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/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) + } +} 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..f27cc06bb 100644 --- a/core/config/config.go +++ b/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: 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 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/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/listener/forwardproxy/server.go b/core/listener/forwardproxy/server.go index 69a0bf2f7..1d5c33820 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 @@ -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 diff --git a/core/listener/forwardproxy/session_affinity_test.go b/core/listener/forwardproxy/session_affinity_test.go new file mode 100644 index 000000000..e7dd0ce92 --- /dev/null +++ b/core/listener/forwardproxy/session_affinity_test.go @@ -0,0 +1,335 @@ +package forwardproxy + +import ( + "bufio" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "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 +} + +// 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() + probe, 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)) + } + // 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 { + 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() + } +} + +// 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) + } +} + +// 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 8319ed611..3a0a3ab64 100644 --- a/core/listener/forwardproxy/transparent.go +++ b/core/listener/forwardproxy/transparent.go @@ -126,13 +126,17 @@ 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 && sessionID != "" { + pctx.OutboundSessionID = sessionID + } + // 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 +195,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/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()) + } +} diff --git a/core/session/affinity.go b/core/session/affinity.go new file mode 100644 index 000000000..8e61206d2 --- /dev/null +++ b/core/session/affinity.go @@ -0,0 +1,189 @@ +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 +} + +// 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 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. +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) || now.Sub(sess.UpdatedAt) > ambiguityWindow { + 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..1c14c9887 --- /dev/null +++ b/core/session/affinity_test.go @@ -0,0 +1,176 @@ +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) + } + // 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) { + 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) } } diff --git a/docs/laptop-service.md b/docs/laptop-service.md index c778ddc18..828fb519f 100644 --- a/docs/laptop-service.md +++ b/docs/laptop-service.md @@ -74,10 +74,22 @@ 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 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. - **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