fix(agent-manager): route mode shortcuts through modal

.changeset/agent-manager-modal-mode-shortcut.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Fix Agent Manager mode shortcuts in the New Worktree dialog so the selected mode and its matching model stay in sync.

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

@@ -0,0 +1,36 @@
+import { describe, expect, it } from "bun:test"
+import { createModeRouter } from "../../webview-ui/agent-manager/mode-router"
+
+describe("Agent Manager mode router", () => {
+  it("dispatches to the active modal handler and reports consumption", () => {
+    const router = createModeRouter()
+    const directions: number[] = []
+
+    router.register((direction) => directions.push(direction))
+
+    expect(router.dispatch(1)).toBe(true)
+    expect(router.dispatch(-1)).toBe(true)
+    expect(directions).toEqual([1, -1])
+  })
+
+  it("restores normal routing after the modal unregisters", () => {
+    const router = createModeRouter()
+    const dispose = router.register(() => undefined)
+
+    dispose()
+
+    expect(router.dispatch(1)).toBe(false)
+  })
+
+  it("does not let an old modal cleanup remove a replacement handler", () => {
+    const router = createModeRouter()
+    const first = router.register(() => undefined)
+    const directions: number[] = []
+
+    router.register((direction) => directions.push(direction))
+    first()
+
+    expect(router.dispatch(1)).toBe(true)
+    expect(directions).toEqual([1])
+  })
+})

packages/kilo-vscode/tests/unit/session-model-store.test.ts

@@ -3,6 +3,7 @@ import {
   type ModelStore,
   type ResolveEnv,
   applyModel,
+  getAgentModel,
   getSessionModel,
   getSelected,
 } from "../../webview-ui/src/context/session-model-store"
