fix(agents): stream forwarded request bodies into sub-agents - #2051
fix(agents): stream forwarded request bodies into sub-agents#2051AntoniTok wants to merge 1 commit into
Conversation
`Agent._cf_forwardToFacet` and `routeSubAgentRequest` both materialised
the whole forwarded body with `await req.arrayBuffer()` before handing
off to the child facet. The read was unbounded and ran in the parent
Durable Object's isolate, ahead of any application-level intake limit —
`Agent.fetch` returns before `onRequest` whenever the path matches
`/sub/{class}/{name}`, so an app could not bound it itself. Nesting
compounded it: every `/sub/` hop re-materialised the same bytes.
Pass `req.body` through as a stream at both call sites instead.
This restores the behaviour from before cloudflare#1443, which switched to an
explicit `RequestInit` so it could set a header on WebSocket upgrades
and re-attached the body with `arrayBuffer()` as a side effect. That
fix's `Upgrade` handling is unchanged.
Measured with `wrangler dev --local` against a handler that never reads
the body, peak RSS across the `workerd` processes for one POST:
body facet before facet after canonical (control)
16 MB +75 MB +4 MB +2 MB
64 MB +268 MB +4 MB +2 MB
128 MB +546 MB +4 MB +2 MB
Adds `sub-agent-body-forwarding.test.ts`, covering both call sites: the
child is dispatched and can read a prefix of a still-open body, and a
2 MB body arrives byte-exact over one hop and two nested hops. The
streaming assertions deadlock against the previous implementation.
Backpressure now reaches the client: a child that returns without
reading the body cancels the remainder of the upload, where the parent
previously drained it in full.
Closes cloudflare#2015
🦋 Changeset detectedLatest commit: a592038 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}
agents
@cloudflare/ai-chat
@cloudflare/codemode
create-think
hono-agents
@cloudflare/shell
@cloudflare/think
@cloudflare/voice
@cloudflare/worker-bundler
commit: |
Could you explain a little more why 1443 needed to rebuild the request in this way, and why reverting to the previous behaviour does not cause the original WebSocket bug not to come back? |
|
Do you think it would be appropriate to add the memory consumption benchmark as a regression test in CI somehow/ |
In terms of behaviour this goes back to how it worked before #1443, meaning bodies stream again. But #1443's code stays the same. The I only tried the literal pre-#1443 form ( #1443 rebuilt the request because It had to set Incoming request headers are immutable in workerd, so setting one means copying the headers and building a fresh The body was re-attached by hand and The tell is So the WS fix stays intact. Upgrades never reach the changed line. |
It could be worthwhile but it would have to be nightly and non-gating and it's not really what protects this. I don't think workers can measure their own memory, so this can't be a And also the streaming tests already catch this regression on their own. "Child reads a prefix of a still-open body" can't pass against a buffering parent, because buffering holds the child back until the upload finishes. On |
cjol
left a comment
cjol
left a comment
There was a problem hiding this comment.
Okay, the key information I was looking for is that the body buffering wasn't actually part of the fix for 1443. That's fine then. Regarding CI, I still think it might be valuable to monitor memory consumption under various scenarios, because the streaming tests cover this specific regression but a more general lens might catch other similar failures. But that doesn't need to block this PR.
2 participants
Closes #2015
The problem
When a request goes to a sub-agent (a URL with
/sub/...in it), the parent used to read the entire upload into memory before passing it on. No limit. A 128 MB upload made the parent balloon by ~546 MB. And your app couldn't stop it, because this happened before your code ever ran.What changed
Two lines. Instead of reading the whole body first, the parent now just hands the body through as a stream — it passes the pipe along instead of filling a bucket and carrying it.
if (req.body && req.method !== "GET" && req.method !== "HEAD") { - forwardedInit.body = await req.arrayBuffer(); + forwardedInit.body = req.body; }Both call sites had it:
Agent._cf_forwardToFacet(src/index.ts) androuteSubAgentRequest(src/sub-routing.ts).Why it broke in the first place
#1443 fixed a WebSocket bug back in May and needed to set a header on upgrade requests. Doing that meant rebuilding the request from an explicit
RequestInit, which doesn't inherit a body — so the body got re-attached by hand, the slow way, as a side effect. It streamed fine before that.I verified this rather than assuming it. Reverting to the pre-#1443 form (
new Request(rewritten, req)) makes all the new streaming tests pass, which confirms where the buffering came from. So this is a genuine regression, shipped inagents@0.12.1. In terms of behaviour this restores the streaming, but #1443's code stays the same. TheRequestInitandUpgradeheader handling are untouched and only the body attachment changes.Worth noting why it mattered beyond raw memory:
Agent.fetchreturns beforeonRequestwhenever the path matches/sub/{class}/{name}, so the unbounded read sat in front of application-level validation. An app that carefully bounded request bodies inonRequeststill had an unbounded read ahead of it and no way to bound it. The cost was also per hop — a nested/sub/.../sub/...address re-materialised the same bytes at every level.Result
Measured with
wrangler dev --localagainst the reporter's repro (a handler that never reads the body, so growth is attributable to the forwarder). Two tarballs built from the same commit, differing only in this fix. Fresh worker per row, peak RSS summed across theworkerdprocesses for a single POST:Flat now, no matter how big the upload. The before column reproduces the reporter's +71/+284/+556 MB; the after column matches the canonical control.
The one caveat
The upload is now connected end to end. If a sub-agent replies without reading the body, the rest of the upload gets cut off instead of being quietly swallowed. That's the correct behaviour and what you'd want — but it is a change, so anyone relying on the old "swallow it all" habit would notice.
Tests
New file
src/tests/sub-agent-body-forwarding.test.ts(9 tests), covering both call sites:routeSubAgentRequest.The suite includes a canonical (non-facet) control that runs first. If the test harness itself couldn't stream a request body, every facet assertion would hang for unrelated reasons, so the control makes a failure interpretable. Fixtures are
BodyProbeSubAgent(facet) andBodyProbeRootAgent(root control), sharing one handler.The three streaming assertions fail against
mainwithexpected 'pending' to be 'settled'after 2.5s, and pass in ~25ms with this change.Verification
pnpm run test:workers— 91 files / 1810 testsagentssuite — 131 files / 2552 tests@cloudflare/ai-chat50/737,@cloudflare/voice12/225,nx affected -t test27/27pnpm run check— sherif, export check, oxfmt, oxlint, typecheck (120 projects)The existing WebSocket sub-agent tests matter here specifically, since this touches the function #1443 changed — they pass.
Not addressed
If an
onBeforeSubAgenthook consumes the body and returnsvoid, forwarding still throws (the stream is disturbed). That was already true witharrayBuffer(), just with a different message — orthogonal to this issue, so I left it alone.