From 8ed0e17855baee32d48fa44881761a8ee6b53f1e Mon Sep 17 00:00:00 2001 From: Marius Date: Wed, 6 May 2026 10:25:55 +0200 Subject: [PATCH] fix(agent-manager): recover worktrees after restart (#9939) --- .changeset/agent-manager-worktree-recovery.md | 5 + .../src/agent-manager/AgentManagerProvider.ts | 90 ++++++++++++++- .../src/agent-manager/WorktreeManager.ts | 2 +- .../src/agent-manager/WorktreeStateManager.ts | 104 ++++++++++++++---- .../src/agent-manager/state-recovery.ts | 39 +++++++ packages/kilo-vscode/src/extension.ts | 6 +- .../tests/unit/agent-manager-arch.test.ts | 28 +++++ .../tests/unit/worktree-manager.test.ts | 3 +- .../tests/unit/worktree-state-manager.test.ts | 88 ++++++++++++++- 9 files changed, 331 insertions(+), 34 deletions(-) create mode 100644 .changeset/agent-manager-worktree-recovery.md create mode 100644 packages/kilo-vscode/src/agent-manager/state-recovery.ts diff --git a/.changeset/agent-manager-worktree-recovery.md b/.changeset/agent-manager-worktree-recovery.md new file mode 100644 index 0000000000..f0479fb492 --- /dev/null +++ b/.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. diff --git a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts b/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts index a511ec63c8..ffcf786856 100644 --- a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts +++ b/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 | 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 { + 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 { + if (!this.closing) this.closing = this.disposeAsync() + return this.closing + } + public dispose(): void { + void this.shutdown() + } + + private async disposeAsync(): Promise { + 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() diff --git a/packages/kilo-vscode/src/agent-manager/WorktreeManager.ts b/packages/kilo-vscode/src/agent-manager/WorktreeManager.ts index 6a64d4bd17..7ed0c1e011 100644 --- a/packages/kilo-vscode/src/agent-manager/WorktreeManager.ts +++ b/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. */ diff --git a/packages/kilo-vscode/src/agent-manager/WorktreeStateManager.ts b/packages/kilo-vscode/src/agent-manager/WorktreeStateManager.ts index 3491a61396..478c2f71da 100644 --- a/packages/kilo-vscode/src/agent-manager/WorktreeStateManager.ts +++ b/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 | 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 { + async load(): Promise { // 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 { + 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 { - if (this.saving) await this.saving + const active = this.saving + if (active) await active + if (this.dirty) await this.save() } async save(): Promise { @@ -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 { + 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 { + 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 + } } } diff --git a/packages/kilo-vscode/src/agent-manager/state-recovery.ts b/packages/kilo-vscode/src/agent-manager/state-recovery.ts new file mode 100644 index 0000000000..b39cbc0f52 --- /dev/null +++ b/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 +} diff --git a/packages/kilo-vscode/src/extension.ts b/packages/kilo-vscode/src/extension.ts index fadee229dc..a393f7843e 100644 --- a/packages/kilo-vscode/src/extension.ts +++ b/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() } 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 96f7785e96..cd51acefe5 100644 --- a/packages/kilo-vscode/tests/unit/agent-manager-arch.test.ts +++ b/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()") + }) }) // --------------------------------------------------------------------------- diff --git a/packages/kilo-vscode/tests/unit/worktree-manager.test.ts b/packages/kilo-vscode/tests/unit/worktree-manager.test.ts index 29a70f1038..96d0737182 100644 --- a/packages/kilo-vscode/tests/unit/worktree-manager.test.ts +++ b/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 () => { diff --git a/packages/kilo-vscode/tests/unit/worktree-state-manager.test.ts b/packages/kilo-vscode/tests/unit/worktree-state-manager.test.ts index e711dd252a..a1c104582f 100644 --- a/packages/kilo-vscode/tests/unit/worktree-state-manager.test.ts +++ b/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 } + const original = api.writeFile + const gate = { + release: () => {}, + promise: Promise.resolve(), + } + gate.promise = new Promise((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 } + 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"])