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
101 changes: 101 additions & 0 deletions packages/cli/src/ui/hooks/useGeminiStream.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -1046,6 +1046,107 @@ describe('useGeminiStream', () => {
});
});

it('should record tool responses in history when the model was switched due to a quota error', async () => {
// Regression test: returning early on a quota-triggered model switch
// without recording the responses leaves the already-recorded
// functionCall unpaired, which corrupts all subsequent requests.
const responseParts: Part[] = [
{
functionResponse: {
name: 'testTool',
id: 'call1',
response: { output: 'tool result' },
},
},
];
const completedToolCalls: TrackedToolCall[] = [
{
request: {
callId: 'call1',
name: 'testTool',
args: {},
isClientInitiated: false,
prompt_id: 'prompt-id-quota',
},
status: CoreToolCallStatus.Success,
responseSubmittedToGemini: false,
response: {
callId: 'call1',
responseParts,
errorType: undefined,
},
tool: { displayName: 'MockTool' },
invocation: {
getDescription: () => `Mock description`,
} as unknown as AnyToolInvocation,
} as TrackedCompletedToolCall,
];

const client = new MockedGeminiClientClass(mockConfig);
const mockConsumeUserHint = vi.fn(() => 'switch to the nprd database');

let capturedOnComplete:
| ((completedTools: TrackedToolCall[]) => Promise<void>)
| null = null;

mockUseToolScheduler.mockImplementation((onComplete) => {
capturedOnComplete = onComplete;
return [
[],
mockScheduleToolCalls,
mockMarkToolsAsSubmitted,
vi.fn(),
mockCancelAllToolCalls,
0,
];
});

await renderHookWithProviders(() =>
useGeminiStream(
client,
[],
mockAddItem,
mockConfig,
mockLoadedSettings,
mockOnDebugMessage,
mockHandleSlashCommand,
false,
() => 'vscode' as EditorType,
() => {},
() => Promise.resolve(),
true, // modelSwitchedFromQuotaError
() => {},
() => {},
() => {},
80,
24,
false,
mockConsumeUserHint,
),
);

await act(async () => {
if (capturedOnComplete) {
await new Promise((resolve) => setTimeout(resolve, 0));
await capturedOnComplete(completedToolCalls);
}
});

await waitFor(() => {
expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['call1']);
// The tool response must be paired with its functionCall in history,
// with no steering-hint text ahead of it...
expect(client.addHistory).toHaveBeenCalledWith({
role: 'user',
parts: responseParts,
});
// ...the turn must NOT auto-continue on the fallback model...
expect(mockSendMessageStream).not.toHaveBeenCalled();
// ...and the pending hint is left for the next real submit.
expect(mockConsumeUserHint).not.toHaveBeenCalled();
});
});

