From 01aeecfc3858ab21997c7c74f93b3ff0230cae6f Mon Sep 17 00:00:00 2001
From: Hai Huang <huang195@gmail.com>
Date: Thu, 14 May 2026 22:08:12 -0400
Subject: [PATCH 1/3] refactor(authlib): proxy listeners session-event parity
 with extproc
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Closes the feature gap between the proxy-sidecar listeners
(forwardproxy, reverseproxy) and the envoy-sidecar listener
(extproc) in their session-event recording. Operators on
proxy-sidecar mode were seeing impoverished abctl rows because the
proxy listeners only carried a subset of what extproc carries on
each event.

What was missing on proxy listeners (now fixed):

  * Plugin Invocations on accept-path events. Previously only the
    deny path threaded Invocations through; happy-path requests +
    responses left the PLUGIN/ACTION columns blank in abctl.

  * Inbound response event. reverseproxy.modifyResponse ran the
    inbound pipeline's response phase but never appended a
    Phase:SessionResponse event. Inbound A2A `message/stream`
    requests showed up as orphan request rows.

  * Rekey default → A2A contextId. The first turn of an A2A
    conversation arrives without a contextId (the agent assigns it
    on response). Without rekey, the inbound request + outbound
    MCP/inference calls during processing landed in `default` while
    only the response went to the contextId bucket. Now mirrors
    extproc.rekeyInboundSession: response phase calls Rekey when
    the parser reveals the contextId, merging the whole turn.

  * Host, StatusCode, Duration fields on every event. Were unset on
    proxy events; abctl's HOST / STATUS / DURATION columns stayed
    blank for proxy-sidecar pods.

  * Pointer aliasing on Extensions.A2A/MCP/Inference. Events stored
    pointers to the live pctx structs; OnResponse mutated the same
    structs and the previously-recorded request events
    retroactively showed response-phase fields (artifact,
    finalStatus, token counts). extproc avoided this with shallow-
    copy Snapshot helpers; proxy listeners did not.

  * Identity (Subject/ClientID/Scopes) and the plugin-public
    Plugins map (Custom entries with /event suffix) were never
    surfaced on proxy listener events.

  * Error on response events derived from Security.Blocked or non-
    2xx StatusCode.

Implementation: lifted the 8 snapshot/derive helpers
(SnapshotA2A, SnapshotMCP, SnapshotInference, SnapshotInvocations,
SnapshotPlugins, SnapshotIdentity, DeriveError, DurationSince) from
the extproc package into authlib/pipeline/snapshot.go as exported
functions. extproc switches to the shared package versions; both
proxy listeners now use the same helpers, single source of truth.

Tier 1 cross-listener dedup also folded in:

  * extractBearer was duplicated byte-for-byte across extauthz,
    forwardproxy, and extproc. Moved to authlib/auth/bearer.go as
    auth.ExtractBearer.

  * writeRejection was duplicated byte-for-byte across forwardproxy
    and reverseproxy. Moved to a new
    authlib/listener/httpx/render.go as httpx.WriteRejection — the
    home for HTTP-listener-shared helpers.

Verified locally: full authlib test suite passes; all three
combined images build clean; weather-agent in proxy-sidecar mode
in a kind cluster produces abctl events with full PLUGIN/ACTION/
HOST/STATUS/DURATION/Identity columns and merges into the A2A
contextId bucket on every turn.

Tier 2 cross-listener dedup (recordInboundReject /
recordOutboundReject / inboundSessionID / Rekey-as-method)
intentionally deferred — needs design care on the abstraction
surface (event-builder taking pctx + action + direction + phase).
Better as a separate PR.

Assisted-By: Claude (Anthropic AI) <noreply@anthropic.com>
Signed-off-by: Hai Huang <huang195@gmail.com>
---
 authbridge/authlib/auth/bearer.go             |  19 ++
 .../authlib/listener/extauthz/server.go       |  10 +-
 authbridge/authlib/listener/extproc/server.go | 203 +++---------------
 .../authlib/listener/extproc/server_test.go   |  12 +-
 .../authlib/listener/forwardproxy/server.go   |  71 +++---
 authbridge/authlib/listener/httpx/render.go   |  30 +++
 .../authlib/listener/reverseproxy/server.go   |  93 +++++---
 authbridge/authlib/pipeline/snapshot.go       | 137 ++++++++++++
 8 files changed, 320 insertions(+), 255 deletions(-)
 create mode 100644 authbridge/authlib/auth/bearer.go
 create mode 100644 authbridge/authlib/listener/httpx/render.go
 create mode 100644 authbridge/authlib/pipeline/snapshot.go

diff --git a/authbridge/authlib/auth/bearer.go b/authbridge/authlib/auth/bearer.go
new file mode 100644
index 000000000..02583f99f
--- /dev/null
+++ b/authbridge/authlib/auth/bearer.go
@@ -0,0 +1,19 @@
+package auth
+
+import "strings"
+
+// ExtractBearer pulls the token out of an HTTP `Authorization` header
+// value. Returns the token (no "Bearer " prefix) on a well-formed
+// bearer header, "" otherwise. The "Bearer" scheme match is
+// case-insensitive per RFC 6750 §2.1.
+//
+// Used by every listener that touches request headers — extauthz,
+// extproc, forwardproxy — so it lives here in authlib/auth (the home
+// of the bearer/JWT composition layer) rather than being duplicated
+// per listener.
+func ExtractBearer(authHeader string) string {
+	if len(authHeader) > 7 && strings.EqualFold(authHeader[:7], "bearer ") {
+		return authHeader[7:]
+	}
+	return ""
+}
diff --git a/authbridge/authlib/listener/extauthz/server.go b/authbridge/authlib/listener/extauthz/server.go
index e3e66c930..4affc90d9 100644
--- a/authbridge/authlib/listener/extauthz/server.go
+++ b/authbridge/authlib/listener/extauthz/server.go
@@ -6,7 +6,6 @@ package extauthz
 import (
 	"context"
 	"net/http"
-	"strings"
 	"time"
 
 	corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
@@ -16,6 +15,7 @@ import (
 
 	rpcstatus "google.golang.org/genproto/googleapis/rpc/status"
 
+	"github.com/kagenti/kagenti-extensions/authbridge/authlib/auth"
 	"github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline"
 )
 
@@ -91,7 +91,7 @@ func (s *Server) Check(ctx context.Context, req *authv3.CheckRequest) (*authv3.C
 
 	newAuth := outPctx.Headers.Get("Authorization")
 	if newAuth != originalAuth {
-		return allowedWithToken(extractBearer(newAuth)), nil
+		return allowedWithToken(auth.ExtractBearer(newAuth)), nil
 	}
 	return allowed(), nil
 }
@@ -121,12 +121,6 @@ func mapToHTTPHeader(m map[string]string) http.Header {
 	return h
 }
 
-func extractBearer(authHeader string) string {
-	if len(authHeader) > 7 && strings.EqualFold(authHeader[:7], "bearer ") {
-		return authHeader[7:]
-	}
-	return ""
-}
 
 // deniedFromAction renders a pipeline Reject as an ext_authz CheckResponse
 // preserving the plugin's status, headers, and body. The flat
diff --git a/authbridge/authlib/listener/extproc/server.go b/authbridge/authlib/listener/extproc/server.go
index 30c48bf1b..ec0e4f466 100644
--- a/authbridge/authlib/listener/extproc/server.go
+++ b/authbridge/authlib/listener/extproc/server.go
@@ -19,6 +19,7 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
+	"github.com/kagenti/kagenti-extensions/authbridge/authlib/auth"
 	"github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline"
 	"github.com/kagenti/kagenti-extensions/authbridge/authlib/session"
 )
@@ -207,7 +208,7 @@ func (s *Server) recordInboundSession(pctx *pipeline.Context) {
 	// Widened gate (was: A2A == nil). Any of A2A / Auth / plugin-public
 	// Custom entries qualify. Keeps traffic with no protocol parser but
 	// meaningful auth state visible in the session stream.
-	plugins := snapshotPlugins(pctx.Extensions.Custom)
+	plugins := pipeline.SnapshotPlugins(pctx.Extensions.Custom)
 	if pctx.Extensions.A2A == nil && pctx.Extensions.Invocations == nil && plugins == nil {
 		return
 	}
@@ -216,10 +217,10 @@ func (s *Server) recordInboundSession(pctx *pipeline.Context) {
 		At:          time.Now(),
 		Direction:   pipeline.Inbound,
 		Phase:       pipeline.SessionRequest,
-		A2A:         snapshotA2A(pctx.Extensions.A2A),
-		Invocations: snapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest),
+		A2A:         pipeline.SnapshotA2A(pctx.Extensions.A2A),
+		Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest),
 		Plugins:     plugins,
