diff --git a/apps/server/src/orchestration-v2/DelegatedCompletionDelivery.test.ts b/apps/server/src/orchestration-v2/DelegatedCompletionDelivery.test.ts index 64ed215ebed5c..7bc7ef37d42de 100644 --- a/apps/server/src/orchestration-v2/DelegatedCompletionDelivery.test.ts +++ b/apps/server/src/orchestration-v2/DelegatedCompletionDelivery.test.ts @@ -1380,6 +1380,145 @@ const seedRestartCancelledChild = (input: { return { taskId, childThreadId, childRunId }; }); +// Parent spawned two always-wake tasks: the incoming one's completion wake is queued behind a busy +// run while the sibling's child run is still going. +const queueCompletionWakeBesideRunningSibling = (key: string) => + Effect.gen(function* () { + const orchestrator = yield* Orchestrator.OrchestratorV2; + const sink = yield* EventSink.EventSinkV2; + const now = yield* DateTime.now; + const threadId = ThreadId.make(`${key}-parent`); + const originRunId = RunId.make(`${key}-origin`); + const originRootId = NodeId.make(`${key}-origin-root`); + const incomingTaskId = NodeId.make(`${key}-incoming`); + yield* seedParentWithTerminalTask({ + threadId, + projectId: ProjectId.make(`${key}-project`), + runId: originRunId, + rootNodeId: originRootId, + taskId: incomingTaskId, + deliveryState: "claimed", + deliveryTaskIds: [incomingTaskId], + completionWake: "always", + now, + }); + yield* sink.write({ + events: [ + { + id: EventId.make(`event:${key}-origin-root`), + type: "node.updated", + threadId, + runId: originRunId, + nodeId: originRootId, + occurredAt: now, + payload: { + id: originRootId, + threadId, + runId: originRunId, + parentNodeId: null, + rootNodeId: originRootId, + kind: "root_turn", + status: "running", + countsForRun: true, + providerThreadId: null, + providerTurnId: null, + nativeItemRef: null, + runtimeRequestId: null, + checkpointScopeId: null, + startedAt: now, + completedAt: null, + }, + }, + ], + }); + yield* orchestrator.dispatch({ + type: "delegated_task.request", + createdBy: "agent", + creationSource: "mcp", + commandId: CommandId.make(`${key}-sibling`), + parentThreadId: threadId, + parentRunId: originRunId, + parentNodeId: originRootId, + task: "Keep working", + modelSelection, + runtimeMode: "full-access", + interactionMode: "default", + completionWake: "always", + }); + const seeded = yield* orchestrator.getThreadProjection(threadId); + const origin = seeded.runs.find((run) => run.id === originRunId)!; + const sibling = seeded.subagents.find((task) => task.id !== incomingTaskId)!; + // The coordinator finished its spawning turn but is busy with another run, so the wake queues. + yield* sink.write({ + events: [ + { + ...runEvent({ threadId, runId: originRunId, ordinal: 1, status: "completed", now }), + payload: { ...origin, status: "completed", completedAt: now }, + }, + runEvent({ + threadId, + runId: RunId.make(`${key}-busy`), + ordinal: 2, + status: "running", + now, + }), + ], + }); + const delivery = origin.delegatedCompletion!.delivery!; + yield* orchestrator.dispatch({ + type: "message.dispatch", + createdBy: "agent", + creationSource: "server", + commandId: CommandId.make(`${key}-wake`), + threadId, + messageId: delivery.messageId, + text: "Delegated task completed", + attachments: [], + modelSelection, + dispatchMode: { type: "queue_after_active" }, + delegatedCompletion: { + parentRunId: originRunId, + generation: delivery.generation, + taskIds: delivery.taskIds, + }, + }); + const wake = (yield* orchestrator.getThreadProjection(threadId)).runs.find( + (run) => run.userMessageId === delivery.messageId, + )!; + assert.equal(wake.status, "queued"); + const completeSibling = Effect.gen(function* () { + const childRun = (yield* orchestrator.getThreadProjection(sibling.childThreadId!)).runs[0]!; + const afterSequence = yield* sink.latestSequence(); + yield* sink.write({ + events: [ + { + ...runEvent({ + threadId: sibling.childThreadId!, + runId: childRun.id, + ordinal: childRun.ordinal, + status: "completed", + now, + }), + payload: { ...childRun, status: "completed", completedAt: now }, + }, + ], + }); + return afterSequence; + }); + return { + orchestrator, + sink, + now, + threadId, + originRunId, + incomingTaskId, + sibling, + delivery, + wake, + completeSibling, + }; + }); + it.layer(layerContinuationWorkerTest)("interrupted completion wakes", (it) => { it.effect.each([ { name: "an interrupt keeps", holdQueue: false, siblingFinishesFirst: false }, @@ -1623,6 +1762,137 @@ it.layer(layerContinuationWorkerTest)("interrupted completion wakes", (it) => { ); }), ); + + it.effect("cancelling a queued completion wake keeps its other siblings", () => + Effect.gen(function* () { + const { + orchestrator, + sink, + threadId, + originRunId, + incomingTaskId, + sibling, + delivery, + wake, + completeSibling, + } = yield* queueCompletionWakeBesideRunningSibling("sibling-queued-cancel"); + const cancelled = yield* orchestrator.dispatch({ + type: "queued-run.cancel", + commandId: CommandId.make("sibling-queued-cancel-cancel"), + threadId, + runId: wake.id, + }); + const cancelDelivery = new Map( + cancelled.storedEvents.flatMap(({ event }) => + event.type === "subagent.updated" + ? [[event.payload.id, event.payload.completionDelivery?.state] as const] + : [], + ), + ); + assert.equal(cancelDelivery.get(incomingTaskId), "disposed"); + assert.notEqual( + cancelDelivery.get(sibling.id), + "disposed", + "cancelling the wake must not dispose a still-running sibling it did not carry", + ); + const afterCancel = yield* orchestrator.getThreadProjection(threadId); + assert.equal( + afterCancel.runs.find((run) => run.id === originRunId)?.delegatedCompletion?.disposition, + "open", + ); + + const afterSequence = yield* completeSibling; + yield* sink.stream({ afterSequence, eventType: "subagent.updated" }).pipe( + Stream.filter( + (stored) => + stored.event.type === "subagent.updated" && + stored.event.payload.id === sibling.id && + stored.event.payload.completionDelivery?.state === "claimed", + ), + Stream.take(1), + Stream.runDrain, + ); + const successor = (yield* orchestrator.getThreadProjection(threadId)).runs.find( + (run) => run.id === originRunId, + )?.delegatedCompletion?.delivery; + assert.deepEqual(successor?.taskIds, [sibling.id]); + assert.notEqual(successor?.messageId, delivery.messageId); + }), + ); + + it.effect("settling with a queued completion wake disposes its running siblings", () => + Effect.gen(function* () { + const { orchestrator, sink, now, threadId, originRunId, sibling, wake, completeSibling } = + yield* queueCompletionWakeBesideRunningSibling("sibling-settle"); + const busy = (yield* orchestrator.getThreadProjection(threadId)).runs.find( + (run) => run.status === "running", + )!; + yield* sink.write({ + events: [ + { + ...runEvent({ + threadId, + runId: busy.id, + ordinal: busy.ordinal, + status: "cancelled", + now, + }), + payload: { ...busy, status: "cancelled", completedAt: now }, + }, + { + ...runEvent({ threadId, runId: wake.id, ordinal: wake.ordinal, status: "queued", now }), + payload: { ...wake, queueHeld: true }, + }, + ], + }); + + const settled = yield* orchestrator.dispatch({ + type: "thread.settle", + commandId: CommandId.make("sibling-settle-settle"), + threadId, + }); + assert.deepEqual( + settled.storedEvents.flatMap(({ event }) => + event.type === "run.updated" && event.payload.id === wake.id + ? [event.payload.status] + : [], + ), + ["cancelled"], + ); + const afterSettle = yield* orchestrator.getThreadProjection(threadId); + assert.equal( + afterSettle.subagents.find((task) => task.id === sibling.id)?.completionDelivery?.state, + "disposed", + ); + + const afterSequence = yield* completeSibling; + yield* sink.stream({ afterSequence, eventType: "subagent.updated" }).pipe( + Stream.filter( + (stored) => + stored.event.type === "subagent.updated" && + stored.event.payload.id === sibling.id && + stored.event.payload.status === "completed", + ), + Stream.take(1), + Stream.runDrain, + ); + const afterLateResult = yield* orchestrator.getThreadProjection(threadId); + assert.equal( + afterLateResult.subagents.find((task) => task.id === sibling.id)?.completionDelivery?.state, + "disposed", + ); + assert.deepEqual( + afterLateResult.runs.map((run) => run.id), + afterSettle.runs.map((run) => run.id), + ); + assert.equal( + afterLateResult.runs.find((run) => run.id === originRunId)?.delegatedCompletion + ?.disposition, + "disposed", + ); + assert.equal(afterLateResult.thread.settledOverride, "settled"); + }), + ); }); it.layer(Layer.merge(layerTest, ProviderContinuationRequests.layer))( diff --git a/apps/server/src/orchestration-v2/Orchestrator.ts b/apps/server/src/orchestration-v2/Orchestrator.ts index a95e698a1fc93..e95fb6ec48656 100644 --- a/apps/server/src/orchestration-v2/Orchestrator.ts +++ b/apps/server/src/orchestration-v2/Orchestrator.ts @@ -3012,8 +3012,8 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio readonly disposition: "stopped" | "disposed"; readonly now: DateTime.Utc; readonly cancelQueuedDelivery?: boolean; - // An interrupted wake disposes only the results it carried. The cohort stays open, and its - // delivery stays the wake's, so the wake's terminal reserves any sibling's result. + // An interrupted or cancelled wake disposes only the results it carried. The cohort stays + // open, and its delivery stays the wake's, so the wake's terminal reserves any sibling's result. readonly interruptedWakeTaskIds?: ReadonlyArray; }) { const parentRun = input.projection.runs.find((candidate) => candidate.id === input.parentRunId); @@ -3752,6 +3752,7 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio runId: run.id, }, events, + "cohort", ); } } @@ -10301,6 +10302,8 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio const dispatchQueuedRunCancel = Effect.fn("Orchestrator.dispatchQueuedRunCancel")(function* ( command: Extract, events: Ref.Ref>, + // Settling disposes a cancelled wake's whole cohort, so a sibling's late result cannot unsettle it. + disposes: "wake" | "cohort" = "wake", ) { const projection = yield* loadProjectionForCommand( command, @@ -10318,8 +10321,8 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio const queuedMessage = projection.messages.find( (candidate) => candidate.id === queuedRun.userMessageId, ); + const now = yield* DateTime.now; if (queuedMessage?.delegatedCompletion !== undefined) { - const now = yield* DateTime.now; yield* disposeDelegatedCompletionCohort({ command, events, @@ -10327,64 +10330,14 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio parentRunId: queuedMessage.delegatedCompletion.parentRunId, disposition: "disposed", now, + cancelQueuedDelivery: false, + ...(disposes === "wake" + ? { interruptedWakeTaskIds: queuedMessage.delegatedCompletion.taskIds } + : {}), }); - return; } - const queuedRootNode = - queuedRun.rootNodeId === null - ? undefined - : projection.nodes.find((candidate) => candidate.id === queuedRun.rootNodeId); - const queuedAttempt = - queuedRun.activeAttemptId === null - ? undefined - : projection.attempts.find((candidate) => candidate.id === queuedRun.activeAttemptId); - - const now = yield* DateTime.now; + yield* emitQueuedRunCancellation({ command, events, projection, run: queuedRun, now }); const emitEvent = emit(events, command); - yield* emitEvent({ - type: "run.updated", - threadId: command.threadId, - runId: queuedRun.id, - ...(queuedRun.rootNodeId === null ? {} : { nodeId: queuedRun.rootNodeId }), - providerInstanceId: queuedRun.providerInstanceId, - occurredAt: now, - payload: { - ...queuedRun, - status: "cancelled", - queuePosition: null, - completedAt: now, - }, - }); - if (queuedAttempt !== undefined && queuedRootNode !== undefined) { - yield* emitEvent({ - type: "run-attempt.updated", - threadId: command.threadId, - runId: queuedRun.id, - nodeId: queuedRootNode.id, - providerInstanceId: queuedRun.providerInstanceId, - occurredAt: now, - payload: { - ...queuedAttempt, - status: "cancelled", - completedAt: now, - }, - }); - } - if (queuedRootNode !== undefined) { - yield* emitEvent({ - type: "node.updated", - threadId: command.threadId, - runId: queuedRun.id, - nodeId: queuedRootNode.id, - providerInstanceId: queuedRun.providerInstanceId, - occurredAt: now, - payload: { - ...queuedRootNode, - status: "cancelled", - completedAt: now, - }, - }); - } const guest = queuedMessage?.context?.guest; if (guest !== undefined) { const thread = guestShareAfterCancelledSend(projection.thread, guest, now);