Skip to content

Commit 7ef9bf5

Browse files
smthfoxysmthfoxy
committed
fix(agents): admit the requester settle wake as tracked gateway root work [AI]
The settle wake was launched as a detached promise from cleanup bookkeeping, so registry cleanup or shutdown could reach quiescence before the wake admitted its gateway turn and the last-child completion could still be lost during restart or teardown. Route the wake through runWithGatewayIndependentRootWorkContinuation: a live cleanup parent reserves the root synchronously, and restart drain now waits for the in-flight wake. Adds a deterministic quiescence-race regression.
1 parent 19bdd5c commit 7ef9bf5

4 files changed

Lines changed: 68 additions & 13 deletions

File tree

‎src/agents/subagent-registry-lifecycle.test.ts‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ import {
66
getActiveGatewayRootWorkCount,
77
markGatewayRestartDraining,
88
resetGatewayWorkAdmission,
9+
runWithGatewayIndependentRootWorkAdmission,
10+
waitForActiveGatewayRootWork,
911
} from "../process/gateway-work-admission.js";
1012
import { SUBAGENT_KILL_TASK_ERROR } from "../tasks/detached-task-runtime-contract.js";
1113
import {
@@ -3643,5 +3645,48 @@ describe("requester settle wake trigger", () => {
36433645
expect(warn).toHaveBeenCalledWith("requester settle wake failed", expect.anything());
36443646
});
36453647
});
3648+
3649+
it("holds the settle wake as tracked root work so restart drain waits for its turn", async () => {
3650+
resetGatewayWorkAdmission();
3651+
try {
3652+
const entry = createRunEntry({ endedAt: 4_000 });
3653+
let releaseWake: (() => void) | undefined;
3654+
const settleWake = vi.fn(
3655+
() =>
3656+
new Promise<boolean>((resolve) => {
3657+
releaseWake = () => resolve(false);
3658+
}),
3659+
);
3660+
const controller = createLifecycleController({
3661+
entry,
3662+
maybeWakeRequesterAfterAllChildrenSettled: settleWake,
3663+
});
3664+
3665+
// Schedule from inside an admitted cleanup parent that finishes before
3666+
// the wake settles — the quiescence window: the wake must reserve its
3667+
// own root before the parent releases.
3668+
await runWithGatewayIndependentRootWorkAdmission(async () => {
3669+
controller.completeCleanupBookkeeping({
3670+
runId: entry.runId,
3671+
entry,
3672+
cleanup: "keep",
3673+
completedAt: 5_000,
3674+
});
3675+
});
3676+
expect(settleWake).toHaveBeenCalledTimes(1);
3677+
await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(1));
3678+
3679+
// A restart drain arriving between scheduling and the wake's gateway
3680+
// turn must wait for the wake instead of reporting quiescence.
3681+
markGatewayRestartDraining();
3682+
expect((await waitForActiveGatewayRootWork(25)).drained).toBe(false);
3683+
3684+
releaseWake?.();
3685+
await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(0));
3686+
expect((await waitForActiveGatewayRootWork(1_000)).drained).toBe(true);
3687+
} finally {
3688+
resetGatewayWorkAdmission();
3689+
}
3690+
});
36463691
});
36473692
/* oxlint-disable max-lines -- TODO: split this grandfathered oversized file. */

‎src/agents/subagent-registry-lifecycle.ts‎

