Skip to content
2 changes: 1 addition & 1 deletion authbridge/authlib/config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -198,7 +198,7 @@ type SessionConfig struct {
// "user didn't say" with "user said false" and silently flip the default.
Enabled *bool `yaml:"enabled" json:"enabled"`
TTL string `yaml:"ttl" json:"ttl"` // duration string; default: 30m
MaxEvents int `yaml:"max_events" json:"max_events"` // max events per session; default: 100
MaxEvents int `yaml:"max_events" json:"max_events"` // max events per session; default: 500
MaxSessions int `yaml:"max_sessions" json:"max_sessions"` // max concurrent sessions; default: 100 (0 = unlimited)
}

Expand Down
106 changes: 106 additions & 0 deletions authbridge/authlib/listener/forwardproxy/mcp_sse_repro_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,106 @@
package forwardproxy

import (
"bytes"
"io"
"net/http"
"net/http/httptest"
"testing"
"time"

"github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline"
"github.com/kagenti/kagenti-extensions/authbridge/authlib/plugins/inferenceparser"
"github.com/kagenti/kagenti-extensions/authbridge/authlib/plugins/mcpparser"
"github.com/kagenti/kagenti-extensions/authbridge/authlib/session"
)

// TestForwardProxy_MCP_SSEResponse_RecordsObserve reproduces the live
// weather-tool-mcp scenario: an MCP tools/list call whose response is
// text/event-stream (Streamable HTTP). The request is parsed (observe), and
// the RESPONSE must also be parsed into a result + recorded as an mcp-parser
// observe on the response event. The live cluster showed result=N,
// invocations=[] on the response — this test pins whether the proxy+parser
// path reproduces that.
func TestForwardProxy_MCP_SSEResponse_RecordsObserve(t *testing.T) {
// Upstream mimics weather-tool-mcp: tools/list -> SSE with one data frame
// carrying the JSON-RPC result, exactly as FastMCP emits it.
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\ndata: {\"jsonrpc\":\"2.0\",\"id\":2,\"result\":{\"tools\":[{\"name\":\"get_weather\"}]}}\n\n")
if f, ok := w.(http.Flusher); ok {
f.Flush()
}
}))
defer upstream.Close()

store := session.New(5*time.Minute, 100, 0)
defer store.Close()

// Production builds Configurable plugins through plugins.Build, which
// wraps them via pipeline.WrapConfigured. mcp-parser is Configurable +
// StreamingResponder, so it MUST stay a StreamingResponder after wrapping
// — that's the bug this test guards (a bare wrapper drops OnResponseFrame
// and RunResponseFrame skips it, leaving the SSE response unparsed).
// inference-parser is not Configurable, so it's added raw, as in production.
pipe, err := pipeline.New([]pipeline.Plugin{
pipeline.WrapConfigured(mcpparser.NewMCPParser(), nil),
inferenceparser.NewInferenceParser(),
})
if err != nil {
t.Fatal(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()

reqBody := `{"jsonrpc":"2.0","id":2,"method":"tools/list"}`
req, _ := http.NewRequest("POST", upstream.URL+"/mcp", bytes.NewReader([]byte(reqBody)))
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)
}
_, _ = io.ReadAll(resp.Body)
resp.Body.Close()

v := store.View(session.DefaultSessionID)
if v == nil {
t.Fatal("no session recorded")
}
var reqEv, respEv *pipeline.SessionEvent
for i := range v.Events {
e := &v.Events[i]
if e.MCP == nil || e.MCP.Method != "tools/list" {
continue
}
switch e.Phase {
case pipeline.SessionRequest:
reqEv = e
case pipeline.SessionResponse:
respEv = e
}
}
if reqEv == nil || respEv == nil {
t.Fatalf("missing req/resp events: req=%v resp=%v", reqEv != nil, respEv != nil)
}

// Request side parsed (sanity).
if reqEv.Invocations == nil || len(reqEv.Invocations.Outbound) == 0 {
t.Errorf("request event has no mcp-parser invocation: %+v", reqEv.Invocations)
}

