Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .changeset/calm-kilos-reconnect.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"kilo-code": patch
---

Show a retryable connection error and preserve unsent prompts when the VS Code background CLI process exits.
18 changes: 14 additions & 4 deletions packages/kilo-vscode/src/KiloProvider.ts
Original file line number Diff line number Diff line change
Expand Up @@ -347,6 +347,16 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
}
}

private postConnectionState(error = this.connectionService.getConnectionError()): void {
this.postMessage({
type: "connectionState",
state: this.connectionState,
...(this.connectionState === "error" && {
error: getErrorMessage(error) || "Connection to CLI backend lost. Retry to reconnect.",
}),
})
}

// 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<T>(part: T): T {
Expand Down Expand Up @@ -386,7 +396,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
}

// Always push connection state first so the UI can render appropriately.
this.postMessage({ type: "connectionState", state: this.connectionState })
this.postConnectionState()
pushTelemetryState((m) => this.postMessage(m))

// Re-send ready so the webview can recover after refresh.
Expand Down Expand Up @@ -1228,9 +1238,9 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
)

// Subscribe to connection state changes
this.unsubscribeState = this.connectionService.onStateChange(async (state) => {
this.unsubscribeState = this.connectionService.onStateChange(async (state, error) => {
this.connectionState = state
this.postMessage({ type: "connectionState", state })
this.postConnectionState(error)

if (state === "connected") {
// Fire config warnings independently so a failure in the
Expand Down Expand Up @@ -1309,7 +1319,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
workspaceDirectory: this.getProjectDirectory(this.currentSession?.id),
})
}
this.postMessage({ type: "connectionState", state: this.connectionState })
this.postConnectionState()

// connect() can resolve after SSE reaches "connected" but before this
// provider subscribes to onStateChange(). In that case the initial
Expand Down
78 changes: 59 additions & 19 deletions packages/kilo-vscode/src/services/cli-backend/connection-service.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ import { resolveEventSessionId as resolveEventSessionIdPure } from "./connection

export type ConnectionState = "connecting" | "connected" | "disconnected" | "error"
type SSEEventListener = (event: Event, directory?: string) => void
type StateListener = (state: ConnectionState) => void
type StateListener = (state: ConnectionState, error?: Error) => void
type SSEEventFilter = (event: Event, directory?: string) => boolean
type NotificationDismissListener = (notificationId: string) => void
type LanguageChangeListener = (locale: string) => void
Expand Down Expand Up @@ -65,6 +65,7 @@ export class KiloConnectionService {
private info: { port: number } | null = null
private config: ServerConfig | null = null
private state: ConnectionState = "disconnected"
private error: Error | null = null
private connectPromise: Promise<void> | null = null
private healthPollTimer: ReturnType<typeof setInterval> | null = null
private remoteService: import("../RemoteStatusService").RemoteStatusService | null = null
Expand Down Expand Up @@ -93,7 +94,7 @@ export class KiloConnectionService {
private unsubRemote: (() => void) | null = null

constructor(context: vscode.ExtensionContext) {
this.serverManager = new ServerManager(context)
this.serverManager = new ServerManager(context, (code) => this.handleServerExit(code))
}

/**
Expand All @@ -115,7 +116,7 @@ export class KiloConnectionService {
await this.connectPromise
} catch (error) {
// If doConnect() fails before SSE can emit a state transition, avoid leaving consumers stuck in "connecting".
this.setState("error")
this.setState("error", this.error ?? (error instanceof Error ? error : new Error(String(error))))
throw error
} finally {
this.connectPromise = null
Expand All @@ -126,7 +127,7 @@ export class KiloConnectionService {
* Get the shared SDK client. Throws if not connected.
*/
getClient(): KiloClient {
if (!this.client) {
if (!this.client || this.state !== "connected") {
throw new Error("Not connected — call connect() first")
}
return this.client
Expand All @@ -139,11 +140,11 @@ export class KiloConnectionService {
* or if the connection fails.
*/
async getClientAsync(dir?: string): Promise<KiloClient> {
if (this.client) return this.client
if (this.client && this.state === "connected") 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!
return this.getClient()
}

/**
Expand Down Expand Up @@ -189,6 +190,13 @@ export class KiloConnectionService {
return this.state
}

/**
* Last connection error. Cleared when a new connection attempt begins.
*/
getConnectionError(): Error | null {
return this.error
}

/**
* Subscribe to SSE events. Returns unsubscribe function.
*/
Expand Down Expand Up @@ -503,12 +511,14 @@ export class KiloConnectionService {
this.config = null
this.info = null
this.state = "disconnected"
this.error = null
}

private setState(state: ConnectionState): void {
private setState(state: ConnectionState, error?: Error): void {
this.state = state
this.error = state === "error" ? (error ?? this.error) : null
for (const listener of this.stateListeners) {
listener(state)
listener(state, this.error ?? undefined)
}
}

Expand Down Expand Up @@ -558,10 +568,28 @@ export class KiloConnectionService {
}
}

private async doConnect(workspaceDir: string): Promise<void> {
// If we reconnect, ensure the previous SSE connection is cleaned up first.
private resetConnection(): void {
this.stopHealthPoll()
this.sseClient?.dispose()
const sse = this.sseClient
this.sseClient = null
sse?.disconnect()
this.client = null
this.config = null
this.info = null
}

private handleServerExit(code: number | null): void {
console.warn("[Kilo New] ConnectionService: CLI background process exited:", code)
this.resetConnection()
this.setState(
"error",
new Error(`CLI background process exited with code ${code ?? "unknown"}. Retry to reconnect.`),
)
}

private async doConnect(workspaceDir: string): Promise<void> {
// Never expose a stale SDK client while its replacement server is starting.
this.resetConnection()

const server = await this.serverManager.getServer()
this.info = { port: server.port }
Expand All @@ -575,14 +603,15 @@ export class KiloConnectionService {

// Create SDK client with Basic Auth header
const authHeader = `Basic ${Buffer.from(`kilo:${server.password}`).toString("base64")}`
this.client = createKiloClient({
const client = createKiloClient({
baseUrl: config.baseUrl,
headers: {
Authorization: authHeader,
},
})

this.sseClient = new SdkSSEAdapter(this.client)
const sse = new SdkSSEAdapter(client)
this.client = client
this.sseClient = sse

// Wait until SSE yields its first server event before resolving connect().
// Initial stream failures are handled by the adapter reconnect loop.
Expand All @@ -596,18 +625,29 @@ export class KiloConnectionService {
let didConnect = false

// Wire SSE events → broadcast to all registered listeners
this.sseClient.onEvent((event, directory) => {
sse.onEvent((event, directory) => {
if (this.sseClient !== sse) return
for (const listener of this.eventListeners) {
listener(event, directory)
}
})

this.sseClient.onError(() => {
this.setState("error")
sse.onError((error) => {
if (this.sseClient !== sse) return
this.setState("error", error)
})

// Wire SSE state → broadcast to all registered state listeners
this.sseClient.onStateChange((sseState) => {
sse.onStateChange((sseState) => {
if (this.sseClient !== sse) {
if (!didConnect && sseState === "disconnected") {
rejectConnected?.(new Error(`SSE connection ended in state: ${sseState}`))
resolveConnected = null
rejectConnected = null
}
return
}

this.setState(sseState)

if (sseState === "connected") {
Expand All @@ -625,7 +665,7 @@ export class KiloConnectionService {
}
})

this.sseClient.connect()
sse.connect()

await connectedPromise

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ export interface ServerInstance {
const STARTUP_TIMEOUT_SECONDS = 30

type WorkspaceFolderLike = { uri: { fsPath: string } }
type ServerExitListener = (code: number | null) => void

export function resolveServerCwd(folders: readonly WorkspaceFolderLike[] | undefined, storage: string): string {
return folders?.[0]?.uri.fsPath ?? storage
Expand All @@ -31,7 +32,10 @@ export class ServerManager {
private instance: ServerInstance | null = null
private startupPromise: Promise<ServerInstance> | null = null

constructor(private readonly context: vscode.ExtensionContext) {}
constructor(
private readonly context: vscode.ExtensionContext,
private readonly onExit?: ServerExitListener,
) {}

/**
* Get or start the server instance
Expand Down Expand Up @@ -171,6 +175,7 @@ export class ServerManager {
console.log("[Kilo New] ServerManager: 🛑 Process exited with code:", code)
if (this.instance?.process === serverProcess) {
this.instance = null
this.onExit?.(code)
}
if (!resolved) {
const { userMessage, userDetails } = toErrorMessage(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@ function connection() {
getServerInfo: () => ({ port: 12345 }),
getServerConfig: () => ({ baseUrl: "http://127.0.0.1:12345", password: "test" }),
getConnectionState: () => "connected" as const,
getConnectionError: () => null,
resolveEventSessionId: (event: Event) => (event.type === "session.created" ? event.properties.info.id : undefined),
recordMessageSessionId: () => undefined,
notifyNotificationDismissed: () => undefined,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,7 @@ function createConnection(client: ReturnType<typeof createClient>) {
registerDirectoryProvider: () => () => undefined,
getServerInfo: () => ({ port: 12345 }),
getConnectionState: () => "connected" as const,
getConnectionError: () => null,
resolveEventSessionId: () => undefined,
recordMessageSessionId: () => undefined,
notifyNotificationDismissed: () => undefined,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@ function createConnection(client: ReturnType<typeof createClient>) {
getServerInfo: () => ({ port: 12345 }),
getServerConfig: () => ({ baseUrl: "http://127.0.0.1:12345", password: "test" }),
getConnectionState: () => "connected" as const,
getConnectionError: () => null,
resolveEventSessionId: () => undefined,
recordMessageSessionId: () => undefined,
notifyNotificationDismissed: () => undefined,
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
import { describe, expect, it } from "bun:test"
import { readFileSync } from "node:fs"
import { join } from "node:path"

describe("PromptInput connection guard", () => {
const path = join(__dirname, "..", "..", "webview-ui", "src", "components", "chat", "PromptInput.tsx")
const src = readFileSync(path, "utf8")

it("rechecks the connection after resolving async attachments and before clearing the draft", () => {
const attachments = src.indexOf("const gitFile = await git.resolveAttachment")
const guard = src.indexOf("if (isDisabled()) return", attachments)
const send = src.indexOf("session.sendMessage(message", guard)
const clear = src.indexOf("drafts.delete(key)", send)

expect(attachments).toBeGreaterThan(-1)
expect(guard).toBeGreaterThan(attachments)
expect(send).toBeGreaterThan(guard)
expect(clear).toBeGreaterThan(send)
})
})
32 changes: 32 additions & 0 deletions packages/kilo-vscode/tests/unit/sdk-sse-adapter.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,38 @@ describe("SdkSSEAdapter", () => {
})
})

describe("KiloConnectionService backend crash", () => {
it("invalidates the stale SDK client and reports a retryable error", () => {
const service = new KiloConnectionService({} as any)
const states: Array<{ state: string; error?: string }> = []
;(service as any).client = {}
;(service as any).config = { baseUrl: "http://127.0.0.1:52512", password: "secret" }
;(service as any).info = { port: 52512 }
;(service as any).state = "connected"
service.onStateChange((state, error) => states.push({ state, error: error?.message }))
;(service as any).handleServerExit(9)

expect(service.getConnectionState()).toBe("error")
expect(service.getConnectionError()?.message).toContain("CLI background process exited with code 9")
expect(service.getServerConfig()).toBeNull()
expect(service.getServerInfo()).toBeNull()
expect(() => service.getClient()).toThrow("Not connected")
expect(states).toEqual([
{ state: "error", error: "CLI background process exited with code 9. Retry to reconnect." },
])
service.dispose()
})

it("does not expose an SDK client while a replacement server is connecting", () => {
const service = new KiloConnectionService({} as any)
;(service as any).client = {}
;(service as any).state = "connecting"

expect(() => service.getClient()).toThrow("Not connected")
service.dispose()
})
})

describe("KiloConnectionService SSE startup", () => {
it("waits through an initial SSE fetch failure until the stream opens", async () => {
const original = globalThis.fetch
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -801,6 +801,7 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
return undefined
})
if (hasGit() && hasGitChangesMention(message) && !gitFile) return
if (isDisabled()) return

const allFiles = [
...mentionFiles,
Expand Down
Loading