Merge pull request #527 from huang195/feat/abctl-one-row-per-message

Feat(abctl): one row per network message, show every message

Author: huang195

authbridge/authlib/config/config.go

@@ -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)
 }
 

authbridge/authlib/listener/forwardproxy/mcp_sse_repro_test.go

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

authbridge/authlib/listener/forwardproxy/server.go

@@ -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")
@@ -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

authbridge/authlib/listener/forwardproxy/server_test.go

@@ -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 {

authbridge/authlib/listener/forwardproxy/transparent.go

@@ -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

authbridge/authlib/listener/forwardproxy/transparent_test.go

@@ -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).

authbridge/authlib/listener/reverseproxy/finisher_test.go

@@ -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 {
@@ -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 {
@@ -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()
@@ -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)")
 	}

authbridge/authlib/listener/reverseproxy/server.go

@@ -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
@@ -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(),
@@ -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.