diff --git a/apps/server/src/orchestration-v2/CheckpointCaptureService.ts b/apps/server/src/orchestration-v2/CheckpointCaptureService.ts index 43cdae25d4bcd..350896a832f93 100644 --- a/apps/server/src/orchestration-v2/CheckpointCaptureService.ts +++ b/apps/server/src/orchestration-v2/CheckpointCaptureService.ts @@ -129,7 +129,9 @@ export const layer: Layer.Layer< }); // Omit delegatedCompletion to preserve cohort writes made after capture loaded the run. const { delegatedCompletion: _delegatedCompletion, ...runWithoutDelegatedCompletion } = run; - const commandId = CommandId.make(`command:effect:checkpoint.capture:${run.id}`); + const commandId = CommandId.make( + `command:effect:checkpoint.capture:${run.activeAttemptId ?? run.id}`, + ); yield* eventSink.commitCommand({ commandId, threadId: input.threadId, diff --git a/apps/server/src/orchestration-v2/Orchestrator.ts b/apps/server/src/orchestration-v2/Orchestrator.ts index dfce0ea977d39..48e51d61c4480 100644 --- a/apps/server/src/orchestration-v2/Orchestrator.ts +++ b/apps/server/src/orchestration-v2/Orchestrator.ts @@ -12286,7 +12286,7 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio yield* Ref.update(input.effects, (effects) => [ ...effects, { - id: `effect:checkpoint.capture:${run.id}`, + id: `effect:checkpoint.capture:${attempt.id}`, commandId: CommandId.make(`command:effect:checkpoint.capture:${run.id}`), threadId: input.command.threadId, request: { diff --git a/apps/server/src/orchestration-v2/RunExecutionService.ts b/apps/server/src/orchestration-v2/RunExecutionService.ts index 8e7947d3147c6..6631c61fa81a0 100644 --- a/apps/server/src/orchestration-v2/RunExecutionService.ts +++ b/apps/server/src/orchestration-v2/RunExecutionService.ts @@ -766,7 +766,7 @@ export const layer: Layer.Layer< input.terminal.status === "cancelled" ? [ { - id: `effect:checkpoint.capture:${input.run.id}`, + id: `effect:checkpoint.capture:${input.attempt.id}`, commandId: checkpointCaptureCommandId, threadId: input.run.threadId, request: { diff --git a/apps/server/src/orchestration-v2/SteeringCompletion.integration.test.ts b/apps/server/src/orchestration-v2/SteeringCompletion.integration.test.ts index acfd20af78749..c4b22046877b3 100644 --- a/apps/server/src/orchestration-v2/SteeringCompletion.integration.test.ts +++ b/apps/server/src/orchestration-v2/SteeringCompletion.integration.test.ts @@ -513,6 +513,7 @@ const nextTurnSelectionHarness = Effect.fn("nextTurnSelectionHarness")(function* turn: ProviderAdapter.ProviderAdapterV2SteerInput, ) => Effect.Effect = () => Effect.void, extraAdapters: ReadonlyArray = [], + beforeSelectionTransition: Effect.Effect = Effect.void, ) { const driver = ProviderDriverKind.make(nativeGoals === undefined ? "codex" : "opencode"); const cwd = yield* checkpointWorkspace(name); @@ -535,7 +536,8 @@ const nextTurnSelectionHarness = Effect.fn("nextTurnSelectionHarness")(function* instanceId, driver, getCapabilities: () => Effect.succeed(capabilities), - planSelectionTransition: () => Effect.succeed({ type: "apply_on_next_turn" }), + planSelectionTransition: () => + beforeSelectionTransition.pipe(Effect.as({ type: "apply_on_next_turn" as const })), openSession: (input) => Effect.gen(function* () { const now = yield* DateTime.now; @@ -1330,6 +1332,125 @@ it.effect("restarts the turn for an image when native steering takes text only", }), ); +it.effect.each(["completed", "interrupted"] as const)( + "checkpoints a run that a steer restarted after its first answer was checkpointed (restart %s)", + (restartEnd) => + Effect.gen(function* () { + const steerReadRunning = yield* Deferred.make(); + const letSteerCommit = yield* Deferred.make(); + const { events, started, layer, startFirstTurn } = yield* nextTurnSelectionHarness( + `steer-restart-after-checkpoint-${restartEnd}`, + true, + false, + Effect.void, + true, + undefined, + undefined, + [], + Deferred.succeed(steerReadRunning, undefined).pipe( + Effect.andThen(Deferred.await(letSteerCommit)), + ), + ); + yield* Effect.gen(function* () { + const orchestrator = yield* Orchestrator.OrchestratorV2; + const worker = yield* EffectWorker.OrchestrationEffectWorkerV2; + const outbox = yield* EffectOutbox.EffectOutboxV2; + const threadId = yield* startFirstTurn; + const runId = started[0]!.runId; + const runReaches = (status: string) => + orchestrator.streamDomainEvents.pipe( + Stream.filter( + (event) => + event.type === "run.updated" && + event.payload.id === runId && + event.payload.status === status, + ), + Stream.take(1), + Stream.runDrain, + Effect.forkScoped, + ); + const endTurn = ( + turn: ProviderAdapter.ProviderAdapterV2TurnInput, + status: "completed" | "interrupted", + ) => + Effect.gen(function* () { + const providerTurn = (yield* orchestrator.getThreadProjection( + threadId, + )).providerTurns.find((candidate) => candidate.runAttemptId === turn.attemptId)!; + const finalized = yield* runReaches(status === "completed" ? "waiting" : status); + yield* Queue.offer(events, { + type: "provider_turn.updated", + driver, + providerTurn: { ...providerTurn, status, completedAt: yield* DateTime.now }, + }); + yield* Queue.offer(events, { + type: "turn.terminal", + driver, + providerThreadId: providerTurn.providerThreadId, + providerTurnId: providerTurn.id, + runOrdinal: turn.runOrdinal, + status, + failure: null, + threadDisposition: "reusable", + }); + yield* Fiber.join(finalized); + }); + + const steer = yield* orchestrator + .dispatch({ + type: "message.dispatch", + commandId: CommandId.make("steer-after-settle"), + threadId, + messageId: MessageId.make("message:steer-after-settle"), + text: "also check the logs", + attachments: [], + modelSelection: composerSelection, + dispatchMode: { type: "steer_active", targetRunId: runId }, + createdBy: "agent", + creationSource: "mcp", + }) + .pipe(Effect.forkScoped); + yield* Deferred.await(steerReadRunning); + const firstCompleted = yield* runReaches("completed"); + yield* endTurn(started[0]!, "completed"); + const firstCapture = yield* worker.drain().pipe(Effect.forkScoped); + yield* Fiber.join(firstCompleted); + const restartRunning = yield* orchestrator.streamDomainEvents.pipe( + Stream.filter( + (event) => + event.type === "provider-turn.updated" && + event.payload.status === "running" && + event.payload.runAttemptId !== started[0]!.attemptId, + ), + Stream.take(1), + Stream.runDrain, + Effect.forkScoped, + ); + yield* Deferred.succeed(letSteerCommit, undefined); + yield* Fiber.join(steer); + yield* Fiber.join(firstCapture); + yield* worker.drain(); + yield* Fiber.join(restartRunning); + yield* (yield* ThreadCommandExecutor.ThreadCommandExecutor).withLock(threadId, Effect.void); + assert.equal(started.length, 2); + assert.equal(started[1]!.runId, runId); + + yield* endTurn(started[1]!, restartEnd); + // Startup recovery keeps a waiting run alive only if it finds a pending capture by this id. + const captures = yield* outbox.listByCommandId( + CommandId.make(`command:effect:checkpoint.capture:${runId}`), + ); + assert.equal(captures.filter((capture) => capture.status === "pending").length, 1); + yield* worker.drain(); + const run = (yield* orchestrator.getThreadProjection(threadId)).runs.find( + (candidate) => candidate.id === runId, + ); + assert.equal(run?.status, restartEnd); + assert.notEqual(run?.checkpointId, null); + }).pipe(Effect.provide(layer)); + }), +); + it.effect("starts a steer that missed the turn on the saved next-turn selection", () => Effect.gen(function* () { const { events, started, steered, layer, startFirstTurn } = yield* nextTurnSelectionHarness(