From 348fc35f176d02403e37d7891bc554a28a71c595 Mon Sep 17 00:00:00 2001 From: Adam Weidman <65992621+adamfweidman@users.noreply.github.com> Date: Wed, 5 Aug 2026 14:53:41 -0400 Subject: [PATCH] fix(core,cli): repair /compress session reload and quota-fallback tool response loss (#28672) Co-authored-by: David Pierce --- .../cli/src/ui/hooks/useGeminiStream.test.tsx | 101 +++++++++++++++ packages/cli/src/ui/hooks/useGeminiStream.ts | 32 +++-- .../src/services/chatRecordingService.test.ts | 119 ++++++++++++++++++ .../core/src/services/chatRecordingService.ts | 78 +++++++++++- 4 files changed, 318 insertions(+), 12 deletions(-) diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 9a35305cb6..abbe933abf 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -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) + | 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[] = [ { diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 9458c6a4ff..7965852dc1 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -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) { @@ -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, diff --git a/packages/core/src/services/chatRecordingService.test.ts b/packages/core/src/services/chatRecordingService.test.ts index 133e9ffe4d..8a63e3e541 100644 --- a/packages/core/src/services/chatRecordingService.test.ts +++ b/packages/core/src/services/chatRecordingService.test.ts @@ -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', () => { diff --git a/packages/core/src/services/chatRecordingService.ts b/packages/core/src/services/chatRecordingService.ts index 18b977bf00..186282eb1d 100644 --- a/packages/core/src/services/chatRecordingService.ts +++ b/packages/core/src/services/chatRecordingService.ts @@ -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 @@ -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): void { if (!this.cachedConversation) return; Object.assign(this.cachedConversation, updates);