fix(agent-manager): clean up failed session moves

.changeset/move-worktrees-safely.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Clean up incomplete worktrees when moving a session fails to transfer its Git changes.

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

@@ -1746,6 +1746,18 @@ export class AgentManagerProvider implements Disposable {
         getClient: () => this.connectionService.getClient(),
         createWorktreeOnDisk: (opts) => this.createWorktreeOnDisk(opts),
         runSetupScript: (p, b, id) => this.runSetupScriptForWorktree(p, b, id),
+        cleanupWorktree: async (id) => {
+          await this.onDeleteWorktree(id)
+        },
+        notifyError: (error, result, id) => {
+          this.postToWebview({
+            type: "agentManager.worktreeSetup",
+            status: "error",
+            message: error,
+            branch: result.branch,
+            worktreeId: id,
+          })
+        },
         getStateManager: () => this.getStateManager(),
         registerWorktreeSession: (sid, dir) => this.registerWorktreeSession(sid, dir),
         registerSession: (session) => this.panel?.sessions.registerSession(session),

packages/kilo-vscode/src/agent-manager/continue-in-worktree.ts

@@ -14,6 +14,8 @@ export interface ContinueContext {
     result: CreateWorktreeResult
   } | null>
   runSetupScript: (path: string, branch: string, worktreeId: string) => Promise<void>
+  cleanupWorktree: (worktreeId: string) => Promise<void>
+  notifyError: (error: string, result: CreateWorktreeResult, worktreeId: string) => void
   getStateManager: () => WorktreeStateManager | undefined
   registerWorktreeSession: (sessionId: string, directory: string) => void
   registerSession: (session: Session) => void
@@ -72,6 +74,23 @@ export async function transferState(
   return { ok: true, value: undefined }
 }
 
+async function rollback(
+  ctx: ContinueContext,
+  prepared: { worktreeId: string; result: CreateWorktreeResult },
+  error: string,
+  progress: (status: string, detail?: string, error?: string) => void,
+): Promise<void> {
+  await ctx.cleanupWorktree(prepared.worktreeId).catch((err) => {
+    ctx.log("Failed to clean up worktree after continue error:", getErrorMessage(err))
+  })
+  try {
+    ctx.notifyError(error, prepared.result, prepared.worktreeId)
+  } catch (err) {
+    ctx.log("Failed to notify Agent Manager about continue error:", getErrorMessage(err))
+  }
+  progress("error", undefined, error)
+}
+
 /** Fork the session into the worktree directory. */
 export async function forkSession(ctx: ContinueContext, sessionId: string, dir: string): Promise<StepResult<Session>> {
   let client: KiloClient
@@ -136,7 +155,7 @@ export async function continueInWorktree(
 
   progress("transferring", "Transferring changes...")
   const transferred = await transferState(ctx, captured.value, prepared.value.result.path)
-  if (!transferred.ok) return progress("error", undefined, transferred.error)
+  if (!transferred.ok) return rollback(ctx, prepared.value, transferred.error, progress)
 
   progress("forking", "Starting session...")
   const forked = await forkSession(ctx, sessionId, prepared.value.result.path)

packages/kilo-vscode/tests/unit/continue-in-worktree.test.ts

@@ -1,18 +1,49 @@
-import { describe, expect, it, mock } from "bun:test"
+import { afterEach, describe, expect, it, mock } from "bun:test"
+import * as fs from "node:fs/promises"
+import * as os from "node:os"
+import * as path from "node:path"
+import simpleGit from "simple-git"
 import {
   abortSession,
-  captureState,
+  continueInWorktree,
   forkSession,
   registerSession,
   type ContinueContext,
-  type StepResult,
 } from "../../src/agent-manager/continue-in-worktree"
-import type { CreateWorktreeResult } from "../../src/agent-manager/WorktreeManager"
+import { WorktreeManager, type CreateWorktreeResult } from "../../src/agent-manager/WorktreeManager"
 import { forkText } from "../../src/agent-manager/fork-handoff"
 import type { Session } from "@kilocode/sdk/v2/client"
 
 const noop = () => {}
 const log = noop as (...args: unknown[]) => void
+const dirs: string[] = []
+
+afterEach(async () => {
+  await Promise.all(dirs.splice(0, dirs.length).map((dir) => fs.rm(dir, { recursive: true, force: true })))
+})
+
+async function repo(): Promise<string> {
+  const dir = await fs.mkdtemp(path.join(os.tmpdir(), "continue-worktree-"))
+  dirs.push(dir)
+  const git = simpleGit(dir)
+  await git.init(["--initial-branch=main"])
+  await git.addConfig("user.email", "test@test.com")
+  await git.addConfig("user.name", "Test")
+  await fs.writeFile(path.join(dir, "state.txt"), "base\n")
+  await git.add("state.txt")
+  await git.commit("initial")
+  return dir
+}
+
+function client(fork = mock(async () => ({ data: session("forked") }))) {
+  return {
+    session: {
+      abort: mock(async () => ({})),
+      fork,
+      promptAsync: mock(async () => ({})),
+    },
+  } as never
+}
 
 function session(id: string): Session {
   return {
@@ -36,6 +67,8 @@ function ctx(overrides: Partial<ContinueContext> = {}): ContinueContext {
     },
     createWorktreeOnDisk: async () => null,
     runSetupScript: async () => {},
+    cleanupWorktree: async () => {},
+    notifyError: noop,
     getStateManager: () => undefined,
     registerWorktreeSession: noop,
     registerSession: noop,
@@ -154,3 +187,47 @@ describe("continue-in-worktree steps", () => {
     })
   })
 })
+
+describe("continueInWorktree", () => {
+  it("rolls back the created worktree when Git transfer fails", async () => {
+    const root = await repo()
+    const git = simpleGit(root)
+    await fs.writeFile(path.join(root, "state.txt"), "local dirty\n")
+
+    const manager = new WorktreeManager(root, noop)
+    const setup = mock(async () => {})
+    const cleanup = mock(async () => {
+      if (created) await manager.removeWorktree(created.path, created.branch)
+    })
+    const notify = mock((_error: string, _result: CreateWorktreeResult, _worktreeId: string) => {})
+    const progress: Array<{ status: string; error?: string }> = []
+    let created: CreateWorktreeResult | undefined
+    const c = ctx({
+      root,
+      getClient: () => client(),
+      createWorktreeOnDisk: async (opts) => {
+        const value = await manager.createWorktree(opts)
+        created = value
+        const target = simpleGit(value.path)
+        await fs.writeFile(path.join(value.path, "state.txt"), "conflict\n")
+        await target.add("state.txt")
+        await target.commit("conflict")
+        return { worktree: { id: "wt-1" }, result: value }
+      },
+      runSetupScript: setup,
+      cleanupWorktree: cleanup,
+      notifyError: notify,
+    })
+
+    await continueInWorktree(c, "source", (status, _detail, error) => progress.push({ status, error }))
+
+    expect(progress.at(-1)?.status).toBe("error")
+    expect(progress.at(-1)?.error).toContain("Unstaged patch failed")
+    expect(setup).toHaveBeenCalledTimes(1)
+    expect(cleanup).toHaveBeenCalledTimes(1)
+    expect(notify).toHaveBeenCalledWith(expect.stringContaining("Unstaged patch failed"), created!, "wt-1")
+    expect(created).toBeDefined()
+    await expect(fs.stat(created!.path)).rejects.toThrow()
+    expect((await git.branchLocal()).all).not.toContain(created!.branch)
+  })
+})