fix(agent-manager): recover worktrees after restart (#9939)

.changeset/agent-manager-worktree-recovery.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Recover Agent Manager worktrees after restart when saved state is missing, corrupt, or stale.

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

@@ -7,7 +7,7 @@ import { resolveLocalDiffTarget } from "../review-utils"
 import { getDiffMarkdownRender, setDiffMarkdownRender } from "../review-settings"
 import { isAbsolutePath } from "../path-utils"
 import { WorktreeManager, type CreateWorktreeResult } from "./WorktreeManager"
-import { remoteRef, WorktreeStateManager } from "./WorktreeStateManager"
+import { remoteRef, WorktreeStateManager, type Worktree } from "./WorktreeStateManager"
 import { handleSection } from "./section-handler"
 import { chooseBaseBranch, normalizeBaseBranch } from "./base-branch"
 import { GitStatsPoller, type LocalStats, type WorktreePresenceResult, type WorktreeStats } from "./GitStatsPoller"
@@ -29,6 +29,7 @@ import { forkSession } from "./fork-session"
 import { continueInWorktree } from "./continue-in-worktree"
 import { WorktreeDiffController } from "./worktree-diff-controller"
 import { WorktreeImporter } from "./worktree-importer"
+import { restoreWorktrees } from "./state-recovery"
 import { diffSummary as localDiffSummary, diffFile as localDiffFile } from "./local-diff"
 import { parseToolRequest, startFromTool, type ToolRequest } from "./tool-start"
 
@@ -68,6 +69,7 @@ export class AgentManagerProvider implements Disposable {
   private cachedWorktreeStats: { type: "agentManager.worktreeStats"; stats: WorktreeStats[] } | undefined
   private cachedLocalStats: { type: "agentManager.localStats"; stats: LocalStats } | undefined
   private unsubTool: (() => void) | undefined
+  private closing: Promise<void> | undefined
 
   /** Session ID most recently loaded via a `loadMessages` message from the webview.
    *  Updated synchronously — unlike the session provider's currentSession which depends on
@@ -248,14 +250,22 @@ export class AgentManagerProvider implements Disposable {
       return
     }
 
-    const migration = await state.load()
+    const loaded = await state.load()
     manager.cleanupOrphanedTempDirs()
 
+    if (loaded.status === "failed" && !(await state.prepareRecovery())) {
+      this.postToWebview({ type: "error", message: "Agent Manager state could not be recovered." })
+      this.pushState()
+      return
+    }
+
+    await this.recoverWorktrees(manager, state)
+
     // When the .kilocode → .kilo migration rewrote git worktree refs, nudge
     // VS Code's git extension to re-discover them. Without this, worktrees
     // won't appear in Source Control until the next VS Code restart.
-    if (migration.refsFixed > 0) {
-      this.log(`Migration fixed ${migration.refsFixed} git worktree ref(s), refreshing git`)
+    if (loaded.refsFixed > 0) {
+      this.log(`Migration fixed ${loaded.refsFixed} git worktree ref(s), refreshing git`)
       this.host.refreshGit()
     }
 
@@ -280,6 +290,20 @@ export class AgentManagerProvider implements Disposable {
     this.panel?.sessions.recoverPendingPrompts()
   }
 
+  private async recoverWorktrees(manager: WorktreeManager, state: WorktreeStateManager): Promise<void> {
+    const infos = await manager.discoverWorktrees().catch((err) => {
+      this.log("Failed to discover worktrees during state recovery:", err)
+      return []
+    })
+    if (infos.length === 0) return
+
+    const result = restoreWorktrees(state, infos)
+    if (result.worktrees === 0 && result.sessions === 0) return
+
+    this.log(`Recovered ${result.worktrees} worktree(s) and ${result.sessions} session(s) from disk`)
+    await state.flush()
+  }
+
   // ---------------------------------------------------------------------------
   // Message interceptor
   // ---------------------------------------------------------------------------
@@ -291,6 +315,7 @@ export class AgentManagerProvider implements Disposable {
     }
     msg = await this.contextMessage(msg)
     const m = msg as unknown as AgentManagerInMessage
+    if (this.shouldWaitForState(m)) await this.waitForStateReady(m.type)
 
     const worktree = await this.onWorktreeMessage(m)
     if (worktree !== undefined) return worktree
@@ -803,6 +828,39 @@ export class AgentManagerProvider implements Disposable {
     await this.stateReady.catch((err) => this.log(`${context}: stateReady rejected, continuing:`, err))
   }
 
+  private shouldWaitForState(m: AgentManagerInMessage): boolean {
+    switch (m.type) {
+      case "agentManager.deleteWorktree":
+      case "agentManager.removeStaleWorktree":
+      case "agentManager.openLocally":
+      case "agentManager.addSessionToWorktree":
+      case "agentManager.closeSession":
+      case "agentManager.persistSession":
+      case "agentManager.forgetSession":
+      case "agentManager.renameWorktree":
+      case "agentManager.requestBranches":
+      case "agentManager.importFromBranch":
+      case "agentManager.importFromPR":
+      case "agentManager.importExternalWorktree":
+      case "agentManager.importAllExternalWorktrees":
+      case "agentManager.setTabOrder":
+      case "agentManager.setWorktreeOrder":
+      case "agentManager.setSessionsCollapsed":
+      case "agentManager.setReviewDiffStyle":
+      case "agentManager.setDefaultBaseBranch":
+      case "agentManager.createSection":
+      case "agentManager.renameSection":
+      case "agentManager.deleteSection":
+      case "agentManager.setSectionColor":
+      case "agentManager.toggleSectionCollapsed":
+      case "agentManager.moveToSection":
+      case "agentManager.moveSection":
+        return true
+      default:
+        return false
+    }
+  }
+
   private onToolEvent(event: unknown, directory?: string): void {
     const properties = (event as { properties?: unknown }).properties
     const req = parseToolRequest(properties)
@@ -1296,6 +1354,9 @@ export class AgentManagerProvider implements Disposable {
   // ---------------------------------------------------------------------------
 
   private registerWorktreeSession(sessionId: string, directory: string): void {
+    const worktree = this.state?.findWorktreeByPath(directory)
+    if (worktree) this.writeMetadata(sessionId, worktree)
+
     if (!this.panel) return
     this.panel.sessions.setSessionDirectory(sessionId, directory)
     this.panel.sessions.trackSession(sessionId)
@@ -1305,6 +1366,14 @@ export class AgentManagerProvider implements Disposable {
     this.panel.sessions.recoverPendingPrompts()
   }
 
+  private writeMetadata(sessionId: string, worktree: Worktree): void {
+    const manager = this.getWorktreeManager()
+    if (!manager) return
+    void manager
+      .writeMetadata(worktree.path, sessionId, worktree.parentBranch, worktree.remote)
+      .catch((err) => this.log(`Failed to write worktree metadata for ${worktree.id}:`, err))
+  }
+
   /** Route a plan follow-up session to its worktree instead of LOCAL. */
   private adoptFollowupInWorktree(session: Session, directory: string): void {
     const state = this.getStateManager()
@@ -1614,7 +1683,18 @@ export class AgentManagerProvider implements Disposable {
     this.panel?.postMessage(message)
   }
 
+  public shutdown(): Promise<void> {
+    if (!this.closing) this.closing = this.disposeAsync()
+    return this.closing
+  }
+
   public dispose(): void {
+    void this.shutdown()
+  }
+
+  private async disposeAsync(): Promise<void> {
+    await this.stateReady?.catch((err) => this.log("dispose: stateReady rejected:", err))
+    await this.state?.flush().catch((err) => this.log("dispose: state flush failed:", err))
     this.unsubTool?.()
     this.connectionService.unregisterFocused("agent-manager")
     this.connectionService.registerOpen("agent-manager", [])
@@ -1624,7 +1704,7 @@ export class AgentManagerProvider implements Disposable {
     this.prBridge.poller.stop()
     this.run.dispose()
     this.terminalManager.dispose()
-    void this.terminalRouter.dispose()
+    await this.terminalRouter.dispose()
     this.panel?.dispose()
     this.outputChannel.dispose()
     this.host.dispose()

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

@@ -31,7 +31,7 @@ import {
 const TEMP_PREFIX = ".kilo-delete-"
 const RM_OPTS: fs.RmOptions = { recursive: true, force: true, maxRetries: 3, retryDelay: 200 }
 
-interface WorktreeInfo {
+export interface WorktreeInfo {
   branch: string
   path: string
   /** Bare branch name (e.g. "main"), without remote prefix. */

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

@@ -75,6 +75,12 @@ interface StateFile {
   defaultBaseBranch?: string
 }
 
+export type StateLoadStatus = "loaded" | "missing" | "failed"
+
+export interface StateLoadResult extends MigrationResult {
+  status: StateLoadStatus
+}
+
 import { KILO_DIR, migrateAgentManagerData, type MigrationResult } from "./constants"
 
 const STATE_FILE = "agent-manager.json"
@@ -97,7 +103,8 @@ export class WorktreeStateManager {
   private defaultBase: string | undefined
   private readonly log: (msg: string) => void
   private saving: Promise<void> | undefined
-  private pendingSave = false
+  private dirty = false
+  private failed = false
 
   private readonly root: string
   private migrated = false
@@ -188,6 +195,31 @@ export class WorktreeStateManager {
     return wt
   }
 
+  restoreWorktree(params: {
+    branch: string
+    path: string
+    parentBranch: string
+    remote?: string
+    createdAt: string
+  }): Worktree {
+    const existing = this.findWorktreeByPath(params.path)
+    if (existing) return existing
+    const id = generateId("wt")
+    const wt: Worktree = {
+      id,
+      branch: params.branch,
+      path: params.path,
+      parentBranch: params.parentBranch,
+      createdAt: params.createdAt,
+    }
+    if (params.remote) wt.remote = params.remote
+    this.worktrees.set(id, wt)
+    this.setNormalizedWorktreeOrder(this.worktreeOrder)
+    this.log(`Restored worktree ${id}: ${params.branch} (${params.path})`)
+    void this.save()
+    return wt
+  }
+
   updateWorktreeBranch(id: string, branch: string): boolean {
     const wt = this.worktrees.get(id)
     if (!wt || wt.branch === branch) return false
@@ -507,7 +539,7 @@ export class WorktreeStateManager {
   // Persistence
   // ---------------------------------------------------------------------------
 
-  async load(): Promise<MigrationResult> {
+  async load(): Promise<StateLoadResult> {
     // Migrate Agent Manager data from .kilocode → .kilo before first read
     let migration: MigrationResult = { refsFixed: 0 }
     if (!this.migrated) {
@@ -518,15 +550,34 @@ export class WorktreeStateManager {
       const content = await fs.promises.readFile(this.file, "utf-8")
       this.apply(content)
       this.loadFailed = false
+      return { ...migration, status: "loaded" }
     } catch (error) {
       const code = (error as NodeJS.ErrnoException).code
-      if (code === "ENOENT") this.loadFailed = false
+      if (code === "ENOENT") {
+        this.loadFailed = false
+        return { ...migration, status: "missing" }
+      }
       if (code !== "ENOENT") {
         this.log(`Failed to load state: ${error}`)
         this.loadFailed = true
       }
     }
-    return migration
+    return { ...migration, status: "failed" }
+  }
+
+  async prepareRecovery(): Promise<boolean> {
+    if (!this.loadFailed) return true
+    const stamp = new Date().toISOString().replace(/[:.]/g, "-")
+    const backup = `${this.file}.corrupt-${stamp}`
+    try {
+      await fs.promises.rename(this.file, backup)
+      this.loadFailed = false
+      this.log(`Backed up unreadable state to ${backup}`)
+      return true
+    } catch (error) {
+      this.log(`Failed to back up unreadable state: ${error}`)
+      return false
+    }
   }
 
   private apply(content: string): void {
@@ -610,7 +661,9 @@ export class WorktreeStateManager {
 
   /** Wait for any in-flight save to complete without triggering a new one. */
   async flush(): Promise<void> {
-    if (this.saving) await this.saving
+    const active = this.saving
+    if (active) await active
+    if (this.dirty) await this.save()
   }
 
   async save(): Promise<void> {
@@ -619,27 +672,32 @@ export class WorktreeStateManager {
       return
     }
 
-    // Serialize concurrent saves — if a save is in-flight, queue one follow-up
-    if (this.saving) {
-      this.pendingSave = true
-      await this.saving
-      // The in-flight save finished but our data may not have been written yet.
-      // If there's a new save already running (the pendingSave follow-up), wait for it.
-      if (this.saving) await this.saving
-      return
+    this.dirty = true
+    this.failed = false
+    while (this.dirty && !this.failed) {
+      await (this.saving ?? this.startSave())
     }
+  }
 
-    this.saving = this.writeToDisk()
-    try {
-      await this.saving
-    } finally {
-      this.saving = undefined
-    }
+  private startSave(): Promise<void> {
+    const run = this.drain().finally(() => {
+      if (this.saving === run) this.saving = undefined
+    })
+    this.saving = run
+    return run
+  }
 
-    // If another save was requested while we were writing, flush it now
-    if (this.pendingSave) {
-      this.pendingSave = false
-      await this.save()
+  private async drain(): Promise<void> {
+    while (this.dirty) {
+      this.dirty = false
+      try {
+        await this.writeToDisk()
+      } catch (error) {
+        this.dirty = true
+        this.failed = true
+        this.log(`Failed to save state: ${error}`)
+        return
+      }
     }
   }
 

packages/kilo-vscode/src/agent-manager/state-recovery.ts

@@ -0,0 +1,39 @@
+import type { WorktreeInfo } from "./WorktreeManager"
+import type { WorktreeStateManager } from "./WorktreeStateManager"
+
+export interface RecoveryResult {
+  worktrees: number
+  sessions: number
+}
+
+export function restoreWorktrees(state: WorktreeStateManager, infos: WorktreeInfo[]): RecoveryResult {
+  const result: RecoveryResult = { worktrees: 0, sessions: 0 }
+
+  for (const info of infos) {
+    const existing = state.findWorktreeByPath(info.path)
+    const wt =
+      existing ??
+      state.restoreWorktree({
+        branch: info.branch,
+        path: info.path,
+        parentBranch: info.parentBranch,
+        remote: info.remote,
+        createdAt: new Date(info.createdAt).toISOString(),
+      })
+
+    if (!existing) result.worktrees++
+    if (!info.sessionId) continue
+
+    const session = state.getSession(info.sessionId)
+    if (!session) {
+      state.addSession(info.sessionId, wt.id)
+      result.sessions++
+      continue
+    }
+    if (session.worktreeId === wt.id) continue
+    state.moveSession(info.sessionId, wt.id)
+    result.sessions++
+  }
+
+  return result
+}

packages/kilo-vscode/src/extension.ts

@@ -21,6 +21,8 @@ import { registerHeapSnapshot } from "./commands/heap-snapshot"
 import { RemoteStatusService } from "./services/RemoteStatusService"
 import { markWorkspace } from "./util/spotlight"
 
+let agentManager: AgentManagerProvider | undefined
+
 // Activated via "onStartupFinished" (package.json) so that commands, code actions, keybindings,
 // autocomplete, commit-message generation, and URI deep links all work immediately — without
 // requiring the user to open a Kilo sidebar or panel first. The CLI backend is NOT spawned here;
@@ -111,6 +113,7 @@ export function activate(context: vscode.ExtensionContext) {
   // Create Agent Manager provider for editor panel
   const agentManagerHost = new VscodeHost(context.extensionUri, connectionService, context)
   const agentManagerProvider = new AgentManagerProvider(agentManagerHost, connectionService)
+  agentManager = agentManagerProvider
   context.subscriptions.push(agentManagerProvider)
 
   // Wire "Continue in Worktree" from sidebar → Agent Manager
@@ -431,7 +434,8 @@ export function activate(context: vscode.ExtensionContext) {
   })
 }
 
-export function deactivate() {
+export async function deactivate() {
+  await agentManager?.shutdown()
   TelemetryProxy.getInstance().shutdown()
 }
 

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

@@ -154,6 +154,34 @@ describe("Agent Manager Provider Messages", () => {
     const body = getMethodBody("onAddSessionToWorktree")
     expect(body).toContain("agentManager.sessionAdded")
   })
+
+  it("state-mutating messages wait for state initialization", () => {
+    const body = getMethodBody("shouldWaitForState")
+    const messages = [
+      "agentManager.setTabOrder",
+      "agentManager.setWorktreeOrder",
+      "agentManager.persistSession",
+      "agentManager.forgetSession",
+      "agentManager.importFromBranch",
+      "agentManager.importFromPR",
+      "agentManager.importExternalWorktree",
+      "agentManager.importAllExternalWorktrees",
+      "agentManager.createSection",
+      "agentManager.moveToSection",
+    ]
+
+    for (const message of messages) {
+      expect(body, `${message} should wait for loaded state`).toContain(message)
+    }
+
+    expect(getMethodBody("onMessage")).toContain("if (this.shouldWaitForState(m)) await this.waitForStateReady(m.type)")
+  })
+
+  it("async shutdown waits for terminal router cleanup", () => {
+    const body = getMethodBody("disposeAsync")
+    expect(body).toContain("await this.terminalRouter.dispose()")
+    expect(body).not.toContain("void this.terminalRouter.dispose()")
+  })
 })
 
 // ---------------------------------------------------------------------------

packages/kilo-vscode/tests/unit/worktree-manager.test.ts

@@ -543,11 +543,12 @@ describe("WorktreeManager metadata", () => {
     const mgr = createManager(root)
     const result = await mgr.createWorktree({ prompt: "session-test" })
 
-    await mgr.writeMetadata(result.path, "sess-abc-123", "feature-branch")
+    await mgr.writeMetadata(result.path, "sess-abc-123", "feature-branch", "origin")
     const meta = await mgr.readMetadata(result.path)
 
     expect(meta?.sessionId).toBe("sess-abc-123")
     expect(meta?.parentBranch).toBe("feature-branch")
+    expect(meta?.remote).toBe("origin")
   })
 
   it("returns undefined when no metadata exists", async () => {

packages/kilo-vscode/tests/unit/worktree-state-manager.test.ts

@@ -3,6 +3,8 @@ import * as fs from "fs"
 import * as path from "path"
 import * as os from "os"
 import { WorktreeStateManager } from "../../src/agent-manager/WorktreeStateManager"
+import { restoreWorktrees } from "../../src/agent-manager/state-recovery"
+import type { WorktreeInfo } from "../../src/agent-manager/WorktreeManager"
 
 describe("WorktreeStateManager", () => {
   let root: string
@@ -160,8 +162,9 @@ describe("WorktreeStateManager", () => {
       await manager.save()
 
       const loaded = new WorktreeStateManager(root, () => {})
-      await loaded.load()
+      const result = await loaded.load()
 
+      expect(result.status).toBe("loaded")
       expect(loaded.getWorktrees()).toHaveLength(1)
       expect(loaded.getWorktrees()[0].branch).toBe("fix")
       expect(loaded.getSessions()).toHaveLength(2)
@@ -171,11 +174,49 @@ describe("WorktreeStateManager", () => {
     })
 
     it("load is a no-op when file does not exist", async () => {
-      await manager.load()
+      const result = await manager.load()
+      expect(result.status).toBe("missing")
       expect(manager.getWorktrees()).toHaveLength(0)
       expect(manager.getSessions()).toHaveLength(0)
     })
 
+    it("flush waits for saves queued during an in-flight write", async () => {
+      const file = path.join(root, ".kilo", "agent-manager.json")
+      const api = fs.promises as unknown as { writeFile: (...args: unknown[]) => Promise<void> }
+      const original = api.writeFile
+      const gate = {
+        release: () => {},
+        promise: Promise.resolve(),
+      }
+      gate.promise = new Promise<void>((resolve) => {
+        gate.release = resolve
+      })
+      const state = { blocked: false }
+
+      api.writeFile = async (...args: unknown[]) => {
+        const target = typeof args[0] === "string" ? args[0] : ""
+        if (!state.blocked && target.includes("agent-manager.json.")) {
+          state.blocked = true
+          await gate.promise
+        }
+        await original(...args)
+      }
+
+      try {
+        manager.addSession("first", null)
+        await Promise.resolve()
+        manager.addSession("second", null)
+        gate.release()
+        await manager.flush()
+      } finally {
+        api.writeFile = original
+      }
+
+      const data = JSON.parse(fs.readFileSync(file, "utf-8")) as { sessions: Record<string, unknown> }
+      expect(data.sessions.first).toBeDefined()
+      expect(data.sessions.second).toBeDefined()
+    })
+
     it("creates .kilo directory if missing", async () => {
       const fresh = path.join(root, "subdir")
       const mgr = new WorktreeStateManager(fresh, () => {})
@@ -190,15 +231,32 @@ describe("WorktreeStateManager", () => {
       const file = path.join(root, ".kilo", "agent-manager.json")
       fs.writeFileSync(file, "{", "utf-8")
 
-      await manager.load()
+      const result = await manager.load()
       manager.addSession("local-after-failure", null)
       await manager.flush()
       await manager.save()
 
+      expect(result.status).toBe("failed")
       expect(fs.readFileSync(file, "utf-8")).toBe("{")
       expect(logs.some((l) => l.includes("Skipping save because state failed to load"))).toBe(true)
     })
 
+    it("backs up a corrupt state file before recovery saves", async () => {
+      const file = path.join(root, ".kilo", "agent-manager.json")
+      fs.writeFileSync(file, "{", "utf-8")
+
+      const result = await manager.load()
+      const recovered = await manager.prepareRecovery()
+      manager.addSession("local-after-recovery", null)
+      await manager.flush()
+
+      const files = fs.readdirSync(path.join(root, ".kilo"))
+      expect(result.status).toBe("failed")
+      expect(recovered).toBe(true)
+      expect(files.some((item) => item.startsWith("agent-manager.json.corrupt-"))).toBe(true)
+      expect(JSON.parse(fs.readFileSync(file, "utf-8")).sessions["local-after-recovery"].worktreeId).toBeNull()
+    })
+
     it("allows saves after a later missing-file reload", async () => {
       const file = path.join(root, ".kilo", "agent-manager.json")
       fs.writeFileSync(file, "{", "utf-8")
@@ -213,6 +271,30 @@ describe("WorktreeStateManager", () => {
     })
   })
 
+  describe("recovery", () => {
+    it("restores discovered worktrees and metadata sessions", async () => {
+      const infos: WorktreeInfo[] = [
+        {
+          branch: "fix-recovered",
+          path: "/tmp/recovered",
+          parentBranch: "main",
+          remote: "origin",
+          createdAt: Date.UTC(2026, 0, 1),
+          sessionId: "sess-recovered",
+        },
+      ]
+
+      const result = restoreWorktrees(manager, infos)
+      await manager.flush()
+
+      const worktree = manager.findWorktreeByPath("/tmp/recovered")
+      expect(result).toEqual({ worktrees: 1, sessions: 1 })
+      expect(worktree?.branch).toBe("fix-recovered")
+      expect(worktree?.remote).toBe("origin")
+      expect(manager.getSession("sess-recovered")?.worktreeId).toBe(worktree?.id)
+    })
+  })
+
   describe("tab order", () => {
     it("sets and gets tab order for a key", () => {
       manager.setTabOrder("wt-1", ["s1", "s2", "s3"])