diff --git a/apps/server/src/checkpointing/CheckpointStore.test.ts b/apps/server/src/checkpointing/CheckpointStore.test.ts index bf332d20d0da..2630fd929b17 100644 --- a/apps/server/src/checkpointing/CheckpointStore.test.ts +++ b/apps/server/src/checkpointing/CheckpointStore.test.ts @@ -114,6 +114,76 @@ it.layer(TestLayer)("CheckpointStore.layer", (it) => { ); }); + describe("readWorkspaceFingerprint", () => { + it.effect("tracks the exact tracked and untracked tree affected by restore", () => + Effect.gen(function* () { + const tmp = yield* makeTmpDir(); + yield* initRepoWithCommit(tmp); + const checkpointStore = yield* CheckpointStore.CheckpointStore; + + const initial = yield* checkpointStore.readWorkspaceFingerprint(tmp); + yield* writeTextFile(NodePath.join(tmp, "README.md"), "# changed\n"); + const trackedChange = yield* checkpointStore.readWorkspaceFingerprint(tmp); + yield* writeTextFile(NodePath.join(tmp, "new-file.txt"), "new\n"); + const untrackedChange = yield* checkpointStore.readWorkspaceFingerprint(tmp); + + expect(trackedChange).not.toBe(initial); + expect(untrackedChange).not.toBe(trackedChange); + expect(yield* git(tmp, ["status", "--short"])).toContain("README.md"); + expect(yield* git(tmp, ["status", "--short"])).toContain("new-file.txt"); + }), + ); + + it.effect("tracks staged-only changes without modifying the user's index", () => + Effect.gen(function* () { + const tmp = yield* makeTmpDir(); + yield* initRepoWithCommit(tmp); + const checkpointStore = yield* CheckpointStore.CheckpointStore; + const readme = NodePath.join(tmp, "README.md"); + + const initial = yield* checkpointStore.readWorkspaceFingerprint(tmp); + yield* writeTextFile(readme, "# staged\n"); + yield* git(tmp, ["add", "README.md"]); + yield* writeTextFile(readme, "# test\n"); + const stagedBefore = yield* git(tmp, ["diff", "--cached", "--", "README.md"]); + const stagedOnly = yield* checkpointStore.readWorkspaceFingerprint(tmp); + const stagedAfter = yield* git(tmp, ["diff", "--cached", "--", "README.md"]); + + expect(stagedOnly).not.toBe(initial); + expect(stagedAfter).toBe(stagedBefore); + expect(yield* git(tmp, ["diff", "--", "README.md"])).not.toBe(""); + }), + ); + + it.effect("limits staged-state fingerprints to a nested restore cwd", () => + Effect.gen(function* () { + const tmp = yield* makeTmpDir(); + yield* initRepoWithCommit(tmp); + const nested = NodePath.join(tmp, "packages", "nested"); + const sibling = NodePath.join(tmp, "sibling.txt"); + const nestedFile = NodePath.join(nested, "file.txt"); + const fileSystem = yield* FileSystem.FileSystem; + yield* fileSystem.makeDirectory(nested, { recursive: true }); + yield* writeTextFile(nestedFile, "nested\n"); + yield* writeTextFile(sibling, "sibling\n"); + yield* git(tmp, ["add", "."]); + yield* git(tmp, ["commit", "-m", "nested files"]); + const checkpointStore = yield* CheckpointStore.CheckpointStore; + + const initial = yield* checkpointStore.readWorkspaceFingerprint(nested); + yield* writeTextFile(sibling, "sibling staged\n"); + yield* git(tmp, ["add", "sibling.txt"]); + yield* writeTextFile(sibling, "sibling\n"); + expect(yield* checkpointStore.readWorkspaceFingerprint(nested)).toBe(initial); + + yield* writeTextFile(nestedFile, "nested staged\n"); + yield* git(tmp, ["add", "packages/nested/file.txt"]); + yield* writeTextFile(nestedFile, "nested\n"); + expect(yield* checkpointStore.readWorkspaceFingerprint(nested)).not.toBe(initial); + }), + ); + }); + describe("diffCheckpoints", () => { it.effect("returns full oversized checkpoint diffs without truncation", () => Effect.gen(function* () { diff --git a/apps/server/src/checkpointing/CheckpointStore.ts b/apps/server/src/checkpointing/CheckpointStore.ts index f13aa4572c17..7c160c7fa4ad 100644 --- a/apps/server/src/checkpointing/CheckpointStore.ts +++ b/apps/server/src/checkpointing/CheckpointStore.ts @@ -53,6 +53,9 @@ export class CheckpointStore extends Context.Service< /** Check whether cwd is inside a Git worktree. */ readonly isGitRepository: (cwd: string) => Effect.Effect; + /** Hash the tracked and untracked workspace state affected by restore. */ + readonly readWorkspaceFingerprint: (cwd: string) => Effect.Effect; + /** * Capture a checkpoint commit and store it at the provided checkpoint ref. * @@ -119,6 +122,15 @@ export const make = Effect.gen(function* () { .detect({ cwd, requestedKind: "git" }) .pipe(Effect.map((repository) => repository !== null)); + const readWorkspaceFingerprint: CheckpointStore["Service"]["readWorkspaceFingerprint"] = + Effect.fn("readWorkspaceFingerprint")(function* (cwd) { + const checkpoints = yield* resolveCheckpoints( + "CheckpointStore.readWorkspaceFingerprint", + cwd, + ); + return yield* checkpoints.readWorkspaceFingerprint(cwd); + }); + const captureCheckpoint: CheckpointStore["Service"]["captureCheckpoint"] = Effect.fn( "captureCheckpoint", )(function* (input) { @@ -159,6 +171,7 @@ export const make = Effect.gen(function* () { return CheckpointStore.of({ isGitRepository, + readWorkspaceFingerprint, captureCheckpoint, hasCheckpointRef, restoreCheckpoint, diff --git a/apps/server/src/mcp/CheckpointMcpRestore.integration.test.ts b/apps/server/src/mcp/CheckpointMcpRestore.integration.test.ts new file mode 100644 index 000000000000..637b2fb5069b --- /dev/null +++ b/apps/server/src/mcp/CheckpointMcpRestore.integration.test.ts @@ -0,0 +1,666 @@ +// @effect-diagnostics nodeBuiltinImport:off +import * as NodePath from "node:path"; + +import * as NodeServices from "@effect/platform-node/NodeServices"; +import { assert, it } from "@effect/vitest"; +import { + CheckpointId, + CheckpointRef, + CheckpointScopeId, + CommandId, + EnvironmentId, + EventId, + MessageId, + NodeId, + type OrchestrationV2DomainEvent, + type OrchestrationV2ProviderSession, + type OrchestrationV2ProviderThread, + type OrchestrationV2Run, + ProjectId, + ProviderDriverKind, + ProviderInstanceId, + ProviderSessionId, + ProviderThreadId, + RunId, + ThreadId, +} from "@t3tools/contracts"; +import * as DateTime from "effect/DateTime"; +import * as Deferred from "effect/Deferred"; +import * as Effect from "effect/Effect"; +import * as FileSystem from "effect/FileSystem"; +import * as Fiber from "effect/Fiber"; +import * as Layer from "effect/Layer"; +import * as Ref from "effect/Ref"; +import * as Schema from "effect/Schema"; +import * as Stream from "effect/Stream"; + +import * as CheckpointStore from "../checkpointing/CheckpointStore.ts"; +import { ServerConfig } from "../config.ts"; +import * as McpSessionRegistryTestkit from "./McpSessionRegistry.testkit.ts"; +import { ProviderInstanceRegistry } from "../provider/Services/ProviderInstanceRegistry.ts"; +import type { ProviderInstance } from "../provider/ProviderDriver.ts"; +import { SqlitePersistenceMemory } from "../persistence/Layers/Sqlite.ts"; +import { ServerSettingsService } from "../serverSettings.ts"; +import * as VcsDriverRegistry from "../vcs/VcsDriverRegistry.ts"; +import * as VcsProcess from "../vcs/VcsProcess.ts"; +import { CodexProviderCapabilitiesV2 } from "../orchestration-v2/Adapters/CodexAdapterV2.ts"; +import { OrchestrationEffectWorkerV2 } from "../orchestration-v2/EffectWorker.ts"; +import { EventSinkV2 } from "../orchestration-v2/EventSink.ts"; +import { + ThreadDispatchLockV2, + threadDispatchLockLayer, +} from "../orchestration-v2/KeyedSerialExecutor.ts"; +import { + OrchestratorCheckpointRollbackTargetUnsupportedError, + OrchestratorV2, +} from "../orchestration-v2/Orchestrator.ts"; +import { + ProviderAdapterRollbackThreadError, + type ProviderAdapterV2SessionRuntime, + type ProviderAdapterV2Shape, +} from "../orchestration-v2/ProviderAdapter.ts"; +import { + OrchestrationV2EventSinkLayerLive, + OrchestrationV2LayerLive, +} from "../orchestration-v2/runtimeLayer.ts"; +import { checkpointWorkspace } from "../orchestration-v2/testkit/ReplayFixtureWorkspace.ts"; +import * as CheckpointMcp from "./CheckpointMcpService.ts"; +import type { McpInvocationScope } from "./McpInvocationContext.ts"; + +const driver = ProviderDriverKind.make("codex"); +const providerInstanceId = ProviderInstanceId.make("codex-checkpoint-restore-test"); +const modelSelection = { instanceId: providerInstanceId, model: "gpt-test" }; +const isRollbackTargetUnsupported = Schema.is(OrchestratorCheckpointRollbackTargetUnsupportedError); +const ServerConfigLayer = ServerConfig.layerTest(process.cwd(), { + prefix: "t3-checkpoint-mcp-restore-", +}); + +function makeAdapter(input: { + readonly rollbackCount: Ref.Ref; + readonly failProviderRollback: boolean; +}): ProviderAdapterV2Shape { + return { + instanceId: providerInstanceId, + driver, + getCapabilities: () => Effect.succeed(CodexProviderCapabilitiesV2), + planSelectionTransition: () => Effect.succeed({ type: "apply_on_next_turn" }), + openSession: (openInput) => + Effect.gen(function* () { + const now = yield* DateTime.now; + const providerSession: OrchestrationV2ProviderSession = { + id: openInput.providerSessionId, + driver, + providerInstanceId, + status: "ready", + cwd: openInput.runtimePolicy.cwd ?? "/repo", + model: openInput.modelSelection.model, + capabilities: CodexProviderCapabilitiesV2, + createdAt: now, + updatedAt: now, + lastError: null, + }; + return { + instanceId: providerInstanceId, + driver, + providerSessionId: openInput.providerSessionId, + providerSession, + events: Stream.empty, + ensureThread: () => Effect.die("ensureThread is unused in checkpoint restore"), + resumeThread: ({ providerThread }) => Effect.succeed(providerThread), + startTurn: () => Effect.die("startTurn is unused in checkpoint restore"), + steerTurn: () => Effect.die("steerTurn is unused in checkpoint restore"), + interruptTurn: () => Effect.die("interruptTurn is unused in checkpoint restore"), + respondToRuntimeRequest: () => + Effect.die("respondToRuntimeRequest is unused in checkpoint restore"), + readThreadSnapshot: () => + Effect.die("readThreadSnapshot is unused in checkpoint restore"), + rollbackThread: ({ providerThread, target }) => + Ref.update(input.rollbackCount, (count) => count + 1).pipe( + Effect.andThen( + input.failProviderRollback + ? Effect.fail( + new ProviderAdapterRollbackThreadError({ + driver, + providerThreadId: providerThread.id, + checkpointId: target.checkpointId, + cause: "simulated provider rollback failure", + }), + ) + : Effect.succeed({ + providerThread, + providerTurns: [], + messages: [], + runtimeRequests: [], + }), + ), + ), + forkThread: () => Effect.die("forkThread is unused in checkpoint restore"), + } satisfies ProviderAdapterV2SessionRuntime; + }), + }; +} + +function makeIntegrationLayer(input: { + readonly rollbackCount: Ref.Ref; + readonly failProviderRollback: boolean; + readonly fingerprintGate?: { + readonly entered: Deferred.Deferred; + readonly release: Deferred.Deferred; + }; +}) { + const adapter = makeAdapter(input); + const providerInstance = { + instanceId: providerInstanceId, + driverKind: driver, + continuationIdentity: { driverKind: driver, continuationKey: "codex:checkpoint-restore" }, + displayName: "Checkpoint restore provider", + enabled: true, + snapshot: {} as ProviderInstance["snapshot"], + orchestrationAdapter: adapter, + textGeneration: {} as ProviderInstance["textGeneration"], + } satisfies ProviderInstance; + const providerRegistry = Layer.succeed(ProviderInstanceRegistry, { + getInstance: (instanceId) => + Effect.succeed(instanceId === providerInstance.instanceId ? providerInstance : undefined), + listInstances: Effect.succeed([providerInstance]), + listUnavailable: Effect.succeed([]), + streamChanges: Stream.empty, + subscribeChanges: Effect.never, + }); + const vcsRegistry = VcsDriverRegistry.layer.pipe( + Layer.provide(VcsProcess.layer), + Layer.provide(ServerConfigLayer), + Layer.provide(NodeServices.layer), + ); + const checkpointStoreLive = CheckpointStore.layer.pipe(Layer.provide(vcsRegistry)); + const fingerprintGate = input.fingerprintGate; + const checkpointStore = + fingerprintGate === undefined + ? checkpointStoreLive + : Layer.effect( + CheckpointStore.CheckpointStore, + Effect.gen(function* () { + const store = yield* CheckpointStore.CheckpointStore; + let fingerprintReads = 0; + return CheckpointStore.CheckpointStore.of({ + ...store, + readWorkspaceFingerprint: (cwd) => + store.readWorkspaceFingerprint(cwd).pipe( + Effect.tap(() => { + fingerprintReads += 1; + return fingerprintReads === 2 + ? Deferred.succeed(fingerprintGate.entered, undefined).pipe( + Effect.andThen(Deferred.await(fingerprintGate.release)), + ) + : Effect.void; + }), + ), + }); + }), + ).pipe(Layer.provide(checkpointStoreLive)); + const runtime = Layer.merge(OrchestrationV2LayerLive, OrchestrationV2EventSinkLayerLive).pipe( + Layer.provide(McpSessionRegistryTestkit.layer), + Layer.provide(SqlitePersistenceMemory), + Layer.provideMerge(checkpointStore), + Layer.provide(ServerConfigLayer), + Layer.provide(ServerSettingsService.layerTest()), + Layer.provide(providerRegistry), + Layer.provide(NodeServices.layer), + ); + const runtimeWithDispatchLock = Layer.merge(runtime, threadDispatchLockLayer); + return CheckpointMcp.layer.pipe( + Layer.provideMerge(runtimeWithDispatchLock), + Layer.provideMerge(NodeServices.layer), + ); +} + +function seedRollbackProjection(input: { + readonly eventSink: EventSinkV2["Service"]; + readonly threadId: ThreadId; + readonly providerSessionId: ProviderSessionId; + readonly providerThreadId: ProviderThreadId; + readonly runId: RunId; + readonly scopeId: CheckpointScopeId; + readonly checkpointId: CheckpointId; + readonly checkpointRef: CheckpointRef; + readonly cwd: string; + readonly checkpointOrdinal?: number; +}) { + return Effect.gen(function* () { + const now = yield* DateTime.now; + const providerSession: OrchestrationV2ProviderSession = { + id: input.providerSessionId, + driver, + providerInstanceId, + status: "ready", + cwd: input.cwd, + model: modelSelection.model, + capabilities: CodexProviderCapabilitiesV2, + createdAt: now, + updatedAt: now, + lastError: null, + }; + const providerThread: OrchestrationV2ProviderThread = { + id: input.providerThreadId, + driver, + providerInstanceId, + providerSessionId: input.providerSessionId, + appThreadId: input.threadId, + ownerNodeId: null, + nativeThreadRef: { driver, nativeId: "native-checkpoint-thread", strength: "strong" }, + nativeConversationHeadRef: null, + status: "idle", + firstRunOrdinal: 1, + lastRunOrdinal: 1, + handoffIds: [], + forkedFrom: null, + pendingBackgroundTasks: [], + createdAt: now, + updatedAt: now, + }; + const completedRun: OrchestrationV2Run = { + id: input.runId, + threadId: input.threadId, + ordinal: 1, + providerInstanceId, + modelSelection, + providerThreadId: input.providerThreadId, + userMessageId: MessageId.make("message:checkpoint-restore:completed"), + rootNodeId: null, + activeAttemptId: null, + status: "completed", + queuePosition: null, + requestedAt: now, + startedAt: now, + completedAt: now, + checkpointId: null, + contextHandoffId: null, + }; + const scopeNodeId = NodeId.make("node:checkpoint-restore:scope"); + const events: ReadonlyArray = [ + { + id: EventId.make("event:checkpoint-restore:provider-session"), + type: "provider-session.attached", + threadId: input.threadId, + providerInstanceId, + occurredAt: now, + payload: providerSession, + }, + { + id: EventId.make("event:checkpoint-restore:provider-thread"), + type: "provider-thread.updated", + threadId: input.threadId, + providerInstanceId, + occurredAt: now, + payload: providerThread, + }, + { + id: EventId.make("event:checkpoint-restore:completed-run"), + type: "run.created", + threadId: input.threadId, + runId: input.runId, + providerInstanceId, + occurredAt: now, + payload: completedRun, + }, + { + id: EventId.make("event:checkpoint-restore:scope"), + type: "checkpoint-scope.created", + threadId: input.threadId, + nodeId: scopeNodeId, + providerInstanceId, + occurredAt: now, + payload: { + id: input.scopeId, + threadId: input.threadId, + runId: null, + nodeId: scopeNodeId, + parentScopeId: null, + providerThreadId: input.providerThreadId, + kind: "root_run", + ordinalWithinParent: 0, + advancesAppRunCount: true, + cwd: input.cwd, + createdAt: now, + }, + }, + { + id: EventId.make("event:checkpoint-restore:checkpoint"), + type: "checkpoint.captured", + threadId: input.threadId, + nodeId: scopeNodeId, + providerInstanceId, + occurredAt: now, + payload: { + id: input.checkpointId, + threadId: input.threadId, + scopeId: input.scopeId, + runId: null, + nodeId: scopeNodeId, + parentCheckpointId: null, + ordinalWithinScope: input.checkpointOrdinal ?? 0, + appRunOrdinal: null, + ref: input.checkpointRef, + status: "ready", + files: [], + capturedAt: now, + }, + }, + ]; + yield* input.eventSink.write({ events }); + }); +} + +function restoreScenario( + input: { + readonly failProviderRollback: boolean; + readonly admitRunBeforeWorker?: boolean; + readonly raceMessageDuringRestore?: boolean; + readonly fingerprintGate?: { + readonly entered: Deferred.Deferred; + readonly release: Deferred.Deferred; + }; + }, + rollbackCount: Ref.Ref, +) { + return Effect.scoped( + Effect.gen(function* () { + const cwd = yield* checkpointWorkspace( + input.admitRunBeforeWorker + ? "mcp-restore-concurrent-run" + : input.failProviderRollback + ? "mcp-restore-partial" + : "mcp-restore-applied", + ); + const fileSystem = yield* FileSystem.FileSystem; + const checkpointStore = yield* CheckpointStore.CheckpointStore; + const orchestrator = yield* OrchestratorV2; + const eventSink = yield* EventSinkV2; + const threadDispatch = yield* ThreadDispatchLockV2; + const worker = yield* OrchestrationEffectWorkerV2; + const service = yield* CheckpointMcp.CheckpointMcpService; + const threadId = ThreadId.make( + input.admitRunBeforeWorker + ? "thread:checkpoint-restore:concurrent-run" + : input.failProviderRollback + ? "thread:checkpoint-restore:partial" + : "thread:checkpoint-restore:applied", + ); + const providerSessionId = ProviderSessionId.make(`provider-session:${threadId}`); + const providerThreadId = ProviderThreadId.make(`provider-thread:${threadId}`); + const runId = RunId.make(`run:${threadId}:completed`); + const scopeId = CheckpointScopeId.make(`scope:${threadId}`); + const checkpointId = CheckpointId.make(`checkpoint:${threadId}`); + const checkpointRef = CheckpointRef.make( + input.admitRunBeforeWorker + ? "refs/t3/test/checkpoint-restore-concurrent-run" + : input.failProviderRollback + ? "refs/t3/test/checkpoint-restore-partial" + : "refs/t3/test/checkpoint-restore-applied", + ); + const readmePath = NodePath.join(cwd, "README.md"); + + yield* fileSystem.writeFileString(readmePath, "checkpoint contents\n"); + yield* checkpointStore.captureCheckpoint({ cwd, checkpointRef }); + yield* fileSystem.writeFileString(readmePath, "current unsaved contents\n"); + yield* fileSystem.writeFileString(NodePath.join(cwd, "untracked.txt"), "remove me\n"); + yield* orchestrator.dispatch({ + type: "thread.create", + createdBy: "user", + creationSource: "web", + commandId: CommandId.make(`command:create:${threadId}`), + threadId, + projectId: ProjectId.make("project:checkpoint-restore"), + title: "Checkpoint restore integration", + modelSelection, + runtimeMode: "full-access", + interactionMode: "default", + branch: "feature/checkpoint-restore", + worktreePath: cwd, + }); + yield* seedRollbackProjection({ + eventSink, + threadId, + providerSessionId, + providerThreadId, + runId, + scopeId, + checkpointId, + checkpointRef, + cwd, + }); + + const invocation: McpInvocationScope = { + environmentId: EnvironmentId.make("environment:checkpoint-restore"), + threadId, + providerSessionId: `mcp-session:${threadId}`, + providerInstanceId, + capabilities: new Set(["orchestration"]), + issuedAt: 1, + }; + const restoreInput = { + scopeId, + checkpointId, + discardChanges: true as const, + clientRequestId: "restore-integration-key", + }; + const requested = yield* service.restore(invocation, restoreInput); + const acceptedRetry = yield* service.restore(invocation, restoreInput); + assert.equal(requested.status, "REQUESTED"); + assert.equal(acceptedRetry.commandId, requested.commandId); + assert.lengthOf(yield* orchestrator.listCommandEffects(requested.commandId), 1); + + if (input.admitRunBeforeWorker === true) { + const now = yield* DateTime.now; + const queuedRunId = RunId.make(`run:concurrent:${threadId}`); + yield* eventSink.write({ + events: [ + { + id: EventId.make(`event:concurrent-run:${threadId}`), + type: "run.created", + threadId, + runId: queuedRunId, + providerInstanceId, + occurredAt: now, + payload: { + id: queuedRunId, + threadId, + ordinal: 2, + providerInstanceId, + modelSelection, + providerThreadId, + userMessageId: MessageId.make(`message:concurrent-run:${threadId}`), + rootNodeId: null, + activeAttemptId: null, + status: "queued", + queuePosition: 1, + requestedAt: now, + startedAt: null, + completedAt: null, + checkpointId: null, + contextHandoffId: null, + }, + }, + ], + }); + } + + if (input.raceMessageDuringRestore === true) { + if (input.fingerprintGate === undefined) { + return yield* Effect.die("fingerprint gate is required for the admission race"); + } + const completionOrder = yield* Ref.make>([]); + const workerFiber = yield* worker.runOnce.pipe( + Effect.tap(() => Ref.update(completionOrder, (order) => [...order, "restore" as const])), + Effect.forkChild, + ); + yield* Deferred.await(input.fingerprintGate.entered); + const lockProbe = yield* threadDispatch + .withLock(threadId, Effect.void) + .pipe(Effect.forkChild); + yield* Effect.yieldNow; + assert.isTrue( + lockProbe.pollUnsafe() === undefined, + "the guarded restore worker must hold the shared thread admission lock", + ); + const dispatchEntered = yield* Deferred.make(); + const messageFiber = yield* Deferred.succeed(dispatchEntered, undefined).pipe( + Effect.andThen( + orchestrator.dispatch({ + type: "message.dispatch", + createdBy: "user", + creationSource: "mcp", + commandId: CommandId.make(`command:message-during-restore:${threadId}`), + threadId, + messageId: MessageId.make(`message:during-restore:${threadId}`), + text: "Run after checkpoint restoration.", + attachments: [], + modelSelection, + dispatchMode: { type: "start_immediately" }, + }), + ), + Effect.tap(() => Ref.update(completionOrder, (order) => [...order, "message" as const])), + Effect.forkChild, + ); + yield* Deferred.await(dispatchEntered); + yield* Effect.yieldNow; + yield* Deferred.succeed(input.fingerprintGate.release, undefined); + assert.isTrue(yield* Fiber.join(workerFiber)); + yield* Fiber.join(lockProbe); + yield* Fiber.join(messageFiber); + assert.deepEqual(yield* Ref.get(completionOrder), ["restore", "message"]); + } else { + assert.isTrue(yield* worker.runOnce); + } + const settled = yield* service.restore(invocation, restoreInput); + assert.equal( + settled.status, + input.admitRunBeforeWorker ? "FAILED" : input.failProviderRollback ? "PARTIAL" : "APPLIED", + ); + assert.equal(yield* Ref.get(rollbackCount), input.admitRunBeforeWorker ? 0 : 1); + assert.equal( + yield* fileSystem.readFileString(readmePath), + input.admitRunBeforeWorker ? "current unsaved contents\n" : "checkpoint contents\n", + ); + assert.equal( + yield* fileSystem.exists(NodePath.join(cwd, "untracked.txt")), + input.admitRunBeforeWorker === true, + ); + assert.lengthOf(yield* orchestrator.listCommandEffects(requested.commandId), 1); + if (input.admitRunBeforeWorker !== true && input.raceMessageDuringRestore !== true) { + assert.isFalse(yield* worker.runOnce); + } + }), + ); +} + +it.effect("restores temporary Git state once through real V2 and reports applied", () => + Effect.gen(function* () { + const rollbackCount = yield* Ref.make(0); + return yield* restoreScenario({ failProviderRollback: false }, rollbackCount).pipe( + Effect.provide(makeIntegrationLayer({ rollbackCount, failProviderRollback: false })), + ); + }), +); + +it.effect("reports provider failure after real filesystem restore as partial without retry", () => + Effect.gen(function* () { + const rollbackCount = yield* Ref.make(0); + return yield* restoreScenario({ failProviderRollback: true }, rollbackCount).pipe( + Effect.provide(makeIntegrationLayer({ rollbackCount, failProviderRollback: true })), + ); + }), +); + +it.effect("rejects work admitted after acceptance at the worker workspace boundary", () => + Effect.gen(function* () { + const rollbackCount = yield* Ref.make(0); + return yield* restoreScenario( + { failProviderRollback: false, admitRunBeforeWorker: true }, + rollbackCount, + ).pipe(Effect.provide(makeIntegrationLayer({ rollbackCount, failProviderRollback: false }))); + }), +); + +it.effect("serializes message admission across the guarded restore critical section", () => + Effect.gen(function* () { + const rollbackCount = yield* Ref.make(0); + const fingerprintGate = { + entered: yield* Deferred.make(), + release: yield* Deferred.make(), + }; + return yield* restoreScenario( + { + failProviderRollback: false, + raceMessageDuringRestore: true, + fingerprintGate, + }, + rollbackCount, + ).pipe( + Effect.provide( + makeIntegrationLayer({ + rollbackCount, + failProviderRollback: false, + fingerprintGate, + }), + ), + ); + }), +); + +it.effect("rejects a nonzero materialized baseline in the serialized command decision", () => + Effect.gen(function* () { + const rollbackCount = yield* Ref.make(0); + return yield* Effect.scoped( + Effect.gen(function* () { + const cwd = yield* checkpointWorkspace("mcp-restore-ambiguous-baseline"); + const orchestrator = yield* OrchestratorV2; + const eventSink = yield* EventSinkV2; + const threadId = ThreadId.make("thread:checkpoint-restore:ambiguous-baseline"); + const providerSessionId = ProviderSessionId.make(`provider-session:${threadId}`); + const providerThreadId = ProviderThreadId.make(`provider-thread:${threadId}`); + const scopeId = CheckpointScopeId.make(`scope:${threadId}`); + const checkpointId = CheckpointId.make(`checkpoint:${threadId}`); + yield* orchestrator.dispatch({ + type: "thread.create", + createdBy: "user", + creationSource: "web", + commandId: CommandId.make(`command:create:${threadId}`), + threadId, + projectId: ProjectId.make("project:checkpoint-restore"), + title: "Ambiguous checkpoint baseline", + modelSelection, + runtimeMode: "full-access", + interactionMode: "default", + branch: null, + worktreePath: cwd, + }); + yield* seedRollbackProjection({ + eventSink, + threadId, + providerSessionId, + providerThreadId, + runId: RunId.make(`run:${threadId}:completed`), + scopeId, + checkpointId, + checkpointRef: CheckpointRef.make("refs/t3/test/ambiguous-baseline"), + cwd, + checkpointOrdinal: 2, + }); + + const error = yield* orchestrator + .dispatch({ + type: "checkpoint.rollback", + commandId: CommandId.make(`command:rollback:${threadId}`), + threadId, + scopeId, + checkpointId, + expectedIdle: true, + expectedWorkspaceFingerprint: "workspace-before", + }) + .pipe(Effect.flip); + assert.isTrue(isRollbackTargetUnsupported(error)); + }), + ).pipe(Effect.provide(makeIntegrationLayer({ rollbackCount, failProviderRollback: false }))); + }), +); diff --git a/apps/server/src/mcp/CheckpointMcpService.test.ts b/apps/server/src/mcp/CheckpointMcpService.test.ts index b577d391d623..ae0359332d6a 100644 --- a/apps/server/src/mcp/CheckpointMcpService.test.ts +++ b/apps/server/src/mcp/CheckpointMcpService.test.ts @@ -3,16 +3,21 @@ import { CheckpointId, CheckpointRef, CheckpointScopeId, + CommandId, EnvironmentId, + EventId, + MessageId, NodeId, type OrchestrationV2AppThread, type OrchestrationV2Checkpoint, type OrchestrationV2CheckpointScope, + type OrchestrationV2Run, type OrchestrationV2ThreadProjection, ProjectId, ProviderInstanceId, ProviderSessionId, ProviderThreadId, + RunId, ThreadId, VcsDriverKind, VcsUnsupportedOperationError, @@ -20,18 +25,21 @@ import { import * as DateTime from "effect/DateTime"; import * as Effect from "effect/Effect"; import * as Layer from "effect/Layer"; +import * as Option from "effect/Option"; import * as CheckpointStore from "../checkpointing/CheckpointStore.ts"; import { CLAUDE_PROVIDER, ClaudeProviderCapabilitiesV2, } from "../orchestration-v2/Adapters/ClaudeAdapterV2.ts"; +import type { OrchestrationEffectV2 } from "../orchestration-v2/EffectOutbox.ts"; +import { CommandReceiptStoreReadError } from "../orchestration-v2/CommandReceiptStore.ts"; +import { OrchestratorProjectionError } from "../orchestration-v2/Orchestrator.ts"; +import { ProjectionStoreReadError } from "../orchestration-v2/ProjectionStore.ts"; import { ThreadManagementService, ThreadManagementThreadNotFoundError, } from "../orchestration-v2/ThreadManagementService.ts"; -import { OrchestratorProjectionError } from "../orchestration-v2/Orchestrator.ts"; -import { ProjectionStoreReadError } from "../orchestration-v2/ProjectionStore.ts"; import { CheckpointMcpService, layer } from "./CheckpointMcpService.ts"; import type { McpInvocationScope } from "./McpInvocationContext.ts"; @@ -186,6 +194,11 @@ function makeHarness( readonly callerError?: OrchestratorProjectionError; readonly hasCheckpointRef?: CheckpointStore.CheckpointStore["Service"]["hasCheckpointRef"]; readonly diffCheckpoints?: CheckpointStore.CheckpointStore["Service"]["diffCheckpoints"]; + readonly readWorkspaceFingerprint?: CheckpointStore.CheckpointStore["Service"]["readWorkspaceFingerprint"]; + readonly effectStatus?: OrchestrationEffectV2["status"]; + readonly effectError?: string | null; + readonly effectFailureCode?: OrchestrationEffectV2["failureCode"]; + readonly receiptReadFailsAfterDispatch?: boolean; } = {}, ) { const projection = input.projection ?? makeProjection(); @@ -193,6 +206,37 @@ function makeHarness( const diffCheckpoints = vi.fn( input.diffCheckpoints ?? (() => Effect.succeed("diff --git a/file b/file")), ); + const readWorkspaceFingerprint = vi.fn( + input.readWorkspaceFingerprint ?? (() => Effect.succeed("tree:checkpoint-mcp")), + ); + let acceptedCommandId: CommandId | undefined; + const dispatch = vi.fn((command: Parameters[0]) => + command.type !== "checkpoint.rollback" + ? Effect.die("Checkpoint MCP test only dispatches checkpoint.rollback") + : Effect.sync(() => { + acceptedCommandId = command.commandId; + return { + sequence: 7, + storedEvents: [ + { + sequence: 7, + commandId: command.commandId, + event: { + id: EventId.make(`event:${command.commandId}`), + type: "checkpoint.rollback-requested" as const, + threadId: command.threadId, + occurredAt: now, + payload: { + scopeId: command.scopeId, + checkpointId: command.checkpointId, + requestedAt: now, + }, + }, + }, + ], + }; + }), + ); const serviceLayer = layer.pipe( Layer.provide( Layer.mergeAll( @@ -210,12 +254,66 @@ function makeHarness( threadId: requested, }), ), + dispatch, + getCommandReceipt: (commandId) => + input.receiptReadFailsAfterDispatch === true && acceptedCommandId !== undefined + ? Effect.fail( + new CommandReceiptStoreReadError({ + commandId, + cause: "simulated receipt read failure", + }), + ) + : Effect.succeed( + acceptedCommandId === commandId + ? Option.some({ + commandId, + threadId: projection.thread.id, + commandType: "checkpoint.rollback", + acceptedAt: now, + resultSequence: 7, + status: "accepted", + error: null, + }) + : Option.none(), + ), + listCommandEffects: (commandId) => + Effect.succeed([ + { + id: `effect:${commandId}:rollback`, + commandId, + threadId: projection.thread.id, + request: { + type: "provider-thread.rollback", + providerThreadId, + checkpointId: makeCheckpoint(0).id, + scopeId, + expectedIdle: true, + expectedWorkspaceFingerprint: "tree:checkpoint-mcp", + }, + status: input.effectStatus ?? "pending", + attemptCount: 0, + availableAt: "2026-08-29T12:00:00.000Z", + leaseOwner: null, + leaseExpiresAt: null, + createdAt: "2026-08-29T12:00:00.000Z", + updatedAt: "2026-08-29T12:00:00.000Z", + completedAt: null, + lastError: input.effectError ?? null, + ...(input.effectFailureCode === undefined + ? {} + : { failureCode: input.effectFailureCode }), + }, + ]), + }), + Layer.mock(CheckpointStore.CheckpointStore)({ + hasCheckpointRef, + diffCheckpoints, + readWorkspaceFingerprint, }), - Layer.mock(CheckpointStore.CheckpointStore)({ hasCheckpointRef, diffCheckpoints }), ), ), ); - return { serviceLayer, hasCheckpointRef, diffCheckpoints }; + return { serviceLayer, hasCheckpointRef, diffCheckpoints, readWorkspaceFingerprint, dispatch }; } it.effect("pages before checking checkpoint refs and bounds file summaries", () => { @@ -393,3 +491,159 @@ it.effect("distinguishes caller projection failures from missing target threads" assert.include(error.message, "Unable to load calling thread"); }).pipe(Effect.provide(harness.serviceLayer)); }); + +it.effect("accepts an exact restore and reuses the same command identity", () => { + const harness = makeHarness(); + const checkpointId = makeCheckpoint(0).id; + return Effect.gen(function* () { + const service = yield* CheckpointMcpService; + const input = { + scopeId, + checkpointId, + discardChanges: true as const, + clientRequestId: "restore-exact-target", + }; + const first = yield* service.restore(invocation, input); + const retry = yield* service.restore(invocation, input); + const conflict = yield* service + .restore(invocation, { + ...input, + checkpointId: CheckpointId.make("checkpoint:checkpoint-mcp:different"), + }) + .pipe(Effect.flip); + + assert.equal(first.status, "REQUESTED"); + assert.equal(first.effectStatus, "pending"); + assert.equal(retry.commandId, first.commandId); + assert.equal(conflict.code, "idempotency_conflict"); + assert.equal(harness.dispatch.mock.calls.length, 1); + assert.equal(harness.dispatch.mock.calls[0]?.[0].type, "checkpoint.rollback"); + assert.deepInclude(harness.dispatch.mock.calls[0]?.[0], { + expectedIdle: true, + expectedWorkspaceFingerprint: "tree:checkpoint-mcp", + }); + }).pipe(Effect.provide(harness.serviceLayer)); +}); + +it.effect("keeps the accepted outcome when post-commit observation fails", () => { + const harness = makeHarness({ receiptReadFailsAfterDispatch: true }); + return Effect.gen(function* () { + const service = yield* CheckpointMcpService; + const result = yield* service.restore(invocation, { + scopeId, + checkpointId: makeCheckpoint(0).id, + discardChanges: true, + clientRequestId: "restore-observation-unavailable", + }); + + assert.equal(result.status, "REQUESTED"); + assert.equal(result.effectStatus, "unavailable"); + assert.equal(result.receipt.sequence, 7); + assert.include(result.detail ?? "", "was accepted"); + assert.include(result.detail ?? "", "same clientRequestId"); + assert.equal(harness.dispatch.mock.calls.length, 1); + }).pipe(Effect.provide(harness.serviceLayer)); +}); + +it.effect("rejects missing checkpoints and unsupported provider rollback", () => { + const missingHarness = makeHarness({ + projection: makeProjection({ + checkpoints: [makeCheckpoint(0, { status: "missing" })], + }), + hasCheckpointRef: () => Effect.succeed(false), + }); + const supportedProjection = makeProjection(); + const unsupportedHarness = makeHarness({ + projection: { + ...supportedProjection, + providerSessions: supportedProjection.providerSessions.map((session) => ({ + ...session, + capabilities: { + ...session.capabilities, + threads: { ...session.capabilities.threads, canRollbackThread: false }, + }, + })), + }, + }); + const restoreInput = { + scopeId, + checkpointId: makeCheckpoint(0).id, + discardChanges: true as const, + clientRequestId: "restore-unavailable", + }; + return Effect.gen(function* () { + const missingService = yield* CheckpointMcpService; + const missing = yield* missingService.restore(invocation, restoreInput).pipe(Effect.flip); + assert.equal(missing.code, "checkpoint_unavailable"); + }).pipe( + Effect.provide(missingHarness.serviceLayer), + Effect.andThen( + Effect.gen(function* () { + const unsupportedService = yield* CheckpointMcpService; + const unsupported = yield* unsupportedService + .restore(invocation, { ...restoreInput, clientRequestId: "restore-unsupported" }) + .pipe(Effect.flip); + assert.equal(unsupported.code, "unsupported"); + }).pipe(Effect.provide(unsupportedHarness.serviceLayer)), + ), + ); +}); + +it.effect("rejects queued work before capturing or dispatching a restore", () => { + const queuedRun: OrchestrationV2Run = { + id: RunId.make("run:checkpoint-mcp:queued"), + threadId, + ordinal: 1, + providerInstanceId, + modelSelection: { instanceId: providerInstanceId, model: "claude-sonnet" }, + providerThreadId, + userMessageId: MessageId.make("message:checkpoint-mcp:queued"), + rootNodeId: null, + activeAttemptId: null, + status: "queued", + queuePosition: 1, + requestedAt: now, + startedAt: null, + completedAt: null, + checkpointId: null, + contextHandoffId: null, + }; + const harness = makeHarness({ + projection: { ...makeProjection(), runs: [queuedRun] }, + }); + return Effect.gen(function* () { + const service = yield* CheckpointMcpService; + const error = yield* service + .restore(invocation, { + scopeId, + checkpointId: makeCheckpoint(0).id, + discardChanges: true, + clientRequestId: "restore-while-queued", + }) + .pipe(Effect.flip); + + assert.equal(error.code, "thread_active"); + assert.equal(harness.readWorkspaceFingerprint.mock.calls.length, 0); + assert.equal(harness.dispatch.mock.calls.length, 0); + }).pipe(Effect.provide(harness.serviceLayer)); +}); + +it.effect("reports provider failure after filesystem restore as partial", () => { + const harness = makeHarness({ + effectStatus: "failed", + effectFailureCode: "checkpoint_restore_partial", + effectError: + "Provider conversation rollback failed after the filesystem checkpoint was restored; the result is partial.", + }); + return Effect.gen(function* () { + const service = yield* CheckpointMcpService; + const result = yield* service.restore(invocation, { + scopeId, + checkpointId: makeCheckpoint(0).id, + discardChanges: true, + clientRequestId: "restore-partial", + }); + assert.equal(result.status, "PARTIAL"); + assert.equal(result.effectStatus, "failed"); + }).pipe(Effect.provide(harness.serviceLayer)); +}); diff --git a/apps/server/src/mcp/CheckpointMcpService.ts b/apps/server/src/mcp/CheckpointMcpService.ts index 69cbff501f34..d8f7689dc96a 100644 --- a/apps/server/src/mcp/CheckpointMcpService.ts +++ b/apps/server/src/mcp/CheckpointMcpService.ts @@ -1,4 +1,5 @@ import { + CommandId, CheckpointMcpFailure, checkpointRollbackAppRunOrdinal, type CheckpointMcpDiffInput, @@ -6,6 +7,8 @@ import { type CheckpointMcpListInput, type CheckpointMcpListResult, type CheckpointMcpRestoreBlocker, + type CheckpointMcpRestoreInput, + type CheckpointMcpRestoreResult, type CheckpointMcpSummary, type OrchestrationV2Checkpoint, type OrchestrationV2CheckpointScope, @@ -15,10 +18,16 @@ import * as Context from "effect/Context"; import * as DateTime from "effect/DateTime"; import * as Effect from "effect/Effect"; import * as Layer from "effect/Layer"; +import * as Option from "effect/Option"; import * as Schema from "effect/Schema"; import * as CheckpointStore from "../checkpointing/CheckpointStore.ts"; -import { OrchestratorProjectionError } from "../orchestration-v2/Orchestrator.ts"; +import { + OrchestratorCheckpointRollbackNotIdleError, + OrchestratorCheckpointRollbackTargetUnsupportedError, + OrchestratorCommandIdConflictError, + OrchestratorProjectionError, +} from "../orchestration-v2/Orchestrator.ts"; import { ProjectionStoreThreadNotFoundError } from "../orchestration-v2/ProjectionStore.ts"; import { ThreadManagementProjectionLoadError, @@ -42,6 +51,10 @@ export class CheckpointMcpService extends Context.Service< scope: McpInvocationScope, input: CheckpointMcpDiffInput, ) => Effect.Effect; + readonly restore: ( + scope: McpInvocationScope, + input: CheckpointMcpRestoreInput, + ) => Effect.Effect; } >()("t3/mcp/CheckpointMcpService") {} @@ -251,6 +264,73 @@ export const make = Effect.gen(function* () { ); }); + const readAcceptedRestore = Effect.fn("CheckpointMcpService.readAcceptedRestore")(function* ( + scope: McpInvocationScope, + input: CheckpointMcpRestoreInput, + commandId: CommandId, + ) { + const receiptOption = yield* threadManagement + .getCommandReceipt(commandId) + .pipe( + Effect.mapError((error) => + failure("operation_failed", `Unable to read restore receipt: ${errorMessage(error)}`), + ), + ); + if (Option.isNone(receiptOption)) return Option.none(); + if (receiptOption.value.status !== "accepted") { + return yield* failure( + "operation_failed", + `Checkpoint restore '${commandId}' has no durable accepted receipt.`, + ); + } + const effects = yield* threadManagement + .listCommandEffects(commandId) + .pipe( + Effect.mapError((error) => + failure("operation_failed", `Unable to read restore effect: ${errorMessage(error)}`), + ), + ); + const rollbackEffect = effects.find( + (effect) => effect.request.type === "provider-thread.rollback", + ); + const intendedThreadId = input.threadId ?? scope.threadId; + if ( + rollbackEffect?.request.type !== "provider-thread.rollback" || + rollbackEffect.threadId !== intendedThreadId || + rollbackEffect.request.scopeId !== input.scopeId || + rollbackEffect.request.checkpointId !== input.checkpointId + ) { + return yield* failure( + "idempotency_conflict", + `clientRequestId '${input.clientRequestId}' was already accepted for a different checkpoint restore.`, + ); + } + const status = + rollbackEffect.status === "succeeded" + ? ("APPLIED" as const) + : rollbackEffect.status === "failed" && + rollbackEffect.failureCode === "checkpoint_restore_partial" + ? ("PARTIAL" as const) + : rollbackEffect.status === "failed" || rollbackEffect.status === "cancelled" + ? ("FAILED" as const) + : ("REQUESTED" as const); + const receipt = receiptOption.value; + return Option.some({ + commandId, + threadId: rollbackEffect.threadId, + scopeId: rollbackEffect.request.scopeId, + checkpointId: rollbackEffect.request.checkpointId, + status, + receipt: { + status: "accepted", + acceptedAt: DateTime.formatIso(receipt.acceptedAt), + sequence: receipt.resultSequence, + }, + effectStatus: rollbackEffect.status, + detail: rollbackEffect.lastError, + }); + }); + const list: CheckpointMcpService["Service"]["list"] = Effect.fn("CheckpointMcpService.list")( function* (scope, input) { yield* requireCapability(scope); @@ -454,7 +534,177 @@ export const make = Effect.gen(function* () { }, ); - return CheckpointMcpService.of({ list, diff }); + const restore: CheckpointMcpService["Service"]["restore"] = Effect.fn( + "CheckpointMcpService.restore", + )(function* (scope, input) { + yield* requireCapability(scope); + const commandId = CommandId.make( + [ + "command", + "mcp", + encodeURIComponent(scope.providerSessionId), + "checkpoint-restore", + encodeURIComponent(input.clientRequestId), + ].join(":"), + ); + const acceptedRestore = yield* readAcceptedRestore(scope, input, commandId); + if (Option.isSome(acceptedRestore)) return acceptedRestore.value; + + const projection = yield* loadThread(scope, input.threadId); + const checkpointScope = projection.checkpointScopes.find( + (candidate) => candidate.id === input.scopeId && candidate.threadId === projection.thread.id, + ); + if (checkpointScope === undefined) { + return yield* failure( + "scope_mismatch", + `Checkpoint scope '${input.scopeId}' does not belong to thread '${projection.thread.id}'.`, + ); + } + const checkpoint = projection.checkpoints.find( + (candidate) => candidate.id === input.checkpointId, + ); + if (checkpoint === undefined || checkpoint.threadId !== projection.thread.id) { + return yield* failure( + "checkpoint_not_found", + `Checkpoint '${input.checkpointId}' was not found on thread '${projection.thread.id}'.`, + ); + } + if (checkpoint.scopeId !== checkpointScope.id) { + return yield* failure( + "scope_mismatch", + `Checkpoint '${checkpoint.id}' belongs to scope '${checkpoint.scopeId}', not '${checkpointScope.id}'.`, + ); + } + const refAvailable = yield* checkpointStore + .hasCheckpointRef({ cwd: checkpointScope.cwd, checkpointRef: checkpoint.ref }) + .pipe( + Effect.mapError((error) => + failure( + "operation_failed", + `Unable to verify checkpoint availability: ${errorMessage(error)}`, + ), + ), + ); + const blockers = restoreBlockers({ + projection, + checkpoint, + scope: checkpointScope, + refAvailable, + }); + if (blockers.includes("thread_active")) { + return yield* failure( + "thread_active", + `Thread '${projection.thread.id}' must be idle with no queued runs before restore.`, + ); + } + if (blockers.length > 0) { + const unsupported = blockers.some((blocker) => + [ + "active_provider_thread_missing", + "provider_thread_mismatch", + "provider_session_missing", + "provider_rollback_unsupported", + "provider_snapshot_unsupported", + "provider_turn_missing", + "rollback_target_ambiguous", + ].includes(blocker), + ); + return yield* failure( + unsupported ? "unsupported" : "checkpoint_unavailable", + `Checkpoint '${checkpoint.id}' cannot be restored: ${blockers.join(", ")}.`, + ); + } + + const expectedWorkspaceFingerprint = yield* checkpointStore + .readWorkspaceFingerprint(checkpointScope.cwd) + .pipe( + Effect.mapError((error) => + failure( + "operation_failed", + `Unable to capture the current workspace guard: ${errorMessage(error)}`, + ), + ), + ); + const dispatched = yield* threadManagement + .dispatch({ + type: "checkpoint.rollback", + commandId, + threadId: projection.thread.id, + scopeId: checkpointScope.id, + checkpointId: checkpoint.id, + expectedIdle: true, + expectedWorkspaceFingerprint, + }) + .pipe( + Effect.catchTags({ + OrchestratorCommandIdConflictError: (error: OrchestratorCommandIdConflictError) => + failure( + "idempotency_conflict", + `Checkpoint restore was not accepted: ${error.message}`, + ), + OrchestratorCheckpointRollbackNotIdleError: ( + error: OrchestratorCheckpointRollbackNotIdleError, + ) => failure("thread_active", `Checkpoint restore was not accepted: ${error.message}`), + OrchestratorCheckpointRollbackTargetUnsupportedError: ( + error: OrchestratorCheckpointRollbackTargetUnsupportedError, + ) => failure("unsupported", `Checkpoint restore was not accepted: ${error.message}`), + }), + Effect.mapError((error) => + isCheckpointMcpFailure(error) + ? error + : failure( + "operation_failed", + `Checkpoint restore was not accepted: ${errorMessage(error)}`, + ), + ), + ); + const rollbackEvent = dispatched.storedEvents.find( + (stored) => stored.event.type === "checkpoint.rollback-requested", + ); + if ( + rollbackEvent?.event.type !== "checkpoint.rollback-requested" || + rollbackEvent.event.payload.scopeId !== checkpointScope.id || + rollbackEvent.event.payload.checkpointId !== checkpoint.id + ) { + return yield* failure( + "idempotency_conflict", + `clientRequestId '${input.clientRequestId}' was already accepted for a different checkpoint restore.`, + ); + } + + const observation = yield* readAcceptedRestore(scope, input, commandId).pipe( + Effect.match({ + onFailure: (error) => ({ type: "unavailable" as const, error }), + onSuccess: (accepted) => ({ type: "available" as const, accepted }), + }), + ); + if (observation.type === "available" && Option.isSome(observation.accepted)) { + return observation.accepted.value; + } + if (observation.type === "unavailable" && observation.error.code === "idempotency_conflict") { + return yield* observation.error; + } + const observationDetail = + observation.type === "unavailable" + ? observation.error.message + : "The durable command receipt was not yet readable."; + return { + commandId, + threadId: projection.thread.id, + scopeId: checkpointScope.id, + checkpointId: checkpoint.id, + status: "REQUESTED", + receipt: { + status: "accepted", + acceptedAt: DateTime.formatIso(rollbackEvent.event.occurredAt), + sequence: dispatched.sequence, + }, + effectStatus: "unavailable", + detail: `Checkpoint restore was accepted, but its effect status is unavailable: ${observationDetail} Reuse the same clientRequestId to observe this request.`, + }; + }); + + return CheckpointMcpService.of({ list, diff, restore }); }); export const layer: Layer.Layer< diff --git a/apps/server/src/mcp/toolkits/checkpoint/handlers.ts b/apps/server/src/mcp/toolkits/checkpoint/handlers.ts index 097c0526f82c..821c9ea82160 100644 --- a/apps/server/src/mcp/toolkits/checkpoint/handlers.ts +++ b/apps/server/src/mcp/toolkits/checkpoint/handlers.ts @@ -17,6 +17,12 @@ const handlers = { const service = yield* CheckpointMcpService; return yield* service.diff(scope, input); }), + t3_checkpoint_restore: (input) => + Effect.gen(function* () { + const scope = yield* McpInvocationContext; + const service = yield* CheckpointMcpService; + return yield* service.restore(scope, input); + }), } satisfies Parameters[0]; export const CheckpointToolkitHandlersLive = CheckpointToolkit.toLayer(handlers); diff --git a/apps/server/src/mcp/toolkits/checkpoint/tools.ts b/apps/server/src/mcp/toolkits/checkpoint/tools.ts index 36b5be0a1af2..a8198e3dbaf2 100644 --- a/apps/server/src/mcp/toolkits/checkpoint/tools.ts +++ b/apps/server/src/mcp/toolkits/checkpoint/tools.ts @@ -4,6 +4,8 @@ import { CheckpointMcpFailure, CheckpointMcpListInput, CheckpointMcpListResult, + CheckpointMcpRestoreInput, + CheckpointMcpRestoreResult, } from "@t3tools/contracts"; import { Tool, Toolkit } from "effect/unstable/ai"; @@ -29,7 +31,7 @@ export const CheckpointListTool = Tool.make("t3_checkpoint_list", { export const CheckpointDiffTool = Tool.make("t3_checkpoint_diff", { description: - "Read a bounded patch between two durable checkpoints in one thread checkpoint scope. Use t3_checkpoint_list first to select stable checkpointId and scopeId values. The thread and both checkpoint identities are validated before an empty diff is returned; arbitrary filesystem paths or Git refs are not accepted. Pagination cursors count UTF-16 code units.", + "Read a bounded patch between two durable checkpoints in one thread checkpoint scope. Use t3_checkpoint_list first to select stable checkpointId and scopeId values. The thread and both checkpoint identities are validated before an empty diff is returned; arbitrary filesystem paths or Git refs are not accepted. Pagination cursors count UTF-16 code units, and a page can exceed its requested limit by one code unit to avoid splitting a surrogate pair.", parameters: CheckpointMcpDiffInput, success: CheckpointMcpDiffResult, failure: CheckpointMcpFailure, @@ -42,4 +44,23 @@ export const CheckpointDiffTool = Tool.make("t3_checkpoint_diff", { .annotate(Tool.Idempotent, true) .annotate(Tool.OpenWorld, false); -export const CheckpointToolkit = Toolkit.make(CheckpointListTool, CheckpointDiffTool); +export const CheckpointRestoreTool = Tool.make("t3_checkpoint_restore", { + description: + "Request restoration of an exact durable checkpoint through the serialized V2 checkpoint.rollback workflow. This discards current tracked, untracked and staged workspace changes covered by Git restore, so discardChanges must be true. The thread must be idle with no queued work, the provider must support conversation rollback, and the guarded workspace/index state must remain unchanged before the locked restore. Reuse the exact clientRequestId to read the original accepted command/effect status without repeating it. REQUESTED means accepted but pending/running or temporarily unobservable; PARTIAL means filesystem/provider state may have changed but the complete outcome was not recorded.", + parameters: CheckpointMcpRestoreInput, + success: CheckpointMcpRestoreResult, + failure: CheckpointMcpFailure, + failureMode: "return", + dependencies, +}) + .annotate(Tool.Title, "Restore a thread checkpoint") + .annotate(Tool.Readonly, false) + .annotate(Tool.Destructive, true) + .annotate(Tool.Idempotent, true) + .annotate(Tool.OpenWorld, true); + +export const CheckpointToolkit = Toolkit.make( + CheckpointListTool, + CheckpointDiffTool, + CheckpointRestoreTool, +); diff --git a/apps/server/src/mcp/toolkits/worktree/registration.test.ts b/apps/server/src/mcp/toolkits/worktree/registration.test.ts index 0b4b44f7eb7a..09159227a7fa 100644 --- a/apps/server/src/mcp/toolkits/worktree/registration.test.ts +++ b/apps/server/src/mcp/toolkits/worktree/registration.test.ts @@ -114,6 +114,7 @@ it.effect("production mcp layer lists worktree tools over http", () => expect(toolNames).toContain("t3_worktree_status"); expect(toolNames).toContain("t3_checkpoint_list"); expect(toolNames).toContain("t3_checkpoint_diff"); + expect(toolNames).toContain("t3_checkpoint_restore"); // The worktree registration merges alongside the other toolkits rather // than replacing them. expect(toolNames).toContain("preview_status"); @@ -131,10 +132,14 @@ it.effect("production mcp layer lists worktree tools over http", () => expect(status?.annotations?.destructiveHint).toBe(false); const checkpointList = tools.find((tool) => tool.name === "t3_checkpoint_list"); const checkpointDiff = tools.find((tool) => tool.name === "t3_checkpoint_diff"); + const checkpointRestore = tools.find((tool) => tool.name === "t3_checkpoint_restore"); expect(checkpointList?.annotations?.readOnlyHint).toBe(true); expect(checkpointList?.annotations?.destructiveHint).toBe(false); expect(checkpointDiff?.annotations?.readOnlyHint).toBe(true); expect(checkpointDiff?.annotations?.destructiveHint).toBe(false); + expect(checkpointRestore?.annotations?.readOnlyHint).toBe(false); + expect(checkpointRestore?.annotations?.destructiveHint).toBe(true); + expect(checkpointRestore?.annotations?.openWorldHint).toBe(true); // MCP requires every tool input schema to be a top-level object schema. // A non-object schema (e.g. the anyOf produced by an empty diff --git a/apps/server/src/orchestration-v2/CheckpointRollbackService.test.ts b/apps/server/src/orchestration-v2/CheckpointRollbackService.test.ts index 5c22663c8ff6..85cbf198b671 100644 --- a/apps/server/src/orchestration-v2/CheckpointRollbackService.test.ts +++ b/apps/server/src/orchestration-v2/CheckpointRollbackService.test.ts @@ -1,27 +1,127 @@ import { assert, it, vi } from "@effect/vitest"; import { CheckpointId, + CheckpointRef, CheckpointScopeId, + MessageId, + NodeId, type OrchestrationV2ThreadProjection, + ProviderDriverKind, ProviderInstanceId, ProviderSessionId, ProviderThreadId, + RunId, ThreadId, } from "@t3tools/contracts"; +import * as DateTime from "effect/DateTime"; import * as Effect from "effect/Effect"; import * as Layer from "effect/Layer"; -import { CheckpointServiceV2 } from "./CheckpointService.ts"; +import { + CheckpointRestoreOutcomeUnknownError, + CheckpointRestorePreflightError, + CheckpointServiceV2, +} from "./CheckpointService.ts"; import { CheckpointRollbackServiceV2, layer as checkpointRollbackServiceLayer, } from "./CheckpointRollbackService.ts"; -import { EventSinkV2 } from "./EventSink.ts"; +import { EventSinkV2, EventSinkWriteError } from "./EventSink.ts"; import { layer as idAllocatorLayer } from "./IdAllocator.ts"; +import { threadDispatchLockLayer } from "./KeyedSerialExecutor.ts"; import { ProjectionStoreReadError, ProjectionStoreV2 } from "./ProjectionStore.ts"; -import { ProviderSessionManagerV2 } from "./ProviderSessionManager.ts"; +import { ProviderSessionManagerV2, ProviderSessionOpenError } from "./ProviderSessionManager.ts"; import { RuntimePolicyV2 } from "./RuntimePolicy.ts"; +const checkpointRollbackServiceTestLayer = checkpointRollbackServiceLayer.pipe( + Layer.provide(threadDispatchLockLayer), +); + +function makeReadyRollbackProjection(input: { + readonly threadId: ThreadId; + readonly providerThreadId: ProviderThreadId; + readonly providerSessionId: ProviderSessionId; + readonly checkpointId: CheckpointId; + readonly scopeId: CheckpointScopeId; + readonly providerInstanceId: ProviderInstanceId; +}): OrchestrationV2ThreadProjection { + const now = DateTime.makeUnsafe("2026-08-29T00:00:00.000Z"); + const driver = ProviderDriverKind.make("codex"); + const nodeId = NodeId.make(`node:${input.threadId}`); + return { + thread: { + id: input.threadId, + activeProviderThreadId: input.providerThreadId, + modelSelection: { instanceId: input.providerInstanceId, model: "test-model" }, + archivedAt: null, + deletedAt: null, + }, + providerThreads: [ + { + id: input.providerThreadId, + providerSessionId: input.providerSessionId, + providerInstanceId: input.providerInstanceId, + driver, + lastRunOrdinal: 1, + }, + ], + providerSessions: [{ id: input.providerSessionId }], + checkpoints: [ + { + id: input.checkpointId, + threadId: input.threadId, + scopeId: input.scopeId, + runId: null, + nodeId, + parentCheckpointId: null, + ordinalWithinScope: 0, + appRunOrdinal: null, + ref: CheckpointRef.make(`refs/t3/${input.checkpointId}`), + status: "ready", + files: [], + capturedAt: now, + }, + ], + checkpointScopes: [ + { + id: input.scopeId, + threadId: input.threadId, + runId: null, + nodeId, + parentScopeId: null, + providerThreadId: input.providerThreadId, + kind: "root_run", + ordinalWithinParent: 0, + advancesAppRunCount: true, + cwd: "/repo", + createdAt: now, + }, + ], + runs: [ + { + id: RunId.make(`run:${input.threadId}`), + threadId: input.threadId, + ordinal: 1, + providerInstanceId: input.providerInstanceId, + modelSelection: { instanceId: input.providerInstanceId, model: "test-model" }, + providerThreadId: input.providerThreadId, + userMessageId: MessageId.make(`message:${input.threadId}`), + rootNodeId: null, + activeAttemptId: null, + status: "completed", + requestedAt: now, + startedAt: now, + completedAt: now, + checkpointId: null, + contextHandoffId: null, + }, + ], + nodes: [], + attempts: [], + providerTurns: [], + } as unknown as OrchestrationV2ThreadProjection; +} + it.effect("rejects a non-ready checkpoint before opening a session or restoring files", () => { const threadId = ThreadId.make("thread:rollback-non-ready"); const providerThreadId = ProviderThreadId.make("provider-thread:rollback-non-ready"); @@ -41,14 +141,15 @@ it.effect("rejects a non-ready checkpoint before opening a session or restoring checkpoints: [{ id: checkpointId, scopeId, status: "stale" }], checkpointScopes: [{ id: scopeId }], } as unknown as OrchestrationV2ThreadProjection; - const testLayer = checkpointRollbackServiceLayer.pipe( + const testLayer = checkpointRollbackServiceTestLayer.pipe( Layer.provide( Layer.mergeAll( Layer.mock(CheckpointServiceV2)({ restore }), Layer.mock(EventSinkV2)({}), idAllocatorLayer, Layer.mock(ProjectionStoreV2)({ - getThreadProjection: () => Effect.succeed(projection), + getThreadSnapshot: () => + Effect.succeed({ schemaVersion: 1, snapshotSequence: 1, projection }), }), Layer.mock(ProviderSessionManagerV2)({ open }), Layer.mock(RuntimePolicyV2)({ resolve: resolveRuntimePolicy }), @@ -67,6 +168,7 @@ it.effect("rejects a non-ready checkpoint before opening a session or restoring }) .pipe(Effect.flip); + assert(error._tag === "CheckpointRollbackRejectedError"); assert.equal(error.reason, "rollback-target-invalid"); assert.equal( error.message, @@ -111,14 +213,15 @@ it.effect("rejects a rollback when another provider thread became active", () => checkpoints: [{ id: checkpointId, scopeId, status: "ready" }], checkpointScopes: [{ id: scopeId }], } as unknown as OrchestrationV2ThreadProjection; - const testLayer = checkpointRollbackServiceLayer.pipe( + const testLayer = checkpointRollbackServiceTestLayer.pipe( Layer.provide( Layer.mergeAll( Layer.mock(CheckpointServiceV2)({ restore }), Layer.mock(EventSinkV2)({}), idAllocatorLayer, Layer.mock(ProjectionStoreV2)({ - getThreadProjection: () => Effect.succeed(projection), + getThreadSnapshot: () => + Effect.succeed({ schemaVersion: 1, snapshotSequence: 1, projection }), }), Layer.mock(ProviderSessionManagerV2)({ open }), Layer.mock(RuntimePolicyV2)({ resolve: resolveRuntimePolicy }), @@ -137,6 +240,7 @@ it.effect("rejects a rollback when another provider thread became active", () => }) .pipe(Effect.flip); + assert(error._tag === "CheckpointRollbackRejectedError"); assert.equal(error.reason, "active-provider-changed"); assert.equal( error.message, @@ -183,14 +287,15 @@ it.effect("rejects a rollback when provider selection changed before execution", checkpoints: [{ id: checkpointId, scopeId, status: "ready" }], checkpointScopes: [{ id: scopeId }], } as unknown as OrchestrationV2ThreadProjection; - const testLayer = checkpointRollbackServiceLayer.pipe( + const testLayer = checkpointRollbackServiceTestLayer.pipe( Layer.provide( Layer.mergeAll( Layer.mock(CheckpointServiceV2)({ restore }), Layer.mock(EventSinkV2)({}), idAllocatorLayer, Layer.mock(ProjectionStoreV2)({ - getThreadProjection: () => Effect.succeed(projection), + getThreadSnapshot: () => + Effect.succeed({ schemaVersion: 1, snapshotSequence: 1, projection }), }), Layer.mock(ProviderSessionManagerV2)({ open }), Layer.mock(RuntimePolicyV2)({ resolve: resolveRuntimePolicy }), @@ -209,6 +314,7 @@ it.effect("rejects a rollback when provider selection changed before execution", }) .pipe(Effect.flip); + assert(error._tag === "CheckpointRollbackRejectedError"); assert.equal(error.reason, "active-provider-changed"); assert.equal( error.message, @@ -246,14 +352,15 @@ it.effect("reports a missing provider turn as a structured rollback failure", () attempts: [], providerTurns: [], } as unknown as OrchestrationV2ThreadProjection; - const testLayer = checkpointRollbackServiceLayer.pipe( + const testLayer = checkpointRollbackServiceTestLayer.pipe( Layer.provide( Layer.mergeAll( Layer.mock(CheckpointServiceV2)({ restore }), Layer.mock(EventSinkV2)({}), idAllocatorLayer, Layer.mock(ProjectionStoreV2)({ - getThreadProjection: () => Effect.succeed(projection), + getThreadSnapshot: () => + Effect.succeed({ schemaVersion: 1, snapshotSequence: 1, projection }), }), Layer.mock(ProviderSessionManagerV2)({ open: () => Effect.succeed({} as never), @@ -276,6 +383,7 @@ it.effect("reports a missing provider turn as a structured rollback failure", () }) .pipe(Effect.flip); + assert(error._tag === "CheckpointRollbackRejectedError"); assert.equal(error.reason, "provider-turn-unavailable"); assert.equal( error.message, @@ -286,7 +394,7 @@ it.effect("reports a missing provider turn as a structured rollback failure", () }).pipe(Effect.provide(testLayer)); }); -it.effect("wraps underlying failures with an unexpected-failure reason and cause", () => { +it.effect("wraps underlying failures with a preflight tag and cause", () => { const threadId = ThreadId.make("thread:rollback-unexpected-failure"); const providerThreadId = ProviderThreadId.make("provider-thread:rollback-unexpected-failure"); const checkpointId = CheckpointId.make("checkpoint:rollback-unexpected-failure"); @@ -295,14 +403,14 @@ it.effect("wraps underlying failures with an unexpected-failure reason and cause threadId, cause: new Error("database read failed"), }); - const testLayer = checkpointRollbackServiceLayer.pipe( + const testLayer = checkpointRollbackServiceTestLayer.pipe( Layer.provide( Layer.mergeAll( Layer.mock(CheckpointServiceV2)({}), Layer.mock(EventSinkV2)({}), idAllocatorLayer, Layer.mock(ProjectionStoreV2)({ - getThreadProjection: () => Effect.fail(projectionError), + getThreadSnapshot: () => Effect.fail(projectionError), }), Layer.mock(ProviderSessionManagerV2)({}), Layer.mock(RuntimePolicyV2)({}), @@ -321,7 +429,7 @@ it.effect("wraps underlying failures with an unexpected-failure reason and cause }) .pipe(Effect.flip); - assert.equal(error.reason, "unexpected-failure"); + assert.equal(error._tag, "CheckpointRollbackPreflightError"); assert.equal( error.message, `Failed to execute rollback target ${checkpointId} on provider thread ${providerThreadId} for thread ${threadId}.`, @@ -329,3 +437,286 @@ it.effect("wraps underlying failures with an unexpected-failure reason and cause assert.strictEqual(error.cause, projectionError); }).pipe(Effect.provide(testLayer)); }); + +it.effect("opens the provider session before restoring files for legacy rollback", () => { + const threadId = ThreadId.make("thread:rollback-open-before-files"); + const providerThreadId = ProviderThreadId.make("provider-thread:rollback-open-before-files"); + const providerSessionId = ProviderSessionId.make("provider-session:rollback-open-before-files"); + const checkpointId = CheckpointId.make("checkpoint:rollback-open-before-files"); + const scopeId = CheckpointScopeId.make("scope:rollback-open-before-files"); + const providerInstanceId = ProviderInstanceId.make("provider_rollback_open_before_files"); + const projection = makeReadyRollbackProjection({ + threadId, + providerThreadId, + providerSessionId, + checkpointId, + scopeId, + providerInstanceId, + }); + const restore = vi.fn(() => Effect.die("session failure must preserve files")); + const openError = new ProviderSessionOpenError({ + instanceId: providerInstanceId, + providerSessionId, + cause: "simulated transient provider open failure", + }); + const testLayer = checkpointRollbackServiceTestLayer.pipe( + Layer.provide( + Layer.mergeAll( + Layer.mock(CheckpointServiceV2)({ restore }), + Layer.mock(EventSinkV2)({}), + idAllocatorLayer, + Layer.mock(ProjectionStoreV2)({ + getThreadSnapshot: () => + Effect.succeed({ schemaVersion: 1, snapshotSequence: 1, projection }), + }), + Layer.mock(ProviderSessionManagerV2)({ open: () => Effect.fail(openError) }), + Layer.mock(RuntimePolicyV2)({ resolve: () => Effect.succeed({} as never) }), + ), + ), + ); + + return Effect.gen(function* () { + const service = yield* CheckpointRollbackServiceV2; + const error = yield* service + .execute({ threadId, providerThreadId, checkpointId, scopeId }) + .pipe(Effect.flip); + + assert.equal(error._tag, "CheckpointRollbackPreflightError"); + assert.strictEqual(error.cause, openError); + assert.equal(restore.mock.calls.length, 0); + }).pipe(Effect.provide(testLayer)); +}); + +it.effect("reports an uncertain filesystem restore as partial before provider rollback", () => { + const threadId = ThreadId.make("thread:rollback-filesystem-partial"); + const providerThreadId = ProviderThreadId.make("provider-thread:rollback-filesystem-partial"); + const providerSessionId = ProviderSessionId.make("provider-session:rollback-filesystem-partial"); + const checkpointId = CheckpointId.make("checkpoint:rollback-filesystem-partial"); + const scopeId = CheckpointScopeId.make("scope:rollback-filesystem-partial"); + const providerInstanceId = ProviderInstanceId.make("provider_rollback_filesystem_partial"); + const projection = makeReadyRollbackProjection({ + threadId, + providerThreadId, + providerSessionId, + checkpointId, + scopeId, + providerInstanceId, + }); + const rollbackThread = vi.fn(() => Effect.void); + const restoreError = new CheckpointRestoreOutcomeUnknownError({ + scopeId, + checkpointId, + cause: "Git restore failed after its first mutating command", + }); + const testLayer = checkpointRollbackServiceTestLayer.pipe( + Layer.provide( + Layer.mergeAll( + Layer.mock(CheckpointServiceV2)({ restore: () => Effect.fail(restoreError) }), + Layer.mock(EventSinkV2)({}), + idAllocatorLayer, + Layer.mock(ProjectionStoreV2)({ + getThreadSnapshot: () => + Effect.succeed({ schemaVersion: 1, snapshotSequence: 1, projection }), + }), + Layer.mock(ProviderSessionManagerV2)({ + open: () => Effect.succeed({ rollbackThread } as never), + }), + Layer.mock(RuntimePolicyV2)({ resolve: () => Effect.succeed({} as never) }), + ), + ), + ); + + return Effect.gen(function* () { + const service = yield* CheckpointRollbackServiceV2; + const error = yield* service + .execute({ + threadId, + providerThreadId, + checkpointId, + scopeId, + expectedIdle: true, + expectedWorkspaceFingerprint: "workspace-before", + }) + .pipe(Effect.flip); + + assert(error._tag === "CheckpointRollbackPartialError"); + assert.equal(error.reason, "post-restore-finalization-failed"); + assert.strictEqual(error.cause, restoreError); + assert.equal(rollbackThread.mock.calls.length, 0); + }).pipe(Effect.provide(testLayer)); +}); + +it.effect("keeps a proven pre-restore failure retryable", () => { + const threadId = ThreadId.make("thread:rollback-restore-preflight"); + const providerThreadId = ProviderThreadId.make("provider-thread:rollback-restore-preflight"); + const providerSessionId = ProviderSessionId.make("provider-session:rollback-restore-preflight"); + const checkpointId = CheckpointId.make("checkpoint:rollback-restore-preflight"); + const scopeId = CheckpointScopeId.make("scope:rollback-restore-preflight"); + const providerInstanceId = ProviderInstanceId.make("provider_rollback_restore_preflight"); + const projection = makeReadyRollbackProjection({ + threadId, + providerThreadId, + providerSessionId, + checkpointId, + scopeId, + providerInstanceId, + }); + const rollbackThread = vi.fn(() => Effect.void); + const restoreError = new CheckpointRestorePreflightError({ + scopeId, + checkpointId, + cause: "Unable to read the workspace fingerprint before restoring files", + }); + const testLayer = checkpointRollbackServiceTestLayer.pipe( + Layer.provide( + Layer.mergeAll( + Layer.mock(CheckpointServiceV2)({ restore: () => Effect.fail(restoreError) }), + Layer.mock(EventSinkV2)({}), + idAllocatorLayer, + Layer.mock(ProjectionStoreV2)({ + getThreadSnapshot: () => + Effect.succeed({ schemaVersion: 1, snapshotSequence: 1, projection }), + }), + Layer.mock(ProviderSessionManagerV2)({ + open: () => Effect.succeed({ rollbackThread } as never), + }), + Layer.mock(RuntimePolicyV2)({ resolve: () => Effect.succeed({} as never) }), + ), + ), + ); + + return Effect.gen(function* () { + const service = yield* CheckpointRollbackServiceV2; + const error = yield* service + .execute({ + threadId, + providerThreadId, + checkpointId, + scopeId, + expectedIdle: true, + expectedWorkspaceFingerprint: "workspace-before", + }) + .pipe(Effect.flip); + + assert.equal(error._tag, "CheckpointRollbackPreflightError"); + assert.strictEqual(error.cause, restoreError); + assert.equal(rollbackThread.mock.calls.length, 0); + }).pipe(Effect.provide(testLayer)); +}); + +it.effect("rejects an ambiguous null-ordinal target inside the worker boundary", () => { + const threadId = ThreadId.make("thread:rollback-ambiguous-target"); + const providerThreadId = ProviderThreadId.make("provider-thread:rollback-ambiguous-target"); + const providerSessionId = ProviderSessionId.make("provider-session:rollback-ambiguous-target"); + const checkpointId = CheckpointId.make("checkpoint:rollback-ambiguous-target"); + const scopeId = CheckpointScopeId.make("scope:rollback-ambiguous-target"); + const providerInstanceId = ProviderInstanceId.make("provider_rollback_ambiguous_target"); + const ready = makeReadyRollbackProjection({ + threadId, + providerThreadId, + providerSessionId, + checkpointId, + scopeId, + providerInstanceId, + }); + const projection = { + ...ready, + checkpoints: ready.checkpoints.map((checkpoint) => ({ + ...checkpoint, + ordinalWithinScope: 2, + appRunOrdinal: null, + })), + }; + const restore = vi.fn(() => Effect.die("ambiguous restore must not touch files")); + const open = vi.fn(() => Effect.die("ambiguous restore must not open a provider session")); + const testLayer = checkpointRollbackServiceTestLayer.pipe( + Layer.provide( + Layer.mergeAll( + Layer.mock(CheckpointServiceV2)({ restore }), + Layer.mock(EventSinkV2)({}), + idAllocatorLayer, + Layer.mock(ProjectionStoreV2)({ + getThreadSnapshot: () => + Effect.succeed({ schemaVersion: 1, snapshotSequence: 1, projection }), + }), + Layer.mock(ProviderSessionManagerV2)({ open }), + Layer.mock(RuntimePolicyV2)({ resolve: () => Effect.succeed({} as never) }), + ), + ), + ); + + return Effect.gen(function* () { + const service = yield* CheckpointRollbackServiceV2; + const error = yield* service + .execute({ threadId, providerThreadId, checkpointId, scopeId }) + .pipe(Effect.flip); + assert(error._tag === "CheckpointRollbackRejectedError"); + assert.equal(error.reason, "rollback-target-ambiguous"); + assert.equal(restore.mock.calls.length, 0); + assert.equal(open.mock.calls.length, 0); + }).pipe(Effect.provide(testLayer)); +}); + +it.effect("reports persistence failure after provider rollback as partial", () => { + const threadId = ThreadId.make("thread:rollback-persistence-partial"); + const providerThreadId = ProviderThreadId.make("provider-thread:rollback-persistence-partial"); + const providerSessionId = ProviderSessionId.make("provider-session:rollback-persistence-partial"); + const checkpointId = CheckpointId.make("checkpoint:rollback-persistence-partial"); + const scopeId = CheckpointScopeId.make("scope:rollback-persistence-partial"); + const providerInstanceId = ProviderInstanceId.make("provider_rollback_persistence_partial"); + const projection = makeReadyRollbackProjection({ + threadId, + providerThreadId, + providerSessionId, + checkpointId, + scopeId, + providerInstanceId, + }); + const providerThread = projection.providerThreads[0]!; + const rollbackThread = vi.fn(() => + Effect.succeed({ providerThread, providerTurns: [], messages: [], runtimeRequests: [] }), + ); + const persistenceError = new EventSinkWriteError({ + eventCount: 2, + cause: "simulated event persistence failure", + }); + const testLayer = checkpointRollbackServiceTestLayer.pipe( + Layer.provide( + Layer.mergeAll( + Layer.mock(CheckpointServiceV2)({ + restore: () => Effect.void, + deleteStaleRefs: () => Effect.void, + }), + Layer.mock(EventSinkV2)({ write: () => Effect.fail(persistenceError) }), + idAllocatorLayer, + Layer.mock(ProjectionStoreV2)({ + getThreadSnapshot: () => + Effect.succeed({ schemaVersion: 1, snapshotSequence: 1, projection }), + }), + Layer.mock(ProviderSessionManagerV2)({ + open: () => Effect.succeed({ rollbackThread } as never), + }), + Layer.mock(RuntimePolicyV2)({ resolve: () => Effect.succeed({} as never) }), + ), + ), + ); + + return Effect.gen(function* () { + const service = yield* CheckpointRollbackServiceV2; + const error = yield* service + .execute({ + threadId, + providerThreadId, + checkpointId, + scopeId, + expectedIdle: true, + expectedWorkspaceFingerprint: "workspace-before", + }) + .pipe(Effect.flip); + + assert(error._tag === "CheckpointRollbackPartialError"); + assert.equal(error.reason, "post-restore-finalization-failed"); + assert.strictEqual(error.cause, persistenceError); + assert.equal(rollbackThread.mock.calls.length, 1); + }).pipe(Effect.provide(testLayer)); +}); diff --git a/apps/server/src/orchestration-v2/CheckpointRollbackService.ts b/apps/server/src/orchestration-v2/CheckpointRollbackService.ts index 1a0acb9007ab..b1e534194aef 100644 --- a/apps/server/src/orchestration-v2/CheckpointRollbackService.ts +++ b/apps/server/src/orchestration-v2/CheckpointRollbackService.ts @@ -1,4 +1,5 @@ import { + checkpointRollbackAppRunOrdinal, CheckpointId, CheckpointScopeId, type OrchestrationV2DomainEvent, @@ -11,44 +12,99 @@ import * as Effect from "effect/Effect"; import * as Layer from "effect/Layer"; import * as Schema from "effect/Schema"; -import { CheckpointServiceV2 } from "./CheckpointService.ts"; +import { + CheckpointRestoreOutcomeUnknownError, + CheckpointRestorePreflightError, + CheckpointRestorePreconditionError, + CheckpointServiceV2, +} from "./CheckpointService.ts"; import { EventSinkV2 } from "./EventSink.ts"; import { IdAllocatorV2 } from "./IdAllocator.ts"; +import { ThreadDispatchLockV2 } from "./KeyedSerialExecutor.ts"; import { ProjectionStoreV2 } from "./ProjectionStore.ts"; import type { ProviderAdapterV2RollbackTarget } from "./ProviderAdapter.ts"; import { ProviderSessionManagerV2 } from "./ProviderSessionManager.ts"; import { RuntimePolicyV2 } from "./RuntimePolicy.ts"; -export class CheckpointRollbackExecutionError extends Schema.TaggedErrorClass()( - "CheckpointRollbackExecutionError", +const CheckpointRollbackErrorFields = { + threadId: ThreadId, + providerThreadId: ProviderThreadId, + checkpointId: CheckpointId, + cause: Schema.optional(Schema.Defect()), +}; + +export class CheckpointRollbackRejectedError extends Schema.TaggedErrorClass()( + "CheckpointRollbackRejectedError", { reason: Schema.Literals([ "rollback-target-invalid", + "rollback-target-ambiguous", "active-provider-changed", "provider-turn-unavailable", - "unexpected-failure", + "thread-not-idle", + "restore-precondition-changed", ]), - threadId: ThreadId, - providerThreadId: ProviderThreadId, - checkpointId: CheckpointId, - cause: Schema.optional(Schema.Defect()), + ...CheckpointRollbackErrorFields, }, ) { override get message(): string { switch (this.reason) { case "rollback-target-invalid": return `Rollback target ${this.checkpointId} for provider thread ${this.providerThreadId} on thread ${this.threadId} is incomplete or invalid.`; + case "rollback-target-ambiguous": + return `Rollback target ${this.checkpointId} does not identify a proven provider-history position.`; case "active-provider-changed": return `Active provider changed before rollback target ${this.checkpointId} could execute on thread ${this.threadId}.`; case "provider-turn-unavailable": return `Provider turn for rollback target ${this.checkpointId} is unavailable on provider thread ${this.providerThreadId}.`; - case "unexpected-failure": - return `Failed to execute rollback target ${this.checkpointId} on provider thread ${this.providerThreadId} for thread ${this.threadId}.`; + case "thread-not-idle": + return `Checkpoint rollback target ${this.checkpointId} requires an idle thread with no queued runs.`; + case "restore-precondition-changed": + return `Checkpoint rollback target ${this.checkpointId} changed before filesystem restoration began.`; } } } +export class CheckpointRollbackPartialError extends Schema.TaggedErrorClass()( + "CheckpointRollbackPartialError", + { + reason: Schema.Literals([ + "provider-rollback-failed-after-restore", + "post-restore-finalization-failed", + ]), + ...CheckpointRollbackErrorFields, + }, +) { + override get message(): string { + return this.reason === "provider-rollback-failed-after-restore" + ? `Provider conversation rollback failed after the filesystem checkpoint ${this.checkpointId} was restored; the result is partial.` + : `Checkpoint ${this.checkpointId} was restored, but rollback finalization failed; the result is partial.`; + } +} + +export class CheckpointRollbackPreflightError extends Schema.TaggedErrorClass()( + "CheckpointRollbackPreflightError", + { + ...CheckpointRollbackErrorFields, + }, +) { + override get message(): string { + return `Failed to execute rollback target ${this.checkpointId} on provider thread ${this.providerThreadId} for thread ${this.threadId}.`; + } +} + +export const CheckpointRollbackExecutionError = Schema.Union([ + CheckpointRollbackRejectedError, + CheckpointRollbackPartialError, + CheckpointRollbackPreflightError, +]); +export type CheckpointRollbackExecutionError = typeof CheckpointRollbackExecutionError.Type; + +const isCheckpointRollbackPartialError = Schema.is(CheckpointRollbackPartialError); const isCheckpointRollbackExecutionError = Schema.is(CheckpointRollbackExecutionError); +const isCheckpointRestoreOutcomeUnknownError = Schema.is(CheckpointRestoreOutcomeUnknownError); +const isCheckpointRestorePreflightError = Schema.is(CheckpointRestorePreflightError); +const isCheckpointRestorePreconditionError = Schema.is(CheckpointRestorePreconditionError); export interface CheckpointRollbackServiceV2Shape { readonly execute: (input: { @@ -56,6 +112,8 @@ export interface CheckpointRollbackServiceV2Shape { readonly providerThreadId: ProviderThreadId; readonly checkpointId: CheckpointId; readonly scopeId: CheckpointScopeId; + readonly expectedIdle?: true; + readonly expectedWorkspaceFingerprint?: string; }) => Effect.Effect; } @@ -70,6 +128,7 @@ export const layer: Layer.Layer< | CheckpointServiceV2 | EventSinkV2 | IdAllocatorV2 + | ThreadDispatchLockV2 | ProjectionStoreV2 | ProviderSessionManagerV2 | RuntimePolicyV2 @@ -79,6 +138,7 @@ export const layer: Layer.Layer< const checkpoints = yield* CheckpointServiceV2; const eventSink = yield* EventSinkV2; const ids = yield* IdAllocatorV2; + const threadDispatch = yield* ThreadDispatchLockV2; const projections = yield* ProjectionStoreV2; const sessions = yield* ProviderSessionManagerV2; const runtimePolicy = yield* RuntimePolicyV2; @@ -88,8 +148,11 @@ export const layer: Layer.Layer< readonly providerThreadId: ProviderThreadId; readonly checkpointId: CheckpointId; readonly scopeId: CheckpointScopeId; + readonly expectedIdle?: true; + readonly expectedWorkspaceFingerprint?: string; }) { - const projection = yield* projections.getThreadProjection(input.threadId); + const initialSnapshot = yield* projections.getThreadSnapshot(input.threadId); + const projection = initialSnapshot.projection; const providerThread = projection.providerThreads.find( (candidate) => candidate.id === input.providerThreadId, ); @@ -105,24 +168,48 @@ export const layer: Layer.Layer< checkpoint.scopeId !== scope.id || checkpoint.status !== "ready" ) { - return yield* new CheckpointRollbackExecutionError({ + return yield* new CheckpointRollbackRejectedError({ reason: "rollback-target-invalid", threadId: input.threadId, providerThreadId: input.providerThreadId, checkpointId: input.checkpointId, }); } + const providerSessionId = providerThread.providerSessionId; if ( providerThread.id !== projection.thread.activeProviderThreadId || providerThread.providerInstanceId !== projection.thread.modelSelection.instanceId ) { - return yield* new CheckpointRollbackExecutionError({ + return yield* new CheckpointRollbackRejectedError({ reason: "active-provider-changed", threadId: input.threadId, providerThreadId: input.providerThreadId, checkpointId: input.checkpointId, }); } + if ( + input.expectedIdle === true && + projection.runs.some((run) => + ["preparing", "queued", "starting", "running", "waiting"].includes(run.status), + ) + ) { + return yield* new CheckpointRollbackRejectedError({ + reason: "thread-not-idle", + threadId: input.threadId, + providerThreadId: input.providerThreadId, + checkpointId: input.checkpointId, + }); + } + + const targetOrdinal = checkpointRollbackAppRunOrdinal(checkpoint, scope); + if (targetOrdinal === null) { + return yield* new CheckpointRollbackRejectedError({ + reason: "rollback-target-ambiguous", + threadId: input.threadId, + providerThreadId: input.providerThreadId, + checkpointId: input.checkpointId, + }); + } const modelSelection = projection.thread.modelSelection; const resolvedRuntimePolicy = yield* runtimePolicy.resolve({ @@ -130,42 +217,106 @@ export const layer: Layer.Layer< modelSelection, }); const existingSession = projection.providerSessions.find( - (candidate) => candidate.id === providerThread.providerSessionId, + (candidate) => candidate.id === providerSessionId, ); const session = yield* sessions.open({ threadId: input.threadId, - providerSessionId: providerThread.providerSessionId, + providerSessionId, modelSelection, runtimePolicy: resolvedRuntimePolicy, ...(existingSession === undefined ? {} : { resumeFromSession: existingSession }), }); - const targetOrdinal = checkpoint.appRunOrdinal ?? 0; - const runsToRollback = projection.runs.filter( + const admittedSnapshot = yield* projections.getThreadSnapshot(input.threadId); + const admittedProjection = admittedSnapshot.projection; + const admittedProviderThread = admittedProjection.providerThreads.find( + (candidate) => candidate.id === input.providerThreadId, + ); + const admittedCheckpoint = admittedProjection.checkpoints.find( + (candidate) => candidate.id === input.checkpointId, + ); + const admittedScope = admittedProjection.checkpointScopes.find( + (candidate) => candidate.id === input.scopeId, + ); + if ( + admittedProviderThread === undefined || + admittedProviderThread.providerSessionId !== providerSessionId || + admittedCheckpoint === undefined || + admittedScope === undefined || + admittedCheckpoint.scopeId !== admittedScope.id || + admittedCheckpoint.status !== "ready" + ) { + return yield* new CheckpointRollbackRejectedError({ + reason: "rollback-target-invalid", + threadId: input.threadId, + providerThreadId: input.providerThreadId, + checkpointId: input.checkpointId, + }); + } + if ( + admittedProviderThread.id !== admittedProjection.thread.activeProviderThreadId || + admittedProviderThread.providerInstanceId !== + admittedProjection.thread.modelSelection.instanceId + ) { + return yield* new CheckpointRollbackRejectedError({ + reason: "active-provider-changed", + threadId: input.threadId, + providerThreadId: input.providerThreadId, + checkpointId: input.checkpointId, + }); + } + if ( + input.expectedIdle === true && + admittedProjection.runs.some((run) => + ["preparing", "queued", "starting", "running", "waiting"].includes(run.status), + ) + ) { + return yield* new CheckpointRollbackRejectedError({ + reason: "thread-not-idle", + threadId: input.threadId, + providerThreadId: input.providerThreadId, + checkpointId: input.checkpointId, + }); + } + if (checkpointRollbackAppRunOrdinal(admittedCheckpoint, admittedScope) !== targetOrdinal) { + return yield* new CheckpointRollbackRejectedError({ + reason: "rollback-target-invalid", + threadId: input.threadId, + providerThreadId: input.providerThreadId, + checkpointId: input.checkpointId, + }); + } + + const runsToRollback = admittedProjection.runs.filter( (run) => run.ordinal > targetOrdinal && run.status === "completed", ); - const providerThreadTurns = projection.providerTurns.filter( - (turn) => turn.providerThreadId === providerThread.id, + const providerThreadTurns = admittedProjection.providerTurns.filter( + (turn) => turn.providerThreadId === admittedProviderThread.id, ); const rollbackTarget: ProviderAdapterV2RollbackTarget = targetOrdinal === 0 ? { type: "thread_start", - checkpointId: checkpoint.id, + checkpointId: admittedCheckpoint.id, appRunOrdinal: 0, } : yield* Effect.gen(function* () { - const targetRun = projection.runs.find((run) => run.ordinal === targetOrdinal); - const targetAttempt = projection.attempts.find( + const targetRun = admittedProjection.runs.find( + (run) => run.ordinal === targetOrdinal, + ); + const targetAttempt = admittedProjection.attempts.find( (attempt) => attempt.id === targetRun?.activeAttemptId, ); - const targetTurn = projection.providerTurns.find( + const targetTurn = admittedProjection.providerTurns.find( (turn) => turn.id === targetAttempt?.providerTurnId || turn.runAttemptId === targetAttempt?.id, ); - if (targetTurn === undefined || targetTurn.providerThreadId !== providerThread.id) { - return yield* new CheckpointRollbackExecutionError({ + if ( + targetTurn === undefined || + targetTurn.providerThreadId !== admittedProviderThread.id + ) { + return yield* new CheckpointRollbackRejectedError({ reason: "provider-turn-unavailable", threadId: input.threadId, providerThreadId: input.providerThreadId, @@ -174,108 +325,220 @@ export const layer: Layer.Layer< } return { type: "provider_turn" as const, - checkpointId: checkpoint.id, + checkpointId: admittedCheckpoint.id, appRunOrdinal: targetOrdinal, providerTurn: targetTurn, }; }); - yield* checkpoints.restore({ scope, checkpoint }); - const snapshot = - runsToRollback.length === 0 - ? { providerThread } - : yield* session.rollbackThread({ - providerThread, - target: rollbackTarget, - providerThreadTurns, + const validateBeforeRestore = + input.expectedIdle === true + ? projections.getThreadSnapshot(input.threadId).pipe( + Effect.flatMap((latestSnapshot) => { + const latest = latestSnapshot.projection; + const latestCheckpoint = latest.checkpoints.find( + (candidate) => candidate.id === input.checkpointId, + ); + const latestProviderThread = latest.providerThreads.find( + (candidate) => candidate.id === input.providerThreadId, + ); + const remainsIdle = !latest.runs.some((run) => + ["preparing", "queued", "starting", "running", "waiting"].includes(run.status), + ); + return remainsIdle && + latestSnapshot.snapshotSequence === admittedSnapshot.snapshotSequence && + latest.thread.archivedAt === null && + latest.thread.deletedAt === null && + latest.thread.activeProviderThreadId === input.providerThreadId && + latest.thread.modelSelection.instanceId === + admittedProviderThread.providerInstanceId && + latestProviderThread?.providerSessionId === providerSessionId && + latestCheckpoint?.scopeId === input.scopeId && + latestCheckpoint.status === "ready" + ? Effect.void + : Effect.fail( + new CheckpointRestorePreconditionError({ + scopeId: input.scopeId, + checkpointId: input.checkpointId, + reason: "precondition-changed", + cause: + "Thread or rollback target changed after admission; current workspace files were preserved.", + }), + ); + }), + Effect.mapError((cause) => + isCheckpointRestorePreconditionError(cause) || + isCheckpointRestorePreflightError(cause) + ? cause + : new CheckpointRestorePreflightError({ + scopeId: input.scopeId, + checkpointId: input.checkpointId, + cause, + }), + ), + ) + : undefined; + yield* checkpoints + .restore({ + scope: admittedScope, + checkpoint: admittedCheckpoint, + ...(input.expectedWorkspaceFingerprint === undefined + ? {} + : { expectedWorkspaceFingerprint: input.expectedWorkspaceFingerprint }), + ...(validateBeforeRestore === undefined ? {} : { validateBeforeRestore }), + }) + .pipe( + Effect.mapError((cause): CheckpointRollbackExecutionError => { + const fields = { + threadId: input.threadId, + providerThreadId: input.providerThreadId, + checkpointId: input.checkpointId, + cause, + }; + if (isCheckpointRestoreOutcomeUnknownError(cause)) { + return new CheckpointRollbackPartialError({ + ...fields, + reason: "post-restore-finalization-failed", + }); + } + if (isCheckpointRestorePreconditionError(cause)) { + return new CheckpointRollbackRejectedError({ + ...fields, + reason: "restore-precondition-changed", + }); + } + return new CheckpointRollbackPreflightError({ + ...fields, }); - const staleCheckpoints = projection.checkpoints.filter( - (candidate) => - candidate.scopeId === scope.id && - candidate.appRunOrdinal !== null && - candidate.appRunOrdinal > targetOrdinal && - candidate.status === "ready", - ); - if (staleCheckpoints.length > 0) { - 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); + yield* Effect.gen(function* () { + const snapshot = + runsToRollback.length === 0 + ? { providerThread: admittedProviderThread } + : yield* session + .rollbackThread({ + providerThread: admittedProviderThread, + target: rollbackTarget, + providerThreadTurns, + }) + .pipe( + Effect.mapError( + (cause) => + new CheckpointRollbackPartialError({ + reason: "provider-rollback-failed-after-restore", + threadId: input.threadId, + providerThreadId: input.providerThreadId, + checkpointId: input.checkpointId, + cause, + }), + ), + ); + const staleCheckpoints = admittedProjection.checkpoints.filter( + (candidate) => + candidate.scopeId === admittedScope.id && + candidate.appRunOrdinal !== null && + candidate.appRunOrdinal > targetOrdinal && + candidate.status === "ready", + ); + if (staleCheckpoints.length > 0) { + yield* checkpoints.deleteStaleRefs({ + scope: admittedScope, + 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: "run.updated", + type: "provider-thread.updated", threadId: input.threadId, - runId: run.id, - ...(rootNode === undefined ? {} : { nodeId: rootNode.id }), - providerInstanceId: run.providerInstanceId, + driver: admittedProviderThread.driver, + providerInstanceId: admittedProviderThread.providerInstanceId, occurredAt: now, - payload: { ...run, status: "rolled_back", completedAt: now }, + payload: { + ...snapshot.providerThread, + lastRunOrdinal: targetOrdinal === 0 ? null : targetOrdinal, + updatedAt: now, + }, }), ); - if (rootNode !== undefined) { + 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: admittedProviderThread.providerInstanceId, + occurredAt: now, + payload: { ...staleCheckpoint, status: "stale" }, + }), + ); + } + for (const run of runsToRollback) { + const rootNode = admittedProjection.nodes.find( + (candidate) => candidate.id === run.rootNodeId, + ); events.push( yield* makeEvent({ - type: "node.updated", + type: "run.updated", threadId: input.threadId, runId: run.id, - nodeId: rootNode.id, + ...(rootNode === undefined ? {} : { nodeId: rootNode.id }), providerInstanceId: run.providerInstanceId, occurredAt: now, - payload: { ...rootNode, status: "rolled_back", completedAt: 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 }); + yield* eventSink.write({ events }); + }).pipe( + Effect.mapError((cause) => + isCheckpointRollbackPartialError(cause) && + cause.reason === "provider-rollback-failed-after-restore" + ? cause + : new CheckpointRollbackPartialError({ + reason: "post-restore-finalization-failed", + threadId: input.threadId, + providerThreadId: input.providerThreadId, + checkpointId: input.checkpointId, + cause, + }), + ), + ); }); return CheckpointRollbackServiceV2.of({ execute: (input) => - execute(input).pipe( + threadDispatch.withLock(input.threadId, execute(input)).pipe( Effect.mapError((cause) => isCheckpointRollbackExecutionError(cause) ? cause - : new CheckpointRollbackExecutionError({ - reason: "unexpected-failure", + : new CheckpointRollbackPreflightError({ threadId: input.threadId, providerThreadId: input.providerThreadId, checkpointId: input.checkpointId, diff --git a/apps/server/src/orchestration-v2/CheckpointService.test.ts b/apps/server/src/orchestration-v2/CheckpointService.test.ts index 1edfe541e3e0..c075c7f21cc6 100644 --- a/apps/server/src/orchestration-v2/CheckpointService.test.ts +++ b/apps/server/src/orchestration-v2/CheckpointService.test.ts @@ -1,11 +1,15 @@ import { assert, it, vi } from "@effect/vitest"; import { + CheckpointId, + CheckpointRef, CheckpointScopeId, NodeId, ProviderThreadId, RunId, ThreadId, + VcsProcessExitError, type OrchestrationV2CheckpointScope, + type OrchestrationV2Checkpoint, } from "@t3tools/contracts"; import * as DateTime from "effect/DateTime"; import * as Effect from "effect/Effect"; @@ -70,3 +74,125 @@ it.effect("materializes the captured baseline at the requested scope ordinal", ( }); }).pipe(Effect.provide(testLayer)); }); + +it.effect("preserves concurrently changed files before locked restore", () => { + const scope: OrchestrationV2CheckpointScope = { + id: CheckpointScopeId.make("checkpoint-scope:guarded-restore"), + threadId: ThreadId.make("thread:guarded-restore"), + runId: null, + nodeId: NodeId.make("node:guarded-restore"), + parentScopeId: null, + providerThreadId: ProviderThreadId.make("provider-thread:guarded-restore"), + kind: "manual", + ordinalWithinParent: 0, + advancesAppRunCount: false, + cwd: "/repo", + createdAt: DateTime.makeUnsafe("2026-08-29T00:00:00.000Z"), + }; + const checkpoint: OrchestrationV2Checkpoint = { + id: CheckpointId.make("checkpoint:guarded-restore"), + threadId: scope.threadId, + scopeId: scope.id, + runId: null, + nodeId: scope.nodeId, + parentCheckpointId: null, + ordinalWithinScope: 0, + appRunOrdinal: null, + ref: CheckpointRef.make("refs/t3/guarded-restore"), + status: "ready", + files: [], + capturedAt: scope.createdAt, + }; + const restoreCheckpoint = vi.fn(() => Effect.succeed(true)); + const testLayer = checkpointServiceLayer.pipe( + Layer.provide( + Layer.mergeAll( + idAllocatorLayer, + Layer.mock(CheckpointStore.CheckpointStore)({ + readWorkspaceFingerprint: () => Effect.succeed("tree:changed"), + restoreCheckpoint, + }), + ), + ), + ); + + return Effect.gen(function* () { + const checkpoints = yield* CheckpointServiceV2; + const error = yield* checkpoints + .restore({ + scope, + checkpoint, + expectedWorkspaceFingerprint: "tree:admitted", + }) + .pipe(Effect.flip); + + assert.equal(error._tag, "CheckpointRestorePreconditionError"); + assert.equal(restoreCheckpoint.mock.calls.length, 0); + }).pipe(Effect.provide(testLayer)); +}); + +it.effect("classifies fingerprint read failures as retryable preflight failures", () => { + const scope: OrchestrationV2CheckpointScope = { + id: CheckpointScopeId.make("checkpoint-scope:fingerprint-read-failure"), + threadId: ThreadId.make("thread:fingerprint-read-failure"), + runId: null, + nodeId: NodeId.make("node:fingerprint-read-failure"), + parentScopeId: null, + providerThreadId: ProviderThreadId.make("provider-thread:fingerprint-read-failure"), + kind: "manual", + ordinalWithinParent: 0, + advancesAppRunCount: false, + cwd: "/repo", + createdAt: DateTime.makeUnsafe("2026-08-29T00:00:00.000Z"), + }; + const checkpoint: OrchestrationV2Checkpoint = { + id: CheckpointId.make("checkpoint:fingerprint-read-failure"), + threadId: scope.threadId, + scopeId: scope.id, + runId: null, + nodeId: scope.nodeId, + parentCheckpointId: null, + ordinalWithinScope: 0, + appRunOrdinal: null, + ref: CheckpointRef.make("refs/t3/fingerprint-read-failure"), + status: "ready", + files: [], + capturedAt: scope.createdAt, + }; + const fingerprintError = new VcsProcessExitError({ + operation: "CheckpointStore.readWorkspaceFingerprint", + command: "git ls-files --stage", + cwd: scope.cwd, + exitCode: 0, + detail: "staged-state output was truncated", + }); + const restoreCheckpoint = vi.fn(() => Effect.succeed(true)); + const testLayer = checkpointServiceLayer.pipe( + Layer.provide( + Layer.mergeAll( + idAllocatorLayer, + Layer.mock(CheckpointStore.CheckpointStore)({ + readWorkspaceFingerprint: () => Effect.fail(fingerprintError), + restoreCheckpoint, + }), + ), + ), + ); + + return Effect.gen(function* () { + const checkpoints = yield* CheckpointServiceV2; + const error = yield* checkpoints + .restore({ + scope, + checkpoint, + expectedWorkspaceFingerprint: "tree:admitted", + }) + .pipe(Effect.flip); + + assert.equal(error._tag, "CheckpointRestorePreflightError"); + if (error._tag === "CheckpointRestorePreflightError") { + assert.strictEqual(error.cause, fingerprintError); + } + assert.equal(restoreCheckpoint.mock.calls.length, 0); + }).pipe(Effect.provide(testLayer)); +}); diff --git a/apps/server/src/orchestration-v2/CheckpointService.ts b/apps/server/src/orchestration-v2/CheckpointService.ts index 80211418308b..5df9b03d3519 100644 --- a/apps/server/src/orchestration-v2/CheckpointService.ts +++ b/apps/server/src/orchestration-v2/CheckpointService.ts @@ -77,16 +77,43 @@ export class CheckpointCaptureError extends Schema.TaggedErrorClass()( - "CheckpointRestoreError", +export class CheckpointRestorePreconditionError extends Schema.TaggedErrorClass()( + "CheckpointRestorePreconditionError", { scopeId: CheckpointScopeId, checkpointId: CheckpointId, + reason: Schema.Literals(["precondition-changed", "target-unavailable"]), cause: Schema.Defect(), }, ) { override get message(): string { - return `Failed to restore checkpoint ${this.checkpointId} for scope ${this.scopeId}.`; + return `Checkpoint ${this.checkpointId} for scope ${this.scopeId} no longer satisfies its restore preconditions.`; + } +} + +export class CheckpointRestorePreflightError extends Schema.TaggedErrorClass()( + "CheckpointRestorePreflightError", + { + scopeId: CheckpointScopeId, + checkpointId: CheckpointId, + cause: Schema.Defect(), + }, +) { + override get message(): string { + return `Failed to verify checkpoint ${this.checkpointId} for scope ${this.scopeId} before restoring files.`; + } +} + +export class CheckpointRestoreOutcomeUnknownError extends Schema.TaggedErrorClass()( + "CheckpointRestoreOutcomeUnknownError", + { + scopeId: CheckpointScopeId, + checkpointId: CheckpointId, + cause: Schema.Defect(), + }, +) { + override get message(): string { + return `Checkpoint ${this.checkpointId} for scope ${this.scopeId} may have been partially restored.`; } } @@ -108,13 +135,13 @@ export const CheckpointServiceV2Error = Schema.Union([ CheckpointScopeEnsureError, CheckpointBaselineCaptureError, CheckpointCaptureError, - CheckpointRestoreError, + CheckpointRestorePreconditionError, + CheckpointRestorePreflightError, + CheckpointRestoreOutcomeUnknownError, CheckpointDeleteStaleRefsError, ]); export type CheckpointServiceV2Error = typeof CheckpointServiceV2Error.Type; -const isCheckpointRestoreError = Schema.is(CheckpointRestoreError); - export interface CheckpointServiceV2Shape { readonly prepareRootRunScope: (input: { readonly threadId: ThreadId; @@ -146,6 +173,11 @@ export interface CheckpointServiceV2Shape { readonly restore: (input: { readonly scope: OrchestrationV2CheckpointScope; readonly checkpoint: OrchestrationV2Checkpoint; + readonly expectedWorkspaceFingerprint?: string; + readonly validateBeforeRestore?: Effect.Effect< + void, + CheckpointRestorePreconditionError | CheckpointRestorePreflightError + >; }) => Effect.Effect; readonly deleteStaleRefs: (input: { readonly scope: OrchestrationV2CheckpointScope; @@ -490,37 +522,66 @@ export const layer: Layer.Layer< withWorkspaceLock( input.scope.cwd, Effect.gen(function* () { + if (input.validateBeforeRestore !== undefined) { + yield* input.validateBeforeRestore; + } if (input.checkpoint.status !== "ready") { - return yield* new CheckpointRestoreError({ + return yield* new CheckpointRestorePreconditionError({ scopeId: input.scope.id, checkpointId: input.checkpoint.id, + reason: "target-unavailable", cause: `Checkpoint status is ${input.checkpoint.status}.`, }); } - const restored = yield* checkpointStore.restoreCheckpoint({ - cwd: input.scope.cwd, - checkpointRef: input.checkpoint.ref, - fallbackToHead: false, - }); + if (input.expectedWorkspaceFingerprint !== undefined) { + const currentFingerprint = yield* checkpointStore + .readWorkspaceFingerprint(input.scope.cwd) + .pipe( + Effect.mapError( + (cause) => + new CheckpointRestorePreflightError({ + scopeId: input.scope.id, + checkpointId: input.checkpoint.id, + cause, + }), + ), + ); + if (currentFingerprint !== input.expectedWorkspaceFingerprint) { + return yield* new CheckpointRestorePreconditionError({ + scopeId: input.scope.id, + checkpointId: input.checkpoint.id, + reason: "precondition-changed", + cause: "Workspace changed after rollback admission; current files were preserved.", + }); + } + } + + const restored = yield* checkpointStore + .restoreCheckpoint({ + cwd: input.scope.cwd, + checkpointRef: input.checkpoint.ref, + fallbackToHead: false, + }) + .pipe( + Effect.mapError( + (cause) => + new CheckpointRestoreOutcomeUnknownError({ + scopeId: input.scope.id, + checkpointId: input.checkpoint.id, + cause, + }), + ), + ); if (!restored) { - return yield* new CheckpointRestoreError({ + return yield* new CheckpointRestorePreconditionError({ scopeId: input.scope.id, checkpointId: input.checkpoint.id, + reason: "target-unavailable", cause: "Checkpoint ref is unavailable.", }); } }), - ).pipe( - Effect.mapError((cause) => - isCheckpointRestoreError(cause) - ? cause - : new CheckpointRestoreError({ - scopeId: input.scope.id, - checkpointId: input.checkpoint.id, - cause, - }), - ), ); const deleteStaleRefs: CheckpointServiceV2Shape["deleteStaleRefs"] = (input) => diff --git a/apps/server/src/orchestration-v2/EffectOutbox.ts b/apps/server/src/orchestration-v2/EffectOutbox.ts index 85180279fc77..a88afde04697 100644 --- a/apps/server/src/orchestration-v2/EffectOutbox.ts +++ b/apps/server/src/orchestration-v2/EffectOutbox.ts @@ -77,6 +77,8 @@ export const OrchestrationEffectRequestV2 = Schema.Union([ providerThreadId: ProviderThreadId, checkpointId: CheckpointId, scopeId: CheckpointScopeId, + expectedIdle: Schema.optional(Schema.Literal(true)), + expectedWorkspaceFingerprint: Schema.optional(Schema.String), }), Schema.Struct({ type: Schema.Literal("checkpoint.capture"), @@ -126,6 +128,12 @@ export const OrchestrationEffectStatusV2 = Schema.Literals([ ]); export type OrchestrationEffectStatusV2 = typeof OrchestrationEffectStatusV2.Type; +export const OrchestrationEffectFailureCodeV2 = Schema.Literals([ + "checkpoint_restore_rejected", + "checkpoint_restore_partial", +]); +export type OrchestrationEffectFailureCodeV2 = typeof OrchestrationEffectFailureCodeV2.Type; + export interface OrchestrationEffectV2 { readonly id: string; readonly commandId: CommandId; @@ -140,6 +148,7 @@ export interface OrchestrationEffectV2 { readonly updatedAt: string; readonly completedAt: string | null; readonly lastError: string | null; + readonly failureCode?: OrchestrationEffectFailureCodeV2; } export interface PendingOrchestrationEffectV2 { @@ -209,6 +218,7 @@ export interface EffectOutboxV2Shape { readonly effectId: string; readonly workerId: string; readonly error: string; + readonly failureCode?: OrchestrationEffectFailureCodeV2; }) => Effect.Effect; } @@ -237,11 +247,41 @@ const encodeRequest = Schema.encodeSync(Schema.fromJsonString(OrchestrationEffec const decodeRequest = Schema.decodeUnknownEffect( Schema.fromJsonString(OrchestrationEffectRequestV2), ); +const StoredEffectFailure = Schema.Struct({ + code: OrchestrationEffectFailureCodeV2, + message: Schema.String, +}); +const encodeStoredEffectFailure = Schema.encodeSync(Schema.fromJsonString(StoredEffectFailure)); +const decodeStoredEffectFailure = Schema.decodeUnknownOption( + Schema.fromJsonString(StoredEffectFailure), +); +const STORED_EFFECT_FAILURE_PREFIX = "t3-effect-failure:"; + +function decodeLastError(value: string | null): { + readonly lastError: string | null; + readonly failureCode?: OrchestrationEffectFailureCodeV2; +} { + if (value === null || !value.startsWith(STORED_EFFECT_FAILURE_PREFIX)) { + return { lastError: value }; + } + const decoded = decodeStoredEffectFailure(value.slice(STORED_EFFECT_FAILURE_PREFIX.length)); + return Option.match(decoded, { + onNone: () => ({ lastError: value }), + onSome: (failure) => ({ lastError: failure.message, failureCode: failure.code }), + }); +} + +function encodeLastError(error: string, failureCode?: OrchestrationEffectFailureCodeV2): string { + return failureCode === undefined + ? error + : `${STORED_EFFECT_FAILURE_PREFIX}${encodeStoredEffectFailure({ code: failureCode, message: error })}`; +} const rowToEffect = (row: EffectRow) => decodeRequest(row.payload_json).pipe( - Effect.map( - (request): OrchestrationEffectV2 => ({ + Effect.map((request): OrchestrationEffectV2 => { + const failure = decodeLastError(row.last_error); + return { id: row.effect_id, commandId: CommandId.make(row.command_id), threadId: ThreadId.make(row.thread_id), @@ -254,9 +294,10 @@ const rowToEffect = (row: EffectRow) => createdAt: row.created_at, updatedAt: row.updated_at, completedAt: row.completed_at, - lastError: row.last_error, - }), - ), + lastError: failure.lastError, + ...(failure.failureCode === undefined ? {} : { failureCode: failure.failureCode }), + }; + }), ); export const layer: Layer.Layer = Layer.effect( @@ -432,6 +473,39 @@ export const layer: Layer.Layer = La }), reconcileAfterProcessLoss: Effect.gen(function* () { const now = DateTime.formatIso(yield* DateTime.now); + const rollbackRows = yield* sql` + SELECT * + FROM orchestration_v2_effect_outbox + WHERE status = 'running' + AND effect_type = 'provider-thread.rollback' + `; + const rollbackEffects = yield* Effect.forEach(rollbackRows, rowToEffect); + const uncertainGuardedRollbackIds = rollbackEffects + .filter( + (effect) => + effect.request.type === "provider-thread.rollback" && + effect.request.expectedIdle === true && + effect.request.expectedWorkspaceFingerprint !== undefined, + ) + .map((effect) => effect.id); + if (uncertainGuardedRollbackIds.length > 0) { + const uncertainError = encodeLastError( + "The server process ended while a guarded checkpoint restore was running; its filesystem/provider outcome is uncertain and the restore was not repeated.", + "checkpoint_restore_partial", + ); + yield* sql` + UPDATE orchestration_v2_effect_outbox + SET + status = 'failed', + lease_owner = NULL, + lease_expires_at = NULL, + completed_at = ${now}, + updated_at = ${now}, + last_error = ${uncertainError} + WHERE effect_id IN ${sql.in(uncertainGuardedRollbackIds)} + AND status = 'running' + `; + } const cancelledRows = yield* sql<{ readonly effect_id: string }>` UPDATE orchestration_v2_effect_outbox SET @@ -580,9 +654,10 @@ export const layer: Layer.Layer = La (cause) => new EffectOutboxError({ operation: "retry", effectId, cause }), ), ), - fail: ({ effectId, workerId, error }) => + fail: ({ effectId, workerId, error, failureCode }) => Effect.gen(function* () { const now = DateTime.formatIso(yield* DateTime.now); + const storedError = encodeLastError(error, failureCode); const rows = yield* sql<{ readonly effect_id: string }>` UPDATE orchestration_v2_effect_outbox SET @@ -591,7 +666,7 @@ export const layer: Layer.Layer = La lease_expires_at = NULL, completed_at = ${now}, updated_at = ${now}, - last_error = ${error} + last_error = ${storedError} WHERE effect_id = ${effectId} AND status = 'running' AND lease_owner = ${workerId} diff --git a/apps/server/src/orchestration-v2/EffectWorker.test.ts b/apps/server/src/orchestration-v2/EffectWorker.test.ts index b39105f36f89..e3e8a729d63a 100644 --- a/apps/server/src/orchestration-v2/EffectWorker.test.ts +++ b/apps/server/src/orchestration-v2/EffectWorker.test.ts @@ -1,5 +1,7 @@ import { assert, it } from "@effect/vitest"; import { + CheckpointId, + CheckpointScopeId, CommandId, MessageId, ProviderSessionId, @@ -20,7 +22,11 @@ import * as Queue from "effect/Queue"; import * as Ref from "effect/Ref"; import * as TestClock from "effect/testing/TestClock"; -import { CheckpointRollbackServiceV2 } from "./CheckpointRollbackService.ts"; +import { + type CheckpointRollbackExecutionError, + CheckpointRollbackPartialError, + CheckpointRollbackServiceV2, +} from "./CheckpointRollbackService.ts"; import { EffectOutboxError, EffectOutboxV2, type OrchestrationEffectV2 } from "./EffectOutbox.ts"; import { executorLayer, @@ -84,9 +90,37 @@ function restartEffect( }; } +function rollbackEffect(now: DateTime.Utc, guarded: boolean): OrchestrationEffectV2 { + const timestamp = DateTime.formatIso(now); + return { + id: `effect:checkpoint-rollback:${guarded ? "guarded" : "legacy"}`, + commandId: CommandId.make(`command:checkpoint-rollback:${guarded ? "guarded" : "legacy"}`), + threadId, + request: { + type: "provider-thread.rollback", + providerThreadId, + checkpointId: CheckpointId.make("checkpoint:effect-worker-rollback"), + scopeId: CheckpointScopeId.make("scope:effect-worker-rollback"), + ...(guarded + ? { expectedIdle: true as const, expectedWorkspaceFingerprint: "workspace-before" } + : {}), + }, + status: "running", + attemptCount: 1, + availableAt: timestamp, + leaseOwner: "test-worker", + leaseExpiresAt: timestamp, + createdAt: timestamp, + updatedAt: timestamp, + completedAt: null, + lastError: null, + }; +} + function makeExecutorLayer(input: { readonly events: Ref.Ref>; readonly failFirstStart?: Ref.Ref; + readonly rollbackError?: CheckpointRollbackExecutionError; }) { const record = (event: string) => Ref.update(input.events, (events) => [...events, event]); const dependencies = Layer.mergeAll( @@ -138,7 +172,10 @@ function makeExecutorLayer(input: { ), Layer.succeed( CheckpointRollbackServiceV2, - CheckpointRollbackServiceV2.of({ execute: () => Effect.void }), + CheckpointRollbackServiceV2.of({ + execute: () => + input.rollbackError === undefined ? Effect.void : Effect.fail(input.rollbackError), + }), ), Layer.succeed( RuntimeRequestServiceV2, @@ -188,6 +225,192 @@ it("does not retry pure interrupt races where the turn is already gone", () => { ); }); +it.effect("classifies retry-unsafe rollback failures only for guarded MCP requests", () => + Effect.gen(function* () { + const now = yield* DateTime.now; + const events = yield* Ref.make>([]); + const rollbackError = new CheckpointRollbackPartialError({ + reason: "provider-rollback-failed-after-restore", + threadId, + providerThreadId, + checkpointId: CheckpointId.make("checkpoint:effect-worker-rollback"), + cause: "simulated provider failure", + }); + const layer = makeExecutorLayer({ events, rollbackError }); + const legacy = yield* OrchestrationEffectExecutorV2.pipe( + Effect.flatMap((executor) => executor.execute(rollbackEffect(now, false))), + Effect.provide(layer), + Effect.flip, + ); + const guarded = yield* OrchestrationEffectExecutorV2.pipe( + Effect.flatMap((executor) => executor.execute(rollbackEffect(now, true))), + Effect.provide(layer), + Effect.flip, + ); + + assert.isUndefined(legacy.failureCode); + assert.equal(guarded.failureCode, "checkpoint_restore_partial"); + }), +); + +it.effect("retries legacy rollback failures while terminalizing guarded partial outcomes", () => + Effect.gen(function* () { + const now = yield* DateTime.now; + + const runCase = ( + guarded: boolean, + failureCode?: "checkpoint_restore_rejected" | "checkpoint_restore_partial", + ) => + Effect.gen(function* () { + const effect = rollbackEffect(now, guarded); + const retries = yield* Ref.make(0); + const failures = yield* Ref.make>([]); + const outboxLayer = Layer.mock(EffectOutboxV2)({ + claimNext: () => Effect.succeed(Option.some(effect)), + get: () => Effect.succeed(Option.some(effect)), + awaitCancellation: () => Effect.never, + clearCancellation: () => Effect.void, + retry: () => Ref.update(retries, (count) => count + 1).pipe(Effect.as(true)), + fail: ({ failureCode }) => + Ref.update(failures, (codes) => [...codes, failureCode]).pipe(Effect.as(true)), + }); + const executorLayer = Layer.succeed( + OrchestrationEffectExecutorV2, + OrchestrationEffectExecutorV2.of({ + execute: () => + Effect.fail( + new OrchestrationEffectExecutionError({ + effectId: effect.id, + effectType: effect.request.type, + ...(failureCode === undefined ? {} : { failureCode }), + cause: "simulated rollback failure", + }), + ), + }), + ); + const workerLayer = effectWorkerLayerWithOptions({ workerId: "test-worker" }).pipe( + Layer.provide(Layer.merge(outboxLayer, executorLayer)), + ); + + assert.isTrue( + yield* OrchestrationEffectWorkerV2.pipe( + Effect.flatMap((worker) => worker.runOnce), + Effect.provide(workerLayer), + ), + ); + return { retries: yield* Ref.get(retries), failures: yield* Ref.get(failures) }; + }); + + assert.deepEqual(yield* runCase(false), { retries: 1, failures: [] }); + assert.deepEqual(yield* runCase(true), { retries: 1, failures: [] }); + assert.deepEqual(yield* runCase(true, "checkpoint_restore_partial"), { + retries: 0, + failures: ["checkpoint_restore_partial"], + }); + }), +); + +it.effect("retains guarded partial classification when the first fail settlement fails", () => + Effect.gen(function* () { + const now = yield* DateTime.now; + const effect = rollbackEffect(now, true); + const failCodes = yield* Ref.make>([]); + const retries = yield* Ref.make(0); + const outboxLayer = Layer.mock(EffectOutboxV2)({ + claimNext: () => Effect.succeed(Option.some(effect)), + get: () => Effect.succeed(Option.some(effect)), + awaitCancellation: () => Effect.never, + clearCancellation: () => Effect.void, + fail: ({ failureCode }) => + Ref.updateAndGet(failCodes, (codes) => [...codes, failureCode]).pipe( + Effect.flatMap((codes) => + codes.length === 1 + ? Effect.fail( + new EffectOutboxError({ + operation: "fail", + effectId: effect.id, + cause: "simulated first fail settlement failure", + }), + ) + : Effect.succeed(true), + ), + ), + retry: () => Ref.update(retries, (count) => count + 1).pipe(Effect.as(true)), + }); + const executorLayer = Layer.succeed( + OrchestrationEffectExecutorV2, + OrchestrationEffectExecutorV2.of({ + execute: () => + Effect.fail( + new OrchestrationEffectExecutionError({ + effectId: effect.id, + effectType: effect.request.type, + failureCode: "checkpoint_restore_partial", + cause: "simulated partial restore", + }), + ), + }), + ); + const workerLayer = effectWorkerLayerWithOptions({ workerId: "test-worker" }).pipe( + Layer.provide(Layer.merge(outboxLayer, executorLayer)), + ); + + yield* OrchestrationEffectWorkerV2.pipe( + Effect.flatMap((worker) => worker.runOnce), + Effect.provide(workerLayer), + Effect.exit, + ); + + assert.deepEqual(yield* Ref.get(failCodes), [ + "checkpoint_restore_partial", + "checkpoint_restore_partial", + ]); + assert.equal(yield* Ref.get(retries), 0); + }), +); + +it.effect("terminalizes an unknown guarded restore defect as partial without replay", () => + Effect.gen(function* () { + const now = yield* DateTime.now; + const effect = rollbackEffect(now, true); + const claims = yield* Ref.make(0); + const executions = yield* Ref.make(0); + const retries = yield* Ref.make(0); + const failCodes = yield* Ref.make>([]); + const outboxLayer = Layer.mock(EffectOutboxV2)({ + claimNext: () => + Ref.getAndUpdate(claims, (count) => count + 1).pipe( + Effect.map((count) => (count === 0 ? Option.some(effect) : Option.none())), + ), + get: () => Effect.succeed(Option.some(effect)), + awaitCancellation: () => Effect.never, + clearCancellation: () => Effect.void, + fail: ({ failureCode }) => + Ref.update(failCodes, (codes) => [...codes, failureCode]).pipe(Effect.as(true)), + retry: () => Ref.update(retries, (count) => count + 1).pipe(Effect.as(true)), + }); + const executorLayer = Layer.succeed( + OrchestrationEffectExecutorV2, + OrchestrationEffectExecutorV2.of({ + execute: () => + Ref.update(executions, (count) => count + 1).pipe( + Effect.andThen(Effect.die("simulated defect after filesystem mutation")), + ), + }), + ); + const workerLayer = effectWorkerLayerWithOptions({ workerId: "test-worker" }).pipe( + Layer.provide(Layer.merge(outboxLayer, executorLayer)), + ); + + const worker = yield* OrchestrationEffectWorkerV2.pipe(Effect.provide(workerLayer)); + assert.isTrue(yield* worker.runOnce); + assert.isFalse(yield* worker.runOnce); + assert.equal(yield* Ref.get(executions), 1); + assert.equal(yield* Ref.get(retries), 0); + assert.deepEqual(yield* Ref.get(failCodes), ["checkpoint_restore_partial"]); + }), +); + it.effect("requeues a claim when a pre-execution worker check fails", () => Effect.gen(function* () { const now = DateTime.formatIso(yield* DateTime.now); diff --git a/apps/server/src/orchestration-v2/EffectWorker.ts b/apps/server/src/orchestration-v2/EffectWorker.ts index 1139d2fea3b1..cba5c65b65b0 100644 --- a/apps/server/src/orchestration-v2/EffectWorker.ts +++ b/apps/server/src/orchestration-v2/EffectWorker.ts @@ -19,10 +19,14 @@ import { RunFinalizationService } from "./RunFinalizationService.ts"; import { ResourceCleanupService } from "./ResourceCleanupService.ts"; import { EffectOutboxV2, + OrchestrationEffectFailureCodeV2, REPLAY_SAFE_EFFECT_TYPES_AFTER_PROCESS_LOSS, type OrchestrationEffectV2, } from "./EffectOutbox.ts"; -import { CheckpointRollbackServiceV2 } from "./CheckpointRollbackService.ts"; +import { + type CheckpointRollbackExecutionError, + CheckpointRollbackServiceV2, +} from "./CheckpointRollbackService.ts"; import { ProviderSessionManagerV2 } from "./ProviderSessionManager.ts"; import { ProviderTurnControlServiceV2 } from "./ProviderTurnControlService.ts"; import { ProviderTurnStartServiceV2 } from "./ProviderTurnStartService.ts"; @@ -34,10 +38,40 @@ export class OrchestrationEffectExecutionError extends Schema.TaggedErrorClass - new OrchestrationEffectExecutionError({ - effectId: effect.id, - effectType: effect.request.type, - cause, - }), - ), + Effect.mapError((cause) => { + const failureCode: OrchestrationEffectFailureCodeV2 | undefined = guardedRestore + ? guardedCheckpointRestoreFailureCode(cause) + : undefined; + return new OrchestrationEffectExecutionError({ + effectId: effect.id, + effectType: effect.request.type, + ...(failureCode === undefined ? {} : { failureCode }), + cause, + }); + }), ); + } case "checkpoint.capture": return runFinalization .finalize({ @@ -402,13 +449,18 @@ export const layerWithOptions = ( }), ), ); - const terminalizeClaim = (effect: OrchestrationEffectV2, cause: Cause.Cause) => { + const terminalizeClaim = ( + effect: OrchestrationEffectV2, + cause: Cause.Cause, + failureCode?: OrchestrationEffectFailureCodeV2, + ) => { if (Cause.hasInterruptsOnly(cause)) return Effect.void; return outbox .fail({ effectId: effect.id, workerId, error: `Worker failed to settle a process-bound effect after execution started: ${Cause.pretty(cause)}`, + ...(failureCode === undefined ? {} : { failureCode }), }) .pipe( Effect.flatMap((failed) => @@ -441,11 +493,13 @@ export const layerWithOptions = ( effect: OrchestrationEffectV2, cause: Cause.Cause, ) => - REPLAY_SAFE_EFFECT_TYPES_AFTER_PROCESS_LOSS.some( - (effectType) => effectType === effect.request.type, - ) - ? requeueClaim(effect, cause) - : terminalizeClaim(effect, cause); + isGuardedCheckpointRestore(effect) + ? terminalizeClaim(effect, cause, "checkpoint_restore_partial") + : REPLAY_SAFE_EFFECT_TYPES_AFTER_PROCESS_LOSS.some( + (effectType) => effectType === effect.request.type, + ) + ? requeueClaim(effect, cause) + : terminalizeClaim(effect, cause); const runOnce = Effect.gen(function* () { const claimExit = yield* Effect.exit(outbox.claimNext({ workerId, leaseDurationMs })); @@ -517,7 +571,21 @@ export const layerWithOptions = ( } const error = Cause.pretty(exit.cause); + const executionError = Cause.findErrorOption(exit.cause).pipe( + Option.filter(isOrchestrationEffectExecutionError), + ); + const uncertainGuardedFailure = + isGuardedCheckpointRestore(effect) && + (Cause.hasDies(exit.cause) || + (!Cause.hasInterruptsOnly(exit.cause) && Option.isNone(executionError))); + const failureCode = uncertainGuardedFailure + ? "checkpoint_restore_partial" + : Option.isSome(executionError) + ? executionError.value.failureCode + : undefined; const nonRetryable = isNonRetryableProviderTurnControlFailure(effect.request.type, error); + const terminalRollbackFailure = + effect.request.type === "provider-thread.rollback" && failureCode !== undefined; yield* Effect.logWarning("Orchestration effect execution failed", { effectId: effect.id, effectType: effect.request.type, @@ -527,22 +595,31 @@ export const layerWithOptions = ( }); // Prefer succeed for terminal interrupt races so the outbox does not // keep a failed interrupt around; fail only when we must not retry. - const updated = nonRetryable + const updated = terminalRollbackFailure ? yield* outbox - .succeed({ effectId: effect.id, workerId }) - .pipe(Effect.onError((cause) => terminalizeClaim(effect, cause))) - : effect.attemptCount >= maxAttempts + .fail({ + effectId: effect.id, + workerId, + error, + ...(failureCode === undefined ? {} : { failureCode }), + }) + .pipe(Effect.onError((cause) => terminalizeClaim(effect, cause, failureCode))) + : nonRetryable ? yield* outbox - .fail({ effectId: effect.id, workerId, error }) + .succeed({ effectId: effect.id, workerId }) .pipe(Effect.onError((cause) => terminalizeClaim(effect, cause))) - : yield* outbox - .retry({ - effectId: effect.id, - workerId, - error, - delayMs: Math.min(30_000, 100 * 2 ** Math.max(0, effect.attemptCount - 1)), - }) - .pipe(Effect.onError((cause) => requeueClaim(effect, cause))); + : effect.attemptCount >= maxAttempts + ? yield* outbox + .fail({ effectId: effect.id, workerId, error }) + .pipe(Effect.onError((cause) => terminalizeClaim(effect, cause))) + : yield* outbox + .retry({ + effectId: effect.id, + workerId, + error, + delayMs: Math.min(30_000, 100 * 2 ** Math.max(0, effect.attemptCount - 1)), + }) + .pipe(Effect.onError((cause) => requeueClaim(effect, cause))); if (!updated) { if (yield* wasCancelled(effect.id)) return true; return yield* new OrchestrationEffectWorkerError({ diff --git a/apps/server/src/orchestration-v2/FoundationPersistence.test.ts b/apps/server/src/orchestration-v2/FoundationPersistence.test.ts index fdb244174846..7660d4c78583 100644 --- a/apps/server/src/orchestration-v2/FoundationPersistence.test.ts +++ b/apps/server/src/orchestration-v2/FoundationPersistence.test.ts @@ -1823,10 +1823,99 @@ it.layer(TestLayer)("orchestration V2 foundation persistence", (it) => { if (Option.isSome(reclaimed)) { assert.equal(reclaimed.value.request.type, "terminal.cleanup"); assert.equal(reclaimed.value.attemptCount, 2); + assert.isTrue( + yield* outbox.succeed({ + effectId: reclaimed.value.id, + workerId: "recovery-worker", + }), + ); } }), ); + it.effect("does not repeat a guarded checkpoint restore after process loss", () => + Effect.gen(function* () { + const outbox = yield* EffectOutboxV2; + const commandId = CommandId.make("command:foundation-rollback-process-loss"); + const rollbackThreadId = ThreadId.make("thread:foundation-rollback-process-loss"); + const rollbackProviderThreadId = ProviderThreadId.make( + "provider-thread:foundation-rollback-process-loss", + ); + const checkpointId = CheckpointId.make("checkpoint:foundation-rollback-process-loss"); + const scopeId = CheckpointScopeId.make("scope:foundation-rollback-process-loss"); + yield* outbox.enqueue([ + { + id: "effect:a-foundation-guarded-rollback", + commandId, + threadId: rollbackThreadId, + request: { + type: "provider-thread.rollback", + providerThreadId: rollbackProviderThreadId, + checkpointId, + scopeId, + expectedIdle: true, + expectedWorkspaceFingerprint: "workspace-before-restore", + }, + }, + { + id: "effect:b-foundation-legacy-rollback", + commandId, + threadId: ThreadId.make("thread:foundation-legacy-rollback-process-loss"), + request: { + type: "provider-thread.rollback", + providerThreadId: rollbackProviderThreadId, + checkpointId, + scopeId, + }, + }, + ]); + assert.isTrue( + Option.isSome( + yield* outbox.claimNext({ workerId: "crashed-worker", leaseDurationMs: 30_000 }), + ), + ); + assert.isTrue( + Option.isSome( + yield* outbox.claimNext({ workerId: "crashed-worker", leaseDurationMs: 30_000 }), + ), + ); + + const runningGuarded = yield* outbox.get("effect:a-foundation-guarded-rollback"); + assert.isTrue(Option.isSome(runningGuarded)); + if ( + Option.isSome(runningGuarded) && + runningGuarded.value.request.type === "provider-thread.rollback" + ) { + assert.equal(runningGuarded.value.status, "running"); + assert.isTrue(runningGuarded.value.request.expectedIdle); + assert.equal( + runningGuarded.value.request.expectedWorkspaceFingerprint, + "workspace-before-restore", + ); + } + + assert.deepEqual(yield* outbox.reconcileAfterProcessLoss, { + cancelled: 0, + requeued: 1, + }); + const guarded = yield* outbox.get("effect:a-foundation-guarded-rollback"); + assert.isTrue(Option.isSome(guarded)); + if (Option.isSome(guarded)) { + assert.equal(guarded.value.status, "failed"); + assert.equal(guarded.value.failureCode, "checkpoint_restore_partial"); + assert.include(guarded.value.lastError ?? "", "outcome is uncertain"); + } + const recovered = yield* outbox.claimNext({ + workerId: "recovery-worker", + leaseDurationMs: 30_000, + }); + assert.isTrue(Option.isSome(recovered)); + if (Option.isSome(recovered)) { + assert.equal(recovered.value.id, "effect:b-foundation-legacy-rollback"); + } + }), + ); + it.effect("atomically cancels stale runs and their process-bound effects", () => Effect.gen(function* () { const eventSink = yield* EventSinkV2; diff --git a/apps/server/src/orchestration-v2/KeyedSerialExecutor.ts b/apps/server/src/orchestration-v2/KeyedSerialExecutor.ts index 93dbe04bbbfd..ef8d6b8b4d48 100644 --- a/apps/server/src/orchestration-v2/KeyedSerialExecutor.ts +++ b/apps/server/src/orchestration-v2/KeyedSerialExecutor.ts @@ -1,4 +1,7 @@ +import { ThreadId } from "@t3tools/contracts"; +import * as Context from "effect/Context"; import * as Effect from "effect/Effect"; +import * as Layer from "effect/Layer"; import * as Ref from "effect/Ref"; import * as Semaphore from "effect/Semaphore"; @@ -53,3 +56,14 @@ export const makeKeyedSerialExecutor = (): Effect.Effect; }); + +/** The single in-process admission boundary for commands targeting a V2 thread. */ +export class ThreadDispatchLockV2 extends Context.Service< + ThreadDispatchLockV2, + KeyedSerialExecutor +>()("t3/orchestration-v2/KeyedSerialExecutor/ThreadDispatchLockV2") {} + +export const threadDispatchLockLayer = Layer.effect( + ThreadDispatchLockV2, + makeKeyedSerialExecutor(), +); diff --git a/apps/server/src/orchestration-v2/Orchestrator.ts b/apps/server/src/orchestration-v2/Orchestrator.ts index 6c5ae7bc7b58..f87af568e300 100644 --- a/apps/server/src/orchestration-v2/Orchestrator.ts +++ b/apps/server/src/orchestration-v2/Orchestrator.ts @@ -1,5 +1,7 @@ import { type ChatAttachment, + checkpointRollbackAppRunOrdinal, + CheckpointId, CommandId, type MessageId, type ModelSelection, @@ -44,9 +46,13 @@ import { CommandPolicyV2 } from "./CommandPolicy.ts"; import { CommandReceiptStoreV2 } from "./CommandReceiptStore.ts"; import { ContextHandoffServiceV2 } from "./ContextHandoffService.ts"; import { EventSinkV2 } from "./EventSink.ts"; -import type { OrchestrationEffectRequestV2, PendingOrchestrationEffectV2 } from "./EffectOutbox.ts"; +import { + EffectOutboxV2, + type OrchestrationEffectRequestV2, + type PendingOrchestrationEffectV2, +} from "./EffectOutbox.ts"; import { IdAllocatorV2 } from "./IdAllocator.ts"; -import { makeKeyedSerialExecutor } from "./KeyedSerialExecutor.ts"; +import { ThreadDispatchLockV2 } from "./KeyedSerialExecutor.ts"; import { applyToProjection, emptyProjection, @@ -143,6 +149,31 @@ export class OrchestratorCommandIdConflictError extends Schema.TaggedErrorClass< } } +export class OrchestratorCheckpointRollbackNotIdleError extends Schema.TaggedErrorClass()( + "OrchestratorCheckpointRollbackNotIdleError", + { + commandId: CommandId, + threadId: ThreadId, + }, +) { + override get message(): string { + return `Checkpoint rollback ${this.commandId} requires idle thread ${this.threadId} with no queued runs.`; + } +} + +export class OrchestratorCheckpointRollbackTargetUnsupportedError extends Schema.TaggedErrorClass()( + "OrchestratorCheckpointRollbackTargetUnsupportedError", + { + commandId: CommandId, + threadId: ThreadId, + checkpointId: CheckpointId, + }, +) { + override get message(): string { + return `Checkpoint ${this.checkpointId} does not identify a proven provider-history target for thread ${this.threadId}.`; + } +} + /** * A command receipt only proves that this exact command already ran for the * thread it was recorded against. Replaying it for a command aimed at another @@ -163,6 +194,8 @@ export const OrchestratorV2Error = Schema.Union([ OrchestratorProviderAdapterError, OrchestratorCommandPreviouslyRejectedError, OrchestratorCommandIdConflictError, + OrchestratorCheckpointRollbackNotIdleError, + OrchestratorCheckpointRollbackTargetUnsupportedError, ]); export type OrchestratorV2Error = typeof OrchestratorV2Error.Type; @@ -196,6 +229,8 @@ export interface OrchestratorV2Shape { readonly getThreadEventSequence: ( threadId: ThreadId, ) => Effect.Effect; + readonly getCommandReceipt: CommandReceiptStoreV2["Service"]["getByCommandId"]; + readonly listCommandEffects: EffectOutboxV2["Service"]["listByCommandId"]; readonly streamStoredEvents: Stream.Stream; readonly streamStoredEventsFrom: (input?: { readonly threadId?: ThreadId; @@ -513,6 +548,7 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio const contextHandoffService = yield* ContextHandoffServiceV2; const eventSink = yield* EventSinkV2; const commandReceipts = yield* CommandReceiptStoreV2; + const effectOutbox = yield* EffectOutboxV2; const idAllocator = yield* IdAllocatorV2; const projectionStore = yield* ProjectionStoreV2; const providerAdapters = yield* ProviderAdapterRegistryV2; @@ -521,7 +557,7 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio const providerSwitchService = yield* ProviderSwitchServiceV2; const runtimePolicy = yield* RuntimePolicyV2; const threadForkService = yield* ThreadForkServiceV2; - const threadDispatch = yield* makeKeyedSerialExecutor(); + const threadDispatch = yield* ThreadDispatchLockV2; const mapDispatchError = (command: OrchestrationV2Command) => @@ -5996,6 +6032,17 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio ) => Effect.gen(function* () { const projection = yield* loadProjectionForCommand(command); + if ( + command.expectedIdle === true && + projection.runs.some((run) => + ["preparing", "queued", "starting", "running", "waiting"].includes(run.status), + ) + ) { + return yield* new OrchestratorCheckpointRollbackNotIdleError({ + commandId: command.commandId, + threadId: command.threadId, + }); + } const providerThread = projection.providerThreads.find( (candidate) => candidate.id === projection.thread.activeProviderThreadId, ); @@ -6062,7 +6109,14 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio cause: `Checkpoint ${command.checkpointId} belongs to scope ${targetScope.id}, not ${command.scopeId}.`, }); } - const targetOrdinal = targetCheckpoint.appRunOrdinal ?? 0; + const targetOrdinal = checkpointRollbackAppRunOrdinal(targetCheckpoint, targetScope); + if (targetOrdinal === null) { + return yield* new OrchestratorCheckpointRollbackTargetUnsupportedError({ + commandId: command.commandId, + threadId: command.threadId, + checkpointId: targetCheckpoint.id, + }); + } if (targetOrdinal > 0) { const targetRun = projection.runs.find((run) => run.ordinal === targetOrdinal); const targetProviderTurn = @@ -6109,6 +6163,10 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio providerThreadId: providerThread.id, checkpointId: targetCheckpoint.id, scopeId: targetScope.id, + ...(command.expectedIdle === undefined ? {} : { expectedIdle: command.expectedIdle }), + ...(command.expectedWorkspaceFingerprint === undefined + ? {} + : { expectedWorkspaceFingerprint: command.expectedWorkspaceFingerprint }), }, } satisfies PendingOrchestrationEffectV2, ]); @@ -7169,6 +7227,8 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio eventSink .latestSequence({ threadId }) .pipe(Effect.mapError((cause) => new OrchestratorProjectionError({ threadId, cause }))), + getCommandReceipt: commandReceipts.getByCommandId, + listCommandEffects: effectOutbox.listByCommandId, streamStoredEvents: eventSink.stream().pipe( Stream.mapError( (cause) => @@ -7214,7 +7274,9 @@ export const layer: Layer.Layer< | CommandReceiptStoreV2 | ContextHandoffServiceV2 | EventSinkV2 + | EffectOutboxV2 | IdAllocatorV2 + | ThreadDispatchLockV2 | ProviderAdapterRegistryV2 | ProviderSessionManagerV2 | ProviderSwitchServiceV2 @@ -7276,6 +7338,8 @@ export const layerUnavailable: Layer.Layer = Layer.succeed( cause: "Orchestration V2 live runtime is not configured.", }), ), + getCommandReceipt: () => Effect.succeed(Option.none()), + listCommandEffects: () => Effect.succeed([]), streamStoredEvents: Stream.fail( new OrchestratorDomainEventStreamError({ cause: "Orchestration V2 live runtime is not configured.", diff --git a/apps/server/src/orchestration-v2/ThreadManagementService.ts b/apps/server/src/orchestration-v2/ThreadManagementService.ts index 77e52bffbbbe..7b14f1a4831a 100644 --- a/apps/server/src/orchestration-v2/ThreadManagementService.ts +++ b/apps/server/src/orchestration-v2/ThreadManagementService.ts @@ -297,6 +297,8 @@ export interface ThreadManagementServiceShape { input: ThreadManagementInterruptInput, ) => Effect.Effect; readonly getThreadEventSequence: OrchestratorV2["Service"]["getThreadEventSequence"]; + readonly getCommandReceipt: OrchestratorV2["Service"]["getCommandReceipt"]; + readonly listCommandEffects: OrchestratorV2["Service"]["listCommandEffects"]; readonly streamStoredEvents: OrchestratorV2["Service"]["streamStoredEvents"]; readonly streamStoredEventsFrom: OrchestratorV2["Service"]["streamStoredEventsFrom"]; readonly streamDomainEvents: OrchestratorV2["Service"]["streamDomainEvents"]; @@ -655,6 +657,8 @@ const make = Effect.gen(function* () { waitForThread, interruptThread, getThreadEventSequence: orchestrator.getThreadEventSequence, + getCommandReceipt: orchestrator.getCommandReceipt, + listCommandEffects: orchestrator.listCommandEffects, streamStoredEvents: orchestrator.streamStoredEvents, streamStoredEventsFrom: orchestrator.streamStoredEventsFrom, streamDomainEvents: orchestrator.streamDomainEvents, diff --git a/apps/server/src/orchestration-v2/runtimeLayer.ts b/apps/server/src/orchestration-v2/runtimeLayer.ts index 6f4df38e1d90..f46a9748808b 100644 --- a/apps/server/src/orchestration-v2/runtimeLayer.ts +++ b/apps/server/src/orchestration-v2/runtimeLayer.ts @@ -22,6 +22,7 @@ import { layerFromStores as eventSinkLayer } from "./EventSink.ts"; import { layerFromOrchestrationEventStore as eventStoreLayer } from "./EventStore.ts"; import { layer as idAllocatorLayer } from "./IdAllocator.ts"; import { layer as legacyV1ThreadImporterLayer } from "./LegacyV1ThreadImporter.ts"; +import { threadDispatchLockLayer } from "./KeyedSerialExecutor.ts"; import { layer as orchestratorLayer } from "./Orchestrator.ts"; import { layer as projectionStoreLayer } from "./ProjectionStore.ts"; import { layer as projectionMaintenanceLayer } from "./ProjectionMaintenance.ts"; @@ -141,6 +142,7 @@ const checkpointRollbackServiceProvided = checkpointRollbackServiceLayer.pipe( projectionStoreLayer, providerSessionManagerProvided, runtimePolicyProvided, + threadDispatchLockLayer, ), ), ); @@ -178,6 +180,7 @@ const orchestratorProvided = orchestratorLayer.pipe( providerSwitchServiceProvided, runExecutionServiceProvided, threadForkServiceLayer, + threadDispatchLockLayer, ), ), ); diff --git a/apps/server/src/orchestration-v2/testkit/ProviderReplayHarness.ts b/apps/server/src/orchestration-v2/testkit/ProviderReplayHarness.ts index 4380f53926c7..492f9d8efb75 100644 --- a/apps/server/src/orchestration-v2/testkit/ProviderReplayHarness.ts +++ b/apps/server/src/orchestration-v2/testkit/ProviderReplayHarness.ts @@ -19,6 +19,7 @@ import * as VcsProcess from "../../vcs/VcsProcess.ts"; import { layer as checkpointCaptureServiceLayer } from "../CheckpointCaptureService.ts"; import { layer as checkpointServiceLayer } from "../CheckpointService.ts"; import { layer as checkpointRollbackServiceLayer } from "../CheckpointRollbackService.ts"; +import { threadDispatchLockLayer } from "../KeyedSerialExecutor.ts"; import { layer as commandPolicyLayer } from "../CommandPolicy.ts"; import { layer as commandReceiptStoreLayer } from "../CommandReceiptStore.ts"; import { layer as contextHandoffServiceLayer } from "../ContextHandoffService.ts"; @@ -343,6 +344,7 @@ export function makeOrchestratorV2ReplayLayerWithRegistry( storesLayer, providerSessionManagerProvided, runtimeLayer, + threadDispatchLockLayer, ), ), ); @@ -387,6 +389,7 @@ export function makeOrchestratorV2ReplayLayerWithRegistry( providerSwitchServiceProvided, runExecutionServiceProvided, threadForkServiceLayer, + threadDispatchLockLayer, ), ), ); diff --git a/apps/server/src/vcs/GitVcsDriver.test.ts b/apps/server/src/vcs/GitVcsDriver.test.ts index 89f7c55d5863..7e93037eb193 100644 --- a/apps/server/src/vcs/GitVcsDriver.test.ts +++ b/apps/server/src/vcs/GitVcsDriver.test.ts @@ -108,3 +108,40 @@ it.effect("GitVcsDriver forwards execute env to the VCS process", () => { ), ); }); + +it.effect("GitVcsDriver rejects a fingerprint built from truncated staged-state output", () => + Effect.gen(function* () { + const driver = yield* GitVcsDriver.makeVcsDriverShape(); + const checkpoints = driver.checkpoints; + assert.ok(checkpoints); + + const error = yield* checkpoints.readWorkspaceFingerprint("/repo").pipe(Effect.flip); + + assert.equal(error._tag, "VcsProcessExitError"); + if (error._tag === "VcsProcessExitError") { + assert.match(error.detail, /incomplete staged-state output/); + } + }).pipe( + Effect.provide( + Layer.mergeAll( + NodeServices.layer, + Layer.mock(VcsProcess.VcsProcess)({ + run: (input) => { + const args = input.args.join(" "); + return Effect.succeed({ + exitCode: ChildProcessSpawner.ExitCode(0), + stdout: args.includes("--git-common-dir") + ? ".git\n" + : args.includes("write-tree") + ? "tree-oid\n" + : "", + stderr: "", + stdoutTruncated: args.includes("ls-files --stage"), + stderrTruncated: false, + }); + }, + }), + ), + ), + ), +); diff --git a/apps/server/src/vcs/GitVcsDriver.ts b/apps/server/src/vcs/GitVcsDriver.ts index 9e8b9615ac9d..22f9eaf17c22 100644 --- a/apps/server/src/vcs/GitVcsDriver.ts +++ b/apps/server/src/vcs/GitVcsDriver.ts @@ -717,7 +717,60 @@ export const makeVcsDriverShape = Effect.fn("makeGitVcsDriverShape")(function* ( return path.isAbsolute(gitCommonDir) ? gitCommonDir : path.resolve(cwd, gitCommonDir); }); + const readWorkspaceFingerprint = Effect.fn("GitVcsDriver.checkpoints.readWorkspaceFingerprint")( + function* (cwd: string, operation = "GitVcsDriver.checkpoints.readWorkspaceFingerprint") { + const gitCommonDir = yield* resolveGitCommonDir(cwd); + const tempIndexPath = path.join( + gitCommonDir, + `t3-checkpoint-index-${NodeCrypto.randomUUID()}`, + ); + const indexEnv: NodeJS.ProcessEnv = { ...process.env, GIT_INDEX_FILE: tempIndexPath }; + const cleanupTempIndex = fileSystem + .remove(tempIndexPath, { force: true }) + .pipe(Effect.ignore); + + return yield* Effect.gen(function* () { + if (yield* hasHeadCommit(cwd)) { + yield* execute({ operation, cwd, args: ["read-tree", "HEAD"], env: indexEnv }); + } + yield* execute({ operation, cwd, args: ["add", "-A", "--", "."], env: indexEnv }); + const result = yield* execute({ operation, cwd, args: ["write-tree"], env: indexEnv }); + const treeOid = result.stdout.trim(); + if (treeOid.length === 0) { + return yield* new VcsProcessExitError({ + operation, + command: "git write-tree", + cwd, + exitCode: 0, + detail: "git write-tree returned an empty tree oid.", + }); + } + const index = yield* execute({ + operation, + cwd, + args: ["ls-files", "--stage", "-z", "--", "."], + }); + if (index.stdoutTruncated) { + return yield* new VcsProcessExitError({ + operation, + command: "git ls-files --stage", + cwd, + exitCode: index.exitCode, + detail: "git ls-files returned incomplete staged-state output.", + }); + } + return NodeCrypto.createHash("sha256") + .update("worktree\0") + .update(treeOid) + .update("\0index\0") + .update(index.stdout) + .digest("hex"); + }).pipe(Effect.ensuring(cleanupTempIndex)); + }, + ); + const checkpoints: VcsDriver.VcsCheckpointOps = { + readWorkspaceFingerprint: (cwd) => readWorkspaceFingerprint(cwd), captureCheckpoint: Effect.fn("GitVcsDriver.checkpoints.captureCheckpoint")(function* (input) { const operation = "GitVcsDriver.checkpoints.captureCheckpoint"; const gitCommonDir = yield* resolveGitCommonDir(input.cwd); diff --git a/apps/server/src/vcs/VcsDriver.ts b/apps/server/src/vcs/VcsDriver.ts index f2daf7935027..7124897a2362 100644 --- a/apps/server/src/vcs/VcsDriver.ts +++ b/apps/server/src/vcs/VcsDriver.ts @@ -39,6 +39,8 @@ export interface VcsDeleteCheckpointRefsInput { } export interface VcsCheckpointOps { + /** Hash the workspace tree and index entries that checkpoint restore can overwrite. */ + readonly readWorkspaceFingerprint: (cwd: string) => Effect.Effect; readonly captureCheckpoint: (input: VcsCaptureCheckpointInput) => Effect.Effect; readonly hasCheckpointRef: ( input: Omit, diff --git a/docs/orchestration-v2/orchestrator-mcp-server.md b/docs/orchestration-v2/orchestrator-mcp-server.md index efe90b1bf1fe..4570d39492df 100644 --- a/docs/orchestration-v2/orchestrator-mcp-server.md +++ b/docs/orchestration-v2/orchestrator-mcp-server.md @@ -331,6 +331,43 @@ code-unit cursors and report total length, truncation, and the next cursor. A page can exceed its requested limit by one code unit to avoid splitting a surrogate pair. This makes pagination match MCP JSON string indexing. +### `t3_checkpoint_restore` + +Requests the ordinary serialized V2 `checkpoint.rollback` command for one +exact checkpoint/scope identity. The tool never runs Git restore directly. +Because the restore discards current tracked and untracked workspace changes, +the caller must pass `discardChanges: true` and a well-formed, exact +`clientRequestId`. + +Admission requires an idle thread with no queued run, a ready checkpoint ref, +the active provider thread/session, and provider conversation rollback with a +returned snapshot. The command carries optional constraints understood by old +servers as absent: `expectedIdle` is enforced inside the thread's serialized +decision. The worker acquires that same per-thread admission boundary, re-reads +idle, archive, provider, checkpoint, and provider-history target state, then +compares the workspace-and-index fingerprint inside the existing per-workspace +checkpoint lock. New same-thread commands wait until guarded restore +finalization releases the admission boundary. An unrelated external process +can still write while Git is executing; T3 does not claim an OS-wide lock. + +The result includes the durable command receipt and current outbox status: + +- `REQUESTED`: accepted, but the rollback effect is pending, running, or its + current observation is temporarily unavailable; +- `APPLIED`: filesystem and required provider rollback completed; +- `FAILED`: a guard rejected the restore before mutation or an unguarded + request failed; and +- `PARTIAL`: filesystem/provider state may have changed, but the complete + durable outcome could not be recorded. + +Reusing the exact key returns the original command/effect state. A key already +accepted for another target is rejected. Guarded MCP rollback failures become +terminal after a retry-unsafe or uncertain phase, so a normal retry cannot +repeat filesystem or provider side effects. If the process ends while such an +effect is running, durable outbox recovery records `PARTIAL` instead of +replaying it. Older unguarded rollback commands retain their existing retry +behavior. + ## Delegated Task Lifecycle The MCP server is a command ingress into V2. It does not call provider adapters @@ -377,6 +414,8 @@ falls back to a terminal-status message when no assistant text exists. - Checkpoint list and diff use the same current-project boundary. They never accept an environment, workspace path, or raw checkpoint ref from the caller. +- Restore uses the same boundary and never expands privileges, substitutes a + different checkpoint, stashes files, or creates a backup workspace. - Provider instances must be enabled, installed, available, authenticated, and backed by a V2 adapter. - A requested model must be advertised by the selected provider when the diff --git a/docs/user/checkpoints.md b/docs/user/checkpoints.md index 7dc3a92779db..e3076584f624 100644 --- a/docs/user/checkpoints.md +++ b/docs/user/checkpoints.md @@ -20,3 +20,21 @@ path, or another project's thread. Inspection is read-only. A checkpoint that is stale, missing, or unavailable is reported honestly rather than substituted with a different snapshot. + +## Restore safety + +`t3_checkpoint_restore` restores one exact checkpoint selected from the list. +It is destructive: current tracked, untracked, and staged changes covered by the +restore are discarded, so the agent must explicitly acknowledge that outcome. +The thread must be idle with no queued work, and the provider must support +rolling its conversation back to the same point. + +Before the locked restore begins, T3 verifies that the covered workspace files +and thread state still match the accepted request. If either changed while the +request was waiting, the restore fails and preserves the newer state. As with +other filesystem commands, an unrelated process can still write while Git is +executing. The result distinguishes a request that is still running, a fully +applied restore, a failure, and a partial result where files were restored but +the provider conversation or durable finalization could not be completed. +Retrying with the same idempotency key reads the original result instead of +starting another restore. diff --git a/packages/contracts/src/checkpointMcp.test.ts b/packages/contracts/src/checkpointMcp.test.ts index e2abba1e50a9..91f77518e9b6 100644 --- a/packages/contracts/src/checkpointMcp.test.ts +++ b/packages/contracts/src/checkpointMcp.test.ts @@ -5,12 +5,14 @@ import { CheckpointId, CheckpointMcpDiffInput, CheckpointMcpListInput, + CheckpointMcpRestoreInput, CheckpointScopeId, ThreadId, } from "./index.ts"; const decodeListInput = Schema.decodeUnknownSync(CheckpointMcpListInput); const decodeDiffInput = Schema.decodeUnknownSync(CheckpointMcpDiffInput); +const decodeRestoreInput = Schema.decodeUnknownSync(CheckpointMcpRestoreInput); describe("checkpoint MCP contracts", () => { it("decodes bounded list and diff inputs from the old empty/default shape", () => { @@ -36,4 +38,22 @@ describe("checkpoint MCP contracts", () => { }), ).toThrow(); }); + + it("rejects malformed UTF-16 idempotency keys without normalizing valid keys", () => { + const base = { + scopeId: CheckpointScopeId.make("scope:checkpoint-contract"), + checkpointId: CheckpointId.make("checkpoint:checkpoint-contract"), + discardChanges: true, + } as const; + expect(() => decodeRestoreInput({ ...base, clientRequestId: "bad\ud800key" })).toThrow(); + expect(() => decodeRestoreInput({ ...base, clientRequestId: "bad\ud800" })).toThrow(); + expect(() => decodeRestoreInput({ ...base, clientRequestId: "bad\udc00key" })).toThrow(); + expect( + decodeRestoreInput({ ...base, clientRequestId: "valid \ud83d\ude80" }).clientRequestId, + ).toBe("valid \ud83d\ude80"); + expect(decodeRestoreInput({ ...base, clientRequestId: " e\u0301 " }).clientRequestId).toBe( + " e\u0301 ", + ); + expect(decodeRestoreInput({ ...base, clientRequestId: " é " }).clientRequestId).toBe(" é "); + }); }); diff --git a/packages/contracts/src/checkpointMcp.ts b/packages/contracts/src/checkpointMcp.ts index 13fbfdd7bb3a..0c477b064c1a 100644 --- a/packages/contracts/src/checkpointMcp.ts +++ b/packages/contracts/src/checkpointMcp.ts @@ -1,8 +1,10 @@ import * as Schema from "effect/Schema"; +import * as SchemaIssue from "effect/SchemaIssue"; import { CheckpointId, CheckpointScopeId, + CommandId, IsoDateTime, NodeId, NonNegativeInt, @@ -146,6 +148,85 @@ export const CheckpointMcpDiffResult = Schema.Struct({ }); export type CheckpointMcpDiffResult = typeof CheckpointMcpDiffResult.Type; +const isWellFormedUtf16 = (value: string) => { + for (let index = 0; index < value.length; index += 1) { + const codeUnit = value.charCodeAt(index); + if (codeUnit >= 0xd800 && codeUnit <= 0xdbff) { + const nextCodeUnit = value.charCodeAt(index + 1); + if (!(nextCodeUnit >= 0xdc00 && nextCodeUnit <= 0xdfff)) return false; + index += 1; + } else if (codeUnit >= 0xdc00 && codeUnit <= 0xdfff) { + return false; + } + } + return true; +}; + +export const CheckpointMcpClientRequestId = Schema.String.check( + Schema.makeFilter( + (value) => + (value.length > 0 && value.length <= 256 && isWellFormedUtf16(value)) || + new SchemaIssue.InvalidValue({ + message: "clientRequestId must contain 1-256 well-formed UTF-16 code units", + }), + { identifier: "CheckpointMcpClientRequestId" }, + ), +).annotate({ + description: + "Exact idempotency key to reuse when retrying this restore. It is not trimmed or Unicode-normalized.", +}); +export type CheckpointMcpClientRequestId = typeof CheckpointMcpClientRequestId.Type; + +export const CheckpointMcpRestoreInput = Schema.Struct({ + threadId: Schema.optional(ThreadId).annotate({ + description: + "Thread to restore. Defaults to the calling thread and must belong to its project.", + }), + scopeId: CheckpointScopeId.annotate({ + description: "Exact durable checkpoint scope returned by t3_checkpoint_list.", + }), + checkpointId: CheckpointId.annotate({ + description: "Exact durable checkpoint returned by t3_checkpoint_list.", + }), + discardChanges: Schema.Literal(true).annotate({ + description: + "Required acknowledgement that applying the checkpoint discards current tracked, untracked and staged workspace changes covered by the checkpoint restore.", + }), + clientRequestId: CheckpointMcpClientRequestId, +}); +export type CheckpointMcpRestoreInput = typeof CheckpointMcpRestoreInput.Type; + +export const CheckpointMcpRestoreStatus = Schema.Literals([ + "REQUESTED", + "APPLIED", + "FAILED", + "PARTIAL", +]); +export type CheckpointMcpRestoreStatus = typeof CheckpointMcpRestoreStatus.Type; + +export const CheckpointMcpRestoreResult = Schema.Struct({ + commandId: CommandId, + threadId: ThreadId, + scopeId: CheckpointScopeId, + checkpointId: CheckpointId, + status: CheckpointMcpRestoreStatus, + receipt: Schema.Struct({ + status: Schema.Literal("accepted"), + acceptedAt: IsoDateTime, + sequence: NonNegativeInt, + }), + effectStatus: Schema.Literals([ + "pending", + "running", + "succeeded", + "failed", + "cancelled", + "unavailable", + ]), + detail: Schema.NullOr(Schema.String), +}); +export type CheckpointMcpRestoreResult = typeof CheckpointMcpRestoreResult.Type; + export class CheckpointMcpFailure extends Schema.TaggedErrorClass()( "CheckpointMcpFailure", { @@ -155,6 +236,8 @@ export class CheckpointMcpFailure extends Schema.TaggedErrorClass { expect(event.payload.id).toBe(RunId.make("run-1")); }); + it("keeps checkpoint rollback commands from older clients decodable", () => { + const command = decodeOrchestrationV2Command({ + type: "checkpoint.rollback", + commandId: "command-checkpoint-rollback-legacy", + threadId: "thread-1", + scopeId: "scope-1", + checkpointId: "checkpoint-1", + }); + + expect(command.type).toBe("checkpoint.rollback"); + if (command.type !== "checkpoint.rollback") { + throw new Error("expected checkpoint.rollback"); + } + expect(command.expectedIdle).toBeUndefined(); + expect(command.expectedWorkspaceFingerprint).toBeUndefined(); + }); + it("decodes app-owned delegated task commands", () => { const command = decodeOrchestrationV2Command({ type: "delegated_task.request", diff --git a/packages/contracts/src/orchestrationV2.ts b/packages/contracts/src/orchestrationV2.ts index 3866758bf5bb..c7910baaca97 100644 --- a/packages/contracts/src/orchestrationV2.ts +++ b/packages/contracts/src/orchestrationV2.ts @@ -2236,6 +2236,9 @@ export const OrchestrationV2Command = Schema.Union([ threadId: ThreadId, scopeId: CheckpointScopeId, checkpointId: CheckpointId, + /** Optional MCP admission constraints. Older clients omit both fields. */ + expectedIdle: Schema.optional(Schema.Literal(true)), + expectedWorkspaceFingerprint: Schema.optional(TrimmedNonEmptyString), }), Schema.Struct({ type: Schema.Literal("thread.fork"), diff --git a/packages/shared/src/t3McpToolPresentation.test.ts b/packages/shared/src/t3McpToolPresentation.test.ts index f3af59729187..d6a7f427cd57 100644 --- a/packages/shared/src/t3McpToolPresentation.test.ts +++ b/packages/shared/src/t3McpToolPresentation.test.ts @@ -44,6 +44,10 @@ describe("resolveT3McpToolPresentation", () => { displayName: "Read a checkpoint diff", logo: "t3-code", }); + expect(resolveT3McpToolPresentation("t3-code.t3_checkpoint_restore")).toEqual({ + displayName: "Restore a thread checkpoint", + logo: "t3-code", + }); }); it("pretty prints preview T3 MCP tool names", () => { diff --git a/packages/shared/src/t3McpToolPresentation.ts b/packages/shared/src/t3McpToolPresentation.ts index e06f671e2cd2..cfc3c0f69809 100644 --- a/packages/shared/src/t3McpToolPresentation.ts +++ b/packages/shared/src/t3McpToolPresentation.ts @@ -25,6 +25,7 @@ const T3_MCP_TOOL_DISPLAY_NAMES: Record = { t3_thread_interrupt: "Interrupt a T3 thread", t3_checkpoint_list: "List thread checkpoints", t3_checkpoint_diff: "Read a checkpoint diff", + t3_checkpoint_restore: "Restore a thread checkpoint", t3_worktree_handoff: "Hand off thread to a git worktree", t3_worktree_status: "Get thread worktree status", preview_status: "Get preview browser status",