From 194c450dd98fcb941b6ba60fffb6538c6a8feea1 Mon Sep 17 00:00:00 2001 From: Kelly Abuelsaad Date: Thu, 4 Jun 2026 11:19:38 -0400 Subject: [PATCH 1/3] fix(authbridge): Stream text/event-stream responses frame-by-frame 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) Signed-off-by: Kelly Abuelsaad --- .../authlib/listener/forwardproxy/server.go | 359 ++++++++++++++++-- .../listener/forwardproxy/streaming_test.go | 321 ++++++++++++++++ .../listener/internal/sseframe/reader.go | 224 +++++++++++ .../listener/internal/sseframe/reader_test.go | 165 ++++++++ .../authlib/listener/reverseproxy/server.go | 258 +++++++++++++ .../listener/reverseproxy/streaming_test.go | 216 +++++++++++ authbridge/authlib/pipeline/holder.go | 18 + authbridge/authlib/pipeline/pipeline.go | 64 ++++ authbridge/authlib/pipeline/plugin.go | 44 +++ authbridge/authlib/pipeline/streaming_test.go | 173 +++++++++ .../authlib/plugins/a2aparser/plugin.go | 196 ++++++---- .../plugins/a2aparser/streaming_test.go | 113 ++++++ .../authlib/plugins/inferenceparser/plugin.go | 131 ++++++- .../plugins/inferenceparser/streaming_test.go | 105 +++++ .../authlib/plugins/mcpparser/plugin.go | 93 ++++- .../plugins/mcpparser/streaming_test.go | 114 ++++++ 16 files changed, 2481 insertions(+), 113 deletions(-) create mode 100644 authbridge/authlib/listener/forwardproxy/streaming_test.go create mode 100644 authbridge/authlib/listener/internal/sseframe/reader.go create mode 100644 authbridge/authlib/listener/internal/sseframe/reader_test.go create mode 100644 authbridge/authlib/listener/reverseproxy/streaming_test.go create mode 100644 authbridge/authlib/pipeline/streaming_test.go create mode 100644 authbridge/authlib/plugins/a2aparser/streaming_test.go create mode 100644 authbridge/authlib/plugins/inferenceparser/streaming_test.go create mode 100644 authbridge/authlib/plugins/mcpparser/streaming_test.go diff --git a/authbridge/authlib/listener/forwardproxy/server.go b/authbridge/authlib/listener/forwardproxy/server.go index 0933cd72b..e08b44e3d 100644 --- a/authbridge/authlib/listener/forwardproxy/server.go +++ b/authbridge/authlib/listener/forwardproxy/server.go @@ -16,6 +16,7 @@ import ( "github.com/kagenti/kagenti-extensions/authbridge/authlib/auth" "github.com/kagenti/kagenti-extensions/authbridge/authlib/listener/httpx" + "github.com/kagenti/kagenti-extensions/authbridge/authlib/listener/internal/sseframe" "github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline" "github.com/kagenti/kagenti-extensions/authbridge/authlib/session" "github.com/kagenti/kagenti-extensions/authbridge/authlib/spiffe" @@ -24,6 +25,31 @@ import ( const maxBodySize = 1 << 20 // 1MB — matches Envoy's default per_stream_buffer_limit_bytes +// responseHeaderTimeout bounds the time-to-first-byte (headers) for an +// upstream response. Replaces the old http.Client.Timeout, which +// covered the entire request lifecycle including the body read and so +// killed slow streaming responses (text/event-stream MCP tools/call, +// A2A message/stream, OpenAI streaming chat completions). With the +// timeout moved to the transport's ResponseHeaderTimeout, slow tools +// can stream as long as they like — only a wedged upstream that fails +// to send any headers within this window aborts. +// +// 30s preserves the prior behavior's time-to-headers ceiling. Tools +// that took longer to *return headers* never worked under the old +// configuration either. +const responseHeaderTimeout = 30 * time.Second + +// streamReadIdleTimeout caps how long the proxy waits for the next +// byte off a streaming response body. The time.Duration is applied +// per ReadFrame iteration (see streamingResponseBody). A wedged +// upstream that goes silent for longer than this aborts the stream, +// rather than hanging the agent indefinitely. Long enough to permit +// slow tool work between SSE heartbeats; short enough to surface a +// dead connection within a few minutes. Tools that need longer idle +// gaps should emit SSE heartbeats — it's what comment lines in SSE +// are for. +const streamReadIdleTimeout = 5 * time.Minute + // Server is an HTTP forward proxy that performs token exchange on outbound requests. // // OutboundPipeline is a holder so the bound pipeline can be hot-swapped @@ -70,6 +96,10 @@ func NewServer(outbound *pipeline.Holder, sessions *session.Store, mtls *MTLSOpt IdleConnTimeout: 90 * time.Second, TLSHandshakeTimeout: 10 * time.Second, ExpectContinueTimeout: 1 * time.Second, + // Time-to-headers ceiling for the upstream response. Replaces + // the old http.Client.Timeout (which covered the body read + // and broke streaming). See responseHeaderTimeout. + ResponseHeaderTimeout: responseHeaderTimeout, } if mtls != nil { @@ -87,7 +117,11 @@ func NewServer(outbound *pipeline.Holder, sessions *session.Store, mtls *MTLSOpt OutboundPipeline: outbound, Sessions: sessions, Client: &http.Client{ - Timeout: 30 * time.Second, + // No Client.Timeout — Go applies that to the entire + // request lifecycle including body read, which kills + // streaming responses. Time-to-headers is enforced via + // transport.ResponseHeaderTimeout above; per-read idle + // behavior on streaming bodies is in streamingResponseBody. Transport: transport, CheckRedirect: func(_ *http.Request, _ []*http.Request) error { return http.ErrUseLastResponse @@ -276,6 +310,25 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) { pctx.StatusCode = resp.StatusCode pctx.ResponseHeaders = resp.Header.Clone() + // Branch on Content-Type per response. The Streamable HTTP transport + // lets the server pick application/json vs text/event-stream per + // response (the client Accepts both), so the same tool may return + // JSON on one call and SSE on the next. Decide here rather than + // negotiating, and don't take the streaming path when a plugin + // declares WritesBody (mutating a body we've already started + // forwarding is incompatible with streaming) — fall back to + // buffered with a warning log instead. + if isEventStream(resp.Header.Get("Content-Type")) && + s.OutboundPipeline.HasStreamingResponders() && + resp.Body != nil { + if s.OutboundPipeline.WritesBody() { + slog.Warn("forward-proxy: text/event-stream response with WritesBody plugin — falling back to buffered path", "host", r.Host) + } else { + s.handleStreamingResponse(w, r, resp, pctx) + return + } + } + if s.OutboundPipeline.NeedsBody() && resp.Body != nil { respBody, err := io.ReadAll(io.LimitReader(resp.Body, maxBodySize+1)) if err != nil { @@ -298,6 +351,19 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) { return } + // Streaming-aware plugins use a single code path for both shapes: + // for the buffered application/json case we deliver the whole body + // as one last=true frame so plugins finalize their running state. + // 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) + if respFrameAction.Type == pipeline.Reject { + httpx.WriteRejection(w, respFrameAction) + return + } + } + // A plugin that called pctx.SetResponseBody flipped the mutation flag. // Use the replaced bytes and rewrite Content-Length so the downstream // client gets a consistent response. Content-Encoding is cleared @@ -310,33 +376,194 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) { resp.Header.Del("Content-Encoding") } - if s.Sessions != nil { - sid := s.Sessions.ActiveSession() - if sid == "" { - sid = session.DefaultSessionID + s.recordOutboundResponseEvent(pctx, resp.StatusCode) + + for key, values := range resp.Header { + for _, value := range values { + w.Header().Add(key, value) } - plugins := pipeline.SnapshotPlugins(pctx.Extensions.Custom) - ev := pipeline.SessionEvent{ - At: time.Now(), - Direction: pipeline.Outbound, - Phase: pipeline.SessionResponse, - MCP: pipeline.SnapshotMCP(pctx.Extensions.MCP), - Inference: pipeline.SnapshotInference(pctx.Extensions.Inference), - Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseResponse), - Plugins: plugins, - Identity: pipeline.SnapshotIdentity(pctx), - Host: pctx.Host, - StatusCode: resp.StatusCode, - Error: pipeline.DeriveError(pctx), - Duration: pipeline.DurationSince(pctx.StartedAt), + } + w.WriteHeader(resp.StatusCode) + if _, err := io.Copy(w, resp.Body); err != nil { + slog.Debug("response copy error", "host", r.Host, "error", err) + } +} + +// recordOutboundResponseEvent emits the SessionResponse event for a +// completed outbound response. Extracted from handleRequest so the +// streaming path can call it once at end-of-stream and the buffered +// path can call it once after RunResponse — both go through the same +// gate and snapshotting logic. +func (s *Server) recordOutboundResponseEvent(pctx *pipeline.Context, statusCode int) { + if s.Sessions == nil { + return + } + sid := s.Sessions.ActiveSession() + if sid == "" { + sid = session.DefaultSessionID + } + plugins := pipeline.SnapshotPlugins(pctx.Extensions.Custom) + ev := pipeline.SessionEvent{ + At: time.Now(), + Direction: pipeline.Outbound, + Phase: pipeline.SessionResponse, + MCP: pipeline.SnapshotMCP(pctx.Extensions.MCP), + Inference: pipeline.SnapshotInference(pctx.Extensions.Inference), + Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseResponse), + Plugins: plugins, + Identity: pipeline.SnapshotIdentity(pctx), + Host: pctx.Host, + StatusCode: statusCode, + Error: pipeline.DeriveError(pctx), + Duration: pipeline.DurationSince(pctx.StartedAt), + } + // Same widened gate as the request side — see the request-phase + // comment for why each clause matters. + if ev.MCP != nil || ev.Inference != nil || ev.Invocations != nil || plugins != nil { + s.Sessions.Append(sid, ev) + } +} + +// isEventStream reports whether a Content-Type header value names the +// SSE media type. Content-Type may carry parameters (charset=, boundary=, +// etc.) so we match on the bare type/subtype prefix and tolerate any +// suffix. Case-insensitive per RFC 9110 §8.3.1. +func isEventStream(contentType string) bool { + if contentType == "" { + return false + } + // Strip parameters: "text/event-stream; charset=utf-8" → "text/event-stream". + if idx := indexByteASCIICaseInsensitive(contentType, ';'); idx >= 0 { + contentType = contentType[:idx] + } + contentType = trimASCIISpace(contentType) + return equalASCIIFold(contentType, "text/event-stream") +} + +// handleStreamingResponse forwards a text/event-stream response to the +// downstream client frame-by-frame. Each parsed SSE event is delivered +// to the pipeline's StreamingResponder plugins (recording-only today) +// and then written + flushed to the client immediately. End-of-stream +// is signaled to plugins with one final last=true call so aggregating +// plugins (inference-parser, a2a-parser) can finalize their running +// state. RunResponse is intentionally NOT invoked on this path — +// streaming-aware plugins move their finalization logic into +// OnResponseFrame(last=true), and legacy non-migrated plugins are +// not called on streaming responses (cleaner contract; no fragmented +// double-dispatch). +func (s *Server) handleStreamingResponse(w http.ResponseWriter, r *http.Request, resp *http.Response, pctx *pipeline.Context) { + flusher, ok := w.(http.Flusher) + if !ok { + // No flusher means the downstream connection can't deliver + // bytes incrementally — fall back to buffered. http.Flusher + // is supported by net/http's default ResponseWriter, so this + // is a defensive guard for exotic wrappers (httptest with a + // custom recorder, embedded servers). + slog.Warn("forward-proxy: ResponseWriter does not support flushing — falling back to buffered for streaming response", "host", r.Host) + s.streamFallbackBuffered(w, r, resp, pctx) + return + } + + // Forward headers and the streaming status code BEFORE the first + // frame is written. Strip Content-Length since we'll be writing + // chunked, and clear hop-by-hop headers as net/http would. + for key, values := range resp.Header { + for _, value := range values { + w.Header().Add(key, value) } - // Same widened gate as the request side — see the request-phase - // comment for why each clause matters. - if ev.MCP != nil || ev.Inference != nil || ev.Invocations != nil || plugins != nil { - s.Sessions.Append(sid, ev) + } + w.Header().Del("Content-Length") + w.WriteHeader(resp.StatusCode) + flusher.Flush() + + reader := sseframe.NewReader(idleReader(resp.Body, streamReadIdleTimeout), maxBodySize) + bytesWritten := 0 + for { + frame, err := reader.ReadFrame() + if err == io.EOF { + break + } + if err != nil { + // Read error or oversized single frame. The client has + // already received some frames; the cleanest signal is to + // close the connection and log. We can't promote this to + // 502 — headers are sent. + slog.Warn("forward-proxy: streaming response read error", "host", r.Host, "error", err, "bytesWritten", bytesWritten) + break + } + + // Record-only dispatch: invoke plugins then write+flush. + // A future enforcement-aware version can inspect-before-forward; + // see StreamingResponder doc. + respAction := s.OutboundPipeline.RunResponseFrame(r.Context(), pctx, frame, false) + if respAction.Type == pipeline.Reject { + // Headers + earlier frames already on the wire — log and + // stop forwarding. The downstream client sees a truncated + // stream, which is the best we can do without inspect- + // before-forward semantics. + slog.Warn("forward-proxy: streaming response rejected mid-stream by plugin", + "host", r.Host, "violation", respAction.Violation) + break + } + + // Write the frame back as a single SSE event. Reconstructing + // "data: ...\n\n" preserves the SSE wire shape regardless of + // whether the upstream used "data: " or "data:" or split across + // multiple data lines. Multi-line data is folded onto one line + // here, which is fine for JSON-RPC payloads (no embedded LF). + if _, err := w.Write([]byte("data: ")); err != nil { + slog.Debug("forward-proxy: streaming write error", "host", r.Host, "error", err) + return } + if _, err := w.Write(frame); err != nil { + slog.Debug("forward-proxy: streaming write error", "host", r.Host, "error", err) + return + } + if _, err := w.Write([]byte("\n\n")); err != nil { + slog.Debug("forward-proxy: streaming write error", "host", r.Host, "error", err) + return + } + flusher.Flush() + bytesWritten += len(frame) + } + + // Final last=true call so aggregating plugins (inference-parser, + // a2a-parser) finalize. Always called exactly once even on a zero- + // frame stream (per StreamingResponder contract). + s.OutboundPipeline.RunResponseFrame(r.Context(), pctx, nil, true) + + s.recordOutboundResponseEvent(pctx, resp.StatusCode) +} + +// streamFallbackBuffered handles the rare case of a streaming +// Content-Type response on a ResponseWriter that doesn't support +// http.Flusher — falls back to the original buffered path. The +// duplicated logic is small enough to keep here rather than refactor +// the main handler around the Flusher fork. +func (s *Server) streamFallbackBuffered(w http.ResponseWriter, r *http.Request, resp *http.Response, pctx *pipeline.Context) { + respBody, err := io.ReadAll(io.LimitReader(resp.Body, maxBodySize+1)) + if err != nil { + slog.Warn("forward-proxy: response body read error", "host", r.Host, "error", err) + http.Error(w, `{"error":"response body read error"}`, http.StatusBadGateway) + return } + if len(respBody) > maxBodySize { + slog.Warn("forward-proxy: response body too large", "host", r.Host, "len", len(respBody)) + http.Error(w, `{"error":"response body too large"}`, http.StatusBadGateway) + return + } + pctx.ResponseBody = respBody + resp.Body = io.NopCloser(bytes.NewReader(respBody)) + respAction := s.OutboundPipeline.RunResponse(r.Context(), pctx) + if respAction.Type == pipeline.Reject { + httpx.WriteRejection(w, respAction) + return + } + if s.OutboundPipeline.HasStreamingResponders() { + _ = s.OutboundPipeline.RunResponseFrame(r.Context(), pctx, respBody, true) + } + s.recordOutboundResponseEvent(pctx, resp.StatusCode) for key, values := range resp.Header { for _, value := range values { w.Header().Add(key, value) @@ -525,6 +752,92 @@ func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) { _ = upstream.Close() } +// idleReader wraps r so each Read enforces an idle deadline. The +// goroutine pattern (timer reset on every Read entry, cancelled on +// every Read exit) is portable across any io.ReadCloser, unlike +// SetReadDeadline which only applies to net.Conn — and for HTTPS +// upstreams the proxy holds the *http.Response.Body, not the +// underlying conn. On idle expiry the reader closes the body, which +// causes the in-flight Read to return an error and unblocks the +// caller. Subsequent Reads return the same close error. +// +// The wrapper does not buffer; bufio's reader inside sseframe.Reader +// continues to do that. The deadline is per-Read, not per-frame, so +// a long-running tool that emits one byte every minute (within the +// idle window) keeps the stream alive. The streamReadIdleTimeout +// constant captures the wall-clock budget. +type idleReadCloser struct { + rc io.ReadCloser + timeout time.Duration +} + +func idleReader(rc io.ReadCloser, timeout time.Duration) io.ReadCloser { + return &idleReadCloser{rc: rc, timeout: timeout} +} + +func (i *idleReadCloser) Read(p []byte) (int, error) { + timer := time.AfterFunc(i.timeout, func() { + // Closing the body unblocks Read with an error. Idempotent on + // http.Response.Body (subsequent Closes are no-ops). + _ = i.rc.Close() + }) + n, err := i.rc.Read(p) + timer.Stop() + return n, err +} + +func (i *idleReadCloser) Close() error { return i.rc.Close() } + +// indexByteASCIICaseInsensitive returns the index of the first +// occurrence of c in s (ASCII), or -1. Case-insensitive in spirit +// but c is treated as exact — the helper exists for ';' lookup in +// Content-Type headers, where the separator is unambiguous. +func indexByteASCIICaseInsensitive(s string, c byte) int { + for i := 0; i < len(s); i++ { + if s[i] == c { + return i + } + } + return -1 +} + +// trimASCIISpace strips leading and trailing ASCII whitespace. +// Avoids the strings package's locale-aware Unicode trim because +// HTTP header tokens are ASCII-only. +func trimASCIISpace(s string) string { + start := 0 + for start < len(s) && (s[start] == ' ' || s[start] == '\t') { + start++ + } + end := len(s) + for end > start && (s[end-1] == ' ' || s[end-1] == '\t') { + end-- + } + return s[start:end] +} + +// equalASCIIFold reports whether s and t are equal under ASCII case +// folding. Used for Content-Type comparison without pulling in a +// Unicode-aware comparator. +func equalASCIIFold(s, t string) bool { + if len(s) != len(t) { + return false + } + for i := 0; i < len(s); i++ { + a, b := s[i], t[i] + if a >= 'A' && a <= 'Z' { + a += 'a' - 'A' + } + if b >= 'A' && b <= 'Z' { + b += 'a' - 'A' + } + if a != b { + return false + } + } + return true +} + // enableKeepalive turns on TCP keepalive with a 30s probe interval on // the underlying *net.TCPConn, if conn unwraps to one. No-op on other // connection types (notably *tls.Conn, which doesn't apply on the diff --git a/authbridge/authlib/listener/forwardproxy/streaming_test.go b/authbridge/authlib/listener/forwardproxy/streaming_test.go new file mode 100644 index 000000000..98cfc2514 --- /dev/null +++ b/authbridge/authlib/listener/forwardproxy/streaming_test.go @@ -0,0 +1,321 @@ +package forwardproxy + +import ( + "bufio" + "bytes" + "context" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline" +) + +// streamingProbe is a minimal Plugin + StreamingResponder used to +// observe how the proxy dispatches frames. It records every frame +// it sees along with the last flag, so tests can assert on order +// and finalization. ReadsBody=true so the listener takes the +// NeedsBody path on requests. +type streamingProbe struct { + mu sync.Mutex + frames [][]byte + lasts []bool + caps pipeline.PluginCapabilities +} + +func newStreamingProbe(writesBody bool) *streamingProbe { + return &streamingProbe{ + caps: pipeline.PluginCapabilities{ + ReadsBody: true, + WritesBody: writesBody, + }, + } +} + +func (p *streamingProbe) Name() string { return "streaming-probe" } +func (p *streamingProbe) Capabilities() pipeline.PluginCapabilities { return p.caps } +func (p *streamingProbe) OnRequest(_ context.Context, _ *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} +func (p *streamingProbe) OnResponse(_ context.Context, _ *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} +func (p *streamingProbe) OnResponseFrame(_ context.Context, _ *pipeline.Context, frame []byte, last bool) pipeline.Action { + p.mu.Lock() + defer p.mu.Unlock() + cp := make([]byte, len(frame)) + copy(cp, frame) + p.frames = append(p.frames, cp) + p.lasts = append(p.lasts, last) + return pipeline.Action{Type: pipeline.Continue} +} + +func (p *streamingProbe) snapshot() ([][]byte, []bool) { + p.mu.Lock() + defer p.mu.Unlock() + frames := make([][]byte, len(p.frames)) + copy(frames, p.frames) + lasts := make([]bool, len(p.lasts)) + copy(lasts, p.lasts) + return frames, lasts +} + +// TestForwardProxy_Streaming_FramesFlowThrough asserts that an upstream +// emitting SSE frames over time has those frames flushed downstream as +// they arrive — no full-body buffering, no 30s timeout firing, and +// each frame reaches the StreamingResponder plugin with last=false. +// A final last=true call is made at end-of-stream. +func TestForwardProxy_Streaming_FramesFlowThrough(t *testing.T) { + // Upstream sends 3 SSE frames with small idle gaps. We use real + // timers (not mock clocks) because the proxy reads off net/http + // internals; the gaps are tens of milliseconds, well under any + // timeout, and far longer than the per-frame work. + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + flusher := w.(http.Flusher) + for i := 1; i <= 3; i++ { + fmt.Fprintf(w, "data: {\"id\":%d}\n\n", i) + flusher.Flush() + time.Sleep(20 * time.Millisecond) + } + })) + defer upstream.Close() + + probe := newStreamingProbe(false) + pipe, err := pipeline.New([]pipeline.Plugin{probe}) + if err != nil { + t.Fatalf("New pipeline: %v", err) + } + srv, err := NewServer(pipeline.NewHolder(pipe), nil, nil) + if err != nil { + t.Fatalf("NewServer: %v", err) + } + proxy := httptest.NewServer(srv.Handler()) + defer proxy.Close() + + proxyClient := &http.Client{ + Transport: &http.Transport{ + Proxy: http.ProxyURL(mustParseURL(proxy.URL)), + }, + } + req, _ := http.NewRequest("GET", upstream.URL+"/stream", nil) + resp, err := proxyClient.Do(req) + if err != nil { + t.Fatalf("Do: %v", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d", resp.StatusCode) + } + if got := resp.Header.Get("Content-Type"); !strings.HasPrefix(got, "text/event-stream") { + t.Errorf("Content-Type = %q, want text/event-stream", got) + } + + body, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("ReadAll: %v", err) + } + // Body should contain 3 SSE events. + if got := bytes.Count(body, []byte("data:")); got != 3 { + t.Errorf("body has %d data: lines, want 3 — body=%q", got, body) + } + + frames, lasts := probe.snapshot() + // We expect 3 frames + 1 last=true call. + if len(frames) != 4 { + t.Fatalf("plugin saw %d calls, want 4 (3 frames + 1 final) — frames=%v lasts=%v", len(frames), framesAsStrings(frames), lasts) + } + for i := 0; i < 3; i++ { + if lasts[i] { + t.Errorf("frame %d last=true, want false", i) + } + if !bytes.Contains(frames[i], []byte(fmt.Sprintf(`"id":%d`, i+1))) { + t.Errorf("frame %d = %q, missing id %d", i, frames[i], i+1) + } + } + if !lasts[3] { + t.Error("final call last=false, want true") + } +} + +// TestForwardProxy_Streaming_WritesBodyFallsBackToBuffered asserts the +// safety guard: a pipeline with a WritesBody plugin can't take the +// streaming path (the plugin can't rewrite a body we've already +// started forwarding). The proxy logs a warning and falls back to +// buffered, so the response is delivered correctly even though it +// loses the streaming property. +func TestForwardProxy_Streaming_WritesBodyFallsBackToBuffered(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + flusher := w.(http.Flusher) + fmt.Fprintf(w, "data: {\"id\":1}\n\n") + flusher.Flush() + })) + defer upstream.Close() + + probe := newStreamingProbe(true) // WritesBody=true → buffered fallback + pipe, err := pipeline.New([]pipeline.Plugin{probe}) + if err != nil { + t.Fatalf("New pipeline: %v", err) + } + srv, err := NewServer(pipeline.NewHolder(pipe), nil, nil) + if err != nil { + t.Fatalf("NewServer: %v", err) + } + proxy := httptest.NewServer(srv.Handler()) + defer proxy.Close() + + proxyClient := &http.Client{ + Transport: &http.Transport{ + Proxy: http.ProxyURL(mustParseURL(proxy.URL)), + }, + } + req, _ := http.NewRequest("GET", upstream.URL+"/stream", nil) + resp, err := proxyClient.Do(req) + if err != nil { + t.Fatalf("Do: %v", err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + if !bytes.Contains(body, []byte(`{"id":1}`)) { + t.Errorf("body did not contain expected payload: %q", body) + } + // Buffered path: streaming-aware plugins still see one last=true + // frame carrying the whole body. Sanity-check. + _, lasts := probe.snapshot() + if len(lasts) == 0 || !lasts[len(lasts)-1] { + t.Errorf("last call lasts = %v; expected final last=true on buffered fallback", lasts) + } +} + +// TestForwardProxy_BufferedDeliversLastTrueFrame asserts the +// application/json one-shot contract: streaming-aware plugins receive +// the buffered body as a single OnResponseFrame call with last=true, +// so the same code path handles streaming and non-streaming responses. +func TestForwardProxy_BufferedDeliversLastTrueFrame(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"ok":true}`)) + })) + defer upstream.Close() + + probe := newStreamingProbe(false) + pipe, err := pipeline.New([]pipeline.Plugin{probe}) + if err != nil { + t.Fatalf("New pipeline: %v", err) + } + srv, err := NewServer(pipeline.NewHolder(pipe), nil, nil) + if err != nil { + t.Fatalf("NewServer: %v", err) + } + proxy := httptest.NewServer(srv.Handler()) + defer proxy.Close() + + proxyClient := &http.Client{ + Transport: &http.Transport{ + Proxy: http.ProxyURL(mustParseURL(proxy.URL)), + }, + } + req, _ := http.NewRequest("GET", upstream.URL+"/oneshot", nil) + resp, err := proxyClient.Do(req) + if err != nil { + t.Fatalf("Do: %v", err) + } + defer resp.Body.Close() + io.Copy(io.Discard, resp.Body) + + frames, lasts := probe.snapshot() + if len(frames) != 1 || !lasts[0] { + t.Fatalf("frames=%v lasts=%v; want exactly one call with last=true", framesAsStrings(frames), lasts) + } + if !bytes.Equal(frames[0], []byte(`{"ok":true}`)) { + t.Errorf("frame[0] = %q, want JSON body", frames[0]) + } +} + +// TestForwardProxy_Streaming_HeadersAndBodyArriveBeforeUpstreamCloses +// is the regression test for #477 — the agent should see frames +// arriving as the upstream produces them, not after the upstream +// finishes. We assert that by reading the first frame from the +// proxy's response BEFORE the upstream writes the second frame, with +// the upstream gated on a channel. +func TestForwardProxy_Streaming_HeadersAndBodyArriveBeforeUpstreamCloses(t *testing.T) { + gate := make(chan struct{}) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + flusher := w.(http.Flusher) + fmt.Fprintf(w, "data: {\"id\":1}\n\n") + flusher.Flush() + <-gate // hold the upstream open until the test releases it + fmt.Fprintf(w, "data: {\"id\":2}\n\n") + flusher.Flush() + })) + defer upstream.Close() + defer close(gate) + + probe := newStreamingProbe(false) + pipe, err := pipeline.New([]pipeline.Plugin{probe}) + if err != nil { + t.Fatalf("New pipeline: %v", err) + } + srv, err := NewServer(pipeline.NewHolder(pipe), nil, nil) + if err != nil { + t.Fatalf("NewServer: %v", err) + } + proxy := httptest.NewServer(srv.Handler()) + defer proxy.Close() + + proxyClient := &http.Client{ + Transport: &http.Transport{ + Proxy: http.ProxyURL(mustParseURL(proxy.URL)), + }, + } + req, _ := http.NewRequest("GET", upstream.URL+"/stream", nil) + resp, err := proxyClient.Do(req) + if err != nil { + t.Fatalf("Do: %v", err) + } + defer resp.Body.Close() + + // Read the first SSE event — must arrive even though upstream is + // still holding open. If buffering had been left in place, this + // read would block until upstream closes (and a 30s timeout would + // likely fire — that was the bug). + br := bufio.NewReader(resp.Body) + deadline := time.Now().Add(2 * time.Second) + var firstFrame []byte + for time.Now().Before(deadline) { + line, err := br.ReadString('\n') + if err != nil { + t.Fatalf("ReadString: %v", err) + } + if bytes.HasPrefix([]byte(line), []byte("data:")) { + firstFrame = []byte(line) + break + } + } + if firstFrame == nil { + t.Fatal("first frame did not arrive within 2s — proxy likely buffered") + } + if !bytes.Contains(firstFrame, []byte(`"id":1`)) { + t.Errorf("first frame = %q, missing id 1", firstFrame) + } +} + +func framesAsStrings(frames [][]byte) []string { + out := make([]string, len(frames)) + for i, f := range frames { + out[i] = string(f) + } + return out +} diff --git a/authbridge/authlib/listener/internal/sseframe/reader.go b/authbridge/authlib/listener/internal/sseframe/reader.go new file mode 100644 index 000000000..7d44623fe --- /dev/null +++ b/authbridge/authlib/listener/internal/sseframe/reader.go @@ -0,0 +1,224 @@ +// Package sseframe reads complete Server-Sent Event frames off a +// streaming HTTP response body. The reader emits one slice per +// "data:"-bearing event, with all data lines of that event +// concatenated and delivered together — exactly the shape MCP's +// Streamable HTTP transport, A2A's message/stream, and OpenAI- +// compatible chat-completions streaming use to carry one JSON-RPC +// message (or one chunk) per event. +// +// This is a deliberately narrow subset of the SSE spec +// (https://html.spec.whatwg.org/multipage/server-sent-events.html): +// we surface only the data payload of each event. Comment lines +// (starting with ":"), event-type lines ("event:"), id, and retry +// fields are skipped. The framing this code cares about is: +// +// data: { ... json-rpc message ... }\n +// \n +// data: { ... another message ... }\n +// \n +// +// Multi-line data within a single event is folded with "\n" +// separators per the spec; an empty line ends the event. CR/LF and +// CRLF terminators are accepted. +package sseframe + +import ( + "bufio" + "bytes" + "errors" + "io" +) + +// ErrFrameTooLarge is returned by Reader.ReadFrame when a single SSE +// event's accumulated data exceeds the configured per-frame cap. +// Streaming responses bound memory to one frame at a time, so this +// signals an upstream that produced a single oversized JSON-RPC +// message — the proxy gives up rather than buffering it. +var ErrFrameTooLarge = errors.New("sseframe: frame exceeds per-frame size cap") + +// Reader scans an io.Reader for complete SSE events and returns the +// concatenated data lines of each. Construct with NewReader; call +// ReadFrame repeatedly until io.EOF. +type Reader struct { + br *bufio.Reader + maxSize int + // scratch holds the in-progress event's accumulated data lines. + // Reused across ReadFrame calls to avoid per-frame allocation in + // the common steady-state case. + scratch []byte +} + +// NewReader wraps r with the given per-frame size cap (bytes). +// A non-positive maxSize falls back to the package default +// (DefaultMaxFrameSize). +func NewReader(r io.Reader, maxSize int) *Reader { + if maxSize <= 0 { + maxSize = DefaultMaxFrameSize + } + return &Reader{ + br: bufio.NewReader(r), + maxSize: maxSize, + } +} + +// DefaultMaxFrameSize bounds a single SSE event's data payload at +// 1 MiB — same as the buffered-path cap so the streaming path +// preserves the per-message memory ceiling without bounding the +// total stream length. +const DefaultMaxFrameSize = 1 << 20 + +// ReadFrame reads from the underlying reader until a complete SSE +// event arrives (i.e. a blank line terminator after one or more +// "data:" lines), then returns the concatenated data payload. +// Empty events (no "data:" lines, or only comment/event-type lines) +// are skipped silently — the returned slice is always non-empty +// unless the stream ends. +// +// At end-of-stream the function returns (nil, io.EOF). If the stream +// ends after a partial event with data already accumulated, that +// trailing event is delivered before EOF (per the spec, an SSE +// stream may end without a final blank-line terminator). +// +// The returned slice is owned by the Reader and reused on the next +// call; callers that need to retain bytes across ReadFrame calls +// must copy. +func (r *Reader) ReadFrame() ([]byte, error) { + r.scratch = r.scratch[:0] + hasData := false + + for { + line, err := r.readLine() + if err == io.EOF { + // Stream ended. If we accumulated data without a final + // blank-line terminator, deliver it before EOF — per spec + // "if the user agent has reached the end of the file, then + // dispatch the event." + if hasData { + return r.scratch, nil + } + return nil, io.EOF + } + if err != nil { + return nil, err + } + + // Blank line — event terminator. Dispatch what we have, or + // keep scanning if this event was empty (comment-only event). + if len(line) == 0 { + if hasData { + return r.scratch, nil + } + continue + } + + // Comment line — silently skipped per spec ("If the line + // starts with a U+003A COLON character, ignore the line"). + if line[0] == ':' { + continue + } + + // Field-bearing line. Split at the first colon; everything + // after it (with one optional leading space stripped) is the + // value. Lines without a colon name a field with an empty + // value, which we don't care about either way. + field, value := splitField(line) + if field != "data" { + // We deliberately ignore "event", "id", "retry", and any + // unknown field — the consumer only needs the data payload. + continue + } + + // Append "\n" before the next data line per spec ("If the + // data buffer's last character is a U+000A LINE FEED (LF) + // character, then remove the last character from the data + // buffer." — equivalently, separator between data lines is + // LF, and trailing LF is stripped at dispatch). + if hasData { + if err := r.appendByte('\n'); err != nil { + return nil, err + } + } + hasData = true + if err := r.appendBytes(value); err != nil { + return nil, err + } + } +} + +// readLine reads one logical SSE line (terminated by LF, CR, or CRLF) +// off the underlying buffered reader. The returned slice excludes the +// terminator and is valid only until the next call. EOF without a +// final newline is reported as io.EOF only when the line is empty; +// otherwise the unterminated tail is returned. +func (r *Reader) readLine() ([]byte, error) { + var line []byte + for { + // Read up to and including the next LF. If the underlying + // reader returns an io.EOF with no data, propagate it; with + // data, treat the final unterminated chunk as a line. + chunk, err := r.br.ReadSlice('\n') + if len(chunk) > 0 { + line = append(line, chunk...) + if chunk[len(chunk)-1] == '\n' { + // Terminated. Strip CRLF or LF. + line = line[:len(line)-1] + if len(line) > 0 && line[len(line)-1] == '\r' { + line = line[:len(line)-1] + } + return line, nil + } + // ReadSlice returned a buffer-full chunk without LF; loop + // to keep accumulating into `line`. This handles SSE + // data lines longer than bufio's buffer. + if err == bufio.ErrBufferFull { + continue + } + } + if err != nil { + if err == io.EOF && len(line) > 0 { + // Strip a possible trailing CR (no LF was seen). + if line[len(line)-1] == '\r' { + line = line[:len(line)-1] + } + return line, nil + } + return nil, err + } + } +} + +// appendByte / appendBytes grow the scratch buffer while enforcing +// the per-frame cap. ErrFrameTooLarge fires the moment we'd exceed +// the cap, before any further reading — bounding memory to maxSize +// regardless of what the upstream sends. +func (r *Reader) appendByte(b byte) error { + if len(r.scratch)+1 > r.maxSize { + return ErrFrameTooLarge + } + r.scratch = append(r.scratch, b) + return nil +} + +func (r *Reader) appendBytes(b []byte) error { + if len(r.scratch)+len(b) > r.maxSize { + return ErrFrameTooLarge + } + r.scratch = append(r.scratch, b...) + return nil +} + +// splitField parses an SSE field line "name:value" or "name:". A +// single optional leading space after the colon is stripped per spec. +// Lines with no colon are treated as a field name with empty value. +func splitField(line []byte) (string, []byte) { + idx := bytes.IndexByte(line, ':') + if idx < 0 { + return string(line), nil + } + name := string(line[:idx]) + value := line[idx+1:] + if len(value) > 0 && value[0] == ' ' { + value = value[1:] + } + return name, value +} diff --git a/authbridge/authlib/listener/internal/sseframe/reader_test.go b/authbridge/authlib/listener/internal/sseframe/reader_test.go new file mode 100644 index 000000000..e593db4f1 --- /dev/null +++ b/authbridge/authlib/listener/internal/sseframe/reader_test.go @@ -0,0 +1,165 @@ +package sseframe + +import ( + "errors" + "io" + "strings" + "testing" +) + +func TestReader_SingleFrame(t *testing.T) { + r := NewReader(strings.NewReader("data: hello\n\n"), 0) + frame, err := r.ReadFrame() + if err != nil { + t.Fatalf("ReadFrame: %v", err) + } + if string(frame) != "hello" { + t.Errorf("frame = %q, want %q", frame, "hello") + } + if _, err := r.ReadFrame(); err != io.EOF { + t.Errorf("second ReadFrame err = %v, want io.EOF", err) + } +} + +func TestReader_MultipleFrames(t *testing.T) { + body := "data: one\n\ndata: two\n\ndata: three\n\n" + r := NewReader(strings.NewReader(body), 0) + want := []string{"one", "two", "three"} + for i, w := range want { + frame, err := r.ReadFrame() + if err != nil { + t.Fatalf("frame %d: %v", i, err) + } + if string(frame) != w { + t.Errorf("frame %d = %q, want %q", i, frame, w) + } + } + if _, err := r.ReadFrame(); err != io.EOF { + t.Errorf("trailing err = %v, want io.EOF", err) + } +} + +func TestReader_MultilineData(t *testing.T) { + // Two data: lines for the same event are joined with \n per spec. + body := "data: line1\ndata: line2\n\n" + r := NewReader(strings.NewReader(body), 0) + frame, err := r.ReadFrame() + if err != nil { + t.Fatalf("ReadFrame: %v", err) + } + if string(frame) != "line1\nline2" { + t.Errorf("frame = %q, want %q", frame, "line1\nline2") + } +} + +func TestReader_CommentLines(t *testing.T) { + body := ": this is a heartbeat\n: another comment\ndata: payload\n\n" + r := NewReader(strings.NewReader(body), 0) + frame, err := r.ReadFrame() + if err != nil { + t.Fatalf("ReadFrame: %v", err) + } + if string(frame) != "payload" { + t.Errorf("frame = %q, want payload", frame) + } +} + +func TestReader_EventAndIDIgnored(t *testing.T) { + body := "event: status-update\nid: 42\ndata: {\"k\":1}\n\n" + r := NewReader(strings.NewReader(body), 0) + frame, err := r.ReadFrame() + if err != nil { + t.Fatalf("ReadFrame: %v", err) + } + if string(frame) != `{"k":1}` { + t.Errorf("frame = %q, want JSON payload", frame) + } +} + +func TestReader_TrailingFrameWithoutBlankLine(t *testing.T) { + // Per spec: end-of-stream dispatches whatever was accumulated + // even without a final blank-line terminator. + body := "data: only\n" + r := NewReader(strings.NewReader(body), 0) + frame, err := r.ReadFrame() + if err != nil { + t.Fatalf("ReadFrame: %v", err) + } + if string(frame) != "only" { + t.Errorf("frame = %q, want only", frame) + } + if _, err := r.ReadFrame(); err != io.EOF { + t.Errorf("trailing err = %v, want EOF", err) + } +} + +func TestReader_EmptyStreamEOF(t *testing.T) { + r := NewReader(strings.NewReader(""), 0) + if _, err := r.ReadFrame(); err != io.EOF { + t.Errorf("err = %v, want io.EOF", err) + } +} + +func TestReader_OnlyComments(t *testing.T) { + r := NewReader(strings.NewReader(": ping\n\n: pong\n\n"), 0) + if _, err := r.ReadFrame(); err != io.EOF { + t.Errorf("err = %v, want io.EOF (no data frames)", err) + } +} + +func TestReader_CRLF(t *testing.T) { + body := "data: a\r\n\r\ndata: b\r\n\r\n" + r := NewReader(strings.NewReader(body), 0) + frame, err := r.ReadFrame() + if err != nil { + t.Fatalf("first frame: %v", err) + } + if string(frame) != "a" { + t.Errorf("first frame = %q, want a", frame) + } + frame, err = r.ReadFrame() + if err != nil { + t.Fatalf("second frame: %v", err) + } + if string(frame) != "b" { + t.Errorf("second frame = %q, want b", frame) + } +} + +func TestReader_FrameTooLarge(t *testing.T) { + // Cap at 4 bytes so even a small payload exceeds. + r := NewReader(strings.NewReader("data: too-long\n\n"), 4) + _, err := r.ReadFrame() + if !errors.Is(err, ErrFrameTooLarge) { + t.Errorf("err = %v, want ErrFrameTooLarge", err) + } +} + +func TestReader_FieldWithoutSpace(t *testing.T) { + // Per spec the leading space after ":" is optional. "data:hi" is valid. + r := NewReader(strings.NewReader("data:hi\n\n"), 0) + frame, err := r.ReadFrame() + if err != nil { + t.Fatalf("ReadFrame: %v", err) + } + if string(frame) != "hi" { + t.Errorf("frame = %q, want hi", frame) + } +} + +func TestReader_LongDataLineExceedsBufioBuffer(t *testing.T) { + // Build a single data line larger than bufio's default 4 KiB + // buffer so the loop accumulates across ReadSlice's + // ErrBufferFull returns. + const n = 8192 + payload := strings.Repeat("x", n) + body := "data: " + payload + "\n\n" + r := NewReader(strings.NewReader(body), 0) + frame, err := r.ReadFrame() + if err != nil { + t.Fatalf("ReadFrame: %v", err) + } + if string(frame) != payload { + t.Errorf("frame len = %d, want %d", len(frame), n) + } +} diff --git a/authbridge/authlib/listener/reverseproxy/server.go b/authbridge/authlib/listener/reverseproxy/server.go index 14ef534a7..9078d0546 100644 --- a/authbridge/authlib/listener/reverseproxy/server.go +++ b/authbridge/authlib/listener/reverseproxy/server.go @@ -17,6 +17,7 @@ import ( "time" "github.com/kagenti/kagenti-extensions/authbridge/authlib/listener/httpx" + "github.com/kagenti/kagenti-extensions/authbridge/authlib/listener/internal/sseframe" "github.com/kagenti/kagenti-extensions/authbridge/authlib/listener/internal/tlssniff" "github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline" "github.com/kagenti/kagenti-extensions/authbridge/authlib/session" @@ -91,6 +92,16 @@ func NewServer(inbound *pipeline.Holder, sessions *session.Store, backendURL str return nil, err } proxy := httputil.NewSingleHostReverseProxy(target) + // FlushInterval -1 makes ReverseProxy flush after every Read of + // the response body. Required for streaming text/event-stream + // responses where each frame must hit the client immediately — + // the default 0 buffers until the client connection's write + // buffer is full. ReverseProxy already auto-flushes on + // text/event-stream Content-Type but only when FlushInterval is + // non-zero, and the explicit -1 makes the streaming behavior + // uniform across content types we install via + // installStreamingResponseBody. + proxy.FlushInterval = -1 s := &Server{ InboundPipeline: inbound, Sessions: sessions, @@ -276,6 +287,33 @@ func (s *Server) modifyResponse(resp *http.Response) error { pctx.StatusCode = resp.StatusCode pctx.ResponseHeaders = resp.Header.Clone() + // Branch on Content-Type per response. Streaming-aware pipelines on + // text/event-stream responses (A2A message/stream, MCP tools/call + // result over Streamable HTTP) replace resp.Body with a streaming + // reader that pulls one frame at a time, dispatches it to the + // pipeline, and rewrites the SSE event back. RunResponse is NOT + // called on this path — streaming-aware plugins finalize via + // OnResponseFrame(last=true). + // + // WritesBody is incompatible with streaming (we can't rewrite a + // body we've already started forwarding) — fall back to buffered + // with a warning. + if isEventStream(resp.Header.Get("Content-Type")) && + s.InboundPipeline.HasStreamingResponders() && + resp.Body != nil { + if s.InboundPipeline.WritesBody() { + slog.Warn("reverse-proxy: text/event-stream response with WritesBody plugin — falling back to buffered path", "host", pctx.Host) + } else { + s.installStreamingResponseBody(resp, pctx) + // Strip Content-Length — the framing reader doesn't know + // the final length and net/http handles chunked encoding + // when Content-Length is unset. + resp.Header.Del("Content-Length") + resp.ContentLength = -1 + return nil + } + } + if s.InboundPipeline.NeedsBody() && resp.Body != nil { body, err := io.ReadAll(io.LimitReader(resp.Body, maxBodySize+1)) if err != nil { @@ -294,6 +332,16 @@ func (s *Server) modifyResponse(resp *http.Response) error { return &responseRejectedError{action: action} } + // Streaming-aware plugins use a single code path for both shapes: + // 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) + if frameAction.Type == pipeline.Reject { + return &responseRejectedError{action: frameAction} + } + } + // A plugin that called pctx.SetResponseBody flipped the mutation flag. // Use the replaced bytes and rewrite Content-Length so the downstream // client gets a consistent response. Content-Encoding is cleared — @@ -427,6 +475,216 @@ func requestScheme(r *http.Request) string { return "http" } +// installStreamingResponseBody replaces resp.Body with a streaming +// reader that pulls SSE frames off the upstream, dispatches each frame +// through the pipeline's StreamingResponder hook, and emits the +// framed bytes downstream. ReverseProxy's normal io.Copy then ferries +// those bytes to the client; FlushInterval=-1 (set in NewServer) +// flushes after each Read so frames hit the client as they arrive. +// +// On end-of-stream the reader emits a final last=true call to +// finalize aggregating plugins, then records the inbound response +// SessionEvent (the call site that buffered responses use is below +// modifyResponse, which doesn't run on the streaming path because +// modifyResponse returned early). +func (s *Server) installStreamingResponseBody(resp *http.Response, pctx *pipeline.Context) { + upstream := resp.Body + resp.Body = &streamingResponseBody{ + upstream: upstream, + reader: sseframe.NewReader(upstream, maxBodySize), + ctx: resp.Request.Context(), + pipeline: s.InboundPipeline, + pctx: pctx, + onClose: func(statusCode int) { + s.recordInboundResponseEvent(pctx, statusCode) + }, + statusCode: resp.StatusCode, + } +} + +// recordInboundResponseEvent emits the SessionResponse event for an +// inbound streaming response. Mirrors the buffered-path block at the +// bottom of modifyResponse; lives here so the streaming body's +// onClose callback can record without holding a reference to the +// status code that close arrived with. +func (s *Server) recordInboundResponseEvent(pctx *pipeline.Context, statusCode int) { + if s.Sessions == nil { + return + } + plugins := pipeline.SnapshotPlugins(pctx.Extensions.Custom) + if !(pctx.Extensions.A2A != nil || pctx.Extensions.Invocations != nil || plugins != nil) { + return + } + sid := inboundSessionID(pctx) + // Rekey default → contextId mirroring the buffered path's behavior; + // streaming A2A message/stream may discover the contextId mid-stream. + if pctx.Extensions.A2A != nil && pctx.Extensions.A2A.SessionID != "" && + pctx.Extensions.A2A.SessionID != session.DefaultSessionID { + s.Sessions.Rekey(session.DefaultSessionID, pctx.Extensions.A2A.SessionID) + } + s.Sessions.Append(sid, pipeline.SessionEvent{ + At: time.Now(), + Direction: pipeline.Inbound, + Phase: pipeline.SessionResponse, + A2A: pipeline.SnapshotA2A(pctx.Extensions.A2A), + Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseResponse), + Plugins: plugins, + Identity: pipeline.SnapshotIdentity(pctx), + Host: pctx.Host, + StatusCode: statusCode, + Error: pipeline.DeriveError(pctx), + Duration: pipeline.DurationSince(pctx.StartedAt), + TLS: eventTLS(pctx), + }) +} + +// streamingResponseBody is the io.ReadCloser ReverseProxy ferries +// downstream. Each Read call pulls one SSE frame from the upstream, +// dispatches it through the pipeline, and writes the re-framed event +// into the caller's buffer. On end-of-stream a final last=true +// dispatch lets aggregating plugins finalize, and onClose records +// the response SessionEvent. +// +// The struct holds an internal buffer (`pending`) that may not +// drain in one Read — large frames are returned across multiple +// Reads, byte-for-byte. Streaming preserves SSE wire framing +// (`data: \n\n`) regardless of how the upstream emitted +// each frame's data lines. +type streamingResponseBody struct { + upstream io.ReadCloser + reader *sseframe.Reader + ctx context.Context + pipeline *pipeline.Holder + pctx *pipeline.Context + onClose func(statusCode int) + statusCode int + + pending []byte + finished bool + closed bool +} + +func (b *streamingResponseBody) Read(p []byte) (int, error) { + if len(b.pending) > 0 { + n := copy(p, b.pending) + b.pending = b.pending[n:] + return n, nil + } + if b.finished { + return 0, io.EOF + } + + frame, err := b.reader.ReadFrame() + if err == io.EOF { + // End of upstream. Finalize aggregating plugins. + b.pipeline.RunResponseFrame(b.ctx, b.pctx, nil, true) + b.finished = true + return 0, io.EOF + } + if err != nil { + // Stream errored mid-flight. Finalize so plugins can record + // what they have, then propagate the error so net/http closes + // the downstream connection. + b.pipeline.RunResponseFrame(b.ctx, b.pctx, nil, true) + b.finished = true + return 0, err + } + + action := b.pipeline.RunResponseFrame(b.ctx, b.pctx, frame, false) + if action.Type == pipeline.Reject { + // Mid-stream reject from a streaming-aware plugin. Headers and + // earlier frames are already on the wire, so the cleanest + // signal is to abort the read; the client sees a truncated + // stream. Finalize first so plugin state is consistent. + b.pipeline.RunResponseFrame(b.ctx, b.pctx, nil, true) + b.finished = true + return 0, fmt.Errorf("reverseproxy: streaming response rejected mid-stream") + } + + // Re-frame as SSE: "data: \n\n". + out := make([]byte, 0, len(frame)+8) + out = append(out, "data: "...) + out = append(out, frame...) + out = append(out, "\n\n"...) + + n := copy(p, out) + if n < len(out) { + b.pending = out[n:] + } + return n, nil +} + +func (b *streamingResponseBody) Close() error { + if b.closed { + return nil + } + b.closed = true + // Ensure plugins finalize even if Read never reached EOF (client + // disconnect, ReverseProxy error). + if !b.finished { + b.pipeline.RunResponseFrame(b.ctx, b.pctx, nil, true) + b.finished = true + } + if b.onClose != nil { + b.onClose(b.statusCode) + } + return b.upstream.Close() +} + +// isEventStream reports whether a Content-Type header value names the +// SSE media type. Tolerates parameters and ASCII case differences. +// Mirrors the forwardproxy helper of the same name. +func isEventStream(contentType string) bool { + if contentType == "" { + return false + } + if idx := indexByteASCII(contentType, ';'); idx >= 0 { + contentType = contentType[:idx] + } + contentType = trimASCIISpace(contentType) + return equalASCIIFold(contentType, "text/event-stream") +} + +func indexByteASCII(s string, c byte) int { + for i := 0; i < len(s); i++ { + if s[i] == c { + return i + } + } + return -1 +} + +func trimASCIISpace(s string) string { + start := 0 + for start < len(s) && (s[start] == ' ' || s[start] == '\t') { + start++ + } + end := len(s) + for end > start && (s[end-1] == ' ' || s[end-1] == '\t') { + end-- + } + return s[start:end] +} + +func equalASCIIFold(s, t string) bool { + if len(s) != len(t) { + return false + } + for i := 0; i < len(s); i++ { + a, b := s[i], t[i] + if a >= 'A' && a <= 'Z' { + a += 'a' - 'A' + } + if b >= 'A' && b <= 'Z' { + b += 'a' - 'A' + } + if a != b { + return false + } + } + return true +} + // inboundSessionID returns the bucket ID for an inbound event. Mirrors // extproc's inboundSessionID: trusts the A2A-stated contextId when // non-empty, otherwise routes to DefaultSessionID. Does NOT fall back diff --git a/authbridge/authlib/listener/reverseproxy/streaming_test.go b/authbridge/authlib/listener/reverseproxy/streaming_test.go new file mode 100644 index 000000000..24fb3eaf5 --- /dev/null +++ b/authbridge/authlib/listener/reverseproxy/streaming_test.go @@ -0,0 +1,216 @@ +package reverseproxy + +import ( + "bufio" + "bytes" + "context" + "fmt" + "io" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline" +) + +// streamingProbe is a Plugin + StreamingResponder for asserting on +// what the reverseproxy dispatches. ReadsBody=true so the listener +// takes the NeedsBody buffering path on requests. +type streamingProbe struct { + mu sync.Mutex + frames [][]byte + lasts []bool + caps pipeline.PluginCapabilities +} + +func newStreamingProbe(writesBody bool) *streamingProbe { + return &streamingProbe{ + caps: pipeline.PluginCapabilities{ReadsBody: true, WritesBody: writesBody}, + } +} + +func (p *streamingProbe) Name() string { return "streaming-probe" } +func (p *streamingProbe) Capabilities() pipeline.PluginCapabilities { return p.caps } +func (p *streamingProbe) OnRequest(_ context.Context, _ *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} +func (p *streamingProbe) OnResponse(_ context.Context, _ *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} +func (p *streamingProbe) OnResponseFrame(_ context.Context, _ *pipeline.Context, frame []byte, last bool) pipeline.Action { + p.mu.Lock() + defer p.mu.Unlock() + cp := make([]byte, len(frame)) + copy(cp, frame) + p.frames = append(p.frames, cp) + p.lasts = append(p.lasts, last) + return pipeline.Action{Type: pipeline.Continue} +} + +func (p *streamingProbe) snapshot() ([][]byte, []bool) { + p.mu.Lock() + defer p.mu.Unlock() + frames := make([][]byte, len(p.frames)) + copy(frames, p.frames) + lasts := make([]bool, len(p.lasts)) + copy(lasts, p.lasts) + return frames, lasts +} + +// TestReverseProxy_Streaming_FramesFlowThrough asserts the inbound +// streaming path: an A2A message/stream-style SSE upstream is +// forwarded frame-by-frame to the downstream client and each frame +// reaches the StreamingResponder hook. Mirrors the forwardproxy +// integration test so the contract is shared across both listeners. +func TestReverseProxy_Streaming_FramesFlowThrough(t *testing.T) { + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + flusher := w.(http.Flusher) + for i := 1; i <= 3; i++ { + fmt.Fprintf(w, "data: {\"event\":%d}\n\n", i) + flusher.Flush() + time.Sleep(20 * time.Millisecond) + } + })) + defer backend.Close() + + probe := newStreamingProbe(false) + pipe, err := pipeline.New([]pipeline.Plugin{probe}) + if err != nil { + t.Fatalf("New pipeline: %v", err) + } + srv, err := NewServer(pipeline.NewHolder(pipe), nil, backend.URL, nil) + if err != nil { + t.Fatalf("NewServer: %v", err) + } + proxy := httptest.NewServer(srv.Handler()) + defer proxy.Close() + + resp, err := http.Get(proxy.URL + "/stream") + if err != nil { + t.Fatalf("Get: %v", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d", resp.StatusCode) + } + body, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("ReadAll: %v", err) + } + if got := bytes.Count(body, []byte("data:")); got != 3 { + t.Errorf("body has %d data: lines, want 3 — body=%q", got, body) + } + + frames, lasts := probe.snapshot() + if len(frames) != 4 { + t.Fatalf("plugin saw %d calls, want 4 (3 frames + 1 final last=true) — frames=%v lasts=%v", len(frames), framesAsStrings(frames), lasts) + } + if !lasts[3] { + t.Error("final call last=false, want true") + } +} + +// TestReverseProxy_Streaming_FlushesBeforeUpstreamCloses asserts that +// an in-flight client read sees the first SSE event before the +// upstream emits the second — i.e. the proxy is not buffering. Same +// regression test as forwardproxy's streaming counterpart, this time +// for the inbound path. +func TestReverseProxy_Streaming_FlushesBeforeUpstreamCloses(t *testing.T) { + gate := make(chan struct{}) + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + flusher := w.(http.Flusher) + fmt.Fprintf(w, "data: {\"id\":1}\n\n") + flusher.Flush() + <-gate + fmt.Fprintf(w, "data: {\"id\":2}\n\n") + flusher.Flush() + })) + defer backend.Close() + defer close(gate) + + probe := newStreamingProbe(false) + pipe, err := pipeline.New([]pipeline.Plugin{probe}) + if err != nil { + t.Fatalf("New pipeline: %v", err) + } + srv, err := NewServer(pipeline.NewHolder(pipe), nil, backend.URL, nil) + if err != nil { + t.Fatalf("NewServer: %v", err) + } + proxy := httptest.NewServer(srv.Handler()) + defer proxy.Close() + + resp, err := http.Get(proxy.URL + "/stream") + if err != nil { + t.Fatalf("Get: %v", err) + } + defer resp.Body.Close() + + br := bufio.NewReader(resp.Body) + deadline := time.Now().Add(2 * time.Second) + var first []byte + for time.Now().Before(deadline) { + line, err := br.ReadString('\n') + if err != nil { + t.Fatalf("ReadString: %v", err) + } + if bytes.HasPrefix([]byte(line), []byte("data:")) { + first = []byte(line) + break + } + } + if first == nil { + t.Fatal("first frame did not arrive within 2s — proxy buffering") + } + if !bytes.Contains(first, []byte(`"id":1`)) { + t.Errorf("first frame = %q, missing id 1", first) + } +} + +// TestReverseProxy_BufferedDeliversLastTrueFrame: streaming-aware +// plugin sees a single OnResponseFrame call with last=true for an +// application/json response, so plugins use one code path. +func TestReverseProxy_BufferedDeliversLastTrueFrame(t *testing.T) { + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"hello":"world"}`)) + })) + defer backend.Close() + + probe := newStreamingProbe(false) + pipe, _ := pipeline.New([]pipeline.Plugin{probe}) + srv, err := NewServer(pipeline.NewHolder(pipe), nil, backend.URL, nil) + if err != nil { + t.Fatal(err) + } + proxy := httptest.NewServer(srv.Handler()) + defer proxy.Close() + + resp, err := http.Get(proxy.URL + "/oneshot") + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + io.Copy(io.Discard, resp.Body) + + frames, lasts := probe.snapshot() + if len(frames) != 1 || !lasts[0] { + t.Fatalf("frames=%v lasts=%v; want single last=true call", framesAsStrings(frames), lasts) + } +} + +func framesAsStrings(frames [][]byte) []string { + out := make([]string, len(frames)) + for i, f := range frames { + out[i] = string(f) + } + return out +} diff --git a/authbridge/authlib/pipeline/holder.go b/authbridge/authlib/pipeline/holder.go index 2ac8df525..b5f8e018c 100644 --- a/authbridge/authlib/pipeline/holder.go +++ b/authbridge/authlib/pipeline/holder.go @@ -56,6 +56,18 @@ func (h *Holder) RunResponse(ctx context.Context, pctx *Context) Action { return h.p.Load().RunResponse(ctx, pctx) } +// RunResponseFrame is equivalent to h.Load().RunResponseFrame(ctx, pctx, frame, last). +// See Pipeline.RunResponseFrame. +func (h *Holder) RunResponseFrame(ctx context.Context, pctx *Context, frame []byte, last bool) Action { + return h.p.Load().RunResponseFrame(ctx, pctx, frame, last) +} + +// HasStreamingResponders is equivalent to h.Load().HasStreamingResponders(). +// Listeners read this to decide whether to take the streaming code path. +func (h *Holder) HasStreamingResponders() bool { + return h.p.Load().HasStreamingResponders() +} + // RunFinish is equivalent to h.Load().RunFinish(ctx, pctx, outcome). // Listeners call this in a defer at request entry to guarantee // Finisher dispatch on every exit path. See Pipeline.RunFinish for @@ -69,6 +81,12 @@ func (h *Holder) RunFinish(ctx context.Context, pctx *Context, outcome Outcome) // that decide whether to buffer the request/response body. func (h *Holder) NeedsBody() bool { return h.p.Load().NeedsBody() } +// WritesBody is equivalent to h.Load().WritesBody(). Listeners read this +// when deciding whether streaming responses are safe — a pipeline with +// a body mutator can't stream because the proxy can't rewrite a body +// it has already started forwarding. +func (h *Holder) WritesBody() bool { return h.p.Load().WritesBody() } + // Ready is equivalent to h.Load().Ready(). func (h *Holder) Ready() bool { return h.p.Load().Ready() } diff --git a/authbridge/authlib/pipeline/pipeline.go b/authbridge/authlib/pipeline/pipeline.go index a741882c8..40ed13dd9 100644 --- a/authbridge/authlib/pipeline/pipeline.go +++ b/authbridge/authlib/pipeline/pipeline.go @@ -163,6 +163,70 @@ func (p *Pipeline) RunResponse(ctx context.Context, pctx *Context) Action { return Action{Type: Continue} } +// RunResponseFrame dispatches a single response frame to every plugin +// implementing StreamingResponder, in reverse declaration order +// (symmetric with RunResponse). Plugins that don't implement the +// interface are skipped — they're handled by the buffered RunResponse +// path the listener still calls when the response is non-streaming. +// +// The off-policy skip and observe-policy shadow conversion are +// applied identically to RunResponse — see Run for the contract. +// +// Frames are dispatched in wire-arrival order: callers invoke this +// once per frame as they arrive off the upstream, then once with +// last=true (typically with an empty frame) at end-of-stream so +// aggregating plugins can finalize. Application/json responses are +// delivered as a single last=true frame so streaming-aware plugins +// have one code path. +// +// A plugin that returns Reject mid-stream causes the listener to +// short-circuit. Today no in-tree plugin returns Reject here (the +// listeners forward+flush before invoking the hook for observability +// only); the contract leaves room for per-message enforcement later. +func (p *Pipeline) RunResponseFrame(ctx context.Context, pctx *Context, frame []byte, last bool) Action { + for i := len(p.plugins) - 1; i >= 0; i-- { + policy := p.policyAt(i) + if policy == ErrorPolicyOff { + continue + } + if ctx.Err() != nil { + slog.Info("pipeline: response frame cancelled", "plugin", p.plugins[i].Name()) + return Deny("pipeline.cancelled", "request cancelled") + } + sr, ok := p.plugins[i].(StreamingResponder) + if !ok { + continue + } + pctx.setCurrent(p.plugins[i].Name(), InvocationPhaseResponse, policy) + action := sr.OnResponseFrame(ctx, pctx, frame, last) + pctx.clearCurrent() + if action.Type == Reject { + stampPluginName(&action, p.plugins[i].Name()) + if policy == ErrorPolicyObserve { + markShadowAndLog(pctx, p.plugins[i].Name(), InvocationPhaseResponse, action, "response-frame") + continue + } + pctx.setRejectingPlugin(p.plugins[i].Name()) + logReject(p.plugins[i].Name(), action, "pipeline: plugin rejected response frame") + return action + } + } + return Action{Type: Continue} +} + +// HasStreamingResponders reports whether any plugin in the pipeline +// implements StreamingResponder. Listeners use this to decide whether +// the streaming code path is worth taking — without any opt-in plugin +// the buffered path delivers the same result for less complexity. +func (p *Pipeline) HasStreamingResponders() bool { + for _, plugin := range p.plugins { + if _, ok := plugin.(StreamingResponder); ok { + return true + } + } + return false +} + // policyAt returns the resolved policy for plugins[i]. The policies // slice is always the same length as plugins (New guarantees this), // but we check defensively so a zero-value Pipeline (constructed diff --git a/authbridge/authlib/pipeline/plugin.go b/authbridge/authlib/pipeline/plugin.go index 094450603..48bee32a7 100644 --- a/authbridge/authlib/pipeline/plugin.go +++ b/authbridge/authlib/pipeline/plugin.go @@ -166,6 +166,50 @@ type Finisher interface { OnFinish(ctx context.Context, pctx *Context) } +// StreamingResponder is an optional interface a plugin may implement +// when its response handling is naturally per-message rather than over +// a fully-buffered body. Listeners that detect a streaming response +// (today: text/event-stream) deliver each complete protocol message +// to the pipeline as a frame; plugins that opted in see one +// OnResponseFrame call per frame and a final empty-frame call with +// last=true so aggregating plugins (inference-parser, a2a-parser) +// can finalize their running state. +// +// Plugins without streaming awareness are unaffected: the listener +// either falls back to the buffered path (Content-Type: application/json +// or any non-streaming response) and runs OnResponse as today, or — for +// streaming responses — skips OnResponse entirely. Aggregating plugins +// that want to support both shapes implement OnResponseFrame and treat +// the buffered application/json case as a single last=true frame; the +// listener delivers it that way for them so one code path covers both +// shapes. +// +// Contract: +// - Per-frame ordering matches the wire: frames are dispatched to +// plugins in the order the listener parsed them off the upstream +// response, in pipeline reverse order (symmetric with RunResponse). +// - The frame slice is owned by the listener and is valid only for +// the duration of the call. Plugins that need to retain bytes +// must copy. +// - 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). +// - last=true is always called exactly once at end-of-stream, even +// for an empty/zero-frame stream. Plugins finalize on last=true. +// - For application/json responses the listener calls +// OnResponseFrame once with the full body and last=true; OnResponse +// is then NOT called for plugins that implement this interface +// (the framework picks one path). +type StreamingResponder interface { + OnResponseFrame(ctx context.Context, pctx *Context, frame []byte, last bool) Action +} + // Readier is an optional interface a plugin may implement when it has // deferred initialization that matters to a /readyz probe. The host // ANDs Ready() across all implementers to decide whether the pipeline diff --git a/authbridge/authlib/pipeline/streaming_test.go b/authbridge/authlib/pipeline/streaming_test.go new file mode 100644 index 000000000..faeace1db --- /dev/null +++ b/authbridge/authlib/pipeline/streaming_test.go @@ -0,0 +1,173 @@ +package pipeline + +import ( + "context" + "testing" +) + +// streamingStubPlugin embeds stubPlugin with a StreamingResponder +// implementation. Tests opt in by setting onFrame. +type streamingStubPlugin struct { + stubPlugin + onFrame func(ctx context.Context, pctx *Context, frame []byte, last bool) Action +} + +func (s *streamingStubPlugin) OnResponseFrame(ctx context.Context, pctx *Context, frame []byte, last bool) Action { + if s.onFrame != nil { + return s.onFrame(ctx, pctx, frame, last) + } + return Action{Type: Continue} +} + +func TestRunResponseFrame_DispatchOrder_ReverseDeclaration(t *testing.T) { + var order []string + p1 := &streamingStubPlugin{ + stubPlugin: stubPlugin{name: "first"}, + onFrame: func(_ context.Context, _ *Context, _ []byte, _ bool) Action { + order = append(order, "first") + return Action{Type: Continue} + }, + } + p2 := &streamingStubPlugin{ + stubPlugin: stubPlugin{name: "second"}, + onFrame: func(_ context.Context, _ *Context, _ []byte, _ bool) Action { + order = append(order, "second") + return Action{Type: Continue} + }, + } + pipe, err := New([]Plugin{p1, p2}) + if err != nil { + t.Fatalf("New: %v", err) + } + pctx := &Context{} + action := pipe.RunResponseFrame(context.Background(), pctx, []byte("frame"), false) + if action.Type != Continue { + t.Errorf("action = %v, want Continue", action.Type) + } + if len(order) != 2 || order[0] != "second" || order[1] != "first" { + t.Errorf("dispatch order = %v, want reverse declaration [second first]", order) + } +} + +func TestRunResponseFrame_NonStreamingPluginsSkipped(t *testing.T) { + called := false + p := &stubPlugin{ + name: "non-streaming", + onResp: func(_ context.Context, _ *Context) Action { + called = true + return Action{Type: Continue} + }, + } + pipe, err := New([]Plugin{p}) + if err != nil { + t.Fatalf("New: %v", err) + } + pctx := &Context{} + pipe.RunResponseFrame(context.Background(), pctx, []byte("frame"), false) + if called { + t.Error("RunResponseFrame called OnResponse on a non-streaming plugin") + } +} + +func TestRunResponseFrame_LastFlag(t *testing.T) { + var seenLast []bool + p := &streamingStubPlugin{ + stubPlugin: stubPlugin{name: "agg"}, + onFrame: func(_ context.Context, _ *Context, _ []byte, last bool) Action { + seenLast = append(seenLast, last) + return Action{Type: Continue} + }, + } + pipe, _ := New([]Plugin{p}) + pctx := &Context{} + pipe.RunResponseFrame(context.Background(), pctx, []byte("a"), false) + pipe.RunResponseFrame(context.Background(), pctx, []byte("b"), false) + pipe.RunResponseFrame(context.Background(), pctx, nil, true) + if len(seenLast) != 3 { + t.Fatalf("got %d calls, want 3", len(seenLast)) + } + if seenLast[0] || seenLast[1] || !seenLast[2] { + t.Errorf("last flags = %v, want [false false true]", seenLast) + } +} + +func TestRunResponseFrame_RejectStops(t *testing.T) { + var dispatched []string + p1 := &streamingStubPlugin{ + stubPlugin: stubPlugin{name: "first"}, + onFrame: func(_ context.Context, _ *Context, _ []byte, _ bool) Action { + dispatched = append(dispatched, "first") + return Action{Type: Continue} + }, + } + p2 := &streamingStubPlugin{ + stubPlugin: stubPlugin{name: "rejecter"}, + onFrame: func(_ context.Context, _ *Context, _ []byte, _ bool) Action { + dispatched = append(dispatched, "rejecter") + return Deny("test.reject", "denied") + }, + } + // Reverse-order dispatch means rejecter runs first (last in slice). + pipe, _ := New([]Plugin{p1, p2}) + pctx := &Context{} + action := pipe.RunResponseFrame(context.Background(), pctx, []byte("x"), false) + if action.Type != Reject { + t.Errorf("action = %v, want Reject", action.Type) + } + if len(dispatched) != 1 || dispatched[0] != "rejecter" { + t.Errorf("dispatched = %v, want [rejecter] only (reject stops chain)", dispatched) + } +} + +func TestHasStreamingResponders(t *testing.T) { + plain := &stubPlugin{name: "plain"} + stream := &streamingStubPlugin{stubPlugin: stubPlugin{name: "stream"}} + + pipe1, _ := New([]Plugin{plain}) + if pipe1.HasStreamingResponders() { + t.Error("pipe with only plain plugin reports HasStreamingResponders=true") + } + pipe2, _ := New([]Plugin{plain, stream}) + if !pipe2.HasStreamingResponders() { + t.Error("pipe with streaming plugin reports HasStreamingResponders=false") + } +} + +func TestRunResponseFrame_ContextCancelled(t *testing.T) { + called := false + p := &streamingStubPlugin{ + stubPlugin: stubPlugin{name: "p"}, + onFrame: func(_ context.Context, _ *Context, _ []byte, _ bool) Action { + called = true + return Action{Type: Continue} + }, + } + pipe, _ := New([]Plugin{p}) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + pctx := &Context{} + action := pipe.RunResponseFrame(ctx, pctx, nil, true) + if action.Type != Reject { + t.Errorf("action = %v, want Reject on cancelled ctx", action.Type) + } + if called { + t.Error("plugin OnResponseFrame called on cancelled ctx") + } +} + +func TestRunResponseFrame_OffPolicySkipped(t *testing.T) { + called := false + p := &streamingStubPlugin{ + stubPlugin: stubPlugin{name: "off"}, + onFrame: func(_ context.Context, _ *Context, _ []byte, _ bool) Action { + called = true + return Action{Type: Continue} + }, + } + pipe, _ := New([]Plugin{p}, WithPolicies(ErrorPolicyOff)) + pctx := &Context{} + pipe.RunResponseFrame(context.Background(), pctx, []byte("x"), false) + if called { + t.Error("off-policy plugin's OnResponseFrame was called") + } +} diff --git a/authbridge/authlib/plugins/a2aparser/plugin.go b/authbridge/authlib/plugins/a2aparser/plugin.go index c40ab130c..d5c86e5b3 100644 --- a/authbridge/authlib/plugins/a2aparser/plugin.go +++ b/authbridge/authlib/plugins/a2aparser/plugin.go @@ -105,22 +105,20 @@ func (p *A2AParser) OnRequest(_ context.Context, pctx *pipeline.Context) pipelin } // OnResponse extracts the server-assigned session/context ID and response summary -// from the response body. The summary includes final status, artifact text, and -// error message — enabling debugging of agent behavior without reading raw SSE. -// -// Handles both JSON-RPC responses (message/send) and SSE event streams (message/stream). +// from a buffered response body. Streaming-aware listeners take the +// OnResponseFrame path instead and this method becomes a no-op when +// the streaming path has already populated the extension's response +// fields. func (p *A2AParser) OnResponse(_ context.Context, pctx *pipeline.Context) pipeline.Action { - // Stay silent when the request side never participated — the parser - // recorded nothing on request, so recording on response would orphan - // the row. if pctx.Extensions.A2A == nil { return pipeline.Action{Type: pipeline.Continue} } - // We DID process the request but the response has no body — record - // a Skip so abctl can pair the response row with the request row. - // Without this, the request invocation orphans and the events table - // shows req+resp as unpaired (no ┌/└ glyphs, plugin/action columns - // blank on the response side). + ext := pctx.Extensions.A2A + // If OnResponseFrame already finalized (any response field set), + // don't re-record on the buffered path. + if ext.FinalStatus != "" || ext.Artifact != "" || ext.ErrorMessage != "" { + return pipeline.Action{Type: pipeline.Continue} + } if len(pctx.ResponseBody) == 0 { pctx.Skip("no_response_body") return pipeline.Action{Type: pipeline.Continue} @@ -128,28 +126,76 @@ func (p *A2AParser) OnResponse(_ context.Context, pctx *pipeline.Context) pipeli // Capture the server-assigned contextId — but only when the request // didn't already carry one. Overwriting would split the inbound - // request and response events into different session buckets (the - // agent's A2A SDK may mint a fresh contextId in its response even - // when the client supplied one, which is legal per A2A but breaks - // telemetry bucketing). The request-side contextId is authoritative - // for session attribution. - if pctx.Extensions.A2A.SessionID == "" { + // request and response events into different session buckets. + if ext.SessionID == "" { if sid := extractSessionID(pctx.ResponseBody); sid != "" { - pctx.Extensions.A2A.SessionID = sid + ext.SessionID = sid } } - // Extract response summary (final status + artifact + error) - extractResponseSummary(pctx.ResponseBody, pctx.Extensions.A2A) + extractResponseSummary(pctx.ResponseBody, ext) + logA2AFinalized(ext) + pctx.Observe("matched_" + ext.Method + "_response") + return pipeline.Action{Type: pipeline.Continue} +} +// OnResponseFrame folds each message/stream SSE event into the +// running A2A response state. Final-status and artifact text are +// accumulated across frames; the public extension fields are +// finalized on last=true so the session event sees a single +// consistent snapshot. +// +// application/json (message/send) responses are delivered as a +// single last=true frame carrying the full envelope — the same path +// extracts the response summary from it. +func (p *A2AParser) OnResponseFrame(_ context.Context, pctx *pipeline.Context, frame []byte, last bool) pipeline.Action { + if pctx.Extensions.A2A == nil { + return pipeline.Action{Type: pipeline.Continue} + } + ext := pctx.Extensions.A2A + + // Per-frame fold. message/stream events arrive one per frame; + // message/send arrives as a single last=true frame. + if len(frame) > 0 { + // Try the message/send envelope first — message/stream events + // don't have the {result: {status: {state: "..."}}} top-level + // shape so extractSendResponse will return false and we fall + // through to the per-event extractor. + if !extractSendResponse(frame, ext) { + extractStreamEvent(frame, ext) + } + // Capture session id from any frame if the request didn't carry one. + if ext.SessionID == "" { + if sid := sessionIDFromJSON(frame); sid != "" { + ext.SessionID = sid + } + } + } + + if last { + // Empty stream and we never recorded anything on the request + // side — record a Skip so the response row is paired with the + // request row in abctl. + if ext.FinalStatus == "" && ext.Artifact == "" && ext.ErrorMessage == "" { + pctx.Skip("no_response_body") + return pipeline.Action{Type: pipeline.Continue} + } + logA2AFinalized(ext) + pctx.Observe("matched_" + ext.Method + "_response") + } + return pipeline.Action{Type: pipeline.Continue} +} + +// logA2AFinalized emits the operator-facing debug log once a response +// is finalized — shared by OnResponse and OnResponseFrame so the two +// paths log identically. +func logA2AFinalized(ext *pipeline.A2AExtension) { slog.Debug("a2a-parser: response parsed", - "sessionId", pctx.Extensions.A2A.SessionID, - "finalStatus", pctx.Extensions.A2A.FinalStatus, - "artifactLen", len(pctx.Extensions.A2A.Artifact), - "error", pctx.Extensions.A2A.ErrorMessage, + "sessionId", ext.SessionID, + "finalStatus", ext.FinalStatus, + "artifactLen", len(ext.Artifact), + "error", ext.ErrorMessage, ) - pctx.Observe("matched_" + pctx.Extensions.A2A.Method + "_response") - return pipeline.Action{Type: pipeline.Continue} } // extractResponseSummary parses the response body for final status, artifact text, @@ -216,7 +262,9 @@ func extractSendResponse(body []byte, ext *pipeline.A2AExtension) bool { return true } -// extractStreamResponse handles message/stream SSE responses. +// extractStreamResponse handles message/stream SSE responses (the +// buffered path). For each "data:" line it extracts a single event +// and folds it into the extension via extractStreamEvent. func extractStreamResponse(body []byte, ext *pipeline.A2AExtension) { for _, line := range bytes.Split(body, []byte("\n")) { line = bytes.TrimSpace(line) @@ -227,61 +275,63 @@ func extractStreamResponse(body []byte, ext *pipeline.A2AExtension) { if len(data) == 0 { continue } + extractStreamEvent(data, ext) + } +} - var event struct { - Result struct { - Kind string `json:"kind"` - Final bool `json:"final"` - TaskID string `json:"taskId"` - Status struct { - State string `json:"state"` - Message struct { - Parts []struct { - Kind string `json:"kind"` - Text string `json:"text"` - } `json:"parts"` - } `json:"message"` - } `json:"status"` - Artifact struct { +// extractStreamEvent folds one A2A message/stream event payload (the +// JSON object after `data: `) into the extension's running response +// state. Used by both the buffered path (extractStreamResponse loops +// over data: lines) and the streaming path (OnResponseFrame already +// has the parsed payload as a frame). +func extractStreamEvent(data []byte, ext *pipeline.A2AExtension) { + var event struct { + Result struct { + Kind string `json:"kind"` + Final bool `json:"final"` + TaskID string `json:"taskId"` + Status struct { + State string `json:"state"` + Message struct { Parts []struct { Kind string `json:"kind"` Text string `json:"text"` } `json:"parts"` - } `json:"artifact"` - } `json:"result"` - } - if json.Unmarshal(data, &event) != nil { - continue - } + } `json:"message"` + } `json:"status"` + Artifact struct { + Parts []struct { + Kind string `json:"kind"` + Text string `json:"text"` + } `json:"parts"` + } `json:"artifact"` + } `json:"result"` + } + if json.Unmarshal(data, &event) != nil { + return + } - // Capture taskId from any event - if ext.TaskID == "" && event.Result.TaskID != "" { - ext.TaskID = event.Result.TaskID - } + if ext.TaskID == "" && event.Result.TaskID != "" { + ext.TaskID = event.Result.TaskID + } - switch event.Result.Kind { - case "status-update": - if event.Result.Final { - ext.FinalStatus = event.Result.Status.State - // Extract error message on failure - if event.Result.Status.State == "failed" { - for _, part := range event.Result.Status.Message.Parts { - if part.Kind == "text" && part.Text != "" { - ext.ErrorMessage = part.Text - break - } + switch event.Result.Kind { + case "status-update": + if event.Result.Final { + ext.FinalStatus = event.Result.Status.State + if event.Result.Status.State == "failed" { + for _, part := range event.Result.Status.Message.Parts { + if part.Kind == "text" && part.Text != "" { + ext.ErrorMessage = part.Text + break } } } - case "artifact-update", "artifact": - // A2A SDKs emit kind="artifact-update" on the stream; older - // samples use "artifact". Accept both. Concatenate text parts - // from the frame; repeated frames for the same artifact carry - // appended text so we accumulate across frames. - for _, part := range event.Result.Artifact.Parts { - if part.Kind == "text" && part.Text != "" { - ext.Artifact += part.Text - } + } + case "artifact-update", "artifact": + for _, part := range event.Result.Artifact.Parts { + if part.Kind == "text" && part.Text != "" { + ext.Artifact += part.Text } } } diff --git a/authbridge/authlib/plugins/a2aparser/streaming_test.go b/authbridge/authlib/plugins/a2aparser/streaming_test.go new file mode 100644 index 000000000..b32f93938 --- /dev/null +++ b/authbridge/authlib/plugins/a2aparser/streaming_test.go @@ -0,0 +1,113 @@ +package a2aparser + +import ( + "context" + "testing" + + "github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline" +) + +func TestA2AParser_OnResponseFrame_FoldsArtifactAndFinalStatus(t *testing.T) { + p := NewA2AParser() + pctx := &pipeline.Context{ + Extensions: pipeline.Extensions{ + A2A: &pipeline.A2AExtension{Method: "message/stream"}, + }, + } + + frames := [][]byte{ + []byte(`{"result":{"kind":"artifact-update","taskId":"task-1","artifact":{"parts":[{"kind":"text","text":"Hello "}]}}}`), + []byte(`{"result":{"kind":"artifact-update","taskId":"task-1","artifact":{"parts":[{"kind":"text","text":"world"}]}}}`), + []byte(`{"result":{"kind":"status-update","final":true,"status":{"state":"completed"}}}`), + } + for _, f := range frames { + if action := p.OnResponseFrame(context.Background(), pctx, f, false); action.Type != pipeline.Continue { + t.Fatalf("frame action = %v, want Continue", action.Type) + } + } + p.OnResponseFrame(context.Background(), pctx, nil, true) + + ext := pctx.Extensions.A2A + if ext.Artifact != "Hello world" { + t.Errorf("Artifact = %q, want %q", ext.Artifact, "Hello world") + } + if ext.FinalStatus != "completed" { + t.Errorf("FinalStatus = %q, want completed", ext.FinalStatus) + } + if ext.TaskID != "task-1" { + t.Errorf("TaskID = %q, want task-1", ext.TaskID) + } +} + +func TestA2AParser_OnResponseFrame_FailedStatusErrorMessage(t *testing.T) { + p := NewA2AParser() + pctx := &pipeline.Context{ + Extensions: pipeline.Extensions{ + A2A: &pipeline.A2AExtension{Method: "message/stream"}, + }, + } + frame := []byte(`{"result":{"kind":"status-update","final":true,"status":{"state":"failed","message":{"parts":[{"kind":"text","text":"upstream timeout"}]}}}}`) + p.OnResponseFrame(context.Background(), pctx, frame, false) + p.OnResponseFrame(context.Background(), pctx, nil, true) + if pctx.Extensions.A2A.FinalStatus != "failed" { + t.Errorf("FinalStatus = %q, want failed", pctx.Extensions.A2A.FinalStatus) + } + if pctx.Extensions.A2A.ErrorMessage != "upstream timeout" { + t.Errorf("ErrorMessage = %q, want upstream timeout", pctx.Extensions.A2A.ErrorMessage) + } +} + +func TestA2AParser_OnResponseFrame_ApplicationJSONOneShot(t *testing.T) { + p := NewA2AParser() + pctx := &pipeline.Context{ + Extensions: pipeline.Extensions{ + A2A: &pipeline.A2AExtension{Method: "message/send"}, + }, + } + body := []byte(`{"result":{"taskId":"t-2","status":{"state":"completed"},"artifacts":[{"parts":[{"kind":"text","text":"answer"}]}]}}`) + action := p.OnResponseFrame(context.Background(), pctx, body, true) + if action.Type != pipeline.Continue { + t.Fatalf("action = %v, want Continue", action.Type) + } + ext := pctx.Extensions.A2A + if ext.FinalStatus != "completed" { + t.Errorf("FinalStatus = %q", ext.FinalStatus) + } + if ext.Artifact != "answer" { + t.Errorf("Artifact = %q, want answer", ext.Artifact) + } + if ext.TaskID != "t-2" { + t.Errorf("TaskID = %q, want t-2", ext.TaskID) + } +} + +func TestA2AParser_OnResponseFrame_CapturesContextID(t *testing.T) { + p := NewA2AParser() + pctx := &pipeline.Context{ + Extensions: pipeline.Extensions{ + A2A: &pipeline.A2AExtension{Method: "message/stream"}, + }, + } + // Request had no contextId; first stream event reveals it. + frame := []byte(`{"result":{"kind":"status-update","contextId":"ctx-7","final":false,"status":{"state":"running"}}}`) + p.OnResponseFrame(context.Background(), pctx, frame, false) + p.OnResponseFrame(context.Background(), pctx, nil, true) + if pctx.Extensions.A2A.SessionID != "ctx-7" { + t.Errorf("SessionID = %q, want ctx-7", pctx.Extensions.A2A.SessionID) + } +} + +func TestA2AParser_OnResponseFrame_RequestContextIDPreserved(t *testing.T) { + p := NewA2AParser() + pctx := &pipeline.Context{ + Extensions: pipeline.Extensions{ + A2A: &pipeline.A2AExtension{Method: "message/stream", SessionID: "from-request"}, + }, + } + frame := []byte(`{"result":{"kind":"status-update","contextId":"different","final":false,"status":{"state":"running"}}}`) + p.OnResponseFrame(context.Background(), pctx, frame, false) + p.OnResponseFrame(context.Background(), pctx, nil, true) + if pctx.Extensions.A2A.SessionID != "from-request" { + t.Errorf("SessionID = %q, want preserved from-request", pctx.Extensions.A2A.SessionID) + } +} diff --git a/authbridge/authlib/plugins/inferenceparser/plugin.go b/authbridge/authlib/plugins/inferenceparser/plugin.go index 876c82d0c..8b59786d6 100644 --- a/authbridge/authlib/plugins/inferenceparser/plugin.go +++ b/authbridge/authlib/plugins/inferenceparser/plugin.go @@ -98,29 +98,138 @@ func (p *InferenceParser) OnRequest(_ context.Context, pctx *pipeline.Context) p } // OnResponse populates the response-side fields (Completion, FinishReason, -// token counts) on pctx.Extensions.Inference. Handles both non-streaming -// JSON responses and SSE streams from OpenAI-compatible servers. +// token counts) on pctx.Extensions.Inference for buffered responses. +// Streaming-aware listeners take the OnResponseFrame path instead and +// this method is a no-op when the streaming path has already +// finalized state. Aggregation across frames lives in OnResponseFrame +// + the inferenceStreamState scratch on the extension. func (p *InferenceParser) OnResponse(_ context.Context, pctx *pipeline.Context) pipeline.Action { - // Stay silent when the request side never participated — the parser - // recorded nothing on request, so recording on response would orphan - // the row. if pctx.Extensions.Inference == nil { return pipeline.Action{Type: pipeline.Continue} } - // We DID process the request but the response has no body — record - // a Skip so abctl can pair the response row with the request row. + // If the streaming path already finalized (Completion populated or + // FinishReason set), don't re-record on the buffered path. + ext := pctx.Extensions.Inference + if ext.Completion != "" || ext.FinishReason != "" || ext.TotalTokens > 0 { + return pipeline.Action{Type: pipeline.Continue} + } if len(pctx.ResponseBody) == 0 { pctx.Skip("no_response_body") return pipeline.Action{Type: pipeline.Continue} } - if pctx.Extensions.Inference.Stream { - parseInferenceSSE(pctx.ResponseBody, pctx.Extensions.Inference) + if ext.Stream { + parseInferenceSSE(pctx.ResponseBody, ext) } else { - parseInferenceJSON(pctx.ResponseBody, pctx.Extensions.Inference) + parseInferenceJSON(pctx.ResponseBody, ext) } + logInferenceFinalized(ext) + pctx.Observe("matched_" + ext.Model + "_response") + return pipeline.Action{Type: pipeline.Continue} +} + +// inferenceStreamState is the scratch state kept on the extension for +// the duration of a streaming response. Lives in pctx.Extensions.Custom +// under a private key — kept off the public InferenceExtension shape so +// the API stays clean. The struct accumulates the in-progress +// completion until last=true triggers finalization. +type inferenceStreamState struct { + completion strings.Builder + usage inferenceUsage +} + +// streamStateKey scopes the scratch state to this plugin in +// pctx.Extensions.Custom. Other plugins see pctx.Extensions.Custom +// keys but won't collide with this one. +const streamStateKey = "inference-parser/stream-state" + +// OnResponseFrame folds each SSE-data chunk into the running +// completion. On last=true the finalized result is written to the +// public InferenceExtension fields (Completion / FinishReason / +// token counts) and the Observe row is recorded. +// +// Application/json responses are delivered as a single last=true +// frame containing the full JSON body — the dual path keeps one +// code path for both shapes. +func (p *InferenceParser) OnResponseFrame(_ context.Context, pctx *pipeline.Context, frame []byte, last bool) pipeline.Action { + if pctx.Extensions.Inference == nil { + return pipeline.Action{Type: pipeline.Continue} + } ext := pctx.Extensions.Inference + + // application/json one-shot: single last=true frame carrying the + // complete envelope. Streaming responses arrive as multiple frames + // where ext.Stream==true; tell them apart by the request-side flag. + if last && !ext.Stream { + if len(frame) == 0 { + pctx.Skip("no_response_body") + return pipeline.Action{Type: pipeline.Continue} + } + parseInferenceJSON(frame, ext) + logInferenceFinalized(ext) + pctx.Observe("matched_" + ext.Model + "_response") + return pipeline.Action{Type: pipeline.Continue} + } + + // Streaming path. Lazily allocate the per-stream scratch. + state := getOrCreateStreamState(pctx) + + if len(frame) > 0 { + // Each frame is one OpenAI streaming chunk: data: { choices: [...], usage: ... } + // or the literal sentinel "[DONE]". Skip the sentinel. + if !bytes.Equal(bytes.TrimSpace(frame), []byte("[DONE]")) { + var chunk inferenceStreamChunk + if err := json.Unmarshal(frame, &chunk); err != nil { + slog.Debug("inference-parser: malformed streaming chunk, skipping", "error", err) + } else { + for _, c := range chunk.Choices { + if c.Delta.Content != "" { + state.completion.WriteString(c.Delta.Content) + } + if c.FinishReason != "" { + ext.FinishReason = c.FinishReason + } + } + if chunk.Usage.TotalTokens > 0 { + state.usage = chunk.Usage + } + } + } + } + + if last { + ext.Completion = state.completion.String() + if state.usage.TotalTokens > 0 { + ext.PromptTokens = state.usage.PromptTokens + ext.CompletionTokens = state.usage.CompletionTokens + ext.TotalTokens = state.usage.TotalTokens + } + // Empty stream with no body and no chunks — record Skip to + // pair the response row with the request row. + if ext.Completion == "" && ext.FinishReason == "" && ext.TotalTokens == 0 { + pctx.Skip("no_response_body") + return pipeline.Action{Type: pipeline.Continue} + } + logInferenceFinalized(ext) + pctx.Observe("matched_" + ext.Model + "_response") + } + return pipeline.Action{Type: pipeline.Continue} +} + +func getOrCreateStreamState(pctx *pipeline.Context) *inferenceStreamState { + if s := pipeline.GetState[inferenceStreamState](pctx, streamStateKey); s != nil { + return s + } + s := &inferenceStreamState{} + pipeline.SetState(pctx, streamStateKey, s) + return s +} + +// logInferenceFinalized emits the operator-facing INFO log + Observe +// once a response is finalized — shared by OnResponse and +// OnResponseFrame so streaming and buffered finalize identically. +func logInferenceFinalized(ext *pipeline.InferenceExtension) { slog.Info("inference-parser: response", "model", ext.Model, "finishReason", ext.FinishReason, @@ -128,8 +237,6 @@ func (p *InferenceParser) OnResponse(_ context.Context, pctx *pipeline.Context) "completionTokens", ext.CompletionTokens, ) slog.Debug("inference-parser: completion", "text", parsercommon.Truncate(ext.Completion, parsercommon.DebugBodyMax)) - pctx.Observe("matched_" + ext.Model + "_response") - return pipeline.Action{Type: pipeline.Continue} } // parseInferenceJSON parses a non-streaming OpenAI chat/completions response. diff --git a/authbridge/authlib/plugins/inferenceparser/streaming_test.go b/authbridge/authlib/plugins/inferenceparser/streaming_test.go new file mode 100644 index 000000000..8957241c4 --- /dev/null +++ b/authbridge/authlib/plugins/inferenceparser/streaming_test.go @@ -0,0 +1,105 @@ +package inferenceparser + +import ( + "context" + "testing" + + "github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline" +) + +func TestInferenceParser_OnResponseFrame_StreamFoldsDeltas(t *testing.T) { + p := NewInferenceParser() + pctx := &pipeline.Context{ + Extensions: pipeline.Extensions{ + Inference: &pipeline.InferenceExtension{ + Model: "gpt-4", + Stream: true, + }, + }, + } + + frames := [][]byte{ + []byte(`{"choices":[{"delta":{"content":"Hello"}}]}`), + []byte(`{"choices":[{"delta":{"content":" world"}}]}`), + []byte(`{"choices":[{"delta":{"content":"!"},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":3,"total_tokens":8}}`), + } + for _, f := range frames { + if action := p.OnResponseFrame(context.Background(), pctx, f, false); action.Type != pipeline.Continue { + t.Fatalf("frame action = %v, want Continue", action.Type) + } + } + // Mid-stream nothing finalized yet. + if pctx.Extensions.Inference.Completion != "" { + t.Errorf("Completion populated mid-stream = %q; finalize should wait for last=true", pctx.Extensions.Inference.Completion) + } + + // Finalize on last=true. + p.OnResponseFrame(context.Background(), pctx, nil, true) + ext := pctx.Extensions.Inference + if ext.Completion != "Hello world!" { + t.Errorf("Completion = %q, want %q", ext.Completion, "Hello world!") + } + if ext.FinishReason != "stop" { + t.Errorf("FinishReason = %q, want stop", ext.FinishReason) + } + if ext.PromptTokens != 5 || ext.CompletionTokens != 3 || ext.TotalTokens != 8 { + t.Errorf("usage = (%d,%d,%d), want (5,3,8)", ext.PromptTokens, ext.CompletionTokens, ext.TotalTokens) + } +} + +func TestInferenceParser_OnResponseFrame_DONESkipped(t *testing.T) { + p := NewInferenceParser() + pctx := &pipeline.Context{ + Extensions: pipeline.Extensions{ + Inference: &pipeline.InferenceExtension{Model: "m", Stream: true}, + }, + } + p.OnResponseFrame(context.Background(), pctx, []byte(`{"choices":[{"delta":{"content":"hi"}}]}`), false) + p.OnResponseFrame(context.Background(), pctx, []byte("[DONE]"), false) + p.OnResponseFrame(context.Background(), pctx, nil, true) + if pctx.Extensions.Inference.Completion != "hi" { + t.Errorf("Completion = %q, want %q ([DONE] should not be parsed as JSON)", pctx.Extensions.Inference.Completion, "hi") + } +} + +func TestInferenceParser_OnResponseFrame_ApplicationJSONOneShot(t *testing.T) { + p := NewInferenceParser() + pctx := &pipeline.Context{ + Extensions: pipeline.Extensions{ + Inference: &pipeline.InferenceExtension{Model: "m", Stream: false}, + }, + } + body := []byte(`{"choices":[{"message":{"role":"assistant","content":"non-streaming reply"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3}}`) + p.OnResponseFrame(context.Background(), pctx, body, true) + ext := pctx.Extensions.Inference + if ext.Completion != "non-streaming reply" { + t.Errorf("Completion = %q", ext.Completion) + } + if ext.TotalTokens != 3 { + t.Errorf("TotalTokens = %d, want 3", ext.TotalTokens) + } +} + +func TestInferenceParser_OnResponseFrame_NoExtensionMeansNoOp(t *testing.T) { + p := NewInferenceParser() + pctx := &pipeline.Context{} + action := p.OnResponseFrame(context.Background(), pctx, []byte(`{}`), true) + if action.Type != pipeline.Continue { + t.Errorf("action = %v, want Continue", action.Type) + } +} + +func TestInferenceParser_OnResponseFrame_EmptyStreamRecordsSkip(t *testing.T) { + p := NewInferenceParser() + pctx := &pipeline.Context{ + Extensions: pipeline.Extensions{ + Inference: &pipeline.InferenceExtension{Model: "m", Stream: true}, + }, + } + pctx.SetCurrentPlugin("inference-parser", pipeline.InvocationPhaseResponse) + p.OnResponseFrame(context.Background(), pctx, nil, true) + pctx.ClearCurrentPlugin() + if pctx.Extensions.Invocations == nil { + t.Fatal("no invocation recorded for empty stream") + } +} diff --git a/authbridge/authlib/plugins/mcpparser/plugin.go b/authbridge/authlib/plugins/mcpparser/plugin.go index e5c3069f1..8011465e1 100644 --- a/authbridge/authlib/plugins/mcpparser/plugin.go +++ b/authbridge/authlib/plugins/mcpparser/plugin.go @@ -211,6 +211,26 @@ func (p *MCPParser) OnRequest(_ context.Context, pctx *pipeline.Context) pipelin return pipeline.Action{Type: pipeline.Continue} } +// OnResponse is the buffered-path response hook. Streaming-aware +// listeners do NOT call it on text/event-stream responses — they +// dispatch through OnResponseFrame instead (see StreamingResponder). +// For application/json responses streaming-aware listeners deliver a +// single last=true frame to OnResponseFrame, so this method exists +// only to support listeners (or pipelines with no streaming-aware +// plugins) that still take the buffered RunResponse path. +// +// To avoid double-recording on streaming-aware listeners that call +// both RunResponse and RunResponseFrame for non-streaming responses, +// the per-frame path bypasses OnResponse via the StreamingResponder +// type assertion in pipeline.RunResponse — actually no, RunResponse +// runs all plugins regardless. So when running a streaming-aware +// pipeline, RunResponseFrame is the one that records; OnResponse +// here is the no-op for the same response. +// +// The compromise: OnResponseFrame populates pctx.Extensions.MCP.Result +// and records the Observe. OnResponse only acts when no streaming- +// aware path has fired, gated on whether the result/error has already +// been populated by OnResponseFrame. func (p *MCPParser) OnResponse(_ context.Context, pctx *pipeline.Context) pipeline.Action { // Stay silent when the request side never participated — the parser // recorded nothing on request, so recording on response would orphan @@ -218,6 +238,15 @@ func (p *MCPParser) OnResponse(_ context.Context, pctx *pipeline.Context) pipeli if pctx.Extensions.MCP == nil { return pipeline.Action{Type: pipeline.Continue} } + // If OnResponseFrame already populated the result/error or recorded + // a skip on this pass, don't re-record. This keeps the buffered and + // streaming paths from emitting duplicate rows when the listener + // runs both (it shouldn't, today — the streaming path skips + // RunResponse entirely — but the guard is cheap and protects + // future listener variants). + if pctx.Extensions.MCP.Result != nil || pctx.Extensions.MCP.Err != nil { + return pipeline.Action{Type: pipeline.Continue} + } // We DID process the request but the response has no body — typical // for JSON-RPC notifications that ack with HTTP 202. Record a Skip // so abctl can pair the response row with the request row in the @@ -235,6 +264,64 @@ func (p *MCPParser) OnResponse(_ context.Context, pctx *pipeline.Context) pipeli return pipeline.Action{Type: pipeline.Continue} } + applyMCPResponseRPC(pctx, rpc) + slog.Debug("mcp-parser: response detail", "method", pctx.Extensions.MCP.Method, "body", parsercommon.Truncate(string(pctx.ResponseBody), parsercommon.DebugBodyMax)) + return pipeline.Action{Type: pipeline.Continue} +} + +// OnResponseFrame is the streaming-aware response hook. Listeners +// invoke it once per SSE frame (text/event-stream) and once with +// last=true at end-of-stream. application/json responses arrive as +// a single last=true frame — so the same code handles both shapes. +// +// Per-message recording: for MCP, each frame's payload is one +// 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 { + // Stay silent when the request side never participated. + if pctx.Extensions.MCP == nil { + return pipeline.Action{Type: pipeline.Continue} + } + + // End-of-stream call with no payload. If we never saw a frame with + // a result/error and the response body was empty, record a Skip + // (matches the buffered path's "no_response_body" semantics so + // abctl pairs request and response rows uniformly across shapes). + // For non-empty streams the per-frame Observe entries are the + // session timeline; nothing to do on last=true. + if len(frame) == 0 { + if last && pctx.Extensions.MCP.Result == nil && pctx.Extensions.MCP.Err == nil { + pctx.Skip("no_response_body") + } + return pipeline.Action{Type: pipeline.Continue} + } + + var rpc jsonRPCResponse + if err := json.Unmarshal(frame, &rpc); err != nil { + // A malformed JSON-RPC message in a stream is unusual but + // recoverable — skip it and keep going. Don't return Reject + // because the listener is forwarding bytes regardless of what + // we say (record-only contract). + slog.Debug("mcp-parser: malformed frame, skipping", "error", err, "frameLen", len(frame)) + return pipeline.Action{Type: pipeline.Continue} + } + if rpc.Result == nil && rpc.Error == nil { + // Notifications, heartbeats, or other JSON-RPC shapes without + // a result/error envelope — silently skip (no observation). + return pipeline.Action{Type: pipeline.Continue} + } + + applyMCPResponseRPC(pctx, rpc) + slog.Debug("mcp-parser: streaming frame", "method", pctx.Extensions.MCP.Method, "body", parsercommon.Truncate(string(frame), parsercommon.DebugBodyMax)) + return pipeline.Action{Type: pipeline.Continue} +} + +// applyMCPResponseRPC mutates pctx.Extensions.MCP from a parsed +// JSON-RPC response and emits the operator-facing log + Observe row. +// Shared by OnResponse (buffered) and OnResponseFrame (streaming) so +// the two paths agree on what gets recorded for one envelope. +func applyMCPResponseRPC(pctx *pipeline.Context, rpc jsonRPCResponse) { if rpc.Error != nil { pctx.Extensions.MCP.Err = &pipeline.MCPError{ Code: rpc.Error.Code, @@ -243,17 +330,13 @@ func (p *MCPParser) OnResponse(_ context.Context, pctx *pipeline.Context) pipeli } slog.Info("mcp-parser: response error", "method", pctx.Extensions.MCP.Method, "code", rpc.Error.Code, "message", rpc.Error.Message) pctx.Observe("response_error") - return pipeline.Action{Type: pipeline.Continue} + return } - if rpc.Result != nil { pctx.Extensions.MCP.Result = rpc.Result slog.Info("mcp-parser: response", "method", pctx.Extensions.MCP.Method, "resultKeys", resultKeys(rpc.Result)) - slog.Debug("mcp-parser: response detail", "method", pctx.Extensions.MCP.Method, "body", parsercommon.Truncate(string(pctx.ResponseBody), parsercommon.DebugBodyMax)) } - pctx.Observe("matched_" + pctx.Extensions.MCP.Method + "_response") - return pipeline.Action{Type: pipeline.Continue} } // parseMCPResponse handles both plain JSON-RPC responses and SSE event streams diff --git a/authbridge/authlib/plugins/mcpparser/streaming_test.go b/authbridge/authlib/plugins/mcpparser/streaming_test.go new file mode 100644 index 000000000..f824ce179 --- /dev/null +++ b/authbridge/authlib/plugins/mcpparser/streaming_test.go @@ -0,0 +1,114 @@ +package mcpparser + +import ( + "context" + "testing" + + "github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline" +) + +func TestMCPParser_OnResponseFrame_PerMessageRecording(t *testing.T) { + p := NewMCPParser() + pctx := &pipeline.Context{ + Extensions: pipeline.Extensions{ + MCP: &pipeline.MCPExtension{Method: "tools/call"}, + }, + } + + // First frame: a complete JSON-RPC result. + frame := []byte(`{"jsonrpc":"2.0","id":1,"result":{"content":[{"type":"text","text":"hello"}]}}`) + action := p.OnResponseFrame(context.Background(), pctx, frame, false) + if action.Type != pipeline.Continue { + t.Fatalf("action = %v, want Continue", action.Type) + } + if pctx.Extensions.MCP.Result == nil { + t.Fatal("Result not populated by OnResponseFrame") + } + // last=true with empty frame finalizes (no-op for already-populated Result). + p.OnResponseFrame(context.Background(), pctx, nil, true) + if pctx.Extensions.MCP.Result == nil { + t.Error("Result cleared on last=true") + } +} + +func TestMCPParser_OnResponseFrame_ApplicationJSONOneShot(t *testing.T) { + p := NewMCPParser() + pctx := &pipeline.Context{ + Extensions: pipeline.Extensions{ + MCP: &pipeline.MCPExtension{Method: "tools/call"}, + }, + } + body := []byte(`{"jsonrpc":"2.0","id":2,"result":{"ok":true}}`) + // application/json: single last=true frame containing whole body. + action := p.OnResponseFrame(context.Background(), pctx, body, true) + if action.Type != pipeline.Continue { + t.Fatalf("action = %v, want Continue", action.Type) + } + if pctx.Extensions.MCP.Result == nil { + t.Fatal("Result not populated") + } + if pctx.Extensions.MCP.Result["ok"] != true { + t.Errorf("Result[ok] = %v, want true", pctx.Extensions.MCP.Result["ok"]) + } +} + +func TestMCPParser_OnResponseFrame_ErrorFrame(t *testing.T) { + p := NewMCPParser() + pctx := &pipeline.Context{ + Extensions: pipeline.Extensions{ + MCP: &pipeline.MCPExtension{Method: "tools/call"}, + }, + } + frame := []byte(`{"jsonrpc":"2.0","id":3,"error":{"code":-32601,"message":"method not found"}}`) + p.OnResponseFrame(context.Background(), pctx, frame, true) + if pctx.Extensions.MCP.Err == nil { + t.Fatal("Err not populated") + } + if pctx.Extensions.MCP.Err.Code != -32601 { + t.Errorf("Err.Code = %d, want -32601", pctx.Extensions.MCP.Err.Code) + } +} + +func TestMCPParser_OnResponseFrame_NoExtensionMeansNoOp(t *testing.T) { + p := NewMCPParser() + pctx := &pipeline.Context{} + action := p.OnResponseFrame(context.Background(), pctx, []byte(`{}`), true) + if action.Type != pipeline.Continue { + t.Errorf("action = %v, want Continue", action.Type) + } + if pctx.Extensions.MCP != nil { + t.Error("MCPExtension created when request side never participated") + } +} + +func TestMCPParser_OnResponseFrame_EmptyStreamRecordsSkip(t *testing.T) { + p := NewMCPParser() + pctx := &pipeline.Context{ + Extensions: pipeline.Extensions{ + MCP: &pipeline.MCPExtension{Method: "tools/call"}, + }, + } + // Streaming response with zero data frames: only the last=true call. + pctx.SetCurrentPlugin("mcp-parser", pipeline.InvocationPhaseResponse) + p.OnResponseFrame(context.Background(), pctx, nil, true) + pctx.ClearCurrentPlugin() + if pctx.Extensions.Invocations == nil { + t.Fatal("no invocation recorded") + } +} + +func TestMCPParser_OnResponseFrame_MalformedFrameSkipped(t *testing.T) { + p := NewMCPParser() + pctx := &pipeline.Context{ + Extensions: pipeline.Extensions{ + MCP: &pipeline.MCPExtension{Method: "tools/call"}, + }, + } + action := p.OnResponseFrame(context.Background(), pctx, []byte("not json"), false) + if action.Type != pipeline.Continue { + t.Errorf("action = %v, want Continue (malformed frame should skip silently)", action.Type) + } + if pctx.Extensions.MCP.Result != nil || pctx.Extensions.MCP.Err != nil { + t.Error("malformed frame populated Result/Err") + } +} From 2724edaec9089d4c8bf64924ba3f1390c7fb77a6 Mon Sep 17 00:00:00 2001 From: Kelly Abuelsaad Date: Thu, 4 Jun 2026 17:10:04 -0400 Subject: [PATCH 2/3] fix(authbridge): Address PR #480 review feedback 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) Signed-off-by: Kelly Abuelsaad --- .../authlib/listener/forwardproxy/server.go | 228 ++++++++++-------- .../listener/forwardproxy/streaming_test.go | 30 ++- .../listener/internal/sseframe/reader.go | 52 ++-- .../authlib/listener/reverseproxy/server.go | 71 ++---- .../listener/reverseproxy/streaming_test.go | 28 ++- authbridge/authlib/pipeline/pipeline.go | 10 + authbridge/authlib/pipeline/plugin.go | 25 +- .../authlib/plugins/a2aparser/plugin.go | 38 ++- .../authlib/plugins/inferenceparser/plugin.go | 15 +- .../authlib/plugins/mcpparser/plugin.go | 102 ++++++-- 10 files changed, 364 insertions(+), 235 deletions(-) diff --git a/authbridge/authlib/listener/forwardproxy/server.go b/authbridge/authlib/listener/forwardproxy/server.go index e08b44e3d..1713c5127 100644 --- a/authbridge/authlib/listener/forwardproxy/server.go +++ b/authbridge/authlib/listener/forwardproxy/server.go @@ -12,6 +12,8 @@ import ( "log/slog" "net" "net/http" + "strings" + "sync" "time" "github.com/kagenti/kagenti-extensions/authbridge/authlib/auth" @@ -25,20 +27,6 @@ import ( const maxBodySize = 1 << 20 // 1MB — matches Envoy's default per_stream_buffer_limit_bytes -// responseHeaderTimeout bounds the time-to-first-byte (headers) for an -// upstream response. Replaces the old http.Client.Timeout, which -// covered the entire request lifecycle including the body read and so -// killed slow streaming responses (text/event-stream MCP tools/call, -// A2A message/stream, OpenAI streaming chat completions). With the -// timeout moved to the transport's ResponseHeaderTimeout, slow tools -// can stream as long as they like — only a wedged upstream that fails -// to send any headers within this window aborts. -// -// 30s preserves the prior behavior's time-to-headers ceiling. Tools -// that took longer to *return headers* never worked under the old -// configuration either. -const responseHeaderTimeout = 30 * time.Second - // streamReadIdleTimeout caps how long the proxy waits for the next // byte off a streaming response body. The time.Duration is applied // per ReadFrame iteration (see streamingResponseBody). A wedged @@ -96,10 +84,15 @@ func NewServer(outbound *pipeline.Holder, sessions *session.Store, mtls *MTLSOpt IdleConnTimeout: 90 * time.Second, TLSHandshakeTimeout: 10 * time.Second, ExpectContinueTimeout: 1 * time.Second, - // Time-to-headers ceiling for the upstream response. Replaces - // the old http.Client.Timeout (which covered the body read - // and broke streaming). See responseHeaderTimeout. - ResponseHeaderTimeout: responseHeaderTimeout, + // No ResponseHeaderTimeout: Streamable HTTP / MCP servers may + // hold response headers open until a slow tool completes, even + // when the eventual response is application/json (the server + // picks JSON vs SSE per call). A fixed time-to-headers ceiling + // reproduces the original 502 — the pre-headers wait is part + // of the same long tool execution we're trying to permit. The + // inbound request context + the per-Read idle timer on + // streaming bodies are the bounds; an unrecoverably wedged + // upstream is closed when the client cancels the request. } if mtls != nil { @@ -433,11 +426,10 @@ func isEventStream(contentType string) bool { return false } // Strip parameters: "text/event-stream; charset=utf-8" → "text/event-stream". - if idx := indexByteASCIICaseInsensitive(contentType, ';'); idx >= 0 { + if idx := strings.IndexByte(contentType, ';'); idx >= 0 { contentType = contentType[:idx] } - contentType = trimASCIISpace(contentType) - return equalASCIIFold(contentType, "text/event-stream") + return strings.EqualFold(strings.TrimSpace(contentType), "text/event-stream") } // handleStreamingResponse forwards a text/event-stream response to the @@ -464,6 +456,23 @@ func (s *Server) handleStreamingResponse(w http.ResponseWriter, r *http.Request, return } + // Defer the final last=true dispatch + session-event recording so + // every exit path (normal EOF, upstream read error, downstream + // client-write error) finalizes aggregating plugins and records + // the response event. Without this, a client disconnect mid-stream + // leaves inference/a2a stuck in an unfinalized state and emits no + // SessionResponse row to abctl. + defer func() { + finalAction := s.OutboundPipeline.RunResponseFrame(r.Context(), pctx, nil, true) + if finalAction.Type == pipeline.Reject { + // Headers already sent; we can't promote to 502, but + // surface the policy violation so operators see it. + slog.Warn("forward-proxy: streaming response rejected on finalization (headers already sent)", + "host", r.Host, "violation", finalAction.Violation) + } + s.recordOutboundResponseEvent(pctx, resp.StatusCode) + }() + // Forward headers and the streaming status code BEFORE the first // frame is written. Strip Content-Length since we'll be writing // chunked, and clear hop-by-hop headers as net/http would. @@ -506,40 +515,31 @@ func (s *Server) handleStreamingResponse(w http.ResponseWriter, r *http.Request, break } - // Write the frame back as a single SSE event. Reconstructing - // "data: ...\n\n" preserves the SSE wire shape regardless of - // whether the upstream used "data: " or "data:" or split across - // multiple data lines. Multi-line data is folded onto one line - // here, which is fine for JSON-RPC payloads (no embedded LF). - if _, err := w.Write([]byte("data: ")); err != nil { - slog.Debug("forward-proxy: streaming write error", "host", r.Host, "error", err) - return - } - if _, err := w.Write(frame); err != nil { - slog.Debug("forward-proxy: streaming write error", "host", r.Host, "error", err) - return - } - if _, err := w.Write([]byte("\n\n")); err != nil { - slog.Debug("forward-proxy: streaming write error", "host", r.Host, "error", err) - return + // Write the frame back as one or more SSE data lines. The + // sseframe reader folds multi-line `data:` events with `\n` + // separators per the spec; re-split here so each original line + // gets its own `data: ` prefix and the downstream parser sees + // the same event boundaries the upstream produced. For the + // single-line JSON-RPC payloads this targets, this loop is + // equivalent to writing `data: \n\n` once. + if !writeSSEFrame(w, frame) { + slog.Debug("forward-proxy: streaming write error", "host", r.Host) + break } flusher.Flush() bytesWritten += len(frame) } - - // Final last=true call so aggregating plugins (inference-parser, - // a2a-parser) finalize. Always called exactly once even on a zero- - // frame stream (per StreamingResponder contract). - s.OutboundPipeline.RunResponseFrame(r.Context(), pctx, nil, true) - - s.recordOutboundResponseEvent(pctx, resp.StatusCode) } // streamFallbackBuffered handles the rare case of a streaming // Content-Type response on a ResponseWriter that doesn't support -// http.Flusher — falls back to the original buffered path. The -// duplicated logic is small enough to keep here rather than refactor -// the main handler around the Flusher fork. +// http.Flusher — buffer the whole SSE body, then re-parse it through +// sseframe so streaming-aware plugins receive one OnResponseFrame call +// per SSE event followed by last=true. Without per-frame dispatch the +// inference parser (and any future fold-and-finalize plugin) would try +// to JSON-decode the whole SSE blob as one chunk, fail, and clobber a +// correctly-parsed completion. Production ResponseWriters implement +// http.Flusher so this path is mostly hit in tests. func (s *Server) streamFallbackBuffered(w http.ResponseWriter, r *http.Request, resp *http.Response, pctx *pipeline.Context) { respBody, err := io.ReadAll(io.LimitReader(resp.Body, maxBodySize+1)) if err != nil { @@ -561,7 +561,30 @@ func (s *Server) streamFallbackBuffered(w http.ResponseWriter, r *http.Request, return } if s.OutboundPipeline.HasStreamingResponders() { - _ = s.OutboundPipeline.RunResponseFrame(r.Context(), pctx, respBody, true) + // Re-parse the buffered SSE body frame-by-frame so plugins see the + // same per-event shape as the real streaming path. A Reject is + // honored here — headers are not yet on the wire. + reader := sseframe.NewReader(bytes.NewReader(respBody), maxBodySize) + for { + frame, ferr := reader.ReadFrame() + if ferr == io.EOF { + break + } + if ferr != nil { + slog.Warn("forward-proxy: streaming response read error in fallback", "host", r.Host, "error", ferr) + break + } + frameAction := s.OutboundPipeline.RunResponseFrame(r.Context(), pctx, frame, false) + if frameAction.Type == pipeline.Reject { + httpx.WriteRejection(w, frameAction) + return + } + } + finalAction := s.OutboundPipeline.RunResponseFrame(r.Context(), pctx, nil, true) + if finalAction.Type == pipeline.Reject { + httpx.WriteRejection(w, finalAction) + return + } } s.recordOutboundResponseEvent(pctx, resp.StatusCode) for key, values := range resp.Header { @@ -752,6 +775,41 @@ func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) { _ = upstream.Close() } +// writeSSEFrame writes one SSE event built from a sseframe-decoded +// frame back to w. The decoder folds multi-line `data:` events with +// `\n` separators; this helper splits on those `\n`s and emits one +// `data: \n` per original line followed by the blank-line +// terminator, so a downstream SSE parser sees the same event +// boundaries the upstream produced. Returns true when every byte +// was written; false on any write error so the caller can stop +// forwarding without re-checking each Write. +func writeSSEFrame(w io.Writer, frame []byte) bool { + for len(frame) > 0 { + nl := bytes.IndexByte(frame, '\n') + var line []byte + if nl < 0 { + line = frame + frame = nil + } else { + line = frame[:nl] + frame = frame[nl+1:] + } + if _, err := w.Write([]byte("data: ")); err != nil { + return false + } + if _, err := w.Write(line); err != nil { + return false + } + if _, err := w.Write([]byte("\n")); err != nil { + return false + } + } + if _, err := w.Write([]byte("\n")); err != nil { + return false + } + return true +} + // idleReader wraps r so each Read enforces an idle deadline. The // goroutine pattern (timer reset on every Read entry, cancelled on // every Read exit) is portable across any io.ReadCloser, unlike @@ -766,9 +824,21 @@ func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) { // a long-running tool that emits one byte every minute (within the // idle window) keeps the stream alive. The streamReadIdleTimeout // constant captures the wall-clock budget. +// +// Race-with-success note: time.AfterFunc + timer.Stop() does NOT +// wait for an already-fired callback. If the timer fires just as a +// Read returns successfully, the close runs after the success and +// would leave the next Read failing under a healthy upstream. The +// closeOnce field makes the close idempotent and Close() also runs +// it, so a stray late timer is harmless: the underlying body is +// closed at most once, and a successful in-flight Read keeps its +// data either way. The wider hazard — closing the body concurrently +// with an active Read — is the documented unblock mechanism the +// stdlib http transport relies on for forced disconnects. type idleReadCloser struct { - rc io.ReadCloser - timeout time.Duration + rc io.ReadCloser + timeout time.Duration + closeOnce sync.Once } func idleReader(rc io.ReadCloser, timeout time.Duration) io.ReadCloser { @@ -776,67 +846,21 @@ func idleReader(rc io.ReadCloser, timeout time.Duration) io.ReadCloser { } func (i *idleReadCloser) Read(p []byte) (int, error) { - timer := time.AfterFunc(i.timeout, func() { - // Closing the body unblocks Read with an error. Idempotent on - // http.Response.Body (subsequent Closes are no-ops). - _ = i.rc.Close() - }) + timer := time.AfterFunc(i.timeout, i.closeIdempotent) n, err := i.rc.Read(p) timer.Stop() return n, err } -func (i *idleReadCloser) Close() error { return i.rc.Close() } - -// indexByteASCIICaseInsensitive returns the index of the first -// occurrence of c in s (ASCII), or -1. Case-insensitive in spirit -// but c is treated as exact — the helper exists for ';' lookup in -// Content-Type headers, where the separator is unambiguous. -func indexByteASCIICaseInsensitive(s string, c byte) int { - for i := 0; i < len(s); i++ { - if s[i] == c { - return i - } - } - return -1 +func (i *idleReadCloser) Close() error { + i.closeIdempotent() + return nil } -// trimASCIISpace strips leading and trailing ASCII whitespace. -// Avoids the strings package's locale-aware Unicode trim because -// HTTP header tokens are ASCII-only. -func trimASCIISpace(s string) string { - start := 0 - for start < len(s) && (s[start] == ' ' || s[start] == '\t') { - start++ - } - end := len(s) - for end > start && (s[end-1] == ' ' || s[end-1] == '\t') { - end-- - } - return s[start:end] +func (i *idleReadCloser) closeIdempotent() { + i.closeOnce.Do(func() { _ = i.rc.Close() }) } -// equalASCIIFold reports whether s and t are equal under ASCII case -// folding. Used for Content-Type comparison without pulling in a -// Unicode-aware comparator. -func equalASCIIFold(s, t string) bool { - if len(s) != len(t) { - return false - } - for i := 0; i < len(s); i++ { - a, b := s[i], t[i] - if a >= 'A' && a <= 'Z' { - a += 'a' - 'A' - } - if b >= 'A' && b <= 'Z' { - b += 'a' - 'A' - } - if a != b { - return false - } - } - return true -} // enableKeepalive turns on TCP keepalive with a 30s probe interval on // the underlying *net.TCPConn, if conn unwraps to one. No-op on other diff --git a/authbridge/authlib/listener/forwardproxy/streaming_test.go b/authbridge/authlib/listener/forwardproxy/streaming_test.go index 98cfc2514..88c5ea9c6 100644 --- a/authbridge/authlib/listener/forwardproxy/streaming_test.go +++ b/authbridge/authlib/listener/forwardproxy/streaming_test.go @@ -20,12 +20,16 @@ import ( // observe how the proxy dispatches frames. It records every frame // it sees along with the last flag, so tests can assert on order // and finalization. ReadsBody=true so the listener takes the -// NeedsBody path on requests. +// NeedsBody path on requests. OnResponse calls are also counted so +// tests can verify the framework's "pick one path" contract for +// StreamingResponder plugins (OnResponse must NOT be called when +// OnResponseFrame is the dispatch path). type streamingProbe struct { - mu sync.Mutex - frames [][]byte - lasts []bool - caps pipeline.PluginCapabilities + mu sync.Mutex + frames [][]byte + lasts []bool + onResponseCalls int + caps pipeline.PluginCapabilities } func newStreamingProbe(writesBody bool) *streamingProbe { @@ -43,6 +47,9 @@ func (p *streamingProbe) OnRequest(_ context.Context, _ *pipeline.Context) pipel return pipeline.Action{Type: pipeline.Continue} } func (p *streamingProbe) OnResponse(_ context.Context, _ *pipeline.Context) pipeline.Action { + p.mu.Lock() + p.onResponseCalls++ + p.mu.Unlock() return pipeline.Action{Type: pipeline.Continue} } func (p *streamingProbe) OnResponseFrame(_ context.Context, _ *pipeline.Context, frame []byte, last bool) pipeline.Action { @@ -65,6 +72,12 @@ func (p *streamingProbe) snapshot() ([][]byte, []bool) { return frames, lasts } +func (p *streamingProbe) onResponseCount() int { + p.mu.Lock() + defer p.mu.Unlock() + return p.onResponseCalls +} + // TestForwardProxy_Streaming_FramesFlowThrough asserts that an upstream // emitting SSE frames over time has those frames flushed downstream as // they arrive — no full-body buffering, no 30s timeout firing, and @@ -240,6 +253,13 @@ func TestForwardProxy_BufferedDeliversLastTrueFrame(t *testing.T) { if !bytes.Equal(frames[0], []byte(`{"ok":true}`)) { t.Errorf("frame[0] = %q, want JSON body", frames[0]) } + // Regression: a StreamingResponder plugin must NOT have OnResponse + // called on the buffered application/json path — the framework + // picks one path and the buffered last=true frame is it. Otherwise + // migrated plugins would record every JSON response twice. + if got := probe.onResponseCount(); got != 0 { + t.Errorf("OnResponse called %d times; want 0 (StreamingResponder picks the OnResponseFrame path)", got) + } } // TestForwardProxy_Streaming_HeadersAndBodyArriveBeforeUpstreamCloses diff --git a/authbridge/authlib/listener/internal/sseframe/reader.go b/authbridge/authlib/listener/internal/sseframe/reader.go index 7d44623fe..e5bd5e0a2 100644 --- a/authbridge/authlib/listener/internal/sseframe/reader.go +++ b/authbridge/authlib/listener/internal/sseframe/reader.go @@ -145,45 +145,43 @@ func (r *Reader) ReadFrame() ([]byte, error) { } } -// readLine reads one logical SSE line (terminated by LF, CR, or CRLF) -// off the underlying buffered reader. The returned slice excludes the +// readLine reads one logical SSE line off the underlying buffered +// reader. Line terminators per the SSE spec are LF, CR, or CRLF — +// any of the three ends a line. The returned slice excludes the // terminator and is valid only until the next call. EOF without a // final newline is reported as io.EOF only when the line is empty; // otherwise the unterminated tail is returned. +// +// Per-line bytes are bounded by r.maxSize so a pathological non-data +// line (huge comment, malformed unknown field) cannot grow memory +// past the per-frame cap. func (r *Reader) readLine() ([]byte, error) { var line []byte for { - // Read up to and including the next LF. If the underlying - // reader returns an io.EOF with no data, propagate it; with - // data, treat the final unterminated chunk as a line. - chunk, err := r.br.ReadSlice('\n') - if len(chunk) > 0 { - line = append(line, chunk...) - if chunk[len(chunk)-1] == '\n' { - // Terminated. Strip CRLF or LF. - line = line[:len(line)-1] - if len(line) > 0 && line[len(line)-1] == '\r' { - line = line[:len(line)-1] - } - return line, nil - } - // ReadSlice returned a buffer-full chunk without LF; loop - // to keep accumulating into `line`. This handles SSE - // data lines longer than bufio's buffer. - if err == bufio.ErrBufferFull { - continue - } - } + b, err := r.br.ReadByte() if err != nil { if err == io.EOF && len(line) > 0 { - // Strip a possible trailing CR (no LF was seen). - if line[len(line)-1] == '\r' { - line = line[:len(line)-1] - } return line, nil } return nil, err } + switch b { + case '\n': + return line, nil + case '\r': + // Coalesce CRLF as a single terminator. If the next byte + // isn't LF, leave it for the next ReadByte; a bare CR is + // itself a valid line terminator. + if next, perr := r.br.Peek(1); perr == nil && len(next) == 1 && next[0] == '\n' { + _, _ = r.br.ReadByte() + } + return line, nil + default: + if len(line)+1 > r.maxSize { + return nil, ErrFrameTooLarge + } + line = append(line, b) + } } } diff --git a/authbridge/authlib/listener/reverseproxy/server.go b/authbridge/authlib/listener/reverseproxy/server.go index 9078d0546..78e9be08b 100644 --- a/authbridge/authlib/listener/reverseproxy/server.go +++ b/authbridge/authlib/listener/reverseproxy/server.go @@ -14,6 +14,7 @@ import ( "net/http" "net/http/httputil" "net/url" + "strings" "time" "github.com/kagenti/kagenti-extensions/authbridge/authlib/listener/httpx" @@ -601,11 +602,28 @@ func (b *streamingResponseBody) Read(p []byte) (int, error) { return 0, fmt.Errorf("reverseproxy: streaming response rejected mid-stream") } - // Re-frame as SSE: "data: \n\n". + // Re-frame as SSE. The decoder folds multi-line `data:` events + // with `\n` separators per spec; split here so each original line + // gets its own `data: ` prefix and the downstream parser sees the + // same event boundaries the upstream produced. For single-line + // JSON-RPC payloads this is equivalent to one `data: \n\n`. out := make([]byte, 0, len(frame)+8) - out = append(out, "data: "...) - out = append(out, frame...) - out = append(out, "\n\n"...) + rest := frame + for len(rest) > 0 { + nl := bytes.IndexByte(rest, '\n') + var line []byte + if nl < 0 { + line = rest + rest = nil + } else { + line = rest[:nl] + rest = rest[nl+1:] + } + out = append(out, "data: "...) + out = append(out, line...) + out = append(out, '\n') + } + out = append(out, '\n') n := copy(p, out) if n < len(out) { @@ -638,51 +656,10 @@ func isEventStream(contentType string) bool { if contentType == "" { return false } - if idx := indexByteASCII(contentType, ';'); idx >= 0 { + if idx := strings.IndexByte(contentType, ';'); idx >= 0 { contentType = contentType[:idx] } - contentType = trimASCIISpace(contentType) - return equalASCIIFold(contentType, "text/event-stream") -} - -func indexByteASCII(s string, c byte) int { - for i := 0; i < len(s); i++ { - if s[i] == c { - return i - } - } - return -1 -} - -func trimASCIISpace(s string) string { - start := 0 - for start < len(s) && (s[start] == ' ' || s[start] == '\t') { - start++ - } - end := len(s) - for end > start && (s[end-1] == ' ' || s[end-1] == '\t') { - end-- - } - return s[start:end] -} - -func equalASCIIFold(s, t string) bool { - if len(s) != len(t) { - return false - } - for i := 0; i < len(s); i++ { - a, b := s[i], t[i] - if a >= 'A' && a <= 'Z' { - a += 'a' - 'A' - } - if b >= 'A' && b <= 'Z' { - b += 'a' - 'A' - } - if a != b { - return false - } - } - return true + return strings.EqualFold(strings.TrimSpace(contentType), "text/event-stream") } // inboundSessionID returns the bucket ID for an inbound event. Mirrors diff --git a/authbridge/authlib/listener/reverseproxy/streaming_test.go b/authbridge/authlib/listener/reverseproxy/streaming_test.go index 24fb3eaf5..17b2fdddb 100644 --- a/authbridge/authlib/listener/reverseproxy/streaming_test.go +++ b/authbridge/authlib/listener/reverseproxy/streaming_test.go @@ -17,12 +17,15 @@ import ( // streamingProbe is a Plugin + StreamingResponder for asserting on // what the reverseproxy dispatches. ReadsBody=true so the listener -// takes the NeedsBody buffering path on requests. +// takes the NeedsBody buffering path on requests. OnResponse calls +// are counted so tests can verify the framework's "pick one path" +// contract for StreamingResponder plugins. type streamingProbe struct { - mu sync.Mutex - frames [][]byte - lasts []bool - caps pipeline.PluginCapabilities + mu sync.Mutex + frames [][]byte + lasts []bool + onResponseCalls int + caps pipeline.PluginCapabilities } func newStreamingProbe(writesBody bool) *streamingProbe { @@ -37,6 +40,9 @@ func (p *streamingProbe) OnRequest(_ context.Context, _ *pipeline.Context) pipel return pipeline.Action{Type: pipeline.Continue} } func (p *streamingProbe) OnResponse(_ context.Context, _ *pipeline.Context) pipeline.Action { + p.mu.Lock() + p.onResponseCalls++ + p.mu.Unlock() return pipeline.Action{Type: pipeline.Continue} } func (p *streamingProbe) OnResponseFrame(_ context.Context, _ *pipeline.Context, frame []byte, last bool) pipeline.Action { @@ -59,6 +65,12 @@ func (p *streamingProbe) snapshot() ([][]byte, []bool) { return frames, lasts } +func (p *streamingProbe) onResponseCount() int { + p.mu.Lock() + defer p.mu.Unlock() + return p.onResponseCalls +} + // TestReverseProxy_Streaming_FramesFlowThrough asserts the inbound // streaming path: an A2A message/stream-style SSE upstream is // forwarded frame-by-frame to the downstream client and each frame @@ -205,6 +217,12 @@ func TestReverseProxy_BufferedDeliversLastTrueFrame(t *testing.T) { if len(frames) != 1 || !lasts[0] { t.Fatalf("frames=%v lasts=%v; want single last=true call", framesAsStrings(frames), lasts) } + // Regression: a StreamingResponder plugin must NOT have OnResponse + // called on the buffered application/json path — the framework + // picks one path so the same body isn't recorded twice. + if got := probe.onResponseCount(); got != 0 { + t.Errorf("OnResponse called %d times; want 0 (StreamingResponder picks the OnResponseFrame path)", got) + } } func framesAsStrings(frames [][]byte) []string { diff --git a/authbridge/authlib/pipeline/pipeline.go b/authbridge/authlib/pipeline/pipeline.go index 40ed13dd9..f46cf6de4 100644 --- a/authbridge/authlib/pipeline/pipeline.go +++ b/authbridge/authlib/pipeline/pipeline.go @@ -133,6 +133,13 @@ func (p *Pipeline) Run(ctx context.Context, pctx *Context) Action { // RunResponse executes the response phase in reverse order. // The last plugin in the chain sees the response first. // +// Plugins implementing StreamingResponder are skipped here — the +// framework picks one path per the StreamingResponder contract: +// streaming-aware plugins receive a final OnResponseFrame(last=true) +// from the listener (single dispatch on the buffered application/json +// path; per-frame + last=true on the SSE path) instead of OnResponse, +// so the same body is never delivered through both hooks. +// // See Run for the pctx attribution stamping, the off-policy skip, and // the observe-policy shadow conversion. Same pattern, phase set to // InvocationPhaseResponse. @@ -142,6 +149,9 @@ func (p *Pipeline) RunResponse(ctx context.Context, pctx *Context) Action { if policy == ErrorPolicyOff { continue } + if _, ok := p.plugins[i].(StreamingResponder); ok { + continue + } if ctx.Err() != nil { slog.Info("pipeline: response cancelled", "plugin", p.plugins[i].Name()) return Deny("pipeline.cancelled", "request cancelled") diff --git a/authbridge/authlib/pipeline/plugin.go b/authbridge/authlib/pipeline/plugin.go index 48bee32a7..a00a1e5ac 100644 --- a/authbridge/authlib/pipeline/plugin.go +++ b/authbridge/authlib/pipeline/plugin.go @@ -192,20 +192,25 @@ type Finisher interface { // the duration of the call. Plugins that need to retain bytes // must copy. // - 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 +// dispatch for that frame and rejects the response. Listener +// ordering varies: forwardproxy invokes OnResponseFrame before +// emitting frame bytes, while reverseproxy emits frame bytes +// first (FlushInterval=-1 ferries each Read straight to the +// client) and then dispatches. Either way previously-emitted +// frames cannot be un-sent, so a mid-stream Reject results in +// a truncated stream rather than a 4xx/5xx response. Today no +// in-tree plugin returns Reject from this hook; the contract // leaves the door open for per-message enforcement to be added -// later (the listener would have to inspect-before-forward at -// that point). +// later (a listener doing enforcement would inspect-before-forward +// at that point). // - last=true is always called exactly once at end-of-stream, even // for an empty/zero-frame stream. Plugins finalize on last=true. // - For application/json responses the listener calls -// OnResponseFrame once with the full body and last=true; OnResponse -// is then NOT called for plugins that implement this interface -// (the framework picks one path). +// OnResponseFrame once with the full body and last=true. +// pipeline.RunResponse skips plugins implementing this interface +// so OnResponse is not called for them — the framework picks one +// path so a single response body is never delivered through both +// hooks. type StreamingResponder interface { OnResponseFrame(ctx context.Context, pctx *Context, frame []byte, last bool) Action } diff --git a/authbridge/authlib/plugins/a2aparser/plugin.go b/authbridge/authlib/plugins/a2aparser/plugin.go index d5c86e5b3..7319b14ce 100644 --- a/authbridge/authlib/plugins/a2aparser/plugin.go +++ b/authbridge/authlib/plugins/a2aparser/plugin.go @@ -104,11 +104,13 @@ func (p *A2AParser) OnRequest(_ context.Context, pctx *pipeline.Context) pipelin return pipeline.Action{Type: pipeline.Continue} } -// OnResponse extracts the server-assigned session/context ID and response summary -// from a buffered response body. Streaming-aware listeners take the -// OnResponseFrame path instead and this method becomes a no-op when -// the streaming path has already populated the extension's response -// fields. +// OnResponse is the legacy buffered-path response hook. Because this +// plugin implements StreamingResponder, pipeline.RunResponse skips it +// and OnResponseFrame is the dispatch path under all listeners — this +// method is unreachable from a normal listener. Kept for tests and +// hypothetical pipelines that call OnResponse directly without going +// through RunResponse, with a defensive guard against re-recording if +// the streaming path has already populated state. func (p *A2AParser) OnResponse(_ context.Context, pctx *pipeline.Context) pipeline.Action { if pctx.Extensions.A2A == nil { return pipeline.Action{Type: pipeline.Continue} @@ -189,6 +191,30 @@ func (p *A2AParser) OnResponseFrame(_ context.Context, pctx *pipeline.Context, f // logA2AFinalized emits the operator-facing debug log once a response // is finalized — shared by OnResponse and OnResponseFrame so the two // paths log identically. +// maxArtifactBytes caps the accumulated A2A artifact text recorded +// for observability. Long agent runs can stream many artifact-update +// events; without a cap, ext.Artifact would grow unbounded across +// frames and live on the response event for the life of the request. +// 64 KiB keeps the most recent text usable in abctl while bounding +// per-request memory. +const maxArtifactBytes = 64 * 1024 + +// appendCapped appends s to dst up to maxArtifactBytes; once the cap +// is reached, further appends are silently dropped (a single +// "…(truncated)" suffix is added the first time we hit the cap so the +// timeline shows truncation explicitly). +func appendCapped(dst, s string) string { + const truncMarker = "…(truncated)" + if len(dst) >= maxArtifactBytes { + return dst + } + remaining := maxArtifactBytes - len(dst) + if len(s) <= remaining { + return dst + s + } + return dst + s[:remaining] + truncMarker +} + func logA2AFinalized(ext *pipeline.A2AExtension) { slog.Debug("a2a-parser: response parsed", "sessionId", ext.SessionID, @@ -331,7 +357,7 @@ func extractStreamEvent(data []byte, ext *pipeline.A2AExtension) { case "artifact-update", "artifact": for _, part := range event.Result.Artifact.Parts { if part.Kind == "text" && part.Text != "" { - ext.Artifact += part.Text + ext.Artifact = appendCapped(ext.Artifact, part.Text) } } } diff --git a/authbridge/authlib/plugins/inferenceparser/plugin.go b/authbridge/authlib/plugins/inferenceparser/plugin.go index 8b59786d6..aee407971 100644 --- a/authbridge/authlib/plugins/inferenceparser/plugin.go +++ b/authbridge/authlib/plugins/inferenceparser/plugin.go @@ -97,18 +97,17 @@ func (p *InferenceParser) OnRequest(_ context.Context, pctx *pipeline.Context) p return pipeline.Action{Type: pipeline.Continue} } -// OnResponse populates the response-side fields (Completion, FinishReason, -// token counts) on pctx.Extensions.Inference for buffered responses. -// Streaming-aware listeners take the OnResponseFrame path instead and -// this method is a no-op when the streaming path has already -// finalized state. Aggregation across frames lives in OnResponseFrame -// + the inferenceStreamState scratch on the extension. +// OnResponse is the legacy buffered-path response hook. Because this +// plugin implements StreamingResponder, pipeline.RunResponse skips it +// and OnResponseFrame is the dispatch path under all listeners — this +// method is unreachable from a normal listener. Kept for tests and +// hypothetical pipelines that call OnResponse directly without going +// through RunResponse, with a defensive guard against re-recording if +// the streaming path has already populated state. func (p *InferenceParser) OnResponse(_ context.Context, pctx *pipeline.Context) pipeline.Action { if pctx.Extensions.Inference == nil { return pipeline.Action{Type: pipeline.Continue} } - // If the streaming path already finalized (Completion populated or - // FinishReason set), don't re-record on the buffered path. ext := pctx.Extensions.Inference if ext.Completion != "" || ext.FinishReason != "" || ext.TotalTokens > 0 { return pipeline.Action{Type: pipeline.Continue} diff --git a/authbridge/authlib/plugins/mcpparser/plugin.go b/authbridge/authlib/plugins/mcpparser/plugin.go index 8011465e1..445560ba5 100644 --- a/authbridge/authlib/plugins/mcpparser/plugin.go +++ b/authbridge/authlib/plugins/mcpparser/plugin.go @@ -211,26 +211,13 @@ func (p *MCPParser) OnRequest(_ context.Context, pctx *pipeline.Context) pipelin return pipeline.Action{Type: pipeline.Continue} } -// OnResponse is the buffered-path response hook. Streaming-aware -// listeners do NOT call it on text/event-stream responses — they -// dispatch through OnResponseFrame instead (see StreamingResponder). -// For application/json responses streaming-aware listeners deliver a -// single last=true frame to OnResponseFrame, so this method exists -// only to support listeners (or pipelines with no streaming-aware -// plugins) that still take the buffered RunResponse path. -// -// To avoid double-recording on streaming-aware listeners that call -// both RunResponse and RunResponseFrame for non-streaming responses, -// the per-frame path bypasses OnResponse via the StreamingResponder -// type assertion in pipeline.RunResponse — actually no, RunResponse -// runs all plugins regardless. So when running a streaming-aware -// pipeline, RunResponseFrame is the one that records; OnResponse -// here is the no-op for the same response. -// -// The compromise: OnResponseFrame populates pctx.Extensions.MCP.Result -// and records the Observe. OnResponse only acts when no streaming- -// aware path has fired, gated on whether the result/error has already -// been populated by OnResponseFrame. +// OnResponse is the legacy buffered-path response hook. Because this +// plugin implements StreamingResponder, pipeline.RunResponse skips it +// and OnResponseFrame is the dispatch path under all listeners — this +// method is unreachable from a normal listener. Kept for tests and +// hypothetical pipelines that call OnResponse directly without going +// through RunResponse, with a defensive guard against re-recording if +// the streaming path has already populated state. func (p *MCPParser) OnResponse(_ context.Context, pctx *pipeline.Context) pipeline.Action { // Stay silent when the request side never participated — the parser // recorded nothing on request, so recording on response would orphan @@ -269,6 +256,17 @@ func (p *MCPParser) OnResponse(_ context.Context, pctx *pipeline.Context) pipeli return pipeline.Action{Type: pipeline.Continue} } +// maxStreamObserves caps the number of per-frame Observe rows a single +// streaming response may emit. Long tools/call streams can produce +// dozens of result envelopes; without a bound, every envelope appends +// an Invocation row that lives on pctx.Extensions.Invocations for the +// life of the request — both noisy in the session timeline and a +// memory growth point. After the cap, additional frames update +// pctx.Extensions.MCP (so the latest result is still observable +// off-stream) but no Invocation row is appended; one final +// "_truncated" Observe at last=true tells operators what happened. +const maxStreamObserves = 50 + // OnResponseFrame is the streaming-aware response hook. Listeners // invoke it once per SSE frame (text/event-stream) and once with // last=true at end-of-stream. application/json responses arrive as @@ -277,22 +275,31 @@ func (p *MCPParser) OnResponse(_ context.Context, pctx *pipeline.Context) pipeli // Per-message recording: for MCP, each frame's payload is one // 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. +// surfaces partial results in the session timeline as they arrive — +// up to maxStreamObserves rows; beyond that the per-frame Observe is +// suppressed and a single _truncated row is emitted at end-of-stream. func (p *MCPParser) OnResponseFrame(_ context.Context, pctx *pipeline.Context, frame []byte, last bool) pipeline.Action { // Stay silent when the request side never participated. if pctx.Extensions.MCP == nil { return pipeline.Action{Type: pipeline.Continue} } + state := getOrCreateMCPStreamState(pctx) + // End-of-stream call with no payload. If we never saw a frame with // a result/error and the response body was empty, record a Skip // (matches the buffered path's "no_response_body" semantics so // abctl pairs request and response rows uniformly across shapes). - // For non-empty streams the per-frame Observe entries are the - // session timeline; nothing to do on last=true. + // If we observed too many frames, emit a single truncation row so + // operators see that records were dropped. if len(frame) == 0 { - if last && pctx.Extensions.MCP.Result == nil && pctx.Extensions.MCP.Err == nil { - pctx.Skip("no_response_body") + if last { + if pctx.Extensions.MCP.Result == nil && pctx.Extensions.MCP.Err == nil && state.observed == 0 { + pctx.Skip("no_response_body") + } + if state.truncated { + pctx.Observe("matched_" + pctx.Extensions.MCP.Method + "_response_truncated") + } } return pipeline.Action{Type: pipeline.Continue} } @@ -312,11 +319,56 @@ func (p *MCPParser) OnResponseFrame(_ context.Context, pctx *pipeline.Context, f return pipeline.Action{Type: pipeline.Continue} } + if state.observed >= maxStreamObserves { + // Past the cap: keep updating ext.MCP so the latest result is + // reflected, but don't append another Invocation row. + state.truncated = true + applyMCPResponseRPCNoObserve(pctx, rpc) + slog.Debug("mcp-parser: streaming frame (suppressed observe)", "method", pctx.Extensions.MCP.Method, "observed", state.observed) + return pipeline.Action{Type: pipeline.Continue} + } + applyMCPResponseRPC(pctx, rpc) + state.observed++ slog.Debug("mcp-parser: streaming frame", "method", pctx.Extensions.MCP.Method, "body", parsercommon.Truncate(string(frame), parsercommon.DebugBodyMax)) return pipeline.Action{Type: pipeline.Continue} } +// mcpStreamState holds per-stream scratch on pctx.Extensions.Custom so +// the per-frame Observe cap can be enforced across calls without +// adding fields to the public MCPExtension shape. +type mcpStreamState struct { + observed int + truncated bool +} + +const mcpStreamStateKey = "mcp-parser/stream-state" + +func getOrCreateMCPStreamState(pctx *pipeline.Context) *mcpStreamState { + if s := pipeline.GetState[mcpStreamState](pctx, mcpStreamStateKey); s != nil { + return s + } + s := &mcpStreamState{} + pipeline.SetState(pctx, mcpStreamStateKey, s) + return s +} + +// applyMCPResponseRPCNoObserve mutates pctx.Extensions.MCP without +// emitting an Observe row. Used past the per-stream Observe cap. +func applyMCPResponseRPCNoObserve(pctx *pipeline.Context, rpc jsonRPCResponse) { + if rpc.Error != nil { + pctx.Extensions.MCP.Err = &pipeline.MCPError{ + Code: rpc.Error.Code, + Message: rpc.Error.Message, + Data: rpc.Error.Data, + } + return + } + if rpc.Result != nil { + pctx.Extensions.MCP.Result = rpc.Result + } +} + // applyMCPResponseRPC mutates pctx.Extensions.MCP from a parsed // JSON-RPC response and emits the operator-facing log + Observe row. // Shared by OnResponse (buffered) and OnResponseFrame (streaming) so From 2fd8c0f651372b306dc519017e0402d1954a5bbf Mon Sep 17 00:00:00 2001 From: Kelly Abuelsaad Date: Thu, 4 Jun 2026 19:52:25 -0400 Subject: [PATCH 3/3] fix(authbridge): Dispatch buffered response through RunResponseFrame in extproc MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) Signed-off-by: Kelly Abuelsaad --- authbridge/authlib/listener/extproc/server.go | 77 ++++++++++++ .../authlib/listener/extproc/server_test.go | 111 ++++++++++++++++++ 2 files changed, 188 insertions(+) diff --git a/authbridge/authlib/listener/extproc/server.go b/authbridge/authlib/listener/extproc/server.go index 7ad533b25..d733d0e4c 100644 --- a/authbridge/authlib/listener/extproc/server.go +++ b/authbridge/authlib/listener/extproc/server.go @@ -4,8 +4,10 @@ package extproc import ( + "bytes" "context" "encoding/json" + "io" "log/slog" "net/http" "strconv" @@ -20,6 +22,7 @@ import ( "google.golang.org/grpc/status" "github.com/kagenti/kagenti-extensions/authbridge/authlib/auth" + "github.com/kagenti/kagenti-extensions/authbridge/authlib/listener/internal/sseframe" "github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline" "github.com/kagenti/kagenti-extensions/authbridge/authlib/session" ) @@ -559,6 +562,16 @@ func (s *Server) handleResponseHeaders(ctx context.Context, headers *corev3.Head return rejectFromAction(action) } + // Body-less response: deliver an empty last=true frame so + // StreamingResponder plugins can finalize (and emit no_response_body + // Skip rows for pairing). Mirrors the buffered-body path's single + // last=true dispatch. + if p.HasStreamingResponders() { + if frameAction := p.RunResponseFrame(ctx, pctx, nil, true); frameAction.Type == pipeline.Reject { + return rejectFromAction(frameAction) + } + } + // No body phase will run; record the response event here. A2A responses // need the body to extract contextId, so the rekey path is body-only; // skip it on this header-only path. @@ -596,6 +609,22 @@ func (s *Server) handleResponseBody(ctx context.Context, body []byte, pctx *pipe return rejectFromAction(action) } + // Streaming-aware plugins use a single code path for both shapes + // (mirrors forwardproxy/reverseproxy). pipeline.RunResponse skips + // StreamingResponder plugins so they wouldn't get a response-phase + // dispatch otherwise; deliver the buffered body via RunResponseFrame + // so mcp/inference/a2a parsers populate their response state and + // the inbound A2A contextId rekey below sees pctx.Extensions.A2A + // fully populated. For text/event-stream bodies (Envoy already + // buffered them at this point), re-parse with sseframe so each + // event arrives as its own frame; otherwise dispatch the whole + // body as one last=true frame. + if p.HasStreamingResponders() { + if frameAction := dispatchBufferedFrames(ctx, p, pctx); frameAction.Type == pipeline.Reject { + return rejectFromAction(frameAction) + } + } + // The server's response may carry the server-assigned A2A contextId. If // the request phase recorded events under DefaultSessionID (because the // client had no contextId yet), migrate them to the real ID so subsequent @@ -849,3 +878,51 @@ func getHeader(headers *corev3.HeaderMap, key string) string { } return "" } + +// dispatchBufferedFrames feeds the buffered response body to +// StreamingResponder plugins via RunResponseFrame, mirroring the +// proxy listeners' single-dispatch contract for buffered bodies. +// Envoy's ext_proc delivers response bodies pre-buffered (we requested +// ResponseBodyMode_BUFFERED via ModeOverride), so we get the whole +// body in one shot regardless of upstream framing. +// +// For application/json the entire body is one last=true frame, so +// non-streaming JSON-RPC responses look the same to plugins as on +// the proxy listeners. 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 — matches the per-message dispatch shape +// streaming-aware plugins expect. +func dispatchBufferedFrames(ctx context.Context, p *pipeline.Holder, pctx *pipeline.Context) pipeline.Action { + contentType := pctx.ResponseHeaders.Get("Content-Type") + if isEventStream(contentType) && len(pctx.ResponseBody) > 0 { + reader := sseframe.NewReader(bytes.NewReader(pctx.ResponseBody), maxBodySize) + for { + frame, err := reader.ReadFrame() + if err == io.EOF { + break + } + if err != nil { + slog.Warn("extproc: SSE re-parse error", "error", err) + break + } + if action := p.RunResponseFrame(ctx, pctx, frame, false); action.Type == pipeline.Reject { + return action + } + } + return p.RunResponseFrame(ctx, pctx, nil, true) + } + return p.RunResponseFrame(ctx, pctx, pctx.ResponseBody, true) +} + +// 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 { + if contentType == "" { + return false + } + if idx := strings.IndexByte(contentType, ';'); idx >= 0 { + contentType = contentType[:idx] + } + return strings.EqualFold(strings.TrimSpace(contentType), "text/event-stream") +} diff --git a/authbridge/authlib/listener/extproc/server_test.go b/authbridge/authlib/listener/extproc/server_test.go index 4d335822a..61feda939 100644 --- a/authbridge/authlib/listener/extproc/server_test.go +++ b/authbridge/authlib/listener/extproc/server_test.go @@ -1570,3 +1570,114 @@ func TestExtProc_PopulatesSchemeFromPseudoHeader(t *testing.T) { }) } } + +// streamingRecorderPlugin is a Plugin + StreamingResponder used to +// verify the extproc listener dispatches the buffered response body +// through OnResponseFrame. The pipeline.RunResponse skip introduced +// for the proxy listeners would otherwise leave streaming-aware +// plugins (mcp/inference/a2a parsers) with no response-phase +// dispatch under envoy-sidecar — the regression this guards. +type streamingRecorderPlugin struct { + frames [][]byte + lasts []bool + onResponseCalls int +} + +func (p *streamingRecorderPlugin) Name() string { return "streaming-recorder" } +func (p *streamingRecorderPlugin) Capabilities() pipeline.PluginCapabilities { + return pipeline.PluginCapabilities{ReadsBody: true} +} +func (p *streamingRecorderPlugin) OnRequest(_ context.Context, _ *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} +func (p *streamingRecorderPlugin) OnResponse(_ context.Context, _ *pipeline.Context) pipeline.Action { + p.onResponseCalls++ + return pipeline.Action{Type: pipeline.Continue} +} +func (p *streamingRecorderPlugin) OnResponseFrame(_ context.Context, _ *pipeline.Context, frame []byte, last bool) pipeline.Action { + cp := make([]byte, len(frame)) + copy(cp, frame) + p.frames = append(p.frames, cp) + p.lasts = append(p.lasts, last) + return pipeline.Action{Type: pipeline.Continue} +} + +// TestExtProc_BufferedJSONResponse_DispatchesToStreamingResponder +// asserts the framework's "pick one path" contract on extproc: +// pipeline.RunResponse skips StreamingResponder plugins, so the +// listener MUST deliver the buffered application/json body via +// RunResponseFrame(last=true) — otherwise mcp/inference/a2a parsers +// lose response observability under envoy-sidecar entirely. +func TestExtProc_BufferedJSONResponse_DispatchesToStreamingResponder(t *testing.T) { + probe := &streamingRecorderPlugin{} + p, err := pipeline.New([]pipeline.Plugin{probe}) + if err != nil { + t.Fatal(err) + } + srv := &Server{InboundPipeline: pipeline.NewHolder(p), OutboundPipeline: pipeline.NewHolder(p)} + + pctx := &pipeline.Context{ + Direction: pipeline.Inbound, + ResponseHeaders: http.Header{ + "Content-Type": []string{"application/json"}, + }, + } + body := []byte(`{"jsonrpc":"2.0","id":1,"result":{"ok":true}}`) + resp := srv.handleResponseBody(context.Background(), body, pctx, "inbound") + if resp.GetImmediateResponse() != nil { + t.Fatalf("unexpected immediate response: %+v", resp.GetImmediateResponse()) + } + + if len(probe.frames) != 1 || !probe.lasts[0] { + t.Fatalf("frames=%d lasts=%v; want exactly one last=true frame", len(probe.frames), probe.lasts) + } + if string(probe.frames[0]) != string(body) { + t.Errorf("frame[0] = %q; want %q", probe.frames[0], body) + } + // OnResponse must NOT fire for a StreamingResponder plugin — + // pipeline.RunResponse skips it on purpose so the body isn't + // delivered through both hooks. + if probe.onResponseCalls != 0 { + t.Errorf("OnResponse called %d times; want 0", probe.onResponseCalls) + } +} + +// TestExtProc_BufferedSSEResponse_DispatchesPerEvent verifies that +// when Envoy buffers a text/event-stream body, the extproc listener +// re-parses with sseframe and dispatches one OnResponseFrame call +// per SSE event followed by a final last=true. Mirrors what the +// proxy listeners do over their wire-streaming dispatch path so +// streaming-aware plugins see the same shape regardless of mode. +func TestExtProc_BufferedSSEResponse_DispatchesPerEvent(t *testing.T) { + probe := &streamingRecorderPlugin{} + p, err := pipeline.New([]pipeline.Plugin{probe}) + if err != nil { + t.Fatal(err) + } + srv := &Server{InboundPipeline: pipeline.NewHolder(p), OutboundPipeline: pipeline.NewHolder(p)} + + pctx := &pipeline.Context{ + Direction: pipeline.Inbound, + ResponseHeaders: http.Header{ + "Content-Type": []string{"text/event-stream"}, + }, + } + body := []byte("data: {\"id\":1}\n\ndata: {\"id\":2}\n\ndata: {\"id\":3}\n\n") + resp := srv.handleResponseBody(context.Background(), body, pctx, "inbound") + if resp.GetImmediateResponse() != nil { + t.Fatalf("unexpected immediate response: %+v", resp.GetImmediateResponse()) + } + + // Three events as non-last + one final last=true call. + if len(probe.frames) != 4 { + t.Fatalf("got %d frames, want 4 (3 events + 1 last=true)", len(probe.frames)) + } + for i := 0; i < 3; i++ { + if probe.lasts[i] { + t.Errorf("frame[%d] last=true, want false", i) + } + } + if !probe.lasts[3] { + t.Errorf("frame[3] last=false, want true (final)") + } +}