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

packages/core/src/agents/a2aUtils.test.ts

@@ -7,12 +7,40 @@
 import { describe, it, expect } from 'vitest';
 import {
   extractMessageText,
-  extractTaskText,
   extractIdsFromResponse,
+  isTerminalState,
+  A2AResultReassembler,
 } from './a2aUtils.js';
-import type { Message, Task, TextPart, DataPart, FilePart } from '@a2a-js/sdk';
+import type { SendMessageResult } from './a2a-client-manager.js';
+import type {
+  Message,
+  Task,
+  TextPart,
+  DataPart,
+  FilePart,
+  TaskStatusUpdateEvent,
+  TaskArtifactUpdateEvent,
+} from '@a2a-js/sdk';
 
 describe('a2aUtils', () => {
+  describe('isTerminalState', () => {
+    it('should return true for completed, failed, canceled, and rejected', () => {
+      expect(isTerminalState('completed')).toBe(true);
+      expect(isTerminalState('failed')).toBe(true);
+      expect(isTerminalState('canceled')).toBe(true);
+      expect(isTerminalState('rejected')).toBe(true);
+    });
+
+    it('should return false for working, submitted, input-required, auth-required, and unknown', () => {
+      expect(isTerminalState('working')).toBe(false);
+      expect(isTerminalState('submitted')).toBe(false);
+      expect(isTerminalState('input-required')).toBe(false);
+      expect(isTerminalState('auth-required')).toBe(false);
+      expect(isTerminalState('unknown')).toBe(false);
+      expect(isTerminalState(undefined)).toBe(false);
+    });
+  });
+
   describe('extractIdsFromResponse', () => {
     it('should extract IDs from a message response', () => {
       const message: Message = {
@@ -25,7 +53,11 @@ describe('a2aUtils', () => {
       };
 
       const result = extractIdsFromResponse(message);
-      expect(result).toEqual({ contextId: 'ctx-1', taskId: 'task-1' });
+      expect(result).toEqual({
+        contextId: 'ctx-1',
+        taskId: 'task-1',
+        clearTaskId: false,
+      });
     });
 
     it('should extract IDs from an in-progress task response', () => {
@@ -37,7 +69,76 @@ describe('a2aUtils', () => {
       };
 
       const result = extractIdsFromResponse(task);
-      expect(result).toEqual({ contextId: 'ctx-2', taskId: 'task-2' });
+      expect(result).toEqual({
+        contextId: 'ctx-2',
+        taskId: 'task-2',
+        clearTaskId: false,
+      });
+    });
+
+    it('should set clearTaskId true for terminal task response', () => {
+      const task: Task = {
+        id: 'task-3',
+        contextId: 'ctx-3',
+        kind: 'task',
+        status: { state: 'completed' },
+      };
+
+      const result = extractIdsFromResponse(task);
+      expect(result.clearTaskId).toBe(true);
+    });
+
+    it('should set clearTaskId true for terminal status update', () => {
+      const update = {
+        kind: 'status-update',
+        contextId: 'ctx-4',
+        taskId: 'task-4',
+        final: true,
+        status: { state: 'failed' },
+      };
+
+      const result = extractIdsFromResponse(
+        update as unknown as TaskStatusUpdateEvent,
+      );
+      expect(result.contextId).toBe('ctx-4');
+      expect(result.taskId).toBe('task-4');
+      expect(result.clearTaskId).toBe(true);
+    });
+
+    it('should extract IDs from an artifact-update event', () => {
+      const update = {
+        kind: 'artifact-update',
+        taskId: 'task-5',
+        contextId: 'ctx-5',
+        artifact: {
+          artifactId: 'art-1',
+          parts: [{ kind: 'text', text: 'artifact content' }],
+        },
+      } as unknown as TaskArtifactUpdateEvent;
+
+      const result = extractIdsFromResponse(update);
+      expect(result).toEqual({
+        contextId: 'ctx-5',
+        taskId: 'task-5',
+        clearTaskId: false,
+      });
+    });
+
+    it('should extract taskId from status update event', () => {
+      const update = {
+        kind: 'status-update',
+        taskId: 'task-6',
+        contextId: 'ctx-6',
+        final: false,
+        status: { state: 'working' },
+      };
+
+      const result = extractIdsFromResponse(
+        update as unknown as TaskStatusUpdateEvent,
+      );
+      expect(result.taskId).toBe('task-6');
+      expect(result.contextId).toBe('ctx-6');
+      expect(result.clearTaskId).toBe(false);
     });
   });
 
@@ -123,49 +224,65 @@ describe('a2aUtils', () => {
     });
   });
 
-  describe('extractTaskText', () => {
-    it('should extract basic task info (clean)', () => {
-      const task: Task = {
-        id: 'task-1',
-        contextId: 'ctx-1',
-        kind: 'task',
+  describe('A2AResultReassembler', () => {
+    it('should reassemble sequential messages and incremental artifacts', () => {
+      const reassembler = new A2AResultReassembler();
+
+      // 1. Initial status
+      reassembler.update({
+        kind: 'status-update',
+        taskId: 't1',
         status: {
           state: 'working',
           message: {
             kind: 'message',
             role: 'agent',
-            messageId: 'm1',
-            parts: [{ kind: 'text', text: 'Processing...' } as TextPart],
-          },
+            parts: [{ kind: 'text', text: 'Analyzing...' }],
+          } as Message,
         },
-      };
+      } as unknown as SendMessageResult);
 
-      const result = extractTaskText(task);
-      expect(result).not.toContain('ID: task-1');
-      expect(result).not.toContain('State: working');
-      expect(result).toBe('Processing...');
-    });
+      // 2. First artifact chunk
+      reassembler.update({
+        kind: 'artifact-update',
+        taskId: 't1',
+        append: false,
+        artifact: {
+          artifactId: 'a1',
+          name: 'Code',
+          parts: [{ kind: 'text', text: 'print(' }],
+        },
+      } as unknown as SendMessageResult);
 
-    it('should extract artifacts with headers', () => {
-      const task: Task = {
-        id: 'task-1',
-        contextId: 'ctx-1',
-        kind: 'task',
-        status: { state: 'completed' },
-        artifacts: [
-          {
-            artifactId: 'art-1',
-            name: 'Report',
-            parts: [{ kind: 'text', text: 'This is the report.' } as TextPart],
-          },
-        ],
-      };
+      // 3. Second status
+      reassembler.update({
+        kind: 'status-update',
+        taskId: 't1',
+        status: {
+          state: 'working',
+          message: {
+            kind: 'message',
+            role: 'agent',
+            parts: [{ kind: 'text', text: 'Processing...' }],
+          } as Message,
+        },
+      } as unknown as SendMessageResult);
 
-      const result = extractTaskText(task);
-      expect(result).toContain('Artifact (Report):');
-      expect(result).toContain('This is the report.');
-      expect(result).not.toContain('Artifacts:');
-      expect(result).not.toContain('  - Name: Report');
+      // 4. Second artifact chunk (append)
+      reassembler.update({
+        kind: 'artifact-update',
+        taskId: 't1',
+        append: true,
+        artifact: {
+          artifactId: 'a1',
+          parts: [{ kind: 'text', text: '"Done")' }],
+        },
+      } as unknown as SendMessageResult);
+
+      const output = reassembler.toString();
+      expect(output).toBe(
+        'Analyzing...\n\nProcessing...\n\nArtifact (Code):\nprint("Done")',
+      );
     });
   });
 });