prevent AI in-editor input from disappearing on rejection

frontend/src/components/editor/ai/add-cell-with-ai.tsx

@@ -231,17 +231,21 @@ 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 = (
     <div className="flex items-center px-3">
       <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>

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 =