Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions packages/core/src/core/geminiChat.ts
Original file line number Diff line number Diff line change
Expand Up @@ -705,6 +705,7 @@ export class GeminiChat {
this.lastPromptTokenCount = estimateTokenCountSync(
this.history.flatMap((c) => c.parts || []),
);
this.chatRecordingService.updateMessagesFromHistory(history);
}

stripThoughtsFromHistory(): void {
Expand Down
143 changes: 143 additions & 0 deletions packages/core/src/services/chatRecordingService.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import type {
ToolCallRecord,
MessageRecord,
} from './chatRecordingService.js';
import type { Content, Part } from '@google/genai';
import { ChatRecordingService } from './chatRecordingService.js';
import type { Config } from '../config/config.js';
import { getProjectHash } from '../utils/paths.js';
Expand Down Expand Up @@ -548,4 +549,146 @@ describe('ChatRecordingService', () => {
writeFileSyncSpy.mockRestore();
});
});

describe('updateMessagesFromHistory', () => {
beforeEach(() => {
chatRecordingService.initialize();
});

it('should update tool results from API history (masking sync)', () => {
// 1. Record an initial message and tool call
chatRecordingService.recordMessage({
type: 'gemini',
content: 'I will list the files.',
model: 'gemini-pro',
});

const callId = 'tool-call-123';
const originalResult = [{ text: 'a'.repeat(1000) }];
chatRecordingService.recordToolCalls('gemini-pro', [
{
id: callId,
name: 'list_files',
args: { path: '.' },
result: originalResult,
status: 'success',
timestamp: new Date().toISOString(),
},
]);

// 2. Prepare mock history with masked content
const maskedSnippet =
'<tool_output_masked>short preview</tool_output_masked>';
const history: Content[] = [
{
role: 'model',
parts: [
{ functionCall: { name: 'list_files', args: { path: '.' } } },
],
},
{
role: 'user',
parts: [
{
functionResponse: {
name: 'list_files',
id: callId,
response: { output: maskedSnippet },
},
},
],
},
];

// 3. Trigger sync
chatRecordingService.updateMessagesFromHistory(history);

// 4. Verify disk content
const sessionFile = chatRecordingService.getConversationFilePath()!;
const conversation = JSON.parse(
fs.readFileSync(sessionFile, 'utf8'),
) as ConversationRecord;

const geminiMsg = conversation.messages[0];
if (geminiMsg.type !== 'gemini')
throw new Error('Expected gemini message');
expect(geminiMsg.toolCalls).toBeDefined();
expect(geminiMsg.toolCalls![0].id).toBe(callId);
// The implementation stringifies the response object
const result = geminiMsg.toolCalls![0].result;
if (!Array.isArray(result)) throw new Error('Expected array result');
const firstPart = result[0] as Part;
expect(firstPart.functionResponse).toBeDefined();
expect(firstPart.functionResponse!.id).toBe(callId);
expect(firstPart.functionResponse!.response).toEqual({
output: maskedSnippet,
});
});
it('should preserve multi-modal sibling parts during sync', () => {
chatRecordingService.initialize();
const callId = 'multi-modal-call';
const originalResult: Part[] = [
{
functionResponse: {
id: callId,
name: 'read_file',
response: { content: '...' },
},
},
{ inlineData: { mimeType: 'image/png', data: 'base64...' } },
];

chatRecordingService.recordMessage({
type: 'gemini',
content: '',
model: 'gemini-pro',
});

chatRecordingService.recordToolCalls('gemini-pro', [
{
id: callId,
name: 'read_file',
args: { path: 'image.png' },
result: originalResult,
status: 'success',
timestamp: new Date().toISOString(),
},
]);

const maskedSnippet = '<masked>';
const history: Content[] = [
{
role: 'user',
parts: [
{
functionResponse: {
name: 'read_file',
id: callId,
response: { output: maskedSnippet },
},
},
{ inlineData: { mimeType: 'image/png', data: 'base64...' } },
],
},
];

chatRecordingService.updateMessagesFromHistory(history);

const sessionFile = chatRecordingService.getConversationFilePath()!;
const conversation = JSON.parse(
fs.readFileSync(sessionFile, 'utf8'),
) as ConversationRecord;

const lastMsg = conversation.messages[0] as MessageRecord & {
type: 'gemini';
};
const result = lastMsg.toolCalls![0].result as Part[];
expect(result).toHaveLength(2);
expect(result[0].functionResponse!.response).toEqual({
output: maskedSnippet,
});
expect(result[1].inlineData).toBeDefined();
expect(result[1].inlineData!.mimeType).toBe('image/png');
});
});
});
57 changes: 57 additions & 0 deletions packages/core/src/services/chatRecordingService.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@ import path from 'node:path';
import fs from 'node:fs';
import { randomUUID } from 'node:crypto';
import type {
Content,
Part,
PartListUnion,
GenerateContentResponseUsageMetadata,
} from '@google/genai';
Expand Down Expand Up @@ -594,4 +596,59 @@ export class ChatRecordingService {
this.writeConversation(conversation, { allowEmpty: true });
return conversation;
}

/**
* Updates the conversation history based on the provided API Content array.
* This is used to persist changes made to the history (like masking) back to disk.
*/
updateMessagesFromHistory(history: Content[]): void {
if (!this.conversationFile) return;

try {
this.updateConversation((conversation) => {
// Create a map of tool results from the API history for quick lookup by call ID.
// We store the full list of parts associated with each tool call ID to preserve
// multi-modal data and proper trajectory structure.
const partsMap = new Map<string, Part[]>();
for (const content of history) {
if (content.role === 'user' && content.parts) {
let currentCallId: string | undefined;
for (const part of content.parts) {
if (part.functionResponse && part.functionResponse.id) {
currentCallId = part.functionResponse.id;
if (!partsMap.has(currentCallId)) {
partsMap.set(currentCallId, []);
}
partsMap.get(currentCallId)!.push(part);
} else if (currentCallId) {
// This handles sibling parts (e.g. inlineData) for the same tool call in Gemini 2.x
partsMap.get(currentCallId)!.push(part);
}
}
}
}

// Update the conversation records tool results if they've changed.
for (const message of conversation.messages) {
if (message.type === 'gemini' && message.toolCalls) {
for (const toolCall of message.toolCalls) {
const newParts = partsMap.get(toolCall.id);
if (newParts !== undefined) {
// Store the results as proper Parts (including functionResponse)
// instead of stringifying them as text parts. This ensures the
// tool trajectory is correctly reconstructed upon session resumption.
toolCall.result = newParts;
}
}
}
}
});
} catch (error) {
debugLogger.error(
'Error updating conversation history from memory.',
error,
);
throw error;
}
}
}
Loading