From 6adf7465bb9197a09d20a1ff051bdd0ea38dc916 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 30 Sep 2026 09:34:55 -0400 Subject: [PATCH 1/8] test: Pin in-cluster session resolution before client affinity lands Adds forwardproxy.Server.ClientAffinity with no behaviour yet, and a table test that runs the in-cluster A2A correlation cases with it off and on and requires the same answers from both: an A2A agent's outbound calls carry no coding-agent header and no recognised User-Agent, so nothing the knob adds may reach them. IBAC and sparc read that correlation. A handler-level twin drives the same case through serveOutbound, so a resolver that agreed in isolation but was bypassed on the request path still fails. Assisted-By: Claude (Anthropic AI) Signed-off-by: Hai Huang --- core/listener/forwardproxy/server.go | 7 + .../forwardproxy/session_affinity_test.go | 125 ++++++++++++++++++ 2 files changed, 132 insertions(+) create mode 100644 core/listener/forwardproxy/session_affinity_test.go diff --git a/core/listener/forwardproxy/server.go b/core/listener/forwardproxy/server.go index 69a0bf2f7..9089273de 100644 --- a/core/listener/forwardproxy/server.go +++ b/core/listener/forwardproxy/server.go @@ -101,6 +101,13 @@ type Server struct { // SessionIDHeaders for the default and how to turn it off. SessionIDHeaders []string + // ClientAffinity files a request that carries no session header under the newest + // session of the SAME coding agent, instead of under ActiveSession() — the one + // global "most recently updated" id, which with two agents running files each + // one's header-less calls into the other's session. See config.SessionConfig + // ClientAffinity; false keeps today's resolution byte for byte. + ClientAffinity bool + // bufferedFallbackOnce keeps the SSE-buffered-path notice to one line per // process; the condition is a supported chain shape, not an error. bufferedFallbackOnce sync.Once diff --git a/core/listener/forwardproxy/session_affinity_test.go b/core/listener/forwardproxy/session_affinity_test.go new file mode 100644 index 000000000..3f201e4c8 --- /dev/null +++ b/core/listener/forwardproxy/session_affinity_test.go @@ -0,0 +1,125 @@ +package forwardproxy + +import ( + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/rossoctl/cortex/core/pipeline" + "github.com/rossoctl/cortex/core/session" +) + +// TestClientAffinity_InClusterResolutionIsUnchanged is a characterization test: it pins +// the in-cluster answers with the knob OFF and asserts the SAME answers with it ON. +// +// In-cluster, an A2A agent's outbound calls carry no coding-agent header and no +// recognised User-Agent, and ActiveSession() is what correlates them with the inbound +// turn that caused them — IBAC and sparc depend on it. ClientAffinity must not reach +// that case at all: every row here is traffic from no known coding agent. +func TestClientAffinity_InClusterResolutionIsUnchanged(t *testing.T) { + cases := []struct { + name string + inbound bool // seed an inbound A2A turn under conv-A + ua string // "" keeps Go's default User-Agent + header string // X-Claude-Code-Session-Id, "" for none + nilHeaders bool // the transparent path: no request to read + wantPlugin string + wantRecord string + }{ + {name: "A2A turn, Go client", inbound: true, wantPlugin: "conv-A", wantRecord: "conv-A"}, + {name: "A2A turn, agent library UA", inbound: true, ua: "python-httpx/0.27.0", wantPlugin: "conv-A", wantRecord: "conv-A"}, + {name: "A2A turn, transparent", inbound: true, nilHeaders: true, wantPlugin: "conv-A", wantRecord: "conv-A"}, + {name: "nothing active", wantPlugin: "", wantRecord: session.DefaultSessionID}, + {name: "header wins over the A2A turn", inbound: true, header: "hdr-1", wantPlugin: "hdr-1", wantRecord: "hdr-1"}, + } + for _, affinity := range []bool{false, true} { + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + store := session.New(5*time.Minute, 100, 0) + defer store.Close() + if tc.inbound { + store.Append("conv-A", pipeline.SessionEvent{ + At: time.Now(), + Direction: pipeline.Inbound, + Phase: pipeline.SessionRequest, + A2A: &pipeline.A2AExtension{Method: "message/send", SessionID: "conv-A"}, + }) + } + s := &Server{ + Sessions: store, + SessionIDHeaders: []string{session.ClaudeCodeSessionHeader}, + ClientAffinity: affinity, + } + var h http.Header + if !tc.nilHeaders { + h = http.Header{} + if tc.ua != "" { + h.Set("User-Agent", tc.ua) + } + if tc.header != "" { + h.Set(session.ClaudeCodeSessionHeader, tc.header) + } + } + if got := s.resolvePluginSessionID(h); got != tc.wantPlugin { + t.Errorf("affinity=%v: plugin identity = %q, want %q", affinity, got, tc.wantPlugin) + } + if got := s.resolveOutboundSessionID(h); got != tc.wantRecord { + t.Errorf("affinity=%v: recording bucket = %q, want %q", affinity, got, tc.wantRecord) + } + }) + } + } +} + +// TestClientAffinity_InClusterPluginIdentityThroughTheHandler drives the same A2A case +// through the real handler with the knob ON, so a resolver that agreed in isolation but +// was bypassed on the request path would still fail here. +func TestClientAffinity_InClusterPluginIdentityThroughTheHandler(t *testing.T) { + store := session.New(5*time.Minute, 100, 0) + defer store.Close() + store.Append("conv-A", pipeline.SessionEvent{ + At: time.Now(), + Direction: pipeline.Inbound, + Phase: pipeline.SessionRequest, + A2A: &pipeline.A2AExtension{Method: "message/send", SessionID: "conv-A"}, + }) + probe, client, backendURL := newAffinityProxy(t, store) + + get(t, client, backendURL, "") + + if len(probe.sawID) != 1 || probe.sawID[0] != "conv-A" { + t.Fatalf("plugin saw %v, want [conv-A] (A2A correlation must survive client_affinity)", probe.sawID) + } + if v := store.View("conv-A"); v == nil || len(v.Events) != 3 { + n := -1 + if v != nil { + n = len(v.Events) + } + t.Fatalf("conv-A holds %d events, want 3 (inbound turn, outbound request, its response)", n) + } +} + +// newAffinityProxy is newProbedProxy with ClientAffinity on. +func newAffinityProxy(t *testing.T, store *session.Store) (*identityProbePlugin, *http.Client, string) { + t.Helper() + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(backend.Close) + probe := &identityProbePlugin{} + p, err := pipeline.New([]pipeline.Plugin{probe}) + if err != nil { + t.Fatal(err) + } + srv := &Server{ + OutboundPipeline: pipeline.NewHolder(p), + Sessions: store, + Client: http.DefaultClient, + SessionIDHeaders: []string{session.ClaudeCodeSessionHeader, session.BobSessionHeader}, + ClientAffinity: true, + } + proxy := httptest.NewServer(srv.Handler()) + t.Cleanup(proxy.Close) + return probe, &http.Client{Transport: &http.Transport{Proxy: http.ProxyURL(mustParseURL(proxy.URL))}}, backend.URL +} From b73512a255ff75d7918063b457f83978a7d02a04 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 30 Sep 2026 09:35:44 -0400 Subject: [PATCH 2/8] feat: Name the coding agent a header-less request belongs to EventClient.AffinityName answers which agent's session a request with no session header should join: Name when the User-Agent was recognised, else a known agent's canonical name found in any token, comments included. Wider than Name on purpose and only here. Claude Code's WebFetch sends "Claude-User (claude-code/2.1.284; ...)" and names its agent only in the comment ParseUserAgent skips. Folding that into Name would move those rows to a different agent in the ledger and the AGENTS pane than their own history; TestAffinityName_LeavesTheLabelAlone pins Label unchanged. Assisted-By: Claude (Anthropic AI) Signed-off-by: Hai Huang --- core/pipeline/client.go | 31 ++++++++++++++++++++++++++ core/pipeline/client_test.go | 43 ++++++++++++++++++++++++++++++++++++ 2 files changed, 74 insertions(+) diff --git a/core/pipeline/client.go b/core/pipeline/client.go index 3193a20ae..d02af0e2b 100644 --- a/core/pipeline/client.go +++ b/core/pipeline/client.go @@ -448,3 +448,34 @@ func (c *EventClient) Label() string { } return c.Name + "/" + c.Version } + +// AffinityName is the coding agent a request belongs to for SESSION ATTRIBUTION: Name when +// the User-Agent was recognised, else a known agent's canonical name found in any token — +// comments included — else "". See session.Store.SessionForClient for what it decides. +// +// WIDER THAN Name, AND ONLY HERE. Claude Code's WebFetch sends +// "Claude-User (claude-code/2.1.284; …)", naming its agent only inside the comment the parser +// deliberately skips. Reading comments into Name would move those rows to a different agent in +// the ledger and the AGENTS pane than the history beside them; for attribution the only +// question is which agent's session a header-less call joins, and there the comment is right. +// +// Canonical names, not knownClients' keys: the comment says "claude-code", while the key is +// the product token "claude-cli". Nil-safe, like Label. +func (c *EventClient) AffinityName() string { + if c == nil { + return "" + } + if c.Name != "" { + return c.Name + } + for _, f := range strings.Fields(c.Raw) { + product, _, _ := strings.Cut(strings.Trim(f, "();,"), "/") + product = strings.ToLower(product) + for _, name := range knownClients { + if product == name { + return name + } + } + } + return "" +} diff --git a/core/pipeline/client_test.go b/core/pipeline/client_test.go index 0b0ee089b..3da999c65 100644 --- a/core/pipeline/client_test.go +++ b/core/pipeline/client_test.go @@ -580,3 +580,46 @@ func TestContextClientInfo_MemoizesTheNilCase(t *testing.T) { t.Errorf("ClientInfo() = %+v after a memoized nil; want the nil to stick", got) } } + +// TestAffinityName pins which coding agent a request is filed with when it carries no +// session header. The Claude-User row is the reason this is not just Name: Claude Code's +// WebFetch announces itself only inside a comment, which the parser deliberately does +// not read. +func TestAffinityName(t *testing.T) { + cases := []struct { + ua string + want string + }{ + {"claude-cli/2.1.284 (external, cli)", "claude-code"}, + {"Claude-User (claude-code/2.1.284; +https://support.anthropic.com/)", "claude-code"}, + {"ai-sdk/5.0.1 openai-compatible/3.0.36 bob-shell/2.0.5", "bob-shell"}, + {"bob-shell/2.0.5", "bob-shell"}, + {"Go-http-client/1.1", ""}, + {"python-httpx/0.27.0", ""}, + {"axios/1.15.2", ""}, + {"node", ""}, + // A comment naming a product that is not a coding agent claims nothing. + {"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7)", ""}, + } + for _, tc := range cases { + if got := ParseUserAgent(tc.ua).AffinityName(); got != tc.want { + t.Errorf("AffinityName(%q) = %q, want %q", tc.ua, got, tc.want) + } + } + if got := (*EventClient)(nil).AffinityName(); got != "" { + t.Errorf("nil client: AffinityName() = %q, want \"\"", got) + } +} + +// TestAffinityName_LeavesTheLabelAlone pins the other half: reading claude-code out of +// the Claude-User comment is for session attribution ONLY. Label is the ledger's agent key +// and the AGENTS pane's row name, and changing it would re-file new cost rows under a +// different agent than the history beside them. +func TestAffinityName_LeavesTheLabelAlone(t *testing.T) { + const ua = "Claude-User (claude-code/2.1.284; +https://support.anthropic.com/)" + c := ParseUserAgent(ua) + if c.Name != "" || c.Label() != ua { + t.Fatalf("ParseUserAgent(%q) = {Name:%q Label:%q}, want an unrecognised client labelled by its raw value", + ua, c.Name, c.Label()) + } +} From 167554c9a52c34aa4790138b8d0274e2c3173e99 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 30 Sep 2026 09:39:54 -0400 Subject: [PATCH 3/8] feat: Let the session store file header-less calls by coding agent Adds the store half of client affinity, inert until a listener calls Claim: - Claim(id, client) records that client named id through its own session header. First claim wins, so a second agent quoting the id does not take it over. - SessionForClient(client) is where a header-less request goes: a known client's newest live session, else its pending bucket (pending:); the default bucket for an unknown client while two or more known clients are live; otherwise "", so ActiveSession() answers exactly as it does today. That last arm is the in-cluster case. - Claim ADOPTS the client's pending bucket into the session its first header names, when that session holds nothing yet. That is what files Bob's /admin/v1/profile, /model/info and task-classifier calls into Bob's session instead of Claude's. The adopted id redirects, because the request in flight at adoption time has its response pinned to the pending id and would otherwise recreate the bucket as a row of orphan responses. - Adopt tells Rekeyer recorders; usage.Aggregator implements it, so /v1/usage?session= and the group=session breakdown follow the rename. Rekey, the A2A default-to-contextId merge, still notifies nobody. Assisted-By: Claude (Anthropic AI) Signed-off-by: Hai Huang --- core/cost/usage/rekey.go | 37 +++++++ core/cost/usage/rekey_test.go | 34 +++++++ core/session/affinity.go | 182 ++++++++++++++++++++++++++++++++++ core/session/affinity_test.go | 168 +++++++++++++++++++++++++++++++ core/session/store.go | 22 +++- 5 files changed, 441 insertions(+), 2 deletions(-) create mode 100644 core/cost/usage/rekey.go create mode 100644 core/cost/usage/rekey_test.go create mode 100644 core/session/affinity.go create mode 100644 core/session/affinity_test.go diff --git a/core/cost/usage/rekey.go b/core/cost/usage/rekey.go new file mode 100644 index 000000000..83b7cd28e --- /dev/null +++ b/core/cost/usage/rekey.go @@ -0,0 +1,37 @@ +package usage + +// Rekeyed moves oldID's figures to newID: the per-session ring, and oldID's row in every +// bucket's session breakdown. It implements session.Rekeyer, which session.Store.Adopt calls +// when a coding agent's pending bucket is renamed to the session its first header named. +// +// Without it the store would list the session under its new id while /v1/usage kept +// answering for the old one — the "listed but zeroed" shape sessionRing exists to prevent. +// A newID that already has a ring keeps it; the store never adopts into a session that holds +// events, so that is a defensive no-op rather than a merge. +func (a *Aggregator) Rekeyed(oldID, newID string) { + if oldID == newID { + return + } + a.mu.Lock() + defer a.mu.Unlock() + if ring, ok := a.sessions[oldID]; ok { + if _, taken := a.sessions[newID]; !taken { + a.sessions[newID] = ring + delete(a.sessions, oldID) + relabelSession(ring.buckets, oldID, newID) + } + } + relabelSession(a.all, oldID, newID) +} + +func relabelSession(ring []bucket, oldID, newID string) { + from, to := ringLabel(oldID), ringLabel(newID) + for i := range ring { + c, ok := ring[i].bySession[from] + if !ok { + continue + } + delete(ring[i].bySession, from) + addLabel(&ring[i].bySession, to, c) + } +} diff --git a/core/cost/usage/rekey_test.go b/core/cost/usage/rekey_test.go new file mode 100644 index 000000000..0d104905a --- /dev/null +++ b/core/cost/usage/rekey_test.go @@ -0,0 +1,34 @@ +package usage + +import ( + "testing" + "time" +) + +// TestRekeyed_MovesTheSessionsFiguresToTheAdoptingSession pins what session.Store.Adopt +// relies on: after a pending bucket is renamed, /v1/usage?session= answers with its +// figures and the all-sessions breakdown names the new id, not the pending one. +func TestRekeyed_MovesTheSessionsFiguresToTheAdoptingSession(t *testing.T) { + a := New() + a.Record("pending:bob-shell", inferenceEvent("m", 13, 0, 0, 5, 0, 0b1001)) + a.Record("other", inferenceEvent("m", 29, 0, 0, 7, 0, 0b1001)) + + a.Rekeyed("pending:bob-shell", "task-1") + + all := mergeSeries(a.Snapshot(10*time.Minute, BucketWidth, "", GroupSession).Buckets) + if _, ok := all["pending:bob-shell"]; ok { + t.Error("all-sessions breakdown still names the adopted pending bucket") + } + if got := all["task-1"].InputTokens; got != 13 { + t.Errorf("task-1 InputTokens = %d, want 13", got) + } + if got := all["other"].InputTokens; got != 29 { + t.Errorf("other InputTokens = %d, want 29 (untouched)", got) + } + if got := a.Snapshot(10*time.Minute, BucketWidth, "task-1", GroupNone).Totals.InputTokens; got != 13 { + t.Errorf("session=task-1 InputTokens = %d, want 13 (the per-session ring must follow)", got) + } + if got := a.Snapshot(10*time.Minute, BucketWidth, "pending:bob-shell", GroupNone).Totals.InputTokens; got != 0 { + t.Errorf("session=pending:bob-shell InputTokens = %d, want 0 after the move", got) + } +} diff --git a/core/session/affinity.go b/core/session/affinity.go new file mode 100644 index 000000000..56050309a --- /dev/null +++ b/core/session/affinity.go @@ -0,0 +1,182 @@ +package session + +import ( + "log/slog" + "strings" + "time" +) + +// Client affinity: filing a request that carries no session header under a session of +// the SAME coding agent, rather than under ActiveSession(), the single global "most +// recently updated" id. With two agents running, ActiveSession() files each one's +// header-less calls — Bob's startup probes and task classifier, Claude Code's WebFetch — +// into whichever session spoke last, which is the other agent's half the time. +// +// Everything here is inert until a listener calls Claim, which only one configured with +// session.client_affinity does. A store nothing claims from answers SessionForClient with +// "" for every unknown client and a pending id for every known one, and holds no state. + +// PendingPrefix marks the bucket a known coding agent's header-less requests collect in +// before its first session header arrives. See Claim for how that bucket is adopted. +const PendingPrefix = "pending:" + +// PendingSessionID is the pending bucket for client. +func PendingSessionID(client string) string { return PendingPrefix + client } + +// Rekeyer is optionally implemented by a Recorder that keeps per-session state, so that +// state follows an Adopt. Called under the store's write lock, like Record. +// +// Only Adopt notifies. Rekey — the A2A path's default→contextId merge — never has, and +// changing that is not this feature's to decide. +type Rekeyer interface { + Rekeyed(oldID, newID string) +} + +// Claim records that client named sessionID through its own session header, which is what +// makes it that client's session for SessionForClient. First claim wins: a session is its +// creator's, and a second agent quoting the same id does not take it over. +// +// It also ADOPTS the client's pending bucket into sessionID when sessionID holds nothing +// yet, so the calls an agent makes before its first header — Bob's /admin/v1/profile, +// /model/info and task classifier — land in the session they belong to. Called at hydration, +// BEFORE the headered request is appended, which is what makes "holds nothing yet" the common +// case rather than a race. When sessionID already exists the pending bucket is left as its +// own row: merging two event histories is not something the store can do honestly. +func (s *Store) Claim(sessionID, client string) { + if sessionID == "" || client == "" || strings.HasPrefix(sessionID, PendingPrefix) { + return + } + if len(sessionID) > MaxSessionIDLen { + sessionID = sessionID[:MaxSessionIDLen] + } + s.mu.Lock() + defer s.mu.Unlock() + if s.owners == nil { + s.owners = make(map[string]string) + } + owner, owned := s.owners[sessionID] + if !owned { + s.pruneOwnersLocked() + s.owners[sessionID], owner = client, client + } + if owner == client { + s.adoptLocked(PendingSessionID(client), sessionID) + } +} + +// Adopt renames pendingID to id the way Rekey does, and additionally tells every Rekeyer +// recorder, so the usage aggregator's per-session figures follow the rename. Reports false, +// changing nothing, when pendingID does not exist or id already does. +func (s *Store) Adopt(pendingID, id string) bool { + s.mu.Lock() + defer s.mu.Unlock() + return s.adoptLocked(pendingID, id) +} + +func (s *Store) adoptLocked(pendingID, id string) bool { + if _, pending := s.sessions[pendingID]; !pending { + return false + } + if !s.rekeyLocked(pendingID, id) { + slog.Info("session: pending bucket not adopted, target already holds events", + "pending", pendingID, "session", id) + return false + } + // The request that was in flight when its client's first header arrived is still pinned + // to the pending id — its response is recorded under pctx.OutboundSessionID. Without the + // redirect that response would recreate the bucket just adopted, as a stray row of orphan + // responses. See followAdoptedLocked. + if s.adopted == nil { + s.adopted = make(map[string]string) + } + s.adopted[pendingID] = id + for _, r := range s.recorders { + if rk, ok := r.(Rekeyer); ok { + rk.Rekeyed(pendingID, id) + } + } + return true +} + +// followAdoptedLocked maps a pending id that has been adopted to the session that adopted +// it, for as long as that session exists. Any other id comes back unchanged. +func (s *Store) followAdoptedLocked(id string) string { + to, ok := s.adopted[id] + if !ok { + return id + } + if _, live := s.sessions[to]; live { + return to + } + delete(s.adopted, id) + return id +} + +// SessionForClient is where a request with no session header goes under client affinity, +// or "" to fall back to ActiveSession(): +// +// 1. A known client (client != ""): its newest live session, else its pending bucket. +// 2. An unknown client while two or more known clients hold live sessions: the default +// bucket, because the owner is ambiguous and guessing files it into one of theirs. +// 3. Anything else: "", so ActiveSession() answers exactly as it does without affinity. +// This is the in-cluster case — an A2A agent is no known coding agent, and +// ActiveSession() is what ties its outbound calls to the inbound turn that caused them. +func (s *Store) SessionForClient(client string) string { + s.mu.RLock() + defer s.mu.RUnlock() + now := time.Now() + if client != "" { + var newest string + var at time.Time + for id, owner := range s.owners { + sess, ok := s.sessions[id] + if owner != client || !ok || s.isExpired(sess, now) { + continue + } + if newest == "" || sess.UpdatedAt.After(at) { + newest, at = id, sess.UpdatedAt + } + } + if newest != "" { + return newest + } + return PendingSessionID(client) + } + seen := make(map[string]struct{}, 2) + for id, sess := range s.sessions { + if s.isExpired(sess, now) { + continue + } + owner := s.owners[id] + if owner == "" { + owner = strings.TrimPrefix(id, PendingPrefix) + if owner == id { + continue + } + } + seen[owner] = struct{}{} + if len(seen) >= 2 { + return DefaultSessionID + } + } + return "" +} + +// pruneOwnersLocked drops claims for sessions the store no longer holds. A claim is made at +// hydration, before anything is appended, so a request rejected there leaves one behind that +// neither cleanup nor eviction will ever see. Swept only past twice the session cap, so the +// common Claim pays nothing. +func (s *Store) pruneOwnersLocked() { + limit := 2 * s.maxSessions + if limit <= 0 { // an uncapped store: sweep at a fixed size instead + limit = 256 + } + if len(s.owners) < limit { + return + } + for id := range s.owners { + if _, ok := s.sessions[id]; !ok { + delete(s.owners, id) + } + } +} diff --git a/core/session/affinity_test.go b/core/session/affinity_test.go new file mode 100644 index 000000000..10ff41521 --- /dev/null +++ b/core/session/affinity_test.go @@ -0,0 +1,168 @@ +package session + +import ( + "testing" + "time" + + "github.com/rossoctl/cortex/core/cost/usage" + "github.com/rossoctl/cortex/core/pipeline" +) + +// Adopt reaches the aggregator through an optional interface, so a signature drift on +// either side would degrade to a silent skip. This makes it a compile error instead. It +// lives here because usage cannot import session. +var _ Rekeyer = (*usage.Aggregator)(nil) + +func ev() pipeline.SessionEvent { + return pipeline.SessionEvent{At: time.Now(), Direction: pipeline.Outbound, Phase: pipeline.SessionRequest} +} + +func TestSessionForClient_ResolutionOrder(t *testing.T) { + s := New(0, 0, 0) + defer s.Close() + + if got := s.SessionForClient(""); got != "" { + t.Fatalf("unknown client, empty store = %q, want \"\" (fall back to ActiveSession)", got) + } + if got := s.SessionForClient("bob-shell"); got != "pending:bob-shell" { + t.Fatalf("known client with no session = %q, want its pending bucket", got) + } + + s.Claim("claude-1", "claude-code") + s.Append("claude-1", ev()) + if got := s.SessionForClient("claude-code"); got != "claude-1" { + t.Fatalf("claude-code = %q, want its own session claude-1", got) + } + if got := s.SessionForClient("bob-shell"); got != "pending:bob-shell" { + t.Fatalf("bob-shell = %q, want its pending bucket, never claude's session", got) + } + if got := s.SessionForClient(""); got != "" { + t.Fatalf("unknown client beside ONE known agent = %q, want \"\" (single-agent behaviour unchanged)", got) + } + + s.Append("pending:bob-shell", ev()) + if got := s.SessionForClient(""); got != DefaultSessionID { + t.Fatalf("unknown client beside two known agents = %q, want the default bucket", got) + } +} + +func TestSessionForClient_PicksTheNewestOfTheClientsSessions(t *testing.T) { + s := New(0, 0, 0) + defer s.Close() + s.Claim("a", "claude-code") + s.Claim("b", "claude-code") + s.Append("a", ev()) + time.Sleep(2 * time.Millisecond) + s.Append("b", ev()) + if got := s.SessionForClient("claude-code"); got != "b" { + t.Fatalf("got %q, want b (updated last)", got) + } + time.Sleep(2 * time.Millisecond) + s.Append("a", ev()) + if got := s.SessionForClient("claude-code"); got != "a" { + t.Fatalf("got %q, want a (updated last)", got) + } +} + +func TestSessionForClient_SkipsExpiredSessions(t *testing.T) { + s := New(20*time.Millisecond, 0, 0) + defer s.Close() + s.Claim("old", "claude-code") + s.Append("old", ev()) + time.Sleep(40 * time.Millisecond) + if got := s.SessionForClient("claude-code"); got != "pending:claude-code" { + t.Fatalf("got %q, want the pending bucket once the only session expired", got) + } +} + +func TestClaim_FirstClaimWins(t *testing.T) { + s := New(0, 0, 0) + defer s.Close() + s.Claim("x", "claude-code") + s.Claim("x", "bob-shell") + s.Append("x", ev()) + if got := s.SessionForClient("bob-shell"); got != "pending:bob-shell" { + t.Fatalf("bob-shell = %q; quoting claude's session id must not make it bob's", got) + } + if got := s.SessionForClient("claude-code"); got != "x" { + t.Fatalf("claude-code = %q, want x", got) + } +} + +func TestClaim_AdoptsThePendingBucketAndRedirectsItsStragglers(t *testing.T) { + s := New(0, 0, 0) + defer s.Close() + s.Append("pending:bob-shell", ev()) + s.Append("pending:bob-shell", ev()) + + s.Claim("task-1", "bob-shell") + + if v := s.View("pending:bob-shell"); v != nil { + t.Fatalf("pending bucket still listed with %d events after its client's first header", len(v.Events)) + } + v := s.View("task-1") + if v == nil || len(v.Events) != 2 { + t.Fatalf("task-1 = %v, want the two adopted events", v) + } + for _, e := range v.Events { + if e.SessionID != "task-1" { + t.Fatalf("adopted event still stamped %q", e.SessionID) + } + } + // A response pinned to the pending id before the adoption follows it. + s.Append("pending:bob-shell", ev()) + if v := s.View("pending:bob-shell"); v != nil { + t.Fatal("a straggler recreated the adopted pending bucket") + } + if v := s.View("task-1"); len(v.Events) != 3 { + t.Fatalf("task-1 holds %d events, want 3 with the straggler", len(v.Events)) + } +} + +func TestClaim_DoesNotAdoptIntoAnotherAgentsSession(t *testing.T) { + s := New(0, 0, 0) + defer s.Close() + s.Claim("x", "claude-code") + s.Append("pending:bob-shell", ev()) + s.Claim("x", "bob-shell") + if v := s.View("pending:bob-shell"); v == nil { + t.Fatal("bob's pending bucket was merged into claude's session") + } +} + +func TestAdopt_LeavesPendingWhenTheTargetHoldsEvents(t *testing.T) { + s := New(0, 0, 0) + defer s.Close() + s.Append("pending:bob-shell", ev()) + s.Append("task-1", ev()) + if s.Adopt("pending:bob-shell", "task-1") { + t.Fatal("Adopt merged into a session that already held events") + } + if s.View("pending:bob-shell") == nil || len(s.View("task-1").Events) != 1 { + t.Fatal("a refused Adopt changed the store") + } +} + +type rekeyRecorder struct{ rekeyed [][2]string } + +func (r *rekeyRecorder) Record(string, *pipeline.SessionEvent) {} +func (r *rekeyRecorder) Rekeyed(o, n string) { r.rekeyed = append(r.rekeyed, [2]string{o, n}) } + +func TestAdopt_NotifiesRekeyersAndRekeyStillDoesNot(t *testing.T) { + s := New(0, 0, 0) + defer s.Close() + r := &rekeyRecorder{} + s.AddRecorder(r) + + s.Append(DefaultSessionID, ev()) + s.Rekey(DefaultSessionID, "ctx-1") + if len(r.rekeyed) != 0 { + t.Fatalf("Rekey notified %v; the A2A merge must keep today's behaviour", r.rekeyed) + } + + s.Append("pending:bob-shell", ev()) + s.Claim("task-1", "bob-shell") + if len(r.rekeyed) != 1 || r.rekeyed[0] != [2]string{"pending:bob-shell", "task-1"} { + t.Fatalf("Rekeyed calls = %v, want exactly [pending:bob-shell task-1]", r.rekeyed) + } +} diff --git a/core/session/store.go b/core/session/store.go index 46496718f..3550e64ff 100644 --- a/core/session/store.go +++ b/core/session/store.go @@ -176,6 +176,11 @@ type Store struct { // recorders are notified of every appended event. Registered at setup via // AddRecorder, before the store serves traffic. recorders []Recorder + + // owners and adopted serve client affinity; see affinity.go. Both nil until + // the first Claim, so a store nothing claims from carries neither. + owners map[string]string + adopted map[string]string } // subscriberChanBuf caps each subscriber's channel depth. 64 absorbs short @@ -343,6 +348,7 @@ func (s *Store) Append(sessionID string, event pipeline.SessionEvent) { defer s.mu.Unlock() now := time.Now() + sessionID = s.followAdoptedLocked(sessionID) sess, ok := s.sessions[sessionID] if !ok { @@ -1004,13 +1010,17 @@ func (s *Store) Rekey(oldID, newID string) { s.mu.Lock() defer s.mu.Unlock() + s.rekeyLocked(oldID, newID) +} +// rekeyLocked is Rekey's body, shared with Adopt. Reports whether it renamed anything. +func (s *Store) rekeyLocked(oldID, newID string) bool { sess, ok := s.sessions[oldID] if !ok { - return + return false } if _, exists := s.sessions[newID]; exists { - return + return false } sess.ID = newID @@ -1024,6 +1034,12 @@ func (s *Store) Rekey(oldID, newID string) { for i := range sess.Events { sess.Events[i].SessionID = newID } + // A claim follows its session. Only client affinity makes one; see affinity.go. + if owner, ok := s.owners[oldID]; ok { + delete(s.owners, oldID) + s.owners[newID] = owner + } + return true } // Cleanup removes expired sessions. Safe for concurrent use. @@ -1037,6 +1053,7 @@ func (s *Store) cleanupLocked(now time.Time) { for id, sess := range s.sessions { if s.isExpired(sess, now) { delete(s.sessions, id) + delete(s.owners, id) if s.activeID == id { s.activeID = "" } @@ -1063,6 +1080,7 @@ func (s *Store) evictOldestLocked() { } if oldestID != "" { delete(s.sessions, oldestID) + delete(s.owners, oldestID) } } From 16a6b874b44cc2911f5724bbc87962f66cdae230 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 30 Sep 2026 09:44:17 -0400 Subject: [PATCH 4/8] feat: Add session.client_affinity to keep coding agents' sessions apart With Claude Code and Bob Shell on one proxy, a request with no session header went to ActiveSession(), the single most-recently-updated session, so each agent's header-less calls landed in whichever session spoke last. Observed on a laptop: Bob's /admin/v1/profile, /model/info and task-classifier completions filed under Claude's session, and Claude Code's WebFetch under Bob's. session.client_affinity (default false) routes those through the store's SessionForClient: the same agent's newest session, its pending bucket before its first header, or the default bucket for traffic from no known agent while two are live. A headered request Claims its session, which adopts that agent's pending bucket. Everything else falls through to ActiveSession() unchanged, so in-cluster A2A correlation is untouched; TestClientAffinity_InClusterResolutionIsUnchanged runs those cases with the knob on. Plugins hear the ambiguous case as "" rather than "default", the same no-identity answer resolvePluginSessionID already gives them, so sessionbudget's DefaultSessionFallback does not start enforcing on it. Under the knob, CONNECT and transparent tunnel rows use the identity the tunnel was gated under instead of re-reading ActiveSession() at record time (#1187). With the knob off that path is unchanged. Not reloadable, like the rest of session.*: the reloader refuses the change and asks for a restart. Assisted-By: Claude (Anthropic AI) Signed-off-by: Hai Huang --- cmd/authbridge-cpex/main.go | 1 + cmd/authbridge-proxy/main.go | 1 + core/config/config.go | 15 +++ core/config/config_test.go | 16 +++ core/listener/forwardproxy/server.go | 41 +++++++ .../forwardproxy/session_affinity_test.go | 102 ++++++++++++++++++ core/listener/forwardproxy/transparent.go | 16 ++- core/session/affinity_test.go | 2 +- 8 files changed, 192 insertions(+), 2 deletions(-) diff --git a/cmd/authbridge-cpex/main.go b/cmd/authbridge-cpex/main.go index a16daaa7f..8a91f02ea 100644 --- a/cmd/authbridge-cpex/main.go +++ b/cmd/authbridge-cpex/main.go @@ -267,6 +267,7 @@ func main() { // off switch wherever attribution is a trust boundary; see // session.IDFromHeaders. fpSrv.SessionIDHeaders = cfg.Session.SessionIDHeaders() + fpSrv.ClientAffinity = cfg.Session.ClientAffinity rpHTTP, err := bootstrap.StartReverseProxyServer("reverse-proxy", rpSrv, cfg.Listener.ReverseProxyAddr) if err != nil { log.Fatalf("reverse-proxy listen: %v", err) diff --git a/cmd/authbridge-proxy/main.go b/cmd/authbridge-proxy/main.go index bb9936133..d604fc54d 100644 --- a/cmd/authbridge-proxy/main.go +++ b/cmd/authbridge-proxy/main.go @@ -894,6 +894,7 @@ func main() { // the supported agents' session headers (config.SessionIDHeaders); // session.id_headers: [] turns it off. fpSrv.SessionIDHeaders = cfg.Session.SessionIDHeaders() + fpSrv.ClientAffinity = cfg.Session.ClientAffinity fpHTTP, herr := bootstrap.StartHTTPServer("forward-proxy", fpSrv.Handler(), cfg.Listener.ForwardProxyAddr) if herr != nil { fatalf("forward-proxy listen: %v", herr) diff --git a/core/config/config.go b/core/config/config.go index 5ffc2847c..9e08501ee 100644 --- a/core/config/config.go +++ b/core/config/config.go @@ -441,6 +441,21 @@ type SessionConfig struct { // boundary rather than a convenience. See session.IDFromHeaders for the full // reasoning. IDHeaders []string `yaml:"id_headers" json:"id_headers"` + + // ClientAffinity files a request that carries NO session header under the newest + // session of the same coding agent (by User-Agent), rather than under the single + // most-recently-updated session. Off by default, which keeps today's attribution. + // + // Turn it on where more than one coding agent shares this proxy — Claude Code and + // Bob Shell on one laptop. Without it each agent's header-less calls (Bob's startup + // probes and task classifier, Claude Code's WebFetch) land in whichever session spoke + // last, which is the other agent's half the time. With it, a known agent's calls + // before its first header collect in a pending: bucket that its first + // session adopts, and traffic from no known agent goes to the default bucket while + // two agents are live. Anything else resolves exactly as before, which is what keeps + // in-cluster A2A correlation intact. Needs id_headers non-empty; see + // session.Store.SessionForClient. + ClientAffinity bool `yaml:"client_affinity" json:"client_affinity"` } // SessionIDHeaders returns the headers to consult for a client-supplied session diff --git a/core/config/config_test.go b/core/config/config_test.go index 3a70cce81..a6973bb24 100644 --- a/core/config/config_test.go +++ b/core/config/config_test.go @@ -947,3 +947,19 @@ listener: t.Errorf("MTLS = %+v, want nil (absent block)", cfg.MTLS) } } + +// TestSessionConfig_ClientAffinity pins the default (off — today's attribution) and the +// YAML key an operator writes to turn it on. +func TestSessionConfig_ClientAffinity(t *testing.T) { + var off, on Config + if err := yaml.Unmarshal([]byte("session: {}\n"), &off); err != nil { + t.Fatal(err) + } + if err := yaml.Unmarshal([]byte("session:\n client_affinity: true\n"), &on); err != nil { + t.Fatal(err) + } + if off.Session.ClientAffinity || !on.Session.ClientAffinity { + t.Fatalf("client_affinity: unset = %v, true = %v; want false, true", + off.Session.ClientAffinity, on.Session.ClientAffinity) + } +} diff --git a/core/listener/forwardproxy/server.go b/core/listener/forwardproxy/server.go index 9089273de..243db51ac 100644 --- a/core/listener/forwardproxy/server.go +++ b/core/listener/forwardproxy/server.go @@ -861,9 +861,25 @@ func (s *Server) recordingSessionID(resolved string, clientHeaders http.Header) // Stated in full under "Decision recorded: client-asserted ids as plugin input" // in #984, which is where to argue with it. func (s *Server) resolvePluginSessionID(clientHeaders http.Header) string { + affinity := s.affinityOn() if sid := session.IDFromHeaders(clientHeaders, s.SessionIDHeaders); sid != "" { + if affinity { + s.Sessions.Claim(sid, affinityClient(clientHeaders)) + } return sid } + if affinity { + // See session.Store.SessionForClient for the order. The default bucket is its + // "ambiguous" answer, and plugins hear that as "" — the no-identity answer this + // function exists to give them — while recording files it under default. + switch sid := s.Sessions.SessionForClient(affinityClient(clientHeaders)); sid { + case "": + case session.DefaultSessionID: + return "" + default: + return sid + } + } if s.Sessions != nil { if sid := s.Sessions.ActiveSession(); sid != "" { return sid @@ -872,6 +888,24 @@ func (s *Server) resolvePluginSessionID(clientHeaders http.Header) string { return "" } +// affinityOn reports whether header-less requests are filed by coding agent. It needs +// header bucketing: a client's session is the one its header named, so with id_headers +// empty no session is ever any client's and every known agent would collect in its +// pending bucket forever. +func (s *Server) affinityOn() bool { + return s.ClientAffinity && s.Sessions != nil && len(s.SessionIDHeaders) > 0 +} + +// affinityClient names the coding agent behind a request for session attribution. Read +// from the headers as RECEIVED, like the session header beside it, so it is the same +// User-Agent pctx.ResolveClient pins. Nil headers — the transparent path — name nobody. +func affinityClient(clientHeaders http.Header) string { + if clientHeaders == nil { + return "" + } + return pipeline.ParseUserAgent(clientHeaders.Get("User-Agent")).AffinityName() +} + // sessionViewFor returns the recorded view for sid, or an empty view carrying // just the id when nothing has been recorded under it yet. // @@ -1373,6 +1407,13 @@ func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) { } } + // Under client affinity the tunnel row joins the session this CONNECT was gated + // under, not ActiveSession() at recording time. Pinned only now, after the + // pipeline, for the reason sessionID is a local. See recordTunnelOpened. + if s.ClientAffinity && !skipped && s.Sessions != nil { + pctx.OutboundSessionID = s.recordingSessionID(sessionID, r.Header) + } + // Verify hijack capability BEFORE dialing upstream. If hijacking // isn't supported the failure mode should be a 500 to the client, // not a half-opened TCP connection to the upstream. The actual diff --git a/core/listener/forwardproxy/session_affinity_test.go b/core/listener/forwardproxy/session_affinity_test.go index 3f201e4c8..6b9f65f3e 100644 --- a/core/listener/forwardproxy/session_affinity_test.go +++ b/core/listener/forwardproxy/session_affinity_test.go @@ -3,6 +3,7 @@ package forwardproxy import ( "net/http" "net/http/httptest" + "strings" "testing" "time" @@ -123,3 +124,104 @@ func newAffinityProxy(t *testing.T, store *session.Store) (*identityProbePlugin, t.Cleanup(proxy.Close) return probe, &http.Client{Transport: &http.Transport{Proxy: http.ProxyURL(mustParseURL(proxy.URL))}}, backend.URL } + +// TestClientAffinity_SeparatesBobFromClaudeCode replays the order observed on a laptop +// running both agents: Bob's startup probes and task classifier before its first +// X-Task-Id, Claude Code's WebFetch with no header, and a library call from neither. +// With ActiveSession() alone, Bob's pre-header calls landed in Claude's session and the +// WebFetch in Bob's. +func TestClientAffinity_SeparatesBobFromClaudeCode(t *testing.T) { + const ( + claudeUA = "claude-cli/2.1.284 (external, cli)" + webFetchUA = "Claude-User (claude-code/2.1.284; +https://support.anthropic.com/)" + bobBareUA = "bob-shell/2.0.5" + bobSDKUA = "ai-sdk/5.0.1 openai-compatible/3.0.36 bob-shell/2.0.5" + ) + store := session.New(5*time.Minute, 100, 0) + defer store.Close() + _, client, backendURL := newAffinityProxy(t, store) + + send := func(path, ua, header, sid string) { + t.Helper() + req, _ := http.NewRequest("POST", backendURL+path, nil) + req.Header.Set("User-Agent", ua) + if header != "" { + req.Header.Set(header, sid) + } + resp, err := client.Do(req) + if err != nil { + t.Fatalf("%s: %v", path, err) + } + _ = resp.Body.Close() + } + send("/v1/messages", claudeUA, session.ClaudeCodeSessionHeader, "claude-1") + send("/admin/v1/profile", bobBareUA, "", "") + send("/inference/v1/chat/completions?classifier", bobSDKUA, "", "") + send("/inference/v1/chat/completions", bobSDKUA, session.BobSessionHeader, "task-1") + send("/web/fetch", webFetchUA, "", "") + send("/api/web/domain_info", "axios/1.15.2", "", "") + + paths := func(id string) []string { + v := store.View(id) + if v == nil { + return nil + } + var out []string + for _, e := range v.Events { + if e.Phase == pipeline.SessionRequest { + out = append(out, e.HTTPPath) + } + } + return out + } + want := map[string][]string{ + "claude-1": {"/v1/messages", "/web/fetch"}, + "task-1": {"/admin/v1/profile", "/inference/v1/chat/completions", "/inference/v1/chat/completions"}, + session.DefaultSessionID: {"/api/web/domain_info"}, + } + for id, w := range want { + if got := paths(id); strings.Join(got, ",") != strings.Join(w, ",") { + t.Errorf("session %s request paths = %v, want %v", id, got, w) + } + } + if v := store.View(session.PendingSessionID("bob-shell")); v != nil { + t.Errorf("pending:bob-shell still holds %d events after Bob's first header", len(v.Events)) + } + // Every response followed its request, including the classifier's, which was + // pinned to the pending id before the adoption. + for id := range want { + v := store.View(id) + var req, resp int + for _, e := range v.Events { + if e.Phase == pipeline.SessionRequest { + req++ + } else { + resp++ + } + } + if req != resp { + t.Errorf("session %s holds %d requests and %d responses", id, req, resp) + } + } +} + +// TestRecordTunnelOpened_UsesTheGatedSessionOnlyUnderAffinity pins #1187's half of this +// change: with the knob on the tunnel row joins the session its CONNECT was gated under; +// with it off, recording still reads ActiveSession() as it always has. +func TestRecordTunnelOpened_UsesTheGatedSessionOnlyUnderAffinity(t *testing.T) { + for _, affinity := range []bool{false, true} { + store := session.New(5*time.Minute, 100, 0) + store.Append("active", pipeline.SessionEvent{At: time.Now(), Direction: pipeline.Inbound, Phase: pipeline.SessionRequest}) + s := &Server{Sessions: store, ClientAffinity: affinity} + pctx := &pipeline.Context{Direction: pipeline.Outbound, Method: "CONNECT", Host: "api.example:443", OutboundSessionID: "gated"} + s.recordTunnelOpened(pctx, "") + want := "active" + if affinity { + want = "gated" + } + if v := store.View(want); v == nil || v.Events[len(v.Events)-1].HTTPMethod != "CONNECT" { + t.Errorf("affinity=%v: tunnel row not recorded under %q", affinity, want) + } + store.Close() + } +} diff --git a/core/listener/forwardproxy/transparent.go b/core/listener/forwardproxy/transparent.go index 8319ed611..aa43c790c 100644 --- a/core/listener/forwardproxy/transparent.go +++ b/core/listener/forwardproxy/transparent.go @@ -133,6 +133,11 @@ func (s *Server) HandleTransparentConn(clientConn net.Conn, dst string) { return } + // See handleConnect's pin: same rule, with no headers to read. + if s.ClientAffinity && s.Sessions != nil { + pctx.OutboundSessionID = s.recordingSessionID(sessionID, nil) + } + // Always dial the original IP (dst), never the sniffed name — the agent // already chose the IP, and re-resolving the name could diverge from it. upstream, err := net.DialTimeout("tcp", dst, connectDialTimeout) @@ -191,7 +196,16 @@ func (s *Server) recordTunnelOpened(pctx *pipeline.Context, reason pipeline.Tunn if s.Sessions == nil { return } - sid := s.Sessions.ActiveSession() + // Without client affinity this reads ActiveSession() at recording time, as it always + // has, ignoring the identity the tunnel was gated under (#1187). With it, the pin set + // after gating wins; see handleConnect. + var sid string + if s.ClientAffinity { + sid = pctx.OutboundSessionID + } + if sid == "" { + sid = s.Sessions.ActiveSession() + } if sid == "" { sid = session.DefaultSessionID } diff --git a/core/session/affinity_test.go b/core/session/affinity_test.go index 10ff41521..6ab4b6443 100644 --- a/core/session/affinity_test.go +++ b/core/session/affinity_test.go @@ -146,7 +146,7 @@ func TestAdopt_LeavesPendingWhenTheTargetHoldsEvents(t *testing.T) { type rekeyRecorder struct{ rekeyed [][2]string } func (r *rekeyRecorder) Record(string, *pipeline.SessionEvent) {} -func (r *rekeyRecorder) Rekeyed(o, n string) { r.rekeyed = append(r.rekeyed, [2]string{o, n}) } +func (r *rekeyRecorder) Rekeyed(o, n string) { r.rekeyed = append(r.rekeyed, [2]string{o, n}) } func TestAdopt_NotifiesRekeyersAndRekeyStillDoesNot(t *testing.T) { s := New(0, 0, 0) From 7d7baf3d2aeff6a85ee90022bd4ec35bf743e375 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 30 Sep 2026 09:45:34 -0400 Subject: [PATCH 5/8] test: Pin client_affinity's off position, its id_headers term, and plugin silence Three mutations of the previous commit survived its tests, each for a reason worth a test of its own: - Turning affinity on regardless of the knob passed every in-cluster test, because those are built to be unaffected by it. Pinned by the knob-off answer it replaces, misfiling included. - Dropping the id_headers term from affinityOn passed; with header bucketing off a known agent would collect in its pending bucket forever. - Handing plugins "default" for the ambiguous case passed, because the scenario asserted only where events were recorded, and recording files it under default either way. Assisted-By: Claude (Anthropic AI) Signed-off-by: Hai Huang --- .../forwardproxy/session_affinity_test.go | 37 ++++++++++++++++++- 1 file changed, 36 insertions(+), 1 deletion(-) diff --git a/core/listener/forwardproxy/session_affinity_test.go b/core/listener/forwardproxy/session_affinity_test.go index 6b9f65f3e..1c9dc01fb 100644 --- a/core/listener/forwardproxy/session_affinity_test.go +++ b/core/listener/forwardproxy/session_affinity_test.go @@ -139,7 +139,7 @@ func TestClientAffinity_SeparatesBobFromClaudeCode(t *testing.T) { ) store := session.New(5*time.Minute, 100, 0) defer store.Close() - _, client, backendURL := newAffinityProxy(t, store) + probe, client, backendURL := newAffinityProxy(t, store) send := func(path, ua, header, sid string) { t.Helper() @@ -187,6 +187,12 @@ func TestClientAffinity_SeparatesBobFromClaudeCode(t *testing.T) { if v := store.View(session.PendingSessionID("bob-shell")); v != nil { t.Errorf("pending:bob-shell still holds %d events after Bob's first header", len(v.Events)) } + // The axios call is recorded under default but its PLUGINS hear no identity, the same + // answer resolvePluginSessionID gives when nothing is known. Handing them "default" + // would satisfy sessionbudget's opt-in DefaultSessionFallback. + if got := probe.sawID[len(probe.sawID)-1]; got != "" { + t.Errorf("plugins saw %q for an unknown client beside two live agents, want \"\"", got) + } // Every response followed its request, including the classifier's, which was // pinned to the pending id before the adoption. for id := range want { @@ -225,3 +231,32 @@ func TestRecordTunnelOpened_UsesTheGatedSessionOnlyUnderAffinity(t *testing.T) { store.Close() } } + +// TestClientAffinity_OffKeepsTodaysAttribution pins the knob's off position to the +// attribution it replaces, misfiling included: Bob's pre-header probe lands in Claude's +// session because Claude's was the one updated last. A change that turned affinity on +// regardless of the knob passes every in-cluster test — those are built to be unaffected — +// and fails this one. +func TestClientAffinity_OffKeepsTodaysAttribution(t *testing.T) { + store := session.New(5*time.Minute, 100, 0) + defer store.Close() + s := &Server{Sessions: store, SessionIDHeaders: []string{session.ClaudeCodeSessionHeader, session.BobSessionHeader}} + store.Append("claude-1", pipeline.SessionEvent{At: time.Now(), Direction: pipeline.Outbound, Phase: pipeline.SessionRequest}) + h := http.Header{"User-Agent": []string{"bob-shell/2.0.5"}} + if got := s.resolveOutboundSessionID(h); got != "claude-1" { + t.Fatalf("knob off: Bob's header-less probe = %q, want claude-1 (today's ActiveSession answer)", got) + } +} + +// TestClientAffinity_NeedsHeaderBucketing pins affinityOn's id_headers term: with header +// bucketing off no session is ever any agent's, so a known agent would collect in its +// pending bucket forever. The knob then does nothing. +func TestClientAffinity_NeedsHeaderBucketing(t *testing.T) { + store := session.New(5*time.Minute, 100, 0) + defer store.Close() + s := &Server{Sessions: store, SessionIDHeaders: []string{}, ClientAffinity: true} + h := http.Header{"User-Agent": []string{"bob-shell/2.0.5"}} + if got := s.resolveOutboundSessionID(h); got != session.DefaultSessionID { + t.Fatalf("id_headers: [] with client_affinity = %q, want %q (affinity inert)", got, session.DefaultSessionID) + } +} From f91d55c61ea494e70158c79085bd28bfa8234de3 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 30 Sep 2026 09:47:40 -0400 Subject: [PATCH 6/8] feat: Turn on client affinity and Bob's session header in the --local preset The built-in laptop config listed X-Claude-Code-Session-Id and X-Session-Id explicitly, and an explicit id_headers REPLACES the built-in default rather than extending it (#1069), so Bob's X-Task-Id was never read under --local and every Bob request fell to ActiveSession(). It is listed now, and the preset turns client_affinity on, since a laptop is where two coding agents share a proxy. Fresh installs only. writeBuiltinConfig never rewrites an existing ~/.cortex/config.yaml, so an upgrade leaves a current install's session block exactly as it is; turning the knob on there is a one-line edit and a restart. Assisted-By: Claude (Anthropic AI) Signed-off-by: Hai Huang --- cmd/authbridge-proxy/local.go | 8 ++++++++ cmd/authbridge-proxy/local_test.go | 23 +++++++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/cmd/authbridge-proxy/local.go b/cmd/authbridge-proxy/local.go index cf7861f93..9b04245ae 100644 --- a/cmd/authbridge-proxy/local.go +++ b/cmd/authbridge-proxy/local.go @@ -279,12 +279,20 @@ cost_ledger: # non-empty value wins. # # X-Claude-Code-Session-Id is set by Claude Code on every inference request. +# X-Task-Id is set by IBM Bob Shell. # X-Session-Id is set by OpenCode, Pi (Inflection AI), and similar frameworks. # # An explicit empty list (id_headers: []) disables header-based bucketing. +# +# client_affinity files a request with none of these headers under the newest +# session of the same coding agent, so with Claude Code and Bob running side by +# side neither one's header-less calls land in the other's session. Not +# hot-reloadable: restart after changing it. session: + client_affinity: true id_headers: - "X-Claude-Code-Session-Id" + - "X-Task-Id" - "X-Session-Id" pipeline: outbound: diff --git a/cmd/authbridge-proxy/local_test.go b/cmd/authbridge-proxy/local_test.go index 4e3fb3e03..34c9abab6 100644 --- a/cmd/authbridge-proxy/local_test.go +++ b/cmd/authbridge-proxy/local_test.go @@ -22,6 +22,7 @@ import ( "github.com/rossoctl/cortex/core/config" "github.com/rossoctl/cortex/core/cost/usage" + "github.com/rossoctl/cortex/core/session" ) // writeBuiltinConfig must produce a config file in cortexDir that loads, presets, @@ -612,3 +613,25 @@ func TestBuiltinConfig_CommentedRetentionDoesNotTruncateAMonth(t *testing.T) { got, usage.Window7dLocalDays) } } + +// TestBuiltinConfig_SeparatesCodingAgentSessions pins the laptop preset's session block: +// Bob's X-Task-Id must be listed — an explicit id_headers REPLACES the built-in default, +// so leaving it out disabled Bob's bucketing under --local — and client_affinity is on, +// since a laptop is where two coding agents share one proxy. +func TestBuiltinConfig_SeparatesCodingAgentSessions(t *testing.T) { + cortexDir := t.TempDir() + p, err := writeBuiltinConfig(cortexDir, filepath.Join(cortexDir, "ca")) + if err != nil { + t.Fatal(err) + } + cfg, err := config.Load(p) + if err != nil { + t.Fatal(err) + } + if !cfg.Session.ClientAffinity { + t.Error("built-in config: session.client_affinity is off, want on") + } + if !slices.Contains(cfg.Session.SessionIDHeaders(), session.BobSessionHeader) { + t.Errorf("built-in config: id_headers = %v, want it to include %s", cfg.Session.SessionIDHeaders(), session.BobSessionHeader) + } +} From a7ac73d61b252b642e2e972c807a853eb554b90a Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 30 Sep 2026 09:48:08 -0400 Subject: [PATCH 7/8] docs: Correct the claim that inference is always attributed exactly docs/laptop-service.md said header-less requests were only MCP tool calls and that "inference traffic ... is always exact". Neither holds with two agents running: Bob's task-classifier completions are inference with tokens and cost, carry no X-Task-Id, and were measured landing in Claude Code's session. Documents session.client_affinity beside the limitation it lifts, including what it does not fix: two sessions of the same agent still share header-less calls by timing. Assisted-By: Claude (Anthropic AI) Signed-off-by: Hai Huang --- docs/laptop-service.md | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/docs/laptop-service.md b/docs/laptop-service.md index c778ddc18..e7df0b2f8 100644 --- a/docs/laptop-service.md +++ b/docs/laptop-service.md @@ -74,10 +74,21 @@ else that egresses through the proxy without announcing a session. Some limitations worth knowing: -- **Tool calls are attributed by timing, not identity.** MCP requests carry no session - header, so they are filed under whichever session was most recently active. That is - right when sessions take turns and can misattribute when two are genuinely - interleaved. Inference traffic — where the tokens and the cost are — is always exact. +- **Header-less requests are attributed by timing, not identity — unless you turn on + `session.client_affinity`.** MCP tool calls, Claude Code's WebFetch, and Bob's startup + probes and task-classifier completions carry no session header, so by default they are + filed under whichever session was most recently active. That is right when sessions + take turns and wrong when two agents run side by side: Bob's classifier completions — + real inference, with tokens and cost — land in Claude Code's session, and Claude Code's + WebFetch in Bob's. Inference is exact only where it carries the header. + + `session.client_affinity: true` files a header-less request under the newest session of + the same agent (by `User-Agent`). An agent's calls before its first header wait in a + `pending:` bucket that its first session then absorbs, and a request from no + known agent goes to `default` while two agents are live. Two sessions of the SAME agent + still share by timing. The built-in `--local` preset turns it on for new installs; an + existing `~/.cortex/config.yaml` is never rewritten, so add the line yourself. It is not + hot-reloadable — restart the proxy after changing it. - **Agents other than Claude Code and Bob need to be named.** Set `session.id_headers` to a list of headers to consult in precedence order if you run a client with its own session header; naming any replaces the built-in list rather than adding to it. An From ea49d8bd2388ee6b3d5ac4de861989406d2f6cdc Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 30 Sep 2026 11:07:58 -0400 Subject: [PATCH 8/8] fix: Keep tunnel rows out of default, and let the ambiguous case lapse Review round 1 of #1194, one class: affinity's ambiguous answer (the default bucket) overrode today's attribution in places it should not, and latched. - Latching: SessionForClient judged "two agents live" by expiry, and session.ttl defaults to never, so after one Bob run every unknown client went to default for the rest of the proxy's life. It now counts an agent only if it sent traffic within ambiguityWindow (5 minutes). - Tunnel rows: a CONNECT rarely carries a User-Agent, so under that answer nearly every tunnel row went to default, apart from the request it carries, which breaks abctl's CONNECT fold. Where affinity has no answer, the tunnel pin is now left unset and the row keeps ActiveSession() at recording time, as without affinity. The two tunnel reject paths get the same rule via tunnelSessionID. Swept every site that records under the answer (recordingSessionID callers and the OutboundSessionID pins): 6 sites, 4 on tunnel paths and fixed; the 2 HTTP sites keep the ambiguous default by design, now bounded by the window. Knob off is unchanged: tunnelSessionID falls through to recordingSessionID. Also drops a transparent-path comment that the knob made false ("empty only when nothing was active"), and corrects "while two agents are live" in config.go and laptop-service.md to what the code checks. Tests drive a real CONNECT with no User-Agent, a transparent connection, and a bob-shell CONNECT through the handler, plus tunnelSessionID (which the reject paths use) and the window lapse. Removing the transparent pin outright is not caught: with no plugin moving ActiveSession() mid-pipeline it is equivalent in that fixture. Assisted-By: Claude (Anthropic AI) Signed-off-by: Hai Huang --- core/config/config.go | 3 +- core/listener/forwardproxy/server.go | 24 +++++- .../forwardproxy/session_affinity_test.go | 73 +++++++++++++++++++ core/listener/forwardproxy/transparent.go | 9 +-- core/session/affinity.go | 13 +++- core/session/affinity_test.go | 8 ++ docs/laptop-service.md | 3 +- 7 files changed, 119 insertions(+), 14 deletions(-) diff --git a/core/config/config.go b/core/config/config.go index 9e08501ee..f27cc06bb 100644 --- a/core/config/config.go +++ b/core/config/config.go @@ -452,7 +452,8 @@ type SessionConfig struct { // last, which is the other agent's half the time. With it, a known agent's calls // before its first header collect in a pending: bucket that its first // session adopts, and traffic from no known agent goes to the default bucket while - // two agents are live. Anything else resolves exactly as before, which is what keeps + // two agents have both sent traffic in the last five minutes; tunnel rows keep + // ActiveSession(). Anything else resolves exactly as before, which is what keeps // in-cluster A2A correlation intact. Needs id_headers non-empty; see // session.Store.SessionForClient. ClientAffinity bool `yaml:"client_affinity" json:"client_affinity"` diff --git a/core/listener/forwardproxy/server.go b/core/listener/forwardproxy/server.go index 243db51ac..1d5c33820 100644 --- a/core/listener/forwardproxy/server.go +++ b/core/listener/forwardproxy/server.go @@ -823,6 +823,21 @@ func (s *Server) recordingSessionID(resolved string, clientHeaders http.Header) return s.resolveOutboundSessionID(clientHeaders) } +// tunnelSessionID is recordingSessionID for a tunnel's own row. Under client affinity +// resolved is "" only for affinity's ambiguous answer or for no identity at all, and +// neither sends a tunnel row to default: a CONNECT rarely carries a User-Agent, so the +// ambiguous answer would file nearly every tunnel away from the request it carries. +// Those rows keep ActiveSession() at recording time, as without affinity. +func (s *Server) tunnelSessionID(resolved string, clientHeaders http.Header) string { + if s.ClientAffinity && resolved == "" && s.Sessions != nil { + if sid := s.Sessions.ActiveSession(); sid != "" { + return sid + } + return session.DefaultSessionID + } + return s.recordingSessionID(resolved, clientHeaders) +} + // resolvePluginSessionID is resolveOutboundSessionID without the default-bucket // fallback: it returns "" when neither a client header nor an active session // names one, and is the identity handed to plugins. @@ -1396,7 +1411,7 @@ func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) { // HTTP body (parsers) see no body, which they handle gracefully. action := s.OutboundPipeline.Run(r.Context(), pctx) if action.Type == pipeline.Reject { - s.recordOutboundReject(pctx, action, s.recordingSessionID(sessionID, r.Header)) + s.recordOutboundReject(pctx, action, s.tunnelSessionID(sessionID, r.Header)) // Render as a JSON-RPC error frame when the rejected // request was MCP JSON-RPC, so the agent's MCP client // surfaces this as one failed tool call rather than a @@ -1409,9 +1424,10 @@ func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) { // Under client affinity the tunnel row joins the session this CONNECT was gated // under, not ActiveSession() at recording time. Pinned only now, after the - // pipeline, for the reason sessionID is a local. See recordTunnelOpened. - if s.ClientAffinity && !skipped && s.Sessions != nil { - pctx.OutboundSessionID = s.recordingSessionID(sessionID, r.Header) + // pipeline, for the reason sessionID is a local; left unpinned where affinity had no + // answer, for the reason in tunnelSessionID. See recordTunnelOpened. + if s.ClientAffinity && !skipped && s.Sessions != nil && sessionID != "" { + pctx.OutboundSessionID = sessionID } // Verify hijack capability BEFORE dialing upstream. If hijacking diff --git a/core/listener/forwardproxy/session_affinity_test.go b/core/listener/forwardproxy/session_affinity_test.go index 1c9dc01fb..e7dd0ce92 100644 --- a/core/listener/forwardproxy/session_affinity_test.go +++ b/core/listener/forwardproxy/session_affinity_test.go @@ -1,6 +1,10 @@ package forwardproxy import ( + "bufio" + "fmt" + "io" + "net" "net/http" "net/http/httptest" "strings" @@ -260,3 +264,72 @@ func TestClientAffinity_NeedsHeaderBucketing(t *testing.T) { t.Fatalf("id_headers: [] with client_affinity = %q, want %q (affinity inert)", got, session.DefaultSessionID) } } + +// TestClientAffinity_TunnelRowsKeepTodaysSessionWhenTheOwnerIsAmbiguous replays the +// review's case: tunnels with no User-Agent — most of a laptop's CONNECTs — while two +// agents are recent. Filing those under default split every bridged call from the +// request it carries. A known agent's tunnel still follows its pin. +func TestClientAffinity_TunnelRowsKeepTodaysSessionWhenTheOwnerIsAmbiguous(t *testing.T) { + backend, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer func() { _ = backend.Close() }() + go func() { + for c, err := backend.Accept(); err == nil; c, err = backend.Accept() { + go func() { _, _ = io.Copy(io.Discard, c); _ = c.Close() }() + } + }() + store := session.New(0, 100, 0) // Bob then Claude Code just spoke: ActiveSession() is claude-1 + defer store.Close() + for _, c := range [][2]string{{"task-1", "bob-shell"}, {"claude-1", "claude-code"}} { + store.Claim(c[0], c[1]) + store.Append(c[0], pipeline.SessionEvent{At: time.Now(), Direction: pipeline.Outbound, Phase: pipeline.SessionRequest}) + } + p, _ := pipeline.New(nil) + srv := &Server{OutboundPipeline: pipeline.NewHolder(p), Sessions: store, Client: http.DefaultClient, + SessionIDHeaders: []string{session.ClaudeCodeSessionHeader, session.BobSessionHeader}, ClientAffinity: true} + proxy := httptest.NewServer(srv.Handler()) + defer proxy.Close() + waitFor := func(id string, want int) { + t.Helper() + for deadline := time.Now().Add(5 * time.Second); ; time.Sleep(5 * time.Millisecond) { + n := 0 + for _, e := range store.View(id).Events { + if e.HTTPMethod == http.MethodConnect { + n++ + } + } + if n >= want { + return + } + if time.Now().After(deadline) { + t.Fatalf("%s holds %d tunnel rows, want %d; default holds %v", id, n, want, store.View(session.DefaultSessionID)) + } + } + } + connect := func(uaLine string) { + t.Helper() + c, err := net.Dial("tcp", proxy.Listener.Addr().String()) + if err != nil { + t.Fatal(err) + } + defer func() { _ = c.Close() }() + _, _ = fmt.Fprintf(c, "CONNECT %[1]s HTTP/1.1\r\nHost: %[1]s\r\n%[2]s\r\n", backend.Addr(), uaLine) + if resp, err := http.ReadResponse(bufio.NewReader(c), nil); err != nil || resp.StatusCode != http.StatusOK { + t.Fatalf("CONNECT: %v %v", resp, err) + } + } + connect("") + waitFor("claude-1", 1) + client, server := net.Pipe() + defer func() { _ = client.Close() }() + go srv.HandleTransparentConn(server, backend.Addr().String()) + waitFor("claude-1", 2) + connect("User-Agent: bob-shell/2.0.5\r\n") + waitFor("task-1", 1) + // The reject paths file through tunnelSessionID rather than the pin. + if got, def := srv.tunnelSessionID("", nil), srv.recordingSessionID("", nil); got != "task-1" || def != session.DefaultSessionID { + t.Errorf("ambiguous denial filed under %q (recordingSessionID %q), want task-1, today's ActiveSession()", got, def) + } +} diff --git a/core/listener/forwardproxy/transparent.go b/core/listener/forwardproxy/transparent.go index aa43c790c..3a0a3ab64 100644 --- a/core/listener/forwardproxy/transparent.go +++ b/core/listener/forwardproxy/transparent.go @@ -126,16 +126,15 @@ func (s *Server) HandleTransparentConn(clientConn net.Conn, dst string) { action := s.OutboundPipeline.Run(ctx, pctx) if action.Type == pipeline.Reject { // Carry the identity this connection was gated under, so the denial lands - // in the session whose traffic it was. Empty only when nothing was active, - // where recording applies the default bucket as it always did. - s.recordOutboundReject(pctx, action, s.recordingSessionID(sessionID, nil)) + // in the session whose traffic it was. + s.recordOutboundReject(pctx, action, s.tunnelSessionID(sessionID, nil)) slog.Warn("transparent-proxy: outbound rejected by policy", "host", host) return } // See handleConnect's pin: same rule, with no headers to read. - if s.ClientAffinity && s.Sessions != nil { - pctx.OutboundSessionID = s.recordingSessionID(sessionID, nil) + if s.ClientAffinity && s.Sessions != nil && sessionID != "" { + pctx.OutboundSessionID = sessionID } // Always dial the original IP (dst), never the sniffed name — the agent diff --git a/core/session/affinity.go b/core/session/affinity.go index 56050309a..8e61206d2 100644 --- a/core/session/affinity.go +++ b/core/session/affinity.go @@ -112,12 +112,19 @@ func (s *Store) followAdoptedLocked(id string) string { return id } +// ambiguityWindow is how recently a known agent must have sent traffic to count toward +// SessionForClient's ambiguous case. Expiry cannot serve: session.ttl defaults to never, +// so one Bob run would otherwise send every unknown client to default for the rest of the +// proxy's life. +const ambiguityWindow = 5 * time.Minute + // SessionForClient is where a request with no session header goes under client affinity, // or "" to fall back to ActiveSession(): // // 1. A known client (client != ""): its newest live session, else its pending bucket. -// 2. An unknown client while two or more known clients hold live sessions: the default -// bucket, because the owner is ambiguous and guessing files it into one of theirs. +// 2. An unknown client while two or more known clients have each sent traffic within +// ambiguityWindow: the default bucket, because the owner is ambiguous and guessing +// files it into one of theirs. // 3. Anything else: "", so ActiveSession() answers exactly as it does without affinity. // This is the in-cluster case — an A2A agent is no known coding agent, and // ActiveSession() is what ties its outbound calls to the inbound turn that caused them. @@ -144,7 +151,7 @@ func (s *Store) SessionForClient(client string) string { } seen := make(map[string]struct{}, 2) for id, sess := range s.sessions { - if s.isExpired(sess, now) { + if s.isExpired(sess, now) || now.Sub(sess.UpdatedAt) > ambiguityWindow { continue } owner := s.owners[id] diff --git a/core/session/affinity_test.go b/core/session/affinity_test.go index 6ab4b6443..1c14c9887 100644 --- a/core/session/affinity_test.go +++ b/core/session/affinity_test.go @@ -44,6 +44,14 @@ func TestSessionForClient_ResolutionOrder(t *testing.T) { if got := s.SessionForClient(""); got != DefaultSessionID { t.Fatalf("unknown client beside two known agents = %q, want the default bucket", got) } + // Recent traffic, not expiry, decides: ttl defaults to never, so one finished Bob run + // would otherwise send every unknown client to default for the proxy's life. + s.mu.Lock() + s.sessions["pending:bob-shell"].UpdatedAt = time.Now().Add(-ambiguityWindow - time.Second) + s.mu.Unlock() + if got := s.SessionForClient(""); got != "" { + t.Fatalf("Bob quiet past ambiguityWindow = %q, want \"\" (fall back to ActiveSession)", got) + } } func TestSessionForClient_PicksTheNewestOfTheClientsSessions(t *testing.T) { diff --git a/docs/laptop-service.md b/docs/laptop-service.md index e7df0b2f8..828fb519f 100644 --- a/docs/laptop-service.md +++ b/docs/laptop-service.md @@ -85,7 +85,8 @@ Some limitations worth knowing: `session.client_affinity: true` files a header-less request under the newest session of the same agent (by `User-Agent`). An agent's calls before its first header wait in a `pending:` bucket that its first session then absorbs, and a request from no - known agent goes to `default` while two agents are live. Two sessions of the SAME agent + known agent goes to `default` while two agents have both sent traffic in the last five + minutes (a tunnel's own row keeps today's attribution). Two sessions of the SAME agent still share by timing. The built-in `--local` preset turns it on for new installs; an existing `~/.cortex/config.yaml` is never rewritten, so add the line yourself. It is not hot-reloadable — restart the proxy after changing it.