1import * as NodeServices from "@effect/platform-node/NodeServices";
2import { assert, it } from "@effect/vitest";
3import {
4 CheckpointId,
5 CheckpointScopeId,
6 ProjectId,
7 ProviderInstanceId,
8 ProviderSessionId,
9 ProviderThreadId,
10 ThreadId,
11} from "@t3tools/contracts";
12import * as Effect from "effect/Effect";
13import * as FileSystem from "effect/FileSystem";
14import * as Layer from "effect/Layer";
15import * as Path from "effect/Path";
16import * as ServerConfig from "../config.ts";
17import * as ServerSettings from "../serverSettings.ts";
18import * as GitWorkflowService from "../git/GitWorkflowService.ts";
19import * as CheckpointRollbackService from "./CheckpointRollbackService.ts";
20import { isCheckpointRestoreIsolated } from "./CheckpointRestoreSafety.ts";
21import * as ThreadCommandExecutor from "./ThreadCommandExecutor.ts";
22import { CheckpointServiceV2 } from "./CheckpointService.ts";
23import { EventSinkV2 } from "./EventSink.ts";
24import * as IdAllocator from "@t3tools/provider-core/server/IdAllocator";
25import { ProjectionStoreV2, type ProjectionCheckpointRollbackContext } from "./ProjectionStore.ts";
26import * as ProjectStore from "./ProjectStore.ts";
27import { ProviderSessionManagerV2 } from "./ProviderSessionManager.ts";
28import { RuntimePolicyV2 } from "./RuntimePolicy.ts";
29
30it.effect.each([
31 "nested",
32 "ancestor",
33 "provider",
34 "sibling",
35 "shared-provider",
36 "conversation",
37] as const)("preserves overlapping workspace files, owner=%s", (owner) =>
38 Effect.gen(function* () {
39 const fs = yield* FileSystem.FileSystem;
40 const path = yield* Path.Path;
41 const parent = yield* fs.makeTempDirectoryScoped({ prefix: "t3-v2-restore-" });
42 const cwd = path.join(parent, "worktree");
43 const nested = path.join(cwd, "nested");
44 const sibling = path.join(parent, "worktree2");
45 yield* fs.makeDirectory(nested, { recursive: true });
46 yield* fs.makeDirectory(sibling);
47 const otherFile = path.join(nested, "other-thread.txt");
48 yield* fs.writeFileString(otherFile, "other thread's uncommitted work");
49 const threadId = ThreadId.make("restore-current");
50 const otherId = ThreadId.make("restore-other");
51 const projectId = ProjectId.make("restore-project");
52 const providerThreadId = ProviderThreadId.make("restore-provider");
53 const providerSessionId = ProviderSessionId.make("restore-session");
54 const instanceId = ProviderInstanceId.make("restore-instance");
55 const checkpointId = CheckpointId.make("restore-checkpoint");
56 const scopeId = CheckpointScopeId.make("restore-scope");
57 const calls: string[] = [];
58 const providerThread = {
59 id: providerThreadId,
60 providerSessionId,
61 providerInstanceId: instanceId,
62 };
63 const otherThread = {
64 id: otherId,
65 projectId,
66 deletedAt: null,
67 lineage: { parentThreadId: null, relationshipToParent: null, rootThreadId: otherId },
68 worktreePath:
69 owner === "ancestor" || owner === "conversation"
70 ? parent
71 : owner === "nested"
72 ? nested
73 : sibling,
74 };
75 const projection = {
76 thread: {
77 id: threadId,
78 projectId,
79 worktreePath: cwd,
80 activeProviderThreadId: providerThreadId,
81 modelSelection: { instanceId, model: "test" },
82 },
83 providerThread,
84 providerSession: undefined,
85 providerTurns: [],
86 rootNodes: [],
87 targetAttempt: undefined,
88 rolledBackAttemptIds: [],
89 staleCheckpoints: [],
90 checkpoint: { id: checkpointId, scopeId, status: "ready", appRunOrdinal: null },
91 scope: { id: scopeId, cwd },
92 runsToRollback: [{ id: "later-run", ordinal: 1, status: "completed", rootNodeId: null }],
93 } as unknown as ProjectionCheckpointRollbackContext;
94 const layerTest = CheckpointRollbackService.layer.pipe(
95 Layer.provide(
96 Layer.mergeAll(
97 NodeServices.layer,
98 Layer.mock(ServerConfig.ServerConfig)({ worktreesDir: "/synthetic-worktrees" } as never),
99 ServerSettings.layerTest(),
100 ThreadCommandExecutor.layer,
101 IdAllocator.layer,
102 Layer.mock(GitWorkflowService.GitWorkflowService)({}),
103 Layer.mock(ProjectStore.ProjectStoreV2)({
104 get: () => Effect.succeedSome({ workspaceRoot: parent } as never),
105 }),
106 Layer.mock(ProjectionStoreV2)({
107 getThread: () => Effect.succeed(projection.thread),
108 getCheckpointRollbackContext: () => Effect.succeed(projection),
109 getShellSnapshot: () =>
110 Effect.succeed({
111 schemaVersion: 1,
112 snapshotSequence: 0,
113 threads: [otherThread] as never,
114 archivedThreads: [],
115 }),
116 getCheckpointContext: () =>
117 Effect.succeed({
118 runs: [],
119 checkpoints: [],
120 checkpointScopes: [],
121 }),
122 getThreadProviderContext: () =>
123 Effect.succeed({
124 thread: otherThread as never,
125 providerThreads: [],
126 providerSessions:
127 owner === "provider" || owner === "shared-provider"
128 ? ([
129 {
130 cwd: owner === "shared-provider" ? cwd : nested,
131 status: "running",
132 capabilities: {
133 sessions: {
134 supportsMultipleProviderThreadsPerSession:
135 owner === "shared-provider",
136 },
137 },
138 },
139 ] as never)
140 : [],
141 }),
142 }),
143 Layer.mock(CheckpointServiceV2)({
144 restore: () =>
145 Effect.gen(function* () {
146 calls.push("files");
147 yield* fs.remove(otherFile).pipe(Effect.orDie);
148 }),
149 }),
150 Layer.mock(EventSinkV2)({ write: () => Effect.succeed([]) }),
151 Layer.mock(ProviderSessionManagerV2)({
152 open: () =>
153 Effect.succeed({
154 rollbackThread: () =>
155 Effect.sync(() => {
156 calls.push("provider");
157 return { providerThread };
158 }),
159 } as never),
160 }),
161 Layer.mock(RuntimePolicyV2)({ resolve: () => Effect.succeed({} as never) }),
162 ),
163 ),
164 );
165 const service = yield* CheckpointRollbackService.CheckpointRollbackServiceV2.pipe(
166 Effect.provide(layerTest),
167 );
168 const restoreFiles = owner !== "conversation";
169 const rejected = !["sibling", "shared-provider", "conversation"].includes(owner);
170 if (rejected) {
171 const error = yield* service
172 .execute({ threadId, providerThreadId, checkpointId, scopeId, restoreFiles })
173 .pipe(Effect.flip);
174 assert.equal(error.reason, "shared-workspace");
175 assert.deepEqual(calls, []);
176 assert.equal(yield* fs.readFileString(otherFile), "other thread's uncommitted work");
177 } else {
178 yield* service.execute({ threadId, providerThreadId, checkpointId, scopeId, restoreFiles });
179 assert.deepEqual(calls, restoreFiles ? ["provider", "files"] : ["provider"]);
180 assert.equal(yield* fs.exists(otherFile), !restoreFiles);
181 }
182 }).pipe(Effect.scoped, Effect.provide(NodeServices.layer)),
183);
184
185it.effect("lets rollback recreate a worktree storage cleanup removed", () =>
186 Effect.gen(function* () {
187 const fileSystem = yield* FileSystem.FileSystem;
188 const path = yield* Path.Path;
189 const parent = yield* fileSystem.makeTempDirectoryScoped({ prefix: "t3-v2-restore-removed-" });
190 const worktreePath = path.join(parent, "worktree");
191 const isolated = yield* isCheckpointRestoreIsolated(
192 { id: ThreadId.make("restore-removed"), worktreePath },
193 { cwd: worktreePath },
194 {
195 fileSystem,
196 path,
197 projections: {
198 getShellSnapshot: () => Effect.succeed({ threads: [], archivedThreads: [] }),
199 } as never,
200 projects: {} as never,
201 },
202 );
203 assert.isTrue(isolated);
204 }).pipe(Effect.scoped, Effect.provide(NodeServices.layer)),
205);
206
207it.effect.each([
208 { sideChat: "answering questions", running: false, toolItems: 0, isolated: true },
209 { sideChat: "running", running: true, toolItems: 0, isolated: false },
210 { sideChat: "after running tools", running: false, toolItems: 1, isolated: false },
211] as const)(
212 "a side chat $sideChat shares the parent's worktree: $isolated",
213 ({ running, toolItems, isolated }) =>
214 Effect.gen(function* () {
215 const fileSystem = yield* FileSystem.FileSystem;
216 const path = yield* Path.Path;
217 const worktreePath = yield* fileSystem.makeTempDirectoryScoped({
218 prefix: "t3-v2-restore-side-chat-",
219 });
220 const parentId = ThreadId.make("restore-side-chat-parent");
221 const sideChat = {
222 id: ThreadId.make("restore-side-chat"),
223 projectId: ProjectId.make("restore-project"),
224 deletedAt: null,
225 worktreePath,
226 activeRunId: running ? "side-chat-run" : null,
227 lineage: {
228 parentThreadId: parentId,
229 relationshipToParent: "fork",
230 rootThreadId: parentId,
231 sideChat: true,
232 },
233 };
234 const result = yield* isCheckpointRestoreIsolated(
235 { id: parentId, worktreePath },
236 { cwd: worktreePath },
237 {
238 fileSystem,
239 path,
240 projections: {
241 getShellSnapshot: () => Effect.succeed({ threads: [sideChat], archivedThreads: [] }),
242 getThreadRecords: () =>
243 Effect.succeed({
244 turnItems: Array.from({ length: toolItems }, () => ({ type: "command_execution" })),
245 }),
246 getCheckpointContext: () => Effect.succeed({ checkpointScopes: [] }),
247 getThreadProviderContext: () => Effect.succeed({ providerSessions: [] }),
248 } as never,
249 projects: {} as never,
250 },
251 );
252 assert.equal(result, isolated);
253 }).pipe(Effect.scoped, Effect.provide(NodeServices.layer)),
254);