// THE ASSERTION: the response was parsed into a result and recorded an
// observe. Live cluster fails both.
if respEv.MCP.Result == nil {
t.Errorf("response MCP.Result not populated (parser never saw the SSE result)")
}
respObserved := respEv.Invocations != nil && len(respEv.Invocations.Outbound) > 0
if !respObserved {
t.Errorf("response event has NO mcp-parser invocation — reproduces the bug. invocations=%+v", respEv.Invocations)
}
}
27 changes: 12 additions & 15 deletions authbridge/authlib/listener/forwardproxy/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -307,16 +307,14 @@ func (s *Server) serveOutbound(w http.ResponseWriter, r *http.Request, isBridge
Identity: pipeline.SnapshotIdentity(pctx),
Host: pctx.Host,
}
// Record whenever ANY protocol-or-plugin context is present —
// MCP/Inference (parser-emitted), Invocations (gate plugins like
// jwt-validation/token-exchange), or plugin-public Plugins
// entries. Earlier the gate was just MCP||Inference; widening
// it ensures auth-only outbound traffic and pure observability
// events show up in abctl. Don't narrow this back without
// understanding why each clause is necessary.
if ev.MCP != nil || ev.Inference != nil || ev.Invocations != nil || plugins != nil {
s.Sessions.Append(sid, ev)
}
// Record EVERY message that reaches the pipeline — even when no
// plugin acted and no parser matched (Invocations/MCP/Inference all
// nil). The session API is an observability surface; a request the
// pipeline saw but no plugin touched is still a network message the
// operator wants to see (it carries Host, and the paired response
// carries StatusCode). skip_hosts traffic never reaches here (the
// !skipped guard above), so it stays suppressed by design.
s.Sessions.Append(sid, ev)
}

newAuth := pctx.Headers.Get("Authorization")
Expand Down Expand Up @@ -519,11 +517,10 @@ func (s *Server) recordOutboundResponseEvent(pctx *pipeline.Context, 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)
}
// Always record — see the request-phase comment. This is what surfaces
// responses no plugin acted on (e.g. a generic 404), carrying StatusCode
// + Error even with empty invocations.
s.Sessions.Append(sid, ev)
}

// isEventStream reports whether a Content-Type header value names the
Expand Down
51 changes: 51 additions & 0 deletions authbridge/authlib/listener/forwardproxy/server_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -567,6 +567,57 @@ func TestRecordOutboundReject_SkipsWithoutInvocations(t *testing.T) {
}
}

// TestForwardProxy_RecordsMessageWithNoPluginActivity locks the Part A
// behavior: a request/response that no plugin acted on (empty pipeline,
// no parser match) is still recorded as two session events so abctl can
// show every network message — not just the ones a plugin touched. The
// response event carries the upstream status even though Invocations is
// nil.
func TestForwardProxy_RecordsMessageWithNoPluginActivity(t *testing.T) {
store := session.New(5*time.Minute, 100, 0)
defer store.Close()

backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusNotFound)
}))
defer backend.Close()

// Empty pipeline: zero plugins, so no Invocations are ever appended.
p, err := pipeline.New([]pipeline.Plugin{})
if err != nil {
t.Fatal(err)
}
srv := &Server{OutboundPipeline: pipeline.NewHolder(p), Sessions: store, Client: http.DefaultClient}
proxy := httptest.NewServer(srv.Handler())
defer proxy.Close()

req, _ := http.NewRequest("GET", backend.URL+"/missing", nil)
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)
}
resp.Body.Close()
if resp.StatusCode != http.StatusNotFound {
t.Errorf("status = %d, want 404", resp.StatusCode)
}

v := store.View(session.DefaultSessionID)
if v == nil || len(v.Events) != 2 {
t.Fatalf("expected 2 events (request + response) with no plugin activity, got %+v", v)
}
reqEv, respEv := v.Events[0], v.Events[1]
if reqEv.Phase != pipeline.SessionRequest || reqEv.Invocations != nil {
t.Errorf("request event = phase %v invocations %+v, want SessionRequest / nil", reqEv.Phase, reqEv.Invocations)
}
if respEv.Phase != pipeline.SessionResponse || respEv.StatusCode != http.StatusNotFound {
t.Errorf("response event = phase %v status %d, want SessionResponse / 404", respEv.Phase, respEv.StatusCode)
}
if respEv.Invocations != nil {
t.Errorf("response event invocations = %+v, want nil (no plugin acted)", respEv.Invocations)
}
}

