feat(core): implement robust A2A streaming reassembly and fix task continuity (#20091)

packages/core/src/agents/remote-invocation.ts

@@ -18,14 +18,12 @@ import type {
 } from './types.js';
 import type { MessageBus } from '../confirmation-bus/message-bus.js';
 import { A2AClientManager } from './a2a-client-manager.js';
-import {
-  extractMessageText,
-  extractTaskText,
-  extractIdsFromResponse,
-} from './a2aUtils.js';
+import { extractIdsFromResponse, A2AResultReassembler } from './a2aUtils.js';
 import { GoogleAuth } from 'google-auth-library';
 import type { AuthenticationHandler } from '@a2a-js/sdk/client';
 import { debugLogger } from '../utils/debugLogger.js';
+import type { AnsiOutput } from '../utils/terminalSerializer.js';
+import type { SendMessageResult } from './a2a-client-manager.js';
 
 /**
  * Authentication handler implementation using Google Application Default Credentials (ADC).
@@ -123,10 +121,14 @@ export class RemoteAgentInvocation extends BaseToolInvocation<
     };
   }
 
-  async execute(_signal: AbortSignal): Promise<ToolResult> {
+  async execute(
+    _signal: AbortSignal,
+    updateOutput?: (output: string | AnsiOutput) => void,
+  ): Promise<ToolResult> {
     // 1. Ensure the agent is loaded (cached by manager)
     // We assume the user has provided an access token via some mechanism (TODO),
     // or we rely on ADC.
+    const reassembler = new A2AResultReassembler();
     try {
       const priorState = RemoteAgentInvocation.sessionState.get(
         this.definition.name,
@@ -146,49 +148,73 @@ export class RemoteAgentInvocation extends BaseToolInvocation<
 
       const message = this.params.query;
 
-      const response = await this.clientManager.sendMessage(
+      const stream = this.clientManager.sendMessageStream(
         this.definition.name,
         message,
         {
           contextId: this.contextId,
           taskId: this.taskId,
+          signal: _signal,
         },
       );
 
-      // Extracts IDs, taskID will be undefined if the task is completed/failed/canceled.
-      const { contextId, taskId } = extractIdsFromResponse(response);
+      let finalResponse: SendMessageResult | undefined;
 
-      this.contextId = contextId ?? this.contextId;
-      this.taskId = taskId;
+      for await (const chunk of stream) {
+        if (_signal.aborted) {
+          throw new Error('Operation aborted');
+        }
+        finalResponse = chunk;
+        reassembler.update(chunk);
 
+        if (updateOutput) {
+          updateOutput(reassembler.toString());
+        }
+
+        const {
+          contextId: newContextId,
+          taskId: newTaskId,
+          clearTaskId,
+        } = extractIdsFromResponse(chunk);
+
+        if (newContextId) {
+          this.contextId = newContextId;
+        }
+
+        this.taskId = clearTaskId ? undefined : (newTaskId ?? this.taskId);
+      }
+
+      if (!finalResponse) {
+        throw new Error('No response from remote agent.');
+      }
+
+      const finalOutput = reassembler.toString();
+
+      debugLogger.debug(
+        `[RemoteAgent] Final response from ${this.definition.name}:\n${JSON.stringify(finalResponse, null, 2)}`,
+      );
+
+      return {
+        llmContent: [{ text: finalOutput }],
+        returnDisplay: finalOutput,
+      };
+    } catch (error: unknown) {
+      const partialOutput = reassembler.toString();
+      const errorMessage = `Error calling remote agent: ${error instanceof Error ? error.message : String(error)}`;
+      const fullDisplay = partialOutput
+        ? `${partialOutput}\n\n${errorMessage}`
+        : errorMessage;
+      return {
+        llmContent: [{ text: fullDisplay }],
+        returnDisplay: fullDisplay,
+        error: { message: errorMessage },
+      };
+    } finally {
+      // Persist state even on partial failures or aborts to maintain conversational continuity.
       RemoteAgentInvocation.sessionState.set(this.definition.name, {
         contextId: this.contextId,
         taskId: this.taskId,
       });
-
-      // Extract the output text
-      const outputText =
-        response.kind === 'task'
-          ? extractTaskText(response)
-          : response.kind === 'message'
-            ? extractMessageText(response)
-            : JSON.stringify(response);
-
-      debugLogger.debug(
-        `[RemoteAgent] Response from ${this.definition.name}:\n${JSON.stringify(response, null, 2)}`,
-      );
-
-      return {
-        llmContent: [{ text: outputText }],
-        returnDisplay: outputText,
-      };
-    } catch (error: unknown) {
-      const errorMessage = `Error calling remote agent: ${error instanceof Error ? error.message : String(error)}`;
-      return {
-        llmContent: [{ text: errorMessage }],
-        returnDisplay: errorMessage,
-        error: { message: errorMessage },
-      };
     }
   }
 }