fix(agent-manager): restore focus to question options and prompt on session switch

.changeset/fix-agent-manager-focus.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Restore keyboard focus to the prompt or pending question when switching Agent Manager worktrees and sessions.

packages/kilo-vscode/tests/unit/agent-manager-focus.test.ts

@@ -0,0 +1,56 @@
+import { describe, expect, it } from "bun:test"
+import { Window } from "happy-dom"
+import { focusQuestionOption, hasQuestionOption } from "../../webview-ui/agent-manager/focus"
+
+describe("Agent Manager focus", () => {
+  it("focuses the first enabled question option", () => {
+    const window = new Window()
+    const root = window.document.createElement("div")
+    const dock = window.document.createElement("div")
+    const disabled = window.document.createElement("button")
+    const option = window.document.createElement("button")
+    disabled.setAttribute("data-slot", "question-option")
+    disabled.disabled = true
+    option.setAttribute("data-slot", "question-option")
+    dock.setAttribute("data-component", "question-dock")
+    dock.append(disabled, option)
+    root.append(dock)
+    window.document.body.append(root)
+
+    expect(focusQuestionOption(root)).toBe(true)
+    expect(root.ownerDocument.activeElement).toBe(option)
+  })
+
+  it("ignores collapsed question bodies", () => {
+    const window = new Window()
+    const root = window.document.createElement("div")
+    const dock = window.document.createElement("div")
+    const body = window.document.createElement("div")
+    const option = window.document.createElement("button")
+    dock.setAttribute("data-component", "question-dock")
+    body.setAttribute("inert", "")
+    option.setAttribute("data-slot", "question-option")
+    body.append(option)
+    dock.append(body)
+    root.append(dock)
+    window.document.body.append(root)
+
+    expect(focusQuestionOption(root)).toBe(false)
+    expect(root.ownerDocument.activeElement).not.toBe(option)
+  })
+
+  it("only reports enabled options outside inert bodies", () => {
+    const window = new Window()
+    const root = window.document.createElement("div")
+    const dock = window.document.createElement("div")
+    const option = window.document.createElement("button")
+    dock.setAttribute("data-component", "question-dock")
+    option.setAttribute("data-slot", "question-option")
+    dock.append(option)
+    root.append(dock)
+
+    expect(hasQuestionOption(root)).toBe(true)
+    dock.setAttribute("inert", "")
+    expect(hasQuestionOption(root)).toBe(false)
+  })
+})

packages/kilo-vscode/webview-ui/agent-manager/AgentManagerApp.tsx

@@ -174,6 +174,7 @@ import { SidebarToggleButton } from "./SidebarToggleButton"
 import { setTabWidths } from "./tab-widths"
 import { buildShortcutCategories } from "./shortcuts"
 import { tracker } from "./telemetry"
+import { createChatFocus, hasQuestionOption } from "./focus"
 import "./agent-manager.css"
 import "./agent-manager-review.css"
 import { cycleAgent as cycle } from "../src/context/session-agent"
@@ -379,6 +380,27 @@ const AgentManagerContent: Component = () => {
     const sel = selection()
     return sel === null ? null : nsKey(sel)
   })
+  const requestChatFocus = createChatFocus({
+    term: () => terms.activeId(),
+    history,
+    review: reviewActive,
+  })
+
+  createEffect(
+    on(
+      () => {
+        const id = session.currentSessionID()
+        return `${id ?? ""}:${session
+          .scopedQuestions(id)
+          .map((question) => question.id)
+          .join(",")}`
+      },
+      () => {
+        requestChatFocus()
+      },
+      { defer: true },
+    ),
+  )
 
   type FocusOwner = "prompt" | { terminal: string }
   const focusMemory = new Map<string, FocusOwner>()
@@ -425,7 +447,7 @@ const AgentManagerContent: Component = () => {
       }
       if (!terminal) focusMemory.delete(key)
     }
