From 05ed3bfe230db21f3baf50c072d2814e02053822 Mon Sep 17 00:00:00 2001 From: Julius Marminge <51714798+juliusmarminge@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:35:44 -0700 Subject: [PATCH] fix(server): avoid loading v2 transcripts for checkpoint effects --- .../CheckpointCaptureService.test.ts | 51 ++- .../CheckpointCaptureService.ts | 2 +- .../CheckpointProjectionReads.test.ts | 352 ++++++++++++++++++ .../CheckpointRestoreSafety.test.ts | 2 +- .../CheckpointRollbackService.test.ts | 12 +- .../CheckpointRollbackService.ts | 2 +- .../src/orchestration-v2/ProjectionStore.ts | 344 +++++++++++++++++ .../ProviderTurnControlService.test.ts | 2 + .../RunFinalizationService.test.ts | 10 +- .../RunFinalizationService.ts | 2 +- 10 files changed, 761 insertions(+), 18 deletions(-) create mode 100644 apps/server/src/orchestration-v2/CheckpointProjectionReads.test.ts diff --git a/apps/server/src/orchestration-v2/CheckpointCaptureService.test.ts b/apps/server/src/orchestration-v2/CheckpointCaptureService.test.ts index 3cf971c1f1d8..5caa612ac44d 100644 --- a/apps/server/src/orchestration-v2/CheckpointCaptureService.test.ts +++ b/apps/server/src/orchestration-v2/CheckpointCaptureService.test.ts @@ -19,6 +19,8 @@ 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 Schema from "effect/Schema"; +import * as SqlClient from "effect/unstable/sql/SqlClient"; import { SqlitePersistenceMemory } from "../persistence/Layers/Sqlite.ts"; import * as CheckpointStore from "../checkpointing/CheckpointStore.ts"; @@ -26,6 +28,7 @@ import { VcsProcessTimeoutError } from "@t3tools/contracts"; import { CheckpointServiceV2, layer as checkpointServiceLayer } from "./CheckpointService.ts"; import * as CheckpointCaptureService from "./CheckpointCaptureService.ts"; import { EventSinkV2 } from "./EventSink.ts"; +import * as RunFinalization from "./RunFinalizationService.ts"; import * as IdAllocator from "./IdAllocator.ts"; import * as ProjectionStore from "./ProjectionStore.ts"; @@ -34,6 +37,8 @@ const ProjectionStoreTestLayer = Layer.mergeAll( SqlitePersistenceMemory, ); +const encodeUnknownJsonString = Schema.encodeSync(Schema.fromJsonString(Schema.Unknown)); + const threadId = ThreadId.make("thread:checkpoint-capture-delegated"); const projectId = ProjectId.make("project:checkpoint-capture-delegated"); const runId = RunId.make("run:checkpoint-capture-delegated"); @@ -252,6 +257,24 @@ it.layer(ProjectionStoreTestLayer)("CheckpointCaptureServiceV2", (it) => { payload: readyBaseline, }); + // Checkpointing must not decode obsolete or large conversation history. + const sql = yield* SqlClient.SqlClient; + yield* sql` + INSERT 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 + ) VALUES ( + ${`turn-item:checkpoint-capture:${refLookupFails}`}, ${threadId}, ${runId}, ${rootNodeId}, NULL, NULL, + NULL, 1, 'assistant_message', 'completed', ${DateTime.formatIso(now)}, + ${encodeUnknownJsonString({ obsolete: "transcript shape", text: "x".repeat(1024 * 1024) })} + ) + `; + yield* sql` + UPDATE orchestration_v2_projection_checkpoints + SET payload_json = json_set(payload_json, '$.files', 'obsolete file summary') + WHERE checkpoint_id = ${readyBaseline.id} + `; + const committed = yield* Ref.make>([]); const captureLayer = CheckpointCaptureService.layer.pipe( Layer.provide( @@ -301,10 +324,20 @@ it.layer(ProjectionStoreTestLayer)("CheckpointCaptureServiceV2", (it) => { ), ); + const finalizationLayer = RunFinalization.layer.pipe( + Layer.provideMerge(captureLayer), + Layer.provide( + Layer.succeed(RunFinalization.RunFinalizationObserver, { + refresh: () => Effect.void, + refreshAfterTurn: () => Effect.void, + }), + ), + ); + yield* Effect.gen(function* () { - const service = yield* CheckpointCaptureService.CheckpointCaptureServiceV2; + const finalization = yield* RunFinalization.RunFinalizationService; // Capture reads the waiting run while the projection still holds the stale cohort. - yield* service.execute({ threadId, runId, scopeId }); + yield* finalization.finalize({ threadId, runId, scopeId }); const events = yield* Ref.get(committed); const runUpdated = events.find((event) => event.type === "run.updated"); @@ -346,15 +379,25 @@ it.layer(ProjectionStoreTestLayer)("CheckpointCaptureServiceV2", (it) => { // Apply the real capture-emitted run.updated through ProjectionStore. yield* projectionStore.apply(runUpdated); - const projection = yield* projectionStore.getThreadProjection(threadId); + const projection = yield* projectionStore.getCheckpointCaptureContext(threadId, { + runId, + scopeId, + }); const projectedRun = projection.runs[0]; assert.isDefined(projectedRun); assert.equal(projectedRun?.status, "completed"); assert.equal(projectedRun?.checkpointId, runUpdated.payload.checkpointId); assert.deepEqual(projectedRun?.delegatedCompletion, newerCohort); assert.equal(projectedRun?.delegatedCompletion?.delivery?.messageId, deliveryMessageId); + yield* Ref.set(committed, []); + yield* finalization.finalize({ threadId, runId, scopeId }); + assert.deepEqual( + yield* Ref.get(committed), + [], + "a repeated capture must remain idempotent", + ); assert.deepEqual(projectedRun?.delegatedCompletion?.delivery?.taskIds, [taskId]); - }).pipe(Effect.provide(captureLayer)); + }).pipe(Effect.provide(finalizationLayer)); }), ); }); diff --git a/apps/server/src/orchestration-v2/CheckpointCaptureService.ts b/apps/server/src/orchestration-v2/CheckpointCaptureService.ts index 6f92331ca6cd..bd2c51ff5c71 100644 --- a/apps/server/src/orchestration-v2/CheckpointCaptureService.ts +++ b/apps/server/src/orchestration-v2/CheckpointCaptureService.ts @@ -62,7 +62,7 @@ export const layer: Layer.Layer< readonly runId: RunId; readonly scopeId: CheckpointScopeId; }) { - const projection = yield* projections.getThreadProjection(input.threadId); + const projection = yield* projections.getCheckpointCaptureContext(input.threadId, input); const run = projection.runs.find((candidate) => candidate.id === input.runId); // The effect is at-least-once. A completed run with a checkpoint proves diff --git a/apps/server/src/orchestration-v2/CheckpointProjectionReads.test.ts b/apps/server/src/orchestration-v2/CheckpointProjectionReads.test.ts new file mode 100644 index 000000000000..c41b5c3fd044 --- /dev/null +++ b/apps/server/src/orchestration-v2/CheckpointProjectionReads.test.ts @@ -0,0 +1,352 @@ +import { assert, it } from "@effect/vitest"; +import { + CheckpointId, + CheckpointRef, + CheckpointScopeId, + EventId, + MessageId, + NodeId, + type OrchestrationV2DomainEvent, + ProjectId, + ProviderDriverKind, + ProviderInstanceId, + ProviderSessionId, + ProviderThreadId, + ProviderTurnId, + RunAttemptId, + RunId, + ThreadId, +} from "@t3tools/contracts"; +import * as DateTime from "effect/DateTime"; +import * as Effect from "effect/Effect"; +import * as Layer from "effect/Layer"; +import * as SqlClient from "effect/unstable/sql/SqlClient"; +import { SqlitePersistenceMemory } from "../persistence/Layers/Sqlite.ts"; +import { CodexProviderCapabilitiesV2 } from "./Adapters/CodexAdapterV2.ts"; +import { ProjectionStoreV2, layer, layerMemory } from "./ProjectionStore.ts"; + +const threadId = ThreadId.make("thread:checkpoint-reads"); +const otherThreadId = ThreadId.make("thread:checkpoint-reads:other"); +const providerThreadId = ProviderThreadId.make("provider-thread:checkpoint-reads"); +const providerSessionId = ProviderSessionId.make("provider-session:checkpoint-reads"); +const providerInstanceId = ProviderInstanceId.make("codex"); +const driver = ProviderDriverKind.make("codex"); +const modelSelection = { instanceId: providerInstanceId, model: "test" }; +const scopeId = CheckpointScopeId.make("scope:checkpoint-reads"); +const otherScopeId = CheckpointScopeId.make("scope:checkpoint-reads:other"); +const runId = (ordinal: number) => RunId.make(`run:checkpoint-reads:${ordinal}`); +const nodeId = (ordinal: number) => NodeId.make(`node:checkpoint-reads:${ordinal}`); +const attemptId = (ordinal: number) => RunAttemptId.make(`attempt:checkpoint-reads:${ordinal}`); +const turnId = (ordinal: number) => ProviderTurnId.make(`turn:checkpoint-reads:${ordinal}`); +const checkpointId = (ordinal: number) => + CheckpointId.make(`checkpoint:checkpoint-reads:${ordinal}`); + +type EventInput = Event extends OrchestrationV2DomainEvent + ? Omit + : never; + +const seed = Effect.gen(function* () { + const store = yield* ProjectionStoreV2; + const now = yield* DateTime.now; + let sequence = 0; + const apply = (event: EventInput) => + store.apply({ + ...event, + id: EventId.make(`event:checkpoint-reads:${sequence++}`), + occurredAt: now, + } as OrchestrationV2DomainEvent); + for (const id of [threadId, otherThreadId]) { + yield* apply({ + type: "thread.created", + threadId: id, + payload: { + createdBy: "user", + creationSource: "web", + id, + projectId: ProjectId.make("project:checkpoint-reads"), + title: "Checkpoint reads", + providerInstanceId, + modelSelection, + runtimeMode: "full-access", + interactionMode: "default", + branch: null, + worktreePath: "/repo/worktree", + activeProviderThreadId: providerThreadId, + lineage: { parentThreadId: null, relationshipToParent: null, rootThreadId: id }, + forkedFrom: null, + createdAt: now, + updatedAt: now, + archivedAt: null, + settledOverride: null, + settledAt: null, + lastVisitedAt: null, + deletedAt: null, + }, + }); + } + yield* apply({ + type: "provider-thread.updated", + threadId, + payload: { + id: providerThreadId, + driver, + providerInstanceId, + providerSessionId, + appThreadId: threadId, + ownerNodeId: null, + nativeThreadRef: null, + nativeConversationHeadRef: null, + status: "idle", + firstRunOrdinal: 1, + lastRunOrdinal: 4, + handoffIds: [], + forkedFrom: null, + createdAt: now, + updatedAt: now, + }, + }); + yield* apply({ + type: "provider-session.attached", + threadId, + payload: { + id: providerSessionId, + driver, + providerInstanceId, + status: "ready", + cwd: "/repo/worktree", + model: "test", + capabilities: CodexProviderCapabilitiesV2, + createdAt: now, + updatedAt: now, + lastError: null, + }, + }); + for (const id of [scopeId, otherScopeId]) { + yield* apply({ + type: "checkpoint-scope.created", + threadId, + payload: { + id, + threadId, + runId: runId(1), + nodeId: nodeId(1), + parentScopeId: null, + providerThreadId, + kind: "root_run", + ordinalWithinParent: 0, + advancesAppRunCount: true, + cwd: "/repo/worktree", + createdAt: now, + }, + }); + } + for (let ordinal = 1; ordinal <= 4; ordinal++) { + const waiting = ordinal === 4; + yield* apply({ + type: "run.created", + threadId, + payload: { + id: runId(ordinal), + threadId, + ordinal, + providerInstanceId, + modelSelection, + providerThreadId, + userMessageId: MessageId.make(`message:checkpoint-reads:${ordinal}`), + rootNodeId: nodeId(ordinal), + activeAttemptId: attemptId(ordinal), + status: waiting ? "waiting" : "completed", + requestedAt: now, + startedAt: now, + completedAt: waiting ? null : now, + checkpointId: waiting ? null : checkpointId(ordinal), + contextHandoffId: null, + }, + }); + yield* apply({ + type: "run-attempt.created", + threadId, + runId: runId(ordinal), + payload: { + id: attemptId(ordinal), + runId: runId(ordinal), + attemptOrdinal: 1, + rootNodeId: nodeId(ordinal), + providerInstanceId, + providerThreadId, + providerTurnId: turnId(ordinal), + reason: "initial", + status: "completed", + startedAt: now, + completedAt: now, + }, + }); + yield* apply({ + type: "node.updated", + threadId, + payload: { + id: nodeId(ordinal), + threadId, + runId: runId(ordinal), + parentNodeId: null, + rootNodeId: nodeId(ordinal), + kind: "root_turn", + status: waiting ? "waiting" : "completed", + countsForRun: true, + providerThreadId, + providerTurnId: turnId(ordinal), + nativeItemRef: null, + runtimeRequestId: null, + checkpointScopeId: scopeId, + startedAt: now, + completedAt: waiting ? null : now, + }, + }); + yield* apply({ + type: "provider-turn.updated", + threadId, + payload: { + id: turnId(ordinal), + providerThreadId, + nodeId: nodeId(ordinal), + runAttemptId: attemptId(ordinal), + nativeTurnRef: null, + ordinal, + status: "completed", + startedAt: now, + completedAt: now, + }, + }); + } + for (let ordinal = 0; ordinal <= 4; ordinal++) { + yield* apply({ + type: "checkpoint.captured", + threadId, + payload: { + id: checkpointId(ordinal), + threadId, + scopeId: ordinal === 4 ? otherScopeId : scopeId, + runId: ordinal === 0 ? null : runId(ordinal), + nodeId: nodeId(Math.max(1, ordinal)), + parentCheckpointId: null, + ordinalWithinScope: ordinal, + appRunOrdinal: ordinal === 0 ? null : ordinal, + ref: CheckpointRef.make(`refs/t3/read/${ordinal}`), + status: ordinal === 3 ? "stale" : "ready", + files: [], + capturedAt: now, + }, + }); + } +}); + +for (const [name, storeLayer] of [ + ["SQLite", layer.pipe(Layer.provideMerge(SqlitePersistenceMemory))], + ["memory", layerMemory], +] as const) { + it.layer(storeLayer)(`checkpoint projection reads (${name})`, (it) => { + it.effect("selects capture readiness and rollback targets without unrelated history", () => + Effect.gen(function* () { + yield* seed; + const store = yield* ProjectionStoreV2; + const capture = yield* store.getCheckpointCaptureContext(threadId, { + runId: runId(4), + scopeId, + }); + assert.deepEqual( + capture.runs.map((run) => run.id), + [runId(4)], + ); + assert.deepEqual( + capture.nodes.map((node) => node.id), + [nodeId(4)], + ); + assert.deepEqual(capture.checkpoints, [ + { scopeId, ordinalWithinScope: 0, status: "ready" }, + { scopeId, ordinalWithinScope: 3, status: "stale" }, + ]); + const rollback = yield* store.getCheckpointRollbackContext(threadId, { + providerThreadId, + checkpointId: checkpointId(1), + scopeId, + }); + assert.deepEqual( + rollback.runs.map((run) => run.id), + [runId(1), runId(2), runId(3)], + ); + assert.deepEqual( + rollback.attempts.map((attempt) => attempt.id), + [attemptId(1)], + ); + assert.deepEqual( + rollback.nodes.map((node) => node.id), + [nodeId(2), nodeId(3)], + ); + assert.deepEqual( + rollback.checkpoints.map((checkpoint) => checkpoint.id), + [checkpointId(1), checkpointId(2)], + ); + assert.deepEqual( + rollback.providerTurns.map((turn) => turn.id), + [turnId(1), turnId(2), turnId(3), turnId(4)], + ); + assert.deepEqual( + rollback.providerSessions.map((session) => session.id), + [providerSessionId], + ); + const retry = yield* store.getCheckpointCaptureContext(threadId, { + runId: runId(1), + scopeId, + }); + assert.deepEqual(retry.nodes, []); + const foreign = yield* store.getCheckpointRollbackContext(otherThreadId, { + providerThreadId, + checkpointId: checkpointId(1), + scopeId, + }); + assert.deepEqual(foreign.providerThreads, []); + assert.deepEqual(foreign.checkpoints, []); + assert.deepEqual(foreign.runs, []); + }), + ); + }); +} + +it.effect("ignores obsolete transcript and unrelated checkpoint/run payloads in SQLite", () => + Effect.gen(function* () { + yield* seed; + const store = yield* ProjectionStoreV2; + const sql = yield* SqlClient.SqlClient; + yield* sql`INSERT 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 + ) VALUES ('obsolete', ${threadId}, NULL, NULL, NULL, NULL, NULL, 1, 'assistant_message', 'completed', '2026-09-17', '{"obsolete":true}')`; + yield* sql`UPDATE orchestration_v2_projection_checkpoints + SET payload_json = json_set(payload_json, '$.files', 'obsolete') + WHERE checkpoint_id IN (${checkpointId(0)}, ${checkpointId(3)}, ${checkpointId(4)})`; + // A waiting later run is not a completed run that rollback invalidates. + yield* sql`UPDATE orchestration_v2_projection_runs SET payload_json = '{"obsolete":true}' WHERE run_id = ${runId(4)}`; + assert.equal((yield* Effect.exit(store.getThreadProjection(threadId)))._tag, "Failure"); + const rollback = yield* store.getCheckpointRollbackContext(threadId, { + providerThreadId, + checkpointId: checkpointId(1), + scopeId, + }); + assert.deepEqual( + rollback.checkpoints.map((checkpoint) => checkpoint.id), + [checkpointId(1), checkpointId(2)], + ); + // Capture readiness reads columns even when the prior file summaries are obsolete. + yield* sql`UPDATE orchestration_v2_projection_runs + SET status = 'waiting', payload_json = json_set(payload_json, '$.status', 'waiting', '$.checkpointId', NULL) + WHERE run_id = ${runId(2)}`; + const capture = yield* store.getCheckpointCaptureContext(threadId, { + runId: runId(2), + scopeId, + }); + assert.deepEqual( + capture.checkpoints.map((checkpoint) => checkpoint.ordinalWithinScope), + [0, 1], + ); + }).pipe(Effect.provide(layer.pipe(Layer.provideMerge(SqlitePersistenceMemory)))), +); diff --git a/apps/server/src/orchestration-v2/CheckpointRestoreSafety.test.ts b/apps/server/src/orchestration-v2/CheckpointRestoreSafety.test.ts index cb3cfef48452..d1efa67e433b 100644 --- a/apps/server/src/orchestration-v2/CheckpointRestoreSafety.test.ts +++ b/apps/server/src/orchestration-v2/CheckpointRestoreSafety.test.ts @@ -102,7 +102,7 @@ it.effect.each([ getById: () => Effect.succeed(Option.some({ workspaceRoot: parent } as never)), }), Layer.mock(ProjectionStoreV2)({ - getThreadProjection: () => Effect.succeed(projection), + getCheckpointRollbackContext: () => Effect.succeed(projection), getShellSnapshot: () => Effect.succeed({ schemaVersion: 1, diff --git a/apps/server/src/orchestration-v2/CheckpointRollbackService.test.ts b/apps/server/src/orchestration-v2/CheckpointRollbackService.test.ts index 196d3821afe3..7436ba29d5ef 100644 --- a/apps/server/src/orchestration-v2/CheckpointRollbackService.test.ts +++ b/apps/server/src/orchestration-v2/CheckpointRollbackService.test.ts @@ -67,7 +67,7 @@ it.effect("rejects a non-ready checkpoint before opening a session or restoring Layer.mock(EventSinkV2)({}), idAllocatorLayer, Layer.mock(ProjectionStoreV2)({ - getThreadProjection: () => Effect.succeed(projection), + getCheckpointRollbackContext: () => Effect.succeed(projection), getShellSnapshot: () => Effect.succeed({ schemaVersion: 1, @@ -144,7 +144,7 @@ it.effect("rejects a rollback when another provider thread became active", () => Layer.mock(EventSinkV2)({}), idAllocatorLayer, Layer.mock(ProjectionStoreV2)({ - getThreadProjection: () => Effect.succeed(projection), + getCheckpointRollbackContext: () => Effect.succeed(projection), getShellSnapshot: () => Effect.succeed({ schemaVersion: 1, @@ -224,7 +224,7 @@ it.effect("rejects a rollback when provider selection changed before execution", Layer.mock(EventSinkV2)({}), idAllocatorLayer, Layer.mock(ProjectionStoreV2)({ - getThreadProjection: () => Effect.succeed(projection), + getCheckpointRollbackContext: () => Effect.succeed(projection), getShellSnapshot: () => Effect.succeed({ schemaVersion: 1, @@ -295,7 +295,7 @@ it.effect("reports a missing provider turn as a structured rollback failure", () Layer.mock(EventSinkV2)({}), idAllocatorLayer, Layer.mock(ProjectionStoreV2)({ - getThreadProjection: () => Effect.succeed(projection), + getCheckpointRollbackContext: () => Effect.succeed(projection), getShellSnapshot: () => Effect.succeed({ schemaVersion: 1, @@ -351,7 +351,7 @@ it.effect("wraps underlying failures with an unexpected-failure reason and cause Layer.mock(EventSinkV2)({}), idAllocatorLayer, Layer.mock(ProjectionStoreV2)({ - getThreadProjection: () => Effect.fail(projectionError), + getCheckpointRollbackContext: () => Effect.fail(projectionError), }), Layer.mock(ProviderSessionManagerV2)({}), Layer.mock(RuntimePolicyV2)({}), @@ -436,7 +436,7 @@ it.effect.each([ }), idAllocatorLayer, Layer.mock(ProjectionStoreV2)({ - getThreadProjection: () => Effect.succeed(projection), + getCheckpointRollbackContext: () => Effect.succeed(projection), getThreadProviderContext: () => Effect.succeed({ providerSessions: [] } as never), getCheckpointContext: () => Effect.succeed({ diff --git a/apps/server/src/orchestration-v2/CheckpointRollbackService.ts b/apps/server/src/orchestration-v2/CheckpointRollbackService.ts index 144d84345576..cf75bb780ef7 100644 --- a/apps/server/src/orchestration-v2/CheckpointRollbackService.ts +++ b/apps/server/src/orchestration-v2/CheckpointRollbackService.ts @@ -108,7 +108,7 @@ export const layer: Layer.Layer< readonly scopeId: CheckpointScopeId; readonly restoreFiles?: boolean; }) { - const projection = yield* projections.getThreadProjection(input.threadId); + const projection = yield* projections.getCheckpointRollbackContext(input.threadId, input); const providerThread = projection.providerThreads.find( (candidate) => candidate.id === input.providerThreadId, ); diff --git a/apps/server/src/orchestration-v2/ProjectionStore.ts b/apps/server/src/orchestration-v2/ProjectionStore.ts index cbef23b92134..56a16616796a 100644 --- a/apps/server/src/orchestration-v2/ProjectionStore.ts +++ b/apps/server/src/orchestration-v2/ProjectionStore.ts @@ -22,6 +22,8 @@ import type { RunAttemptId, RuntimeRequestId, MessageId, + CheckpointId, + CheckpointScopeId, } from "@t3tools/contracts"; import { OrchestrationV2AppThreadJson as OrchestrationV2AppThreadJsonSchema, @@ -226,6 +228,52 @@ export interface ProjectionRuntimeResponseContext { readonly session: OrchestrationV2ThreadProjection["providerSessions"][number] | undefined; } +type ProjectionCheckpointCaptureContext = Pick< + OrchestrationV2ThreadProjection, + "runs" | "nodes" | "checkpointScopes" | "providerThreads" +> & { + readonly checkpoints: ReadonlyArray< + Pick< + OrchestrationV2ThreadProjection["checkpoints"][number], + "scopeId" | "ordinalWithinScope" | "status" + > + >; +}; + +type ProjectionCheckpointRollbackContext = Pick< + OrchestrationV2ThreadProjection, + | "thread" + | "runs" + | "attempts" + | "nodes" + | "providerThreads" + | "providerSessions" + | "providerTurns" + | "checkpointScopes" + | "checkpoints" +>; + +interface CheckpointCaptureTarget { + readonly runId: RunId; + readonly scopeId: CheckpointScopeId; +} + +interface CheckpointRollbackTarget { + readonly providerThreadId: ProviderThreadId; + readonly checkpointId: CheckpointId; + readonly scopeId: CheckpointScopeId; +} + +const decodeCaptureCheckpointReadiness = Schema.decodeUnknownEffect( + Schema.Array( + OrchestrationV2CheckpointJsonSchema.mapFields(({ scopeId, ordinalWithinScope, status }) => ({ + scopeId, + ordinalWithinScope, + status, + })), + ), +); + export interface ProjectionStoreV2Shape { readonly apply: ( event: OrchestrationV2DomainEvent, @@ -285,6 +333,14 @@ export interface ProjectionStoreV2Shape { readonly getCheckpointContext: ( threadId: ThreadId, ) => Effect.Effect; + readonly getCheckpointCaptureContext: ( + threadId: ThreadId, + target: CheckpointCaptureTarget, + ) => Effect.Effect; + readonly getCheckpointRollbackContext: ( + threadId: ThreadId, + target: CheckpointRollbackTarget, + ) => Effect.Effect; readonly getRecoveryThreadIds: ( kind: ProjectionRecoveryKind, ) => Effect.Effect, ProjectionStoreV2Error>; @@ -3788,6 +3844,174 @@ export const layer: Layer.Layer = ) .pipe(Effect.mapError(controlReadError(threadId))); + const getCheckpointCaptureContext: ProjectionStoreV2Shape["getCheckpointCaptureContext"] = ( + threadId, + target, + ) => + sql + .withTransaction( + Effect.gen(function* () { + yield* requireThread(threadId); + const runRows = yield* sql` + SELECT payload_json FROM orchestration_v2_projection_runs + WHERE thread_id = ${threadId} AND run_id = ${target.runId}`; + const runs = yield* decodeRows(decodeRunPayload, threadId)(runRows); + const run = runs[0]; + if (run === undefined || (run.status === "completed" && run.checkpointId !== null)) { + return { + runs, + nodes: [], + checkpointScopes: [], + providerThreads: [], + checkpoints: [], + }; + } + const [nodeRows, scopeRows, providerRows, readinessRows] = yield* Effect.all([ + sql`SELECT payload_json FROM orchestration_v2_projection_nodes + WHERE thread_id = ${threadId} AND node_id = ${run.rootNodeId}`, + sql`SELECT payload_json FROM orchestration_v2_projection_checkpoint_scopes + WHERE thread_id = ${threadId} AND scope_id = ${target.scopeId}`, + sql`SELECT payload_json FROM orchestration_v2_projection_provider_threads + WHERE provider_thread_id = ${run.providerThreadId} AND ( + thread_id = ${threadId} + OR owner_node_id IN (SELECT node_id FROM orchestration_v2_projection_nodes WHERE thread_id = ${threadId}) + OR provider_thread_id IN (SELECT provider_thread_id FROM orchestration_v2_projection_subagents WHERE thread_id = ${threadId}) + )`, + sql`SELECT scope_id AS "scopeId", ordinal_within_scope AS "ordinalWithinScope", status + FROM orchestration_v2_projection_checkpoints + WHERE thread_id = ${threadId} AND scope_id = ${target.scopeId} + AND ordinal_within_scope IN (0, ${Math.max(0, run.ordinal - 1)})`, + ]); + return { + runs, + nodes: yield* decodeRows(decodeNodePayload, threadId)(nodeRows), + checkpointScopes: yield* decodeRows( + decodeCheckpointScopePayload, + threadId, + )(scopeRows), + providerThreads: yield* decodeRows( + decodeProviderThreadPayload, + threadId, + )(providerRows), + checkpoints: yield* decodeCaptureCheckpointReadiness(readinessRows), + }; + }), + ) + .pipe(Effect.mapError(controlReadError(threadId))); + + const getCheckpointRollbackContext: ProjectionStoreV2Shape["getCheckpointRollbackContext"] = ( + threadId, + target, + ) => + sql + .withTransaction( + Effect.gen(function* () { + const thread = yield* getThread(threadId); + const [providerRows, checkpointRows, scopeRows] = yield* Effect.all([ + sql`SELECT payload_json FROM orchestration_v2_projection_provider_threads + WHERE provider_thread_id = ${target.providerThreadId} AND ( + thread_id = ${threadId} + OR owner_node_id IN (SELECT node_id FROM orchestration_v2_projection_nodes WHERE thread_id = ${threadId}) + OR provider_thread_id IN (SELECT provider_thread_id FROM orchestration_v2_projection_subagents WHERE thread_id = ${threadId}) + )`, + sql`SELECT payload_json FROM orchestration_v2_projection_checkpoints + WHERE thread_id = ${threadId} AND checkpoint_id = ${target.checkpointId}`, + sql`SELECT payload_json FROM orchestration_v2_projection_checkpoint_scopes + WHERE thread_id = ${threadId} AND scope_id = ${target.scopeId}`, + ]); + const providerThreads = yield* decodeRows( + decodeProviderThreadPayload, + threadId, + )(providerRows); + const checkpointScopes = yield* decodeRows( + decodeCheckpointScopePayload, + threadId, + )(scopeRows); + const checkpoints = yield* decodeRows( + decodeCheckpointPayload, + threadId, + )(checkpointRows); + const providerThread = providerThreads[0]; + const checkpoint = checkpoints[0]; + const scope = checkpointScopes[0]; + if ( + providerThread === undefined || + checkpoint === undefined || + scope === undefined || + checkpoint.scopeId !== scope.id || + checkpoint.status !== "ready" || + providerThread.providerSessionId === null || + providerThread.id !== thread.activeProviderThreadId || + providerThread.providerInstanceId !== thread.modelSelection.instanceId + ) { + return { + thread, + providerThreads, + checkpointScopes, + checkpoints, + runs: [], + attempts: [], + nodes: [], + providerSessions: [], + providerTurns: [], + }; + } + const targetOrdinal = checkpoint.appRunOrdinal ?? 0; + const runRows = + yield* sql`SELECT payload_json FROM orchestration_v2_projection_runs + WHERE thread_id = ${threadId} + AND (ordinal = ${targetOrdinal} OR (ordinal > ${targetOrdinal} AND status = 'completed')) + ORDER BY ordinal ASC`; + const runs = yield* decodeRows(decodeRunPayload, threadId)(runRows); + const targetRun = runs.find((run) => run.ordinal === targetOrdinal); + const attemptRows = + yield* sql`SELECT payload_json FROM orchestration_v2_projection_run_attempts + WHERE thread_id = ${threadId} AND attempt_id = ${targetRun?.activeAttemptId ?? null}`; + const attempts = yield* decodeRows(decodeRunAttemptPayload, threadId)(attemptRows); + const attempt = attempts[0]; + const [nodeRows, sessionRows, turnRows, staleRows] = yield* Effect.all([ + sql`SELECT payload_json FROM orchestration_v2_projection_nodes + WHERE thread_id = ${threadId} AND node_id IN ( + SELECT json_extract(payload_json, '$.rootNodeId') FROM orchestration_v2_projection_runs + WHERE thread_id = ${threadId} AND ordinal > ${targetOrdinal} AND status = 'completed' + ) ORDER BY node_id ASC`, + sql`SELECT session.payload_json FROM orchestration_v2_projection_provider_sessions AS session + JOIN orchestration_v2_projection_provider_session_bindings AS binding + ON binding.provider_session_id = session.provider_session_id + WHERE binding.thread_id = ${threadId} AND session.provider_session_id = ${providerThread.providerSessionId}`, + sql`SELECT payload_json FROM orchestration_v2_projection_provider_turns + WHERE thread_id = ${threadId} AND ( + provider_thread_id = ${providerThread.id} + OR provider_turn_id = ${attempt?.providerTurnId ?? null} + OR run_attempt_id = ${attempt?.id ?? null} + ) ORDER BY provider_thread_id ASC, ordinal ASC`, + sql`SELECT payload_json FROM orchestration_v2_projection_checkpoints + WHERE thread_id = ${threadId} AND scope_id = ${scope.id} AND status = 'ready' + AND app_run_ordinal > ${targetOrdinal} + AND checkpoint_id != ${checkpoint.id} + ORDER BY scope_id ASC, ordinal_within_scope ASC`, + ]); + return { + thread, + providerThreads, + checkpointScopes, + runs, + attempts, + checkpoints: [ + ...checkpoints, + ...(yield* decodeRows(decodeCheckpointPayload, threadId)(staleRows)), + ], + nodes: yield* decodeRows(decodeNodePayload, threadId)(nodeRows), + providerSessions: yield* decodeRows( + decodeProviderSessionPayload, + threadId, + )(sessionRows), + providerTurns: yield* decodeRows(decodeProviderTurnPayload, threadId)(turnRows), + }; + }), + ) + .pipe(Effect.mapError(controlReadError(threadId))); + const getCheckpointContext: ProjectionStoreV2Shape["getCheckpointContext"] = (threadId) => sql .withTransaction( @@ -4524,6 +4748,8 @@ export const layer: Layer.Layer = getThreadProviderContext, getRuntimeResponseContext, getCheckpointContext, + getCheckpointCaptureContext, + getCheckpointRollbackContext, getPendingNativeUserInputs, getRuntimeRequest, getPlan, @@ -4774,6 +5000,124 @@ export const layerMemory: Layer.Layer = Layer.effect( ), }; }), + getCheckpointCaptureContext: (threadId, target) => + Effect.gen(function* () { + const projection = (yield* Ref.get(replayState)).projections.get(threadId); + if (projection === undefined) + return yield* new ProjectionStoreThreadNotFoundError({ threadId }); + const runs = projection.runs.filter((run) => run.id === target.runId); + const run = runs[0]; + if (run === undefined || (run.status === "completed" && run.checkpointId !== null)) { + return { runs, nodes: [], checkpointScopes: [], providerThreads: [], checkpoints: [] }; + } + return { + runs, + nodes: projection.nodes.filter((node) => node.id === run.rootNodeId), + checkpointScopes: projection.checkpointScopes.filter( + (scope) => scope.id === target.scopeId, + ), + providerThreads: projection.providerThreads.filter( + (provider) => provider.id === run.providerThreadId, + ), + checkpoints: projection.checkpoints + .filter( + (checkpoint) => + checkpoint.scopeId === target.scopeId && + (checkpoint.ordinalWithinScope === 0 || + checkpoint.ordinalWithinScope === Math.max(0, run.ordinal - 1)), + ) + .map(({ scopeId, ordinalWithinScope, status }) => ({ + scopeId, + ordinalWithinScope, + status, + })), + }; + }), + getCheckpointRollbackContext: (threadId, target) => + Effect.gen(function* () { + const projection = (yield* Ref.get(replayState)).projections.get(threadId); + if (projection === undefined) + return yield* new ProjectionStoreThreadNotFoundError({ threadId }); + const thread = projection.thread; + const providerThreads = projection.providerThreads.filter( + (provider) => provider.id === target.providerThreadId, + ); + const checkpointScopes = projection.checkpointScopes.filter( + (scope) => scope.id === target.scopeId, + ); + const checkpoints = projection.checkpoints.filter( + (checkpoint) => checkpoint.id === target.checkpointId, + ); + const providerThread = providerThreads[0]; + const checkpoint = checkpoints[0]; + const scope = checkpointScopes[0]; + if ( + providerThread === undefined || + checkpoint === undefined || + scope === undefined || + checkpoint.scopeId !== scope.id || + checkpoint.status !== "ready" || + providerThread.providerSessionId === null || + providerThread.id !== thread.activeProviderThreadId || + providerThread.providerInstanceId !== thread.modelSelection.instanceId + ) { + return { + thread, + providerThreads, + checkpointScopes, + checkpoints, + runs: [], + attempts: [], + nodes: [], + providerSessions: [], + providerTurns: [], + }; + } + const targetOrdinal = checkpoint.appRunOrdinal ?? 0; + const runs = projection.runs.filter( + (run) => + run.ordinal === targetOrdinal || + (run.ordinal > targetOrdinal && run.status === "completed"), + ); + const targetRun = runs.find((run) => run.ordinal === targetOrdinal); + const attempts = projection.attempts.filter( + (attempt) => attempt.id === targetRun?.activeAttemptId, + ); + const attempt = attempts[0]; + const nodeIds = new Set( + runs + .filter((run) => run.ordinal > targetOrdinal && run.status === "completed") + .map((run) => run.rootNodeId), + ); + return { + thread, + providerThreads, + checkpointScopes, + runs, + attempts, + nodes: projection.nodes.filter((node) => nodeIds.has(node.id)), + checkpoints: [ + ...checkpoints, + ...projection.checkpoints.filter( + (candidate) => + candidate.scopeId === scope.id && + candidate.status === "ready" && + candidate.appRunOrdinal !== null && + candidate.appRunOrdinal > targetOrdinal && + candidate.id !== checkpoint.id, + ), + ], + providerSessions: projection.providerSessions.filter( + (session) => session.id === providerThread.providerSessionId, + ), + providerTurns: projection.providerTurns.filter( + (turn) => + turn.providerThreadId === providerThread.id || + turn.id === attempt?.providerTurnId || + (attempt !== undefined && turn.runAttemptId === attempt.id), + ), + }; + }), getCheckpointContext: (threadId) => Effect.gen(function* () { const projection = (yield* Ref.get(replayState)).projections.get(threadId); diff --git a/apps/server/src/orchestration-v2/ProviderTurnControlService.test.ts b/apps/server/src/orchestration-v2/ProviderTurnControlService.test.ts index 839a5992ef65..bbe89b5d4e3d 100644 --- a/apps/server/src/orchestration-v2/ProviderTurnControlService.test.ts +++ b/apps/server/src/orchestration-v2/ProviderTurnControlService.test.ts @@ -241,6 +241,8 @@ it.effect( })), ), getCheckpointContext: () => Effect.die("not used"), + getCheckpointCaptureContext: () => Effect.die("not used"), + getCheckpointRollbackContext: () => Effect.die("not used"), getRecoveryThreadIds: () => Effect.die("unused getRecoveryThreadIds"), getUnreadableThreadIds: () => Effect.die("unused getUnreadableThreadIds"), getThreadSnapshot: () => Effect.die("unused getThreadSnapshot"), diff --git a/apps/server/src/orchestration-v2/RunFinalizationService.test.ts b/apps/server/src/orchestration-v2/RunFinalizationService.test.ts index cdc330f3c9ff..e162626e8a09 100644 --- a/apps/server/src/orchestration-v2/RunFinalizationService.test.ts +++ b/apps/server/src/orchestration-v2/RunFinalizationService.test.ts @@ -4,7 +4,6 @@ import { RunId, ThreadId, type OrchestrationV2ThreadShell, - type OrchestrationV2ThreadProjection, } from "@t3tools/contracts"; import * as Effect from "effect/Effect"; import * as Layer from "effect/Layer"; @@ -23,14 +22,17 @@ it.effect("captures the root checkpoint and refreshes workspace state", () => { const capture = vi.fn(() => Effect.void); const refresh = vi.fn(() => Effect.void); const projection = { - checkpointScopes: [{ id: scopeId, cwd: "/repo" }], - } as unknown as OrchestrationV2ThreadProjection; + runs: [], + checkpoints: [], + checkpointScopes: [{ id: scopeId, runId, kind: "root_run" as const, cwd: "/repo" }], + }; const layer = RunFinalization.layer.pipe( Layer.provide( Layer.mergeAll( Layer.mock(CheckpointCapture.CheckpointCaptureServiceV2)({ execute: capture }), Layer.mock(ProjectionStore.ProjectionStoreV2)({ - getThreadProjection: () => Effect.succeed(projection), + getCheckpointContext: () => Effect.succeed(projection), + getThreadProjection: () => Effect.die("finalization must not load transcript"), }), Layer.succeed(RunFinalization.RunFinalizationObserver, { refresh, diff --git a/apps/server/src/orchestration-v2/RunFinalizationService.ts b/apps/server/src/orchestration-v2/RunFinalizationService.ts index 13c9fdd386c3..d4a5dc342005 100644 --- a/apps/server/src/orchestration-v2/RunFinalizationService.ts +++ b/apps/server/src/orchestration-v2/RunFinalizationService.ts @@ -64,7 +64,7 @@ const make = Effect.gen(function* () { ), ); const projection = yield* projections - .getThreadProjection(input.threadId) + .getCheckpointContext(input.threadId) .pipe( Effect.mapError( (cause) => new RunFinalizationError({ ...input, operation: "refresh-workspace", cause }),