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..6f91a46c69d 100644
--- a/frontend/src/components/editor/ai/add-cell-with-ai.tsx
+++ b/frontend/src/components/editor/ai/add-cell-with-ai.tsx
@@ -231,6 +231,13 @@ export const AddCellWithAI: React.FC<{
 
   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..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 =