fix: Restore user input when the user cancels response (#5601)

Co-authored-by: Shi Shu <[email protected]>
Co-authored-by: Jacob Richman <[email protected]>

1 files changed, 19 insertions, 10 deletions

diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx
index 66396c36..f2dcc79e 100644
--- a/packages/cli/src/ui/App.tsx
+++ b/packages/cli/src/ui/App.tsx
@@ -486,6 +486,24 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => {
     setGeminiMdFileCount,
   );
 
+  const buffer = useTextBuffer({
+    initialText: '',
+    viewport: { height: 10, width: inputWidth },
+    stdin,
+    setRawMode,
+    isValidPath,
+    shellModeActive,
+  });
+
+  const [userMessages, setUserMessages] = useState<string[]>([]);
+
+  const handleUserCancel = useCallback(() => {
+    const lastUserMessage = userMessages.at(-1);
+    if (lastUserMessage) {
+      buffer.setText(lastUserMessage);
+    }
+  }, [buffer, userMessages]);
+
   const {
     streamingState,
     submitQuery,
@@ -506,6 +524,7 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => {
     modelSwitchedFromQuotaError,
     setModelSwitchedFromQuotaError,
     refreshStatic,
+    handleUserCancel,
   );
 
   // Input handling
@@ -519,15 +538,6 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => {
     [submitQuery],
   );
 
-  const buffer = useTextBuffer({
-    initialText: '',
-    viewport: { height: 10, width: inputWidth },
-    stdin,
-    setRawMode,
-    isValidPath,
-    shellModeActive,
-  });
-
   const { handleInput: vimHandleInput } = useVim(buffer, handleFinalSubmit);
   const pendingHistoryItems = [...pendingSlashCommandHistoryItems];
   pendingHistoryItems.push(...pendingGeminiHistoryItems);
@@ -607,7 +617,6 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => {
   }, [config, config.getGeminiMdFileCount]);
 
   const logger = useLogger();
-  const [userMessages, setUserMessages] = useState<string[]>([]);
 
   useEffect(() => {
     const fetchUserMessages = async () => {