Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
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
196 changes: 196 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,199 @@ 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');
});

it('should handle parts appearing BEFORE the functionResponse in a content block', () => {
chatRecordingService.initialize();
const callId = 'prefix-part-call';

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

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

const history: Content[] = [
{
role: 'user',
parts: [
{ text: 'Prefix metadata or text' },
{
functionResponse: {
name: 'read_file',
id: callId,
response: { output: 'file content' },
},
},
],
},
];

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].text).toBe('Prefix metadata or text');
expect(result[1].functionResponse!.id).toBe(callId);
});
});
});
64 changes: 64 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,66 @@ 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) {
// Find all unique call IDs in this message
const callIds = content.parts
.map((p) => p.functionResponse?.id)
.filter((id): id is string => !!id);

if (callIds.length === 0) continue;

// Use the first ID as a seed to capture any "leading" non-ID parts
// in this specific content block.
let currentCallId = callIds[0];
for (const part of content.parts) {
if (part.functionResponse?.id) {
currentCallId = part.functionResponse.id;
}

if (!partsMap.has(currentCallId)) {
partsMap.set(currentCallId, []);
}
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