it('should NOT stop responding when only update_topic is called', async () => {
const topicToolCalls: TrackedToolCall[] = [
{
Expand Down
32 changes: 21 additions & 11 deletions packages/cli/src/ui/hooks/useGeminiStream.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2110,6 +2110,27 @@ export const useGeminiStream = (
(toolCall) => toolCall.response.responseParts,
);

const callIdsToMarkAsSubmitted = geminiTools.map(
(toolCall) => toolCall.request.callId,
);

markToolsAsSubmitted(callIdsToMarkAsSubmitted);

// Don't continue if model was switched due to quota error, but still
// record the responses: the matching functionCall is already in history,
// and leaving it unpaired corrupts every subsequent request. Any pending
// steering hint is deliberately left unconsumed so it rides along with
// the next query the user actually submits.
if (modelSwitchedFromQuotaError) {
if (geminiClient && responsesToSend.length > 0) {
await geminiClient.addHistory({
role: 'user',
parts: responsesToSend,
});
}
return;
}

if (consumeUserHint) {
const userHint = consumeUserHint();
if (userHint && userHint.trim().length > 0) {
Expand All @@ -2120,21 +2141,10 @@ export const useGeminiStream = (
}
}

const callIdsToMarkAsSubmitted = geminiTools.map(
(toolCall) => toolCall.request.callId,
);

const prompt_ids = geminiTools.map(
(toolCall) => toolCall.request.prompt_id,
);

markToolsAsSubmitted(callIdsToMarkAsSubmitted);

// Don't continue if model was switched due to quota error
if (modelSwitchedFromQuotaError) {
return;
}

// eslint-disable-next-line @typescript-eslint/no-floating-promises
submitQuery(
responsesToSend,
Expand Down
119 changes: 119 additions & 0 deletions packages/core/src/services/chatRecordingService.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -308,6 +308,125 @@ describe('ChatRecordingService', () => {
)) as ConversationRecord;
expect(conversation.sessionId).toBe('old-session-id');
});

it('should fall back to the in-memory conversation when the file cannot be reloaded', async () => {
// Regression test for the `/compress` "Failed to load resumed session
// data from file" bug: when resuming with a filePath that cannot be
// loaded from disk, initialize must NOT throw. It should adopt the
// in-memory conversation it was handed and rewrite a clean file.
const chatsDir = path.join(testTempDir, 'chats');
fs.mkdirSync(chatsDir, { recursive: true });
const missingFile = path.join(chatsDir, 'missing-session.jsonl');
expect(fs.existsSync(missingFile)).toBe(false);

const inMemoryConversation = {
sessionId: 'resumed-session-id',
projectHash: 'resumed-project-hash',
startTime: new Date().toISOString(),
lastUpdated: new Date().toISOString(),
messages: [
{
id: 'msg-1',
type: 'user',
timestamp: new Date().toISOString(),
content: 'hello from memory',
},
],
} as unknown as ConversationRecord;

await expect(
chatRecordingService.initialize({
filePath: missingFile,
conversation: inMemoryConversation,
}),
).resolves.not.toThrow();

// The in-memory conversation is adopted.
expect(chatRecordingService.getConversation()?.sessionId).toBe(
'resumed-session-id',
);

// A clean, loadable file is rewritten from the in-memory copy so future
// loads and appends succeed.
const reloaded = (await loadConversationRecord(
missingFile,
)) as ConversationRecord;
expect(reloaded).not.toBeNull();
expect(reloaded.sessionId).toBe('resumed-session-id');
expect(reloaded.projectHash).toBe('resumed-project-hash');
expect(reloaded.messages).toHaveLength(1);
});

it('should preserve an unreadable session file instead of destroying it', async () => {
// The reload may have failed only transiently, so the original bytes
// must survive the recovery rewrite.
const chatsDir = path.join(testTempDir, 'chats');
fs.mkdirSync(chatsDir, { recursive: true });
const sessionFile = path.join(chatsDir, 'unreadable.jsonl');

// No usable metadata line => loadConversationRecord() returns null.
const originalBytes = '{"not":"a valid metadata line"}\n';
fs.writeFileSync(sessionFile, originalBytes);

await chatRecordingService.initialize({
filePath: sessionFile,
conversation: {
sessionId: 'recovered-session-id',
projectHash: 'recovered-project-hash',
startTime: new Date().toISOString(),
lastUpdated: new Date().toISOString(),
messages: [],
} as unknown as ConversationRecord,
});

// The rewritten file is loadable again...
const reloaded = (await loadConversationRecord(
sessionFile,
)) as ConversationRecord;
expect(reloaded.sessionId).toBe('recovered-session-id');

// ...and the original bytes were kept alongside it.
const preserved = fs
.readdirSync(chatsDir)
.filter((f) => f.startsWith('unreadable.jsonl.unreadable-'));
expect(preserved).toHaveLength(1);
expect(fs.readFileSync(path.join(chatsDir, preserved[0]), 'utf-8')).toBe(
originalBytes,
);
});

it('should not leave a temp file behind when the rewrite fails', async () => {
const chatsDir = path.join(testTempDir, 'chats');
fs.mkdirSync(chatsDir, { recursive: true });
const sessionFile = path.join(chatsDir, 'rewrite-fails.jsonl');

// Fail the rename that publishes the temp file, leaving it orphaned.
const realRename = fs.renameSync;
vi.spyOn(fs, 'renameSync').mockImplementation((from, to) => {
if (String(from).includes('.tmp-')) {
throw new Error('simulated rename failure');
}
return realRename(from, to);
});

await expect(
chatRecordingService.initialize({
filePath: sessionFile,
conversation: {
sessionId: 'temp-cleanup-session',
projectHash: 'temp-cleanup-hash',
startTime: new Date().toISOString(),
lastUpdated: new Date().toISOString(),
messages: [],
} as unknown as ConversationRecord,
}),
).rejects.toThrow('simulated rename failure');

const leftovers = fs
.readdirSync(chatsDir)
.filter((f) => f.includes('.tmp-'));
expect(leftovers).toEqual([]);
});
});

describe('recordMessage', () => {
Expand Down
78 changes: 77 additions & 1 deletion packages/core/src/services/chatRecordingService.ts
Original file line number Diff line number Diff line change
Expand Up @@ -462,7 +462,16 @@ export class ChatRecordingService {
// Update the session ID in the existing file
this.updateMetadata({ sessionId: this.sessionId });
} else {
throw new Error('Failed to load resumed session data from file');
// The file could not be reloaded (missing, corrupt metadata, or an
// I/O error). Fall back to the in-memory conversation we were handed
// rather than failing the caller, and rewrite a clean file from it.
debugLogger.warn(
'Failed to reload resumed session data from file; falling back ' +
'to the in-memory conversation.',
);
this.cachedConversation = resumedSessionData.conversation;
this.projectHash = this.cachedConversation.projectHash;
this.rewriteConversationFile(this.cachedConversation);
}
} else {
// Create new session
Expand Down Expand Up @@ -563,6 +572,73 @@ export class ChatRecordingService {
}
}

/**
* Rewrites the session file from an in-memory record. Any existing
* (unreadable) file is preserved alongside rather than destroyed, and the
* new file is written atomically (temp file + rename).
*/
private rewriteConversationFile(conversation: ConversationRecord): void {
if (!this.conversationFile) return;

// Normalize legacy `.json` paths to the `.jsonl` format we write.
if (this.conversationFile.endsWith('.json')) {
this.conversationFile = this.conversationFile + 'l';
}

const { messages, memoryScratchpad, ...metadata } = conversation;
const lines: string[] = [JSON.stringify(metadata)];
for (const msg of messages) {
lines.push(JSON.stringify(msg));
}
if (memoryScratchpad) {
lines.push(JSON.stringify({ $set: { memoryScratchpad } }));
}
const content = lines.join('\n') + '\n';

try {
fs.mkdirSync(path.dirname(this.conversationFile), { recursive: true });

// The existing file was unreadable, but it may have been only
// transiently so (a lock or I/O blip) rather than truly corrupt. Keep
// its bytes rather than destroying them.
if (fs.existsSync(this.conversationFile)) {
const backup = `${this.conversationFile}.unreadable-${Date.now()}`;
try {
fs.renameSync(this.conversationFile, backup);
debugLogger.warn(
`Preserved the unreadable session file at ${backup}.`,
);
} catch (backupError) {
debugLogger.error(
'Failed to preserve the unreadable session file.',
backupError,
);
}
}

const tempFile = `${this.conversationFile}.tmp-${process.pid}`;
try {
fs.writeFileSync(tempFile, content);
fs.renameSync(tempFile, this.conversationFile);
} catch (error) {
// The rename did not complete, so the temp file would be left behind.
try {
fs.unlinkSync(tempFile);
} catch {
// Ignore cleanup errors so the original failure still surfaces.
}
throw error;
}
} catch (error) {
if (isNodeError(error) && error.code === 'ENOSPC') {
this.conversationFile = null;
debugLogger.warn(ENOSPC_WARNING_MESSAGE);
} else {
throw error;
}
}
}

private updateMetadata(updates: Partial<ConversationRecord>): void {
if (!this.cachedConversation) return;
Object.assign(this.cachedConversation, updates);
Expand Down
Loading