From 090198a7d644f24c617bd35db6a287b930729280 Mon Sep 17 00:00:00 2001
From: Taylor Mullen <ntaylormullen@google.com>
Date: Fri, 9 May 2025 22:47:18 -0700
Subject: [PATCH] Make cancel not explode.

- We were console.erroring, throwing and early aborting. Instead we now treat cancels like a normal user message and show an indicator in the UI

Fixes https://b.corp.google.com/issues/416515841
---
 .../cli/src/ui/components/messages/InfoMessage.tsx |  2 +-
 packages/cli/src/ui/hooks/useGeminiStream.ts       | 14 ++++++++++++--
 packages/server/src/core/turn.ts                   | 13 +++++--------
 3 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/packages/cli/src/ui/components/messages/InfoMessage.tsx b/packages/cli/src/ui/components/messages/InfoMessage.tsx
index b30e4473..c9129999 100644
--- a/packages/cli/src/ui/components/messages/InfoMessage.tsx
+++ b/packages/cli/src/ui/components/messages/InfoMessage.tsx
@@ -17,7 +17,7 @@ export const InfoMessage: React.FC<InfoMessageProps> = ({ text }) => {
   const prefixWidth = prefix.length;
 
   return (
-    <Box flexDirection="row">
+    <Box flexDirection="row" marginTop={1}>
       <Box width={prefixWidth}>
         <Text color={Colors.AccentYellow}>{prefix}</Text>
       </Box>
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts
index de073373..3f8cee40 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.ts
+++ b/packages/cli/src/ui/hooks/useGeminiStream.ts
@@ -177,8 +177,6 @@ export const useGeminiStream = (
         let geminiMessageBuffer = '';
 
         for await (const event of stream) {
-          if (signal.aborted) break;
-
           if (event.type === ServerGeminiEventType.Content) {
             if (
               pendingHistoryItemRef.current?.type !== 'gemini' &&
@@ -293,6 +291,18 @@ export const useGeminiStream = (
             );
             setStreamingState(StreamingState.WaitingForConfirmation);
             return; // Wait for user confirmation
+          } else if (event.type === ServerGeminiEventType.UserCancelled) {
+            // Flush out existing pending history item.
+            if (pendingHistoryItemRef.current) {
+              addItem(pendingHistoryItemRef.current, userMessageTimestamp);
+              setPendingHistoryItem(null);
+            }
+            addItem(
+              { type: 'info', text: 'User cancelled the request.' },
+              userMessageTimestamp,
+            );
+            setStreamingState(StreamingState.Idle);
+            return; // Stop processing the stream
           }
         } // End stream loop
 
diff --git a/packages/server/src/core/turn.ts b/packages/server/src/core/turn.ts
index ad461d74..7d8bf7b6 100644
--- a/packages/server/src/core/turn.ts
+++ b/packages/server/src/core/turn.ts
@@ -47,6 +47,7 @@ export enum GeminiEventType {
   ToolCallRequest = 'tool_call_request',
   ToolCallResponse = 'tool_call_response',
   ToolCallConfirmation = 'tool_call_confirmation',
+  UserCancelled = 'user_cancelled',
 }
 
 export interface ToolCallRequestInfo {
@@ -74,7 +75,8 @@ export type ServerGeminiStreamEvent =
   | {
       type: GeminiEventType.ToolCallConfirmation;
       value: ServerToolCallConfirmationDetails;
-    };
+    }
+  | { type: GeminiEventType.UserCancelled };
 
 // A turn manages the agentic loop turn within the server context.
 export class Turn {
@@ -108,7 +110,8 @@ export class Turn {
     for await (const resp of responseStream) {
       this.debugResponses.push(resp);
       if (signal?.aborted) {
-        throw this.abortError();
+        yield { type: GeminiEventType.UserCancelled };
+        return;
       }
 
       const text = getResponseText(resp);
@@ -240,12 +243,6 @@ export class Turn {
     };
   }
 
-  private abortError(): Error {
-    const error = new Error('Request cancelled by user during stream.');
-    error.name = 'AbortError';
-    return error; // Return instead of throw, let caller handle
-  }
-
   getConfirmationDetails(): ToolCallConfirmationDetails[] {
     return this.confirmationDetails;
   }