diff --git a/.changeset/stop-moved-agent-sessions.md b/.changeset/stop-moved-agent-sessions.md new file mode 100644 index 00000000000..ee125bf2d8f --- /dev/null +++ b/.changeset/stop-moved-agent-sessions.md @@ -0,0 +1,5 @@ +--- +"kilo-code": patch +--- + +Keep the Stop button working when an active session moves between the workspace and an Agent Manager worktree. diff --git a/packages/kilo-vscode/src/KiloProvider.ts b/packages/kilo-vscode/src/KiloProvider.ts index 3ef1f944182..da0ea7d794b 100644 --- a/packages/kilo-vscode/src/KiloProvider.ts +++ b/packages/kilo-vscode/src/KiloProvider.ts @@ -72,7 +72,7 @@ import { fetchMessagePage, MESSAGE_PAGE_LIMIT } from "./kilo-provider/message-pa import { childID } from "./kilo-provider/task-session" import { VisibleTaskStreams } from "./kilo-provider/visible-task-streams" import { handleNetworkEvent, clearNetworkWaits } from "./kilo-provider/network" -import { abortSession } from "./kilo-provider/abort" +import { SessionAbort } from "./kilo-provider/abort" import { buildAutocompleteSettingsMessage, validAutocompleteSetting, @@ -310,6 +310,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper private syncedChildSessions: Set = new Set() private sessionStatusMap = new Map() // Latest status used for destructive config warnings. private sessionDirectories = new Map() // Per-session directory overrides, such as Agent Manager worktrees. + private readonly aborts = new SessionAbort() private projectID: string | undefined // Current workspace project ID used to filter sessions. private loadMessagesAbort: AbortController | null = null // Current load request cancellation. private lastReconciledAt = new Map() // Per-session focus-mode reconcile timestamp. @@ -657,10 +658,12 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper * When set, all operations for this session use this directory instead of the workspace root. */ public setSessionDirectory(sessionId: string, directory: string): void { + this.aborts.preserve(sessionId, this.sessionStatusMap.get(sessionId), this.getWorkspaceDirectory(sessionId)) this.sessionDirectories.set(sessionId, directory) } public clearSessionDirectory(sessionId: string): void { + this.aborts.preserve(sessionId, this.sessionStatusMap.get(sessionId), this.getWorkspaceDirectory(sessionId)) this.sessionDirectories.delete(sessionId) } @@ -1755,6 +1758,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper this.visibleTaskStreams.delete(sessionID) this.syncedChildSessions.delete(sessionID) this.sessionDirectories.delete(sessionID) + this.aborts.delete(sessionID) this.lastReconciledAt.delete(sessionID) this.connectionService.pruneSession(sessionID) if (this.currentSession?.id === sessionID) { @@ -2754,24 +2758,11 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper } private async handleAbort(sessionID?: string): Promise { - if (!this.client) { - return - } - - const targetSessionID = sessionID || this.currentSession?.id - if (!targetSessionID) { - return - } - - try { - await abortSession({ - client: this.client, - sessionID: targetSessionID, - dir: this.getWorkspaceDirectory(targetSessionID), - }) - } catch (error) { - console.error("[Kilo New] KiloProvider: Failed to abort session:", error) - } + const sid = sessionID || this.currentSession?.id + if (!this.client || !sid || !(await this.aborts.stop(this.client, sid, this.getWorkspaceDirectory(sid)))) return + this.sessionStatusMap.set(sid, "idle") + this.streams.flush(sid) + this.postMessage({ type: "sessionStatus", sessionID: sid, status: "idle" }) } private async handleRevertSession(sessionID: string, messageID: string, partID?: string): Promise { @@ -3156,6 +3147,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper if (event.type === "session.status") { const sid = event.properties.sessionID 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) @@ -3188,6 +3180,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper if (event.type === "server.instance.disposed") { const props = event.properties as Record | null const dir = typeof props?.directory === "string" ? props.directory : undefined + if (dir) for (const sid of this.aborts.dispose(dir)) this.sessionStatusMap.set(sid, "idle") if (dir && !sameDirectory(dir, this.getWorkspaceDirectory())) return void this.reloadAfterAuthChange() return @@ -3600,6 +3593,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper this.trackedSessionIds.clear() this.syncedChildSessions.clear() this.sessionDirectories.clear() + this.aborts.clear() this.sessionStatusMap.clear() this.ignoreController?.dispose() this.chatAutocomplete?.dispose() diff --git a/packages/kilo-vscode/src/kilo-provider/abort.ts b/packages/kilo-vscode/src/kilo-provider/abort.ts index 8d3d00676e9..7bb355b5b26 100644 --- a/packages/kilo-vscode/src/kilo-provider/abort.ts +++ b/packages/kilo-vscode/src/kilo-provider/abort.ts @@ -1,4 +1,69 @@ -import type { KiloClient } from "@kilocode/sdk/v2/client" +import type { KiloClient, SessionStatus } from "@kilocode/sdk/v2/client" +import { sameDirectory } from "../kilo-provider-utils" + +export class SessionAbort { + private active = new Map>() + + observe(sessionID: string, status: SessionStatus["type"], dir?: string) { + if (!dir) return + const dirs = this.active.get(sessionID) + if (status === "idle") { + if (!dirs) return + for (const entry of dirs) { + if (sameDirectory(entry, dir)) dirs.delete(entry) + } + if (dirs.size === 0) this.active.delete(sessionID) + return + } + if (!dirs) { + this.active.set(sessionID, new Set([dir])) + return + } + if (![...dirs].some((entry) => sameDirectory(entry, dir))) dirs.add(dir) + } + + preserve(sessionID: string, status: SessionStatus["type"] | undefined, dir: string) { + if (!status || status === "idle" || this.active.has(sessionID)) return + this.observe(sessionID, status, dir) + } + + async stop(client: KiloClient, sessionID: string, fallback: string) { + const known = this.active.has(sessionID) + const dirs = [...(this.active.get(sessionID) ?? [])] + if (!dirs.some((dir) => sameDirectory(dir, fallback))) dirs.push(fallback) + const results = await Promise.allSettled(dirs.map((dir) => abortSession({ client, sessionID, dir }))) + const failures = results.flatMap((result, index) => + result.status === "rejected" ? [{ dir: dirs[index], error: result.reason }] : [], + ) + if (failures.length > 0) { + console.error("[Kilo New] KiloProvider: Failed to abort session in one or more directories:", failures) + return false + } + if (known) this.active.delete(sessionID) + return known + } + + dispose(dir: string) { + const idle: string[] = [] + for (const [sessionID, dirs] of this.active) { + for (const entry of dirs) { + if (sameDirectory(entry, dir)) dirs.delete(entry) + } + if (dirs.size > 0) continue + this.active.delete(sessionID) + idle.push(sessionID) + } + return idle + } + + delete(sessionID: string) { + this.active.delete(sessionID) + } + + clear() { + this.active.clear() + } +} export async function abortSession(input: { client: KiloClient; sessionID: string; dir: string }) { await input.client.session.abort({ sessionID: input.sessionID, directory: input.dir }, { throwOnError: true }) diff --git a/packages/kilo-vscode/tests/unit/abort.test.ts b/packages/kilo-vscode/tests/unit/abort.test.ts index 126333e6a4e..e971a1211b8 100644 --- a/packages/kilo-vscode/tests/unit/abort.test.ts +++ b/packages/kilo-vscode/tests/unit/abort.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it } from "bun:test" import type { KiloClient } from "@kilocode/sdk/v2/client" -import { abortSession } from "../../src/kilo-provider/abort" +import { abortSession, SessionAbort } from "../../src/kilo-provider/abort" function client(calls: unknown[], fail = false) { return { @@ -14,6 +14,53 @@ function client(calls: unknown[], fail = false) { } as unknown as KiloClient } +describe("SessionAbort", () => { + it("stops the active owner and current mapped directory", async () => { + const calls: unknown[] = [] + const aborts = new SessionAbort() + aborts.observe("session_1", "busy", "/repo") + + expect(await aborts.stop(client(calls), "session_1", "/repo/worktree")).toBe(true) + expect(calls).toEqual([ + { + type: "abort", + params: { sessionID: "session_1", directory: "/repo" }, + opts: { throwOnError: true }, + }, + { + type: "abort", + params: { sessionID: "session_1", directory: "/repo/worktree" }, + opts: { throwOnError: true }, + }, + ]) + }) + + it("forgets an owner when its instance becomes idle", async () => { + const calls: unknown[] = [] + const aborts = new SessionAbort() + aborts.observe("session_1", "busy", "/repo") + aborts.observe("session_1", "idle", "/repo") + + expect(await aborts.stop(client(calls), "session_1", "/repo/worktree")).toBe(false) + expect(calls).toEqual([ + { + type: "abort", + params: { sessionID: "session_1", directory: "/repo/worktree" }, + opts: { throwOnError: true }, + }, + ]) + }) + + it("deduplicates equivalent directory paths", async () => { + const calls: unknown[] = [] + const aborts = new SessionAbort() + aborts.observe("session_1", "busy", "/repo/worktree") + + expect(await aborts.stop(client(calls), "session_1", "/repo/worktree/.")).toBe(true) + expect(calls).toHaveLength(1) + }) +}) + describe("abortSession", () => { it("calls session.abort with the session id and directory", async () => { const calls: unknown[] = [] 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 3145932a2ca..a909cbe407c 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,4 @@ -import { describe, it, expect } from "bun:test" +import { describe, it, expect, spyOn } from "bun:test" import type { PartUpdate } from "../../src/shared/stream-messages" // vscode mock is provided by the shared preload (tests/setup/vscode-mock.ts) @@ -44,12 +44,15 @@ function createClient(options?: { deleteDeferred?: Deferred sessionData?: unknown sessionGet?: (params: { sessionID: string; directory?: string }) => Promise<{ data: unknown }> + abortFailures?: string[] }) { const calls: { before?: string; limit?: number }[] = [] const stopped: { sessionID: string; directory?: string }[] = [] + const aborted: { sessionID: string; directory?: string }[] = [] return { calls, stopped, + aborted, session: { list: async () => ({ data: [] }), get: async (params: { sessionID: string; directory?: string }) => { @@ -57,6 +60,11 @@ function createClient(options?: { return { data: options?.sessionData ?? null } }, status: async () => ({ data: {} }), + abort: async (params: { sessionID: string; directory?: string }) => { + aborted.push(params) + if (params.directory && options?.abortFailures?.includes(params.directory)) throw new Error("abort failed") + return { data: true } + }, messages: async (params: { before?: string; limit?: number }) => { calls.push({ before: params.before, limit: params.limit }) if (options?.messagesDeferred) return options.messagesDeferred.promise @@ -118,7 +126,8 @@ type ProviderInternals = { trackedSessionIds: Set streams: { push: (msg: PartUpdate) => void } stopCurrentSessionProcesses: (next?: string) => void - handleEvent: (event: unknown) => void + handleEvent: (event: unknown, directory?: string) => void + handleAbort: (sid?: string) => Promise handleLoadMessages: (sid: string, opts?: { mode?: string; before?: string; limit?: number }) => Promise handleDeleteSession: (sid: string) => Promise } @@ -137,6 +146,70 @@ function makeProvider(client: ReturnType) { return { provider, internal, sent } } +describe("KiloProvider.handleAbort", () => { + it("aborts the original owner after a running session moves to a worktree", async () => { + const client = createClient() + const { provider, internal, sent } = makeProvider(client) + internal.handleEvent( + { + type: "session.status", + properties: { sessionID: "s1", status: { type: "busy" } }, + }, + "/repo", + ) + provider.setSessionDirectory("s1", "/repo/worktree") + + await internal.handleAbort("s1") + + expect(client.aborted).toEqual([ + { sessionID: "s1", directory: "/repo" }, + { sessionID: "s1", directory: "/repo/worktree" }, + ]) + expect(sent.at(-1)).toMatchObject({ type: "sessionStatus", sessionID: "s1", status: "idle" }) + }) + + it("preserves the original owner when the status event lacks a directory", async () => { + const client = createClient() + const { provider, internal } = makeProvider(client) + internal.handleEvent({ + type: "session.status", + properties: { sessionID: "s1", status: { type: "busy" } }, + }) + provider.setSessionDirectory("s1", "/repo/worktree") + + await internal.handleAbort("s1") + + expect(client.aborted).toEqual([ + { sessionID: "s1", directory: "/repo" }, + { sessionID: "s1", directory: "/repo/worktree" }, + ]) + }) + + it("attempts every owner and stays busy when one abort fails", async () => { + const error = spyOn(console, "error").mockImplementation(() => {}) + const client = createClient({ abortFailures: ["/repo"] }) + const { provider, internal, sent } = makeProvider(client) + internal.handleEvent( + { + type: "session.status", + properties: { sessionID: "s1", status: { type: "busy" } }, + }, + "/repo", + ) + provider.setSessionDirectory("s1", "/repo/worktree") + + await internal.handleAbort("s1") + + expect(client.aborted).toEqual([ + { sessionID: "s1", directory: "/repo" }, + { sessionID: "s1", directory: "/repo/worktree" }, + ]) + expect(sent.at(-1)).toMatchObject({ type: "sessionStatus", sessionID: "s1", status: "busy" }) + expect(error).toHaveBeenCalledTimes(1) + error.mockRestore() + }) +}) + describe("KiloProvider.handleLoadMessages / focus mode freshness", () => { it("stops background processes for the previous session when switching sessions", async () => { const client = createClient({