diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go index 328be59c7..80cc89b48 100644 --- a/sdk/cliproxy/auth/selector.go +++ b/sdk/cliproxy/auth/selector.go @@ -1126,6 +1126,12 @@ func extractExplicitSessionIDs(headers http.Header, payload []byte, metadata map var parentIDCandidate string if len(payload) > 0 { root = util.ParseGJSONBytesNoCopy(payload) + reqRoot := root + req := root.Get("request") + hasNestedReq := req.Exists() && !root.Get("contents").Exists() + if hasNestedReq { + reqRoot = req + } for _, parentPath := range []string{ "parent_session_id", "parentSessionId", "parent_thread_id", "parentThreadId", @@ -1138,6 +1144,12 @@ func extractExplicitSessionIDs(headers http.Header, payload []byte, metadata map parentIDCandidate = psid break } + if hasNestedReq { + if psid := normalizedSessionCandidate(reqRoot.Get(parentPath).String()); psid != "" { + parentIDCandidate = psid + break + } + } } if parentIDCandidate == "" { parentIDCandidate = cliproxysession.ClaudeMetadataParentSessionID(payload) @@ -1306,9 +1318,20 @@ func extractExplicitSessionIDs(headers http.Header, payload []byte, metadata map // 5. Body payload inspection if len(payload) > 0 && root.Exists() { + reqRoot := root + req := root.Get("request") + hasNestedReq := req.Exists() && !root.Get("contents").Exists() + if hasNestedReq { + reqRoot = req + } + // Google Gemini Context Caching for _, cachePath := range []string{"cachedContent", "cached_content"} { - if cacheID := normalizedSessionCandidate(root.Get(cachePath).String()); cacheID != "" { + cacheID := normalizedSessionCandidate(root.Get(cachePath).String()) + if cacheID == "" && hasNestedReq { + cacheID = normalizedSessionCandidate(reqRoot.Get(cachePath).String()) + } + if cacheID != "" { if parentIDCandidate != "" && parentIDCandidate != cacheID { return "geminicache:" + cacheID, "geminicache:" + parentIDCandidate } @@ -1318,7 +1341,11 @@ func extractExplicitSessionIDs(headers http.Header, payload []byte, metadata map // OpenAI Assistants / Threads for _, threadPath := range []string{"thread_id", "threadId", "metadata.thread_id"} { - if tid := normalizedSessionCandidate(root.Get(threadPath).String()); tid != "" { + tid := normalizedSessionCandidate(root.Get(threadPath).String()) + if tid == "" && hasNestedReq { + tid = normalizedSessionCandidate(reqRoot.Get(threadPath).String()) + } + if tid != "" { if parentIDCandidate != "" && parentIDCandidate != tid { return "thread:" + tid, "thread:" + parentIDCandidate } @@ -1332,7 +1359,11 @@ func extractExplicitSessionIDs(headers http.Header, payload []byte, metadata map agentID = normalizedSessionCandidate(root.Get("metadata.subagent_id").String()) } for _, path := range []string{"session_id", "sessionId", "sessionID", "metadata.session_id", "extra_body.session_id"} { - if sid := normalizedSessionCandidate(root.Get(path).String()); sid != "" { + sid := normalizedSessionCandidate(root.Get(path).String()) + if sid == "" && hasNestedReq { + sid = normalizedSessionCandidate(reqRoot.Get(path).String()) + } + if sid != "" { if agentID != "" && agentID != "main" { primary = "session:" + sid + ":agent:" + agentID fallback = "session:" + sid @@ -1350,6 +1381,9 @@ func extractExplicitSessionIDs(headers http.Header, payload []byte, metadata map conversationID := "" conversation := root.Get("conversation") + if !conversation.Exists() && hasNestedReq { + conversation = reqRoot.Get("conversation") + } if sid := normalizedSessionCandidate(conversation.Get("id").String()); sid != "" { conversationID = "conv:" + sid } else if conversation.Type == gjson.String { @@ -1357,7 +1391,11 @@ func extractExplicitSessionIDs(headers http.Header, payload []byte, metadata map conversationID = "conv:" + sid } } - if sid := normalizedSessionCandidate(root.Get("prompt_cache_key").String()); sid != "" { + pck := root.Get("prompt_cache_key") + if !pck.Exists() { + pck = root.Get("promptCacheKey") + } + if sid := normalizedSessionCandidate(pck.String()); sid != "" { return "pck:" + sid, conversationID } if conversationID != "" { @@ -1370,7 +1408,11 @@ func extractExplicitSessionIDs(headers http.Header, payload []byte, metadata map return "user:" + userID, "" } for _, convPath := range []string{"conversation_id", "conversationId", "chat_id", "chatId", "metadata.conversation_id", "extra_body.conversation_id"} { - if cid := normalizedSessionCandidate(root.Get(convPath).String()); cid != "" { + cid := normalizedSessionCandidate(root.Get(convPath).String()) + if cid == "" && hasNestedReq { + cid = normalizedSessionCandidate(reqRoot.Get(convPath).String()) + } + if cid != "" { if parentIDCandidate != "" && ("conv:"+parentIDCandidate) != ("conv:"+cid) { return "conv:" + cid, "conv:" + parentIDCandidate } diff --git a/sdk/cliproxy/auth/selector_lcp_test.go b/sdk/cliproxy/auth/selector_lcp_test.go index 34b8633c5..82f9316db 100644 --- a/sdk/cliproxy/auth/selector_lcp_test.go +++ b/sdk/cliproxy/auth/selector_lcp_test.go @@ -761,6 +761,38 @@ func TestSessionAffinitySelectorNilFallbackNoPanic(t *testing.T) { } } +func TestSessionAffinitySelectorPromptCacheKeyCamelCase(t *testing.T) { + t.Parallel() + + payload := []byte(`{"promptCacheKey":"camel-pck-123","input":"hello"}`) + primary, fallback := extractExplicitSessionIDs(nil, payload, nil) + if primary != "pck:camel-pck-123" { + t.Fatalf("primary = %q, want pck:camel-pck-123", primary) + } + if fallback != "" { + t.Fatalf("fallback = %q, want empty", fallback) + } +} + +func TestSessionAffinitySelectorNestedAntigravityPayload(t *testing.T) { + t.Parallel() + + payload := []byte(`{ + "project_id": "proj-123", + "request": { + "parentSessionId": "parent-456", + "sessionId": "child-789" + } + }`) + primary, fallback := extractExplicitSessionIDs(nil, payload, nil) + if primary != "session:child-789" { + t.Fatalf("primary = %q, want session:child-789", primary) + } + if fallback != "session:parent-456" { + t.Fatalf("fallback = %q, want session:parent-456", fallback) + } +} + func TestSessionCacheTinyTTLNoPanic(t *testing.T) { t.Parallel() diff --git a/sdk/cliproxy/session/lcp.go b/sdk/cliproxy/session/lcp.go index e0c056dfe..57d8c3198 100644 --- a/sdk/cliproxy/session/lcp.go +++ b/sdk/cliproxy/session/lcp.go @@ -12,6 +12,7 @@ import ( "sync" "time" + "github.com/router-for-me/CLIProxyAPI/v7/internal/util" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" "github.com/tidwall/gjson" ) @@ -83,10 +84,13 @@ func writeFingerprintField(hash interface{ Write([]byte) (int, error) }, value s // ExtractCanonicalTurns extracts all logical turns from the five supported inbound protocols. // Invalid or empty payloads return an empty slice. func ExtractCanonicalTurns(format sdktranslator.Format, payload []byte) []CanonicalTurn { - if len(payload) == 0 || !gjson.ValidBytes(payload) { + if len(payload) == 0 { + return nil + } + root := util.ParseGJSONBytesNoCopy(payload) + if !root.Exists() { return nil } - root := gjson.ParseBytes(payload) if format == "" { format = inferCanonicalFormat(root) } @@ -715,16 +719,37 @@ func (m *MerklePrefixMatcher) Match(namespace string, turns []CanonicalTurn) (Me return m.MatchFingerprints(namespace, fingerprints, minPrefixLength) } +func (m *MerklePrefixMatcher) sanitizeFingerprints(fingerprints []string, minPrefixLength int) ([]string, int, bool) { + if len(fingerprints) == 0 || minPrefixLength <= 0 || minPrefixLength > len(fingerprints) { + return nil, 0, false + } + maxTurns := m.maxTurns + if maxTurns <= 0 { + maxTurns = defaultMatcherMaxTurns + } + if len(fingerprints) > maxTurns { + fingerprints = fingerprints[:maxTurns] + if minPrefixLength > len(fingerprints) { + return nil, 0, false + } + } + return fingerprints, minPrefixLength, true +} + // MatchFingerprints returns the longest known prefix match without reparsing turns. func (m *MerklePrefixMatcher) MatchFingerprints(namespace string, fingerprints []string, minPrefixLength int) (MerklePrefixMatch, bool) { - if m == nil || namespace == "" || len(fingerprints) == 0 || minPrefixLength <= 0 || minPrefixLength > len(fingerprints) { + if m == nil || namespace == "" { + return MerklePrefixMatch{}, false + } + var ok bool + if fingerprints, minPrefixLength, ok = m.sanitizeFingerprints(fingerprints, minPrefixLength); !ok { return MerklePrefixMatch{}, false } m.mu.Lock() defer m.mu.Unlock() m.prepareLocked() - match, ok := m.matchLocked(namespace, fingerprints, minPrefixLength, time.Now()) - if !ok { + match, matchOK := m.matchLocked(namespace, fingerprints, minPrefixLength, time.Now()) + if !matchOK { return MerklePrefixMatch{}, false } return match, true @@ -738,7 +763,11 @@ func (m *MerklePrefixMatcher) Bind(namespace string, turns []CanonicalTurn, auth // BindFingerprints records a precomputed request sequence for an auth. func (m *MerklePrefixMatcher) BindFingerprints(namespace string, fingerprints []string, minPrefixLength int, authID string) string { - if m == nil || strings.TrimSpace(namespace) == "" || strings.TrimSpace(authID) == "" || len(fingerprints) == 0 || minPrefixLength <= 0 || minPrefixLength > len(fingerprints) { + if m == nil || strings.TrimSpace(namespace) == "" || strings.TrimSpace(authID) == "" { + return "" + } + var ok bool + if fingerprints, minPrefixLength, ok = m.sanitizeFingerprints(fingerprints, minPrefixLength); !ok { return "" } m.mu.Lock() @@ -755,7 +784,11 @@ func (m *MerklePrefixMatcher) Touch(namespace string, turns []CanonicalTurn, aut // TouchFingerprints refreshes or binds a precomputed request sequence. func (m *MerklePrefixMatcher) TouchFingerprints(namespace string, fingerprints []string, minPrefixLength int, authID string) bool { - if m == nil || strings.TrimSpace(namespace) == "" || strings.TrimSpace(authID) == "" || len(fingerprints) == 0 || minPrefixLength <= 0 || minPrefixLength > len(fingerprints) { + if m == nil || strings.TrimSpace(namespace) == "" || strings.TrimSpace(authID) == "" { + return false + } + var ok bool + if fingerprints, minPrefixLength, ok = m.sanitizeFingerprints(fingerprints, minPrefixLength); !ok { return false } m.mu.Lock() @@ -775,6 +808,13 @@ func (m *MerklePrefixMatcher) RemoveFingerprints(namespace string, fingerprints if m == nil || namespace == "" || authID == "" || len(fingerprints) == 0 { return false } + maxTurns := m.maxTurns + if maxTurns <= 0 { + maxTurns = defaultMatcherMaxTurns + } + if len(fingerprints) > maxTurns { + fingerprints = fingerprints[:maxTurns] + } m.mu.Lock() defer m.mu.Unlock() m.prepareLocked() diff --git a/sdk/cliproxy/session/lcp_test.go b/sdk/cliproxy/session/lcp_test.go index 7a07c7381..1262d9f8f 100644 --- a/sdk/cliproxy/session/lcp_test.go +++ b/sdk/cliproxy/session/lcp_test.go @@ -546,6 +546,40 @@ func TestExtractCanonicalTurnsAntigravityNestedRequest(t *testing.T) { } } +func TestMerklePrefixMatcherFingerprintsBounding(t *testing.T) { + t.Parallel() + + matcher := NewMerklePrefixMatcherWithConfig(MerklePrefixMatcherConfig{ + MaxTurns: 10, + TTL: time.Hour, + }) + + fps := make([]string, 50) + for i := range fps { + fps[i] = fmt.Sprintf("fp-%d", i) + } + + // Binding 50 fingerprints with MaxTurns=10 should truncate to 10 + sid := matcher.BindFingerprints("ns", fps, 2, "auth-1") + if sid == "" { + t.Fatal("BindFingerprints failed") + } + + // Match should succeed with 50 fingerprints (truncated to 10) + match, ok := matcher.MatchFingerprints("ns", fps, 2) + if !ok || match.PrefixLength != 10 { + t.Fatalf("expected 10 matched turns, got %d (ok=%v)", match.PrefixLength, ok) + } + + // Touch and Remove should also handle 50 fingerprints cleanly + if !matcher.TouchFingerprints("ns", fps, 2, "auth-1") { + t.Fatal("TouchFingerprints failed") + } + if !matcher.RemoveFingerprints("ns", fps, "auth-1") { + t.Fatal("RemoveFingerprints failed") + } +} + func BenchmarkExtractCanonicalTurns(b *testing.B) { payload := []byte(`{ "messages": [ diff --git a/sdk/cliproxy/session/tree.go b/sdk/cliproxy/session/tree.go index 9c492d8b7..3821736de 100644 --- a/sdk/cliproxy/session/tree.go +++ b/sdk/cliproxy/session/tree.go @@ -49,21 +49,30 @@ func (n *SessionTreeNode) Clone() *SessionTreeNode { return &res } +const maxMetadataCloneDepth = 16 + func cloneTreeMetadata(metadata map[string]any) map[string]any { - if metadata == nil { + return cloneTreeMetadataWithDepth(metadata, 0) +} + +func cloneTreeMetadataWithDepth(metadata map[string]any, depth int) map[string]any { + if metadata == nil || depth > maxMetadataCloneDepth { return nil } cloned := make(map[string]any, len(metadata)) for key, value := range metadata { - cloned[key] = cloneTreeMetadataValue(value) + cloned[key] = cloneTreeMetadataValueWithDepth(value, depth+1) } return cloned } -func cloneTreeMetadataValue(value any) any { +func cloneTreeMetadataValueWithDepth(value any, depth int) any { + if depth > maxMetadataCloneDepth { + return nil + } switch typed := value.(type) { case map[string]any: - return cloneTreeMetadata(typed) + return cloneTreeMetadataWithDepth(typed, depth+1) case map[string]string: cloned := make(map[string]string, len(typed)) for k, v := range typed { @@ -73,7 +82,7 @@ func cloneTreeMetadataValue(value any) any { case []any: cloned := make([]any, len(typed)) for index, item := range typed { - cloned[index] = cloneTreeMetadataValue(item) + cloned[index] = cloneTreeMetadataValueWithDepth(item, depth+1) } return cloned case []string: @@ -231,6 +240,12 @@ func ExtractTreeInfo(headers http.Header, payload []byte, metadata map[string]an // 2. Payload Inspection if headers didn't fully resolve if len(payload) > 0 { root := util.ParseGJSONBytesNoCopy(payload) + reqRoot := root + req := root.Get("request") + hasNestedReq := req.Exists() && !root.Get("contents").Exists() + if hasNestedReq { + reqRoot = req + } var parentCandidate string for _, p := range []string{ "parent_session_id", "parentSessionId", @@ -244,6 +259,12 @@ func ExtractTreeInfo(headers http.Header, payload []byte, metadata map[string]an parentCandidate = val break } + if hasNestedReq { + if val := normalizedSessionCandidate(reqRoot.Get(p).String()); val != "" { + parentCandidate = val + break + } + } } if parentCandidate == "" { parentCandidate = ClaudeMetadataParentSessionID(payload) @@ -277,7 +298,11 @@ func ExtractTreeInfo(headers http.Header, payload []byte, metadata map[string]an // Gemini context caching if info.SessionID == "" { for _, cachePath := range []string{"cachedContent", "cached_content"} { - if cacheID := normalizedSessionCandidate(root.Get(cachePath).String()); cacheID != "" { + cacheID := normalizedSessionCandidate(root.Get(cachePath).String()) + if cacheID == "" && hasNestedReq { + cacheID = normalizedSessionCandidate(reqRoot.Get(cachePath).String()) + } + if cacheID != "" { info.ClientType = "gemini" info.SessionID = "geminicache:" + cacheID if parentCandidate != "" && parentCandidate != cacheID { @@ -294,7 +319,11 @@ func ExtractTreeInfo(headers http.Header, payload []byte, metadata map[string]an // OpenAI thread in payload if info.SessionID == "" { for _, threadPath := range []string{"thread_id", "threadId", "metadata.thread_id"} { - if tid := normalizedSessionCandidate(root.Get(threadPath).String()); tid != "" { + tid := normalizedSessionCandidate(root.Get(threadPath).String()) + if tid == "" && hasNestedReq { + tid = normalizedSessionCandidate(reqRoot.Get(threadPath).String()) + } + if tid != "" { info.ClientType = "openai-thread" info.SessionID = "thread:" + tid if parentCandidate != "" && parentCandidate != tid { @@ -315,7 +344,11 @@ func ExtractTreeInfo(headers http.Header, payload []byte, metadata map[string]an agentID = normalizedSessionCandidate(root.Get("metadata.subagent_id").String()) } for _, path := range []string{"session_id", "sessionId", "sessionID", "metadata.session_id", "extra_body.session_id"} { - if sid := normalizedSessionCandidate(root.Get(path).String()); sid != "" { + sid := normalizedSessionCandidate(root.Get(path).String()) + if sid == "" && hasNestedReq { + sid = normalizedSessionCandidate(reqRoot.Get(path).String()) + } + if sid != "" { info.ClientType = "generic" if agentID != "" && agentID != "main" { info.SessionID = "session:" + sid + ":agent:" + agentID @@ -341,7 +374,11 @@ func ExtractTreeInfo(headers http.Header, payload []byte, metadata map[string]an // Conversation paths if info.SessionID == "" { for _, convPath := range []string{"conversation_id", "conversationId", "chat_id", "chatId", "metadata.conversation_id", "extra_body.conversation_id"} { - if cid := normalizedSessionCandidate(root.Get(convPath).String()); cid != "" { + cid := normalizedSessionCandidate(root.Get(convPath).String()) + if cid == "" && hasNestedReq { + cid = normalizedSessionCandidate(reqRoot.Get(convPath).String()) + } + if cid != "" { info.ClientType = "conv" info.SessionID = "conv:" + cid if parentCandidate != "" && parentCandidate != cid { @@ -775,6 +812,7 @@ func (s *InMemorySessionTreeStore) computeNodeLineageLocked(node *SessionTreeNod newPath = parent.TreePath + "/" + node.SessionID newDepth = parent.TreeDepth + 1 } else { + s.updateParentIndexLocked(node.SessionID, node.ParentSessionID, "") node.ParentSessionID = "" newRoot = node.SessionID newPath = node.SessionID diff --git a/sdk/cliproxy/session/tree_test.go b/sdk/cliproxy/session/tree_test.go index e630569ec..c92e5b296 100644 --- a/sdk/cliproxy/session/tree_test.go +++ b/sdk/cliproxy/session/tree_test.go @@ -637,6 +637,64 @@ func TestSessionTreeStoreMaxDepthCapping(t *testing.T) { } current = next } + + // Verify that the capped node's parent index was cleaned up + cappedParent := fmt.Sprintf("node-%d", 128) + cappedNode := fmt.Sprintf("node-%d", 129) + store.mu.Lock() + if children, ok := store.parentIndex[cappedParent]; ok { + if _, exists := children[cappedNode]; exists { + store.mu.Unlock() + t.Fatalf("capped node %s still found in parentIndex[%s]", cappedNode, cappedParent) + } + } + store.mu.Unlock() +} + +func TestSessionTreeStoreMetadataDeepRecursionProtection(t *testing.T) { + t.Parallel() + + store := NewInMemorySessionTreeStore(100, time.Hour) + + // Create deeply nested map + nested := map[string]any{"level": 0} + curr := nested + for i := 1; i <= 30; i++ { + next := map[string]any{"level": i} + curr["child"] = next + curr = next + } + + // Recording node should succeed without stack overflow + node := store.RecordNode(SessionTreeInfo{ + SessionID: "sess-deep-meta", + Metadata: nested, + }) + if node.SessionID != "sess-deep-meta" { + t.Fatalf("expected sess-deep-meta, got %s", node.SessionID) + } +} + +func TestExtractTreeInfoNestedAntigravityRequest(t *testing.T) { + t.Parallel() + + payload := []byte(`{ + "project_id": "proj-123", + "request": { + "parentSessionId": "parent-sess-456", + "sessionId": "child-sess-789" + } + }`) + info, ok := ExtractTreeInfo(nil, payload, nil) + if !ok { + t.Fatal("ExtractTreeInfo failed on nested Antigravity payload") + } + if info.SessionID != "session:child-sess-789" { + t.Fatalf("SessionID = %q, want session:child-sess-789", info.SessionID) + } + if info.ParentSessionID != "session:parent-sess-456" { + t.Fatalf("ParentSessionID = %q, want session:parent-sess-456", info.ParentSessionID) + } } func BenchmarkSessionTreeRecordAndCascade(b *testing.B) {