From 76e34bdcf094a3490228b87f80d52761bcd5cd42 Mon Sep 17 00:00:00 2001 From: marius-kilocode Date: Mon, 3 Aug 2026 11:57:15 +0200 Subject: [PATCH] fix(agent-manager): align modal model resolution --- .../tests/unit/session-model-store.test.ts | 19 +++++++++----- .../tests/unit/session-variant-store.test.ts | 8 ++++++ .../agent-manager/NewWorktreeDialog.tsx | 2 +- .../src/context/session-model-store.ts | 11 ++++++++ .../src/context/session-variant-store.ts | 10 ++++++++ .../webview-ui/src/context/session.tsx | 25 +++++++++++++++---- 6 files changed, 63 insertions(+), 12 deletions(-) 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 bed7b4c6a09..0c99900dcb1 100644 --- a/packages/kilo-vscode/tests/unit/session-model-store.test.ts +++ b/packages/kilo-vscode/tests/unit/session-model-store.test.ts @@ -3,6 +3,7 @@ import { type ModelStore, type ResolveEnv, applyModel, + getAgentModel, getSessionModel, getSelected, } from "../../webview-ui/src/context/session-model-store" @@ -149,14 +150,20 @@ describe("per-session model selection", () => { }) describe("per-mode model memory", () => { - it("resolves a specific mode's remembered model without a session", () => { - let store = emptyStore() - const e = env() + it("uses remembered model selections for modes without configured models", () => { + const store = { ...emptyStore(), modelSelections: { ask: gpt } } - const result = applyModel(store, "ask", gpt, undefined) - store = { ...store, ...result } + expect(getAgentModel(store, env(), "ask")).toEqual(gpt) + }) - expect(getSelected(store, e, undefined, "ask")).toEqual(gpt) + it("ignores stale remembered selections when a configured mode model is user-set", () => { + const configured: ResolveEnv = { + ...env(), + getModeModel: (name) => (name === "code" ? claude : null), + } + const store = { ...emptyStore(), modelSelections: { code: gpt } } + + expect(getAgentModel(store, configured, "code", true)).toEqual(claude) }) it("applyModel in a session writes only to sessionOverrides", () => { 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 c6c27294b16..a2b678009c3 100644 --- a/packages/kilo-vscode/tests/unit/session-variant-store.test.ts +++ b/packages/kilo-vscode/tests/unit/session-variant-store.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it } from "bun:test" import { cycleVariant, + getAgentVariant, getVariant, sessionVariantKeys, sessionVariants, @@ -43,6 +44,13 @@ describe("per-session variant selection", () => { expect(getVariant(store, model, variants, "ask")).toBe("high") }) + it("resolves the effective variant for a mode and model", () => { + const store: Record = {} + store[variantKey(model, "ask")] = "high" + + expect(getAgentVariant(store, model, { variants: { low: {}, high: {} } }, "ask")).toBe("high") + }) + it("carries the pre-submit agent variant into a newly created session", () => { const store: Record = {} diff --git a/packages/kilo-vscode/webview-ui/agent-manager/NewWorktreeDialog.tsx b/packages/kilo-vscode/webview-ui/agent-manager/NewWorktreeDialog.tsx index 164f3c7a9d8..65217564e36 100644 --- a/packages/kilo-vscode/webview-ui/agent-manager/NewWorktreeDialog.tsx +++ b/packages/kilo-vscode/webview-ui/agent-manager/NewWorktreeDialog.tsx @@ -107,7 +107,7 @@ export const NewWorktreeDialog: Component<{ const [prompt, setPrompt] = createSignal((cached?.advancedDialogPrompt as string) ?? "") const [versions, setVersions] = createSignal(1) const initialAgent = session.selectedAgent() - const initialModel = session.configModelForAgent(initialAgent) + const initialModel = session.modelForAgent(initialAgent) const [model, setModel] = createSignal<{ providerID: string; modelID: string } | null>(initialModel) const [compareMode, setCompareMode] = createSignal(false) const [modelAllocations, setModelAllocations] = createSignal(new Map()) 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 424d5bf51d8..704e03fc648 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 @@ -78,6 +78,17 @@ export function getSelected( return resolveModel(env, agentName, store.modelSelections[agentName], store.recentModels) } +/** Returns the effective model for a mode outside a session scope. */ +export function getAgentModel( + store: ModelStore, + env: ResolveEnv, + agentName: string, + userSet = false, +): ModelSelection | null { + const override = env.getModeModel(agentName) && userSet ? null : store.modelSelections[agentName] + return resolveModel(env, agentName, override, store.recentModels) +} + export interface ApplyResult { modelSelections: Record sessionOverrides: Record 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 a997f251d7d..bbbb5ca5a8a 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 @@ -24,6 +24,16 @@ export function getVariant( return stored && variants.includes(stored) ? stored : variants[0] } +export function getAgentVariant( + store: Record, + sel: ModelSelection, + model: { variants?: Record } | undefined, + agent: string, +) { + if (!model?.variants) return undefined + return getVariant(store, sel, Object.keys(model.variants), agent) +} + /** * Next variant in the list, wrapping back to the first after the last. * An unknown or missing current value starts at the first variant. diff --git a/packages/kilo-vscode/webview-ui/src/context/session.tsx b/packages/kilo-vscode/webview-ui/src/context/session.tsx index dec3809626e..54917021876 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session.tsx +++ b/packages/kilo-vscode/webview-ui/src/context/session.tsx @@ -67,12 +67,13 @@ import { } from "./session-utils" import { Identifier } from "../utils/id" import { resolveModelSelection } from "./model-selection" +import { getAgentModel } from "./session-model-store" import { resolveMessagePrefs } from "./session-preferences" import { errorIDs } from "./session-errors" import { PartStash } from "./part-stash" import { mergeParts, sameParts } from "./session-parts" import { state as todoState } from "./todo-revert" -import { getVariant, sessionVariantKeys, transferVariants, variantKey } from "./session-variant-store" +import { getAgentVariant, getVariant, sessionVariantKeys, transferVariants, variantKey } from "./session-variant-store" import { KILO_AUTO, KILO_PROVIDER_ID, parseModelString } from "../../../src/shared/provider-model" import { reviewMetadata, type ReviewMessageData } from "../../../src/shared/review-comments" import { visibleMessages as filterVisibleMessages } from "./session-queue" @@ -739,8 +740,23 @@ export const SessionProvider: ParentComponent = (props) => { } function modelForAgent(agentName: string): ModelSelection | null { - const override = shouldClearModeModelSelection(agentName) ? null : store.modelSelections[agentName] - return resolveModel(agentName, override) + return getAgentModel( + { + modelSelections: store.modelSelections, + sessionOverrides: store.sessionOverrides, + agentSelections: store.agentSelections, + recentModels: store.recentModels, + }, + { + providers: provider.providers(), + connected: provider.connected(), + getModeModel, + getGlobalModel, + fallback: KILO_AUTO, + }, + agentName, + userSetAgents()[agentName] === true, + ) } function configModelForAgent(agentName: string): ModelSelection | null { @@ -893,8 +909,7 @@ export const SessionProvider: ParentComponent = (props) => { function variantForAgent(agentName: string, sel: ModelSelection | null) { if (!sel) return undefined const model = provider.findModel(sel) - if (!model?.variants) return undefined - return getVariant(store.variantSelections, sel, Object.keys(model.variants), agentName) + return getAgentVariant(store.variantSelections, sel, model, agentName) } const currentVariant = (sessionID?: string) => {