From e899f0e53985605a96fb6dd49d5204ec4845acb8 Mon Sep 17 00:00:00 2001 From: sususu98 <33882693+sususu98@users.noreply.github.com> Date: Thu, 3 Sep 2026 17:52:40 +0800 Subject: [PATCH] feat(session): derive distinct branch session ID, parent lineage on Merkle LCP forks, and enhance Codex fork/subagent affinity (#5418) (#5454) --- sdk/cliproxy/auth/home_session_alias.go | 16 +- sdk/cliproxy/auth/home_session_alias_test.go | 24 + sdk/cliproxy/auth/selector.go | 295 ++++++++- sdk/cliproxy/auth/selector_lcp_test.go | 624 ++++++++++++++++++ .../auth/session_overhead_benchmark_test.go | 278 ++++++++ sdk/cliproxy/executor/types.go | 7 + sdk/cliproxy/session/info.go | 219 +++++- sdk/cliproxy/session/info_test.go | 73 ++ sdk/cliproxy/session/lcp.go | 141 +++- sdk/cliproxy/session/lcp_test.go | 269 +++++++- 10 files changed, 1854 insertions(+), 92 deletions(-) create mode 100644 sdk/cliproxy/auth/session_overhead_benchmark_test.go diff --git a/sdk/cliproxy/auth/home_session_alias.go b/sdk/cliproxy/auth/home_session_alias.go index dce16b54..e9e69029 100644 --- a/sdk/cliproxy/auth/home_session_alias.go +++ b/sdk/cliproxy/auth/home_session_alias.go @@ -240,7 +240,7 @@ func isHierarchyParent(primary, fallback string) bool { if strings.Contains(primary, ":agent:") { return true } - for _, prefix := range []string{"codex:", "header:", "affinity:", "slot:", "thread:", "conv:", "session:", "claude:", "agy:", "geminicache:"} { + for _, prefix := range []string{"codex:", "header:", "affinity:", "slot:", "thread:", "conv:", "session:", "claude:", "agy:", "geminicache:", "lcp:"} { if strings.HasPrefix(primary, prefix) && strings.HasPrefix(fallback, prefix) && primary != fallback { return true } @@ -272,13 +272,25 @@ func (m *Manager) homeDispatchSessionIDs(opts cliproxyexecutor.Options) (string, aliasFallback = fallback } } + if parentSessionID == "" && opts.Metadata != nil { + if metaParent, ok := opts.Metadata[cliproxyexecutor.ParentSessionIDMetadataKey].(string); ok && strings.TrimSpace(metaParent) != "" { + parentSessionID = strings.TrimSpace(metaParent) + } + } cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) ttl := homeSessionAliasTTL(cfg) now := time.Now() canonical := m.homeSessionAliases.canonical(primary, aliasFallback, ttl, now) if parentSessionID != "" { - parentSessionID = m.homeSessionAliases.canonical(parentSessionID, "", ttl, now) + if parentSessionID == canonical || parentSessionID == primary || (aliasFallback != "" && parentSessionID == aliasFallback) { + parentSessionID = "" + } else { + parentSessionID = m.homeSessionAliases.canonical(parentSessionID, "", ttl, now) + if parentSessionID == canonical { + parentSessionID = "" + } + } } return canonical, parentSessionID } diff --git a/sdk/cliproxy/auth/home_session_alias_test.go b/sdk/cliproxy/auth/home_session_alias_test.go index b1788ac0..5c13d835 100644 --- a/sdk/cliproxy/auth/home_session_alias_test.go +++ b/sdk/cliproxy/auth/home_session_alias_test.go @@ -545,6 +545,30 @@ func TestHomeDispatchSessionIDsExtractsParentFromHeaderPlusBody(t *testing.T) { t.Fatalf("lcpParentID = %q, want empty", lcpParentID) } + // 3b. LCP with metadata parent + lcpForkOpts := cliproxyexecutor.Options{ + Metadata: map[string]any{ + cliproxyexecutor.CanonicalSessionIDMetadataKey: "lcp:v1:fork-child-xyz", + cliproxyexecutor.ParentSessionIDMetadataKey: "lcp:v1:parent-root-abc", + }, + } + lcpForkSessionID, lcpForkParentID := manager.homeDispatchSessionIDs(lcpForkOpts) + if lcpForkSessionID != "lcp:v1:fork-child-xyz" || lcpForkParentID != "lcp:v1:parent-root-abc" { + t.Fatalf("lcpFork = (%q, %q), want (lcp:v1:fork-child-xyz, lcp:v1:parent-root-abc)", lcpForkSessionID, lcpForkParentID) + } + + // 3c. Self-referential parent is suppressed + lcpSelfParentOpts := cliproxyexecutor.Options{ + Metadata: map[string]any{ + cliproxyexecutor.CanonicalSessionIDMetadataKey: "lcp:v1:same-session", + cliproxyexecutor.ParentSessionIDMetadataKey: "lcp:v1:same-session", + }, + } + selfSessionID, selfParentID := manager.homeDispatchSessionIDs(lcpSelfParentOpts) + if selfSessionID != "lcp:v1:same-session" || selfParentID != "" { + t.Fatalf("self-referential parent = (%q, %q), want (%q, empty)", selfSessionID, selfParentID, "lcp:v1:same-session") + } + // 4. Antigravity hierarchy agyOpts := cliproxyexecutor.Options{ Headers: http.Header{ diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go index c472c715..75a3f228 100644 --- a/sdk/cliproxy/auth/selector.go +++ b/sdk/cliproxy/auth/selector.go @@ -957,6 +957,17 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri if auth, handled, errLCP := s.pickLCP(ctx, provider, model, opts, auths, entry); handled || errLCP != nil { return auth, errLCP } + } else if opts.Metadata != nil { + delete(opts.Metadata, cliproxyexecutor.LCPAffinitySessionIDMetadataKey) + delete(opts.Metadata, cliproxyexecutor.LCPAccessGenerationMetadataKey) + if explicitFallbackID != "" { + opts.Metadata[cliproxyexecutor.ParentSessionIDMetadataKey] = explicitFallbackID + } else { + delete(opts.Metadata, cliproxyexecutor.ParentSessionIDMetadataKey) + } + if isFork, ok := opts.Metadata[cliproxyexecutor.IsForkMetadataKey].(bool); !ok || !isFork { + delete(opts.Metadata, cliproxyexecutor.IsForkMetadataKey) + } } primaryID, fallbackID := explicitID, explicitFallbackID @@ -990,13 +1001,19 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri modelKey := canonicalModelKey(model) cacheKey := provider + "::" + primaryID + "::" + modelKey - isSubagent := isSubagentSession(primaryID, fallbackID) + isFork := false + if opts.Metadata != nil { + if forkFlag, ok := opts.Metadata[cliproxyexecutor.IsForkMetadataKey].(bool); ok && forkFlag { + isFork = true + } + } + isSubagent := !isFork && isSubagentSession(primaryID, fallbackID) fallbackKey := "" if fallbackID != "" && fallbackID != primaryID { fallbackKey = provider + "::" + fallbackID + "::" + modelKey } bind := func(authID string) { - if fallbackKey != "" && !isSubagent { + if fallbackKey != "" && !isSubagent && !isFork { s.cache.SetAliases(authID, cacheKey, fallbackKey) } else { s.cache.Set(cacheKey, authID) @@ -1030,7 +1047,11 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri if auth.ID == cachedAuthID { if !isSubagent || s.subagentAffinity { bind(auth.ID) - entry.Infof("session-affinity: fallback cache hit | session=%s fallback=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), truncateSessionID(fallbackID), auth.ID, provider, model) + if isFork { + entry.Infof("session-affinity: fork cache hit | session=%s parent=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), truncateSessionID(fallbackID), auth.ID, provider, model) + } else { + entry.Infof("session-affinity: fallback cache hit | session=%s fallback=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), truncateSessionID(fallbackID), auth.ID, provider, model) + } return auth, nil } } @@ -1046,7 +1067,11 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri return nil, nil } bind(auth.ID) - entry.Infof("session-affinity: cache miss, new binding | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) + if isFork && fallbackID != "" { + entry.Infof("session-affinity: fork bound to new auth | session=%s parent=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), truncateSessionID(fallbackID), auth.ID, provider, model) + } else { + entry.Infof("session-affinity: cache miss, new binding | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) + } return auth, nil } @@ -1085,12 +1110,29 @@ func (s *SessionAffinitySelector) pickLCP(ctx context.Context, provider, model s if auth == nil || auth.ID != match.AuthID { continue } - s.matcher.TouchFingerprints(namespace, fingerprints, minPrefixLength, auth.ID) if match.SessionID != "" { opts.Metadata[cliproxyexecutor.LCPAffinitySessionIDMetadataKey] = match.SessionID opts.Metadata[cliproxyexecutor.CanonicalSessionIDMetadataKey] = match.SessionID } - entry.Infof("session-affinity: LCP cache hit | session=%s prefix=%d auth=%s provider=%s model=%s", truncateSessionID(match.SessionID), match.PrefixLength, auth.ID, provider, model) + if match.ParentSessionID != "" { + opts.Metadata[cliproxyexecutor.ParentSessionIDMetadataKey] = match.ParentSessionID + } else if opts.Metadata != nil { + delete(opts.Metadata, cliproxyexecutor.ParentSessionIDMetadataKey) + } + if match.AccessNumber > 0 && opts.Metadata != nil { + opts.Metadata[cliproxyexecutor.LCPAccessGenerationMetadataKey] = match.AccessNumber + } + if match.IsFork { + if opts.Metadata != nil { + opts.Metadata[cliproxyexecutor.IsForkMetadataKey] = true + } + entry.Infof("session-affinity: LCP fork hit | session=%s parent=%s prefix=%d auth=%s provider=%s model=%s", truncateSessionID(match.SessionID), truncateSessionID(match.ParentSessionID), match.PrefixLength, auth.ID, provider, model) + } else { + if opts.Metadata != nil { + delete(opts.Metadata, cliproxyexecutor.IsForkMetadataKey) + } + entry.Infof("session-affinity: LCP cache hit | session=%s prefix=%d auth=%s provider=%s model=%s", truncateSessionID(match.SessionID), match.PrefixLength, auth.ID, provider, model) + } return auth, true, nil } } @@ -1103,10 +1145,28 @@ func (s *SessionAffinitySelector) pickLCP(ctx context.Context, provider, model s if auth == nil { return nil, true, &Error{Code: "auth_not_found", Message: "selector returned no auth"} } - if sessionID := s.matcher.BindFingerprints(namespace, fingerprints, minPrefixLength, auth.ID); sessionID != "" { - opts.Metadata[cliproxyexecutor.LCPAffinitySessionIDMetadataKey] = sessionID - opts.Metadata[cliproxyexecutor.CanonicalSessionIDMetadataKey] = sessionID - entry.Infof("session-affinity: LCP cache miss, new binding | session=%s auth=%s provider=%s model=%s", truncateSessionID(sessionID), auth.ID, provider, model) + if bindRes := s.matcher.BindFingerprintsWithResult(namespace, fingerprints, minPrefixLength, auth.ID); bindRes.SessionID != "" { + opts.Metadata[cliproxyexecutor.LCPAffinitySessionIDMetadataKey] = bindRes.SessionID + opts.Metadata[cliproxyexecutor.CanonicalSessionIDMetadataKey] = bindRes.SessionID + if bindRes.ParentSessionID != "" { + opts.Metadata[cliproxyexecutor.ParentSessionIDMetadataKey] = bindRes.ParentSessionID + } else if opts.Metadata != nil { + delete(opts.Metadata, cliproxyexecutor.ParentSessionIDMetadataKey) + } + if bindRes.AccessNumber > 0 && opts.Metadata != nil { + opts.Metadata[cliproxyexecutor.LCPAccessGenerationMetadataKey] = bindRes.AccessNumber + } + if bindRes.IsFork { + if opts.Metadata != nil { + opts.Metadata[cliproxyexecutor.IsForkMetadataKey] = true + } + entry.Infof("session-affinity: LCP fork bound to new auth | session=%s parent=%s auth=%s provider=%s model=%s", truncateSessionID(bindRes.SessionID), truncateSessionID(bindRes.ParentSessionID), auth.ID, provider, model) + } else { + if opts.Metadata != nil { + delete(opts.Metadata, cliproxyexecutor.IsForkMetadataKey) + } + entry.Infof("session-affinity: LCP cache miss, new binding | session=%s auth=%s provider=%s model=%s", truncateSessionID(bindRes.SessionID), auth.ID, provider, model) + } } return auth, true, nil } @@ -1236,7 +1296,13 @@ func (s *SessionAffinitySelector) OnResult(res Result) { if res.Success { s.matcher.TouchFingerprints(namespace, fingerprints, minPrefixLength, res.AuthID) } else { - s.matcher.RemoveFingerprints(namespace, fingerprints, res.AuthID) + var generation uint64 + if res.Options.Metadata != nil { + if gen, ok := res.Options.Metadata[cliproxyexecutor.LCPAccessGenerationMetadataKey].(uint64); ok { + generation = gen + } + } + s.matcher.RemoveFingerprintsBefore(namespace, fingerprints, res.AuthID, generation) } } } @@ -1245,6 +1311,11 @@ func (s *SessionAffinitySelector) OnResult(res Result) { if s.cache == nil { return } + if explicitID == "" && s.matcher != nil && res.Options.Metadata != nil { + if _, isLCP := res.Options.Metadata[cliproxyexecutor.LCPAffinitySessionIDMetadataKey]; isLCP { + return + } + } primaryID, fallbackID := explicitID, explicitFallbackID if primaryID == "" { primaryID, fallbackID = extractSessionIDs(res.Options.Headers, res.Options.OriginalRequest, res.Options.Metadata) @@ -1272,6 +1343,27 @@ func (s *SessionAffinitySelector) OnResult(res Result) { } } +func isBodyForkCandidate(root, reqRoot gjson.Result, hasNestedReq bool) bool { + if !root.Exists() { + return false + } + for _, k := range []string{ + "forked_from_thread_id", "forked_from_id", + "metadata.forked_from_thread_id", "metadata.forked_from_id", + "extra_body.forked_from_thread_id", "extra_body.forked_from_id", + } { + if val := normalizedSessionCandidate(root.Get(k).String()); val != "" { + return true + } + if hasNestedReq { + if val := normalizedSessionCandidate(reqRoot.Get(k).String()); val != "" { + return true + } + } + } + return false +} + // normalizedSessionCandidate validates an explicit client-provided session signal. // It keeps opaque printable IDs intact while rejecting values that are unsafe or // implausibly large for routing keys and logs. @@ -1369,7 +1461,9 @@ func extractExplicitSessionIDs(headers http.Header, payload []byte, metadata map "forked_from_thread_id", "forked_from_id", "parent_conversation_id", "parentConversationId", "metadata.parent_session_id", "metadata.parent_thread_id", + "metadata.forked_from_thread_id", "metadata.forked_from_id", "extra_body.parent_session_id", "extra_body.parent_thread_id", + "extra_body.forked_from_thread_id", "extra_body.forked_from_id", } { if psid := normalizedSessionCandidate(root.Get(parentPath).String()); psid != "" { parentIDCandidate = psid @@ -1461,31 +1555,161 @@ func extractExplicitSessionIDs(headers http.Header, payload []byte, metadata map } // 2. OpenAI / Codex CLI - if sid := sessionHeaderValue(headers, "Session-Id"); sid != "" { - parentThreadID := sessionHeaderValue(headers, "x-codex-parent-thread-id") - if parentThreadID == "" { - parentThreadID = sessionHeaderValue(headers, "X-Codex-Parent-Thread-Id") - } - if parentThreadID != "" && parentThreadID != sid { - return "codex:" + sid, "codex:" + parentThreadID + sid := sessionHeaderValue(headers, "Session-Id") + if sid == "" { + sid = sessionHeaderValue(headers, "Session_id") + } + tid := sessionHeaderValue(headers, "Thread-Id") + if tid == "" { + tid = sessionHeaderValue(headers, "Thread_id") + } + + var codexTurnMeta string + for k, v := range headers { + if strings.EqualFold(k, "X-Codex-Turn-Metadata") && len(v) > 0 { + codexTurnMeta = strings.TrimSpace(v[0]) + break } - if parentIDCandidate != "" && parentIDCandidate != sid { - return "codex:" + sid, "codex:" + parentIDCandidate + } + var codexTurnMetaJSON gjson.Result + if codexTurnMeta != "" { + codexTurnMetaJSON = gjson.Parse(codexTurnMeta) + } + + if sid == "" && codexTurnMetaJSON.Exists() { + sid = normalizedSessionCandidate(codexTurnMetaJSON.Get("session_id").String()) + } + if tid == "" && codexTurnMetaJSON.Exists() { + tid = normalizedSessionCandidate(codexTurnMetaJSON.Get("thread_id").String()) + } + if tid == "" && sid != "" && root.Exists() { + for _, path := range []string{"thread_id", "threadId", "metadata.thread_id"} { + if tid = normalizedSessionCandidate(root.Get(path).String()); tid != "" { + break + } + if hasNestedReq { + if tid = normalizedSessionCandidate(reqRoot.Get(path).String()); tid != "" { + break + } + } } - return "codex:" + sid, "" } - if sid := sessionHeaderValue(headers, "Session_id"); sid != "" { + + if sid != "" || tid != "" { parentThreadID := sessionHeaderValue(headers, "x-codex-parent-thread-id") if parentThreadID == "" { parentThreadID = sessionHeaderValue(headers, "X-Codex-Parent-Thread-Id") } - if parentThreadID != "" && parentThreadID != sid { - return "codex:" + sid, "codex:" + parentThreadID + if parentThreadID == "" && codexTurnMetaJSON.Exists() { + parentThreadID = normalizedSessionCandidate(codexTurnMetaJSON.Get("parent_thread_id").String()) } - if parentIDCandidate != "" && parentIDCandidate != sid { - return "codex:" + sid, "codex:" + parentIDCandidate + + forkedFrom := "" + if codexTurnMetaJSON.Exists() { + forkedFrom = normalizedSessionCandidate(codexTurnMetaJSON.Get("forked_from_thread_id").String()) + if forkedFrom == "" { + forkedFrom = normalizedSessionCandidate(codexTurnMetaJSON.Get("forked_from_id").String()) + } + } + if forkedFrom == "" && root.Exists() { + for _, forkPath := range []string{ + "forked_from_thread_id", "forked_from_id", + "metadata.forked_from_thread_id", "metadata.forked_from_id", + "extra_body.forked_from_thread_id", "extra_body.forked_from_id", + } { + if forkedFrom = normalizedSessionCandidate(root.Get(forkPath).String()); forkedFrom != "" { + break + } + if hasNestedReq { + if forkedFrom = normalizedSessionCandidate(reqRoot.Get(forkPath).String()); forkedFrom != "" { + break + } + } + } + } + + cleanAgentName := "" + if codexTurnMetaJSON.Exists() { + rawName := codexTurnMetaJSON.Get("agent_name").String() + rawName = strings.TrimPrefix(rawName, "/root/") + rawName = strings.TrimPrefix(rawName, "/") + rawName = strings.TrimSpace(rawName) + rawName = normalizedSessionCandidate(rawName) + if rawName != "" && rawName != "root" && rawName != "main" { + cleanAgentName = rawName + } + } + + subVal := sessionHeaderValue(headers, "X-Openai-Subagent") + subagentSignal := subVal != "" && !strings.EqualFold(subVal, "false") && subVal != "0" + if codexTurnMetaJSON.Exists() && codexTurnMetaJSON.Get("subagent_kind").String() == "thread_spawn" { + subagentSignal = true + } + + // 1. Fork detection + if forkedFrom != "" { + forkSessionID := tid + if forkSessionID == "" { + forkSessionID = sid + } + if forkSessionID == forkedFrom && sid != "" && sid != forkedFrom { + forkSessionID = sid + } + primary = "codex:" + forkSessionID + fallback = "codex:" + forkedFrom + if metadata != nil { + metadata[cliproxyexecutor.IsForkMetadataKey] = true + metadata[cliproxyexecutor.ParentSessionIDMetadataKey] = fallback + } + return primary, fallback + } + + // 2. Multi-Agent / Subagent detection + if subagentSignal || (tid != "" && sid != "" && tid != sid) || (parentThreadID != "" && parentThreadID != tid && parentThreadID != sid) { + childSessionID := tid + if childSessionID == "" { + childSessionID = sid + } + parentSID := parentThreadID + if parentSID == "" { + parentSID = sid + } + if cleanAgentName != "" && sid != "" { + primary = "codex:" + sid + ":agent:" + cleanAgentName + if parentSID != "" { + fallback = "codex:" + parentSID + } else if parentIDCandidate != "" && parentIDCandidate != sid { + fallback = "codex:" + parentIDCandidate + } + } else { + primary = "codex:" + childSessionID + if parentSID != "" && parentSID != childSessionID { + fallback = "codex:" + parentSID + } else if parentIDCandidate != "" && parentIDCandidate != childSessionID { + fallback = "codex:" + parentIDCandidate + } + } + if metadata != nil && fallback != "" { + metadata[cliproxyexecutor.ParentSessionIDMetadataKey] = fallback + } + return primary, fallback } - return "codex:" + sid, "" + + // 3. Normal session + sessionID := sid + if sessionID == "" { + sessionID = tid + } + primary = "codex:" + sessionID + if parentThreadID != "" && parentThreadID != sessionID { + fallback = "codex:" + parentThreadID + } else if parentIDCandidate != "" && parentIDCandidate != sessionID { + fallback = "codex:" + parentIDCandidate + } + if metadata != nil && fallback != "" { + metadata[cliproxyexecutor.ParentSessionIDMetadataKey] = fallback + } + return primary, fallback } // 3. Antigravity CLI (agy) / Google Cloud Code @@ -1567,12 +1791,6 @@ func extractExplicitSessionIDs(headers http.Header, payload []byte, metadata map } return "thread:" + sid, "" } - if sid := sessionHeaderValue(headers, "Thread-Id"); sid != "" { - if parentIDCandidate != "" && parentIDCandidate != sid { - return "thread:" + sid, "thread:" + parentIDCandidate - } - return "thread:" + sid, "" - } if sid := sessionHeaderValue(headers, "X-Client-Request-Id"); sid != "" { return "clientreq:" + sid, "" } @@ -1608,6 +1826,10 @@ func extractExplicitSessionIDs(headers http.Header, payload []byte, metadata map } if tid != "" { if parentIDCandidate != "" && parentIDCandidate != tid { + if isBodyForkCandidate(root, reqRoot, hasNestedReq) && metadata != nil { + metadata[cliproxyexecutor.IsForkMetadataKey] = true + metadata[cliproxyexecutor.ParentSessionIDMetadataKey] = "thread:" + parentIDCandidate + } return "thread:" + tid, "thread:" + parentIDCandidate } return "thread:" + tid, "" @@ -1646,7 +1868,12 @@ func extractExplicitSessionIDs(headers http.Header, payload []byte, metadata map return primary, fallback } if parentIDCandidate != "" && parentIDCandidate != sid { - return "session:" + sid, "session:" + parentIDCandidate + fallback = "session:" + parentIDCandidate + if isBodyForkCandidate(root, reqRoot, hasNestedReq) && metadata != nil { + metadata[cliproxyexecutor.IsForkMetadataKey] = true + metadata[cliproxyexecutor.ParentSessionIDMetadataKey] = fallback + } + return "session:" + sid, fallback } return "session:" + sid, "" } diff --git a/sdk/cliproxy/auth/selector_lcp_test.go b/sdk/cliproxy/auth/selector_lcp_test.go index 616bcd37..531d9392 100644 --- a/sdk/cliproxy/auth/selector_lcp_test.go +++ b/sdk/cliproxy/auth/selector_lcp_test.go @@ -480,6 +480,322 @@ func TestSessionAffinityCodexSubagentInheritsParentThread(t *testing.T) { if got := childOpts.Metadata[cliproxyexecutor.CanonicalSessionIDMetadataKey]; got != "codex:thread-child-888" { t.Fatalf("child canonical session ID = %v, want codex:thread-child-888", got) } + + // 3. Forked thread carries X-Codex-Turn-Metadata with forked_from_thread_id + forkOpts := cliproxyexecutor.Options{ + Headers: http.Header{ + "Session-Id": []string{"thread-fork-777"}, + "X-Codex-Turn-Metadata": []string{ + `{"session_id":"thread-fork-777","forked_from_thread_id":"thread-parent-999","request_kind":"turn"}`, + }, + }, + OriginalRequest: []byte(`{"input":[{"role":"user","content":"fork turn"}]}`), + Metadata: map[string]any{}, + } + forkAuth, errFork := selector.Pick(context.Background(), "openai", "model", forkOpts, auths) + if errFork != nil { + t.Fatalf("fork Pick() error = %v", errFork) + } + if forkAuth.ID != parentAuth.ID { + t.Fatalf("fork thread did not inherit parent auth: got %q, want %q", forkAuth.ID, parentAuth.ID) + } + if got := forkOpts.Metadata[cliproxyexecutor.CanonicalSessionIDMetadataKey]; got != "codex:thread-fork-777" { + t.Fatalf("fork canonical session ID = %v, want codex:thread-fork-777", got) + } +} + +func TestSessionAffinityCodexForkInheritsParentOnGeminiModel(t *testing.T) { + t.Parallel() + + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + + auths := []*Auth{{ID: "auth-1", Provider: "antigravity"}, {ID: "auth-2", Provider: "antigravity"}} + + // 1. Parent thread binds to an antigravity auth + parentOpts := cliproxyexecutor.Options{ + Headers: http.Header{"Session-Id": []string{"parent-gemini-thread"}}, + OriginalRequest: []byte(`{"input":[{"role":"user","content":"parent query"}]}`), + Metadata: map[string]any{}, + } + parentAuth, errParent := selector.Pick(context.Background(), "mixed", "gemini-3.8-flash-high", parentOpts, auths) + if errParent != nil { + t.Fatalf("parent Pick() error = %v", errParent) + } + + // 2. Forked thread carries X-Codex-Turn-Metadata with forked_from_thread_id + // Even on Gemini/Antigravity where subagents are non-inheriting, conversational FORKS must inherit parent auth! + forkOpts := cliproxyexecutor.Options{ + Headers: http.Header{ + "Session-Id": []string{"fork-gemini-thread"}, + "X-Codex-Turn-Metadata": []string{ + `{"session_id":"fork-gemini-thread","forked_from_thread_id":"parent-gemini-thread","request_kind":"turn"}`, + }, + }, + OriginalRequest: []byte(`{"input":[{"role":"user","content":"fork query"}]}`), + Metadata: map[string]any{}, + } + forkAuth, errFork := selector.Pick(context.Background(), "mixed", "gemini-3.8-flash-high", forkOpts, auths) + if errFork != nil { + t.Fatalf("fork Pick() error = %v", errFork) + } + if forkAuth.ID != parentAuth.ID { + t.Fatalf("conversational fork did not inherit parent auth on Gemini: got %q, want %q", forkAuth.ID, parentAuth.ID) + } + if got := forkOpts.Metadata[cliproxyexecutor.CanonicalSessionIDMetadataKey]; got != "codex:fork-gemini-thread" { + t.Fatalf("fork canonical session ID = %v, want codex:fork-gemini-thread", got) + } + if got := forkOpts.Metadata[cliproxyexecutor.ParentSessionIDMetadataKey]; got != "codex:parent-gemini-thread" { + t.Fatalf("fork parent session ID = %v, want codex:parent-gemini-thread", got) + } + if isFork, ok := forkOpts.Metadata[cliproxyexecutor.IsForkMetadataKey].(bool); !ok || !isFork { + t.Fatalf("expected is_fork=true in metadata, got %v", forkOpts.Metadata[cliproxyexecutor.IsForkMetadataKey]) + } +} + +func TestSessionAffinityCodexMultiAgentV2CollabSpawn(t *testing.T) { + t.Parallel() + + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + + auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}} + + // 1. Parent session in Codex CLI + parentOpts := cliproxyexecutor.Options{ + Headers: http.Header{ + "Session-Id": []string{"parent-root-100"}, + "Thread-Id": []string{"parent-root-100"}, + "X-Codex-Turn-Metadata": []string{ + `{"session_id":"parent-root-100","thread_id":"parent-root-100","agent_name":"/root","request_kind":"turn"}`, + }, + }, + OriginalRequest: []byte(`{"input":[{"role":"user","content":"parent task"}]}`), + Metadata: map[string]any{}, + } + parentAuth, errParent := selector.Pick(context.Background(), "openai", "model", parentOpts, auths) + if errParent != nil { + t.Fatalf("parent Pick() error = %v", errParent) + } + + // 2. Multi-Agent v2 child subagent (collab_spawn) + childOpts := cliproxyexecutor.Options{ + Headers: http.Header{ + "Session-Id": []string{"parent-root-100"}, + "Thread-Id": []string{"subagent-thread-200"}, + "X-Codex-Parent-Thread-Id": []string{"parent-root-100"}, + "X-Openai-Subagent": []string{"collab_spawn"}, + "X-Codex-Turn-Metadata": []string{ + `{"session_id":"parent-root-100","thread_id":"subagent-thread-200","agent_name":"/root/check_readme","parent_thread_id":"parent-root-100","subagent_kind":"thread_spawn"}`, + }, + }, + OriginalRequest: []byte(`{"input":[{"role":"user","content":"child check readme"}]}`), + Metadata: map[string]any{}, + } + childAuth, errChild := selector.Pick(context.Background(), "openai", "model", childOpts, auths) + if errChild != nil { + t.Fatalf("child Pick() error = %v", errChild) + } + if childAuth.ID != parentAuth.ID { + t.Fatalf("child subagent did not inherit parent auth: got %q, want %q", childAuth.ID, parentAuth.ID) + } + if got := childOpts.Metadata[cliproxyexecutor.CanonicalSessionIDMetadataKey]; got != "codex:parent-root-100:agent:check_readme" { + t.Fatalf("child canonical session ID = %v, want codex:parent-root-100:agent:check_readme", got) + } + if got := childOpts.Metadata[cliproxyexecutor.ParentSessionIDMetadataKey]; got != "codex:parent-root-100" { + t.Fatalf("child parent session ID = %v, want codex:parent-root-100", got) + } + + // 3. Failure on child subagent does not evict parent binding + selector.OnResult(Result{ + Provider: "openai", + Model: "model", + AuthID: childAuth.ID, + Success: false, + Options: childOpts, + }) + + // Parent session should still be bound to parentAuth + parentOpts2 := cliproxyexecutor.Options{ + Headers: http.Header{ + "Session-Id": []string{"parent-root-100"}, + "Thread-Id": []string{"parent-root-100"}, + }, + OriginalRequest: []byte(`{"input":[{"role":"user","content":"parent task 2"}]}`), + Metadata: map[string]any{}, + } + parentAuth2, errParent2 := selector.Pick(context.Background(), "openai", "model", parentOpts2, auths) + if errParent2 != nil { + t.Fatalf("parent Pick() 2 error = %v", errParent2) + } + if parentAuth2.ID != parentAuth.ID { + t.Fatalf("parent auth binding was evicted by subagent failure: got %q, want %q", parentAuth2.ID, parentAuth.ID) + } +} + +func TestSessionAffinityBodyOnlyCodexFork(t *testing.T) { + t.Parallel() + + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + + auths := []*Auth{{ID: "auth-1", Provider: "antigravity"}, {ID: "auth-2", Provider: "antigravity"}} + + // 1. Parent thread in body + parentOpts := cliproxyexecutor.Options{ + OriginalRequest: []byte(`{"thread_id":"body-parent-100","messages":[{"role":"user","content":"parent"}]}`), + Metadata: map[string]any{}, + } + parentAuth, errParent := selector.Pick(context.Background(), "mixed", "gemini-3.8-flash-high", parentOpts, auths) + if errParent != nil { + t.Fatalf("parent Pick() error = %v", errParent) + } + + // 2. Forked thread in body without headers (e.g. {"thread_id":"child","forked_from_thread_id":"parent"}) + forkOpts := cliproxyexecutor.Options{ + OriginalRequest: []byte(`{"thread_id":"body-fork-200","forked_from_thread_id":"body-parent-100","messages":[{"role":"user","content":"fork"}]}`), + Metadata: map[string]any{}, + } + forkAuth, errFork := selector.Pick(context.Background(), "mixed", "gemini-3.8-flash-high", forkOpts, auths) + if errFork != nil { + t.Fatalf("fork Pick() error = %v", errFork) + } + if forkAuth.ID != parentAuth.ID { + t.Fatalf("body-only fork did not inherit parent auth on Gemini: got %q, want %q", forkAuth.ID, parentAuth.ID) + } + if isFork, ok := forkOpts.Metadata[cliproxyexecutor.IsForkMetadataKey].(bool); !ok || !isFork { + t.Fatalf("expected is_fork=true for body-only fork, got %v", forkOpts.Metadata[cliproxyexecutor.IsForkMetadataKey]) + } +} + +func TestSessionAffinityForkRebindDoesNotMutateParentBinding(t *testing.T) { + t.Parallel() + + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + + auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}} + + // 1. Parent thread binds to auth-a + parentOpts := cliproxyexecutor.Options{ + Headers: http.Header{"Session-Id": []string{"parent-thread-alpha"}}, + OriginalRequest: []byte(`{"input":[{"role":"user","content":"parent"}]}`), + Metadata: map[string]any{}, + } + parentAuth, errParent := selector.Pick(context.Background(), "openai", "model", parentOpts, auths) + if errParent != nil || parentAuth.ID != "auth-a" { + t.Fatalf("parent Pick() error = %v, auth = %v", errParent, parentAuth) + } + + // 2. Fork thread inherits parent auth-a + forkOpts := cliproxyexecutor.Options{ + Headers: http.Header{ + "Session-Id": []string{"fork-thread-beta"}, + "X-Codex-Turn-Metadata": []string{ + `{"session_id":"fork-thread-beta","forked_from_thread_id":"parent-thread-alpha","request_kind":"turn"}`, + }, + }, + OriginalRequest: []byte(`{"input":[{"role":"user","content":"fork"}]}`), + Metadata: map[string]any{}, + } + forkAuth, errFork := selector.Pick(context.Background(), "openai", "model", forkOpts, auths) + if errFork != nil || forkAuth.ID != "auth-a" { + t.Fatalf("fork Pick() error = %v, auth = %v", errFork, forkAuth) + } + + // 3. Fork thread fails and rebinds to auth-b (failover rebind) + selector.OnResult(Result{ + Provider: "openai", + Model: "model", + AuthID: "auth-a", + Success: false, + Options: forkOpts, + }) + + // Re-pick fork with auth-b + forkOpts2 := cliproxyexecutor.Options{ + Headers: http.Header{ + "Session-Id": []string{"fork-thread-beta"}, + "X-Codex-Turn-Metadata": []string{ + `{"session_id":"fork-thread-beta","forked_from_thread_id":"parent-thread-alpha","request_kind":"turn"}`, + }, + }, + OriginalRequest: []byte(`{"input":[{"role":"user","content":"fork retry"}]}`), + Metadata: map[string]any{}, + } + forkAuth2, errFork2 := selector.Pick(context.Background(), "openai", "model", forkOpts2, []*Auth{{ID: "auth-b"}}) + if errFork2 != nil || forkAuth2.ID != "auth-b" { + t.Fatalf("fork retry Pick() error = %v, auth = %v", errFork2, forkAuth2) + } + + // 4. Parent session MUST STILL BE BOUND to auth-a, not mutated to auth-b! + parentOpts2 := cliproxyexecutor.Options{ + Headers: http.Header{"Session-Id": []string{"parent-thread-alpha"}}, + OriginalRequest: []byte(`{"input":[{"role":"user","content":"parent turn 2"}]}`), + Metadata: map[string]any{}, + } + parentAuth2, errParent2 := selector.Pick(context.Background(), "openai", "model", parentOpts2, auths) + if errParent2 != nil || parentAuth2.ID != "auth-a" { + t.Fatalf("parent auth binding was mutated by fork failover: got %q, want auth-a", parentAuth2.ID) + } +} + +func TestSessionAffinityCodexSubagentWithOmittedThreadIdRetainsParent(t *testing.T) { + t.Parallel() + + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + + auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}} + + // Parent binds + parentOpts := cliproxyexecutor.Options{ + Headers: http.Header{"Session-Id": []string{"sess-main-999"}}, + OriginalRequest: []byte(`{"input":[{"role":"user","content":"parent"}]}`), + Metadata: map[string]any{}, + } + parentAuth, _ := selector.Pick(context.Background(), "openai", "model", parentOpts, auths) + + // Subagent sends Session-Id, but Thread-Id is omitted, while agent_name is set in turn metadata + subOpts := cliproxyexecutor.Options{ + Headers: http.Header{ + "Session-Id": []string{"sess-main-999"}, + "X-Openai-Subagent": []string{"collab_spawn"}, + "X-Codex-Turn-Metadata": []string{ + `{"session_id":"sess-main-999","agent_name":"/root/worker","subagent_kind":"thread_spawn"}`, + }, + }, + OriginalRequest: []byte(`{"input":[{"role":"user","content":"subagent"}]}`), + Metadata: map[string]any{}, + } + subAuth, errSub := selector.Pick(context.Background(), "openai", "model", subOpts, auths) + if errSub != nil { + t.Fatalf("subagent Pick() error = %v", errSub) + } + if subAuth.ID != parentAuth.ID { + t.Fatalf("subagent did not inherit parent auth: got %q, want %q", subAuth.ID, parentAuth.ID) + } + if got := subOpts.Metadata[cliproxyexecutor.ParentSessionIDMetadataKey]; got != "codex:sess-main-999" { + t.Fatalf("expected ParentSessionID=codex:sess-main-999, got %v", got) + } + if got := subOpts.Metadata[cliproxyexecutor.CanonicalSessionIDMetadataKey]; got != "codex:sess-main-999:agent:worker" { + t.Fatalf("expected CanonicalSessionID=codex:sess-main-999:agent:worker, got %v", got) + } } func TestSessionAffinityPayloadParentSessionInheritance(t *testing.T) { @@ -856,6 +1172,314 @@ func TestSessionAffinityClaudeMetadataSubagentNonInheritingGeminiModel(t *testin } } +func TestSessionAffinitySelectorLCPForkDerivesDistinctSessionIDAndParentLineage(t *testing.T) { + t.Parallel() + + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Hour, + }) + defer selector.Stop() + + auths := []*Auth{{ID: "auth-1"}, {ID: "auth-2"}} + + // 1. Root conversation: 3 user turns + 2 assistant answers + rootOpts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + OriginalRequest: []byte(`{"messages":[` + + `{"role":"user","content":"turn 1"},` + + `{"role":"assistant","content":"ans 1"},` + + `{"role":"user","content":"turn 2 trunk"},` + + `{"role":"assistant","content":"ans 2 trunk"},` + + `{"role":"user","content":"turn 3 trunk"}` + + `]}`), + Metadata: map[string]any{ + cliproxyexecutor.CallerScopeMetadataKey: "caller-user-1", + }, + } + rootAuth, errRoot := selector.Pick(context.Background(), "openai", "model", rootOpts, auths) + if errRoot != nil { + t.Fatalf("root Pick() error = %v", errRoot) + } + rootSessionID, ok := rootOpts.Metadata[cliproxyexecutor.CanonicalSessionIDMetadataKey].(string) + if !ok || rootSessionID == "" { + t.Fatalf("expected non-empty canonical root session ID, got %#v", rootOpts.Metadata[cliproxyexecutor.CanonicalSessionIDMetadataKey]) + } + if parentID := rootOpts.Metadata[cliproxyexecutor.ParentSessionIDMetadataKey]; parentID != nil { + t.Fatalf("expected nil parent ID on root session, got %#v", parentID) + } + + // 2. Fork request: shares turn 1 & ans 1, but diverges on turn 2 + forkOpts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + OriginalRequest: []byte(`{"messages":[` + + `{"role":"user","content":"turn 1"},` + + `{"role":"assistant","content":"ans 1"},` + + `{"role":"user","content":"turn 2 fork branch B"}` + + `]}`), + Metadata: map[string]any{ + cliproxyexecutor.CallerScopeMetadataKey: "caller-user-1", + }, + } + forkAuth, errFork := selector.Pick(context.Background(), "openai", "model", forkOpts, auths) + if errFork != nil { + t.Fatalf("fork Pick() error = %v", errFork) + } + // Routing MUST keep the exact same auth for hardware KV cache reuse + if forkAuth.ID != rootAuth.ID { + t.Fatalf("fork routed to %q, want same auth %q as root session", forkAuth.ID, rootAuth.ID) + } + forkSessionID, ok := forkOpts.Metadata[cliproxyexecutor.CanonicalSessionIDMetadataKey].(string) + if !ok || forkSessionID == "" { + t.Fatalf("expected non-empty canonical fork session ID, got %#v", forkOpts.Metadata[cliproxyexecutor.CanonicalSessionIDMetadataKey]) + } + // Branch identity MUST differ from root + if forkSessionID == rootSessionID { + t.Fatalf("fork session ID %q should differ from root session ID %q", forkSessionID, rootSessionID) + } + forkParentID, ok := forkOpts.Metadata[cliproxyexecutor.ParentSessionIDMetadataKey].(string) + if !ok || forkParentID == "" { + t.Fatalf("expected non-empty parent session ID on fork, got %#v", forkOpts.Metadata[cliproxyexecutor.ParentSessionIDMetadataKey]) + } + + // 3. Linear continuation on the fork branch (turn 3 on branch B) + forkContOpts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + OriginalRequest: []byte(`{"messages":[` + + `{"role":"user","content":"turn 1"},` + + `{"role":"assistant","content":"ans 1"},` + + `{"role":"user","content":"turn 2 fork branch B"},` + + `{"role":"assistant","content":"ans 2 fork branch B"},` + + `{"role":"user","content":"turn 3 fork branch B"}` + + `]}`), + Metadata: map[string]any{ + cliproxyexecutor.CallerScopeMetadataKey: "caller-user-1", + }, + } + forkContAuth, errForkCont := selector.Pick(context.Background(), "openai", "model", forkContOpts, auths) + if errForkCont != nil { + t.Fatalf("forkCont Pick() error = %v", errForkCont) + } + if forkContAuth.ID != forkAuth.ID { + t.Fatalf("fork continuation routed to %q, want same auth %q", forkContAuth.ID, forkAuth.ID) + } + contSessionID := forkContOpts.Metadata[cliproxyexecutor.CanonicalSessionIDMetadataKey] + if contSessionID != forkSessionID { + t.Fatalf("fork continuation session ID = %q, want identical to fork session %q", contSessionID, forkSessionID) + } + contParentID := forkContOpts.Metadata[cliproxyexecutor.ParentSessionIDMetadataKey] + if contParentID != forkParentID { + t.Fatalf("fork continuation parent ID = %q, want identical to fork parent %q", contParentID, forkParentID) + } +} + +func TestSessionAffinitySelectorLCPFailureEvictionOnCacheHit(t *testing.T) { + t.Parallel() + + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Hour, + }) + defer selector.Stop() + + auths := []*Auth{{ID: "auth-1"}, {ID: "auth-2"}} + + opts1 := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + OriginalRequest: []byte(`{"messages":[{"role":"user","content":"test failure eviction"}]}`), + Metadata: map[string]any{ + cliproxyexecutor.CallerScopeMetadataKey: "test-caller", + }, + } + + // 1. Initial request binds to auth-1 and succeeds + auth1, err1 := selector.Pick(context.Background(), "openai", "model", opts1, auths) + if err1 != nil { + t.Fatalf("Pick 1 error: %v", err1) + } + selector.OnResult(Result{ + Provider: "openai", + AuthID: auth1.ID, + Options: opts1, + Success: true, + }) + + // 2. Second request hits LCP cache for auth-1 + opts2 := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + OriginalRequest: []byte(`{"messages":[{"role":"user","content":"test failure eviction"}]}`), + Metadata: map[string]any{ + cliproxyexecutor.CallerScopeMetadataKey: "test-caller", + }, + } + auth2, err2 := selector.Pick(context.Background(), "openai", "model", opts2, auths) + if err2 != nil { + t.Fatalf("Pick 2 error: %v", err2) + } + if auth2.ID != auth1.ID { + t.Fatalf("Pick 2 did not hit LCP cache: got %q, want %q", auth2.ID, auth1.ID) + } + + // Second request fails upstream (e.g. 500 error) + selector.OnResult(Result{ + Provider: "openai", + AuthID: auth2.ID, + Options: opts2, + Success: false, + Error: &Error{Code: "upstream_500", Message: "500 internal server error"}, + }) + + // 3. Third request should see the failed binding evicted, and fall back to round-robin + opts3 := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAI, + OriginalRequest: []byte(`{"messages":[{"role":"user","content":"test failure eviction"}]}`), + Metadata: map[string]any{ + cliproxyexecutor.CallerScopeMetadataKey: "test-caller", + }, + } + auth3, err3 := selector.Pick(context.Background(), "openai", "model", opts3, auths) + if err3 != nil { + t.Fatalf("Pick 3 error: %v", err3) + } + // Since auth-1 was evicted from LCP on failure, round-robin picks auth-2 + if auth3.ID != "auth-2" { + t.Fatalf("Pick 3 was not evicted on failure: got %q, want round-robin fallback to auth-2", auth3.ID) + } +} + +func TestSessionAffinityCodexForkWithBothSessionAndThreadIDs(t *testing.T) { + t.Parallel() + + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + + auths := []*Auth{{ID: "auth-codex-1"}, {ID: "auth-codex-2"}} + + // Parent session + parentOpts := cliproxyexecutor.Options{ + Headers: http.Header{ + "Session-Id": []string{"parent-thread-100"}, + "Thread-Id": []string{"parent-thread-100"}, + }, + Metadata: map[string]any{}, + } + parentAuth, _ := selector.Pick(context.Background(), "openai", "model", parentOpts, auths) + + // Child fork has Session-Id: parent-thread-100, Thread-Id: child-thread-200, and forked_from_thread_id: parent-thread-100 + forkOpts := cliproxyexecutor.Options{ + Headers: http.Header{ + "Session-Id": []string{"parent-thread-100"}, + "Thread-Id": []string{"child-thread-200"}, + "X-Codex-Turn-Metadata": []string{ + `{"session_id":"parent-thread-100","thread_id":"child-thread-200","forked_from_thread_id":"parent-thread-100"}`, + }, + }, + Metadata: map[string]any{}, + } + forkAuth, errFork := selector.Pick(context.Background(), "openai", "model", forkOpts, auths) + if errFork != nil { + t.Fatalf("fork Pick() error = %v", errFork) + } + if forkAuth.ID != parentAuth.ID { + t.Fatalf("fork did not inherit parent auth: got %q, want %q", forkAuth.ID, parentAuth.ID) + } + // Child must NOT collapse onto parent! + if got := forkOpts.Metadata[cliproxyexecutor.CanonicalSessionIDMetadataKey]; got != "codex:child-thread-200" { + t.Fatalf("child fork session ID collapsed onto parent: got %v, want codex:child-thread-200", got) + } + if got := forkOpts.Metadata[cliproxyexecutor.ParentSessionIDMetadataKey]; got != "codex:parent-thread-100" { + t.Fatalf("child fork parent ID = %v, want codex:parent-thread-100", got) + } +} + +func TestSessionAffinityNestedMetadataForkedFromThreadID(t *testing.T) { + t.Parallel() + + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + + auths := []*Auth{{ID: "auth-1"}, {ID: "auth-2"}} + + // Parent + parentOpts := cliproxyexecutor.Options{ + OriginalRequest: []byte(`{"thread_id":"parent-t-1"}`), + Metadata: map[string]any{}, + } + parentAuth, _ := selector.Pick(context.Background(), "openai", "model", parentOpts, auths) + + // Child fork with nested metadata.forked_from_thread_id + forkOpts := cliproxyexecutor.Options{ + OriginalRequest: []byte(`{"thread_id":"child-t-2","metadata":{"forked_from_thread_id":"parent-t-1"}}`), + Metadata: map[string]any{}, + } + forkAuth, errFork := selector.Pick(context.Background(), "openai", "model", forkOpts, auths) + if errFork != nil { + t.Fatalf("nested fork Pick() error = %v", errFork) + } + if forkAuth.ID != parentAuth.ID { + t.Fatalf("nested fork did not inherit parent auth: got %q, want %q", forkAuth.ID, parentAuth.ID) + } + if isFork, ok := forkOpts.Metadata[cliproxyexecutor.IsForkMetadataKey].(bool); !ok || !isFork { + t.Fatalf("expected is_fork=true for nested metadata fork, got %v", forkOpts.Metadata[cliproxyexecutor.IsForkMetadataKey]) + } + if got := forkOpts.Metadata[cliproxyexecutor.ParentSessionIDMetadataKey]; got != "thread:parent-t-1" { + t.Fatalf("expected ParentSessionID=thread:parent-t-1, got %v", got) + } +} + +func TestSessionAffinityCodexForkWithSessionIdHeaderAndBodyThreadId(t *testing.T) { + t.Parallel() + + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + + auths := []*Auth{{ID: "auth-1", Provider: "antigravity"}, {ID: "auth-2", Provider: "antigravity"}} + + // Parent binds with Session-Id header + parentOpts := cliproxyexecutor.Options{ + Headers: http.Header{"Session-Id": []string{"parent-sess-uuid"}}, + Metadata: map[string]any{}, + } + parentAuth, _ := selector.Pick(context.Background(), "mixed", "gemini-3.8-flash-high", parentOpts, auths) + + // Fork carries Session-Id header (parent-sess-uuid), but body contains thread_id (child-thread-uuid) and metadata.forked_from_thread_id + forkOpts := cliproxyexecutor.Options{ + Headers: http.Header{"Session-Id": []string{"parent-sess-uuid"}}, + OriginalRequest: []byte(`{ + "thread_id": "child-thread-uuid", + "metadata": { + "forked_from_thread_id": "parent-sess-uuid" + } + }`), + Metadata: map[string]any{}, + } + forkAuth, errFork := selector.Pick(context.Background(), "mixed", "gemini-3.8-flash-high", forkOpts, auths) + if errFork != nil { + t.Fatalf("fork Pick() error = %v", errFork) + } + if forkAuth.ID != parentAuth.ID { + t.Fatalf("fork did not inherit parent auth on Gemini: got %q, want %q", forkAuth.ID, parentAuth.ID) + } + if got := forkOpts.Metadata[cliproxyexecutor.CanonicalSessionIDMetadataKey]; got != "codex:child-thread-uuid" { + t.Fatalf("child fork collapsed onto parent: got %v, want codex:child-thread-uuid", got) + } + if got := forkOpts.Metadata[cliproxyexecutor.ParentSessionIDMetadataKey]; got != "codex:parent-sess-uuid" { + t.Fatalf("child fork parent ID = %v, want codex:parent-sess-uuid", got) + } + if isFork, ok := forkOpts.Metadata[cliproxyexecutor.IsForkMetadataKey].(bool); !ok || !isFork { + t.Fatalf("expected is_fork=true, got %v", forkOpts.Metadata[cliproxyexecutor.IsForkMetadataKey]) + } +} + func BenchmarkSessionAffinitySelectorPickLCP(b *testing.B) { log.SetLevel(log.WarnLevel) defer log.SetLevel(log.InfoLevel) diff --git a/sdk/cliproxy/auth/session_overhead_benchmark_test.go b/sdk/cliproxy/auth/session_overhead_benchmark_test.go new file mode 100644 index 00000000..239a3f25 --- /dev/null +++ b/sdk/cliproxy/auth/session_overhead_benchmark_test.go @@ -0,0 +1,278 @@ +package auth + +import ( + "context" + "fmt" + "net/http" + "testing" + "time" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + cliproxysession "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/session" + log "github.com/sirupsen/logrus" +) + +// 1. Explicit Header Fast Path (Pi, Claude Code, OpenCode) +func BenchmarkSessionExplicitHeaderFastPath(b *testing.B) { + log.SetLevel(log.WarnLevel) + defer log.SetLevel(log.InfoLevel) + + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + TTL: time.Hour, + }) + defer selector.Stop() + + auths := []*Auth{ + {ID: "auth-1", Provider: "openai", Status: StatusActive}, + {ID: "auth-2", Provider: "openai", Status: StatusActive}, + } + + opts := cliproxyexecutor.Options{ + Headers: http.Header{ + "Session-Id": []string{"pi-interactive-session-abc-123"}, + }, + OriginalRequest: []byte(`{"messages":[{"role":"user","content":"hello"}]}`), + Metadata: map[string]any{}, + } + + // Warmup and bind + ctx := context.Background() + _, _ = selector.Pick(ctx, "openai", "gpt-5", opts, auths) + + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + // Clear transient metadata + delete(opts.Metadata, cliproxyexecutor.CanonicalSessionIDMetadataKey) + delete(opts.Metadata, cliproxyexecutor.ParentSessionIDMetadataKey) + delete(opts.Metadata, cliproxyexecutor.IsForkMetadataKey) + + _, _ = selector.Pick(ctx, "openai", "gpt-5", opts, auths) + } +} + +// 2. Codex Multi-Agent v2 Path (Header + JSON Turn Metadata + Subagent Isolation) +func BenchmarkSessionCodexMultiAgentV2Path(b *testing.B) { + log.SetLevel(log.WarnLevel) + defer log.SetLevel(log.InfoLevel) + + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + TTL: time.Hour, + }) + defer selector.Stop() + + auths := []*Auth{ + {ID: "auth-codex-1", Provider: "openai", Status: StatusActive}, + {ID: "auth-codex-2", Provider: "openai", Status: StatusActive}, + } + + // Parent binding + parentOpts := cliproxyexecutor.Options{ + Headers: http.Header{ + "Session-Id": []string{"parent-session-100"}, + }, + Metadata: map[string]any{}, + } + ctx := context.Background() + _, _ = selector.Pick(ctx, "openai", "codex-5", parentOpts, auths) + + // Child subagent request + childOpts := cliproxyexecutor.Options{ + Headers: http.Header{ + "Session-Id": []string{"parent-session-100"}, + "Thread-Id": []string{"subagent-thread-200"}, + "X-Codex-Parent-Thread-Id": []string{"parent-session-100"}, + "X-Openai-Subagent": []string{"collab_spawn"}, + "X-Codex-Turn-Metadata": []string{ + `{"session_id":"parent-session-100","thread_id":"subagent-thread-200","agent_name":"/root/check_readme","parent_thread_id":"parent-session-100","subagent_kind":"thread_spawn"}`, + }, + }, + OriginalRequest: []byte(`{"input":[{"role":"user","content":"inspect"}]}`), + Metadata: map[string]any{}, + } + + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + delete(childOpts.Metadata, cliproxyexecutor.CanonicalSessionIDMetadataKey) + delete(childOpts.Metadata, cliproxyexecutor.ParentSessionIDMetadataKey) + delete(childOpts.Metadata, cliproxyexecutor.IsForkMetadataKey) + + _, _ = selector.Pick(ctx, "openai", "codex-5", childOpts, auths) + } +} + +// 3. Codex Fork Path (Header + Turn Metadata Fork Derivation) +func BenchmarkSessionCodexForkPath(b *testing.B) { + log.SetLevel(log.WarnLevel) + defer log.SetLevel(log.InfoLevel) + + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + TTL: time.Hour, + }) + defer selector.Stop() + + auths := []*Auth{ + {ID: "auth-codex-1", Provider: "openai", Status: StatusActive}, + {ID: "auth-codex-2", Provider: "openai", Status: StatusActive}, + } + + // Parent + parentOpts := cliproxyexecutor.Options{ + Headers: http.Header{ + "Session-Id": []string{"parent-thread-alpha"}, + }, + Metadata: map[string]any{}, + } + ctx := context.Background() + _, _ = selector.Pick(ctx, "openai", "codex-5", parentOpts, auths) + + forkOpts := cliproxyexecutor.Options{ + Headers: http.Header{ + "Session-Id": []string{"fork-thread-beta"}, + "X-Codex-Turn-Metadata": []string{ + `{"session_id":"fork-thread-beta","forked_from_thread_id":"parent-thread-alpha","request_kind":"turn"}`, + }, + }, + OriginalRequest: []byte(`{"input":[{"role":"user","content":"fork turn"}]}`), + Metadata: map[string]any{}, + } + + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + delete(forkOpts.Metadata, cliproxyexecutor.CanonicalSessionIDMetadataKey) + delete(forkOpts.Metadata, cliproxyexecutor.ParentSessionIDMetadataKey) + delete(forkOpts.Metadata, cliproxyexecutor.IsForkMetadataKey) + + _, _ = selector.Pick(ctx, "openai", "codex-5", forkOpts, auths) + } +} + +// 4. Body-Only Fallback Path (GJSON payload parsing for thread_id & forked_from_thread_id) +func BenchmarkSessionBodyOnlyForkPath(b *testing.B) { + log.SetLevel(log.WarnLevel) + defer log.SetLevel(log.InfoLevel) + + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + TTL: time.Hour, + }) + defer selector.Stop() + + auths := []*Auth{ + {ID: "auth-1", Provider: "openai", Status: StatusActive}, + } + + // Parent in body + parentOpts := cliproxyexecutor.Options{ + OriginalRequest: []byte(`{"thread_id":"body-parent-100","messages":[{"role":"user","content":"hello"}]}`), + Metadata: map[string]any{}, + } + ctx := context.Background() + _, _ = selector.Pick(ctx, "openai", "gpt-5", parentOpts, auths) + + forkOpts := cliproxyexecutor.Options{ + OriginalRequest: []byte(`{"thread_id":"body-fork-200","forked_from_thread_id":"body-parent-100","messages":[{"role":"user","content":"fork"}]}`), + Metadata: map[string]any{}, + } + + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + delete(forkOpts.Metadata, cliproxyexecutor.CanonicalSessionIDMetadataKey) + delete(forkOpts.Metadata, cliproxyexecutor.ParentSessionIDMetadataKey) + delete(forkOpts.Metadata, cliproxyexecutor.IsForkMetadataKey) + + _, _ = selector.Pick(ctx, "openai", "gpt-5", forkOpts, auths) + } +} + +// 5. SessionCache Concurrent High QPS Stress (simulating multi-threaded worker pool) +func BenchmarkSessionCacheConcurrentHighQPS(b *testing.B) { + cache := NewSessionCache(time.Hour) + defer cache.Stop() + + // Pre-populate 1000 active sessions + for i := 0; i < 1000; i++ { + cache.Set(fmt.Sprintf("session-%04d", i), fmt.Sprintf("auth-%d", i%10)) + } + + b.ReportAllocs() + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + idx := 0 + for pb.Next() { + key := fmt.Sprintf("session-%04d", idx%1000) + // 90% read/touch, 10% write + if idx%10 == 0 { + cache.Set(key, fmt.Sprintf("auth-%d", idx%10)) + } else { + _, _ = cache.Get(key) + cache.Touch(key, fmt.Sprintf("auth-%d", idx%10)) + } + idx++ + } + }) +} + +// 6. MerklePrefixMatcher Concurrent High QPS (simulating parallel LCP lookups) +func BenchmarkMerklePrefixMatcherConcurrentHighQPS(b *testing.B) { + matcher := cliproxysession.NewMerklePrefixMatcher(time.Hour) + + // Pre-populate trunk conversation + texts := make([]string, 16) + for i := 0; i < 16; i++ { + texts[i] = fmt.Sprintf("system instruction turn %d", i) + } + turns := make([]cliproxysession.CanonicalTurn, 16) + for i, t := range texts { + turns[i] = cliproxysession.CanonicalTurn{ + Role: "user", + Parts: []cliproxysession.CanonicalPart{ + {Kind: "text", Value: t}, + }, + } + } + matcher.Bind("lcp:v1:bench:group-0", turns, "auth-1") + + b.ReportAllocs() + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + _, _ = matcher.Match("lcp:v1:bench:group-0", turns) + } + }) +} + +// 7. Large Scale Memory & LRU Eviction Benchmark (Simulates continuous session churn) +func BenchmarkSessionCacheScale100kEviction(b *testing.B) { + cache := NewSessionCacheWithCapacity(time.Hour, 10000) + defer cache.Stop() + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + key := fmt.Sprintf("session-scale-%d", i) + cache.Set(key, "auth-target") + } +} + +// 8. MerklePrefixMatcher Large Scale LRU Eviction (Simulates continuous conversation churn with maxGroups=1000) +func BenchmarkMerkleMatcherScaleLRUEviction(b *testing.B) { + matcher := cliproxysession.NewMerklePrefixMatcherWithConfig(cliproxysession.MerklePrefixMatcherConfig{ + TTL: time.Hour, + MaxTurns: 32, + MaxGroups: 1000, + }) + + turns := []cliproxysession.CanonicalTurn{ + {Role: "user", Parts: []cliproxysession.CanonicalPart{{Kind: "text", Value: "hello world"}}}, + } + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + ns := fmt.Sprintf("lcp:v1:group-%d", i) + matcher.Bind(ns, turns, "auth-target") + } +} diff --git a/sdk/cliproxy/executor/types.go b/sdk/cliproxy/executor/types.go index 7b6b1632..a1a648f0 100644 --- a/sdk/cliproxy/executor/types.go +++ b/sdk/cliproxy/executor/types.go @@ -55,6 +55,13 @@ const ( // across explicit harness headers, body fields, execution sessions, LCP inference, // and fallback context derivation for unified debugging and cross-subsystem tracing. CanonicalSessionIDMetadataKey = "canonical_session_id" + // ParentSessionIDMetadataKey stores the parent session identity for hierarchical sessions and forks. + // For top-level Merkle LCP forks, it represents the deterministic Merkle prefix hash at the divergence point. + ParentSessionIDMetadataKey = "parent_session_id" + // IsForkMetadataKey indicates whether the request represents a conversational branch or fork. + IsForkMetadataKey = "is_fork" + // LCPAccessGenerationMetadataKey stores the monotonic access generation when an LCP entry was touched or bound. + LCPAccessGenerationMetadataKey = "lcp_access_generation" // LCPFingerprintMetadataKey stores bounded request-scoped turn fingerprints so // SessionAffinitySelector.OnResult can avoid reparsing the original payload. LCPFingerprintMetadataKey = "lcp_fingerprints" diff --git a/sdk/cliproxy/session/info.go b/sdk/cliproxy/session/info.go index 79efc50f..4eef2483 100644 --- a/sdk/cliproxy/session/info.go +++ b/sdk/cliproxy/session/info.go @@ -20,6 +20,8 @@ type SessionInfo struct { AuthID string `json:"auth_id,omitempty"` Provider string `json:"provider,omitempty"` Model string `json:"model,omitempty"` + IsFork bool `json:"is_fork,omitempty"` + IsSubagent bool `json:"is_subagent,omitempty"` Metadata map[string]any `json:"metadata,omitempty"` } @@ -67,7 +69,9 @@ func ExtractSessionInfo(headers http.Header, payload []byte, metadata map[string "forked_from_thread_id", "forked_from_id", "parent_conversation_id", "parentConversationId", "metadata.parent_session_id", "metadata.parent_thread_id", + "metadata.forked_from_thread_id", "metadata.forked_from_id", "extra_body.parent_session_id", "extra_body.parent_thread_id", + "extra_body.forked_from_thread_id", "extra_body.forked_from_id", } { if val := normalizedSessionCandidate(root.Get(p).String()); val != "" { parentCandidate = val @@ -173,37 +177,164 @@ func ExtractSessionInfo(headers http.Header, payload []byte, metadata map[string } // 3. OpenAI / Codex CLI Headers - if sid := sessionHeaderValue(headers, "Session-Id"); sid != "" { - info.ClientType = "codex" - info.SessionID = "codex:" + sid - parentThread := sessionHeaderValue(headers, "x-codex-parent-thread-id") - if parentThread == "" { - parentThread = sessionHeaderValue(headers, "X-Codex-Parent-Thread-Id") + sid := sessionHeaderValue(headers, "Session-Id") + if sid == "" { + sid = sessionHeaderValue(headers, "Session_id") + } + tid := sessionHeaderValue(headers, "Thread-Id") + if tid == "" { + tid = sessionHeaderValue(headers, "Thread_id") + } + + var codexTurnMeta string + for k, v := range headers { + if strings.EqualFold(k, "X-Codex-Turn-Metadata") && len(v) > 0 { + codexTurnMeta = strings.TrimSpace(v[0]) + break } - if parentThread != "" && parentThread != sid { - info.ParentSessionID = "codex:" + parentThread - info.AgentName = "subagent" - } else if parentCandidate != "" && parentCandidate != sid { - info.ParentSessionID = "codex:" + parentCandidate - info.AgentName = "subagent" - } else { - info.AgentName = "main" + } + var codexTurnMetaJSON gjson.Result + if codexTurnMeta != "" { + codexTurnMetaJSON = gjson.Parse(codexTurnMeta) + } + + if sid == "" && codexTurnMetaJSON.Exists() { + sid = normalizedSessionCandidate(codexTurnMetaJSON.Get("session_id").String()) + } + if tid == "" && codexTurnMetaJSON.Exists() { + tid = normalizedSessionCandidate(codexTurnMetaJSON.Get("thread_id").String()) + } + if tid == "" && sid != "" && root.Exists() { + for _, path := range []string{"thread_id", "threadId", "metadata.thread_id"} { + if tid = normalizedSessionCandidate(root.Get(path).String()); tid != "" { + break + } + if hasNestedReq { + if tid = normalizedSessionCandidate(reqRoot.Get(path).String()); tid != "" { + break + } + } } - return finalizeSessionInfo(info) } - if sid := sessionHeaderValue(headers, "Session_id"); sid != "" { + + if sid != "" || tid != "" { info.ClientType = "codex" - info.SessionID = "codex:" + sid parentThread := sessionHeaderValue(headers, "x-codex-parent-thread-id") if parentThread == "" { parentThread = sessionHeaderValue(headers, "X-Codex-Parent-Thread-Id") } - if parentThread != "" && parentThread != sid { + if parentThread == "" && codexTurnMetaJSON.Exists() { + parentThread = normalizedSessionCandidate(codexTurnMetaJSON.Get("parent_thread_id").String()) + } + + forkedFrom := "" + if codexTurnMetaJSON.Exists() { + forkedFrom = normalizedSessionCandidate(codexTurnMetaJSON.Get("forked_from_thread_id").String()) + if forkedFrom == "" { + forkedFrom = normalizedSessionCandidate(codexTurnMetaJSON.Get("forked_from_id").String()) + } + } + if forkedFrom == "" && root.Exists() { + for _, forkPath := range []string{ + "forked_from_thread_id", "forked_from_id", + "metadata.forked_from_thread_id", "metadata.forked_from_id", + "extra_body.forked_from_thread_id", "extra_body.forked_from_id", + } { + if forkedFrom = normalizedSessionCandidate(root.Get(forkPath).String()); forkedFrom != "" { + break + } + if hasNestedReq { + if forkedFrom = normalizedSessionCandidate(reqRoot.Get(forkPath).String()); forkedFrom != "" { + break + } + } + } + } + + cleanAgentName := "" + if codexTurnMetaJSON.Exists() { + rawName := codexTurnMetaJSON.Get("agent_name").String() + rawName = strings.TrimPrefix(rawName, "/root/") + rawName = strings.TrimPrefix(rawName, "/") + rawName = strings.TrimSpace(rawName) + rawName = normalizedSessionCandidate(rawName) + if rawName != "" && rawName != "root" && rawName != "main" { + cleanAgentName = rawName + } + } + + subVal := sessionHeaderValue(headers, "X-Openai-Subagent") + subagentSignal := subVal != "" && !strings.EqualFold(subVal, "false") && subVal != "0" + if codexTurnMetaJSON.Exists() && codexTurnMetaJSON.Get("subagent_kind").String() == "thread_spawn" { + subagentSignal = true + } + + // 1. Fork detection + if forkedFrom != "" { + forkSessionID := tid + if forkSessionID == "" { + forkSessionID = sid + } + if forkSessionID == forkedFrom && sid != "" && sid != forkedFrom { + forkSessionID = sid + } + info.SessionID = "codex:" + forkSessionID + info.ParentSessionID = "codex:" + forkedFrom + info.AgentName = "main" + info.IsFork = true + info.IsSubagent = false + return finalizeSessionInfo(info) + } + + // 2. Subagent detection (Multi-Agent v2) + if subagentSignal || (tid != "" && sid != "" && tid != sid) || (parentThread != "" && parentThread != tid && parentThread != sid) { + childSessionID := tid + if childSessionID == "" { + childSessionID = sid + } + parentSID := parentThread + if parentSID == "" { + parentSID = sid + } + if cleanAgentName != "" && sid != "" { + info.SessionID = "codex:" + sid + ":agent:" + cleanAgentName + info.AgentName = cleanAgentName + if parentSID != "" { + info.ParentSessionID = "codex:" + parentSID + } else if parentCandidate != "" && parentCandidate != sid { + info.ParentSessionID = "codex:" + parentCandidate + } + } else { + info.SessionID = "codex:" + childSessionID + if cleanAgentName != "" { + info.AgentName = cleanAgentName + } else { + info.AgentName = "subagent" + } + if parentSID != "" && parentSID != childSessionID { + info.ParentSessionID = "codex:" + parentSID + } else if parentCandidate != "" && parentCandidate != childSessionID { + info.ParentSessionID = "codex:" + parentCandidate + } + } + info.IsSubagent = true + return finalizeSessionInfo(info) + } + + // 3. Normal interactive session + sessionID := sid + if sessionID == "" { + sessionID = tid + } + info.SessionID = "codex:" + sessionID + if parentThread != "" && parentThread != sessionID { info.ParentSessionID = "codex:" + parentThread info.AgentName = "subagent" - } else if parentCandidate != "" && parentCandidate != sid { + info.IsSubagent = true + } else if parentCandidate != "" && parentCandidate != sessionID { info.ParentSessionID = "codex:" + parentCandidate info.AgentName = "subagent" + info.IsSubagent = true } else { info.AgentName = "main" } @@ -329,17 +460,6 @@ func ExtractSessionInfo(headers http.Header, payload []byte, metadata map[string } return finalizeSessionInfo(info) } - if sid := sessionHeaderValue(headers, "Thread-Id"); sid != "" { - info.ClientType = "openai-thread" - info.SessionID = "thread:" + sid - if parentCandidate != "" && parentCandidate != sid { - info.ParentSessionID = "thread:" + parentCandidate - info.AgentName = "subagent" - } else { - info.AgentName = "main" - } - return finalizeSessionInfo(info) - } if sid := sessionHeaderValue(headers, "X-Client-Request-Id"); sid != "" { info.ClientType = "generic" info.SessionID = "clientreq:" + sid @@ -379,7 +499,14 @@ func ExtractSessionInfo(headers http.Header, payload []byte, metadata map[string info.SessionID = "thread:" + tid if parentCandidate != "" && parentCandidate != tid { info.ParentSessionID = "thread:" + parentCandidate - info.AgentName = "subagent" + if isBodyForkCandidate(root, reqRoot, hasNestedReq) { + info.IsFork = true + info.IsSubagent = false + info.AgentName = "main" + } else { + info.AgentName = "subagent" + info.IsSubagent = true + } } else { info.AgentName = "main" } @@ -423,7 +550,14 @@ func ExtractSessionInfo(headers http.Header, payload []byte, metadata map[string info.SessionID = "session:" + sid if parentCandidate != "" && parentCandidate != sid { info.ParentSessionID = "session:" + parentCandidate - info.AgentName = "subagent" + if isBodyForkCandidate(root, reqRoot, hasNestedReq) { + info.IsFork = true + info.IsSubagent = false + info.AgentName = "main" + } else { + info.AgentName = "subagent" + info.IsSubagent = true + } } else { info.AgentName = "main" } @@ -523,6 +657,27 @@ func ExtractSessionInfo(headers http.Header, payload []byte, metadata map[string return SessionInfo{}, false } +func isBodyForkCandidate(root, reqRoot gjson.Result, hasNestedReq bool) bool { + if !root.Exists() { + return false + } + for _, k := range []string{ + "forked_from_thread_id", "forked_from_id", + "metadata.forked_from_thread_id", "metadata.forked_from_id", + "extra_body.forked_from_thread_id", "extra_body.forked_from_id", + } { + if val := normalizedSessionCandidate(root.Get(k).String()); val != "" { + return true + } + if hasNestedReq { + if val := normalizedSessionCandidate(reqRoot.Get(k).String()); val != "" { + return true + } + } + } + return false +} + func finalizeSessionInfo(info SessionInfo) (SessionInfo, bool) { if info.SessionID == "" { return SessionInfo{}, false diff --git a/sdk/cliproxy/session/info_test.go b/sdk/cliproxy/session/info_test.go index 00294a38..35b98fa2 100644 --- a/sdk/cliproxy/session/info_test.go +++ b/sdk/cliproxy/session/info_test.go @@ -45,6 +45,44 @@ func TestExtractSessionInfoAllClients(t *testing.T) { t.Errorf("Codex session ids mismatch: %+v", info) } + // 2b. Codex CLI fork with X-Codex-Turn-Metadata + codexForkHeaders := http.Header{ + "Session-Id": []string{"codex-fork-666"}, + "X-Codex-Turn-Metadata": []string{ + `{"session_id":"codex-fork-666","forked_from_thread_id":"codex-parent-111","request_kind":"turn"}`, + }, + } + info, ok = ExtractSessionInfo(codexForkHeaders, nil, nil) + if !ok || info.ClientType != "codex" { + t.Fatalf("ExtractSessionInfo failed for Codex fork") + } + if info.SessionID != "codex:codex-fork-666" || info.ParentSessionID != "codex:codex-parent-111" { + t.Errorf("Codex fork session ids mismatch: %+v", info) + } + if !info.IsFork { + t.Errorf("expected IsFork=true for Codex fork: %+v", info) + } + + // 2c. Codex CLI Multi-Agent v2 (collab_spawn) + codexSubagentHeaders := http.Header{ + "Session-Id": []string{"codex-root-001"}, + "Thread-Id": []string{"codex-sub-thread-002"}, + "X-Openai-Subagent": []string{"collab_spawn"}, + "X-Codex-Turn-Metadata": []string{ + `{"session_id":"codex-root-001","thread_id":"codex-sub-thread-002","agent_name":"/root/check_readme","parent_thread_id":"codex-root-001","subagent_kind":"thread_spawn"}`, + }, + } + info, ok = ExtractSessionInfo(codexSubagentHeaders, nil, nil) + if !ok || info.ClientType != "codex" { + t.Fatalf("ExtractSessionInfo failed for Codex Multi-Agent v2") + } + if info.SessionID != "codex:codex-root-001:agent:check_readme" || info.ParentSessionID != "codex:codex-root-001" { + t.Errorf("Codex Multi-Agent session ids mismatch: %+v", info) + } + if info.AgentName != "check_readme" || !info.IsSubagent { + t.Errorf("Codex Multi-Agent metadata mismatch: AgentName=%q, IsSubagent=%v", info.AgentName, info.IsSubagent) + } + // 3. Pi Slot Session piHeaders := http.Header{ "X-Slot-Session-Id": []string{"pi-slot-777"}, @@ -64,6 +102,41 @@ func TestExtractSessionInfoAllClients(t *testing.T) { t.Errorf("OpenCode mismatch: %+v", info) } + // 4b. Body-only fork in thread_id + bodyForkPayload := []byte(`{"thread_id":"child-thread-01","forked_from_thread_id":"parent-thread-00"}`) + info, ok = ExtractSessionInfo(nil, bodyForkPayload, nil) + if !ok || info.SessionID != "thread:child-thread-01" || info.ParentSessionID != "thread:parent-thread-00" || !info.IsFork || info.IsSubagent { + t.Errorf("body-only fork mismatch: %+v", info) + } + + // 4c. Codex fork with both Session-Id and Thread-Id + codexForkBothHeaders := http.Header{ + "Session-Id": []string{"parent-thread-00"}, + "Thread-Id": []string{"child-thread-01"}, + "X-Codex-Turn-Metadata": []string{ + `{"session_id":"parent-thread-00","thread_id":"child-thread-01","forked_from_thread_id":"parent-thread-00"}`, + }, + } + info, ok = ExtractSessionInfo(codexForkBothHeaders, nil, nil) + if !ok || info.SessionID != "codex:child-thread-01" || info.ParentSessionID != "codex:parent-thread-00" || !info.IsFork { + t.Errorf("Codex fork with both sid and tid mismatch: %+v", info) + } + + // 4d. Nested metadata forked_from_thread_id in body + nestedForkPayload := []byte(`{"thread_id":"child-t-99","metadata":{"forked_from_thread_id":"parent-t-88"}}`) + info, ok = ExtractSessionInfo(nil, nestedForkPayload, nil) + if !ok || info.SessionID != "thread:child-t-99" || info.ParentSessionID != "thread:parent-t-88" || !info.IsFork { + t.Errorf("nested body-only fork mismatch: %+v", info) + } + + // 4e. Codex fork with Session-Id in header and thread_id + metadata.forked_from_thread_id in body + codexHeaderBodyForkHeaders := http.Header{"Session-Id": []string{"parent-sess-uuid"}} + codexHeaderBodyForkPayload := []byte(`{"thread_id":"child-thread-uuid","metadata":{"forked_from_thread_id":"parent-sess-uuid"}}`) + info, ok = ExtractSessionInfo(codexHeaderBodyForkHeaders, codexHeaderBodyForkPayload, nil) + if !ok || info.SessionID != "codex:child-thread-uuid" || info.ParentSessionID != "codex:parent-sess-uuid" || !info.IsFork { + t.Errorf("Codex fork with Session-Id header and body thread_id mismatch: %+v", info) + } + // 5. Antigravity X-Http-Session-Id agyHeaders := http.Header{ "X-Http-Session-Id": []string{"agy-sess-888"}, diff --git a/sdk/cliproxy/session/lcp.go b/sdk/cliproxy/session/lcp.go index bb77afa1..e895cdf2 100644 --- a/sdk/cliproxy/session/lcp.go +++ b/sdk/cliproxy/session/lcp.go @@ -627,9 +627,12 @@ func formatEqual(left, right sdktranslator.Format) bool { // MerklePrefixMatch describes the best known affinity match for a request. type MerklePrefixMatch struct { - AuthID string - SessionID string - PrefixLength int + AuthID string + SessionID string + ParentSessionID string + PrefixLength int + IsFork bool + AccessNumber uint64 } // MerklePrefixMatcherConfig controls the bounded in-memory LCP index. @@ -639,6 +642,8 @@ type MerklePrefixMatcherConfig struct { MaxGroups int // MaxPrefixes bounds group-to-prefix index entries, not just groups. MaxPrefixes int + // NowFunc provides a mockable clock for deterministic TTL/expiration tests. + NowFunc func() time.Time } // MerklePrefixMatcher stores rolling Merkle prefixes and their selected auth bindings. @@ -651,6 +656,7 @@ type MerklePrefixMatcher struct { maxTurns int maxGroups int maxPrefixes int + nowFunc func() time.Time groups map[string]*lcpNamespace lru *list.List lruElements map[*lcpGroup]*list.Element @@ -670,6 +676,7 @@ type lcpGroup struct { namespace string authID string sessionID string + parentSessionID string minPrefixLength int fingerprints []string prefixKeys []string @@ -699,17 +706,29 @@ func NewMerklePrefixMatcherWithConfig(cfg MerklePrefixMatcherConfig) *MerklePref if cfg.MaxPrefixes < cfg.MaxTurns { cfg.MaxPrefixes = cfg.MaxTurns } + nowFunc := cfg.NowFunc + if nowFunc == nil { + nowFunc = time.Now + } return &MerklePrefixMatcher{ ttl: cfg.TTL, maxTurns: cfg.MaxTurns, maxGroups: cfg.MaxGroups, maxPrefixes: cfg.MaxPrefixes, + nowFunc: nowFunc, groups: make(map[string]*lcpNamespace), lru: list.New(), lruElements: make(map[*lcpGroup]*list.Element), } } +func (m *MerklePrefixMatcher) now() time.Time { + if m != nil && m.nowFunc != nil { + return m.nowFunc() + } + return time.Now() +} + // Prepare returns bounded turn fingerprints and the first eligible prefix boundary. // The returned fingerprints can be retained in request-scoped metadata. func (m *MerklePrefixMatcher) Prepare(turns []CanonicalTurn) ([]string, int) { @@ -754,7 +773,7 @@ func (m *MerklePrefixMatcher) MatchFingerprints(namespace string, fingerprints [ m.mu.Lock() defer m.mu.Unlock() m.prepareLocked() - match, matchOK := m.matchLocked(namespace, fingerprints, minPrefixLength, time.Now()) + match, matchOK := m.matchLocked(namespace, fingerprints, minPrefixLength, m.now()) if !matchOK { return MerklePrefixMatch{}, false } @@ -767,19 +786,38 @@ func (m *MerklePrefixMatcher) Bind(namespace string, turns []CanonicalTurn, auth return m.BindFingerprints(namespace, fingerprints, minPrefixLength, authID) } +// BindWithResult records a request sequence for an auth and returns detailed session identities. +func (m *MerklePrefixMatcher) BindWithResult(namespace string, turns []CanonicalTurn, authID string) MerklePrefixBindResult { + fingerprints, minPrefixLength := m.Prepare(turns) + return m.BindFingerprintsWithResult(namespace, fingerprints, minPrefixLength, authID) +} + +// MerklePrefixBindResult describes the session identities produced by binding an LCP sequence. +type MerklePrefixBindResult struct { + SessionID string + ParentSessionID string + IsFork bool + AccessNumber uint64 +} + // BindFingerprints records a precomputed request sequence for an auth. func (m *MerklePrefixMatcher) BindFingerprints(namespace string, fingerprints []string, minPrefixLength int, authID string) string { + return m.BindFingerprintsWithResult(namespace, fingerprints, minPrefixLength, authID).SessionID +} + +// BindFingerprintsWithResult records a precomputed request sequence for an auth and returns detailed session identities. +func (m *MerklePrefixMatcher) BindFingerprintsWithResult(namespace string, fingerprints []string, minPrefixLength int, authID string) MerklePrefixBindResult { if m == nil || strings.TrimSpace(namespace) == "" || strings.TrimSpace(authID) == "" { - return "" + return MerklePrefixBindResult{} } var ok bool if fingerprints, minPrefixLength, ok = m.sanitizeFingerprints(fingerprints, minPrefixLength); !ok { - return "" + return MerklePrefixBindResult{} } m.mu.Lock() defer m.mu.Unlock() m.prepareLocked() - return m.bindLocked(namespace, fingerprints, minPrefixLength, strings.TrimSpace(authID), time.Now()) + return m.bindLocked(namespace, fingerprints, minPrefixLength, strings.TrimSpace(authID), m.now()) } // Touch refreshes an existing sequence or binds it to authID when it is a new extension. @@ -800,7 +838,7 @@ func (m *MerklePrefixMatcher) TouchFingerprints(namespace string, fingerprints [ m.mu.Lock() defer m.mu.Unlock() m.prepareLocked() - return m.touchLocked(namespace, fingerprints, minPrefixLength, strings.TrimSpace(authID), time.Now()) + return m.touchLocked(namespace, fingerprints, minPrefixLength, strings.TrimSpace(authID), m.now()) } // Remove removes the exact request sequence when it is still bound to authID. @@ -811,6 +849,12 @@ func (m *MerklePrefixMatcher) Remove(namespace string, turns []CanonicalTurn, au // RemoveFingerprints removes an exact precomputed request sequence. func (m *MerklePrefixMatcher) RemoveFingerprints(namespace string, fingerprints []string, authID string) bool { + return m.RemoveFingerprintsBefore(namespace, fingerprints, authID, 0) +} + +// RemoveFingerprintsBefore removes an exact precomputed request sequence only if it has not +// been refreshed after maxGeneration. If maxGeneration is 0, it removes the sequence unconditionally. +func (m *MerklePrefixMatcher) RemoveFingerprintsBefore(namespace string, fingerprints []string, authID string, maxGeneration uint64) bool { if m == nil || namespace == "" || authID == "" || len(fingerprints) == 0 { return false } @@ -828,10 +872,18 @@ func (m *MerklePrefixMatcher) RemoveFingerprints(namespace string, fingerprints if ns == nil { return false } - group := ns.groups[sequenceKey(fingerprints)] + prefixKeys := rollingPrefixKeys(fingerprints) + if len(prefixKeys) == 0 { + return false + } + group := ns.groups[prefixKeys[len(prefixKeys)-1]] if group == nil || group.authID != authID { return false } + if maxGeneration > 0 && group.lastAccessNumber > maxGeneration { + // Entry was refreshed/touched by a newer concurrent request; preserve the active binding. + return false + } m.removeGroupLocked(group) return true } @@ -863,7 +915,9 @@ func (m *MerklePrefixMatcher) Clear() { m.lruElements = make(map[*lcpGroup]*list.Element) m.groupCount = 0 m.prefixCount = 0 - m.accessCounter = 0 + // Do NOT reset accessCounter to 0. Keeping accessCounter monotonically increasing + // across Clear() ensures that in-flight requests with pre-clear generations cannot + // accidentally evict newly created post-clear bindings. m.mu.Unlock() } @@ -894,7 +948,7 @@ func (m *MerklePrefixMatcher) prepareLocked() { } m.operations++ if m.operations%128 == 0 { - m.cleanupLocked(time.Now()) + m.cleanupLocked(m.now()) } } @@ -936,30 +990,44 @@ func (m *MerklePrefixMatcher) touchLocked(namespace string, fingerprints []strin return true } -func (m *MerklePrefixMatcher) bindLocked(namespace string, fingerprints []string, minPrefixLength int, authID string, now time.Time) string { +func (m *MerklePrefixMatcher) bindLocked(namespace string, fingerprints []string, minPrefixLength int, authID string, now time.Time) MerklePrefixBindResult { ns := m.namespaceLocked(namespace) key := sequenceKey(fingerprints) if existing := ns.groups[key]; existing != nil { if now.Before(existing.expiresAt) { sessionID := existing.sessionID + parentSessionID := existing.parentSessionID + isFork := parentSessionID != "" m.removeGroupLocked(existing) - m.addGroupLocked(&lcpGroup{ + reboundGroup := &lcpGroup{ key: key, namespace: namespace, authID: authID, sessionID: sessionID, + parentSessionID: parentSessionID, minPrefixLength: minPrefixLength, fingerprints: append([]string(nil), fingerprints...), + prefixKeys: existing.prefixKeys, expiresAt: now.Add(m.ttl), - }) - return sessionID + } + m.addGroupLocked(reboundGroup) + return MerklePrefixBindResult{ + SessionID: sessionID, + ParentSessionID: parentSessionID, + IsFork: isFork, + AccessNumber: reboundGroup.lastAccessNumber, + } } m.removeGroupLocked(existing) } sessionID := "" + parentSessionID := "" + isFork := false if match, ok := m.matchLocked(namespace, fingerprints, minPrefixLength, now); ok { sessionID = match.SessionID + parentSessionID = match.ParentSessionID + isFork = match.IsFork } prefixKeys := rollingPrefixKeys(fingerprints) if sessionID == "" { @@ -973,17 +1041,24 @@ func (m *MerklePrefixMatcher) bindLocked(namespace string, fingerprints []string } sessionID = newLCPSessionID(namespace, firstKey) } - m.addGroupLocked(&lcpGroup{ + createdGroup := &lcpGroup{ key: key, namespace: namespace, authID: authID, sessionID: sessionID, + parentSessionID: parentSessionID, minPrefixLength: minPrefixLength, fingerprints: append([]string(nil), fingerprints...), prefixKeys: prefixKeys, expiresAt: now.Add(m.ttl), - }) - return sessionID + } + m.addGroupLocked(createdGroup) + return MerklePrefixBindResult{ + SessionID: sessionID, + ParentSessionID: parentSessionID, + IsFork: isFork, + AccessNumber: createdGroup.lastAccessNumber, + } } func (m *MerklePrefixMatcher) addGroupLocked(group *lcpGroup) { @@ -1084,7 +1159,28 @@ func (m *MerklePrefixMatcher) matchLocked(namespace string, fingerprints []strin if element := m.lruElements[best]; element != nil { m.lru.MoveToBack(element) } - return MerklePrefixMatch{AuthID: best.authID, SessionID: best.sessionID, PrefixLength: bestLength}, true + + sessionID := best.sessionID + parentSessionID := best.parentSessionID + isFork := false + + // Divergence check: + // A request represents a true fork if the longest matched common prefix is strictly + // shorter than the matched group's trajectory, and the request extends past that prefix. + if bestLength < len(best.fingerprints) && len(fingerprints) > bestLength { + isFork = true + parentSessionID = newLCPSessionID(namespace, prefixKeys[bestLength-1]) + sessionID = newLCPSessionID(namespace, prefixKeys[bestLength]) + } + + return MerklePrefixMatch{ + AuthID: best.authID, + SessionID: sessionID, + ParentSessionID: parentSessionID, + PrefixLength: bestLength, + IsFork: isFork, + AccessNumber: best.lastAccessNumber, + }, true } func newestMatchingGroup(bucket map[string]*lcpGroup, fingerprints []string, now time.Time) *lcpGroup { @@ -1093,7 +1189,12 @@ func newestMatchingGroup(bucket map[string]*lcpGroup, fingerprints []string, now if group == nil || !now.Before(group.expiresAt) || group.minPrefixLength > len(fingerprints) || len(group.fingerprints) < len(fingerprints) || !equalStrings(group.fingerprints[:len(fingerprints)], fingerprints) { continue } - if best == nil || group.lastAccessNumber > best.lastAccessNumber || (group.lastAccessNumber == best.lastAccessNumber && group.expiresAt.After(best.expiresAt)) { + // Prefer the longest known trajectory so an exact prefix match on an earlier turn + // does not mask a deeper divergent fork. Break ties by recency of access. + if best == nil || len(group.fingerprints) > len(best.fingerprints) || + (len(group.fingerprints) == len(best.fingerprints) && + (group.lastAccessNumber > best.lastAccessNumber || + (group.lastAccessNumber == best.lastAccessNumber && group.expiresAt.After(best.expiresAt)))) { best = group } } diff --git a/sdk/cliproxy/session/lcp_test.go b/sdk/cliproxy/session/lcp_test.go index af6c2a61..80f5b1c6 100644 --- a/sdk/cliproxy/session/lcp_test.go +++ b/sdk/cliproxy/session/lcp_test.go @@ -372,14 +372,22 @@ func TestMerklePrefixMatcherPrefixEntryBound(t *testing.T) { func TestMerklePrefixMatcherTTLAndAuthInvalidation(t *testing.T) { t.Parallel() - matcher := NewMerklePrefixMatcher(5 * time.Millisecond) + current := time.Now() + matcher := NewMerklePrefixMatcherWithConfig(MerklePrefixMatcherConfig{ + TTL: 5 * time.Minute, + NowFunc: func() time.Time { + return current + }, + }) namespace := "lcp:v1:ttl:model:caller" turns := turnsFromTexts("one") matcher.Bind(namespace, turns, "auth-a") if match, ok := matcher.Match(namespace, turns); !ok || match.AuthID != "auth-a" { t.Fatalf("initial match = %#v, %v", match, ok) } - time.Sleep(10 * time.Millisecond) + + // Advance mock clock past TTL without wall-clock sleep + current = current.Add(10 * time.Minute) if _, ok := matcher.Match(namespace, turns); ok { t.Fatal("expired matcher entry remained available") } @@ -571,7 +579,15 @@ func TestMerklePrefixMatcherTouchFingerprintsDelayedSuccessProtection(t *testing } func TestMerklePrefixMatcherTouchExpiredEntry(t *testing.T) { - matcher := NewMerklePrefixMatcher(10 * time.Millisecond) + t.Parallel() + + current := time.Now() + matcher := NewMerklePrefixMatcherWithConfig(MerklePrefixMatcherConfig{ + TTL: 10 * time.Minute, + NowFunc: func() time.Time { + return current + }, + }) namespace := "lcp:v1:test:model:caller" turns := turnsFromTexts("alpha", "beta") fingerprints, minPrefixLength := matcher.Prepare(turns) @@ -580,7 +596,8 @@ func TestMerklePrefixMatcherTouchExpiredEntry(t *testing.T) { } matcher.BindFingerprints(namespace, fingerprints, minPrefixLength, "auth-A") - time.Sleep(25 * time.Millisecond) + // Advance mock clock past TTL without wall-clock sleep + current = current.Add(25 * time.Minute) // After TTL expires, Match should not find the expired binding if _, ok := matcher.MatchFingerprints(namespace, fingerprints, minPrefixLength); ok { @@ -678,6 +695,250 @@ func BenchmarkMerklePrefixMatcherMatch(b *testing.B) { } } +func TestMerklePrefixMatcherForkLineageTree(t *testing.T) { + t.Parallel() + + matcher := NewMerklePrefixMatcher(time.Hour) + defer matcher.Clear() + namespace := "lcp:v1:test-fork:model:caller" + + // 4 test sequences: + // Seq 1: 1 2 3 4 5 6 7 A + // Seq 2: 1 2 3 4 5 6 7 8 + // Seq 3: 1 2 3 C 12 + // Seq 4: 1 2 3 C D + seq1 := turnsFromTexts("1", "2", "3", "4", "5", "6", "7", "A") + seq2 := turnsFromTexts("1", "2", "3", "4", "5", "6", "7", "8") + seq3 := turnsFromTexts("1", "2", "3", "C", "12") + seq4 := turnsFromTexts("1", "2", "3", "C", "D") + + // 1. Seq 1: Initial root establishment + res1 := matcher.BindWithResult(namespace, seq1, "auth-root") + if res1.SessionID == "" { + t.Fatal("seq1 bind returned empty session ID") + } + if res1.IsFork { + t.Fatal("seq1 unexpectedly marked as fork") + } + if res1.ParentSessionID != "" { + t.Fatalf("seq1 parentSessionID = %q, want empty", res1.ParentSessionID) + } + + // 2. Seq 2: Forks from Seq 1 at depth 7 ("8" vs "A") + match2, ok2 := matcher.Match(namespace, seq2) + if !ok2 { + t.Fatal("seq2 match failed") + } + if !match2.IsFork { + t.Fatal("seq2 should be recognized as a true fork") + } + if match2.PrefixLength != 7 { + t.Fatalf("seq2 prefix length = %d, want 7", match2.PrefixLength) + } + if match2.AuthID != "auth-root" { + t.Fatalf("seq2 authID = %q, want auth-root", match2.AuthID) + } + if match2.SessionID == res1.SessionID { + t.Fatalf("seq2 sessionID = %q should differ from seq1 root %q", match2.SessionID, res1.SessionID) + } + if match2.ParentSessionID == "" { + t.Fatal("seq2 parentSessionID should not be empty") + } + res2 := matcher.BindWithResult(namespace, seq2, "auth-root") + if res2.SessionID != match2.SessionID || res2.ParentSessionID != match2.ParentSessionID { + t.Fatalf("seq2 bind result mismatch: bind=%+v, match=%+v", res2, match2) + } + + // 3. Seq 3: Forks from root tree at depth 3 ("C" vs "4") + match3, ok3 := matcher.Match(namespace, seq3) + if !ok3 { + t.Fatal("seq3 match failed") + } + if !match3.IsFork { + t.Fatal("seq3 should be recognized as a true fork") + } + if match3.PrefixLength != 3 { + t.Fatalf("seq3 prefix length = %d, want 3", match3.PrefixLength) + } + if match3.AuthID != "auth-root" { + t.Fatalf("seq3 authID = %q, want auth-root", match3.AuthID) + } + if match3.SessionID == res1.SessionID || match3.SessionID == res2.SessionID { + t.Fatalf("seq3 sessionID = %q collided with seq1=%q or seq2=%q", match3.SessionID, res1.SessionID, res2.SessionID) + } + if match3.ParentSessionID == "" || match3.ParentSessionID == match2.ParentSessionID { + t.Fatalf("seq3 parentSessionID = %q, should point to prefix 3 (distinct from seq2 prefix 7 parent %q)", match3.ParentSessionID, match2.ParentSessionID) + } + res3 := matcher.BindWithResult(namespace, seq3, "auth-root") + if res3.SessionID != match3.SessionID || res3.ParentSessionID != match3.ParentSessionID { + t.Fatalf("seq3 bind result mismatch: bind=%+v, match=%+v", res3, match3) + } + + // 4. Seq 4: Nested fork from Seq 3 at depth 4 ("D" vs "12") + match4, ok4 := matcher.Match(namespace, seq4) + if !ok4 { + t.Fatal("seq4 match failed") + } + if !match4.IsFork { + t.Fatal("seq4 should be recognized as a true fork") + } + if match4.PrefixLength != 4 { + t.Fatalf("seq4 prefix length = %d, want 4", match4.PrefixLength) + } + if match4.AuthID != "auth-root" { + t.Fatalf("seq4 authID = %q, want auth-root", match4.AuthID) + } + if match4.SessionID == res1.SessionID || match4.SessionID == res2.SessionID || match4.SessionID == res3.SessionID { + t.Fatalf("seq4 sessionID = %q collided with earlier sessions", match4.SessionID) + } + // Crucial: seq4's parent should be seq3's branch session ID (prefix 4: 1 2 3 C) + if match4.ParentSessionID != res3.SessionID { + t.Fatalf("seq4 parentSessionID = %q, want seq3 session ID %q", match4.ParentSessionID, res3.SessionID) + } + res4 := matcher.BindWithResult(namespace, seq4, "auth-root") + if res4.SessionID != match4.SessionID || res4.ParentSessionID != match4.ParentSessionID { + t.Fatalf("seq4 bind result mismatch: bind=%+v, match=%+v", res4, match4) + } + + // 5. Linear continuation of Seq 4 (Seq 4.2: 1 2 3 C D E) + seq4Cont := turnsFromTexts("1", "2", "3", "C", "D", "E") + match4Cont, ok4Cont := matcher.Match(namespace, seq4Cont) + if !ok4Cont { + t.Fatal("seq4Cont match failed") + } + if match4Cont.IsFork { + t.Fatal("seq4Cont is a linear continuation, should NOT be marked as fork") + } + if match4Cont.PrefixLength != 5 { + t.Fatalf("seq4Cont prefix length = %d, want 5", match4Cont.PrefixLength) + } + if match4Cont.SessionID != res4.SessionID { + t.Fatalf("seq4Cont sessionID = %q, want identical to seq4 %q across linear growth", match4Cont.SessionID, res4.SessionID) + } + if match4Cont.ParentSessionID != res4.ParentSessionID { + t.Fatalf("seq4Cont parentSessionID = %q, want %q", match4Cont.ParentSessionID, res4.ParentSessionID) + } +} + +func BenchmarkMerklePrefixMatcherFork(b *testing.B) { + matcher := NewMerklePrefixMatcher(time.Hour) + trunk := turnsFromTexts("1", "2", "3", "4", "5", "6", "7", "A") + fork := turnsFromTexts("1", "2", "3", "4", "5", "6", "7", "8") + matcher.Bind("lcp:v1:benchmark:fork", trunk, "auth-a") + + b.ReportAllocs() + b.ResetTimer() + for b.Loop() { + _, _ = matcher.Match("lcp:v1:benchmark:fork", fork) + } +} + +func TestMerklePrefixMatcherShorterGroupDoesNotMaskLongerTrajectory(t *testing.T) { + t.Parallel() + + matcher := NewMerklePrefixMatcher(time.Hour) + defer matcher.Clear() + namespace := "lcp:v1:mask-test:model:caller" + + // 1. Bind long trajectory [1, 2, 3] + longTrunk := turnsFromTexts("1", "2", "3") + matcher.Bind(namespace, longTrunk, "auth-root") + + // 2. Bind shorter exact-prefix group [1, 2] afterwards (more recently accessed) + shortPrefix := turnsFromTexts("1", "2") + matcher.Bind(namespace, shortPrefix, "auth-root") + + // 3. Query a fork [1, 2, 4] diverging at turn 3 from the long trunk + fork := turnsFromTexts("1", "2", "4") + match, ok := matcher.Match(namespace, fork) + if !ok { + t.Fatal("fork match failed") + } + if !match.IsFork { + t.Fatal("fork must NOT be masked by shorter exact-prefix group [1, 2]") + } + if match.PrefixLength != 2 { + t.Fatalf("prefix length = %d, want 2", match.PrefixLength) + } +} + +func TestMerklePrefixMatcherRemoveFingerprintsBeforeGenerationGuard(t *testing.T) { + t.Parallel() + + matcher := NewMerklePrefixMatcher(time.Hour) + defer matcher.Clear() + namespace := "lcp:v1:gen-test:model:caller" + + turns := turnsFromTexts("1", "2") + fps, minPrefix := matcher.Prepare(turns) + + // 1. Initial bind at generation G1 + res1 := matcher.BindFingerprintsWithResult(namespace, fps, minPrefix, "auth-1") + g1 := res1.AccessNumber + if g1 == 0 { + t.Fatal("expected non-zero access generation") + } + + // 2. Concurrent success touches entry and advances to generation G2 > G1 + if !matcher.TouchFingerprints(namespace, fps, minPrefix, "auth-1") { + t.Fatal("TouchFingerprints failed") + } + + // 3. Stale failure from request 1 attempts to remove with generation G1 + removedStale := matcher.RemoveFingerprintsBefore(namespace, fps, "auth-1", g1) + if removedStale { + t.Fatal("stale RemoveFingerprintsBefore should not delete entry refreshed by newer touch") + } + + // Entry must remain active + if _, ok := matcher.MatchFingerprints(namespace, fps, minPrefix); !ok { + t.Fatal("entry should still be present after stale removal attempt") + } + + // 4. Current failure with generation 0 (unconditional) removes it + removedCurrent := matcher.RemoveFingerprints(namespace, fps, "auth-1") + if !removedCurrent { + t.Fatal("unconditional RemoveFingerprints should succeed") + } + if _, ok := matcher.MatchFingerprints(namespace, fps, minPrefix); ok { + t.Fatal("entry should be removed after unconditional removal") + } +} + +func TestMerklePrefixMatcherClearPreservesMonotonicGeneration(t *testing.T) { + t.Parallel() + matcher := NewMerklePrefixMatcher(time.Hour) + namespace := "lcp:v1:test:model:caller" + fps := []string{"fp-1", "fp-2"} + minPrefix := 1 + + // Bind initial sequence and record generation + res1 := matcher.BindFingerprintsWithResult(namespace, fps, minPrefix, "auth-old") + if res1.AccessNumber == 0 { + t.Fatal("expected non-zero generation for initial bind") + } + + // Clear matcher + matcher.Clear() + + // Rebind same sequence under new auth post-clear + res2 := matcher.BindFingerprintsWithResult(namespace, fps, minPrefix, "auth-new") + if res2.AccessNumber <= res1.AccessNumber { + t.Fatalf("expected generation to increase monotonically across Clear(), got %d <= %d", res2.AccessNumber, res1.AccessNumber) + } + + // Pre-clear generation should NOT be able to evict post-clear binding + if matcher.RemoveFingerprintsBefore(namespace, fps, "auth-new", res1.AccessNumber) { + t.Fatal("stale pre-clear generation should not evict post-clear binding") + } + + // The binding must remain intact + match, ok := matcher.MatchFingerprints(namespace, fps, minPrefix) + if !ok || match.AuthID != "auth-new" { + t.Fatalf("binding should remain intact, got ok=%v match=%+v", ok, match) + } +} + func BenchmarkMerklePrefixMatcherMatch_100Turns(b *testing.B) { matcher := NewMerklePrefixMatcher(time.Hour) texts := make([]string, 100) -- 2.51.2