From c987b99394a318e55095a18b8c1066a66edd91c7 Mon Sep 17 00:00:00 2001 From: Adam Weidman <65992621+adamfweidman@users.noreply.github.com> Date: Tue, 12 May 2026 14:58:25 -0400 Subject: [PATCH] refactor(core): introduce SubagentState enum for progress (#26934) --- .../messages/SubagentGroupDisplay.test.tsx | 14 ++++-- .../messages/SubagentGroupDisplay.tsx | 25 +++++----- .../messages/SubagentHistoryMessage.test.tsx | 7 +-- .../messages/SubagentProgressDisplay.test.tsx | 20 ++++---- .../messages/SubagentProgressDisplay.tsx | 27 ++++++----- .../ToolGroupMessageRegression.test.tsx | 5 +- .../cli/src/ui/hooks/useToolScheduler.test.ts | 13 ++++-- packages/core/src/agents/a2aUtils.ts | 8 ++-- .../agents/browser/browserAgentInvocation.ts | 46 +++++++++++-------- .../core/src/agents/local-invocation.test.ts | 17 +++---- packages/core/src/agents/local-invocation.ts | 44 ++++++++++-------- .../core/src/agents/remote-invocation.test.ts | 40 ++++++++++------ packages/core/src/agents/remote-invocation.ts | 11 +++-- .../src/agents/remote-subagent-protocol.ts | 5 +- packages/core/src/agents/types.ts | 11 ++++- 15 files changed, 172 insertions(+), 121 deletions(-) diff --git a/packages/cli/src/ui/components/messages/SubagentGroupDisplay.test.tsx b/packages/cli/src/ui/components/messages/SubagentGroupDisplay.test.tsx index 484ca8a8ed..1a3572a82a 100644 --- a/packages/cli/src/ui/components/messages/SubagentGroupDisplay.test.tsx +++ b/packages/cli/src/ui/components/messages/SubagentGroupDisplay.test.tsx @@ -6,7 +6,11 @@ import { waitFor } from '../../../test-utils/async.js'; import { renderWithProviders } from '../../../test-utils/render.js'; import { SubagentGroupDisplay } from './SubagentGroupDisplay.js'; -import { Kind, CoreToolCallStatus } from '@google/gemini-cli-core'; +import { + Kind, + CoreToolCallStatus, + SubagentState, +} from '@google/gemini-cli-core'; import type { IndividualToolCallDisplay } from '../../types.js'; import { describe, it, expect, vi } from 'vitest'; import { Text } from 'ink'; @@ -27,12 +31,12 @@ describe('', () => { resultDisplay: { isSubagentProgress: true, agentName: 'api-monitor', - state: 'running', + state: SubagentState.RUNNING, recentActivity: [ { id: 'act-1', type: 'tool_call', - status: 'running', + status: SubagentState.RUNNING, content: '', displayName: 'Action Required', description: 'Verify server is running', @@ -50,13 +54,13 @@ describe('', () => { resultDisplay: { isSubagentProgress: true, agentName: 'db-manager', - state: 'completed', + state: SubagentState.COMPLETED, result: 'Database schema validated', recentActivity: [ { id: 'act-2', type: 'thought', - status: 'completed', + status: SubagentState.COMPLETED, content: 'Database schema validated', }, ], diff --git a/packages/cli/src/ui/components/messages/SubagentGroupDisplay.tsx b/packages/cli/src/ui/components/messages/SubagentGroupDisplay.tsx index b57160966b..02ff8d461b 100644 --- a/packages/cli/src/ui/components/messages/SubagentGroupDisplay.tsx +++ b/packages/cli/src/ui/components/messages/SubagentGroupDisplay.tsx @@ -13,6 +13,7 @@ import { isSubagentProgress, checkExhaustive, type SubagentActivityItem, + SubagentState, } from '@google/gemini-cli-core'; import { SubagentProgressDisplay, @@ -66,13 +67,13 @@ export const SubagentGroupDisplay: React.FC = ({ const singleAgent = toolCalls[0].resultDisplay; if (isSubagentProgress(singleAgent)) { switch (singleAgent.state) { - case 'completed': + case SubagentState.COMPLETED: headerText = 'Agent Completed'; break; - case 'cancelled': + case SubagentState.CANCELLED: headerText = 'Agent Cancelled'; break; - case 'error': + case SubagentState.ERROR: headerText = 'Agent Error'; break; default: @@ -88,8 +89,8 @@ export const SubagentGroupDisplay: React.FC = ({ for (const tc of toolCalls) { const progress = tc.resultDisplay; if (isSubagentProgress(progress)) { - if (progress.state === 'completed') completedCount++; - else if (progress.state === 'running') runningCount++; + if (progress.state === SubagentState.COMPLETED) completedCount++; + else if (progress.state === SubagentState.RUNNING) runningCount++; } else { // It hasn't emitted progress yet, but it is "running" runningCount++; @@ -200,7 +201,7 @@ export const SubagentGroupDisplay: React.FC = ({ let content = 'Starting...'; let formattedArgs: string | undefined; - if (progress.state === 'completed') { + if (progress.state === SubagentState.COMPLETED) { if ( progress.terminateReason && progress.terminateReason !== 'GOAL' @@ -223,18 +224,18 @@ export const SubagentGroupDisplay: React.FC = ({ } const displayArgs = - progress.state === 'completed' ? '' : formattedArgs; + progress.state === SubagentState.COMPLETED ? '' : formattedArgs; const renderStatusIcon = () => { - const state = progress.state ?? 'running'; + const state = progress.state ?? SubagentState.RUNNING; switch (state) { - case 'running': + case SubagentState.RUNNING: return !; - case 'completed': + case SubagentState.COMPLETED: return ✓; - case 'cancelled': + case SubagentState.CANCELLED: return ℹ; - case 'error': + case SubagentState.ERROR: return ✗; default: return checkExhaustive(state); diff --git a/packages/cli/src/ui/components/messages/SubagentHistoryMessage.test.tsx b/packages/cli/src/ui/components/messages/SubagentHistoryMessage.test.tsx index 20a86cb5a9..9db757b240 100644 --- a/packages/cli/src/ui/components/messages/SubagentHistoryMessage.test.tsx +++ b/packages/cli/src/ui/components/messages/SubagentHistoryMessage.test.tsx @@ -8,6 +8,7 @@ import { describe, it, expect } from 'vitest'; import { renderWithProviders } from '../../../test-utils/render.js'; import { SubagentHistoryMessage } from './SubagentHistoryMessage.js'; import type { HistoryItemSubagent } from '../../types.js'; +import { SubagentState } from '@google/gemini-cli-core'; describe('SubagentHistoryMessage', () => { const mockItem: HistoryItemSubagent = { @@ -18,19 +19,19 @@ describe('SubagentHistoryMessage', () => { id: '1', type: 'thought', content: 'Thinking about the problem', - status: 'completed', + status: SubagentState.COMPLETED, }, { id: '2', type: 'tool_call', content: 'Calling search_web', - status: 'running', + status: SubagentState.RUNNING, }, { id: '3', type: 'tool_call', content: 'Calling read_file fail', - status: 'error', + status: SubagentState.ERROR, }, ], }; diff --git a/packages/cli/src/ui/components/messages/SubagentProgressDisplay.test.tsx b/packages/cli/src/ui/components/messages/SubagentProgressDisplay.test.tsx index fcafa4ed28..d1f2d70f0e 100644 --- a/packages/cli/src/ui/components/messages/SubagentProgressDisplay.test.tsx +++ b/packages/cli/src/ui/components/messages/SubagentProgressDisplay.test.tsx @@ -6,7 +6,7 @@ import { render, cleanup } from '../../../test-utils/render.js'; import { SubagentProgressDisplay } from './SubagentProgressDisplay.js'; -import type { SubagentProgress } from '@google/gemini-cli-core'; +import { type SubagentProgress, SubagentState } from '@google/gemini-cli-core'; import { describe, it, expect, vi, afterEach } from 'vitest'; describe('', () => { @@ -25,7 +25,7 @@ describe('', () => { type: 'tool_call', content: 'run_shell_command', args: '{"command": "echo hello", "description": "Say hello"}', - status: 'running', + status: SubagentState.RUNNING, }, ], }; @@ -48,7 +48,7 @@ describe('', () => { displayName: 'RunShellCommand', description: 'Executing echo hello', args: '{"command": "echo hello"}', - status: 'running', + status: SubagentState.RUNNING, }, ], }; @@ -69,7 +69,7 @@ describe('', () => { type: 'tool_call', content: 'run_shell_command', args: '{"command": "echo hello"}', - status: 'running', + status: SubagentState.RUNNING, }, ], }; @@ -90,7 +90,7 @@ describe('', () => { type: 'tool_call', content: 'write_file', args: '{"file_path": "/tmp/test.txt", "content": "foo"}', - status: 'completed', + status: SubagentState.COMPLETED, }, ], }; @@ -113,7 +113,7 @@ describe('', () => { type: 'tool_call', content: 'run_shell_command', args: JSON.stringify({ description: longDesc }), - status: 'running', + status: SubagentState.RUNNING, }, ], }; @@ -133,7 +133,7 @@ describe('', () => { id: '5', type: 'thought', content: 'Thinking about life', - status: 'running', + status: SubagentState.RUNNING, }, ], }; @@ -149,7 +149,7 @@ describe('', () => { isSubagentProgress: true, agentName: 'TestAgent', recentActivity: [], - state: 'cancelled', + state: SubagentState.CANCELLED, }; const { lastFrame } = await render( @@ -167,7 +167,7 @@ describe('', () => { id: '6', type: 'thought', content: 'Request cancelled.', - status: 'error', + status: SubagentState.ERROR, }, ], }; @@ -188,7 +188,7 @@ describe('', () => { type: 'tool_call', content: 'run_shell_command', args: '{"command": "echo hello"}', - status: 'error', + status: SubagentState.ERROR, }, ], }; diff --git a/packages/cli/src/ui/components/messages/SubagentProgressDisplay.tsx b/packages/cli/src/ui/components/messages/SubagentProgressDisplay.tsx index 995c404d9d..b46756c5d3 100644 --- a/packages/cli/src/ui/components/messages/SubagentProgressDisplay.tsx +++ b/packages/cli/src/ui/components/messages/SubagentProgressDisplay.tsx @@ -9,9 +9,10 @@ import { Box, Text } from 'ink'; import { theme } from '../../semantic-colors.js'; import Spinner from 'ink-spinner'; import { MarkdownDisplay } from '../../utils/MarkdownDisplay.js'; -import type { - SubagentProgress, - SubagentActivityItem, +import { + type SubagentProgress, + type SubagentActivityItem, + SubagentState, } from '@google/gemini-cli-core'; import { TOOL_STATUS } from '../../constants.js'; import { STATUS_INDICATOR_WIDTH } from './ToolShared.js'; @@ -62,13 +63,13 @@ export const SubagentProgressDisplay: React.FC< let headerText: string | undefined; let headerColor = theme.text.secondary; - if (progress.state === 'cancelled') { + if (progress.state === SubagentState.CANCELLED) { headerText = `Subagent ${progress.agentName} was cancelled.`; headerColor = theme.status.warning; - } else if (progress.state === 'error') { + } else if (progress.state === SubagentState.ERROR) { headerText = `Subagent ${progress.agentName} failed.`; headerColor = theme.status.error; - } else if (progress.state === 'completed') { + } else if (progress.state === SubagentState.COMPLETED) { headerText = `Subagent ${progress.agentName} completed.`; headerColor = theme.status.success; } else { @@ -107,13 +108,13 @@ export const SubagentProgressDisplay: React.FC< ); } else if (item.type === 'tool_call') { const statusSymbol = - item.status === 'running' ? ( + item.status === SubagentState.RUNNING ? ( - ) : item.status === 'completed' ? ( + ) : item.status === SubagentState.COMPLETED ? ( {TOOL_STATUS.SUCCESS} - ) : item.status === 'cancelled' ? ( + ) : item.status === SubagentState.CANCELLED ? ( {TOOL_STATUS.CANCELED} @@ -135,7 +136,7 @@ export const SubagentProgressDisplay: React.FC< {item.displayName || item.content} @@ -144,7 +145,9 @@ export const SubagentProgressDisplay: React.FC< {displayArgs} @@ -170,7 +173,7 @@ export const SubagentProgressDisplay: React.FC< )} diff --git a/packages/cli/src/ui/components/messages/ToolGroupMessageRegression.test.tsx b/packages/cli/src/ui/components/messages/ToolGroupMessageRegression.test.tsx index 96239fb720..5206145c9e 100644 --- a/packages/cli/src/ui/components/messages/ToolGroupMessageRegression.test.tsx +++ b/packages/cli/src/ui/components/messages/ToolGroupMessageRegression.test.tsx @@ -13,6 +13,7 @@ import { ApprovalMode, WRITE_FILE_DISPLAY_NAME, Kind, + SubagentState, } from '@google/gemini-cli-core'; import os from 'node:os'; import { createMockSettings } from '../../../test-utils/settings.js'; @@ -76,7 +77,7 @@ describe('ToolGroupMessage Regression Tests', () => { resultDisplay: { isSubagentProgress: true, agentName: 'TestAgent', - state: 'running', + state: SubagentState.RUNNING, recentActivity: [], }, }), @@ -112,7 +113,7 @@ describe('ToolGroupMessage Regression Tests', () => { resultDisplay: { isSubagentProgress: true, agentName: 'TestAgent', - state: 'completed', + state: SubagentState.COMPLETED, recentActivity: [], }, }), diff --git a/packages/cli/src/ui/hooks/useToolScheduler.test.ts b/packages/cli/src/ui/hooks/useToolScheduler.test.ts index efb9b8a6fd..e9665ec63b 100644 --- a/packages/cli/src/ui/hooks/useToolScheduler.test.ts +++ b/packages/cli/src/ui/hooks/useToolScheduler.test.ts @@ -21,6 +21,7 @@ import { ROOT_SCHEDULER_ID, CoreToolCallStatus, type WaitingToolCall, + SubagentState, } from '@google/gemini-cli-core'; import { createMockMessageBus } from '@google/gemini-cli-core/src/test-utils/mock-message-bus.js'; @@ -630,7 +631,7 @@ describe('useToolScheduler', () => { id: '1', type: 'thought', content: 'Thinking...', - status: 'running', + status: SubagentState.RUNNING, }, }); }); @@ -648,7 +649,7 @@ describe('useToolScheduler', () => { id: '2', type: 'tool_call', content: 'Calling tool', - status: 'completed', + status: SubagentState.COMPLETED, }, }); }); @@ -697,7 +698,7 @@ describe('useToolScheduler', () => { id: '1', type: 'thought', content: 'Thinking...', - status: 'running', + status: SubagentState.RUNNING, }, }); }); @@ -716,7 +717,7 @@ describe('useToolScheduler', () => { id: '1', type: 'thought', content: 'Thinking... Done!', - status: 'completed', + status: SubagentState.COMPLETED, }, }); }); @@ -726,6 +727,8 @@ describe('useToolScheduler', () => { expect(result.current[0][0].subagentHistory![0].content).toBe( 'Thinking... Done!', ); - expect(result.current[0][0].subagentHistory![0].status).toBe('completed'); + expect(result.current[0][0].subagentHistory![0].status).toBe( + SubagentState.COMPLETED, + ); }); }); diff --git a/packages/core/src/agents/a2aUtils.ts b/packages/core/src/agents/a2aUtils.ts index 2d146fc420..876f623911 100644 --- a/packages/core/src/agents/a2aUtils.ts +++ b/packages/core/src/agents/a2aUtils.ts @@ -16,7 +16,7 @@ import type { AgentInterface, } from '@a2a-js/sdk'; import type { SendMessageResult } from './a2a-client-manager.js'; -import type { SubagentActivityItem } from './types.js'; +import { type SubagentActivityItem, SubagentState } from './types.js'; export const AUTH_REQUIRED_MSG = `[Authorization Required] The agent has indicated it requires authorization to proceed. Please follow the agent's instructions.`; @@ -143,7 +143,7 @@ export class A2AResultReassembler { id: 'auth-required', type: 'thought', content: AUTH_REQUIRED_MSG, - status: 'running', + status: SubagentState.RUNNING, }); } @@ -152,7 +152,7 @@ export class A2AResultReassembler { id: `msg-${index}`, type: 'thought', content: msg.trim(), - status: 'completed', + status: SubagentState.COMPLETED, }); }); @@ -161,7 +161,7 @@ export class A2AResultReassembler { id: 'pending', type: 'thought', content: 'Working...', - status: 'running', + status: SubagentState.RUNNING, }); } diff --git a/packages/core/src/agents/browser/browserAgentInvocation.ts b/packages/core/src/agents/browser/browserAgentInvocation.ts index a59ffc25b5..a27a8d29ed 100644 --- a/packages/core/src/agents/browser/browserAgentInvocation.ts +++ b/packages/core/src/agents/browser/browserAgentInvocation.ts @@ -32,6 +32,7 @@ import { type SubagentActivityItem, AgentTerminateMode, isToolActivityError, + SubagentState, } from '../types.js'; import type { MessageBus } from '../../confirmation-bus/message-bus.js'; import { createBrowserAgentDefinition } from './browserAgentFactory.js'; @@ -123,7 +124,7 @@ export class BrowserAgentInvocation extends BaseToolInvocation< isSubagentProgress: true, agentName: this.agentName, recentActivity: [], - state: 'running', + state: SubagentState.RUNNING, }; updateOutput(initialProgress); } @@ -137,7 +138,7 @@ export class BrowserAgentInvocation extends BaseToolInvocation< id: randomUUID(), type: 'thought', content: sanitizedMsg, - status: 'completed', + status: SubagentState.COMPLETED, }); if (recentActivity.length > MAX_RECENT_ACTIVITY) { recentActivity = recentActivity.slice(-MAX_RECENT_ACTIVITY); @@ -146,7 +147,7 @@ export class BrowserAgentInvocation extends BaseToolInvocation< isSubagentProgress: true, agentName: this.agentName, recentActivity: [...recentActivity], - state: 'running', + state: SubagentState.RUNNING, } as SubagentProgress); } : undefined; @@ -175,7 +176,7 @@ export class BrowserAgentInvocation extends BaseToolInvocation< if ( lastItem && lastItem.type === 'thought' && - lastItem.status === 'running' + lastItem.status === SubagentState.RUNNING ) { lastItem.content = sanitizeThoughtContent(text); } else { @@ -183,7 +184,7 @@ export class BrowserAgentInvocation extends BaseToolInvocation< id: randomUUID(), type: 'thought', content: sanitizeThoughtContent(text), - status: 'running', + status: SubagentState.RUNNING, }); } updated = true; @@ -210,7 +211,7 @@ export class BrowserAgentInvocation extends BaseToolInvocation< displayName, description, args, - status: 'running', + status: SubagentState.RUNNING, }); updated = true; break; @@ -227,9 +228,11 @@ export class BrowserAgentInvocation extends BaseToolInvocation< recentActivity[i].type === 'tool_call' && callId != null && recentActivity[i].id === callId && - recentActivity[i].status === 'running' + recentActivity[i].status === SubagentState.RUNNING ) { - recentActivity[i].status = isError ? 'error' : 'completed'; + recentActivity[i].status = isError + ? SubagentState.ERROR + : SubagentState.COMPLETED; updated = true; break; } @@ -242,7 +245,9 @@ export class BrowserAgentInvocation extends BaseToolInvocation< const callId = activity.data['callId'] ? String(activity.data['callId']) : undefined; - const newStatus = isCancellation ? 'cancelled' : 'error'; + const newStatus = isCancellation + ? SubagentState.CANCELLED + : SubagentState.ERROR; if (callId) { // Mark the specific tool as error/cancelled @@ -250,7 +255,7 @@ export class BrowserAgentInvocation extends BaseToolInvocation< if ( recentActivity[i].type === 'tool_call' && recentActivity[i].id === callId && - recentActivity[i].status === 'running' + recentActivity[i].status === SubagentState.RUNNING ) { recentActivity[i].status = newStatus; updated = true; @@ -260,7 +265,10 @@ export class BrowserAgentInvocation extends BaseToolInvocation< } else { // No specific tool — mark ALL running tool_call items for (const item of recentActivity) { - if (item.type === 'tool_call' && item.status === 'running') { + if ( + item.type === 'tool_call' && + item.status === SubagentState.RUNNING + ) { item.status = newStatus; updated = true; } @@ -293,7 +301,7 @@ export class BrowserAgentInvocation extends BaseToolInvocation< isSubagentProgress: true, agentName: this.agentName, recentActivity: [...recentActivity], - state: 'running', + state: SubagentState.RUNNING, }; updateOutput(progress); } @@ -330,13 +338,13 @@ ${output.result}`; // GOAL = agent completed its task normally. // ABORTED = user cancelled. // Others (ERROR, MAX_TURNS, ERROR_NO_COMPLETE_TASK_CALL) = error. - let progressState: SubagentProgress['state']; + let progressState: SubagentState; if (output.terminate_reason === AgentTerminateMode.ABORTED) { - progressState = 'cancelled'; + progressState = SubagentState.CANCELLED; } else if (output.terminate_reason === AgentTerminateMode.GOAL) { - progressState = 'completed'; + progressState = SubagentState.COMPLETED; } else { - progressState = 'error'; + progressState = SubagentState.ERROR; } const progress: SubagentProgress = { @@ -366,8 +374,8 @@ ${output.result}`; // Mark any running items as error/cancelled for (const item of recentActivity) { - if (item.status === 'running') { - item.status = isAbort ? 'cancelled' : 'error'; + if (item.status === SubagentState.RUNNING) { + item.status = isAbort ? SubagentState.CANCELLED : SubagentState.ERROR; } } @@ -375,7 +383,7 @@ ${output.result}`; isSubagentProgress: true, agentName: this.agentName, recentActivity: [...recentActivity], - state: isAbort ? 'cancelled' : 'error', + state: isAbort ? SubagentState.CANCELLED : SubagentState.ERROR, }; if (updateOutput) { diff --git a/packages/core/src/agents/local-invocation.test.ts b/packages/core/src/agents/local-invocation.test.ts index eaea2b9ffa..297b46592e 100644 --- a/packages/core/src/agents/local-invocation.test.ts +++ b/packages/core/src/agents/local-invocation.test.ts @@ -21,6 +21,7 @@ import { type SubagentProgress, SubagentActivityErrorType, SUBAGENT_REJECTED_ERROR_PREFIX, + SubagentState, } from './types.js'; import { LocalSubagentInvocation } from './local-invocation.js'; import { LocalAgentExecutor } from './local-executor.js'; @@ -215,7 +216,7 @@ describe('LocalSubagentInvocation', () => { ]); const display = result.returnDisplay as SubagentProgress; expect(display.isSubagentProgress).toBe(true); - expect(display.state).toBe('completed'); + expect(display.state).toBe(SubagentState.COMPLETED); expect(display.result).toBe('Analysis complete.'); expect(display.terminateReason).toBe(AgentTerminateMode.GOAL); }); @@ -234,7 +235,7 @@ describe('LocalSubagentInvocation', () => { const display = result.returnDisplay as SubagentProgress; expect(display.isSubagentProgress).toBe(true); - expect(display.state).toBe('completed'); + expect(display.state).toBe(SubagentState.COMPLETED); expect(display.result).toBe('Partial progress...'); expect(display.terminateReason).toBe(AgentTerminateMode.TIMEOUT); }); @@ -340,7 +341,7 @@ describe('LocalSubagentInvocation', () => { expect.objectContaining({ type: 'thought', content: 'Error: Failed', - status: 'error', + status: SubagentState.ERROR, }), ); }); @@ -376,7 +377,7 @@ describe('LocalSubagentInvocation', () => { expect.objectContaining({ type: 'tool_call', content: 'ls', - status: 'error', + status: SubagentState.ERROR, }), ); }); @@ -418,7 +419,7 @@ describe('LocalSubagentInvocation', () => { expect.objectContaining({ type: 'tool_call', content: 'ls', - status: 'cancelled', + status: SubagentState.CANCELLED, }), ); }); @@ -443,7 +444,7 @@ describe('LocalSubagentInvocation', () => { expect(result.error).toBeUndefined(); const display = result.returnDisplay as SubagentProgress; expect(display.isSubagentProgress).toBe(true); - expect(display.state).toBe('completed'); + expect(display.state).toBe(SubagentState.COMPLETED); expect(display.result).toBe('Done'); }); @@ -466,7 +467,7 @@ describe('LocalSubagentInvocation', () => { expect.objectContaining({ type: 'thought', content: `Error: ${error.message}`, - status: 'error', + status: SubagentState.ERROR, }), ); }); @@ -488,7 +489,7 @@ describe('LocalSubagentInvocation', () => { expect(display.recentActivity).toContainEqual( expect.objectContaining({ content: `Error: ${creationError.message}`, - status: 'error', + status: SubagentState.ERROR, }), ); }); diff --git a/packages/core/src/agents/local-invocation.ts b/packages/core/src/agents/local-invocation.ts index 186f015979..f4d3153d79 100644 --- a/packages/core/src/agents/local-invocation.ts +++ b/packages/core/src/agents/local-invocation.ts @@ -23,6 +23,7 @@ import { SUBAGENT_REJECTED_ERROR_PREFIX, SUBAGENT_CANCELLED_ERROR_MESSAGE, isToolActivityError, + SubagentState, } from './types.js'; import { randomUUID } from 'node:crypto'; import type { z } from 'zod'; @@ -117,7 +118,7 @@ export class LocalSubagentInvocation extends BaseToolInvocation< isSubagentProgress: true, agentName: this.definition.name, recentActivity: [], - state: 'running', + state: SubagentState.RUNNING, }; updateOutput(initialProgress); } @@ -137,7 +138,7 @@ export class LocalSubagentInvocation extends BaseToolInvocation< if ( lastItem && lastItem.type === 'thought' && - lastItem.status === 'running' + lastItem.status === SubagentState.RUNNING ) { lastItem.content = sanitizeThoughtContent(text); } else { @@ -145,7 +146,7 @@ export class LocalSubagentInvocation extends BaseToolInvocation< id: randomUUID(), type: 'thought', content: sanitizeThoughtContent(text), - status: 'running', + status: SubagentState.RUNNING, }); } updated = true; @@ -174,7 +175,7 @@ export class LocalSubagentInvocation extends BaseToolInvocation< displayName, description, args, - status: 'running', + status: SubagentState.RUNNING, }); updated = true; @@ -193,9 +194,11 @@ export class LocalSubagentInvocation extends BaseToolInvocation< if ( recentActivity[i].type === 'tool_call' && recentActivity[i].content === name && - recentActivity[i].status === 'running' + recentActivity[i].status === SubagentState.RUNNING ) { - recentActivity[i].status = isError ? 'error' : 'completed'; + recentActivity[i].status = isError + ? SubagentState.ERROR + : SubagentState.COMPLETED; updated = true; this.publishActivity(recentActivity[i]); @@ -224,9 +227,9 @@ export class LocalSubagentInvocation extends BaseToolInvocation< if ( recentActivity[i].type === 'tool_call' && recentActivity[i].content === toolName && - recentActivity[i].status === 'running' + recentActivity[i].status === SubagentState.RUNNING ) { - recentActivity[i].status = 'cancelled'; + recentActivity[i].status = SubagentState.CANCELLED; updated = true; break; } @@ -237,9 +240,9 @@ export class LocalSubagentInvocation extends BaseToolInvocation< if ( recentActivity[i].type === 'tool_call' && recentActivity[i].content === toolName && - recentActivity[i].status === 'running' + recentActivity[i].status === SubagentState.RUNNING ) { - recentActivity[i].status = 'error'; + recentActivity[i].status = SubagentState.ERROR; updated = true; break; } @@ -253,7 +256,10 @@ export class LocalSubagentInvocation extends BaseToolInvocation< isCancellation || isRejection ? sanitizedError : `Error: ${sanitizedError}`, - status: isCancellation || isRejection ? 'cancelled' : 'error', + status: + isCancellation || isRejection + ? SubagentState.CANCELLED + : SubagentState.ERROR, }); updated = true; break; @@ -267,7 +273,7 @@ export class LocalSubagentInvocation extends BaseToolInvocation< isSubagentProgress: true, agentName: this.definition.name, recentActivity: [...recentActivity], // Copy to avoid mutation issues - state: 'running', + state: SubagentState.RUNNING, }; updateOutput(progress); @@ -287,7 +293,7 @@ export class LocalSubagentInvocation extends BaseToolInvocation< isSubagentProgress: true, agentName: this.definition.name, recentActivity: [...recentActivity], - state: 'cancelled', + state: SubagentState.CANCELLED, }; if (updateOutput) { @@ -303,7 +309,7 @@ export class LocalSubagentInvocation extends BaseToolInvocation< isSubagentProgress: true, agentName: this.definition.name, recentActivity: [...recentActivity], - state: 'completed', + state: SubagentState.COMPLETED, result: output.result, terminateReason: output.terminate_reason, }; @@ -334,8 +340,8 @@ ${output.result}`; // Mark any running items as error/cancelled for (const item of recentActivity) { - if (item.status === 'running') { - item.status = isAbort ? 'cancelled' : 'error'; + if (item.status === SubagentState.RUNNING) { + item.status = isAbort ? SubagentState.CANCELLED : SubagentState.ERROR; } } @@ -343,12 +349,12 @@ ${output.result}`; // But only if it's NOT an abort, or if we want to show "Cancelled" as a thought if (!isAbort) { const lastActivity = recentActivity[recentActivity.length - 1]; - if (!lastActivity || lastActivity.status !== 'error') { + if (!lastActivity || lastActivity.status !== SubagentState.ERROR) { recentActivity.push({ id: randomUUID(), type: 'thought', content: `Error: ${errorMessage}`, - status: 'error', + status: SubagentState.ERROR, }); // Maintain size limit // No limit on UI events sent via bus @@ -359,7 +365,7 @@ ${output.result}`; isSubagentProgress: true, agentName: this.definition.name, recentActivity: [...recentActivity], - state: isAbort ? 'cancelled' : 'error', + state: isAbort ? SubagentState.CANCELLED : SubagentState.ERROR, }; if (updateOutput) { diff --git a/packages/core/src/agents/remote-invocation.test.ts b/packages/core/src/agents/remote-invocation.test.ts index 0ec7774192..c2b89f49df 100644 --- a/packages/core/src/agents/remote-invocation.test.ts +++ b/packages/core/src/agents/remote-invocation.test.ts @@ -20,7 +20,11 @@ import { type A2AClientManager, } from './a2a-client-manager.js'; -import type { RemoteAgentDefinition, SubagentProgress } from './types.js'; +import { + type RemoteAgentDefinition, + type SubagentProgress, + SubagentState, +} from './types.js'; import { createMockMessageBus } from '../test-utils/mock-message-bus.js'; import { A2AAuthProviderFactory } from './auth-provider/factory.js'; import type { A2AAuthProvider } from './auth-provider/types.js'; @@ -268,7 +272,9 @@ describe('RemoteAgentInvocation', () => { abortSignal: new AbortController().signal, }); - expect(result.returnDisplay).toMatchObject({ state: 'error' }); + expect(result.returnDisplay).toMatchObject({ + state: SubagentState.ERROR, + }); expect((result.returnDisplay as SubagentProgress).result).toContain( "Failed to create auth provider for agent 'test-agent'", ); @@ -461,7 +467,7 @@ describe('RemoteAgentInvocation', () => { expect(updateOutput).toHaveBeenCalledWith( expect.objectContaining({ isSubagentProgress: true, - state: 'running', + state: SubagentState.RUNNING, recentActivity: expect.arrayContaining([ expect.objectContaining({ content: 'Working...' }), ]), @@ -470,7 +476,7 @@ describe('RemoteAgentInvocation', () => { expect(updateOutput).toHaveBeenCalledWith( expect.objectContaining({ isSubagentProgress: true, - state: 'completed', + state: SubagentState.COMPLETED, result: 'HelloHello World', }), ); @@ -508,7 +514,9 @@ describe('RemoteAgentInvocation', () => { abortSignal: controller.signal, }); - expect(result.returnDisplay).toMatchObject({ state: 'error' }); + expect(result.returnDisplay).toMatchObject({ + state: SubagentState.ERROR, + }); }); it('should handle errors gracefully', async () => { @@ -533,7 +541,7 @@ describe('RemoteAgentInvocation', () => { }); expect(result.returnDisplay).toMatchObject({ - state: 'error', + state: SubagentState.ERROR, result: expect.stringContaining('Network error'), }); }); @@ -616,7 +624,7 @@ describe('RemoteAgentInvocation', () => { expect(updateOutput).toHaveBeenCalledWith( expect.objectContaining({ isSubagentProgress: true, - state: 'running', + state: SubagentState.RUNNING, recentActivity: expect.arrayContaining([ expect.objectContaining({ content: 'Working...' }), ]), @@ -625,7 +633,7 @@ describe('RemoteAgentInvocation', () => { expect(updateOutput).toHaveBeenCalledWith( expect.objectContaining({ isSubagentProgress: true, - state: 'completed', + state: SubagentState.COMPLETED, result: 'Thinking...Final Answer', }), ); @@ -693,7 +701,7 @@ describe('RemoteAgentInvocation', () => { expect(updateOutput).toHaveBeenCalledWith( expect.objectContaining({ isSubagentProgress: true, - state: 'running', + state: SubagentState.RUNNING, recentActivity: expect.arrayContaining([ expect.objectContaining({ content: 'Working...' }), ]), @@ -702,7 +710,7 @@ describe('RemoteAgentInvocation', () => { expect(updateOutput).toHaveBeenCalledWith( expect.objectContaining({ isSubagentProgress: true, - state: 'completed', + state: SubagentState.COMPLETED, result: 'Generating...\n\nArtifact (Result):\nPart 1 Part 2', }), ); @@ -760,7 +768,9 @@ describe('RemoteAgentInvocation', () => { abortSignal: new AbortController().signal, }); - expect(result.returnDisplay).toMatchObject({ state: 'error' }); + expect(result.returnDisplay).toMatchObject({ + state: SubagentState.ERROR, + }); expect((result.returnDisplay as SubagentProgress).result).toContain( a2aError.userMessage, ); @@ -782,7 +792,9 @@ describe('RemoteAgentInvocation', () => { abortSignal: new AbortController().signal, }); - expect(result.returnDisplay).toMatchObject({ state: 'error' }); + expect(result.returnDisplay).toMatchObject({ + state: SubagentState.ERROR, + }); expect((result.returnDisplay as SubagentProgress).result).toContain( 'Error calling remote agent: something unexpected', ); @@ -813,7 +825,9 @@ describe('RemoteAgentInvocation', () => { abortSignal: new AbortController().signal, }); - expect(result.returnDisplay).toMatchObject({ state: 'error' }); + expect(result.returnDisplay).toMatchObject({ + state: SubagentState.ERROR, + }); // Should contain both the partial output and the error message expect(result.returnDisplay).toMatchObject({ result: expect.stringContaining('Partial response'), diff --git a/packages/core/src/agents/remote-invocation.ts b/packages/core/src/agents/remote-invocation.ts index e0869603fe..1510849683 100644 --- a/packages/core/src/agents/remote-invocation.ts +++ b/packages/core/src/agents/remote-invocation.ts @@ -17,6 +17,7 @@ import { type RemoteAgentDefinition, type AgentInputs, type SubagentProgress, + SubagentState, getAgentCardLoadOptions, getRemoteAgentTargetUrl, } from './types.js'; @@ -138,13 +139,13 @@ export class RemoteAgentInvocation extends BaseToolInvocation< updateOutput({ isSubagentProgress: true, agentName, - state: 'running', + state: SubagentState.RUNNING, recentActivity: [ { id: 'pending', type: 'thought', content: 'Working...', - status: 'running', + status: SubagentState.RUNNING, }, ], }); @@ -193,7 +194,7 @@ export class RemoteAgentInvocation extends BaseToolInvocation< updateOutput({ isSubagentProgress: true, agentName, - state: 'running', + state: SubagentState.RUNNING, recentActivity: reassembler.toActivityItems(), result: reassembler.toString(), }); @@ -225,7 +226,7 @@ export class RemoteAgentInvocation extends BaseToolInvocation< const finalProgress: SubagentProgress = { isSubagentProgress: true, agentName, - state: 'completed', + state: SubagentState.COMPLETED, result: finalOutput, recentActivity: reassembler.toActivityItems(), }; @@ -249,7 +250,7 @@ export class RemoteAgentInvocation extends BaseToolInvocation< const errorProgress: SubagentProgress = { isSubagentProgress: true, agentName, - state: 'error', + state: SubagentState.ERROR, result: fullDisplay, recentActivity: reassembler.toActivityItems(), }; diff --git a/packages/core/src/agents/remote-subagent-protocol.ts b/packages/core/src/agents/remote-subagent-protocol.ts index 4179e5587b..1231b0f068 100644 --- a/packages/core/src/agents/remote-subagent-protocol.ts +++ b/packages/core/src/agents/remote-subagent-protocol.ts @@ -28,6 +28,7 @@ import { DEFAULT_QUERY_STRING, type RemoteAgentDefinition, type SubagentProgress, + SubagentState, getRemoteAgentTargetUrl, getAgentCardLoadOptions, } from './types.js'; @@ -233,7 +234,7 @@ class RemoteSubagentProtocol implements AgentProtocol { this._latestProgress = { isSubagentProgress: true, agentName: this._agentName, - state: 'running', + state: SubagentState.RUNNING, recentActivity: reassembler.toActivityItems(), result: currentText, }; @@ -259,7 +260,7 @@ class RemoteSubagentProtocol implements AgentProtocol { const finalProgress: SubagentProgress = { isSubagentProgress: true, agentName: this._agentName, - state: 'completed', + state: SubagentState.COMPLETED, result: finalOutput, recentActivity: reassembler.toActivityItems(), }; diff --git a/packages/core/src/agents/types.ts b/packages/core/src/agents/types.ts index 86c9bec63b..7d99c10933 100644 --- a/packages/core/src/agents/types.ts +++ b/packages/core/src/agents/types.ts @@ -88,6 +88,13 @@ export interface SubagentActivityEvent { data: Record; } +export enum SubagentState { + RUNNING = 'running', + COMPLETED = 'completed', + ERROR = 'error', + CANCELLED = 'cancelled', +} + export interface SubagentActivityItem { id: string; type: 'thought' | 'tool_call'; @@ -95,14 +102,14 @@ export interface SubagentActivityItem { displayName?: string; description?: string; args?: string; - status: 'running' | 'completed' | 'error' | 'cancelled'; + status: SubagentState; } export interface SubagentProgress { isSubagentProgress: true; agentName: string; recentActivity: SubagentActivityItem[]; - state?: 'running' | 'completed' | 'error' | 'cancelled'; + state?: SubagentState; result?: string; terminateReason?: AgentTerminateMode; }