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

import { describe, expect, it } from "bun:test"
import {
  type ModelStore,
  type ResolveEnv,
  applyModel,
  getAgentModel,
  getSessionModel,
  getSelected,
} from "../../webview-ui/src/context/session-model-store"
import type { ModelSelection, Provider } from "../../webview-ui/src/types/messages"

function makeProvider(id: string, models: string[]): Provider {
  const result: Provider = { id, name: id, models: {} }
  for (const m of models) {
    result.models[m] = { id: m, name: m }
  }
  return result
}

const KILO_AUTO: ModelSelection = { providerID: "kilo", modelID: "kilo-auto/free" }

const providers: Record<string, Provider> = {
  kilo: makeProvider("kilo", ["kilo-auto/free"]),
  anthropic: makeProvider("anthropic", ["claude-sonnet-4"]),
  openai: makeProvider("openai", ["gpt-4.1"]),
}

function env(): ResolveEnv {
  return {
    providers,
    connected: ["kilo", "anthropic", "openai"],
    fallback: KILO_AUTO,
    getModeModel: () => null,
    getGlobalModel: () => null,
  }
}

function emptyStore(): ModelStore {
  return {
    modelSelections: {},
    sessionOverrides: {},
    agentSelections: {},
    recentModels: [],
  }
}

const claude: ModelSelection = { providerID: "anthropic", modelID: "claude-sonnet-4" }
const gpt: ModelSelection = { providerID: "openai", modelID: "gpt-4.1" }

describe("per-session model selection", () => {
  it("selecting a model in session A does not write per-mode globally", () => {
    const store = emptyStore()
    const e = env()

    // User picks claude in session A
    const after = applyModel(store, "code", claude, "session-a")
    const updated: ModelStore = { ...store, ...after }

    // Session A should see claude (via session override)
    expect(getSessionModel(updated, e, "session-a", "code")).toEqual(claude)

    // Session B (no override) keeps the default model.
    const sessionB = getSessionModel(updated, e, "session-b", "code")
    expect(sessionB).toEqual(KILO_AUTO)
  })

  it("each session preserves its own model independently", () => {
    let store = emptyStore()
    const e = env()

    // User picks claude in session A
    const a = applyModel(store, "code", claude, "session-a")
    store = { ...store, ...a }

    // User picks gpt in session B
    const b = applyModel(store, "code", gpt, "session-b")
    store = { ...store, ...b }

    // Both sessions should keep their own model
    expect(getSessionModel(store, e, "session-a", "code")).toEqual(claude)
    expect(getSessionModel(store, e, "session-b", "code")).toEqual(gpt)
  })

  it("getSelected returns per-session override when session is active", () => {
    let store = emptyStore()
    const e = env()

    const a = applyModel(store, "code", claude, "session-a")
    store = { ...store, ...a }

    expect(getSelected(store, e, "session-a", "code")).toEqual(claude)
  })

  it("getSelected returns global model when no session is active", () => {
    let store = emptyStore()
    const e = env()

    // Sidebar mode (no session) — writes globally
    const result = applyModel(store, "code", claude, undefined)
    store = { ...store, ...result }

    expect(getSelected(store, e, undefined, "code")).toEqual(claude)
  })

  it("sidebar model selection writes globally and is visible to new sessions without overrides", () => {
    let store = emptyStore()
    const e = env()

    // User picks claude in sidebar (no session)
    const result = applyModel(store, "code", claude, undefined)
    store = { ...store, ...result }

    // A new session without an override should see the global model
    expect(getSessionModel(store, e, "session-new", "code")).toEqual(claude)
  })

  it("setSessionModel (compare mode) only writes per-session override", () => {
    const store = emptyStore()

    // Simulate setSessionModel — writes only to sessionOverrides
    store.sessionOverrides["session-a"] = claude
    store.sessionOverrides["session-b"] = gpt

    const e = env()
    expect(getSessionModel(store, e, "session-a", "code")).toEqual(claude)
    expect(getSessionModel(store, e, "session-b", "code")).toEqual(gpt)
  })

  it("switching sessions preserves model selection after multiple changes", () => {
    let store = emptyStore()
    const e = env()

    // Simulate: user in session A picks claude
    let result = applyModel(store, "code", claude, "session-a")
    store = { ...store, ...result }

    // Switch to session B — picks gpt
    result = applyModel(store, "code", gpt, "session-b")
    store = { ...store, ...result }

    // Switch back to session A — picks gpt this time
    result = applyModel(store, "code", gpt, "session-a")
    store = { ...store, ...result }

    // Switch back to session B — should still have gpt
    expect(getSessionModel(store, e, "session-b", "code")).toEqual(gpt)
    // Session A was updated to gpt
    expect(getSessionModel(store, e, "session-a", "code")).toEqual(gpt)
  })
})

