fix(vscode): restore per-mode model memory to match CLI TUI behavior

packages/kilo-docs/pages/code-with-ai/agents/model-selection.md

@@ -26,7 +26,8 @@ While the specifics change constantly, some principles stay consistent:
 
 - Use the **model selector** in the chat prompt area to pick a model for the current session. You can also type `/models` to open the model picker.
 - Set per-agent defaults and a global default in the **Settings** panel (Models tab), or directly in the `kilo.jsonc` config file.
-- **Model precedence:** Session override → Per-agent config → Global config → Recent models → Kilo Auto (free).
+- **Model precedence:** Session override → Last picked per agent → Per-agent config → Global config → Kilo Auto (free).
+- The model selector remembers the last model you picked for each agent — switching agents restores your previous choice. A manual pick always beats config settings; use the **reset button** (visible when your active model differs from config) to go back to the config default.
 
 {% /tab %}
 {% tab label="CLI" %}

packages/kilo-docs/pages/customize/custom-modes.md

@@ -161,6 +161,8 @@ Pin a specific model using the `provider/model` format:
 model: anthropic/claude-sonnet-4-20250514
 ```
 
+The model selector also **remembers the last model you picked for each agent** across sessions. A config-pinned `model` acts as the default when no manual pick exists. To reset a pick and let the config take over, use the **reset button** in the model selector (visible when your active model differs from what the config specifies).
+
 ### `steps`
 
 Limits the number of agentic iterations (tool call rounds) before the agent is forced to respond with text only. Useful for preventing runaway agents:
@@ -379,6 +381,8 @@ Pin a specific model using the `provider/model` format:
 model: anthropic/claude-sonnet-4-20250514
 ```
 
+The TUI also **remembers the last model you picked for each agent** across sessions. A config-pinned `model` acts as the default when no manual pick exists. To reset a pick and let the config take over, use the model picker (`Ctrl+X m`) and select a different model, or remove the saved pick from `~/.local/state/kilo/model.json`.
+
 ### `steps`
 
 Limits the number of agentic iterations (tool call rounds) before the agent is forced to respond with text only. Useful for preventing runaway agents:

packages/kilo-vscode/eslint.config.mjs

