fix(agent-manager): safely clean up deleted worktree snapshots

.changeset/prune-orphaned-worktree-snapshots.md

@@ -0,0 +1,6 @@
+---
+"@kilocode/cli": patch
+"kilo-code": patch
+---
+
+Remove deleted worktree checkpoints without losing conversation history and stop showing activity for deleted sessions.

packages/kilo-vscode/src/KiloProvider.ts

@@ -379,6 +379,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
   private promptRecoveryQueued = false
   private promptRecovery: Promise<void> | null = null
   private trackedSessionIds: Set<string> = new Set()
+  private readonly removedSessionIds = new Set<string>()
   private readonly openSessionIds = new Set<string>()
   private modelUsageSessionIds: Set<string> = new Set()
   private syncedChildSessions: Set<string> = new Set()
@@ -801,6 +802,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
 
   /** Register a session created externally and notify the webview. */
   public registerSession(session: Session, activate = false): void {
+    this.removedSessionIds.delete(session.id)
     this.stopCurrentSessionProcesses(session.id)
     this.setCurrentSession(session)
     this.contextSessionID = session.id
@@ -2340,6 +2342,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
    * a session the backend has already deleted.
    */
   private pruneDeletedSession(sessionID: string): void {
+    this.removedSessionIds.add(sessionID)
     this.trackedSessionIds.delete(sessionID)
     this.openSessionIds.delete(sessionID)
     for (const [key, session] of this.draftSessions) {
@@ -4724,6 +4727,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
     // busy-session warning on Save.
     if (event.type === "session.status") {
       const sid = event.properties.sessionID
+      if (this.removedSessionIds.has(sid)) return
       const status = event.properties.status
       this.mark(sid, directory)
       this.aborts.observe(sid, status.type, directory)
@@ -5440,6 +5444,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
     this.promptRecoveryQueued = false
     clearNetworkWaits(this.trackedSessionIds)
     this.trackedSessionIds.clear()
+    this.removedSessionIds.clear()
     this.openSessionIds.clear()
     this.syncedChildSessions.clear()
     this.inspectorSessionIds.clear()

packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts

@@ -118,6 +118,7 @@ export class AgentManagerProvider implements Disposable {
   private onVisibilityChange: ((visible: boolean) => void) | undefined
   private panelSessions = new Set<string>()
   private busySessions = new Set<string>()
+  private removedSessions = new Set<string>()
   readonly settings: ProjectWiring["settings"]
   /** Session ID most recently loaded via `loadMessages`; updated synchronously. */
   private activeSessionId: string | undefined
@@ -308,6 +309,7 @@ export class AgentManagerProvider implements Disposable {
     if (ev.type === "session.deleted") {
       const id = ev.properties?.sessionID
       if (!id) return
+      this.removedSessions.add(id)
       this.busySessions.delete(id)
       const ctx = this.contexts.byLiveSession(id)
       if (!ctx) return
@@ -316,7 +318,8 @@ export class AgentManagerProvider implements Disposable {
       return
     }
     const info = ev.properties?.info
-    const dir = info?.directory
+    if (ev.type === "session.created" && info) this.removedSessions.delete(info.id)
+    const dir = info && !this.removedSessions.has(info.id) ? info.directory : undefined
     // Session events from sync or older backends can lack time/directory; a throw
     // would escape into the SSE dispatch loop and starve the other listeners.
     if (!info?.time || !dir || (info.parentID !== undefined && info.parentID !== null)) return
@@ -332,12 +335,11 @@ export class AgentManagerProvider implements Disposable {
     ctx.invalidateSessions()
     this.postToWebview({ type: "agentManager.projectSessions", projectId: ctx.id, sessions: [...ctx.sessions()] })
   }
-
   private onSessionStatus(event: unknown): void {
     const props = (event as { properties?: { sessionID?: string; status?: { type?: string } } }).properties
     const sid = props?.sessionID
     const type = props?.status?.type
-    if (!sid || !type) return
+    if (!sid || !type || this.removedSessions.has(sid)) return
     if (type === "idle") {
       this.busySessions.delete(sid)
       this.naming.idle(sid)
@@ -346,12 +348,10 @@ export class AgentManagerProvider implements Disposable {
     this.busySessions.add(sid)
     this.naming.busy(sid)
   }
-
   private log(...args: unknown[]) {
     const msg = args.map((a) => (typeof a === "string" ? a : JSON.stringify(a))).join(" ")
     this.outputChannel.appendLine(`${new Date().toISOString()} ${msg}`)
   }
-
   public openPanel(preserveFocus?: boolean): void {
     if (this.panel) {
       this.log("Panel already open, revealing")
@@ -1123,7 +1123,6 @@ export class AgentManagerProvider implements Disposable {
       req,
     )
   }
-
   // Worktree actions
 
   /** Create a new worktree with an auto-created first session. */
@@ -1452,9 +1451,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 {
     return {
       createOnDisk: (opts) => this.createWorktreeOnDisk(opts),
@@ -1464,6 +1460,8 @@ export class AgentManagerProvider implements Disposable {
       sessions: {
         register: (session) => this.panel?.sessions.registerSession(session),
         clearDirectory: (sid) => this.panel?.sessions.clearSessionDirectory(sid),
+        setSessionDirectory: (sid, dir) => this.panel?.sessions.setSessionDirectory(sid, dir),
+        registerSessionRoute: (ref, dir, gen) => this.panel?.sessions.registerSessionRoute?.(ref, dir, gen),
         directories: () => this.panel?.sessions.getSessionDirectories(),
         abort: (ids) => this.panel?.sessions.abortSessions(ids) ?? Promise.resolve(),
         forget: (sid) => void this.panelSessions.delete(sid),
@@ -1485,6 +1483,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),
     }
   }

packages/kilo-vscode/src/agent-manager/provider-lifecycle.ts

@@ -9,6 +9,7 @@ import type { CreateWorktreeResult, WorktreeManager } from "./WorktreeManager"
 import type { CreateWorktreeOnDiskOptions, CreateWorktreeOnDiskResult } from "./worktree-create"
 import { recordPromotionHandoff } from "./promotion-handoff"
 import { stopSessionProcesses } from "../kilo-provider/background-process"
+import { routeProjectSession } from "./project/messages"
 
 /**
  * Provider capabilities the worktree lifecycle needs beyond project state.
@@ -25,6 +26,12 @@ export interface LifecycleHost {
   sessions: {
     register: (session: Session) => void
     clearDirectory: (sessionId: string) => void
+    setSessionDirectory: (sessionId: string, directory: string) => void
+    registerSessionRoute?: (
+      ref: { projectId: string; sessionId: string },
+      directory: string,
+      generation: number,
+    ) => void
     directories: () => ReadonlyMap<string, string> | undefined
     abort: (sessionIds: string[]) => Promise<void>
     forget: (sessionId: string) => void
@@ -45,6 +52,7 @@ export interface LifecycleHost {
   acquirePtyCleanup: (directory: string) => Promise<() => void>
   metadata: (client: KiloClient, dir: string) => Promise<Record<string, unknown>>
   post: (message: AgentManagerOutMessage) => void
+  notify: (message: string) => void
   log: (...args: unknown[]) => void
 }
 
@@ -114,14 +122,55 @@ export async function deleteLifecycleWorktree(
     host.log(`Worktree ${worktreeId} not found in state`)
     return null
   }
+  const fail = (message: string) => {
+    host.post({ type: "error", code: "agentManager.worktreeDeleteFailed", projectId: ctx.id, worktreeId, message })
+    return null
+  }
+  const retained = new Set(state.getSessions(worktreeId).map((session) => session.id))
+  let client: KiloClient
+  try {
+    client = host.client()
+    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 ||
+      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)
+      return fail("Cannot delete a worktree while a session is active or waiting for input")
+  } catch (error) {
+    host.log(`Failed to verify worktree deletion safety: ${error}`)
+    return fail("Cannot verify worktree sessions before deletion")
+  }
   // Stop pollers before cleanup. State is removed only after PTYs and disk are gone so a failed
   // process cleanup cannot leave a live shell rooted in an untracked worktree.
-  host.skipStats(worktreeId)
-  await host.removeRun(worktreeId)
-  if (!(await host.clearRun(worktreeId))) {
+  try {
+    host.skipStats(worktreeId)
+    await host.removeRun(worktreeId)
+  } catch (error) {
     host.unskipStats(worktreeId)
-    host.post({ type: "error", message: "Failed to stop the Run script before deleting the worktree" })
-    return null
+    host.log(`Failed to stop worktree services: ${error}`)
+    return fail("Failed to stop worktree services before deletion")
+  }
+  const cleared = await host.clearRun(worktreeId).catch((error) => {
+    host.log(`Failed to stop the Run script: ${error}`)
+    return false
+  })
+  if (!cleared) {
+    host.unskipStats(worktreeId)
+    return fail("Failed to stop the Run script before deleting the worktree")
   }
   const branch = worktree.branchOwned === false ? undefined : (worktree.originalBranch ?? worktree.branch)
   let releasePtyCleanup: () => void
@@ -130,17 +179,37 @@ export async function deleteLifecycleWorktree(
   } catch (error) {
     host.log(`Failed to remove worktree from disk: ${error}`)
     host.unskipStats(worktreeId)
-    return null
+    return fail("Failed to remove worktree PTYs before deletion")
   }
   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.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) host.sessions.clearDirectory(s.id)
+    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) {
+    host.unskipStats(worktreeId)
+    host.log(`Failed to delete worktree ${worktreeId}: ${error}`)
+    return fail("Failed to delete the worktree")
   } finally {
     releasePtyCleanup()
   }

packages/kilo-vscode/src/agent-manager/types.ts

@@ -263,6 +263,9 @@ interface ScriptTerminalsMessage {
 interface ErrorOutMessage {
   type: "error"
   message: string
+  code?: string
+  projectId?: string
+  worktreeId?: string
 }
 
 interface SessionAddedMessage {

packages/kilo-vscode/tests/unit/agent-manager-arch.test.ts

@@ -637,12 +637,23 @@ describe("Agent Manager Provider — onMessage routing", () => {
    * Regression: deletion must clean up both disk (manager) and state, then
    * push to webview. Missing any step leaves ghost worktrees or stale UI.
    */
-  it("onDeleteWorktree removes from disk, state, clears orphans, and pushes", () => {
+  it("does not restore running indicators after a session is deleted", () => {
+    const lifecycle = body("onSessionLifecycle")
+    const status = body("onSessionStatus")
+
+    expect(lifecycle).toContain("this.removedSessions.add(id)")
+    expect(lifecycle).toContain("this.busySessions.delete(id)")
+    expect(lifecycle).toContain("info && !this.removedSessions.has(info.id) ? info.directory : undefined")
+    expect(status).toContain("this.removedSessions.has(sid)")
+  })
+
+  it("limits snapshot cleanup to explicit worktree deletion without deleting sessions", () => {
     const text = body("onDeleteWorktree")
-    expect(text).toContain("worktreeManager().removeWorktree")
-    expect(text).toContain("state.removeWorktree")
-    expect(text).toContain("sessions.clearDirectory")
-    expect(text).toContain("host.push()")
+    expect(text).toContain(".kilocode.removeSnapshot")
+    expect(text).not.toContain("session.delete")
+    for (const name of ["onCreateWorktree", "onCreateMultiVersion", "onRemoveStaleWorktree"]) {
+      expect(body(name)).not.toContain("removeSnapshot")
+    }
   })
 
   // -- onCreateWorktree invariants -------------------------------------------