-		Identity:    snapshotIdentity(pctx),
+		Identity:    pipeline.SnapshotIdentity(pctx),
 		Host:        pctx.Host,
 	}
 	s.Sessions.Append(sid, ev)
@@ -254,9 +255,9 @@ func (s *Server) recordInboundReject(pctx *pipeline.Context, action pipeline.Act
 		At:          time.Now(),
 		Direction:   pipeline.Inbound,
 		Phase:       pipeline.SessionDenied,
-		Invocations: snapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest),
-		Plugins:     snapshotPlugins(pctx.Extensions.Custom),
-		Identity:    snapshotIdentity(pctx),
+		Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest),
+		Plugins:     pipeline.SnapshotPlugins(pctx.Extensions.Custom),
+		Identity:    pipeline.SnapshotIdentity(pctx),
 		Host:        pctx.Host,
 		StatusCode:  status,
 		Error: &pipeline.EventError{
@@ -264,7 +265,7 @@ func (s *Server) recordInboundReject(pctx *pipeline.Context, action pipeline.Act
 			Code:    code,
 			Message: message,
 		},
-		Duration: durationSince(pctx.StartedAt),
+		Duration: pipeline.DurationSince(pctx.StartedAt),
 	}
 	s.Sessions.Append(inboundSessionID(pctx), ev)
 }
@@ -307,9 +308,9 @@ func (s *Server) recordOutboundReject(pctx *pipeline.Context, action pipeline.Ac
 		At:          time.Now(),
 		Direction:   pipeline.Outbound,
 		Phase:       pipeline.SessionDenied,
-		Invocations: snapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest),
-		Plugins:     snapshotPlugins(pctx.Extensions.Custom),
-		Identity:    snapshotIdentity(pctx),
+		Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest),
+		Plugins:     pipeline.SnapshotPlugins(pctx.Extensions.Custom),
+		Identity:    pipeline.SnapshotIdentity(pctx),
 		Host:        pctx.Host,
 		StatusCode:  status,
 		Error: &pipeline.EventError{
@@ -317,58 +318,12 @@ func (s *Server) recordOutboundReject(pctx *pipeline.Context, action pipeline.Ac
 			Code:    code,
 			Message: message,
 		},
-		Duration: durationSince(pctx.StartedAt),
+		Duration: pipeline.DurationSince(pctx.StartedAt),
 	}
 	s.Sessions.Append(sid, ev)
 }
 
-// snapshotInvocations returns a shallow copy of the Invocations extension
-// filtered by phase. Plugins append to pctx.Extensions.Invocations as
-// both OnRequest and OnResponse fire; the full list lives there for
-// cross-phase inspection. At record time each SessionEvent should carry
-// only the invocations from its own phase, so request events don't
-// double-report request-phase entries AFTER the response phase has
-// already added its own. Request events pass InvocationPhaseRequest,
-// response events pass InvocationPhaseResponse, denied events pass
-// InvocationPhaseRequest (denial terminates the pass before response
-// runs). Returns nil when no matching entry exists, so the recording
-// gate can check for "no invocations on this phase" cleanly.
-//
-// Thin wrapper over (*Invocations).FilteredByPhase — kept as a named
-// function so recordXxxSession call sites stay grep-friendly.
-func snapshotInvocations(ext *pipeline.Invocations, phase pipeline.InvocationPhase) *pipeline.Invocations {
-	return ext.FilteredByPhase(phase)
-}
 
-// snapshotPlugins collects plugin-public observability events from
-// pctx.Extensions.Custom entries whose keys end in PluginEventSuffix.
-// Each matching value is json.Marshaled into the wire-form map under
-// the plugin name (suffix stripped). Marshal errors downgrade to slog
-// Debug and skip the entry rather than aborting recording — that keeps
-// a misbehaving plugin from taking out the whole session stream.
-func snapshotPlugins(custom map[string]any) map[string]json.RawMessage {
-	if len(custom) == 0 {
-		return nil
-	}
-	var out map[string]json.RawMessage
-	for k, v := range custom {
-		if !strings.HasSuffix(k, pipeline.PluginEventSuffix) {
-			continue
-		}
-		raw, err := json.Marshal(v)
-		if err != nil {
-			slog.Debug("session: skipping non-marshalable plugin event",
-				"key", k, "error", err)
-			continue
-		}
-		if out == nil {
-			out = make(map[string]json.RawMessage)
-		}
-		pluginName := strings.TrimSuffix(k, pipeline.PluginEventSuffix)
-		out[pluginName] = raw
-	}
-	return out
-}
 
 // recordInboundResponseSession appends a Phase:SessionResponse event for the
 // inbound direction. Called after RunResponse completes so the event carries
@@ -385,7 +340,7 @@ func (s *Server) recordInboundResponseSession(pctx *pipeline.Context) {
 	if s.Sessions == nil {
 		return
 	}
-	plugins := snapshotPlugins(pctx.Extensions.Custom)
+	plugins := pipeline.SnapshotPlugins(pctx.Extensions.Custom)
 	if pctx.Extensions.A2A == nil && pctx.Extensions.Invocations == nil && plugins == nil {
 		return
 	}
@@ -394,14 +349,14 @@ func (s *Server) recordInboundResponseSession(pctx *pipeline.Context) {
 		At:          time.Now(),
 		Direction:   pipeline.Inbound,
 		Phase:       pipeline.SessionResponse,
-		A2A:         snapshotA2A(pctx.Extensions.A2A),
-		Invocations: snapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseResponse),
+		A2A:         pipeline.SnapshotA2A(pctx.Extensions.A2A),
+		Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseResponse),
 		Plugins:     plugins,
-		Identity:    snapshotIdentity(pctx),
+		Identity:    pipeline.SnapshotIdentity(pctx),
 		StatusCode:  pctx.StatusCode,
-		Error:       deriveError(pctx),
+		Error:       pipeline.DeriveError(pctx),
 		Host:        pctx.Host,
-		Duration:    durationSince(pctx.StartedAt),
+		Duration:    pipeline.DurationSince(pctx.StartedAt),
 	}
 	s.Sessions.Append(sid, ev)
 }
@@ -417,20 +372,20 @@ func (s *Server) recordOutboundResponseSession(pctx *pipeline.Context) {
 	if sid == "" {
 		sid = session.DefaultSessionID
 	}
-	plugins := snapshotPlugins(pctx.Extensions.Custom)
+	plugins := pipeline.SnapshotPlugins(pctx.Extensions.Custom)
 	ev := pipeline.SessionEvent{
 		At:          time.Now(),
 		Direction:   pipeline.Outbound,
 		Phase:       pipeline.SessionResponse,
-		MCP:         snapshotMCP(pctx.Extensions.MCP),
-		Inference:   snapshotInference(pctx.Extensions.Inference),
-		Invocations: snapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseResponse),
+		MCP:         pipeline.SnapshotMCP(pctx.Extensions.MCP),
+		Inference:   pipeline.SnapshotInference(pctx.Extensions.Inference),
+		Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseResponse),
 		Plugins:     plugins,
-		Identity:    snapshotIdentity(pctx),
+		Identity:    pipeline.SnapshotIdentity(pctx),
 		StatusCode:  pctx.StatusCode,
-		Error:       deriveError(pctx),
+		Error:       pipeline.DeriveError(pctx),
 		Host:        pctx.Host,
-		Duration:    durationSince(pctx.StartedAt),
+		Duration:    pipeline.DurationSince(pctx.StartedAt),
 	}
 	// Auth / Plugins alone qualify for recording; matches the widened
 	// gate in recordInboundSession so outbound denials and plugin-public
@@ -441,57 +396,8 @@ func (s *Server) recordOutboundResponseSession(pctx *pipeline.Context) {
 	}
 }
 
-// snapshotIdentity copies Claims + Agent identity off pctx so the session event
-// stays valid after pctx is discarded. Returns nil when no identity information
-// is available (e.g., jwt-validation didn't run on this path).
-func snapshotIdentity(pctx *pipeline.Context) *pipeline.EventIdentity {
-	// Read identity via the pipeline.Identity interface. Whichever auth
-	// plugin ran (jwt-validation today; SAML / mTLS / custom later)
-	// published an adapter onto pctx.Identity — we don't need to know
-	// its concrete type to snapshot the subject/client/scopes.
-	if pctx.Identity == nil && pctx.Agent == nil {
-		return nil
-	}
-	id := &pipeline.EventIdentity{}
-	if pctx.Identity != nil {
-		id.Subject = pctx.Identity.Subject()
-		id.ClientID = pctx.Identity.ClientID()
-		if scopes := pctx.Identity.Scopes(); len(scopes) > 0 {
-			id.Scopes = append([]string(nil), scopes...)
-		}
-	}
-	if pctx.Agent != nil {
-		id.AgentID = pctx.Agent.WorkloadID
-	}
-	return id
-}
 
