Fix: forward-proxy streams SSE responses without a StreamingResponder

The forward proxy only took its flushing SSE path (handleStreamingResponse)
when a StreamingResponder plugin was configured. A plain pipeline (e.g. just
token-exchange) fell through to an unflushed io.Copy, so intermittent SSE
events from an MCP server never reached the agent until the upstream closed
the connection — the agent timed out.

Add streamPassthrough: a byte-faithful flushing relay used when no
StreamingResponder is present. It copies raw chunks and flushes each write,
preserving the event:/id:/retry: lines that generic SSE clients (MCP
Streamable HTTP) rely on — re-framing via sseframe would drop them. The
WritesBody buffered-fallback guard and the responder re-framing path are
unchanged.

Add a regression test whose upstream holds the SSE connection open after the
first event (the shape the existing repro misses), asserting the event is
delivered before the upstream closes and that event:/id:/data: survive
verbatim.

Fixes #642

Assisted-By: Claude (Anthropic AI) <noreply@anthropic.com>
Signed-off-by: Hai Huang <huang195@gmail.com>

Author: huang195

authbridge/authlib/listener/forwardproxy/mcp_sse_stream_test.go

@@ -0,0 +1,135 @@
+package forwardproxy
+
+import (
+	"bufio"
+	"bytes"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"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
+}

authbridge/authlib/listener/forwardproxy/server.go

@@ -380,14 +380,23 @@ 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.
+				s.streamPassthrough(w, r, resp, pctx)
+				return
 			}
 		}
 
@@ -651,6 +660,69 @@ 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).
+//
+// Response-phase plugins (RunResponse / RunResponseFrame) are intentionally not
+// invoked here — consistent with the streaming contract that legacy/body
+// plugins don't run on streamed responses — but the response event is still
+// recorded for the session API / abctl.
+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
+	}
+
+	// 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