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)")
+ }
+}
diff --git a/authbridge/authlib/listener/forwardproxy/server.go b/authbridge/authlib/listener/forwardproxy/server.go
index 0933cd72b..1713c5127 100644
--- a/authbridge/authlib/listener/forwardproxy/server.go
+++ b/authbridge/authlib/listener/forwardproxy/server.go
@@ -12,10 +12,13 @@ import (
"log/slog"
"net"
"net/http"
+ "strings"
+ "sync"
"time"
"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 +27,17 @@ import (
const maxBodySize = 1 << 20 // 1MB — matches Envoy's default per_stream_buffer_limit_bytes
+// 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 +84,15 @@ func NewServer(outbound *pipeline.Holder, sessions *session.Store, mtls *MTLSOpt
IdleConnTimeout: 90 * time.Second,
TLSHandshakeTimeout: 10 * time.Second,
ExpectContinueTimeout: 1 * time.Second,
+ // 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 {
@@ -87,7 +110,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 +303,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 +344,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 +369,224 @@ 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 := strings.IndexByte(contentType, ';'); idx >= 0 {
+ contentType = contentType[:idx]
+ }
+ return strings.EqualFold(strings.TrimSpace(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
+ }
+
+ // 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)
}
- // 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)
+ 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.
+ for key, values := range resp.Header {
+ for _, value := range values {
+ w.Header().Add(key, value)
}
}
+ 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 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)
+ }
+}
+
+// streamFallbackBuffered handles the rare case of a streaming
+// Content-Type response on a ResponseWriter that doesn't support
+// 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 {
+ 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() {
+ // 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 {
for _, value := range values {
w.Header().Add(key, value)
@@ -525,6 +775,93 @@ 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
+// 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.
+//
+// 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
+ closeOnce sync.Once
+}
+
+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, i.closeIdempotent)
+ n, err := i.rc.Read(p)
+ timer.Stop()
+ return n, err
+}
+
+func (i *idleReadCloser) Close() error {
+ i.closeIdempotent()
+ return nil
+}
+
+func (i *idleReadCloser) closeIdempotent() {
+ i.closeOnce.Do(func() { _ = i.rc.Close() })
+}
+
+
// 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..88c5ea9c6
--- /dev/null
+++ b/authbridge/authlib/listener/forwardproxy/streaming_test.go
@@ -0,0 +1,341 @@
+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. 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
+ onResponseCalls int
+ 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 {
+ 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 {
+ 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
+}
+
+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
+// 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])
+ }
+ // 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
+// 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..e5bd5e0a2
--- /dev/null
+++ b/authbridge/authlib/listener/internal/sseframe/reader.go
@@ -0,0 +1,222 @@
+// 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 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 {
+ b, err := r.br.ReadByte()
+ if err != nil {
+ if err == io.EOF && len(line) > 0 {
+ 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)
+ }
+ }
+}
+
+// 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..78e9be08b 100644
--- a/authbridge/authlib/listener/reverseproxy/server.go
+++ b/authbridge/authlib/listener/reverseproxy/server.go
@@ -14,9 +14,11 @@ import (
"net/http"
"net/http/httputil"
"net/url"
+ "strings"
"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 +93,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 +288,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 +333,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 +476,192 @@ 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. 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)
+ 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) {
+ 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 := strings.IndexByte(contentType, ';'); idx >= 0 {
+ contentType = contentType[:idx]
+ }
+ return strings.EqualFold(strings.TrimSpace(contentType), "text/event-stream")
+}
+
// 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..17b2fdddb
--- /dev/null
+++ b/authbridge/authlib/listener/reverseproxy/streaming_test.go
@@ -0,0 +1,234 @@
+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. 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
+ onResponseCalls int
+ 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 {
+ 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 {
+ 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
+}
+
+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
+// 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)
+ }
+ // 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 {
+ 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..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")
@@ -163,6 +173,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..a00a1e5ac 100644
--- a/authbridge/authlib/pipeline/plugin.go
+++ b/authbridge/authlib/pipeline/plugin.go
@@ -166,6 +166,55 @@ 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. 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 (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.
+// 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
+}
+
// 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..7319b14ce 100644
--- a/authbridge/authlib/plugins/a2aparser/plugin.go
+++ b/authbridge/authlib/plugins/a2aparser/plugin.go
@@ -104,23 +104,23 @@ 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 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).
+// 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 {
- // 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 +128,100 @@ 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.
+// 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", 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 +288,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 +301,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 = appendCapped(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..aee407971 100644
--- a/authbridge/authlib/plugins/inferenceparser/plugin.go
+++ b/authbridge/authlib/plugins/inferenceparser/plugin.go
@@ -97,30 +97,138 @@ 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. Handles both non-streaming
-// JSON responses and SSE streams from OpenAI-compatible servers.
+// 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 {
- // 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.
+ 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 +236,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..445560ba5 100644
--- a/authbridge/authlib/plugins/mcpparser/plugin.go
+++ b/authbridge/authlib/plugins/mcpparser/plugin.go
@@ -211,6 +211,13 @@ func (p *MCPParser) OnRequest(_ context.Context, pctx *pipeline.Context) pipelin
return pipeline.Action{Type: pipeline.Continue}
}
+// 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
@@ -218,6 +225,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 +251,129 @@ 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}
+}
+
+// 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
+// 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 —
+// 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).
+ // If we observed too many frames, emit a single truncation row so
+ // operators see that records were dropped.
+ if len(frame) == 0 {
+ 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}
+ }
+
+ 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}
+ }
+
+ 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
+// 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 +382,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")
+ }
+}