Subagent activity UX. (#17570)

packages/core/src/agents/local-invocation.test.ts

@@ -4,17 +4,25 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect, vi, beforeEach, type Mocked } from 'vitest';
+import {
+  describe,
+  it,
+  expect,
+  vi,
+  beforeEach,
+  afterEach,
+  type Mocked,
+} from 'vitest';
 import type {
   LocalAgentDefinition,
   SubagentActivityEvent,
   AgentInputs,
+  SubagentProgress,
 } from './types.js';
 import { LocalSubagentInvocation } from './local-invocation.js';
 import { LocalAgentExecutor } from './local-executor.js';
 import { AgentTerminateMode } from './types.js';
 import { makeFakeConfig } from '../test-utils/config.js';
-import { ToolErrorType } from '../tools/tool-error.js';
 import type { Config } from '../config/config.js';
 import type { MessageBus } from '../confirmation-bus/message-bus.js';
 import { type z } from 'zod';
@@ -29,6 +37,7 @@ let mockConfig: Config;
 const testDefinition: LocalAgentDefinition<z.ZodUnknown> = {
   kind: 'local',
   name: 'MockAgent',
+  displayName: 'Mock Agent',
   description: 'A mock agent.',
   inputConfig: {
     inputSchema: {
@@ -70,6 +79,10 @@ describe('LocalSubagentInvocation', () => {
     );
   });
 
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
   it('should pass the messageBus to the parent constructor', () => {
     const params = { task: 'Analyze data' };
     const invocation = new LocalSubagentInvocation(
@@ -173,7 +186,12 @@ describe('LocalSubagentInvocation', () => {
         mockConfig,
         expect.any(Function),
       );
-      expect(updateOutput).toHaveBeenCalledWith('Subagent starting...\n');
+      expect(updateOutput).toHaveBeenCalledWith(
+        expect.objectContaining({
+          isSubagentProgress: true,
+          agentName: 'MockAgent',
+        }),
+      );
 
       expect(mockExecutorInstance.run).toHaveBeenCalledWith(params, signal);
 
@@ -211,13 +229,17 @@ describe('LocalSubagentInvocation', () => {
 
       await invocation.execute(signal, updateOutput);
 
-      expect(updateOutput).toHaveBeenCalledWith('Subagent starting...\n');
-      expect(updateOutput).toHaveBeenCalledWith('🤖💭 Analyzing...');
-      expect(updateOutput).toHaveBeenCalledWith('🤖💭  Still thinking.');
-      expect(updateOutput).toHaveBeenCalledTimes(3); // Initial message + 2 thoughts
+      expect(updateOutput).toHaveBeenCalledTimes(3); // Initial + 2 updates
+      const lastCall = updateOutput.mock.calls[2][0] as SubagentProgress;
+      expect(lastCall.recentActivity).toContainEqual(
+        expect.objectContaining({
+          type: 'thought',
+          content: 'Analyzing... Still thinking.',
+        }),
+      );
     });
 
-    it('should NOT stream other activities (e.g., TOOL_CALL_START, ERROR)', async () => {
+    it('should stream other activities (e.g., TOOL_CALL_START, ERROR)', async () => {
       mockExecutorInstance.run.mockImplementation(async () => {
         const onActivity = MockLocalAgentExecutor.create.mock.calls[0][2];
 
@@ -226,7 +248,7 @@ describe('LocalSubagentInvocation', () => {
             isSubagentActivityEvent: true,
             agentName: 'MockAgent',
             type: 'TOOL_CALL_START',
-            data: { name: 'ls' },
+            data: { name: 'ls', args: {} },
           } as SubagentActivityEvent);
           onActivity({
             isSubagentActivityEvent: true,
@@ -240,9 +262,15 @@ describe('LocalSubagentInvocation', () => {
 
       await invocation.execute(signal, updateOutput);
 
-      // Should only contain the initial "Subagent starting..." message
-      expect(updateOutput).toHaveBeenCalledTimes(1);
-      expect(updateOutput).toHaveBeenCalledWith('Subagent starting...\n');
+      expect(updateOutput).toHaveBeenCalledTimes(3);
+      const lastCall = updateOutput.mock.calls[2][0] as SubagentProgress;
+      expect(lastCall.recentActivity).toContainEqual(
+        expect.objectContaining({
+          type: 'thought',
+          content: 'Error: Failed',
+          status: 'error',
+        }),
+      );
     });
 
     it('should run successfully without an updateOutput callback', async () => {
@@ -272,16 +300,19 @@ describe('LocalSubagentInvocation', () => {
 
       const result = await invocation.execute(signal, updateOutput);
 
-      expect(result.error).toEqual({
-        message: error.message,
-        type: ToolErrorType.EXECUTION_FAILED,
-      });
-      expect(result.returnDisplay).toBe(
-        `Subagent Failed: MockAgent\nError: ${error.message}`,
-      );
+      expect(result.error).toBeUndefined();
       expect(result.llmContent).toBe(
         `Subagent 'MockAgent' failed. Error: ${error.message}`,
       );
+      const display = result.returnDisplay as SubagentProgress;
+      expect(display.isSubagentProgress).toBe(true);
+      expect(display.recentActivity).toContainEqual(
+        expect.objectContaining({
+          type: 'thought',
+          content: `Error: ${error.message}`,
+          status: 'error',
+        }),
+      );
     });
 
     it('should handle executor creation failure', async () => {
@@ -291,19 +322,21 @@ describe('LocalSubagentInvocation', () => {
       const result = await invocation.execute(signal, updateOutput);
 
       expect(mockExecutorInstance.run).not.toHaveBeenCalled();
-      expect(result.error).toEqual({
-        message: creationError.message,
-        type: ToolErrorType.EXECUTION_FAILED,
-      });
-      expect(result.returnDisplay).toContain(`Error: ${creationError.message}`);
+      expect(result.error).toBeUndefined();
+      expect(result.llmContent).toContain(creationError.message);
+
+      const display = result.returnDisplay as SubagentProgress;
+      expect(display.recentActivity).toContainEqual(
+        expect.objectContaining({
+          content: `Error: ${creationError.message}`,
+          status: 'error',
+        }),
+      );
     });
 
-    /**
-     * This test verifies that the AbortSignal is correctly propagated and
-     * that a rejection from the executor due to abortion is handled gracefully.
-     */
     it('should handle abortion signal during execution', async () => {
       const abortError = new Error('Aborted');
+      abortError.name = 'AbortError';
       mockExecutorInstance.run.mockRejectedValue(abortError);
 
       const controller = new AbortController();
@@ -312,14 +345,24 @@ describe('LocalSubagentInvocation', () => {
         updateOutput,
       );
       controller.abort();
-      const result = await executePromise;
+      await expect(executePromise).rejects.toThrow('Aborted');
 
       expect(mockExecutorInstance.run).toHaveBeenCalledWith(
         params,
         controller.signal,
       );
-      expect(result.error?.message).toBe('Aborted');
-      expect(result.error?.type).toBe(ToolErrorType.EXECUTION_FAILED);
+    });
+
+    it('should throw an error and bubble cancellation when execution returns ABORTED', async () => {
+      const mockOutput = {
+        result: 'Cancelled by user',
+        terminate_reason: AgentTerminateMode.ABORTED,
+      };
+      mockExecutorInstance.run.mockResolvedValue(mockOutput);
+
+      await expect(invocation.execute(signal, updateOutput)).rejects.toThrow(
+        'Operation cancelled by user',
+      );
     });
   });
 });