diff --git a/config.example.yaml b/config.example.yaml index c38c565f..4bba36ac 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -229,6 +229,13 @@ routing: session-affinity: false # default: false # How long session-to-auth bindings are retained. Default: 1h session-affinity-ttl: "1h" + # When true (default), subagents (child sessions with parent references) inherit and bind + # to the parent's upstream credential across all providers (Claude, Codex, Antigravity, Gemini), + # maximizing Prompt/KV Cache reuse and minimizing TTFT. + # When false, subagents are distributed across the credential pool via the fallback selector + # for maximum parallel concurrency. + # Ignored when routing.session-affinity is false. + session-affinity-subagents: true # Codex provider behavior. codex: diff --git a/internal/config/config_types.go b/internal/config/config_types.go index 6228284c..05b21e71 100644 --- a/internal/config/config_types.go +++ b/internal/config/config_types.go @@ -245,6 +245,13 @@ type RoutingConfig struct { // SessionAffinityTTL specifies how long session-to-auth bindings are retained. // Default: 1h. Accepts duration strings like "30m", "1h", "2h30m". SessionAffinityTTL string `yaml:"session-affinity-ttl,omitempty" json:"session-affinity-ttl,omitempty"` + + // SessionAffinitySubagents controls whether subagents (child sessions with parent references) + // inherit and bind to the parent's upstream credential across all providers (Claude, Codex, + // Antigravity, Gemini), maximizing prompt and KV cache reuse. + // When false, subagents are distributed across the credential pool via the fallback selector. + // Default: true. Ignored when SessionAffinity is false. + SessionAffinitySubagents *bool `yaml:"session-affinity-subagents,omitempty" json:"session-affinity-subagents,omitempty"` } // OAuthModelAlias defines a model ID alias for a specific channel. diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go index 04f93da5..510a8a52 100644 --- a/sdk/cliproxy/auth/selector.go +++ b/sdk/cliproxy/auth/selector.go @@ -874,15 +874,17 @@ func availabilityBlock(unavailable, quotaExceeded bool, nextRetryAfter, nextReco // It extracts session ID from multiple sources and maintains session-to-auth // mappings with automatic failover when the bound auth becomes unavailable. type SessionAffinitySelector struct { - fallback Selector - cache *SessionCache - matcher *cliproxysession.MerklePrefixMatcher + fallback Selector + cache *SessionCache + matcher *cliproxysession.MerklePrefixMatcher + subagentAffinity bool } // SessionAffinityConfig configures the session affinity selector. type SessionAffinityConfig struct { - Fallback Selector - TTL time.Duration + Fallback Selector + TTL time.Duration + SubagentAffinity *bool } // NewSessionAffinitySelector creates a new session-aware selector. @@ -901,10 +903,15 @@ func NewSessionAffinitySelectorWithConfig(cfg SessionAffinityConfig) *SessionAff if cfg.TTL <= 0 { cfg.TTL = time.Hour } + subagentAffinity := true + if cfg.SubagentAffinity != nil { + subagentAffinity = *cfg.SubagentAffinity + } return &SessionAffinitySelector{ - fallback: cfg.Fallback, - cache: NewSessionCache(cfg.TTL), - matcher: cliproxysession.NewMerklePrefixMatcher(cfg.TTL), + fallback: cfg.Fallback, + cache: NewSessionCache(cfg.TTL), + matcher: cliproxysession.NewMerklePrefixMatcher(cfg.TTL), + subagentAffinity: subagentAffinity, } } @@ -1013,7 +1020,7 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri if cachedAuthID, ok := s.cache.Get(fallbackKey); ok { for _, auth := range available { if auth.ID == cachedAuthID { - if !isSubagent || allowsSubagentAuthInheritance(auth, model) { + 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) return auth, nil @@ -1274,31 +1281,6 @@ func isSubagentSession(primaryID, fallbackID string) bool { return isHierarchyParent(primaryID, fallbackID) } -func isNonInheritingProvider(provider string) bool { - p := strings.ToLower(strings.TrimSpace(provider)) - return p == "antigravity" || p == "gemini" || p == "vertex" || p == "gemini-vertex" || p == "aistudio" || p == "gemini-interactions" || strings.Contains(p, "antigravity") || strings.Contains(p, "gemini") -} - -func isNonInheritingModel(model string) bool { - baseModel := strings.ToLower(strings.TrimSpace(model)) - if parsed := thinking.ParseSuffix(baseModel); parsed.ModelName != "" { - baseModel = strings.ToLower(strings.TrimSpace(parsed.ModelName)) - } - baseModel = strings.TrimPrefix(baseModel, "models/") - baseModel = strings.TrimPrefix(baseModel, "google/") - return strings.HasPrefix(baseModel, "gemini-") || strings.HasPrefix(baseModel, "antigravity-") || strings.Contains(baseModel, "gemini") || strings.Contains(baseModel, "antigravity") -} - -func allowsSubagentAuthInheritance(auth *Auth, model string) bool { - if auth != nil && isNonInheritingProvider(auth.Provider) { - return false - } - if isNonInheritingModel(model) { - return false - } - return true -} - func sessionHeaderValue(headers http.Header, name string) string { if headers == nil { return "" diff --git a/sdk/cliproxy/auth/selector_antigravity_subagent_test.go b/sdk/cliproxy/auth/selector_antigravity_subagent_test.go index 9c6c72d7..da1984b3 100644 --- a/sdk/cliproxy/auth/selector_antigravity_subagent_test.go +++ b/sdk/cliproxy/auth/selector_antigravity_subagent_test.go @@ -9,12 +9,17 @@ import ( cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" ) +func boolPointer(b bool) *bool { + return &b +} + func TestSessionAffinityAntigravitySubagentDoesNotInheritParentBinding(t *testing.T) { t.Parallel() selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ - Fallback: &RoundRobinSelector{}, - TTL: time.Minute, + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + SubagentAffinity: boolPointer(false), }) defer selector.Stop() @@ -111,8 +116,9 @@ func TestSessionAffinityMixedProviderAntigravitySubagentIsolation(t *testing.T) t.Parallel() selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ - Fallback: &RoundRobinSelector{}, - TTL: time.Minute, + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + SubagentAffinity: boolPointer(false), }) defer selector.Stop() @@ -242,7 +248,7 @@ func TestSessionAffinityClaudeAndCodexStillInheritParentBinding(t *testing.T) { } } -func TestSessionAffinityMixedPoolClaudeInheritsWhileAntigravityIsolates(t *testing.T) { +func TestSessionAffinityMixedPoolSubagentInheritsParentAcrossProviders(t *testing.T) { t.Parallel() selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ @@ -316,8 +322,8 @@ func TestSessionAffinityMixedPoolClaudeInheritsWhileAntigravityIsolates(t *testi t.Fatalf("parent Ag auth = %q, want auth-ag-primary", parentAgAuth.ID) } - // 4. Subagent of Antigravity parent in mixed pool where Claude is also available - // MUST NOT inherit Antigravity auth, and independently picking Claude must NOT overwrite parent binding + // 4. Subagent of Antigravity parent in mixed pool under unified inheritance + // MUST inherit parent auth-ag-primary for prompt cache reuse mixedCandidates := []*Auth{ {ID: "auth-ag-primary", Provider: "antigravity"}, {ID: "auth-claude-primary", Provider: "claude"}, @@ -337,25 +343,22 @@ func TestSessionAffinityMixedPoolClaudeInheritsWhileAntigravityIsolates(t *testi if errSubAg != nil { t.Fatalf("subagent Ag Pick() error = %v", errSubAg) } - if subAgAuth.ID == parentAgAuth.ID { - t.Fatalf("subagent with Antigravity parent incorrectly inherited parent auth %q", parentAgAuth.ID) - } - if subAgAuth.ID != "auth-claude-primary" { - t.Fatalf("subagent Ag auth = %q, want auth-claude-primary", subAgAuth.ID) + if subAgAuth.ID != parentAgAuth.ID { + t.Fatalf("subagent with Antigravity parent did not inherit parent auth: got %q, want %q", subAgAuth.ID, parentAgAuth.ID) } - // 5. Verify parent binding in cache was NOT overwritten by subagent picking Claude + // 5. Verify parent binding in cache remains intact if bound, ok := selector.cache.Get("mixed::claude:mixed-sess-ag::claude-3-7-sonnet"); !ok || bound != "auth-ag-primary" { - t.Fatalf("parent cache binding was overwritten by subagent: got (%q, %v), want (auth-ag-primary, true)", bound, ok) + t.Fatalf("parent cache binding was lost: got (%q, %v), want (auth-ag-primary, true)", bound, ok) } - // 6. Subagent turn 2 must stay sticky to auth-claude-primary + // 6. Subagent turn 2 must stay sticky to auth-ag-primary subAgTurn2Auth, errSubAg2 := selector.Pick(context.Background(), "mixed", "claude-3-7-sonnet", subAgOpts, mixedCandidates) if errSubAg2 != nil { t.Fatalf("subagent turn 2 Pick() error = %v", errSubAg2) } - if subAgTurn2Auth.ID != "auth-claude-primary" { - t.Fatalf("subagent turn 2 auth = %q, want auth-claude-primary", subAgTurn2Auth.ID) + if subAgTurn2Auth.ID != "auth-ag-primary" { + t.Fatalf("subagent turn 2 auth = %q, want auth-ag-primary", subAgTurn2Auth.ID) } // 7. Parent turn 2 must stay sticky to auth-ag-primary @@ -450,8 +453,9 @@ func TestSessionAffinityOtherGoogleProvidersSubagentIsolation(t *testing.T) { for _, provider := range []string{"gemini", "vertex", "aistudio", "gemini-interactions"} { t.Run(provider, func(t *testing.T) { selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ - Fallback: &RoundRobinSelector{}, - TTL: time.Minute, + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + SubagentAffinity: boolPointer(false), }) defer selector.Stop() @@ -499,8 +503,9 @@ func TestSessionAffinityNestedAntigravitySubagentIsolation(t *testing.T) { t.Parallel() selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ - Fallback: &RoundRobinSelector{}, - TTL: time.Minute, + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + SubagentAffinity: boolPointer(false), }) defer selector.Stop() @@ -617,8 +622,9 @@ func TestSessionAffinityNestedGeminiProviderSubagentIsolation(t *testing.T) { t.Parallel() selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ - Fallback: &RoundRobinSelector{}, - TTL: time.Minute, + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + SubagentAffinity: boolPointer(false), }) defer selector.Stop() diff --git a/sdk/cliproxy/auth/selector_lcp_test.go b/sdk/cliproxy/auth/selector_lcp_test.go index 7e349d17..616bcd37 100644 --- a/sdk/cliproxy/auth/selector_lcp_test.go +++ b/sdk/cliproxy/auth/selector_lcp_test.go @@ -792,8 +792,9 @@ func TestSessionAffinityClaudeMetadataSubagentNonInheritingGeminiModel(t *testin t.Parallel() selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ - Fallback: &RoundRobinSelector{}, - TTL: time.Minute, + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + SubagentAffinity: boolPointer(false), }) defer selector.Stop() @@ -830,9 +831,9 @@ func TestSessionAffinityClaudeMetadataSubagentNonInheritingGeminiModel(t *testin if errSub1 != nil { t.Fatalf("subagent 1 Pick() error = %v", errSub1) } - // For Gemini/Antigravity, subagents must NOT inherit the parent's auth; they should balance to auth-b + // When SubagentAffinity is false, subagents must NOT inherit the parent's auth; they should balance to auth-b if subagent1Auth.ID != "auth-b" { - t.Fatalf("subagent 1 should not inherit parent auth for Gemini, got %q, want auth-b", subagent1Auth.ID) + t.Fatalf("subagent 1 should not inherit parent auth when subagent affinity is disabled, got %q, want auth-b", subagent1Auth.ID) } if got := subagent1Opts.Metadata[cliproxyexecutor.CanonicalSessionIDMetadataKey]; got != "claude:sess-main-1:agent:subagent-001" { t.Fatalf("subagent 1 canonical session ID = %v, want claude:sess-main-1:agent:subagent-001", got) diff --git a/sdk/cliproxy/auth/selector_subagent_affinity_test.go b/sdk/cliproxy/auth/selector_subagent_affinity_test.go new file mode 100644 index 00000000..b8e81fc8 --- /dev/null +++ b/sdk/cliproxy/auth/selector_subagent_affinity_test.go @@ -0,0 +1,312 @@ +package auth + +import ( + "context" + "net/http" + "testing" + "time" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +// TestSessionAffinityAntigravitySubagentInheritsParentByDefault reproduces issue #5417: +// Subagents across all providers (including Antigravity and Gemini) should inherit +// the parent session's credential by default to maximize prompt and KV cache reuse. +func TestSessionAffinityAntigravitySubagentInheritsParentByDefault(t *testing.T) { + t.Parallel() + + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + }) + defer selector.Stop() + + auths := []*Auth{ + {ID: "auth-ag-1", Provider: "antigravity"}, + {ID: "auth-ag-2", Provider: "antigravity"}, + } + + // 1. Parent request binds to auth-ag-1. + parentOpts := cliproxyexecutor.Options{ + Headers: http.Header{"X-Claude-Code-Session-Id": []string{"claude-root-100"}}, + OriginalRequest: []byte(`{"messages":[{"role":"user","content":"parent task"}]}`), + Metadata: map[string]any{}, + } + parentAuth, errParent := selector.Pick(context.Background(), "antigravity", "gemini-3.7-flash-high", parentOpts, auths) + if errParent != nil { + t.Fatalf("parent Pick() error = %v", errParent) + } + if parentAuth.ID != "auth-ag-1" { + t.Fatalf("parent auth = %q, want auth-ag-1", parentAuth.ID) + } + + // 2. Subagent request carries child agent ID and should inherit parent auth-ag-1 by default. + subagentOpts := cliproxyexecutor.Options{ + Headers: http.Header{ + "X-Claude-Code-Session-Id": []string{"claude-root-100"}, + "X-Claude-Code-Agent-Id": []string{"subagent-001"}, + }, + OriginalRequest: []byte(`{"messages":[{"role":"user","content":"subagent task"}]}`), + Metadata: map[string]any{}, + } + subagentAuth, errSub := selector.Pick(context.Background(), "antigravity", "gemini-3.7-flash-high", subagentOpts, auths) + if errSub != nil { + t.Fatalf("subagent Pick() error = %v", errSub) + } + if subagentAuth.ID != parentAuth.ID { + t.Fatalf("subagent did not inherit parent auth %q, got %q", parentAuth.ID, subagentAuth.ID) + } + + // 3. Subsequent turn of subagent must stay sticky to the inherited auth. + subagentTurn2Auth, errSub2 := selector.Pick(context.Background(), "antigravity", "gemini-3.7-flash-high", subagentOpts, auths) + if errSub2 != nil { + t.Fatalf("subagent turn 2 Pick() error = %v", errSub2) + } + if subagentTurn2Auth.ID != "auth-ag-1" { + t.Fatalf("subagent turn 2 did not stay sticky: got %q, want auth-ag-1", subagentTurn2Auth.ID) + } +} + +// TestSessionAffinityGeminiSubagentInheritsParentWhenExplicitlyTrue verifies that when +// SubagentAffinity is explicitly set to true, Gemini subagents inherit the parent's auth. +func TestSessionAffinityGeminiSubagentInheritsParentWhenExplicitlyTrue(t *testing.T) { + t.Parallel() + + enabled := true + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + SubagentAffinity: &enabled, + }) + defer selector.Stop() + + auths := []*Auth{ + {ID: "auth-gem-1", Provider: "gemini"}, + {ID: "auth-gem-2", Provider: "gemini"}, + } + + // 1. Parent request binds to auth-gem-1. + parentOpts := cliproxyexecutor.Options{ + Headers: http.Header{"X-Claude-Code-Session-Id": []string{"gem-root-200"}}, + OriginalRequest: []byte(`{"messages":[{"role":"user","content":"parent gemini task"}]}`), + Metadata: map[string]any{}, + } + parentAuth, errParent := selector.Pick(context.Background(), "gemini", "gemini-2.5-pro", parentOpts, auths) + if errParent != nil { + t.Fatalf("parent Pick() error = %v", errParent) + } + if parentAuth.ID != "auth-gem-1" { + t.Fatalf("parent auth = %q, want auth-gem-1", parentAuth.ID) + } + + // 2. Subagent request inherits parent auth-gem-1. + subagentOpts := cliproxyexecutor.Options{ + Headers: http.Header{ + "X-Claude-Code-Session-Id": []string{"gem-root-200"}, + "X-Claude-Code-Agent-Id": []string{"gem-sub-001"}, + }, + OriginalRequest: []byte(`{"messages":[{"role":"user","content":"subagent gemini task"}]}`), + Metadata: map[string]any{}, + } + subagentAuth, errSub := selector.Pick(context.Background(), "gemini", "gemini-2.5-pro", subagentOpts, auths) + if errSub != nil { + t.Fatalf("subagent Pick() error = %v", errSub) + } + if subagentAuth.ID != parentAuth.ID { + t.Fatalf("subagent did not inherit parent auth %q, got %q", parentAuth.ID, subagentAuth.ID) + } +} + +// TestSessionAffinitySubagentIsolatesWhenExplicitlyFalse verifies that when +// SubagentAffinity is false, subagents do not inherit parent credentials and balance across pool. +func TestSessionAffinitySubagentIsolatesWhenExplicitlyFalse(t *testing.T) { + t.Parallel() + + disabled := false + selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ + Fallback: &RoundRobinSelector{}, + TTL: time.Minute, + SubagentAffinity: &disabled, + }) + defer selector.Stop() + + auths := []*Auth{ + {ID: "auth-1", Provider: "claude"}, + {ID: "auth-2", Provider: "claude"}, + } + + // 1. Parent request binds to auth-1. + parentOpts := cliproxyexecutor.Options{ + Headers: http.Header{"X-Claude-Code-Session-Id": []string{"sess-iso-300"}}, + OriginalRequest: []byte(`{"messages":[{"role":"user","content":"parent task"}]}`), + Metadata: map[string]any{}, + } + parentAuth, errParent := selector.Pick(context.Background(), "claude", "claude-3-7-sonnet", parentOpts, auths) + if errParent != nil { + t.Fatalf("parent Pick() error = %v", errParent) + } + if parentAuth.ID != "auth-1" { + t.Fatalf("parent auth = %q, want auth-1", parentAuth.ID) + } + + // 2. Subagent request should NOT inherit auth-1; it picks auth-2 via fallback round-robin. + subagentOpts := cliproxyexecutor.Options{ + Headers: http.Header{ + "X-Claude-Code-Session-Id": []string{"sess-iso-300"}, + "X-Claude-Code-Agent-Id": []string{"sub-iso-001"}, + }, + OriginalRequest: []byte(`{"messages":[{"role":"user","content":"subagent task"}]}`), + Metadata: map[string]any{}, + } + subagentAuth, errSub := selector.Pick(context.Background(), "claude", "claude-3-7-sonnet", subagentOpts, auths) + if errSub != nil { + t.Fatalf("subagent Pick() error = %v", errSub) + } + if subagentAuth.ID == parentAuth.ID { + t.Fatalf("subagent incorrectly inherited parent auth %q when SubagentAffinity is false", parentAuth.ID) + } + if subagentAuth.ID != "auth-2" { + t.Fatalf("subagent auth = %q, want auth-2", subagentAuth.ID) + } +} + +// TestSessionAffinitySubagentFailureIsolation verifies that failure on a subagent session +// only deletes the subagent's own binding, leaving the parent's binding intact. +func TestSessionAffinitySubagentFailureIsolation(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 request binds to auth-1. + parentOpts := cliproxyexecutor.Options{ + Headers: http.Header{"X-Claude-Code-Session-Id": []string{"sess-fail-400"}}, + OriginalRequest: []byte(`{"messages":[{"role":"user","content":"parent task"}]}`), + Metadata: map[string]any{}, + } + parentAuth, errParent := selector.Pick(context.Background(), "antigravity", "gemini-3.7-flash-high", parentOpts, auths) + if errParent != nil { + t.Fatalf("parent Pick() error = %v", errParent) + } + if parentAuth.ID != "auth-1" { + t.Fatalf("parent auth = %q, want auth-1", parentAuth.ID) + } + + // Parent cache key must be bound to auth-1. + parentCacheKey := "antigravity::claude:sess-fail-400::gemini-3.7-flash-high" + if bound, ok := selector.cache.Get(parentCacheKey); !ok || bound != "auth-1" { + t.Fatalf("parent cache binding = (%q, %v), want (auth-1, true)", bound, ok) + } + + // 2. Subagent inherits auth-1 and binds its own cache key. + subagentOpts := cliproxyexecutor.Options{ + Headers: http.Header{ + "X-Claude-Code-Session-Id": []string{"sess-fail-400"}, + "X-Claude-Code-Agent-Id": []string{"sub-fail-001"}, + }, + OriginalRequest: []byte(`{"messages":[{"role":"user","content":"subagent task"}]}`), + Metadata: map[string]any{}, + } + subagentAuth, errSub := selector.Pick(context.Background(), "antigravity", "gemini-3.7-flash-high", subagentOpts, auths) + if errSub != nil { + t.Fatalf("subagent Pick() error = %v", errSub) + } + if subagentAuth.ID != "auth-1" { + t.Fatalf("subagent auth = %q, want auth-1", subagentAuth.ID) + } + + subagentCacheKey := "antigravity::claude:sess-fail-400:agent:sub-fail-001::gemini-3.7-flash-high" + if bound, ok := selector.cache.Get(subagentCacheKey); !ok || bound != "auth-1" { + t.Fatalf("subagent cache binding = (%q, %v), want (auth-1, true)", bound, ok) + } + + // 3. Subagent fails with an upstream error via OnResult. + selector.OnResult(Result{ + AuthID: "auth-1", + Success: false, + Options: subagentOpts, + }) + + // 4. Subagent's cache key must be cleared. + if bound, ok := selector.cache.Get(subagentCacheKey); ok { + t.Fatalf("subagent cache binding was not cleared after failure: got %q", bound) + } + + // 5. Parent session binding in cache MUST still remain intact on auth-1. + if bound, ok := selector.cache.Get(parentCacheKey); !ok || bound != "auth-1" { + t.Fatalf("parent binding was damaged by subagent failure: got (%q, %v), want (auth-1, true)", bound, ok) + } + + // 6. Parent request turn 2 must still hit auth-1. + parentTurn2Auth, errParent2 := selector.Pick(context.Background(), "antigravity", "gemini-3.7-flash-high", parentOpts, auths) + if errParent2 != nil { + t.Fatalf("parent turn 2 Pick() error = %v", errParent2) + } + if parentTurn2Auth.ID != "auth-1" { + t.Fatalf("parent turn 2 auth = %q, want auth-1", parentTurn2Auth.ID) + } +} + +// TestSessionAffinitySubagentAliasIsolation verifies that subagents bind only their +// own cacheKey and never register an alias that overrides the parent's fallbackKey. +func TestSessionAffinitySubagentAliasIsolation(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 binds auth-1. + parentOpts := cliproxyexecutor.Options{ + Headers: http.Header{"X-Claude-Code-Session-Id": []string{"sess-alias-500"}}, + OriginalRequest: []byte(`{"messages":[{"role":"user","content":"parent task"}]}`), + Metadata: map[string]any{}, + } + _, errParent := selector.Pick(context.Background(), "antigravity", "gemini-3.7-flash-high", parentOpts, auths) + if errParent != nil { + t.Fatalf("parent Pick() error = %v", errParent) + } + + // 2. Subagent picks and inherits auth-1. + subagentOpts := cliproxyexecutor.Options{ + Headers: http.Header{ + "X-Claude-Code-Session-Id": []string{"sess-alias-500"}, + "X-Claude-Code-Agent-Id": []string{"sub-alias-001"}, + }, + OriginalRequest: []byte(`{"messages":[{"role":"user","content":"sub task"}]}`), + Metadata: map[string]any{}, + } + _, errSub := selector.Pick(context.Background(), "antigravity", "gemini-3.7-flash-high", subagentOpts, auths) + if errSub != nil { + t.Fatalf("subagent Pick() error = %v", errSub) + } + + // 3. Invalidate subagent's auth by calling OnResult(Success: false). + selector.OnResult(Result{ + AuthID: "auth-1", + Success: false, + Options: subagentOpts, + }) + + // Parent fallbackKey in cache MUST still resolve to auth-1, proving subagent did not alias it. + parentCacheKey := "antigravity::claude:sess-alias-500::gemini-3.7-flash-high" + val, ok := selector.cache.Get(parentCacheKey) + if !ok || val != "auth-1" { + t.Fatalf("parent cache key %s = (%q, %v), want (auth-1, true)", parentCacheKey, val, ok) + } +} diff --git a/sdk/cliproxy/service_config.go b/sdk/cliproxy/service_config.go index 40c08d93..d04374af 100644 --- a/sdk/cliproxy/service_config.go +++ b/sdk/cliproxy/service_config.go @@ -25,15 +25,17 @@ type configCommit struct { } type routingRuntimeState struct { - strategy string - sessionAffinity bool - sessionAffinityTTL time.Duration + strategy string + sessionAffinity bool + sessionAffinityTTL time.Duration + sessionAffinitySubagents bool } func normalizedRoutingRuntimeState(cfg *config.Config) routingRuntimeState { state := routingRuntimeState{ - strategy: "round-robin", - sessionAffinityTTL: time.Hour, + strategy: "round-robin", + sessionAffinityTTL: time.Hour, + sessionAffinitySubagents: true, } if cfg == nil { return state @@ -54,6 +56,9 @@ func normalizedRoutingRuntimeState(cfg *config.Config) routingRuntimeState { state.sessionAffinityTTL = parsed } } + if state.sessionAffinity && cfg.Routing.SessionAffinitySubagents != nil { + state.sessionAffinitySubagents = *cfg.Routing.SessionAffinitySubagents + } return state } @@ -68,9 +73,11 @@ func newRoutingSelector(state routingRuntimeState) coreauth.Selector { selector = &coreauth.RoundRobinSelector{} } if state.sessionAffinity { + subagents := state.sessionAffinitySubagents selector = coreauth.NewSessionAffinitySelectorWithConfig(coreauth.SessionAffinityConfig{ - Fallback: selector, - TTL: state.sessionAffinityTTL, + Fallback: selector, + TTL: state.sessionAffinityTTL, + SubagentAffinity: &subagents, }) } return selector diff --git a/sdk/cliproxy/service_executionregistry_test.go b/sdk/cliproxy/service_executionregistry_test.go index 8219d939..38c7e815 100644 --- a/sdk/cliproxy/service_executionregistry_test.go +++ b/sdk/cliproxy/service_executionregistry_test.go @@ -2915,6 +2915,70 @@ func TestServiceApplyConfigRuntimePreservesSelectorForUnchangedRouting(t *testin } } +func TestServiceApplyConfigRuntimeSessionAffinitySubagentsChangeRecreatesSelector(t *testing.T) { + manager := coreauth.NewManager(nil, &coreauth.RoundRobinSelector{}, nil) + service := &Service{cfg: &config.Config{}, coreManager: manager} + + subagentsEnabled := true + initial := service.commitConfigUpdate(&config.Config{Routing: internalconfig.RoutingConfig{ + Strategy: "round-robin", + SessionAffinity: true, + SessionAffinityTTL: "1h", + SessionAffinitySubagents: &subagentsEnabled, + }}) + if !service.applyConfigRuntime(context.Background(), initial, false) { + t.Fatal("initial config runtime apply failed") + } + initialSelector := manager.Selector() + + subagentsDisabled := false + changed := service.commitConfigUpdate(&config.Config{Routing: internalconfig.RoutingConfig{ + Strategy: "round-robin", + SessionAffinity: true, + SessionAffinityTTL: "1h", + SessionAffinitySubagents: &subagentsDisabled, + }}) + if !service.applyConfigRuntime(context.Background(), changed, false) { + t.Fatal("changed config runtime apply failed") + } + changedSelector := manager.Selector() + if changedSelector == initialSelector { + t.Fatal("changing SessionAffinitySubagents should recreate selector") + } +} + +func TestServiceApplyConfigRuntimeSessionAffinityDisabledSubagentsChangeIsNoOp(t *testing.T) { + manager := coreauth.NewManager(nil, &coreauth.RoundRobinSelector{}, nil) + service := &Service{cfg: &config.Config{}, coreManager: manager} + + subagentsEnabled := true + initial := service.commitConfigUpdate(&config.Config{Routing: internalconfig.RoutingConfig{ + Strategy: "round-robin", + SessionAffinity: false, + SessionAffinityTTL: "1h", + SessionAffinitySubagents: &subagentsEnabled, + }}) + if !service.applyConfigRuntime(context.Background(), initial, false) { + t.Fatal("initial config runtime apply failed") + } + initialSelector := manager.Selector() + + subagentsDisabled := false + changed := service.commitConfigUpdate(&config.Config{Routing: internalconfig.RoutingConfig{ + Strategy: "round-robin", + SessionAffinity: false, + SessionAffinityTTL: "1h", + SessionAffinitySubagents: &subagentsDisabled, + }}) + if !service.applyConfigRuntime(context.Background(), changed, false) { + t.Fatal("changed config runtime apply failed") + } + changedSelector := manager.Selector() + if changedSelector != initialSelector { + t.Fatal("toggling SessionAffinitySubagents when SessionAffinity is false must be a no-op") + } +} + func TestServiceSerializesHomeAndWatcherConfigRuntimeApply(t *testing.T) { baseCfg := &config.Config{} baseCfg.Home.Enabled = true