From a09dafe61b757e2e737d3f935bd1fddc61835b55 Mon Sep 17 00:00:00 2001 From: marius-kilocode Date: Fri, 12 Jun 2026 08:54:35 +0200 Subject: [PATCH] fix(vscode): abort moved active sessions --- .changeset/abort-moved-sessions.md | 5 + packages/kilo-vscode/src/KiloProvider.ts | 77 ++++++++++-- .../kilo-vscode/src/kilo-provider/abort.ts | 35 +++++- packages/kilo-vscode/tests/unit/abort.test.ts | 59 ++++++++- .../unit/kilo-provider-load-messages.test.ts | 116 +++++++++++++++++- 5 files changed, 276 insertions(+), 16 deletions(-) create mode 100644 .changeset/abort-moved-sessions.md diff --git a/.changeset/abort-moved-sessions.md b/.changeset/abort-moved-sessions.md new file mode 100644 index 0000000000..287b8a8823 --- /dev/null +++ b/.changeset/abort-moved-sessions.md @@ -0,0 +1,5 @@ +--- +"kilo-code": patch +--- + +Stop active Agent Manager turns after moving or continuing their sessions in another worktree. diff --git a/packages/kilo-vscode/src/KiloProvider.ts b/packages/kilo-vscode/src/KiloProvider.ts index c393d27912..4a89ebc917 100644 --- a/packages/kilo-vscode/src/KiloProvider.ts +++ b/packages/kilo-vscode/src/KiloProvider.ts @@ -71,7 +71,12 @@ 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 { + abortSession, + resolveAbortDirectories, + updateActiveSessionDirectory, + type ActiveSessionDirectories, +} from "./kilo-provider/abort" import { buildAutocompleteSettingsMessage, validAutocompleteSetting, @@ -281,6 +286,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 activeSessionDirectories: ActiveSessionDirectories = new Map() private permissionDirectories = new Map() private projectID: string | undefined // Current workspace project ID used to filter sessions. private loadMessagesAbort: AbortController | null = null // Current load request cancellation. @@ -628,10 +634,28 @@ 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 { + const status = this.sessionStatusMap.get(sessionId) + if (status && status !== "idle" && !this.activeSessionDirectories.has(sessionId)) { + updateActiveSessionDirectory({ + active: this.activeSessionDirectories, + sessionID: sessionId, + status, + dir: this.getWorkspaceDirectory(sessionId), + }) + } this.sessionDirectories.set(sessionId, directory) } public clearSessionDirectory(sessionId: string): void { + const status = this.sessionStatusMap.get(sessionId) + if (status && status !== "idle" && !this.activeSessionDirectories.has(sessionId)) { + updateActiveSessionDirectory({ + active: this.activeSessionDirectories, + sessionID: sessionId, + status, + dir: this.getWorkspaceDirectory(sessionId), + }) + } this.sessionDirectories.delete(sessionId) } @@ -1708,6 +1732,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper this.visibleTaskStreams.delete(sessionID) this.syncedChildSessions.delete(sessionID) this.sessionDirectories.delete(sessionID) + this.activeSessionDirectories.delete(sessionID) this.lastReconciledAt.delete(sessionID) this.connectionService.pruneSession(sessionID) if (this.currentSession?.id === sessionID) { @@ -2705,23 +2730,33 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper } private async handleAbort(sessionID?: string): Promise { - if (!this.client) { - return - } + const client = this.client + if (!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 dirs = resolveAbortDirectories( + this.activeSessionDirectories, + targetSessionID, + this.getWorkspaceDirectory(targetSessionID), + ) + const results = await Promise.allSettled( + dirs.map((dir) => + abortSession({ + client, + sessionID: targetSessionID, + 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) } } @@ -3096,6 +3131,12 @@ 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) + updateActiveSessionDirectory({ + active: this.activeSessionDirectories, + sessionID: sid, + status: event.properties.status.type, + dir: directory, + }) const msg = mapSSEEventToWebviewMessage(event, sid) if (msg) { this.streams.flush(sid) @@ -3128,6 +3169,17 @@ 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 sessionID of [...this.activeSessionDirectories.keys()]) { + updateActiveSessionDirectory({ + active: this.activeSessionDirectories, + sessionID, + status: "idle", + dir, + }) + if (!this.activeSessionDirectories.has(sessionID)) this.sessionStatusMap.set(sessionID, "idle") + } + } if (dir && !sameDirectory(dir, this.getWorkspaceDirectory())) return void this.reloadAfterAuthChange() return @@ -3527,6 +3579,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper this.trackedSessionIds.clear() this.syncedChildSessions.clear() this.sessionDirectories.clear() + this.activeSessionDirectories.clear() this.permissionDirectories.clear() this.sessionStatusMap.clear() this.ignoreController?.dispose() diff --git a/packages/kilo-vscode/src/kilo-provider/abort.ts b/packages/kilo-vscode/src/kilo-provider/abort.ts index 8d3d00676e..fa0eeb3edf 100644 --- a/packages/kilo-vscode/src/kilo-provider/abort.ts +++ b/packages/kilo-vscode/src/kilo-provider/abort.ts @@ -1,4 +1,37 @@ -import type { KiloClient } from "@kilocode/sdk/v2/client" +import type { KiloClient, SessionStatus } from "@kilocode/sdk/v2/client" +import { sameDirectory } from "../kilo-provider-utils" + +export type ActiveSessionDirectories = Map> + +export function updateActiveSessionDirectory(input: { + active: ActiveSessionDirectories + sessionID: string + status: SessionStatus["type"] + dir?: string +}) { + const target = input.dir + if (!target) return + const dirs = input.active.get(input.sessionID) + if (input.status === "idle") { + if (!dirs) return + for (const dir of dirs) { + if (sameDirectory(dir, target)) dirs.delete(dir) + } + if (dirs.size === 0) input.active.delete(input.sessionID) + return + } + if (!dirs) { + input.active.set(input.sessionID, new Set([target])) + return + } + if (![...dirs].some((dir) => sameDirectory(dir, target))) dirs.add(target) +} + +export function resolveAbortDirectories(active: ActiveSessionDirectories, sessionID: string, fallback: string) { + const dirs = [...(active.get(sessionID) ?? [])] + if (!dirs.some((dir) => sameDirectory(dir, fallback))) dirs.push(fallback) + return dirs +} 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 126333e6a4..2c182cbfe8 100644 --- a/packages/kilo-vscode/tests/unit/abort.test.ts +++ b/packages/kilo-vscode/tests/unit/abort.test.ts @@ -1,6 +1,11 @@ import { describe, expect, it } from "bun:test" import type { KiloClient } from "@kilocode/sdk/v2/client" -import { abortSession } from "../../src/kilo-provider/abort" +import { + abortSession, + resolveAbortDirectories, + updateActiveSessionDirectory, + type ActiveSessionDirectories, +} from "../../src/kilo-provider/abort" function client(calls: unknown[], fail = false) { return { @@ -14,6 +19,58 @@ function client(calls: unknown[], fail = false) { } as unknown as KiloClient } +describe("active session directories", () => { + function update(active: ActiveSessionDirectories, status: "busy" | "retry" | "idle", dir?: string) { + updateActiveSessionDirectory({ active, sessionID: "session_1", status, dir }) + } + + it("includes the active owner and current session directory", () => { + const active: ActiveSessionDirectories = new Map() + update(active, "busy", "/repo/source") + + expect(resolveAbortDirectories(active, "session_1", "/repo/worktree")).toEqual(["/repo/source", "/repo/worktree"]) + }) + + it("falls back to the current session directory after the active turn becomes idle", () => { + const active: ActiveSessionDirectories = new Map() + update(active, "busy", "/repo/source") + update(active, "idle", "/repo/source") + + expect(resolveAbortDirectories(active, "session_1", "/repo/worktree")).toEqual(["/repo/worktree"]) + }) + + it("retains active directories when an unrelated instance reports idle", () => { + const active: ActiveSessionDirectories = new Map() + update(active, "retry", "/repo/source") + update(active, "idle", "/repo/worktree") + update(active, "idle") + + expect(resolveAbortDirectories(active, "session_1", "/repo/worktree")).toEqual(["/repo/source", "/repo/worktree"]) + }) + + it("tracks concurrent instances and ignores delayed idle events from the old one", () => { + const active: ActiveSessionDirectories = new Map() + update(active, "busy", "/repo/source") + update(active, "busy", "/repo/worktree") + + expect(resolveAbortDirectories(active, "session_1", "/repo/fallback")).toEqual([ + "/repo/source", + "/repo/worktree", + "/repo/fallback", + ]) + + update(active, "idle", "/repo/source") + expect(resolveAbortDirectories(active, "session_1", "/repo/fallback")).toEqual(["/repo/worktree", "/repo/fallback"]) + }) + + it("deduplicates the current directory when it is already active", () => { + const active: ActiveSessionDirectories = new Map() + update(active, "busy", "/repo/worktree") + + expect(resolveAbortDirectories(active, "session_1", "/repo/worktree/.")).toEqual(["/repo/worktree"]) + }) +}) + 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 a63808d556..f8f44116d8 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" // vscode mock is provided by the shared preload (tests/setup/vscode-mock.ts) const { KiloProvider } = await import("../../src/KiloProvider") @@ -43,12 +43,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 }) => { @@ -56,6 +59,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 @@ -114,9 +122,11 @@ type ProviderInternals = { currentSession: { id: string; directory?: string } | null contextSessionID: string | undefined sessionDirectories: Map + activeSessionDirectories: Map> trackedSessionIds: Set 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 } @@ -135,6 +145,108 @@ function makeProvider(client: ReturnType) { return { provider, internal, sent } } +describe("KiloProvider.handleAbort", () => { + it("keeps the active owner when a running session moves to a worktree", 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("preserves a directory-qualified owner instead of inferring the current directory", async () => { + const client = createClient() + const { provider, internal } = makeProvider(client) + internal.handleEvent( + { + type: "session.status", + properties: { sessionID: "s1", status: { type: "busy" } }, + }, + "/repo/source", + ) + provider.setSessionDirectory("s1", "/repo/worktree") + + await internal.handleAbort("s1") + + expect(client.aborted).toEqual([ + { sessionID: "s1", directory: "/repo/source" }, + { sessionID: "s1", directory: "/repo/worktree" }, + ]) + }) + + it("falls back to the current directory after the old owner becomes idle", 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") + internal.handleEvent( + { + type: "session.status", + properties: { sessionID: "s1", status: { type: "idle" } }, + }, + "/repo", + ) + + await internal.handleAbort("s1") + + expect(client.aborted).toEqual([{ sessionID: "s1", directory: "/repo/worktree" }]) + }) + + it("deduplicates the current directory when it is also an active owner", async () => { + const client = createClient() + const { provider, internal } = makeProvider(client) + internal.trackedSessionIds.add("s1") + provider.setSessionDirectory("s1", "/repo/worktree") + internal.handleEvent( + { + type: "session.status", + properties: { sessionID: "s1", status: { type: "busy" } }, + }, + "/repo/worktree", + ) + + await internal.handleAbort("s1") + + expect(client.aborted).toEqual([{ sessionID: "s1", directory: "/repo/worktree" }]) + }) + + it("attempts every owner when one abort request fails", async () => { + const error = spyOn(console, "error").mockImplementation(() => {}) + const client = createClient({ abortFailures: ["/repo"] }) + const { provider, internal } = makeProvider(client) + internal.trackedSessionIds.add("s1") + 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(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({