diff --git a/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.ts b/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.ts index e9b41ac8a5b05..747b8e4056931 100644 --- a/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.ts +++ b/apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.ts @@ -266,6 +266,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 85babb8a869bf..b98a5d544cf00 100644 --- a/apps/server/src/orchestration-v2/CommandPolicy.test.ts +++ b/apps/server/src/orchestration-v2/CommandPolicy.test.ts @@ -271,6 +271,7 @@ layer("CommandPolicyV2", (it) => { capabilities: CursorProviderCapabilitiesV2, sameProvider: true, hasStrongNativeSource: true, + hasNativeTurnBoundary: false, sourceRunStatus: "completed", fromSpecificTurn: true, sourceBusy: false, @@ -280,24 +281,38 @@ 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: 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 cd56460369bd6..b54c994cd1e6d 100644 --- a/apps/server/src/orchestration-v2/CommandPolicy.ts +++ b/apps/server/src/orchestration-v2/CommandPolicy.ts @@ -219,9 +219,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; @@ -386,7 +387,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 c7cd7a9ff288e..ce0a0f1e44661 100644 --- a/apps/server/src/orchestration-v2/Orchestrator.ts +++ b/apps/server/src/orchestration-v2/Orchestrator.ts @@ -8017,6 +8017,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 17d40cf05162e..fe761106acfa6 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, @@ -306,6 +307,9 @@ function makeLocalCommandHarness(input: { readonly sideChat?: { readonly contextRun: "current" | "failed" | "cancelled" | "delivered"; readonly nativeForkFrom?: "idle" | "busy"; + readonly canForkWhileRunning?: boolean; + readonly sourceBoundary?: "native" | "synthetic"; + readonly sourceRolledBack?: boolean; readonly switchedProvider?: boolean; readonly threadRuntimeMode?: OrchestrationV2ThreadProjection["thread"]["runtimeMode"]; }; @@ -510,11 +514,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, @@ -537,7 +544,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, }, { @@ -787,6 +797,12 @@ function makeLocalCommandHarness(input: { model: null, capabilities: { ...CodexProviderCapabilitiesV2, + threads: { + ...CodexProviderCapabilitiesV2.threads, + canForkWhileRunning: + input.sideChat?.canForkWhileRunning ?? + CodexProviderCapabilitiesV2.threads.canForkWhileRunning, + }, turns: { ...CodexProviderCapabilitiesV2.turns, ...(input.nativeGoals === undefined ? {} : { supportsNativeGoals: input.nativeGoals }), @@ -805,7 +821,7 @@ function makeLocalCommandHarness(input: { strength: "strong" as const, }, }; - const forkThread = vi.fn((_fork: ProviderAdapter.ProviderAdapterV2ForkThreadInput) => + const forkThread = vi.fn(() => Effect.succeed(sideChatThread), ); const sideChatSession = { @@ -915,6 +931,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, @@ -948,18 +979,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, @@ -1052,6 +1072,7 @@ function makeLocalCommandHarness(input: { startRootRun, startTurn, forkThread, + prepareProviderHandoff, interruptTurn, tryHandlePromptCommand, events, @@ -1336,20 +1357,45 @@ 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: "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"], }, ] 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; @@ -1359,6 +1405,11 @@ effectIt.effect.each([ expect( harness.forkThread.mock.calls[0]?.[0].sourceProviderTurns?.map((turn) => turn.id), ).toEqual(forks ? ["side-chat-parent-turn-1"] : undefined); + if (forks) { + expect(harness.forkThread).toHaveBeenCalledWith( + expect.objectContaining({ providerTurnId: "side-chat-parent-turn-1" }), + ); + } expect(harness.projection().contextTransfers).toMatchObject([ { id: "side-chat-fork", status: transferStatus }, ]); @@ -1392,6 +1443,109 @@ effectIt.effect.each([ }), ); +effectIt.effect.each(["before-start", "during-fork"] as const)( + "recovers a planned native fork after the parent rolls back %s", + (rollback) => + Effect.gen(function* () { + const harness = makeLocalCommandHarness({ + text: "Continue from the selected turn", + sideChat: { + contextRun: "current", + nativeForkFrom: "idle", + 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.start; + + 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.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 648e12e5a2aa8..7ec476cfd29b1 100644 --- a/apps/server/src/orchestration-v2/ProviderTurnStartService.ts +++ b/apps/server/src/orchestration-v2/ProviderTurnStartService.ts @@ -1145,7 +1145,7 @@ export const layer: Layer.Layer< let recoveryTurnItem: OrchestrationV2TurnItem | 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"], ); @@ -1169,13 +1169,17 @@ 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), ); let nativeForkFailure: string | undefined; - 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 @@ -1218,6 +1222,13 @@ export const layer: Layer.Layer< message: startFailureMessage(forked.failure), class: "provider_error", }).message; + sourceProjection = { + ...sourceProjection, + ...(yield* projectionStore.getThreadRecords(nativeForkTransfer.sourceThreadId, [ + "runs", + "attempts", + ])), + }; } const nativeForkFailed = nativeForkFailure !== undefined; const sourceItems = yield* ContextHandoffService.readForkHistoryItems( diff --git a/packages/contracts/src/orchestrationV2.ts b/packages/contracts/src/orchestrationV2.ts index 38046a17e8b4a..518cd9aea9d5d 100644 --- a/packages/contracts/src/orchestrationV2.ts +++ b/packages/contracts/src/orchestrationV2.ts @@ -285,6 +285,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, });