diff --git a/apps/server/src/orchestration-v2/DelegatedCompletionDelivery.test.ts b/apps/server/src/orchestration-v2/DelegatedCompletionDelivery.test.ts index cabc322d31331..3f0f9cf0342ee 100644 --- a/apps/server/src/orchestration-v2/DelegatedCompletionDelivery.test.ts +++ b/apps/server/src/orchestration-v2/DelegatedCompletionDelivery.test.ts @@ -158,7 +158,7 @@ const seedParentWithTerminalTask = (input: { readonly runId: RunId; readonly rootNodeId: NodeId; readonly taskId: NodeId; - readonly deliveryState: "delivered" | "claimed" | "acknowledged" | "disposed"; + readonly deliveryState: "pending" | "delivered" | "claimed" | "acknowledged" | "disposed"; readonly completionWake?: "always" | "settled_only"; readonly deliveryTaskIds?: ReadonlyArray; readonly parentRunStatus?: "running" | "completed"; @@ -1049,6 +1049,7 @@ it.layer(layerTest)("delegated completion delivery repairs", (it) => { { command: "acknowledge", before: "running", after: "failed", capture: false }, { command: "acknowledge", before: "waiting", after: "completed", capture: true }, { command: "accept", before: "waiting", after: "completed", capture: true }, + { command: "wake-policy", before: "running", after: "completed", capture: true }, ] as const)( "keeps a $after run that settled between $command's read and commit", ({ command, before, after, capture }) => @@ -1062,14 +1063,20 @@ it.layer(layerTest)("delegated completion delivery repairs", (it) => { const taskId = NodeId.make(`node:${race}-task`); const checkpointId = CheckpointId.make(`checkpoint:${race}`); const messageId = MessageId.make(`message:delegated-delivery:${threadId}`); + const commandType = { + acknowledge: "delegated_task.completion-delivery.acknowledge", + accept: "notification.delivery.accept", + "wake-policy": "delegated_task.wake-policy", + }[command]; yield* seedParentWithTerminalTask({ threadId, projectId: ProjectId.make(`project:${race}`), runId, rootNodeId: NodeId.make(`node:${race}-root`), taskId, - deliveryState: "claimed", - deliveryTaskIds: [taskId], + ...(command === "wake-policy" + ? { deliveryState: "pending" } + : { deliveryState: "claimed", deliveryTaskIds: [taskId] }), now, }); const parent = (yield* orchestrator.getThreadProjection(threadId)).runs[0]!; @@ -1110,18 +1117,18 @@ it.layer(layerTest)("delegated completion delivery repairs", (it) => { const { delegatedCompletion: _cohort, ...settledRun } = parent; const commit = sink.commitCommand; let settled = false; + let planned: OrchestrationV2Run["delegatedCompletion"]; yield* Effect.acquireUseRelease( Effect.sync(() => vi.spyOn(sink, "commitCommand").mockImplementation((input) => Effect.gen(function* () { - if ( - !settled && - input.commandType === - (command === "accept" - ? "notification.delivery.accept" - : "delegated_task.completion-delivery.acknowledge") - ) { + if (!settled && input.commandType === commandType) { settled = true; + const plannedRun = input.events.find((event) => event.type === "run.updated"); + planned = + plannedRun?.type === "run.updated" + ? plannedRun.payload.delegatedCompletion + : undefined; // The provider turn ends (or its checkpoint lands) after the command read the run. yield* commit({ commandId: CommandId.make(`command:${race}`), @@ -1159,13 +1166,21 @@ it.layer(layerTest)("delegated completion delivery repairs", (it) => { threadId, messageId, } - : { - type: "delegated_task.completion-delivery.acknowledge", - commandId: CommandId.make(`command:${race}:task-status`), - parentThreadId: threadId, - taskId, - observedByRunId: runId, - }, + : command === "wake-policy" + ? { + type: "delegated_task.wake-policy", + commandId: CommandId.make(`command:${race}:wake-policy`), + parentThreadId: threadId, + taskId, + completionWake: "always", + } + : { + type: "delegated_task.completion-delivery.acknowledge", + commandId: CommandId.make(`command:${race}:task-status`), + parentThreadId: threadId, + taskId, + observedByRunId: runId, + }, ), (spy) => Effect.sync(() => spy.mockRestore()), ); @@ -1173,12 +1188,22 @@ it.layer(layerTest)("delegated completion delivery repairs", (it) => { const finished = yield* orchestrator.getThreadProjection(threadId); assert.equal( finished.subagents.find((row) => row.id === taskId)?.completionDelivery?.state, - command === "accept" ? "delivered" : "acknowledged", + { acknowledge: "acknowledged", accept: "delivered", "wake-policy": "claimed" }[command], ); const run = finished.runs.find((candidate) => candidate.id === runId)!; assert.deepEqual( - { status: run.status, completedAt: run.completedAt, checkpointId: run.checkpointId }, - { status: after, completedAt: now, checkpointId: capture ? checkpointId : null }, + { + status: run.status, + completedAt: run.completedAt, + checkpointId: run.checkpointId, + delegatedCompletion: run.delegatedCompletion, + }, + { + status: after, + completedAt: now, + checkpointId: capture ? checkpointId : null, + delegatedCompletion: planned, + }, ); }), ); diff --git a/apps/server/src/orchestration-v2/Orchestrator.ts b/apps/server/src/orchestration-v2/Orchestrator.ts index 59b708f5ba1cf..1059c88230788 100644 --- a/apps/server/src/orchestration-v2/Orchestrator.ts +++ b/apps/server/src/orchestration-v2/Orchestrator.ts @@ -8586,6 +8586,7 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio payload: completionPlan.message, }); } + return completionPlan.parentRun?.id; }); const dispatchCreatedThreadRecord = Effect.fn("orchestrationV2.dispatch.createdThreadRecord")( @@ -13413,7 +13414,10 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio yield* dispatchRemoteDelegatedTaskStop(command, events); break; case "delegated_task.wake-policy": - yield* dispatchDelegatedTaskWakePolicy(command, events); + updateDelegatedCompletionOnlyRunId = yield* dispatchDelegatedTaskWakePolicy( + command, + events, + ); break; case "delegated_task.completion-delivery.acknowledge": case "delegated_task.completion-delivery.dispose":