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

packages/core/src/agents/a2aUtils.ts

@@ -6,12 +6,120 @@
 
 import type {
   Message,
-  Task,
   Part,
   TextPart,
   DataPart,
   FilePart,
+  Artifact,
+  TaskState,
+  TaskStatusUpdateEvent,
 } from '@a2a-js/sdk';
+import type { SendMessageResult } from './a2a-client-manager.js';
+
+/**
+ * Reassembles incremental A2A streaming updates into a coherent result.
+ * Shows sequential status/messages followed by all reassembled artifacts.
+ */
+export class A2AResultReassembler {
+  private messageLog: string[] = [];
+  private artifacts = new Map<string, Artifact>();
+  private artifactChunks = new Map<string, string[]>();
+
+  /**
+   * Processes a new chunk from the A2A stream.
+   */
+  update(chunk: SendMessageResult) {
+    if (!('kind' in chunk)) return;
+
+    switch (chunk.kind) {
+      case 'status-update':
+        this.pushMessage(chunk.status?.message);
+        break;
+
+      case 'artifact-update':
+        if (chunk.artifact) {
+          const id = chunk.artifact.artifactId;
+          const existing = this.artifacts.get(id);
+
+          if (chunk.append && existing) {
+            for (const part of chunk.artifact.parts) {
+              existing.parts.push(structuredClone(part));
+            }
+          } else {
+            this.artifacts.set(id, structuredClone(chunk.artifact));
+          }
+
+          const newText = extractPartsText(chunk.artifact.parts, '');
+          let chunks = this.artifactChunks.get(id);
+          if (!chunks) {
+            chunks = [];
+            this.artifactChunks.set(id, chunks);
+          }
+          if (chunk.append) {
+            chunks.push(newText);
+          } else {
+            chunks.length = 0;
+            chunks.push(newText);
+          }
+        }
+        break;
+
+      case 'task':
+        this.pushMessage(chunk.status?.message);
+        if (chunk.artifacts) {
+          for (const art of chunk.artifacts) {
+            this.artifacts.set(art.artifactId, structuredClone(art));
+            this.artifactChunks.set(art.artifactId, [
+              extractPartsText(art.parts, ''),
+            ]);
+          }
+        }
+        break;
+
+      case 'message': {
+        this.pushMessage(chunk);
+        break;
+      }
+
+      default:
+        break;
+    }
+  }
+
+  private pushMessage(message: Message | undefined) {
+    if (!message) return;
+    const text = extractPartsText(message.parts, '\n');
+    if (text && this.messageLog[this.messageLog.length - 1] !== text) {
+      this.messageLog.push(text);
+    }
+  }
+
+  /**
+   * Returns a human-readable string representation of the current reassembled state.
+   */
+  toString(): string {
+    const joinedMessages = this.messageLog.join('\n\n');
+
+    const artifactsOutput = Array.from(this.artifacts.keys())
+      .map((id) => {
+        const chunks = this.artifactChunks.get(id);
+        const artifact = this.artifacts.get(id);
+        if (!chunks || !artifact) return '';
+        const content = chunks.join('');
+        const header = artifact.name
+          ? `Artifact (${artifact.name}):`
+          : 'Artifact:';
+        return `${header}\n${content}`;
+      })
+      .filter(Boolean)
+      .join('\n\n');
+
+    if (joinedMessages && artifactsOutput) {
+      return `${joinedMessages}\n\n${artifactsOutput}`;
+    }
+    return joinedMessages || artifactsOutput;
+  }
+}
 
 /**
  * Extracts a human-readable text representation from a Message object.
@@ -22,7 +130,23 @@ export function extractMessageText(message: Message | undefined): string {
     return '';
   }
 
-  return extractPartsText(message.parts);
+  return extractPartsText(message.parts, '\n');
+}
+
+/**
+ * Extracts text from an array of parts, joining them with the specified separator.
+ */
+function extractPartsText(
+  parts: Part[] | undefined,
+  separator: string,
+): string {
+  if (!parts || parts.length === 0) {
+    return '';
+  }
+  return parts
+    .map((p) => extractPartText(p))
+    .filter(Boolean)
+    .join(separator);
 }
 
 /**
@@ -52,50 +176,6 @@ function extractPartText(part: Part): string {
   return '';
 }
 
-/**
- * Extracts a clean, human-readable text summary from a Task object.
- * Includes the status message and any artifact content with context headers.
- * Technical metadata like ID and State are omitted for better clarity and token efficiency.
- */
-export function extractTaskText(task: Task): string {
-  const parts: string[] = [];
-
-  // Status Message
-  const statusMessageText = extractMessageText(task.status?.message);
-  if (statusMessageText) {
-    parts.push(statusMessageText);
-  }
-
-  // Artifacts
-  if (task.artifacts) {
-    for (const artifact of task.artifacts) {
-      const artifactContent = extractPartsText(artifact.parts);
-
-      if (artifactContent) {
-        const header = artifact.name
-          ? `Artifact (${artifact.name}):`
-          : 'Artifact:';
-        parts.push(`${header}\n${artifactContent}`);
-      }
-    }
-  }
-
-  return parts.join('\n\n');
-}
-
-/**
- * Extracts text from an array of parts.
- */
-function extractPartsText(parts: Part[] | undefined): string {
-  if (!parts || parts.length === 0) {
-    return '';
-  }
-  return parts
-    .map((p) => extractPartText(p))
-    .filter(Boolean)
-    .join('\n');
-}
-
 // Type Guards
 
 function isTextPart(part: Part): part is TextPart {
@@ -110,36 +190,58 @@ function isFilePart(part: Part): part is FilePart {
   return part.kind === 'file';
 }
 
+function isStatusUpdateEvent(
+  result: SendMessageResult,
+): result is TaskStatusUpdateEvent {
+  return result.kind === 'status-update';
+}
+
 /**
- * Extracts contextId and taskId from a Message or Task response.
+ * Returns true if the given state is a terminal state for a task.
+ */
+export function isTerminalState(state: TaskState | undefined): boolean {
+  return (
+    state === 'completed' ||
+    state === 'failed' ||
+    state === 'canceled' ||
+    state === 'rejected'
+  );
+}
+
+/**
+ * Extracts contextId and taskId from a Message, Task, or Update response.
  * Follows the pattern from the A2A CLI sample to maintain conversational continuity.
  */
-export function extractIdsFromResponse(result: Message | Task): {
+export function extractIdsFromResponse(result: SendMessageResult): {
   contextId?: string;
   taskId?: string;
+  clearTaskId?: boolean;
 } {
   let contextId: string | undefined;
   let taskId: string | undefined;
+  let clearTaskId = false;
 
-  if (result.kind === 'message') {
-    taskId = result.taskId;
-    contextId = result.contextId;
-  } else if (result.kind === 'task') {
-    taskId = result.id;
-    contextId = result.contextId;
-
-    // If the task is in a final state (and not input-required), we clear the taskId
-    // so that the next interaction starts a fresh task (or keeps context without being bound to the old task).
-    if (
-      result.status &&
-      result.status.state !== 'input-required' &&
-      (result.status.state === 'completed' ||
-        result.status.state === 'failed' ||
-        result.status.state === 'canceled')
-    ) {
-      taskId = undefined;
+  if ('kind' in result) {
+    const kind = result.kind;
+    if (kind === 'message' || kind === 'artifact-update') {
+      taskId = result.taskId;
+      contextId = result.contextId;
+    } else if (kind === 'task') {
+      taskId = result.id;
+      contextId = result.contextId;
+      if (isTerminalState(result.status?.state)) {
+        clearTaskId = true;
+      }
+    } else if (isStatusUpdateEvent(result)) {
+      taskId = result.taskId;
+      contextId = result.contextId;
+      // Note: We ignore the 'final' flag here per A2A protocol best practices,
+      // as a stream can close while a task is still in a 'working' state.
+      if (isTerminalState(result.status?.state)) {
+        clearTaskId = true;
+      }
     }
   }
 
-  return { contextId, taskId };
+  return { contextId, taskId, clearTaskId };
 }