From 1f0ad865444c07481385c39b272f9ec2b94d41b9 Mon Sep 17 00:00:00 2001
From: shishu314 <shishu_1998@yahoo.com>
Date: Wed, 6 Aug 2025 15:19:10 -0400
Subject: fix: Restore user input when the user cancels response (#5601)

Co-authored-by: Shi Shu <shii@google.com>
Co-authored-by: Jacob Richman <jacob314@gmail.com>
---
 packages/cli/src/ui/hooks/useGeminiStream.test.tsx | 66 ++++++++++++++++++++++
 packages/cli/src/ui/hooks/useGeminiStream.ts       |  2 +
 2 files changed, 68 insertions(+)

(limited to 'packages/cli/src/ui/hooks')

diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
index dd2428bb..751b869e 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
+++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
@@ -406,6 +406,8 @@ describe('useGeminiStream', () => {
           () => Promise.resolve(),
           false,
           () => {},
+          () => {},
+          () => {},
         );
       },
       {
@@ -560,6 +562,8 @@ describe('useGeminiStream', () => {
         () => Promise.resolve(),
         false,
         () => {},
+        () => {},
+        () => {},
       ),
     );
 
@@ -633,6 +637,8 @@ describe('useGeminiStream', () => {
         () => Promise.resolve(),
         false,
         () => {},
+        () => {},
+        () => {},
       ),
     );
 
@@ -737,6 +743,8 @@ describe('useGeminiStream', () => {
         () => Promise.resolve(),
         false,
         () => {},
+        () => {},
+        () => {},
       ),
     );
 
@@ -843,6 +851,8 @@ describe('useGeminiStream', () => {
         () => Promise.resolve(),
         false,
         () => {},
+        () => {},
+        () => {},
       ),
     );
 
@@ -943,6 +953,44 @@ describe('useGeminiStream', () => {
       expect(result.current.streamingState).toBe(StreamingState.Idle);
     });
 
+    it('should call onCancelSubmit handler when escape is pressed', async () => {
+      const cancelSubmitSpy = vi.fn();
+      const mockStream = (async function* () {
+        yield { type: 'content', value: 'Part 1' };
+        // Keep the stream open
+        await new Promise(() => {});
+      })();
+      mockSendMessageStream.mockReturnValue(mockStream);
+
+      const { result } = renderHook(() =>
+        useGeminiStream(
+          mockConfig.getGeminiClient(),
+          [],
+          mockAddItem,
+          mockConfig,
+          mockOnDebugMessage,
+          mockHandleSlashCommand,
+          false,
+          () => 'vscode' as EditorType,
+          () => {},
+          () => Promise.resolve(),
+          false,
+          () => {},
+          () => {},
+          cancelSubmitSpy,
+        ),
+      );
+
+      // Start a query
+      await act(async () => {
+        result.current.submitQuery('test query');
+      });
+
+      simulateEscapeKeyPress();
+
+      expect(cancelSubmitSpy).toHaveBeenCalled();
+    });
+
     it('should not do anything if escape is pressed when not responding', () => {
       const { result } = renderTestHook();
 
@@ -1202,6 +1250,8 @@ describe('useGeminiStream', () => {
           mockPerformMemoryRefresh,
           false,
           () => {},
+          () => {},
+          () => {},
         ),
       );
 
@@ -1253,6 +1303,8 @@ describe('useGeminiStream', () => {
           () => Promise.resolve(),
           false,
           () => {},
+          () => {},
+          () => {},
         ),
       );
 
@@ -1301,6 +1353,8 @@ describe('useGeminiStream', () => {
           () => Promise.resolve(),
           false,
           () => {},
+          () => {},
+          () => {},
         ),
       );
 
@@ -1347,6 +1401,8 @@ describe('useGeminiStream', () => {
           () => Promise.resolve(),
           false,
           () => {},
+          () => {},
+          () => {},
         ),
       );
 
@@ -1394,6 +1450,8 @@ describe('useGeminiStream', () => {
           () => Promise.resolve(),
           false,
           () => {},
+          () => {},
+          () => {},
         ),
       );
 
@@ -1481,6 +1539,8 @@ describe('useGeminiStream', () => {
             () => Promise.resolve(),
             false,
             () => {},
+            () => {},
+            () => {},
           ),
         );
 
@@ -1535,6 +1595,8 @@ describe('useGeminiStream', () => {
           () => Promise.resolve(),
           false,
           () => {},
+          () => {},
+          () => {},
         ),
       );
 
@@ -1611,6 +1673,8 @@ describe('useGeminiStream', () => {
           () => Promise.resolve(),
           false,
           () => {},
+          () => {},
+          () => {},
         ),
       );
 
@@ -1663,6 +1727,8 @@ describe('useGeminiStream', () => {
           () => Promise.resolve(),
           false,
           () => {},
+          () => {},
+          () => {},
         ),
       );
 
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts
index 63ba961f..58bec431 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.ts
+++ b/packages/cli/src/ui/hooks/useGeminiStream.ts
@@ -94,6 +94,7 @@ export const useGeminiStream = (
   modelSwitchedFromQuotaError: boolean,
   setModelSwitchedFromQuotaError: React.Dispatch<React.SetStateAction<boolean>>,
   onEditorClose: () => void,
+  onCancelSubmit: () => void,
 ) => {
   const [initError, setInitError] = useState<string | null>(null);
   const abortControllerRef = useRef<AbortController | null>(null);
@@ -200,6 +201,7 @@ export const useGeminiStream = (
         Date.now(),
       );
       setPendingHistoryItem(null);
+      onCancelSubmit();
       setIsResponding(false);
     }
   });
-- 
cgit v1.2.3