-// durationSince returns the elapsed time since StartedAt, or 0 when StartedAt
-// is zero (pctx constructed without wall-clock stamping, e.g. in unit tests).
-func durationSince(start time.Time) time.Duration {
-	if start.IsZero() {
-		return 0
-	}
-	return time.Since(start)
-}
 
-// deriveError constructs an EventError from response-side signals. Returns nil
-// for 2xx / no guardrail block / no parser error.
-func deriveError(pctx *pipeline.Context) *pipeline.EventError {
-	if pctx.Extensions.Security != nil && pctx.Extensions.Security.Blocked {
-		return &pipeline.EventError{
-			Kind:    "blocked",
-			Message: pctx.Extensions.Security.BlockReason,
-		}
-	}
-	if pctx.StatusCode >= 400 {
-		return &pipeline.EventError{
-			Kind: "backend_error",
-			Code: strconv.Itoa(pctx.StatusCode),
-		}
-	}
-	return nil
-}
 
 // rekeyInboundSession renames the DefaultSessionID bucket to the
 // server-assigned A2A contextId when the response reveals one, so events
@@ -516,16 +422,16 @@ func (s *Server) recordOutboundSession(pctx *pipeline.Context) {
 	if sid == "" {
 		sid = session.DefaultSessionID
 	}
-	plugins := snapshotPlugins(pctx.Extensions.Custom)
+	plugins := pipeline.SnapshotPlugins(pctx.Extensions.Custom)
 	ev := pipeline.SessionEvent{
 		At:          time.Now(),
 		Direction:   pipeline.Outbound,
 		Phase:       pipeline.SessionRequest,
-		MCP:         snapshotMCP(pctx.Extensions.MCP),
-		Inference:   snapshotInference(pctx.Extensions.Inference),
-		Invocations: snapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest),
+		MCP:         pipeline.SnapshotMCP(pctx.Extensions.MCP),
+		Inference:   pipeline.SnapshotInference(pctx.Extensions.Inference),
+		Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest),
 		Plugins:     plugins,
-		Identity:    snapshotIdentity(pctx),
+		Identity:    pipeline.SnapshotIdentity(pctx),
 		Host:        pctx.Host,
 	}
 	if ev.MCP != nil || ev.Inference != nil || ev.Invocations != nil || plugins != nil {
@@ -533,45 +439,8 @@ func (s *Server) recordOutboundSession(pctx *pipeline.Context) {
 	}
 }
 
-// snapshotA2A returns a shallow copy of ext. The record helpers attach
-// the snapshot to the SessionEvent rather than the live pointer so
-// response-phase mutations on pctx.Extensions.A2A (e.g. the parser
-// stamping the server-assigned contextId onto SessionID during OnResponse)
-// don't retroactively rewrite request-phase events that were already
-// appended. Slice fields are reused intentionally — they are only
-// assigned, never mutated in place, after the parser completes.
-func snapshotA2A(ext *pipeline.A2AExtension) *pipeline.A2AExtension {
-	if ext == nil {
-		return nil
-	}
-	c := *ext
-	return &c
-}
 
-// snapshotMCP returns a shallow copy of ext. Important for outbound
-// request events: the same pctx.Extensions.MCP pointer receives Result
-// or Err on the response side, so without snapshotting, the
-// already-recorded request event would display the future response's
-// result map.
-func snapshotMCP(ext *pipeline.MCPExtension) *pipeline.MCPExtension {
-	if ext == nil {
-		return nil
-	}
-	c := *ext
-	return &c
-}
 
