Skip to content

Commit 937b301

Browse files
rynfarrynfar
committed
feat(prime): add exact checkpoint rollback
1 parent d1227e0 commit 937b301

26 files changed

Lines changed: 1717 additions & 110 deletions

‎apps/server/src/git/GitWorkflowService.test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -218,6 +218,8 @@ describe("GitWorkflowService", () => {
218218
workspaceCwd: "/repo",
219219
sourceRevision: 2,
220220
targetRevision: 1,
221+
sourceTurnId: null,
222+
targetTurnId: null,
221223
sourceCheckpointRef: "refs/source" as never,
222224
sourceCheckpointOid: "a".repeat(40),
223225
targetCheckpointRef: "refs/target" as never,

‎apps/server/src/orchestration/Layers/CheckpointReactor.ts‎

Lines changed: 42 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -247,6 +247,7 @@ export const make = Effect.gen(function* () {
247247
readonly threadId: ThreadId;
248248
readonly cwd: string;
249249
readonly checkpointTurnCount: number;
250+
readonly turnId: TurnId | null;
250251
readonly checkpointRef: ReturnType<typeof checkpointRefForThreadTurn>;
251252
readonly capturedAt: string;
252253
}) {
@@ -273,10 +274,22 @@ export const make = Effect.gen(function* () {
273274
cwd: input.cwd,
274275
checkpointRef: input.checkpointRef,
275276
});
276-
const anchor = yield* providerService.captureConversationAnchor(input.threadId);
277+
const anchor = yield* providerService.captureConversationAnchor({
278+
threadId: input.threadId,
279+
binding: {
280+
kind: "checkpoint",
281+
checkpointTurnCount: input.checkpointTurnCount,
282+
turnId: input.turnId,
283+
checkpointRef: input.checkpointRef,
284+
checkpointOid: checkpoint.oid,
285+
sourceRevision: input.checkpointTurnCount,
286+
},
287+
});
277288
yield* rollbackRepository.value.putCheckpointAnchor({
278289
threadId: input.threadId,
279290
checkpointTurnCount: input.checkpointTurnCount,
291+
turnId: input.turnId,
292+
sourceRevision: input.checkpointTurnCount,
280293
providerInstanceId: session.value.providerInstanceId,
281294
sessionIncarnationId: session.value.sessionIncarnationId,
282295
checkpointRef: input.checkpointRef,
@@ -331,6 +344,7 @@ export const make = Effect.gen(function* () {
331344
threadId: input.threadId,
332345
cwd: input.cwd,
333346
checkpointTurnCount: input.turnCount,
347+
turnId: input.turnId,
334348
checkpointRef: targetCheckpointRef,
335349
capturedAt: input.createdAt,
336350
});
@@ -590,14 +604,24 @@ export const make = Effect.gen(function* () {
590604
cwd: checkpointCwd,
591605
checkpointRef: baselineCheckpointRef,
592606
});
593-
if (baselineExists) {
594-
return;
607+
if (!baselineExists) {
608+
yield* checkpointStore.captureCheckpoint({
609+
cwd: checkpointCwd,
610+
checkpointRef: baselineCheckpointRef,
611+
});
595612
}
596-
597-
yield* checkpointStore.captureCheckpoint({
613+
yield* capturePrivateCheckpointAnchor({
614+
threadId: thread.id,
598615
cwd: checkpointCwd,
616+
checkpointTurnCount: currentTurnCount,
617+
turnId:
618+
thread.checkpoints.find(
619+
(checkpoint) => checkpoint.checkpointTurnCount === currentTurnCount,
620+
)?.turnId ?? null,
599621
checkpointRef: baselineCheckpointRef,
622+
capturedAt: event.createdAt,
600623
});
624+
if (baselineExists) return;
601625
yield* receiptBus.publish({
602626
type: "checkpoint.baseline.captured",
603627
threadId: thread.id,
@@ -749,14 +773,23 @@ export const make = Effect.gen(function* () {
749773
cwd: checkpointCwd,
750774
checkpointRef: baselineCheckpointRef,
751775
});
752-
if (baselineExists) {
753-
return;
776+
if (!baselineExists) {
777+
yield* checkpointStore.captureCheckpoint({
778+
cwd: checkpointCwd,
779+
checkpointRef: baselineCheckpointRef,
780+
});
754781
}
755-
756-
yield* checkpointStore.captureCheckpoint({
782+
yield* capturePrivateCheckpointAnchor({
783+
threadId,
757784
cwd: checkpointCwd,
785+
checkpointTurnCount: currentTurnCount,
786+
turnId:
787+
thread.checkpoints.find((checkpoint) => checkpoint.checkpointTurnCount === currentTurnCount)
788+
?.turnId ?? null,
758789
checkpointRef: baselineCheckpointRef,
790+
capturedAt: event.occurredAt,
759791
});
792+
if (baselineExists) return;
760793
yield* receiptBus.publish({
761794
type: "checkpoint.baseline.captured",
762795
threadId,

‎apps/server/src/orchestration/Layers/RollbackAdmissionAtomic.test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,8 @@ const admission = Layer.succeed(RollbackAdmission, {
5353
workspaceCwd: "/workspace/atomic",
5454
sourceRevision: 2,
5555
targetRevision: 1,
56+
sourceTurnId: null,
57+
targetTurnId: null,
5658
sourceCheckpointRef: checkpointRefForThreadTurn(threadId, 2),
5759
sourceCheckpointOid: "2".repeat(40),
5860
targetCheckpointRef: checkpointRefForThreadTurn(threadId, 1),

‎apps/server/src/orchestration/Layers/RollbackReconciliation.test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,8 @@ const pending = {
4141
workspaceCwd: "/startup/workspace",
4242
sourceRevision: 2,
4343
targetRevision: 1,
44+
sourceTurnId: null,
45+
targetTurnId: null,
4446
sourceCheckpointRef: "refs/t3/checkpoints/source" as never,
4547
sourceCheckpointOid: "a".repeat(40),
4648
targetCheckpointRef: "refs/t3/checkpoints/target" as never,

‎apps/server/src/persistence/Layers/RollbackSagas.test.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@ const makeState = (
3737
workspaceCwd: "/private/workspace/canary",
3838
sourceRevision: 2,
3939
targetRevision: 1,
40+
sourceTurnId: null,
41+
targetTurnId: null,
4042
sourceCheckpointRef: CheckpointRef.make("refs/t3/checkpoints/thread-rollback-a/turn/2"),
4143
sourceCheckpointOid: "a".repeat(40),
4244
targetCheckpointRef: CheckpointRef.make("refs/t3/checkpoints/thread-rollback-a/turn/1"),
@@ -160,6 +162,8 @@ layer("RollbackSagaRepository", (it) => {
160162
yield* repository.putCheckpointAnchor({
161163
threadId: threadA,
162164
checkpointTurnCount,
165+
turnId: null,
166+
sourceRevision: checkpointTurnCount,
163167
providerInstanceId,
164168
sessionIncarnationId,
165169
checkpointRef: CheckpointRef.make(
@@ -175,6 +179,8 @@ layer("RollbackSagaRepository", (it) => {
175179
yield* repository.putCheckpointAnchor({
176180
threadId: threadA,
177181
checkpointTurnCount: 1,
182+
turnId: null,
183+
sourceRevision: 1,
178184
providerInstanceId,
179185
sessionIncarnationId,
180186
checkpointRef: CheckpointRef.make("refs/t3/checkpoints/thread-rollback-a/turn/1"),
@@ -188,6 +194,8 @@ layer("RollbackSagaRepository", (it) => {
188194
.putCheckpointAnchor({
189195
threadId: threadA,
190196
checkpointTurnCount: 1,
197+
turnId: null,
198+
sourceRevision: 1,
191199
providerInstanceId,
192200
sessionIncarnationId,
193201
checkpointRef: CheckpointRef.make("refs/t3/checkpoints/thread-rollback-a/turn/1"),

‎apps/server/src/persistence/Layers/RollbackSagas.ts‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import * as Option from "effect/Option";
55
import * as Schema from "effect/Schema";
66
import * as SqlClient from "effect/unstable/sql/SqlClient";
77
import type { SqlError } from "effect/unstable/sql/SqlError";
8-
import { NonNegativeInt, ProjectId, ThreadId } from "@t3tools/contracts";
8+
import { NonNegativeInt, ProjectId, ThreadId, TurnId } from "@t3tools/contracts";
99
import {
1010
PersistenceDecodeError,
1111
toPersistenceDecodeError,
@@ -41,6 +41,8 @@ const LeaseDbRow = Schema.Struct({
4141
const AnchorDbRow = Schema.Struct({
4242
threadId: ThreadId,
4343
checkpointTurnCount: NonNegativeInt,
44+
turnId: Schema.NullOr(TurnId),
45+
sourceRevision: NonNegativeInt,
4446
providerInstanceId: Schema.String,
4547
sessionIncarnationId: Schema.String,
4648
checkpointRef: Schema.String,
@@ -299,16 +301,19 @@ const make = Effect.gen(function* () {
299301
});
300302
const rows = yield* sql<{ readonly anchorDigest: string }>`
301303
INSERT INTO rollback_checkpoint_anchors (
302-
thread_id, checkpoint_turn_count, provider_instance_id, session_incarnation_id,
303-
checkpoint_ref, checkpoint_oid, anchor_json, anchor_digest, captured_at
304+
thread_id, checkpoint_turn_count, turn_id, source_revision, provider_instance_id,
305+
session_incarnation_id, checkpoint_ref, checkpoint_oid, anchor_json, anchor_digest, captured_at
304306
) VALUES (
305-
${anchor.threadId}, ${anchor.checkpointTurnCount}, ${anchor.providerInstanceId},
307+
${anchor.threadId}, ${anchor.checkpointTurnCount}, ${anchor.turnId}, ${anchor.sourceRevision},
308+
${anchor.providerInstanceId},
306309
${anchor.sessionIncarnationId}, ${anchor.checkpointRef}, ${anchor.checkpointOid},
307310
${anchorJson}, ${anchor.anchorDigest}, ${anchor.capturedAt}
308311
)
309312
ON CONFLICT (thread_id, checkpoint_turn_count, provider_instance_id, session_incarnation_id)
310313
DO UPDATE SET captured_at = rollback_checkpoint_anchors.captured_at
311-
WHERE rollback_checkpoint_anchors.checkpoint_ref = excluded.checkpoint_ref
314+
WHERE rollback_checkpoint_anchors.turn_id IS excluded.turn_id
315+
AND rollback_checkpoint_anchors.source_revision = excluded.source_revision
316+
AND rollback_checkpoint_anchors.checkpoint_ref = excluded.checkpoint_ref
312317
AND rollback_checkpoint_anchors.checkpoint_oid = excluded.checkpoint_oid
313318
AND rollback_checkpoint_anchors.anchor_digest = excluded.anchor_digest
314319
RETURNING anchor_digest AS "anchorDigest"
@@ -324,6 +329,7 @@ const make = Effect.gen(function* () {
324329
const getCheckpointAnchor: RollbackSagaRepositoryShape["getCheckpointAnchor"] = (input) =>
325330
sql`
326331
SELECT thread_id AS "threadId", checkpoint_turn_count AS "checkpointTurnCount",
332+
turn_id AS "turnId", source_revision AS "sourceRevision",
327333
provider_instance_id AS "providerInstanceId", session_incarnation_id AS "sessionIncarnationId",
328334
checkpoint_ref AS "checkpointRef", checkpoint_oid AS "checkpointOid",
329335
anchor_json AS "anchorJson", anchor_digest AS "anchorDigest", captured_at AS "capturedAt"
@@ -353,6 +359,8 @@ const make = Effect.gen(function* () {
353359
decodeAnchor({
354360
threadId: row.threadId,
355361
checkpointTurnCount: row.checkpointTurnCount,
362+
turnId: row.turnId,
363+
sourceRevision: row.sourceRevision,
356364
providerInstanceId: row.providerInstanceId,
357365
sessionIncarnationId: row.sessionIncarnationId,
358366
checkpointRef: row.checkpointRef,

‎apps/server/src/persistence/Migrations/051_DurableRollbackSagas.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,8 @@ export default Effect.gen(function* () {
5656
CREATE TABLE IF NOT EXISTS rollback_checkpoint_anchors (
5757
thread_id TEXT NOT NULL,
5858
checkpoint_turn_count INTEGER NOT NULL,
59+
turn_id TEXT,
60+
source_revision INTEGER NOT NULL,
5961
provider_instance_id TEXT NOT NULL,
6062
session_incarnation_id TEXT NOT NULL,
6163
checkpoint_ref TEXT NOT NULL,

‎apps/server/src/persistence/Services/RollbackSagas.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import {
1010
ProviderInstanceId,
1111
RuntimeSessionId,
1212
ThreadId,
13+
TurnId,
1314
} from "@t3tools/contracts";
1415
import type { PersistenceDecodeError, PersistenceSqlError } from "../Errors.ts";
1516

@@ -43,6 +44,8 @@ export const RollbackSagaState = Schema.Struct({
4344
workspaceCwd: Schema.String,
4445
sourceRevision: NonNegativeInt,
4546
targetRevision: NonNegativeInt,
47+
sourceTurnId: Schema.NullOr(TurnId),
48+
targetTurnId: Schema.NullOr(TurnId),
4649
sourceCheckpointRef: CheckpointRef,
4750
sourceCheckpointOid: Schema.String,
4851
targetCheckpointRef: CheckpointRef,
@@ -87,6 +90,8 @@ export type RollbackSagaRecord = typeof RollbackSagaRecord.Type;
8790
export const RollbackCheckpointAnchor = Schema.Struct({
8891
threadId: ThreadId,
8992
checkpointTurnCount: NonNegativeInt,
93+
turnId: Schema.NullOr(TurnId),
94+
sourceRevision: NonNegativeInt,
9095
providerInstanceId: ProviderInstanceId,
9196
sessionIncarnationId: RuntimeSessionId,
9297
checkpointRef: CheckpointRef,

‎apps/server/src/provider/Layers/ProviderService.test.ts‎

Lines changed: 103 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2720,6 +2720,79 @@ fanout.layer("ProviderServiceLive fanout", (it) => {
27202720
}),
27212721
);
27222722

2723+
it.effect("rejects an exact anchor result after its provider generation retires", () =>
2724+
Effect.gen(function* () {
2725+
const provider = yield* ProviderService.ProviderService;
2726+
const threadId = asThreadId("thread-generation-replaced-absolute-anchor");
2727+
yield* provider.startSession(threadId, {
2728+
provider: CODEX_DRIVER,
2729+
providerInstanceId: codexInstanceId,
2730+
threadId,
2731+
runtimeMode: "full-access",
2732+
});
2733+
2734+
const originalCapabilities = fanout.codex.adapter.capabilities;
2735+
const originalAbsoluteRollback = fanout.codex.adapter.absoluteConversationRollback;
2736+
const originalRuntimeFence = fanout.codex.adapter.runtimeFence;
2737+
const current = yield* Ref.make(true);
2738+
const captureEntered = yield* Deferred.make<void>();
2739+
const releaseCapture = yield* Deferred.make<void>();
2740+
Object.assign(fanout.codex.adapter, {
2741+
capabilities: { ...originalCapabilities, conversationRollback: "absolute" },
2742+
absoluteConversationRollback: {
2743+
isAvailable: () => Effect.succeed(true),
2744+
captureAnchor: () =>
2745+
Deferred.succeed(captureEntered, undefined).pipe(
2746+
Effect.andThen(Deferred.await(releaseCapture)),
2747+
Effect.as({ anchor: { privateLeaf: "stale-private" }, digest: "stale-digest" }),
2748+
),
2749+
inspectAnchor: () => Effect.succeed({ anchor: {}, digest: "unused" }),
2750+
applyAnchor: () => Effect.void,
2751+
releaseAnchor: () => Effect.void,
2752+
},
2753+
runtimeFence: {
2754+
generation: {},
2755+
configRevision: "private-anchor-test-revision",
2756+
isCurrent: Ref.get(current),
2757+
},
2758+
});
2759+
yield* Effect.addFinalizer(() =>
2760+
Effect.sync(() => {
2761+
Object.assign(fanout.codex.adapter, { capabilities: originalCapabilities });
2762+
if (originalAbsoluteRollback === undefined) {
2763+
delete (fanout.codex.adapter as { absoluteConversationRollback?: unknown })
2764+
.absoluteConversationRollback;
2765+
} else {
2766+
Object.assign(fanout.codex.adapter, {
2767+
absoluteConversationRollback: originalAbsoluteRollback,
2768+
});
2769+
}
2770+
if (originalRuntimeFence === undefined) {
2771+
delete (fanout.codex.adapter as { runtimeFence?: unknown }).runtimeFence;
2772+
} else {
2773+
Object.assign(fanout.codex.adapter, { runtimeFence: originalRuntimeFence });
2774+
}
2775+
}),
2776+
);
2777+
2778+
const anchorFiber = yield* provider.captureConversationAnchor!({
2779+
threadId,
2780+
binding: {
2781+
kind: "source",
2782+
sourceRevision: 1,
2783+
checkpointRef: "refs/t3/checkpoints/private/source" as never,
2784+
checkpointOid: "a".repeat(40),
2785+
turnId: null,
2786+
},
2787+
}).pipe(Effect.forkChild);
2788+
yield* Deferred.await(captureEntered);
2789+
yield* Ref.set(current, false);
2790+
yield* Deferred.succeed(releaseCapture, undefined);
2791+
2792+
assert.isTrue(Exit.isFailure(yield* Fiber.await(anchorFiber)));
2793+
}),
2794+
);
2795+
27232796
it.effect("keeps Stop authoritative while an account transition is starting", () =>
27242797
Effect.gen(function* () {
27252798
const provider = yield* ProviderService.ProviderService;
@@ -3883,6 +3956,23 @@ describe("agent browser access", () => {
38833956
let recoveredSession: ProviderSession | undefined;
38843957
const recoveryAdapter: ProviderAdapterShape<ProviderAdapterError> = {
38853958
...codex.adapter,
3959+
capabilities: {
3960+
...codex.adapter.capabilities,
3961+
conversationRollback: "absolute",
3962+
},
3963+
absoluteConversationRollback: {
3964+
isAvailable: () => Effect.succeed(true),
3965+
captureAnchor: () => Effect.succeed({ anchor: {}, digest: "unused" }),
3966+
inspectAnchor: () => Effect.succeed({ anchor: {}, digest: "unused" }),
3967+
applyAnchor: () => Effect.void,
3968+
releaseAnchor: () => Effect.void,
3969+
prepareRecovery: (input) =>
3970+
Effect.sync(() => {
3971+
assert.deepEqual(input.sourceAnchor, { privateLeaf: "restart-source-private" });
3972+
assert.deepEqual(input.desiredAnchor, { privateLeaf: "restart-target-private" });
3973+
order.push("quarantine");
3974+
}),
3975+
},
38863976
recoverSession: (input) =>
38873977
Effect.gen(function* () {
38883978
assert.isDefined(McpProviderSession.readMcpProviderSession(threadId));
@@ -3943,10 +4033,21 @@ describe("agent browser access", () => {
39434033
);
39444034
yield* Effect.yieldNow;
39454035

3946-
yield* provider.recoverRestartSessions!();
4036+
yield* provider.recoverRestartSessions!({
4037+
pendingAbsoluteRollbacks: new Map([
4038+
[
4039+
threadId,
4040+
{
4041+
sourceAnchor: { privateLeaf: "restart-source-private" },
4042+
desiredAnchor: { privateLeaf: "restart-target-private" },
4043+
expectedAnchor: { privateLeaf: "restart-source-private" },
4044+
},
4045+
],
4046+
]),
4047+
});
39474048
yield* Fiber.join(consumer);
39484049

3949-
assert.deepEqual(order, ["mcp", "recover", "activate"]);
4050+
assert.deepEqual(order, ["mcp", "recover", "quarantine", "activate"]);
39504051
const [recoveredEvent] = yield* Ref.get(recoveredEvents);
39514052
assert.equal(recoveredEvent?.eventId, asEventId("evt-restart-adopted-output"));
39524053
assert.equal(recoveredEvent?.turnId, recoveredTurnId);

0 commit comments

Comments
 (0)