// schemeCapturePlugin captures pctx.Scheme for the scheme-wiring
// test below.
type schemeCapturePlugin struct {
Expand Down
10 changes: 7 additions & 3 deletions authbridge/authlib/listener/forwardproxy/transparent.go
Original file line number Diff line number Diff line change
Expand Up @@ -179,10 +179,14 @@ func (s *Server) recordTunnelOpened(pctx *pipeline.Context) {
Plugins: plugins,
Identity: pipeline.SnapshotIdentity(pctx),
Host: pctx.Host,
// Explicit opaque-tunnel marker so abctl can fold this CONNECT into
// the decrypted inner request without inferring "tunnel" from shape.
Tunnel: true,
}
if ev.Invocations != nil || plugins != nil {
s.Sessions.Append(sid, ev)
}
// Always record the tunnel-open so passthrough/non-bridged tunnels (no
// plugin activity) are still visible. For a TLS-bridged call abctl folds
// this CONNECT event into the decrypted inner-request row.
s.Sessions.Append(sid, ev)
}

// tunnel bidirectionally copies between two connections until either side
Expand Down
25 changes: 25 additions & 0 deletions authbridge/authlib/listener/forwardproxy/transparent_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,33 @@ import (

"github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline"
"github.com/kagenti/kagenti-extensions/authbridge/authlib/plugins/plugintesting"
"github.com/kagenti/kagenti-extensions/authbridge/authlib/session"
)

// TestRecordTunnelOpened_SetsTunnelMarker locks the explicit producer marker:
// a CONNECT / transparent-redirect tunnel-open is recorded with Tunnel=true so
// consumers (abctl) fold it into the decrypted inner request without inferring
// "tunnel" from host/extension shape.
func TestRecordTunnelOpened_SetsTunnelMarker(t *testing.T) {
store := session.New(5*time.Minute, 100, 0)
defer store.Close()
s := &Server{Sessions: store}

s.recordTunnelOpened(&pipeline.Context{Direction: pipeline.Outbound, Host: "example.com:443"})

v := store.View(session.DefaultSessionID)
if v == nil || len(v.Events) != 1 {
t.Fatalf("expected 1 tunnel-open event, got %+v", v)
}
ev := v.Events[0]
if !ev.Tunnel {
t.Error("tunnel-open event must have Tunnel=true")
}
if ev.Direction != pipeline.Outbound || ev.Phase != pipeline.SessionRequest {
t.Errorf("tunnel-open = %v/%v, want Outbound/SessionRequest", ev.Direction, ev.Phase)
}
}

// HandleTransparentConn gates then blind-tunnels: with an allow-all pipeline it
// must dial the recovered destination and copy bytes both ways, emitting no
// proxy-protocol bytes of its own (the agent thinks it's talking to dst).
Expand Down
29 changes: 25 additions & 4 deletions authbridge/authlib/listener/reverseproxy/finisher_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,27 @@ import (
"net/http/httptest"
"sync/atomic"
"testing"
"time"

"github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline"
"github.com/kagenti/kagenti-extensions/authbridge/authlib/plugins/plugintesting"
)

// waitSeen polls b until it's true, up to ~1s. OnFinish runs in the
// reverse-proxy handler's deferred RunFinish, which fires only after the
// streamed response (NewServer sets FlushInterval=-1) has already reached the
// client — so http.Get can return before OnFinish has run. Poll rather than
// read immediately, which otherwise races (flaky under -race).
func waitSeen(b *atomic.Bool) bool {
for i := 0; i < 1000; i++ {
if b.Load() {
return true
}
time.Sleep(time.Millisecond)
}
return false
}

