Fix Agent Manager session selector state (#9922)

* tmp

* fix(agent-manager): recover session preferences from history

* fix(agent-manager): use scoped prompt selections

* Delete packages/kilo-vscode/docs/agent-manager-session-preferences.md
This commit is contained in:
Marius
2026-05-06 15:46:43 +00:00
committed by GitHub
parent e5e9d0ba37
commit 5fcc4231c0
21 changed files with 610 additions and 141 deletions
+5
View File
@@ -0,0 +1,5 @@
---
"kilo-code": patch
---
Keep Agent Manager model, mode, and reasoning effort selections independent per session and preserve them after restart.
@@ -600,8 +600,9 @@ export class WorktreeStateManager {
let pruned = 0
for (const [id, s] of Object.entries(data.sessions ?? {})) {
const ref = s.worktreeId
const session: ManagedSession = { id, worktreeId: s.worktreeId, createdAt: s.createdAt }
if (ref === null) {
this.sessions.set(id, { id, ...s })
this.sessions.set(id, session)
continue
}
// Skip orphaned sessions referencing a deleted worktree.
@@ -609,7 +610,7 @@ export class WorktreeStateManager {
pruned++
continue
}
this.sessions.set(id, { id, ...s })
this.sessions.set(id, session)
}
for (const [id, sec] of Object.entries(data.sections ?? {})) {
this.sections.set(id, { id, ...sec })
@@ -0,0 +1,96 @@
import { describe, expect, it } from "bun:test"
import { initialMessage, initialVariant, seedInitialVariant } from "../../webview-ui/agent-manager/initial-message"
describe("Agent Manager initial message", () => {
it("forwards the selected variant to sendMessage", () => {
const msg = initialMessage({
type: "agentManager.sendInitialMessage",
sessionId: "session-a",
worktreeId: "wt-a",
text: "Fix it",
providerID: "anthropic",
modelID: "claude-sonnet-4",
agent: "code",
variant: "high",
})
expect(msg).toEqual({
type: "sendMessage",
text: "Fix it",
sessionID: "session-a",
providerID: "anthropic",
modelID: "claude-sonnet-4",
agent: "code",
variant: "high",
files: undefined,
})
})
it("does not create an empty sendMessage payload", () => {
expect(
initialMessage({
type: "agentManager.sendInitialMessage",
sessionId: "session-a",
worktreeId: "wt-a",
}),
).toBeUndefined()
})
it("builds the initial session variant state", () => {
const state = initialVariant(
{
type: "agentManager.sendInitialMessage",
sessionId: "session-a",
worktreeId: "wt-a",
providerID: "anthropic",
modelID: "claude-sonnet-4",
variant: "medium",
},
"code",
)
expect(state).toEqual({
sessionID: "session-a",
providerID: "anthropic",
modelID: "claude-sonnet-4",
agent: "code",
value: "medium",
})
})
it("does not build variant state without a complete model variant", () => {
expect(
initialVariant(
{
type: "agentManager.sendInitialMessage",
sessionId: "session-a",
worktreeId: "wt-a",
providerID: "anthropic",
modelID: "claude-sonnet-4",
},
"code",
),
).toBeUndefined()
})
it("seeds initial variant state into the session store", () => {
const calls: unknown[] = []
seedInitialVariant(
{
getSessionAgent: () => "code",
setSessionVariant: (...args) => calls.push(args),
},
{
type: "agentManager.sendInitialMessage",
sessionId: "session-a",
worktreeId: "wt-a",
providerID: "anthropic",
modelID: "claude-sonnet-4",
variant: "medium",
},
)
expect(calls).toEqual([["session-a", "anthropic", "claude-sonnet-4", "medium", "code"]])
})
})
@@ -47,7 +47,7 @@ const claude: ModelSelection = { providerID: "anthropic", modelID: "claude-sonne
const gpt: ModelSelection = { providerID: "openai", modelID: "gpt-4.1" }
describe("per-session model selection", () => {
it("selecting a model in session A writes per-mode globally", () => {
it("selecting a model in session A does not write per-mode globally", () => {
const store = emptyStore()
const e = env()
@@ -58,10 +58,9 @@ describe("per-session model selection", () => {
// Session A should see claude (via session override)
expect(getSessionModel(updated, e, "session-a", "code")).toEqual(claude)
// Session B (no override) inherits the per-mode global selection.
// This matches CLI behavior: per-mode model is global, not per-session.
// Session B (no override) keeps the default model.
const sessionB = getSessionModel(updated, e, "session-b", "code")
expect(sessionB).toEqual(claude)
expect(sessionB).toEqual(KILO_AUTO)
})
it("each session preserves its own model independently", () => {
@@ -150,15 +149,15 @@ describe("per-session model selection", () => {
})
describe("per-mode model memory", () => {
it("applyModel in a session writes to both sessionOverrides and modelSelections", () => {
it("applyModel in a session writes only to sessionOverrides", () => {
const store = emptyStore()
const result = applyModel(store, "code", claude, "session-a")
expect(result.sessionOverrides["session-a"]).toEqual(claude)
expect(result.modelSelections["code"]).toEqual(claude)
expect(result.modelSelections["code"]).toBeUndefined()
})
it("switching modes restores per-mode model after session override is cleared", () => {
it("switching modes falls back to default after session override is cleared", () => {
let store = emptyStore()
const e = env()
@@ -169,20 +168,19 @@ describe("per-mode model memory", () => {
// Simulate mode switch: clear session override (like selectAgent does)
const cleared = { ...store, sessionOverrides: {} }
// The global modelSelections["code"] still has claude
expect(getSelected(cleared, e, "session-a", "code")).toEqual(claude)
expect(getSelected(cleared, e, "session-a", "code")).toEqual(KILO_AUTO)
})
it("different modes remember their own model independently", () => {
let store = emptyStore()
const e = env()
// User picks claude for "code" in session A
let result = applyModel(store, "code", claude, "session-a")
// User picks claude for "code" globally
let result = applyModel(store, "code", claude, undefined)
store = { ...store, ...result }
// User switches to "ask" mode and picks gpt
result = applyModel(store, "ask", gpt, "session-a")
// User switches to "ask" mode and picks gpt globally
result = applyModel(store, "ask", gpt, undefined)
store = { ...store, ...result }
// Clear session overrides (simulating mode switch)
@@ -207,8 +205,8 @@ describe("per-mode model memory", () => {
// Session A sees gpt (its override), not the global claude
expect(getSelected(store, e, "session-a", "code")).toEqual(gpt)
// Global modelSelections was updated to gpt (last write wins)
expect(store.modelSelections["code"]).toEqual(gpt)
// Global modelSelections stays at the sidebar/default choice.
expect(store.modelSelections["code"]).toEqual(claude)
})
it("applyModel without session only writes to modelSelections, not sessionOverrides", () => {
@@ -0,0 +1,72 @@
import { describe, expect, it } from "bun:test"
import { resolveMessagePrefs } from "../../webview-ui/src/context/session-preferences"
import type { Message } from "../../webview-ui/src/types/messages"
function msg(input: Partial<Message>): Message {
return {
id: input.id ?? "msg",
sessionID: input.sessionID ?? "session-a",
role: input.role ?? "user",
createdAt: input.createdAt ?? "2026-01-01T00:00:00.000Z",
...input,
}
}
const agents = new Set(["code", "ask"])
describe("session preference recovery", () => {
it("recovers model, variant, and agent from the latest user message", () => {
const prefs = resolveMessagePrefs(
[
msg({
id: "old",
agent: "ask",
model: { providerID: "anthropic", modelID: "claude-sonnet-4", variant: "low" },
}),
msg({
id: "new",
agent: "code",
model: { providerID: "openai", modelID: "gpt-5.5", variant: "medium" },
}),
],
agents,
)
expect(prefs).toEqual({
agent: "code",
model: { providerID: "openai", modelID: "gpt-5.5" },
variant: "medium",
})
})
it("ignores assistant-only model data and invalid agents", () => {
const prefs = resolveMessagePrefs(
[
msg({
role: "assistant",
agent: "task",
model: { providerID: "openai", modelID: "gpt-5.5", variant: "high" },
}),
],
agents,
)
expect(prefs).toEqual({})
})
it("can recover the latest valid agent separately from the latest user model", () => {
const prefs = resolveMessagePrefs(
[
msg({ agent: "ask", model: { providerID: "anthropic", modelID: "claude-sonnet-4" } }),
msg({ role: "assistant", agent: "code" }),
],
agents,
)
expect(prefs).toEqual({
agent: "code",
model: { providerID: "anthropic", modelID: "claude-sonnet-4" },
variant: undefined,
})
})
})
@@ -0,0 +1,94 @@
import { describe, expect, it } from "bun:test"
import {
getVariant,
sessionVariantKeys,
sessionVariants,
transferVariants,
variantKey,
} from "../../webview-ui/src/context/session-variant-store"
import type { ModelSelection } from "../../webview-ui/src/types/messages"
const model: ModelSelection = { providerID: "anthropic", modelID: "claude-sonnet-4" }
const variants = ["low", "medium", "high"]
describe("per-session variant selection", () => {
it("keeps reasoning effort independent for each Agent Manager session", () => {
const store: Record<string, string> = {}
store[variantKey(model, "code", "session-a")] = "low"
store[variantKey(model, "code", "session-b")] = "high"
expect(getVariant(store, model, variants, "code", "session-a")).toBe("low")
expect(getVariant(store, model, variants, "code", "session-b")).toBe("high")
})
it("keeps reasoning effort independent for each pending local tab", () => {
const store: Record<string, string> = {}
store[variantKey(model, "code", "pending-local-1")] = "medium"
store[variantKey(model, "code", "pending-local-2")] = "high"
expect(getVariant(store, model, variants, "code", "pending-local-1")).toBe("medium")
expect(getVariant(store, model, variants, "code", "pending-local-2")).toBe("high")
})
it("keeps no-session reasoning effort independent per agent", () => {
const store: Record<string, string> = {}
store[variantKey(model, "code")] = "medium"
store[variantKey(model, "ask")] = "high"
expect(getVariant(store, model, variants, "code")).toBe("medium")
expect(getVariant(store, model, variants, "ask")).toBe("high")
})
it("carries the pre-submit agent variant into a newly created session", () => {
const store: Record<string, string> = {}
store[variantKey(model, "code")] = "medium"
expect(getVariant(store, model, variants, "code", "session-a")).toBe("medium")
})
it("prefers a session variant over the pre-submit agent variant", () => {
const store: Record<string, string> = {}
store[variantKey(model, "code")] = "medium"
store[variantKey(model, "code", "session-a")] = "high"
expect(getVariant(store, model, variants, "code", "session-a")).toBe("high")
})
it("falls back to the legacy provider/model variant key", () => {
const store: Record<string, string> = { "anthropic/claude-sonnet-4": "medium" }
expect(getVariant(store, model, variants, "code", "session-a")).toBe("medium")
})
it("transfers a pending local tab variant to the created session", () => {
const store: Record<string, string> = {}
store[variantKey(model, "code", "pending-local-1")] = "medium"
Object.assign(store, transferVariants(store, "pending-local-1", "session-a"))
expect(getVariant(store, model, variants, "code", "session-a")).toBe("medium")
})
it("extracts persisted session variant preferences", () => {
const store: Record<string, string> = {}
store[variantKey(model, "code", "session-a")] = "medium"
store[variantKey(model, "code", "session-b")] = "high"
expect(sessionVariants(store, "session-a")).toEqual({ "anthropic/claude-sonnet-4": "medium" })
})
it("finds only variant keys for the requested session", () => {
const store: Record<string, string> = {}
store[variantKey(model, "code", "pending-local-1")] = "medium"
store[variantKey(model, "code", "pending-local-2")] = "high"
expect(sessionVariantKeys(store, "pending-local-1")).toEqual(["session/pending-local-1/anthropic/claude-sonnet-4"])
})
})
@@ -77,6 +77,27 @@ describe("WorktreeStateManager", () => {
expect(s.worktreeId).toBeNull()
})
it("drops obsolete session prefs while loading state", async () => {
const file = path.join(root, ".kilo", "agent-manager.json")
fs.writeFileSync(
file,
JSON.stringify({
worktrees: {},
sessions: {
"local-1": {
worktreeId: null,
createdAt: new Date().toISOString(),
prefs: { agent: "code" },
},
},
}),
)
await manager.load()
expect(manager.getSession("local-1")).not.toHaveProperty("prefs")
})
it("filters sessions by worktreeId", () => {
const wt1 = manager.addWorktree({ branch: "a", path: "/tmp/a", parentBranch: "main" })
const wt2 = manager.addWorktree({ branch: "b", path: "/tmp/b", parentBranch: "main" })
@@ -114,6 +114,7 @@ import {
import { sectionAwareDetector } from "./section-dnd"
import { ConstrainDragXAxis } from "./constrain-drag-x"
import { mergeWorktreeDiffs } from "./diff-state"
import { initialMessage, seedInitialVariant } from "./initial-message"
import { createMarkdownRender } from "./review-preferences"
import "./agent-manager.css"
import "./agent-manager-review.css"
@@ -1407,18 +1408,12 @@ const AgentManagerContent: Component = () => {
if (ev.providerID && ev.modelID) {
session.setSessionModel(ev.sessionId, ev.providerID, ev.modelID)
}
seedInitialVariant(session, ev)
// Only send a message if there's text — otherwise just clear busy state
if (ev.text) {
vscode.postMessage({
type: "sendMessage",
text: ev.text,
sessionID: ev.sessionId,
providerID: ev.providerID,
modelID: ev.modelID,
agent: ev.agent,
files: ev.files,
})
const init = initialMessage(ev)
if (init) {
vscode.postMessage(init)
}
// Clear busy state — use worktreeId from the message directly
// to avoid race condition where managedSessions() hasn't updated yet
@@ -0,0 +1,37 @@
import type { AgentManagerSendInitialMessage, SendMessageRequest } from "../src/types/messages"
interface VariantSession {
getSessionAgent: (sessionID: string) => string
setSessionVariant: (sessionID: string, providerID: string, modelID: string, value: string, agent?: string) => void
}
export function initialMessage(ev: AgentManagerSendInitialMessage): SendMessageRequest | undefined {
if (!ev.text) return undefined
return {
type: "sendMessage",
text: ev.text,
sessionID: ev.sessionId,
providerID: ev.providerID,
modelID: ev.modelID,
agent: ev.agent,
variant: ev.variant,
files: ev.files,
}
}
export function initialVariant(ev: AgentManagerSendInitialMessage, agent: string) {
if (!ev.providerID || !ev.modelID || !ev.variant) return undefined
return {
sessionID: ev.sessionId,
providerID: ev.providerID,
modelID: ev.modelID,
agent: ev.agent ?? agent,
value: ev.variant,
}
}
export function seedInitialVariant(session: VariantSession, ev: AgentManagerSendInitialMessage) {
const state = initialVariant(ev, session.getSessionAgent(ev.sessionId))
if (!state) return
session.setSessionVariant(state.sessionID, state.providerID, state.modelID, state.value, state.agent)
}
@@ -82,7 +82,9 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
const mention = useFileMention(vscode, sid, hasGit)
const terminal = useTerminalContext(vscode)
const git = useGitChangesContext(vscode, ctx, hasGit)
const slash = useSlashCommand(vscode, () => (session.variantList().length > 0 ? new Set() : new Set(["variant"])))
const slash = useSlashCommand(vscode, () =>
session.variantList(sid()).length > 0 ? new Set() : new Set(["variant"]),
)
const imageAttach = useImageAttachments()
imageAttach.setFilePathDropHandler((paths) => {
const cwd = server.workspaceDirectory()
@@ -315,7 +317,7 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
const onCompact = () => {
if (session.status() === "busy") return
if (session.messages().length === 0) return
if (!session.selected()) return
if (!session.selected(sid())) return
session.compact()
}
window.addEventListener("compactSession", onCompact)
@@ -387,7 +389,7 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
if (message.type === "triggerTask") {
if (isDisabled()) return
const sel = session.selected()
const sel = session.selected(sid())
session.sendMessage(message.text, sel?.providerID, sel?.modelID, undefined, undefined, ctx())
}
@@ -679,9 +681,9 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
const mentionFiles = mention.parseFileAttachments(draft)
const imgFiles = imgs.map((img) => ({ mime: img.mime, url: img.dataUrl, filename: img.filename }))
const sel = session.selected()
const pendingId = props.pendingSessionID ?? session.draftSessionID()
const id = sid()
const sel = session.selected(id)
const terminalFile = await terminal.resolveAttachment(message, id).catch((err: Error) => {
showToast({ variant: "error", title: "Terminal context unavailable", description: err.message })
@@ -955,15 +957,15 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
</div>
<div class="prompt-input-hint">
<div class="prompt-input-hint-selectors">
<ModeSwitcher />
<ModelSelector />
<ThinkingSelector />
<Show when={session.hasModelOverride()}>
<ModeSwitcher sessionID={sid} />
<ModelSelector sessionID={sid} />
<ThinkingSelector sessionID={sid} />
<Show when={session.hasModelOverride(sid())}>
<Tooltip value={language.t("prompt.action.resetModel")} placement="top">
<Button
variant="ghost"
size="small"
onClick={() => session.clearModelOverride()}
onClick={() => session.clearModelOverride(sid())}
aria-label={language.t("prompt.action.resetModel")}
>
<svg width="12" height="12" viewBox="0 0 16 16" fill="currentColor">
@@ -7,7 +7,7 @@
* ModeSwitcher — thin wrapper wired to session context for chat usage.
*/
import { Component, createSignal, onCleanup, For, Show } from "solid-js"
import { type Accessor, Component, createSignal, onCleanup, For, Show } from "solid-js"
import { PopupSelector } from "./PopupSelector"
import { Button } from "@kilocode/kilo-ui/button"
import { useSession } from "../../context/session"
@@ -167,15 +167,20 @@ export const ModeSwitcherBase: Component<ModeSwitcherBaseProps> = (props) => {
// Chat-specific wrapper (backwards-compatible)
// ---------------------------------------------------------------------------
export const ModeSwitcher: Component = () => {
interface ModeSwitcherProps {
sessionID?: Accessor<string | undefined>
}
export const ModeSwitcher: Component<ModeSwitcherProps> = (props) => {
const session = useSession()
const id = () => props.sessionID?.()
return (
<ModeSwitcherBase
agents={session.agents()}
value={session.selectedAgent()}
value={session.selectedAgent(id())}
onSelect={(name) => {
session.selectAgent(name)
session.selectAgent(name, id())
requestAnimationFrame(() => window.dispatchEvent(new Event("focusPrompt")))
}}
/>
@@ -8,7 +8,7 @@
*/
import { createSignal, createMemo, createEffect, onCleanup, For, Show, createSelector, useContext } from "solid-js"
import type { Component } from "solid-js"
import type { Accessor, Component } from "solid-js"
import { PopupSelector } from "./PopupSelector"
import { Button } from "@kilocode/kilo-ui/button"
import { IconButton } from "@kilocode/kilo-ui/icon-button"
@@ -730,14 +730,19 @@ export const ModelSelectorBase: Component<ModelSelectorBaseProps> = (props) => {
// Chat-specific wrapper (backwards-compatible default export)
// ---------------------------------------------------------------------------
export const ModelSelector: Component = () => {
interface ModelSelectorProps {
sessionID?: Accessor<string | undefined>
}
export const ModelSelector: Component<ModelSelectorProps> = (props) => {
const session = useSession()
const id = () => props.sessionID?.()
return (
<ModelSelectorBase
value={session.selected()}
value={session.selected(id())}
onSelect={(providerID, modelID) => {
session.selectModel(providerID, modelID)
session.selectModel(providerID, modelID, id())
}}
onPick={() => {
requestAnimationFrame(() => window.dispatchEvent(new CustomEvent("focusPrompt", { detail: { restore: true } })))
@@ -7,7 +7,7 @@
* ThinkingSelector — thin wrapper wired to session context for chat usage.
*/
import { Component, createSignal, For, onCleanup, Show } from "solid-js"
import { type Accessor, Component, createSignal, For, onCleanup, Show } from "solid-js"
import { PopupSelector } from "./PopupSelector"
import { Button } from "@kilocode/kilo-ui/button"
import { useSession } from "../../context/session"
@@ -157,14 +157,19 @@ export const ThinkingSelectorBase: Component<ThinkingSelectorBaseProps> = (props
// Chat-specific wrapper (backwards-compatible)
// ---------------------------------------------------------------------------
export const ThinkingSelector: Component = () => {
interface ThinkingSelectorProps {
sessionID?: Accessor<string | undefined>
}
export const ThinkingSelector: Component<ThinkingSelectorProps> = (props) => {
const session = useSession()
const id = () => props.sessionID?.()
return (
<ThinkingSelectorBase
variants={session.variantList()}
value={session.currentVariant()}
onSelect={(value) => session.selectVariant(value)}
variants={session.variantList(id())}
value={session.currentVariant(id())}
onSelect={(value) => session.selectVariant(value, id())}
/>
)
}
@@ -86,10 +86,9 @@ export interface ApplyResult {
/**
* Apply a user-initiated model selection.
*
* Always writes to the global modelSelections map so switching modes
* restores the last-used model (mirrors CLI TUI's model.json behavior).
* When a session is active, also writes to the per-session override so
* the active session uses the chosen model immediately.
* Session-scoped selections write only to the per-session override.
* No-session selections write to the global modelSelections map so sidebar
* default picks still mirror CLI TUI's model.json behavior.
*/
export function applyModel(
store: ModelStore,
@@ -97,7 +96,9 @@ export function applyModel(
selection: ModelSelection,
sessionID: string | undefined,
): ApplyResult {
const modelSelections = { ...store.modelSelections, [agentName]: selection }
const modelSelections = sessionID
? { ...store.modelSelections }
: { ...store.modelSelections, [agentName]: selection }
const sessionOverrides = { ...store.sessionOverrides }
if (sessionID) {
@@ -0,0 +1,25 @@
import type { Message, ModelSelection } from "../types/messages"
export interface MessagePrefs {
agent?: string
model?: ModelSelection
variant?: string
}
export function resolveMessagePrefs(messages: Message[], names: Set<string>): MessagePrefs {
const prefs: MessagePrefs = {}
for (let i = messages.length - 1; i >= 0; i--) {
const msg = messages[i]
if (!msg) continue
if (!prefs.agent) {
const agent = msg.agent?.trim()
if (agent && names.has(agent)) prefs.agent = agent
}
if (!prefs.model && msg.role === "user" && msg.model?.providerID && msg.model.modelID) {
prefs.model = { providerID: msg.model.providerID, modelID: msg.model.modelID }
prefs.variant = msg.model.variant
}
if (prefs.agent && prefs.model) break
}
return prefs
}
@@ -0,0 +1,48 @@
import type { ModelSelection } from "../types/messages"
export function legacyVariantKey(sel: ModelSelection) {
return `${sel.providerID}/${sel.modelID}`
}
export function variantKey(sel: ModelSelection, agent: string, session?: string) {
const base = legacyVariantKey(sel)
if (session) return `session/${session}/${base}`
return `agent/${agent}/${base}`
}
export function getVariant(
store: Record<string, string>,
sel: ModelSelection,
variants: string[],
agent: string,
session?: string,
) {
if (variants.length === 0) return undefined
const key = variantKey(sel, agent, session)
const fallback = session ? store[variantKey(sel, agent)] : undefined
const stored = store[key] ?? fallback ?? store[legacyVariantKey(sel)]
return stored && variants.includes(stored) ? stored : variants[0]
}
export function transferVariants(store: Record<string, string>, from: string, to: string) {
const prefix = `session/${from}/`
return Object.fromEntries(
Object.entries(store)
.filter(([key]) => key.startsWith(prefix))
.map(([key, value]) => [`session/${to}/${key.slice(prefix.length)}`, value]),
)
}
export function sessionVariantKeys(store: Record<string, string>, session: string) {
const prefix = `session/${session}/`
return Object.keys(store).filter((key) => key.startsWith(prefix))
}
export function sessionVariants(store: Record<string, string>, session: string) {
const prefix = `session/${session}/`
return Object.fromEntries(
Object.entries(store)
.filter(([key]) => key.startsWith(prefix))
.map(([key, value]) => [key.slice(prefix.length), value]),
)
}
@@ -45,9 +45,10 @@ import {
} from "./session-utils"
import { Identifier } from "../utils/id"
import { resolveModelSelection } from "./model-selection"
import { resolveSessionAgent } from "./session-agent"
import { resolveMessagePrefs } from "./session-preferences"
import { errorIDs } from "./session-errors"
import { PartStash } from "./part-stash"
import { getVariant, sessionVariantKeys, transferVariants, variantKey } from "./session-variant-store"
import { KILO_AUTO, parseModelString } from "../../../src/shared/provider-model"
const RECENT_LIMIT = 5
@@ -80,7 +81,7 @@ interface SessionStore {
modelSelections: Record<string, ModelSelection | null> // agentName -> model (global, extension-lifetime)
sessionOverrides: Record<string, ModelSelection> // sessionID -> per-session model override (compare mode)
agentSelections: Record<string, string> // sessionID -> agent name
variantSelections: Record<string, string> // "providerID/modelID" -> variant name
variantSelections: Record<string, string> // session/agent scoped variant key -> variant name
recentModels: ModelSelection[]
favoriteModels: ModelSelection[]
}
@@ -150,10 +151,10 @@ interface SessionContextValue {
scopedSuggestions: (sessionID: string | undefined) => SuggestionRequest[]
// Model selection (global, extension-lifetime)
selected: Accessor<ModelSelection | null>
selectModel: (providerID: string, modelID: string) => void
hasModelOverride: Accessor<boolean>
clearModelOverride: () => void
selected: (sessionID?: string) => ModelSelection | null
selectModel: (providerID: string, modelID: string, sessionID?: string) => void
hasModelOverride: (sessionID?: string) => boolean
clearModelOverride: (sessionID?: string) => void
// Cost and context usage for the current session
costBreakdown: Accessor<Array<{ label: string; cost: number }>>
@@ -177,17 +178,18 @@ interface SessionContextValue {
disconnectMcp: (name: string) => void
authenticateMcp: (name: string) => void
refreshMcpStatus: () => void
selectedAgent: Accessor<string>
selectAgent: (name: string) => void
selectedAgent: (sessionID?: string) => string
selectAgent: (name: string, sessionID?: string) => void
getSessionAgent: (sessionID: string) => string
getSessionModel: (sessionID: string) => ModelSelection | null
setSessionModel: (sessionID: string, providerID: string, modelID: string) => void
setSessionAgent: (sessionID: string, name: string) => void
setSessionVariant: (sessionID: string, providerID: string, modelID: string, value: string, agent?: string) => void
// Thinking variant for the selected model
variantList: () => string[]
currentVariant: () => string | undefined
selectVariant: (value: string) => void
variantList: (sessionID?: string) => string[]
currentVariant: (sessionID?: string) => string | undefined
selectVariant: (value: string, sessionID?: string) => void
// Model favorites
favoriteModels: Accessor<ModelSelection[]>
@@ -407,6 +409,11 @@ export const SessionProvider: ParentComponent = (props) => {
return pendingAgentSelection() ?? defaultAgent()
})
function agentForScope(sessionID?: string) {
if (sessionID) return store.agentSelections[sessionID] ?? defaultAgent()
return selectedAgentName()
}
const agentNames = createMemo(() => new Set(agents().map((agent) => agent.name)))
/** Per-mode model from config (e.g. config.agent.code.model). */
@@ -440,10 +447,7 @@ export const SessionProvider: ParentComponent = (props) => {
setStore("modelSelections", agentName, sel)
})
// Global model selection per agent/mode
// Precedence: per-session override > user override > per-mode config > global config model > VS Code default > kilo-auto/free
// Each candidate is validated against the provider catalog; invalid models fall through.
const selected = createMemo<ModelSelection | null>(() => {
const currentSelected = createMemo<ModelSelection | null>(() => {
const sid = currentSessionID()
if (sid) {
const session = store.sessionOverrides[sid]
@@ -453,6 +457,15 @@ export const SessionProvider: ParentComponent = (props) => {
return resolveModel(agentName, store.modelSelections[agentName])
})
// Precedence: scoped override > per-agent global/default > config/default.
function selected(sessionID?: string): ModelSelection | null {
if (!sessionID) return currentSelected()
const session = store.sessionOverrides[sessionID]
if (session) return session
const agentName = agentForScope(sessionID)
return resolveModel(agentName, store.modelSelections[agentName])
}
function pushRecent(selection: ModelSelection) {
const key = `${selection.providerID}/${selection.modelID}`
const filtered = store.recentModels.filter((r) => `${r.providerID}/${r.modelID}` !== key)
@@ -461,8 +474,12 @@ export const SessionProvider: ParentComponent = (props) => {
vscode.postMessage({ type: "persistRecents", recents: updated })
}
function applyModel(agentName: string, selection: ModelSelection) {
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 }))
@@ -474,20 +491,21 @@ export const SessionProvider: ParentComponent = (props) => {
providerID: selection.providerID,
modelID: selection.modelID,
})
const sid = currentSessionID()
if (sid) {
setStore("sessionOverrides", sid, selection)
}
}
function selectModel(providerID: string, modelID: string) {
applyModel(selectedAgentName(), { providerID, modelID })
const sid = currentSessionID()
function selectModel(providerID: string, modelID: string, sessionID?: string) {
const sid = sessionID ?? currentSessionID()
applyModel(agentForScope(sid), { providerID, modelID }, sid)
if (sid) {
hideErrors(sid)
}
}
function promptAgent(sessionID?: string) {
const name = agentForScope(sessionID)
return name !== defaultAgent() ? name : undefined
}
function hideErrors(sid: string) {
const ids = errorIDs(store.messages[sid] ?? [])
if (ids.length === 0) return
@@ -508,22 +526,32 @@ export const SessionProvider: ParentComponent = (props) => {
})
}
/** The config/default model for the current mode (what settings says). */
const configModel = createMemo<ModelSelection | null>(() => {
const agentName = selectedAgentName()
function configModel(sessionID?: string): ModelSelection | null {
const agentName = agentForScope(sessionID)
return resolveModel(agentName)
})
}
/** True when the active model differs from what the config dictates. */
const hasModelOverride = createMemo<boolean>(() => {
const sel = selected()
const cfg = configModel()
function hasModelOverride(sessionID?: string) {
const sel = selected(sessionID)
const cfg = configModel(sessionID)
if (!sel || !cfg) return false
return sel.providerID !== cfg.providerID || sel.modelID !== cfg.modelID
})
}
/** Clear the per-mode model override, falling back to config default. */
function clearModelOverride() {
function clearModelOverride(sessionID?: string) {
const sid = sessionID ?? currentSessionID()
if (sid) {
setStore(
"sessionOverrides",
produce((overrides) => {
delete overrides[sid]
}),
)
hideErrors(sid)
return
}
const agentName = selectedAgentName()
setUserSetAgents((prev) => {
const next = { ...prev }
@@ -538,17 +566,6 @@ export const SessionProvider: ParentComponent = (props) => {
)
// Clear from model.json via extension host
vscode.postMessage({ type: "clearModelSelection", agent: agentName })
// Also clear per-session override so the session falls back to config default
const sid = currentSessionID()
if (sid) {
setStore(
"sessionOverrides",
produce((overrides) => {
delete overrides[sid]
}),
)
hideErrors(sid)
}
}
// Handle agentsLoaded immediately (not in onMount) so we never miss
@@ -585,9 +602,7 @@ export const SessionProvider: ParentComponent = (props) => {
// backfilled now that we know the valid agent names.
batch(() => {
for (const [sid, msgs] of Object.entries(store.messages)) {
if (store.agentSelections[sid]) continue
const agent = resolveSessionAgent(msgs, names)
if (agent) setStore("agentSelections", sid, agent)
recoverPrefs(sid, msgs, names)
}
})
})
@@ -662,38 +677,37 @@ export const SessionProvider: ParentComponent = (props) => {
clearTimeout(fallback)
})
// Variant (thinking effort) selection — keyed by "providerID/modelID"
const variantKey = (sel: ModelSelection) => `${sel.providerID}/${sel.modelID}`
const variantList = () => {
const sel = selected()
const variantList = (sessionID?: string) => {
const sel = selected(sessionID)
if (!sel) return []
const model = provider.findModel(sel)
if (!model?.variants) return []
return Object.keys(model.variants)
}
const currentVariant = () => {
const sel = selected()
const currentVariant = (sessionID?: string) => {
const sid = sessionID ?? currentSessionID()
const sel = selected(sid)
if (!sel) return undefined
const list = variantList()
const list = variantList(sid)
if (list.length === 0) return undefined
const stored = store.variantSelections[variantKey(sel)]
return stored && list.includes(stored) ? stored : list[0]
return getVariant(store.variantSelections, sel, list, agentForScope(sid), sid)
}
const selectVariant = (value: string) => {
const sel = selected()
const selectVariant = (value: string, sessionID?: string) => {
const sid = sessionID ?? currentSessionID()
const sel = selected(sid)
if (!sel) return
const key = variantKey(sel)
const key = variantKey(sel, agentForScope(sid), sid)
setStore("variantSelections", key, value)
vscode.postMessage({ type: "persistVariant", key, value })
if (!sid) vscode.postMessage({ type: "persistVariant", key, value })
}
// Load persisted variants from extension globalState
const unsubVariants = vscode.onMessage((message: ExtensionMessage) => {
if (message.type !== "variantsLoaded") return
for (const [k, v] of Object.entries(message.variants)) {
if (k.startsWith("session/")) continue
setStore("variantSelections", k, v)
}
})
@@ -909,9 +923,35 @@ export const SessionProvider: ParentComponent = (props) => {
setStore("messages", session.id, [])
}
// Transfer pending agent selection to the new session
const pendingAgent = pendingAgentSelection()
if (pendingAgent && !store.agentSelections[session.id]) {
const pendingAgent = draftID ? store.agentSelections[draftID] : pendingAgentSelection()
const pendingModel = draftID ? store.sessionOverrides[draftID] : undefined
if (draftID) {
const entries = transferVariants(store.variantSelections, draftID, session.id)
for (const [key, value] of Object.entries(entries)) {
setStore("variantSelections", key, value)
vscode.postMessage({ type: "persistVariant", key, value })
}
if (pendingAgent) setStore("agentSelections", session.id, pendingAgent)
if (pendingModel) setStore("sessionOverrides", session.id, pendingModel)
setStore(
"agentSelections",
produce((agents) => {
delete agents[draftID]
}),
)
setStore(
"sessionOverrides",
produce((models) => {
delete models[draftID]
}),
)
setStore(
"variantSelections",
produce((variants) => {
for (const key of sessionVariantKeys(variants, draftID)) delete variants[key]
}),
)
} else if (pendingAgent && !store.agentSelections[session.id]) {
setStore("agentSelections", session.id, pendingAgent)
setPendingAgentSelection(null)
}
@@ -951,6 +991,21 @@ export const SessionProvider: ParentComponent = (props) => {
})
}
function recoverPrefs(sessionID: string, messages: Message[], names = agentNames()) {
const prefs = resolveMessagePrefs(messages, names)
if (prefs.agent && !store.agentSelections[sessionID]) {
setStore("agentSelections", sessionID, prefs.agent)
}
if (prefs.model && !store.sessionOverrides[sessionID]) {
setStore("sessionOverrides", sessionID, prefs.model)
}
if (prefs.model && prefs.variant) {
const agent = prefs.agent ?? store.agentSelections[sessionID] ?? defaultAgent()
const key = variantKey(prefs.model, agent, sessionID)
if (!store.variantSelections[key]) setStore("variantSelections", key, prefs.variant)
}
}
function withPending(sessionID: string, messages: Message[]) {
const pending = pendingOptimistic.get(sessionID)
if (!pending || pending.size === 0) return messages
@@ -1047,10 +1102,7 @@ export const SessionProvider: ParentComponent = (props) => {
})
}
const agent = resolveSessionAgent(merged, agentNames())
if (agent) {
setStore("agentSelections", sessionID, agent)
}
recoverPrefs(sessionID, merged)
})
if (reset) requestAnimationFrame(() => patchPage(sessionID, { lastMutation: undefined }))
}
@@ -1088,13 +1140,7 @@ export const SessionProvider: ParentComponent = (props) => {
})
patchPage(message.sessionID, { initialLoaded: true, lastMutation: exists ? "update" : "append" })
// Sync mode picker from any message role (user or assistant).
// agentNames() already excludes subagent/hidden agents, so subtask
// assistant messages (e.g. "task" agent) are silently ignored.
const agent = message.agent?.trim()
if (agent && agentNames().has(agent)) {
setStore("agentSelections", message.sessionID, agent)
}
recoverPrefs(message.sessionID, [message])
if (message.parts && message.parts.length > 0) {
stash.remove(message.id)
@@ -1634,8 +1680,8 @@ export const SessionProvider: ParentComponent = (props) => {
}
// Actions
function selectAgent(name: string) {
const id = currentSessionID()
function selectAgent(name: string, sessionID?: string) {
const id = sessionID ?? currentSessionID()
if (id) {
setStore("agentSelections", id, name)
// Clear per-session model override so the new mode's configured/default
@@ -1709,7 +1755,8 @@ export const SessionProvider: ParentComponent = (props) => {
const preview = cloudPreviewId()
if (preview) {
const agent = selectedAgentName() !== defaultAgent() ? selectedAgentName() : undefined
const scope = draftID ?? currentSessionID()
const agent = promptAgent(scope)
vscode.postMessage({
type: "importAndSend",
cloudSessionId: preview,
@@ -1718,7 +1765,7 @@ export const SessionProvider: ParentComponent = (props) => {
providerID,
modelID,
agent,
variant: currentVariant(),
variant: currentVariant(scope),
files,
})
return
@@ -1732,7 +1779,8 @@ export const SessionProvider: ParentComponent = (props) => {
}
if (sid) addOptimistic(sid, messageID, text, files)
const agent = selectedAgentName() !== defaultAgent() ? selectedAgentName() : undefined
const scope = draftID ?? sid
const agent = promptAgent(scope)
vscode.postMessage({
type: "sendMessage",
@@ -1743,7 +1791,7 @@ export const SessionProvider: ParentComponent = (props) => {
providerID,
modelID,
agent,
variant: currentVariant(),
variant: currentVariant(scope),
files,
agentManagerContext: context,
})
@@ -1766,7 +1814,8 @@ export const SessionProvider: ParentComponent = (props) => {
// Cloud previews need import-then-command; post importAndSend with command metadata
const preview = cloudPreviewId()
if (preview) {
const agent = selectedAgentName() !== defaultAgent() ? selectedAgentName() : undefined
const scope = draftID ?? currentSessionID()
const agent = promptAgent(scope)
vscode.postMessage({
type: "importAndSend",
cloudSessionId: preview,
@@ -1775,7 +1824,7 @@ export const SessionProvider: ParentComponent = (props) => {
providerID,
modelID,
agent,
variant: currentVariant(),
variant: currentVariant(scope),
files,
command,
commandArgs: args,
@@ -1793,7 +1842,8 @@ export const SessionProvider: ParentComponent = (props) => {
if (sid) addOptimistic(sid, messageID, `/${command} ${args}`.trim(), files)
const agent = selectedAgentName() !== defaultAgent() ? selectedAgentName() : undefined
const scope = draftID ?? sid
const agent = promptAgent(scope)
vscode.postMessage({
type: "sendCommand",
@@ -1805,7 +1855,7 @@ export const SessionProvider: ParentComponent = (props) => {
providerID,
modelID,
agent,
variant: currentVariant(),
variant: currentVariant(scope),
files,
agentManagerContext: context,
})
@@ -2248,7 +2298,7 @@ export const SessionProvider: ParentComponent = (props) => {
disconnectMcp,
authenticateMcp,
refreshMcpStatus,
selectedAgent: selectedAgentName,
selectedAgent: agentForScope,
selectAgent,
getSessionAgent: (sessionID: string) => store.agentSelections[sessionID] ?? defaultAgent(),
getSessionModel: (sessionID: string) => {
@@ -2264,11 +2314,17 @@ export const SessionProvider: ParentComponent = (props) => {
// agent may not yet be assigned (sendInitialMessage calls setSessionModel
// before setSessionAgent), so the write would land on defaultAgent() and
// corrupt the default mode's model for later sessions.
setStore("sessionOverrides", sessionID, { providerID, modelID })
const model = { providerID, modelID }
setStore("sessionOverrides", sessionID, model)
},
setSessionAgent: (sessionID: string, name: string) => {
setStore("agentSelections", sessionID, name)
},
setSessionVariant: (sessionID: string, providerID: string, modelID: string, value: string, agent?: string) => {
const name = agent ?? store.agentSelections[sessionID] ?? defaultAgent()
const key = variantKey({ providerID, modelID }, name, sessionID)
setStore("variantSelections", key, value)
},
allMessages,
allParts,
allStatusMap,
@@ -214,6 +214,7 @@ export function mockSessionValue(overrides?: {
getSessionModel: () => ({ providerID: "kilo", modelID: "anthropic/claude-sonnet-4-6" }),
setSessionModel: noop,
setSessionAgent: noop,
setSessionVariant: noop,
revert: () => undefined,
revertedCount: () => 0,
summary: () => undefined,
@@ -81,6 +81,7 @@ const WithSessions: ParentComponent<{ sessions?: typeof mockSessions }> = (props
getSessionModel: () => ({ providerID: "kilo", modelID: "anthropic/claude-sonnet-4-6" }),
setSessionModel: noop,
setSessionAgent: noop,
setSessionVariant: noop,
variantList: () => [],
currentVariant: () => undefined,
selectVariant: noop,
@@ -686,6 +686,7 @@ export interface AgentManagerSendInitialMessage {
providerID?: string
modelID?: string
agent?: string
variant?: string
files?: Array<{ mime: string; url: string }>
}
@@ -10,7 +10,7 @@ export interface Message {
createdAt: string
time?: { created: number; completed?: number }
agent?: string
model?: { providerID: string; modelID: string }
model?: { providerID: string; modelID: string; variant?: string }
providerID?: string
modelID?: string
mode?: string