diff --git a/packages/core/src/agents/remote-session-invocation.test.ts b/packages/core/src/agents/remote-session-invocation.test.ts index d611dd7399..555af74ad0 100644 --- a/packages/core/src/agents/remote-session-invocation.test.ts +++ b/packages/core/src/agents/remote-session-invocation.test.ts @@ -1,6 +1,6 @@ /** * @license - * Copyright 2025 Google LLC + * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ @@ -354,7 +354,22 @@ describe('RemoteSessionInvocation', () => { it('should handle abort gracefully', async () => { const controller = new AbortController(); - const { mockSession } = setupMockSession(); + const partialProgress: SubagentProgress = { + isSubagentProgress: true, + agentName: 'Test Agent', + state: SubagentState.RUNNING, + result: '', + recentActivity: [ + { + id: 'a1', + type: 'thought', + content: 'Thinking...', + status: SubagentState.RUNNING, + }, + ], + }; + + const { mockSession } = setupMockSession({ progress: partialProgress }); // When getResult resolves, the signal will already be aborted mockSession.getResult.mockImplementation(async () => { @@ -378,7 +393,10 @@ describe('RemoteSessionInvocation', () => { updateOutput, }); - expect(result.returnDisplay).toMatchObject({ state: 'error' }); + expect(result.returnDisplay).toMatchObject({ state: 'cancelled' }); + expect( + (result.returnDisplay as SubagentProgress).recentActivity[0].status, + ).toBe(SubagentState.CANCELLED); expect(result.llmContent).toEqual([ { text: 'Operation cancelled by user' }, ]); @@ -452,9 +470,10 @@ describe('RemoteSessionInvocation', () => { // Should contain both the partial output and the error expect(display.result).toContain('Partial work so far'); expect(display.result).toContain('mid-stream error'); - // Should preserve partial activity + // Should preserve and update partial activity status to ERROR expect(display.recentActivity).toHaveLength(1); expect(display.recentActivity[0].content).toBe('Thinking...'); + expect(display.recentActivity[0].status).toBe(SubagentState.ERROR); }); it('should clean up listeners in finally', async () => { diff --git a/packages/core/src/agents/remote-session-invocation.ts b/packages/core/src/agents/remote-session-invocation.ts index 2434f6fc30..6ff32ba41a 100644 --- a/packages/core/src/agents/remote-session-invocation.ts +++ b/packages/core/src/agents/remote-session-invocation.ts @@ -1,6 +1,6 @@ /** * @license - * Copyright 2025 Google LLC + * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ @@ -17,6 +17,7 @@ import { type RemoteAgentDefinition, type AgentInputs, type SubagentProgress, + type SubagentActivityItem, SubagentState, getRemoteAgentTargetUrl, } from './types.js'; @@ -116,6 +117,7 @@ export class RemoteSessionInvocation extends BaseToolInvocation< async execute(options: ExecuteOptions): Promise { const { abortSignal: _signal, updateOutput } = options; const agentName = this.definition.displayName ?? this.definition.name; + const emptyActivity: SubagentActivityItem[] = []; // Seed session with prior A2A conversation state const stateKey = RemoteSessionInvocation.sessionKey(this.definition); @@ -172,15 +174,19 @@ export class RemoteSessionInvocation extends BaseToolInvocation< // rejecting. Detect this and surface proper error state. if (_signal?.aborted) { const partialProgress = session.getLatestProgress(); + const recentActivity = this.stopRunningActivities( + partialProgress?.recentActivity ?? emptyActivity, + SubagentState.CANCELLED, + ); const errorProgress: SubagentProgress = { isSubagentProgress: true, agentName, - state: SubagentState.ERROR, + state: SubagentState.CANCELLED, result: typeof partialProgress?.result === 'string' ? partialProgress.result : '', - recentActivity: partialProgress?.recentActivity ?? [], + recentActivity, }; if (updateOutput) updateOutput(errorProgress); return { @@ -207,12 +213,22 @@ export class RemoteSessionInvocation extends BaseToolInvocation< ? `${partialOutput}\n\n${errorMessage}` : errorMessage; + const isAbort = + (error instanceof Error && error.name === 'AbortError') || + errorMessage.includes('Aborted'); + + const status = isAbort ? SubagentState.CANCELLED : SubagentState.ERROR; + const recentActivity = this.stopRunningActivities( + partialProgress?.recentActivity ?? emptyActivity, + status, + ); + const errorProgress: SubagentProgress = { isSubagentProgress: true, agentName, - state: SubagentState.ERROR, + state: status, result: fullDisplay, - recentActivity: partialProgress?.recentActivity ?? [], + recentActivity, }; if (updateOutput) { @@ -235,6 +251,19 @@ export class RemoteSessionInvocation extends BaseToolInvocation< } } + private stopRunningActivities( + activity: SubagentActivityItem[], + status: SubagentState, + ): SubagentActivityItem[] { + const result: SubagentActivityItem[] = []; + for (const item of activity) { + result.push( + item.status === SubagentState.RUNNING ? { ...item, status } : item, + ); + } + return result; + } + /** * Formats an execution error into a user-friendly message. * Recognizes typed A2AAgentError subclasses and falls back to