diff --git a/authbridge/authlib/listener/forwardproxy/mcp_sse_stream_test.go b/authbridge/authlib/listener/forwardproxy/mcp_sse_stream_test.go new file mode 100644 index 000000000..a68da382f --- /dev/null +++ b/authbridge/authlib/listener/forwardproxy/mcp_sse_stream_test.go @@ -0,0 +1,325 @@ +package forwardproxy + +import ( + "bufio" + "bytes" + "context" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline" + "github.com/kagenti/kagenti-extensions/authbridge/authlib/session" +) + +// TestForwardProxy_SSE_StreamsWithoutResponder is the regression test for +// issue #642: a generic (e.g. MCP Streamable HTTP) upstream returns +// text/event-stream, but the outbound pipeline has NO StreamingResponder +// plugin. Before the fix, such a response fell through to an unflushed +// io.Copy, so intermittent SSE events never reached the client until the +// upstream closed the connection — the agent timed out. +// +// Unlike TestForwardProxy_MCP_SSEResponse_RecordsObserve (which uses an +// upstream that writes one frame then RETURNS, so io.Copy sees EOF +// immediately and never exposes the buffering), this upstream flushes one +// event and then holds the connection OPEN. A buffering proxy delivers +// nothing until release; a flushing proxy delivers the event at once. +func TestForwardProxy_SSE_StreamsWithoutResponder(t *testing.T) { + release := make(chan struct{}) + closed := false + defer func() { + if !closed { + close(release) + } + }() + + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + f, ok := w.(http.Flusher) + if !ok { + t.Error("upstream ResponseWriter lacks http.Flusher") + return + } + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + // event: + id: exercise byte-faithful framing — the re-framing path + // (handleStreamingResponse) would drop a non-allowlisted event name + // and the id: line entirely. + io.WriteString(w, "event: message\nid: 42\ndata: {\"jsonrpc\":\"2.0\",\"id\":2,\"result\":{\"ok\":true}}\n\n") + f.Flush() + <-release // hold the stream open, like a live MCP server between events + })) + defer upstream.Close() + + store := session.New(5*time.Minute, 100, 0) + defer store.Close() + + // Empty pipeline: HasStreamingResponders()==false and WritesBody()==false, + // so serveOutbound routes to streamPassthrough — the reporter's plain-proxy + // shape (their only outbound plugin, token-exchange, is likewise not a + // StreamingResponder). + pipe, err := pipeline.New(nil) + if err != nil { + t.Fatalf("pipeline.New: %v", err) + } + srv, err := NewServer(pipeline.NewHolder(pipe), store, nil) + if err != nil { + t.Fatalf("NewServer: %v", err) + } + proxy := httptest.NewServer(srv.Handler()) + defer proxy.Close() + + req, _ := http.NewRequest("POST", upstream.URL+"/mcp", bytes.NewReader([]byte(`{"jsonrpc":"2.0","id":2,"method":"tools/list"}`))) + req.Header.Set("Content-Type", "application/json") + proxyClient := &http.Client{Transport: &http.Transport{Proxy: http.ProxyURL(mustParseURL(proxy.URL))}} + resp, err := proxyClient.Do(req) + if err != nil { + t.Fatalf("request failed: %v", err) + } + defer resp.Body.Close() + + if ct := resp.Header.Get("Content-Type"); ct != "text/event-stream" { + t.Errorf("Content-Type = %q, want text/event-stream", ct) + } + + // Read the first full SSE frame in a goroutine, WHILE the upstream is still + // blocked on <-release. A buffering proxy delivers nothing until release, so + // this read blocks and the select hits the timeout — that is the #642 + // regression. + type frameResult struct { + data []byte + err error + } + got := make(chan frameResult, 1) + go func() { + br := bufio.NewReader(resp.Body) + var acc []byte + for { + line, err := br.ReadBytes('\n') + acc = append(acc, line...) + if err != nil { + got <- frameResult{acc, err} + return + } + if bytes.HasSuffix(acc, []byte("\n\n")) { + got <- frameResult{acc, nil} + return + } + } + }() + + select { + case fr := <-got: + if fr.err != nil && fr.err != io.EOF { + t.Fatalf("reading first SSE frame: %v", fr.err) + } + frame := string(fr.data) + // Byte-faithful framing: event: and id: must survive verbatim. + if !strings.Contains(frame, "event: message") { + t.Errorf("first frame missing 'event: message' (framing not preserved): %q", frame) + } + if !strings.Contains(frame, "id: 42") { + t.Errorf("first frame missing 'id: 42' (framing not preserved): %q", frame) + } + if !strings.Contains(frame, `data: {"jsonrpc":"2.0"`) { + t.Errorf("first frame missing/garbled data line: %q", frame) + } + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for the first SSE event while upstream held the connection open — proxy buffered the stream (regression of #642)") + } + + // Let the upstream handler exit cleanly. + close(release) + closed = true +} + +// responseProbePlugin is a minimal non-StreamingResponder plugin whose +// OnResponse uses only status/headers (like opa / litellm-budgettrack). It +// records that it ran and can deny — used to prove streamPassthrough still runs +// the response-phase pipeline before streaming. When readsBody is set it also +// declares ReadsBody, so it exercises the ReadsBody-but-not-StreamingResponder +// footgun (see TestForwardProxy_SSE_ReadsBodyPluginWarnsAndStreams). +type responseProbePlugin struct { + deny bool + readsBody bool + ranResp atomic.Bool + // sawBodyLen is the length of pctx.ResponseBody observed in OnResponse. + sawBodyLen atomic.Int64 +} + +func (p *responseProbePlugin) Name() string { return "response-probe" } +func (p *responseProbePlugin) Capabilities() pipeline.PluginCapabilities { + // Never a StreamingResponder → streamPassthrough path. ReadsBody is opt-in + // so most callers stay on the status/header-only shape. + return pipeline.PluginCapabilities{ReadsBody: p.readsBody} +} +func (p *responseProbePlugin) OnRequest(context.Context, *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} +func (p *responseProbePlugin) OnResponse(_ context.Context, pctx *pipeline.Context) pipeline.Action { + p.ranResp.Store(true) + p.sawBodyLen.Store(int64(len(pctx.ResponseBody))) + if p.deny { + return pipeline.DenyStatus(403, "test.denied", "denied by response probe") + } + return pipeline.Action{Type: pipeline.Continue} +} + +// syncBuffer is a goroutine-safe io.Writer for capturing slog output written +// from the proxy's request-handler goroutine while the test reads it. +type syncBuffer struct { + mu sync.Mutex + buf bytes.Buffer +} + +func (b *syncBuffer) Write(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.Write(p) +} + +func (b *syncBuffer) String() string { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.String() +} + +// sseProxy wires an SSE upstream + a forward proxy with the given pipeline and +// returns a client bound to the proxy. The upstream writes one event and +// returns (closes) — enough to exercise the response-phase decision. +func sseProxy(t *testing.T, plugins []pipeline.Plugin) (*http.Client, string) { + t.Helper() + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + io.WriteString(w, "event: message\nid: 7\ndata: {\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{}}\n\n") + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + })) + t.Cleanup(upstream.Close) + + store := session.New(5*time.Minute, 100, 0) + t.Cleanup(store.Close) + pipe, err := pipeline.New(plugins) + if err != nil { + t.Fatalf("pipeline.New: %v", err) + } + srv, err := NewServer(pipeline.NewHolder(pipe), store, nil) + if err != nil { + t.Fatalf("NewServer: %v", err) + } + proxy := httptest.NewServer(srv.Handler()) + t.Cleanup(proxy.Close) + + client := &http.Client{Transport: &http.Transport{Proxy: http.ProxyURL(mustParseURL(proxy.URL))}} + return client, upstream.URL + "/mcp" +} + +// TestForwardProxy_SSE_ResponsePhaseDenyShortCircuits proves the fix keeps +// response-phase enforcement on streamed responses: a plugin that denies in +// OnResponse must short-circuit BEFORE any SSE byte is written. +func TestForwardProxy_SSE_ResponsePhaseDenyShortCircuits(t *testing.T) { + probe := &responseProbePlugin{deny: true} + client, url := sseProxy(t, []pipeline.Plugin{probe}) + + req, _ := http.NewRequest("POST", url, bytes.NewReader([]byte(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`))) + req.Header.Set("Content-Type", "application/json") + resp, err := client.Do(req) + if err != nil { + t.Fatalf("request failed: %v", err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + + if !probe.ranResp.Load() { + t.Error("OnResponse did not run on the SSE response (RunResponse skipped)") + } + if resp.StatusCode != http.StatusForbidden { + t.Errorf("status = %d, want 403 (response-phase deny not honored)", resp.StatusCode) + } + if strings.Contains(string(body), "event: message") { + t.Errorf("SSE body was forwarded despite the deny: %q", body) + } +} + +// TestForwardProxy_SSE_HeaderOnResponseRuns proves a non-denying, header-level +// OnResponse (e.g. cost accounting) still fires on an SSE response and the +// stream is delivered. +func TestForwardProxy_SSE_HeaderOnResponseRuns(t *testing.T) { + probe := &responseProbePlugin{deny: false} + client, url := sseProxy(t, []pipeline.Plugin{probe}) + + req, _ := http.NewRequest("POST", url, bytes.NewReader([]byte(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`))) + req.Header.Set("Content-Type", "application/json") + resp, err := client.Do(req) + if err != nil { + t.Fatalf("request failed: %v", err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + + if !probe.ranResp.Load() { + t.Error("OnResponse did not run on the SSE response") + } + if resp.StatusCode != http.StatusOK { + t.Errorf("status = %d, want 200", resp.StatusCode) + } + if !strings.Contains(string(body), "event: message") || !strings.Contains(string(body), "id: 7") { + t.Errorf("SSE body not delivered verbatim: %q", body) + } +} + +// TestForwardProxy_SSE_ReadsBodyPluginWarnsAndStreams covers the ReadsBody +// footgun raised in review: a plugin that declares ReadsBody but is NOT a +// StreamingResponder falls into streamPassthrough. The proxy must NOT buffer to +// feed it a body (that would reintroduce the #642 timeout on a live stream), so +// the plugin's OnResponse sees an empty pctx.ResponseBody. The stream must still +// be delivered verbatim, and the misconfiguration must be surfaced via a warning +// rather than silently starving the plugin of the body. +func TestForwardProxy_SSE_ReadsBodyPluginWarnsAndStreams(t *testing.T) { + // Capture slog output. The warning is emitted from the proxy's handler + // goroutine, so use a mutex-guarded buffer to stay race-clean; tests in this + // package don't run in parallel, so swapping the default logger is safe. + logBuf := &syncBuffer{} + prev := slog.Default() + slog.SetDefault(slog.New(slog.NewTextHandler(logBuf, &slog.HandlerOptions{Level: slog.LevelWarn}))) + t.Cleanup(func() { slog.SetDefault(prev) }) + + probe := &responseProbePlugin{readsBody: true} + client, url := sseProxy(t, []pipeline.Plugin{probe}) + + req, _ := http.NewRequest("POST", url, bytes.NewReader([]byte(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`))) + req.Header.Set("Content-Type", "application/json") + resp, err := client.Do(req) + if err != nil { + t.Fatalf("request failed: %v", err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + + if !probe.ranResp.Load() { + t.Error("OnResponse did not run on the SSE response") + } + // The plugin declared ReadsBody, but the stream is passed through unbuffered, + // so it sees no body — the documented limitation this warning surfaces. + if n := probe.sawBodyLen.Load(); n != 0 { + t.Errorf("ReadsBody plugin saw %d body bytes, want 0 (stream is not buffered for it)", n) + } + // The stream must still be delivered verbatim — no regression to the #642 timeout. + if resp.StatusCode != http.StatusOK { + t.Errorf("status = %d, want 200", resp.StatusCode) + } + if !strings.Contains(string(body), "event: message") || !strings.Contains(string(body), "id: 7") { + t.Errorf("SSE body not delivered verbatim: %q", body) + } + // The misconfiguration must be logged, not silent. + if got := logBuf.String(); !strings.Contains(got, "ReadsBody plugin that is not a StreamingResponder") { + t.Errorf("expected a warning about the ReadsBody+SSE misconfiguration, got logs: %q", got) + } +} diff --git a/authbridge/authlib/listener/forwardproxy/server.go b/authbridge/authlib/listener/forwardproxy/server.go index 5615133c1..8d2a6472a 100644 --- a/authbridge/authlib/listener/forwardproxy/server.go +++ b/authbridge/authlib/listener/forwardproxy/server.go @@ -380,14 +380,38 @@ func (s *Server) serveOutbound(w http.ResponseWriter, r *http.Request, isBridge // 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 isEventStream(resp.Header.Get("Content-Type")) && resp.Body != nil { if s.OutboundPipeline.WritesBody() { + // A body mutator needs the whole body to rewrite it, so it + // can't stream — fall back to the buffered path with a warning. slog.Warn("forward-proxy: text/event-stream response with WritesBody plugin — falling back to buffered path", "host", r.Host) - } else { + } else if s.OutboundPipeline.HasStreamingResponders() { + // Streaming-aware plugins (inference-parser, a2a-parser) parse + // each SSE frame; handleStreamingResponse re-frames via sseframe. s.handleStreamingResponse(w, r, resp, pctx) return + } else { + // No streaming responder: relay the SSE stream byte-for-byte + // with per-write flushing. Re-framing (handleStreamingResponse) + // would drop the event:/id:/retry: lines that generic SSE + // clients (e.g. an MCP Streamable HTTP client) depend on. Fixes #642. + // + // A plugin that declares ReadsBody (but not WritesBody, and is + // not a StreamingResponder) also lands here, and its OnResponse + // runs against an empty pctx.ResponseBody: streamPassthrough + // forwards the stream without buffering it. We deliberately don't + // buffer to satisfy such a plugin — that would reintroduce the + // #642 timeout on a live stream. A plugin that must inspect a + // streamed body should implement StreamingResponder. Warn + // (mirroring the WritesBody fallback above) so the + // misconfiguration surfaces instead of the plugin silently seeing + // no body. WritesBody is already false in this branch, so + // NeedsBody() here implies ReadsBody. + if s.OutboundPipeline.NeedsBody() { + slog.Warn("forward-proxy: text/event-stream response with a ReadsBody plugin that is not a StreamingResponder — streaming byte-for-byte; its OnResponse will see an empty body (implement StreamingResponder to inspect a streamed body)", "host", r.Host) + } + s.streamPassthrough(w, r, resp, pctx) + return } } @@ -651,6 +675,84 @@ func (s *Server) handleStreamingResponse(w http.ResponseWriter, r *http.Request, } } +// streamPassthrough forwards a text/event-stream response to the downstream +// client byte-for-byte with per-write flushing. It is the streaming path when +// no StreamingResponder plugin is configured (a plain proxy pipeline). Unlike +// handleStreamingResponse it does NOT parse or re-frame the stream through +// sseframe — it relays the exact upstream bytes so that event:, id:, retry:, +// and comment lines survive, which generic SSE consumers such as an MCP +// Streamable HTTP client require. Without this path such a response fell +// through to an unflushed io.Copy and never reached the client until the +// upstream closed the connection (issue #642). +// +// Unlike handleStreamingResponse, the response-phase pipeline (RunResponse) IS +// run before the first byte. That is safe only because this path is reached +// exclusively when no StreamingResponder is configured, so RunResponse cannot +// double-dispatch a plugin that also handles OnResponseFrame. Running it lets +// header/status-based response gates fire on streamed responses too — e.g. +// opa's response-phase deny (status + headers) and litellm-budgettrack's cost +// accounting (a response header) — and a deny is honored before any byte is +// written. The plugins reachable here do not read pctx.ResponseBody in +// OnResponse, so leaving the body unbuffered is fine; body-level response +// inspection on a stream requires implementing StreamingResponder. +func (s *Server) streamPassthrough(w http.ResponseWriter, r *http.Request, resp *http.Response, pctx *pipeline.Context) { + flusher, ok := w.(http.Flusher) + if !ok { + // No incremental delivery possible on this ResponseWriter — buffer. + // http.Flusher is implemented by net/http's default writer; this is a + // defensive guard for test recorders and exotic wrappers. + 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 + } + + // Run the response-phase pipeline before the first byte so header/status + // gates still fire on a streamed response and a deny short-circuits before + // anything is written. streamFallbackBuffered runs its own RunResponse, so + // this is done only on the flushing path to avoid double-dispatch. + if respAction := s.OutboundPipeline.RunResponse(r.Context(), pctx); respAction.Type == pipeline.Reject { + httpx.WriteRejection(w, respAction) + return + } + + // Record the response event on every exit path (normal EOF, upstream read + // error, downstream write error) so a SessionResponse row still lands. + defer s.recordOutboundResponseEvent(pctx, resp.StatusCode) + + // Forward headers + status before the first byte. Drop Content-Length since + // we relay an open-ended chunked stream. + 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() + + // Copy raw chunks and flush each so intermittent SSE events reach the client + // immediately. idleReader bounds a wedged upstream; total size stays + // unbounded so long-lived streams aren't cut off. + body := idleReader(resp.Body, streamReadIdleTimeout) + buf := make([]byte, 32*1024) + for { + n, readErr := body.Read(buf) + if n > 0 { + if _, writeErr := w.Write(buf[:n]); writeErr != nil { + slog.Debug("forward-proxy: streaming write error", "host", r.Host, "error", writeErr) + return + } + flusher.Flush() + } + if readErr != nil { + if readErr != io.EOF { + slog.Warn("forward-proxy: streaming response read error", "host", r.Host, "error", readErr) + } + return + } + } +} + // 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