diff --git a/apps/server/src/checkpointing/CheckpointDiffQuery.test.ts b/apps/server/src/checkpointing/CheckpointDiffQuery.test.ts index bdb3e02785d8..1a4d1092972a 100644 --- a/apps/server/src/checkpointing/CheckpointDiffQuery.test.ts +++ b/apps/server/src/checkpointing/CheckpointDiffQuery.test.ts @@ -37,6 +37,7 @@ function makeProjection(): ProjectionCheckpointContext { { scopeId: secondScopeId, runId: secondRunId, + ordinalWithinScope: 1, appRunOrdinal: 2, status: "ready", ref: secondRef, @@ -200,6 +201,90 @@ it.effect("excludes ready checkpoints from rolled-back runs", () => { }).pipe(Effect.provide(layer)); }); +function makeProjectionAfterStoppedRun(stoppedRunCheckpoint: boolean) { + const thirdRunId = RunId.make("run:checkpoint-diff-v2:3"); + const checkpoint = ( + ordinalWithinScope: number, + runId: RunId | null, + ): ProjectionCheckpointContext["checkpoints"][number] => ({ + scopeId: firstScopeId, + runId, + ordinalWithinScope, + appRunOrdinal: runId === null ? null : ordinalWithinScope, + status: "ready", + ref: CheckpointRef.make(`refs/t3/test/ordinal/${ordinalWithinScope}`), + }); + return { + runs: [ + { id: firstRunId, ordinal: 1, status: "completed" }, + { id: secondRunId, ordinal: 2, status: "interrupted" }, + { id: thirdRunId, ordinal: 3, status: "completed" }, + ], + checkpointScopes: [ + { id: firstScopeId, runId: thirdRunId, kind: "root_run", cwd: "/repo" }, + { id: secondScopeId, runId: secondRunId, kind: "tool", cwd: "/repo" }, + ], + checkpoints: [ + { + ...checkpoint(2, null), + scopeId: secondScopeId, + ref: CheckpointRef.make("refs/t3/test/tool/2"), + }, + checkpoint(0, null), + checkpoint(1, firstRunId), + checkpoint(2, stoppedRunCheckpoint ? secondRunId : null), + checkpoint(3, thirdRunId), + ], + } satisfies ProjectionCheckpointContext; +} + +it.effect.each([false, true])( + "diffs the turn after a stopped run from its parent checkpoint (stopped run captured: %s)", + (stoppedRunCheckpoint) => { + const diffCheckpoints = vi.fn((_input: CheckpointStore.DiffCheckpointsInput) => + Effect.succeed("diff --git a/file b/file"), + ); + const layer = layerFor({ + projection: Effect.succeed(makeProjectionAfterStoppedRun(stoppedRunCheckpoint)), + diffCheckpoints, + }); + + return Effect.gen(function* () { + const query = yield* CheckpointDiffQuery.CheckpointDiffQuery; + const result = yield* query.getTurnDiff({ threadId, fromTurnCount: 2, toTurnCount: 3 }); + + assert.strictEqual(result.diff, "diff --git a/file b/file"); + yield* query.getTurnDiff({ threadId, fromTurnCount: 1, toTurnCount: 3 }); + yield* query.getFullThreadDiff({ threadId, toTurnCount: 3 }); + assert.deepEqual( + diffCheckpoints.mock.calls.map(([input]) => ({ + from: input.fromCheckpointRef, + to: input.toCheckpointRef, + })), + [ + { from: "refs/t3/test/ordinal/2", to: "refs/t3/test/ordinal/3" }, + { from: "refs/t3/test/ordinal/1", to: "refs/t3/test/ordinal/3" }, + { + from: yield* checkpointRefForScopeOrdinal({ + scopeId: firstScopeId, + ordinalWithinScope: 0, + }), + to: "refs/t3/test/ordinal/3", + }, + ], + ); + const stoppedTurnError = yield* query + .getTurnDiff({ threadId, fromTurnCount: 1, toTurnCount: 2 }) + .pipe(Effect.flip); + assert.instanceOf(stoppedTurnError, CheckpointRefUnavailableError); + assert.deepEqual( + { checkpoint: stoppedTurnError.checkpoint, turnCount: stoppedTurnError.turnCount }, + { checkpoint: "to", turnCount: 2 }, + ); + }).pipe(Effect.provide(layer)); + }, +); + it.effect("preserves the typed missing-baseline-ref error contract", () => { const projection = makeProjection(); const layer = layerFor({ diff --git a/apps/server/src/checkpointing/CheckpointDiffQuery.ts b/apps/server/src/checkpointing/CheckpointDiffQuery.ts index 08ed4abca0da..3cf37fa38343 100644 --- a/apps/server/src/checkpointing/CheckpointDiffQuery.ts +++ b/apps/server/src/checkpointing/CheckpointDiffQuery.ts @@ -161,6 +161,18 @@ export const make = Effect.gen(function* () { }); } + // Capture compares against the preceding scope ordinal, even when that + // checkpoint belongs to a stopped run or has no run. + const parentCheckpoint = + input.fromTurnCount === input.toTurnCount - 1 + ? projection.checkpoints.find( + (checkpoint) => + checkpoint.status === "ready" && + checkpoint.scopeId === toCheckpoint.scopeId && + checkpoint.ordinalWithinScope === toCheckpoint.ordinalWithinScope - 1, + ) + : undefined; + const firstScope = input.fromTurnCount === 0 ? projection.checkpointScopes.find((scope) => scope.kind === "root_run") @@ -173,8 +185,12 @@ export const make = Effect.gen(function* () { scopeId: firstScope.id, ordinalWithinScope: 0, }).pipe(Effect.provideService(Crypto.Crypto, crypto)) - : readyCheckpoints.find((checkpoint) => checkpoint.appRunOrdinal === input.fromTurnCount) - ?.ref; + : ( + parentCheckpoint ?? + readyCheckpoints.find( + (checkpoint) => checkpoint.appRunOrdinal === input.fromTurnCount, + ) + )?.ref; if (fromCheckpointRef === undefined) { return yield* new CheckpointRefUnavailableError({ operation, diff --git a/apps/server/src/orchestration-v2/ProjectionStore.test.ts b/apps/server/src/orchestration-v2/ProjectionStore.test.ts index 74377160a820..39fce5c078ba 100644 --- a/apps/server/src/orchestration-v2/ProjectionStore.test.ts +++ b/apps/server/src/orchestration-v2/ProjectionStore.test.ts @@ -4060,6 +4060,11 @@ it.layer(layerTest)("ProjectionStoreV2", (it) => { () => Effect.gen(function* () { const projectionStore = yield* ProjectionStore.ProjectionStoreV2; + const memoryStore = yield* ProjectionStore.ProjectionStoreV2.pipe( + Effect.provide(ProjectionStore.layerMemory), + ); + const apply = (event: Parameters[0]) => + Effect.all([projectionStore.apply(event), memoryStore.apply(event)]); const sql = yield* SqlClient.SqlClient; const now = yield* DateTime.now; const threadId = ThreadId.make("thread:checkpoint-context"); @@ -4068,7 +4073,7 @@ it.layer(layerTest)("ProjectionStoreV2", (it) => { const scopeId = CheckpointScopeId.make("scope:checkpoint-context"); const checkpointId = CheckpointId.make("checkpoint:checkpoint-context"); const ref = CheckpointRef.make("refs/t3/checkpoint-context/1"); - yield* projectionStore.apply({ + yield* apply({ id: EventId.make("event:checkpoint-context:thread"), type: "thread.created", threadId, @@ -4097,7 +4102,7 @@ it.layer(layerTest)("ProjectionStoreV2", (it) => { deletedAt: null, }, }); - yield* projectionStore.apply({ + yield* apply({ id: EventId.make("event:checkpoint-context:run"), type: "run.created", threadId, @@ -4120,7 +4125,7 @@ it.layer(layerTest)("ProjectionStoreV2", (it) => { contextHandoffId: null, }, }); - yield* projectionStore.apply({ + yield* apply({ id: EventId.make("event:checkpoint-context:scope"), type: "checkpoint-scope.created", threadId, @@ -4139,7 +4144,7 @@ it.layer(layerTest)("ProjectionStoreV2", (it) => { createdAt: now, }, }); - yield* projectionStore.apply({ + yield* apply({ id: EventId.make("event:checkpoint-context:checkpoint"), type: "checkpoint.captured", threadId, @@ -4178,15 +4183,19 @@ it.layer(layerTest)("ProjectionStoreV2", (it) => { SET payload_json = json_set(payload_json, '$.files', 'obsolete file summary') WHERE checkpoint_id = ${checkpointId} `; - assert.deepEqual(yield* projectionStore.getCheckpointContext(threadId), { - runs: [{ id: runId, ordinal: 1, status: "completed" }], - checkpointScopes: [{ id: scopeId, runId, kind: "root_run", cwd: "/repo/worktree" }], - checkpoints: [{ scopeId, runId, appRunOrdinal: 1, status: "ready", ref }], - }); - const missing = yield* projectionStore - .getCheckpointContext(ThreadId.make("thread:checkpoint-context:missing")) - .pipe(Effect.flip); - assert.instanceOf(missing, ProjectionStore.ProjectionStoreThreadNotFoundError); + for (const store of [projectionStore, memoryStore]) { + assert.deepEqual(yield* store.getCheckpointContext(threadId), { + runs: [{ id: runId, ordinal: 1, status: "completed" }], + checkpointScopes: [{ id: scopeId, runId, kind: "root_run", cwd: "/repo/worktree" }], + checkpoints: [ + { scopeId, runId, ordinalWithinScope: 1, appRunOrdinal: 1, status: "ready", ref }, + ], + }); + const missing = yield* store + .getCheckpointContext(ThreadId.make("thread:checkpoint-context:missing")) + .pipe(Effect.flip); + assert.instanceOf(missing, ProjectionStore.ProjectionStoreThreadNotFoundError); + } }), ); diff --git a/apps/server/src/orchestration-v2/ProjectionStore.ts b/apps/server/src/orchestration-v2/ProjectionStore.ts index 002ef7bafa95..885c763b9981 100644 --- a/apps/server/src/orchestration-v2/ProjectionStore.ts +++ b/apps/server/src/orchestration-v2/ProjectionStore.ts @@ -220,9 +220,10 @@ const ProjectionCheckpointContext = Schema.Struct({ ), checkpoints: Schema.Array( OrchestrationV2CheckpointJsonSchema.mapFields( - ({ scopeId, runId, appRunOrdinal, status, ref }) => ({ + ({ scopeId, runId, ordinalWithinScope, appRunOrdinal, status, ref }) => ({ scopeId, runId, + ordinalWithinScope, appRunOrdinal, status, ref, @@ -4721,6 +4722,7 @@ export const layer: Layer.Layer = `, sql` SELECT scope_id AS "scopeId", run_id AS "runId", + ordinal_within_scope AS "ordinalWithinScope", app_run_ordinal AS "appRunOrdinal", status, json_extract(payload_json, '$.ref') AS ref FROM orchestration_v2_projection_checkpoints @@ -6583,9 +6585,10 @@ export const layerMemory: Layer.Layer = Layer.effect( cwd, })), checkpoints: projection.checkpoints.map( - ({ scopeId, runId, appRunOrdinal, status, ref }) => ({ + ({ scopeId, runId, ordinalWithinScope, appRunOrdinal, status, ref }) => ({ scopeId, runId, + ordinalWithinScope, appRunOrdinal, status, ref,