fix: restore prompt focus after footer selection (#20841)

2 files changed, 38 insertions, 17 deletions

diff --git a/packages/app/src/components/dialog-select-model.tsx b/packages/app/src/components/dialog-select-model.tsx
index cb688c30a..fdef866a7 100644
--- a/packages/app/src/components/dialog-select-model.tsx
+++ b/packages/app/src/components/dialog-select-model.tsx
@@ -86,6 +86,7 @@ const ModelList: Component<{
 }
 
 type ModelSelectorTriggerProps = Omit<ComponentProps<typeof Kobalte.Trigger>, "as" | "ref">
+type Dismiss = "escape" | "outside" | "select" | "manage" | "provider"
 
 export function ModelSelectorPopover(props: {
   provider?: string
@@ -93,25 +94,31 @@ export function ModelSelectorPopover(props: {
   children?: JSX.Element
   triggerAs?: ValidComponent
   triggerProps?: ModelSelectorTriggerProps
+  onClose?: (cause: "escape" | "select") => void
 }) {
   const [store, setStore] = createStore<{
     open: boolean
-    dismiss: "escape" | "outside" | null
+    dismiss: Dismiss | null
   }>({
     open: false,
     dismiss: null,
   })
   const dialog = useDialog()
 
-  const handleManage = () => {
+  const close = (dismiss: Dismiss) => {
+    setStore("dismiss", dismiss)
     setStore("open", false)
+  }
+
+  const handleManage = () => {
+    close("manage")
     void import("./dialog-manage-models").then((x) => {
       dialog.show(() => <x.DialogManageModels />)
     })
   }
 
   const handleConnectProvider = () => {
-    setStore("open", false)
+    close("provider")
     void import("./dialog-select-provider").then((x) => {
       dialog.show(() => <x.DialogSelectProvider />)
     })
@@ -136,21 +143,19 @@ export function ModelSelectorPopover(props: {
         <Kobalte.Content
           class="w-72 h-80 flex flex-col p-2 rounded-md border border-border-base bg-surface-raised-stronger-non-alpha shadow-md z-50 outline-none overflow-hidden"
           onEscapeKeyDown={(event) => {
-            setStore("dismiss", "escape")
-            setStore("open", false)
+            close("escape")
             event.preventDefault()
             event.stopPropagation()
           }}
-          onPointerDownOutside={() => {
-            setStore("dismiss", "outside")
-            setStore("open", false)
-          }}
-          onFocusOutside={() => {
-            setStore("dismiss", "outside")
-            setStore("open", false)
-          }}
+          onPointerDownOutside={() => close("outside")}
+          onFocusOutside={() => close("outside")}
           onCloseAutoFocus={(event) => {
-            if (store.dismiss === "outside") event.preventDefault()
+            const dismiss = store.dismiss
+            if (dismiss === "outside") event.preventDefault()
+            if (dismiss === "escape" || dismiss === "select") {
+              event.preventDefault()
+              props.onClose?.(dismiss)
+            }
             setStore("dismiss", null)
           }}
         >
@@ -158,7 +163,7 @@ export function ModelSelectorPopover(props: {
           <ModelList
             provider={props.provider}
             model={props.model}
-            onSelect={() => setStore("open", false)}
+            onSelect={() => close("select")}
             class="p-1"
             action={
               <div class="flex items-center gap-1">
diff --git a/packages/app/src/components/prompt-input.tsx b/packages/app/src/components/prompt-input.tsx
index 653e89f51..ff31c8c2d 100644
--- a/packages/app/src/components/prompt-input.tsx
+++ b/packages/app/src/components/prompt-input.tsx
@@ -502,6 +502,15 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
     return getCursorPosition(editorRef)
   }
 
+  const restoreFocus = () => {
+    requestAnimationFrame(() => {
+      const cursor = prompt.cursor() ?? promptLength(prompt.current())
+      editorRef.focus()
+      setCursorPosition(editorRef, cursor)
+      queueScroll()
+    })
+  }
+
   const renderEditorWithCursor = (parts: Prompt) => {
     const cursor = currentCursor()
     renderEditor(parts)
@@ -1471,7 +1480,10 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
                       size="normal"
                       options={agentNames()}
                       current={local.agent.current()?.name ?? ""}
-                      onSelect={local.agent.set}
+                      onSelect={(value) => {
+                        local.agent.set(value)
+                        restoreFocus()
+                      }}
                       class="capitalize max-w-[160px] text-text-base"
                       valueClass="truncate text-13-regular text-text-base"
                       triggerStyle={control()}
@@ -1535,6 +1547,7 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
                             class: "min-w-0 max-w-[320px] text-13-regular text-text-base group",
                             "data-action": "prompt-model",
                           }}
+                          onClose={restoreFocus}
                         >
                           <Show when={local.model.current()?.provider?.id}>
                             <ProviderIcon
@@ -1563,7 +1576,10 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
                         options={variants()}
                         current={local.model.variant.current() ?? "default"}
                         label={(x) => (x === "default" ? language.t("common.default") : x)}
-                        onSelect={(x) => local.model.variant.set(x === "default" ? undefined : x)}
+                        onSelect={(value) => {
+                          local.model.variant.set(value === "default" ? undefined : value)
+                          restoreFocus()
+                        }}
                         class="capitalize max-w-[160px] text-text-base"
                         valueClass="truncate text-13-regular text-text-base"
                         triggerStyle={control()}