From 100c564313bbb12e07afe12e9828111e452ab09b Mon Sep 17 00:00:00 2001 From: sususu Date: Mon, 31 Aug 2026 14:00:10 +0800 Subject: [PATCH] fix(session): harden nil fallback, normalize agent id, and support nested antigravity payloads (#5202) --- sdk/cliproxy/auth/home_session_alias.go | 7 ++- sdk/cliproxy/auth/selector.go | 6 +++ sdk/cliproxy/auth/selector_lcp_test.go | 68 +++++++++++++++++++++++++ sdk/cliproxy/auth/session_cache.go | 31 +++++------ sdk/cliproxy/session/identity.go | 30 ++++++++--- sdk/cliproxy/session/identity_test.go | 65 ++++++++++++++++++----- sdk/cliproxy/session/lcp.go | 6 +++ sdk/cliproxy/session/lcp_test.go | 21 ++++++++ sdk/cliproxy/session/tree.go | 14 +++-- sdk/cliproxy/session/tree_test.go | 28 ++++++++++ 10 files changed, 234 insertions(+), 42 deletions(-) diff --git a/sdk/cliproxy/auth/home_session_alias.go b/sdk/cliproxy/auth/home_session_alias.go index 027bac0a..dce16b54 100644 --- a/sdk/cliproxy/auth/home_session_alias.go +++ b/sdk/cliproxy/auth/home_session_alias.go @@ -274,7 +274,12 @@ func (m *Manager) homeDispatchSessionIDs(opts cliproxyexecutor.Options) (string, } cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config) - canonical := m.homeSessionAliases.canonical(primary, aliasFallback, homeSessionAliasTTL(cfg), time.Now()) + ttl := homeSessionAliasTTL(cfg) + now := time.Now() + canonical := m.homeSessionAliases.canonical(primary, aliasFallback, ttl, now) + if parentSessionID != "" { + parentSessionID = m.homeSessionAliases.canonical(parentSessionID, "", ttl, now) + } return canonical, parentSessionID } diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go index 74819db0..328be59c 100644 --- a/sdk/cliproxy/auth/selector.go +++ b/sdk/cliproxy/auth/selector.go @@ -796,6 +796,9 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri if err != nil { return nil, err } + if auth == nil { + return nil, nil + } bind(auth.ID) entry.Infof("session-affinity: cache hit but auth unavailable, reselected | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) return auth, nil @@ -817,6 +820,9 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri if err != nil { return nil, err } + if auth == nil { + 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) return auth, nil diff --git a/sdk/cliproxy/auth/selector_lcp_test.go b/sdk/cliproxy/auth/selector_lcp_test.go index 631913f1..34b8633c 100644 --- a/sdk/cliproxy/auth/selector_lcp_test.go +++ b/sdk/cliproxy/auth/selector_lcp_test.go @@ -693,6 +693,74 @@ func TestSessionAffinitySelectorSessionTreeIntegration(t *testing.T) { } } +func TestSessionCacheCompareAndDeleteMultipleAliases(t *testing.T) { + t.Parallel() + + cache := NewSessionCache(time.Hour) + defer cache.Stop() + + // Set 3 aliases in the same group + cache.SetAliases("auth-1", "s1", "s2", "s3") + + if authID, ok := cache.Get("s1"); !ok || authID != "auth-1" { + t.Fatalf("s1 = %q, %v", authID, ok) + } + if authID, ok := cache.Get("s2"); !ok || authID != "auth-1" { + t.Fatalf("s2 = %q, %v", authID, ok) + } + if authID, ok := cache.Get("s3"); !ok || authID != "auth-1" { + t.Fatalf("s3 = %q, %v", authID, ok) + } + + // Compare and delete s1 + if !cache.CompareAndDelete("s1", "auth-1") { + t.Fatal("CompareAndDelete s1 failed") + } + + // s1 should be gone + if _, ok := cache.Get("s1"); ok { + t.Fatal("s1 still exists after CompareAndDelete") + } + + // s2 and s3 should still exist and point to auth-1 + if authID, ok := cache.Get("s2"); !ok || authID != "auth-1" { + t.Fatalf("s2 = %q, %v", authID, ok) + } + if authID, ok := cache.Get("s3"); !ok || authID != "auth-1" { + t.Fatalf("s3 = %q, %v", authID, ok) + } +} + +type nilSelector struct{} + +func (nilSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, candidates []*Auth) (*Auth, error) { + return nil, nil +} + +func TestSessionAffinitySelectorNilFallbackNoPanic(t *testing.T) { + t.Parallel() + + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: nilSelector{}, + TTL: time.Hour, + }) + defer selector.Stop() + + opts := cliproxyexecutor.Options{ + Headers: http.Header{ + "X-Session-ID": []string{"sess-nil-test"}, + }, + } + candidates := []*Auth{{ID: "auth-1", Status: StatusActive}} + auth, err := selector.Pick(context.Background(), "openai", "gpt-4o", opts, candidates) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if auth != nil { + t.Fatalf("expected nil auth, got %+v", auth) + } +} + func TestSessionCacheTinyTTLNoPanic(t *testing.T) { t.Parallel() diff --git a/sdk/cliproxy/auth/session_cache.go b/sdk/cliproxy/auth/session_cache.go index 8927d15f..b493bc81 100644 --- a/sdk/cliproxy/auth/session_cache.go +++ b/sdk/cliproxy/auth/session_cache.go @@ -302,21 +302,19 @@ func (c *SessionCache) CompareAndDelete(sessionID, expectedAuthID string) bool { return false } delete(c.entries, sessionID) + + surviving := make([]string, 0, len(entry.aliases)) for _, alias := range entry.aliases { - if alias == sessionID { - continue + if alias != sessionID { + surviving = append(surviving, alias) } + } + for _, alias := range surviving { current, exists := c.entries[alias] if !exists || current.authID != entry.authID { continue } - filtered := make([]string, 0, len(current.aliases)) - for _, candidate := range current.aliases { - if candidate != sessionID { - filtered = append(filtered, candidate) - } - } - current.aliases = filtered + current.aliases = append([]string(nil), surviving...) c.entries[alias] = current } return true @@ -332,21 +330,18 @@ func (c *SessionCache) Invalidate(sessionID string) { entry, ok := c.entries[sessionID] delete(c.entries, sessionID) if ok { + surviving := make([]string, 0, len(entry.aliases)) for _, alias := range entry.aliases { - if alias == sessionID { - continue + if alias != sessionID { + surviving = append(surviving, alias) } + } + for _, alias := range surviving { current, exists := c.entries[alias] if !exists || current.authID != entry.authID { continue } - filtered := make([]string, 0, len(current.aliases)) - for _, candidate := range current.aliases { - if candidate != sessionID { - filtered = append(filtered, candidate) - } - } - current.aliases = filtered + current.aliases = append([]string(nil), surviving...) c.entries[alias] = current } } diff --git a/sdk/cliproxy/session/identity.go b/sdk/cliproxy/session/identity.go index 0fce505c..e092272e 100644 --- a/sdk/cliproxy/session/identity.go +++ b/sdk/cliproxy/session/identity.go @@ -69,7 +69,7 @@ func ClaudeMetadataIdentities(payload []byte) (sessionID, parentSessionID, agent parsed := gjson.Parse(userID) sessionID = NormalizeExplicitID(parsed.Get("session_id").String()) parentSessionID = NormalizeExplicitID(parsed.Get("parent_session_id").String()) - agentID = strings.TrimSpace(parsed.Get("agent_id").String()) + agentID = NormalizeExplicitID(parsed.Get("agent_id").String()) return sessionID, parentSessionID, agentID } if matches := legacyClaudeSessionPattern.FindStringSubmatch(userID); len(matches) >= 2 { @@ -254,12 +254,16 @@ func DeriveID(format sdktranslator.Format, payload []byte, callerScope string) s Format: format.String(), CallerScope: strings.TrimSpace(callerScope), } - if sourceFormatEqual(format, sdktranslator.FormatGemini) { - root.Resource = stringField(body, "cachedContent", "cached_content") + if sourceFormatEqual(format, sdktranslator.FormatGemini) || sourceFormatEqual(format, sdktranslator.FormatAntigravity) { + reqBody := body + if req, ok := body["request"].(map[string]any); ok { + reqBody = req + } + root.Resource = stringField(reqBody, "cachedContent", "cached_content") } switch { - case sourceFormatEqual(format, sdktranslator.FormatGemini): + case sourceFormatEqual(format, sdktranslator.FormatGemini), sourceFormatEqual(format, sdktranslator.FormatAntigravity): root.Instructions, root.User = geminiRoot(body) case sourceFormatEqual(format, sdktranslator.FormatInteractions): root.Instructions, root.User = interactionsRoot(body) @@ -294,7 +298,10 @@ func messagesRoot(body map[string]any, includeTopLevelSystem bool) ([]string, [] case "system", "developer": instructions = appendInstruction(instructions, message["content"]) case "user": - return instructions, canonicalParts(message["content"]) + parts := canonicalParts(message["content"]) + if len(parts) > 0 { + return instructions, parts + } } } return instructions, nil @@ -323,13 +330,19 @@ func responsesRoot(body map[string]any) ([]string, []canonicalPart) { case "system", "developer": instructions = appendInstruction(instructions, item["content"]) case "user": - return instructions, canonicalParts(item["content"]) + parts := canonicalParts(item["content"]) + if len(parts) > 0 { + return instructions, parts + } } } return instructions, nil } func geminiRoot(body map[string]any) ([]string, []canonicalPart) { + if req, ok := body["request"].(map[string]any); ok { + body = req + } instructions := make([]string, 0) if value, ok := firstField(body, "systemInstruction", "system_instruction"); ok { instructions = appendInstruction(instructions, contentValue(value)) @@ -340,7 +353,10 @@ func geminiRoot(body map[string]any) ([]string, []canonicalPart) { if !okContent || normalizedString(content["role"]) != "user" { continue } - return instructions, canonicalParts(contentValue(content)) + parts := canonicalParts(contentValue(content)) + if len(parts) > 0 { + return instructions, parts + } } return instructions, nil } diff --git a/sdk/cliproxy/session/identity_test.go b/sdk/cliproxy/session/identity_test.go index 7211d379..6d441c3f 100644 --- a/sdk/cliproxy/session/identity_test.go +++ b/sdk/cliproxy/session/identity_test.go @@ -1,7 +1,6 @@ package session import ( - "bytes" "net/http" "strings" "testing" @@ -327,22 +326,62 @@ func TestEnrichCopiesDerivedIdentityToRequestAndOptions(t *testing.T) { } } -func TestEnrichCarriesRequestPayloadIntoSelectionOptions(t *testing.T) { +func TestDeriveIDAntigravityNestedRequestAndEmptyFirstUser(t *testing.T) { t.Parallel() - payload := []byte(`{"conversation":{"id":"request-only-conversation"},"input":"hello"}`) - _, enrichedOpts := Enrich( - cliproxyexecutor.Request{Payload: payload}, - cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse}, - ) + nestedAntigravity := []byte(`{ + "project_id": "test-project", + "request": { + "systemInstruction": {"parts":[{"text":"system prompt"}]}, + "contents": [ + {"role":"user","parts":[{"text":""}]}, + {"role":"user","parts":[{"text":"actual user prompt"}]} + ] + } + }`) + id := DeriveID(sdktranslator.FormatAntigravity, nestedAntigravity, "caller-a") + if id == "" { + t.Fatal("DeriveID returned empty for nested Antigravity request with empty first turn") + } + + directAntigravity := []byte(`{ + "systemInstruction": {"parts":[{"text":"system prompt"}]}, + "contents": [ + {"role":"user","parts":[{"text":"actual user prompt"}]} + ] + }`) + directID := DeriveID(sdktranslator.FormatAntigravity, directAntigravity, "caller-a") + if id != directID { + t.Fatalf("DeriveID mismatch: nested=%s, direct=%s", id, directID) + } +} - if !bytes.Equal(enrichedOpts.OriginalRequest, payload) { - t.Fatalf("OriginalRequest = %q, want request payload %q", enrichedOpts.OriginalRequest, payload) +func TestClaudeMetadataIdentitiesNormalizesAgentID(t *testing.T) { + t.Parallel() + + validPayload := []byte(`{ + "metadata": { + "user_id": "{\"session_id\":\"sess-123\",\"parent_session_id\":\"parent-456\",\"agent_id\":\" subagent-alpha \"}" + } + }`) + sessionID, parentSessionID, agentID := ClaudeMetadataIdentities(validPayload) + if sessionID != "sess-123" { + t.Fatalf("sessionID = %q, want sess-123", sessionID) } - if len(enrichedOpts.OriginalRequest) > 0 && &enrichedOpts.OriginalRequest[0] == &payload[0] { - t.Fatal("OriginalRequest aliases Request.Payload instead of preserving a snapshot") + if parentSessionID != "parent-456" { + t.Fatalf("parentSessionID = %q, want parent-456", parentSessionID) } - if got := DerivedID(enrichedOpts.Metadata); got != "" { - t.Fatalf("DerivedSessionID = %q, want explicit conversation to remain authoritative", got) + if agentID != "subagent-alpha" { + t.Fatalf("agentID = %q, want subagent-alpha", agentID) + } + + invalidPayload := []byte(`{ + "metadata": { + "user_id": "{\"session_id\":\"sess-123\",\"agent_id\":\"bad\nagent\"}" + } + }`) + _, _, badAgentID := ClaudeMetadataIdentities(invalidPayload) + if badAgentID != "" { + t.Fatalf("expected badAgentID to be empty for control character, got %q", badAgentID) } } diff --git a/sdk/cliproxy/session/lcp.go b/sdk/cliproxy/session/lcp.go index 09831a27..e0c056df 100644 --- a/sdk/cliproxy/session/lcp.go +++ b/sdk/cliproxy/session/lcp.go @@ -108,6 +108,9 @@ func ExtractCanonicalTurns(format sdktranslator.Format, payload []byte) []Canoni } func inferCanonicalFormat(root gjson.Result) sdktranslator.Format { + if req := root.Get("request"); req.Exists() && !root.Get("contents").Exists() { + root = req + } if root.Get("contents").IsArray() || root.Get("systemInstruction").Exists() || root.Get("system_instruction").Exists() { return sdktranslator.FormatGemini } @@ -207,6 +210,9 @@ func appendResponsesTurns(turns *[]CanonicalTurn, root gjson.Result) { } func appendGeminiTurns(turns *[]CanonicalTurn, root gjson.Result) { + if req := root.Get("request"); req.Exists() && !root.Get("contents").Exists() { + root = req + } cached := root.Get("cachedContent") if !cached.Exists() { cached = root.Get("cached_content") diff --git a/sdk/cliproxy/session/lcp_test.go b/sdk/cliproxy/session/lcp_test.go index 2a1a9c81..7a07c738 100644 --- a/sdk/cliproxy/session/lcp_test.go +++ b/sdk/cliproxy/session/lcp_test.go @@ -525,6 +525,27 @@ func TestMerklePrefixMatcherTouchExpiredEntry(t *testing.T) { } } +func TestExtractCanonicalTurnsAntigravityNestedRequest(t *testing.T) { + t.Parallel() + + nested := []byte(`{ + "project_id": "proj-123", + "request": { + "systemInstruction": {"parts":[{"text":"system prompt"}]}, + "contents": [ + {"role":"user","parts":[{"text":"hello antigravity"}]} + ] + } + }`) + turns := ExtractCanonicalTurns(sdktranslator.FormatAntigravity, nested) + if len(turns) != 2 { + t.Fatalf("len(turns) = %d, want 2", len(turns)) + } + if turns[0].Role != "system" || turns[1].Role != "user" { + t.Fatalf("unexpected turn roles: %+v", turns) + } +} + func BenchmarkExtractCanonicalTurns(b *testing.B) { payload := []byte(`{ "messages": [ diff --git a/sdk/cliproxy/session/tree.go b/sdk/cliproxy/session/tree.go index 7b106808..9c492d8b 100644 --- a/sdk/cliproxy/session/tree.go +++ b/sdk/cliproxy/session/tree.go @@ -16,6 +16,7 @@ import ( const ( defaultMaxTreeNodes = 65536 defaultTreeTTL = 24 * time.Hour + maxTreeDepth = 128 ) // SessionTreeNode represents a node in the hierarchical session tree. @@ -769,9 +770,16 @@ func (s *InMemorySessionTreeStore) computeNodeLineageLocked(node *SessionTreeNod if node.ParentSessionID != "" && node.ParentSessionID != node.SessionID && !s.isDescendantLocked(node.ParentSessionID, node.SessionID) { if parent, ok := s.nodes[node.ParentSessionID]; ok { - newRoot = parent.RootSessionID - newPath = parent.TreePath + "/" + node.SessionID - newDepth = parent.TreeDepth + 1 + if parent.TreeDepth < maxTreeDepth { + newRoot = parent.RootSessionID + newPath = parent.TreePath + "/" + node.SessionID + newDepth = parent.TreeDepth + 1 + } else { + node.ParentSessionID = "" + newRoot = node.SessionID + newPath = node.SessionID + newDepth = 0 + } } else { newRoot = node.ParentSessionID newPath = node.ParentSessionID + "/" + node.SessionID diff --git a/sdk/cliproxy/session/tree_test.go b/sdk/cliproxy/session/tree_test.go index 1786d6c2..e630569e 100644 --- a/sdk/cliproxy/session/tree_test.go +++ b/sdk/cliproxy/session/tree_test.go @@ -611,6 +611,34 @@ func TestExtractTreeInfoGeminiAndAntigravityHierarchy(t *testing.T) { } } +func TestSessionTreeStoreMaxDepthCapping(t *testing.T) { + t.Parallel() + + store := NewInMemorySessionTreeStore(500, time.Hour) + current := "node-0" + store.RecordNode(SessionTreeInfo{SessionID: current}) + + // Create a chain exceeding maxTreeDepth (128) + for i := 1; i <= 150; i++ { + next := fmt.Sprintf("node-%d", i) + node := store.RecordNode(SessionTreeInfo{ + SessionID: next, + ParentSessionID: current, + }) + if i <= 128 { + if node.TreeDepth != i { + t.Fatalf("expected depth %d, got %d", i, node.TreeDepth) + } + } else { + // Once exceeding maxTreeDepth, the chain is capped and the node becomes its own root + if node.TreeDepth > maxTreeDepth { + t.Fatalf("depth %d exceeded maxTreeDepth %d", node.TreeDepth, maxTreeDepth) + } + } + current = next + } +} + func BenchmarkSessionTreeRecordAndCascade(b *testing.B) { store := NewInMemorySessionTreeStore(10000, time.Hour) // Pre-populate root and parent -- 2.54.0