Repository navigation
test(agents): regression tests for nested sub-agent WebSockets (#2026) - #2096
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.
|
![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.
Co-authored-by: Cursor <cursoragent@cursor.com>
threepointone
changed the title
|
Thanks @AntoniTok. The same root cause was fixed on main in #2400 while this was open, after the forwarding code moved to |
| it("establishes a WebSocket through two nested facet hops", async () => { | ||
| // Regression for #2026. The root owns the physical socket and |
There was a problem hiding this comment.
🔍 Description describes fixes already merged
The current diff adds only regression tests; the routing and delivery fixes described in the PR have already landed on the base branch. Update the description to reflect this PR's remaining scope.
Was this helpful? React with 👍 or 👎 to provide feedback.
There was a problem hiding this comment.
Done. The title and description now say the fix landed in #2400 and this PR adds only the two regression tests.
|
✅ agents import sizes: no significant changes ( |
2 participants
Closes #2026
The routing fix for #2026 landed on main in #2400 (descendants no longer receive the root's outer-URL header), after the WebSocket forwarding code moved into
dynamic-agents.ts. This PR is now reduced to the regression tests from the original work, which pass on main:/sub/spike-sub-child/{middle}/sub/spike-sub-child/{leaf}upgrades, round-trips a message, and stays open.broadcast()(sync, so its caller can't await it) reaches the root-owned socket.The original
index.tschanges are dropped. Main tracks facet deliveries per stream inbridges.ts, so it doesn't need the agent-wide delivery drain, or the per-frame scoping and background-delivery test built on it. Test-only, so no changeset.