fix(core): rollback entire multi-turn request on cancellation or abort

packages/core/src/core/geminiChat.test.ts

@@ -1628,6 +1628,89 @@
       expect(chat.agentHistory.length).toBe(initialHistoryLength);
     });
 
+    it('should roll back the entire multi-turn request including function responses when a continuation stream is aborted/cancelled', async () => {
+      const initialHistoryLength = chat.agentHistory.length;
+      const abortController = new AbortController();
+
+      // 1. Send the first message of the prompt. This will succeed and register prompt-id-multi-turn-abort.
+      const streamFirst = (async function* () {
+        yield {
+          candidates: [
+            {
+              content: {
+                role: 'model',
+                parts: [{ text: 'model first response' }],
+              },
+              finishReason: 'STOP',
+            },
+          ],
+        } as unknown as GenerateContentResponse;
+      })();
+      vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(
+        streamFirst,
+      );
+
+      const s1 = await chat.sendMessageStream(
+        { model: 'gemini-2.0-flash' },
+        'user original prompt',
+        'prompt-id-multi-turn-abort',
+        new AbortController().signal,
+        LlmRole.MAIN,
+      );
+      for await (const _ of s1) {
+        // consume the stream
+      }
+
+      // Expect history to contain: user, model
+      expect(chat.agentHistory.length).toBe(initialHistoryLength + 2);
+
+      // 2. Send a continuation (functionResponse), which is cancelled mid-stream.
+      const streamSecond = (async function* () {
+        yield {
+          candidates: [
+            {
+              content: {
+                role: 'model',
+                parts: [{ text: 'partial model response before abort' }],
+              },
+            },
+          ],
+        } as unknown as GenerateContentResponse;
+        abortController.abort();
+        throw new Error('User aborted a continuation stream.');
+      })();
+      vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(
+        streamSecond,
+      );
+
+      const s2 = await chat.sendMessageStream(
+        { model: 'gemini-2.0-flash' },
+        [
+          {
+            functionResponse: {
+              id: 'call_id_1',
+              name: 'my_tool',
+              response: { result: 'success' },
+            },
+          },
+        ],
+        'prompt-id-multi-turn-abort',
+        abortController.signal,
+        LlmRole.MAIN,
+      );
+
+      await expect(
+        (async () => {
+          for await (const _ of s2) {
+            // consume the stream to trigger abort
+          }
+        })(),
+      ).rejects.toThrow();
+
+      // Verify history has been rolled back entirely to initialHistoryLength (before the original prompt started)!
+      expect(chat.agentHistory.length).toBe(initialHistoryLength);
+    });
+
     it('should roll back the un-responded user turn from history when an ApiError is thrown', async () => {
       const initialHistoryLength = chat.agentHistory.length;
 
@@ -3330,7 +3413,7 @@
       });
     });
 
     it('should completely filter out thought parts from getHistoryTurns when context management is disabled but model is gemini-2/modern', () => {
       vi.mocked(mockConfig.isContextManagementEnabled).mockReturnValue(false);
       vi.mocked(mockConfig.getModel).mockReturnValue('gemini-2.5-pro');
 

packages/core/src/core/geminiChat.ts

@@ -38,6 +38,7 @@ import {
 import { hasCycleInSchema } from '../tools/tools.js';
 import type { StructuredError } from './turn.js';
 import type { CompletedToolCall } from '../scheduler/types.js';
+import { isAbortError } from '../utils/errors.js';
 import {
   logContentRetry,
   logContentRetryFailure,
@@ -296,6 +297,9 @@ export class GeminiChat {
   private lastPromptTokenCount: number;
   private callCounter = 0;
   agentHistory: AgentChatHistory;
+  private lastPromptId?: string;
+  private promptOriginalHistoryLength?: number;
+  private promptOriginalTokenCount?: number;
 
   constructor(
     readonly context: AgentLoopContext,
@@ -407,6 +411,17 @@ export class GeminiChat {
     const historyLengthBefore = this.agentHistory.length;
     const baselinePromptTokenCount = this.lastPromptTokenCount;
 
+    if (this.lastPromptId && this.lastPromptId !== prompt_id) {
+      this.promptOriginalHistoryLength = undefined;
+      this.promptOriginalTokenCount = undefined;
+    }
+    this.lastPromptId = prompt_id;
+
+    if (this.promptOriginalHistoryLength === undefined) {
+      this.promptOriginalHistoryLength = historyLengthBefore;
+      this.promptOriginalTokenCount = baselinePromptTokenCount;
+    }
+
     let streamDoneResolver: () => void;
     const streamDonePromise = new Promise<void>((resolve) => {
       streamDoneResolver = resolve;
@@ -684,7 +699,26 @@ export class GeminiChat {
           }
         }
       } catch (error) {
-        if (!isOriginalFunctionResponse) {
+        const isAborted =
+          signal?.aborted ||
+          isAbortError(error) ||
+          (error instanceof Error &&
+            (error.name === 'CanceledError' ||
+              error.name === 'FatalCancellationError'));
+        const originalLength = this.promptOriginalHistoryLength;
+        const originalTokenCount = this.promptOriginalTokenCount;
+        if (isAborted && originalLength !== undefined) {
+          this.agentHistory.rollback(originalLength);
+          this.chatRecordingService.updateMessagesFromHistory(
+            this.agentHistory.get(),
+          );
+          if (originalTokenCount !== undefined) {
+            this.lastPromptTokenCount = originalTokenCount;
+          }
+          this.promptOriginalHistoryLength = undefined;
+          this.promptOriginalTokenCount = undefined;
+          this.lastPromptId = undefined;
+        } else if (!isOriginalFunctionResponse) {
           this.agentHistory.rollback(historyLengthBefore);
           this.chatRecordingService.updateMessagesFromHistory(
             this.agentHistory.get(),