diff --git a/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.ts b/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.ts index 688de092c03d0..6eb19d23c150d 100644 --- a/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.ts +++ b/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.ts @@ -250,6 +250,7 @@ export const CodexProviderCapabilitiesV2 = { canRollbackThread: true, canForkThread: true, canForkFromTurn: true, + canForkWhileRunning: true, canForkFromSubagentThread: true, exposesNativeThreadId: true, }, diff --git a/apps/server/src/orchestration-v2/CommandPolicy.test.ts b/apps/server/src/orchestration-v2/CommandPolicy.test.ts index 8ed4de3e9dea0..a0902c244aa6b 100644 --- a/apps/server/src/orchestration-v2/CommandPolicy.test.ts +++ b/apps/server/src/orchestration-v2/CommandPolicy.test.ts @@ -258,6 +258,7 @@ layer("CommandPolicyV2", (it) => { capabilities: CursorProviderCapabilitiesV2, sameProvider: true, hasStrongNativeSource: true, + hasNativeTurnBoundary: false, sourceRunStatus: "completed", fromSpecificTurn: true, sourceBusy: false, @@ -267,24 +268,39 @@ layer("CommandPolicyV2", (it) => { }), ); - it.effect("forks a busy source through portable context", () => - Effect.gen(function* () { - const policy = yield* CommandPolicy.CommandPolicyV2; - const decide = (sourceBusy: boolean) => - policy.decideForkExecution({ - commandId, - threadId, - providerInstanceId: ProviderInstanceId.make("codex"), - capabilities: CodexProviderCapabilitiesV2, - sameProvider: true, - hasStrongNativeSource: true, - sourceRunStatus: "completed", - fromSpecificTurn: true, - sourceBusy, - }); + it.effect.each([ + { canForkWhileRunning: true, hasNativeTurnBoundary: true, busyStrategy: "native_fork" }, + { canForkWhileRunning: false, hasNativeTurnBoundary: true, busyStrategy: "portable_context" }, + { + canForkWhileRunning: undefined, + hasNativeTurnBoundary: true, + busyStrategy: "portable_context", + }, + { canForkWhileRunning: true, hasNativeTurnBoundary: false, busyStrategy: "portable_context" }, + ] as const)( + "forks a busy source with canForkWhileRunning=$canForkWhileRunning and hasNativeTurnBoundary=$hasNativeTurnBoundary", + ({ canForkWhileRunning, hasNativeTurnBoundary, busyStrategy }) => + Effect.gen(function* () { + const policy = yield* CommandPolicy.CommandPolicyV2; + const decide = (sourceBusy: boolean) => + policy.decideForkExecution({ + commandId, + threadId, + providerInstanceId: ProviderInstanceId.make("codex"), + capabilities: { + ...CodexProviderCapabilitiesV2, + threads: { ...CodexProviderCapabilitiesV2.threads, canForkWhileRunning }, + }, + sameProvider: true, + hasStrongNativeSource: true, + hasNativeTurnBoundary, + sourceRunStatus: "completed", + fromSpecificTurn: true, + sourceBusy, + }); - assert.equal(yield* decide(false), "native_fork"); - assert.equal(yield* decide(true), "portable_context"); - }), + assert.equal(yield* decide(false), "native_fork"); + assert.equal(yield* decide(true), busyStrategy); + }), ); }); diff --git a/apps/server/src/orchestration-v2/CommandPolicy.ts b/apps/server/src/orchestration-v2/CommandPolicy.ts index 02f329e536feb..77624bb32fe61 100644 --- a/apps/server/src/orchestration-v2/CommandPolicy.ts +++ b/apps/server/src/orchestration-v2/CommandPolicy.ts @@ -209,9 +209,10 @@ export interface CommandPolicyV2Shape { input: CapabilityCheckInput & { readonly sameProvider: boolean; readonly hasStrongNativeSource: boolean; + readonly hasNativeTurnBoundary: boolean; readonly sourceRunStatus: OrchestrationV2Run["status"]; readonly fromSpecificTurn: boolean; - /** Later source runs are still active; providers cannot fork a thread mid-turn. */ + /** Later source runs are still active. */ readonly sourceBusy: boolean; }, ) => Effect.Effect; @@ -376,7 +377,8 @@ const decideForkExecution: CommandPolicyV2Shape["decideForkExecution"] = (input) (input.sourceRunStatus === "completed" || input.sourceRunStatus === "waiting") && input.sameProvider && input.hasStrongNativeSource && - !input.sourceBusy && + (!input.sourceBusy || + (input.capabilities.threads.canForkWhileRunning === true && input.hasNativeTurnBoundary)) && input.capabilities.threads.canForkThread && (!input.fromSpecificTurn || input.capabilities.threads.canForkFromTurn) && input.capabilities.identity.nativeThreadIds === "strong"; diff --git a/apps/server/src/orchestration-v2/Orchestrator.ts b/apps/server/src/orchestration-v2/Orchestrator.ts index 1059c88230788..61e4acd545dd9 100644 --- a/apps/server/src/orchestration-v2/Orchestrator.ts +++ b/apps/server/src/orchestration-v2/Orchestrator.ts @@ -6775,6 +6775,9 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio sameProvider: pendingForkTransfer.sourceProviderInstanceId === modelSelection.instanceId, hasStrongNativeSource: sourceProviderThread?.nativeThreadRef?.strength === "strong", + hasNativeTurnBoundary: + sourceProjection?.providerTurns.find((turn) => turn.id === sourceProviderTurnId) + ?.nativeTurnRef != null, sourceRunStatus: sourceRun.status, fromSpecificTurn: sourceRun !== null, sourceBusy: sourceRunsAfterFork.some(isActiveRun), diff --git a/apps/server/src/orchestration-v2/ProviderTurnStartService.test.ts b/apps/server/src/orchestration-v2/ProviderTurnStartService.test.ts index f7f9dc25ddc92..3843c807e7869 100644 --- a/apps/server/src/orchestration-v2/ProviderTurnStartService.test.ts +++ b/apps/server/src/orchestration-v2/ProviderTurnStartService.test.ts @@ -18,6 +18,7 @@ import { RunAttemptId, RunId, ThreadId, + TurnItemId, ProjectId, type OrchestrationV2ThreadProjection, OrchestrationV2DomainEvent, @@ -300,6 +301,9 @@ function makeLocalCommandHarness(input: { readonly sideChat?: { readonly contextRun: "current" | "failed" | "cancelled" | "delivered"; readonly nativeForkFrom?: "idle" | "busy"; + readonly canForkWhileRunning?: boolean; + readonly sourceBoundary?: "native" | "synthetic" | "missing"; + readonly sourceRolledBack?: boolean; readonly switchedProvider?: boolean; }; }) { @@ -501,11 +505,14 @@ function makeLocalCommandHarness(input: { providerThreadId: parentProviderThread.id, nodeId: rootNodeId, runAttemptId: RunAttemptId.make(`side-chat-parent-attempt-${ordinal}`), - nativeTurnRef: { - driver: providerThread.driver, - nativeId: `side-chat-parent-native-turn-${ordinal}`, - strength: "strong" as const, - }, + nativeTurnRef: + ordinal === 1 && input.sideChat?.sourceBoundary === "synthetic" + ? null + : { + driver: providerThread.driver, + nativeId: `side-chat-parent-native-turn-${ordinal}`, + strength: "strong" as const, + }, ordinal, status: "completed" as const, startedAt: now, @@ -528,7 +535,10 @@ function makeLocalCommandHarness(input: { ordinal: 1, providerThreadId: parentProviderThread.id, activeAttemptId: parentTurns[0]!.attempt.id, - status: "completed" as const, + status: + input.sideChat?.sourceRolledBack === true + ? ("rolled_back" as const) + : ("completed" as const), completedAt: now, }, { @@ -556,7 +566,9 @@ function makeLocalCommandHarness(input: { ], providerThreads: [parentProviderThread], attempts: parentTurns.map(({ attempt }) => attempt), - providerTurns: parentTurns.map(({ turn }) => turn), + providerTurns: parentTurns + .map(({ turn }) => turn) + .filter((turn) => turn.ordinal !== 1 || input.sideChat?.sourceBoundary !== "missing"), turnItems: [], }; if (input.sideChat !== undefined) { @@ -774,7 +786,15 @@ function makeLocalCommandHarness(input: { status: "ready", cwd: "/tmp/native-account-command", model: null, - capabilities: CodexProviderCapabilitiesV2, + capabilities: { + ...CodexProviderCapabilitiesV2, + threads: { + ...CodexProviderCapabilitiesV2.threads, + canForkWhileRunning: + input.sideChat?.canForkWhileRunning ?? + CodexProviderCapabilitiesV2.threads.canForkWhileRunning, + }, + }, createdAt: now, updatedAt: now, lastError: null, @@ -788,13 +808,14 @@ function makeLocalCommandHarness(input: { strength: "strong" as const, }, }; - const forkThread = vi.fn((_fork: ProviderAdapter.ProviderAdapterV2ForkThreadInput) => + const forkThread = vi.fn(() => Effect.succeed(sideChatThread), ); const sideChatSession = { driver: providerThread.driver, providerSession, - ensureThread: () => Effect.succeed(sideChatThread), + ensureThread: () => + "ensureThreadFailure" in input ? ensureThread() : Effect.succeed(sideChatThread), resumeThread: () => Effect.succeed(sideChatThread), forkThread, startTurn, @@ -898,6 +919,21 @@ function makeLocalCommandHarness(input: { const getThreadRecords = vi.fn((id) => Effect.succeed(id === parentThreadId ? parentRecords : projection), ); + const prepareProviderHandoff = vi.fn< + ContextHandoffService.ContextHandoffServiceV2["Service"]["prepareProviderHandoff"] + >((handoff) => + input.sideChat?.nativeForkFrom === undefined + ? Effect.die("history read must fail first") + : Effect.succeed({ + ...handoff, + id: ContextHandoffId.make("side-chat-fork-context"), + createdByProviderInstanceId: handoff.toProviderInstanceId, + status: "ready", + summaryMessageId: null, + summaryText: "PARENT_HISTORY", + updatedAt: handoff.createdAt, + }), + ); const getTurnStartContext = vi.fn(() => Effect.succeed({ ...projection, @@ -912,18 +948,7 @@ function makeLocalCommandHarness(input: { Layer.provide( Layer.mergeAll( Layer.mock(ContextHandoffService.ContextHandoffServiceV2)({ - prepareProviderHandoff: (handoff) => - input.sideChat?.nativeForkFrom === undefined - ? Effect.die("history read must fail first") - : Effect.succeed({ - ...handoff, - id: ContextHandoffId.make("side-chat-fork-context"), - createdByProviderInstanceId: handoff.toProviderInstanceId, - status: "ready", - summaryMessageId: null, - summaryText: "PARENT_HISTORY", - updatedAt: handoff.createdAt, - }), + prepareProviderHandoff, }), Layer.mock(EventSink.EventSinkV2)({ writeIfRunCurrent, @@ -1016,6 +1041,7 @@ function makeLocalCommandHarness(input: { startRootRun, startTurn, forkThread, + prepareProviderHandoff, interruptTurn, tryHandlePromptCommand, events, @@ -1295,20 +1321,61 @@ effectIt.effect.each([ ); effectIt.effect.each([ - { nativeForkFrom: "idle", forks: true, transferStatus: "consumed", handoffIds: [] }, + { + nativeForkFrom: "idle", + canForkWhileRunning: false, + sourceBoundary: "native", + forks: true, + transferStatus: "consumed", + handoffIds: [], + }, + { + nativeForkFrom: "idle", + canForkWhileRunning: false, + sourceBoundary: "missing", + forks: true, + transferStatus: "consumed", + handoffIds: [], + }, + { + nativeForkFrom: "busy", + canForkWhileRunning: true, + sourceBoundary: "native", + forks: true, + transferStatus: "consumed", + handoffIds: [], + }, { nativeForkFrom: "busy", + canForkWhileRunning: false, + sourceBoundary: "native", + forks: false, + transferStatus: "resolved_portable", + handoffIds: ["side-chat-fork-context"], + }, + { + nativeForkFrom: "busy", + canForkWhileRunning: true, + sourceBoundary: "synthetic", + forks: false, + transferStatus: "resolved_portable", + handoffIds: ["side-chat-fork-context"], + }, + { + nativeForkFrom: "busy", + canForkWhileRunning: true, + sourceBoundary: "missing", forks: false, transferStatus: "resolved_portable", handoffIds: ["side-chat-fork-context"], }, ] as const)( - "runs a planned native side chat fork when the parent is $nativeForkFrom", - ({ nativeForkFrom, forks, transferStatus, handoffIds }) => + "runs a planned native side chat fork when the parent is $nativeForkFrom with canForkWhileRunning=$canForkWhileRunning and sourceBoundary=$sourceBoundary", + ({ nativeForkFrom, canForkWhileRunning, sourceBoundary, forks, transferStatus, handoffIds }) => Effect.gen(function* () { const harness = makeLocalCommandHarness({ text: "What is the main thread doing?", - sideChat: { contextRun: "current", nativeForkFrom }, + sideChat: { contextRun: "current", nativeForkFrom, canForkWhileRunning, sourceBoundary }, }); yield* harness.start; @@ -1317,7 +1384,14 @@ effectIt.effect.each([ // The parent's rolled-back turn is gone from its native conversation. expect( harness.forkThread.mock.calls[0]?.[0].sourceProviderTurns?.map((turn) => turn.id), - ).toEqual(forks ? ["side-chat-parent-turn-1"] : undefined); + ).toEqual( + forks ? (sourceBoundary === "missing" ? [] : ["side-chat-parent-turn-1"]) : undefined, + ); + if (forks && sourceBoundary === "native") { + expect(harness.forkThread).toHaveBeenCalledWith( + expect.objectContaining({ providerTurnId: "side-chat-parent-turn-1" }), + ); + } expect(harness.projection().contextTransfers).toMatchObject([ { id: "side-chat-fork", status: transferStatus }, ]); @@ -1328,6 +1402,147 @@ effectIt.effect.each([ }), ); +effectIt.effect.each([ + { nativeForkFrom: "idle", rollback: "before-start" }, + { nativeForkFrom: "busy", rollback: "before-start" }, + { nativeForkFrom: "idle", rollback: "during-fork" }, + { nativeForkFrom: "busy", rollback: "during-fork" }, +] as const)( + "recovers a planned $nativeForkFrom fork after rollback $rollback", + ({ nativeForkFrom, rollback }) => + Effect.gen(function* () { + const harness = makeLocalCommandHarness({ + text: "Continue from the selected turn", + sideChat: { + contextRun: "current", + nativeForkFrom, + sourceRolledBack: rollback === "before-start", + }, + }); + const parentThreadId = ThreadId.make("side-chat-parent"); + let parent = { + ...harness.projection(), + ...(yield* harness.getThreadRecords(parentThreadId, [ + "runs", + "attempts", + "providerThreads", + "providerTurns", + "turnItems", + ])), + }; + const selectedRun = parent.runs[0]!; + const earlierRun = { + ...selectedRun, + id: RunId.make("side-chat-kept-run"), + ordinal: 1, + activeAttemptId: null, + status: "completed" as const, + }; + const items = [earlierRun, selectedRun].map((sourceRun, ordinal) => ({ + type: "user_message" as const, + id: TurnItemId.make(`parent-item-${ordinal}`), + threadId: parentThreadId, + runId: sourceRun.id, + nodeId: null, + providerThreadId: sourceRun.providerThreadId, + providerTurnId: null, + nativeItemRef: null, + parentItemId: null, + ordinal, + status: "completed" as const, + title: null, + startedAt: parent.updatedAt, + completedAt: parent.updatedAt, + updatedAt: parent.updatedAt, + createdBy: "user" as const, + creationSource: "server" as const, + messageId: MessageId.make(`parent-message-${ordinal}`), + text: ordinal === 0 ? "KEPT_HISTORY" : "ROLLED_BACK_HISTORY", + inputIntent: "turn_start" as const, + attachments: [], + })); + parent = { + ...parent, + runs: [ + earlierRun, + ...parent.runs.map((sourceRun) => ({ ...sourceRun, ordinal: sourceRun.ordinal + 1 })), + ], + turnItems: items, + }; + harness.getThreadRecords.mockImplementation((id) => + Effect.succeed(id === parentThreadId ? parent : harness.projection()), + ); + harness.forkThread.mockImplementation(() => + Effect.sync(() => { + parent = { + ...parent, + runs: parent.runs.map((sourceRun) => + sourceRun.id === selectedRun.id ? { ...sourceRun, status: "rolled_back" } : sourceRun, + ), + }; + }).pipe( + Effect.andThen( + Effect.fail( + new ProviderAdapter.ProviderAdapterForkThreadError({ + driver: ProviderDriverKind.make("codex"), + providerThreadId: ProviderThreadId.make("side-chat-parent-provider-thread"), + cause: "Selected native turn was removed by rollback", + }), + ), + ), + ), + ); + + yield* harness.startWithRetry; + + expect(harness.forkThread).toHaveBeenCalledTimes(rollback === "before-start" ? 0 : 1); + expect(harness.projection().contextTransfers).toMatchObject([ + { id: "side-chat-fork", status: "resolved_portable" }, + ]); + expect(harness.startTurn.mock.calls[0]?.[0].message.text).toContain("PARENT_HISTORY"); + expect(harness.projection().providerThreads[0]?.handoffIds).toEqual([ + "side-chat-fork-context", + ]); + expect(harness.settledRuns).toEqual([]); + expect( + harness.prepareProviderHandoff.mock.calls[0]?.[0].items.map((item) => item.id), + ).toEqual([items[0]!.id]); + }), +); + +effectIt.effect("preserves portable initialization failures after a native fork fails", () => + Effect.gen(function* () { + const harness = makeLocalCommandHarness({ + text: "Continue from the selected turn", + sideChat: { contextRun: "current", nativeForkFrom: "busy" }, + ensureThreadFailure: "Portable session unavailable", + }); + harness.forkThread.mockImplementation(() => + Effect.fail( + new ProviderAdapter.ProviderAdapterForkThreadError({ + driver: ProviderDriverKind.make("codex"), + providerThreadId: ProviderThreadId.make("side-chat-parent-provider-thread"), + cause: "Selected native turn was removed by rollback", + }), + ), + ); + + const result = yield* Effect.result(harness.startWithRetry); + + expect(result._tag).toBe("Failure"); + if (result._tag === "Failure") { + expect(result.failure).toMatchObject({ + _tag: "ProviderTurnStartError", + cause: { + _tag: "ProviderAdapterEventStreamError", + cause: "Portable session unavailable", + }, + }); + } + expect(harness.startTurn).not.toHaveBeenCalled(); + }), +); + effectIt.effect.each([ { runStatus: "failed", turnStatus: "completed" }, { runStatus: "completed", turnStatus: "completed" }, diff --git a/apps/server/src/orchestration-v2/ProviderTurnStartService.ts b/apps/server/src/orchestration-v2/ProviderTurnStartService.ts index cd121c7115a2a..7ac29e2feb402 100644 --- a/apps/server/src/orchestration-v2/ProviderTurnStartService.ts +++ b/apps/server/src/orchestration-v2/ProviderTurnStartService.ts @@ -1120,7 +1120,7 @@ export const layer: Layer.Layer< let portableForkFallback: OrchestrationV2ContextHandoff | undefined; const loadedProviderThread = yield* Effect.gen(function* () { if (nativeForkTransfer !== undefined) { - const sourceProjection = yield* projectionStore.getThreadRecords( + let sourceProjection = yield* projectionStore.getThreadRecords( nativeForkTransfer.sourceThreadId, ["runs", "providerThreads", "attempts", "providerTurns"], ); @@ -1144,12 +1144,16 @@ export const layer: Layer.Layer< cause: `Native fork transfer ${nativeForkTransfer.id} has no source provider execution.`, }); } - // Providers refuse to fork a busy thread, so a source that started a turn after - // dispatch planned this fork gets the portable context dispatch gives a busy source. + // The source may have started a turn after dispatch planned this fork. const sourceBusy = sourceProjection.runs.some( (candidate) => candidate.ordinal > sourceRun.ordinal && isActiveRun(candidate), ); - if (!sourceBusy) { + if ( + (sourceRun.status === "completed" || sourceRun.status === "waiting") && + (!sourceBusy || + (session.providerSession.capabilities.threads.canForkWhileRunning === true && + sourceProviderTurn?.nativeTurnRef != null)) + ) { // Rolled-back turns are no longer in the source conversation, so they cannot bound a fork. const rolledBackAttemptIds = new Set( sourceProjection.attempts @@ -1161,7 +1165,7 @@ export const layer: Layer.Layer< ) .map((attempt) => attempt.id), ); - return yield* loadFromProvider( + const forked = yield* Effect.result( session.forkThread({ sourceProviderThread, sourceProviderTurns: sourceProjection.providerTurns.filter( @@ -1176,6 +1180,23 @@ export const layer: Layer.Layer< : { providerTurnId: sourceProviderTurn.id }), }), ); + if (forked._tag === "Success") return forked.success; + if (forked.failure._tag !== "ProviderAdapterForkThreadError") { + return yield* loadFromProvider(Effect.fail(forked.failure)); + } + // A parent rollback can remove the boundary after the source snapshot was read. + yield* Effect.logWarning("Recovering a failed native fork with portable context", { + sourceThreadId: nativeForkTransfer.sourceThreadId, + runId, + error: forked.failure, + }); + sourceProjection = { + ...sourceProjection, + ...(yield* projectionStore.getThreadRecords(nativeForkTransfer.sourceThreadId, [ + "runs", + "attempts", + ])), + }; } const sourceItems = yield* ContextHandoffService.readForkHistoryItems( projectionStore, diff --git a/packages/contracts/src/orchestrationV2.ts b/packages/contracts/src/orchestrationV2.ts index 785f41e3aa0fb..d3431e9a5fe3b 100644 --- a/packages/contracts/src/orchestrationV2.ts +++ b/packages/contracts/src/orchestrationV2.ts @@ -272,6 +272,7 @@ export const OrchestrationV2ThreadCapabilities = Schema.Struct({ canRollbackThread: Schema.Boolean, canForkThread: Schema.Boolean, canForkFromTurn: Schema.Boolean, + canForkWhileRunning: Schema.optional(Schema.Boolean), canForkFromSubagentThread: Schema.Boolean, exposesNativeThreadId: Schema.Boolean, });