fix(agents): route nested facet WebSockets without recursing - #2096
fix(agents): route nested facet WebSockets without recursing#2096AntoniTok wants to merge 2 commits into
Conversation
A sub-agent route two or more hops deep upgraded with HTTP 101 and then closed with 1011. The root's private route header was copied into every descendant's forwarded request, so a leaf read its own ancestor as one of its children and created facets recursively until workerd rejected the chain at the depth limit. One hop passed only because the single-hop self-strip happened to consume the leaf's own segment. Descendants now route from their already-stripped connection URI, and the header is dropped when forwarding. Fixing the routing exposed a second defect on the reply path: each hop's bridge is an RpcTarget whose stub is disposed when its inbound call returns, and every hop was fire-and-forget, so with two hops the inner delivery was cut off with "RPC stub used after being disposed". Deliveries are now awaited at each hop, with the sync `send`/`broadcast` contracts tracked and drained before the frame returns.
🦋 Changeset detectedLatest commit: 8e18b89 The changes in this PR will be included in the next version bump. This PR includes changesets to release 2 packages
Not sure what this means? Click here to learn what changesets are. Click here if you're a maintainer who wants to add another changeset to this PR |
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
| private async _cf_drainSubAgentConnection(id: string): Promise<void> { | ||
| const awaited = new Set<Promise<unknown>>(); | ||
| // Terminates because each pass only awaits promises it has not seen, | ||
| // and a frame can only queue finitely many. | ||
| while (true) { | ||
| const inFlight = [ | ||
| ...(this._cf_virtualSubAgentConnections.get(id)?.pending ?? []), | ||
| ...this._cf_pendingSubAgentDeliveries | ||
| ].filter((promise) => !awaited.has(promise)); | ||
| if (inFlight.length === 0) return; | ||
| for (const promise of inFlight) awaited.add(promise); | ||
| await Promise.allSettled(inFlight); | ||
| } |
There was a problem hiding this comment.
🟡 A sub-agent's message handling can stay open indefinitely when the agent keeps sending in the background
Every incoming frame now waits for all of the agent's outstanding outgoing messages, including ones belonging to other clients or to background work (_cf_drainSubAgentConnection at packages/agents/src/index.ts:7410-7422), and the wait restarts whenever a new outgoing message appears, so handling a single frame can stay open for as long as the agent keeps sending.
Impact: A sub-agent that streams or periodically pushes messages (e.g. a chat agent emitting chunks) keeps the frame's call open for the whole stream, and one slow client's delivery can hold up an unrelated client's frame.
Why the drain loop's termination assumption does not hold
The loop mixes two sources: the per-connection pending set and the agent-wide _cf_pendingSubAgentDeliveries set, which is populated by broadcast() on a facet (packages/agents/src/index.ts:7091-7093) regardless of which connection or which turn started it.
The inline comment claims termination because "a frame can only queue finitely many" deliveries. That holds only if the deliveries are all started by the frame being drained. In practice the agent-wide set is also fed by background/streaming code (e.g. AIChatAgent._broadcastChatMessage emitting chunks after onMessage returns) and by concurrent frames on other connections. Each loop pass re-reads both sets and filters out only promises it has already awaited, so a steady supply of new broadcasts keeps inFlight non-empty and the loop keeps going.
The callers are the finally blocks of _cf_handleSubAgentWebSocketConnect (packages/agents/src/index.ts:7559-7564), _cf_handleSubAgentWebSocketMessage (:7609-7615) and _cf_handleSubAgentWebSocketClose (:7627-7634), plus SubAgentConnectionBridge's #flush on every send/close/setState, so the coupling applies to each hop.
Scoping the drain to the deliveries actually started by the current frame (for example by recording the connection id alongside each tracked broadcast, and snapshotting the set at frame entry) would keep the #2026 fix while bounding the wait.
Prompt for agents
_cf_drainSubAgentConnection in packages/agents/src/index.ts waits on two sources: the per-connection `pending` set of the virtual sub-agent connection, and the agent-wide `_cf_pendingSubAgentDeliveries` set fed by Agent.broadcast() on a facet. The loop re-reads both sets each pass and only skips promises it has already awaited, so it terminates only if no new deliveries keep appearing. The comment asserts a frame can only queue finitely many deliveries, but the agent-wide set is also fed by background work (streaming chat chunks broadcast after onMessage returns) and by concurrent frames on other connections. As a result the finally-block drains in _cf_handleSubAgentWebSocketConnect/Message/Close, and the #flush awaited inside SubAgentConnectionBridge.send/close/setState, can stay pending for the whole duration of an unrelated stream, keeping the inbound RPC (and the lent bridge stub) open far longer than intended and coupling unrelated connections. Consider scoping tracked deliveries to the frame/connection that started them — e.g. record the connection id (or a per-frame token) when tracking a broadcast, and have the drain only await deliveries belonging to the frame being completed, or snapshot the pending set at frame entry rather than re-reading it in a loop.
Was this helpful? React with 👍 or 👎 to provide feedback.
There was a problem hiding this comment.
Good catch — this was a real bug and the termination comment was wrong. Fixed in 8e18b89.
You were right on both counts. _cf_pendingSubAgentDeliveries was agent-wide, and the loop re-read it every pass, so any steady source of broadcasts kept inFlight non-empty. My "a frame can only queue finitely many" claim only held for the per-connection set, not the shared one.
Two changes:
Deliveries from a sync API are now collected per frame. _cf_runSubAgentFrame installs a fresh sink around each connect/message/close handler and drains only that. The shared set is gone, so unrelated work cannot be in scope by construction.
A broadcast started outside a frame is no longer tracked at all. That turned out to be the more useful observation: those travel via _rootAlarmOwner() on a stub the agent owns, not on a borrowed bridge, so they were never exposed to the disposal race this PR is about. Waiting on them was pure cost. This covers your AIChatAgent._broadcastChatMessage case — chunks emitted after onMessage returns take the owned-stub path.
The drain snapshots once instead of polling. send and broadcast register synchronously, so everything a frame started is already present when it returns, and a completing delivery never queues another. Removes the loop entirely.
Also added a regression test. Worth noting how it observes the problem, because my first attempt was wrong: asserting that the echo still arrives passes either way, since the message is delivered before the drain. It is frame completion that stalls. So the test stalls a background delivery, closes the socket, and asserts the facet drops the connection — _cf_handleSubAgentWebSocketClose only deletes the entry once its frame finishes. Verified it hangs and fails against the agent-wide version, and passes with the fix.
pnpm run test:workers 90 files / 1804 tests, pnpm run check clean across 121 projects.
agents
@cloudflare/ai-chat
@cloudflare/codemode
hono-agents
@cloudflare/shell
@cloudflare/think
@cloudflare/voice
@cloudflare/worker-bundler
commit: |
Review feedback on cloudflare#2096: deliveries were tracked in one agent-wide set, and the drain re-read it each pass. Background broadcasts and concurrent frames on other connections therefore kept it non-empty, so a frame could stay open for the length of an unrelated stream and one connection could stall another. The termination comment was wrong for the same reason. Deliveries started from a sync API are now collected per frame, and the drain snapshots once instead of polling. A broadcast started outside a frame is left untracked: it travels on a stub the agent owns rather than a borrowed bridge, so it was never exposed to the disposal race. Adds a regression test that stalls a background delivery and asserts the close frame still completes, observed through the facet dropping the connection. It hangs against the agent-wide version.
1 participant
Closes #2026
The problem
A WebSocket two or more hops deep (
/sub/{class}/{name}/sub/{class}/{name}) upgraded with HTTP 101 and then immediately died with code1011. One hop was fine.Only the root agent owns the real socket. It remembers the whole route in a private header,
x-cf-agents-subagent-url, and strips one/sub/{class}/{name}hop each time it passes the connection down. But it also copied every header down — including that private one.So the leaf ended up holding a route written from the root's point of view. It read the first hop in it, saw
middle, and concludedmiddlemust be one of its own children.middleis its grandparent. It created a new facet, which read the same note, and so on until workerd stopped it:One hop only passed by accident: the single-hop self-strip in
_cf_resolveSubAgentConnectionhappened to consume the leaf's own segment, leaving nothing to route.What changed
Descendants no longer receive the header, and no longer read it even if they somehow get one:
A facet's connection URI is already stripped one hop per level, so routing from it is correct at any depth. The invariant — only the root may hold this route — is now written down at the declaration.
_cf_getForwardedSubAgentStatealready stripped the same key from forwarded state; the header channel was the remaining hole.The second defect
Fixing the routing exposed a bug underneath it. The nested socket now reached the leaf, but the reply never came back:
Each hop hands the next one a
SubAgentConnectionBridge, anRpcTargetwhose stub is disposed the moment its inbound call returns. Every delivery was fire-and-forget:The bridge is effectively on loan for the duration of the call. Returning without waiting means hanging up while your reply is still travelling:
Step 7 is only the first place this is visible, not the only place it is wrong. Every link had the same flaw: if the reply had survived step 7, Middle would then have sent it over bridge-A and returned without waiting, and Root would have disposed bridge-A out from under it. The message dies at whichever link loses the race first.
That is also why one hop passed. With no middle there is a single link and a single short trip, so the send usually landed before the bridge was reclaimed. Two hops means winning two races instead of one, and it lost consistently. Depth 1 was never correct, only lucky.
Deliveries are now awaited at each hop, so no agent finishes while it is still holding someone else's message:
Connection.sendandbroadcastare synchronous by contract, so their promises have nowhere to be returned. Those are recorded instead, and drained before the frame's RPC returns. Deliveries flow strictly rootward, so the drain cannot cycle; on the root it is a no-op, since a real socket sends natively.I kept these together because #2026's expected behaviour — the client receiving
pong:{leaf}:hello— is unreachable with only the routing fix, and the delivery bug is unobservable without it.Tests
Two additions to
src/tests/spike-sub-agent-routing.test.ts, using the existingSpikeSubParent/SpikeSubChildfixtures:mainwith the depth-limit error.broadcast(), the sync-contract path that cannot await itself. Also fails onmain.Verified red→green rather than assumed: with only the header fix applied, the first test still fails with
RPC stub used after being disposed. That is the evidence the delivery fix is load-bearing rather than speculative.Verification
pnpm run test:workers— 90 files / 1803 testspnpm run check— sherif, export check, oxfmt, oxlint, 121 projects typecheck@cloudflare/ai-chat50/737,@cloudflare/voice12/227,@cloudflare/think2/5nx affected -t test— 16 projectsAIChatAgentandThinkboth wraponConnect/onMessageand branch on_cf_connectionTargetsSubAgent, which readsconnection.uri. That semantics is unchanged, and their suites pass.Not addressed
_cf_broadcastToSubAgentstill falls through tosuper.getConnections()when a facet has no ambient bridge, which is the Cannot perform I/O on behalf of a different Durable Object. #1677 cross-DO hazard. Reachable, but a separate fix._cf_currentSubAgentBridgeis an instance field with LIFO save/restore and is not concurrency-safe. That is RPC reply silently dropped for a facet @callable that awaits before returning, under a burst of concurrent useAgent().call() frames #1991's territory and Fix dropped async callable replies on facets #2027 is already in that code.sub-agent-rpc-bridge.test.tsis not actionable here — that file only exists on Fix dropped async callable replies on facets #2027's branch, and the test was removed there in9a66232c. It should be restored on that branch now this has landed.