@@ -149,6 +150,22 @@ describe("per-session model selection", () => {
 })
 
 describe("per-mode model memory", () => {
+  it("uses remembered model selections for modes without configured models", () => {
+    const store = { ...emptyStore(), modelSelections: { ask: gpt } }
+
+    expect(getAgentModel(store, env(), "ask")).toEqual(gpt)
+  })
+
+  it("ignores stale remembered selections when a configured mode model is user-set", () => {
+    const configured: ResolveEnv = {
+      ...env(),
+      getModeModel: (name) => (name === "code" ? claude : null),
+    }
+    const store = { ...emptyStore(), modelSelections: { code: gpt } }
+
+    expect(getAgentModel(store, configured, "code", true)).toEqual(claude)
+  })
+
   it("applyModel in a session writes only to sessionOverrides", () => {
     const store = emptyStore()
     const result = applyModel(store, "code", claude, "session-a")

packages/kilo-vscode/tests/unit/session-variant-store.test.ts

@@ -1,6 +1,7 @@
 import { describe, expect, it } from "bun:test"
 import {
   cycleVariant,
+  getAgentVariant,
   getVariant,
   sessionVariantKeys,
   sessionVariants,
@@ -43,6 +44,13 @@ describe("per-session variant selection", () => {
     expect(getVariant(store, model, variants, "ask")).toBe("high")
   })
 
+  it("resolves the effective variant for a mode and model", () => {
+    const store: Record<string, string> = {}
+    store[variantKey(model, "ask")] = "high"
+
+    expect(getAgentVariant(store, model, { variants: { low: {}, high: {} } }, "ask")).toBe("high")
+  })
+
   it("carries the pre-submit agent variant into a newly created session", () => {
     const store: Record<string, string> = {}
 

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

@@ -76,6 +76,7 @@ import { ProviderShell } from "../src/context/provider-shell"
 import { ChatView } from "../src/components/chat"
 import HistoryView from "../src/components/history/HistoryView"
 import { NewWorktreeDialog } from "./NewWorktreeDialog"
+import { createModeRouter } from "./mode-router"
 import { ProjectList } from "./ProjectList"
 import { SidebarBody } from "./SidebarBody"
 import { TabBar } from "./TabBar"
@@ -227,6 +228,7 @@ const AgentManagerContent: Component = () => {
   const session = useSession()
   const vscode = useVSCode()
   const dialog = useDialog()
+  const mode = createModeRouter()
   let sidebarSearchMenu: SidebarSearchMenuRef | undefined
 
   const [kb, setKb] = createSignal<Record<string, string>>(defaultBindings)
@@ -1111,9 +1113,11 @@ const AgentManagerContent: Component = () => {
       else if (msg.action === "focusSearch")
         focusChatSearch({ history: setHistory, review: setReviewActive, terminal: () => terms.setActiveId(undefined) })
       else if (msg.action === "newTerminal") termHandlers.requestNew()
-      else if (msg.action === "cycleAgentMode" && document.hasFocus()) cycleAgent(1)
-      else if (msg.action === "cyclePreviousAgentMode" && document.hasFocus()) cycleAgent(-1)
-      else {
+      else if (msg.action === "cycleAgentMode" && document.hasFocus()) {
+        if (!mode.dispatch(1)) cycleAgent(1)
+      } else if (msg.action === "cyclePreviousAgentMode" && document.hasFocus()) {
+        if (!mode.dispatch(-1)) cycleAgent(-1)
+      } else {
         // Handle jumpTo1 through jumpTo9
         const match = /^jumpTo([1-9])$/.exec(msg.action ?? "")
         if (match) projectNav.jump(parseInt(match[1]!) - 1)
@@ -1790,7 +1794,9 @@ const AgentManagerContent: Component = () => {
   const showNewWorktreeDialog = () => {
     if (!loaded()) return
     expandSidebar()
-    dialog.show(() => <NewWorktreeDialog onClose={() => dialog.close()} defaultBaseBranch={repoDefaultBranch()} />)
+    dialog.show(() => (
+      <NewWorktreeDialog mode={mode} onClose={() => dialog.close()} defaultBaseBranch={repoDefaultBranch()} />
+    ))
   }
 
   const confirmDeleteWorktree = (worktreeId: string) => {
@@ -2256,6 +2262,7 @@ const AgentManagerContent: Component = () => {
             selectedProject={activeProjectId()}
             selection={selection() ?? undefined}
             currentSessionID={session.currentSessionID}
+            mode={mode}
             bindings={kb()}
             t={t}
             onSearchRef={(ref) => (sidebarSearchMenu = ref)}

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

@@ -44,6 +44,8 @@ import { insertSpacedText } from "../src/components/chat/prompt-input-utils"
 import { WandSparkles } from "@kilocode/kilo-ui/lucide"
 import { BranchSelect, BranchSelectPopover } from "../src/components/shared/BranchSelect"
 import { tracker } from "./telemetry"
+import { cycleAgent } from "../src/context/session-agent"
+import type { ModeRouter } from "./mode-router"
 
 type VersionCount = 1 | 2 | 3 | 4
 const VERSION_OPTIONS: VersionCount[] = [1, 2, 3, 4]
@@ -74,9 +76,12 @@ function sanitizeBranchName(name: string): string {
     .join("/")
 }
 
-export const NewWorktreeDialog: Component<{ onClose: () => void; defaultBaseBranch?: string; projectId?: string }> = (
-  props,
-) => {
+export const NewWorktreeDialog: Component<{
+  onClose: () => void
+  defaultBaseBranch?: string
+  projectId?: string
+  mode: ModeRouter
+}> = (props) => {
   const { t } = useLanguage()
   const vscode = useVSCode()
   const server = useServer()
@@ -101,10 +106,12 @@ export const NewWorktreeDialog: Component<{ onClose: () => void; defaultBaseBran
   const cached = vscode.getState<Record<string, unknown>>()
   const [prompt, setPrompt] = createSignal((cached?.advancedDialogPrompt as string) ?? "")
   const [versions, setVersions] = createSignal<VersionCount>(1)
-  const [model, setModel] = createSignal<{ providerID: string; modelID: string } | null>(session.configModel())
+  const initialAgent = session.selectedAgent()
+  const initialModel = session.modelForAgent(initialAgent)
+  const [model, setModel] = createSignal<{ providerID: string; modelID: string } | null>(initialModel)
   const [compareMode, setCompareMode] = createSignal(false)
   const [modelAllocations, setModelAllocations] = createSignal<ModelAllocations>(new Map())
-  const [agent, setAgent] = createSignal(session.selectedAgent())
+  const [agent, setAgent] = createSignal(initialAgent)
   const [starting, setStarting] = createSignal(false)
   const [enhancing, setEnhancing] = createSignal(false)
   const [showAdvanced, setShowAdvanced] = createSignal(false)
@@ -113,7 +120,7 @@ export const NewWorktreeDialog: Component<{ onClose: () => void; defaultBaseBran
   const [baseBranchOpen, setBaseBranchOpen] = createSignal(false)
   const [compareOpen, setCompareOpen] = createSignal(false)
   const [highlightedIndex, setHighlightedIndex] = createSignal(0)
-  const [variant, setVariant] = createSignal<string | undefined>(session.currentVariant())
+  const [variant, setVariant] = createSignal<string | undefined>(session.variantForAgent(initialAgent, initialModel))
   const [sandbox, setSandbox] = createSignal<boolean | undefined>()
   const [sandboxDefault, setSandboxDefault] = createSignal<boolean | undefined>()
   const [sandboxOverride, setSandboxOverride] = createSignal<boolean | undefined>()
@@ -133,6 +140,34 @@ export const NewWorktreeDialog: Component<{ onClose: () => void; defaultBaseBran
     setEnhancing(false)
   }
 
+  const selectAgent = (name: string) => {
+    setAgent(name)
+    const sel = session.modelForAgent(name)
+    setModel(sel)
+    setVariant(session.variantForAgent(name, sel))
+  }
+
+  const resetModel = () => {
+    const sel = session.configModelForAgent(agent())
+    setModel(sel)
+    setVariant(session.variantForAgent(agent(), sel))
+  }
+
+  const cycle = (direction: 1 | -1) => {
+    cycleAgent({
+      agents: session.agents(),
+      direction,
+      selected: () => agent(),
+      select: selectAgent,
+    })
+  }
+
+  createEffect(() => {
+    if (tab() !== "new") return
+    const dispose = props.mode.register(cycle)
+    onCleanup(dispose)
+  })
+
   // Variant list for the currently selected model
   const variants = createMemo(() => {
     const sel = model()
@@ -153,7 +188,7 @@ export const NewWorktreeDialog: Component<{ onClose: () => void; defaultBaseBran
   // True when the user has changed the model from the session/config default
   const overridden = createMemo(() => {
     const sel = model()
-    const cfg = session.configModel()
+    const cfg = session.configModelForAgent(agent())
     if (!sel || !cfg) return false
     return sel.providerID !== cfg.providerID || sel.modelID !== cfg.modelID
   })
@@ -583,7 +618,7 @@ export const NewWorktreeDialog: Component<{ onClose: () => void; defaultBaseBran
                     <ModeSwitcherBase
                       agents={session.agents()}
                       value={agent()}
-                      onSelect={setAgent}
+                      onSelect={selectAgent}
                       portal={false}
                       deferDismiss
                     />
@@ -611,7 +646,7 @@ export const NewWorktreeDialog: Component<{ onClose: () => void; defaultBaseBran
                         <Button
                           variant="ghost"
                           size="small"
-                          onClick={() => setModel(session.configModel())}
+                          onClick={resetModel}
                           aria-label={t("prompt.action.resetModel")}
                         >
                           <svg width="12" height="12" viewBox="0 0 16 16" fill="currentColor">

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

@@ -19,6 +19,7 @@ import type { SidebarSearchItem } from "./sidebar-search"
 import { LOCAL } from "./navigate"
 import { NewWorktreeDialog } from "./NewWorktreeDialog"
 import { ProjectBranchDialog } from "./ProjectBranchDialog"
+import type { ModeRouter } from "./mode-router"
 
 interface Props {
   projects: AgentProjectSnapshot[]
@@ -30,6 +31,7 @@ interface Props {
   selectedProject?: string
   selection?: string
   currentSessionID?: () => string | undefined
+  mode: ModeRouter
   busy?: (id: string) => boolean
   bindings: Record<string, string>
   t: LanguageContextValue["t"]
@@ -131,6 +133,7 @@ export const ProjectList: Component<Props> = (props) => {
     dialog.show(() => (
       <NewWorktreeDialog
         projectId={projectId}
+        mode={props.mode}
         defaultBaseBranch={state?.defaultBaseBranch ?? props.local[projectId]?.branch}
         onClose={() => dialog.close()}
       />

packages/kilo-vscode/webview-ui/agent-manager/mode-router.ts

@@ -0,0 +1,26 @@
+export type ModeDirection = 1 | -1
+export type ModeHandler = (direction: ModeDirection) => void
+
+export interface ModeRouter {
+  register: (handler: ModeHandler) => () => void
+  dispatch: (direction: ModeDirection) => boolean
+}
+
+export function createModeRouter(): ModeRouter {
+  const state: { handler?: ModeHandler } = {}
+
+  return {
+    register(handler) {
+      state.handler = handler
+      return () => {
+        if (state.handler === handler) state.handler = undefined
+      }
+    },
+    dispatch(direction) {
+      const handler = state.handler
+      if (!handler) return false
+      handler(direction)
+      return true
+    },
+  }
+}

packages/kilo-vscode/webview-ui/src/context/session-model-store.ts

@@ -78,6 +78,17 @@ export function getSelected(
   return resolveModel(env, agentName, store.modelSelections[agentName], store.recentModels)
 }
 
+/** Returns the effective model for a mode outside a session scope. */
+export function getAgentModel(
+  store: ModelStore,
+  env: ResolveEnv,
+  agentName: string,
+  userSet = false,
+): ModelSelection | null {
+  const override = env.getModeModel(agentName) && userSet ? null : store.modelSelections[agentName]
+  return resolveModel(env, agentName, override, store.recentModels)
+}
+
 export interface ApplyResult {
   modelSelections: Record<string, ModelSelection | null>
   sessionOverrides: Record<string, ModelSelection>

packages/kilo-vscode/webview-ui/src/context/session-variant-store.ts

@@ -24,6 +24,16 @@ export function getVariant(
   return stored && variants.includes(stored) ? stored : variants[0]
 }
 
+export function getAgentVariant(
+  store: Record<string, string>,
+  sel: ModelSelection,
+  model: { variants?: Record<string, unknown> } | undefined,
+  agent: string,
+) {
+  if (!model?.variants) return undefined
+  return getVariant(store, sel, Object.keys(model.variants), agent)
+}
+
 /**
  * Next variant in the list, wrapping back to the first after the last.
  * An unknown or missing current value starts at the first variant.