@@ -38,7 +38,7 @@ export default [
   // New code must stay ≤ 20. Do not raise these caps; refactor instead.
   {
     files: ["src/KiloProvider.ts"],
-    rules: { complexity: ["error", 140], "max-lines": ["error", 3350] },
+    rules: { complexity: ["error", 140], "max-lines": ["error", 3353] },
   },
   {
     files: ["webview-ui/agent-manager/AgentManagerApp.tsx"],

packages/kilo-vscode/src/KiloProvider.ts

@@ -55,6 +55,7 @@ import { fetchMessagePage, MESSAGE_PAGE_LIMIT } from "./kilo-provider/message-pa
 import { childID } from "./kilo-provider/task-session"
 import { handleNetworkEvent, clearNetworkWaits } from "./kilo-provider/network"
 import { abortSession, parseQueued } from "./kilo-provider/abort"
+import * as ModelState from "./kilo-provider/model-state"
 import { handleForkSession } from "./kilo-provider/fork-session"
 import { retryable, backoff, MAX_RETRIES } from "./util/retry"
 import { hasGit } from "./kilo-provider/git-status"
@@ -569,6 +570,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
       }
 
       await routeSuggestionWebviewMessage(this.questionCtx, message)
+      if (await ModelState.handleMessage(message.type, message, this.client, (msg) => this.postMessage(msg))) return
       switch (message.type) {
         case "webviewReady":
           console.log("[Kilo New] KiloProvider: ✅ webviewReady received")
@@ -2836,6 +2838,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
     this.sendBrowserSettings()
     this.sendNotificationSettings()
     this.sendTimelineSetting()
+    await ModelState.reset(this.client, (msg) => this.postMessage(msg))
 
     // Re-send globalState items to the webview
     this.postMessage({ type: "variantsLoaded", variants: {} })

packages/kilo-vscode/src/kilo-provider/model-state.ts

@@ -0,0 +1,94 @@
+/**
+ * Per-mode model selection persistence via the CLI's model.json.
+ *
+ * Reads/writes ~/.local/state/kilo/model.json (same file the CLI TUI uses)
+ * so per-mode model choices are shared between CLI and extension.
+ */
+
+import * as fs from "fs"
+import * as path from "path"
+import type { KiloClient } from "@kilocode/sdk/v2/client"
+import { validateModelSelections } from "../provider-actions"
+
+type PostMessage = (msg: unknown) => void
+
+let cached: string | undefined
+let queue: Promise<void> = Promise.resolve()
+
+async function resolve(client: KiloClient | null): Promise<string | undefined> {
+  if (cached) return cached
+  try {
+    const resp = await client?.path.get()
+    if (!resp?.data?.state) return undefined
+    cached = path.join(resp.data.state, "model.json")
+    return cached
+  } catch {
+    return undefined
+  }
+}
+
+async function read(client: KiloClient | null): Promise<Record<string, unknown>> {
+  const p = await resolve(client)
+  if (!p) return {}
+  try {
+    const raw = await fs.promises.readFile(p, "utf-8")
+    const parsed = JSON.parse(raw)
+    return typeof parsed === "object" && parsed !== null && !Array.isArray(parsed)
+      ? (parsed as Record<string, unknown>)
+      : {}
+  } catch {
+    return {}
+  }
+}
+
+function write(client: KiloClient | null, key: string, value: unknown): Promise<void> {
+  const op = queue.then(async () => {
+    const p = await resolve(client)
+    if (!p) return
+    const existing = await read(client)
+    existing[key] = value
+    await fs.promises.writeFile(p, JSON.stringify(existing, null, 2))
+  })
+  queue = op.catch(() => {})
+  return op
+}
+
+/**
+ * Handle a model-state webview message. Returns true if handled.
+ */
+export async function handleMessage(
+  type: string,
+  message: Record<string, unknown>,
+  client: KiloClient | null,
+  post: PostMessage,
+): Promise<boolean> {
+  if (type === "persistModelSelection") {
+    const data = await read(client)
+    const model = validateModelSelections(data.model)
+    model[message.agent as string] = {
+      providerID: message.providerID as string,
+      modelID: message.modelID as string,
+    }
+    await write(client, "model", model)
+    return true
+  }
+  if (type === "clearModelSelection") {
+    const data = await read(client)
+    const model = validateModelSelections(data.model)
+    delete model[message.agent as string]
+    await write(client, "model", model)
+    return true
+  }
+  if (type === "requestModelSelections") {
+    const data = await read(client)
+    const selections = validateModelSelections(data.model)
+    post({ type: "modelSelectionsLoaded", selections })
+    return true
+  }
+  return false
+}
+
+export async function reset(client: KiloClient | null, post: PostMessage): Promise<void> {
+  await write(client, "model", {})
+  post({ type: "modelSelectionsLoaded", selections: {} })
+}

packages/kilo-vscode/src/provider-actions.ts

@@ -93,6 +93,18 @@ export function validateFavorites(raw: unknown): Array<{ providerID: string; mod
   return raw.filter(isModelSelection).map((r) => ({ providerID: r.providerID, modelID: r.modelID }))
 }
 
+/** Validate and sanitize per-mode model selections from untrusted sources. */
+export function validateModelSelections(raw: unknown): Record<string, { providerID: string; modelID: string }> {
+  if (!raw || typeof raw !== "object" || Array.isArray(raw)) return {}
+  const result: Record<string, { providerID: string; modelID: string }> = {}
+  for (const [key, val] of Object.entries(raw as Record<string, unknown>)) {
+    if (isModelSelection(val)) {
+      result[key] = { providerID: val.providerID, modelID: val.modelID }
+    }
+  }
+  return result
+}
+
 export function computeDefaultSelection(
   cachedConfig: { config?: { model?: string } } | null,
   vscodePID: string,

packages/kilo-vscode/tests/unit/provider-actions-validate.test.ts

@@ -0,0 +1,75 @@
+import { describe, expect, it } from "bun:test"
+import { validateModelSelections, validateRecents, validateFavorites } from "../../src/provider-actions"
+
+describe("validateModelSelections", () => {
+  it("returns empty object for null", () => {
+    expect(validateModelSelections(null)).toEqual({})
+  })
+
+  it("returns empty object for undefined", () => {
+    expect(validateModelSelections(undefined)).toEqual({})
+  })
+
+  it("returns empty object for array", () => {
+    expect(validateModelSelections([{ providerID: "a", modelID: "b" }])).toEqual({})
+  })
+
+  it("returns empty object for non-object primitives", () => {
+    expect(validateModelSelections("string")).toEqual({})
+    expect(validateModelSelections(42)).toEqual({})
+    expect(validateModelSelections(true)).toEqual({})
+  })
+
+  it("passes through valid selections", () => {
+    const input = {
+      code: { providerID: "anthropic", modelID: "claude-sonnet-4" },
+      ask: { providerID: "openai", modelID: "gpt-4.1" },
+    }
+    expect(validateModelSelections(input)).toEqual(input)
+  })
+
+  it("filters out entries with non-string providerID", () => {
+    const input = {
+      code: { providerID: "anthropic", modelID: "claude-sonnet-4" },
+      broken: { providerID: 42, modelID: "model" },
+    }
+    expect(validateModelSelections(input)).toEqual({
+      code: { providerID: "anthropic", modelID: "claude-sonnet-4" },
+    })
+  })
+
+  it("filters out entries with non-string modelID", () => {
+    const input = {
+      code: { providerID: "anthropic", modelID: "claude-sonnet-4" },
+      broken: { providerID: "openai", modelID: null },
+    }
+    expect(validateModelSelections(input)).toEqual({
+      code: { providerID: "anthropic", modelID: "claude-sonnet-4" },
+    })
+  })
+
+  it("filters out null and non-object entries", () => {
+    const input = {
+      code: { providerID: "anthropic", modelID: "claude-sonnet-4" },
+      empty: null,
+      str: "not-an-object",
+      num: 123,
+    }
+    expect(validateModelSelections(input)).toEqual({
+      code: { providerID: "anthropic", modelID: "claude-sonnet-4" },
+    })
+  })
+
+  it("strips extra properties from valid entries", () => {
+    const input = {
+      code: { providerID: "anthropic", modelID: "claude-sonnet-4", extra: true, nested: { x: 1 } },
+    }
+    const result = validateModelSelections(input)
+    expect(result).toEqual({ code: { providerID: "anthropic", modelID: "claude-sonnet-4" } })
+    expect(Object.keys(result.code!)).toEqual(["providerID", "modelID"])
+  })
+
+  it("returns empty object for empty input object", () => {
+    expect(validateModelSelections({})).toEqual({})
+  })
+})

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

@@ -47,20 +47,21 @@ const claude: ModelSelection = { providerID: "anthropic", modelID: "claude-sonne
 const gpt: ModelSelection = { providerID: "openai", modelID: "gpt-4.1" }
 
 describe("per-session model selection", () => {
-  it("selecting a model in session A does not affect session B", () => {
+  it("selecting a model in session A writes per-mode globally", () => {
     const store = emptyStore()
     const e = env()
 
     // User picks claude in session A
     const after = applyModel(store, "code", claude, "session-a")
     const updated: ModelStore = { ...store, ...after }
 
-    // Session A should see claude
+    // Session A should see claude (via session override)
     expect(getSessionModel(updated, e, "session-a", "code")).toEqual(claude)
 
-    // Session B should NOT see claude — it should fall back to the default
+    // Session B (no override) inherits the per-mode global selection.
+    // This matches CLI behavior: per-mode model is global, not per-session.
     const sessionB = getSessionModel(updated, e, "session-b", "code")
-    expect(sessionB).toEqual(KILO_AUTO)
+    expect(sessionB).toEqual(claude)
   })
 
   it("each session preserves its own model independently", () => {
@@ -147,3 +148,74 @@ describe("per-session model selection", () => {
     expect(getSessionModel(store, e, "session-a", "code")).toEqual(gpt)
   })
 })
+
+describe("per-mode model memory", () => {
+  it("applyModel in a session writes to both sessionOverrides and modelSelections", () => {
+    const store = emptyStore()
+    const result = applyModel(store, "code", claude, "session-a")
+
+    expect(result.sessionOverrides["session-a"]).toEqual(claude)
+    expect(result.modelSelections["code"]).toEqual(claude)
+  })
+
+  it("switching modes restores per-mode model after session override is cleared", () => {
+    let store = emptyStore()
+    const e = env()
+
+    // User picks claude for "code" mode in session A
+    const result = applyModel(store, "code", claude, "session-a")
+    store = { ...store, ...result }
+
+    // Simulate mode switch: clear session override (like selectAgent does)
+    const cleared = { ...store, sessionOverrides: {} }
+
+    // The global modelSelections["code"] still has claude
+    expect(getSelected(cleared, e, "session-a", "code")).toEqual(claude)
+  })
+
+  it("different modes remember their own model independently", () => {
+    let store = emptyStore()
+    const e = env()
+
+    // User picks claude for "code" in session A
+    let result = applyModel(store, "code", claude, "session-a")
+    store = { ...store, ...result }
+
+    // User switches to "ask" mode and picks gpt
+    result = applyModel(store, "ask", gpt, "session-a")
+    store = { ...store, ...result }
+
+    // Clear session overrides (simulating mode switch)
+    const cleared: ModelStore = { ...store, sessionOverrides: {} }
+
+    // Each mode should have its own saved model
+    expect(getSelected(cleared, e, undefined, "code")).toEqual(claude)
+    expect(getSelected(cleared, e, undefined, "ask")).toEqual(gpt)
+  })
+
+  it("per-session override still takes priority over global modelSelections", () => {
+    let store = emptyStore()
+    const e = env()
+
+    // User picks claude globally for "code"
+    let result = applyModel(store, "code", claude, undefined)
+    store = { ...store, ...result }
+
+    // Session A overrides with gpt
+    result = applyModel(store, "code", gpt, "session-a")
+    store = { ...store, ...result }
+
+    // Session A sees gpt (its override), not the global claude
+    expect(getSelected(store, e, "session-a", "code")).toEqual(gpt)
+    // Global modelSelections was updated to gpt (last write wins)
+    expect(store.modelSelections["code"]).toEqual(gpt)
+  })
+
+  it("applyModel without session only writes to modelSelections, not sessionOverrides", () => {
+    const store = emptyStore()
+    const result = applyModel(store, "code", claude, undefined)
+
+    expect(result.modelSelections["code"]).toEqual(claude)
+    expect(Object.keys(result.sessionOverrides)).toHaveLength(0)
+  })
+})

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

@@ -86,27 +86,22 @@ export interface ApplyResult {
 /**
  * Apply a user-initiated model selection.
  *
- * When a session is active, writes ONLY to the per-session override so other
- * sessions are not affected. When no session is active (sidebar), writes to
- * the global modelSelections map.
+ * Always writes to the global modelSelections map so switching modes
+ * restores the last-used model (mirrors CLI TUI's model.json behavior).
+ * When a session is active, also writes to the per-session override so
+ * the active session uses the chosen model immediately.
  */
 export function applyModel(
   store: ModelStore,
   agentName: string,
   selection: ModelSelection,
   sessionID: string | undefined,
 ): ApplyResult {
-  const modelSelections = { ...store.modelSelections }
+  const modelSelections = { ...store.modelSelections, [agentName]: selection }
   const sessionOverrides = { ...store.sessionOverrides }
 
   if (sessionID) {
-    // Per-session only — do NOT mutate the global map.  Writing globally
-    // here would cause every other session (that hasn't set its own
-    // override) to inherit this session's model.
     sessionOverrides[sessionID] = selection
-  } else {
-    // No active session (sidebar) — write globally
-    modelSelections[agentName] = selection
   }
 
   return { modelSelections, sessionOverrides }