Lines changed: 17 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import { formatErrorMessage, readErrorName } from "../infra/errors.js";
1212
import {
1313
isGatewayRestartDraining,
1414
runWithGatewayIndependentRootWorkAdmission,
15+
runWithGatewayIndependentRootWorkContinuation,
1516
} from "../process/gateway-work-admission.js";
1617
import { defaultRuntime } from "../runtime.js";
1718
import { emitSessionLifecycleEvent } from "../sessions/session-lifecycle-events.js";
@@ -756,10 +757,13 @@ export function createSubagentRegistryLifecycleController(params: {
756757
return true;
757758
};
758759

759-
// Fire-and-forget: once a child reaches a terminal settle, let the announce
760-
// layer decide whether its requester's batch has fully drained and, if so,
761-
// wake the registry-less top-level requester to synthesize. Failures are
762-
// logged only — settle bookkeeping must never block on the wake.
760+
// Once a child reaches a terminal settle, let the announce layer decide
761+
// whether its requester's batch has fully drained and, if so, wake the
762+
// registry-less top-level requester to synthesize. Settle bookkeeping never
763+
// blocks on the wake, but the wake must run as tracked root work: a live
764+
// cleanup parent reserves the root synchronously, so restart or suspend
765+
// cannot reach quiescence between scheduling and the wake's gateway turn.
766+
// Failures are logged only.
763767
const scheduleRequesterSettleWake = (
764768
runId: string,
765769
entry: SubagentRunRecord,
@@ -769,20 +773,20 @@ export function createSubagentRegistryLifecycleController(params: {
769773
if (!requesterSessionKey) {
770774
return;
771775
}
772-
void params
773-
.maybeWakeRequesterAfterAllChildrenSettled({
776+
void runWithGatewayIndependentRootWorkContinuation(() =>
777+
params.maybeWakeRequesterAfterAllChildrenSettled({
774778
requesterSessionKey,
775779
requesterOrigin: entry.requesterOrigin,
776780
settledEntry: entry,
777781
settledRowRetired: options?.rowRetired === true,
778-
})
779-
.catch((error: unknown) => {
780-
params.warn("requester settle wake failed", {
781-
error: buildSafeLifecycleErrorMeta(error),
782-
runId: maskRunId(runId),
783-
requesterSessionKey: maskSessionKey(requesterSessionKey),
784-
});
782+
}),
783+
).catch((error: unknown) => {
784+
params.warn("requester settle wake failed", {
785+
error: buildSafeLifecycleErrorMeta(error),
786+
runId: maskRunId(runId),
787+
requesterSessionKey: maskSessionKey(requesterSessionKey),
785788
});
789+
});
786790
};
787791

788792
const suspendPendingFinalDelivery = (args: {

‎src/agents/subagent-registry.test-helpers.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ type RegistryDeps = {
4646
resolveAgentTimeoutMs: typeof import("./timeout.js").resolveAgentTimeoutMs;
4747
restoreSubagentRunsFromDisk: typeof import("./subagent-registry-state.js").restoreSubagentRunsFromDisk;
4848
runSubagentAnnounceFlow: typeof import("./subagent-announce.js").runSubagentAnnounceFlow;
49+
maybeWakeRequesterAfterAllChildrenSettled: typeof import("./subagent-announce.requester-settle-wake.js").maybeWakeRequesterAfterAllChildrenSettled;
4950
ensureContextEnginesInitialized?: () => void;
5051
ensureRuntimePluginsLoaded?: typeof import("./runtime-plugins.js").ensureRuntimePluginsLoaded;
5152
resolveContextEngine?: typeof import("../context-engine/registry.js").resolveContextEngine;

‎src/agents/subagent-registry.test.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,7 @@ const mocks = vi.hoisted(() => ({
177177
captureSubagentCompletionReply: vi.fn(async () => "final completion reply"),
178178
cleanupBrowserSessionsForLifecycleEnd: vi.fn(async () => {}),
179179
runSubagentAnnounceFlow: vi.fn(async () => true),
180+
maybeWakeRequesterAfterAllChildrenSettled: vi.fn(async () => false),
180181
getGlobalHookRunner: vi.fn(() => null),
181182
ensureRuntimePluginsLoaded: vi.fn(),
182183
ensureContextEnginesInitialized: vi.fn(),
@@ -320,6 +321,10 @@ describe("subagent registry seam flow", () => {
320321
resolveAgentTimeoutMs: mocks.resolveAgentTimeoutMs,
321322
restoreSubagentRunsFromDisk: mocks.restoreSubagentRunsFromDisk,
322323
runSubagentAnnounceFlow: mocks.runSubagentAnnounceFlow,
324+
// Registry seam tests must not run the real settle wake: it holds
325+
// tracked root work through its own async gating, which races the
326+
// root-count drain assertions here. Wake behavior has its own suites.
327+
maybeWakeRequesterAfterAllChildrenSettled: mocks.maybeWakeRequesterAfterAllChildrenSettled,
323328
ensureContextEnginesInitialized: mocks.ensureContextEnginesInitialized,
324329
ensureRuntimePluginsLoaded: mocks.ensureRuntimePluginsLoaded,
325330
resolveContextEngine: mocks.resolveContextEngine,

0 commit comments

Comments
 (0)