Feat: Keep each coding agent's header-less requests in its own session

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)

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:

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

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)

core/config/config.go

@@ -441,6 +441,22 @@ 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:<agent> bucket that its first
+	// session adopts, and traffic from no known agent goes to the default bucket while
+	// 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"`
 }
 
 // SessionIDHeaders returns the headers to consult for a client-supplied session

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

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

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=<new> 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)
+	}
+}

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
@@ -816,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.
@@ -854,9 +876,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
@@ -865,6 +903,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.
 //
@@ -1355,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
@@ -1366,6 +1422,14 @@ 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; 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
 	// isn't supported the failure mode should be a 500 to the client,
 	// not a half-opened TCP connection to the upstream. The actual