Fix: run response-phase pipeline before SSE passthrough streaming

streamPassthrough previously skipped RunResponse (matching
handleStreamingResponse), which silently bypassed header/status response
gates on SSE responses — e.g. opa's response-phase deny and
litellm-budgettrack's cost accounting would not run for a streamed
response.

Run RunResponse before the first byte in streamPassthrough and honor a
deny. This is safe only here: the path is reached exclusively when no
StreamingResponder is configured, so RunResponse cannot double-dispatch a
plugin that also implements OnResponseFrame, and the plugins reachable
here use status/headers only (never pctx.ResponseBody). Body-level
response inspection on a stream still requires implementing
StreamingResponder.

Add tests: a response-phase deny short-circuits before any SSE byte is
written, and a non-denying header-level OnResponse still runs while the
stream is delivered verbatim.

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

@@ -3,10 +3,12 @@ package forwardproxy
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"io"
 	"net/http"
 	"net/http/httptest"
 	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -133,3 +135,113 @@ func TestForwardProxy_SSE_StreamsWithoutResponder(t *testing.T) {
 	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.
+type responseProbePlugin struct {
+	deny    bool
+	ranResp atomic.Bool
+}
+
+func (p *responseProbePlugin) Name() string { return "response-probe" }
+func (p *responseProbePlugin) Capabilities() pipeline.PluginCapabilities {
+	return pipeline.PluginCapabilities{} // no ReadsBody/WritesBody/StreamingResponder → streamPassthrough path
+}
+func (p *responseProbePlugin) OnRequest(context.Context, *pipeline.Context) pipeline.Action {
+	return pipeline.Action{Type: pipeline.Continue}
+}
+func (p *responseProbePlugin) OnResponse(_ context.Context, _ *pipeline.Context) pipeline.Action {
+	p.ranResp.Store(true)
+	if p.deny {
+		return pipeline.DenyStatus(403, "test.denied", "denied by response probe")
+	}
+	return pipeline.Action{Type: pipeline.Continue}
+}
+
+// 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)
+	}
+}

authbridge/authlib/listener/forwardproxy/server.go

@@ -670,10 +670,16 @@ func (s *Server) handleStreamingResponse(w http.ResponseWriter, r *http.Request,
 // 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.
+// 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 {
@@ -685,6 +691,15 @@ func (s *Server) streamPassthrough(w http.ResponseWriter, r *http.Request, resp
 		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)