Skip to content
Source
Sign inSign up

CheckpointRollbackService.test.ts

781 lines29.6 KB

1import * as NodeServices from "@effect/platform-node/NodeServices";
2import { assert, it, vi } from "@effect/vitest";
3import {
4 CheckpointId,
5 CheckpointRef,
6 EventId,
7 MessageId,
8 NodeId,
9 type OrchestrationV2DomainEvent,
10 ProviderDriverKind,
11 ProviderTurnId,
12 RunAttemptId,
13 RunId,
14 CheckpointScopeId,
15 ProjectId,
16 ProviderInstanceId,
17 ProviderSessionId,
18 ProviderThreadId,
19 ThreadId,
20} from "@t3tools/contracts";
21import * as Effect from "effect/Effect";
22import * as DateTime from "effect/DateTime";
23import * as SqlClient from "effect/sql/SqlClient";
24import * as Deferred from "effect/Deferred";
25import * as Fiber from "effect/Fiber";
26import * as FileSystem from "effect/FileSystem";
27import * as Path from "effect/Path";
28import * as Option from "effect/Option";
29import * as Layer from "effect/Layer";
30​
31import * as SqlitePersistence from "../persistence/Sqlite.ts";
32import * as ServerConfig from "../config.ts";
33import * as ServerSettings from "../serverSettings.ts";
34import * as GitWorkflow from "../git/GitWorkflowService.ts";
35import { withWorkspaceLease } from "../workspace/workspaceLease.ts";
36import { ProviderWorkspaceMissingError } from "../provider/Errors.ts";
37import { resolveCodexRollbackTurnCount } from "./Adapters/CodexAdapterV2.ts";
38import * as CheckpointService from "./CheckpointService.ts";
39import * as CheckpointRollbackService from "./CheckpointRollbackService.ts";
40import * as ThreadCommandExecutor from "./ThreadCommandExecutor.ts";
41import * as EventSink from "./EventSink.ts";
42import * as IdAllocator from "@t3tools/provider-core/server/IdAllocator";
43import * as ProjectionStore from "./ProjectionStore.ts";
44import * as ProjectStore from "./ProjectStore.ts";
45import type * as ProviderAdapter from "@t3tools/provider-core/server/ProviderAdapter";
46import * as ProviderSessionManager from "./ProviderSessionManager.ts";
47import * as RuntimePolicy from "./RuntimePolicy.ts";
48​
49const unrelatedProject = Option.some({
50 workspaceRoot: "/nonexistent/t3-rollback-project",
51} as never);
52const layerCheckpointRollbackService = CheckpointRollbackService.layer.pipe(
53 Layer.provide(
54 Layer.mergeAll(
55 NodeServices.layer,
56 Layer.mock(ServerConfig.ServerConfig)({ worktreesDir: "/synthetic-worktrees" } as never),
57 ServerSettings.layerTest(),
58 Layer.mock(GitWorkflow.GitWorkflowService)({}),
59 Layer.mock(ProjectStore.ProjectStoreV2)({ get: () => Effect.succeed(unrelatedProject) }),
60 ),
61 ),
62);
63​
64it.effect.each(["initial", "waiting"])(
65 "rejects a non-ready checkpoint from %s before opening a session or restoring files",
66 (state) => {
67 const threadId = ThreadId.make("thread:rollback-non-ready");
68 const providerThreadId = ProviderThreadId.make("provider-thread:rollback-non-ready");
69 const providerSessionId = ProviderSessionId.make("provider-session:rollback-non-ready");
70 const checkpointId = CheckpointId.make("checkpoint:rollback-non-ready");
71 const scopeId = CheckpointScopeId.make("checkpoint-scope:rollback-non-ready");
72 const providerInstanceId = ProviderInstanceId.make("provider_rollback_non_ready");
73 const restore = vi.fn(() => Effect.die("checkpoint restore must not run"));
74 const open = vi.fn(() => Effect.die("provider session open must not run"));
75 const resolveRuntimePolicy = vi.fn(() => Effect.die("runtime policy resolution must not run"));
76 const readEntered = Deferred.makeUnsafe<void>();
77 const release = Deferred.makeUnsafe<void>();
78 let projection = {
79 thread: {
80 worktreePath: process.cwd(),
81 activeProviderThreadId: providerThreadId,
82 modelSelection: { instanceId: providerInstanceId, model: "test-model" },
83 },
84 providerThread: { id: providerThreadId, providerSessionId, providerInstanceId },
85 checkpoint: {
86 id: checkpointId,
87 scopeId,
88 status: state === "waiting" ? "ready" : "stale",
89 appRunOrdinal: 0,
90 },
91 scope: { id: scopeId, cwd: process.cwd() },
92 } as unknown as ProjectionStore.ProjectionCheckpointRollbackContext;
93 const getThread = vi.fn<ProjectionStore.ProjectionStoreV2Shape["getThread"]>(() =>
94 Deferred.succeed(readEntered, undefined).pipe(Effect.as(projection.thread)),
95 );
96 const getCheckpointRollbackContext = vi.fn(() => Effect.succeed(projection));
97 const testLayer = layerCheckpointRollbackService.pipe(
98 Layer.provide(
99 Layer.mergeAll(
100 ThreadCommandExecutor.layer,
101 Layer.mock(CheckpointService.CheckpointServiceV2)({ restore }),
102 Layer.mock(EventSink.EventSinkV2)({}),
103 IdAllocator.layer,
104 Layer.mock(ProjectionStore.ProjectionStoreV2)({
105 getThread,
106 getCheckpointRollbackContext,
107 getShellSnapshot: () =>
108 Effect.succeed({
109 schemaVersion: 1,
110 snapshotSequence: 0,
111 threads: [],
112 archivedThreads: [],
113 }),
114 }),
115 Layer.mock(ProviderSessionManager.ProviderSessionManagerV2)({ open }),
116 Layer.mock(RuntimePolicy.RuntimePolicyV2)({ resolve: resolveRuntimePolicy }),
117 ),
118 ),
119 );
120​
121 return Effect.gen(function* () {
122 const service = yield* CheckpointRollbackService.CheckpointRollbackServiceV2;
123 const holder =
124 state === "waiting"
125 ? yield* withWorkspaceLease(process.cwd(), Deferred.await(release)).pipe(
126 Effect.forkChild({ startImmediately: true }),
127 )
128 : undefined;
129 const rollback = yield* service
130 .execute({
131 threadId,
132 providerThreadId,
133 checkpointId,
134 scopeId,
135 })
136 .pipe(Effect.flip, Effect.forkChild({ startImmediately: true }));
137 if (state === "waiting") {
138 yield* Deferred.await(readEntered);
139 projection = {
140 ...projection,
141 checkpoint: { ...projection.checkpoint!, status: "stale" },
142 };
143 yield* Deferred.succeed(release, undefined);
144 if (holder !== undefined) yield* Fiber.join(holder);
145 }
146 const error = yield* Fiber.join(rollback);
147​
148 assert.deepEqual(getThread.mock.calls, [[threadId]]);
149 assert.deepEqual(getCheckpointRollbackContext.mock.calls, [
150 [
151 threadId,
152 {
153 threadId,
154 providerThreadId,
155 checkpointId,
156 scopeId,
157 },
158 ],
159 ]);
160 assert.equal(error.reason, "rollback-target-invalid");
161 assert.equal(
162 error.message,
163 `Rollback target ${checkpointId} for provider thread ${providerThreadId} on thread ${threadId} is incomplete or invalid.`,
164 );
165 assert.equal(error.cause, undefined);
166 assert.equal(resolveRuntimePolicy.mock.calls.length, 0);
167 assert.equal(open.mock.calls.length, 0);
168 assert.equal(restore.mock.calls.length, 0);
169 }).pipe(Effect.provide(testLayer));
170 },
171);
172​
173it.effect.each([
174 { order: "cleanup-first", restoreFiles: true },
175 { order: "rollback-first", restoreFiles: true },
176 { order: "cleanup-first", restoreFiles: false, unmanaged: true },
177 { order: "cleanup-first", restoreFiles: true, leavesFolder: true },
178 { order: "cleanup-first", restoreFiles: false, leavesFolder: true },
179 { order: "cleanup-first", restoreFiles: true, leavesFolder: true, shared: true },
180 { order: "cleanup-first", restoreFiles: true, leavesFolder: true, createFails: true },
181])("locks rollback with %s", (row) =>
182 Effect.gen(function* () {
183 const { order, restoreFiles, unmanaged, leavesFolder, shared, createFails } = row;
184 const fileSystem = yield* FileSystem.FileSystem;
185 const path = yield* Path.Path;
186 const workspaceRoot = yield* fileSystem.realPath(
187 yield* fileSystem.makeTempDirectoryScoped({ prefix: "t3-rollback-" }),
188 );
189 const worktreesDir = path.join(workspaceRoot, "worktrees");
190 const worktreePath = unmanaged
191 ? path.join(workspaceRoot, "elsewhere", "feature")
192 : path.join(worktreesDir, "repo", "feature");
193 yield* fileSystem.makeDirectory(path.join(worktreePath, ".git"), { recursive: true });
194 const threadId = ThreadId.make("rollback-worktree");
195 const projectId = ProjectId.make("rollback-worktree-project");
196 const providerThreadId = ProviderThreadId.make("rollback-worktree-provider-thread");
197 const providerSessionId = ProviderSessionId.make("rollback-worktree-provider-session");
198 const providerInstanceId = ProviderInstanceId.make("rollback-worktree-provider-instance");
199 const checkpointId = CheckpointId.make("rollback-worktree-checkpoint");
200 const scopeId = CheckpointScopeId.make("rollback-worktree-scope");
201 const branch = "feature/rollback";
202 const release = Deferred.makeUnsafe<void>();
203 const openEntered = Deferred.makeUnsafe<void>();
204 const calls: string[] = [];
205 let cleanupEntered = false;
206 const providerThread = { id: providerThreadId, providerSessionId, providerInstanceId };
207 const projection = {
208 thread: {
209 id: threadId,
210 projectId,
211 worktreePath,
212 branch,
213 activeProviderThreadId: providerThreadId,
214 modelSelection: { instanceId: providerInstanceId, model: "test-model" },
215 },
216 providerThread,
217 providerSession: undefined,
218 checkpoint: { id: checkpointId, scopeId, status: "ready", appRunOrdinal: 0 },
219 scope: { id: scopeId, cwd: worktreePath },
220 runsToRollback: [],
221 targetAttempt: undefined,
222 rolledBackAttemptIds: [],
223 staleCheckpoints: [],
224 rootNodes: [],
225 providerTurns: [],
226 } as unknown as ProjectionStore.ProjectionCheckpointRollbackContext;
227 const testLayer = CheckpointRollbackService.layer.pipe(
228 Layer.provide(
229 Layer.mergeAll(
230 ThreadCommandExecutor.layer,
231 NodeServices.layer,
232 Layer.mock(ServerConfig.ServerConfig)({ worktreesDir } as never),
233 ServerSettings.layerTest(),
234 Layer.mock(ProjectStore.ProjectStoreV2)({
235 get: (id) => {
236 assert.equal(id, projectId);
237 return Effect.succeedSome({ workspaceRoot } as never);
238 },
239 }),
240 Layer.mock(GitWorkflow.GitWorkflowService)({
241 pruneWorktrees: (input) =>
242 Effect.sync(() => {
243 assert.deepEqual(input, { cwd: workspaceRoot });
244 calls.push("prune");
245 }),
246 createWorktree: (input) =>
247 Effect.gen(function* () {
248 assert.deepEqual(input, {
249 cwd: workspaceRoot,
250 refName: branch,
251 path: worktreePath,
252 });
253 calls.push("create");
254 if (createFails) return yield* Effect.die("the branch is gone");
255 yield* fileSystem
256 .makeDirectory(path.join(worktreePath, ".git"), { recursive: true })
257 .pipe(Effect.orDie);
258 return {} as never;
259 }),
260 }),
261 Layer.mock(CheckpointService.CheckpointServiceV2)({
262 restore: () =>
263 Effect.gen(function* () {
264 assert.isTrue(
265 yield* fileSystem.exists(path.join(worktreePath, ".git")).pipe(Effect.orDie),
266 );
267 calls.push("files");
268 }),
269 }),
270 Layer.mock(EventSink.EventSinkV2)({
271 write: () =>
272 Effect.sync(() => {
273 calls.push("projection");
274 return [];
275 }),
276 }),
277 IdAllocator.layer,
278 Layer.mock(ProjectionStore.ProjectionStoreV2)({
279 getThread: () => Effect.succeed(projection.thread),
280 getCheckpointRollbackContext: () => Effect.succeed(projection),
281 getShellSnapshot: () =>
282 Effect.succeed({
283 schemaVersion: 1,
284 snapshotSequence: 0,
285 threads: shared
286 ? [
287 {
288 id: ThreadId.make("rollback-worktree-sharer"),
289 projectId,
290 worktreePath,
291 deletedAt: null,
292 lineage: {},
293 },
294 ]
295 : [],
296 archivedThreads: [],
297 } as never),
298 getCheckpointContext: () => Effect.succeed({ checkpointScopes: [] } as never),
299 getThreadProviderContext: () => Effect.succeed({ providerSessions: [] } as never),
300 }),
301 Layer.mock(ProviderSessionManager.ProviderSessionManagerV2)({
302 open: (input) =>
303 Effect.gen(function* () {
304 assert.equal(input.runtimePolicy.cwd, worktreePath);
305 if (!(yield* fileSystem.exists(worktreePath).pipe(Effect.orDie))) {
306 return yield* new ProviderWorkspaceMissingError({
307 threadId,
308 cwd: worktreePath,
309 });
310 }
311 calls.push("session");
312 yield* Deferred.succeed(openEntered, undefined);
313 if (order === "rollback-first") yield* Deferred.await(release);
314 return {} as never;
315 }),
316 }),
317 Layer.mock(RuntimePolicy.RuntimePolicyV2)({
318 resolve: ({ thread }) => {
319 assert.equal(thread.worktreePath, worktreePath);
320 return Effect.succeed({ cwd: thread.worktreePath } as never);
321 },
322 }),
323 ),
324 ),
325 );
326 yield* Effect.gen(function* () {
327 const service = yield* CheckpointRollbackService.CheckpointRollbackServiceV2;
328 const cleanupEffect = withWorkspaceLease(
329 worktreePath,
330 Effect.gen(function* () {
331 cleanupEntered = true;
332 if (order === "cleanup-first") yield* Deferred.await(release);
333 yield* fileSystem.remove(worktreePath, { recursive: true, force: true });
334 // A failed `git worktree remove` can leave the folder without its checkout.
335 if (leavesFolder) {
336 yield* fileSystem.makeDirectory(worktreePath);
337 yield* fileSystem.writeFileString(path.join(worktreePath, "unsaved.txt"), "keep");
338 }
339 calls.push("cleanup");
340 }),
341 );
342 const cleanup =
343 order === "cleanup-first"
344 ? yield* cleanupEffect.pipe(Effect.forkChild({ startImmediately: true }))
345 : undefined;
346 const rollback = yield* service
347 .execute({ threadId, providerThreadId, checkpointId, scopeId, restoreFiles })
348 .pipe(Effect.forkChild({ startImmediately: true }));
349 if (order === "cleanup-first") assert.deepEqual(calls, []);
350 if (order === "rollback-first") yield* Deferred.await(openEntered);
351 const laterCleanup =
352 order === "rollback-first"
353 ? yield* cleanupEffect.pipe(Effect.forkChild({ startImmediately: true }))
354 : undefined;
355 if (order === "rollback-first") assert.isFalse(cleanupEntered);
356 yield* Deferred.succeed(release, undefined);
357 if (cleanup !== undefined) yield* Fiber.join(cleanup);
358 if (unmanaged || shared) {
359 const error = yield* Fiber.join(rollback).pipe(Effect.flip);
360 assert.equal(error._tag, "CheckpointRollbackExecutionError");
361 assert.deepEqual(calls, ["cleanup"]);
362 if (shared) {
363 assert.equal(error.reason, "shared-workspace");
364 assert.isTrue(yield* fileSystem.exists(path.join(worktreePath, "unsaved.txt")));
365 }
366 return;
367 }
368 if (createFails) {
369 yield* Fiber.join(rollback).pipe(Effect.exit);
370 assert.deepEqual(calls, ["cleanup", "prune", "create"]);
371 assert.equal(
372 yield* fileSystem.readFileString(path.join(worktreePath, "unsaved.txt")),
373 "keep",
374 );
375 assert.deepEqual(yield* fileSystem.readDirectory(path.dirname(worktreePath)), ["feature"]);
376 return;
377 }
378 yield* Fiber.join(rollback);
379 if (laterCleanup !== undefined) yield* Fiber.join(laterCleanup);
380 const rollbackCalls = restoreFiles
381 ? ["session", "files", "projection"]
382 : ["session", "projection"];
383 const recreates = !leavesFolder || restoreFiles;
384 assert.deepEqual(
385 calls,
386 order === "cleanup-first"
387 ? ["cleanup", ...(recreates ? ["prune", "create"] : []), ...rollbackCalls]
388 : [...rollbackCalls, "cleanup"],
389 );
390 if (!recreates) {
391 assert.isTrue(yield* fileSystem.exists(path.join(worktreePath, "unsaved.txt")));
392 }
393 }).pipe(Effect.provide(testLayer));
394 }).pipe(Effect.scoped, Effect.provide(NodeServices.layer)),
395);
396​
397const seedRollbackHistory = Effect.gen(function* () {
398 const store = yield* ProjectionStore.ProjectionStoreV2;
399 const now = yield* DateTime.now;
400 const threadId = ThreadId.make("rollback-history");
401 const providerThreadId = ProviderThreadId.make("rollback-history-provider");
402 const providerSessionId = ProviderSessionId.make("rollback-history-session");
403 const providerInstanceId = ProviderInstanceId.make("codex");
404 const driver = ProviderDriverKind.make("codex");
405 const modelSelection = { instanceId: providerInstanceId, model: "test" };
406 const scopeId = CheckpointScopeId.make("rollback-history-root");
407 const otherScopeId = CheckpointScopeId.make("rollback-history-child");
408 const runId = (ordinal: number) => RunId.make(`rollback-history-run-${ordinal}`);
409 const nodeId = (ordinal: number) => NodeId.make(`rollback-history-node-${ordinal}`);
410 const attemptId = (ordinal: number) => RunAttemptId.make(`rollback-history-attempt-${ordinal}`);
411 const turnId = (ordinal: number) => ProviderTurnId.make(`rollback-history-turn-${ordinal}`);
412 const checkpointId = (ordinal: number) =>
413 CheckpointId.make(`rollback-history-checkpoint-${ordinal}`);
414 type EventInput<Event = OrchestrationV2DomainEvent> = Event extends OrchestrationV2DomainEvent
415 ? Omit<Event, "id" | "occurredAt" | "threadId">
416 : never;
417 let sequence = 0;
418 const apply = (event: EventInput) =>
419 store.apply({
420 ...event,
421 id: EventId.make(`rollback-history-event-${sequence++}`),
422 threadId,
423 occurredAt: now,
424 } as OrchestrationV2DomainEvent);
425 yield* apply({
426 type: "thread.created",
427 payload: {
428 id: threadId,
429 createdBy: "user",
430 creationSource: "web",
431 projectId: ProjectId.make("rollback-history-project"),
432 title: "Rollback history",
433 providerInstanceId,
434 modelSelection,
435 runtimeMode: "full-access",
436 interactionMode: "default",
437 branch: null,
438 worktreePath: process.cwd(),
439 activeProviderThreadId: providerThreadId,
440 lineage: { parentThreadId: null, relationshipToParent: null, rootThreadId: threadId },
441 forkedFrom: null,
442 createdAt: now,
443 updatedAt: now,
444 archivedAt: null,
445 settledOverride: null,
446 settledAt: null,
447 lastVisitedAt: null,
448 deletedAt: null,
449 },
450 });
451 const providerThread = {
452 id: providerThreadId,
453 driver,
454 providerInstanceId,
455 providerSessionId,
456 appThreadId: threadId,
457 ownerNodeId: null,
458 nativeThreadRef: null,
459 nativeConversationHeadRef: null,
460 status: "idle" as const,
461 firstRunOrdinal: 1,
462 lastRunOrdinal: 6,
463 handoffIds: [],
464 forkedFrom: null,
465 createdAt: now,
466 updatedAt: now,
467 };
468 yield* apply({ type: "provider-thread.updated", payload: providerThread });
469 for (const id of [scopeId, otherScopeId]) {
470 yield* apply({
471 type: "checkpoint-scope.created",
472 payload: {
473 id,
474 threadId,
475 runId: runId(1),
476 nodeId: nodeId(1),
477 parentScopeId: id === scopeId ? null : scopeId,
478 providerThreadId,
479 kind: id === scopeId ? "root_run" : "subagent",
480 ordinalWithinParent: 0,
481 advancesAppRunCount: id === scopeId,
482 cwd: process.cwd(),
483 createdAt: now,
484 },
485 });
486 }
487 const statuses = [
488 "completed",
489 "completed",
490 "completed",
491 "interrupted",
492 "failed",
493 "cancelled",
494 "rolled_back",
495 ] as const;
496 for (const [index, status] of statuses.entries()) {
497 const ordinal = index + 1;
498 yield* apply({
499 type: "run.created",
500 payload: {
501 id: runId(ordinal),
502 threadId,
503 ordinal,
504 providerInstanceId,
505 modelSelection,
506 providerThreadId,
507 userMessageId: MessageId.make(`rollback-history-message-${ordinal}`),
508 rootNodeId: nodeId(ordinal),
509 activeAttemptId: attemptId(ordinal),
510 status,
511 requestedAt: now,
512 startedAt: now,
513 completedAt: now,
514 checkpointId: checkpointId(ordinal),
515 contextHandoffId: null,
516 },
517 });
518 yield* apply({
519 type: "node.updated",
520 payload: {
521 id: nodeId(ordinal),
522 threadId,
523 runId: runId(ordinal),
524 parentNodeId: null,
525 rootNodeId: nodeId(ordinal),
526 kind: "root_turn",
527 status,
528 countsForRun: true,
529 providerThreadId,
530 providerTurnId: turnId(ordinal),
531 nativeItemRef: null,
532 runtimeRequestId: null,
533 checkpointScopeId: scopeId,
534 startedAt: now,
535 completedAt: now,
536 },
537 });
538 // Rollback must exclude every attempt of discarded runs, including an inactive retry.
539 for (const attemptOrdinal of ordinal === 7 ? [6.5, 7] : [ordinal]) {
540 yield* apply({
541 type: "run-attempt.created",
542 runId: runId(ordinal),
543 payload: {
544 id: attemptId(attemptOrdinal),
545 runId: runId(ordinal),
546 attemptOrdinal: attemptOrdinal === ordinal ? 2 : 1,
547 rootNodeId: nodeId(ordinal),
548 providerInstanceId,
549 providerThreadId,
550 providerTurnId: turnId(attemptOrdinal),
551 reason: "initial",
552 status: "completed",
553 startedAt: now,
554 completedAt: now,
555 },
556 });
557 yield* apply({
558 type: "provider-turn.updated",
559 payload: {
560 id: turnId(attemptOrdinal),
561 providerThreadId,
562 nodeId: nodeId(ordinal),
563 runAttemptId: attemptId(attemptOrdinal),
564 nativeTurnRef: { driver, nativeId: `native-turn-${attemptOrdinal}`, strength: "strong" },
565 ordinal: attemptOrdinal * 2,
566 status: status === "rolled_back" ? "completed" : status,
567 startedAt: now,
568 completedAt: now,
569 },
570 });
571 }
572 }
573 // A continuation shares the target attempt; its first provider turn is no longer the boundary.
574 yield* apply({
575 type: "provider-turn.updated",
576 payload: {
577 id: turnId(2.5),
578 providerThreadId,
579 nodeId: nodeId(2),
580 runAttemptId: attemptId(2),
581 nativeTurnRef: { driver, nativeId: "native-turn-2-continuation", strength: "strong" },
582 ordinal: 5,
583 status: "completed",
584 startedAt: now,
585 completedAt: now,
586 },
587 });
588 for (let ordinal = 0; ordinal <= 8; ordinal++) {
589 yield* apply({
590 type: "checkpoint.captured",
591 payload: {
592 id: checkpointId(ordinal),
593 threadId,
594 scopeId: ordinal === 8 ? otherScopeId : scopeId,
595 runId: ordinal === 0 ? null : runId(Math.min(ordinal, 7)),
596 nodeId: nodeId(Math.max(1, Math.min(ordinal, 7))),
597 parentCheckpointId: null,
598 ordinalWithinScope: ordinal,
599 appRunOrdinal: ordinal === 0 ? null : ordinal,
600 ref: CheckpointRef.make(`refs/t3/rollback-history/${ordinal}`),
601 status: ordinal === 7 ? "stale" : "ready",
602 files: [],
603 capturedAt: now,
604 },
605 });
606 }
607​
608 const rollbackCounts: number[] = [];
609 const restored: CheckpointId[] = [];
610 const deletedRefs: CheckpointId[] = [];
611 const committed: OrchestrationV2DomainEvent[] = [];
612 const serviceLayer = layerCheckpointRollbackService.pipe(
613 Layer.provide(
614 Layer.mergeAll(
615 ThreadCommandExecutor.layer,
616 IdAllocator.layer,
617 Layer.mock(CheckpointService.CheckpointServiceV2)({
618 restore: ({ checkpoint }) =>
619 Effect.sync(() => {
620 restored.push(checkpoint.id);
621 }),
622 deleteStaleRefs: ({ checkpoints }) =>
623 Effect.sync(() => {
624 deletedRefs.push(...checkpoints.map((checkpoint) => checkpoint.id));
625 }),
626 }),
627 Layer.mock(EventSink.EventSinkV2)({
628 write: ({ events }) =>
629 Effect.gen(function* () {
630 for (const event of events) yield* store.apply(event);
631 committed.push(...events);
632 return [];
633 }).pipe(
634 Effect.mapError(
635 (cause) =>
636 new EventSink.EventSinkWriteError({
637 eventCount: events.length,
638 cause,
639 }),
640 ),
641 ),
642 }),
643 Layer.mock(ProviderSessionManager.ProviderSessionManagerV2)({
644 open: () =>
645 Effect.succeed({
646 rollbackThread: (input: ProviderAdapter.ProviderAdapterV2RollbackThreadInput) =>
647 Effect.gen(function* () {
648 rollbackCounts.push(yield* resolveCodexRollbackTurnCount(input));
649 return { providerThread: input.providerThread };
650 }),
651 } as never),
652 }),
653 Layer.mock(RuntimePolicy.RuntimePolicyV2)({ resolve: () => Effect.succeed({} as never) }),
654 ),
655 ),
656 );
657 const rollback = (ordinal: number) =>
658 Effect.gen(function* () {
659 const service = yield* CheckpointRollbackService.CheckpointRollbackServiceV2;
660 yield* service.execute({
661 threadId,
662 providerThreadId,
663 scopeId,
664 checkpointId: checkpointId(ordinal),
665 });
666 }).pipe(Effect.provide(serviceLayer));
667​
668 return {
669 store,
670 now,
671 threadId,
672 runId,
673 nodeId,
674 checkpointId,
675 rollback,
676 rollbackCounts,
677 restored,
678 deletedRefs,
679 committed,
680 };
681});
682​
683it.layer(ProjectionStore.layer.pipe(Layer.provideMerge(SqlitePersistence.layerMemory)))(
684 "rollback history reads",
685 (it) => {
686 it.effect("rolls back later terminal runs despite obsolete unrelated SQLite payloads", () =>
687 Effect.gen(function* () {
688 const history = yield* seedRollbackHistory;
689 const sql = yield* SqlClient.SqlClient;
690 const { threadId, runId, nodeId, checkpointId } = history;
691 yield* sql`
692 INSERT INTO orchestration_v2_projection_turn_items (
693 turn_item_id, thread_id, run_id, node_id, provider_thread_id, provider_turn_id,
694 parent_item_id, ordinal, type, status, updated_at, payload_json
695 ) VALUES ('rollback-obsolete-transcript', ${threadId}, ${runId(1)}, ${nodeId(1)},
696 NULL, NULL, NULL, 1, 'assistant_message', 'completed', ${DateTime.formatIso(history.now)},
697 '{"obsolete":"transcript"}')`;
698 yield* sql`UPDATE orchestration_v2_projection_runs
699 SET payload_json = '{"obsolete":"run"}' WHERE run_id = ${runId(1)}`;
700 yield* sql`UPDATE orchestration_v2_projection_nodes
701 SET payload_json = '{"obsolete":"node"}' WHERE node_id = ${nodeId(1)}`;
702 yield* sql`UPDATE orchestration_v2_projection_checkpoints
703 SET payload_json = json_set(payload_json, '$.files', 'obsolete file summary')
704 WHERE checkpoint_id IN (${checkpointId(0)}, ${checkpointId(1)}, ${checkpointId(7)}, ${checkpointId(8)})`;
705 assert.equal(
706 (yield* Effect.exit(history.store.getThreadProjection(threadId)))._tag,
707 "Failure",
708 );
709​
710 yield* history.rollback(2);
711​
712 assert.deepEqual(history.rollbackCounts, [4]);
713 assert.deepEqual(history.restored, [checkpointId(2)]);
714 assert.sameMembers(history.deletedRefs, [3, 4, 5, 6].map(checkpointId));
715 const runs = yield* sql<{ ordinal: number; status: string }>`
716 SELECT ordinal, status FROM orchestration_v2_projection_runs
717 WHERE thread_id = ${threadId} ORDER BY ordinal`;
718 assert.deepEqual(
719 runs.map((run) => run.status),
720 [
721 "completed",
722 "completed",
723 "rolled_back",
724 "rolled_back",
725 "rolled_back",
726 "rolled_back",
727 "rolled_back",
728 ],
729 );
730 const nodes = yield* sql<{ status: string }>`
731 SELECT status FROM orchestration_v2_projection_nodes
732 WHERE thread_id = ${threadId} ORDER BY node_id`;
733 assert.deepEqual(
734 nodes.map((node) => node.status),
735 runs.map((run) => run.status),
736 );
737 const checkpoints = yield* sql<{ status: string }>`
738 SELECT status FROM orchestration_v2_projection_checkpoints
739 WHERE thread_id = ${threadId} ORDER BY ordinal_within_scope`;
740 assert.deepEqual(
741 checkpoints.map((checkpoint) => checkpoint.status),
742 ["ready", "ready", "ready", "stale", "stale", "stale", "stale", "stale", "ready"],
743 );
744 }),
745 );
746​
747 it.effect(
748 "rewinds stopped runs repeatedly without counting discarded provider turns again",
749 () =>
750 Effect.gen(function* () {
751 const history = yield* seedRollbackHistory;
752 yield* history.rollback(5);
753 yield* history.rollback(2);
754 yield* history.rollback(0);
755 yield* history.rollback(0);
756​
757 assert.deepEqual(history.rollbackCounts, [1, 3, 3]);
758 assert.sameMembers(history.deletedRefs, [1, 2, 3, 4, 5, 6].map(history.checkpointId));
759 assert.sameMembers(
760 history.committed
761 .filter((event) => event.type === "run.updated")
762 .map((event) => event.payload.id),
763 [1, 2, 3, 4, 5, 6].map(history.runId),
764 );
765 const records = yield* history.store.getThreadRecords(history.threadId, [
766 "runs",
767 "nodes",
768 "checkpoints",
769 ]);
770 assert.isTrue(records.runs.every((run) => run.status === "rolled_back"));
771 assert.isTrue(records.nodes.every((node) => node.status === "rolled_back"));
772 assert.sameMembers(
773 records.checkpoints
774 .filter((checkpoint) => checkpoint.status === "ready")
775 .map((checkpoint) => checkpoint.id),
776 [history.checkpointId(0), history.checkpointId(8)],
777 );
778 }),
779 );
780 },
781);