From b06b42d2e3c15f27fe276c6c6ac0df61a8acaece Mon Sep 17 00:00:00 2001 From: Josh Lambert Date: Tue, 2 Jun 2026 12:15:02 -0400 Subject: [PATCH 1/6] feat(kilo-vscode): support promoted model deep links --- packages/kilo-vscode/package.json | 3 +- packages/kilo-vscode/src/KiloProvider.ts | 17 ++++++ packages/kilo-vscode/src/extension.ts | 29 ++++++---- .../src/kilo-provider/model-state.ts | 32 ++++++----- packages/kilo-vscode/webview-ui/src/App.tsx | 16 ++++++ .../webview-ui/src/context/session.tsx | 54 +++++++++++++++---- .../src/types/messages/extension-messages.ts | 6 +++ 7 files changed, 124 insertions(+), 33 deletions(-) diff --git a/packages/kilo-vscode/package.json b/packages/kilo-vscode/package.json index 36e2d728f1c..96fc4b60845 100644 --- a/packages/kilo-vscode/package.json +++ b/packages/kilo-vscode/package.json @@ -52,7 +52,8 @@ "zoo code" ], "activationEvents": [ - "onStartupFinished" + "onStartupFinished", + "onUri" ], "main": "./dist/extension.js", "contributes": { diff --git a/packages/kilo-vscode/src/KiloProvider.ts b/packages/kilo-vscode/src/KiloProvider.ts index 89b8c8d254b..e76acd6be32 100644 --- a/packages/kilo-vscode/src/KiloProvider.ts +++ b/packages/kilo-vscode/src/KiloProvider.ts @@ -200,6 +200,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper private configWarningsShown = false /** Cached notificationsLoaded payload */ private cachedNotificationsMessage: unknown = null + private pendingKiloModelID: string | null = null private pendingReviewComments: { comments: unknown[]; autoSend: boolean }[] = [] private readyResolvers: (() => void)[] = [] private promptRecoveryQueued = false @@ -595,6 +596,11 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper this.postMessage({ type: "openCloudSession", sessionId }) } + public selectKiloModel(modelID: string): void { + this.pendingKiloModelID = modelID + this.flushPendingKiloModel() + } + public setContinueInWorktreeHandler( handler: (sessionId: string, progress: (status: string, detail?: string, error?: string) => void) => Promise, ): void { @@ -668,6 +674,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper this.isWebviewReady = true this.visibleTaskStreams.clear() await this.syncWebviewState("webviewReady") + this.flushPendingKiloModel() this.flushPendingReviewComments() this.recoverPendingPrompts() this.readyResolvers.splice(0).forEach((r) => r()) @@ -1195,6 +1202,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper // Connect the shared service (no-op if already connected) await this.connectionService.connect(workspaceDir) + this.flushPendingKiloModel() // Subscribe to SSE events for this webview (filtered by tracked sessions) this.unsubscribeEvent = this.connectionService.onEventFiltered( @@ -1229,6 +1237,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper this.postConnectionState(error) if (state === "connected") { + this.flushPendingKiloModel() // Fire config warnings independently so a failure in the // sequential await chain doesn't prevent warnings from being shown void this.checkConfigWarnings("state") @@ -3189,6 +3198,14 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper }) } + private flushPendingKiloModel(): void { + if (!this.webview || !this.isWebviewReady || !this.client || !this.pendingKiloModelID) return + + const modelID = this.pendingKiloModelID + this.pendingKiloModelID = null + this.postMessage({ type: "selectKiloModel", modelID }) + } + public async appendReviewComments(comments: unknown[], autoSend = false): Promise { this.pendingReviewComments.push({ comments, autoSend }) diff --git a/packages/kilo-vscode/src/extension.ts b/packages/kilo-vscode/src/extension.ts index e45d49bbe33..2c1f9e471d1 100644 --- a/packages/kilo-vscode/src/extension.ts +++ b/packages/kilo-vscode/src/extension.ts @@ -26,6 +26,7 @@ let agentManager: AgentManagerProvider | undefined let shuttingDown = false const RESTORE_KEY = "kilo.workbench.restore" +const PROMOTED_KILO_MODEL_IDS = new Set(["stealth/claude-opus-4.8"]) type RestoreState = { sidebar?: boolean @@ -36,9 +37,9 @@ const panelTitleHandler = (panel: vscode.WebviewPanel) => (title: string) => { panel.title = title || EXTENSION_DISPLAY_NAME } -// Activated via "onStartupFinished" (package.json) so that commands, code actions, keybindings, -// autocomplete, commit-message generation, and URI deep links all work immediately — without -// requiring the user to open a Kilo sidebar or panel first. The CLI backend is NOT spawned here; +// Activated via "onStartupFinished" and "onUri" (package.json) so that commands, code actions, +// keybindings, autocomplete, commit-message generation, and URI deep links all work immediately — +// without requiring the user to open a Kilo sidebar or panel first. The CLI backend is NOT spawned here; // it starts lazily when a webview connects or when ensureBackendForAutocomplete() triggers it. export function activate(context: vscode.ExtensionContext) { console.log("Kilo Code extension is now active") @@ -471,17 +472,25 @@ export function activate(context: vscode.ExtensionContext) { ), ) - // Register URI handler for session imports (vscode://kilocode.kilo-code/kilocode/s/{sessionId}) + // Register URI handler for extension deep links (vscode://kilocode.kilo-code/kilocode/...) context.subscriptions.push( vscode.window.registerUriHandler({ async handleUri(uri: vscode.Uri) { - const match = uri.path.match(/^\/kilocode\/s\/([a-zA-Z0-9_-]+)$/) - if (!match) return - const sessionId = match[1] - if (!sessionId) return - console.log("[Kilo New] URI handler: opening cloud session:", sessionId) + const sessionMatch = uri.path.match(/^\/kilocode\/s\/([a-zA-Z0-9_-]+)$/) + const sessionId = sessionMatch?.[1] + if (sessionId) { + console.log("[Kilo New] URI handler: opening cloud session:", sessionId) + await vscode.commands.executeCommand(`${KiloProvider.viewType}.focus`) + provider.openCloudSession(sessionId) + return + } + + if (uri.path !== "/kilocode/model") return + const modelID = new URLSearchParams(uri.query).get("model") + if (!modelID || !PROMOTED_KILO_MODEL_IDS.has(modelID)) return + console.log("[Kilo New] URI handler: selecting promoted model:", modelID) await vscode.commands.executeCommand(`${KiloProvider.viewType}.focus`) - provider.openCloudSession(sessionId) + provider.selectKiloModel(modelID) }, }), ) diff --git a/packages/kilo-vscode/src/kilo-provider/model-state.ts b/packages/kilo-vscode/src/kilo-provider/model-state.ts index c97775bb2d1..e79ed94467d 100644 --- a/packages/kilo-vscode/src/kilo-provider/model-state.ts +++ b/packages/kilo-vscode/src/kilo-provider/model-state.ts @@ -41,18 +41,22 @@ async function read(client: KiloClient | null): Promise> } } -function write(client: KiloClient | null, key: string, value: unknown): Promise { +function update(client: KiloClient | null, key: string, updater: (value: unknown) => unknown): Promise { const op = queue.then(async () => { const p = await resolve(client) if (!p) return const existing = await read(client) - existing[key] = value + existing[key] = updater(existing[key]) await fs.promises.writeFile(p, JSON.stringify(existing, null, 2)) }) queue = op.catch(() => {}) return op } +function write(client: KiloClient | null, key: string, value: unknown): Promise { + return update(client, key, () => value) +} + /** * Handle a model-state webview message. Returns true if handled. */ @@ -63,20 +67,22 @@ 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) + await update(client, "model", (value) => { + const model = validateModelSelections(value) + model[message.agent as string] = { + providerID: message.providerID as string, + modelID: message.modelID as string, + } + return model + }) return true } if (type === "clearModelSelection") { - const data = await read(client) - const model = validateModelSelections(data.model) - delete model[message.agent as string] - await write(client, "model", model) + await update(client, "model", (value) => { + const model = validateModelSelections(value) + delete model[message.agent as string] + return model + }) return true } if (type === "requestModelSelections") { diff --git a/packages/kilo-vscode/webview-ui/src/App.tsx b/packages/kilo-vscode/webview-ui/src/App.tsx index 9843ad7d7d9..8d17087cdda 100644 --- a/packages/kilo-vscode/webview-ui/src/App.tsx +++ b/packages/kilo-vscode/webview-ui/src/App.tsx @@ -10,6 +10,7 @@ import { Diff } from "@kilocode/kilo-ui/diff" import { File } from "@kilocode/kilo-ui/file" import { DataProvider } from "@kilocode/kilo-ui/context/data" import { Toast } from "@kilocode/kilo-ui/toast" +import { KILO_PROVIDER_ID } from "../../src/shared/provider-model" import Settings from "./components/settings/Settings" import ProfileView from "./components/profile/ProfileView" import { VSCodeProvider, useVSCode } from "./context/vscode" @@ -36,11 +37,18 @@ import { NotificationsProvider } from "./context/notifications" import { FeedbackProvider } from "./context/feedback" import { KiloEmbeddingModelsProvider } from "./context/kilo-embedding-models" import type { Message as SDKMessage, Part as SDKPart } from "@kilocode/sdk/v2" +import type { SelectKiloModelMessage } from "./types/messages" import "./styles/chat.css" type ViewType = "newTask" | "marketplace" | "history" | "profile" | "settings" | "subAgentViewer" const VALID_VIEWS = new Set(["newTask", "marketplace", "history", "profile", "settings", "subAgentViewer"]) +function isSelectKiloModelMessage(message: unknown): message is SelectKiloModelMessage { + if (!message || typeof message !== "object") return false + if (!("type" in message) || !("modelID" in message)) return false + return message.type === "selectKiloModel" && typeof message.modelID === "string" +} + /** * Bridge our session store to the DataProvider's expected Data shape. * @@ -261,6 +269,13 @@ const AppContent: Component = () => { setCurrentView("newTask") } + const handleSelectKiloModel = (message: unknown) => { + if (!isSelectKiloModelMessage(message)) return + console.log("[Kilo New] App: selecting promoted Kilo model:", message.modelID) + session.selectPersistedModel(KILO_PROVIDER_ID, message.modelID) + setCurrentView("newTask") + } + onMount(() => { const handler = (event: MessageEvent) => { const message = event.data @@ -279,6 +294,7 @@ const AppContent: Component = () => { session.selectCloudSession(message.sessionId) setCurrentView("newTask") } + handleSelectKiloModel(message) handleForked(message) if (message?.type === "viewSubAgentSession" && message.sessionID) { console.log("[Kilo New] App: 🔍 viewSubAgentSession:", message.sessionID) diff --git a/packages/kilo-vscode/webview-ui/src/context/session.tsx b/packages/kilo-vscode/webview-ui/src/context/session.tsx index d93c05ffce5..73e7d2fa968 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session.tsx +++ b/packages/kilo-vscode/webview-ui/src/context/session.tsx @@ -172,6 +172,7 @@ interface SessionContextValue { // Model selection (global, extension-lifetime) selected: (sessionID?: string) => ModelSelection | null selectModel: (providerID: string, modelID: string, sessionID?: string) => void + selectPersistedModel: (providerID: string, modelID: string) => void hasModelOverride: (sessionID?: string) => boolean clearModelOverride: (sessionID?: string) => void @@ -345,6 +346,7 @@ export const SessionProvider: ParentComponent = (props) => { const [agents, setAgents] = createSignal([]) const [allAgents, setAllAgents] = createSignal([]) const [defaultAgent, setDefaultAgent] = createSignal("code") + const [pendingPersistedModel, setPendingPersistedModel] = createSignal(null) // Skills loaded from the CLI backend const [skills, setSkills] = createSignal([]) @@ -506,17 +508,9 @@ export const SessionProvider: ParentComponent = (props) => { vscode.postMessage({ type: "persistRecents", recents: updated }) } - 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). + function persistModelSelection(agentName: string, selection: ModelSelection) { setUserSetAgents((prev) => ({ ...prev, [agentName]: true })) setStore("modelSelections", agentName, selection) - // Persist to model.json via the extension host vscode.postMessage({ type: "persistModelSelection", agent: agentName, @@ -525,6 +519,17 @@ export const SessionProvider: ParentComponent = (props) => { }) } + 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). + persistModelSelection(agentName, selection) + } + function selectModel(providerID: string, modelID: string, sessionID?: string) { const sid = sessionID ?? currentSessionID() applyModel(agentForScope(sid), { providerID, modelID }, sid) @@ -533,6 +538,30 @@ export const SessionProvider: ParentComponent = (props) => { } } + function applyPersistedModel(selection: ModelSelection) { + const sid = currentSessionID() + const defaultAgentName = defaultAgent() + const activeAgentName = agentForScope() + pushRecent(selection) + persistModelSelection(defaultAgentName, selection) + if (activeAgentName !== defaultAgentName) { + persistModelSelection(activeAgentName, selection) + } + if (sid) { + setStore("sessionOverrides", sid, selection) + hideErrors(sid) + } + } + + function selectPersistedModel(providerID: string, modelID: string) { + const selection = { providerID, modelID } + if (agents().length === 0) { + setPendingPersistedModel(selection) + return + } + applyPersistedModel(selection) + } + function promptAgent(sessionID?: string) { const name = agentForScope(sessionID) return name !== defaultAgent() ? name : undefined @@ -637,6 +666,12 @@ export const SessionProvider: ParentComponent = (props) => { }), ) + const pendingModel = pendingPersistedModel() + if (pendingModel) { + setPendingPersistedModel(null) + applyPersistedModel(pendingModel) + } + // Rescan already-loaded message history so sessions whose messagesLoaded // arrived before agentsLoaded (and therefore got no agent selection) are // backfilled now that we know the valid agent names. @@ -2478,6 +2513,7 @@ export const SessionProvider: ParentComponent = (props) => { scopedSuggestions, selected, selectModel, + selectPersistedModel, hasModelOverride, clearModelOverride, costBreakdown, 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 6d4b5d9a0b4..63a83ff3097 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 @@ -229,6 +229,11 @@ export interface OpenCloudSessionMessage { sessionId: string } +export interface SelectKiloModelMessage { + type: "selectKiloModel" + modelID: string +} + export interface ActionMessage { type: "action" action: string @@ -1011,6 +1016,7 @@ export type ExtensionMessage = | CloudSessionImportedMessage | CloudSessionImportFailedMessage | OpenCloudSessionMessage + | SelectKiloModelMessage | AgentManagerBranchesMessage | AgentManagerExternalWorktreesMessage | AgentManagerImportResultMessage From d1b29009dd8e362b9a59b5cfeb724fbe4c113234 Mon Sep 17 00:00:00 2001 From: Josh Lambert Date: Wed, 3 Jun 2026 23:07:24 -0400 Subject: [PATCH 2/6] fix(vscode): harden promoted model deep links --- .changeset/promoted-vscode-model-links.md | 5 + packages/kilo-vscode/src/KiloProvider.ts | 20 ++-- .../src/kilo-provider/model-state.ts | 10 +- .../unit/kilo-provider-promoted-model.test.ts | 71 ++++++++++++++ .../tests/unit/model-selection.test.ts | 37 +++++++- .../tests/unit/model-state.test.ts | 93 +++++++++++++++++++ .../webview-ui/src/context/model-selection.ts | 15 +++ .../webview-ui/src/context/provider.tsx | 4 + .../webview-ui/src/context/session.tsx | 40 +++++--- .../webview-ui/src/stories/StoryProviders.tsx | 1 + .../src/types/messages/extension-messages.ts | 1 + .../src/types/messages/webview-messages.ts | 1 + 12 files changed, 274 insertions(+), 24 deletions(-) create mode 100644 .changeset/promoted-vscode-model-links.md create mode 100644 packages/kilo-vscode/tests/unit/kilo-provider-promoted-model.test.ts create mode 100644 packages/kilo-vscode/tests/unit/model-state.test.ts diff --git a/.changeset/promoted-vscode-model-links.md b/.changeset/promoted-vscode-model-links.md new file mode 100644 index 00000000000..4e01463f762 --- /dev/null +++ b/.changeset/promoted-vscode-model-links.md @@ -0,0 +1,5 @@ +--- +"kilo-code": minor +--- + +Support selecting promoted Kilo models from VS Code deep links. diff --git a/packages/kilo-vscode/src/KiloProvider.ts b/packages/kilo-vscode/src/KiloProvider.ts index e76acd6be32..db1d0e6ec82 100644 --- a/packages/kilo-vscode/src/KiloProvider.ts +++ b/packages/kilo-vscode/src/KiloProvider.ts @@ -670,14 +670,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper this.visibleTaskStreams.handle(message) switch (message.type) { case "webviewReady": - console.log("[Kilo New] KiloProvider: ✅ webviewReady received") - this.isWebviewReady = true - this.visibleTaskStreams.clear() - await this.syncWebviewState("webviewReady") - this.flushPendingKiloModel() - this.flushPendingReviewComments() - this.recoverPendingPrompts() - this.readyResolvers.splice(0).forEach((r) => r()) + await this.handleWebviewReady() break case "sendMessage": { const msg = message as typeof message & ContextMessage @@ -1150,6 +1143,17 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper this.webviewMessageDisposable = watchFontSizeConfig((msg) => this.postMessage(msg), this.webviewMessageDisposable) } + private async handleWebviewReady(): Promise { + console.log("[Kilo New] KiloProvider: ✅ webviewReady received") + this.isWebviewReady = true + this.visibleTaskStreams.clear() + this.flushPendingKiloModel() + await this.syncWebviewState("webviewReady") + this.flushPendingReviewComments() + this.recoverPendingPrompts() + this.readyResolvers.splice(0).forEach((r) => r()) + } + private handleEditorOpenMessage(message: Parameters[0]): boolean { return handleEditorAction(message, { dir: () => this.getWorkspaceDirectory(this.currentSession?.id), diff --git a/packages/kilo-vscode/src/kilo-provider/model-state.ts b/packages/kilo-vscode/src/kilo-provider/model-state.ts index e79ed94467d..2529dc46bd1 100644 --- a/packages/kilo-vscode/src/kilo-provider/model-state.ts +++ b/packages/kilo-vscode/src/kilo-provider/model-state.ts @@ -41,6 +41,11 @@ async function read(client: KiloClient | null): Promise> } } +async function latest(client: KiloClient | null): Promise> { + await queue + return read(client) +} + function update(client: KiloClient | null, key: string, updater: (value: unknown) => unknown): Promise { const op = queue.then(async () => { const p = await resolve(client) @@ -86,9 +91,10 @@ export async function handleMessage( return true } if (type === "requestModelSelections") { - const data = await read(client) + const data = await latest(client) const selections = validateModelSelections(data.model) - post({ type: "modelSelectionsLoaded", selections }) + const revision = typeof message.revision === "number" ? message.revision : undefined + post({ type: "modelSelectionsLoaded", selections, ...(revision !== undefined && { revision }) }) return true } return false diff --git a/packages/kilo-vscode/tests/unit/kilo-provider-promoted-model.test.ts b/packages/kilo-vscode/tests/unit/kilo-provider-promoted-model.test.ts new file mode 100644 index 00000000000..74429ec0b00 --- /dev/null +++ b/packages/kilo-vscode/tests/unit/kilo-provider-promoted-model.test.ts @@ -0,0 +1,71 @@ +import { describe, expect, it } from "bun:test" + +// vscode mock is provided by the shared preload (tests/setup/vscode-mock.ts) +const { KiloProvider } = await import("../../src/KiloProvider") + +type Internals = { + webview: { postMessage: (message: unknown) => Promise } | null + isWebviewReady: boolean + syncWebviewState: () => Promise + handleWebviewReady: () => Promise + flushPendingKiloModel: () => void +} + +function connection() { + const state = { connected: true } + return { + state, + getClient: () => { + if (!state.connected) throw new Error("Not connected") + return {} + }, + } +} + +describe("KiloProvider promoted model selection", () => { + it("flushes a queued model before ancillary webview sync can fail", async () => { + const service = connection() + const provider = new KiloProvider({} as never, service as never) + const internal = provider as unknown as Internals + const sent: unknown[] = [] + + internal.webview = { + postMessage: async (message) => { + sent.push(message) + return true + }, + } + internal.syncWebviewState = async () => { + throw new Error("profile unavailable") + } + + provider.selectKiloModel("stealth/claude-opus-4.8") + expect(sent).toEqual([]) + + await expect(internal.handleWebviewReady()).rejects.toThrow("profile unavailable") + expect(sent).toEqual([{ type: "selectKiloModel", modelID: "stealth/claude-opus-4.8" }]) + }) + + it("keeps a queued model until the backend client is available", () => { + const service = connection() + service.state.connected = false + const provider = new KiloProvider({} as never, service as never) + const internal = provider as unknown as Internals + const sent: unknown[] = [] + + internal.webview = { + postMessage: async (message) => { + sent.push(message) + return true + }, + } + internal.isWebviewReady = true + + provider.selectKiloModel("stealth/claude-opus-4.8") + expect(sent).toEqual([]) + + service.state.connected = true + internal.flushPendingKiloModel() + expect(sent).toEqual([{ type: "selectKiloModel", modelID: "stealth/claude-opus-4.8" }]) + }) +}) diff --git a/packages/kilo-vscode/tests/unit/model-selection.test.ts b/packages/kilo-vscode/tests/unit/model-selection.test.ts index c61d1df1db6..61ef7d7dd7c 100644 --- a/packages/kilo-vscode/tests/unit/model-selection.test.ts +++ b/packages/kilo-vscode/tests/unit/model-selection.test.ts @@ -1,5 +1,9 @@ import { describe, expect, it } from "bun:test" -import { resolveModelSelection } from "../../webview-ui/src/context/model-selection" +import { + isCurrentModelSelections, + promotion, + resolveModelSelection, +} from "../../webview-ui/src/context/model-selection" import { KILO_AUTO, parseModelString } from "../../src/shared/provider-model" import type { Provider } from "../../webview-ui/src/types/messages" @@ -38,6 +42,37 @@ describe("parseModelString", () => { }) }) +describe("promotion", () => { + it("waits for agents and provider metadata", () => { + expect(promotion({ agents: false, loaded: true, providers, connected: [], selection: KILO_AUTO })).toBe("pending") + expect(promotion({ agents: true, loaded: false, providers, connected: [], selection: KILO_AUTO })).toBe("pending") + }) + + it("rejects models missing from the loaded Kilo catalog", () => { + expect( + promotion({ + agents: true, + loaded: true, + providers, + connected: [], + selection: { providerID: "kilo", modelID: "stealth/claude-opus-4.8" }, + }), + ).toBe("invalid") + }) + + it("applies models exposed by the loaded Kilo catalog", () => { + expect(promotion({ agents: true, loaded: true, providers, connected: [], selection: KILO_AUTO })).toBe("apply") + }) +}) + +describe("isCurrentModelSelections", () => { + it("rejects stale hydration responses after local selection changes", () => { + expect(isCurrentModelSelections(0, 1)).toBe(false) + expect(isCurrentModelSelections(1, 1)).toBe(true) + expect(isCurrentModelSelections(undefined, 1)).toBe(true) + }) +}) + describe("resolveModelSelection", () => { it("prefers a valid override", () => { const result = resolveModelSelection({ 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..ed0d4458840 --- /dev/null +++ b/packages/kilo-vscode/tests/unit/model-state.test.ts @@ -0,0 +1,93 @@ +import { afterAll, beforeEach, describe, expect, it } from "bun:test" +import * as fs from "fs" +import * as os from "os" +import * as path from "path" +import type { KiloClient } from "@kilocode/sdk/v2/client" +import { handleMessage } from "../../src/kilo-provider/model-state" + +const root = fs.mkdtempSync(path.join(os.tmpdir(), "model-state-")) +const file = path.join(root, "model.json") +const client = { + path: { + get: async () => ({ data: { state: root } }), + }, +} as unknown as KiloClient + +afterAll(() => { + fs.rmSync(root, { recursive: true, force: true }) +}) + +beforeEach(() => { + fs.writeFileSync(file, JSON.stringify({ model: {} })) +}) + +describe("model state", () => { + it("preserves concurrent selections for different agents", async () => { + await Promise.all([ + handleMessage( + "persistModelSelection", + { agent: "code", providerID: "kilo", modelID: "stealth/claude-opus-4.8" }, + client, + () => {}, + ), + handleMessage( + "persistModelSelection", + { agent: "plan", providerID: "kilo", modelID: "kilo-auto/free" }, + client, + () => {}, + ), + ]) + + const data = JSON.parse(fs.readFileSync(file, "utf-8")) as { model: Record } + expect(data.model).toEqual({ + code: { providerID: "kilo", modelID: "stealth/claude-opus-4.8" }, + plan: { providerID: "kilo", modelID: "kilo-auto/free" }, + }) + }) + + it("waits for queued writes before hydrating selections", async () => { + const api = fs.promises as unknown as { writeFile: (...args: unknown[]) => Promise } + const original = api.writeFile + const gate = Promise.withResolvers() + const started = Promise.withResolvers() + const sent: unknown[] = [] + + api.writeFile = async (...args: unknown[]) => { + if (args[0] === file) { + started.resolve() + await gate.promise + } + await original(...args) + } + + try { + const pending = handleMessage( + "persistModelSelection", + { agent: "code", providerID: "kilo", modelID: "stealth/claude-opus-4.8" }, + client, + () => {}, + ) + await started.promise + + const hydration = handleMessage("requestModelSelections", { revision: 3 }, client, (message) => + sent.push(message), + ) + await Promise.resolve() + expect(sent).toEqual([]) + + gate.resolve() + await Promise.all([pending, hydration]) + } finally { + gate.resolve() + api.writeFile = original + } + + expect(sent).toEqual([ + { + type: "modelSelectionsLoaded", + selections: { code: { providerID: "kilo", modelID: "stealth/claude-opus-4.8" } }, + revision: 3, + }, + ]) + }) +}) 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 86d9a255b98..c679b429afd 100644 --- a/packages/kilo-vscode/webview-ui/src/context/model-selection.ts +++ b/packages/kilo-vscode/webview-ui/src/context/model-selection.ts @@ -41,3 +41,18 @@ export function resolveModelSelection(input: { null ) } + +export function promotion(input: { + agents: boolean + loaded: boolean + providers: Record + connected: string[] + selection: ModelSelection +}): "pending" | "invalid" | "apply" { + if (!input.agents || !input.loaded) return "pending" + return isModelValid(input.providers, input.connected, input.selection) ? "apply" : "invalid" +} + +export function isCurrentModelSelections(revision: number | undefined, current: number): boolean { + return revision === undefined || revision === current +} diff --git a/packages/kilo-vscode/webview-ui/src/context/provider.tsx b/packages/kilo-vscode/webview-ui/src/context/provider.tsx index 4b173b8baa0..723b50ba51d 100644 --- a/packages/kilo-vscode/webview-ui/src/context/provider.tsx +++ b/packages/kilo-vscode/webview-ui/src/context/provider.tsx @@ -17,6 +17,7 @@ export type EnrichedModel = ProviderModel & { providerID: string; providerName: interface ProviderContextValue { providers: Accessor> connected: Accessor + loaded: Accessor defaults: Accessor> defaultSelection: Accessor models: Accessor @@ -33,6 +34,7 @@ export const ProviderProvider: ParentComponent = (props) => { const [providers, setProviders] = createSignal>({}) const [connected, setConnected] = createSignal([]) + const [loaded, setLoaded] = createSignal(false) const [defaults, setDefaults] = createSignal>({}) const [defaultSelection, setDefaultSelection] = createSignal(KILO_AUTO) const [authMethods, setAuthMethods] = createSignal>({}) @@ -57,6 +59,7 @@ export const ProviderProvider: ParentComponent = (props) => { setProviders(message.providers) setConnected(message.connected) + setLoaded(true) setDefaults(message.defaults) setDefaultSelection(message.defaultSelection) setAuthMethods(message.authMethods) @@ -92,6 +95,7 @@ export const ProviderProvider: ParentComponent = (props) => { const value: ProviderContextValue = { providers, connected, + loaded, defaults, defaultSelection, models, diff --git a/packages/kilo-vscode/webview-ui/src/context/session.tsx b/packages/kilo-vscode/webview-ui/src/context/session.tsx index 73e7d2fa968..af40b6e3b4a 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session.tsx +++ b/packages/kilo-vscode/webview-ui/src/context/session.tsx @@ -51,7 +51,7 @@ import { upsertSessionToolPart, } from "./session-utils" import { Identifier } from "../utils/id" -import { resolveModelSelection } from "./model-selection" +import { isCurrentModelSelections, promotion, resolveModelSelection } from "./model-selection" import { resolveMessagePrefs } from "./session-preferences" import { errorIDs } from "./session-errors" import { PartStash } from "./part-stash" @@ -347,6 +347,7 @@ export const SessionProvider: ParentComponent = (props) => { const [allAgents, setAllAgents] = createSignal([]) const [defaultAgent, setDefaultAgent] = createSignal("code") const [pendingPersistedModel, setPendingPersistedModel] = createSignal(null) + let revision = 0 // Skills loaded from the CLI backend const [skills, setSkills] = createSignal([]) @@ -509,6 +510,7 @@ export const SessionProvider: ParentComponent = (props) => { } function persistModelSelection(agentName: string, selection: ModelSelection) { + revision++ setUserSetAgents((prev) => ({ ...prev, [agentName]: true })) setStore("modelSelections", agentName, selection) vscode.postMessage({ @@ -554,13 +556,27 @@ export const SessionProvider: ParentComponent = (props) => { } function selectPersistedModel(providerID: string, modelID: string) { - const selection = { providerID, modelID } - if (agents().length === 0) { - setPendingPersistedModel(selection) + setPendingPersistedModel({ providerID, modelID }) + } + + createEffect(() => { + const selection = pendingPersistedModel() + if (!selection) return + const status = promotion({ + agents: agents().length > 0, + loaded: provider.loaded(), + providers: provider.providers(), + connected: provider.connected(), + selection, + }) + if (status === "pending") return + setPendingPersistedModel(null) + if (status === "invalid") { + console.warn("[Kilo New] Ignoring unavailable promoted Kilo model:", selection.modelID) return } applyPersistedModel(selection) - } + }) function promptAgent(sessionID?: string) { const name = agentForScope(sessionID) @@ -589,7 +605,10 @@ export const SessionProvider: ParentComponent = (props) => { delete selections[agentName] }), ) - if (persist) vscode.postMessage({ type: "clearModelSelection", agent: agentName }) + if (persist) { + revision++ + vscode.postMessage({ type: "clearModelSelection", agent: agentName }) + } } function shouldClearModeModelSelection(agentName: string) { @@ -666,12 +685,6 @@ export const SessionProvider: ParentComponent = (props) => { }), ) - const pendingModel = pendingPersistedModel() - if (pendingModel) { - setPendingPersistedModel(null) - applyPersistedModel(pendingModel) - } - // Rescan already-loaded message history so sessions whose messagesLoaded // arrived before agentsLoaded (and therefore got no agent selection) are // backfilled now that we know the valid agent names. @@ -795,6 +808,7 @@ export const SessionProvider: ParentComponent = (props) => { // Uses replace semantics so a reset (empty payload) clears old entries. const unsubSelections = vscode.onMessage((message: ExtensionMessage) => { if (message.type !== "modelSelectionsLoaded") return + if (!isCurrentModelSelections(message.revision, revision)) return setStore("modelSelections", reconcile(message.selections)) const flags: Record = {} for (const name of Object.keys(message.selections)) { @@ -802,7 +816,7 @@ export const SessionProvider: ParentComponent = (props) => { } setUserSetAgents(flags) }) - vscode.postMessage({ type: "requestModelSelections" }) + vscode.postMessage({ type: "requestModelSelections", revision }) onCleanup(unsubSelections) // Load persisted recent models from extension globalState diff --git a/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx b/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx index 9fca0126250..9b069b7938d 100644 --- a/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx +++ b/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx @@ -95,6 +95,7 @@ const MockProviderProvider: ParentComponent<{ kiloAuth?: boolean }> = (props) => const value = { providers: () => MOCK_PROVIDERS as any, connected: () => ["kilo"], + loaded: () => true, defaults: () => ({}), defaultSelection: () => ({ providerID: "kilo", modelID: "anthropic/claude-sonnet-4-6" }), models: () => MOCK_MODELS, 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 63a83ff3097..b1b3272fbe2 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 @@ -635,6 +635,7 @@ export interface FavoritesLoadedMessage { export interface ModelSelectionsLoadedMessage { type: "modelSelectionsLoaded" selections: Record + revision?: number } 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 f0c0f315c0a..a064f604d26 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 @@ -989,6 +989,7 @@ export interface ClearModelSelectionRequest { export interface RequestModelSelectionsMessage { type: "requestModelSelections" + revision: number } // Continue in Worktree: transfer sidebar session + git state to an isolated worktree From 235e769fb359531da4716999ef5feac839affdcc Mon Sep 17 00:00:00 2001 From: Josh Lambert Date: Wed, 3 Jun 2026 23:31:59 -0400 Subject: [PATCH 3/6] fix(vscode): generalize model deep links --- .changeset/promoted-vscode-model-links.md | 5 --- .changeset/vscode-kilo-model-links.md | 5 +++ packages/kilo-vscode/src/extension.ts | 9 +++-- .../src/kilo-provider/model-uri.ts | 9 +++++ ...est.ts => kilo-provider-model-uri.test.ts} | 10 +++--- .../tests/unit/model-selection.test.ts | 36 ++++++++++--------- .../tests/unit/model-state.test.ts | 8 ++--- .../kilo-vscode/tests/unit/model-uri.test.ts | 25 +++++++++++++ packages/kilo-vscode/webview-ui/src/App.tsx | 5 ++- .../webview-ui/src/context/model-selection.ts | 8 ++--- .../webview-ui/src/context/provider-utils.ts | 3 +- .../webview-ui/src/context/session.tsx | 29 ++++++++------- 12 files changed, 94 insertions(+), 58 deletions(-) delete mode 100644 .changeset/promoted-vscode-model-links.md create mode 100644 .changeset/vscode-kilo-model-links.md create mode 100644 packages/kilo-vscode/src/kilo-provider/model-uri.ts rename packages/kilo-vscode/tests/unit/{kilo-provider-promoted-model.test.ts => kilo-provider-model-uri.test.ts} (83%) create mode 100644 packages/kilo-vscode/tests/unit/model-uri.test.ts diff --git a/.changeset/promoted-vscode-model-links.md b/.changeset/promoted-vscode-model-links.md deleted file mode 100644 index 4e01463f762..00000000000 --- a/.changeset/promoted-vscode-model-links.md +++ /dev/null @@ -1,5 +0,0 @@ ---- -"kilo-code": minor ---- - -Support selecting promoted Kilo models from VS Code deep links. diff --git a/.changeset/vscode-kilo-model-links.md b/.changeset/vscode-kilo-model-links.md new file mode 100644 index 00000000000..2f71e28221a --- /dev/null +++ b/.changeset/vscode-kilo-model-links.md @@ -0,0 +1,5 @@ +--- +"kilo-code": minor +--- + +Support selecting Kilo Gateway catalog models from VS Code deep links. diff --git a/packages/kilo-vscode/src/extension.ts b/packages/kilo-vscode/src/extension.ts index 2c1f9e471d1..4ba283a8d60 100644 --- a/packages/kilo-vscode/src/extension.ts +++ b/packages/kilo-vscode/src/extension.ts @@ -21,12 +21,12 @@ import { registerToggleAutoApprove } from "./commands/toggle-auto-approve" import { registerHeapSnapshot } from "./commands/heap-snapshot" import { RemoteStatusService } from "./services/RemoteStatusService" import { markWorkspace } from "./util/spotlight" +import { kiloModelFromURI } from "./kilo-provider/model-uri" let agentManager: AgentManagerProvider | undefined let shuttingDown = false const RESTORE_KEY = "kilo.workbench.restore" -const PROMOTED_KILO_MODEL_IDS = new Set(["stealth/claude-opus-4.8"]) type RestoreState = { sidebar?: boolean @@ -485,10 +485,9 @@ export function activate(context: vscode.ExtensionContext) { return } - if (uri.path !== "/kilocode/model") return - const modelID = new URLSearchParams(uri.query).get("model") - if (!modelID || !PROMOTED_KILO_MODEL_IDS.has(modelID)) return - console.log("[Kilo New] URI handler: selecting promoted model:", modelID) + const modelID = kiloModelFromURI(uri) + if (!modelID) return + console.log("[Kilo New] URI handler: selecting linked Kilo model:", modelID) await vscode.commands.executeCommand(`${KiloProvider.viewType}.focus`) provider.selectKiloModel(modelID) }, diff --git a/packages/kilo-vscode/src/kilo-provider/model-uri.ts b/packages/kilo-vscode/src/kilo-provider/model-uri.ts new file mode 100644 index 00000000000..c97492e039e --- /dev/null +++ b/packages/kilo-vscode/src/kilo-provider/model-uri.ts @@ -0,0 +1,9 @@ +type URI = { + path: string + query: string +} + +export function kiloModelFromURI(uri: URI): string | undefined { + if (uri.path !== "/kilocode/model") return undefined + return new URLSearchParams(uri.query).get("model") || undefined +} diff --git a/packages/kilo-vscode/tests/unit/kilo-provider-promoted-model.test.ts b/packages/kilo-vscode/tests/unit/kilo-provider-model-uri.test.ts similarity index 83% rename from packages/kilo-vscode/tests/unit/kilo-provider-promoted-model.test.ts rename to packages/kilo-vscode/tests/unit/kilo-provider-model-uri.test.ts index 74429ec0b00..ff4ff60220e 100644 --- a/packages/kilo-vscode/tests/unit/kilo-provider-promoted-model.test.ts +++ b/packages/kilo-vscode/tests/unit/kilo-provider-model-uri.test.ts @@ -22,7 +22,7 @@ function connection() { } } -describe("KiloProvider promoted model selection", () => { +describe("KiloProvider model URI selection", () => { it("flushes a queued model before ancillary webview sync can fail", async () => { const service = connection() const provider = new KiloProvider({} as never, service as never) @@ -39,11 +39,11 @@ describe("KiloProvider promoted model selection", () => { throw new Error("profile unavailable") } - provider.selectKiloModel("stealth/claude-opus-4.8") + provider.selectKiloModel("vendor/new-live-model") expect(sent).toEqual([]) await expect(internal.handleWebviewReady()).rejects.toThrow("profile unavailable") - expect(sent).toEqual([{ type: "selectKiloModel", modelID: "stealth/claude-opus-4.8" }]) + expect(sent).toEqual([{ type: "selectKiloModel", modelID: "vendor/new-live-model" }]) }) it("keeps a queued model until the backend client is available", () => { @@ -61,11 +61,11 @@ describe("KiloProvider promoted model selection", () => { } internal.isWebviewReady = true - provider.selectKiloModel("stealth/claude-opus-4.8") + provider.selectKiloModel("vendor/new-live-model") expect(sent).toEqual([]) service.state.connected = true internal.flushPendingKiloModel() - expect(sent).toEqual([{ type: "selectKiloModel", modelID: "stealth/claude-opus-4.8" }]) + expect(sent).toEqual([{ type: "selectKiloModel", modelID: "vendor/new-live-model" }]) }) }) diff --git a/packages/kilo-vscode/tests/unit/model-selection.test.ts b/packages/kilo-vscode/tests/unit/model-selection.test.ts index 61ef7d7dd7c..f39c4157c05 100644 --- a/packages/kilo-vscode/tests/unit/model-selection.test.ts +++ b/packages/kilo-vscode/tests/unit/model-selection.test.ts @@ -1,7 +1,7 @@ import { describe, expect, it } from "bun:test" import { isCurrentModelSelections, - promotion, + kiloCatalogModelStatus, resolveModelSelection, } from "../../webview-ui/src/context/model-selection" import { KILO_AUTO, parseModelString } from "../../src/shared/provider-model" @@ -16,7 +16,7 @@ function makeProvider(id: string, name: string, modelIds: string[]): Provider { } const providers = { - kilo: makeProvider("kilo", "Kilo Gateway", ["kilo-auto/free"]), + kilo: makeProvider("kilo", "Kilo Gateway", ["kilo-auto/free", "vendor/new-live-model"]), anthropic: makeProvider("anthropic", "Anthropic", ["claude-sonnet-4"]), openai: makeProvider("openai", "OpenAI", ["gpt-4.1"]), } @@ -42,26 +42,30 @@ describe("parseModelString", () => { }) }) -describe("promotion", () => { +describe("Kilo catalog model URI validation", () => { it("waits for agents and provider metadata", () => { - expect(promotion({ agents: false, loaded: true, providers, connected: [], selection: KILO_AUTO })).toBe("pending") - expect(promotion({ agents: true, loaded: false, providers, connected: [], selection: KILO_AUTO })).toBe("pending") + expect(kiloCatalogModelStatus({ agents: false, loaded: true, providers, modelID: "vendor/new-live-model" })).toBe( + "pending", + ) + expect(kiloCatalogModelStatus({ agents: true, loaded: false, providers, modelID: "vendor/new-live-model" })).toBe( + "pending", + ) }) it("rejects models missing from the loaded Kilo catalog", () => { - expect( - promotion({ - agents: true, - loaded: true, - providers, - connected: [], - selection: { providerID: "kilo", modelID: "stealth/claude-opus-4.8" }, - }), - ).toBe("invalid") + expect(kiloCatalogModelStatus({ agents: true, loaded: true, providers, modelID: "vendor/missing-model" })).toBe( + "invalid", + ) }) - it("applies models exposed by the loaded Kilo catalog", () => { - expect(promotion({ agents: true, loaded: true, providers, connected: [], selection: KILO_AUTO })).toBe("apply") + it("applies arbitrary models exposed by the loaded Kilo catalog", () => { + expect(kiloCatalogModelStatus({ agents: true, loaded: true, providers, modelID: "vendor/new-live-model" })).toBe( + "apply", + ) + }) + + it("does not accept models exposed only by non-Kilo providers", () => { + expect(kiloCatalogModelStatus({ agents: true, loaded: true, providers, modelID: "gpt-4.1" })).toBe("invalid") }) }) diff --git a/packages/kilo-vscode/tests/unit/model-state.test.ts b/packages/kilo-vscode/tests/unit/model-state.test.ts index ed0d4458840..54ebc98c148 100644 --- a/packages/kilo-vscode/tests/unit/model-state.test.ts +++ b/packages/kilo-vscode/tests/unit/model-state.test.ts @@ -26,7 +26,7 @@ describe("model state", () => { await Promise.all([ handleMessage( "persistModelSelection", - { agent: "code", providerID: "kilo", modelID: "stealth/claude-opus-4.8" }, + { agent: "code", providerID: "kilo", modelID: "vendor/new-live-model" }, client, () => {}, ), @@ -40,7 +40,7 @@ describe("model state", () => { const data = JSON.parse(fs.readFileSync(file, "utf-8")) as { model: Record } expect(data.model).toEqual({ - code: { providerID: "kilo", modelID: "stealth/claude-opus-4.8" }, + code: { providerID: "kilo", modelID: "vendor/new-live-model" }, plan: { providerID: "kilo", modelID: "kilo-auto/free" }, }) }) @@ -63,7 +63,7 @@ describe("model state", () => { try { const pending = handleMessage( "persistModelSelection", - { agent: "code", providerID: "kilo", modelID: "stealth/claude-opus-4.8" }, + { agent: "code", providerID: "kilo", modelID: "vendor/new-live-model" }, client, () => {}, ) @@ -85,7 +85,7 @@ describe("model state", () => { expect(sent).toEqual([ { type: "modelSelectionsLoaded", - selections: { code: { providerID: "kilo", modelID: "stealth/claude-opus-4.8" } }, + selections: { code: { providerID: "kilo", modelID: "vendor/new-live-model" } }, revision: 3, }, ]) diff --git a/packages/kilo-vscode/tests/unit/model-uri.test.ts b/packages/kilo-vscode/tests/unit/model-uri.test.ts new file mode 100644 index 00000000000..f37f97e7c6f --- /dev/null +++ b/packages/kilo-vscode/tests/unit/model-uri.test.ts @@ -0,0 +1,25 @@ +import { describe, expect, it } from "bun:test" +import { kiloModelFromURI } from "../../src/kilo-provider/model-uri" + +describe("kiloModelFromURI", () => { + it("accepts arbitrary Kilo catalog model ids", () => { + expect(kiloModelFromURI({ path: "/kilocode/model", query: "model=vendor%2Fnew-live-model" })).toBe( + "vendor/new-live-model", + ) + }) + + it("preserves additional slashes inside model ids", () => { + expect(kiloModelFromURI({ path: "/kilocode/model", query: "model=vendor%2Ffamily%2Fmodel" })).toBe( + "vendor/family/model", + ) + }) + + it("rejects missing and empty model ids", () => { + expect(kiloModelFromURI({ path: "/kilocode/model", query: "" })).toBeUndefined() + expect(kiloModelFromURI({ path: "/kilocode/model", query: "model=" })).toBeUndefined() + }) + + it("rejects unrelated URI paths", () => { + expect(kiloModelFromURI({ path: "/kilocode/s/session-id", query: "model=vendor%2Fmodel" })).toBeUndefined() + }) +}) diff --git a/packages/kilo-vscode/webview-ui/src/App.tsx b/packages/kilo-vscode/webview-ui/src/App.tsx index 8d17087cdda..3b9a26cbc03 100644 --- a/packages/kilo-vscode/webview-ui/src/App.tsx +++ b/packages/kilo-vscode/webview-ui/src/App.tsx @@ -10,7 +10,6 @@ import { Diff } from "@kilocode/kilo-ui/diff" import { File } from "@kilocode/kilo-ui/file" import { DataProvider } from "@kilocode/kilo-ui/context/data" import { Toast } from "@kilocode/kilo-ui/toast" -import { KILO_PROVIDER_ID } from "../../src/shared/provider-model" import Settings from "./components/settings/Settings" import ProfileView from "./components/profile/ProfileView" import { VSCodeProvider, useVSCode } from "./context/vscode" @@ -271,8 +270,8 @@ const AppContent: Component = () => { const handleSelectKiloModel = (message: unknown) => { if (!isSelectKiloModelMessage(message)) return - console.log("[Kilo New] App: selecting promoted Kilo model:", message.modelID) - session.selectPersistedModel(KILO_PROVIDER_ID, message.modelID) + console.log("[Kilo New] App: selecting linked Kilo model:", message.modelID) + session.selectKiloModel(message.modelID) setCurrentView("newTask") } 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 c679b429afd..898dc1f9504 100644 --- a/packages/kilo-vscode/webview-ui/src/context/model-selection.ts +++ b/packages/kilo-vscode/webview-ui/src/context/model-selection.ts @@ -1,4 +1,5 @@ import type { ModelSelection, Provider } from "../types/messages" +import { KILO_PROVIDER_ID } from "../../../src/shared/provider-model" import { isModelValid } from "./provider-utils" function validate( @@ -42,15 +43,14 @@ export function resolveModelSelection(input: { ) } -export function promotion(input: { +export function kiloCatalogModelStatus(input: { agents: boolean loaded: boolean providers: Record - connected: string[] - selection: ModelSelection + modelID: string }): "pending" | "invalid" | "apply" { if (!input.agents || !input.loaded) return "pending" - return isModelValid(input.providers, input.connected, input.selection) ? "apply" : "invalid" + return input.providers[KILO_PROVIDER_ID]?.models[input.modelID] ? "apply" : "invalid" } export function isCurrentModelSelections(revision: number | undefined, current: number): boolean { diff --git a/packages/kilo-vscode/webview-ui/src/context/provider-utils.ts b/packages/kilo-vscode/webview-ui/src/context/provider-utils.ts index 6afa2f6b5e3..5c1bca22e25 100644 --- a/packages/kilo-vscode/webview-ui/src/context/provider-utils.ts +++ b/packages/kilo-vscode/webview-ui/src/context/provider-utils.ts @@ -1,3 +1,4 @@ +import { KILO_PROVIDER_ID } from "../../../src/shared/provider-model" import type { Provider, ProviderModel, ModelSelection } from "../types/messages" export type EnrichedModel = ProviderModel & { providerID: string; providerName: string } @@ -41,6 +42,6 @@ export function isModelValid( if (!selection) return false const provider = providers[selection.providerID] if (!provider) return false - if (selection.providerID !== "kilo" && !connected.includes(selection.providerID)) return false + if (selection.providerID !== KILO_PROVIDER_ID && !connected.includes(selection.providerID)) return false return !!provider.models[selection.modelID] } diff --git a/packages/kilo-vscode/webview-ui/src/context/session.tsx b/packages/kilo-vscode/webview-ui/src/context/session.tsx index af40b6e3b4a..be784a20aac 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session.tsx +++ b/packages/kilo-vscode/webview-ui/src/context/session.tsx @@ -51,14 +51,14 @@ import { upsertSessionToolPart, } from "./session-utils" import { Identifier } from "../utils/id" -import { isCurrentModelSelections, promotion, resolveModelSelection } from "./model-selection" +import { isCurrentModelSelections, kiloCatalogModelStatus, resolveModelSelection } from "./model-selection" 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 { KILO_AUTO, parseModelString } from "../../../src/shared/provider-model" +import { KILO_AUTO, KILO_PROVIDER_ID, parseModelString } from "../../../src/shared/provider-model" import { visibleMessages as filterVisibleMessages } from "./session-queue" const RECENT_LIMIT = 5 @@ -172,7 +172,7 @@ interface SessionContextValue { // Model selection (global, extension-lifetime) selected: (sessionID?: string) => ModelSelection | null selectModel: (providerID: string, modelID: string, sessionID?: string) => void - selectPersistedModel: (providerID: string, modelID: string) => void + selectKiloModel: (modelID: string) => void hasModelOverride: (sessionID?: string) => boolean clearModelOverride: (sessionID?: string) => void @@ -346,7 +346,7 @@ export const SessionProvider: ParentComponent = (props) => { const [agents, setAgents] = createSignal([]) const [allAgents, setAllAgents] = createSignal([]) const [defaultAgent, setDefaultAgent] = createSignal("code") - const [pendingPersistedModel, setPendingPersistedModel] = createSignal(null) + const [pendingKiloModelID, setPendingKiloModelID] = createSignal(null) let revision = 0 // Skills loaded from the CLI backend @@ -555,27 +555,26 @@ export const SessionProvider: ParentComponent = (props) => { } } - function selectPersistedModel(providerID: string, modelID: string) { - setPendingPersistedModel({ providerID, modelID }) + function selectKiloModel(modelID: string) { + setPendingKiloModelID(modelID) } createEffect(() => { - const selection = pendingPersistedModel() - if (!selection) return - const status = promotion({ + const modelID = pendingKiloModelID() + if (!modelID) return + const status = kiloCatalogModelStatus({ agents: agents().length > 0, loaded: provider.loaded(), providers: provider.providers(), - connected: provider.connected(), - selection, + modelID, }) if (status === "pending") return - setPendingPersistedModel(null) + setPendingKiloModelID(null) if (status === "invalid") { - console.warn("[Kilo New] Ignoring unavailable promoted Kilo model:", selection.modelID) + console.warn("[Kilo New] Ignoring unavailable Kilo catalog model:", modelID) return } - applyPersistedModel(selection) + applyPersistedModel({ providerID: KILO_PROVIDER_ID, modelID }) }) function promptAgent(sessionID?: string) { @@ -2527,7 +2526,7 @@ export const SessionProvider: ParentComponent = (props) => { scopedSuggestions, selected, selectModel, - selectPersistedModel, + selectKiloModel, hasModelOverride, clearModelOverride, costBreakdown, From 7691373d86951833765c39034b73a21be4b30c08 Mon Sep 17 00:00:00 2001 From: Josh Lambert Date: Thu, 4 Jun 2026 08:59:45 -0400 Subject: [PATCH 4/6] fix(vscode): refresh catalog for model deep links --- .../src/kilo-provider/model-state.ts | 24 ++++-------- .../tests/unit/model-selection.test.ts | 37 ++++-------------- packages/kilo-vscode/webview-ui/src/App.tsx | 16 ++------ .../webview-ui/src/context/model-selection.ts | 20 ++++------ .../webview-ui/src/context/provider-utils.ts | 3 +- .../webview-ui/src/context/provider.tsx | 4 -- .../webview-ui/src/context/session.tsx | 38 ++++++++++--------- .../webview-ui/src/stories/StoryProviders.tsx | 1 - 8 files changed, 49 insertions(+), 94 deletions(-) diff --git a/packages/kilo-vscode/src/kilo-provider/model-state.ts b/packages/kilo-vscode/src/kilo-provider/model-state.ts index 2529dc46bd1..7f0d242d1f8 100644 --- a/packages/kilo-vscode/src/kilo-provider/model-state.ts +++ b/packages/kilo-vscode/src/kilo-provider/model-state.ts @@ -41,27 +41,18 @@ async function read(client: KiloClient | null): Promise> } } -async function latest(client: KiloClient | null): Promise> { - await queue - return read(client) -} - -function update(client: KiloClient | null, key: string, updater: (value: unknown) => unknown): Promise { +function write(client: KiloClient | null, key: string, update: (value: unknown) => unknown): Promise { const op = queue.then(async () => { const p = await resolve(client) if (!p) return const existing = await read(client) - existing[key] = updater(existing[key]) + existing[key] = update(existing[key]) await fs.promises.writeFile(p, JSON.stringify(existing, null, 2)) }) queue = op.catch(() => {}) return op } -function write(client: KiloClient | null, key: string, value: unknown): Promise { - return update(client, key, () => value) -} - /** * Handle a model-state webview message. Returns true if handled. */ @@ -72,7 +63,7 @@ export async function handleMessage( post: PostMessage, ): Promise { if (type === "persistModelSelection") { - await update(client, "model", (value) => { + await write(client, "model", (value) => { const model = validateModelSelections(value) model[message.agent as string] = { providerID: message.providerID as string, @@ -83,7 +74,7 @@ export async function handleMessage( return true } if (type === "clearModelSelection") { - await update(client, "model", (value) => { + await write(client, "model", (value) => { const model = validateModelSelections(value) delete model[message.agent as string] return model @@ -91,16 +82,17 @@ export async function handleMessage( return true } if (type === "requestModelSelections") { - const data = await latest(client) + await queue + const data = await read(client) const selections = validateModelSelections(data.model) const revision = typeof message.revision === "number" ? message.revision : undefined - post({ type: "modelSelectionsLoaded", selections, ...(revision !== undefined && { revision }) }) + post({ type: "modelSelectionsLoaded", selections, revision }) return true } return false } export async function reset(client: KiloClient | null, post: PostMessage): Promise { - await write(client, "model", {}) + await write(client, "model", () => ({})) post({ type: "modelSelectionsLoaded", selections: {} }) } diff --git a/packages/kilo-vscode/tests/unit/model-selection.test.ts b/packages/kilo-vscode/tests/unit/model-selection.test.ts index f39c4157c05..4c2cf1e51b0 100644 --- a/packages/kilo-vscode/tests/unit/model-selection.test.ts +++ b/packages/kilo-vscode/tests/unit/model-selection.test.ts @@ -1,9 +1,5 @@ import { describe, expect, it } from "bun:test" -import { - isCurrentModelSelections, - kiloCatalogModelStatus, - resolveModelSelection, -} from "../../webview-ui/src/context/model-selection" +import { kiloCatalogModelStatus, resolveModelSelection } from "../../webview-ui/src/context/model-selection" import { KILO_AUTO, parseModelString } from "../../src/shared/provider-model" import type { Provider } from "../../webview-ui/src/types/messages" @@ -43,37 +39,20 @@ describe("parseModelString", () => { }) describe("Kilo catalog model URI validation", () => { - it("waits for agents and provider metadata", () => { - expect(kiloCatalogModelStatus({ agents: false, loaded: true, providers, modelID: "vendor/new-live-model" })).toBe( - "pending", - ) - expect(kiloCatalogModelStatus({ agents: true, loaded: false, providers, modelID: "vendor/new-live-model" })).toBe( - "pending", - ) + it("waits for a catalog refresh after the link arrives", () => { + expect(kiloCatalogModelStatus(providers, "vendor/new-live-model", 1, 1)).toBe("pending") }) - it("rejects models missing from the loaded Kilo catalog", () => { - expect(kiloCatalogModelStatus({ agents: true, loaded: true, providers, modelID: "vendor/missing-model" })).toBe( - "invalid", - ) + it("accepts arbitrary models exposed by a refreshed Kilo catalog", () => { + expect(kiloCatalogModelStatus(providers, "vendor/new-live-model", 1, 2)).toBe("apply") }) - it("applies arbitrary models exposed by the loaded Kilo catalog", () => { - expect(kiloCatalogModelStatus({ agents: true, loaded: true, providers, modelID: "vendor/new-live-model" })).toBe( - "apply", - ) + it("rejects models missing from the refreshed Kilo catalog", () => { + expect(kiloCatalogModelStatus(providers, "vendor/missing-model", 1, 2)).toBe("invalid") }) it("does not accept models exposed only by non-Kilo providers", () => { - expect(kiloCatalogModelStatus({ agents: true, loaded: true, providers, modelID: "gpt-4.1" })).toBe("invalid") - }) -}) - -describe("isCurrentModelSelections", () => { - it("rejects stale hydration responses after local selection changes", () => { - expect(isCurrentModelSelections(0, 1)).toBe(false) - expect(isCurrentModelSelections(1, 1)).toBe(true) - expect(isCurrentModelSelections(undefined, 1)).toBe(true) + expect(kiloCatalogModelStatus(providers, "gpt-4.1", 1, 2)).toBe("invalid") }) }) diff --git a/packages/kilo-vscode/webview-ui/src/App.tsx b/packages/kilo-vscode/webview-ui/src/App.tsx index 3b9a26cbc03..9c240f1d411 100644 --- a/packages/kilo-vscode/webview-ui/src/App.tsx +++ b/packages/kilo-vscode/webview-ui/src/App.tsx @@ -36,18 +36,11 @@ import { NotificationsProvider } from "./context/notifications" import { FeedbackProvider } from "./context/feedback" import { KiloEmbeddingModelsProvider } from "./context/kilo-embedding-models" import type { Message as SDKMessage, Part as SDKPart } from "@kilocode/sdk/v2" -import type { SelectKiloModelMessage } from "./types/messages" import "./styles/chat.css" type ViewType = "newTask" | "marketplace" | "history" | "profile" | "settings" | "subAgentViewer" const VALID_VIEWS = new Set(["newTask", "marketplace", "history", "profile", "settings", "subAgentViewer"]) -function isSelectKiloModelMessage(message: unknown): message is SelectKiloModelMessage { - if (!message || typeof message !== "object") return false - if (!("type" in message) || !("modelID" in message)) return false - return message.type === "selectKiloModel" && typeof message.modelID === "string" -} - /** * Bridge our session store to the DataProvider's expected Data shape. * @@ -268,11 +261,8 @@ const AppContent: Component = () => { setCurrentView("newTask") } - const handleSelectKiloModel = (message: unknown) => { - if (!isSelectKiloModelMessage(message)) return - console.log("[Kilo New] App: selecting linked Kilo model:", message.modelID) - session.selectKiloModel(message.modelID) - setCurrentView("newTask") + const handleKiloModel = (message: { type?: string }) => { + if (message.type === "selectKiloModel") setCurrentView("newTask") } onMount(() => { @@ -293,7 +283,7 @@ const AppContent: Component = () => { session.selectCloudSession(message.sessionId) setCurrentView("newTask") } - handleSelectKiloModel(message) + handleKiloModel(message) handleForked(message) if (message?.type === "viewSubAgentSession" && message.sessionID) { console.log("[Kilo New] App: 🔍 viewSubAgentSession:", message.sessionID) 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 898dc1f9504..a30ba6164ed 100644 --- a/packages/kilo-vscode/webview-ui/src/context/model-selection.ts +++ b/packages/kilo-vscode/webview-ui/src/context/model-selection.ts @@ -43,16 +43,12 @@ export function resolveModelSelection(input: { ) } -export function kiloCatalogModelStatus(input: { - agents: boolean - loaded: boolean - providers: Record - modelID: string -}): "pending" | "invalid" | "apply" { - if (!input.agents || !input.loaded) return "pending" - return input.providers[KILO_PROVIDER_ID]?.models[input.modelID] ? "apply" : "invalid" -} - -export function isCurrentModelSelections(revision: number | undefined, current: number): boolean { - return revision === undefined || revision === current +export function kiloCatalogModelStatus( + providers: Record, + modelID: string, + after: number, + current: number, +): "pending" | "invalid" | "apply" { + if (current <= after) return "pending" + return providers[KILO_PROVIDER_ID]?.models[modelID] ? "apply" : "invalid" } diff --git a/packages/kilo-vscode/webview-ui/src/context/provider-utils.ts b/packages/kilo-vscode/webview-ui/src/context/provider-utils.ts index 5c1bca22e25..6afa2f6b5e3 100644 --- a/packages/kilo-vscode/webview-ui/src/context/provider-utils.ts +++ b/packages/kilo-vscode/webview-ui/src/context/provider-utils.ts @@ -1,4 +1,3 @@ -import { KILO_PROVIDER_ID } from "../../../src/shared/provider-model" import type { Provider, ProviderModel, ModelSelection } from "../types/messages" export type EnrichedModel = ProviderModel & { providerID: string; providerName: string } @@ -42,6 +41,6 @@ export function isModelValid( if (!selection) return false const provider = providers[selection.providerID] if (!provider) return false - if (selection.providerID !== KILO_PROVIDER_ID && !connected.includes(selection.providerID)) return false + if (selection.providerID !== "kilo" && !connected.includes(selection.providerID)) return false return !!provider.models[selection.modelID] } diff --git a/packages/kilo-vscode/webview-ui/src/context/provider.tsx b/packages/kilo-vscode/webview-ui/src/context/provider.tsx index 723b50ba51d..4b173b8baa0 100644 --- a/packages/kilo-vscode/webview-ui/src/context/provider.tsx +++ b/packages/kilo-vscode/webview-ui/src/context/provider.tsx @@ -17,7 +17,6 @@ export type EnrichedModel = ProviderModel & { providerID: string; providerName: interface ProviderContextValue { providers: Accessor> connected: Accessor - loaded: Accessor defaults: Accessor> defaultSelection: Accessor models: Accessor @@ -34,7 +33,6 @@ export const ProviderProvider: ParentComponent = (props) => { const [providers, setProviders] = createSignal>({}) const [connected, setConnected] = createSignal([]) - const [loaded, setLoaded] = createSignal(false) const [defaults, setDefaults] = createSignal>({}) const [defaultSelection, setDefaultSelection] = createSignal(KILO_AUTO) const [authMethods, setAuthMethods] = createSignal>({}) @@ -59,7 +57,6 @@ export const ProviderProvider: ParentComponent = (props) => { setProviders(message.providers) setConnected(message.connected) - setLoaded(true) setDefaults(message.defaults) setDefaultSelection(message.defaultSelection) setAuthMethods(message.authMethods) @@ -95,7 +92,6 @@ export const ProviderProvider: ParentComponent = (props) => { const value: ProviderContextValue = { providers, connected, - loaded, defaults, defaultSelection, models, diff --git a/packages/kilo-vscode/webview-ui/src/context/session.tsx b/packages/kilo-vscode/webview-ui/src/context/session.tsx index be784a20aac..c15908e9662 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session.tsx +++ b/packages/kilo-vscode/webview-ui/src/context/session.tsx @@ -51,7 +51,7 @@ import { upsertSessionToolPart, } from "./session-utils" import { Identifier } from "../utils/id" -import { isCurrentModelSelections, kiloCatalogModelStatus, resolveModelSelection } from "./model-selection" +import { kiloCatalogModelStatus, resolveModelSelection } from "./model-selection" import { resolveMessagePrefs } from "./session-preferences" import { errorIDs } from "./session-errors" import { PartStash } from "./part-stash" @@ -172,7 +172,6 @@ interface SessionContextValue { // Model selection (global, extension-lifetime) selected: (sessionID?: string) => ModelSelection | null selectModel: (providerID: string, modelID: string, sessionID?: string) => void - selectKiloModel: (modelID: string) => void hasModelOverride: (sessionID?: string) => boolean clearModelOverride: (sessionID?: string) => void @@ -346,7 +345,8 @@ export const SessionProvider: ParentComponent = (props) => { const [agents, setAgents] = createSignal([]) const [allAgents, setAllAgents] = createSignal([]) const [defaultAgent, setDefaultAgent] = createSignal("code") - const [pendingKiloModelID, setPendingKiloModelID] = createSignal(null) + const [pendingKiloModel, setPendingKiloModel] = createSignal<{ modelID: string; after: number } | null>(null) + const [catalog, setCatalog] = createSignal(0) let revision = 0 // Skills loaded from the CLI backend @@ -556,25 +556,30 @@ export const SessionProvider: ParentComponent = (props) => { } function selectKiloModel(modelID: string) { - setPendingKiloModelID(modelID) + setPendingKiloModel({ modelID, after: catalog() }) + vscode.postMessage({ type: "requestProviders" }) } + const unsubKiloModel = vscode.onMessage((message: ExtensionMessage) => { + if (message.type === "providersLoaded") { + setCatalog((value) => value + 1) + return + } + if (message.type === "selectKiloModel") selectKiloModel(message.modelID) + }) + onCleanup(unsubKiloModel) + createEffect(() => { - const modelID = pendingKiloModelID() - if (!modelID) return - const status = kiloCatalogModelStatus({ - agents: agents().length > 0, - loaded: provider.loaded(), - providers: provider.providers(), - modelID, - }) + const pending = pendingKiloModel() + if (!pending || agents().length === 0) return + const status = kiloCatalogModelStatus(provider.providers(), pending.modelID, pending.after, catalog()) if (status === "pending") return - setPendingKiloModelID(null) + setPendingKiloModel(null) if (status === "invalid") { - console.warn("[Kilo New] Ignoring unavailable Kilo catalog model:", modelID) + console.warn("[Kilo New] Ignoring unavailable Kilo catalog model:", pending.modelID) return } - applyPersistedModel({ providerID: KILO_PROVIDER_ID, modelID }) + applyPersistedModel({ providerID: KILO_PROVIDER_ID, modelID: pending.modelID }) }) function promptAgent(sessionID?: string) { @@ -807,7 +812,7 @@ export const SessionProvider: ParentComponent = (props) => { // Uses replace semantics so a reset (empty payload) clears old entries. const unsubSelections = vscode.onMessage((message: ExtensionMessage) => { if (message.type !== "modelSelectionsLoaded") return - if (!isCurrentModelSelections(message.revision, revision)) return + if (message.revision !== undefined && message.revision !== revision) return setStore("modelSelections", reconcile(message.selections)) const flags: Record = {} for (const name of Object.keys(message.selections)) { @@ -2526,7 +2531,6 @@ export const SessionProvider: ParentComponent = (props) => { scopedSuggestions, selected, selectModel, - selectKiloModel, hasModelOverride, clearModelOverride, costBreakdown, diff --git a/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx b/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx index 9b069b7938d..9fca0126250 100644 --- a/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx +++ b/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx @@ -95,7 +95,6 @@ const MockProviderProvider: ParentComponent<{ kiloAuth?: boolean }> = (props) => const value = { providers: () => MOCK_PROVIDERS as any, connected: () => ["kilo"], - loaded: () => true, defaults: () => ({}), defaultSelection: () => ({ providerID: "kilo", modelID: "anthropic/claude-sonnet-4-6" }), models: () => MOCK_MODELS, From 4c5b17a0e598ab739074fbf79a9cbefba64623a1 Mon Sep 17 00:00:00 2001 From: Josh Lambert Date: Thu, 11 Jun 2026 10:33:06 -0400 Subject: [PATCH 5/6] refactor(vscode): simplify model deep link handling --- packages/kilo-vscode/src/KiloProvider.ts | 20 +++--- packages/kilo-vscode/src/extension.ts | 4 +- .../src/kilo-provider/model-uri.ts | 9 --- .../unit/kilo-provider-model-uri.test.ts | 71 ------------------- .../tests/unit/model-selection.test.ts | 22 +----- .../tests/unit/model-state.test.ts | 46 ------------ .../kilo-vscode/tests/unit/model-uri.test.ts | 25 ------- .../webview-ui/src/context/model-selection.ts | 11 --- .../webview-ui/src/context/session.tsx | 13 ++-- 9 files changed, 19 insertions(+), 202 deletions(-) delete mode 100644 packages/kilo-vscode/src/kilo-provider/model-uri.ts delete mode 100644 packages/kilo-vscode/tests/unit/kilo-provider-model-uri.test.ts delete mode 100644 packages/kilo-vscode/tests/unit/model-uri.test.ts diff --git a/packages/kilo-vscode/src/KiloProvider.ts b/packages/kilo-vscode/src/KiloProvider.ts index de401802069..ea9de174e07 100644 --- a/packages/kilo-vscode/src/KiloProvider.ts +++ b/packages/kilo-vscode/src/KiloProvider.ts @@ -762,7 +762,14 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper this.visibleTaskStreams.handle(message) switch (message.type) { case "webviewReady": - await this.handleWebviewReady() + console.log("[Kilo New] KiloProvider: ✅ webviewReady received") + this.isWebviewReady = true + this.visibleTaskStreams.clear() + this.flushPendingKiloModel() + await this.syncWebviewState("webviewReady") + this.flushPendingReviewComments() + this.recoverPendingPrompts() + this.readyResolvers.splice(0).forEach((r) => r()) break case "sendMessage": { const msg = message as typeof message & ContextMessage @@ -1197,17 +1204,6 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper this.webviewMessageDisposable = watchFontSizeConfig((msg) => this.postMessage(msg), this.webviewMessageDisposable) } - private async handleWebviewReady(): Promise { - console.log("[Kilo New] KiloProvider: ✅ webviewReady received") - this.isWebviewReady = true - this.visibleTaskStreams.clear() - this.flushPendingKiloModel() - await this.syncWebviewState("webviewReady") - this.flushPendingReviewComments() - this.recoverPendingPrompts() - this.readyResolvers.splice(0).forEach((r) => r()) - } - private handleEditorOpenMessage(message: Parameters[0]): boolean { return handleEditorAction(message, { dir: () => this.getWorkspaceDirectory(this.currentSession?.id), diff --git a/packages/kilo-vscode/src/extension.ts b/packages/kilo-vscode/src/extension.ts index d0c2b058a27..9d19f9aab36 100644 --- a/packages/kilo-vscode/src/extension.ts +++ b/packages/kilo-vscode/src/extension.ts @@ -22,7 +22,6 @@ import { registerToggleAutoApprove } from "./commands/toggle-auto-approve" import { registerHeapSnapshot } from "./commands/heap-snapshot" import { RemoteStatusService } from "./services/RemoteStatusService" import { markWorkspace } from "./util/spotlight" -import { kiloModelFromURI } from "./kilo-provider/model-uri" let agentManager: AgentManagerProvider | undefined let shuttingDown = false @@ -490,7 +489,8 @@ export function activate(context: vscode.ExtensionContext) { return } - const modelID = kiloModelFromURI(uri) + if (uri.path !== "/kilocode/model") return + const modelID = new URLSearchParams(uri.query).get("model") if (!modelID) return console.log("[Kilo New] URI handler: selecting linked Kilo model:", modelID) await vscode.commands.executeCommand(`${KiloProvider.viewType}.focus`) diff --git a/packages/kilo-vscode/src/kilo-provider/model-uri.ts b/packages/kilo-vscode/src/kilo-provider/model-uri.ts deleted file mode 100644 index c97492e039e..00000000000 --- a/packages/kilo-vscode/src/kilo-provider/model-uri.ts +++ /dev/null @@ -1,9 +0,0 @@ -type URI = { - path: string - query: string -} - -export function kiloModelFromURI(uri: URI): string | undefined { - if (uri.path !== "/kilocode/model") return undefined - return new URLSearchParams(uri.query).get("model") || undefined -} diff --git a/packages/kilo-vscode/tests/unit/kilo-provider-model-uri.test.ts b/packages/kilo-vscode/tests/unit/kilo-provider-model-uri.test.ts deleted file mode 100644 index ff4ff60220e..00000000000 --- a/packages/kilo-vscode/tests/unit/kilo-provider-model-uri.test.ts +++ /dev/null @@ -1,71 +0,0 @@ -import { describe, expect, it } from "bun:test" - -// vscode mock is provided by the shared preload (tests/setup/vscode-mock.ts) -const { KiloProvider } = await import("../../src/KiloProvider") - -type Internals = { - webview: { postMessage: (message: unknown) => Promise } | null - isWebviewReady: boolean - syncWebviewState: () => Promise - handleWebviewReady: () => Promise - flushPendingKiloModel: () => void -} - -function connection() { - const state = { connected: true } - return { - state, - getClient: () => { - if (!state.connected) throw new Error("Not connected") - return {} - }, - } -} - -describe("KiloProvider model URI selection", () => { - it("flushes a queued model before ancillary webview sync can fail", async () => { - const service = connection() - const provider = new KiloProvider({} as never, service as never) - const internal = provider as unknown as Internals - const sent: unknown[] = [] - - internal.webview = { - postMessage: async (message) => { - sent.push(message) - return true - }, - } - internal.syncWebviewState = async () => { - throw new Error("profile unavailable") - } - - provider.selectKiloModel("vendor/new-live-model") - expect(sent).toEqual([]) - - await expect(internal.handleWebviewReady()).rejects.toThrow("profile unavailable") - expect(sent).toEqual([{ type: "selectKiloModel", modelID: "vendor/new-live-model" }]) - }) - - it("keeps a queued model until the backend client is available", () => { - const service = connection() - service.state.connected = false - const provider = new KiloProvider({} as never, service as never) - const internal = provider as unknown as Internals - const sent: unknown[] = [] - - internal.webview = { - postMessage: async (message) => { - sent.push(message) - return true - }, - } - internal.isWebviewReady = true - - provider.selectKiloModel("vendor/new-live-model") - expect(sent).toEqual([]) - - service.state.connected = true - internal.flushPendingKiloModel() - expect(sent).toEqual([{ type: "selectKiloModel", modelID: "vendor/new-live-model" }]) - }) -}) diff --git a/packages/kilo-vscode/tests/unit/model-selection.test.ts b/packages/kilo-vscode/tests/unit/model-selection.test.ts index 4c2cf1e51b0..c61d1df1db6 100644 --- a/packages/kilo-vscode/tests/unit/model-selection.test.ts +++ b/packages/kilo-vscode/tests/unit/model-selection.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "bun:test" -import { kiloCatalogModelStatus, resolveModelSelection } from "../../webview-ui/src/context/model-selection" +import { resolveModelSelection } from "../../webview-ui/src/context/model-selection" import { KILO_AUTO, parseModelString } from "../../src/shared/provider-model" import type { Provider } from "../../webview-ui/src/types/messages" @@ -12,7 +12,7 @@ function makeProvider(id: string, name: string, modelIds: string[]): Provider { } const providers = { - kilo: makeProvider("kilo", "Kilo Gateway", ["kilo-auto/free", "vendor/new-live-model"]), + kilo: makeProvider("kilo", "Kilo Gateway", ["kilo-auto/free"]), anthropic: makeProvider("anthropic", "Anthropic", ["claude-sonnet-4"]), openai: makeProvider("openai", "OpenAI", ["gpt-4.1"]), } @@ -38,24 +38,6 @@ describe("parseModelString", () => { }) }) -describe("Kilo catalog model URI validation", () => { - it("waits for a catalog refresh after the link arrives", () => { - expect(kiloCatalogModelStatus(providers, "vendor/new-live-model", 1, 1)).toBe("pending") - }) - - it("accepts arbitrary models exposed by a refreshed Kilo catalog", () => { - expect(kiloCatalogModelStatus(providers, "vendor/new-live-model", 1, 2)).toBe("apply") - }) - - it("rejects models missing from the refreshed Kilo catalog", () => { - expect(kiloCatalogModelStatus(providers, "vendor/missing-model", 1, 2)).toBe("invalid") - }) - - it("does not accept models exposed only by non-Kilo providers", () => { - expect(kiloCatalogModelStatus(providers, "gpt-4.1", 1, 2)).toBe("invalid") - }) -}) - describe("resolveModelSelection", () => { it("prefers a valid override", () => { const result = resolveModelSelection({ diff --git a/packages/kilo-vscode/tests/unit/model-state.test.ts b/packages/kilo-vscode/tests/unit/model-state.test.ts index 54ebc98c148..6b42a142ff5 100644 --- a/packages/kilo-vscode/tests/unit/model-state.test.ts +++ b/packages/kilo-vscode/tests/unit/model-state.test.ts @@ -44,50 +44,4 @@ describe("model state", () => { plan: { providerID: "kilo", modelID: "kilo-auto/free" }, }) }) - - it("waits for queued writes before hydrating selections", async () => { - const api = fs.promises as unknown as { writeFile: (...args: unknown[]) => Promise } - const original = api.writeFile - const gate = Promise.withResolvers() - const started = Promise.withResolvers() - const sent: unknown[] = [] - - api.writeFile = async (...args: unknown[]) => { - if (args[0] === file) { - started.resolve() - await gate.promise - } - await original(...args) - } - - try { - const pending = handleMessage( - "persistModelSelection", - { agent: "code", providerID: "kilo", modelID: "vendor/new-live-model" }, - client, - () => {}, - ) - await started.promise - - const hydration = handleMessage("requestModelSelections", { revision: 3 }, client, (message) => - sent.push(message), - ) - await Promise.resolve() - expect(sent).toEqual([]) - - gate.resolve() - await Promise.all([pending, hydration]) - } finally { - gate.resolve() - api.writeFile = original - } - - expect(sent).toEqual([ - { - type: "modelSelectionsLoaded", - selections: { code: { providerID: "kilo", modelID: "vendor/new-live-model" } }, - revision: 3, - }, - ]) - }) }) diff --git a/packages/kilo-vscode/tests/unit/model-uri.test.ts b/packages/kilo-vscode/tests/unit/model-uri.test.ts deleted file mode 100644 index f37f97e7c6f..00000000000 --- a/packages/kilo-vscode/tests/unit/model-uri.test.ts +++ /dev/null @@ -1,25 +0,0 @@ -import { describe, expect, it } from "bun:test" -import { kiloModelFromURI } from "../../src/kilo-provider/model-uri" - -describe("kiloModelFromURI", () => { - it("accepts arbitrary Kilo catalog model ids", () => { - expect(kiloModelFromURI({ path: "/kilocode/model", query: "model=vendor%2Fnew-live-model" })).toBe( - "vendor/new-live-model", - ) - }) - - it("preserves additional slashes inside model ids", () => { - expect(kiloModelFromURI({ path: "/kilocode/model", query: "model=vendor%2Ffamily%2Fmodel" })).toBe( - "vendor/family/model", - ) - }) - - it("rejects missing and empty model ids", () => { - expect(kiloModelFromURI({ path: "/kilocode/model", query: "" })).toBeUndefined() - expect(kiloModelFromURI({ path: "/kilocode/model", query: "model=" })).toBeUndefined() - }) - - it("rejects unrelated URI paths", () => { - expect(kiloModelFromURI({ path: "/kilocode/s/session-id", query: "model=vendor%2Fmodel" })).toBeUndefined() - }) -}) 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 a30ba6164ed..86d9a255b98 100644 --- a/packages/kilo-vscode/webview-ui/src/context/model-selection.ts +++ b/packages/kilo-vscode/webview-ui/src/context/model-selection.ts @@ -1,5 +1,4 @@ import type { ModelSelection, Provider } from "../types/messages" -import { KILO_PROVIDER_ID } from "../../../src/shared/provider-model" import { isModelValid } from "./provider-utils" function validate( @@ -42,13 +41,3 @@ export function resolveModelSelection(input: { null ) } - -export function kiloCatalogModelStatus( - providers: Record, - modelID: string, - after: number, - current: number, -): "pending" | "invalid" | "apply" { - if (current <= after) return "pending" - return providers[KILO_PROVIDER_ID]?.models[modelID] ? "apply" : "invalid" -} diff --git a/packages/kilo-vscode/webview-ui/src/context/session.tsx b/packages/kilo-vscode/webview-ui/src/context/session.tsx index 4126641c186..ec47ab50f3e 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session.tsx +++ b/packages/kilo-vscode/webview-ui/src/context/session.tsx @@ -62,7 +62,7 @@ import { upsertSessionToolPart, } from "./session-utils" import { Identifier } from "../utils/id" -import { kiloCatalogModelStatus, resolveModelSelection } from "./model-selection" +import { resolveModelSelection } from "./model-selection" import { resolveMessagePrefs } from "./session-preferences" import { errorIDs } from "./session-errors" import { PartStash } from "./part-stash" @@ -583,11 +583,9 @@ export const SessionProvider: ParentComponent = (props) => { createEffect(() => { const pending = pendingKiloModel() - if (!pending || agents().length === 0) return - const status = kiloCatalogModelStatus(provider.providers(), pending.modelID, pending.after, catalog()) - if (status === "pending") return + if (!pending || agents().length === 0 || catalog() <= pending.after) return setPendingKiloModel(null) - if (status === "invalid") { + if (!provider.providers()[KILO_PROVIDER_ID]?.models[pending.modelID]) { console.warn("[Kilo New] Ignoring unavailable Kilo catalog model:", pending.modelID) return } @@ -824,7 +822,10 @@ export const SessionProvider: ParentComponent = (props) => { // Uses replace semantics so a reset (empty payload) clears old entries. const unsubSelections = vscode.onMessage((message: ExtensionMessage) => { if (message.type !== "modelSelectionsLoaded") return - if (message.revision !== undefined && message.revision !== revision) return + if (message.revision !== undefined && message.revision !== revision) { + vscode.postMessage({ type: "requestModelSelections", revision }) + return + } setStore("modelSelections", reconcile(message.selections)) const flags: Record = {} for (const name of Object.keys(message.selections)) { From 14d0cc2b9661f40f68445cd92f699fb3cf826f50 Mon Sep 17 00:00:00 2001 From: Josh Lambert Date: Thu, 11 Jun 2026 17:19:49 -0400 Subject: [PATCH 6/6] refactor(vscode): reuse model selection for deep links --- .../src/kilo-provider/model-state.ts | 34 ++++++------- .../tests/unit/model-state.test.ts | 47 ----------------- .../webview-ui/src/context/session.tsx | 51 +++++-------------- .../src/types/messages/extension-messages.ts | 1 - .../src/types/messages/webview-messages.ts | 1 - 5 files changed, 27 insertions(+), 107 deletions(-) delete mode 100644 packages/kilo-vscode/tests/unit/model-state.test.ts diff --git a/packages/kilo-vscode/src/kilo-provider/model-state.ts b/packages/kilo-vscode/src/kilo-provider/model-state.ts index 7f0d242d1f8..c97775bb2d1 100644 --- a/packages/kilo-vscode/src/kilo-provider/model-state.ts +++ b/packages/kilo-vscode/src/kilo-provider/model-state.ts @@ -41,12 +41,12 @@ async function read(client: KiloClient | null): Promise> } } -function write(client: KiloClient | null, key: string, update: (value: unknown) => unknown): Promise { +function write(client: KiloClient | null, key: string, value: unknown): Promise { const op = queue.then(async () => { const p = await resolve(client) if (!p) return const existing = await read(client) - existing[key] = update(existing[key]) + existing[key] = value await fs.promises.writeFile(p, JSON.stringify(existing, null, 2)) }) queue = op.catch(() => {}) @@ -63,36 +63,32 @@ export async function handleMessage( post: PostMessage, ): Promise { if (type === "persistModelSelection") { - await write(client, "model", (value) => { - const model = validateModelSelections(value) - model[message.agent as string] = { - providerID: message.providerID as string, - modelID: message.modelID as string, - } - return model - }) + const data = await read(client) + const model = validateModelSelections(data.model) + model[message.agent as string] = { + providerID: message.providerID as string, + modelID: message.modelID as string, + } + await write(client, "model", model) return true } if (type === "clearModelSelection") { - await write(client, "model", (value) => { - const model = validateModelSelections(value) - delete model[message.agent as string] - return model - }) + const data = await read(client) + const model = validateModelSelections(data.model) + delete model[message.agent as string] + await write(client, "model", model) return true } if (type === "requestModelSelections") { - await queue const data = await read(client) const selections = validateModelSelections(data.model) - const revision = typeof message.revision === "number" ? message.revision : undefined - post({ type: "modelSelectionsLoaded", selections, revision }) + post({ type: "modelSelectionsLoaded", selections }) return true } return false } export async function reset(client: KiloClient | null, post: PostMessage): Promise { - await write(client, "model", () => ({})) + await write(client, "model", {}) post({ type: "modelSelectionsLoaded", selections: {} }) } diff --git a/packages/kilo-vscode/tests/unit/model-state.test.ts b/packages/kilo-vscode/tests/unit/model-state.test.ts deleted file mode 100644 index 6b42a142ff5..00000000000 --- a/packages/kilo-vscode/tests/unit/model-state.test.ts +++ /dev/null @@ -1,47 +0,0 @@ -import { afterAll, beforeEach, describe, expect, it } from "bun:test" -import * as fs from "fs" -import * as os from "os" -import * as path from "path" -import type { KiloClient } from "@kilocode/sdk/v2/client" -import { handleMessage } from "../../src/kilo-provider/model-state" - -const root = fs.mkdtempSync(path.join(os.tmpdir(), "model-state-")) -const file = path.join(root, "model.json") -const client = { - path: { - get: async () => ({ data: { state: root } }), - }, -} as unknown as KiloClient - -afterAll(() => { - fs.rmSync(root, { recursive: true, force: true }) -}) - -beforeEach(() => { - fs.writeFileSync(file, JSON.stringify({ model: {} })) -}) - -describe("model state", () => { - it("preserves concurrent selections for different agents", async () => { - await Promise.all([ - handleMessage( - "persistModelSelection", - { agent: "code", providerID: "kilo", modelID: "vendor/new-live-model" }, - client, - () => {}, - ), - handleMessage( - "persistModelSelection", - { agent: "plan", providerID: "kilo", modelID: "kilo-auto/free" }, - client, - () => {}, - ), - ]) - - const data = JSON.parse(fs.readFileSync(file, "utf-8")) as { model: Record } - expect(data.model).toEqual({ - code: { providerID: "kilo", modelID: "vendor/new-live-model" }, - plan: { providerID: "kilo", modelID: "kilo-auto/free" }, - }) - }) -}) diff --git a/packages/kilo-vscode/webview-ui/src/context/session.tsx b/packages/kilo-vscode/webview-ui/src/context/session.tsx index ec47ab50f3e..39d5dfaca24 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session.tsx +++ b/packages/kilo-vscode/webview-ui/src/context/session.tsx @@ -359,7 +359,6 @@ export const SessionProvider: ParentComponent = (props) => { const [defaultAgent, setDefaultAgent] = createSignal("code") const [pendingKiloModel, setPendingKiloModel] = createSignal<{ modelID: string; after: number } | null>(null) const [catalog, setCatalog] = createSignal(0) - let revision = 0 // Skills loaded from the CLI backend const [skills, setSkills] = createSignal([]) @@ -521,18 +520,6 @@ export const SessionProvider: ParentComponent = (props) => { vscode.postMessage({ type: "persistRecents", recents: updated }) } - function persistModelSelection(agentName: string, selection: ModelSelection) { - revision++ - setUserSetAgents((prev) => ({ ...prev, [agentName]: true })) - setStore("modelSelections", agentName, selection) - vscode.postMessage({ - type: "persistModelSelection", - agent: agentName, - providerID: selection.providerID, - modelID: selection.modelID, - }) - } - function applyModel(agentName: string, selection: ModelSelection, sessionID?: string) { pushRecent(selection) if (sessionID) { @@ -541,7 +528,15 @@ export const SessionProvider: ParentComponent = (props) => { } // Always remember the per-mode model choice so switching modes restores // the last-used model (mirrors CLI TUI's model.json behavior). - persistModelSelection(agentName, selection) + 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, + }) } function selectModel(providerID: string, modelID: string, sessionID?: string) { @@ -552,21 +547,6 @@ export const SessionProvider: ParentComponent = (props) => { } } - function applyPersistedModel(selection: ModelSelection) { - const sid = currentSessionID() - const defaultAgentName = defaultAgent() - const activeAgentName = agentForScope() - pushRecent(selection) - persistModelSelection(defaultAgentName, selection) - if (activeAgentName !== defaultAgentName) { - persistModelSelection(activeAgentName, selection) - } - if (sid) { - setStore("sessionOverrides", sid, selection) - hideErrors(sid) - } - } - function selectKiloModel(modelID: string) { setPendingKiloModel({ modelID, after: catalog() }) vscode.postMessage({ type: "requestProviders" }) @@ -589,7 +569,7 @@ export const SessionProvider: ParentComponent = (props) => { console.warn("[Kilo New] Ignoring unavailable Kilo catalog model:", pending.modelID) return } - applyPersistedModel({ providerID: KILO_PROVIDER_ID, modelID: pending.modelID }) + selectModel(KILO_PROVIDER_ID, pending.modelID) }) function promptAgent(sessionID?: string) { @@ -619,10 +599,7 @@ export const SessionProvider: ParentComponent = (props) => { delete selections[agentName] }), ) - if (persist) { - revision++ - vscode.postMessage({ type: "clearModelSelection", agent: agentName }) - } + if (persist) vscode.postMessage({ type: "clearModelSelection", agent: agentName }) } function shouldClearModeModelSelection(agentName: string) { @@ -822,10 +799,6 @@ export const SessionProvider: ParentComponent = (props) => { // Uses replace semantics so a reset (empty payload) clears old entries. const unsubSelections = vscode.onMessage((message: ExtensionMessage) => { if (message.type !== "modelSelectionsLoaded") return - if (message.revision !== undefined && message.revision !== revision) { - vscode.postMessage({ type: "requestModelSelections", revision }) - return - } setStore("modelSelections", reconcile(message.selections)) const flags: Record = {} for (const name of Object.keys(message.selections)) { @@ -833,7 +806,7 @@ export const SessionProvider: ParentComponent = (props) => { } setUserSetAgents(flags) }) - vscode.postMessage({ type: "requestModelSelections", revision }) + vscode.postMessage({ type: "requestModelSelections" }) onCleanup(unsubSelections) // Load persisted recent models from extension globalState 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 492ee38b305..7fc8d4e4f4b 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 @@ -642,7 +642,6 @@ export interface FavoritesLoadedMessage { export interface ModelSelectionsLoadedMessage { type: "modelSelectionsLoaded" selections: Record - revision?: number } 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 a064f604d26..f0c0f315c0a 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 @@ -989,7 +989,6 @@ export interface ClearModelSelectionRequest { export interface RequestModelSelectionsMessage { type: "requestModelSelections" - revision: number } // Continue in Worktree: transfer sidebar session + git state to an isolated worktree