Merge pull request #480 from kagenti/fix/477-streaming-mcp-responses

fix(authbridge): Stream text/event-stream responses frame-by-frame

Author: huang195

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")
+}

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)")
+	}
+}