diff --git a/packages/kilo-vscode/src/KiloProvider.ts b/packages/kilo-vscode/src/KiloProvider.ts index b18192e38f4..d0d3027b985 100644 --- a/packages/kilo-vscode/src/KiloProvider.ts +++ b/packages/kilo-vscode/src/KiloProvider.ts @@ -106,13 +106,13 @@ import { } from "./kilo-provider/handlers/question" import { fetchAndSendPendingSuggestions, routeSuggestionWebviewMessage } from "./kilo-provider/handlers/suggestion" import { nativeTitle } from "./kilo-provider/native-tab-title" +import * as Preferences from "./kilo-provider/model-preferences" +import { handleEnhancePrompt } from "./kilo-provider/enhance-prompt" import { buildActionContext, computeDefaultSelection, fetchProviderData, - validateRecents, - validateFavorites, connectProvider as connectProviderAction, authorizeProviderOAuth as authorizeOAuthAction, completeProviderOAuth as completeOAuthAction, @@ -319,6 +319,25 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper } } + private get preferenceCtx(): Parameters[0] { + return { + extensionContext: this.extensionContext, + postMessage: (msg) => this.postMessage(msg), + notifyFavoritesChanged: (favorites) => this.connectionService.notifyFavoritesChanged(favorites), + } + } + + private get enhancePromptCtx(): Parameters[0] { + return { + client: this.client, + postMessage: (msg) => this.postMessage(msg), + getErrorMessage, + showErrorMessage: (msg) => { + void vscode.window.showErrorMessage(msg) + }, + } + } + // Strip edit-tool metadata.filediff.before/after (multi-MB for edit-heavy // sessions) to keep session switches fast. Logic in kilo-provider/slim-metadata.ts. private slimPart(part: T): T { @@ -971,41 +990,29 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper TelemetryProxy.capture(message.event, message.properties) break case "persistVariant": { - const stored = this.extensionContext?.globalState.get>("variantSelections") ?? {} - stored[message.key] = message.value - await this.extensionContext?.globalState.update("variantSelections", stored) + await Preferences.persistVariant(this.preferenceCtx, message.key, message.value) break } case "requestVariants": { - const variants = this.extensionContext?.globalState.get>("variantSelections") ?? {} - this.postMessage({ type: "variantsLoaded", variants }) + Preferences.requestVariants(this.preferenceCtx) break } case "persistRecents": - await this.extensionContext?.globalState.update("recentModels", validateRecents(message.recents)) + await Preferences.persistRecents(this.preferenceCtx, message.recents) break case "requestRecents": { - const recents = validateRecents(this.extensionContext?.globalState.get("recentModels")) - this.postMessage({ type: "recentsLoaded", recents }) + Preferences.requestRecents(this.preferenceCtx) break } case "toggleFavorite": { - const current = validateFavorites(this.extensionContext?.globalState.get("favoriteModels")) - const key = `${message.providerID}/${message.modelID}` - const exists = current.some((f) => `${f.providerID}/${f.modelID}` === key) - const favorites = - message.action === "add" && !exists - ? [...current, { providerID: message.providerID, modelID: message.modelID }] - : message.action === "remove" && exists - ? current.filter((f) => `${f.providerID}/${f.modelID}` !== key) - : current - await this.extensionContext?.globalState.update("favoriteModels", favorites) - this.connectionService.notifyFavoritesChanged(favorites) + await Preferences.toggleFavorite(this.preferenceCtx, message.action, { + providerID: message.providerID, + modelID: message.modelID, + }) break } case "requestFavorites": { - const favorites = validateFavorites(this.extensionContext?.globalState.get("favoriteModels")) - this.postMessage({ type: "favoritesLoaded", favorites }) + Preferences.requestFavorites(this.preferenceCtx) break } // legacy-migration start @@ -1026,30 +1033,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper break // legacy-migration end case "enhancePrompt": { - const sdkClient = this.client - if (!sdkClient) { - this.postMessage({ - type: "enhancePromptError", - error: "Not connected to CLI backend", - requestId: message.requestId, - }) - break - } - void sdkClient.enhancePrompt - .enhance({ text: message.text }, { throwOnError: true }) - .then(({ data }) => { - this.postMessage({ type: "enhancePromptResult", text: data.text, requestId: message.requestId }) - }) - .catch((err: unknown) => { - const msg = getErrorMessage(err) || "Failed to enhance prompt" - console.error("[Kilo New] KiloProvider: Failed to enhance prompt:", err) - vscode.window.showErrorMessage(`Enhance prompt failed: ${msg}`) - this.postMessage({ - type: "enhancePromptError", - error: msg, - requestId: message.requestId, - }) - }) + handleEnhancePrompt(this.enhancePromptCtx, message.text, message.requestId) break } case "fetchMarketplaceData": { diff --git a/packages/kilo-vscode/src/kilo-provider/enhance-prompt.ts b/packages/kilo-vscode/src/kilo-provider/enhance-prompt.ts new file mode 100644 index 00000000000..43ae66706be --- /dev/null +++ b/packages/kilo-vscode/src/kilo-provider/enhance-prompt.ts @@ -0,0 +1,38 @@ +import type { KiloClient } from "@kilocode/sdk/v2/client" + +type Post = (msg: unknown) => void + +interface Context { + readonly client: KiloClient | null + postMessage: Post + getErrorMessage(error: unknown): string + showErrorMessage(message: string): void +} + +export function handleEnhancePrompt(ctx: Context, text: string, requestId: string): void { + const client = ctx.client + if (!client) { + ctx.postMessage({ + type: "enhancePromptError", + error: "Not connected to CLI backend", + requestId, + }) + return + } + + void client.enhancePrompt + .enhance({ text }, { throwOnError: true }) + .then(({ data }) => { + ctx.postMessage({ type: "enhancePromptResult", text: data.text, requestId }) + }) + .catch((err: unknown) => { + const msg = ctx.getErrorMessage(err) || "Failed to enhance prompt" + console.error("[Kilo New] KiloProvider: Failed to enhance prompt:", err) + ctx.showErrorMessage(`Enhance prompt failed: ${msg}`) + ctx.postMessage({ + type: "enhancePromptError", + error: msg, + requestId, + }) + }) +} diff --git a/packages/kilo-vscode/src/kilo-provider/model-preferences.ts b/packages/kilo-vscode/src/kilo-provider/model-preferences.ts new file mode 100644 index 00000000000..975bc03a8ad --- /dev/null +++ b/packages/kilo-vscode/src/kilo-provider/model-preferences.ts @@ -0,0 +1,76 @@ +import type { ExtensionContext } from "vscode" + +type Model = { providerID: string; modelID: string } +type Post = (msg: unknown) => void + +interface Context { + readonly extensionContext?: ExtensionContext + postMessage: Post + notifyFavoritesChanged(favorites: Model[]): void +} + +function selection(value: unknown): value is Model { + return ( + !!value && + typeof value === "object" && + typeof (value as Record).providerID === "string" && + typeof (value as Record).modelID === "string" + ) +} + +export function validateRecents(raw: unknown): Model[] { + if (!Array.isArray(raw)) return [] + return raw + .filter(selection) + .slice(0, 5) + .map((item) => ({ providerID: item.providerID, modelID: item.modelID })) +} + +export function validateFavorites(raw: unknown): Model[] { + if (!Array.isArray(raw)) return [] + return raw.filter(selection).map((item) => ({ providerID: item.providerID, modelID: item.modelID })) +} + +export function updateFavorites(current: Model[], action: "add" | "remove", model: Model): Model[] { + const key = `${model.providerID}/${model.modelID}` + const exists = current.some((item) => `${item.providerID}/${item.modelID}` === key) + if (action === "add" && !exists) return [...current, model] + if (action === "remove" && exists) return current.filter((item) => `${item.providerID}/${item.modelID}` !== key) + return current +} + +export async function persistVariant(ctx: Context, key: string, value: string): Promise { + const stored = ctx.extensionContext?.globalState.get>("variantSelections") ?? {} + stored[key] = value + await ctx.extensionContext?.globalState.update("variantSelections", stored) +} + +export function requestVariants(ctx: Context): void { + const variants = ctx.extensionContext?.globalState.get>("variantSelections") ?? {} + ctx.postMessage({ type: "variantsLoaded", variants }) +} + +export async function persistRecents(ctx: Context, recents: unknown): Promise { + const valid = validateRecents(recents) + await ctx.extensionContext?.globalState.update("recentModels", valid) +} + +export function requestRecents(ctx: Context): void { + const stored = ctx.extensionContext?.globalState.get("recentModels") + const recents = validateRecents(stored) + ctx.postMessage({ type: "recentsLoaded", recents }) +} + +export async function toggleFavorite(ctx: Context, action: "add" | "remove", model: Model): Promise { + const stored = ctx.extensionContext?.globalState.get("favoriteModels") + const current = validateFavorites(stored) + const favorites = updateFavorites(current, action, model) + await ctx.extensionContext?.globalState.update("favoriteModels", favorites) + ctx.notifyFavoritesChanged(favorites) +} + +export function requestFavorites(ctx: Context): void { + const stored = ctx.extensionContext?.globalState.get("favoriteModels") + const favorites = validateFavorites(stored) + ctx.postMessage({ type: "favoritesLoaded", favorites }) +} diff --git a/packages/kilo-vscode/src/provider-actions.ts b/packages/kilo-vscode/src/provider-actions.ts index 3d8e4089224..8a5052091ec 100644 --- a/packages/kilo-vscode/src/provider-actions.ts +++ b/packages/kilo-vscode/src/provider-actions.ts @@ -105,21 +105,6 @@ function isModelSelection(r: unknown): r is { providerID: string; modelID: strin ) } -/** Validate and sanitize recent model selections from untrusted sources. */ -export function validateRecents(raw: unknown): Array<{ providerID: string; modelID: string }> { - if (!Array.isArray(raw)) return [] - return raw - .filter(isModelSelection) - .slice(0, 5) - .map((r) => ({ providerID: r.providerID, modelID: r.modelID })) -} - -/** Validate and sanitize favorite model selections from untrusted sources. */ -export function validateFavorites(raw: unknown): Array<{ providerID: string; modelID: string }> { - if (!Array.isArray(raw)) return [] - return raw.filter(isModelSelection).map((r) => ({ providerID: r.providerID, modelID: r.modelID })) -} - /** Validate and sanitize per-mode model selections from untrusted sources. */ export function validateModelSelections(raw: unknown): Record { if (!raw || typeof raw !== "object" || Array.isArray(raw)) return {} diff --git a/packages/kilo-vscode/tests/unit/enhance-prompt.test.ts b/packages/kilo-vscode/tests/unit/enhance-prompt.test.ts new file mode 100644 index 00000000000..f9ab10d86d8 --- /dev/null +++ b/packages/kilo-vscode/tests/unit/enhance-prompt.test.ts @@ -0,0 +1,83 @@ +import { describe, expect, it, spyOn } from "bun:test" +import { handleEnhancePrompt } from "../../src/kilo-provider/enhance-prompt" + +type Context = Parameters[0] +type Client = NonNullable + +async function tick() { + await new Promise((resolve) => setTimeout(resolve, 0)) +} + +function createCtx(client: Client | null, error = "boom") { + const calls = { + posts: [] as unknown[], + errors: [] as string[], + } + const ctx: Context = { + client, + postMessage: (msg) => calls.posts.push(msg), + getErrorMessage: () => error, + showErrorMessage: (msg) => calls.errors.push(msg), + } + return { calls, ctx } +} + +describe("handleEnhancePrompt", () => { + it("posts an error when the client is unavailable", () => { + const { calls, ctx } = createCtx(null) + + handleEnhancePrompt(ctx, "draft", "req-1") + + expect(calls.posts).toEqual([ + { + type: "enhancePromptError", + error: "Not connected to CLI backend", + requestId: "req-1", + }, + ]) + expect(calls.errors).toEqual([]) + }) + + it("posts enhanced text on success", async () => { + const client = { + enhancePrompt: { + enhance: async (input: { text?: string }) => ({ data: { text: `better ${input.text}` } }), + }, + } as unknown as Client + const { calls, ctx } = createCtx(client) + + handleEnhancePrompt(ctx, "draft", "req-2") + await tick() + + expect(calls.posts).toEqual([{ type: "enhancePromptResult", text: "better draft", requestId: "req-2" }]) + expect(calls.errors).toEqual([]) + }) + + it("posts and shows the error message on failure", async () => { + const log = spyOn(console, "error").mockImplementation(() => {}) + try { + const client = { + enhancePrompt: { + enhance: async () => { + throw new Error("server failed") + }, + }, + } as unknown as Client + const { calls, ctx } = createCtx(client, "server failed") + + handleEnhancePrompt(ctx, "draft", "req-3") + await tick() + + expect(calls.errors).toEqual(["Enhance prompt failed: server failed"]) + expect(calls.posts).toEqual([ + { + type: "enhancePromptError", + error: "server failed", + requestId: "req-3", + }, + ]) + } finally { + log.mockRestore() + } + }) +}) diff --git a/packages/kilo-vscode/tests/unit/model-preferences.test.ts b/packages/kilo-vscode/tests/unit/model-preferences.test.ts new file mode 100644 index 00000000000..f5cf231cde9 --- /dev/null +++ b/packages/kilo-vscode/tests/unit/model-preferences.test.ts @@ -0,0 +1,36 @@ +import { describe, expect, it } from "bun:test" +import { updateFavorites } from "../../src/kilo-provider/model-preferences" + +describe("updateFavorites", () => { + it("adds a favorite that is not present", () => { + const current = [{ providerID: "anthropic", modelID: "claude-sonnet-4" }] + + expect(updateFavorites(current, "add", { providerID: "openai", modelID: "gpt-5" })).toEqual([ + { providerID: "anthropic", modelID: "claude-sonnet-4" }, + { providerID: "openai", modelID: "gpt-5" }, + ]) + }) + + it("does not add a duplicate favorite", () => { + const current = [{ providerID: "openai", modelID: "gpt-5" }] + + expect(updateFavorites(current, "add", { providerID: "openai", modelID: "gpt-5" })).toBe(current) + }) + + it("removes an existing favorite", () => { + const current = [ + { providerID: "anthropic", modelID: "claude-sonnet-4" }, + { providerID: "openai", modelID: "gpt-5" }, + ] + + expect(updateFavorites(current, "remove", { providerID: "openai", modelID: "gpt-5" })).toEqual([ + { providerID: "anthropic", modelID: "claude-sonnet-4" }, + ]) + }) + + it("keeps favorites unchanged when removing a missing favorite", () => { + const current = [{ providerID: "anthropic", modelID: "claude-sonnet-4" }] + + expect(updateFavorites(current, "remove", { providerID: "openai", modelID: "gpt-5" })).toBe(current) + }) +}) diff --git a/packages/kilo-vscode/tests/unit/provider-actions-validate.test.ts b/packages/kilo-vscode/tests/unit/provider-actions-validate.test.ts index 688f917d8d6..7c82d7904a2 100644 --- a/packages/kilo-vscode/tests/unit/provider-actions-validate.test.ts +++ b/packages/kilo-vscode/tests/unit/provider-actions-validate.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "bun:test" -import { validateModelSelections, validateRecents, validateFavorites } from "../../src/provider-actions" +import { validateModelSelections } from "../../src/provider-actions" describe("validateModelSelections", () => { it("returns empty object for null", () => {