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

packages/app/e2e/prompt/prompt-footer-focus.spec.ts

@@ -0,0 +1,88 @@
+import type { Locator, Page } from "@playwright/test"
+import { test, expect } from "../fixtures"
+import { promptAgentSelector, promptModelSelector, promptSelector } from "../selectors"
+
+type Probe = {
+  agent?: string
+  model?: { providerID: string; modelID: string; name?: string }
+  models?: Array<{ providerID: string; modelID: string; name: string }>
+  agents?: Array<{ name: string }>
+}
+
+async function probe(page: Page): Promise<Probe | null> {
+  return page.evaluate(() => {
+    const win = window as Window & {
+      __opencode_e2e?: {
+        model?: {
+          current?: Probe
+        }
+      }
+    }
+    return win.__opencode_e2e?.model?.current ?? null
+  })
+}
+
+async function state(page: Page) {
+  const value = await probe(page)
+  if (!value) throw new Error("Failed to resolve model selection probe")
+  return value
+}
+
+async function ready(page: Page) {
+  const prompt = page.locator(promptSelector)
+  await prompt.click()
+  await expect(prompt).toBeFocused()
+  await prompt.pressSequentially("focus")
+  return prompt
+}
+
+async function body(prompt: Locator) {
+  return prompt.evaluate((el) => (el as HTMLElement).innerText)
+}
+
+test("agent select returns focus to the prompt", async ({ page, gotoSession }) => {
+  await gotoSession()
+
+  const prompt = await ready(page)
+
+  const info = await state(page)
+  const next = info.agents?.map((item) => item.name).find((name) => name !== info.agent)
+  test.skip(!next, "only one agent available")
+  if (!next) return
+
+  await page.locator(`${promptAgentSelector} [data-slot="select-select-trigger"]`).first().click()
+
+  const item = page.locator('[data-slot="select-select-item"]').filter({ hasText: next }).first()
+  await expect(item).toBeVisible()
+  await item.click({ force: true })
+
+  await expect(page.locator(`${promptAgentSelector} [data-slot="select-select-trigger-value"]`).first()).toHaveText(
+    next,
+  )
+  await expect(prompt).toBeFocused()
+  await prompt.pressSequentially(" agent")
+  await expect.poll(() => body(prompt)).toContain("focus agent")
+})
+
+test("model select returns focus to the prompt", async ({ page, gotoSession }) => {
+  await gotoSession()
+
+  const prompt = await ready(page)
+
+  const info = await state(page)
+  const key = info.model ? `${info.model.providerID}:${info.model.modelID}` : null
+  const next = info.models?.find((item) => `${item.providerID}:${item.modelID}` !== key)
+  test.skip(!next, "only one model available")
+  if (!next) return
+
+  await page.locator(`${promptModelSelector} [data-action="prompt-model"]`).first().click()
+
+  const item = page.locator(`[data-slot="list-item"][data-key="${next.providerID}:${next.modelID}"]`).first()
+  await expect(item).toBeVisible()
+  await item.click({ force: true })
+
+  await expect(page.locator(`${promptModelSelector} [data-action="prompt-model"] span`).first()).toHaveText(next.name)
+  await expect(prompt).toBeFocused()
+  await prompt.pressSequentially(" model")
+  await expect.poll(() => body(prompt)).toContain("focus model")
+})

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">

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()}