fix(session): add keyboard support to question dock (#20439)

3 files changed, 183 insertions, 5 deletions

diff --git a/packages/app/e2e/session/session-composer-dock.spec.ts b/packages/app/e2e/session/session-composer-dock.spec.ts
index f083bf359..c56079337 100644
--- a/packages/app/e2e/session/session-composer-dock.spec.ts
+++ b/packages/app/e2e/session/session-composer-dock.spec.ts
@@ -13,6 +13,7 @@ import {
   sessionComposerDockSelector,
   sessionTodoToggleButtonSelector,
 } from "../selectors"
+import { modKey } from "../utils"
 
 type Sdk = Parameters<typeof clearSessionDockSeed>[0]
 type PermissionRule = { permission: string; pattern: string; action: "allow" | "deny" | "ask" }
@@ -310,6 +311,73 @@ test("blocked question flow unblocks after submit", async ({ page, sdk, gotoSess
   })
 })
 
+test("blocked question flow supports keyboard shortcuts", async ({ page, sdk, gotoSession }) => {
+  await withDockSession(sdk, "e2e composer dock question keyboard", async (session) => {
+    await withDockSeed(sdk, session.id, async () => {
+      await gotoSession(session.id)
+
+      await seedSessionQuestion(sdk, {
+        sessionID: session.id,
+        questions: [
+          {
+            header: "Need input",
+            question: "Pick one option",
+            options: [
+              { label: "Continue", description: "Continue now" },
+              { label: "Stop", description: "Stop here" },
+            ],
+          },
+        ],
+      })
+
+      const dock = page.locator(questionDockSelector)
+      const first = dock.locator('[data-slot="question-option"]').first()
+      const second = dock.locator('[data-slot="question-option"]').nth(1)
+
+      await expectQuestionBlocked(page)
+      await expect(first).toBeFocused()
+
+      await page.keyboard.press("ArrowDown")
+      await expect(second).toBeFocused()
+
+      await page.keyboard.press("Space")
+      await page.keyboard.press(`${modKey}+Enter`)
+      await expectQuestionOpen(page)
+    })
+  })
+})
+
+test("blocked question flow supports escape dismiss", async ({ page, sdk, gotoSession }) => {
+  await withDockSession(sdk, "e2e composer dock question escape", async (session) => {
+    await withDockSeed(sdk, session.id, async () => {
+      await gotoSession(session.id)
+
+      await seedSessionQuestion(sdk, {
+        sessionID: session.id,
+        questions: [
+          {
+            header: "Need input",
+            question: "Pick one option",
+            options: [
+              { label: "Continue", description: "Continue now" },
+              { label: "Stop", description: "Stop here" },
+            ],
+          },
+        ],
+      })
+
+      const dock = page.locator(questionDockSelector)
+      const first = dock.locator('[data-slot="question-option"]').first()
+
+      await expectQuestionBlocked(page)
+      await expect(first).toBeFocused()
+
+      await page.keyboard.press("Escape")
+      await expectQuestionOpen(page)
+    })
+  })
+})
+
 test("blocked permission flow supports allow once", async ({ page, sdk, gotoSession }) => {
   await withDockSession(sdk, "e2e composer dock permission once", async (session) => {
     await gotoSession(session.id)
diff --git a/packages/app/src/pages/session/composer/session-question-dock.tsx b/packages/app/src/pages/session/composer/session-question-dock.tsx
index ef1e52d26..38974b246 100644
--- a/packages/app/src/pages/session/composer/session-question-dock.tsx
+++ b/packages/app/src/pages/session/composer/session-question-dock.tsx
@@ -29,16 +29,20 @@ function Option(props: {
   label: string
   description?: string
   disabled: boolean
+  ref?: (el: HTMLButtonElement) => void
+  onFocus?: VoidFunction
   onClick: VoidFunction
 }) {
   return (
     <button
       type="button"
+      ref={props.ref}
       data-slot="question-option"
       data-picked={props.picked}
       role={props.multi ? "checkbox" : "radio"}
       aria-checked={props.picked}
       disabled={props.disabled}
+      onFocus={props.onFocus}
       onClick={props.onClick}
     >
       <Mark multi={props.multi} picked={props.picked} />
@@ -66,16 +70,21 @@ export const SessionQuestionDock: Component<{ request: QuestionRequest; onSubmit
     custom: cached?.custom ?? ([] as string[]),
     customOn: cached?.customOn ?? ([] as boolean[]),
     editing: false,
+    focus: 0,
   })
 
   let root: HTMLDivElement | undefined
+  let customRef: HTMLButtonElement | undefined
+  let optsRef: HTMLButtonElement[] = []
   let replied = false
+  let focusFrame: number | undefined
 
   const question = createMemo(() => questions()[store.tab])
   const options = createMemo(() => question()?.options ?? [])
   const input = createMemo(() => store.custom[store.tab] ?? "")
   const on = createMemo(() => store.customOn[store.tab] === true)
   const multi = createMemo(() => question()?.multiple === true)
+  const count = createMemo(() => options().length + 1)
 
   const summary = createMemo(() => {
     const n = Math.min(store.tab + 1, total())
@@ -129,6 +138,29 @@ export const SessionQuestionDock: Component<{ request: QuestionRequest; onSubmit
     root.style.setProperty("--question-prompt-max-height", `${max}px`)
   }
 
+  const clamp = (i: number) => Math.max(0, Math.min(count() - 1, i))
+
+  const pickFocus = (tab: number = store.tab) => {
+    const list = questions()[tab]?.options ?? []
+    if (store.customOn[tab] === true) return list.length
+    return Math.max(
+      0,
+      list.findIndex((item) => store.answers[tab]?.includes(item.label) ?? false),
+    )
+  }
+
+  const focus = (i: number) => {
+    const next = clamp(i)
+    setStore("focus", next)
+    if (store.editing) return
+    if (focusFrame !== undefined) cancelAnimationFrame(focusFrame)
+    focusFrame = requestAnimationFrame(() => {
+      focusFrame = undefined
+      const el = next === options().length ? customRef : optsRef[next]
+      el?.focus()
+    })
+  }
+
   onMount(() => {
     let raf: number | undefined
     const update = () => {
@@ -153,9 +185,12 @@ export const SessionQuestionDock: Component<{ request: QuestionRequest; onSubmit
       observer.disconnect()
       if (raf !== undefined) cancelAnimationFrame(raf)
     })
+
+    focus(pickFocus())
   })
 
   onCleanup(() => {
+    if (focusFrame !== undefined) cancelAnimationFrame(focusFrame)
     if (replied) return
     cache.set(props.request.id, {
       tab: store.tab,
@@ -231,6 +266,7 @@ export const SessionQuestionDock: Component<{ request: QuestionRequest; onSubmit
 
   const customToggle = () => {
     if (sending()) return
+    setStore("focus", options().length)
 
     if (!multi()) {
       setStore("customOn", store.tab, true)
@@ -250,15 +286,68 @@ export const SessionQuestionDock: Component<{ request: QuestionRequest; onSubmit
     const value = input().trim()
     if (value) setStore("answers", store.tab, (current = []) => current.filter((item) => item.trim() !== value))
     setStore("editing", false)
+    focus(options().length)
   }
 
   const customOpen = () => {
     if (sending()) return
+    setStore("focus", options().length)
     if (!on()) setStore("customOn", store.tab, true)
     setStore("editing", true)
     customUpdate(input(), true)
   }
 
+  const move = (step: number) => {
+    if (store.editing || sending()) return
+    focus(store.focus + step)
+  }
+
+  const nav = (event: KeyboardEvent) => {
+    if (event.defaultPrevented) return
+
+    if (event.key === "Escape") {
+      event.preventDefault()
+      void reject()
+      return
+    }
+
+    const mod = (event.metaKey || event.ctrlKey) && !event.altKey
+    if (mod && event.key === "Enter") {
+      if (event.repeat) return
+      event.preventDefault()
+      next()
+      return
+    }
+
+    const target =
+      event.target instanceof HTMLElement ? event.target.closest('[data-slot="question-options"]') : undefined
+    if (store.editing) return
+    if (!(target instanceof HTMLElement)) return
+    if (event.altKey || event.ctrlKey || event.metaKey) return
+
+    if (event.key === "ArrowDown" || event.key === "ArrowRight") {
+      event.preventDefault()
+      move(1)
+      return
+    }
+
+    if (event.key === "ArrowUp" || event.key === "ArrowLeft") {
+      event.preventDefault()
+      move(-1)
+      return
+    }
+
+    if (event.key === "Home") {
+      event.preventDefault()
+      focus(0)
+      return
+    }
+
+    if (event.key !== "End") return
+    event.preventDefault()
+    focus(count() - 1)
+  }
+
   const selectOption = (optIndex: number) => {
     if (sending()) return
 
@@ -270,6 +359,7 @@ export const SessionQuestionDock: Component<{ request: QuestionRequest; onSubmit
     const opt = options()[optIndex]
     if (!opt) return
     if (multi()) {
+      setStore("editing", false)
       toggle(opt.label)
       return
     }
@@ -279,6 +369,7 @@ export const SessionQuestionDock: Component<{ request: QuestionRequest; onSubmit
   const commitCustom = () => {
     setStore("editing", false)
     customUpdate(input())
+    focus(options().length)
   }
 
   const resizeInput = (el: HTMLTextAreaElement) => {
@@ -308,27 +399,33 @@ export const SessionQuestionDock: Component<{ request: QuestionRequest; onSubmit
       return
     }
 
-    setStore("tab", store.tab + 1)
+    const tab = store.tab + 1
+    setStore("tab", tab)
     setStore("editing", false)
+    focus(pickFocus(tab))
   }
 
   const back = () => {
     if (sending()) return
     if (store.tab <= 0) return
-    setStore("tab", store.tab - 1)
+    const tab = store.tab - 1
+    setStore("tab", tab)
     setStore("editing", false)
+    focus(pickFocus(tab))
   }
 
   const jump = (tab: number) => {
     if (sending()) return
     setStore("tab", tab)
     setStore("editing", false)
+    focus(pickFocus(tab))
   }
 
   return (
     <DockPrompt
       kind="question"
       ref={(el) => (root = el)}
+      onKeyDown={nav}
       header={
         <>
           <div data-slot="question-header-title">{summary()}</div>
@@ -351,7 +448,7 @@ export const SessionQuestionDock: Component<{ request: QuestionRequest; onSubmit
       }
       footer={
         <>
-          <Button variant="ghost" size="large" disabled={sending()} onClick={reject}>
+          <Button variant="ghost" size="large" disabled={sending()} onClick={reject} aria-keyshortcuts="Escape">
             {language.t("ui.common.dismiss")}
           </Button>
           <div data-slot="question-footer-actions">
@@ -360,7 +457,13 @@ export const SessionQuestionDock: Component<{ request: QuestionRequest; onSubmit
                 {language.t("ui.common.back")}
               </Button>
             </Show>
-            <Button variant={last() ? "primary" : "secondary"} size="large" disabled={sending()} onClick={next}>
+            <Button
+              variant={last() ? "primary" : "secondary"}
+              size="large"
+              disabled={sending()}
+              onClick={next}
+              aria-keyshortcuts="Meta+Enter Control+Enter"
+            >
               {last() ? language.t("ui.common.submit") : language.t("ui.common.next")}
             </Button>
           </div>
@@ -380,6 +483,8 @@ export const SessionQuestionDock: Component<{ request: QuestionRequest; onSubmit
               label={opt.label}
               description={opt.description}
               disabled={sending()}
+              ref={(el) => (optsRef[i()] = el)}
+              onFocus={() => setStore("focus", i())}
               onClick={() => selectOption(i())}
             />
           )}
@@ -390,12 +495,14 @@ export const SessionQuestionDock: Component<{ request: QuestionRequest; onSubmit
           fallback={
             <button
               type="button"
+              ref={customRef}
               data-slot="question-option"
               data-custom="true"
               data-picked={on()}
               role={multi() ? "checkbox" : "radio"}
               aria-checked={on()}
               disabled={sending()}
+              onFocus={() => setStore("focus", options().length)}
               onClick={customOpen}
             >
               <Mark multi={multi()} picked={on()} onClick={toggleCustomMark} />
@@ -440,8 +547,10 @@ export const SessionQuestionDock: Component<{ request: QuestionRequest; onSubmit
                   if (e.key === "Escape") {
                     e.preventDefault()
                     setStore("editing", false)
+                    focus(options().length)
                     return
                   }
+                  if ((e.metaKey || e.ctrlKey) && !e.altKey) return
                   if (e.key !== "Enter" || e.shiftKey) return
                   e.preventDefault()
                   commitCustom()
diff --git a/packages/ui/src/components/dock-prompt.tsx b/packages/ui/src/components/dock-prompt.tsx
index d774e7f17..01f1848c8 100644
--- a/packages/ui/src/components/dock-prompt.tsx
+++ b/packages/ui/src/components/dock-prompt.tsx
@@ -7,11 +7,12 @@ export function DockPrompt(props: {
   children: JSX.Element
   footer: JSX.Element
   ref?: (el: HTMLDivElement) => void
+  onKeyDown?: JSX.EventHandlerUnion<HTMLDivElement, KeyboardEvent>
 }) {
   const slot = (name: string) => `${props.kind}-${name}`
 
   return (
-    <div data-component="dock-prompt" data-kind={props.kind} ref={props.ref}>
+    <div data-component="dock-prompt" data-kind={props.kind} ref={props.ref} onKeyDown={props.onKeyDown}>
       <DockShell data-slot={slot("body")}>
         <div data-slot={slot("header")}>{props.header}</div>
         <div data-slot={slot("content")}>{props.children}</div>