diff --git a/apps/server/src/orchestration-v2/CheckpointRollbackService.test.ts b/apps/server/src/orchestration-v2/CheckpointRollbackService.test.ts index c61f9bc175edc..dea2c0151d38f 100644 --- a/apps/server/src/orchestration-v2/CheckpointRollbackService.test.ts +++ b/apps/server/src/orchestration-v2/CheckpointRollbackService.test.ts @@ -174,7 +174,10 @@ it.effect.each([ { order: "cleanup-first", restoreFiles: true }, { order: "rollback-first", restoreFiles: true }, { order: "cleanup-first", restoreFiles: false, unmanaged: true }, -])("locks rollback with %s", ({ order, restoreFiles, unmanaged = false }) => + { order: "cleanup-first", restoreFiles: true, leavesFolder: true }, + { order: "cleanup-first", restoreFiles: false, leavesFolder: true }, + { order: "cleanup-first", restoreFiles: true, leavesFolder: true, shared: true }, +])("locks rollback with %s", ({ order, restoreFiles, unmanaged, leavesFolder, shared }) => Effect.gen(function* () { const fileSystem = yield* FileSystem.FileSystem; const path = yield* Path.Path; @@ -185,7 +188,7 @@ it.effect.each([ const worktreePath = unmanaged ? path.join(workspaceRoot, "elsewhere", "feature") : path.join(worktreesDir, "repo", "feature"); - yield* fileSystem.makeDirectory(worktreePath, { recursive: true }); + yield* fileSystem.makeDirectory(path.join(worktreePath, ".git"), { recursive: true }); const threadId = ThreadId.make("rollback-worktree"); const projectId = ProjectId.make("rollback-worktree-project"); const providerThreadId = ProviderThreadId.make("rollback-worktree-provider-thread"); @@ -246,14 +249,18 @@ it.effect.each([ path: worktreePath, }); calls.push("create"); - yield* fileSystem.makeDirectory(worktreePath).pipe(Effect.orDie); + yield* fileSystem + .makeDirectory(path.join(worktreePath, ".git"), { recursive: true }) + .pipe(Effect.orDie); return {} as never; }), }), Layer.mock(CheckpointService.CheckpointServiceV2)({ restore: () => Effect.gen(function* () { - assert.isTrue(yield* fileSystem.exists(worktreePath).pipe(Effect.orDie)); + assert.isTrue( + yield* fileSystem.exists(path.join(worktreePath, ".git")).pipe(Effect.orDie), + ); calls.push("files"); }), }), @@ -272,16 +279,31 @@ it.effect.each([ Effect.succeed({ schemaVersion: 1, snapshotSequence: 0, - threads: [], + threads: shared + ? [ + { + id: ThreadId.make("rollback-worktree-sharer"), + projectId, + worktreePath, + deletedAt: null, + lineage: {}, + }, + ] + : [], archivedThreads: [], - }), + } as never), + getCheckpointContext: () => Effect.succeed({ checkpointScopes: [] } as never), + getThreadProviderContext: () => Effect.succeed({ providerSessions: [] } as never), }), Layer.mock(ProviderSessionManager.ProviderSessionManagerV2)({ open: (input) => Effect.gen(function* () { assert.equal(input.runtimePolicy.cwd, worktreePath); if (!(yield* fileSystem.exists(worktreePath).pipe(Effect.orDie))) { - return yield* new ProviderWorkspaceMissingError({ threadId, cwd: worktreePath }); + return yield* new ProviderWorkspaceMissingError({ + threadId, + cwd: worktreePath, + }); } calls.push("session"); yield* Deferred.succeed(openEntered, undefined); @@ -306,6 +328,11 @@ it.effect.each([ cleanupEntered = true; if (order === "cleanup-first") yield* Deferred.await(release); yield* fileSystem.remove(worktreePath, { recursive: true, force: true }); + // A failed `git worktree remove` can leave the folder without its checkout. + if (leavesFolder) { + yield* fileSystem.makeDirectory(worktreePath); + yield* fileSystem.writeFileString(path.join(worktreePath, "unsaved.txt"), "keep"); + } calls.push("cleanup"); }), ); @@ -325,10 +352,14 @@ it.effect.each([ if (order === "rollback-first") assert.isFalse(cleanupEntered); yield* Deferred.succeed(release, undefined); if (cleanup !== undefined) yield* Fiber.join(cleanup); - if (unmanaged) { + if (unmanaged || shared) { const error = yield* Fiber.join(rollback).pipe(Effect.flip); assert.equal(error._tag, "CheckpointRollbackExecutionError"); assert.deepEqual(calls, ["cleanup"]); + if (shared) { + assert.equal(error.reason, "shared-workspace"); + assert.isTrue(yield* fileSystem.exists(path.join(worktreePath, "unsaved.txt"))); + } return; } yield* Fiber.join(rollback); @@ -336,12 +367,16 @@ it.effect.each([ const rollbackCalls = restoreFiles ? ["session", "files", "projection"] : ["session", "projection"]; + const recreates = !leavesFolder || restoreFiles; assert.deepEqual( calls, order === "cleanup-first" - ? ["cleanup", "prune", "create", ...rollbackCalls] + ? ["cleanup", ...(recreates ? ["prune", "create"] : []), ...rollbackCalls] : [...rollbackCalls, "cleanup"], ); + if (!recreates) { + assert.isTrue(yield* fileSystem.exists(path.join(worktreePath, "unsaved.txt"))); + } }).pipe(Effect.provide(testLayer)); }).pipe(Effect.scoped, Effect.provide(NodeServices.layer)), ); diff --git a/apps/server/src/orchestration-v2/CheckpointRollbackService.ts b/apps/server/src/orchestration-v2/CheckpointRollbackService.ts index 4d40641fa54db..a3604c278e3dd 100644 --- a/apps/server/src/orchestration-v2/CheckpointRollbackService.ts +++ b/apps/server/src/orchestration-v2/CheckpointRollbackService.ts @@ -217,42 +217,51 @@ export const layer: Layer.Layer< }; }); - const { worktreePath, branch } = context.thread; - // Recreated only where the server would have made it, as turn start does. - if ( - worktreePath != null && - branch != null && - !(yield* fileSystem.exists(worktreePath)) && - (yield* isManagedWorktreePath(worktreePath)) - ) { - const project = yield* projects.get(context.thread.projectId); - if (Option.isSome(project)) { - yield* gitWorkflow.pruneWorktrees({ cwd: project.value.workspaceRoot }); - yield* gitWorkflow.createWorktree({ - cwd: project.value.workspaceRoot, - refName: branch, - path: worktreePath, - }); - } - } - - if ( - input.restoreFiles !== false && - !(yield* isCheckpointRestoreIsolated(context.thread, scope, { - fileSystem, - projections, - projects, - path, - })) - ) { - return yield* new CheckpointRollbackExecutionError({ + const restoreFiles = input.restoreFiles !== false; + const isolated = isCheckpointRestoreIsolated(context.thread, scope, { + fileSystem, + projections, + projects, + path, + }); + const sharedWorkspace = () => + new CheckpointRollbackExecutionError({ reason: "shared-workspace", threadId: input.threadId, providerThreadId: input.providerThreadId, checkpointId: input.checkpointId, }); + + const { worktreePath, branch } = context.thread; + // Recreated only where the server would have made it, as turn start does. + if (worktreePath != null && branch != null) { + // A failed removal can leave the folder without its checkout; only a file restore replaces it. + const leftover = + restoreFiles && + (yield* fileSystem.exists(worktreePath)) && + !(yield* fileSystem.exists(path.join(worktreePath, ".git"))); + if ( + (leftover || !(yield* fileSystem.exists(worktreePath))) && + (yield* isManagedWorktreePath(worktreePath)) + ) { + const project = yield* projects.get(context.thread.projectId); + if (Option.isSome(project)) { + if (leftover) { + if (!(yield* isolated)) return yield* sharedWorkspace(); + yield* fileSystem.remove(worktreePath, { recursive: true, force: true }); + } + yield* gitWorkflow.pruneWorktrees({ cwd: project.value.workspaceRoot }); + yield* gitWorkflow.createWorktree({ + cwd: project.value.workspaceRoot, + refName: branch, + path: worktreePath, + }); + } + } } + if (restoreFiles && !(yield* isolated)) return yield* sharedWorkspace(); + const modelSelection = context.thread.modelSelection; const resolvedRuntimePolicy = yield* runtimePolicy.resolve({ thread: context.thread,