fix(core,cli): repair /compress session reload and quota-fallback tool response loss (#28672)

Co-authored-by: David Pierce <davidapierce@google.com>

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<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[] = [
       {

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,

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', () => {

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<ConversationRecord>): void {
     if (!this.cachedConversation) return;
     Object.assign(this.cachedConversation, updates);