-// snapshotInference returns a shallow copy of ext. Scalar response
-// fields (Completion, FinishReason, *Tokens) get assigned on the live
-// extension during OnResponse; without snapshotting, the request event's
-// view would contain the eventual response's token counts and completion.
-func snapshotInference(ext *pipeline.InferenceExtension) *pipeline.InferenceExtension {
-	if ext == nil {
-		return nil
-	}
-	c := *ext
-	return &c
-}
 
 func (s *Server) handleOutbound(stream extprocv3.ExternalProcessor_ProcessServer, headers *corev3.HeaderMap, body []byte) (*extprocv3.ProcessingResponse, *pipeline.Context) {
 	ctx := stream.Context()
@@ -605,7 +474,7 @@ func (s *Server) handleOutbound(stream extprocv3.ExternalProcessor_ProcessServer
 
 	newAuth := pctx.Headers.Get("Authorization")
 	if newAuth != originalAuth {
-		return replaceTokenResponse(extractBearer(newAuth)), pctx
+		return replaceTokenResponse(auth.ExtractBearer(newAuth)), pctx
 	}
 	return passResponse(), pctx
 }
@@ -642,7 +511,7 @@ func (s *Server) handleOutboundBody(stream extprocv3.ExternalProcessor_ProcessSe
 
 	newAuth := pctx.Headers.Get("Authorization")
 	if newAuth != originalAuth {
-		return withBodyMutation(replaceTokenBodyResponse(extractBearer(newAuth)), pctx), pctx
+		return withBodyMutation(replaceTokenBodyResponse(auth.ExtractBearer(newAuth)), pctx), pctx
 	}
 	return withBodyMutation(passBodyResponse(), pctx), pctx
 }
@@ -772,12 +641,6 @@ func headerMapToHTTP(headers *corev3.HeaderMap) http.Header {
 	return h
 }
 
-func extractBearer(authHeader string) string {
-	if len(authHeader) > 7 && strings.EqualFold(authHeader[:7], "bearer ") {
-		return authHeader[7:]
-	}
-	return ""
-}
 
 func requestBodyResponse() *extprocv3.ProcessingResponse {
 	return &extprocv3.ProcessingResponse{
diff --git a/authbridge/authlib/listener/extproc/server_test.go b/authbridge/authlib/listener/extproc/server_test.go
index a977f2fc8..a60948bb1 100644
--- a/authbridge/authlib/listener/extproc/server_test.go
+++ b/authbridge/authlib/listener/extproc/server_test.go
@@ -1068,7 +1068,7 @@ func TestSnapshotIdentity(t *testing.T) {
 
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			got := snapshotIdentity(&pipeline.Context{Identity: tc.identity, Agent: tc.agent})
+			got := pipeline.SnapshotIdentity(&pipeline.Context{Identity: tc.identity, Agent: tc.agent})
 			if tc.want == nil {
 				if got != nil {
 					t.Errorf("got %+v, want nil", got)
@@ -1091,7 +1091,7 @@ func TestSnapshotIdentity(t *testing.T) {
 func TestSnapshotIdentity_ScopesDeepCopy(t *testing.T) {
 	// Mutating the source slice after snapshot must not mutate the event.
 	scopes := []string{"a", "b"}
-	id := snapshotIdentity(&pipeline.Context{Identity: stubIdentity{subject: "alice", scopes: scopes}})
+	id := pipeline.SnapshotIdentity(&pipeline.Context{Identity: stubIdentity{subject: "alice", scopes: scopes}})
 	scopes[0] = "x"
 	if id.Scopes[0] != "a" {
 		t.Errorf("snapshot scopes aliased original: got %v", id.Scopes)
@@ -1124,7 +1124,7 @@ func TestDeriveError(t *testing.T) {
 
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			got := deriveError(tc.pctx)
+			got := pipeline.DeriveError(tc.pctx)
 			if tc.wantNil {
 				if got != nil {
 					t.Errorf("got %+v, want nil", got)
@@ -1176,11 +1176,11 @@ func TestRecordOutboundResponseSession_CapturesStatusAndError(t *testing.T) {
 }
 
 func TestDurationSince(t *testing.T) {
-	if d := durationSince(time.Time{}); d != 0 {
+	if d := pipeline.DurationSince(time.Time{}); d != 0 {
 		t.Errorf("zero StartedAt should yield 0, got %v", d)
 	}
 	start := time.Now().Add(-50 * time.Millisecond)
-	if d := durationSince(start); d < 50*time.Millisecond {
+	if d := pipeline.DurationSince(start); d < 50*time.Millisecond {
 		t.Errorf("elapsed = %v, want >= 50ms", d)
 	}
 }
@@ -1456,7 +1456,7 @@ func TestSnapshotPlugins_FiltersByEventSuffix(t *testing.T) {
 			Allowed: true, TokensLeft: 42,
 		},
 	}
-	out := snapshotPlugins(custom)
+	out := pipeline.SnapshotPlugins(custom)
 	if _, private := out["rate-limiter"]; !private {
 		// Key exists in out because the /event entry WAS promoted.
 		// Clarifying: we want exactly one entry, keyed "rate-limiter".
diff --git a/authbridge/authlib/listener/forwardproxy/server.go b/authbridge/authlib/listener/forwardproxy/server.go
index 4866d3f9c..e9fb67ddf 100644
--- a/authbridge/authlib/listener/forwardproxy/server.go
+++ b/authbridge/authlib/listener/forwardproxy/server.go
@@ -9,9 +9,10 @@ import (
 	"io"
 	"log/slog"
 	"net/http"
-	"strings"
 	"time"
 
+	"github.com/kagenti/kagenti-extensions/authbridge/authlib/auth"
+	"github.com/kagenti/kagenti-extensions/authbridge/authlib/listener/httpx"
 	"github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline"
 	"github.com/kagenti/kagenti-extensions/authbridge/authlib/session"
 )
@@ -95,7 +96,7 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) {
 
 	if action.Type == pipeline.Reject {
 		s.recordOutboundReject(pctx, action)
-		writeRejection(w, action)
+		httpx.WriteRejection(w, action)
 		return
 	}
 
@@ -104,21 +105,29 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) {
 		if sid == "" {
 			sid = session.DefaultSessionID
 		}
+		// Snapshot-copy the protocol extension so the request event
+		// doesn't see response-phase mutations on the same MCP/Inference
+		// struct (e.g. token counts assigned in OnResponse).
+		plugins := pipeline.SnapshotPlugins(pctx.Extensions.Custom)
 		ev := pipeline.SessionEvent{
-			At:        time.Now(),
-			Direction: pipeline.Outbound,
-			Phase:     pipeline.SessionRequest,
-			MCP:       pctx.Extensions.MCP,
-			Inference: pctx.Extensions.Inference,
+			At:          time.Now(),
+			Direction:   pipeline.Outbound,
+			Phase:       pipeline.SessionRequest,
+			MCP:         pipeline.SnapshotMCP(pctx.Extensions.MCP),
+			Inference:   pipeline.SnapshotInference(pctx.Extensions.Inference),
+			Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest),
+			Plugins:     plugins,
+			Identity:    pipeline.SnapshotIdentity(pctx),
+			Host:        pctx.Host,
 		}
-		if ev.MCP != nil || ev.Inference != nil {
+		if ev.MCP != nil || ev.Inference != nil || ev.Invocations != nil || plugins != nil {
 			s.Sessions.Append(sid, ev)
 		}
 	}
 
 	newAuth := pctx.Headers.Get("Authorization")
 	if newAuth != originalAuth {
-		r.Header.Set("Authorization", "Bearer "+extractBearer(newAuth))
+		r.Header.Set("Authorization", "Bearer "+auth.ExtractBearer(newAuth))
 	}
 
 	// If a WritesBody plugin rewrote pctx.Body, ship the new bytes
@@ -174,7 +183,7 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) {
 
 	respAction := s.OutboundPipeline.RunResponse(r.Context(), pctx)
 	if respAction.Type == pipeline.Reject {
-		writeRejection(w, respAction)
+		httpx.WriteRejection(w, respAction)
 		return
 	}
 
@@ -195,14 +204,22 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) {
 		if sid == "" {
 			sid = session.DefaultSessionID
 		}
+		plugins := pipeline.SnapshotPlugins(pctx.Extensions.Custom)
 		ev := pipeline.SessionEvent{
-			At:        time.Now(),
-			Direction: pipeline.Outbound,
-			Phase:     pipeline.SessionResponse,
-			MCP:       pctx.Extensions.MCP,
-			Inference: pctx.Extensions.Inference,
+			At:          time.Now(),
+			Direction:   pipeline.Outbound,
+			Phase:       pipeline.SessionResponse,
+			MCP:         pipeline.SnapshotMCP(pctx.Extensions.MCP),
+			Inference:   pipeline.SnapshotInference(pctx.Extensions.Inference),
+			Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseResponse),
+			Plugins:     plugins,
+			Identity:    pipeline.SnapshotIdentity(pctx),
+			Host:        pctx.Host,
+			StatusCode:  resp.StatusCode,
+			Error:       pipeline.DeriveError(pctx),
+			Duration:    pipeline.DurationSince(pctx.StartedAt),
 		}
-		if ev.MCP != nil || ev.Inference != nil {
+		if ev.MCP != nil || ev.Inference != nil || ev.Invocations != nil || plugins != nil {
 			s.Sessions.Append(sid, ev)
 		}
 	}
@@ -250,7 +267,7 @@ func (s *Server) recordOutboundReject(pctx *pipeline.Context, action pipeline.Ac
 		At:          time.Now(),
 		Direction:   pipeline.Outbound,
 		Phase:       pipeline.SessionDenied,
-		Invocations: pctx.Extensions.Invocations.FilteredByPhase(pipeline.InvocationPhaseRequest),
+		Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest),
 		Host:        pctx.Host,
 		StatusCode:  status,
 		Error: &pipeline.EventError{
@@ -262,24 +279,4 @@ func (s *Server) recordOutboundReject(pctx *pipeline.Context, action pipeline.Ac
 	s.Sessions.Append(sid, ev)
 }
 
-func extractBearer(authHeader string) string {
-	if len(authHeader) > 7 && strings.EqualFold(authHeader[:7], "bearer ") {
-		return authHeader[7:]
-	}
-	return ""
-}
 
-// writeRejection renders a pipeline Reject to the http.ResponseWriter.
-// Preserves the plugin's status, headers, and body — the old flat
-// {"error":reason} is gone; plugins now control the wire shape via
-// action.Violation.
-func writeRejection(w http.ResponseWriter, action pipeline.Action) {
-	status, headers, body := action.Violation.Render()
-	for k, vs := range headers {
-		for _, v := range vs {
-			w.Header().Add(k, v)
-		}
-	}
-	w.WriteHeader(status)
-	_, _ = w.Write(body)
-}
diff --git a/authbridge/authlib/listener/httpx/render.go b/authbridge/authlib/listener/httpx/render.go
new file mode 100644
index 000000000..740d24d9f
--- /dev/null
+++ b/authbridge/authlib/listener/httpx/render.go
@@ -0,0 +1,30 @@
+// Package httpx contains HTTP-listener helpers shared between the
+// forwardproxy and reverseproxy listeners. extproc and extauthz speak
+// gRPC and don't use this package.
+package httpx
+
+import (
+	"net/http"
+
+	"github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline"
+)
+
+// WriteRejection renders a pipeline.Reject Action onto an HTTP
+// ResponseWriter. Status, headers, and body all come from the action's
+// Violation — listeners hand the writer + action over and let the
+// pipeline-defined contract drive the response shape.
+//
+// Safe to call only when action.Violation is non-nil (i.e. the action
+// was Type=Reject). The forward/reverse proxy listeners only invoke
+// this on the Reject branch of their action switch, matching that
+// invariant.
+func WriteRejection(w http.ResponseWriter, action pipeline.Action) {
+	status, headers, body := action.Violation.Render()
+	for k, vs := range headers {
+		for _, v := range vs {
+			w.Header().Add(k, v)
+		}
+	}
+	w.WriteHeader(status)
+	_, _ = w.Write(body)
+}
diff --git a/authbridge/authlib/listener/reverseproxy/server.go b/authbridge/authlib/listener/reverseproxy/server.go
index 21c521108..da4c8ce0b 100644
--- a/authbridge/authlib/listener/reverseproxy/server.go
+++ b/authbridge/authlib/listener/reverseproxy/server.go
@@ -14,6 +14,7 @@ import (
 	"net/url"
 	"time"
 
+	"github.com/kagenti/kagenti-extensions/authbridge/authlib/listener/httpx"
 	"github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline"
 	"github.com/kagenti/kagenti-extensions/authbridge/authlib/session"
 )
@@ -107,7 +108,7 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) {
 	action := s.InboundPipeline.Run(r.Context(), pctx)
 	if action.Type == pipeline.Reject {
 		s.recordInboundReject(pctx, action)
-		writeRejection(w, action)
+		httpx.WriteRejection(w, action)
 		return
 	}
 
@@ -129,11 +130,19 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) {
 		if sid == "" {
 			sid = session.DefaultSessionID
 		}
+		// Snapshot-copy the protocol extension and use the shared helpers
+		// for plugin invocations / observability / identity. Mirrors what
+		// extproc does so request events don't pick up response-phase
+		// mutations on the same pctx.Extensions.A2A struct.
 		s.Sessions.Append(sid, pipeline.SessionEvent{
-			At:        time.Now(),
-			Direction: pipeline.Inbound,
-			Phase:     pipeline.SessionRequest,
-			A2A:       pctx.Extensions.A2A,
+			At:          time.Now(),
+			Direction:   pipeline.Inbound,
+			Phase:       pipeline.SessionRequest,
+			A2A:         pipeline.SnapshotA2A(pctx.Extensions.A2A),
+			Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest),
+			Plugins:     pipeline.SnapshotPlugins(pctx.Extensions.Custom),
+			Identity:    pipeline.SnapshotIdentity(pctx),
+			Host:        pctx.Host,
 		})
 	}
 
@@ -178,12 +187,54 @@ func (s *Server) modifyResponse(resp *http.Response) error {
 		resp.Header.Set("Content-Length", fmt.Sprintf("%d", len(pctx.ResponseBody)))
 		resp.Header.Del("Content-Encoding")
 	}
+
+	// Rekey the default bucket → A2A contextId when the response
+	// reveals one. The first turn of an A2A conversation arrives
+	// without a contextId (the agent assigns it on response), so the
+	// inbound request + any outbound MCP/inference calls during
+	// processing land in `default`. Without rekey those events stay
+	// orphaned while only the response goes to the contextId bucket.
+	// Mirrors extproc.rekeyInboundSession.
+	if s.Sessions != nil && pctx.Extensions.A2A != nil &&
+		pctx.Extensions.A2A.SessionID != "" &&
+		pctx.Extensions.A2A.SessionID != session.DefaultSessionID {
+		s.Sessions.Rekey(session.DefaultSessionID, pctx.Extensions.A2A.SessionID)
+	}
+
+	// Mirror forwardproxy's response-phase event so abctl pairs every
+	// inbound request with a response row. Without this, A2A
+	// `message/stream` requests show up as orphan request events.
+	// SSE responses still get recorded — the body is whatever the
+	// pipeline saw at this point (may be empty for streamed bodies),
+	// but the status code and plugin invocations are always meaningful.
+	if s.Sessions != nil && pctx.Extensions.A2A != nil {
+		sid := pctx.Extensions.A2A.SessionID
+		if sid == "" {
+			sid = s.Sessions.ActiveSession()
+		}
+		if sid == "" {
+			sid = session.DefaultSessionID
+		}
+		s.Sessions.Append(sid, pipeline.SessionEvent{
+			At:          time.Now(),
+			Direction:   pipeline.Inbound,
+			Phase:       pipeline.SessionResponse,
+			A2A:         pipeline.SnapshotA2A(pctx.Extensions.A2A),
+			Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseResponse),
+			Plugins:     pipeline.SnapshotPlugins(pctx.Extensions.Custom),
+			Identity:    pipeline.SnapshotIdentity(pctx),
+			Host:        pctx.Host,
+			StatusCode:  resp.StatusCode,
+			Error:       pipeline.DeriveError(pctx),
+			Duration:    pipeline.DurationSince(pctx.StartedAt),
+		})
+	}
 	return nil
 }
 
 func (s *Server) errorHandler(w http.ResponseWriter, _ *http.Request, err error) {
 	if rErr, ok := err.(*responseRejectedError); ok {
-		writeRejection(w, rErr.action)
+		httpx.WriteRejection(w, rErr.action)
 		return
 	}
 	http.Error(w, `{"error":"bad gateway"}`, http.StatusBadGateway)
@@ -229,7 +280,7 @@ func (s *Server) recordInboundReject(pctx *pipeline.Context, action pipeline.Act
 		At:          time.Now(),
 		Direction:   pipeline.Inbound,
 		Phase:       pipeline.SessionDenied,
-		Invocations: pctx.Extensions.Invocations.FilteredByPhase(pipeline.InvocationPhaseRequest),
+		Invocations: pipeline.SnapshotInvocations(pctx.Extensions.Invocations, pipeline.InvocationPhaseRequest),
 		Host:        pctx.Host,
 		StatusCode:  status,
 		Error: &pipeline.EventError{
@@ -241,40 +292,14 @@ func (s *Server) recordInboundReject(pctx *pipeline.Context, action pipeline.Act
 	s.Sessions.Append(sid, ev)
 }
 
-// writeRejection renders a pipeline Reject to the http.ResponseWriter,
-// preserving the plugin's status, headers, and body.
 // requestScheme derives the URL scheme for an incoming server-side
 // request. On server requests Go does not populate r.URL.Scheme (it's
 // only set for client-side / proxy requests where the full URL is on
 // the request line), so we read it from r.TLS instead: TLS present =>
-// https, absent => http.
-//
-// Contract note: this listener intentionally diverges from the
-// Context.Scheme godoc's "empty when undetermined" convention — it
-// always returns "http" or "https" based on r.TLS. The fallback is
-// confidently wrong when reverseproxy sits behind a TLS-terminating
-// upstream (LB, ingress): r.TLS is nil on the inner hop even though
-// the caller's actual scheme was https. Consumers that need the
-// caller's scheme in that topology should plumb X-Forwarded-Proto
-// once a trusted-upstream policy exists (not in this PR).
-//
-// Does not consult X-Forwarded-Proto. Honoring that header is only
-// safe when the upstream proxy is trusted; wiring a trust policy is
-// deferred until we have a concrete multi-hop deployment story.
+// https, absent => http. Reverse-proxy-specific helper.
 func requestScheme(r *http.Request) string {
 	if r.TLS != nil {
 		return "https"
 	}
 	return "http"
 }
-
-func writeRejection(w http.ResponseWriter, action pipeline.Action) {
-	status, headers, body := action.Violation.Render()
-	for k, vs := range headers {
-		for _, v := range vs {
-			w.Header().Add(k, v)
-		}
-	}
-	w.WriteHeader(status)
-	_, _ = w.Write(body)
-}
diff --git a/authbridge/authlib/pipeline/snapshot.go b/authbridge/authlib/pipeline/snapshot.go
new file mode 100644
index 000000000..079cc36d0
--- /dev/null
+++ b/authbridge/authlib/pipeline/snapshot.go
@@ -0,0 +1,137 @@
+package pipeline
+
+import (
+	"encoding/json"
+	"log/slog"
+	"strconv"
+	"strings"
+	"time"
+)
+
+// SnapshotA2A returns a shallow copy of ext. The record helpers attach
+// the snapshot to the SessionEvent rather than the live pointer so
+// response-phase mutations on pctx.Extensions.A2A (e.g. the parser
+// stamping the server-assigned contextId onto SessionID during
+// OnResponse) don't retroactively rewrite request-phase events that
+// were already appended. Slice fields are reused intentionally — they
+// are only assigned, never mutated in place, after the parser
+// completes.
+func SnapshotA2A(ext *A2AExtension) *A2AExtension {
+	if ext == nil {
+		return nil
+	}
+	c := *ext
+	return &c
+}
+
+// SnapshotMCP returns a shallow copy of ext. Important for outbound
+// request events: the same pctx.Extensions.MCP pointer receives Result
+// or Err on the response side, so without snapshotting, the
+// already-recorded request event would display the future response's
+// result map.
+func SnapshotMCP(ext *MCPExtension) *MCPExtension {
+	if ext == nil {
+		return nil
+	}
+	c := *ext
+	return &c
+}
+
+// SnapshotInference returns a shallow copy of ext. Scalar response
+// fields (Completion, FinishReason, *Tokens) get assigned on the live
+// extension during OnResponse; without snapshotting, the request event's
+// view would contain the eventual response's token counts and completion.
+func SnapshotInference(ext *InferenceExtension) *InferenceExtension {
+	if ext == nil {
+		return nil
+	}
+	c := *ext
+	return &c
+}
+
+// SnapshotInvocations is an alias for FilteredByPhase that participates
+// in the Snapshot* family for symmetry with the other shallow-copy
+// helpers. The underlying call already returns a fresh slice.
+func SnapshotInvocations(ext *Invocations, phase InvocationPhase) *Invocations {
+	return ext.FilteredByPhase(phase)
+}
+
+// SnapshotPlugins collects plugin-public observability events from
+// pctx.Extensions.Custom entries whose keys end in PluginEventSuffix.
+// Each matching value is json.Marshaled into the wire-form map under
+// the plugin name (suffix stripped). Marshal errors downgrade to slog
+// Debug and skip the entry rather than aborting recording — that keeps
+// a misbehaving plugin from taking out the whole session stream.
+func SnapshotPlugins(custom map[string]any) map[string]json.RawMessage {
+	if len(custom) == 0 {
+		return nil
+	}
+	var out map[string]json.RawMessage
+	for k, v := range custom {
+		if !strings.HasSuffix(k, PluginEventSuffix) {
+			continue
+		}
+		raw, err := json.Marshal(v)
+		if err != nil {
+			slog.Debug("session: skipping non-marshalable plugin event",
+				"key", k, "error", err)
+			continue
+		}
+		if out == nil {
+			out = make(map[string]json.RawMessage)
+		}
+		pluginName := strings.TrimSuffix(k, PluginEventSuffix)
+		out[pluginName] = raw
+	}
+	return out
+}
+
+// SnapshotIdentity copies the caller identity off pctx so the session
+// event stays valid after pctx is discarded. Returns nil when no
+// identity information is available (e.g., jwt-validation didn't run
+// on this path and no agent identity was attached).
+func SnapshotIdentity(pctx *Context) *EventIdentity {
+	if pctx.Identity == nil && pctx.Agent == nil {
+		return nil
+	}
+	id := &EventIdentity{}
+	if pctx.Identity != nil {
+		id.Subject = pctx.Identity.Subject()
+		id.ClientID = pctx.Identity.ClientID()
+		if scopes := pctx.Identity.Scopes(); len(scopes) > 0 {
+			id.Scopes = append([]string(nil), scopes...)
+		}
+	}
+	if pctx.Agent != nil {
+		id.AgentID = pctx.Agent.WorkloadID
+	}
+	return id
+}
+
+// DurationSince returns the elapsed time since start, or 0 when start
+// is zero (pctx constructed without wall-clock stamping, e.g. in unit
+// tests).
+func DurationSince(start time.Time) time.Duration {
+	if start.IsZero() {
+		return 0
+	}
+	return time.Since(start)
+}
+
+// DeriveError constructs an EventError from response-side signals.
+// Returns nil for 2xx / no guardrail block / no parser error.
+func DeriveError(pctx *Context) *EventError {
+	if pctx.Extensions.Security != nil && pctx.Extensions.Security.Blocked {
+		return &EventError{
+			Kind:    "blocked",
+			Message: pctx.Extensions.Security.BlockReason,
+		}
+	}
+	if pctx.StatusCode >= 400 {
+		return &EventError{
+			Kind: "backend_error",
+			Code: strconv.Itoa(pctx.StatusCode),
+		}
+	}
+	return nil
+}

From cbc3150343783bbd8febe7fe8984429488d966f0 Mon Sep 17 00:00:00 2001
From: Hai Huang <huang195@gmail.com>
Date: Thu, 14 May 2026 22:30:39 -0400
Subject: [PATCH 2/3] =?UTF-8?q?test:=20address=20PR=20#412=20review=20feed?=
 =?UTF-8?q?back=20=E2=80=94=20snapshot=20+=20rekey=20unit=20tests,=20comme?=
 =?UTF-8?q?nts?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Picked up suggestions from the PR #412 review:

* New authlib/pipeline/snapshot_test.go covers the helpers that
  this PR newly exported. Specifically:
    - SnapshotA2A / SnapshotMCP / SnapshotInference: protect the
      no-pointer-aliasing contract by mutating the source after
      snapshot and asserting the snapshot is unchanged.
    - SnapshotIdentity: nil case + Identity-populated case +
      Agent-populated case, plus a check that the snapshot's Scopes
      slice is independently owned (mutating it can't bleed back to
      the source).
    - SnapshotPlugins: empty input, /event-suffix filter +
      suffix-stripping, isolation against non-marshalable values
      (one bad plugin event must not drop the rest).
    - DeriveError: Blocked branch, non-2xx branch, 2xx no-op, and
      Blocked-takes-precedence-over-StatusCode policy.
    - DurationSince: zero-time short-circuit + valid start.

* New TestReverseProxy_ModifyResponse_RekeyAndResponseEvent in the
  reverseproxy test suite locks in the modifyResponse behavior
  introduced by this PR — without it, future refactors of the
  event-construction path could silently regress to "request lands
  in default, response lands in a fresh contextId bucket, no merge."
  The test drives a real httptest server, uses an inline a2aStamp
  plugin to mimic the parser stamping a SessionID during OnResponse,
  and asserts (a) Sessions.Rekey moved events out of `default` into
  the contextId bucket, and (b) a SessionResponse event was
  appended carrying StatusCode + Invocations.

* Comments added at:
    - forwardproxy.go widened append guard
      (MCP || Inference || Invocations || plugins != nil) — noted
      that each clause is intentional and removing one silently
      regresses event coverage.
    - reverseproxy.go modifyResponse rekey condition — noted why
      empty SessionID and "default" SessionID are skipped.

The Tier 2 dedup work for recordInboundReject / recordOutboundReject
is filed as kagenti/kagenti-extensions#413 so it doesn't fall off
the radar.

Assisted-By: Claude (Anthropic AI) <noreply@anthropic.com>
Signed-off-by: Hai Huang <huang195@gmail.com>
---
 .../authlib/listener/forwardproxy/server.go   |   9 +
 .../authlib/listener/reverseproxy/server.go   |   4 +
 .../listener/reverseproxy/server_test.go      |  97 +++++++
 authbridge/authlib/pipeline/snapshot_test.go  | 266 ++++++++++++++++++
 4 files changed, 376 insertions(+)
 create mode 100644 authbridge/authlib/pipeline/snapshot_test.go

diff --git a/authbridge/authlib/listener/forwardproxy/server.go b/authbridge/authlib/listener/forwardproxy/server.go
index e9fb67ddf..f53b704a0 100644
--- a/authbridge/authlib/listener/forwardproxy/server.go
+++ b/authbridge/authlib/listener/forwardproxy/server.go
@@ -120,6 +120,13 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) {
 			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)
 		}
@@ -219,6 +226,8 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) {
 			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)
 		}
diff --git a/authbridge/authlib/listener/reverseproxy/server.go b/authbridge/authlib/listener/reverseproxy/server.go
index da4c8ce0b..76ae3e10c 100644
--- a/authbridge/authlib/listener/reverseproxy/server.go
+++ b/authbridge/authlib/listener/reverseproxy/server.go
@@ -195,6 +195,10 @@ func (s *Server) modifyResponse(resp *http.Response) error {
 	// processing land in `default`. Without rekey those events stay
 	// orphaned while only the response goes to the contextId bucket.
 	// Mirrors extproc.rekeyInboundSession.
+	//
+	// Skip when SessionID is empty (auth-only or non-A2A response —
+	// no contextId to merge against) or already "default" (a no-op
+	// that would also collide with the source bucket name).
 	if s.Sessions != nil && pctx.Extensions.A2A != nil &&
 		pctx.Extensions.A2A.SessionID != "" &&
 		pctx.Extensions.A2A.SessionID != session.DefaultSessionID {
diff --git a/authbridge/authlib/listener/reverseproxy/server_test.go b/authbridge/authlib/listener/reverseproxy/server_test.go
index a45d81dde..c2ee2ce5d 100644
--- a/authbridge/authlib/listener/reverseproxy/server_test.go
+++ b/authbridge/authlib/listener/reverseproxy/server_test.go
@@ -465,3 +465,100 @@ func TestReverseProxy_SchemeFromTLS(t *testing.T) {
 		})
 	}
 }
+
+// a2aStampPlugin is a tiny inline plugin that simulates a2a-parser's
+// two-phase behavior:
+//   - OnRequest populates pctx.Extensions.A2A so the request-side
+//     session-event gate fires (no SessionID yet — first turn).
+//   - OnResponse stamps the server-assigned SessionID, mimicking
+//     extractSessionID(pctx.ResponseBody) in the real parser. This is
+//     what triggers the rekey path in modifyResponse.
+type a2aStampPlugin struct {
+	method      string
+	responseSID string // SessionID to stamp during OnResponse
+}
+
+func (p *a2aStampPlugin) Name() string { return "a2a-stamp" }
+func (p *a2aStampPlugin) Capabilities() pipeline.PluginCapabilities {
+	return pipeline.PluginCapabilities{}
+}
+func (p *a2aStampPlugin) OnRequest(_ context.Context, pctx *pipeline.Context) pipeline.Action {
+	pctx.Extensions.A2A = &pipeline.A2AExtension{Method: p.method}
+	pctx.Observe("matched_request")
+	return pipeline.Action{Type: pipeline.Continue}
+}
+func (p *a2aStampPlugin) OnResponse(_ context.Context, pctx *pipeline.Context) pipeline.Action {
+	if pctx.Extensions.A2A != nil {
+		pctx.Extensions.A2A.SessionID = p.responseSID
+	}
+	pctx.Observe("matched_response")
+	return pipeline.Action{Type: pipeline.Continue}
+}
+
+// TestReverseProxy_ModifyResponse_RekeyAndResponseEvent locks in the
+// behavior that an A2A first-turn (request without contextId, agent
+// assigns contextId in response) ends up with all the turn's events
+// merged into the contextId bucket, plus a SessionResponse event
+// recorded. Without rekey + the response-side append, abctl users see
+// orphan request rows in `default` and no inbound response row at all.
+func TestReverseProxy_ModifyResponse_RekeyAndResponseEvent(t *testing.T) {
+	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte(`{"ok":true}`))
+	}))
+	defer backend.Close()
+
+	const ctxID = "ctx-1234"
+	stamp := &a2aStampPlugin{method: "message/stream", responseSID: ctxID}
+	p, err := plugintesting.BuildPipeline([]pipeline.Plugin{stamp})
+	if err != nil {
+		t.Fatalf("BuildPipeline: %v", err)
+	}
+
+	store := session.New(30*time.Minute, 100, 100)
+	srv, err := NewServer(pipeline.NewHolder(p), store, backend.URL)
+	if err != nil {
+		t.Fatalf("NewServer: %v", err)
+	}
+	proxy := httptest.NewServer(srv.Handler())
+	defer proxy.Close()
+
+	resp, err := http.DefaultClient.Get(proxy.URL + "/api")
+	if err != nil {
+		t.Fatalf("Get: %v", err)
+	}
+	resp.Body.Close()
+
+	// After the response: the `default` bucket should be gone (rekeyed
+	// to ctxID) and the ctxID bucket should hold both events from this
+	// turn (request + response).
+	if v := store.View(session.DefaultSessionID); v != nil {
+		t.Errorf("default bucket still exists with %d events; expected rekey to %q", len(v.Events), ctxID)
+	}
+	view := store.View(ctxID)
+	if view == nil {
+		t.Fatalf("session %q not found; rekey did not run", ctxID)
+	}
+	if len(view.Events) != 2 {
+		t.Fatalf("session %q has %d events, want 2 (request + response)", ctxID, len(view.Events))
+	}
+
+	gotPhases := []pipeline.SessionPhase{view.Events[0].Phase, view.Events[1].Phase}
+	wantPhases := []pipeline.SessionPhase{pipeline.SessionRequest, pipeline.SessionResponse}
+	for i := range wantPhases {
+		if gotPhases[i] != wantPhases[i] {
+			t.Errorf("event %d phase = %v, want %v", i, gotPhases[i], wantPhases[i])
+		}
+	}
+
+	// Response event must carry the StatusCode + populated Invocations
+	// (the response-phase "observe" we appended in OnResponse) — those
+	// are the bits that previously dropped out.
+	respEvent := view.Events[1]
+	if respEvent.StatusCode != http.StatusOK {
+		t.Errorf("response StatusCode = %d, want 200", respEvent.StatusCode)
+	}
+	if respEvent.Invocations == nil || len(respEvent.Invocations.Inbound) == 0 {
+		t.Errorf("response event has no Invocations; expected one a2a-stamp observe entry")
+	}
+}
diff --git a/authbridge/authlib/pipeline/snapshot_test.go b/authbridge/authlib/pipeline/snapshot_test.go
new file mode 100644
index 000000000..8c870a124
--- /dev/null
+++ b/authbridge/authlib/pipeline/snapshot_test.go
@@ -0,0 +1,266 @@
+package pipeline
+
+import (
+	"encoding/json"
+	"testing"
+	"time"
+)
+
+// TestSnapshotA2A_PointerAliasing protects the no-aliasing contract:
+// after Snapshot, mutating the source must not affect the snapshot. The
+// proxy listeners depend on this so a request-phase event isn't
+// retroactively rewritten when OnResponse mutates the same A2A struct
+// (the symptom that triggered this work).
+func TestSnapshotA2A_PointerAliasing(t *testing.T) {
+	src := &A2AExtension{
+		Method:      "message/stream",
+		SessionID:   "ctx-1",
+		FinalStatus: "",
+		Artifact:    "",
+	}
+	snap := SnapshotA2A(src)
+
+	// Mutate source after snapshot — common pattern: parser stamps
+	// SessionID + final fields on the live ext during OnResponse.
+	src.SessionID = "ctx-2"
+	src.FinalStatus = "completed"
+	src.Artifact = "the answer is 42"
+
+	if snap.SessionID != "ctx-1" {
+		t.Errorf("SessionID = %q, want %q (snapshot retroactively mutated)", snap.SessionID, "ctx-1")
+	}
+	if snap.FinalStatus != "" {
+		t.Errorf("FinalStatus = %q, want empty (snapshot picked up response-phase mutation)", snap.FinalStatus)
+	}
+	if snap.Artifact != "" {
+		t.Errorf("Artifact = %q, want empty (snapshot picked up response-phase mutation)", snap.Artifact)
+	}
+}
+
+func TestSnapshotA2A_Nil(t *testing.T) {
+	if got := SnapshotA2A(nil); got != nil {
+		t.Errorf("SnapshotA2A(nil) = %v, want nil", got)
+	}
+}
+
+func TestSnapshotMCP_PointerAliasing(t *testing.T) {
+	src := &MCPExtension{Method: "tools/call"}
+	snap := SnapshotMCP(src)
+
+	// Result map is assigned on the live ext during OnResponse.
+	src.Result = map[string]any{"content": "tool output"}
+	src.Err = &MCPError{Code: -32000, Message: "boom"}
+
+	if snap.Result != nil {
+		t.Errorf("Result = %v, want nil (snapshot picked up response-phase Result)", snap.Result)
+	}
+	if snap.Err != nil {
+		t.Errorf("Err = %v, want nil (snapshot picked up response-phase Err)", snap.Err)
+	}
+}
+
+func TestSnapshotMCP_Nil(t *testing.T) {
+	if got := SnapshotMCP(nil); got != nil {
+		t.Errorf("SnapshotMCP(nil) = %v, want nil", got)
+	}
+}
+
+func TestSnapshotInference_PointerAliasing(t *testing.T) {
+	src := &InferenceExtension{Model: "llama3.2"}
+	snap := SnapshotInference(src)
+
+	// Completion + token counts are assigned on the live ext during OnResponse.
+	src.Completion = "the answer is 42"
+	src.PromptTokens = 100
+	src.CompletionTokens = 25
+
+	if snap.Completion != "" {
+		t.Errorf("Completion = %q, want empty (snapshot picked up response-phase mutation)", snap.Completion)
+	}
+	if snap.PromptTokens != 0 {
+		t.Errorf("PromptTokens = %d, want 0 (snapshot picked up response-phase mutation)", snap.PromptTokens)
+	}
+	if snap.CompletionTokens != 0 {
+		t.Errorf("CompletionTokens = %d, want 0 (snapshot picked up response-phase mutation)", snap.CompletionTokens)
+	}
+}
+
+func TestSnapshotInference_Nil(t *testing.T) {
+	if got := SnapshotInference(nil); got != nil {
+		t.Errorf("SnapshotInference(nil) = %v, want nil", got)
+	}
+}
+
+// stubIdentity is a minimal pipeline.Identity implementation for tests.
+type stubIdentity struct {
+	subject  string
+	clientID string
+	scopes   []string
+}
+
+func (s stubIdentity) Subject() string  { return s.subject }
+func (s stubIdentity) ClientID() string { return s.clientID }
+func (s stubIdentity) Scopes() []string { return s.scopes }
+
+func TestSnapshotIdentity_NoIdentity(t *testing.T) {
+	pctx := &Context{}
+	if got := SnapshotIdentity(pctx); got != nil {
+		t.Errorf("SnapshotIdentity(empty pctx) = %v, want nil", got)
+	}
+}
+
+func TestSnapshotIdentity_FromIdentity(t *testing.T) {
+	pctx := &Context{
+		Identity: stubIdentity{
+			subject:  "alice",
+			clientID: "kagenti-ui",
+			scopes:   []string{"openid", "agent-aud"},
+		},
+	}
+	got := SnapshotIdentity(pctx)
+	if got == nil {
+		t.Fatal("SnapshotIdentity returned nil despite populated Identity")
+	}
+	if got.Subject != "alice" || got.ClientID != "kagenti-ui" {
+		t.Errorf("Subject/ClientID = %q/%q, want alice/kagenti-ui", got.Subject, got.ClientID)
+	}
+	if len(got.Scopes) != 2 || got.Scopes[0] != "openid" {
+		t.Errorf("Scopes = %v, want [openid agent-aud]", got.Scopes)
+	}
+
+	// Mutating the snapshot's Scopes slice must not bleed back to the
+	// caller's identity (we copy the slice).
+	got.Scopes[0] = "ZAP"
+	if pctx.Identity.Scopes()[0] != "openid" {
+		t.Error("snapshot mutation bled back to source Identity.Scopes")
+	}
+}
+
+func TestSnapshotIdentity_FromAgent(t *testing.T) {
+	pctx := &Context{Agent: &AgentIdentity{WorkloadID: "team1/weather-agent"}}
+	got := SnapshotIdentity(pctx)
+	if got == nil {
+		t.Fatal("SnapshotIdentity returned nil despite populated Agent")
+	}
+	if got.AgentID != "team1/weather-agent" {
+		t.Errorf("AgentID = %q, want team1/weather-agent", got.AgentID)
+	}
+}
+
+func TestSnapshotPlugins_Empty(t *testing.T) {
+	if got := SnapshotPlugins(nil); got != nil {
+		t.Errorf("SnapshotPlugins(nil) = %v, want nil", got)
+	}
+	if got := SnapshotPlugins(map[string]any{}); got != nil {
+		t.Errorf("SnapshotPlugins(empty) = %v, want nil", got)
+	}
+}
+
+func TestSnapshotPlugins_FilterAndStripSuffix(t *testing.T) {
+	custom := map[string]any{
+		"rate-limiter" + PluginEventSuffix:  map[string]any{"remaining": 42},
+		"audit" + PluginEventSuffix:         "logged",
+		"some-internal-key":                 "should-not-appear", // no /event suffix
+		"more.internal.state":               struct{ X int }{X: 1},
+	}
+	got := SnapshotPlugins(custom)
+
+	if _, ok := got["some-internal-key"]; ok {
+		t.Error("non-/event key leaked into Plugins map")
+	}
+	if _, ok := got["more.internal.state"]; ok {
+		t.Error("non-/event key leaked into Plugins map")
+	}
+	if _, ok := got["rate-limiter"]; !ok {
+		t.Error("rate-limiter key missing (suffix should be stripped)")
+	}
+	if _, ok := got["audit"]; !ok {
+		t.Error("audit key missing (suffix should be stripped)")
+	}
+
+	// Round-trip the rate-limiter value and verify it preserved structure.
+	var rl struct {
+		Remaining int `json:"remaining"`
+	}
+	if err := json.Unmarshal(got["rate-limiter"], &rl); err != nil {
+		t.Fatalf("unmarshal rate-limiter: %v", err)
+	}
+	if rl.Remaining != 42 {
+		t.Errorf("Remaining = %d, want 42", rl.Remaining)
+	}
+}
+
+func TestSnapshotPlugins_SkipsUnmarshalable(t *testing.T) {
+	// channels can't json.Marshal — should be skipped, not abort recording.
+	custom := map[string]any{
+		"goodplugin" + PluginEventSuffix: map[string]string{"ok": "yes"},
+		"badplugin" + PluginEventSuffix:  make(chan int),
+	}
+	got := SnapshotPlugins(custom)
+	if _, ok := got["badplugin"]; ok {
+		t.Error("non-marshalable value should have been skipped")
+	}
+	if _, ok := got["goodplugin"]; !ok {
+		t.Error("good plugin event was dropped along with the bad one — should be independent")
+	}
+}
+
+func TestDeriveError_Blocked(t *testing.T) {
+	pctx := &Context{
+		Extensions: Extensions{
+			Security: &SecurityExtension{Blocked: true, BlockReason: "PII detected"},
+		},
+	}
+	got := DeriveError(pctx)
+	if got == nil {
+		t.Fatal("DeriveError returned nil for Blocked=true")
+	}
+	if got.Kind != "blocked" || got.Message != "PII detected" {
+		t.Errorf("Kind/Message = %q/%q, want blocked/PII detected", got.Kind, got.Message)
+	}
+}
+
+func TestDeriveError_Non2xx(t *testing.T) {
+	pctx := &Context{StatusCode: 502}
+	got := DeriveError(pctx)
+	if got == nil {
+		t.Fatal("DeriveError returned nil for StatusCode=502")
+	}
+	if got.Kind != "backend_error" || got.Code != "502" {
+		t.Errorf("Kind/Code = %q/%q, want backend_error/502", got.Kind, got.Code)
+	}
+}
+
+func TestDeriveError_2xx(t *testing.T) {
+	pctx := &Context{StatusCode: 200}
+	if got := DeriveError(pctx); got != nil {
+		t.Errorf("DeriveError(200) = %+v, want nil", got)
+	}
+}
+
+func TestDeriveError_BlockedTakesPrecedenceOverStatusCode(t *testing.T) {
+	// A guardrail block returning 200 still produces a "blocked" error —
+	// the failure mode is policy, not the eventual HTTP status.
+	pctx := &Context{
+		StatusCode: 200,
+		Extensions: Extensions{Security: &SecurityExtension{Blocked: true, BlockReason: "policy"}},
+	}
+	got := DeriveError(pctx)
+	if got == nil || got.Kind != "blocked" {
+		t.Errorf("expected blocked error to win over 200 StatusCode, got %+v", got)
+	}
+}
+
+func TestDurationSince_ZeroStart(t *testing.T) {
+	if got := DurationSince(time.Time{}); got != 0 {
+		t.Errorf("DurationSince(zero) = %v, want 0", got)
+	}
+}
+
+func TestDurationSince_ValidStart(t *testing.T) {
+	start := time.Now().Add(-50 * time.Millisecond)
+	got := DurationSince(start)
+	if got < 40*time.Millisecond || got > 5*time.Second {
+		t.Errorf("DurationSince(50ms ago) = %v, want roughly ≥40ms", got)
+	}
+}

From 96533a3eb2a5d485134904c89beb1586d89f3cd8 Mon Sep 17 00:00:00 2001
From: Hai Huang <huang195@gmail.com>
Date: Fri, 15 May 2026 08:49:42 -0400
Subject: [PATCH 3/3] =?UTF-8?q?docs:=20address=20PR=20#412=20review=20?=
 =?UTF-8?q?=E2=80=94=20restore=20requestScheme=20contract=20notes,=20docum?=
 =?UTF-8?q?ent=20A2A-only=20inbound=20gate?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two non-blocking suggestions from esnible's PR #412 review:

1. requestScheme godoc — restore the contract note paragraphs that the
   refactor's writeRejection-comment dedup swept up by accident:
   - Intentional divergence from Context.Scheme's "empty when
     undetermined" convention (this listener always returns http/https
     based on r.TLS).
   - Failure mode behind a TLS-terminating ingress: r.TLS is nil on the
     inner hop even though the caller's actual scheme was https.
   - Deferred X-Forwarded-Proto trust-policy plan.

   These are design notes future readers need when deciding whether to
   trust this scheme value; pruning was unintentional.

2. Inbound A2A gate — add a one-line contract comment at the
   `pctx.Extensions.A2A != nil` gate explaining the asymmetry with
   forwardproxy. Reverseproxy is A2A-only by design (session keying
   and modifyResponse rekey logic are A2A-specific). Forwardproxy
   widens its analogous gate to MCP/Inference/Invocations/plugins
   because outbound traffic is not A2A-only. Non-A2A inbound is
   intentionally not recorded here.

Comment-only changes; no behavior change. Build + vet + tests all
pass on authlib/listener/reverseproxy and authlib/pipeline.

Assisted-By: Claude (Anthropic AI) <noreply@anthropic.com>
Signed-off-by: Hai Huang <huang195@gmail.com>
---
 .../authlib/listener/reverseproxy/server.go   | 21 ++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/authbridge/authlib/listener/reverseproxy/server.go b/authbridge/authlib/listener/reverseproxy/server.go
index 76ae3e10c..b28150fe1 100644
--- a/authbridge/authlib/listener/reverseproxy/server.go
+++ b/authbridge/authlib/listener/reverseproxy/server.go
@@ -122,6 +122,12 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) {
 		r.Header.Del("Content-Encoding")
 	}
 
+	// Inbound recording is gated on A2A by design: reverseproxy is the
+	// A2A-only listener (its session keying and rekey logic are A2A-specific
+	// — see modifyResponse). Forwardproxy widens the analogous gate to
+	// cover MCP/Inference/Invocations/plugins because outbound traffic is
+	// not A2A-only. A non-A2A inbound, or an A2A request that fails to
+	// parse, is intentionally not recorded here.
 	if s.Sessions != nil && pctx.Extensions.A2A != nil {
 		sid := pctx.Extensions.A2A.SessionID
 		if sid == "" {
@@ -300,7 +306,20 @@ func (s *Server) recordInboundReject(pctx *pipeline.Context, action pipeline.Act
 // request. On server requests Go does not populate r.URL.Scheme (it's
 // only set for client-side / proxy requests where the full URL is on
 // the request line), so we read it from r.TLS instead: TLS present =>
-// https, absent => http. Reverse-proxy-specific helper.
+// https, absent => http.
+//
+// Contract note: this listener intentionally diverges from the
+// Context.Scheme godoc's "empty when undetermined" convention — it
+// always returns "http" or "https" based on r.TLS. The fallback is
+// confidently wrong when reverseproxy sits behind a TLS-terminating
+// upstream (LB, ingress): r.TLS is nil on the inner hop even though
+// the caller's actual scheme was https. Consumers that need the
+// caller's scheme in that topology should plumb X-Forwarded-Proto
+// once a trusted-upstream policy exists (not in this PR).
+//
+// Does not consult X-Forwarded-Proto. Honoring that header is only
+// safe when the upstream proxy is trusted; wiring a trust policy is
+// deferred until we have a concrete multi-hop deployment story.
 func requestScheme(r *http.Request) string {
 	if r.TLS != nil {
 		return "https"