diff --git a/.changeset/fix-stale-agent-session-status.md b/.changeset/fix-stale-agent-session-status.md new file mode 100644 index 00000000000..d6e569b25ce --- /dev/null +++ b/.changeset/fix-stale-agent-session-status.md @@ -0,0 +1,5 @@ +--- +"kilo-code": patch +--- + +Prevent completed sessions from staying stuck in the working state. diff --git a/packages/kilo-vscode/src/KiloProvider.ts b/packages/kilo-vscode/src/KiloProvider.ts index 206972ce9e6..215ff9fd259 100644 --- a/packages/kilo-vscode/src/KiloProvider.ts +++ b/packages/kilo-vscode/src/KiloProvider.ts @@ -391,6 +391,9 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper private readonly refreshes = new Map() private readonly anacondaDesktop = new AnacondaDesktopBridge() private sessionStatusMap = new Map() // Latest status used for destructive config warnings. + private readonly epochs = new Map>() + private readonly requests = new Map() + private epoch = 0 private sessionDirectories = new Map() // Per-session directory overrides, such as Agent Manager worktrees. private sessionGitDirectories = new Map() // Stable Git root resolved for each session. private sessionGitRecoveries = new Set() // Sessions whose older history was scanned for a Git root. @@ -717,14 +720,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper if (this.cachedStats) this.postMessage(this.cachedStats) this.postMessage({ type: "gitStatus", repo: this.cachedGitRepo }) - // Seed session status map so the Settings panel knows about already-running sessions. - // Must run after webview is ready (postMessage is a no-op before that). - // Only reconcile (reset missing busy→idle) when the map is empty, i.e. - // on the very first seed before any real-time SSE events have arrived. - // On SSE reconnects or webview recreations the live SSE data is - // authoritative and reconciliation risks race-resetting busy sessions. - const reconcile = this.sessionStatusMap.size === 0 - void this.seedSessionStatusMap(reconcile) + void this.seedSessionStatusMap() this.sendRemoteStatus() } @@ -2006,21 +2002,30 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper }) .catch((e: unknown) => console.warn("[Kilo New] KiloProvider: getSession failed (non-critical):", e)) this.postMessage({ type: "workspaceDirectoryChanged", directory: this.getWorkspaceDirectory(sessionID) }) - this.client.session + this.sync(sessionID, dir, signal, refresh) + } + + private sync(sessionID: string, dir: string, signal?: AbortSignal, refresh?: number): void { + const client = this.client + if (!client) return + const epoch = this.epoch + const request = this.begin(dir) + void client.session .status({ directory: dir }) - .then((r) => { - if (!r.data || signal?.aborted) return - for (const [sid, info] of Object.entries(r.data) as [string, SessionStatus][]) { - if (!this.trackedSessionIds.has(sid)) continue - this.postMessage({ - type: "sessionStatus", - sessionID: sid, - status: info.type, - ...(info.type === "retry" ? { attempt: info.attempt, message: info.message, next: info.next } : {}), - }) + .then((result) => { + if (!result.data || signal?.aborted || !this.latest(dir, request)) return + if (refresh !== undefined && this.refreshes.get(sessionID) !== refresh) return + for (const [sid, status] of Object.entries(result.data) as [string, SessionStatus][]) { + if (!this.trackedSessionIds.has(sid) || !this.accept(sid, status, dir, epoch)) continue + this.publish(sid, status) + } + for (const [sid, current] of this.sessionStatusMap) { + if (result.data[sid] || current === "idle" || !this.trackedSessionIds.has(sid)) continue + const status = { type: "idle" as const } + if (this.accept(sid, status, dir, epoch)) this.publish(sid, status) } }) - .catch((e: unknown) => console.error("[Kilo New] KiloProvider: Failed to fetch session statuses:", e)) + .catch((error: unknown) => console.error("[Kilo New] KiloProvider: Failed to fetch session statuses:", error)) } private fetchAndSendSessionModelUsage(sessionID: string, requestID: string): Promise { @@ -2352,6 +2357,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper this.checkpoints.delete(sessionID) this.revisions.delete(sessionID) this.refreshes.delete(sessionID) + this.epochs.delete(sessionID) this.sessionStatusMap.delete(sessionID) this.costs.onSessionDeleted(sessionID) const deletedAlertLimit = this.activeAlerts.get(sessionID) @@ -2950,19 +2956,65 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper } } - /** - * Seed sessionStatusMap with current session statuses on connect. - * Without this, the Settings panel (which has no tracked sessions) would see - * busyCount() = 0 for sessions that were already running before it opened. - * - * @param reconcile When true, reset locally-busy sessions absent from the - * server response to idle (crash recovery). Set to false on SSE reconnects - * to avoid a race where a brief HTTP fetch gap causes the spinner to vanish. - */ + private begin(dir: string): number { + const key = [...this.requests.keys()].find((entry) => sameDirectory(entry, dir)) ?? dir + const request = (this.requests.get(key) ?? 0) + 1 + this.requests.set(key, request) + return request + } + + private latest(dir: string, request: number): boolean { + return [...this.requests].some(([entry, value]) => value === request && sameDirectory(entry, dir)) + } + + private mark(sessionID: string, dir?: string): void { + const entries = this.epochs.get(sessionID) ?? new Map() + const key = dir ? ([...entries.keys()].find((entry) => entry && sameDirectory(entry, dir)) ?? dir) : "" + entries.set(key, ++this.epoch) + this.epochs.set(sessionID, entries) + } + + private stale(sessionID: string, dir: string, epoch: number): boolean { + const entries = this.epochs.get(sessionID) + if (!entries) return false + return [...entries].some(([entry, value]) => value > epoch && (!entry || sameDirectory(entry, dir))) + } + + private accept(sessionID: string, status: SessionStatus, dir: string, epoch: number): boolean { + if (this.stale(sessionID, dir, epoch)) return false + if (status.type === "idle" && !sameDirectory(this.getWorkspaceDirectory(sessionID), dir)) return false + this.mark(sessionID, dir) + this.aborts.observe(sessionID, status.type, dir) + return true + } + + private publish(sessionID: string, status: SessionStatus): void { + const previous = this.sessionStatusMap.get(sessionID) + if ((previous === undefined || previous === "idle") && status.type !== "idle") this.costs.rearm(sessionID) + this.sessionStatusMap.set(sessionID, status.type) + this.streams.flush(sessionID) + this.postMessage({ + type: "sessionStatus", + sessionID, + status: status.type, + ...(status.type === "retry" ? { attempt: status.attempt, message: status.message, next: status.next } : {}), + ...(status.type === "offline" ? { message: status.message } : {}), + }) + } + private async seedSessionStatusMap(reconcile = true): Promise { if (!this.client || this.connectionState !== "connected") return const dir = this.getWorkspaceDirectory() - await seedSessionStatuses(this.client, dir, this.sessionStatusMap, (msg) => this.postMessage(msg), reconcile) + const epoch = this.epoch + const request = this.begin(dir) + await seedSessionStatuses( + this.client, + dir, + this.sessionStatusMap, + (message) => this.postMessage(message), + reconcile, + (sessionID, status) => this.latest(dir, request) && this.accept(sessionID, status, dir, epoch), + ) } /** @@ -4141,9 +4193,8 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper private async handleAbort(sessionID?: string): Promise { const sid = sessionID || this.currentSession?.id if (!sid || !(await this.stopSession(sid))) return - this.sessionStatusMap.set(sid, "idle") - this.streams.flush(sid) - this.postMessage({ type: "sessionStatus", sessionID: sid, status: "idle" }) + this.mark(sid) + this.publish(sid, { type: "idle" }) } private async handleRevertSession(sessionID: string, messageID: string, partID?: string): Promise { @@ -4677,17 +4728,10 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper // busy-session warning on Save. if (event.type === "session.status") { const sid = event.properties.sessionID - const prev = this.sessionStatusMap.get(sid) - if ((prev === undefined || prev === "idle") && event.properties.status.type !== "idle") { - this.costs.rearm(sid) - } - this.sessionStatusMap.set(sid, event.properties.status.type) - this.aborts.observe(sid, event.properties.status.type, directory) - const msg = mapSSEEventToWebviewMessage(event, sid) - if (msg) { - this.streams.flush(sid) - this.postMessage(msg) - } + const status = event.properties.status + this.mark(sid, directory) + this.aborts.observe(sid, status.type, directory) + this.publish(sid, status) return } @@ -4845,6 +4889,19 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper } this.streams.flush(sessionID) this.postMessage(next) + const sid = event.type === "session.turn.close" ? event.properties.sessionID : sessionID + if (!sid) return + const status = this.sessionStatusMap.get(sid) + if (!status || status === "idle") return + if ( + event.type === "session.turn.close" || + (event.type === "message.updated" && + event.properties.info.role === "assistant" && + event.properties.info.finish === "stop" && + event.properties.info.time.completed !== undefined) + ) { + this.sync(sid, directory ?? this.getWorkspaceDirectory(sid)) + } } /** Wait until the webview has sent "webviewReady". Resolves immediately when already ready. */ @@ -5387,6 +5444,8 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper this.sessionDirectories.clear() this.anacondaDesktop.dispose() this.aborts.clear() + this.requests.clear() + this.epochs.clear() this.sessionStatusMap.clear() this.ignoreController?.dispose() this.chatAutocomplete?.dispose() diff --git a/packages/kilo-vscode/src/session-status.ts b/packages/kilo-vscode/src/session-status.ts index 67856f81f54..0f83d93ad2d 100644 --- a/packages/kilo-vscode/src/session-status.ts +++ b/packages/kilo-vscode/src/session-status.ts @@ -4,11 +4,6 @@ import type { KiloClient, SessionStatus } from "@kilocode/sdk/v2/client" * Fetch all current session statuses and seed the provided map + webview. * Called on connect so the Settings panel knows about already-running sessions * without waiting for the next session.status SSE event. - * - * When `reconcile` is true (default: first seed), locally-busy sessions absent - * from the server response are reset to idle — covering server crash/restart. - * On SSE reconnects set `reconcile: false` to avoid a race where the HTTP - * fetch briefly returns stale data and the spinner disappears mid-stream. */ export async function seedSessionStatuses( client: KiloClient, @@ -16,6 +11,7 @@ export async function seedSessionStatuses( map: Map, post: (msg: unknown) => void, reconcile = true, + accept?: (sessionID: string, status: SessionStatus) => boolean, ): Promise { try { const result = await client.session.status({ directory: dir }) @@ -24,6 +20,7 @@ export async function seedSessionStatuses( // Seed/update entries the server knows about for (const [sid, info] of Object.entries(active) as [string, SessionStatus][]) { + if (accept && !accept(sid, info)) continue map.set(sid, info.type) post({ type: "sessionStatus", @@ -35,11 +32,10 @@ export async function seedSessionStatuses( // Reconcile: any locally non-idle session absent from the server response // means the server lost its in-memory state (crash/restart). Reset to idle. - // Skipped on SSE reconnects — the real-time SSE events are authoritative - // for status transitions and the brief HTTP fetch can race with them. if (reconcile) { for (const [sid, status] of map) { if (status !== "idle" && !active[sid]) { + if (accept && !accept(sid, { type: "idle" })) continue map.set(sid, "idle") post({ type: "sessionStatus", sessionID: sid, status: "idle" }) } diff --git a/packages/kilo-vscode/tests/unit/kilo-provider-load-messages.test.ts b/packages/kilo-vscode/tests/unit/kilo-provider-load-messages.test.ts index 24be2aaf257..36d3d4390c5 100644 --- a/packages/kilo-vscode/tests/unit/kilo-provider-load-messages.test.ts +++ b/packages/kilo-vscode/tests/unit/kilo-provider-load-messages.test.ts @@ -1,4 +1,5 @@ import { describe, it, expect, spyOn } from "bun:test" +import type { SessionStatus } from "@kilocode/sdk/v2/client" import * as vscode from "vscode" import type { PartUpdate } from "../../src/shared/stream-messages" @@ -65,6 +66,7 @@ function createClient(options?: { revertDeferred?: Deferred<{ data?: unknown; error?: unknown }> sessionData?: unknown sessionGet?: (params: { sessionID: string; directory?: string }) => Promise<{ data: unknown }> + status?: (params: { directory?: string }) => Promise<{ data: Record | null }> createDeferred?: Deferred<{ data: ReturnType }> abortFailures?: string[] abortDeferred?: Deferred @@ -105,7 +107,7 @@ function createClient(options?: { if (options?.sessionGet) return options.sessionGet(params) return { data: options?.sessionData ?? null } }, - status: async () => ({ data: {} }), + status: async (params: { directory?: string }) => options?.status?.(params) ?? { data: {} }, revert: async (params: Record) => { reverted.push(params) if (options?.revertDeferred) return options.revertDeferred.promise @@ -223,6 +225,7 @@ type ProviderInternals = { currentSession: { id: string; directory?: string; cost?: number; revert?: { messageID: string } } | null contextSessionID: string | undefined sessionDirectories: Map + sessionStatusMap: Map trackedSessionIds: Set openSessionIds: Set draftSessions: Map @@ -232,6 +235,7 @@ type ProviderInternals = { checkpoint: (sid: string, run: () => Promise) => void gatherEditorContext: () => Promise> refreshSessionDetails: (sid: string, dir: string) => void + seedSessionStatusMap: (reconcile?: boolean) => Promise stopCurrentSessionProcesses: (next?: string) => void handleEvent: (event: unknown, directory?: string) => void handleAbort: (sid?: string) => Promise @@ -263,6 +267,10 @@ function makeProvider(client: ReturnType) { return { provider, internal, sent } } +function status(internal: ProviderInternals, type: "busy" | "idle", directory = "/repo", sessionID = "s1") { + internal.handleEvent({ type: "session.status", properties: { sessionID, status: { type } } }, directory) +} + function mockMaxCost(internal: ProviderInternals, value: number) { internal.setMaxCost(value) } @@ -287,6 +295,7 @@ describe("KiloProvider.handleAbort", () => { { sessionID: "s1", directory: "/repo/worktree" }, ]) expect(sent.at(-1)).toMatchObject({ type: "sessionStatus", sessionID: "s1", status: "idle" }) + expect(sent).not.toContainEqual({ type: "sessionTurnClosed", sessionID: "s1", reason: "interrupted" }) }) it("preserves the original owner when the status event lacks a directory", async () => { @@ -412,6 +421,110 @@ describe("KiloProvider.handleAbort", () => { }) }) +describe("KiloProvider session status reconciliation", () => { + it("rejects a stale busy snapshot after a newer idle event", async () => { + const pending = defer<{ data: Record }>() + const client = createClient({ status: async () => pending.promise }) + const { internal } = makeProvider(client) + internal.trackedSessionIds.add("s1") + status(internal, "busy") + + internal.refreshSessionDetails("s1", "/repo") + status(internal, "idle") + pending.resolve({ data: { s1: { type: "busy" } } }) + await Bun.sleep(0) + + expect(internal.sessionStatusMap.get("s1")).toBe("idle") + }) + + it("preserves newer busy status when a stale snapshot omits the session", async () => { + const pending = defer<{ data: Record }>() + const client = createClient({ status: async () => pending.promise }) + const { internal } = makeProvider(client) + internal.trackedSessionIds.add("s1") + internal.refreshSessionDetails("s1", "/repo") + status(internal, "busy") + pending.resolve({ data: {} }) + await Bun.sleep(0) + + expect(internal.sessionStatusMap.get("s1")).toBe("busy") + }) + + it("recovers missing idle after a completed assistant response", async () => { + const client = createClient() + const { internal, sent } = makeProvider(client) + internal.trackedSessionIds.add("s1") + status(internal, "busy") + + internal.handleEvent( + { + type: "message.updated", + properties: { + sessionID: "s1", + info: { id: "m1", sessionID: "s1", role: "assistant", finish: "stop", time: { created: 1, completed: 2 } }, + }, + }, + "/repo", + ) + await Bun.sleep(0) + + expect(internal.sessionStatusMap.get("s1")).toBe("idle") + expect(sent).toContainEqual({ type: "sessionStatus", sessionID: "s1", status: "idle" }) + }) + + it("accepts the latest overlapping directory snapshot", async () => { + const first = defer<{ data: Record }>() + const second = defer<{ data: Record }>() + const pending = [first, second] + const client = createClient({ status: async () => pending.shift()!.promise }) + const { internal } = makeProvider(client) + internal.trackedSessionIds.add("s1") + internal.trackedSessionIds.add("s2") + status(internal, "busy", "/repo", "s1") + status(internal, "busy", "/repo", "s2") + + internal.refreshSessionDetails("s1", "/repo") + internal.refreshSessionDetails("s2", "/repo") + second.resolve({ data: { s2: { type: "busy" } } }) + await Bun.sleep(0) + first.resolve({ data: { s1: { type: "busy" }, s2: { type: "busy" } } }) + await Bun.sleep(0) + + expect(internal.sessionStatusMap.get("s1")).toBe("idle") + expect(internal.sessionStatusMap.get("s2")).toBe("busy") + }) + + it("accepts the latest overlapping seeded snapshot", async () => { + const first = defer<{ data: Record }>() + const second = defer<{ data: Record }>() + const pending = [first, second] + const client = createClient({ status: async () => pending.shift()!.promise }) + const { internal } = makeProvider(client) + status(internal, "busy") + + const older = internal.seedSessionStatusMap() + const newer = internal.seedSessionStatusMap() + first.resolve({ data: { s1: { type: "busy" } } }) + await older + second.resolve({ data: {} }) + await newer + + expect(internal.sessionStatusMap.get("s1")).toBe("idle") + }) + + it("does not idle a worktree session from a root snapshot", async () => { + const client = createClient() + const { provider, internal } = makeProvider(client) + provider.setSessionDirectory("s1", "/repo/worktree") + internal.trackedSessionIds.add("s1") + status(internal, "busy", "/repo/worktree") + + await internal.seedSessionStatusMap() + + expect(internal.sessionStatusMap.get("s1")).toBe("busy") + }) +}) + describe("KiloProvider sandbox status", () => { it("ignores events from another directory for the same session", () => { const client = createClient()