diff --git a/.changeset/preserve-explicit-model-choice.md b/.changeset/preserve-explicit-model-choice.md new file mode 100644 index 00000000000..2dde8c881ae --- /dev/null +++ b/.changeset/preserve-explicit-model-choice.md @@ -0,0 +1,5 @@ +--- +"kilo-code": patch +--- + +Keep model and reasoning choices when switching modes, and remember explicit new-session choices as the default without changing open sessions. diff --git a/packages/kilo-vscode/src/agent-manager/multi-version.ts b/packages/kilo-vscode/src/agent-manager/multi-version.ts index 5d0704f5522..c38a1ae57cf 100644 --- a/packages/kilo-vscode/src/agent-manager/multi-version.ts +++ b/packages/kilo-vscode/src/agent-manager/multi-version.ts @@ -96,12 +96,12 @@ export function buildInitialMessages( worktreeId: entry.worktreeId, providerID: pid, modelID: mid, + agent, + // A per-allocation effort pick wins even when preparing an empty session. + variant: model?.variant ?? variant, } if (prompt) { msg.text = prompt - msg.agent = agent - // A per-allocation effort pick wins over the dialog-level variant. - msg.variant = model?.variant ?? variant msg.files = files } return msg diff --git a/packages/kilo-vscode/src/kilo-provider/model-state.ts b/packages/kilo-vscode/src/kilo-provider/model-state.ts index 655370aa1bb..815d85a2e23 100644 --- a/packages/kilo-vscode/src/kilo-provider/model-state.ts +++ b/packages/kilo-vscode/src/kilo-provider/model-state.ts @@ -1,12 +1,13 @@ /** - * Per-mode model selection persistence via the CLI's model.json. + * Model selection persistence in the CLI state directory. * - * Reads/writes ~/.local/state/kilo/model.json (same file the CLI TUI uses) - * so per-mode model choices are shared between CLI and extension. + * Per-mode choices use the shared model.json. The explicit preferred combo uses + * vscode-model.json so CLI/TUI writers cannot discard extension-only state. */ import * as fs from "fs" import * as path from "path" +import { randomUUID } from "crypto" import type { KiloClient } from "@kilocode/sdk/v2/client" import { validateModelSelections } from "../provider-actions" @@ -17,21 +18,20 @@ let queue: Promise = Promise.resolve() async function resolve(client: KiloClient | null): Promise { if (cached) return cached + if (!client) return undefined try { - const resp = await client?.path.get() - if (!resp?.data?.state) return undefined - cached = path.join(resp.data.state, "model.json") + const resp = await client.path.get() + if (typeof resp?.data?.state !== "string" || !resp.data.state) return undefined + cached = resp.data.state return cached } catch { return undefined } } -async function read(client: KiloClient | null): Promise> { - const p = await resolve(client) - if (!p) return {} +async function read(file: string): Promise> { try { - const raw = await fs.promises.readFile(p, "utf-8") + const raw = await fs.promises.readFile(file, "utf-8") const parsed = JSON.parse(raw) return typeof parsed === "object" && parsed !== null && !Array.isArray(parsed) ? (parsed as Record) @@ -41,15 +41,51 @@ async function read(client: KiloClient | null): Promise> } } -function write(client: KiloClient | null, key: string, value: unknown): Promise { +function selection(raw: unknown) { + if (!raw || typeof raw !== "object" || Array.isArray(raw)) return undefined + const value = raw as Record + if ( + typeof value.providerID !== "string" || + !value.providerID || + typeof value.modelID !== "string" || + !value.modelID || + (value.variant !== undefined && typeof value.variant !== "string") + ) { + return undefined + } + return { + providerID: value.providerID, + modelID: value.modelID, + ...(typeof value.variant === "string" ? { variant: value.variant } : {}), + } +} + +function write( + client: KiloClient | null, + update: (data: Record) => void, + preferred?: ReturnType, +): Promise { 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)) + const dir = await resolve(client) + if (!dir) return + const file = path.join(dir, "model.json") + const existing = await read(file) + update(existing) + await fs.promises.writeFile(file, JSON.stringify(existing, null, 2)) + const target = path.join(dir, "vscode-model.json") + const temp = `${target}.${randomUUID()}.tmp` + try { + await fs.promises.writeFile(temp, JSON.stringify({ preferred }, null, 2)) + await fs.promises.rename(temp, target) + } finally { + await fs.promises.rm(temp, { force: true }).catch((err) => { + console.error("[Kilo New] Failed to remove temporary model preferences:", err) + }) + } + }) + queue = op.catch((err) => { + console.error("[Kilo New] Failed to persist model selections:", err) }) - queue = op.catch(() => {}) return op } @@ -63,25 +99,39 @@ export async function handleMessage( post: PostMessage, ): Promise { 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) + const preferred = selection(message) + const agent = message.agent + if (!preferred || typeof agent !== "string" || !agent) return true + await write( + client, + (data) => { + const model = validateModelSelections(data.model) + model[agent] = { providerID: preferred.providerID, modelID: preferred.modelID } + data.model = model + }, + { ...preferred, variant: preferred.variant ?? "" }, + ) return true } if (type === "requestModelSelections") { - const data = await read(client) + await queue + const dir = await resolve(client) + if (!dir) return true + const [data, prefs] = await Promise.all([ + read(path.join(dir, "model.json")), + read(path.join(dir, "vscode-model.json")), + ]) const selections = validateModelSelections(data.model) - post({ type: "modelSelectionsLoaded", selections }) + const preferred = selection(prefs.preferred) + post({ type: "modelSelectionsLoaded", selections, ...(preferred ? { preferred } : {}) }) return true } return false } export async function reset(client: KiloClient | null, post: PostMessage): Promise { - await write(client, "model", {}) + await write(client, (data) => { + data.model = {} + }) post({ type: "modelSelectionsLoaded", selections: {} }) } diff --git a/packages/kilo-vscode/tests/fixtures/session-provider-activity.tsx b/packages/kilo-vscode/tests/fixtures/session-provider-activity.tsx index c5c4db95d78..42ab8addf0c 100644 --- a/packages/kilo-vscode/tests/fixtures/session-provider-activity.tsx +++ b/packages/kilo-vscode/tests/fixtures/session-provider-activity.tsx @@ -350,7 +350,13 @@ try { const writes = () => sent.filter((item) => item.type === "persistModelSelection" || item.type === "persistRecents") const requests = () => sent.filter((item) => ["sendMessage", "sendCommand", "importAndSend", "compact"].includes(item.type)) - const catalog = async (organizationId: string | null, ids: string[], model?: string, ready = true) => { + const catalog = async ( + organizationId: string | null, + ids: string[], + model?: string, + ready = true, + variants = ["low", "high"], + ) => { await emit({ type: "providersLoaded", organizationId, @@ -359,7 +365,12 @@ try { kilo: { id: "kilo", name: "Kilo", - models: Object.fromEntries(ids.map((id) => [id, { id, name: id, variants: { low: {}, high: {} } }])), + models: Object.fromEntries( + ids.map((id) => [ + id, + { id, name: id, variants: Object.fromEntries(variants.map((variant) => [variant, {}])) }, + ]), + ), }, openai: { id: "openai", name: "OpenAI", models: { external: { id: "external", name: "External" } } }, }, @@ -656,13 +667,43 @@ try { choice(value.selected(), auto) setSettings({ agent: { code: { model: "kilo/z-first" } } }) await settle() - choice(value.modelForAgent("code"), first) + choice(value.modelForAgent("code"), auto) choice(value.selected(), auto) setSettings({}) await emit({ type: "modelSelectionsLoaded", selections: {} }) choice(value.selected(), recommended) assert.deepEqual(writes(), remembered) + // Mode changes preserve both explicit and inherited session choices, including Default. + await catalog("org-a", [personal.modelID, first.modelID, recommended.modelID], recommended.modelID) + setSettings({ agent: { ask: { model: "kilo/z-first", variant: "low" } } }) + for (const effort of ["high", ""]) { + for (const scope of [`selection-${effort}`, `pending:${effort}`, `sidebar-pending:${effort}`]) { + value.setSessionAgent(scope, "code") + value.setSessionModel(scope, personal.providerID, personal.modelID) + value.selectVariant(effort, scope) + const before = sent.length + for (const agent of ["ask", "code", "ask"]) { + value.selectAgent(agent, scope) + await settle() + assert.deepEqual(value.submission(scope), { model: personal, variant: effort, agent }) + } + assert.equal(sent.length, before, "Switching mode must not persist a different preference") + await emit({ type: "providersLoading" }) + value.selectAgent("code", scope) + await catalog("org-a", [personal.modelID, first.modelID, recommended.modelID], recommended.modelID) + assert.deepEqual(value.submission(scope), { model: personal, variant: effort, agent: "code" }) + } + } + setSettings({ + agent: { code: { model: "kilo/personal", variant: "high" }, ask: { model: "kilo/z-first", variant: "low" } }, + }) + value.setSessionAgent("inherited-mode", "code") + const inherited = value.submission("inherited-mode") + value.selectAgent("ask", "inherited-mode") + assert.deepEqual(value.submission("inherited-mode"), { ...inherited, agent: "ask" }) + setSettings({}) + const snapshot = (scope?: string) => JSON.stringify({ session: value.currentSessionID(), @@ -743,6 +784,7 @@ try { value.setSessionAgent(scope, "code") value.setSessionModel(scope, personal.providerID, personal.modelID) await settle() + const preserved = value.submission(scope) assert.equal( value.sendCommand( "review-test", @@ -761,10 +803,10 @@ try { assert(request?.type === "sendCommand") assert.equal(request.sessionID, scope) assert.equal(request.agent, "ask") - assert.equal(request.modelID, configured ? first.modelID : recommended.modelID) - assert.equal(request.variant, configured ? "high" : "low") + assert.equal(request.modelID, personal.modelID) + assert.equal(request.variant, preserved.variant) assert.equal(value.selectedAgent(scope), "ask") - choice(value.selected(scope), configured ? first : recommended) + choice(value.selected(scope), personal) } setSettings({}) await catalog(null, [auto.modelID, personal.modelID, first.modelID, recommended.modelID]) @@ -777,6 +819,7 @@ try { setSettings({ agent: { ask: { model: "kilo/a-recommended", variant: "high" } } }) value.setCurrentSessionID("ses_command-cached") await settle() + const inheritedCommand = value.submission("ses_command-cached") assert.equal( value.sendCommand( "review-test", @@ -793,15 +836,37 @@ try { ) const configured = requests().at(-1) assert(configured?.type === "sendCommand") - assert.equal(configured.modelID, recommended.modelID) - assert.equal(configured.variant, "high") - choice(value.selected(), recommended) + assert.equal(configured.modelID, inheritedCommand.model?.modelID) + assert.equal(configured.variant, inheritedCommand.variant) + assert.deepEqual(value.submission("ses_command-cached"), { ...inheritedCommand, agent: "ask" }) + + assert.equal( + value.sendCommand( + "review-test", + "default effort", + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + { variant: "" }, + ), + true, + ) + const commandDefault = sent.findLast((item) => item.type === "sendCommand") + assert(commandDefault) + assert.equal(commandDefault.variant, "") + assert.equal(value.currentVariant("ses_command-cached"), undefined) setSettings({}) await catalog("org-a", [personal.modelID, first.modelID, recommended.modelID], recommended.modelID) await emit({ type: "modelSelectionsLoaded", selections: { code: first, ask: recommended } }) value.setCurrentSessionID(undefined) value.selectAgent("ask") + const pendingModel = value.selected() + assert(pendingModel) + const pendingVariant = value.variantForAgent("ask", pendingModel) value.setCurrentSessionID("selection") assert.equal( value.sendCommand( @@ -822,12 +887,12 @@ try { assert(pending.draftID) assert.equal(pending.sessionID, undefined) assert.equal(pending.agent, "ask") - assert.equal(pending.modelID, recommended.modelID) + assert.equal(pending.modelID, pendingModel.modelID) assert.equal(pending.variant, "high") assert.equal(value.selectedAgent(pending.draftID), "ask") - choice(value.selected(pending.draftID), recommended) + choice(value.selected(pending.draftID), pendingModel) assert.equal(value.currentVariant(pending.draftID), "high") - assert.equal(value.variantForAgent("ask", recommended), "low") + assert.equal(value.variantForAgent("ask", pendingModel), pendingVariant) const persisted = sent.length assert.equal( value.sendCommand( @@ -855,7 +920,7 @@ try { choice(value.selected(accepted.draftID), personal) assert.equal(value.selectedAgent(accepted.draftID), "ask") assert.equal(value.currentVariant(accepted.draftID), "high") - choice(value.modelForAgent("ask"), recommended) + choice(value.modelForAgent("ask"), pendingModel) assert.equal( sent.slice(persisted).some((message) => message.type === "persistModelSelection"), false, @@ -900,7 +965,7 @@ try { value.clearCurrentSession() value.selectAgent("ask") assert.equal(value.selectedAgent(), "ask") - choice(value.selected(), recommended) + choice(value.selected(), pendingModel) const usage = JSON.stringify(value.modelUsageHistory()) const recent = JSON.stringify(value.recentModels()) const start = sent.length @@ -920,7 +985,7 @@ try { await emit({ type: "sessionCreated", session: info("ses_goal-draft"), draftID: command.draftID }) assert.equal(value.currentSessionID(), "ses_goal-draft") assert.equal(value.selectedAgent(), "ask") - choice(value.selected(), recommended) + choice(value.selected(), pendingModel) await emit({ type: "sessionCommandCompleted", messageID: command.messageID }) assert.equal(value.submitting(), false) assert.equal(JSON.stringify(value.modelUsageHistory()), usage) @@ -2011,6 +2076,117 @@ try { mode: "focus", }, ) + // An explicit next-session combo becomes the shared default, not a change to open sessions. + await catalog("org-a", [personal.modelID, first.modelID, recommended.modelID], recommended.modelID) + await emit({ type: "modelSelectionsLoaded", selections: {} }) + setSettings({ + agent: { code: { model: "kilo/z-first", variant: "low" }, ask: { model: "kilo/a-recommended", variant: "low" } }, + }) + value.setCurrentSessionID("preference-active") + value.setSessionAgent("preference-active", "ask") + const combo = value.submission("preference-active") + for (const scope of ["pending:preferred", "sidebar-pending:preferred"]) { + value.setSessionAgent(scope, "code") + value.selectModel(personal.providerID, personal.modelID, scope) + value.selectVariant("high", scope) + assert.deepEqual(value.preferredSelection(), { ...personal, variant: "high" }) + assert.deepEqual(value.submission("preference-active"), combo) + const saved = sent.findLast((item) => item.type === "persistModelSelection") + assert.deepEqual(saved, { type: "persistModelSelection", agent: "code", ...personal, variant: "high" }) + for (const agent of ["code", "ask"]) { + choice(value.modelForAgent(agent), personal) + assert.equal(value.variantForAgent(agent, personal), "high") + } + value.selectAgent("ask", scope) + assert.deepEqual(value.submission(scope), { model: personal, variant: "high", agent: "ask" }) + value.selectVariant(undefined, scope) + assert.deepEqual(value.preferredSelection(), { ...personal, variant: "" }) + assert.equal(value.variantForAgent("code", personal), undefined) + await emit({ type: "providersLoading" }) + value.selectAgent("code", scope) + await catalog("org-b", [first.modelID], first.modelID) + choice(value.selected(scope), first) + await catalog("org-a", [personal.modelID, first.modelID, recommended.modelID], recommended.modelID) + assert.deepEqual(value.submission(scope), { model: personal, variant: "", agent: "code" }) + assert.deepEqual(value.submission("preference-active"), combo) + } + value.rememberSelection("ask", first, "high") + assert.deepEqual(value.submission("preference-active"), combo) + assert.deepEqual(value.submission("pending:preferred"), { model: personal, variant: "", agent: "code" }) + choice(value.modelForAgent("code"), first) + assert.equal(value.variantForAgent("code", first), "high") + value.setSessionModel("comparison-only", personal.providerID, personal.modelID) + value.setSessionVariant("comparison-only", personal.providerID, personal.modelID, "low") + assert.deepEqual(value.preferredSelection(), { ...first, variant: "high" }) + + // Restore the persisted combo through the real host-message boundary. + await emit({ type: "modelSelectionsLoaded", selections: { ask: first }, preferred: { ...first, variant: "high" } }) + value.setCurrentSessionID(undefined) + value.selectAgent("code") + choice(value.selected(), first) + assert.equal(value.currentVariant(), "high") + value.selectModel(personal.providerID, personal.modelID) + choice(value.selected(), personal) + assert.equal(value.currentVariant(), "high") + assert.deepEqual(value.preferredSelection(), { ...personal, variant: "high" }) + assert.deepEqual(value.submission("preference-active"), combo) + + // Retention must not assign defaults to unopened or still-loading historical sessions. + await emit({ + type: "sessionsLoaded", + sessions: [...unwrap(value.sessions()), info("unopened-preference"), info("loading-preference")], + }) + value.setCurrentSessionID("loading-preference") + value.rememberSelection("code", first, "low") + for (const id of ["unopened-preference", "loading-preference"]) { + await emit({ + type: "messagesLoaded", + sessionID: id, + messages: [ + { + id: `${id}-message`, + sessionID: id, + role: "user", + agent: "ask", + model: { ...personal, variant: "high" }, + createdAt: info(id).createdAt, + }, + ], + }) + assert.deepEqual(value.submission(id), { model: personal, variant: "high", agent: "ask" }) + } + + // Pin raw inherited preferences, never a temporary catalog fallback or mapped effort. + value.rememberSelection("code", personal, "high") + for (const missing of ["model", "effort"]) { + const id = `pending:inherited-${missing}` + value.setSessionAgent(id, "code") + if (missing === "model") await catalog("org-b", [first.modelID], first.modelID) + if (missing === "effort") await catalog("org-a", [personal.modelID], personal.modelID, true, ["low"]) + value.selectAgent("ask", id) + await catalog("org-a", [personal.modelID, first.modelID, recommended.modelID], recommended.modelID) + assert.deepEqual(value.submission(id), { model: personal, variant: "high", agent: "ask" }) + } + + // A delayed startup response cannot erase a choice made while preferences load. + setSharing(false) + await settle() + setSharing(true) + await settle() + const fresh = peer.value + assert(fresh) + assert.equal(fresh.preferencesReady(), false) + const loadingPreferences = sent.filter((item) => item.type === "requestModelSelections").length + await emit({ type: "extensionDataReady" }) + assert.equal(sent.filter((item) => item.type === "requestModelSelections").length, loadingPreferences + 1) + fresh.selectModel(personal.providerID, personal.modelID) + fresh.selectVariant("high") + await emit({ type: "modelSelectionsLoaded", selections: { code: first }, preferred: { ...first, variant: "low" } }) + await emit({ type: "variantsLoaded", variants: { "agent/code/kilo/personal": "low" } }) + choice(fresh.selected(), personal) + assert.equal(fresh.currentVariant(), "high") + assert.deepEqual(fresh.preferredSelection(), { ...personal, variant: "high" }) + assert.equal(fresh.preferencesReady(), true) assert.deepEqual(failures, []) } finally { const before = state("background") diff --git a/packages/kilo-vscode/tests/unit/model-state.test.ts b/packages/kilo-vscode/tests/unit/model-state.test.ts new file mode 100644 index 00000000000..f238dd06d2c --- /dev/null +++ b/packages/kilo-vscode/tests/unit/model-state.test.ts @@ -0,0 +1,343 @@ +import { afterEach, describe, expect, it } from "bun:test" +import { link, mkdtemp, rm } from "node:fs/promises" +import { tmpdir } from "node:os" +import { join } from "node:path" +import type { KiloClient } from "@kilocode/sdk/v2/client" +import type { PersistModelSelectionRequest } from "../../webview-ui/src/types/messages/webview-messages" + +const dirs: string[] = [] +const model = { providerID: "kilo", modelID: "model-a" } +const other = { providerID: "other", modelID: "model-b" } +const unrelated = { + recent: [other], + favorite: [model], + variant: { "kilo/model-a": "low" }, + custom: { nested: [1, "keep", null] }, +} + +function load(): Promise { + // Reload module-level path caching without replacing the persistence implementation. + return import(`../../src/kilo-provider/model-state.ts?${crypto.randomUUID()}`) +} + +async function fixture(data?: unknown, prefs?: unknown) { + const dir = await mkdtemp(join(tmpdir(), "kilo-model-state-")) + dirs.push(dir) + const file = join(dir, "model.json") + const preference = join(dir, "vscode-model.json") + if (data !== undefined) await Bun.write(file, JSON.stringify(data)) + if (prefs !== undefined) await Bun.write(preference, JSON.stringify(prefs)) + const client = { path: { get: async () => ({ data: { state: dir } }) } } as unknown as KiloClient + const host = await load() + const messages: unknown[] = [] + const post = (message: unknown) => { + messages.push(message) + } + return { + dir, + get file() { + return Bun.file(file) + }, + get preference() { + return Bun.file(preference) + }, + client, + host, + post, + messages, + save: (selection: Omit) => + host.handleMessage("persistModelSelection", selection, client, post), + request: async () => { + expect(await host.handleMessage("requestModelSelections", {}, client, post)).toBe(true) + return messages.at(-1) + }, + } +} + +afterEach(async () => { + await Promise.all(dirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true }))) +}) + +describe("model-state", () => { + it("persists a preferred model and effort together while preserving other modes and unrelated fields", async () => { + const state = await fixture({ ...unrelated, model: { code: other, plan: other } }) + + expect(await state.save({ agent: "code", ...model, variant: "high" })).toBe(true) + + expect(await state.file.json()).toEqual({ + ...unrelated, + model: { code: model, plan: other }, + }) + expect(await state.preference.json()).toEqual({ preferred: { ...model, variant: "high" } }) + expect(await state.request()).toEqual({ + type: "modelSelectionsLoaded", + selections: { code: model, plan: other }, + preferred: { ...model, variant: "high" }, + }) + }) + + it.each([undefined, ""])("saves an explicit default effort for variant %j", async (variant) => { + const state = await fixture({ model: { plan: other } }, { preferred: { ...other, variant: "high" } }) + + await state.save({ agent: "code", ...model, ...(variant === undefined ? {} : { variant }) }) + + expect(await state.file.json()).toEqual({ + model: { plan: other, code: model }, + }) + expect(await state.preference.json()).toEqual({ preferred: { ...model, variant: "" } }) + }) + + it("loads the saved combo after the host module reloads", async () => { + const state = await fixture() + await state.save({ agent: "code", ...model, variant: "high" }) + const host = await load() + + expect(await host.handleMessage("requestModelSelections", {}, state.client, state.post)).toBe(true) + + expect(state.messages).toEqual([ + { type: "modelSelectionsLoaded", selections: { code: model }, preferred: { ...model, variant: "high" } }, + ]) + }) + + it("atomically replaces the complete preferred combo without modifying the previous file", async () => { + const prefs = { preferred: { ...other, variant: "low" } } + const state = await fixture({ model: { plan: other } }, prefs) + const previous = join(state.dir, "previous.json") + await link(join(state.dir, "vscode-model.json"), previous) + + await state.save({ agent: "code", ...model, variant: "high" }) + + expect(await Bun.file(previous).json()).toEqual(prefs) + expect(await state.preference.json()).toEqual({ preferred: { ...model, variant: "high" } }) + }) + + it("loads legacy model memory without inventing a preferred combo", async () => { + const state = await fixture({ model: { code: model, bad: { providerID: 42, modelID: "invalid" } } }) + + expect(await state.request()).toEqual({ type: "modelSelectionsLoaded", selections: { code: model } }) + }) + + it.each([undefined, "", "high"])( + "validates and sanitizes a loaded preferred combo with variant %j", + async (variant) => { + const preferred = { ...model, ...(variant === undefined ? {} : { variant }) } + const data = { model: { plan: other } } + const prefs = { preferred: { ...preferred, extra: "discard" } } + const state = await fixture(data, prefs) + + expect(await state.request()).toEqual({ type: "modelSelectionsLoaded", selections: { plan: other }, preferred }) + expect(await state.file.json()).toEqual(data) + expect(await state.preference.json()).toEqual(prefs) + }, + ) + + it.each( + [ + null, + true, + 42, + "invalid", + [], + {}, + { providerID: "kilo" }, + { modelID: "model-a" }, + { ...model, providerID: 42 }, + { ...model, providerID: "" }, + { ...model, modelID: null }, + { ...model, modelID: "" }, + { ...model, variant: null }, + { ...model, variant: 42 }, + { ...model, variant: false }, + { ...model, variant: {} }, + { ...model, variant: [] }, + ].map((preferred) => [preferred] as const), + )("omits invalid preferred data %j without discarding model memory", async (preferred) => { + const state = await fixture({ model: { code: model } }, { preferred }) + + expect(await state.request()).toEqual({ type: "modelSelectionsLoaded", selections: { code: model } }) + }) + + it.each(["{", "null", "[]", '"invalid"'])("recovers from invalid model.json contents %s", async (raw) => { + const preferred = { ...model, variant: "high" } + const state = await fixture(undefined, { preferred }) + await Bun.write(state.file, raw) + + expect(await state.request()).toEqual({ type: "modelSelectionsLoaded", selections: {}, preferred }) + await state.save({ agent: "code", ...model, variant: "high" }) + expect(await state.file.json()).toEqual({ model: { code: model } }) + expect(await state.preference.json()).toEqual({ preferred }) + }) + + it.each(["{", "null", "[]", '"invalid"'])("ignores invalid vscode-model.json contents %s", async (raw) => { + const state = await fixture({ model: { code: model } }) + await Bun.write(state.preference, raw) + + expect(await state.request()).toEqual({ type: "modelSelectionsLoaded", selections: { code: model } }) + await state.save({ agent: "code", ...model, variant: "high" }) + expect(await state.preference.json()).toEqual({ preferred: { ...model, variant: "high" } }) + }) + + it.each([ + { agent: "", ...model }, + { agent: 42, ...model }, + { agent: "code", modelID: "model-a" }, + { agent: "code", ...model, modelID: null }, + { agent: "code", ...model, variant: null }, + { agent: "code", ...model, variant: 42 }, + ])("ignores invalid explicit selections %j without overwriting saved preferences", async (message) => { + const data = { ...unrelated, model: { plan: other } } + const prefs = { preferred: { ...other, variant: "high" } } + const state = await fixture(data, prefs) + + expect(await state.host.handleMessage("persistModelSelection", message, state.client, state.post)).toBe(true) + expect(await state.file.json()).toEqual(data) + expect(await state.preference.json()).toEqual(prefs) + }) + + it("returns empty memory for a missing file and creates it on an explicit save", async () => { + const state = await fixture() + + expect(await state.request()).toEqual({ type: "modelSelectionsLoaded", selections: {} }) + expect(await state.file.exists()).toBe(false) + expect(await state.preference.exists()).toBe(false) + await state.save({ agent: "code", ...model }) + expect(await state.file.json()).toEqual({ model: { code: model } }) + expect(await state.preference.json()).toEqual({ preferred: { ...model, variant: "" } }) + }) + + it("clears model memory and preferred together while retaining unrelated fields", async () => { + const state = await fixture({ ...unrelated, model: { code: model } }, { preferred: { ...model, variant: "high" } }) + + await state.host.reset(state.client, state.post) + + expect(await state.file.json()).toEqual({ ...unrelated, model: {} }) + expect(await state.preference.json()).toEqual({}) + expect(state.messages).toEqual([{ type: "modelSelectionsLoaded", selections: {} }]) + const host = await load() + await host.handleMessage("requestModelSelections", {}, state.client, state.post) + expect(state.messages.at(-1)).toEqual({ type: "modelSelectionsLoaded", selections: {} }) + }) + + it("preserves every model entry during concurrent saves and keeps the latest preferred combo", async () => { + const state = await fixture({ ...unrelated, model: { existing: other } }) + const choices = Array.from({ length: 12 }, (_, index) => ({ + agent: `mode-${index}`, + providerID: "kilo", + modelID: `model-${index}`, + variant: `effort-${index}`, + })) + + await Promise.all([ + ...choices.map((choice) => state.save(choice)), + state.save({ agent: "mode-0", ...other, variant: "" }), + ]) + + expect(await state.file.json()).toEqual({ + ...unrelated, + model: { + existing: other, + ...Object.fromEntries( + choices.map((choice) => [choice.agent, { providerID: choice.providerID, modelID: choice.modelID }]), + ), + "mode-0": other, + }, + }) + expect(await state.preference.json()).toEqual({ preferred: { ...other, variant: "" } }) + }) + + it("waits for pending persistence before returning model selections", async () => { + const state = await fixture({ model: { plan: other } }) + + const saving = state.save({ agent: "code", ...model, variant: "high" }) + const loaded = await state.request() + await saving + + expect(loaded).toEqual({ + type: "modelSelectionsLoaded", + selections: { plan: other, code: model }, + preferred: { ...model, variant: "high" }, + }) + }) + + it("orders reset with pending saves without resurrecting cleared model memory", async () => { + const state = await fixture({ ...unrelated, model: { existing: other } }) + + await Promise.all([ + state.save({ agent: "code", ...model, variant: "high" }), + state.host.reset(state.client, state.post), + state.save({ agent: "plan", ...other }), + ]) + + expect(await state.file.json()).toEqual({ + ...unrelated, + model: { plan: other }, + }) + expect(await state.preference.json()).toEqual({ preferred: { ...other, variant: "" } }) + }) + + it.each(["high", ""])("retains preferred effort %j after a TUI rewrite of shared model state", async (variant) => { + const state = await fixture() + const preferred = { ...model, variant } + await state.save({ agent: "code", ...preferred }) + const before = await state.preference.text() + const shared = { model: { plan: other }, recent: [other], favorite: [], variant: { "other/model-b": "low" } } + + await Bun.write(state.file, JSON.stringify(shared)) + const host = await load() + await host.handleMessage("requestModelSelections", {}, state.client, state.post) + + expect(state.messages).toEqual([{ type: "modelSelectionsLoaded", selections: { plan: other }, preferred }]) + expect(await state.preference.text()).toBe(before) + expect(await state.file.json()).toEqual(shared) + }) + + it("does not read preferred metadata from the shared TUI file", async () => { + const state = await fixture({ model: { code: model }, preferred: { ...other, variant: "high" } }) + + expect(await state.request()).toEqual({ type: "modelSelectionsLoaded", selections: { code: model } }) + }) + + it("keeps saving to the resolved state directory during a temporary disconnect", async () => { + const state = await fixture() + await state.request() + await state.host.handleMessage( + "persistModelSelection", + { agent: "code", ...model, variant: "high" }, + null, + state.post, + ) + expect(await state.preference.json()).toEqual({ preferred: { ...model, variant: "high" } }) + await state.host.handleMessage("requestModelSelections", {}, null, state.post) + expect(state.messages.at(-1)).toEqual({ + type: "modelSelectionsLoaded", + selections: { code: model }, + preferred: { ...model, variant: "high" }, + }) + }) + + it("defers a cold-start request without a client until a ready-client retry", async () => { + const preferred = { ...model, variant: "high" } + const state = await fixture({ model: { code: model } }, { preferred }) + + expect(await state.host.handleMessage("requestModelSelections", {}, null, state.post)).toBe(true) + expect(state.messages).toEqual([]) + expect(await state.request()).toEqual({ type: "modelSelectionsLoaded", selections: { code: model }, preferred }) + }) + + it.each(["missing", "error"])("defers loading when the CLI state path is unavailable (%s)", async (reason) => { + const preferred = { ...model, variant: "high" } + const state = await fixture({ model: { code: model } }, { preferred }) + const client = { + path: { + get: async () => { + if (reason === "error") throw new Error("CLI is not ready") + return { data: {} } + }, + }, + } as unknown as KiloClient + + expect(await state.host.handleMessage("requestModelSelections", {}, client, state.post)).toBe(true) + expect(state.messages).toEqual([]) + expect(await state.request()).toEqual({ type: "modelSelectionsLoaded", selections: { code: model }, preferred }) + }) +}) diff --git a/packages/kilo-vscode/tests/unit/multi-version.test.ts b/packages/kilo-vscode/tests/unit/multi-version.test.ts index 9d5a1882ef1..6f082f0196e 100644 --- a/packages/kilo-vscode/tests/unit/multi-version.test.ts +++ b/packages/kilo-vscode/tests/unit/multi-version.test.ts @@ -36,6 +36,14 @@ describe("resolveVersionModels", () => { }) describe("buildInitialMessages", () => { + test.each(["high", ""])("preserves model, agent, and effort %s without an initial prompt", (variant) => { + const model = { providerID: "a", modelID: "m1" } + const messages = buildInitialMessages(created(1), [], model, undefined, "plan", variant) + expect(messages.at(0)).toEqual({ sessionId: "ses-0", worktreeId: "wt-0", ...model, agent: "plan", variant }) + const comparisons = buildInitialMessages(created(1), [{ ...model, variant }], {}, "", "plan", "low") + expect(comparisons.at(0)).toEqual(messages.at(0)) + }) + test("per-allocation variant wins over the dialog-level variant", () => { const models = resolveVersionModels( [ diff --git a/packages/kilo-vscode/tests/unit/new-worktree-dialog-sandbox.test.ts b/packages/kilo-vscode/tests/unit/new-worktree-dialog-sandbox.test.ts index dc6a5d4a6e6..e48c6354e6c 100644 --- a/packages/kilo-vscode/tests/unit/new-worktree-dialog-sandbox.test.ts +++ b/packages/kilo-vscode/tests/unit/new-worktree-dialog-sandbox.test.ts @@ -7,6 +7,14 @@ const providerPath = join(__dirname, "..", "..", "src", "KiloProvider.ts") const src = readFileSync(path, "utf8") const provider = readFileSync(providerPath, "utf8") +function fragment(start: string, end: string) { + const from = src.indexOf(start) + const to = src.indexOf(end, from) + expect(from).toBeGreaterThanOrEqual(0) + expect(to).toBeGreaterThan(from) + return new Bun.Transpiler({ loader: "tsx" }).transformSync(src.slice(from, to)) +} + describe("NewWorktreeDialog sandbox toggle", () => { it("uses the persisted default and only sends explicit modal overrides", () => { expect(src).toContain('vscode.postMessage({ type: "requestSandboxDefault", requestID: sandboxRequestID })') @@ -47,6 +55,7 @@ function check(code: string) { import { plugin } from "bun" import { isModelValid } from "./src/context/provider-utils.ts" import { toggleModel, setAllocationVariant } from "./agent-manager/multi-model-utils.ts" + import { DEFAULT_VARIANT, preserveVariant } from "./src/context/session-variant-store.ts" const solid = join(dirname(require.resolve("solid-js")), "solid.js") plugin({ @@ -55,7 +64,7 @@ function check(code: string) { build.onResolve({ filter: /^solid-js$/ }, () => ({ path: solid })) }, }) - const { batch, createComputed, createRoot, createSignal } = await import("solid-js") + const { batch, createComputed, createEffect, createMemo, createRoot, createSignal, onCleanup } = await import("solid-js") const { createDialogModels } = await import("./agent-manager/new-worktree-models.ts") const x = { providerID: "kilo", modelID: "x" } @@ -87,7 +96,50 @@ function check(code: string) { createComputed(() => seen.push(state.model())) return { state, snapshot, refresh: (update) => refresh((current) => ({ ...current, ...update })), switchAgent, seen } } - createRoot((dispose) => { + + // Run the dialog's actual model/variant setup, handlers and persistence effect, without rendering unrelated UI. + function dialog(saved = {}, initial = { providers: catalog(x, y), fallback: y, alternate: y, ready: true, connected: [] }) { + const result = createRoot((dispose) => { + const [snapshot, refresh] = createSignal(initial) + const [agent, setAgent] = createSignal(saved.agent ?? "code") + const [compareMode, setCompareMode] = createSignal(false) + const [modelAllocations, setModelAllocations] = createSignal(new Map()) + const [sandbox] = createSignal(undefined) + const preferences = [] + const provider = { + ready: () => snapshot().ready, + isModelValid: (value) => isModelValid(snapshot().providers, snapshot().connected, value), + findModel: (value) => snapshot().providers[value.providerID]?.models[value.modelID], + } + const session = { + modelForAgent: (name) => name === "code" ? snapshot().fallback : snapshot().alternate ?? null, + variantForAgent: (name) => snapshot().efforts?.[name], + preferredSelection: () => snapshot().preferred, + preferencesReady: () => snapshot().hydrated ?? true, + rememberSelection: (...args) => preferences.push(args), + } + let cached = {} + const vscode = { getState: () => cached, setState: (value) => { cached = value } } + ${fragment("const preferred =", "const [versions")} + ${fragment("const selection = createDialogModels({", "const [compareMode")} + ${fragment("const [variant, setVariant] =", "const [sandbox, setSandbox] =")} + ${fragment("const selectAgent =", "const cycle =")} + ${fragment("// Variant list for the currently selected model", " createEffect(() => {\n if (!sandboxVisible())")} + ${fragment(" createEffect(() => {\n const state = vscode.getState", "// Auto-persist images")} + const pick = ${src.match(/onSelect=\{(\(pid, mid\) => \{[\s\S]*?\n\s*\})\}/)?.[1]} + const choose = ${src.match(/ cached.advancedDialogSelections, + refresh: (update) => refresh((current) => ({ ...current, ...update })), dispose, + } + }) + onCleanup(result.dispose) + return result + } + await createRoot(async (dispose) => { try { ${code} } finally { @@ -106,14 +158,17 @@ function check(code: string) { describe("NewWorktreeDialog models", () => { it("persists only the saved choice and wires the effective model to display, variants, and guarded submission", () => { expect(src).toContain("saved: saved.model,") + expect(src).toContain("const preferred = session.preferredSelection()") expect(src).toContain("fallback: () => session.modelForAgent(agent()),") expect(src).toContain("ready: provider.ready,") expect(src).toContain("const model = selection.model") expect(src).toContain("model: selection.choice(),") expect(src).not.toContain("model: model(),") - expect(src).toContain("selection.select(undefined)") + expect(src).not.toContain("selection.select(undefined)") + expect(src).toContain("selection.retain()") expect(src).not.toContain("setModel(") expect(src).toContain("selection.select(next)") + expect(src).toContain("selectVariant(next)") expect(src).toContain("value={model()}") expect(src).toContain("const sel = model()") expect(src).toContain("session.variantForAgent(agent(), model())") @@ -137,6 +192,226 @@ describe("NewWorktreeDialog models", () => { `) }) + it("keeps an explicit model and effort through real dialog mode switches", () => { + check(` + const state = dialog({ model: x, variant: "high" }) + await Promise.resolve() + state.selectAgent("plan") + assert.deepEqual(state.model(), x) + assert.equal(state.effectiveVariant(), "high") + assert.deepEqual(state.cached().model, x) + assert.equal(state.cached().variant, "high") + state.selectAgent("code") + assert.deepEqual(state.model(), x) + assert.equal(state.effectiveVariant(), "high") + assert.deepEqual(state.preferences, []) + `) + }) + + it("pins the displayed inherited model and effort on mode switch without saving a model preference", () => { + check(` + const state = dialog({}, { + providers: catalog(x, y), fallback: x, alternate: y, ready: true, connected: [], + efforts: { code: "high", plan: undefined }, + }) + await Promise.resolve() + state.selectAgent("plan") + assert.deepEqual(state.model(), x) + assert.equal(state.effectiveVariant(), "high") + state.refresh({ providers: catalog(y) }) + assert.deepEqual(state.model(), y) + state.selectAgent("code") + state.refresh({ providers: catalog(x, y) }) + assert.deepEqual(state.model(), x) + assert.equal(state.selection.choice(), undefined) + assert.equal(state.cached().model, undefined) + assert.deepEqual(state.preferences, []) + `) + }) + + it("restores the saved effort after an empty catalog refresh and reopening", () => { + check(` + const state = dialog({ model: x, variant: "high" }) + await Promise.resolve() + state.refresh({ providers: {}, ready: false }) + assert.equal(state.model(), null) + assert.equal(state.effectiveVariant(), undefined) + assert.equal(state.cached().variant, "high") + const reopened = dialog(state.cached(), { providers: {}, fallback: y, ready: false, connected: [] }) + await Promise.resolve() + assert.equal(reopened.cached().variant, "high") + reopened.refresh({ providers: catalog(x, y), ready: true }) + assert.deepEqual(reopened.model(), x) + assert.equal(reopened.effectiveVariant(), "high") + assert.deepEqual(reopened.preferences, []) + `) + }) + + it("does not overwrite saved effort with a temporary catalog fallback's nearest effort", () => { + check(` + const state = dialog({ model: x, variant: "high" }) + await Promise.resolve() + const providers = catalog(y) + providers.kilo.models.y.variants = { low: {} } + state.refresh({ providers }) + assert.deepEqual(state.model(), y) + assert.equal(state.effectiveVariant(), "low") + assert.equal(state.cached().variant, "high") + state.selectAgent("plan") + state.refresh({ providers: catalog(x, y) }) + assert.deepEqual(state.model(), x) + assert.equal(state.effectiveVariant(), "high") + assert.deepEqual(state.preferences, []) + `) + }) + + it("carries the saved effort when explicitly picking a model during a temporary fallback", () => { + check(` + const state = dialog({ model: x, variant: "high" }) + await Promise.resolve() + const providers = catalog(y, z) + providers.kilo.models.y.variants = { low: {} } + providers.kilo.models.z.variants = { low: {}, high: {} } + state.refresh({ providers }) + assert.equal(state.effectiveVariant(), "low") + state.pick("kilo", "z") + assert.deepEqual(state.model(), z) + assert.equal(state.effectiveVariant(), "high") + assert.deepEqual(state.preferences, [["code", z, "high"]]) + `) + }) + + it("shares only explicit single-model picks and keeps explicit Default across mode switches", () => { + check(` + const state = dialog({}, { + providers: catalog(x, y), fallback: x, alternate: y, ready: true, connected: [], + efforts: { code: "high", plan: "high" }, + }) + await Promise.resolve() + assert.deepEqual(state.preferences, []) + state.pick("kilo", "y") + assert.deepEqual(state.preferences, [["code", y, "high"]]) + state.clear() + assert.deepEqual(state.preferences.at(-1), ["code", y, ""]) + state.selectAgent("plan") + assert.equal(state.effectiveVariant(), undefined) + assert.equal(state.variant(), "") + assert.equal(state.preferences.length, 2) + state.choose("high") + assert.deepEqual(state.preferences.at(-1), ["plan", y, "high"]) + const reopened = dialog(state.cached()) + await Promise.resolve() + assert.deepEqual(reopened.model(), y) + assert.equal(reopened.effectiveVariant(), "high") + assert.deepEqual(reopened.preferences, []) + `) + }) + + it("starts with the latest shared preference rather than stale dialog choices without following later updates", () => { + check(` + const state = dialog({ model: x, variant: "high" }, { + providers: catalog(x, y), fallback: x, ready: true, connected: [], preferred: { ...y, variant: "" }, + }) + await Promise.resolve() + assert.deepEqual(state.model(), y) + assert.equal(state.effectiveVariant(), undefined) + assert.equal(state.variant(), "") + state.refresh({ preferred: { ...x, variant: "high" } }) + assert.deepEqual(state.model(), y) + assert.equal(state.variant(), "") + assert.deepEqual(state.preferences, []) + `) + }) + + it("adopts a delayed first hydrated preference once without following later shared updates", () => { + check(` + const providers = catalog(x, y, z) + providers.kilo.models.y.variants = { low: {}, high: {} } + const state = dialog({ model: x, variant: "high" }, { + providers, fallback: x, ready: true, connected: [], hydrated: false, + }) + await Promise.resolve() + state.refresh({ preferred: { ...y, variant: "low" } }) + assert.deepEqual(state.model(), x) + assert.equal(state.effectiveVariant(), "high") + state.refresh({ hydrated: true }) + assert.deepEqual(state.model(), y) + assert.equal(state.effectiveVariant(), "low") + assert.deepEqual(state.cached().model, y) + assert.equal(state.cached().variant, "low") + state.refresh({ preferred: { ...z, variant: "high" } }) + assert.deepEqual(state.model(), y) + assert.equal(state.effectiveVariant(), "low") + assert.deepEqual(state.preferences, []) + `) + }) + + it("retains cached choices when first hydration has no preference and ignores a later preference", () => { + check(` + const state = dialog({ model: x, variant: "high" }, { + providers: catalog(x, y), fallback: y, ready: true, connected: [], hydrated: false, + }) + await Promise.resolve() + state.refresh({ hydrated: true }) + assert.deepEqual(state.model(), x) + assert.equal(state.effectiveVariant(), "high") + state.refresh({ preferred: { ...y, variant: "" } }) + assert.deepEqual(state.model(), x) + assert.equal(state.effectiveVariant(), "high") + assert.deepEqual(state.cached().model, x) + assert.equal(state.cached().variant, "high") + assert.deepEqual(state.preferences, []) + `) + }) + + it.each(["model", "variant", "default", "mode"])( + "preserves a %s interaction before initial preferences arrive", + (action) => { + check(` + const state = dialog({ model: x, variant: "high" }, { + providers: catalog(x, y, z), fallback: x, alternate: y, ready: true, connected: [], hydrated: false, + }) + const actions = { + model: () => state.pick("kilo", "z"), + variant: () => state.choose("high"), + default: state.clear, + mode: () => state.selectAgent("plan"), + } + actions[${JSON.stringify(action)}]() + await Promise.resolve() + const expected = state.model() + const effort = state.variant() + const writes = state.preferences.length + state.refresh({ hydrated: true, preferred: { ...y, variant: "" } }) + assert.deepEqual(state.model(), expected) + assert.equal(state.variant(), effort) + assert.deepEqual(state.cached().model, expected) + assert.equal(state.cached().variant, effort) + assert.equal(state.preferences.length, writes) + `) + }, + ) + + it("pins a model on explicit effort selection but never shares comparison choices", () => { + check(` + const state = dialog() + await Promise.resolve() + state.choose("high") + assert.deepEqual(state.preferences, [["code", y, "high"]]) + state.refresh({ fallback: x }) + assert.deepEqual(state.model(), y) + state.setCompareMode(true) + state.allocate(setAllocationVariant(toggleModel(new Map(), "kilo", "x", "X"), "kilo", "x", "high")) + state.choose(undefined) + assert.equal(state.preferences.length, 1) + state.selectAgent("plan") + state.refresh({ ready: false }) + state.refresh({ ready: true }) + assert.deepEqual(state.preferences, [["code", y, "high"]]) + assert.deepEqual(state.selection.choice(), y) + `) + }) + it("restores an initially unavailable cached X without replacing it with Y", () => { check(` const { state, refresh } = scene(x, { providers: catalog(y), fallback: y, ready: true, connected: [] }) @@ -158,17 +433,15 @@ describe("NewWorktreeDialog models", () => { const { state, refresh, switchAgent, seen } = scene(undefined) assert.deepEqual(state.model(), y) assert.equal(state.choice(), undefined) - state.select(y) - assert.deepEqual(state.choice(), y) refresh({ providers: catalog(y, z), alternate: z }) batch(() => { + state.retain() switchAgent("plan") - state.select(undefined) }) - assert.deepEqual(state.model(), z) + assert.deepEqual(state.model(), y) assert.equal(state.choice(), undefined) refresh({ providers: catalog(x), alternate: x }) - assert.deepEqual(seen, [y, z, x]) + assert.deepEqual(seen, [y, x]) assert.equal(state.choice(), undefined) `) }) diff --git a/packages/kilo-vscode/tests/unit/prompt-send-contract.test.ts b/packages/kilo-vscode/tests/unit/prompt-send-contract.test.ts index 6cfd9dbc161..1d1244e1ae8 100644 --- a/packages/kilo-vscode/tests/unit/prompt-send-contract.test.ts +++ b/packages/kilo-vscode/tests/unit/prompt-send-contract.test.ts @@ -90,7 +90,7 @@ describe("sendCommand dismisses pending tool requests", () => { expect(body).toContain("selectAgent(overrides.agent, scope)") expect(body).toContain("if (overrides?.model)") expect(body).toContain("selectModel(effectiveSelection.providerID, effectiveSelection.modelID, scope)") - expect(body).toContain("if (overrides?.variant)") + expect(body).toContain("if (overrides?.variant !== undefined)") expect(body).toContain("selectVariant(overrides.variant, scope)") }) }) diff --git a/packages/kilo-vscode/tests/unit/session-model-store.test.ts b/packages/kilo-vscode/tests/unit/session-model-store.test.ts index e866efc19ca..d8329d74dd7 100644 --- a/packages/kilo-vscode/tests/unit/session-model-store.test.ts +++ b/packages/kilo-vscode/tests/unit/session-model-store.test.ts @@ -156,14 +156,14 @@ describe("per-mode model memory", () => { expect(getAgentModel(store, env(), "ask")).toEqual(gpt) }) - it("ignores stale remembered selections when a configured mode model is user-set", () => { + it("keeps an explicit remembered model ahead of the configured mode default", () => { const configured: ResolveEnv = { ...env(), getModeModel: (name) => (name === "code" ? claude : null), } const store = { ...emptyStore(), modelSelections: { code: gpt } } - expect(getAgentModel(store, configured, "code", true)).toEqual(claude) + expect(getAgentModel(store, configured, "code", true)).toEqual(gpt) }) it("applyModel in a session writes only to sessionOverrides", () => { @@ -174,7 +174,7 @@ describe("per-mode model memory", () => { expect(result.modelSelections["code"]).toBeUndefined() }) - it("switching modes falls back to default after session override is cleared", () => { + it("keeps a session override when its selected mode changes", () => { let store = emptyStore() const e = env() @@ -182,10 +182,8 @@ describe("per-mode model memory", () => { const result = applyModel(store, "code", claude, "session-a") store = { ...store, ...result } - // Simulate mode switch: clear session override (like selectAgent does) - const cleared = { ...store, sessionOverrides: {} } - - expect(getSelected(cleared, e, "session-a", "code")).toEqual(KILO_AUTO) + const switched = { ...store, agentSelections: { "session-a": "ask" } } + expect(getSessionModel(switched, e, "session-a", "code")).toEqual(claude) }) it("different modes remember their own model independently", () => { @@ -234,7 +232,7 @@ describe("per-mode model memory", () => { expect(Object.keys(result.sessionOverrides)).toHaveLength(0) }) - it("switching from plan to implementation uses implementation config after clearing stale memory", () => { + it("keeps a plan session's explicit model when switching to configured implementation mode", () => { let store = emptyStore() const configured: ResolveEnv = { ...env(), @@ -252,11 +250,9 @@ describe("per-mode model memory", () => { const switched: ModelStore = { ...store, agentSelections: { "session-a": "code" }, - sessionOverrides: {}, - modelSelections: { ...store.modelSelections, code: null }, } - expect(getSelected(switched, configured, "session-a", "code")).toEqual(gpt) + expect(getSelected(switched, configured, "session-a", "code")).toEqual(claude) }) }) @@ -338,10 +334,10 @@ describe("organization model store", () => { expect(getSessionModel(store, { ...organization, connected: [] }, "session-a", "code")).toEqual(recommendation) }) - it("keeps Agent Manager mode configuration precedence without destroying the manual choice", () => { + it("honors the same explicit choice in Agent Manager and the chat picker", () => { const store = { ...emptyStore(), modelSelections: { code: KILO_AUTO }, userSetAgents: { code: true } } const configured = { ...organization, getModeModel: () => first, getGlobalModel: () => gpt } - expect(getAgentModel(store, configured, "code")).toEqual(first) + expect(getAgentModel(store, configured, "code")).toEqual(KILO_AUTO) expect(getSelected(store, configured, undefined, "code")).toEqual(KILO_AUTO) expect(store.modelSelections.code).toEqual(KILO_AUTO) expect(getAgentModel(store, organization, "code")).toEqual(KILO_AUTO) @@ -352,4 +348,24 @@ describe("organization model store", () => { expect(getAgentModel(store, { ...organization, getModeModel: () => first }, "code")).toEqual(first) expect(getSelected(store, { ...organization, getGlobalModel: () => gpt }, undefined, "code")).toEqual(gpt) }) + + it("uses the latest explicit default across modes without changing a session override", () => { + const store: ModelStore = { + ...emptyStore(), + preferred: gpt, + modelSelections: { code: KILO_AUTO, ask: first }, + userSetAgents: { code: true, ask: true }, + sessionOverrides: { active: claude }, + } + const configured = { ...organization, getModeModel: () => first, getGlobalModel: () => KILO_AUTO } + for (const agent of ["code", "ask"]) { + expect(getAgentModel(store, configured, agent)).toEqual(gpt) + expect(getSelected(store, configured, undefined, agent)).toEqual(gpt) + expect(getSelected(store, configured, "active", agent)).toEqual(claude) + } + const unavailable = { ...configured, connected: ["kilo"] } + expect(getAgentModel(store, unavailable, "code")).toEqual(KILO_AUTO) + expect(getAgentModel(store, configured, "code")).toEqual(gpt) + expect(store.preferred).toEqual(gpt) + }) }) diff --git a/packages/kilo-vscode/tests/unit/session-variant-store.test.ts b/packages/kilo-vscode/tests/unit/session-variant-store.test.ts index 148acc32ced..dfa6afa7f3e 100644 --- a/packages/kilo-vscode/tests/unit/session-variant-store.test.ts +++ b/packages/kilo-vscode/tests/unit/session-variant-store.test.ts @@ -57,11 +57,11 @@ describe("per-session variant selection", () => { }) it.each(["anthropic/claude-sonnet-4", variantKey(model, "code")])( - "prefers the configured variant over the remembered preference %s", + "preserves the remembered preference %s above the configured variant", (key) => { const store = { [key]: "high" } - expect(getVariant(store, model, ["high", "max"], "code", "pending-new", "max")).toBe("max") - expect(getAgentVariant(store, model, { variants: { high: {}, max: {} } }, "code", "max")).toBe("max") + expect(getVariant(store, model, ["high", "max"], "code", "pending-new", "max")).toBe("high") + expect(getAgentVariant(store, model, { variants: { high: {}, max: {} } }, "code", "max")).toBe("high") }, ) diff --git a/packages/kilo-vscode/tests/unit/session-variants.test.ts b/packages/kilo-vscode/tests/unit/session-variants.test.ts index 670c870aaf3..1014e82a987 100644 --- a/packages/kilo-vscode/tests/unit/session-variants.test.ts +++ b/packages/kilo-vscode/tests/unit/session-variants.test.ts @@ -8,6 +8,7 @@ function setup(session?: string, configured?: string) { const config = { model: "anthropic/claude-sonnet-4", variant: configured } const selections: Record = {} const messages: Array<{ type: string; key?: string; value?: string }> = [] + const remembered: Array<{ agent: string; model: ModelSelection; variant: string }> = [] const order: string[] = [] let handler: ((message: ExtensionMessage) => void) | undefined const variants = createSessionVariants({ @@ -20,6 +21,7 @@ function setup(session?: string, configured?: string) { agent: () => "code", config: () => config, find: () => ({ variants: { low: {}, high: {}, max: {} } }), + remember: (agent, model, variant) => remembered.push({ agent, model, variant }), post: (message) => { order.push("post") messages.push(message) @@ -30,7 +32,15 @@ function setup(session?: string, configured?: string) { return () => order.push("unsub") }, }) - return { variants, config, selections, messages, order, dispatch: (message: ExtensionMessage) => handler?.(message) } + return { + variants, + config, + selections, + messages, + remembered, + order, + dispatch: (message: ExtensionMessage) => handler?.(message), + } } describe("session variants", () => { @@ -60,14 +70,14 @@ describe("session variants", () => { expect(state.variants.request()).toBe("max") }) - it("uses updated configuration ahead of remembered defaults for new tabs", () => { + it("keeps remembered effort when configuration changes for new tabs", () => { const state = setup("pending-new", "high") state.selections["agent/code/anthropic/claude-sonnet-4"] = "low" - expect(state.variants.current()).toBe("high") + expect(state.variants.current()).toBe("low") state.config.variant = "max" - expect(state.variants.current()).toBe("max") - expect(state.variants.request()).toBe("max") - expect(state.variants.agent("code", model)).toBe("max") + expect(state.variants.current()).toBe("low") + expect(state.variants.request()).toBe("low") + expect(state.variants.agent("code", model)).toBe("low") }) it("does not apply a configured variant to another model", () => { @@ -86,26 +96,25 @@ describe("session variants", () => { expect(state.variants.request("session-b")).toBe("max") }) - it.each(["sidebar-pending:new", "pending:new"])("keeps a pre-submit Default choice scoped to %s", (id) => { + it.each(["sidebar-pending:new", "pending:new"])("remembers a pre-submit Default choice from %s", (id) => { const state = setup(undefined, "max") state.variants.select(undefined, id) expect(state.variants.current(id)).toBeUndefined() expect(state.variants.request(id)).toBe("") - expect(state.variants.current("another-draft")).toBe("max") + expect(state.remembered).toEqual([{ agent: "code", model, variant: "" }]) expect(state.messages).toEqual([]) }) it("persists global selections but keeps session selections local", () => { const global = setup() global.variants.select("high") - expect(global.messages).toEqual([ - { type: "persistVariant", key: "agent/code/anthropic/claude-sonnet-4", value: "high" }, - ]) + expect(global.remembered).toEqual([{ agent: "code", model, variant: "high" }]) const scoped = setup("session-a") scoped.variants.select("low") expect(scoped.selections).toEqual({ "session/session-a/anthropic/claude-sonnet-4": "low" }) expect(scoped.messages).toEqual([]) + expect(scoped.remembered).toEqual([]) }) it("persists an explicit default selection", () => { @@ -114,7 +123,7 @@ describe("session variants", () => { state.variants.select(undefined) expect(state.selections).toEqual({ "agent/code/anthropic/claude-sonnet-4": "" }) expect(state.variants.current()).toBeUndefined() - expect(state.messages).toEqual([{ type: "persistVariant", key: "agent/code/anthropic/claude-sonnet-4", value: "" }]) + expect(state.remembered).toEqual([{ agent: "code", model, variant: "" }]) }) it("does not shadow a cached variant when carrying the model default", () => { diff --git a/packages/kilo-vscode/webview-ui/agent-manager/NewWorktreeDialog.tsx b/packages/kilo-vscode/webview-ui/agent-manager/NewWorktreeDialog.tsx index ffbe4e695da..c3cadf0c29d 100644 --- a/packages/kilo-vscode/webview-ui/agent-manager/NewWorktreeDialog.tsx +++ b/packages/kilo-vscode/webview-ui/agent-manager/NewWorktreeDialog.tsx @@ -136,6 +136,12 @@ export const NewWorktreeDialog: Component<{ const cached = vscode.getState>() const [prompt, setPrompt] = createSignal((cached?.advancedDialogPrompt as string) ?? "") const saved = readDialogSelections(cached?.advancedDialogSelections) + const preferred = session.preferredSelection() + let pending = !session.preferencesReady() + if (preferred) { + saved.model = { providerID: preferred.providerID, modelID: preferred.modelID } + saved.variant = preferred.variant + } const [versions, setVersions] = createSignal(1) const initialAgent = restoreAgent(saved.agent, session.agents(), session.selectedAgent()) const [agent, setAgent] = createSignal(initialAgent) @@ -179,9 +185,10 @@ export const NewWorktreeDialog: Component<{ } const selectAgent = (name: string) => { + pending = false + selection.retain() + setVariant(variant() ?? effectiveVariant() ?? (variants().length > 0 ? DEFAULT_VARIANT : undefined)) setAgent(name) - selection.select(undefined) - setVariant(undefined) } const cycle = (direction: 1 | -1) => { @@ -212,18 +219,28 @@ export const NewWorktreeDialog: Component<{ const list = variants() if (list.length === 0) return undefined const stored = variant() ?? session.variantForAgent(agent(), model()) - return stored && list.includes(stored) ? stored : undefined + // Catalog refreshes may temporarily hide a model or effort. Never rewrite the saved choice. + return preserveVariant(stored, list) }) - // Reset variant when model changes and stored variant is not in new list + const selectVariant = (value: string | undefined) => { + pending = false + const next = value ?? DEFAULT_VARIANT + setVariant(next) + const sel = model() + if (!sel || compareMode()) return + selection.select(sel) + session.rememberSelection(agent(), sel, next) + } + createEffect(() => { - const list = variants() - if (list.length === 0) { - setVariant(undefined) - return - } - const stored = variant() - if (stored && !list.includes(stored)) setVariant(preserveVariant(stored, list)) + if (!pending || !session.preferencesReady()) return + // Initial host preferences may arrive after opening, but never replace an in-progress choice. + pending = false + const preferred = session.preferredSelection() + if (!preferred) return + selection.select({ providerID: preferred.providerID, modelID: preferred.modelID }) + setVariant(preferred.variant) }) createEffect(() => { @@ -483,7 +500,7 @@ export const NewWorktreeDialog: Component<{ if (list.length === 0) return const next = cycleVariant(effectiveVariant(), list) e.preventDefault() - setVariant(next ?? DEFAULT_VARIANT) + selectVariant(next) return } undo(e) @@ -826,11 +843,14 @@ export const NewWorktreeDialog: Component<{ value={model()} onSelect={(pid, mid) => { if (!pid || !mid) return - const current = effectiveVariant() + pending = false + const current = variant() ?? effectiveVariant() const next = { providerID: pid, modelID: mid } const list = Object.keys(provider.findModel(next)?.variants ?? {}) + const effort = preserveVariant(current, list) ?? DEFAULT_VARIANT selection.select(next) - setVariant(preserveVariant(current, list) ?? DEFAULT_VARIANT) + setVariant(effort) + session.rememberSelection(agent(), next, effort) }} onPick={restorePrompt} onCancel={restorePrompt} @@ -842,8 +862,8 @@ export const NewWorktreeDialog: Component<{ setVariant(DEFAULT_VARIANT)} + onSelect={selectVariant} + onClear={() => selectVariant(DEFAULT_VARIANT)} allowClear clearLabel={t("common.default")} trigger={WORKTREE_PROMPT_SCOPE} diff --git a/packages/kilo-vscode/webview-ui/agent-manager/new-worktree-models.ts b/packages/kilo-vscode/webview-ui/agent-manager/new-worktree-models.ts index 7fafbdd6111..9aaf32aa2c1 100644 --- a/packages/kilo-vscode/webview-ui/agent-manager/new-worktree-models.ts +++ b/packages/kilo-vscode/webview-ui/agent-manager/new-worktree-models.ts @@ -10,9 +10,10 @@ export function createDialogModels(opts: { variants: (model: ModelSelection) => string[] }) { const [choice, select] = createSignal(opts.saved) + const [held, hold] = createSignal(null) const valid = (value: ModelSelection) => (value.providerID !== "kilo" || opts.ready()) && opts.valid(value) const model = createMemo(() => { - const saved = choice() + const saved = choice() ?? held() if (saved && valid(saved)) return saved const fallback = opts.fallback() return fallback && valid(fallback) ? fallback : null @@ -29,5 +30,9 @@ export function createDialogModels(opts: { (entry.variant === undefined || opts.variants(entry).includes(entry.variant)), ) } - return { choice, select, model, canSubmit } + const retain = () => { + // Mode switches retain the displayed default without turning it into a saved user preference. + if (!choice() && !held()) hold(model()) + } + return { choice, select, model, canSubmit, retain } } diff --git a/packages/kilo-vscode/webview-ui/src/context/model-selection.ts b/packages/kilo-vscode/webview-ui/src/context/model-selection.ts index efea49c065d..2ef584e1ef0 100644 --- a/packages/kilo-vscode/webview-ui/src/context/model-selection.ts +++ b/packages/kilo-vscode/webview-ui/src/context/model-selection.ts @@ -8,6 +8,7 @@ export function resolveModelSelection(input: { organizationId?: string | null defaults?: Record session?: ModelSelection | null + preferred?: ModelSelection | null override?: ModelSelection | null mode?: ModelSelection | null global?: ModelSelection | null @@ -20,7 +21,11 @@ export function resolveModelSelection(input: { return isModelValid(input.providers, input.connected, selection) ? selection : null } const preference = - validate(input.session) ?? validate(input.override) ?? validate(input.mode) ?? validate(input.global) + validate(input.session) ?? + validate(input.preferred) ?? + validate(input.override) ?? + validate(input.mode) ?? + validate(input.global) if (preference) return preference if (pending) return null if (input.organizationId) { diff --git a/packages/kilo-vscode/webview-ui/src/context/session-model-preferences.ts b/packages/kilo-vscode/webview-ui/src/context/session-model-preferences.ts new file mode 100644 index 00000000000..1d840613a56 --- /dev/null +++ b/packages/kilo-vscode/webview-ui/src/context/session-model-preferences.ts @@ -0,0 +1,72 @@ +import { batch } from "solid-js" +import type { ModelSelection, WebviewMessage } from "../types/messages" +import { variantKey } from "./session-variant-store" + +interface Store { + modelSelections: Record + sessionOverrides: Record + agentSelections: Record + variantSelections: Record +} + +export function createModelPreferences(options: { + store: Store + model: (scope: "modelSelections" | "sessionOverrides", id: string, model: ModelSelection) => void + set: (key: string, variant: string) => void + scopes: () => (string | undefined)[] + initialized: (id: string) => boolean + selected: (id: string) => ModelSelection | null + defaults: (agent: string) => ModelSelection | null | undefined + agent: (id: string) => string + variant: (id: string, model: ModelSelection) => string | undefined + recent: (model: ModelSelection) => void + update: (agent: string, model: ModelSelection, variant: string) => void + post: (message: WebviewMessage) => void +}) { + function pin(id: string) { + if (!options.store.sessionOverrides[id] && !options.initialized(id)) return + const model = options.store.sessionOverrides[id] ?? options.defaults(options.agent(id)) ?? options.selected(id) + if (!model) return + const key = variantKey(model, options.agent(id), id) + const value = options.variant(id, model) + if (options.store.variantSelections[key] === undefined && value !== undefined) options.set(key, value) + // Copy inherited models so updates to a mode's store cannot mutate the session. + if (!options.store.sessionOverrides[id]) options.model("sessionOverrides", id, { ...model }) + } + + function retain() { + const ids = new Set([...options.scopes(), ...Object.keys(options.store.agentSelections)]) + for (const id of ids) if (id) pin(id) + } + + function apply(agent: string, model: ModelSelection, id?: string) { + if (id) { + if (!/^(?:sidebar-)?pending:/.test(id)) options.recent(model) + options.model("sessionOverrides", id, model) + return + } + retain() + options.model("modelSelections", agent, model) + } + + function remember(agent: string, model: ModelSelection, variant = "") { + batch(() => { + // New-session defaults must not change other open sessions or drafts. + retain() + options.recent(model) + options.model("modelSelections", agent, { ...model }) + options.set(variantKey(model, agent), variant) + options.update(agent, model, variant) + }) + options.post({ + type: "persistModelSelection", + agent, + providerID: model.providerID, + modelID: model.modelID, + variant, + }) + options.post({ type: "persistVariant", key: variantKey(model, agent), value: variant }) + } + + return { apply, pin, remember } +} diff --git a/packages/kilo-vscode/webview-ui/src/context/session-model-store.ts b/packages/kilo-vscode/webview-ui/src/context/session-model-store.ts index 105d60fa373..3b27eb61991 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session-model-store.ts +++ b/packages/kilo-vscode/webview-ui/src/context/session-model-store.ts @@ -17,6 +17,7 @@ export interface ModelStore { agentSelections: Record recentModels: ModelSelection[] userSetAgents?: Record + preferred?: ModelSelection } export interface ResolveEnv { @@ -36,6 +37,7 @@ function resolveModel( override?: ModelSelection | null, recents?: ModelSelection[], session?: ModelSelection, + preferred?: ModelSelection, ): ModelSelection | null { return resolveModelSelection({ providers: env.providers, @@ -44,6 +46,7 @@ function resolveModel( organizationId: env.organizationId, defaults: env.defaults, session, + preferred: preferred && { providerID: preferred.providerID, modelID: preferred.modelID }, override, mode: env.getModeModel(agentName), global: env.getGlobalModel(), @@ -85,6 +88,7 @@ export function getSelected( override, store.recentModels, sessionID ? store.sessionOverrides[sessionID] : undefined, + store.preferred, ) } @@ -95,11 +99,8 @@ export function getAgentModel( agentName: string, userSet = store.userSetAgents?.[agentName] === true, ): ModelSelection | null { - const override = - (env.getModeModel(agentName) && userSet) || (env.organizationId && !userSet) - ? null - : store.modelSelections[agentName] - return resolveModel(env, agentName, override, store.recentModels) + const override = env.organizationId && !userSet ? null : store.modelSelections[agentName] + return resolveModel(env, agentName, override, store.recentModels, undefined, store.preferred) } export interface ApplyResult { diff --git a/packages/kilo-vscode/webview-ui/src/context/session-types.ts b/packages/kilo-vscode/webview-ui/src/context/session-types.ts index 6bf42703c8b..7d7daec36a1 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session-types.ts +++ b/packages/kilo-vscode/webview-ui/src/context/session-types.ts @@ -112,6 +112,9 @@ export interface SessionContextValue { selected: (sessionID?: string) => ModelSelection | null modelForAgent: (agent: string) => ModelSelection | null selectModel: (providerID: string, modelID: string, sessionID?: string) => void + preferredSelection: Accessor<(ModelSelection & { variant?: string }) | undefined> + preferencesReady: Accessor + rememberSelection: (agent: string, model: ModelSelection, variant?: string) => void // Cost and context usage for the current session costBreakdown: Accessor> diff --git a/packages/kilo-vscode/webview-ui/src/context/session-variant-store.ts b/packages/kilo-vscode/webview-ui/src/context/session-variant-store.ts index 757e0f63e6d..d2ede4e2aeb 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session-variant-store.ts +++ b/packages/kilo-vscode/webview-ui/src/context/session-variant-store.ts @@ -40,11 +40,12 @@ export function getVariant( agent: string, session?: string, configured?: string, + preferred?: string, ) { if (variants.length === 0) return undefined const scoped = session ? store[variantKey(sel, agent, session)] : undefined const preset = configured && variants.includes(configured) ? configured : undefined - const stored = scoped ?? preset ?? store[variantKey(sel, agent)] ?? store[legacyVariantKey(sel)] + const stored = scoped ?? preferred ?? store[variantKey(sel, agent)] ?? store[legacyVariantKey(sel)] ?? preset if (stored === undefined || stored === DEFAULT_VARIANT) return undefined return preserveVariant(stored, variants) } @@ -55,9 +56,10 @@ export function getAgentVariant( model: { variants?: Record } | undefined, agent: string, configured?: string, + preferred?: string, ) { if (!model?.variants) return undefined - return getVariant(store, sel, Object.keys(model.variants), agent, undefined, configured) + return getVariant(store, sel, Object.keys(model.variants), agent, undefined, configured, preferred) } /** diff --git a/packages/kilo-vscode/webview-ui/src/context/session-variants.ts b/packages/kilo-vscode/webview-ui/src/context/session-variants.ts index 5565b6426a2..e5be9564de0 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session-variants.ts +++ b/packages/kilo-vscode/webview-ui/src/context/session-variants.ts @@ -1,6 +1,13 @@ import type { Accessor } from "solid-js" import type { AgentConfig, ExtensionMessage, ModelSelection } from "../types/messages" -import { DEFAULT_VARIANT, getAgentVariant, getVariant, preserveVariant, variantKey } from "./session-variant-store" +import { + DEFAULT_VARIANT, + getAgentVariant, + getVariant, + legacyVariantKey, + preserveVariant, + variantKey, +} from "./session-variant-store" interface Model { variants?: Record @@ -18,6 +25,8 @@ interface Options { find: (selection: ModelSelection) => Model | undefined post: (message: Message) => void listen: (handler: (message: ExtensionMessage) => void) => () => void + preferred?: Accessor<(ModelSelection & { variant?: string }) | undefined> + remember: (agent: string, model: ModelSelection, variant: string) => void } export function createSessionVariants(options: Options) { @@ -33,9 +42,22 @@ export function createSessionVariants(options: Options) { return config.variant ?? undefined } + const preferred = (selection: ModelSelection) => { + const value = options.preferred?.() + if (value?.providerID !== selection.providerID || value.modelID !== selection.modelID) return undefined + return value.variant ?? DEFAULT_VARIANT + } + const agent = (name: string, selection: ModelSelection | null) => { if (!selection) return undefined - return getAgentVariant(options.selections(), selection, options.find(selection), name, configured(name, selection)) + return getAgentVariant( + options.selections(), + selection, + options.find(selection), + name, + configured(name, selection), + preferred(selection), + ) } const current = (sessionID?: string) => { @@ -45,12 +67,28 @@ export function createSessionVariants(options: Options) { const variants = list(sid) if (variants.length === 0) return undefined const name = options.agent(sid) - return getVariant(options.selections(), selection, variants, name, sid, configured(name, selection)) + return getVariant( + options.selections(), + selection, + variants, + name, + sid, + configured(name, selection), + preferred(selection), + ) } const request = (sessionID?: string) => current(sessionID) ?? (list(sessionID).length > 0 ? DEFAULT_VARIANT : undefined) + const saved = (selection: ModelSelection, name: string, sessionID?: string) => + (sessionID ? options.selections()[variantKey(selection, name, sessionID)] : undefined) ?? + preferred(selection) ?? + options.selections()[variantKey(selection, name)] ?? + options.selections()[legacyVariantKey(selection)] ?? + configured(name, selection) ?? + DEFAULT_VARIANT + const select = (value: string | undefined, sessionID?: string) => { const sid = sessionID ?? options.session() const selection = options.selected(sid) @@ -58,7 +96,9 @@ export function createSessionVariants(options: Options) { const key = variantKey(selection, options.agent(sid), sid) const next = value ?? DEFAULT_VARIANT options.set(key, next) - if (!sid) options.post({ type: "persistVariant", key, value: next }) + if (!sid || /^(?:sidebar-)?pending:/.test(sid)) { + options.remember(options.agent(sid), selection, next) + } } const carry = (selection: ModelSelection, value: string | undefined, name: string, sessionID?: string) => { @@ -67,7 +107,7 @@ export function createSessionVariants(options: Options) { // An absent value means the model default, not an explicit user choice. // Do not write a default sentinel here because it would shadow a cached // agent-level variant when this selection is resolved for a new session. - const next = preserveVariant(value, list) + const next = value === DEFAULT_VARIANT ? DEFAULT_VARIANT : preserveVariant(value, list) if (next === undefined) return const key = variantKey(selection, name, sessionID) options.set(key, next) @@ -86,5 +126,5 @@ export function createSessionVariants(options: Options) { return unsub } - return { carry, list, agent, current, request, select, load } + return { carry, list, agent, current, request, saved, select, load } } diff --git a/packages/kilo-vscode/webview-ui/src/context/session.tsx b/packages/kilo-vscode/webview-ui/src/context/session.tsx index fec11a0632e..71a8a9ad2fe 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session.tsx +++ b/packages/kilo-vscode/webview-ui/src/context/session.tsx @@ -86,7 +86,7 @@ import { PartStash } from "./part-stash" import { isolate, mergeOptimisticPart, mergeOptimisticParts, mergeParts } from "./session-parts" import { mergeMessages, sameReconcileShape } from "./session-merge" import { state as todoState } from "./todo-revert" -import { sessionVariantKeys, transferVariants, variantKey } from "./session-variant-store" +import { preserveVariant, sessionVariantKeys, transferVariants, variantKey } from "./session-variant-store" import { createSessionVariants } from "./session-variants" import { KILO_AUTO, KILO_PROVIDER_ID, parseModelString } from "../../../src/shared/provider-model" import { type ReviewMessageData } from "../../../src/shared/review-comments" @@ -100,6 +100,7 @@ import { clearIfOn, createCloudPrune } from "./session-cloud-prune" import { isSameSessionTree } from "./model-usage" import { createDraftAgentSeed, resolvePromptAgent } from "./session-agent" import { createModelSelector } from "./session-model-selector" +import { createModelPreferences } from "./session-model-preferences" import { activities, type Activity } from "../utils/session-activity" import { active as activeTiming, hold, type Timing } from "./session-timing" import type { SessionContextValue } from "./session-types" @@ -248,6 +249,9 @@ export const SessionProvider: ParentComponent = (props) => { // Tracks whether the user has explicitly set a model override per agent (to // prevent the default-sync effect from overwriting it). const [userSetAgents, setUserSetAgents] = createSignal>({}) + const [preferredSelection, setPreferredSelection] = createSignal() + const [preferencesReady, setPreferencesReady] = createSignal(false) + let remembered = false // Agents (modes) loaded from the CLI backend const [agents, setAgents] = createSignal([]) @@ -471,25 +475,21 @@ export const SessionProvider: ParentComponent = (props) => { agentSelections: store.agentSelections, recentModels: store.recentModels, userSetAgents: userSetAgents(), + preferred: preferredSelection(), } } - function resolveModel(agentName: string): ModelSelection | null { - return resolveModelSelection({ + // Keep automatic defaults in sync until the user chooses a model. + createEffect(() => { + const agentName = selectedAgentName() + if (userSetAgents()[agentName]) return + const sel = resolveModelSelection({ ...environment(), mode: getModeModel(agentName), global: getGlobalModel(), recent: store.recentModels, }) - } - - // Keep model selection in sync with provider/mode default until the user - // explicitly overrides it. - createEffect(() => { - const agentName = selectedAgentName() - if (userSetAgents()[agentName]) return - const sel = resolveModel(agentName) - setStore("modelSelections", agentName, sel) + setStore("modelSelections", agentName, sel && { ...sel }) }) const currentSelected = createMemo(() => @@ -518,24 +518,29 @@ export const SessionProvider: ParentComponent = (props) => { vscode.postMessage({ type: "recordModelUsage", providerID, modelID }) } - function applyModel(agentName: string, selection: ModelSelection, sessionID?: string) { - pushRecent(selection) - if (sessionID) { - setStore("sessionOverrides", sessionID, selection) - return - } - // Always remember the per-mode model choice so switching modes restores - // the last-used model (mirrors CLI TUI's model.json behavior). - setUserSetAgents((prev) => ({ ...prev, [agentName]: true })) - setStore("modelSelections", agentName, selection) - // Persist to model.json via the extension host - vscode.postMessage({ - type: "persistModelSelection", - agent: agentName, - providerID: selection.providerID, - modelID: selection.modelID, - }) - } + const memory = createModelPreferences({ + store, + model: (scope, id, model) => setStore(scope, id, model), + set: (key, value) => setStore("variantSelections", key, value), + scopes: () => [...Object.keys(store.messages), currentSessionID(), draftSessionID()], + initialized: (id) => !store.sessions[id] || store.messages[id] !== undefined, + selected, + defaults: (agent) => + preferredSelection() ?? + (userSetAgents()[agent] ? store.modelSelections[agent] : undefined) ?? + getModeModel(agent) ?? + getGlobalModel(), + agent: agentForScope, + variant: (id, model) => variants.saved(model, agentForScope(id), id), + recent: pushRecent, + update: (agent, model, variant) => { + remembered = true + setUserSetAgents((prev) => ({ ...prev, [agent]: true })) + setPreferredSelection({ ...model, variant }) + }, + post: vscode.postMessage, + }) + const rememberSelection = memory.remember const variants = createSessionVariants({ selections: () => store.variantSelections, @@ -547,6 +552,8 @@ export const SessionProvider: ParentComponent = (props) => { find: provider.findModel, post: vscode.postMessage, listen: vscode.onMessage, + preferred: preferredSelection, + remember: rememberSelection, }) const { carry: carryVariant, list: variantList, agent: variantForAgent, current: currentVariant } = variants const selectVariant = variants.select @@ -554,13 +561,25 @@ export const SessionProvider: ParentComponent = (props) => { current: currentSessionID, agent: agentForScope, selected, - variant: currentVariant, - apply: applyModel, + variant: variants.request, + apply: memory.apply, set: (id, selection) => setStore("sessionOverrides", id, selection), carry: carryVariant, hide: hideErrors, }) - const selectModel = models.select + function selectModel(providerID: string, modelID: string, sessionID?: string) { + const id = sessionID ?? currentSessionID() + batch(() => { + models.select(providerID, modelID, id) + if (!id || /^(?:sidebar-)?pending:/.test(id)) { + const model = { providerID, modelID } + const agent = agentForScope(id) + const value = store.variantSelections[variantKey(model, agent, id)] ?? variantForAgent(agent, model) + const list = Object.keys(provider.findModel(model)?.variants ?? {}) + rememberSelection(agent, model, value === "" ? "" : preserveVariant(value, list)) + } + }) + } function selectKiloModel(modelID?: string, agent?: string) { if (!modelID && !agent) return @@ -615,24 +634,6 @@ export const SessionProvider: ParentComponent = (props) => { }) } - function clearModeModelSelection(agentName: string) { - setUserSetAgents((prev) => { - const next = { ...prev } - delete next[agentName] - return next - }) - setStore( - "modelSelections", - produce((selections) => { - delete selections[agentName] - }), - ) - } - - function shouldClearModeModelSelection(agentName: string) { - return getModeModel(agentName) !== null && userSetAgents()[agentName] === true - } - function clearHiddenErrors(ids: string[]) { if (ids.length === 0) return setHiddenErrors((prev) => { @@ -736,6 +737,7 @@ export const SessionProvider: ParentComponent = (props) => { vscode.postMessage({ type: "requestMcpStatus" }) const fallback = setTimeout(() => { + if (!preferencesReady()) vscode.postMessage({ type: "requestModelSelections" }) if (agents().length === 0) vscode.postMessage({ type: "requestAgents" }) if (Object.keys(mcpStatus()).length === 0) vscode.postMessage({ type: "requestMcpStatus" }) }, 3000) @@ -744,6 +746,7 @@ export const SessionProvider: ParentComponent = (props) => { if (message.type !== "extensionDataReady") return unsubReady() clearTimeout(fallback) + if (!preferencesReady()) vscode.postMessage({ type: "requestModelSelections" }) if (agents().length === 0) vscode.postMessage({ type: "requestAgents" }) if (Object.keys(mcpStatus()).length === 0) vscode.postMessage({ type: "requestMcpStatus" }) }) @@ -763,12 +766,19 @@ export const SessionProvider: ParentComponent = (props) => { const unsubSelections = vscode.onMessage((message: ExtensionMessage) => { if (message.type !== "modelSelectionsLoaded") return batch(() => { - setStore("modelSelections", reconcile(message.selections)) + const local = + !preferencesReady() && remembered + ? Object.fromEntries(Object.entries(store.modelSelections).filter(([name]) => userSetAgents()[name])) + : {} + const selections = { ...message.selections, ...local } + setStore("modelSelections", reconcile(selections)) const flags: Record = {} - for (const name of Object.keys(message.selections)) { + for (const name of Object.keys(selections)) { flags[name] = true } setUserSetAgents(flags) + if (preferencesReady() || !remembered) setPreferredSelection(message.preferred) + setPreferencesReady(true) }) }) vscode.postMessage({ type: "requestModelSelections" }) @@ -2122,30 +2132,16 @@ export const SessionProvider: ParentComponent = (props) => { function selectAgent(name: string, sessionID?: string) { const id = sessionID ?? currentSessionID() if (id) { + memory.pin(id) setStore("agentSelections", id, name) - // Clear per-session model override so the new mode's configured/default - // model takes effect instead of the previous mode's override. - setStore( - "sessionOverrides", - produce((overrides) => { - delete overrides[id] - }), - ) - if (shouldClearModeModelSelection(name)) { - clearModeModelSelection(name) - } - } else { - setPendingAgentSelection(name) - if (shouldClearModeModelSelection(name)) { - clearModeModelSelection(name) - return - } - // When switching mode, initialize model for the new mode if the user - // hasn't explicitly set one for it - if (!userSetAgents()[name] && !store.modelSelections[name]) { - setStore("modelSelections", name, resolveModel(name)) - } + return } + const agent = selectedAgentName() + const model = store.modelSelections[agent] + if (!preferredSelection() && userSetAgents()[agent] && model) { + setPreferredSelection({ ...model, variant: variants.saved(model, agent) }) + } + setPendingAgentSelection(name) } /** Create an optimistic user message + parts in the store so the UI updates instantly. */ @@ -2303,11 +2299,9 @@ export const SessionProvider: ParentComponent = (props) => { if (control) return null if (overrides?.model) return parseModelString(overrides.model) const scope = draftID ?? sid - const model = overrides?.agent - ? modelForAgent(overrides.agent) - : scope - ? selected(scope) - : getSelected(preferences(), environment(), undefined, pendingAgentSelection() ?? defaultAgent()) + const model = scope + ? selected(scope) + : getSelected(preferences(), environment(), undefined, pendingAgentSelection() ?? defaultAgent()) return model ?? (providerID && modelID ? { providerID, modelID } : null) })() if (!control && !available(effectiveSelection)) return false @@ -2323,7 +2317,7 @@ export const SessionProvider: ParentComponent = (props) => { if (overrides?.model) { selectModel(effectiveSelection.providerID, effectiveSelection.modelID, scope) } - if (overrides?.variant) { + if (overrides?.variant !== undefined) { selectVariant(overrides.variant, scope) } recordModelUsage(effectiveSelection.providerID, effectiveSelection.modelID) @@ -2991,6 +2985,9 @@ export const SessionProvider: ParentComponent = (props) => { selected, modelForAgent, selectModel, + preferredSelection, + preferencesReady, + rememberSelection, costBreakdown, contextUsage, modelUsage, diff --git a/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx b/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx index e526d3b5d29..cfbc5aa5fe5 100644 --- a/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx +++ b/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx @@ -237,6 +237,9 @@ export function mockSessionValue(overrides?: { selected: () => ({ providerID: "kilo", modelID: "anthropic/claude-sonnet-4-6" }), modelForAgent: () => ({ providerID: "kilo", modelID: "anthropic/claude-sonnet-4-6" }), selectModel: noop, + preferredSelection: () => undefined, + preferencesReady: () => true, + rememberSelection: noop, costBreakdown: () => [], contextUsage: () => undefined, modelUsage: () => undefined, diff --git a/packages/kilo-vscode/webview-ui/src/types/messages/extension-messages.ts b/packages/kilo-vscode/webview-ui/src/types/messages/extension-messages.ts index 00616725669..d99c177c505 100644 --- a/packages/kilo-vscode/webview-ui/src/types/messages/extension-messages.ts +++ b/packages/kilo-vscode/webview-ui/src/types/messages/extension-messages.ts @@ -1105,10 +1105,11 @@ export interface FavoritesLoadedMessage { favorites: ModelSelection[] } -// Per-mode model selections loaded from model.json (extension → webview) +// Preferred and per-mode model selections loaded from persisted state (extension → webview) export interface ModelSelectionsLoadedMessage { type: "modelSelectionsLoaded" selections: Record + preferred?: ModelSelection & { variant?: string } } export interface AgentManagerBranchesMessage { diff --git a/packages/kilo-vscode/webview-ui/src/types/messages/webview-messages.ts b/packages/kilo-vscode/webview-ui/src/types/messages/webview-messages.ts index 963fdcbc6a5..1ed42908be9 100644 --- a/packages/kilo-vscode/webview-ui/src/types/messages/webview-messages.ts +++ b/packages/kilo-vscode/webview-ui/src/types/messages/webview-messages.ts @@ -1468,12 +1468,13 @@ export interface RequestFavoritesMessage { type: "requestFavorites" } -// Per-mode model selection persistence (webview → extension) +// Explicit preferred and per-mode model selection persistence (webview → extension) export interface PersistModelSelectionRequest { type: "persistModelSelection" agent: string providerID: string modelID: string + variant?: string } export interface RequestModelSelectionsMessage {