diff --git a/apps/server/src/orchestration-v2/CheckpointRestoreSafety.test.ts b/apps/server/src/orchestration-v2/CheckpointRestoreSafety.test.ts index 5f79aefdd0b7bd598851f18a5acc5098f6415d35..34bf5cce6a5375b37a3d48f4285ee7fc4bc9d273 100644 --- a/apps/server/src/orchestration-v2/CheckpointRestoreSafety.test.ts +++ b/apps/server/src/orchestration-v2/CheckpointRestoreSafety.test.ts @@ -17,6 +17,7 @@ import * as Option from "effect/Option"; import * as Path from "effect/Path"; import { symlinksSupported } from "@t3tools/shared/testing/symlinks"; import * as CheckpointRollbackService from "./CheckpointRollbackService.ts"; +import * as ThreadCommandExecutor from "./ThreadCommandExecutor.ts"; import { CheckpointServiceV2 } from "./CheckpointService.ts"; import { EventSinkV2 } from "./EventSink.ts"; import * as IdAllocator from "./IdAllocator.ts"; @@ -102,6 +103,7 @@ it.effect.each([ Layer.provide( Layer.mergeAll( NodeServices.layer, + ThreadCommandExecutor.layer, IdAllocator.layer, Layer.mock(ProjectStore.ProjectStoreV2)({ get: () => Effect.succeed(Option.some({ workspaceRoot: parent } as never)), diff --git a/apps/server/src/orchestration-v2/CheckpointRollbackService.test.ts b/apps/server/src/orchestration-v2/CheckpointRollbackService.test.ts index 433a3be43ac3cd07ec094b4deb0628387c9fb758..bb37290363ac7f59690b61caf2f669f3905f9ffe 100644 --- a/apps/server/src/orchestration-v2/CheckpointRollbackService.test.ts +++ b/apps/server/src/orchestration-v2/CheckpointRollbackService.test.ts @@ -20,6 +20,8 @@ import { resolveCodexRollbackTurnCount } from "./Adapters/CodexAdapterV2.ts"; import { isCheckpointRestoreIsolated } from "./CheckpointRestoreSafety.ts"; import * as CheckpointService from "./CheckpointService.ts"; import * as CheckpointRollbackService from "./CheckpointRollbackService.ts"; +import * as ThreadCommandExecutor from "./ThreadCommandExecutor.ts"; +import * as KeyedLock from "@t3tools/shared/KeyedLock"; import * as EventSink from "./EventSink.ts"; import * as IdAllocator from "./IdAllocator.ts"; import * as ProjectionStore from "./ProjectionStore.ts"; @@ -64,6 +66,7 @@ it.effect("rejects a non-ready checkpoint before opening a session or restoring const layerTest = layerCheckpointRollbackService.pipe( Layer.provide( Layer.mergeAll( + ThreadCommandExecutor.layer, Layer.mock(CheckpointService.CheckpointServiceV2)({ restore }), Layer.mock(EventSink.EventSinkV2)({}), IdAllocator.layer, @@ -141,6 +144,7 @@ it.effect("rejects a rollback when another provider thread became active", () => const layerTest = layerCheckpointRollbackService.pipe( Layer.provide( Layer.mergeAll( + ThreadCommandExecutor.layer, Layer.mock(CheckpointService.CheckpointServiceV2)({ restore }), Layer.mock(EventSink.EventSinkV2)({}), IdAllocator.layer, @@ -221,6 +225,7 @@ it.effect("rejects a rollback when provider selection changed before execution", const layerTest = layerCheckpointRollbackService.pipe( Layer.provide( Layer.mergeAll( + ThreadCommandExecutor.layer, Layer.mock(CheckpointService.CheckpointServiceV2)({ restore }), Layer.mock(EventSink.EventSinkV2)({}), IdAllocator.layer, @@ -292,6 +297,7 @@ it.effect("reports a missing provider turn as a structured rollback failure", () const layerTest = layerCheckpointRollbackService.pipe( Layer.provide( Layer.mergeAll( + ThreadCommandExecutor.layer, Layer.mock(CheckpointService.CheckpointServiceV2)({ restore }), Layer.mock(EventSink.EventSinkV2)({}), IdAllocator.layer, @@ -352,6 +358,15 @@ it.effect.each([ const checkpointId = CheckpointId.make("rewind-start"); const scopeId = CheckpointScopeId.make("rewind-scope"); const calls: string[] = []; + let lockedThreads = Effect.succeed>([]); + const layerThreadCommands = Layer.effect( + ThreadCommandExecutor.ThreadCommandExecutor, + Effect.tap(KeyedLock.make(), (lock) => + Effect.sync(() => { + lockedThreads = lock.activeKeys; + }), + ), + ); const providerThread = { id: providerThreadId, providerSessionId, @@ -390,6 +405,7 @@ it.effect.each([ const layerTest = layerCheckpointRollbackService.pipe( Layer.provide( Layer.mergeAll( + layerThreadCommands, Layer.mock(CheckpointService.CheckpointServiceV2)({ restore: () => Effect.sync(() => { @@ -398,7 +414,10 @@ it.effect.each([ }), Layer.mock(EventSink.EventSinkV2)({ write: ({ events }) => - Effect.sync(() => { + Effect.map(lockedThreads, (locked) => { + // A thread command that planned against the old provider thread + // must not commit after this and undo the rollback. + assert.include(locked, threadId); assert.ok( events.some( (event) => event.type === "run.updated" && event.payload.status === "rolled_back", diff --git a/apps/server/src/orchestration-v2/CheckpointRollbackService.ts b/apps/server/src/orchestration-v2/CheckpointRollbackService.ts index 1710e7a8cd844336f69feab4e2a99c7da12a5aba..b1532a74c6dc7d7e12cf7c8da1f1c1f3549568cf 100644 --- a/apps/server/src/orchestration-v2/CheckpointRollbackService.ts +++ b/apps/server/src/orchestration-v2/CheckpointRollbackService.ts @@ -26,6 +26,7 @@ import * as ProjectStore from "./ProjectStore.ts"; import type { ProviderAdapterV2RollbackTarget } from "./ProviderAdapter.ts"; import { ProviderSessionManagerV2 } from "./ProviderSessionManager.ts"; import { RuntimePolicyV2 } from "./RuntimePolicy.ts"; +import * as ThreadCommandExecutor from "./ThreadCommandExecutor.ts"; export const ROLLBACK_FAILED_MESSAGE = "The provider could not roll back this conversation. Try again; if it keeps failing, check the provider and server logs."; @@ -91,6 +92,7 @@ export const layer: Layer.Layer< | FileSystem.FileSystem | Path.Path | ProjectStore.ProjectStoreV2 + | ThreadCommandExecutor.ThreadCommandExecutor > = Layer.effect( CheckpointRollbackServiceV2, Effect.gen(function* () { @@ -103,6 +105,7 @@ export const layer: Layer.Layer< const fileSystem = yield* FileSystem.FileSystem; const projects = yield* ProjectStore.ProjectStoreV2; const path = yield* Path.Path; + const threadCommands = yield* ThreadCommandExecutor.ThreadCommandExecutor; const execute = Effect.fn("orchestrationV2.checkpointRollback.execute")(function* (input: { readonly threadId: ThreadId; @@ -274,72 +277,80 @@ export const layer: Layer.Layer< yield* checkpoints.deleteStaleRefs({ scope, checkpoints: staleCheckpoints }); } - const now = yield* DateTime.now; - const makeEvent = (event: Omit) => - Effect.map( - ids.allocate.event({ threadId: event.threadId }), - (id) => - ({ - ...event, - id, - }) as Event, - ); - const events: Array = []; - events.push( - yield* makeEvent({ - type: "provider-thread.updated", - threadId: input.threadId, - driver: providerThread.driver, - providerInstanceId: providerThread.providerInstanceId, - occurredAt: now, - payload: { - ...snapshot.providerThread, - lastRunOrdinal: targetOrdinal === 0 ? null : targetOrdinal, - updatedAt: now, - }, - }), - ); - for (const staleCheckpoint of staleCheckpoints) { - events.push( - yield* makeEvent({ - type: "checkpoint.captured", - threadId: input.threadId, - ...(staleCheckpoint.runId === null ? {} : { runId: staleCheckpoint.runId }), - nodeId: staleCheckpoint.nodeId, - providerInstanceId: providerThread.providerInstanceId, - occurredAt: now, - payload: { ...staleCheckpoint, status: "stale" }, - }), - ); - } - for (const run of runsToRollback) { - const rootNode = projection.nodes.find((candidate) => candidate.id === run.rootNodeId); - events.push( - yield* makeEvent({ - type: "run.updated", - threadId: input.threadId, - runId: run.id, - ...(rootNode === undefined ? {} : { nodeId: rootNode.id }), - providerInstanceId: run.providerInstanceId, - occurredAt: now, - payload: { ...run, status: "rolled_back", completedAt: now }, - }), - ); - if (rootNode !== undefined) { + // Thread commands write full provider thread rows under this lock. Without + // it, a message sent during the rollback can plan against the old native + // session and commit after this, undoing the provider's fork. + yield* threadCommands.withLock( + input.threadId, + Effect.gen(function* () { + const now = yield* DateTime.now; + const makeEvent = (event: Omit) => + Effect.map( + ids.allocate.event({ threadId: event.threadId }), + (id) => + ({ + ...event, + id, + }) as Event, + ); + const events: Array = []; events.push( yield* makeEvent({ - type: "node.updated", + type: "provider-thread.updated", threadId: input.threadId, - runId: run.id, - nodeId: rootNode.id, - providerInstanceId: run.providerInstanceId, + driver: providerThread.driver, + providerInstanceId: providerThread.providerInstanceId, occurredAt: now, - payload: { ...rootNode, status: "rolled_back", completedAt: now }, + payload: { + ...snapshot.providerThread, + lastRunOrdinal: targetOrdinal === 0 ? null : targetOrdinal, + updatedAt: now, + }, }), ); - } - } - yield* eventSink.write({ events }); + for (const staleCheckpoint of staleCheckpoints) { + events.push( + yield* makeEvent({ + type: "checkpoint.captured", + threadId: input.threadId, + ...(staleCheckpoint.runId === null ? {} : { runId: staleCheckpoint.runId }), + nodeId: staleCheckpoint.nodeId, + providerInstanceId: providerThread.providerInstanceId, + occurredAt: now, + payload: { ...staleCheckpoint, status: "stale" }, + }), + ); + } + for (const run of runsToRollback) { + const rootNode = projection.nodes.find((candidate) => candidate.id === run.rootNodeId); + events.push( + yield* makeEvent({ + type: "run.updated", + threadId: input.threadId, + runId: run.id, + ...(rootNode === undefined ? {} : { nodeId: rootNode.id }), + providerInstanceId: run.providerInstanceId, + occurredAt: now, + payload: { ...run, status: "rolled_back", completedAt: now }, + }), + ); + if (rootNode !== undefined) { + events.push( + yield* makeEvent({ + type: "node.updated", + threadId: input.threadId, + runId: run.id, + nodeId: rootNode.id, + providerInstanceId: run.providerInstanceId, + occurredAt: now, + payload: { ...rootNode, status: "rolled_back", completedAt: now }, + }), + ); + } + } + yield* eventSink.write({ events }); + }), + ); }); return CheckpointRollbackServiceV2.of({ diff --git a/apps/server/src/orchestration-v2/runtimeLayer.ts b/apps/server/src/orchestration-v2/runtimeLayer.ts index 98e05de9a5a264ab997d1b53027bf8c03753920d..9cf3351f5d252200f9fc273bc2b3cc87816e6043 100644 --- a/apps/server/src/orchestration-v2/runtimeLayer.ts +++ b/apps/server/src/orchestration-v2/runtimeLayer.ts @@ -182,6 +182,7 @@ const layerCheckpointRollbackServiceProvided = CheckpointRollbackService.layer.p IdAllocator.layer, ProjectionStore.layer, layerProviderSessionManagerProvided, + ThreadCommandExecutor.layer, layerRuntimePolicyProvided, ), ), diff --git a/apps/server/src/orchestration-v2/testkit/ProviderReplayHarness.ts b/apps/server/src/orchestration-v2/testkit/ProviderReplayHarness.ts index 83762ef06ceabb1baa9393a186c026cc88917355..230ad7aa0feb8e8fcc5c2db14a5e98d711610f6f 100644 --- a/apps/server/src/orchestration-v2/testkit/ProviderReplayHarness.ts +++ b/apps/server/src/orchestration-v2/testkit/ProviderReplayHarness.ts @@ -395,6 +395,7 @@ export function layerWithRegistry( IdAllocator.layer, layerStores, layerProviderSessionManagerProvided, + ThreadCommandExecutor.layer, layerRuntime, ), ),