describe("per-mode model memory", () => {
  it("uses remembered model selections for modes without configured models", () => {
    const store = { ...emptyStore(), modelSelections: { ask: gpt } }

    expect(getAgentModel(store, env(), "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", () => {
    const store = emptyStore()
    const result = applyModel(store, "code", claude, "session-a")

    expect(result.sessionOverrides["session-a"]).toEqual(claude)
    expect(result.modelSelections["code"]).toBeUndefined()
  })

  it("switching modes falls back to default after session override is cleared", () => {
    let store = emptyStore()
    const e = env()

    // User picks claude for "code" mode in session A
    const result = applyModel(store, "code", claude, "session-a")
    store = { ...store, ...result }

    // Simulate mode switch: clear session override (like selectAgent does)
    const cleared = { ...store, sessionOverrides: {} }

    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" globally
    let result = applyModel(store, "code", claude, undefined)
    store = { ...store, ...result }

    // User switches to "ask" mode and picks gpt globally
    result = applyModel(store, "ask", gpt, undefined)
    store = { ...store, ...result }

    // Clear session overrides (simulating mode switch)
    const cleared: ModelStore = { ...store, sessionOverrides: {} }

    // Each mode should have its own saved model
    expect(getSelected(cleared, e, undefined, "code")).toEqual(claude)
    expect(getSelected(cleared, e, undefined, "ask")).toEqual(gpt)
  })

  it("per-session override still takes priority over global modelSelections", () => {
    let store = emptyStore()
    const e = env()

    // User picks claude globally for "code"
    let result = applyModel(store, "code", claude, undefined)
    store = { ...store, ...result }

    // Session A overrides with gpt
    result = applyModel(store, "code", gpt, "session-a")
    store = { ...store, ...result }

    // Session A sees gpt (its override), not the global claude
    expect(getSelected(store, e, "session-a", "code")).toEqual(gpt)
    // Global modelSelections stays at the sidebar/default choice.
    expect(store.modelSelections["code"]).toEqual(claude)
  })

  it("applyModel without session only writes to modelSelections, not sessionOverrides", () => {
    const store = emptyStore()
    const result = applyModel(store, "code", claude, undefined)

    expect(result.modelSelections["code"]).toEqual(claude)
    expect(Object.keys(result.sessionOverrides)).toHaveLength(0)
  })

  it("clearing both session override and per-mode selection falls back to config default", () => {
    let store = emptyStore()
    // Simulate mode model set in config (getModeModel returns claude).
    const configured: ResolveEnv = {
      ...env(),
      getModeModel: (name) => (name === "code" ? claude : null),
    }

    // User picked gpt globally for "code"
    let result = applyModel(store, "code", gpt, undefined)
    store = { ...store, ...result }

    // User then overrode the session with claude
    result = applyModel(store, "code", claude, "session-a")
    store = { ...store, ...result }

    // Simulate clearModelOverride: clear both session override and per-mode selection
    const reset: ModelStore = {
      ...store,
      sessionOverrides: {},
      modelSelections: { ...store.modelSelections, code: null },
    }

    // Should fall through to the configured per-mode model (claude from config)
    expect(getSelected(reset, configured, "session-a", "code")).toEqual(claude)
  })

  it("clearing only session override but not per-mode selection leaves persisted pick visible", () => {
    let store = emptyStore()
    const e = env()

    // User picked claude globally for "code"
    const result = applyModel(store, "code", claude, undefined)
    store = { ...store, ...result }

    // User then overrode the session with gpt
    const r2 = applyModel(store, "code", gpt, "session-a")
    store = { ...store, ...r2 }

    // Simulate OLD behaviour: only clear session override, leave modelSelections intact
    const partial: ModelStore = { ...store, sessionOverrides: {} }

    // The persisted per-mode selection (claude) is still returned — this is
    // why the reset appeared to "do nothing" when config also resolved to claude.
    expect(getSelected(partial, e, "session-a", "code")).toEqual(claude)
  })

  it("switching from plan to implementation uses implementation config after clearing stale memory", () => {
    let store = emptyStore()
    const configured: ResolveEnv = {
      ...env(),
      getModeModel: (name) => (name === "code" ? gpt : name === "plan" ? claude : null),
    }

    // Old manual memory says implementation/code should use claude.
    let result = applyModel(store, "code", claude, undefined)
    store = { ...store, ...result }

    // Current plan session is using its own model.
    result = applyModel(store, "plan", claude, "session-a")
    store = { ...store, ...result, agentSelections: { "session-a": "plan" } }

    const switched: ModelStore = {
      ...store,
      agentSelections: { "session-a": "code" },
      sessionOverrides: {},
      modelSelections: { ...store.modelSelections, code: null },
    }

    expect(getSelected(switched, configured, "session-a", "code")).toEqual(gpt)
  })
})