Merge pull request #8593 from Kilo-Org/fix/commit-message-auto-connect

fix(vscode): auto-connect backend when generating commit message without sidebar open

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<boolean> {
-    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
   }
 }

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 }),
       )
     })
   })

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<KiloClient> {
+    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.
    */

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

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()