feat: foundation for subagent trajectories (Stage 1)

This commit is contained in:
Aishanee Shah
2026-05-14 16:39:47 +00:00
parent 488d71b8c9
commit a9648de39f
12 changed files with 298 additions and 39 deletions

View File

@@ -173,6 +173,8 @@ describe('bugCommand', () => {
geminiClient: {
getChat: () => ({
getHistory: () => history,
getSubagentTrajectories: vi.fn().mockResolvedValue({}),
getConversation: vi.fn().mockReturnValue({ messages: [] }),
}),
},
},
@@ -187,8 +189,10 @@ describe('bugCommand', () => {
'bug-report-history-1704067200000.json',
);
expect(exportHistoryToFile).toHaveBeenCalledWith({
history,
messages: [],
filePath: expectedPath,
trajectories: {},
history,
});
const addItemCall = vi.mocked(mockContext.ui.addItem).mock.calls[0];

View File

@@ -88,7 +88,14 @@ export const bugCommand: SlashCommand = {
const historyFileName = `bug-report-history-${Date.now()}.json`;
const historyFilePath = path.join(tempDir, historyFileName);
try {
await exportHistoryToFile({ history, filePath: historyFilePath });
const trajectories = await chat?.getSubagentTrajectories();
const messages = chat?.getConversation()?.messages ?? [];
await exportHistoryToFile({
messages,
filePath: historyFilePath,
trajectories,
history,
});
historyFileMessage = `\n\n--------------------------------------------------------------------------------\n\n📄 **Chat History Exported**\nTo help us debug, we've exported your current chat history to:\n${historyFilePath}\n\nPlease consider attaching this file to your GitHub issue if you feel comfortable doing so.\n\n**Privacy Disclaimer:** Please do not upload any logs containing sensitive or private information that you are not comfortable sharing publicly.`;
problemValue += `\n\n[ACTION REQUIRED] 📎 PLEASE ATTACH THE EXPORTED CHAT HISTORY JSON FILE TO THIS ISSUE IF YOU FEEL COMFORTABLE SHARING IT.`;
} catch (err) {

View File

@@ -63,6 +63,8 @@ describe('chatCommand', () => {
mockGetHistory = vi.fn().mockReturnValue([]);
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: [] });
@@ -191,6 +193,15 @@ describe('chatCommand', () => {
{ role: 'user', parts: [{ text: 'Hello, how are you?' }] },
]);
result = await saveCommand?.action?.(mockContext, tag);
expect(mockSaveCheckpoint).toHaveBeenCalledWith(
{
history: expect.any(Array),
authType: AuthType.LOGIN_WITH_GOOGLE,
trajectories: {},
messages: [],
},
tag,
);
expect(result).toEqual({
type: 'message',
messageType: 'info',
@@ -230,7 +241,12 @@ describe('chatCommand', () => {
expect(mockCheckpointExists).not.toHaveBeenCalled(); // Should skip existence check
expect(mockSaveCheckpoint).toHaveBeenCalledWith(
{ history, authType: AuthType.LOGIN_WITH_GOOGLE },
{
history,
authType: AuthType.LOGIN_WITH_GOOGLE,
trajectories: {},
messages: [],
},
tag,
);
expect(result).toEqual({
@@ -292,6 +308,8 @@ describe('chatCommand', () => {
{ type: 'gemini', text: 'hello world' },
] as HistoryItemWithoutId[],
clientHistory: conversation,
messages: undefined,
version: undefined,
});
});
@@ -332,6 +350,8 @@ describe('chatCommand', () => {
{ type: 'gemini', text: 'hello world' },
] as HistoryItemWithoutId[],
clientHistory: conversation,
messages: undefined,
version: undefined,
});
});
@@ -463,8 +483,10 @@ describe('chatCommand', () => {
'gemini-conversation-1234567890.json',
);
expect(mockExport).toHaveBeenCalledWith({
history: mockHistory,
messages: [],
filePath: expectedPath,
trajectories: {},
history: mockHistory,
});
expect(result).toEqual({
type: 'message',
@@ -478,8 +500,10 @@ describe('chatCommand', () => {
const result = await shareCommand?.action?.(mockContext, filePath);
const expectedPath = path.join(process.cwd(), 'my-chat.json');
expect(mockExport).toHaveBeenCalledWith({
history: mockHistory,
messages: [],
filePath: expectedPath,
trajectories: {},
history: mockHistory,
});
expect(result).toEqual({
type: 'message',
@@ -493,8 +517,10 @@ describe('chatCommand', () => {
const result = await shareCommand?.action?.(mockContext, filePath);
const expectedPath = path.join(process.cwd(), 'my-chat.md');
expect(mockExport).toHaveBeenCalledWith({
history: mockHistory,
messages: [],
filePath: expectedPath,
trajectories: {},
history: mockHistory,
});
expect(result).toEqual({
type: 'message',
@@ -543,8 +569,10 @@ describe('chatCommand', () => {
await shareCommand?.action?.(mockContext, filePath);
const expectedPath = path.join(process.cwd(), 'my-chat.json');
expect(mockExport).toHaveBeenCalledWith({
history: mockHistory,
messages: [],
filePath: expectedPath,
trajectories: {},
history: mockHistory,
});
});
@@ -553,8 +581,10 @@ describe('chatCommand', () => {
await shareCommand?.action?.(mockContext, filePath);
const expectedPath = path.join(process.cwd(), 'my-chat.md');
expect(mockExport).toHaveBeenCalledWith({
history: mockHistory,
messages: [],
filePath: expectedPath,
trajectories: {},
history: mockHistory,
});
});
});

View File

@@ -139,7 +139,12 @@ const saveCommand: SlashCommand = {
const history = chat.getHistory();
if (history.length > INITIAL_HISTORY_LENGTH) {
const authType = config?.getContentGeneratorConfig()?.authType;
await logger.saveCheckpoint({ history, authType }, tag);
const trajectories = await chat.getSubagentTrajectories();
const messages = chat.getConversation()?.messages ?? [];
await logger.saveCheckpoint(
{ history, authType, trajectories, messages },
tag,
);
return {
type: 'message',
messageType: 'info',
@@ -324,7 +329,9 @@ const shareCommand: SlashCommand = {
}
try {
await exportHistoryToFile({ history, filePath });
const trajectories = await chat.getSubagentTrajectories();
const messages = chat.getConversation()?.messages ?? [];
await exportHistoryToFile({ messages, filePath, trajectories, history });
return {
type: 'message',
messageType: 'info',

View File

@@ -7,6 +7,11 @@
import * as fsPromises from 'node:fs/promises';
import path from 'node:path';
import type { Content } from '@google/genai';
import {
type ConversationRecord,
type MessageRecord,
reconstructHistory,
} from '@google/gemini-cli-core';
/**
* Serializes chat history to a Markdown string.
@@ -51,8 +56,20 @@ export function serializeHistoryToMarkdown(
* Options for exporting chat history.
*/
export interface ExportHistoryOptions {
history: readonly Content[];
/**
* Optional full message records which contain metadata like agentId for tool calls,
* providing the link between history and trajectories.
*/
messages?: MessageRecord[];
/** The file path to export to. */
filePath: string;
/** Optional subagent trajectories to include. */
trajectories?: Record<string, ConversationRecord>;
/**
* Optional standard history array used for model requests.
* If provided, it is used for Markdown export to avoid reconstruction.
*/
history?: readonly Content[];
}
/**
@@ -61,13 +78,21 @@ export interface ExportHistoryOptions {
export async function exportHistoryToFile(
options: ExportHistoryOptions,
): Promise<void> {
const { history, filePath } = options;
const {
messages,
filePath,
trajectories: _trajectories, // Collected but not yet included in Stage 1 JSON output
history: providedHistory,
} = options;
const extension = path.extname(filePath).toLowerCase();
let content: string;
if (extension === '.json') {
content = JSON.stringify(history, null, 2);
// Stage 1 & 2: Maintain legacy behavior - only export the raw history array.
// Trajectories and messages are collected but not yet included in Stage 2 JSON output.
content = JSON.stringify(providedHistory ?? [], null, 2);
} else if (extension === '.md') {
const history = providedHistory ?? reconstructHistory(messages ?? []);
content = serializeHistoryToMarkdown(history);
} else {
throw new Error(

View File

@@ -40,6 +40,7 @@ import { tokenLimit } from './tokenLimits.js';
import type {
ChatRecordingService,
ResumedSessionData,
MessageRecord,
} from '../services/chatRecordingService.js';
import type { ContentGenerator } from './contentGenerator.js';
import { LoopDetectionService } from '../services/loopDetectionService.js';
@@ -337,8 +338,9 @@ export class GeminiClient {
async resumeChat(
history: Content[],
resumedSessionData?: ResumedSessionData,
messages?: MessageRecord[],
): Promise<void> {
this.chat = await this.startChat(history, resumedSessionData);
this.chat = await this.startChat(history, resumedSessionData, messages);
this.updateTelemetryTokenCount();
}
@@ -378,6 +380,7 @@ export class GeminiClient {
async startChat(
extraHistory?: Content[],
resumedSessionData?: ResumedSessionData,
messages?: MessageRecord[],
): Promise<GeminiChat> {
this.forceFullIdeContext = true;
this.hasFailedCompressionAttempt = false;
@@ -407,8 +410,9 @@ export class GeminiClient {
toolRegistry.getFunctionDeclarations(modelId);
return [{ functionDeclarations: toolDeclarations }];
},
messages,
);
await chat.initialize(resumedSessionData, 'main');
await chat.initialize(resumedSessionData, 'main', messages);
this.contextManager = await initializeContextManager(
this.config,
chat,

View File

@@ -39,6 +39,8 @@ import {
import {
ChatRecordingService,
type ResumedSessionData,
type MessageRecord,
type ConversationRecord,
} from '../services/chatRecordingService.js';
import {
ContentRetryEvent,
@@ -266,6 +268,7 @@ export class GeminiChat {
private readonly chatRecordingService: ChatRecordingService;
private lastPromptTokenCount: number;
private callCounter = 0;
private initialMessages?: MessageRecord[];
agentHistory: AgentChatHistory;
constructor(
@@ -275,8 +278,10 @@ export class GeminiChat {
history: Content[] = [],
resumedSessionData?: ResumedSessionData,
private readonly onModelChanged?: (modelId: string) => Promise<Tool[]>,
messages?: MessageRecord[],
) {
validateHistory(history);
this.initialMessages = messages;
this.agentHistory = new AgentChatHistory(history);
this.chatRecordingService = new ChatRecordingService(context);
this.lastPromptTokenCount = estimateTokenCountSync(
@@ -291,14 +296,31 @@ export class GeminiChat {
async initialize(
resumedSessionData?: ResumedSessionData,
kind: 'main' | 'subagent' = 'main',
messages?: MessageRecord[],
) {
const messagesToUse = messages ?? this.initialMessages;
await this.chatRecordingService.initialize(resumedSessionData, kind);
if (messagesToUse) {
this.chatRecordingService.resetMessages(messagesToUse);
}
}
setSystemInstruction(sysInstr: string) {
this.systemInstruction = sysInstr;
}
getConversation(): ConversationRecord | null {
return this.chatRecordingService.getConversation();
}
getChatRecordingService(): ChatRecordingService {
return this.chatRecordingService;
}
async getSubagentTrajectories(): Promise<Record<string, ConversationRecord>> {
return this.chatRecordingService.getSubagentTrajectories();
}
/**
* Sends a message to the model and returns the response in chunks.
*
@@ -1215,13 +1237,6 @@ export class GeminiChat {
return this.lastPromptTokenCount;
}
/**
* Gets the chat recording service instance.
*/
getChatRecordingService(): ChatRecordingService {
return this.chatRecordingService;
}
/**
* Records completed tool calls with full metadata.
* This is called by external components when tool calls complete, before sending responses to Gemini.

View File

@@ -437,7 +437,11 @@ describe('Logger', () => {
},
])('should save a checkpoint', async ({ tag, encodedTag }) => {
await logger.saveCheckpoint(
{ history: conversation, authType: AuthType.LOGIN_WITH_GOOGLE },
{
history: conversation,
messages: [],
authType: AuthType.LOGIN_WITH_GOOGLE,
},
tag,
);
const taggedFilePath = path.join(
@@ -447,6 +451,7 @@ describe('Logger', () => {
const fileContent = await fs.readFile(taggedFilePath, 'utf-8');
expect(JSON.parse(fileContent)).toEqual({
history: conversation,
messages: [],
authType: AuthType.LOGIN_WITH_GOOGLE,
});
});
@@ -462,7 +467,10 @@ describe('Logger', () => {
.mockImplementation(() => {});
await expect(
uninitializedLogger.saveCheckpoint({ history: conversation }, 'tag'),
uninitializedLogger.saveCheckpoint(
{ history: conversation, messages: [] },
'tag',
),
).resolves.not.toThrow();
expect(consoleErrorSpy).toHaveBeenCalledWith(
'Logger not initialized or checkpoint file path not set. Cannot save a checkpoint.',
@@ -507,6 +515,7 @@ describe('Logger', () => {
...conversation,
{ role: 'user', parts: [{ text: 'hello' }] },
],
messages: [],
authType: AuthType.USE_GEMINI,
};
const taggedFilePath = path.join(
@@ -534,21 +543,21 @@ describe('Logger', () => {
await fs.writeFile(taggedFilePath, JSON.stringify(conversation, null, 2));
const loaded = await logger.loadCheckpoint(tag);
expect(loaded).toEqual({ history: conversation });
expect(loaded).toEqual({ history: conversation, messages: [] });
});
it('should return an empty history if a tagged checkpoint file does not exist', async () => {
it('should return an empty message list if a tagged checkpoint file does not exist', async () => {
const loaded = await logger.loadCheckpoint('nonexistent-tag');
expect(loaded).toEqual({ history: [] });
expect(loaded).toEqual({ history: [], messages: [] });
});
it('should return an empty history if the checkpoint file does not exist', async () => {
it('should return an empty message list if the checkpoint file does not exist', async () => {
await fs.unlink(TEST_CHECKPOINT_FILE_PATH); // Ensure it's gone
const loaded = await logger.loadCheckpoint('missing');
expect(loaded).toEqual({ history: [] });
expect(loaded).toEqual({ history: [], messages: [] });
});
it('should return an empty history if the file contains invalid JSON', async () => {
it('should return an empty message list if the file contains invalid JSON', async () => {
const tag = 'invalid-json-tag';
const encodedTag = 'invalid-json-tag';
const taggedFilePath = path.join(
@@ -560,14 +569,14 @@ describe('Logger', () => {
.spyOn(debugLogger, 'error')
.mockImplementation(() => {});
const loadedCheckpoint = await logger.loadCheckpoint(tag);
expect(loadedCheckpoint).toEqual({ history: [] });
expect(loadedCheckpoint).toEqual({ history: [], messages: [] });
expect(consoleErrorSpy).toHaveBeenCalledWith(
expect.stringContaining('Failed to read or parse checkpoint file'),
expect.any(Error),
);
});
it('should return an empty history if logger is not initialized', async () => {
it('should return an empty message list if logger is not initialized', async () => {
const uninitializedLogger = new Logger(
testSessionId,
new Storage(process.cwd()),
@@ -577,7 +586,7 @@ describe('Logger', () => {
.spyOn(debugLogger, 'error')
.mockImplementation(() => {});
const loadedCheckpoint = await uninitializedLogger.loadCheckpoint('tag');
expect(loadedCheckpoint).toEqual({ history: [] });
expect(loadedCheckpoint).toEqual({ history: [], messages: [] });
expect(consoleErrorSpy).toHaveBeenCalledWith(
'Logger not initialized or checkpoint file path not set. Cannot load checkpoint.',
);

View File

@@ -12,6 +12,11 @@ import type { Storage } from '../config/storage.js';
import { debugLogger } from '../utils/debugLogger.js';
import { coreEvents } from '../utils/events.js';
import {
type ConversationRecord,
type MessageRecord,
} from '../services/chatRecordingService.js';
const LOG_FILE_NAME = 'logs.json';
export enum MessageSenderType {
@@ -29,6 +34,17 @@ export interface LogEntry {
export interface Checkpoint {
history: readonly Content[];
authType?: AuthType;
/**
* The rich message records which are the source of truth for the session.
* Optional in Stage 1 to maintain backward compatibility.
*/
messages?: MessageRecord[];
/**
* The version of the checkpoint format.
*/
version?: '2.0';
/** Optional subagent trajectories to include. */
trajectories?: Record<string, ConversationRecord>;
}
// This regex matches any character that is NOT a letter (a-z, A-Z),
@@ -347,7 +363,7 @@ export class Logger {
debugLogger.error(
'Logger not initialized or checkpoint file path not set. Cannot load checkpoint.',
);
return { history: [] };
return { history: [], messages: [] };
}
const path = await this._getCheckpointPath(tag);
@@ -359,7 +375,7 @@ export class Logger {
// Handle legacy format (just an array of Content)
if (Array.isArray(parsedContent)) {
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
return { history: parsedContent as Content[] };
return { history: parsedContent as Content[], messages: [] };
}
if (
@@ -368,25 +384,29 @@ export class Logger {
'history' in parsedContent
) {
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
return parsedContent as Checkpoint;
const checkpoint = parsedContent as Checkpoint;
return {
...checkpoint,
messages: checkpoint.messages ?? [],
};
}
debugLogger.warn(
`Checkpoint file at ${path} has an unknown format. Returning empty checkpoint.`,
);
return { history: [] };
return { history: [], messages: [] };
} catch (error) {
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
const nodeError = error as NodeJS.ErrnoException;
if (nodeError.code === 'ENOENT') {
// This is okay, it just means the checkpoint doesn't exist in either format.
return { history: [] };
return { history: [], messages: [] };
}
debugLogger.error(
`Failed to read or parse checkpoint file ${path}:`,
error,
);
return { history: [] };
return { history: [], messages: [] };
}
}

View File

@@ -128,6 +128,7 @@ export * from './utils/channel.js';
export * from './utils/constants.js';
export * from './utils/sessionUtils.js';
export * from './utils/cache.js';
export * from './utils/history-reconstruction.js';
export * from './utils/markdownUtils.js';
// Export services

View File

@@ -673,6 +673,63 @@ export class ChatRecordingService {
return this.conversationFile;
}
/**
* Recursively finds all subagent trajectories called from this session.
*/
async getSubagentTrajectories(): Promise<Record<string, ConversationRecord>> {
const trajectories: Record<string, ConversationRecord> = {};
const conversation = this.getConversation();
if (!conversation) return trajectories;
const agentIds = new Set<string>();
for (const message of conversation.messages) {
if (message.type === 'gemini' && message.toolCalls) {
for (const toolCall of message.toolCalls) {
if (toolCall.agentId) {
agentIds.add(toolCall.agentId);
}
}
}
}
if (agentIds.size === 0) return trajectories;
const tempDir = this.context.config.storage.getProjectTempDir();
const chatsDir = path.join(tempDir, 'chats');
for (const agentId of agentIds) {
const subagentFilePath = path.join(chatsDir, `${agentId}.json`);
try {
const content = await fs.promises.readFile(subagentFilePath, 'utf-8');
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
const subagentRecord = JSON.parse(content) as ConversationRecord;
trajectories[agentId] = subagentRecord;
// Recursive discovery: Create a temporary service to find nested trajectories
const subService = new ChatRecordingService(this.context);
// Manual override of cached state for discovery
subService.cachedConversation = subagentRecord;
subService.conversationFile = subagentFilePath;
const nested = await subService.getSubagentTrajectories();
Object.assign(trajectories, nested);
} catch (err) {
debugLogger.warn(`Failed to load subagent trajectory ${agentId}:`, err);
}
}
return trajectories;
}
/**
* Resets the current message history. Used during session resumption.
*/
resetMessages(messages: MessageRecord[]): void {
if (!this.cachedConversation) return;
this.cachedConversation.messages = [...messages];
// We don't append to the log here, as we are resetting the in-memory state
// to match a loaded checkpoint.
}
/**
* Deletes a session file by sessionId, filename, or basename.
* Derives an 8-character shortId to find and delete all associated files

View File

@@ -0,0 +1,80 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { Content, Part } from '@google/genai';
import type { MessageRecord } from '../services/chatRecordingTypes.js';
/**
* Reconstructs the model-compatible history array from rich message records.
* This allows us to treat MessageRecord as the source of truth and generate
* the API-specific Content array on-the-fly.
*/
export function reconstructHistory(messages: MessageRecord[]): Content[] {
const history: Content[] = [];
for (const msg of messages) {
const parts: Part[] = [];
if (Array.isArray(msg.content)) {
// Map PartUnion to Part
for (const p of msg.content) {
if (typeof p === 'string') {
parts.push({ text: p });
} else {
parts.push(p);
}
}
} else if (typeof msg.content === 'string') {
parts.push({ text: msg.content });
}
if (msg.type === 'user') {
history.push({ role: 'user', parts });
} else if (msg.type === 'gemini') {
// 1. Add model-generated tool calls if present
if (msg.toolCalls && msg.toolCalls.length > 0) {
msg.toolCalls.forEach((tc) => {
parts.push({
functionCall: {
name: tc.name,
args: tc.args,
id: tc.id,
},
});
});
}
history.push({ role: 'model', parts });
// 2. Add the tool responses as a following user turn if results exist
const toolResponseParts: Part[] = [];
if (msg.toolCalls) {
for (const tc of msg.toolCalls) {
if (tc.result) {
if (Array.isArray(tc.result)) {
for (const r of tc.result) {
if (typeof r === 'string') {
toolResponseParts.push({ text: r });
} else {
toolResponseParts.push(r);
}
}
} else if (typeof tc.result === 'string') {
toolResponseParts.push({ text: tc.result });
} else {
toolResponseParts.push(tc.result);
}
}
}
}
if (toolResponseParts.length > 0) {
history.push({ role: 'user', parts: toolResponseParts });
}
}
}
return history;
}