diff --git a/apps/server/src/orchestration-v2/Adapters/AcpAdapterV2.test.ts b/apps/server/src/orchestration-v2/Adapters/AcpAdapterV2.test.ts index a58a83bffa349..b47e60ebfdb86 100644 --- a/apps/server/src/orchestration-v2/Adapters/AcpAdapterV2.test.ts +++ b/apps/server/src/orchestration-v2/Adapters/AcpAdapterV2.test.ts @@ -4612,6 +4612,19 @@ describe("AcpAdapterV2", () => { if (activeRollback._tag === "Failure") { assert.include(String(activeRollback.failure.cause), "while turn"); } + const activeAdmission = yield* runtime.withRollbackAdmission!( + { + providerThread, + providerThreadTurns: [], + target: { + type: "thread_start", + checkpointId: CheckpointId.make("checkpoint-acp-active-admission"), + appRunOrdinal: 0, + }, + }, + () => Effect.die("active ACP turn must not admit restore"), + ).pipe(Effect.result); + assert.equal(activeAdmission._tag, "Failure"); const firstProviderTurnId = idAllocator.derive.providerTurn({ driver: ACP_TEST_DRIVER, nativeTurnId: acpScopedNativeId(instanceId, `${firstSessionId}:turn:1`), @@ -4798,15 +4811,24 @@ describe("AcpAdapterV2", () => { runtimePolicy, }); - yield* runtime.rollbackThread({ - providerThread, - providerThreadTurns: [], - target: { - type: "thread_start", - checkpointId: CheckpointId.make("checkpoint-acp-rollback-retry-generation"), - appRunOrdinal: 0, + assert.isFalse(yield* hasPendingBackgroundWork); + const fs = yield* FileSystem.FileSystem; + const directory = yield* fs.makeTempDirectoryScoped({ prefix: "t3-acp-rollback-admission-" }); + const file = path.join(directory, "workspace.txt"); + yield* fs.writeFileString(file, "current"); + yield* runtime.withRollbackAdmission!( + { + providerThread, + providerThreadTurns: [], + target: { + type: "thread_start", + checkpointId: CheckpointId.make("checkpoint-acp-rollback-retry-generation"), + appRunOrdinal: 0, + }, }, - }); + (rollback) => fs.writeFileString(file, "restored").pipe(Effect.andThen(rollback())), + ); + assert.equal(yield* fs.readFileString(file), "restored"); assert.equal(runtimeOrdinalSeen, 3); assert.deepEqual( @@ -8637,6 +8659,27 @@ describe("AcpAdapterV2", () => { if (Exit.isSuccess(interruptExit)) assert.fail("missing hard teardown must fail interrupt"); assert.include(Cause.pretty(interruptExit.cause), "session is poisoned"); + let admitted = false; + const admissionExit = yield* runtime.withRollbackAdmission!( + { + providerThread, + target: { + type: "thread_start", + checkpointId: CheckpointId.make("poisoned-rollback"), + appRunOrdinal: 0, + }, + providerThreadTurns: [], + }, + () => + Effect.sync(() => { + admitted = true; + }), + ).pipe(Effect.exit); + if (Exit.isSuccess(admissionExit)) + assert.fail("poisoned session must reject rollback admission"); + assert.include(Cause.pretty(admissionExit.cause), "session is poisoned"); + assert.isFalse(admitted); + const startExit = yield* runtime .startTurn( makeTurnInput({ threadId, providerThread, instanceId, runtimePolicy, now, ordinal: 2 }), diff --git a/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.test.ts b/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.test.ts index 45d21b5944e24..d1948d0dce2e7 100644 --- a/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.test.ts +++ b/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.test.ts @@ -3081,6 +3081,7 @@ describe("ClaudeAdapterV2 background wake turns", () => { }); const makeWakeHarnessWithOptions = (options?: { + readonly afterTerminalEvent?: Effect.Effect; readonly close?: (sdkMessages: Queue.Queue) => Effect.Effect; readonly interrupt?: Effect.Effect; readonly setPermissionMode?: ClaudeAdapterV2.ClaudeAgentSdkQuerySession["setPermissionMode"]; @@ -3152,6 +3153,9 @@ describe("ClaudeAdapterV2 background wake turns", () => { >(); let openedOptions: ClaudeAdapterV2.ClaudeAgentSdkQueryOptions | undefined; const adapter = yield* ClaudeAdapterV2.makeClaudeAdapterV2({ + ...(options?.afterTerminalEvent === undefined + ? {} + : { testHooks: { afterTerminalEvent: options.afterTerminalEvent } }), instanceId: ClaudeAdapterV2.CLAUDE_DEFAULT_INSTANCE_ID, settings: options?.settings ?? DEFAULT_CLAUDE_SETTINGS, environment: options?.environment ?? {}, @@ -6043,6 +6047,193 @@ describe("ClaudeAdapterV2 background wake turns", () => { close: (sdkMessages) => Queue.shutdown(sdkMessages), }); + it.effect("rejects rollback while a Claude provider turn is active", () => + Effect.gen(function* () { + const harness = yield* makeRollbackHarness; + yield* harness.runtime.startTurn( + makeClaudeTestTurnInput({ + threadId: harness.threadId, + providerThread: harness.providerThread, + now: yield* DateTime.now, + attemptId: RunAttemptId.make("active-rollback-proof"), + text: "Keep the turn active.", + attachments: [], + }), + ); + const error = yield* harness.runtime + .rollbackThread({ + providerThread: harness.providerThread, + target: { + type: "thread_start", + checkpointId: CheckpointId.make("active-rollback-checkpoint"), + appRunOrdinal: 0, + }, + providerThreadTurns: [], + }) + .pipe(Effect.flip); + assert.equal(error._tag, "ProviderAdapterRollbackThreadError"); + if (error._tag !== "ProviderAdapterRollbackThreadError") return; + assert.instanceOf(error.cause, ProviderAdapter.ProviderAdapterProtocolError); + assert.include(String(error.cause), "while provider turn"); + const admissionError = yield* harness.runtime.withRollbackAdmission!( + { + providerThread: harness.providerThread, + target: { + type: "thread_start", + checkpointId: CheckpointId.make("active-admission-checkpoint"), + appRunOrdinal: 0, + }, + providerThreadTurns: [], + }, + () => Effect.die("active turns must not enter rollback admission"), + ).pipe(Effect.flip); + assert.equal(admissionError._tag, "ProviderAdapterProtocolError"); + }).pipe( + Effect.provide( + Layer.mergeAll( + IdAllocator.layer, + McpProviderSessions.layer, + NodeServices.layer, + NodeCrypto.layer, + ), + ), + ), + ); + + it.effect("keeps admission closed after terminal until Claude clears the active turn", () => + Effect.gen(function* () { + const terminalEmitted = yield* Deferred.make(); + const releaseTerminal = yield* Deferred.make(); + const harness = yield* makeWakeHarnessWithOptions({ + close: (messages) => Queue.shutdown(messages), + afterTerminalEvent: Deferred.succeed(terminalEmitted, undefined).pipe( + Effect.andThen(Deferred.await(releaseTerminal)), + ), + }); + yield* harness.runtime.startTurn( + makeClaudeTestTurnInput({ + threadId: harness.threadId, + providerThread: harness.providerThread, + now: yield* DateTime.now, + attemptId: RunAttemptId.make("terminal-admission-turn"), + text: "Finish this turn.", + attachments: [], + }), + ); + yield* Queue.offer( + harness.sdkMessages, + makeResultFrame({ uuid: "terminal-admission-result", result: "Done." }), + ); + yield* Deferred.await(terminalEmitted); + assert.equal((yield* Queue.take(harness.terminalReceipts)).status, "completed"); + const rejected = yield* harness.runtime.withRollbackAdmission!( + { + providerThread: harness.providerThread, + target: { + type: "thread_start", + checkpointId: CheckpointId.make("terminal-admission"), + appRunOrdinal: 0, + }, + providerThreadTurns: [], + }, + () => Effect.die("terminal is not native cleanup"), + ).pipe(Effect.flip); + assert.equal(rejected._tag, "ProviderAdapterProtocolError"); + yield* Deferred.succeed(releaseTerminal, undefined); + }).pipe( + Effect.provide( + Layer.mergeAll( + IdAllocator.layer, + McpProviderSessions.layer, + NodeServices.layer, + NodeCrypto.layer, + ), + ), + ), + ); + + it.effect.each(["complete", "interrupt"] as const)( + "waits for rollback admission to %s before starting Claude", + (outcome) => + Effect.gen(function* () { + const harness = yield* makeRollbackHarness; + const entered = yield* Deferred.make(); + const release = yield* Deferred.make(); + const admitted = yield* harness.runtime.withRollbackAdmission!( + { + providerThread: harness.providerThread, + target: { + type: "thread_start", + checkpointId: CheckpointId.make("exclusive-admission"), + appRunOrdinal: 0, + }, + providerThreadTurns: [], + }, + () => Deferred.succeed(entered, undefined).pipe(Effect.andThen(Deferred.await(release))), + ).pipe(Effect.forkChild({ startImmediately: true })); + yield* Deferred.await(entered); + const starting = yield* harness.runtime + .startTurn( + makeClaudeTestTurnInput({ + threadId: harness.threadId, + providerThread: harness.providerThread, + now: yield* DateTime.now, + attemptId: RunAttemptId.make("exclusive-admission-turn"), + text: "Start after the rollback.", + attachments: [], + }), + ) + .pipe(Effect.forkChild({ startImmediately: true })); + assert.isUndefined(harness.getOpenedOptions()); + if (outcome === "interrupt") { + yield* Fiber.interrupt(admitted); + } else { + yield* Deferred.succeed(release, undefined); + yield* Fiber.join(admitted); + } + yield* Fiber.join(starting); + assert.isDefined(harness.getOpenedOptions()); + }).pipe( + Effect.provide( + Layer.mergeAll( + IdAllocator.layer, + McpProviderSessions.layer, + NodeServices.layer, + NodeCrypto.layer, + ), + ), + ), + ); + + it.effect("preserves failures from inside rollback admission", () => + Effect.gen(function* () { + const harness = yield* makeRollbackHarness; + const expected = { _tag: "host-restore-failure" as const }; + const error = yield* harness.runtime.withRollbackAdmission!( + { + providerThread: harness.providerThread, + target: { + type: "thread_start", + checkpointId: CheckpointId.make("admission-host-failure"), + appRunOrdinal: 0, + }, + providerThreadTurns: [], + }, + () => Effect.fail(expected), + ).pipe(Effect.flip); + assert.strictEqual(error, expected); + }).pipe( + Effect.provide( + Layer.mergeAll( + IdAllocator.layer, + McpProviderSessions.layer, + NodeServices.layer, + NodeCrypto.layer, + ), + ), + ), + ); + it.effect.each([ { kind: "fork with inherited head", fork: true, headMessageId: "remapped-inherited-head" }, { kind: "fork without recorded head", fork: true, headMessageId: undefined }, @@ -11161,6 +11352,19 @@ describe("ClaudeAdapterV2 background wake turns", () => { assert.lengthOf(harness.continuationRequests, 0); yield* harness.offerAndWait(wakeTurnInit); assert.lengthOf(harness.continuationRequests, 1); + const admissionError = yield* harness.runtime.withRollbackAdmission!( + { + providerThread: harness.providerThread, + providerThreadTurns: [], + target: { + type: "thread_start", + checkpointId: CheckpointId.make("wake-init-admission"), + appRunOrdinal: 0, + }, + }, + () => Effect.die("an initialized native wake must exclude restoration"), + ).pipe(Effect.flip); + assert.equal(admissionError._tag, "ProviderAdapterProtocolError"); assert.equal(harness.continuationRequests[0]?.detail, WAKE_SUMMARY); yield* harness.runtime.startTurn( @@ -18013,6 +18217,28 @@ describe("ClaudeAdapterV2 background wake turns", () => { event.type === "provider_thread.updated" && event.providerThread.goal === null, ), ); + const clearedThread = harness.events.findLast( + (event) => event.type === "provider_thread.updated" && event.providerThread.goal === null, + ); + assert.equal(clearedThread?.type, "provider_thread.updated"); + if (clearedThread?.type !== "provider_thread.updated") return; + assert.isTrue( + yield* harness.runtime.withRollbackAdmission!( + { + providerThread: { + ...clearedThread.providerThread, + goal: { objective: condition, status: "active", checks: 1 }, + }, + target: { + type: "thread_start", + checkpointId: CheckpointId.make("cleared-goal-admission"), + appRunOrdinal: 0, + }, + providerThreadTurns: [], + }, + () => Effect.succeed(true), + ), + ); assert.isFalse( harness.events.some( (event) => event.type === "message.updated" && event.message.role === "assistant", diff --git a/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.ts b/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.ts index 674f55c524d75..b2e8cd4d22a7b 100644 --- a/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.ts +++ b/apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.ts @@ -4591,6 +4591,7 @@ function claudeProposedPlan(input: unknown): string | null { } export interface ClaudeAdapterV2Options { + readonly testHooks?: { readonly afterTerminalEvent?: Effect.Effect }; readonly instanceId: ProviderInstanceId; readonly settings: ClaudeSettings; readonly environment: NodeJS.ProcessEnv; @@ -4804,6 +4805,7 @@ export const makeClaudeAdapterV2 = Effect.fn("makeClaudeAdapterV2")(function* ( }); const events = yield* Queue.unbounded(); const activeTurn = yield* Ref.make(null); + const rollbackAdmission = yield* Semaphore.make(1); const taskOutputTails = new Map< string, { readonly fiber: Fiber.Fiber; readonly stop: Deferred.Deferred } @@ -7814,6 +7816,7 @@ export const makeClaudeAdapterV2 = Effect.fn("makeClaudeAdapterV2")(function* ( ], { concurrency: 1 }, ); + yield* adapterOptions.testHooks?.afterTerminalEvent ?? Effect.void; yield* Ref.update(activeTurn, (current) => current?.providerTurnId === input.context.providerTurnId ? null : current, ); @@ -11474,12 +11477,14 @@ export const makeClaudeAdapterV2 = Effect.fn("makeClaudeAdapterV2")(function* ( ), ), ), - startTurn, + startTurn: (turnInput) => rollbackAdmission.withPermit(startTurn(turnInput)), compactThread: (turnInput) => - startTurn({ - ...turnInput, - message: { ...turnInput.message, text: "/compact" }, - }), + rollbackAdmission.withPermit( + startTurn({ + ...turnInput, + message: { ...turnInput.message, text: "/compact" }, + }), + ), steerTurn, stopTask: Effect.fn("ClaudeAdapterV2.stopTask")(function* ({ providerThread, taskId }) { const liveQuery = yield* Ref.get(queryContext); @@ -11665,6 +11670,31 @@ export const makeClaudeAdapterV2 = Effect.fn("makeClaudeAdapterV2")(function* ( cause: "Claude V2 adapter does not implement snapshots.", }), ), + withRollbackAdmission: (rollbackInput, use) => + rollbackAdmission.withPermit( + Effect.gen(function* () { + const currentTurn = yield* Ref.get(activeTurn); + const nativeThreadId = rollbackInput.providerThread.nativeThreadRef?.nativeId; + const goal = + nativeThreadId != null && goalsByNativeThread.has(nativeThreadId) + ? goalsByNativeThread.get(nativeThreadId) + : rollbackInput.providerThread.goal; + if ( + currentTurn !== null || + [...(yield* Ref.get(wakeBuffers)).values()].some((entry) => + entry.messages.some(isClaudeTurnStartMessage), + ) || + goal?.status === "active" || + (yield* runtime.hasPendingBackgroundWork!) + ) { + return yield* new ProviderAdapter.ProviderAdapterProtocolError({ + driver: CLAUDE_PROVIDER, + detail: `Cannot roll back Claude provider thread ${rollbackInput.providerThread.id} while provider work is active.`, + }); + } + return yield* use(() => runtime.rollbackThread(rollbackInput)); + }), + ), rollbackThread: Effect.fn("ClaudeAdapterV2.rollbackThread")( function* (rollbackInput) { const currentTurn = yield* Ref.get(activeTurn); diff --git a/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.test.ts b/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.test.ts index 4471be814bf39..a7247eda4c954 100644 --- a/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.test.ts +++ b/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.test.ts @@ -1174,6 +1174,7 @@ describe("CodexAdapterV2 post-settle continuation", () => { onRequest: (method: string, params: unknown) => Effect.Effect = () => Effect.void, readChildMetadata?: Parameters[2], beforeEmitInbound?: CodexReplay.CodexAppServerReplayDriver["beforeEmitInbound"], + testHooks?: CodexAdapterV2.CodexAdapterV2Options["testHooks"], ) => Effect.gen(function* () { const fileSystem = yield* FileSystem.FileSystem; @@ -1231,6 +1232,7 @@ describe("CodexAdapterV2 post-settle continuation", () => { settings: DEFAULT_CODEX_SETTINGS, environment: {}, clientFactory, + ...(testHooks === undefined ? {} : { testHooks }), digest: yield* Digest.Digest, fileSystem, idAllocator, @@ -2823,6 +2825,7 @@ describe("CodexAdapterV2 post-settle continuation", () => { it.effect.each([ "recover", + "admission", "completion-only", "exhaust", "stop", @@ -2975,6 +2978,7 @@ describe("CodexAdapterV2 post-settle continuation", () => { if (outcome === "exhaust") entries.push(...failureEntries(attempt)); if ( outcome === "recover" || + outcome === "admission" || outcome === "completion-only" || outcome === "started-first" || outcome === "stop-starting" @@ -3033,6 +3037,10 @@ describe("CodexAdapterV2 post-settle continuation", () => { const milestones = yield* Queue.unbounded(); const starts: Array = []; const retryStarted = yield* Deferred.make(); + const retryScheduled = yield* Deferred.make(); + const retryFinalized = yield* Deferred.make(); + const releaseRetry = yield* Deferred.make(); + const admitted = yield* Deferred.make(); const harness = yield* makeCodexReplayHarness( makeCodexReplayTranscript({ scenario, entries }), (event) => @@ -3049,6 +3057,18 @@ describe("CodexAdapterV2 post-settle continuation", () => { if (starts.length === 2) yield* Deferred.succeed(retryStarted, undefined); } }), + undefined, + undefined, + outcome === "admission" + ? { + afterCapacityRetryScheduled: Deferred.succeed(retryScheduled, undefined).pipe( + Effect.asVoid, + ), + afterCapacityRetryFinalization: Deferred.succeed(retryFinalized, undefined).pipe( + Effect.andThen(Deferred.await(releaseRetry)), + ), + } + : undefined, ); const turnInput = makeCodexTestTurnInput({ threadId: harness.threadId, @@ -3087,7 +3107,27 @@ describe("CodexAdapterV2 post-settle continuation", () => { }) .pipe(Effect.forkScoped) : undefined; + if (outcome === "admission") yield* Deferred.await(retryScheduled); yield* TestClock.adjust(Math.min(2_000 * 2 ** (attempt - 1), 30_000)); + if (outcome === "admission") { + yield* Deferred.await(retryFinalized); + const admission = yield* harness.runtime.withRollbackAdmission!( + { + providerThread: harness.providerThread, + providerThreadTurns: [], + target: { + type: "thread_start", + checkpointId: CheckpointId.make("capacity-admission"), + appRunOrdinal: 0, + }, + }, + () => Deferred.succeed(admitted, undefined), + ).pipe(Effect.result, Effect.forkChild({ startImmediately: true })); + assert.isFalse(yield* Deferred.isDone(admitted)); + yield* Deferred.succeed(releaseRetry, undefined); + yield* Fiber.join(admission); + assert.equal(starts.length, 2); + } if (outcome === "stop-starting") { yield* Deferred.await(retryStarted); const stopping = yield* harness.runtime @@ -3129,6 +3169,7 @@ describe("CodexAdapterV2 post-settle continuation", () => { assert.equal( terminal?.status, outcome === "recover" || + outcome === "admission" || outcome === "completion-only" || outcome === "steer" || outcome === "started-first" || @@ -6683,6 +6724,37 @@ describe("CodexAdapterV2 post-settle continuation", () => { assert.equal(harness.terminalEvents()[0]?.status, "completed"); assert.lengthOf(harness.continuationRequests, 0); + assert.isTrue( + yield* harness.runtime.hasPendingBackgroundWorkForThread!(harness.providerThread), + ); + const rollbackInput = { + providerThread: harness.providerThread, + target: { + type: "thread_start" as const, + checkpointId: CheckpointId.make("late-child-admission"), + appRunOrdinal: 0 as const, + }, + providerThreadTurns: [], + }; + const admissionError = yield* harness.runtime.withRollbackAdmission!(rollbackInput, () => + Effect.die("a running descendant must exclude restoration"), + ).pipe(Effect.flip); + assert.equal(admissionError._tag, "ProviderAdapterProtocolError"); + const unrelated = { + ...rollbackInput, + providerThread: { + ...harness.providerThread, + id: ProviderThreadId.make("unrelated-admission-thread"), + nativeThreadRef: null, + }, + }; + assert.isFalse( + yield* harness.runtime.hasPendingBackgroundWorkForThread!(unrelated.providerThread), + ); + assert.isTrue( + yield* harness.runtime.withRollbackAdmission!(unrelated, () => Effect.succeed(true)), + ); + yield* TestClock.adjust("60 seconds"); yield* harness.until( () => harness.subagentUpdates().at(-1)?.subagent.result === "NATIVE_CHILD_16900_DONE", @@ -7807,6 +7879,23 @@ describe("CodexAdapterV2 post-settle continuation", () => { now, }); + const activeGoalAdmission = yield* harness.runtime.withRollbackAdmission!( + { + providerThread: { + ...harness.providerThread, + goal: { objective: "Continue autonomously", status: "active" }, + }, + target: { + type: "thread_start", + checkpointId: CheckpointId.make("checkpoint-codex-active-goal"), + appRunOrdinal: 0, + }, + providerThreadTurns: [], + }, + () => Effect.die("an active goal must not admit restoration"), + ).pipe(Effect.result); + assert.equal(activeGoalAdmission._tag, "Failure"); + const error = yield* Effect.flip( harness.runtime.rollbackThread({ providerThread: harness.providerThread, @@ -8175,6 +8264,67 @@ describe("CodexAdapterV2 post-settle continuation", () => { event.type === "provider_thread.updated" ? [event.providerThread.goal?.status ?? null] : [], ); + it.effect("admits rollback after a native goal is cleared despite stale input", () => + Effect.gen(function* () { + const harness = yield* makeCodexReplayHarness( + makeCodexReplayTranscript({ + scenario: "goal-clear-admission", + entries: [ + ...codexReplayPreamble({ + nativeThreadId, + nativeTurnId: "clear-turn", + prompt: "clear test", + }), + notification("turn done", "turn/completed", { + threadId: nativeThreadId, + turn: makeCodexReplayTurn({ id: "clear-turn", status: "completed" }), + }), + notification("goal active", "thread/goal/updated", { + threadId: nativeThreadId, + turnId: null, + goal: codexGoal("active", 0), + }), + notification("goal cleared", "thread/goal/cleared", { + threadId: nativeThreadId, + turnId: null, + }), + ], + }), + ); + yield* harness.runtime.startTurn(yield* goalTurnInput(harness, "clear test")); + yield* harness.firstTerminal; + yield* harness.until(() => + harness.events.some( + (event) => + event.type === "provider_thread.updated" && event.providerThread.goal === null, + ), + ); + const result = yield* harness.runtime.withRollbackAdmission!( + { + providerThread: { ...harness.providerThread, goal: { objective, status: "active" } }, + providerThreadTurns: [], + target: { + type: "thread_start", + checkpointId: CheckpointId.make("goal-cleared"), + appRunOrdinal: 0, + }, + }, + () => Effect.succeed(true), + ); + assert.isTrue(result); + }).pipe( + Effect.scoped, + Effect.provide( + Layer.mergeAll( + IdAllocator.layer, + McpProviderSessions.layer, + NodeServices.layer, + NodeCrypto.layer, + ), + ), + ), + ); + it.effect("keeps a /goal run open across the turns Codex continues on its own", () => Effect.gen(function* () { const transcript = makeCodexReplayTranscript({ diff --git a/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.ts b/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.ts index 688de092c03d0..98dbd084b7636 100644 --- a/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.ts +++ b/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.ts @@ -1654,6 +1654,10 @@ export const createCodexAdapterV2 = Effect.fn("CodexAdapterV2.createCodexAdapter }); export interface CodexAdapterV2Options { + readonly testHooks?: { + readonly afterCapacityRetryFinalization?: Effect.Effect; + readonly afterCapacityRetryScheduled?: Effect.Effect; + }; readonly instanceId: ProviderInstanceId; readonly settings: CodexSettings; readonly environment: NodeJS.ProcessEnv; @@ -1836,6 +1840,7 @@ export const makeCodexAdapterV2 = Effect.fn("makeCodexAdapterV2")(function* ( new Map>(), ); const activeTurns = yield* Ref.make(new Map()); + const rollbackAdmission = yield* Semaphore.make(1); const turnTokenUsageByThread = new Map(); const usageStateForThread = (nativeThreadId: string) => { let state = turnTokenUsageByThread.get(nativeThreadId); @@ -6461,66 +6466,83 @@ export const makeCodexAdapterV2 = Effect.fn("makeCodexAdapterV2")(function* ( yield* Effect.sleep(retryDelayMs).pipe( Effect.andThen( Effect.gen(function* () { - const nativeThreadId = yield* getNativeThreadId(context.providerThread); - const claimed = yield* turnTerminalizationPermit.withPermits(1)( + // Release admission once pending startup is visible; notifications may precede the RPC response. + const pendingInstalled = yield* Deferred.make(); + const retryFiber = yield* rollbackAdmission.withPermit( Effect.gen(function* () { - if ( - (yield* Ref.get(activeTurns)).get(terminal.nativeTurnId) !== context || - goalStops.has(nativeThreadId) || - (yield* Ref.get(interruptingNativeTurns)).has(terminal.nativeTurnId) - ) - return false; - capacityContinuations.set(nativeThreadId, { context, retry: state }); - yield* finalizeCodexTurn({ ...terminal, continueRun: true }); - return true; + const nativeThreadId = yield* getNativeThreadId(context.providerThread); + const claimed = yield* turnTerminalizationPermit.withPermits(1)( + Effect.gen(function* () { + if ( + (yield* Ref.get(activeTurns)).get(terminal.nativeTurnId) !== + context || + goalStops.has(nativeThreadId) || + (yield* Ref.get(interruptingNativeTurns)).has(terminal.nativeTurnId) + ) + return false; + capacityContinuations.set(nativeThreadId, { context, retry: state }); + yield* finalizeCodexTurn({ ...terminal, continueRun: true }); + return true; + }), + ); + if (!claimed) return undefined; + yield* ( + adapterOptions.testHooks?.afterCapacityRetryFinalization ?? Effect.void + ); + // Notifications before the RPC response also take the terminalization permit. + const retryFiber = yield* startNativeTurn( + { + ...context.input, + providerTurnOrdinal: context.providerTurnOrdinal + 1, + }, + [], + Deferred.succeed(pendingInstalled, undefined).pipe(Effect.asVoid), + ).pipe( + Effect.catch((cause) => + Effect.gen(function* () { + if (yield* Deferred.isDone(pending)) return; + yield* Effect.logWarning( + "orchestration-v2.codex-capacity-retry-start-failed", + { + providerTurnId: context.providerTurnId, + cause, + }, + ); + yield* emitProviderEvent({ + type: "turn_item.updated", + driver: CODEX_PROVIDER, + turnItem: codexRetryTurnItem( + context, + state, + "failed", + yield* DateTime.now, + ), + }); + yield* Deferred.succeed(pending, undefined); + yield* turnTerminalizationPermit.withPermits(1)( + emitOrDeferRootTerminal({ ...terminal, providerRetry: state }), + ); + }), + ), + Effect.ensuring(Deferred.succeed(pendingInstalled, undefined)), + Effect.ensuring( + Effect.sync(() => { + if (capacityContinuations.get(nativeThreadId)?.context === context) + capacityContinuations.delete(nativeThreadId); + }), + ), + Effect.forkChild({ startImmediately: true }), + ); + yield* Deferred.await(pendingInstalled); + return retryFiber; }), ); - if (!claimed) return; - // Notifications before the RPC response also take the terminalization permit. - yield* startNativeTurn( - { - ...context.input, - providerTurnOrdinal: context.providerTurnOrdinal + 1, - }, - [], - ).pipe( - Effect.catch((cause) => - Effect.gen(function* () { - if (yield* Deferred.isDone(pending)) return; - yield* Effect.logWarning( - "orchestration-v2.codex-capacity-retry-start-failed", - { - providerTurnId: context.providerTurnId, - cause, - }, - ); - yield* emitProviderEvent({ - type: "turn_item.updated", - driver: CODEX_PROVIDER, - turnItem: codexRetryTurnItem( - context, - state, - "failed", - yield* DateTime.now, - ), - }); - yield* Deferred.succeed(pending, undefined); - yield* turnTerminalizationPermit.withPermits(1)( - emitOrDeferRootTerminal({ ...terminal, providerRetry: state }), - ); - }), - ), - Effect.ensuring( - Effect.sync(() => { - if (capacityContinuations.get(nativeThreadId)?.context === context) - capacityContinuations.delete(nativeThreadId); - }), - ), - ); + if (retryFiber !== undefined) yield* Fiber.join(retryFiber); }), ), Effect.forkIn(scope, { startImmediately: true }), ); + yield* adapterOptions.testHooks?.afterCapacityRetryScheduled ?? Effect.void; return true; }), ); @@ -6744,6 +6766,7 @@ export const makeCodexAdapterV2 = Effect.fn("makeCodexAdapterV2")(function* ( function* ( turnInput: ProviderAdapter.ProviderAdapterV2TurnInput, codexInput: ReadonlyArray, + onPendingInstalled: Effect.Effect = Effect.void, ) { const threadId = yield* getNativeThreadId(turnInput.providerThread); const mcpSession = yield* mcpSessions.read(turnInput.threadId); @@ -6770,6 +6793,7 @@ export const makeCodexAdapterV2 = Effect.fn("makeCodexAdapterV2")(function* ( else next.delete(threadId); return next; }); + yield* onPendingInstalled; // Completion can remove the old context while turn/start is in flight. const previousTurns = new Map([ ...(yield* Ref.get(settledTurns)), @@ -7041,12 +7065,13 @@ export const makeCodexAdapterV2 = Effect.fn("makeCodexAdapterV2")(function* ( ...(yield* Ref.get(activeTurns)).values(), ...(yield* Ref.get(settledTurns)).values(), ]; - const roots = contexts.filter( - (context) => context.providerThread.id === providerThread.id, - ); for (const context of contexts) { - if (!roots.some((root) => context === root || isDescendantCodexTurn(context, root))) - continue; + // Descendants retain their ancestry after the root leaves the runtime maps. + let ancestor: ActiveCodexTurnContext | undefined = context; + while (ancestor !== undefined && ancestor.providerThread.id !== providerThread.id) { + ancestor = ancestor.subagent?.parentContext; + } + if (ancestor === undefined) continue; if (yield* turnHasRetainedBackgroundWork(context.nativeTurnId)) return true; if ( context.subagent !== null && @@ -7156,6 +7181,7 @@ export const makeCodexAdapterV2 = Effect.fn("makeCodexAdapterV2")(function* ( }, (effect, turnInput) => effect.pipe( + (effect) => rollbackAdmission.withPermit(effect), Effect.mapError( (cause) => new ProviderAdapter.ProviderAdapterTurnStartError({ @@ -7199,24 +7225,27 @@ export const makeCodexAdapterV2 = Effect.fn("makeCodexAdapterV2")(function* ( turnInput.restartContinuationOfRunId === undefined ? parseGoalCommand(turnInput.message) : null; - if (goalCommand !== null) return runGoalCommand(turnInput, goalCommand); - return Effect.gen(function* () { - yield* startNativeTurn( - turnInput, - turnInput.restartContinuationOfRunId === undefined - ? yield* toCodexInput(turnInput) - : [], - ); - }).pipe( - Effect.mapError( - (cause) => - new ProviderAdapter.ProviderAdapterTurnStartError({ - driver: CODEX_PROVIDER, - threadId: turnInput.threadId, - providerThreadId: turnInput.providerThread.id, - runId: turnInput.runId, - cause, - }), + if (goalCommand !== null) + return rollbackAdmission.withPermit(runGoalCommand(turnInput, goalCommand)); + return rollbackAdmission.withPermit( + Effect.gen(function* () { + yield* startNativeTurn( + turnInput, + turnInput.restartContinuationOfRunId === undefined + ? yield* toCodexInput(turnInput) + : [], + ); + }).pipe( + Effect.mapError( + (cause) => + new ProviderAdapter.ProviderAdapterTurnStartError({ + driver: CODEX_PROVIDER, + threadId: turnInput.threadId, + providerThreadId: turnInput.providerThread.id, + runId: turnInput.runId, + cause, + }), + ), ), ); }, @@ -7829,6 +7858,37 @@ export const makeCodexAdapterV2 = Effect.fn("makeCodexAdapterV2")(function* ( ), ), ), + withRollbackAdmission: (rollbackInput, use) => + rollbackAdmission.withPermit( + Effect.gen(function* () { + const nativeThreadId = rollbackInput.providerThread.nativeThreadRef?.nativeId; + const active = [...(yield* Ref.get(activeTurns)).values()].some( + (turn) => turn.providerThread.id === rollbackInput.providerThread.id, + ); + const pending = + nativeThreadId != null && (yield* Ref.get(pendingRootTurns)).has(nativeThreadId); + const goalRunning = + (nativeThreadId == null && + rollbackInput.providerThread.goal?.status === "active") || + (nativeThreadId != null && + (goalActivations.has(nativeThreadId) || + goalHolds.has(nativeThreadId) || + (goalsByNativeThread.has(nativeThreadId) + ? goalsByNativeThread.get(nativeThreadId) + : rollbackInput.providerThread.goal + )?.status === "active")); + const background = yield* runtime.hasPendingBackgroundWorkForThread!( + rollbackInput.providerThread, + ); + if (active || pending || background || goalRunning) { + return yield* new ProviderAdapter.ProviderAdapterProtocolError({ + driver: CODEX_PROVIDER, + detail: "Cannot roll back while the provider thread has active work.", + }); + } + return yield* use(() => runtime.rollbackThread(rollbackInput)); + }), + ), rollbackThread: Effect.fn("CodexAdapterV2.rollbackThread")( function* (threadInput) { const threadId = yield* getNativeThreadId(threadInput.providerThread); diff --git a/packages/provider-acp/src/server/adapter.ts b/packages/provider-acp/src/server/adapter.ts index 985f6c8ee9887..cea07220b1054 100644 --- a/packages/provider-acp/src/server/adapter.ts +++ b/packages/provider-acp/src/server/adapter.ts @@ -8347,6 +8347,61 @@ export const makeAcpAdapterV2 = Effect.fn("makeAcpAdapterV2")(function* ( }), ); + const rollbackUnlocked = Effect.fn("AcpAdapterV2.rollbackThread.unlocked")( + function* (rollbackInput: ProviderAdapter.ProviderAdapterV2RollbackThreadInput) { + const currentTurn = yield* Ref.get(activeTurn); + if (currentTurn !== null) { + return yield* new ProviderAdapter.ProviderAdapterProtocolError({ + driver, + detail: `Cannot roll back ACP provider thread ${rollbackInput.providerThread.id} while turn ${currentTurn.providerTurnId} is active`, + }); + } + yield* awaitRuntimeTeardown(); + itemIdentityVersion = 2; + prepareTerminalEnvironment( + yield* readMcpContext(rollbackInput.providerThread.appThreadId), + ); + const replacement = yield* startReplacementAcpRuntime( + rollbackInput.providerThread.appThreadId, + (candidate) => + commitReplacementSession(rollbackInput.providerThread.appThreadId, candidate), + ).pipe(Effect.retry({ times: 1 })); + const now = yield* DateTime.now; + return { + providerThread: { + ...rollbackInput.providerThread, + nativeThreadRef: { + driver, + nativeId: replacement.sessionId, + strength: "strong" as const, + }, + nativeConversationHeadRef: null, + nativeMetadata: { + ...rollbackInput.providerThread.nativeMetadata, + itemIdentityVersion: 2 as const, + }, + status: "idle" as const, + updatedAt: now, + }, + providerTurns: [], + messages: [], + runtimeRequests: [], + }; + }, + (effect, rollbackInput) => + effect.pipe( + Effect.mapError( + (cause) => + new ProviderAdapter.ProviderAdapterRollbackThreadError({ + driver, + providerThreadId: rollbackInput.providerThread.id, + checkpointId: rollbackInput.target.checkpointId, + cause, + }), + ), + ), + ); + const sessionRuntime: ProviderAdapter.ProviderAdapterV2SessionRuntime = { instanceId: options.instanceId, driver, @@ -9050,61 +9105,35 @@ export const makeAcpAdapterV2 = Effect.fn("makeAcpAdapterV2")(function* ( ), ), ), - rollbackThread: (rollbackInput) => - runtimeTransitionPermit - .withPermit( - Effect.gen(function* () { - const currentTurn = yield* Ref.get(activeTurn); - if (currentTurn !== null) { - return yield* new ProviderAdapter.ProviderAdapterProtocolError({ - driver, - detail: `Cannot roll back ACP provider thread ${rollbackInput.providerThread.id} while turn ${currentTurn.providerTurnId} is active`, - }); - } - yield* awaitRuntimeTeardown(); - itemIdentityVersion = 2; - prepareTerminalEnvironment( - yield* readMcpContext(rollbackInput.providerThread.appThreadId), - ); - const replacement = yield* startReplacementAcpRuntime( - rollbackInput.providerThread.appThreadId, - (candidate) => - commitReplacementSession(rollbackInput.providerThread.appThreadId, candidate), - ).pipe(Effect.retry({ times: 1 })); - const now = yield* DateTime.now; - return { - providerThread: { - ...rollbackInput.providerThread, - nativeThreadRef: { + withRollbackAdmission: (rollbackInput, use) => + runtimeTransitionPermit.withPermit( + Effect.gen(function* () { + yield* awaitRuntimeTeardown().pipe( + Effect.mapError( + (cause) => + new ProviderAdapter.ProviderAdapterRollbackThreadError({ driver, - nativeId: replacement.sessionId, - strength: "strong" as const, - }, - nativeConversationHeadRef: null, - nativeMetadata: { - ...rollbackInput.providerThread.nativeMetadata, - itemIdentityVersion: 2 as const, - }, - status: "idle" as const, - updatedAt: now, - }, - providerTurns: [], - messages: [], - runtimeRequests: [], - }; - }), - ) - .pipe( - Effect.mapError( - (cause) => - new ProviderAdapter.ProviderAdapterRollbackThreadError({ - driver, - providerThreadId: rollbackInput.providerThread.id, - checkpointId: rollbackInput.target.checkpointId, - cause, - }), - ), - ), + providerThreadId: rollbackInput.providerThread.id, + checkpointId: rollbackInput.target.checkpointId, + cause, + }), + ), + ); + const currentTurn = yield* Ref.get(activeTurn); + if ( + currentTurn !== null || + (yield* sessionRuntime.hasPendingBackgroundWork ?? Effect.succeed(false)) + ) { + return yield* new ProviderAdapter.ProviderAdapterProtocolError({ + driver, + detail: `Cannot roll back ACP provider thread ${rollbackInput.providerThread.id} while provider work is active`, + }); + } + return yield* use(() => rollbackUnlocked(rollbackInput)); + }), + ), + rollbackThread: (rollbackInput) => + runtimeTransitionPermit.withPermit(rollbackUnlocked(rollbackInput)), forkThread: Effect.fn("AcpAdapterV2.forkThread")( function* (forkInput) { return yield* runtimeTransitionPermit.withPermit( diff --git a/packages/provider-core/src/server/ProviderAdapter.ts b/packages/provider-core/src/server/ProviderAdapter.ts index 9d83da489b306..eeeb26140366f 100644 --- a/packages/provider-core/src/server/ProviderAdapter.ts +++ b/packages/provider-core/src/server/ProviderAdapter.ts @@ -658,6 +658,13 @@ export interface ProviderAdapterV2SessionRuntime { readonly rollbackThread: ( input: ProviderAdapterV2RollbackThreadInput, ) => Effect.Effect; + /** Reject active native work, then hold admission through `use` and its unlocked rollback. */ + readonly withRollbackAdmission?: ( + input: ProviderAdapterV2RollbackThreadInput, + use: ( + rollback: () => Effect.Effect, + ) => Effect.Effect, + ) => Effect.Effect; readonly forkThread: ( input: ProviderAdapterV2ForkThreadInput, ) => Effect.Effect;