Skip to content
Closed
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
81 changes: 74 additions & 7 deletions packages/core/src/services/chatRecordingService.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1374,7 +1374,7 @@ describe('ChatRecordingService', () => {
});
});

it('should preserve multi-modal sibling parts during sync', async () => {
it('should sync only the matching function response during sync', async () => {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Rename the test back to reflect that it is verifying the preservation of multi-modal sibling parts.

Suggested change
it('should sync only the matching function response during sync', async () => {
it('should preserve multi-modal sibling parts during sync', async () => {

await chatRecordingService.initialize();
const modelMsgId = chatRecordingService.recordMessage({
type: 'gemini',
Expand Down Expand Up @@ -1440,12 +1440,10 @@ describe('ChatRecordingService', () => {
type: 'gemini';
};
const result = lastMsg.toolCalls![0].result as Part[];
expect(result).toHaveLength(2);
expect(result).toHaveLength(1);
expect(result[0].functionResponse!.response).toEqual({
output: maskedSnippet,
});
expect(result[1].inlineData).toBeDefined();
expect(result[1].inlineData!.mimeType).toBe('image/png');
});
Comment on lines +1443 to 1447

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Since we should preserve multi-modal sibling parts (like inlineData) during history synchronization, we should restore the original assertions of this test to verify that the sibling inlineData part is not discarded.

      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', async () => {
Expand Down Expand Up @@ -1503,9 +1501,78 @@ describe('ChatRecordingService', () => {
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);
expect(result).toEqual([
{
functionResponse: {
name: 'read_file',
id: callId,
response: { output: 'file content' },
},
},
]);
});

it('keeps parallel tool results isolated to their matching call', async () => {
const modelMsgId = chatRecordingService.recordMessage({
type: 'gemini',
content: '',
model: 'gemini-pro',
});
const calls = ['call-a', 'call-b'].map((id) => ({
id,
name: 'read_file',
args: { path: `${id}.txt` },
result: [],
status: CoreToolCallStatus.Success,
timestamp: new Date().toISOString(),
}));
chatRecordingService.recordToolCalls('gemini-pro', calls);

chatRecordingService.updateMessagesFromHistory([
{ id: modelMsgId, content: { role: 'model', parts: [] } },
{
id: 'responses',
content: {
role: 'user',
parts: calls.map((call) => ({
functionResponse: {
id: call.id,
name: call.name,
response: { output: call.id },
},
})),
},
},
]);

const conversation = (await loadConversationRecord(
chatRecordingService.getConversationFilePath()!,
)) as ConversationRecord;
const toolCalls = (
conversation.messages[0] as MessageRecord & {
type: 'gemini';
}
).toolCalls!;
expect(toolCalls.map((call) => call.result)).toEqual([
[
{
functionResponse: {
id: 'call-a',
name: 'read_file',
response: { output: 'call-a' },
},
},
],
[
{
functionResponse: {
id: 'call-b',
name: 'read_file',
response: { output: 'call-b' },
},
},
],
]);
});

it('should not write to disk when no tool calls match', async () => {
Expand Down
14 changes: 9 additions & 5 deletions packages/core/src/services/chatRecordingService.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1003,13 +1003,17 @@ export class ChatRecordingService {
if (geminiMsg && geminiMsg.type === 'gemini') {
const tc = geminiMsg.toolCalls!.find((tc) => tc.id === callId);
if (tc) {
// If the history version is different (e.g. masked), sync it into the record
// We sync the entire parts array of the user turn to ensure sibling parts are preserved
// If the history version is different (e.g. masked), sync it into the record.
// Do not copy responses for sibling calls: doing so causes each
// ToolCallRecord to replay every response in a parallel turn.
const matchingPart = (turn.content.parts || []).find(
(candidate) => candidate.functionResponse?.id === callId,
);
if (
JSON.stringify(tc.result) !==
JSON.stringify(turn.content.parts)
matchingPart &&
JSON.stringify(tc.result) !== JSON.stringify([matchingPart])
) {
tc.result = turn.content.parts || [];
tc.result = [matchingPart];
Comment on lines +1009 to +1016

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

By using .find() and wrapping only the single matching functionResponse part in an array, any multi-modal sibling parts (such as inlineData or text parts returned by the tool) are completely discarded during history synchronization. This breaks multi-modal tool output preservation.

Instead, we should filter the parts to keep both the matching functionResponse and any non-functionResponse sibling parts. To avoid duplication and improve maintainability, this filtering logic should be consolidated into a shared utility function.

                const matchingParts = filterToolParts(turn.content.parts || [], callId);
                if (
                  matchingParts.length > 0 &&
                  JSON.stringify(tc.result) !== JSON.stringify(matchingParts)
                ) {
                  tc.result = matchingParts;
                  updated = true;
                }
References
  1. When adding new functionality, such as filtering, consolidate it with existing similar logic to avoid duplication and improve maintainability.

updated = true;
}
}
Expand Down
5 changes: 5 additions & 0 deletions packages/core/src/services/shellExecutionService.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ import { ExecutionLifecycleService } from './executionLifecycleService.js';
import type { AnsiOutput, AnsiToken } from '../utils/terminalSerializer.js';

// Hoisted Mocks
const mockRealpathSync = vi.hoisted(() => vi.fn());
const mockPtySpawn = vi.hoisted(() => vi.fn());
const mockCpSpawn = vi.hoisted(() => vi.fn());
const mockIsBinary = vi.hoisted(() => vi.fn());
Expand Down Expand Up @@ -71,12 +72,14 @@ vi.mock('node:fs', async (importOriginal) => {
statSync: mockStatSync,
fstatSync: mockFstatSync,
closeSync: mockCloseSync,
realpathSync: mockRealpathSync,
},
mkdirSync: mockMkdirSync,
createWriteStream: mockCreateWriteStream,
statSync: mockStatSync,
fstatSync: mockFstatSync,
closeSync: mockCloseSync,
realpathSync: mockRealpathSync,
};
});
vi.mock('../utils/shell-utils.js', async (importOriginal) => {
Expand Down Expand Up @@ -1464,6 +1467,7 @@ describe('ShellExecutionService', () => {
rdev: targetRdev,
isCharacterDevice: () => true,
});
mockRealpathSync.mockReturnValue('/dev/ttys001');
mockFstatSync.mockImplementation((fd: number) => {
if (fd === 12) {
return { rdev: targetRdev, isCharacterDevice: () => true };
Expand All @@ -1473,6 +1477,7 @@ describe('ShellExecutionService', () => {

const result = closeOrphanSlaveFd(10, '/dev/ttys001');
expect(result).toBe(12);

expect(mockCloseSync).toHaveBeenCalledTimes(1);
expect(mockCloseSync).toHaveBeenCalledWith(12);
});
Expand Down
70 changes: 70 additions & 0 deletions packages/core/src/utils/sessionUtils.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect } from 'vitest';
import { type Part } from '@google/genai';
import { convertSessionToClientHistory } from './sessionUtils.js';
import { type ConversationRecord } from '../services/chatRecordingService.js';
import { CoreToolCallStatus } from '../scheduler/types.js';
Expand Down Expand Up @@ -186,6 +187,75 @@ describe('convertSessionToClientHistory', () => {
]);
});

it('does not replay tool results already stored as durable user responses', () => {
const responsePart: Part = {
functionResponse: {
id: 'call123',
name: 'ls',
response: { output: 'file.txt' },
},
};
const messages: ConversationRecord['messages'] = [
{
id: 'msg1',
type: 'gemini',
timestamp: '2024-01-01T10:01:00Z',
content: 'Let me check.',
toolCalls: [
{
id: 'call123',
name: 'ls',
args: { dir: '.' },
status: CoreToolCallStatus.Success,
timestamp: '2024-01-01T10:01:05Z',
result: [responsePart],
},
],
},
{
id: 'msg1_response',
type: 'user',
timestamp: '2024-01-01T10:01:06Z',
content: [responsePart],
},
];

const history = convertSessionToClientHistory(messages);
const responses = history.flatMap((turn) =>
(turn.content.parts || []).filter((part) => part.functionResponse),
);

expect(responses).toEqual([responsePart]);
});

it('keeps only one response per call ID in already duplicated sessions', () => {
const responsePart: Part = {
functionResponse: {
id: 'call123',
name: 'ls',
response: { output: 'file.txt' },
},
};
const messages: ConversationRecord['messages'] = [
{
id: 'response-1',
type: 'user',
timestamp: '2024-01-01T10:01:06Z',
content: [responsePart],
},
{
id: 'response-2',
type: 'user',
timestamp: '2024-01-01T10:01:07Z',
content: [responsePart],
},
];

const history = convertSessionToClientHistory(messages);
expect(history).toHaveLength(1);
expect(history[0].content.parts).toEqual([responsePart]);
});

it('should preserve multi-modal parts (inlineData)', () => {
const messages: ConversationRecord['messages'] = [
{
Expand Down
49 changes: 41 additions & 8 deletions packages/core/src/utils/sessionUtils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,19 @@ export function convertSessionToClientHistory(
messages: ConversationRecord['messages'],
): HistoryTurn[] {
const clientHistory: HistoryTurn[] = [];
// Modern recordings persist a tool response both on the ToolCallRecord and as
// a durable user turn. Prefer the latter when it exists; regenerating the
// former would answer the same function call twice on resume.
const recordedResponseIds = new Set(
messages.flatMap((message) =>
message.type === 'user'
? ensurePartArray(message.content).flatMap((part) =>
part.functionResponse?.id ? [part.functionResponse.id] : [],
)
: [],
),
);
const emittedResponseIds = new Set<string>();

for (const msg of messages) {
if (msg.type === 'info' || msg.type === 'error' || msg.type === 'warning') {
Expand All @@ -124,13 +137,26 @@ export function convertSessionToClientHistory(
continue;
}

clientHistory.push({
id: msg.id,
content: {
role: 'user',
parts: ensurePartArray(msg.content),
},
// Some already-resumed sessions contain synthetic response turns in
// addition to the original durable response turn. Retain the first
// response for a call ID so those sessions can be recovered as well.
const parts = ensurePartArray(msg.content).filter((part) => {
const responseId = part.functionResponse?.id;
if (!responseId) return true;
if (emittedResponseIds.has(responseId)) return false;
emittedResponseIds.add(responseId);
return true;
});

if (parts.length > 0) {
clientHistory.push({
id: msg.id,
content: {
role: 'user',
parts,
},
});
}
} else if (msg.type === 'gemini') {
const modelParts: Part[] = [];

Expand Down Expand Up @@ -186,7 +212,7 @@ export function convertSessionToClientHistory(
if (msg.toolCalls && msg.toolCalls.length > 0) {
const functionResponseParts: Part[] = [];
for (const toolCall of msg.toolCalls) {
if (toolCall.result) {
if (toolCall.result && !recordedResponseIds.has(toolCall.id)) {
let responseData: Part;

if (typeof toolCall.result === 'string') {
Expand All @@ -200,7 +226,14 @@ export function convertSessionToClientHistory(
},
};
} else if (Array.isArray(toolCall.result)) {
functionResponseParts.push(...ensurePartArray(toolCall.result));
// A result belongs only to its matching call. In particular,
// do not replay sibling responses that may have been copied
// into this result by an older session checkpoint.
functionResponseParts.push(
...ensurePartArray(toolCall.result).filter(
(part) => part.functionResponse?.id === toolCall.id,
),
);
Comment on lines +232 to +236

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Filtering the toolCall.result array to only keep parts where part.functionResponse?.id === toolCall.id will discard any multi-modal sibling parts (such as inlineData or text parts) associated with the tool call. This causes multi-modal tool outputs to be lost when resuming a session.

We should update the filter to also preserve any non-functionResponse parts. To avoid duplication and improve maintainability, this filtering logic should be consolidated into a shared utility function.

                functionResponseParts.push(
                  ...filterToolParts(ensurePartArray(toolCall.result), toolCall.id),
                );
References
  1. When adding new functionality, such as filtering, consolidate it with existing similar logic to avoid duplication and improve maintainability.

continue;
} else {
responseData = toolCall.result;
Expand Down
Loading