fix(vscode): sync mode picker with agent mode changes from backend

packages/kilo-vscode/tests/unit/kilo-provider-utils.test.ts

@@ -363,11 +363,21 @@ describe("mapSSEEventToWebviewMessage", () => {
       properties: {
         id: "q1",
         sessionID: "sess-1",
-        questions: [],
+        questions: [
+          {
+            question: "Ready to implement?",
+            header: "Implement",
+            options: [{ label: "Implement", description: "Switch to code", mode: "code" }],
+          },
+        ],
       },
     }
     const msg = mapSSEEventToWebviewMessage(event, "sess-1")
     expect(msg?.type).toBe("questionRequest")
+    if (msg?.type === "questionRequest") {
+      const questions = msg.question.questions as Array<{ options?: Array<{ mode?: string }> }>
+      expect(questions[0]?.options?.[0]?.mode).toBe("code")
+    }
   })
 
   it("maps question.replied to questionResolved", () => {

packages/kilo-vscode/tests/unit/question-dock-utils.test.ts

@@ -1,5 +1,10 @@
 import { describe, it, expect } from "bun:test"
-import { toggleAnswer } from "../../webview-ui/src/components/chat/question-dock-utils"
+import {
+  resolveOptimisticQuestionAgent,
+  resolveQuestionMode,
+  resolveSelectedQuestionMode,
+  toggleAnswer,
+} from "../../webview-ui/src/components/chat/question-dock-utils"
 
 describe("toggleAnswer", () => {
   it("adds answer when not present", () => {
@@ -36,3 +41,90 @@ describe("toggleAnswer", () => {
     expect(result).toEqual(["a"])
   })
 })
+
+describe("resolveQuestionMode", () => {
+  it("returns mode for matching predefined option", () => {
+    const result = resolveQuestionMode(
+      [
+        { label: "Implement", description: "Switch to code", mode: "code" },
+        { label: "Stay", description: "Remain here" },
+      ],
+      "Implement",
+    )
+
+    expect(result).toBe("code")
+  })
+
+  it("returns undefined for unknown answer", () => {
+    const result = resolveQuestionMode([{ label: "Implement", description: "Switch", mode: "code" }], "Custom")
+    expect(result).toBeUndefined()
+  })
+
+  it("returns undefined when option has no mode", () => {
+    const result = resolveQuestionMode([{ label: "Stay", description: "Remain here" }], "Stay")
+    expect(result).toBeUndefined()
+  })
+})
+
+describe("resolveSelectedQuestionMode", () => {
+  it("returns the selected mode from predefined answers", () => {
+    const result = resolveSelectedQuestionMode(
+      [
+        [
+          { label: "Implement", description: "Switch to code", mode: "code" },
+          { label: "Stay", description: "Remain here" },
+        ],
+      ].map((options) => ({ options })),
+      [["Implement"]],
+    )
+
+    expect(result).toBe("code")
+  })
+
+  it("ignores custom answers that replace a mode option", () => {
+    const result = resolveSelectedQuestionMode(
+      [{ options: [{ label: "Implement", description: "Switch to code", mode: "code" }] }],
+      [["Implement custom flow"]],
+    )
+
+    expect(result).toBeUndefined()
+  })
+
+  it("keeps mode answers from other questions", () => {
+    const result = resolveSelectedQuestionMode(
+      [
+        { options: [{ label: "Implement", description: "Switch to code", mode: "code" }] },
+        { options: [{ label: "Stay", description: "Remain here" }] },
+      ],
+      [["Implement"], ["Stay"]],
+    )
+
+    expect(result).toBe("code")
+  })
+})
+
+describe("resolveOptimisticQuestionAgent", () => {
+  it("stores the previous agent when applying an optimistic mode", () => {
+    const result = resolveOptimisticQuestionAgent(undefined, "ask", "code")
+
+    expect(result).toEqual({ base: "ask", agent: "code" })
+  })
+
+  it("reverts to the stored previous agent when the mode is cleared", () => {
+    const result = resolveOptimisticQuestionAgent("ask", "code", undefined)
+
+    expect(result).toEqual({ base: undefined, agent: "ask" })
+  })
+
+  it("avoids switching when the selected mode already matches the current agent", () => {
+    const result = resolveOptimisticQuestionAgent(undefined, "code", "code")
+
+    expect(result).toEqual({ base: undefined, agent: undefined })
+  })
+
+  it("keeps the original base agent while changing between mode answers", () => {
+    const result = resolveOptimisticQuestionAgent("ask", "code", "architect")
+
+    expect(result).toEqual({ base: "ask", agent: "architect" })
+  })
+})

packages/kilo-vscode/tests/unit/session-agent.test.ts

@@ -0,0 +1,60 @@
+import { describe, it, expect } from "bun:test"
+import { resolveSessionAgent } from "../../webview-ui/src/context/session-agent"
+import type { Message } from "../../webview-ui/src/types/messages"
+
+function makeMessage(overrides: Partial<Message> = {}): Message {
+  return {
+    id: "msg-1",
+    sessionID: "sess-1",
+    role: "user",
+    createdAt: new Date(0).toISOString(),
+    ...overrides,
+  }
+}
+
+describe("resolveSessionAgent", () => {
+  it("returns the latest valid user agent", () => {
+    const result = resolveSessionAgent(
+      [
+        makeMessage({ id: "1", agent: "plan" }),
+        makeMessage({ id: "2", role: "assistant", agent: "ask" }),
+        makeMessage({ id: "3", agent: "code" }),
+      ],
+      new Set(["plan", "code", "ask"]),
+    )
+
+    expect(result).toBe("code")
+  })
+
+  it("ignores assistant messages", () => {
+    const result = resolveSessionAgent(
+      [makeMessage({ role: "assistant", agent: "code" }), makeMessage({ agent: "plan" })],
+      new Set(["plan", "code"]),
+    )
+
+    expect(result).toBe("plan")
+  })
+
+  it("ignores unknown agent names", () => {
+    const result = resolveSessionAgent(
+      [makeMessage({ agent: "missing" }), makeMessage({ agent: "code" })],
+      new Set(["code"]),
+    )
+
+    expect(result).toBe("code")
+  })
+
+  it("ignores empty agent values", () => {
+    const result = resolveSessionAgent([makeMessage({ agent: "  " })], new Set(["code"]))
+    expect(result).toBeUndefined()
+  })
+
+  it("returns undefined when no valid user agent exists", () => {
+    const result = resolveSessionAgent(
+      [makeMessage({ role: "assistant", agent: "code" }), makeMessage({ agent: undefined })],
+      new Set(["code"]),
+    )
+
+    expect(result).toBeUndefined()
+  })
+})

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

@@ -4,14 +4,15 @@
  * Uses kilo-ui's DockPrompt component for proper surface styling.
  */
 
-import { Component, For, Show, createMemo, createEffect } from "solid-js"
+import { For, Show, createMemo, createEffect } from "solid-js"
+import type { Component } from "solid-js"
 import { createStore } from "solid-js/store"
 import { Button } from "@kilocode/kilo-ui/button"
 import { Icon } from "@kilocode/kilo-ui/icon"
 import { useSession } from "../../context/session"
 import { useLanguage } from "../../context/language"
 import type { QuestionRequest } from "../../types/messages"
-import { toggleAnswer } from "./question-dock-utils"
+import { resolveOptimisticQuestionAgent, resolveSelectedQuestionMode, toggleAnswer } from "./question-dock-utils"
 
 export const QuestionDock: Component<{ request: QuestionRequest }> = (props) => {
   const session = useSession()
@@ -24,17 +25,23 @@ export const QuestionDock: Component<{ request: QuestionRequest }> = (props) =>
     tab: 0,
     answers: [] as string[][],
     custom: [] as string[],
+    kinds: [] as Record<string, "option" | "custom">[],
     editing: false,
     sending: false,
     collapsed: false,
   })
 
   let root!: HTMLDivElement
+  let prevAgent: string | undefined
 
-  // Reset sending state when an error occurs for this question
+  // Reset sending state and roll back the optimistic agent change on error
   createEffect(() => {
     if (session.questionErrors().has(props.request.id)) {
       setStore("sending", false)
+      if (prevAgent !== undefined) {
+        session.selectAgent(prevAgent)
+        prevAgent = undefined
+      }
     }
   })
 
@@ -64,6 +71,8 @@ export const QuestionDock: Component<{ request: QuestionRequest }> = (props) =>
     setStore("sending", true)
     session.replyToQuestion(props.request.id, answers)
     focusPrompt()
+    // prevAgent is intentionally left set until either questionError (rollback)
+    // or the question is dismissed (success — the question unmounts, so no cleanup needed)
   }
 
   const reject = () => {
@@ -83,17 +92,33 @@ export const QuestionDock: Component<{ request: QuestionRequest }> = (props) =>
     setStore("editing", false)
   }
 
+  const syncAgent = (answers: string[][], kinds: Record<string, "option" | "custom">[] = store.kinds) => {
+    const mode = resolveSelectedQuestionMode(questions(), answers, kinds)
+    const next = resolveOptimisticQuestionAgent(prevAgent, session.selectedAgent(), mode)
+
+    prevAgent = next.base
+    if (!next.agent) return
+    if (next.agent === session.selectedAgent()) return
+    session.selectAgent(next.agent)
+  }
+
   const pick = (answer: string, custom = false) => {
     const answers = [...store.answers]
     answers[store.tab] = [answer]
     setStore("answers", answers)
 
+    const kinds = [...store.kinds]
+    kinds[store.tab] = { [answer]: custom ? "custom" : "option" }
+    setStore("kinds", kinds)
+
     if (custom) {
       const inputs = [...store.custom]
       inputs[store.tab] = answer
       setStore("custom", inputs)
     }
 
+    syncAgent(answers, kinds)
+
     if (!single() && !multi()) {
       setStore("tab", store.tab + 1)
     }
@@ -104,6 +129,13 @@ export const QuestionDock: Component<{ request: QuestionRequest }> = (props) =>
     const answers = [...store.answers]
     answers[store.tab] = next
     setStore("answers", answers)
+    const kinds = [...store.kinds]
+    const current = { ...(kinds[store.tab] ?? {}) }
+    if (next.includes(answer)) current[answer] = "option"
+    else delete current[answer]
+    kinds[store.tab] = current
+    setStore("kinds", kinds)
+    syncAgent(answers, kinds)
   }
 
   const selectTab = (index: number) => {
@@ -167,6 +199,12 @@ export const QuestionDock: Component<{ request: QuestionRequest }> = (props) =>
       const answers = [...store.answers]
       answers[store.tab] = next
       setStore("answers", answers)
+      const kinds = [...store.kinds]
+      const current = { ...(kinds[store.tab] ?? {}) }
+      current[value] = "custom"
+      kinds[store.tab] = current
+      setStore("kinds", kinds)
+      syncAgent(answers, kinds)
       setStore("editing", false)
       return
     }

packages/kilo-vscode/webview-ui/src/components/chat/question-dock-utils.ts

@@ -1,7 +1,53 @@
+import type { QuestionOption } from "../../types/messages"
+
 export function toggleAnswer(existing: string[], answer: string): string[] {
   const next = [...existing]
   const index = next.indexOf(answer)
   if (index === -1) next.push(answer)
   if (index !== -1) next.splice(index, 1)
   return next
 }
+
+export function resolveQuestionMode(options: QuestionOption[], answer: string): string | undefined {
+  return options.find((item) => item.label === answer)?.mode
+}
+
+export function resolveSelectedQuestionMode(
+  questions: Array<{ options?: QuestionOption[] }>,
+  answers: string[][],
+  kinds: Record<string, "option" | "custom">[] = [],
+): string | undefined {
+  let mode: string | undefined
+
+  for (const [i, list] of answers.entries()) {
+    const options = questions[i]?.options ?? []
+    for (const answer of list) {
+      if (kinds[i]?.[answer] === "custom") continue
+      const next = resolveQuestionMode(options, answer)
+      if (next) mode = next
+    }
+  }
+
+  return mode
+}
+
+export function resolveOptimisticQuestionAgent(base: string | undefined, current: string, mode: string | undefined) {
+  if (!mode) {
+    return {
+      base: undefined,
+      agent: base,
+    }
+  }
+
+  if (base === undefined && current === mode) {
+    return {
+      base: undefined,
+      agent: undefined,
+    }
+  }
+
+  return {
+    base: base ?? current,
+    agent: mode,
+  }
+}

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

@@ -0,0 +1,12 @@
+import type { Message } from "../types/messages"
+
+export function resolveSessionAgent(messages: Message[], names: Set<string>): string | undefined {
+  for (let i = messages.length - 1; i >= 0; i--) {
+    const msg = messages[i]
+    if (msg.role !== "user") continue
+    const name = msg.agent?.trim()
+    if (!name) continue
+    if (!names.has(name)) continue
+    return name
+  }
+}