From 57ac9958d1f0861a7b80136acf625aa548a5157b Mon Sep 17 00:00:00 2001
From: Shahmir Varqha <sham9871@gmail.com>
Date: Fri, 17 Jul 2026 17:24:03 +0800
Subject: [PATCH 1/5] prevent input from disappearing on rejection

---
 .../components/editor/ai/add-cell-with-ai.tsx | 21 +++++++++++++------
 .../editor/ai/ai-completion-editor.tsx        |  4 +++-
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/frontend/src/components/editor/ai/add-cell-with-ai.tsx b/frontend/src/components/editor/ai/add-cell-with-ai.tsx
index 13e79609eee..33ac0d71f2a 100644
--- a/frontend/src/components/editor/ai/add-cell-with-ai.tsx
+++ b/frontend/src/components/editor/ai/add-cell-with-ai.tsx
@@ -226,11 +226,18 @@ export const AddCellWithAI: React.FC<{
 
   const handleAcceptCompletion = () => {
     clearStagedCells();
-    onClose();
+    setInput("");
   };
 
   const handleDeclineCompletion = () => {
     deleteAllStagedCells();
+    // Focus the input so the user can refine the prompt.
+    inputRef.current?.view?.focus();
+  };
+
+  const handleClose = () => {
+    deleteAllStagedCells();
+    onClose();
   };
 
   const inputComponent = (
@@ -238,10 +245,7 @@ export const AddCellWithAI: React.FC<{
       <SparklesIcon className="size-4 text-(--blue-11) mr-2" />
       <PromptInput
         inputRef={inputRef}
-        onClose={() => {
-          deleteAllStagedCells();
-          onClose();
-        }}
+        onClose={handleClose}
         value={input}
         onChange={(newValue) => {
           setInput(newValue);
@@ -254,7 +258,12 @@ export const AddCellWithAI: React.FC<{
           hasCompletion,
         })}
       />
-      <Button variant="text" size="sm" className="mb-0 px-1" onClick={onClose}>
+      <Button
+        variant="text"
+        size="sm"
+        className="mb-0 px-1"
+        onClick={handleClose}
+      >
         <XIcon className="size-4" />
       </Button>
     </div>
diff --git a/frontend/src/components/editor/ai/ai-completion-editor.tsx b/frontend/src/components/editor/ai/ai-completion-editor.tsx
index ca85b52706c..fd8f0452391 100644
--- a/frontend/src/components/editor/ai/ai-completion-editor.tsx
+++ b/frontend/src/components/editor/ai/ai-completion-editor.tsx
@@ -194,9 +194,11 @@ export const AiCompletionEditor: React.FC<Props> = ({
     setCompletion("");
   };
 
+  // Reject discards the suggestion but keeps the prompt open for refinement.
   const handleDeclineCompletion = () => {
-    declineChange();
+    stop();
     setCompletion("");
+    inputRef.current?.view?.focus();
   };
 
   const showCompletionBanner =

From 1962bffc7624375d13a74abc92e8f2cb3fa09fcd Mon Sep 17 00:00:00 2001
From: Shahmir Varqha <sham9871@gmail.com>
Date: Fri, 17 Jul 2026 17:27:40 +0800
Subject: [PATCH 2/5] onClose

---
 frontend/src/components/editor/ai/add-cell-with-ai.tsx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/frontend/src/components/editor/ai/add-cell-with-ai.tsx b/frontend/src/components/editor/ai/add-cell-with-ai.tsx
index 33ac0d71f2a..6f91a46c69d 100644
--- a/frontend/src/components/editor/ai/add-cell-with-ai.tsx
+++ b/frontend/src/components/editor/ai/add-cell-with-ai.tsx
@@ -226,7 +226,7 @@ export const AddCellWithAI: React.FC<{
 
   const handleAcceptCompletion = () => {
     clearStagedCells();
-    setInput("");
+    onClose();
   };
 
   const handleDeclineCompletion = () => {

From 1620733b4d0e6d7cb45b21bad11f816201261631 Mon Sep 17 00:00:00 2001
From: Shahmir Varqha <sham9871@gmail.com>
Date: Fri, 17 Jul 2026 17:57:03 +0800
Subject: [PATCH 3/5] address review

---
 frontend/src/components/editor/ai/ai-completion-editor.tsx | 1 +
 1 file changed, 1 insertion(+)

diff --git a/frontend/src/components/editor/ai/ai-completion-editor.tsx b/frontend/src/components/editor/ai/ai-completion-editor.tsx
index fd8f0452391..109eda05831 100644
--- a/frontend/src/components/editor/ai/ai-completion-editor.tsx
+++ b/frontend/src/components/editor/ai/ai-completion-editor.tsx
@@ -198,6 +198,7 @@ export const AiCompletionEditor: React.FC<Props> = ({
   const handleDeclineCompletion = () => {
     stop();
     setCompletion("");
+    setShowInputPrompt(true);
     inputRef.current?.view?.focus();
   };
 

From 1622043a640215e91eed956be9aa41962c8a8ae8 Mon Sep 17 00:00:00 2001
From: Shahmir Varqha <sham9871@gmail.com>
Date: Fri, 17 Jul 2026 20:19:49 +0800
Subject: [PATCH 4/5] revert to original

---
 frontend/src/components/editor/ai/ai-completion-editor.tsx | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/frontend/src/components/editor/ai/ai-completion-editor.tsx b/frontend/src/components/editor/ai/ai-completion-editor.tsx
index 109eda05831..ca85b52706c 100644
--- a/frontend/src/components/editor/ai/ai-completion-editor.tsx
+++ b/frontend/src/components/editor/ai/ai-completion-editor.tsx
@@ -194,12 +194,9 @@ export const AiCompletionEditor: React.FC<Props> = ({
     setCompletion("");
   };
 
-  // Reject discards the suggestion but keeps the prompt open for refinement.
   const handleDeclineCompletion = () => {
-    stop();
+    declineChange();
     setCompletion("");
-    setShowInputPrompt(true);
-    inputRef.current?.view?.focus();
   };
 
   const showCompletionBanner =

From 6ecd6e246aa9e239b07a99e4f57906f84a0d86a5 Mon Sep 17 00:00:00 2001
From: Shahmir Varqha <sham9871@gmail.com>
Date: Fri, 17 Jul 2026 20:22:56 +0800
Subject: [PATCH 5/5] revert back

---
 frontend/src/components/editor/ai/ai-completion-editor.tsx | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/frontend/src/components/editor/ai/ai-completion-editor.tsx b/frontend/src/components/editor/ai/ai-completion-editor.tsx
index ca85b52706c..109eda05831 100644
--- a/frontend/src/components/editor/ai/ai-completion-editor.tsx
+++ b/frontend/src/components/editor/ai/ai-completion-editor.tsx
@@ -194,9 +194,12 @@ export const AiCompletionEditor: React.FC<Props> = ({
     setCompletion("");
   };
 
+  // Reject discards the suggestion but keeps the prompt open for refinement.
   const handleDeclineCompletion = () => {
-    declineChange();
+    stop();
     setCompletion("");
+    setShowInputPrompt(true);
+    inputRef.current?.view?.focus();
   };
 
   const showCompletionBanner =