Skip to content

Commit aaa3dba

Browse files
juliusmarmingejuliusmarminge
authored
fix(server): a follow-up to a delegated child waits for its awaited background work (#36411)
1 parent c8db8be commit aaa3dba

4 files changed

Lines changed: 312 additions & 14 deletions

File tree

‎apps/server/src/mcp/OrchestratorMcpService.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,7 @@ import type {
102102
ProjectionRecordFilter,
103103
} from "../orchestration-v2/ProjectionStore.ts";
104104
import {
105-
delegatedTaskAwaitsSubagents,
105+
delegatedTaskAwaitsWork,
106106
delegatedTaskResultRun,
107107
delegatedTaskRuns,
108108
runOutcome,
@@ -1640,7 +1640,7 @@ const make = Effect.gen(function* () {
16401640
? progress.state
16411641
: !isTerminalTaskStatus(taskStatusForRun(taskResultRun))
16421642
? "working"
1643-
: delegatedTaskAwaitsSubagents(followupRuns, childControls.subagents)
1643+
: delegatedTaskAwaitsWork(followupRuns, childControls)
16441644
? "waiting_for_children"
16451645
: "result_available";
16461646
const resultRunIds = [

‎apps/server/src/orchestration-v2/DelegatedCompletionDelivery.test.ts‎

Lines changed: 240 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3190,6 +3190,246 @@ it.layer(Layer.merge(layerTest, ProviderContinuationRequests.layer))(
31903190
}),
31913191
);
31923192

3193+
it.effect.each(["woken", "cancelled"] as const)(
3194+
"holds a follow-up open while the child's awaited background command runs (%s)",
3195+
(ending) =>
3196+
Effect.gen(function* () {
3197+
const orchestrator = yield* Orchestrator.OrchestratorV2;
3198+
const eventSink = yield* EventSink.EventSinkV2;
3199+
const now = yield* DateTime.now;
3200+
const parentThreadId = ThreadId.make(`thread:followup-command-parent-${ending}`);
3201+
const projectId = ProjectId.make(`project:followup-command-${ending}`);
3202+
const parentRunId = RunId.make(`run:followup-command-parent-${ending}`);
3203+
const rootNodeId = NodeId.make(`node:followup-command-parent-${ending}`);
3204+
yield* seedParentWithTerminalTask({
3205+
threadId: parentThreadId,
3206+
projectId,
3207+
runId: parentRunId,
3208+
rootNodeId,
3209+
taskId: NodeId.make(`node:followup-command-seed-${ending}`),
3210+
deliveryState: "acknowledged",
3211+
now,
3212+
});
3213+
const child = yield* seedRestartCancelledChild({
3214+
parentThreadId,
3215+
projectId,
3216+
parentRunId,
3217+
rootNodeId,
3218+
name: `followup-command-child-${ending}`,
3219+
completionWake: "always",
3220+
continuationPending: false,
3221+
runStatus: "completed",
3222+
now,
3223+
});
3224+
yield* orchestrator.recoverDelegatedTasks;
3225+
3226+
const followupMessageId = MessageId.make(`message:followup-command-send-${ending}`);
3227+
yield* orchestrator.dispatch({
3228+
type: "message.dispatch",
3229+
createdBy: "agent",
3230+
creationSource: "mcp",
3231+
commandId: CommandId.make(`command:followup-command-send-${ending}`),
3232+
threadId: child.childThreadId,
3233+
senderThreadId: parentThreadId,
3234+
messageId: followupMessageId,
3235+
text: "Start the check in the background, wait on it, and reply when it ends.",
3236+
attachments: [],
3237+
dispatchMode: { type: "queue_after_active" },
3238+
});
3239+
const followupRun = (yield* orchestrator.getThreadProjection(
3240+
child.childThreadId,
3241+
)).runs.find((run) => run.userMessageId === followupMessageId)!;
3242+
const followupTaskId = followupRun.delegatedTaskId!;
3243+
const parent = yield* orchestrator.getThreadProjection(parentThreadId);
3244+
3245+
// The child awaits a background command, replies with an interim note and ends its turn.
3246+
yield* eventSink.write({
3247+
commandId: reconcileCommandId(`followup-command-turn-ended-${ending}`),
3248+
events: [
3249+
{
3250+
id: EventId.make(`event:followup-command-awaited-${ending}`),
3251+
type: "provider-thread.updated",
3252+
threadId: child.childThreadId,
3253+
driver,
3254+
providerInstanceId: modelSelection.instanceId,
3255+
occurredAt: now,
3256+
payload: {
3257+
...parent.providerThreads[0]!,
3258+
id: ProviderThreadId.make(`provider-thread:followup-command-child-${ending}`),
3259+
appThreadId: child.childThreadId,
3260+
pendingBackgroundTasks: [
3261+
{ taskId: "sleep-90", kind: "command" as const, awaited: true as const },
3262+
],
3263+
},
3264+
},
3265+
{
3266+
id: EventId.make(`event:followup-command-interim-${ending}`),
3267+
type: "message.updated",
3268+
threadId: child.childThreadId,
3269+
runId: followupRun.id,
3270+
occurredAt: now,
3271+
payload: {
3272+
id: MessageId.make(`message:followup-command-interim-${ending}`),
3273+
threadId: child.childThreadId,
3274+
runId: followupRun.id,
3275+
nodeId: null,
3276+
role: "assistant",
3277+
text: "INTERIM",
3278+
attachments: [],
3279+
streaming: false,
3280+
createdBy: "agent",
3281+
creationSource: "server",
3282+
createdAt: now,
3283+
updatedAt: now,
3284+
},
3285+
},
3286+
{
3287+
id: EventId.make(`event:followup-command-run-completed-${ending}`),
3288+
type: "run.updated",
3289+
threadId: child.childThreadId,
3290+
runId: followupRun.id,
3291+
occurredAt: now,
3292+
payload: { ...followupRun, status: "completed", startedAt: now, completedAt: now },
3293+
},
3294+
],
3295+
});
3296+
yield* orchestrator.recoverDelegatedTasks;
3297+
3298+
const held = (yield* orchestrator.getThreadProjection(parentThreadId)).subagents.find(
3299+
(task) => task.id === followupTaskId,
3300+
);
3301+
assert.isNull(
3302+
held?.result ?? null,
3303+
"the follow-up closed with its interim reply while the awaited command still ran",
3304+
);
3305+
assert.equal(held?.status, "running");
3306+
3307+
if (ending === "cancelled") {
3308+
yield* orchestrator.cancelDelegatedTask({
3309+
parentThreadId,
3310+
taskId: followupTaskId,
3311+
command: {
3312+
type: "thread.stop",
3313+
commandId: CommandId.make(`command:followup-command-cancel-${ending}`),
3314+
threadId: child.childThreadId,
3315+
stopDelegatedTasks: false,
3316+
},
3317+
stopDescendants: () => Effect.void,
3318+
});
3319+
yield* eventSink.write({
3320+
commandId: reconcileCommandId(`followup-command-stopped-${ending}`),
3321+
events: [
3322+
{
3323+
id: EventId.make(`event:followup-command-stopped-${ending}`),
3324+
type: "provider-thread.updated",
3325+
threadId: child.childThreadId,
3326+
driver,
3327+
providerInstanceId: modelSelection.instanceId,
3328+
occurredAt: now,
3329+
payload: {
3330+
...parent.providerThreads[0]!,
3331+
id: ProviderThreadId.make(`provider-thread:followup-command-child-${ending}`),
3332+
appThreadId: child.childThreadId,
3333+
pendingBackgroundTasks: [],
3334+
},
3335+
},
3336+
],
3337+
});
3338+
yield* orchestrator.recoverDelegatedTasks;
3339+
const stopped = (yield* orchestrator.getThreadProjection(
3340+
parentThreadId,
3341+
)).subagents.find((task) => task.id === followupTaskId);
3342+
assert.equal(stopped?.cancelledRunId, followupRun.id);
3343+
assert.equal(stopped?.status, "interrupted");
3344+
return;
3345+
}
3346+
3347+
// The finished command wakes the child; that turn continues the same follow-up.
3348+
const wakeMessageId = MessageId.make(`message:followup-command-wake-${ending}`);
3349+
yield* orchestrator.dispatch({
3350+
type: "message.dispatch",
3351+
createdBy: "agent",
3352+
creationSource: "server",
3353+
commandId: CommandId.make(`command:followup-command-wake-${ending}`),
3354+
threadId: child.childThreadId,
3355+
messageId: wakeMessageId,
3356+
text: "Background activity updated",
3357+
notification: {
3358+
source: { kind: "command" },
3359+
outcome: "completed",
3360+
summary: 'Command "sleep 90" finished',
3361+
},
3362+
attachments: [],
3363+
dispatchMode: { type: "queue_after_active" },
3364+
});
3365+
const wakeRun = (yield* orchestrator.getThreadProjection(child.childThreadId)).runs.find(
3366+
(run) => run.userMessageId === wakeMessageId,
3367+
)!;
3368+
assert.equal(wakeRun.delegatedTaskId, followupTaskId);
3369+
yield* eventSink.write({
3370+
commandId: reconcileCommandId(`followup-command-wake-ended-${ending}`),
3371+
events: [
3372+
{
3373+
id: EventId.make(`event:followup-command-settled-${ending}`),
3374+
type: "provider-thread.updated",
3375+
threadId: child.childThreadId,
3376+
driver,
3377+
providerInstanceId: modelSelection.instanceId,
3378+
occurredAt: now,
3379+
payload: {
3380+
...parent.providerThreads[0]!,
3381+
id: ProviderThreadId.make(`provider-thread:followup-command-child-${ending}`),
3382+
appThreadId: child.childThreadId,
3383+
pendingBackgroundTasks: [],
3384+
},
3385+
},
3386+
{
3387+
id: EventId.make(`event:followup-command-final-${ending}`),
3388+
type: "message.updated",
3389+
threadId: child.childThreadId,
3390+
runId: wakeRun.id,
3391+
occurredAt: now,
3392+
payload: {
3393+
id: MessageId.make(`message:followup-command-final-${ending}`),
3394+
threadId: child.childThreadId,
3395+
runId: wakeRun.id,
3396+
nodeId: null,
3397+
role: "assistant",
3398+
text: "FINAL CHECK-OK",
3399+
attachments: [],
3400+
streaming: false,
3401+
createdBy: "agent",
3402+
creationSource: "server",
3403+
createdAt: now,
3404+
updatedAt: now,
3405+
},
3406+
},
3407+
{
3408+
id: EventId.make(`event:followup-command-wake-completed-${ending}`),
3409+
type: "run.updated",
3410+
threadId: child.childThreadId,
3411+
runId: wakeRun.id,
3412+
occurredAt: now,
3413+
payload: { ...wakeRun, status: "completed", startedAt: now, completedAt: now },
3414+
},
3415+
],
3416+
});
3417+
yield* orchestrator.recoverDelegatedTasks;
3418+
3419+
const finished = yield* orchestrator.getThreadProjection(parentThreadId);
3420+
const childTasks = finished.subagents.filter(
3421+
(task) => task.childThreadId === child.childThreadId,
3422+
);
3423+
assert.deepEqual(
3424+
childTasks.map((task) => task.id).toSorted(),
3425+
[child.taskId, followupTaskId].toSorted(),
3426+
);
3427+
const followup = childTasks.find((task) => task.id === followupTaskId);
3428+
assert.equal(followup?.status, "completed");
3429+
assert.equal(followup?.result, "FINAL CHECK-OK");
3430+
}),
3431+
);
3432+
31933433
it.effect("waits for a run the follow-up's child queued on its own thread", () =>
31943434
Effect.gen(function* () {
31953435
const requests = yield* ProviderContinuationRequests.ProviderContinuationRequests;

‎apps/server/src/orchestration-v2/Orchestrator.ts‎

Lines changed: 49 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -128,7 +128,7 @@ import {
128128
cancelledBeforeStart,
129129
continuationFailedBeforePrompt,
130130
continuationStartRetryDelayMs,
131-
delegatedTaskAwaitsSubagents,
131+
delegatedTaskAwaitsWork,
132132
delegatedTaskResultRun,
133133
delegatedTaskRuns,
134134
isRestartContinuationSource,
@@ -9489,9 +9489,21 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio
94899489
(row) => row.id === message.runId,
94909490
);
94919491
if (run === undefined || run.delegatedTaskId !== undefined) return;
9492+
// A wake continues the latest follow-up while it is open, as it continues the original task.
9493+
const latestFollowupId =
9494+
command.notification === undefined
9495+
? undefined
9496+
: (yield* projectionStore
9497+
.getThreadRecords(command.threadId, ["runs"])
9498+
.pipe(mapDispatchError(command))).runs
9499+
.filter((row) => row.id !== run.id && row.delegatedTaskId !== undefined)
9500+
.toSorted((a, b) => b.ordinal - a.ordinal)[0]?.delegatedTaskId;
94929501
const parent = yield* projectionStore
94939502
.getThreadRecords(parentThreadId, ["runs", "subagents", "providerTurns", "attempts"], {
9494-
subagentIds: child.forkedFrom?.type === "node" ? [child.forkedFrom.nodeId] : [],
9503+
subagentIds: [
9504+
...(child.forkedFrom?.type === "node" ? [child.forkedFrom.nodeId] : []),
9505+
...(latestFollowupId === undefined ? [] : [latestFollowupId]),
9506+
],
94959507
})
94969508
.pipe(mapDispatchError(command));
94979509
const original = parent.subagents.find(
@@ -9501,6 +9513,27 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio
95019513
task.id === (child.forkedFrom?.type === "node" ? child.forkedFrom.nodeId : undefined),
95029514
);
95039515
if (original === undefined || original.result === null) return;
9516+
const openFollowup = parent.subagents.find(
9517+
(task) =>
9518+
task.id === latestFollowupId &&
9519+
task.origin === "app_owned" &&
9520+
task.childThreadId === child.id &&
9521+
task.cancelledRunId === undefined &&
9522+
task.result === null,
9523+
);
9524+
if (openFollowup !== undefined) {
9525+
yield* emit(
9526+
events,
9527+
command,
9528+
)({
9529+
type: "run.updated",
9530+
threadId: child.id,
9531+
runId: run.id,
9532+
occurredAt: message.createdAt,
9533+
payload: { ...run, delegatedTaskId: openFollowup.id },
9534+
});
9535+
return;
9536+
}
95049537
const owner = followUpOwnerRun(parent.runs);
95059538
if (owner === null || owner.rootNodeId === null) return;
95069539
const now = message.createdAt;
@@ -14115,8 +14148,17 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio
1411514148
]);
1411614149
return;
1411714150
}
14118-
// A follow-up ends with the subagents its turns started, as the original task does.
14119-
if (delegatedTaskAwaitsSubagents(taskRuns, childControls.subagents)) return;
14151+
// A follow-up ends with the work it awaits, as the original task does. A cancelled task
14152+
// does not wait for the background work its cancel is stopping.
14153+
if (
14154+
delegatedTaskAwaitsWork(
14155+
taskRuns,
14156+
options.task.cancelledRunId === undefined
14157+
? childControls
14158+
: { subagents: childControls.subagents, providerThreads: [] },
14159+
)
14160+
)
14161+
return;
1412014162
if (
1412114163
yield* childAwaitsRestartContinuation(
1412214164
taskRuns.length === 0 ? childControls.runs : taskRuns,
@@ -15691,7 +15733,7 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio
1569115733
if (
1569215734
run === undefined ||
1569315735
delegatedTaskTerminalStatus(run.status) === null ||
15694-
delegatedTaskAwaitsSubagents(runs, child.subagents)
15736+
delegatedTaskAwaitsWork(runs, child)
1569515737
)
1569615738
return true;
1569715739
return yield* childAwaitsRestartContinuation(runs, run);
@@ -15786,7 +15828,8 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio
1578615828
? child.thread.forkedFrom?.type === "node" &&
1578715829
child.thread.forkedFrom.nodeId === input.taskId &&
1578815830
progressAfterTurn(child).state !== "result_available"
15789-
: delegatedTaskTerminalStatus(runOutcome(followUpRun)) === null;
15831+
: delegatedTaskTerminalStatus(runOutcome(followUpRun)) === null ||
15832+
delegatedTaskAwaitsWork(delegatedTaskRuns(child, input.taskId), child);
1579015833
yield* threadDispatch.withLock(
1579115834
command.threadId,
1579215835
Effect.gen(function* () {

0 commit comments

Comments
 (0)