// finisherStub is a minimal Plugin + Finisher used by these tests to
// observe OnFinish dispatch and the Outcome the listener derived.
type finisherStub struct {
Expand All @@ -34,11 +50,13 @@ func (p *finisherStub) OnResponse(context.Context, *pipeline.Context) pipeline.A
return pipeline.Action{Type: pipeline.Continue}
}
func (p *finisherStub) OnFinish(_ context.Context, pctx *pipeline.Context) {
p.seen.Store(true)
if o := pctx.Outcome(); o != nil {
cp := *o
p.outcome.Store(&cp)
}
// Store seen LAST so a reader that observes seen==true is guaranteed to
// also see the outcome that was set just above.
p.seen.Store(true)
}

func pipelineWith(t *testing.T, plugins ...pipeline.Plugin) *pipeline.Holder {
Expand Down Expand Up @@ -73,7 +91,7 @@ func TestReverseProxy_Finisher_Allow(t *testing.T) {
}
resp.Body.Close()

if !f.seen.Load() {
if !waitSeen(&f.seen) {
t.Fatal("OnFinish did not fire")
}
o := f.outcome.Load()
Expand Down Expand Up @@ -133,12 +151,15 @@ func TestReverseProxy_Finisher_Deny(t *testing.T) {
t.Errorf("HTTP status = %d, want 403", resp.StatusCode)
}

if !before.seen.Load() {
if !waitSeen(&before.seen) {
t.Error("before-deny.OnFinish should have fired (OnRequest ran before denial)")
}
if !denier.seen.Load() {
if !waitSeen(&denier.seen) {
t.Error("denier.OnFinish should have fired (OnRequest ran and produced the deny)")
}
// before/denier have fired, so the single RunFinish dispatch is complete;
// after-deny was never dispatched (its OnRequest never ran), so its
// OnFinish must not have fired.
if after.seen.Load() {
t.Error("after-deny.OnFinish should NOT have fired (OnRequest never ran)")
}
Expand Down
12 changes: 7 additions & 5 deletions authbridge/authlib/listener/reverseproxy/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -256,7 +256,9 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) {
// as denials already do via recordInboundReject). The A2A-specific session
// rekey in modifyResponse stays A2A-gated.
plugins := pipeline.SnapshotPlugins(pctx.Extensions.Custom)
if s.Sessions != nil && (pctx.Extensions.A2A != nil || pctx.Extensions.Invocations != nil || plugins != nil) {
// Record every inbound request the pipeline saw, even with no plugin
// activity (skip_hosts is N/A inbound).
if s.Sessions != nil {
sid := inboundSessionID(pctx)
// Snapshot-copy the protocol extension and use the shared helpers
// for plugin invocations / observability / identity. Mirrors what
Expand Down Expand Up @@ -378,7 +380,8 @@ func (s *Server) modifyResponse(resp *http.Response) error {
// pipeline saw at this point (may be empty for streamed bodies),
// but the status code and plugin invocations are always meaningful.
plugins := pipeline.SnapshotPlugins(pctx.Extensions.Custom)
if s.Sessions != nil && (pctx.Extensions.A2A != nil || pctx.Extensions.Invocations != nil || plugins != nil) {
// Always pair every inbound request with a response row (carries StatusCode).
if s.Sessions != nil {
sid := inboundSessionID(pctx)
s.Sessions.Append(sid, pipeline.SessionEvent{
At: time.Now(),
Expand Down Expand Up @@ -513,9 +516,8 @@ func (s *Server) recordInboundResponseEvent(pctx *pipeline.Context, statusCode i
return
}
plugins := pipeline.SnapshotPlugins(pctx.Extensions.Custom)
if !(pctx.Extensions.A2A != nil || pctx.Extensions.Invocations != nil || plugins != nil) {
return
}
// Always record the streaming response (carries StatusCode), even with
// no plugin activity.
sid := inboundSessionID(pctx)
// Rekey default → contextId mirroring the buffered path's behavior;
// streaming A2A message/stream may discover the contextId mid-stream.
Expand Down
Loading
Loading