From 67064b5dfcc461ab8eec943828729c016d490594 Mon Sep 17 00:00:00 2001 From: marius-kilocode Date: Thu, 27 Aug 2026 11:30:14 +0200 Subject: [PATCH] fix(agent-manager): persist retained session relocation after worktree deletion --- .../src/agent-manager/AgentManagerProvider.ts | 2 +- .../src/agent-manager/provider-lifecycle.ts | 37 +++++++--- .../tests/unit/agent-manager-arch.test.ts | 3 +- .../agent-manager-provider-lifecycle.test.ts | 73 ++++++++++++++++--- 4 files changed, 93 insertions(+), 22 deletions(-) diff --git a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts b/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts index abfc4467fc..05f887fc62 100644 --- a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts +++ b/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts @@ -1454,7 +1454,6 @@ export class AgentManagerProvider implements Disposable { runScriptConfigured: false, }) } - // Manager accessors — repository-bound services are owned by the active ProjectContext (immutable per root). /** Provider capabilities for the worktree lifecycle handlers (state stays in ProjectContext). */ private get lifecycleHost(): LifecycleHost { @@ -1489,6 +1488,7 @@ export class AgentManagerProvider implements Disposable { acquirePtyCleanup: (directory) => this.acquirePtyCleanup(directory), metadata: (client, dir) => sandboxSessionMetadata(this.connectionService.sandboxPreference, client, dir), post: (msg) => this.postToWebview(msg), + notify: (message) => this.host.showError(message), log: (...args) => this.log(...args), } } diff --git a/packages/kilo-vscode/src/agent-manager/provider-lifecycle.ts b/packages/kilo-vscode/src/agent-manager/provider-lifecycle.ts index c5a1378522..509bf9a47b 100644 --- a/packages/kilo-vscode/src/agent-manager/provider-lifecycle.ts +++ b/packages/kilo-vscode/src/agent-manager/provider-lifecycle.ts @@ -53,6 +53,7 @@ export interface LifecycleHost { acquirePtyCleanup: (directory: string) => Promise<() => void> metadata: (client: KiloClient, dir: string) => Promise> post: (message: AgentManagerOutMessage) => void + notify: (message: string) => void log: (...args: unknown[]) => void } @@ -122,16 +123,27 @@ export async function deleteLifecycleWorktree( host.log(`Worktree ${worktreeId} not found in state`) return null } + const retained = new Set(state.getSessions(worktreeId).map((session) => session.id)) let client: KiloClient try { client = host.client() - const [status, permissions, questions] = await Promise.all([ + const [status, permissions, questions, sessions] = await Promise.all([ client.session.status({ directory: worktree.path }, { throwOnError: true }), client.permission.list({ directory: worktree.path }, { throwOnError: true }), client.question.list({ directory: worktree.path }, { throwOnError: true }), + client.experimental.session.list( + { directory: worktree.path, archived: true, roots: false, limit: Number.MAX_SAFE_INTEGER }, + { throwOnError: true }, + ), ]) - if (status.data === undefined || permissions.data === undefined || questions.data === undefined) + if ( + status.data === undefined || + permissions.data === undefined || + questions.data === undefined || + sessions.data === undefined + ) throw new Error("Deletion safety checks returned no data") + sessions.data.forEach((session) => retained.add(session.id)) const active = Object.values(status.data).some((value) => value.type !== "idle") if (active || permissions.data.length > 0 || questions.data.length > 0) { host.post({ @@ -204,24 +216,27 @@ export async function deleteLifecycleWorktree( } try { await ctx.worktreeManager().removeWorktree(worktree.path, branch) + await Promise.all( + [...retained].map((sessionID) => + client.experimental.controlPlane.moveSession( + { sessionID, destination: { directory: ctx.root }, moveChanges: false }, + { throwOnError: true }, + ), + ), + ) try { await client.kilocode.removeSnapshot({ directory: ctx.root, worktree: worktree.path }, { throwOnError: true }) } catch (error) { host.log(`Failed to remove worktree snapshots: ${error}`) - host.post({ - type: "error", - code: "agentManager.snapshotCleanupFailed", - projectId: ctx.id, - worktreeId, - message: - "The worktree was deleted, but its checkpoint data could not be removed. Conversation history is preserved.", - }) + host.notify( + "The worktree was deleted, but its checkpoint data could not be removed. Conversation history is preserved.", + ) } const orphaned = state.removeWorktree(worktreeId) host.removePR(worktreeId) host.forgetName(worktreeId) host.stopDiffs(worktree.path, orphaned) - for (const s of orphaned) routeProjectSession(host.sessions, ctx.id, s.id, ctx.root, ctx.generation) + for (const sessionID of retained) routeProjectSession(host.sessions, ctx.id, sessionID, ctx.root, ctx.generation) host.push() host.log(`Deleted worktree ${worktreeId}${branch ? ` (${branch})` : ""}`) } catch (error) { diff --git a/packages/kilo-vscode/tests/unit/agent-manager-arch.test.ts b/packages/kilo-vscode/tests/unit/agent-manager-arch.test.ts index b3911f7ea0..cdb1f2f1cf 100644 --- a/packages/kilo-vscode/tests/unit/agent-manager-arch.test.ts +++ b/packages/kilo-vscode/tests/unit/agent-manager-arch.test.ts @@ -663,7 +663,8 @@ describe("Agent Manager Provider — onMessage routing", () => { expect(text).toContain("worktree: worktree.path") expect(text).toContain("throwOnError: true") expect(text).not.toContain("session.delete") - expect(text).toContain("routeProjectSession(host.sessions, ctx.id, s.id, ctx.root, ctx.generation)") + expect(text).toContain("client.experimental.controlPlane.moveSession") + expect(text).toContain("routeProjectSession(host.sessions, ctx.id, sessionID, ctx.root, ctx.generation)") expect(text).not.toContain("sessions.clearDirectory(s.id)") expect(text).toContain("host.push()") for (const name of ["onCreateWorktree", "onCreateMultiVersion", "onRemoveStaleWorktree"]) { diff --git a/packages/kilo-vscode/tests/unit/agent-manager-provider-lifecycle.test.ts b/packages/kilo-vscode/tests/unit/agent-manager-provider-lifecycle.test.ts index 58a21e0b7d..7d9b4d32a8 100644 --- a/packages/kilo-vscode/tests/unit/agent-manager-provider-lifecycle.test.ts +++ b/packages/kilo-vscode/tests/unit/agent-manager-provider-lifecycle.test.ts @@ -18,6 +18,7 @@ describe("Agent Manager worktree deletion lifecycle", () => { session: { status: ReturnType; delete: ReturnType } permission: { list: ReturnType } question: { list: ReturnType } + experimental: { session: { list: ReturnType }; controlPlane: { moveSession: ReturnType } } kilocode: { removeSnapshot: ReturnType } } let host: LifecycleHost @@ -46,7 +47,24 @@ describe("Agent Manager worktree deletion lifecycle", () => { }, permission: { list: mock(async () => ({ data: [] })) }, question: { list: mock(async () => ({ data: [] })) }, - kilocode: { removeSnapshot: mock(async () => ({ data: true })) }, + experimental: { + session: { + list: mock(async () => ({ + data: state.getSessions().map((session) => ({ id: session.id, directory: worktree })), + })), + }, + controlPlane: { + moveSession: mock(async ({ sessionID }: { sessionID: string }) => { + calls.push(`move:${sessionID}`) + }), + }, + }, + kilocode: { + removeSnapshot: mock(async () => { + calls.push("snapshots") + return { data: true } + }), + }, } host = { createOnDisk: async () => null, @@ -84,6 +102,7 @@ describe("Agent Manager worktree deletion lifecycle", () => { }, metadata: async () => ({}), post: (message) => calls.push(`post:${message.type}`), + notify: (message) => calls.push(`notify:${message}`), log: () => undefined, } }) @@ -134,18 +153,14 @@ describe("Agent Manager worktree deletion lifecycle", () => { it("reports checkpoint cleanup failures while preserving and retargeting sessions", async () => { const session = state.addSession("retained", state.getWorktrees()[0]!.id) - const post = mock(host.post) - host.post = post + const notify = mock(host.notify) + host.notify = notify client.kilocode.removeSnapshot.mockRejectedValue(new Error("checkpoint cleanup failed")) await deleteWorktree() - expect(post).toHaveBeenCalledWith( - expect.objectContaining({ - type: "error", - code: "agentManager.snapshotCleanupFailed", - projectId: ctx.id, - }), + expect(notify).toHaveBeenCalledWith( + "The worktree was deleted, but its checkpoint data could not be removed. Conversation history is preserved.", ) expect(state.getWorktrees()).toHaveLength(0) expect(routes).toContainEqual({ @@ -157,6 +172,34 @@ describe("Agent Manager worktree deletion lifecycle", () => { expect(client.session.delete).not.toHaveBeenCalled() }) + it("relocates archived and child sessions not present in Agent Manager state", async () => { + client.experimental.session.list.mockResolvedValue({ + data: [ + { id: "archived", directory: worktree, time: { archived: 1 } }, + { id: "child", directory: worktree, parentID: "parent" }, + ], + }) + + await deleteWorktree() + + expect(routes.map((route) => route.sessionID)).toEqual(["archived", "child"]) + expect(client.experimental.controlPlane.moveSession).toHaveBeenCalledTimes(2) + expect(client.session.delete).not.toHaveBeenCalled() + }) + + it("does not discard checkpoints when persistent session relocation fails", async () => { + state.addSession("retained", state.getWorktrees()[0]!.id) + client.experimental.controlPlane.moveSession.mockRejectedValue(new Error("move failed")) + + await deleteWorktree() + + expect(calls).toContain("disk") + expect(calls).toContain("post:error") + expect(client.kilocode.removeSnapshot).not.toHaveBeenCalled() + expect(state.getWorktrees()).toHaveLength(1) + expect(client.session.delete).not.toHaveBeenCalled() + }) + it("retargets orphaned sessions to the exact project root without deleting them", async () => { const first = state.addSession("first", state.getWorktrees()[0]!.id) const second = state.addSession("second", state.getWorktrees()[0]!.id) @@ -170,6 +213,18 @@ describe("Agent Manager worktree deletion lifecycle", () => { expect(calls).not.toContain(`clear:${first.id}`) expect(calls).not.toContain(`clear:${second.id}`) expect(client.session.delete).not.toHaveBeenCalled() + expect(client.experimental.session.list).toHaveBeenCalledWith( + { directory: worktree, archived: true, roots: false, limit: Number.MAX_SAFE_INTEGER }, + { throwOnError: true }, + ) + for (const session of [first, second]) { + expect(client.experimental.controlPlane.moveSession).toHaveBeenCalledWith( + { sessionID: session.id, destination: { directory: ctx.root }, moveChanges: false }, + { throwOnError: true }, + ) + expect(calls.indexOf(`move:${session.id}`)).toBeGreaterThan(calls.indexOf("disk")) + expect(calls.indexOf(`move:${session.id}`)).toBeLessThan(calls.indexOf("snapshots")) + } expect(client.kilocode.removeSnapshot).toHaveBeenCalledWith( { directory: ctx.root, worktree }, { throwOnError: true },