fix(vscode): clear stale optimistic question mode

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

@@ -1,5 +1,10 @@
 import { describe, it, expect } from "bun:test"
-import { resolveQuestionMode, 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", () => {
@@ -60,3 +65,66 @@ describe("resolveQuestionMode", () => {
     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/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 { resolveQuestionMode, toggleAnswer } from "./question-dock-utils"
+import { resolveOptimisticQuestionAgent, resolveSelectedQuestionMode, toggleAnswer } from "./question-dock-utils"
 
 export const QuestionDock: Component<{ request: QuestionRequest }> = (props) => {
   const session = useSession()
@@ -24,6 +25,7 @@ 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,
@@ -90,22 +92,32 @@ 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 mode = resolveQuestionMode(options(), answer)
     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)
     }
 
-    if (mode && !custom) {
-      prevAgent = session.selectedAgent()
-      session.selectAgent(mode)
-    }
+    syncAgent(answers, kinds)
 
     if (!single() && !multi()) {
       setStore("tab", store.tab + 1)
@@ -117,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) => {
@@ -180,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

@@ -11,3 +11,43 @@ export function toggleAnswer(existing: string[], answer: string): string[] {
 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,
+  }
+}