fix(authbridge): Stream text/event-stream responses frame-by-frame#480
Conversation
Fixes #477. The forward proxy fully buffered outbound response bodies via io.ReadAll whenever any plugin in the outbound pipeline declared ReadsBody -- which today is every realistic pipeline (mcp-parser, inference-parser, ibac all set ReadsBody=true). For MCP servers using Streamable HTTP, tools/call responses arrive as text/event-stream that may stay open for the full duration of slow tool execution; buffering converted those streams into blocking Client.Do reads bounded by the hard-coded 30s http.Client.Timeout, producing 502s or apparent hangs. The reverse proxy had the same buffering for inbound A2A message/stream responses. This stream-aware path detects text/event-stream at the response header, then forwards SSE frames to the client and dispatches each JSON-RPC message to plugins as it arrives. application/json keeps the buffered path. Per-frame memory is capped at 1 MiB; total stream length is unbounded. - Add an optional pipeline.StreamingResponder interface with OnResponseFrame(ctx, pctx, frame, last). Plumbed via pipeline.Pipeline.RunResponseFrame and pipeline.Holder. - Drop the outbound http.Client.Timeout (which covered the body read and broke streaming); replace with Transport.ResponseHeaderTimeout (30s, time-to-headers ceiling) and a per-read idle deadline on streaming bodies. - forwardproxy: branch on Content-Type per response. SSE streams with HasStreamingResponders take a flush-per-frame path that invokes OnResponseFrame for each event and a final last=true call; RunResponse is not called on this path. application/json and any other Content-Type keep the existing buffered path, plus a single last=true OnResponseFrame call so streaming-aware plugins use one code path for both shapes. WritesBody=true forces buffered fallback with a warning (a body mutator can't rewrite a body already on the wire). - reverseproxy: mirror the same branch in modifyResponse via a streaming response body that pulls one SSE frame per Read, dispatches it to OnResponseFrame, and re-emits the SSE event; ReverseProxy.FlushInterval=-1 ferries each frame to the client immediately. Same WritesBody fallback contract. - Convert the aggregating parsers to fold-and-finalize: - mcp-parser: per-message recording (one JSON-RPC result per frame); OnResponse stays as a fallback for non-streaming-aware listeners. - inference-parser: accumulates content deltas + usage on a private SetState scratch across frames; finalizes on last=true. - a2a-parser: folds artifact text and final-status across message/stream events; finalizes on last=true. New extractStreamEvent helper shared by the buffered and streaming paths so one envelope shape is parsed in one place. - New authlib/listener/internal/sseframe package: a narrow SSE reader emitting one ([]byte, error) per data-bearing event. Per-frame size cap (default 1 MiB), CRLF tolerance, multi-line-data folding, comment lines silently skipped, trailing unterminated event delivered before EOF. - New tests cover: the SSE reader (multi-frame, CRLF, oversize, long lines, only-comments, multi-line data); RunResponseFrame dispatch ordering, off-policy skip, reject-stops-chain, and cancel; per-plugin OnResponseFrame for all three parsers (per-message MCP, fold-and-finalize inference + A2A, application/json one-shot, no-extension no-op, empty-stream skip pairing); end-to-end streaming through the forward and reverse proxies (frames arrive before upstream closes -- the regression check; WritesBody buffered fallback; application/json single last=true frame). The OnResponseFrame hook leaves room for per-message response-side enforcement later. Today's plugins are observability-only on the response side (IBAC's OnResponse is a no-op; it gates at request time), so the listener forwards-then-dispatches each frame; an inspect-before-forward variant would be a clean follow-up. Assisted-By: Claude (Anthropic AI) <noreply@anthropic.com> Signed-off-by: Kelly Abuelsaad <kaymar@gmail.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughAdds SSE frame parsing, a StreamingResponder pipeline hook and RunResponseFrame dispatch, streaming-aware forward/reverse proxy paths, and streaming adapters for A2A, Inference, and MCP parsers with tests. ChangesStreaming Response Support
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 golangci-lint (2.12.2)level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies" Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (5)
authbridge/authlib/listener/forwardproxy/server.go (2)
778-787: 💤 Low valueMinor race between timer firing and
timer.Stop().If
Readreturns just as the timer fires (narrow window), the body could be closed after a successful read, causing the nextReadto fail unexpectedly. The 5-minute timeout makes this extremely unlikely, and the failure mode is a clean error rather than data corruption, so this is low risk. If you want to eliminate it, async.Oncearound the close would ensure idempotency at the cost of slight complexity.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@authbridge/authlib/listener/forwardproxy/server.go` around lines 778 - 787, idleReadCloser.Read has a tiny race where the timer callback can call i.rc.Close() just after a successful Read, causing the next Read to fail; wrap the close in a sync.Once to make the timer callback idempotent so Close() is only executed once. Add a sync.Once field (e.g. closeOnce) to the idleReadCloser struct and change the timer callback to call closeOnce.Do(func(){ _ = i.rc.Close() }), leaving the rest of Read (timer creation, stop, return) unchanged.
791-802: 💤 Low valueFunction name is misleading.
indexByteASCIICaseInsensitiveperforms an exact byte match (as noted in the comment), not a case-insensitive search. Consider renaming toindexByteorindexASCIIByteto avoid confusion, since the "CaseInsensitive" suffix suggests behavior that doesn't apply to the byte parameter.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@authbridge/authlib/listener/forwardproxy/server.go` around lines 791 - 802, The function indexByteASCIICaseInsensitive is misnamed because it performs an exact byte match; rename it to a clear name like indexByte or indexASCIIByte (choose one) and update all call sites to use the new name (e.g., replace indexByteASCIICaseInsensitive(...) with indexByte(...)); also update the function comment to reflect exact ASCII byte lookup semantics and run the build/tests to ensure no references remain to the old name.authbridge/authlib/plugins/a2aparser/plugin.go (1)
331-337: 💤 Low valueUnbounded artifact concatenation could accumulate large strings.
For very long streaming responses with many artifact-update events,
ext.Artifact += part.Textaccumulates all text parts into one string without bounds. While the SSE frame reader has per-frame limits (1 MiB default per PR summary), a stream could have many frames. Consider whether a truncation or size cap is appropriate for observability purposes.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@authbridge/authlib/plugins/a2aparser/plugin.go` around lines 331 - 337, The ext.Artifact concatenation in the artifact handling branch can grow unbounded: in the loop over event.Result.Artifact.Parts (checking part.Kind == "text") the code uses ext.Artifact += part.Text; introduce a fixed cap (e.g., const maxArtifactSize) and track ext.Artifact length before appending so you only append up to the remaining allowed bytes, and if you hit the cap append a sentinel like "…(truncated)" or stop further concatenation; update the loop that processes event.Result.Artifact.Parts to enforce this cap and ensure subsequent artifact-update events respect the same limit.authbridge/authlib/listener/reverseproxy/server.go (1)
648-686: 💤 Low valueConsider stdlib string functions for clarity.
strings.IndexByte,strings.TrimSpace, andstrings.EqualFoldwould reduce code and improve readability. The ASCII-only behavior here matches stdlib for the expected inputs (HTTP header values).That said, these custom implementations are correct and the duplication with forwardproxy is acceptable to keep the packages decoupled.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@authbridge/authlib/listener/reverseproxy/server.go` around lines 648 - 686, Replace the custom implementations of indexByteASCII, trimASCIISpace, and equalASCIIFold with calls to the stdlib string helpers: use strings.IndexByte for indexByteASCII, strings.TrimSpace (or strings.Trim(s, " \t") if you want to strictly preserve only space and tab) for trimASCIISpace, and strings.EqualFold for equalASCIIFold; also add the "strings" import and update the three function bodies to delegate to these stdlib functions (keeping the same function names so callers in this file stay unchanged).authbridge/authlib/pipeline/plugin.go (1)
194-202: ⚡ Quick winUpdate the
StreamingRespondercontract comment to match current proxy ordering.The comment says current listeners forward+flush before
OnResponseFrame, but reverse proxy flow invokesRunResponseFramebefore emitting the current frame bytes. Please align this contract text to avoid misleading plugin implementers about mid-stream reject behavior.📝 Suggested doc tweak
-// - A non-Continue Action returned mid-stream stops further -// dispatch for that frame and rejects the response. The current -// forwardproxy/reverseproxy implementations forward+flush before -// calling OnResponseFrame for record-only observability, so a -// mid-stream Reject from a plugin would not un-send already-sent -// bytes; today no plugin returns Reject from this hook. The hook -// leaves the door open for per-message enforcement to be added -// later (the listener would have to inspect-before-forward at -// that point). +// - A non-Continue Action returned mid-stream stops further +// dispatch for that frame and rejects the response. Depending on +// listener implementation, previously forwarded bytes may already +// be on the wire, so a mid-stream Reject can result in a truncated +// stream. Today no in-tree plugin returns Reject from this hook.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@authbridge/authlib/pipeline/plugin.go` around lines 194 - 202, Update the StreamingResponder contract comment to reflect the actual proxy ordering: state that reverse proxy flow invokes RunResponseFrame (and thus emits frame bytes) before calling OnResponseFrame, so a mid-stream non-Continue Action cannot un-send already-emitted bytes; remove or correct the sentence that claims listeners forward+flush before OnResponseFrame. Edit the comment block near the StreamingResponder type and the OnResponseFrame / RunResponseFrame mentions in plugin.go to explicitly describe current behavior and its implication for mid-stream Reject handling.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@authbridge/authlib/listener/forwardproxy/server.go`:
- Around line 530-536: Capture the return Action from the finalization call to
OutboundPipeline.RunResponseFrame and handle Reject the same way the per-frame
loop does: assign the result (e.g., action :=
s.OutboundPipeline.RunResponseFrame(r.Context(), pctx, nil, true)), then if
action == pipeline.ActionReject (or the package's Reject constant) log the
rejection via s.recordOutboundResponseEvent or the existing logger so the policy
violation is recorded even though headers are already sent; keep the call to
s.recordOutboundResponseEvent(pctx, resp.StatusCode) afterward as before.
- Around line 563-565: In streamFallbackBuffered, do not ignore the return value
from s.OutboundPipeline.RunResponseFrame; call RunResponseFrame(r.Context(),
pctx, respBody, true) and inspect its error/response result (same way as the
buffered path that checks lines 359-364) and if it indicates a rejection or
error (e.g., instructing a 502/403) stop further handling and return that result
instead of proceeding to write headers/body; update the logic around
s.OutboundPipeline.HasStreamingResponders and RunResponseFrame so rejections are
honored before headers are written.
In `@authbridge/authlib/listener/internal/sseframe/reader.go`:
- Around line 153-188: readLine currently only treats '\n' as a terminator and
allows unbounded accumulation for long non-`data:` lines; update Reader.readLine
to also accept a lone '\r' as a line terminator (handle chunks that end with
'\r' returned by r.br.ReadSlice and treat them like terminated lines, stripping
the CR) and enforce a cap on accumulated non-`data:` lines to prevent unbounded
growth (introduce a maxLineBytes constant and when appending to line check if
the line does not start with "data:" and its length exceeds maxLineBytes, return
a clear error). Use the existing r.br.ReadSlice and bufio.ErrBufferFull loop but
add the CR-termination branch and the non-data length check to safely bound
memory.
---
Nitpick comments:
In `@authbridge/authlib/listener/forwardproxy/server.go`:
- Around line 778-787: idleReadCloser.Read has a tiny race where the timer
callback can call i.rc.Close() just after a successful Read, causing the next
Read to fail; wrap the close in a sync.Once to make the timer callback
idempotent so Close() is only executed once. Add a sync.Once field (e.g.
closeOnce) to the idleReadCloser struct and change the timer callback to call
closeOnce.Do(func(){ _ = i.rc.Close() }), leaving the rest of Read (timer
creation, stop, return) unchanged.
- Around line 791-802: The function indexByteASCIICaseInsensitive is misnamed
because it performs an exact byte match; rename it to a clear name like
indexByte or indexASCIIByte (choose one) and update all call sites to use the
new name (e.g., replace indexByteASCIICaseInsensitive(...) with indexByte(...));
also update the function comment to reflect exact ASCII byte lookup semantics
and run the build/tests to ensure no references remain to the old name.
In `@authbridge/authlib/listener/reverseproxy/server.go`:
- Around line 648-686: Replace the custom implementations of indexByteASCII,
trimASCIISpace, and equalASCIIFold with calls to the stdlib string helpers: use
strings.IndexByte for indexByteASCII, strings.TrimSpace (or strings.Trim(s, "
\t") if you want to strictly preserve only space and tab) for trimASCIISpace,
and strings.EqualFold for equalASCIIFold; also add the "strings" import and
update the three function bodies to delegate to these stdlib functions (keeping
the same function names so callers in this file stay unchanged).
In `@authbridge/authlib/pipeline/plugin.go`:
- Around line 194-202: Update the StreamingResponder contract comment to reflect
the actual proxy ordering: state that reverse proxy flow invokes
RunResponseFrame (and thus emits frame bytes) before calling OnResponseFrame, so
a mid-stream non-Continue Action cannot un-send already-emitted bytes; remove or
correct the sentence that claims listeners forward+flush before OnResponseFrame.
Edit the comment block near the StreamingResponder type and the OnResponseFrame
/ RunResponseFrame mentions in plugin.go to explicitly describe current behavior
and its implication for mid-stream Reject handling.
In `@authbridge/authlib/plugins/a2aparser/plugin.go`:
- Around line 331-337: The ext.Artifact concatenation in the artifact handling
branch can grow unbounded: in the loop over event.Result.Artifact.Parts
(checking part.Kind == "text") the code uses ext.Artifact += part.Text;
introduce a fixed cap (e.g., const maxArtifactSize) and track ext.Artifact
length before appending so you only append up to the remaining allowed bytes,
and if you hit the cap append a sentinel like "…(truncated)" or stop further
concatenation; update the loop that processes event.Result.Artifact.Parts to
enforce this cap and ensure subsequent artifact-update events respect the same
limit.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 5923a09c-defe-4843-9c2c-81ad62f311bc
📒 Files selected for processing (16)
authbridge/authlib/listener/forwardproxy/server.goauthbridge/authlib/listener/forwardproxy/streaming_test.goauthbridge/authlib/listener/internal/sseframe/reader.goauthbridge/authlib/listener/internal/sseframe/reader_test.goauthbridge/authlib/listener/reverseproxy/server.goauthbridge/authlib/listener/reverseproxy/streaming_test.goauthbridge/authlib/pipeline/holder.goauthbridge/authlib/pipeline/pipeline.goauthbridge/authlib/pipeline/plugin.goauthbridge/authlib/pipeline/streaming_test.goauthbridge/authlib/plugins/a2aparser/plugin.goauthbridge/authlib/plugins/a2aparser/streaming_test.goauthbridge/authlib/plugins/inferenceparser/plugin.goauthbridge/authlib/plugins/inferenceparser/streaming_test.goauthbridge/authlib/plugins/mcpparser/plugin.goauthbridge/authlib/plugins/mcpparser/streaming_test.go
|
One concern on the timeout design — I think the hard-coded The breakStreamable HTTP lets the server pick Why a stream-declared timeout can't save usNo layer here lets the sender declare "expect my response in N seconds": HTTP has no such header, SSE's Proposal — remove a clock, don't add configRather than tuning or configuring the header timeout, collapse to a single proxy-imposed timeout:
Result: connect timeout (unchanged) + one idle timeout (now used on both paths) + caller context. No new config; a constant removed, not added. "What if the server never terminates the stream?"Every realistic broken-server case goes silent → the idle timer catches it (hang after headers, hang mid-stream, dropped conn). The only case it misses is a server looping forever emitting valid frames/heartbeats — and that's bounded by the agent's own request deadline, which the proxy already carries. The proxy shouldn't out-guess the caller on total runtime. If we ever want a hard proxy-side stop even when the agent sets no deadline, that's one in-code constant ( Net: delete the header timeout, reuse the idle timer on both paths, lean on the caller's context — simpler than the current two-clock setup and it closes the slow-JSON regression. |
huang195
left a comment
huang195
left a comment
There was a problem hiding this comment.
Review summary
The streaming design is sound — the true text/event-stream path is single-dispatch and correct, per-request state lives on pctx (no shared-instance race across concurrent requests), nil-safety and empty-stream pairing are handled, and RunResponseFrame mirrors RunResponse's skip/gating/cancel semantics. Nicely done.
One root-cause must-fix blocks it: the buffered application/json finalize runs both RunResponse and RunResponseFrame(last=true) on the same body, so migrated plugins record every JSON response twice. Details inline (#1/#2). The remaining items are suggestions/nits.
Areas reviewed: Go (forward/reverse proxy streaming, SSE reader, pipeline plumbing, mcp/inference/a2a parsers). Commits: 1, signed-off ✓. CI: all green.
Fix for #1/#2/#3 is one change: have RunResponse skip plugins implementing StreamingResponder (the buffered path then finalizes them via the single RunResponseFrame(last=true) call you already make), or conversely only deliver the buffered finalize through one hook. That makes the documented "framework picks one path" contract actually true.
| // Plugins that didn't migrate — i.e. don't implement | ||
| // StreamingResponder — are unaffected (RunResponseFrame skips them). | ||
| if s.OutboundPipeline.HasStreamingResponders() && resp.Body != nil { | ||
| respFrameAction := s.OutboundPipeline.RunResponseFrame(r.Context(), pctx, pctx.ResponseBody, true) |
There was a problem hiding this comment.
must-fix — buffered application/json double-dispatches. This runs after RunResponse (line 348), so for a plugin implementing StreamingResponder both hooks fire on the same buffered body. Traced it: OnResponse runs first, its "already-finalized" guard is false on a fresh response, so it does the full parse + Observe; then OnResponseFrame(last=true) re-parses the same bytes and Observes again → two identical matched_<method>_response rows per JSON response (confirmed in mcp-parser and inference-parser). The guards in each OnResponse check fields that OnResponseFrame populates, but since OnResponse runs first they never fire — the guard is in the wrong order relative to the actual call order. The doc on StreamingResponder says "OnResponse is then NOT called for plugins that implement this interface," but nothing in the framework enforces that on the buffered path.
Fix: skip StreamingResponder plugins inside RunResponse (then this single RunResponseFrame(last=true) finalizes them), so JSON and SSE both go through exactly one hook.
There was a problem hiding this comment.
Fixed in 2724eda. Pipeline.RunResponse now skips plugins implementing StreamingResponder so the buffered application/json path delivers the body through OnResponseFrame(last=true) only — single dispatch on both shapes. Added a regression test in TestForwardProxy_BufferedDeliversLastTrueFrame that asserts OnResponse is not called for a StreamingResponder plugin.
| // for buffered application/json we deliver the body as one | ||
| // last=true frame so plugins can finalize via OnResponseFrame. | ||
| if s.InboundPipeline.HasStreamingResponders() && resp.Body != nil { | ||
| frameAction := s.InboundPipeline.RunResponseFrame(resp.Request.Context(), pctx, pctx.ResponseBody, true) |
There was a problem hiding this comment.
must-fix — same buffered double-dispatch inbound. RunResponse (330) + RunResponseFrame(last=true) (339) both run on the buffered application/json body, so A2A message/send (non-stream) responses through this listener double-record the same way as the forward proxy (see the forwardproxy:360 comment). Same one-line fix resolves both listeners.
There was a problem hiding this comment.
Fixed in 2724eda by the same Pipeline.RunResponse change — reverseproxy.modifyResponse no longer double-dispatches buffered application/json bodies. Regression test added in TestReverseProxy_BufferedDeliversLastTrueFrame.
| } | ||
|
|
||
| if last { | ||
| ext.Completion = state.completion.String() |
There was a problem hiding this comment.
suggestion — fallback can clobber a correctly-parsed completion to empty. On the non-flusher fallback (streamFallbackBuffered) with Stream==true: OnResponse runs parseInferenceSSE and sets Completion correctly; then OnResponseFrame(whole-SSE-blob, last=true) takes the streaming branch, json.Unmarshal of the multi-line blob as one chunk fails, and this line then overwrites ext.Completion with the empty state.completion. Production ResponseWriters implement http.Flusher so this is mostly a test-surface hazard, but it's the same root cause as the buffered double-dispatch — fixing RunResponse to skip StreamingResponder plugins removes it too.
There was a problem hiding this comment.
Fixed in 2724eda. Two parts: (1) Pipeline.RunResponse skipping StreamingResponder plugins removes the OnResponse pre-population path entirely; (2) streamFallbackBuffered now re-parses the buffered SSE body frame-by-frame through sseframe.NewReader and dispatches each event as a non-last frame, so the streaming branch sees properly-shaped chunks instead of one giant blob that would fail json.Unmarshal.
| return | ||
| } | ||
| if _, err := w.Write(frame); err != nil { | ||
| slog.Debug("forward-proxy: streaming write error", "host", r.Host, "error", err) |
There was a problem hiding this comment.
suggestion — client-write error skips finalize + record. These w.Write error branches return directly, bypassing the bottom RunResponseFrame(nil, true) and recordOutboundResponseEvent. So when the downstream client disconnects mid-stream, aggregating plugins (inference/a2a) never get last=true to finalize, and no response SessionEvent is recorded for the request. (Upstream read errors take the break path and finalize correctly — only the client-write path returns early.) Consider a defer-based finalize or routing these through the same end-of-stream tail.
There was a problem hiding this comment.
Fixed in 2724eda. handleStreamingResponse now uses a defer for the final RunResponseFrame(nil, true) and recordOutboundResponseEvent so every exit path — normal EOF, upstream read error, and downstream client-write error — finalizes plugins and records the SessionResponse event. Client-write errors break out of the loop instead of returning early.
| // Closing the body unblocks Read with an error. Idempotent on | ||
| // http.Response.Body (subsequent Closes are no-ops). | ||
| _ = i.rc.Close() | ||
| }) |
There was a problem hiding this comment.
suggestion — idle-timeout close races with a healthy Read. time.AfterFunc + i.rc.Close(), with timer.Stop() after the Read: Stop() doesn't wait for an already-fired callback, so if the timer fires just as a Read returns successfully, the body is closed under a live stream and the next ReadFrame aborts it. Window is small per read, but the timer is re-armed on every Read across a multi-minute stream, so the cumulative chance of a spurious truncation is nonzero. Also: closing the body concurrently with an in-flight Read is relied-upon-but-undocumented for HTTP/2 response bodies. Consider one long-lived timer reset between reads (and a guard flag) rather than per-Read allocation, or at least document the close-to-unblock assumption.
There was a problem hiding this comment.
Fixed in 2724eda. idleReadCloser now wraps the close in a sync.Once and Close() routes through the same idempotent path, so a late-firing timer after a successful Read is harmless: the underlying body is closed at most once, and a successful in-flight Read keeps its bytes either way. Added a doc note about the close-to-unblock assumption (which is what the stdlib http transport relies on for forced disconnects).
| // complete JSON-RPC response message. We parse + record per message | ||
| // rather than waiting for end-of-stream, so a long-running tools/call | ||
| // surfaces partial results in the session timeline as they arrive. | ||
| func (p *MCPParser) OnResponseFrame(_ context.Context, pctx *pipeline.Context, frame []byte, last bool) pipeline.Action { |
There was a problem hiding this comment.
suggestion — per-frame Observe is unbounded and asymmetric. OnResponseFrame records one invocation per result frame, so a long tools/call that streams many envelopes appends an unbounded number of matched_<method>_response rows onto pctx.Extensions.Invocations for a single request — both noisy in the session timeline and a memory-growth point. inference/a2a fold and emit once at last=true; mcp is the odd one out. If per-message rows are intended, fine — but worth confirming, and maybe capping.
There was a problem hiding this comment.
Fixed in 2724eda. Capped at 50 per-frame Observe rows per stream; further frames update pctx.Extensions.MCP (so the latest result is still observable off-stream) but no Invocation row is appended. A single matched__response_truncated row is emitted at last=true so operators see the truncation explicitly. Per-message rows are still the design — abctl benefits from partial visibility on long tools/call streams — but the cap removes the unboundedness.
| @@ -70,6 +96,10 @@ func NewServer(outbound *pipeline.Holder, sessions *session.Store, mtls *MTLSOpt | |||
| IdleConnTimeout: 90 * time.Second, | |||
There was a problem hiding this comment.
suggestion — hard-coded ResponseHeaderTimeout=30s re-creates #477 for application/json. Streamable HTTP lets the server return JSON instead of SSE per call; for JSON the server holds headers until the tool finishes, so a slow tool 502s at 30s — the same failure #477 is fixing, on the other response shape. (Raised in a standalone comment already; folding into the review for completeness.) Suggestion there: drop this and let the inbound request context + the existing idle timer bound things, so no new config is added.
There was a problem hiding this comment.
Fixed in 2724eda. Dropped ResponseHeaderTimeout entirely. Streamable HTTP servers can hold headers open on either response shape, so any fixed time-to-headers ceiling reproduces #477 on the JSON path. The inbound request context and the per-Read idle timer (streamReadIdleTimeout, 5min) remain as bounds; an unrecoverably wedged upstream is closed when the client cancels.
| if _, err := w.Write([]byte("data: ")); err != nil { | ||
| slog.Debug("forward-proxy: streaming write error", "host", r.Host, "error", err) | ||
| return | ||
| } |
There was a problem hiding this comment.
nit — multi-line data: not round-tripped faithfully. sseframe folds multi-line data: with \n; re-emitting as one data: <frame>\n\n embeds those \ns, which a downstream SSE parser reads as line boundaries. Harmless for the single-line JSON-RPC payloads this targets (as the comment notes), but it'd corrupt genuine multi-line SSE. A guard or a re-split-per-line on emit would make it shape-faithful.
There was a problem hiding this comment.
Fixed in 2724eda. Both listeners now re-split the frame on \n at emit time so each original data line gets its own data: prefix, with one trailing blank line as the event terminator. For single-line JSON-RPC payloads this is equivalent to the old data: \n\n; for genuine multi-line data the wire shape is now faithful to what the upstream produced.
Addresses review items from PR #480: - Fix double-dispatch on buffered application/json: Pipeline.RunResponse now skips plugins implementing StreamingResponder so the same body is never delivered through both OnResponse and OnResponseFrame. The framework picks one path per the StreamingResponder contract. - Fix inference clobber on streamFallbackBuffered: re-parse the buffered SSE body frame-by-frame (matches the real streaming dispatch shape) instead of delivering the whole blob as one frame, which would fail json.Unmarshal in the streaming branch and zero out a parsed completion. - Defer-based finalize + record on every exit path of handleStreamingResponse so a client disconnect mid-stream still finalizes aggregating plugins and records the SessionResponse event. - Capture and log the final RunResponseFrame Reject (headers already on the wire, but the policy violation now surfaces). - streamFallbackBuffered: honor RunResponseFrame Reject (headers not yet sent, so a 502/403 is still possible). - Drop hard-coded 30s ResponseHeaderTimeout. Streamable HTTP servers hold headers open until tool completion even for application/json, so a fixed time-to-headers ceiling reproduces #477 on the JSON shape. Request context + per-Read idle timer remain the bounds. - idleReadCloser: sync.Once around close so a stray late timer cannot fire after a successful Read; document the close-to-unblock contract. - SSE reader: accept lone CR as a line terminator and bound non-data line growth at maxSize so pathological comment / unknown-field lines cannot grow memory past the per-frame cap. - mcp-parser: cap per-frame Observe at 50 to bound noisy session timelines on long tools/call streams; emit a single _truncated row past the cap so operators see truncation explicitly. - a2a-parser: cap accumulated artifact text at 64 KiB with a single truncation marker, mirroring the per-stream bounds elsewhere. - Fix multi-line data round-trip: re-split frame on LF at emit time so each original line gets its own data prefix instead of embedding LFs inside one event. - Replace custom indexByteASCII / trimASCIISpace / equalASCIIFold with stdlib strings.IndexByte / TrimSpace / EqualFold in both listeners. - Update StreamingResponder contract comment to reflect the actual proxy ordering (reverseproxy emits frame bytes before OnResponseFrame; forwardproxy emits after) and the new RunResponse skip behavior. Regression tests: assert OnResponse is NOT called on the buffered application/json path for a StreamingResponder plugin in both listeners. All existing tests continue to pass. Assisted-By: Claude (Anthropic AI) <noreply@anthropic.com> Signed-off-by: Kelly Abuelsaad <kaymar@gmail.com>
|
Pushed 2724eda addressing all review feedback. Must-fix (huang195):
Suggestions (huang195):
CodeRabbit:
All tests pass; Assisted-By: Claude (Anthropic AI) noreply@anthropic.com |
huang195
left a comment
huang195
left a comment
There was a problem hiding this comment.
Re-review (commit 2724eda)
Thanks for the thorough follow-up — all eight points from the previous round are addressed, and well:
- #1/#2 double-dispatch (must-fix) →
RunResponsenow skipsStreamingResponderplugins, so the bufferedapplication/jsonpath single-dispatches. ✔ - #3 inference clobber →
streamFallbackBufferedre-parses the buffered SSE body frame-by-frame instead of dumping the whole blob. ✔ - #4 client-write finalize → moved to a
defercovering every exit path; the old explicit tail calls are removed, so it's still exactly one finalize. ✔ - #5 idle race →
sync.Once(closeIdempotent) makes close safe; residual deadline-boundary race documented. ✔ - #6 mcp per-frame Observe →
maxStreamObserves=50cap +_truncatedmarker. ✔ - #7 30s
ResponseHeaderTimeout→ removed; relies on the idle timer + inbound request context. ✔ - #8 multi-line
data:→writeSSEFramere-splits per original line. ✔
One new must-fix that the #1 fix introduced — details inline on pipeline.go. It's a consequence of the skip landing framework-wide while only the proxy listeners dispatch frames; the envoy-sidecar (extproc) path now drops streaming-parser response handling entirely. Once that's resolved this is good to go.
CI: green.
| if policy == ErrorPolicyOff { | ||
| continue | ||
| } | ||
| if _, ok := p.plugins[i].(StreamingResponder); ok { |
There was a problem hiding this comment.
must-fix — this skip regresses the envoy-sidecar (extproc) path. Skipping StreamingResponder plugins here is correct for forwardproxy/reverseproxy because they finalize those plugins via RunResponseFrame. But extproc calls only RunResponse (extproc/server.go:557 and :594) and never RunResponseFrame — so for the mcp/inference/a2a parsers (all of which now implement StreamingResponder), their response-phase logic runs nowhere under envoy-sidecar: skipped here, never reached as frames.
Two concrete regressions in envoy mode:
- Lost response observability — MCP results, inference completion/usage, and A2A artifacts/final-status are no longer recorded on the response side.
- Broken A2A contextId rekey —
extproc/server.go:~599rekeys the session fromDefaultSessionIDto the server-assigned contextId, which depends on a2a-parser populatingpctx.Extensions.A2A.SessionIDduringRunResponse. With a2a-parser now skipped, that field isn't set on the response and inbound A2A sessions fragment.
Envoy-sidecar is a supported, tested deployment path, so this needs to be closed before merge. Two options:
- (a) Have extproc call
RunResponseFrame(pctx.ResponseBody, last=true)afterRunResponseon its buffered body (envoy delivers the response BUFFERED already — seehandleResponseHeadersModeOverride), mirroring what the proxy listeners do forapplication/json. For an SSE body buffered by envoy, re-parse viasseframelikestreamFallbackBuffereddoes. - (b) Make the skip listener-aware (e.g. a flag/variant so
RunResponseonly skips when the caller will also dispatch frames), so extproc's plainRunResponsekeeps runningOnResponsefor these plugins.
(a) is the more consistent fix — it gives extproc the same single-dispatch contract as the proxies.
There was a problem hiding this comment.
Fixed in 2fd8c0f via option (a). After RunResponse, extproc now dispatches the buffered body through RunResponseFrame whenever the pipeline has any StreamingResponder — application/json as one last=true frame, text/event-stream re-parsed via sseframe (one non-last frame per event + a final last=true). Header-only path also gets a single empty last=true so plugins finalize. The A2A contextId rekey now sees pctx.Extensions.A2A.SessionID populated, restoring single-session behavior. Regression tests added: TestExtProc_BufferedJSONResponse_DispatchesToStreamingResponder and TestExtProc_BufferedSSEResponse_DispatchesPerEvent.
…in extproc Fixes a regression introduced in 2724eda: pipeline.RunResponse now skips StreamingResponder plugins so the proxy listeners single-dispatch on the buffered application/json path, but extproc only calls RunResponse — never RunResponseFrame — so under envoy-sidecar mode mcp/inference/a2a parsers lost response-phase dispatch entirely. That broke MCP / inference / A2A response observability and the inbound A2A contextId rekey (which depends on a2a-parser populating pctx.Extensions.A2A.SessionID during response handling). extproc now mirrors the proxy listeners' single-dispatch contract: after RunResponse, when any plugin in the pipeline implements StreamingResponder, dispatch the buffered body via RunResponseFrame. For application/json the entire body is one last=true frame; for text/event-stream we re-parse with sseframe so each event arrives as its own non-last frame followed by a final last=true — same shape the proxy wire-streaming path produces. The header-only response path also delivers a single empty last=true frame so plugins finalize. Regression tests: - TestExtProc_BufferedJSONResponse_DispatchesToStreamingResponder: buffered JSON delivers exactly one OnResponseFrame(last=true) and OnResponse is not called for the StreamingResponder plugin. - TestExtProc_BufferedSSEResponse_DispatchesPerEvent: a 3-event SSE body produces 3 non-last frames + 1 final last=true. Assisted-By: Claude (Anthropic AI) <noreply@anthropic.com> Signed-off-by: Kelly Abuelsaad <kaymar@gmail.com>
huang195
left a comment
huang195
left a comment
There was a problem hiding this comment.
Re-review (commit 2fd8c0f) — LGTM
The envoy-sidecar regression from the previous round is correctly fixed. Verified:
handleResponseBodycallsdispatchBufferedFramesafterRunResponseand before the A2A contextId rekey, so a2a-parser'sOnResponseFramepopulatespctx.Extensions.A2A.SessionIDin time — the rekey works again.- Single dispatch is preserved:
RunResponsestill skipsStreamingResponderplugins, anddispatchBufferedFramesis the only place they're dispatched (no return of the original double-dispatch). - Both shapes covered:
application/json→ onelast=trueframe;text/event-stream→ re-parsed per-event viasseframe+ finallast=true; header-only path → emptylast=trueso plugins finalize. - Two regression tests added (buffered JSON single
last=true+OnResponsenot called; 3-event SSE → 3 frames + final).
All eight points from the first round remain resolved. Two non-blocking notes inline. Nice work across the revisions.
CI: green.
| // sseframe so each event arrives as its own non-last frame followed | ||
| // by a final last=true — matches the per-message dispatch shape | ||
| // streaming-aware plugins expect. | ||
| func dispatchBufferedFrames(ctx context.Context, p *pipeline.Holder, pctx *pipeline.Context) pipeline.Action { |
There was a problem hiding this comment.
suggestion (non-blocking, scope note). Because envoy ext_proc delivers the response ResponseBodyMode_BUFFERED, this re-parses an already-fully-buffered SSE body. So #477's actual streaming benefit — not holding a slow/large SSE tools/call until complete — applies to proxy-sidecar only; in envoy-sidecar mode the whole stream is still buffered (bounded by maxBodySize = 1 MiB), so the original slow-tool/large-response risk persists there. This is a pre-existing extproc limitation and the right behavior for plugin dispatch — not something to fix in this PR — but worth a doc line or a tracked follow-up so operators know envoy mode isn't streamed end-to-end (it would need ext_proc STREAMED body mode).
| // isEventStream reports whether a Content-Type header value names the | ||
| // SSE media type. Tolerates parameters and ASCII case differences. | ||
| // Mirrors the helpers in forwardproxy/reverseproxy. | ||
| func isEventStream(contentType string) bool { |
There was a problem hiding this comment.
nit. isEventStream now exists in three listeners (forwardproxy, reverseproxy, extproc) with the same logic. Consider hoisting to a shared internal helper (e.g. next to sseframe) so the three copies can't drift.
github-project-automation
Bot
moved this from New /:ToDo
to Done
in Kagenti Issue Prioritization
3 participants
Fixes #477.
The forward proxy fully buffered outbound response bodies via io.ReadAll whenever any plugin in the outbound pipeline declared ReadsBody -- which today is every realistic pipeline (mcp-parser, inference-parser, ibac all set ReadsBody=true). For MCP servers using Streamable HTTP, tools/call responses arrive as text/event-stream that may stay open for the full duration of slow tool execution; buffering converted those streams into blocking Client.Do reads bounded by the hard-coded 30s http.Client.Timeout, producing 502s or apparent hangs. The reverse proxy had the same buffering for inbound A2A message/stream responses.
This stream-aware path detects text/event-stream at the response header, then forwards SSE frames to the client and dispatches each JSON-RPC message to plugins as it arrives. application/json keeps the buffered path. Per-frame memory is capped at 1 MiB; total stream length is unbounded.
The OnResponseFrame hook leaves room for per-message response-side enforcement later. Today's plugins are observability-only on the response side (IBAC's OnResponse is a no-op; it gates at request time), so the listener forwards-then-dispatches each frame; an inspect-before-forward variant would be a clean follow-up.
Assisted-By: Claude (Anthropic AI) noreply@anthropic.com
Summary
Related issue(s)
(Optional) Testing Instructions
Fixes #
Summary by CodeRabbit
New Features
Bug Fixes / Reliability
Tests