diff --git a/apps/server/src/orchestration-v2/CheckpointCaptureService.test.ts b/apps/server/src/orchestration-v2/CheckpointCaptureService.test.ts index 9d34049c3b164..16c1bdcf2df1b 100644 --- a/apps/server/src/orchestration-v2/CheckpointCaptureService.test.ts +++ b/apps/server/src/orchestration-v2/CheckpointCaptureService.test.ts @@ -20,6 +20,7 @@ import * as DateTime from "effect/DateTime"; import * as Effect from "effect/Effect"; import * as Layer from "effect/Layer"; import * as Ref from "effect/Ref"; +import * as Tracer from "effect/Tracer"; import * as SqlClient from "effect/sql/SqlClient"; import * as SqlitePersistence from "../persistence/Sqlite.ts"; @@ -52,9 +53,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 +82,7 @@ it.layer(layerProjectionStoreTest)("CheckpointCaptureServiceV2", (it) => { const staleRun: OrchestrationV2Run = { id: runId, threadId, - ordinal: 1, + ordinal, providerInstanceId, modelSelection, providerThreadId, @@ -160,8 +165,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 +256,24 @@ 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 fullContext = yield* projectionStore.getCheckpointCaptureContext(threadId, { + runId, + scopeId, + }); + assert.deepEqual( + [...fullContext.readyCheckpointOrdinals].sort((a, b) => a - b), + Array.from({ length: ordinal }, (_, index) => index), + ); 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) @@ -321,7 +344,43 @@ it.layer(layerProjectionStoreTest)("CheckpointCaptureServiceV2", (it) => { .execute({ threadId, runId, scopeId: CheckpointScopeId.make("missing-scope") }) .pipe(Effect.flip); assert.instanceOf(incomplete, CheckpointCaptureService.CheckpointCaptureExecutionError); - yield* service.execute({ threadId, runId, scopeId }); + let readyQuery: string | undefined; + let readyRows: unknown; + const tracer = Tracer.make({ + span(options) { + const span = new Tracer.NativeSpan(options); + const end = span.end.bind(span); + span.end = (endTime, exit) => { + end(endTime, exit); + const query = span.attributes.get("db.query.text"); + if ( + typeof query === "string" && + query.includes("SELECT ordinal_within_scope AS ordinal") && + exit._tag === "Success" + ) { + readyQuery = query; + readyRows = exit.value; + } + }; + return span; + }, + }); + yield* service.execute({ threadId, runId, scopeId }).pipe(Effect.withTracer(tracer)); + assert.isDefined(readyQuery); + assert.deepEqual( + readyRows, + ordinal === 1 ? [{ ordinal: 0 }] : [{ ordinal: 0 }, { ordinal: ordinal - 1 }], + ); + const plan = yield* sql.unsafe<{ readonly detail: string }>( + `EXPLAIN QUERY PLAN ${readyQuery}`, + ); + assert.isTrue( + plan.some(({ detail }) => + detail.includes( + "checkpoints_scope_ordinal_idx (scope_id=? AND ordinal_within_scope=?)", + ), + ), + ); const events = yield* Ref.get(committed); const runUpdated = events.find((event) => event.type === "run.updated"); diff --git a/apps/server/src/orchestration-v2/CheckpointCaptureService.ts b/apps/server/src/orchestration-v2/CheckpointCaptureService.ts index 6a2d9b2062b35..2f5914e254b30 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, readyCheckpointOrdinals, sideChat } = - yield* projections.getCheckpointCaptureContext(input.threadId, input); + yield* projections.getCheckpointCaptureContext(input.threadId, input, { + readyCheckpoints: "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 2f8c972e7ded9..f982a2b84f130 100644 --- a/apps/server/src/orchestration-v2/ProjectionStore.ts +++ b/apps/server/src/orchestration-v2/ProjectionStore.ts @@ -659,6 +659,7 @@ export interface ProjectionStoreV2Shape { readonly getCheckpointCaptureContext: ( threadId: ThreadId, target: { readonly runId: RunId; readonly scopeId: CheckpointScopeId }, + options?: { readonly readyCheckpoints: "baselines" }, ) => Effect.Effect; readonly getCheckpointRollbackContext: ( threadId: ThreadId, @@ -5300,6 +5301,7 @@ export const layer: Layer.Layer = const getCheckpointCaptureContext: ProjectionStoreV2Shape["getCheckpointCaptureContext"] = ( threadId, target, + options, ) => sql .withTransaction( @@ -5329,7 +5331,12 @@ export const layer: Layer.Layer = OR provider_thread_id IN (SELECT provider_thread_id FROM orchestration_v2_projection_subagents WHERE thread_id = ${threadId}))`, 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'`, + WHERE thread_id = ${threadId} AND scope_id = ${target.scopeId} AND status = 'ready' + ${ + options?.readyCheckpoints === "baselines" + ? sql`AND ordinal_within_scope IN (0, ${run === undefined ? null : Math.max(0, run.ordinal - 1)})` + : sql`` + }`, ]); return { run,