diff --git a/packages/kilo-vscode/src/services/autocomplete/AutocompleteModel.ts b/packages/kilo-vscode/src/services/autocomplete/AutocompleteModel.ts index ea0570c216f..8fd77bd4a49 100644 --- a/packages/kilo-vscode/src/services/autocomplete/AutocompleteModel.ts +++ b/packages/kilo-vscode/src/services/autocomplete/AutocompleteModel.ts @@ -54,12 +54,7 @@ export class AutocompleteModel { throw new Error("Connection service is not available") } - const state = this.connectionService.getConnectionState() - if (state !== "connected") { - throw new Error(`CLI backend is not connected (state: ${state})`) - } - - const client = this.connectionService.getClient() + const client = await this.connectionService.getClientAsync() let cost = 0 let inputTokens = 0 @@ -146,13 +141,13 @@ export class AutocompleteModel { * Returns false on any error (not connected, fetch failed, etc.). */ public async hasBalance(): Promise { - if (!this.connectionService || this.connectionService.getConnectionState() !== "connected") { + if (!this.connectionService) return false + try { + const client = await this.connectionService.getClientAsync() + const result = await client.kilo.profile().catch(() => null) + return (result?.data?.balance?.balance ?? 0) > 0 + } catch { return false } - const result = await this.connectionService - .getClient() - .kilo.profile() - .catch(() => null) - return (result?.data?.balance?.balance ?? 0) > 0 } } diff --git a/packages/kilo-vscode/src/services/autocomplete/__tests__/AutocompleteModel.spec.ts b/packages/kilo-vscode/src/services/autocomplete/__tests__/AutocompleteModel.spec.ts index d94bcd0612e..e98fc60cc67 100644 --- a/packages/kilo-vscode/src/services/autocomplete/__tests__/AutocompleteModel.spec.ts +++ b/packages/kilo-vscode/src/services/autocomplete/__tests__/AutocompleteModel.spec.ts @@ -1,20 +1,30 @@ -import { describe, it, expect, vi } from "vitest" +import { describe, it, expect, vi, beforeEach } from "vitest" import { AutocompleteModel } from "../AutocompleteModel" import type { KiloConnectionService } from "../../cli-backend" +const mockClient = { + kilo: { + fim: vi.fn(), + }, +} + function createMockConnectionService(state: "connecting" | "connected" | "disconnected" | "error" = "connected") { return { getConnectionState: vi.fn().mockReturnValue(state), - getClient: vi.fn().mockReturnValue({ - kilo: { - fim: vi.fn(), - }, - }), + getClient: vi.fn().mockReturnValue(mockClient), + getClientAsync: + state === "connected" + ? vi.fn().mockResolvedValue(mockClient) + : vi.fn().mockRejectedValue(new Error(`CLI backend is not connected (state: ${state})`)), onStateChange: vi.fn().mockReturnValue(() => {}), } as unknown as KiloConnectionService } describe("AutocompleteModel", () => { + beforeEach(() => { + mockClient.kilo.fim.mockReset() + }) + describe("constructor", () => { it("defaults profileName and profileType to null", () => { const model = new AutocompleteModel() @@ -98,7 +108,7 @@ describe("AutocompleteModel", () => { const connection = createMockConnectionService("disconnected") const model = new AutocompleteModel(connection) await expect(model.generateFimResponse("prefix", "suffix", vi.fn())).rejects.toThrow( - "CLI backend is not connected (state: disconnected)", + "CLI backend is not connected", ) }) @@ -113,8 +123,7 @@ describe("AutocompleteModel", () => { ] const connection = createMockConnectionService("connected") - const client = (connection as any).getClient() - client.kilo.fim.mockResolvedValue({ + mockClient.kilo.fim.mockResolvedValue({ stream: (async function* () { for (const chunk of chunks) yield chunk })(), @@ -136,8 +145,7 @@ describe("AutocompleteModel", () => { it("passes model parameters to fim call", async () => { const connection = createMockConnectionService("connected") - const client = (connection as any).getClient() - client.kilo.fim.mockResolvedValue({ + mockClient.kilo.fim.mockResolvedValue({ stream: (async function* () {})(), }) @@ -145,7 +153,7 @@ describe("AutocompleteModel", () => { const signal = new AbortController().signal await model.generateFimResponse("pre", "suf", vi.fn(), signal) - expect(client.kilo.fim).toHaveBeenCalledWith( + expect(mockClient.kilo.fim).toHaveBeenCalledWith( { prefix: "pre", suffix: "suf", @@ -153,7 +161,7 @@ describe("AutocompleteModel", () => { maxTokens: 256, temperature: 0.2, }, - { signal }, + expect.objectContaining({ signal }), ) }) }) diff --git a/packages/kilo-vscode/src/services/cli-backend/connection-service.ts b/packages/kilo-vscode/src/services/cli-backend/connection-service.ts index 22d8306eda2..d3fdf70415e 100644 --- a/packages/kilo-vscode/src/services/cli-backend/connection-service.ts +++ b/packages/kilo-vscode/src/services/cli-backend/connection-service.ts @@ -121,6 +121,20 @@ export class KiloConnectionService { return this.client } + /** + * Get the shared SDK client, auto-connecting if not yet started. + * Accepts an optional directory to use as the workspace root; falls back + * to the first VS Code workspace folder. Throws if neither is available + * or if the connection fails. + */ + async getClientAsync(dir?: string): Promise { + if (this.client) return this.client + const root = dir ?? vscode.workspace.workspaceFolders?.[0]?.uri.fsPath + if (!root) throw new Error("No workspace folder open") + await this.connect(root) + return this.client! + } + /** * Get server info (port). Returns null if not connected. */ diff --git a/packages/kilo-vscode/src/services/commit-message/__tests__/index.spec.ts b/packages/kilo-vscode/src/services/commit-message/__tests__/index.spec.ts index 124002d8028..f35427f1513 100644 --- a/packages/kilo-vscode/src/services/commit-message/__tests__/index.spec.ts +++ b/packages/kilo-vscode/src/services/commit-message/__tests__/index.spec.ts @@ -54,7 +54,7 @@ describe("commit-message service", () => { } mockConnectionService = { - getClient: vi.fn().mockReturnValue(mockClient), + getClientAsync: vi.fn().mockResolvedValue(mockClient), } as any }) @@ -115,7 +115,7 @@ describe("commit-message service", () => { expect(vscode.window.showErrorMessage).toHaveBeenCalledWith("No Git repository found") }) - it("shows error when backend is not connected", async () => { + it("shows error when backend fails to connect", async () => { ;(vscode.extensions.getExtension as Mock).mockReturnValue({ isActive: true, activate: vi.fn().mockResolvedValue(undefined), @@ -125,17 +125,38 @@ describe("commit-message service", () => { }), }, }) - ;(mockConnectionService.getClient as Mock).mockImplementation(() => { - throw new Error("Not connected") - }) + ;(mockConnectionService.getClientAsync as Mock).mockRejectedValue(new Error("Connect failed")) await commandCallback() expect(vscode.window.showErrorMessage).toHaveBeenCalledWith( - "Kilo backend is not connected. Please wait for the connection to establish.", + "Failed to connect to Kilo backend. Please try again.", ) }) + it("auto-connects backend and generates message when client not yet ready", async () => { + const mockInputBox = { value: "" } + ;(vscode.extensions.getExtension as Mock).mockReturnValue({ + isActive: true, + activate: vi.fn().mockResolvedValue(undefined), + exports: { + getAPI: () => ({ + repositories: [{ inputBox: mockInputBox, rootUri: { fsPath: "/auto-connect-repo" } }], + }), + }, + }) + + const mockToken = { onCancellationRequested: vi.fn() } + ;(vscode.window.withProgress as Mock).mockImplementation(async (_options: any, task: any) => { + await task({}, mockToken) + }) + + await commandCallback() + + expect(mockConnectionService.getClientAsync).toHaveBeenCalled() + expect(mockInputBox.value).toBe("feat: add new feature") + }) + it("calls commitMessage.generate on the SDK client with repository root path", async () => { const mockInputBox = { value: "" } ;(vscode.extensions.getExtension as Mock).mockReturnValue({ @@ -229,7 +250,7 @@ describe("commit-message service", () => { } as any) vi.mocked(vscode.window.withProgress).mockImplementation(async (_options, task) => { - await task({} as any, {} as any) + await task({} as any, { onCancellationRequested: vi.fn() } as any) }) // Simulate SCM title/input passing the SourceControl for the worktree repo @@ -254,7 +275,7 @@ describe("commit-message service", () => { } as any) vi.mocked(vscode.window.withProgress).mockImplementation(async (_options, task) => { - await task({} as any, {} as any) + await task({} as any, { onCancellationRequested: vi.fn() } as any) }) const scmArg = { rootUri: { fsPath: "/nonexistent-repo" } } as vscode.SourceControl diff --git a/packages/kilo-vscode/src/services/commit-message/index.ts b/packages/kilo-vscode/src/services/commit-message/index.ts index 4c57712edcc..bdd397327b4 100644 --- a/packages/kilo-vscode/src/services/commit-message/index.ts +++ b/packages/kilo-vscode/src/services/commit-message/index.ts @@ -1,5 +1,4 @@ import * as vscode from "vscode" -import type { KiloClient } from "@kilocode/sdk/v2/client" import type { KiloConnectionService } from "../cli-backend/connection-service" import { getErrorMessage } from "../../kilo-provider-utils" @@ -53,21 +52,17 @@ export function registerCommitMessageService( return } - let client: KiloClient | undefined - try { - client = connectionService.getClient() - } catch (err) { - console.error("[Kilo New] Failed to get client:", err) - vscode.window.showErrorMessage("Kilo backend is not connected. Please wait for the connection to establish.") - return - } - if (!client) { - vscode.window.showErrorMessage("Kilo backend is not connected. Please wait for the connection to establish.") - return - } - const path = repository.rootUri.fsPath + let client + try { + client = await connectionService.getClientAsync(path) + } catch (err) { + console.error("[Kilo New] Failed to connect to Kilo backend:", err) + vscode.window.showErrorMessage("Failed to connect to Kilo backend. Please try again.") + return + } + const previousMessage = lastWorkspacePath === path ? lastGeneratedMessage : undefined const controller = new AbortController()