From 5dda532573ec98d12d2db0f819d84ee18c532721 Mon Sep 17 00:00:00 2001 From: Aishanee Shah Date: Mon, 11 May 2026 22:14:59 +0000 Subject: [PATCH] feat: include full message records to link subagent trajectories --- packages/cli/src/ui/commands/bugCommand.test.ts | 6 ++++++ packages/cli/src/ui/commands/bugCommand.ts | 2 ++ .../cli/src/ui/commands/chatCommand.test.ts | 13 ++++++++++++- packages/cli/src/ui/commands/chatCommand.ts | 9 +++++++-- packages/cli/src/ui/utils/historyExportUtils.ts | 17 ++++++++++++++--- packages/core/src/core/geminiChat.ts | 8 ++++++++ packages/core/src/core/logger.ts | 6 +++++- 7 files changed, 54 insertions(+), 7 deletions(-) diff --git a/packages/cli/src/ui/commands/bugCommand.test.ts b/packages/cli/src/ui/commands/bugCommand.test.ts index f47ba0716c..a9c868639c 100644 --- a/packages/cli/src/ui/commands/bugCommand.test.ts +++ b/packages/cli/src/ui/commands/bugCommand.test.ts @@ -162,6 +162,7 @@ describe('bugCommand', () => { { role: 'model', parts: [{ text: 'hi' }] }, ]; const mockGetSubagentTrajectories = vi.fn().mockResolvedValue({}); + const mockGetConversation = vi.fn().mockReturnValue({ messages: [] }); const mockContext = createMockCommandContext({ services: { agentContext: { @@ -179,6 +180,7 @@ describe('bugCommand', () => { getChat: () => ({ getHistory: () => history, getSubagentTrajectories: mockGetSubagentTrajectories, + getConversation: mockGetConversation, }), }, }, @@ -194,6 +196,7 @@ describe('bugCommand', () => { ); expect(exportHistoryToFile).toHaveBeenCalledWith({ history, + messages: [], filePath: expectedPath, trajectories: {}, }); @@ -222,6 +225,7 @@ describe('bugCommand', () => { } as unknown as ConversationRecord, }; const mockGetSubagentTrajectories = vi.fn().mockResolvedValue(trajectories); + const mockGetConversation = vi.fn().mockReturnValue({ messages: [] }); const mockContext = createMockCommandContext({ services: { @@ -240,6 +244,7 @@ describe('bugCommand', () => { getChat: () => ({ getHistory: () => history, getSubagentTrajectories: mockGetSubagentTrajectories, + getConversation: mockGetConversation, }), }, }, @@ -256,6 +261,7 @@ describe('bugCommand', () => { expect(mockGetSubagentTrajectories).toHaveBeenCalled(); expect(exportHistoryToFile).toHaveBeenCalledWith({ history, + messages: [], filePath: expectedPath, trajectories, }); diff --git a/packages/cli/src/ui/commands/bugCommand.ts b/packages/cli/src/ui/commands/bugCommand.ts index ed226636cc..5ab6d4ad9f 100644 --- a/packages/cli/src/ui/commands/bugCommand.ts +++ b/packages/cli/src/ui/commands/bugCommand.ts @@ -89,8 +89,10 @@ export const bugCommand: SlashCommand = { const historyFilePath = path.join(tempDir, historyFileName); try { const trajectories = await chat?.getSubagentTrajectories(); + const messages = chat?.getConversation()?.messages; await exportHistoryToFile({ history, + messages, filePath: historyFilePath, trajectories, }); diff --git a/packages/cli/src/ui/commands/chatCommand.test.ts b/packages/cli/src/ui/commands/chatCommand.test.ts index 5a7e53c690..b735990827 100644 --- a/packages/cli/src/ui/commands/chatCommand.test.ts +++ b/packages/cli/src/ui/commands/chatCommand.test.ts @@ -64,6 +64,7 @@ describe('chatCommand', () => { mockGetChat = vi.fn().mockReturnValue({ getHistory: mockGetHistory, getSubagentTrajectories: vi.fn().mockResolvedValue({}), + getConversation: vi.fn().mockReturnValue({ messages: [] }), }); mockSaveCheckpoint = vi.fn().mockResolvedValue(undefined); mockLoadCheckpoint = vi.fn().mockResolvedValue({ history: [] }); @@ -231,7 +232,12 @@ describe('chatCommand', () => { expect(mockCheckpointExists).not.toHaveBeenCalled(); // Should skip existence check expect(mockSaveCheckpoint).toHaveBeenCalledWith( - { history, authType: AuthType.LOGIN_WITH_GOOGLE, trajectories: {} }, + { + history, + authType: AuthType.LOGIN_WITH_GOOGLE, + trajectories: {}, + messages: [], + }, tag, ); expect(result).toEqual({ @@ -465,6 +471,7 @@ describe('chatCommand', () => { ); expect(mockExport).toHaveBeenCalledWith({ history: mockHistory, + messages: [], filePath: expectedPath, trajectories: {}, }); @@ -481,6 +488,7 @@ describe('chatCommand', () => { const expectedPath = path.join(process.cwd(), 'my-chat.json'); expect(mockExport).toHaveBeenCalledWith({ history: mockHistory, + messages: [], filePath: expectedPath, trajectories: {}, }); @@ -497,6 +505,7 @@ describe('chatCommand', () => { const expectedPath = path.join(process.cwd(), 'my-chat.md'); expect(mockExport).toHaveBeenCalledWith({ history: mockHistory, + messages: [], filePath: expectedPath, trajectories: {}, }); @@ -548,6 +557,7 @@ describe('chatCommand', () => { const expectedPath = path.join(process.cwd(), 'my-chat.json'); expect(mockExport).toHaveBeenCalledWith({ history: mockHistory, + messages: [], filePath: expectedPath, trajectories: {}, }); @@ -559,6 +569,7 @@ describe('chatCommand', () => { const expectedPath = path.join(process.cwd(), 'my-chat.md'); expect(mockExport).toHaveBeenCalledWith({ history: mockHistory, + messages: [], filePath: expectedPath, trajectories: {}, }); diff --git a/packages/cli/src/ui/commands/chatCommand.ts b/packages/cli/src/ui/commands/chatCommand.ts index dd67ae1e35..cfb2be3961 100644 --- a/packages/cli/src/ui/commands/chatCommand.ts +++ b/packages/cli/src/ui/commands/chatCommand.ts @@ -140,7 +140,11 @@ const saveCommand: SlashCommand = { if (history.length > INITIAL_HISTORY_LENGTH) { const authType = config?.getContentGeneratorConfig()?.authType; const trajectories = await chat.getSubagentTrajectories(); - await logger.saveCheckpoint({ history, authType, trajectories }, tag); + const messages = chat.getConversation()?.messages; + await logger.saveCheckpoint( + { history, authType, trajectories, messages }, + tag, + ); return { type: 'message', messageType: 'info', @@ -326,7 +330,8 @@ const shareCommand: SlashCommand = { try { const trajectories = await chat.getSubagentTrajectories(); - await exportHistoryToFile({ history, filePath, trajectories }); + const messages = chat.getConversation()?.messages; + await exportHistoryToFile({ history, filePath, trajectories, messages }); return { type: 'message', messageType: 'info', diff --git a/packages/cli/src/ui/utils/historyExportUtils.ts b/packages/cli/src/ui/utils/historyExportUtils.ts index cff446d83f..2c20bcb4af 100644 --- a/packages/cli/src/ui/utils/historyExportUtils.ts +++ b/packages/cli/src/ui/utils/historyExportUtils.ts @@ -7,7 +7,10 @@ import * as fsPromises from 'node:fs/promises'; import path from 'node:path'; import type { Content } from '@google/genai'; -import type { ConversationRecord } from '@google/gemini-cli-core'; +import type { + ConversationRecord, + MessageRecord, +} from '@google/gemini-cli-core'; /** * Serializes chat history to a Markdown string. @@ -52,9 +55,17 @@ export function serializeHistoryToMarkdown( * Options for exporting chat history. */ export interface ExportHistoryOptions { + /** The standard history array used for model requests. */ history: readonly Content[]; + /** The file path to export to. */ filePath: string; + /** Optional subagent trajectories to include. */ trajectories?: Record; + /** + * Optional full message records which contain metadata like agentId for tool calls, + * providing the link between history and trajectories. + */ + messages?: MessageRecord[]; } /** @@ -63,13 +74,13 @@ export interface ExportHistoryOptions { export async function exportHistoryToFile( options: ExportHistoryOptions, ): Promise { - const { history, filePath, trajectories } = options; + const { history, filePath, trajectories, messages } = options; const extension = path.extname(filePath).toLowerCase(); let content: string; if (extension === '.json') { if (trajectories && Object.keys(trajectories).length > 0) { - content = JSON.stringify({ history, trajectories }, null, 2); + content = JSON.stringify({ history, messages, trajectories }, null, 2); } else { content = JSON.stringify(history, null, 2); } diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 0052ac3609..dccbd282e2 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -1238,8 +1238,16 @@ export class GeminiChat { return this.chatRecordingService.getSubagentTrajectories(); } + /** + * Gets the current conversation record. + */ + getConversation(): ConversationRecord | null { + return this.chatRecordingService.getConversation(); + } + /** * Records completed tool calls with full metadata. + * This is called by external components when tool calls complete, before sending responses to Gemini. */ recordCompletedToolCalls( diff --git a/packages/core/src/core/logger.ts b/packages/core/src/core/logger.ts index 58a73fdd68..81852ab4c2 100644 --- a/packages/core/src/core/logger.ts +++ b/packages/core/src/core/logger.ts @@ -11,7 +11,10 @@ import type { AuthType } from './contentGenerator.js'; import type { Storage } from '../config/storage.js'; import { debugLogger } from '../utils/debugLogger.js'; import { coreEvents } from '../utils/events.js'; -import type { ConversationRecord } from '../services/chatRecordingService.js'; +import { + type ConversationRecord, + type MessageRecord, +} from '../services/chatRecordingService.js'; const LOG_FILE_NAME = 'logs.json'; @@ -31,6 +34,7 @@ export interface Checkpoint { history: readonly Content[]; authType?: AuthType; trajectories?: Record; + messages?: MessageRecord[]; } // This regex matches any character that is NOT a letter (a-z, A-Z),