diff --git a/apps/server/src/orchestration-v2/CheckpointCaptureService.test.ts b/apps/server/src/orchestration-v2/CheckpointCaptureService.test.ts index 828d438646cd..c2dfd8609e80 100644 --- a/apps/server/src/orchestration-v2/CheckpointCaptureService.test.ts +++ b/apps/server/src/orchestration-v2/CheckpointCaptureService.test.ts @@ -52,9 +52,13 @@ const modelSelection = { } as const; it.layer(layerProjectionStoreTest)("CheckpointCaptureServiceV2", (it) => { - it.effect.each([false, true])( - "captures without decoding history or losing newer delegated completion, ref lookup fails=%s", - (refLookupFails) => + it.effect.each([ + { ordinal: 1, refLookupFails: false }, + { ordinal: 1, refLookupFails: true }, + { ordinal: 128, refLookupFails: false }, + ])( + "captures without decoding history or losing newer delegated completion, ordinal=$ordinal, ref lookup fails=$refLookupFails", + ({ ordinal, refLookupFails }) => Effect.gen(function* () { const projectionStore = yield* ProjectionStore.ProjectionStoreV2; const now = yield* DateTime.now; @@ -77,7 +81,7 @@ it.layer(layerProjectionStoreTest)("CheckpointCaptureServiceV2", (it) => { const staleRun: OrchestrationV2Run = { id: runId, threadId, - ordinal: 1, + ordinal, providerInstanceId, modelSelection, providerThreadId, @@ -160,8 +164,8 @@ it.layer(layerProjectionStoreTest)("CheckpointCaptureServiceV2", (it) => { runId, nodeId: rootNodeId, parentCheckpointId: readyBaseline.id, - ordinalWithinScope: 1, - appRunOrdinal: 1, + ordinalWithinScope: ordinal, + appRunOrdinal: ordinal, ref: CheckpointRef.make("checkpoint-ref:captured-1"), status: "ready" as const, files: [], @@ -251,6 +255,25 @@ it.layer(layerProjectionStoreTest)("CheckpointCaptureServiceV2", (it) => { payload: readyBaseline, }); const sql = yield* SqlClient.SqlClient; + if (ordinal > 1) { + yield* sql`WITH RECURSIVE history(ordinal) AS ( + SELECT 1 UNION ALL SELECT ordinal + 1 FROM history WHERE ordinal < ${ordinal - 1} + ) + INSERT INTO orchestration_v2_projection_checkpoints + (checkpoint_id, thread_id, scope_id, node_id, ordinal_within_scope, status, + captured_at, payload_json) + SELECT 'checkpoint:history:' || ordinal, ${threadId}, ${scopeId}, ${rootNodeId}, + ordinal, 'ready', ${DateTime.formatIso(now)}, '{}' FROM history`; + } + const baselineContext = yield* projectionStore.getCheckpointCaptureContext( + threadId, + { runId, scopeId }, + { preservedCheckpoints: "baselines" }, + ); + assert.deepEqual( + [...baselineContext.preservedCheckpointOrdinals].toSorted((a, b) => a - b), + ordinal === 1 ? [0] : [0, ordinal - 1], + ); yield* sql`INSERT OR REPLACE INTO orchestration_v2_projection_turn_items (turn_item_id, thread_id, run_id, node_id, provider_thread_id, provider_turn_id, parent_item_id, ordinal, type, status, updated_at, payload_json) diff --git a/apps/server/src/orchestration-v2/CheckpointCaptureService.ts b/apps/server/src/orchestration-v2/CheckpointCaptureService.ts index 43cdae25d4bc..9f665f55df46 100644 --- a/apps/server/src/orchestration-v2/CheckpointCaptureService.ts +++ b/apps/server/src/orchestration-v2/CheckpointCaptureService.ts @@ -66,7 +66,9 @@ export const layer: Layer.Layer< readonly scopeId: CheckpointScopeId; }) { const { run, rootNode, scope, providerThread, preservedCheckpointOrdinals, sideChat } = - yield* projections.getCheckpointCaptureContext(input.threadId, input); + yield* projections.getCheckpointCaptureContext(input.threadId, input, { + preservedCheckpoints: "baselines", + }); // Stopped runs remain terminal; their checkpoint is the next message's rollback point. const stopped = run?.status === "interrupted" || run?.status === "cancelled"; diff --git a/apps/server/src/orchestration-v2/ProjectionStore.ts b/apps/server/src/orchestration-v2/ProjectionStore.ts index a003ade1ae8f..cfeaf36196eb 100644 --- a/apps/server/src/orchestration-v2/ProjectionStore.ts +++ b/apps/server/src/orchestration-v2/ProjectionStore.ts @@ -748,6 +748,7 @@ export interface ProjectionStoreV2Shape { readonly getCheckpointCaptureContext: ( threadId: ThreadId, target: { readonly runId: RunId; readonly scopeId: CheckpointScopeId }, + options?: { readonly preservedCheckpoints: "baselines" }, ) => Effect.Effect; readonly getCheckpointRollbackContext: ( threadId: ThreadId, @@ -5962,6 +5963,7 @@ export const layer: Layer.Layer = const getCheckpointCaptureContext: ProjectionStoreV2Shape["getCheckpointCaptureContext"] = ( threadId, target, + options, ) => sql .withTransaction( @@ -5992,7 +5994,12 @@ export const layer: Layer.Layer = sql<{ readonly ordinal: number }>`SELECT ordinal_within_scope AS ordinal FROM orchestration_v2_projection_checkpoints WHERE thread_id = ${threadId} AND scope_id = ${target.scopeId} - AND (status = 'ready' OR (status = 'missing' AND run_id IS NOT NULL))`, + AND (status = 'ready' OR (status = 'missing' AND run_id IS NOT NULL)) + ${ + options?.preservedCheckpoints === "baselines" + ? sql`AND ordinal_within_scope IN (0, ${run === undefined ? null : Math.max(0, run.ordinal - 1)})` + : sql`` + }`, ]); return { run,