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: 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
---
 packages/cli/src/ui/components/messages/InfoMessage.tsx |  2 +-
 packages/cli/src/ui/hooks/useGeminiStream.ts            | 14 ++++++++++++--
 2 files changed, 13 insertions(+), 3 deletions(-)

(limited to 'packages/cli/src')

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
 
-- 
cgit v1.2.3