@@ -49,6 +49,16 @@ struct ToolCallRequest {
4949 let completion : ( AnyJSONRPCResponse ) -> Void
5050}
5151
52+ struct ConversationTurnTrackingState {
53+ var turnParentMap : [ String : String ] = [ : ] // Maps subturn ID to parent turn ID
54+ var validConversationIds : Set < String > = [ ] // Tracks all valid conversation IDs including subagents
55+
56+ mutating func reset ( ) {
57+ turnParentMap . removeAll ( )
58+ validConversationIds . removeAll ( )
59+ }
60+ }
61+
5262public final class ChatService : ChatServiceType , ObservableObject {
5363
5464 public var memory : ContextAwareAutoManagedChatMemory
@@ -69,6 +79,9 @@ public final class ChatService: ChatServiceType, ObservableObject {
6979 private var lastUserRequest : ConversationRequest ?
7080 private var isRestored : Bool = false
7181 private var pendingToolCallRequests : [ String : ToolCallRequest ] = [ : ]
82+ // Workaround: toolConfirmation request does not have parent turnId
83+ private var conversationTurnTracking = ConversationTurnTrackingState ( )
84+
7285 init ( provider : any ConversationServiceProvider ,
7386 memory : ContextAwareAutoManagedChatMemory = ContextAwareAutoManagedChatMemory ( ) ,
7487 conversationProgressHandler : ConversationProgressHandler = ConversationProgressHandlerImpl . shared ,
@@ -136,7 +149,15 @@ public final class ChatService: ChatServiceType, ObservableObject {
136149
137150 private func subscribeToClientToolConfirmationEvent ( ) {
138151 ClientToolHandlerImpl . shared . onClientToolConfirmationEvent . sink ( receiveValue : { [ weak self ] ( request , completion ) in
139- guard let params = request . params , params . conversationId == self ? . conversationId else { return }
152+ guard let params = request . params else { return }
153+
154+ // Check if this conversationId is valid (main conversation or subagent conversation)
155+ guard let validIds = self ? . conversationTurnTracking . validConversationIds , validIds . contains ( params . conversationId ) else {
156+ return
157+ }
158+
159+ let parentTurnId = self ? . conversationTurnTracking . turnParentMap [ params . turnId ]
160+
140161 let editAgentRounds : [ AgentRound ] = [
141162 AgentRound ( roundId : params . roundId ,
142163 reply : " " ,
@@ -145,7 +166,7 @@ public final class ChatService: ChatServiceType, ObservableObject {
145166 ]
146167 )
147168 ]
148- self ? . appendToolCallHistory ( turnId : params . turnId , editAgentRounds : editAgentRounds )
169+ self ? . appendToolCallHistory ( turnId : params . turnId , editAgentRounds : editAgentRounds , parentTurnId : parentTurnId )
149170 self ? . pendingToolCallRequests [ params . toolCallId ] = ToolCallRequest (
150171 requestId : request . id ,
151172 turnId : params . turnId ,
@@ -157,7 +178,13 @@ public final class ChatService: ChatServiceType, ObservableObject {
157178
158179 private func subscribeToClientToolInvokeEvent ( ) {
159180 ClientToolHandlerImpl . shared . onClientToolInvokeEvent . sink ( receiveValue : { [ weak self ] ( request , completion ) in
160- guard let params = request . params , params . conversationId == self ? . conversationId else { return }
181+ guard let params = request . params else { return }
182+
183+ // Check if this conversationId is valid (main conversation or subagent conversation)
184+ guard let validIds = self ? . conversationTurnTracking . validConversationIds , validIds . contains ( params . conversationId ) else {
185+ return
186+ }
187+
161188 guard let copilotTool = CopilotToolRegistry . shared . getTool ( name : params . name ) else {
162189 completion ( AnyJSONRPCResponse ( id : request . id ,
163190 result : JSONValue . array ( [
@@ -173,11 +200,11 @@ public final class ChatService: ChatServiceType, ObservableObject {
173200 return
174201 }
175202
176- copilotTool . invokeTool ( request , completion : completion , contextProvider : self )
203+ _ = copilotTool . invokeTool ( request , completion : completion , contextProvider : self )
177204 } ) . store ( in : & cancellables )
178205 }
179206
180- func appendToolCallHistory ( turnId : String , editAgentRounds : [ AgentRound ] , fileEdits : [ FileEdit ] = [ ] ) {
207+ func appendToolCallHistory ( turnId : String , editAgentRounds : [ AgentRound ] , fileEdits : [ FileEdit ] = [ ] , parentTurnId : String ? = nil ) {
181208 let chatTabId = self . chatTabInfo . id
182209 Task {
183210 let turnStatus : ChatMessage . TurnStatus ? = {
@@ -196,6 +223,7 @@ public final class ChatService: ChatServiceType, ObservableObject {
196223 assistantMessageWithId : turnId ,
197224 chatTabID : chatTabId ,
198225 editAgentRounds : editAgentRounds ,
226+ parentTurnId : parentTurnId ,
199227 fileEdits : fileEdits ,
200228 turnStatus : turnStatus
201229 )
@@ -229,75 +257,64 @@ public final class ChatService: ChatServiceType, ObservableObject {
229257 self . isRestored = true
230258 }
231259
260+ /// Updates the status of a tool call (accepted, cancelled, etc.) and notifies the server
261+ ///
262+ /// This method handles two key responsibilities:
263+ /// 1. Sends confirmation response back to the server when user accepts/cancels
264+ /// 2. Updates the tool call status in chat history UI (including subagent tool calls)
232265 public func updateToolCallStatus ( toolCallId : String , status : AgentToolCall . ToolCallStatus , payload : Any ? = nil ) {
233- // Send the tool call result back to the server
234- if let toolCallRequest = self . pendingToolCallRequests [ toolCallId ] , status == . accepted || status == . cancelled {
266+ // Capture the pending request info before removing it from the dictionary
267+ let toolCallRequest = self . pendingToolCallRequests [ toolCallId ]
268+
269+ // Step 1: Send confirmation response to server (for accept/cancel actions only)
270+ if let toolCallRequest = toolCallRequest , status == . accepted || status == . cancelled {
235271 self . pendingToolCallRequests . removeValue ( forKey : toolCallId )
236- let toolResult = LanguageModelToolConfirmationResult (
237- result : status == . accepted ? . Accept : . Dismiss
238- )
239- let jsonResult = try? JSONEncoder ( ) . encode ( toolResult )
240- let jsonValue = ( try? JSONDecoder ( ) . decode ( JSONValue . self , from : jsonResult ?? Data ( ) ) ) ?? JSONValue . null
241- toolCallRequest . completion (
242- AnyJSONRPCResponse (
243- id : toolCallRequest . requestId ,
244- result : JSONValue . array ( [
245- jsonValue ,
246- JSONValue . null
247- ] )
248- )
249- )
272+ sendToolConfirmationResponse ( toolCallRequest , accepted : status == . accepted )
250273 }
251274
252- // Update the tool call status in the chat history
275+ // Step 2: Update the tool call status in chat history UI
253276 Task {
254- guard let lastMessage = await memory . history . last , lastMessage . role == . assistant else {
277+ guard let targetMessage = await ToolCallStatusUpdater . findMessageContainingToolCall (
278+ toolCallRequest ,
279+ conversationTurnTracking : conversationTurnTracking ,
280+ history : await memory . history
281+ ) else {
255282 return
256283 }
257-
258- var updatedAgentRounds : [ AgentRound ] = [ ]
259- for i in 0 ..< lastMessage . editAgentRounds . count {
260- if lastMessage . editAgentRounds [ i ] . toolCalls == nil {
261- continue
262- }
263- for j in 0 ..< lastMessage . editAgentRounds [ i ] . toolCalls ! . count {
264- if lastMessage . editAgentRounds [ i ] . toolCalls ! [ j ] . id == toolCallId {
265- updatedAgentRounds . append (
266- AgentRound ( roundId : lastMessage . editAgentRounds [ i ] . roundId ,
267- reply : " " ,
268- toolCalls : [
269- AgentToolCall ( id : toolCallId ,
270- name : lastMessage . editAgentRounds [ i ] . toolCalls ! [ j ] . name ,
271- status : status )
272- ]
273- )
274- )
275- break
276- }
277- }
278- if ! updatedAgentRounds . isEmpty {
279- break
280- }
281- }
282-
283- if ! updatedAgentRounds . isEmpty {
284- let message = ChatMessage (
285- id : lastMessage . id ,
286- chatTabID : lastMessage . chatTabID ,
287- clsTurnID : lastMessage . clsTurnID ,
288- role : . assistant ,
289- content : " " ,
290- references : [ ] ,
291- steps : [ ] ,
292- editAgentRounds : updatedAgentRounds ,
293- turnStatus : . inProgress
284+
285+ // Search for the tool call in main rounds or subagent rounds
286+ if let updatedRound = ToolCallStatusUpdater . findAndUpdateToolCall (
287+ toolCallId : toolCallId ,
288+ newStatus : status ,
289+ in : targetMessage . editAgentRounds
290+ ) {
291+ let message = ToolCallStatusUpdater . createMessageUpdate (
292+ targetMessage : targetMessage ,
293+ updatedRound : updatedRound
294294 )
295-
296- await self . memory . appendMessage ( message )
295+ await memory . appendMessage ( message )
297296 }
298297 }
299298 }
300299
300+ // MARK: - Helper Methods for Tool Call Status Updates
301+
302+ /// Sends the confirmation response (accept/dismiss) back to the server
303+ private func sendToolConfirmationResponse ( _ request : ToolCallRequest , accepted : Bool ) {
304+ let toolResult = LanguageModelToolConfirmationResult (
305+ result : accepted ? . Accept : . Dismiss
306+ )
307+ let jsonResult = try? JSONEncoder ( ) . encode ( toolResult )
308+ let jsonValue = ( try? JSONDecoder ( ) . decode ( JSONValue . self , from : jsonResult ?? Data ( ) ) ) ?? JSONValue . null
309+
310+ request . completion (
311+ AnyJSONRPCResponse (
312+ id : request . requestId ,
313+ result : JSONValue . array ( [ jsonValue , JSONValue . null ] )
314+ )
315+ )
316+ }
317+
301318 public enum ChatServiceError : Error , LocalizedError {
302319 case conflictingImageFormats ( String )
303320
@@ -676,9 +693,18 @@ public final class ChatService: ChatServiceType, ObservableObject {
676693 if progress . parentTurnId == nil {
677694 conversationId = progress . conversationId
678695 }
696+
697+ // Track all valid conversation IDs for the current turn (main conversation + its subturns)
698+ conversationTurnTracking . validConversationIds . insert ( progress . conversationId )
699+
679700 let turnId = progress . turnId
680701 let parentTurnId = progress . parentTurnId
681702
703+ // Track parent-subturn relationship
704+ if let parentTurnId = parentTurnId {
705+ conversationTurnTracking . turnParentMap [ turnId ] = parentTurnId
706+ }
707+
682708 Task {
683709 if var lastUserMessage = await memory . history . last ( where : { $0 . role == . user } ) {
684710
@@ -870,6 +896,9 @@ public final class ChatService: ChatServiceType, ObservableObject {
870896 activeRequestId = nil
871897 isReceivingMessage = false
872898 requestType = nil
899+
900+ // Clear turn tracking data
901+ conversationTurnTracking . reset ( )
873902
874903 // cancel all pending tool call requests
875904 for ( _, request ) in pendingToolCallRequests {
0 commit comments