From 7d1de3bccc42a387def2bf169d9b4f256b9e6480 Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Fri, 10 Apr 2026 12:47:25 -0400 Subject: [PATCH] feat(core): persist subagent agentId in tool call records (#25092) --- packages/core/src/agents/local-executor.ts | 2 +- .../core/src/agents/local-invocation.test.ts | 1 + packages/core/src/agents/local-invocation.ts | 9 +++++- packages/core/src/core/geminiChat.ts | 4 +++ .../src/services/chatRecordingService.test.ts | 28 +++++++++++++++++++ .../core/src/services/chatRecordingTypes.ts | 1 + 6 files changed, 43 insertions(+), 2 deletions(-) diff --git a/packages/core/src/agents/local-executor.ts b/packages/core/src/agents/local-executor.ts index cfae92c870..e7d8078579 100644 --- a/packages/core/src/agents/local-executor.ts +++ b/packages/core/src/agents/local-executor.ts @@ -114,7 +114,7 @@ export function createUnauthorizedToolError(toolName: string): string { export class LocalAgentExecutor { readonly definition: LocalAgentDefinition; - private readonly agentId: string; + readonly agentId: string; private readonly toolRegistry: ToolRegistry; private readonly promptRegistry: PromptRegistry; private readonly resourceRegistry: ResourceRegistry; diff --git a/packages/core/src/agents/local-invocation.test.ts b/packages/core/src/agents/local-invocation.test.ts index 592bcb59e8..854a32ec64 100644 --- a/packages/core/src/agents/local-invocation.test.ts +++ b/packages/core/src/agents/local-invocation.test.ts @@ -79,6 +79,7 @@ describe('LocalSubagentInvocation', () => { mockExecutorInstance = { run: vi.fn(), definition: testDefinition, + agentId: 'test-agent-id', } as unknown as Mocked>; MockLocalAgentExecutor.create.mockResolvedValue( diff --git a/packages/core/src/agents/local-invocation.ts b/packages/core/src/agents/local-invocation.ts index 771be7b68a..228d5010ec 100644 --- a/packages/core/src/agents/local-invocation.ts +++ b/packages/core/src/agents/local-invocation.ts @@ -25,12 +25,14 @@ import { isToolActivityError, } from './types.js'; import { randomUUID } from 'node:crypto'; +import type { z } from 'zod'; import type { MessageBus } from '../confirmation-bus/message-bus.js'; import { sanitizeThoughtContent, sanitizeToolArgs, sanitizeErrorMessage, } from '../utils/agent-sanitization-utils.js'; +import { debugLogger } from '../utils/debugLogger.js'; const INPUT_PREVIEW_MAX_LENGTH = 50; const DESCRIPTION_MAX_LENGTH = 200; @@ -108,6 +110,7 @@ export class LocalSubagentInvocation extends BaseToolInvocation< updateOutput?: (output: ToolLiveOutput) => void, ): Promise { const recentActivity: SubagentActivityItem[] = []; + let executor: LocalAgentExecutor | undefined; try { if (updateOutput) { @@ -273,7 +276,7 @@ export class LocalSubagentInvocation extends BaseToolInvocation< } }; - const executor = await LocalAgentExecutor.create( + executor = await LocalAgentExecutor.create( this.definition, this.context, onActivity, @@ -319,11 +322,14 @@ ${output.result}`; return { llmContent: [{ text: resultContent }], returnDisplay: progress, + data: { agentId: executor.agentId }, }; } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error); + debugLogger.error(`Subagent '${this.definition.name}' failed:`, error); + const isAbort = (error instanceof Error && error.name === 'AbortError') || errorMessage.includes('Aborted'); @@ -369,6 +375,7 @@ ${output.result}`; return { llmContent: `Subagent '${this.definition.name}' failed. Error: ${errorMessage}`, returnDisplay: progress, + data: executor ? { agentId: executor.agentId } : undefined, // We omit the 'error' property so that the UI renders our rich returnDisplay // instead of the raw error message. The llmContent still informs the agent of the failure. }; diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index f5ee37e565..c480c3800b 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -1050,6 +1050,10 @@ export class GeminiChat { result: call.response?.responseParts || null, status: call.status, timestamp: new Date().toISOString(), + agentId: + typeof call.response?.data?.['agentId'] === 'string' + ? call.response.data['agentId'] + : undefined, resultDisplay, description: 'invocation' in call ? call.invocation?.getDescription() : undefined, diff --git a/packages/core/src/services/chatRecordingService.test.ts b/packages/core/src/services/chatRecordingService.test.ts index 22ba6c2c03..94b9c61c7a 100644 --- a/packages/core/src/services/chatRecordingService.test.ts +++ b/packages/core/src/services/chatRecordingService.test.ts @@ -536,6 +536,34 @@ describe('ChatRecordingService', () => { .toolCalls, ).toHaveLength(1); }); + + it('should record agentId when provided', async () => { + chatRecordingService.recordMessage({ + type: 'gemini', + content: '', + model: 'gemini-pro', + }); + + const toolCall: ToolCallRecord = { + id: 'tool-1', + name: 'testTool', + args: {}, + status: CoreToolCallStatus.Success, + timestamp: new Date().toISOString(), + agentId: 'test-agent-id', + }; + chatRecordingService.recordToolCalls('gemini-pro', [toolCall]); + + const sessionFile = chatRecordingService.getConversationFilePath()!; + const conversation = (await loadConversationRecord( + sessionFile, + )) as ConversationRecord; + const geminiMsg = conversation.messages[0] as MessageRecord & { + type: 'gemini'; + }; + expect(geminiMsg.toolCalls).toHaveLength(1); + expect(geminiMsg.toolCalls![0].agentId).toBe('test-agent-id'); + }); }); describe('deleteSession', () => { diff --git a/packages/core/src/services/chatRecordingTypes.ts b/packages/core/src/services/chatRecordingTypes.ts index c2564c0eec..2ddc218bdc 100644 --- a/packages/core/src/services/chatRecordingTypes.ts +++ b/packages/core/src/services/chatRecordingTypes.ts @@ -45,6 +45,7 @@ export interface ToolCallRecord { result?: PartListUnion | null; status: Status; timestamp: string; + agentId?: string; // UI-specific fields for display purposes displayName?: string; description?: string;