-    window.dispatchEvent(new Event("focusPrompt"))
+    requestChatFocus()
   }
   createEffect(
     on(
@@ -913,12 +935,14 @@ const AgentManagerContent: Component = () => {
     setSelection(null)
     setReviewActive(false)
     session.selectSession(id)
+    requestChatFocus(true)
   }
 
   const focusSidebarItem = (item: { type: string; id: string }) => {
     if (item.type === "local") selectLocal()
     else if (item.type === "wt") selectWorktree(item.id)
     else selectUnassigned(item.id)
+    requestChatFocus(true)
     const el = document.querySelector(`[data-sidebar-id="${item.id}"]`)
     if (el instanceof HTMLElement) scrollIntoView(el)
   }
@@ -951,6 +975,7 @@ const AgentManagerContent: Component = () => {
     const next = direction === "left" ? idx - 1 : idx + 1
     if (next < 0 || next >= ids.length) return
     focusTab(ids[next]!)
+    requestChatFocus(true)
   }
 
   const selectionDeps = {
@@ -971,10 +996,15 @@ const AgentManagerContent: Component = () => {
       remembered === REVIEW_TAB_ID && reviewOpenByContext()[sel] === true,
   }
 
-  const selectLocal = () => selectLocalAction(selectionDeps, localSessions())
+  const selectLocal = () => {
+    selectLocalAction(selectionDeps, localSessions())
+    requestChatFocus()
+  }
 
-  const selectWorktree = (worktreeId: string) =>
+  const selectWorktree = (worktreeId: string) => {
     selectWorktreeAction(selectionDeps, worktreeId, sessionsForWorktree(worktreeId))
+    requestChatFocus()
+  }
 
   const addSessionToCurrentWorktree = (sid: string) => {
     const sel = selection()
@@ -994,6 +1024,7 @@ const AgentManagerContent: Component = () => {
     selectWorktree(worktreeId)
     setHistory(false)
     session.selectSession(sid)
+    requestChatFocus()
     return true
   }
 
@@ -1113,6 +1144,7 @@ const AgentManagerContent: Component = () => {
         setSelection,
         setActivePendingId,
       })
+      requestChatFocus()
     }
     // Recover sidebar collapsed state and mark hydrated so transitions enable
     sidebar.hydrate(state.sidebarCollapsed)
@@ -1177,7 +1209,7 @@ const AgentManagerContent: Component = () => {
       else if (msg.action === "advancedWorktree") showNewWorktreeDialog()
       else if (msg.action === "closeWorktree") closeSelectedWorktree()
       else if (msg.action === "showShortcuts") handleShowKeyboardShortcuts()
-      else if (msg.action === "focusInput") window.dispatchEvent(new Event("focusPrompt"))
+      else if (msg.action === "focusInput") requestChatFocus(true)
       else if (msg.action === "focusSearch")
         focusChatSearch({ history: setHistory, review: setReviewActive, terminal: () => terms.setActiveId(undefined) })
       else if (msg.action === "newTerminal") termHandlers.requestNew()
@@ -1382,6 +1414,7 @@ const AgentManagerContent: Component = () => {
             const ms = managedSessions().find((s) => s.id === ev.sessionId)
             if (ms?.worktreeId) setSelection(ms.worktreeId)
             evictLocal(ev.sessionId)
+            requestChatFocus(true)
           }
         } else {
           // Track this worktree as setting up and auto-select it in the sidebar
@@ -1410,6 +1443,7 @@ const AgentManagerContent: Component = () => {
         evictLocal(ev.sessionId)
         drafts.apply(ev.worktreeId, ev.sessionId)
         session.selectSession(ev.sessionId)
+        requestChatFocus(true)
       }
 
       if (msg.type === "agentManager.sessionForked") {
@@ -1429,6 +1463,7 @@ const AgentManagerContent: Component = () => {
           evictLocal(ev.sessionId)
         }
         session.selectSession(ev.sessionId)
+        requestChatFocus(true)
       }
 
       if (msg.type === "agentManager.keybindings") {
@@ -1969,6 +2004,7 @@ const AgentManagerContent: Component = () => {
     setSelection(LOCAL)
     setReviewActive(false)
     session.selectSession(sid)
+    requestChatFocus()
     vscode.postMessage({ type: "agentManager.openLocally", sessionId: sid })
   }
 
@@ -2077,7 +2113,7 @@ const AgentManagerContent: Component = () => {
       cancelAmbientSetup()
       setSidePanel(null)
     },
-    refocus: () => window.dispatchEvent(new Event("focusPrompt")),
+    refocus: requestChatFocus,
     postMessage: (msg) => vscode.postMessage(msg as never),
     track: (button, surface, properties) => metrics.track(button, surface, properties),
     // Panel-local pick, immune to cross-window setting echoes (see side.ts).
@@ -2173,7 +2209,7 @@ const AgentManagerContent: Component = () => {
     return activeTabs().find((s) => s.id === id)
   })
 
-  const focusTab = (id: string) =>
+  const focusTab = (id: string) => {
     focusCurrentTab({
       id,
       terms,
@@ -2189,6 +2225,7 @@ const AgentManagerContent: Component = () => {
       selectSession: session.selectSession,
       activateTerminal: termHandlers.activate,
     })
+  }
   const tabFocus = createTabFocus({ ids: () => tabIds(), select: focusTab })
 
   // Close the currently active tab via keyboard shortcut.
@@ -2498,13 +2535,15 @@ const AgentManagerContent: Component = () => {
                 saveTabMemory()
                 session.selectSession(id)
                 setSelection(LOCAL)
+                requestChatFocus(true)
                 return
               }
               const ms = worktreeSessionIds().has(id) ? managedSessions().find((s) => s.id === id) : undefined
               if (ms?.worktreeId) {
                 selectWorktree(ms.worktreeId)
                 session.selectSession(id)
                 setReviewActive(false)
+                requestChatFocus()
                 return
               }
               openLocally(id)
@@ -2559,6 +2598,7 @@ const AgentManagerContent: Component = () => {
                         if (localSessionIDs().includes(id)) {
                           session.selectSession(id)
                           if (selection() === null) setSelection(LOCAL)
+                          requestChatFocus()
                           return
                         }
                         // Navigate to owning worktree instead of forcing into local mode
@@ -2568,6 +2608,7 @@ const AgentManagerContent: Component = () => {
                             selectWorktree(ms.worktreeId)
                             session.selectSession(id)
                             setReviewActive(false)
+                            requestChatFocus()
                             return
                           }
                         }
@@ -2579,6 +2620,7 @@ const AgentManagerContent: Component = () => {
                       readonly={readOnly()}
                       continueInWorktree={selection() === LOCAL}
                       promptBoxId={`agent-manager:${selection() ?? "unassigned"}`}
+                      deferFocusToQuestion={hasQuestionOption}
                       pendingSessionID={selection() === LOCAL ? activePendingId() : undefined}
                       focusOnDraftChange={focusOnDraftChange}
                       onFocusChange={rememberPromptFocus}

packages/kilo-vscode/webview-ui/agent-manager/focus.ts

@@ -0,0 +1,48 @@
+const OPTION = '[data-component="question-dock"] button[data-slot="question-option"]'
+
+export function createChatFocus(deps: {
+  term: () => string | undefined
+  history: () => boolean
+  review: () => boolean
+}) {
+  const focus = (force: boolean) => {
+    if ((!force && !document.hasFocus()) || deps.term() || deps.history() || deps.review()) return
+    if (!force && document.activeElement?.matches('[role="tab"]')) return
+    if (!force && document.activeElement?.closest('[data-component="question-dock"]')) return
+    if (focusQuestionOption()) return
+    const defer = hasQuestionOption()
+    window.dispatchEvent(
+      new CustomEvent("focusPrompt", {
+        detail: { restore: !defer, deferFocusToQuestion: defer },
+      }),
+    )
+  }
+  return (force = false) => {
+    queueMicrotask(() => focus(force))
+    requestAnimationFrame(() => {
+      focus(force)
+      requestAnimationFrame(() => {
+        focus(force)
+        requestAnimationFrame(() => focus(force))
+      })
+    })
+  }
+}
+
+/** Return whether the visible question dock has an enabled option to focus. */
+export function hasQuestionOption(root: ParentNode = document): boolean {
+  for (const option of root.querySelectorAll<HTMLButtonElement>(OPTION)) {
+    if (!option.disabled && !option.closest("[inert]")) return true
+  }
+  return false
+}
+
+/** Focus the first enabled option in the visible question dock, if one exists. */
+export function focusQuestionOption(root: ParentNode = document): boolean {
+  for (const option of root.querySelectorAll<HTMLButtonElement>(OPTION)) {
+    if (option.disabled || option.closest("[inert]")) continue
+    option.focus({ preventScroll: true })
+    return true
+  }
+  return false
+}

packages/kilo-vscode/webview-ui/src/components/chat/ChatView.tsx

@@ -39,6 +39,7 @@ interface ChatViewProps {
   /** When true, show the "Continue in Worktree" button. Defaults to true in the sidebar. */
   continueInWorktree?: boolean
   promptBoxId?: string
+  deferFocusToQuestion?: () => boolean
   pendingSessionID?: string
   focusOnDraftChange?: () => boolean
   onFocusChange?: (focused: boolean) => void
@@ -386,6 +387,7 @@ export const ChatView: Component<ChatViewProps> = (props) => {
                 suggesting={suggesting}
                 questioning={questioning}
                 boxId={props.promptBoxId}
+                deferFocusToQuestion={props.deferFocusToQuestion}
                 pendingSessionID={pendingSessionID()}
                 focusOnDraftChange={props.focusOnDraftChange}
                 onFocusChange={props.onFocusChange}

packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx

@@ -109,6 +109,8 @@ interface PromptInputProps {
   suggesting?: () => boolean
   /** When true, session is busy only because a question is pending — treat as idle for input */
   questioning?: () => boolean
+  /** When true, defer prompt focus while switching to a pending question */
+  deferFocusToQuestion?: () => boolean
   boxId?: string
   pendingSessionID?: string
   /** Agent Manager can suppress automatic prompt focus when this session last
@@ -386,7 +388,9 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
         textareaRef.scrollTop = scroll
         if (highlightRef) highlightRef.scrollTop = scroll
       }
-      if (props.focusOnDraftChange?.() ?? true) window.dispatchEvent(new Event("focusPrompt"))
+      if (!props.deferFocusToQuestion?.() && (props.focusOnDraftChange?.() ?? true)) {
+        window.dispatchEvent(new Event("focusPrompt"))
+      }
     }),
   )
 
@@ -410,14 +414,18 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
 
   // Focus textarea when any part of the app requests it
   const onFocusPrompt = (event: Event) => {
+    const defer = () =>
+      event instanceof CustomEvent && event.detail?.deferFocusToQuestion && props.deferFocusToQuestion?.()
     const focus = () => {
+      if (defer()) return
       const ref = textareaRef
       if (!ref) return
       ref.focus({ preventScroll: true })
     }
     focus()
     if (!(event instanceof CustomEvent) || !event.detail?.restore) return
     const restore = () => {
+      if (defer()) return
       window.focus()
       focus()
     }