diff --git a/backend/internal/handler/dto/mappers.go b/backend/internal/handler/dto/mappers.go index d3fc51ed66..816550ab04 100644 --- a/backend/internal/handler/dto/mappers.go +++ b/backend/internal/handler/dto/mappers.go @@ -250,6 +250,12 @@ func AccountFromServiceShallow(a *service.Account) *Account { } } + // 客户端亲和调度(适用于 Anthropic/Antigravity 平台) + if a.IsClientAffinityEnabled() { + enabled := true + out.ClientAffinityEnabled = &enabled + } + // 提取 API Key 账号配额限制(仅 apikey 类型有效) if a.Type == service.AccountTypeAPIKey { if limit := a.GetQuotaLimit(); limit > 0 { diff --git a/backend/internal/handler/dto/types.go b/backend/internal/handler/dto/types.go index f0d13d3f02..e5310a21f5 100644 --- a/backend/internal/handler/dto/types.go +++ b/backend/internal/handler/dto/types.go @@ -188,6 +188,10 @@ type Account struct { CacheTTLOverrideEnabled *bool `json:"cache_ttl_override_enabled,omitempty"` CacheTTLOverrideTarget *string `json:"cache_ttl_override_target,omitempty"` + // 客户端亲和调度(仅 Anthropic/Antigravity 平台有效) + // 启用后新会话会优先调度到客户端之前使用过的账号 + ClientAffinityEnabled *bool `json:"client_affinity_enabled,omitempty"` + // API Key 账号配额限制 QuotaLimit *float64 `json:"quota_limit,omitempty"` QuotaUsed *float64 `json:"quota_used,omitempty"` diff --git a/backend/internal/repository/gateway_cache.go b/backend/internal/repository/gateway_cache.go index 58291b6652..8fa6d0c401 100644 --- a/backend/internal/repository/gateway_cache.go +++ b/backend/internal/repository/gateway_cache.go @@ -3,13 +3,17 @@ package repository import ( "context" "fmt" + "strconv" "time" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/redis/go-redis/v9" ) -const stickySessionPrefix = "sticky_session:" +const ( + stickySessionPrefix = "sticky_session:" + clientAffinityPrefix = "client_affinity:" +) type gatewayCache struct { rdb *redis.Client @@ -51,3 +55,70 @@ func (c *gatewayCache) DeleteSessionAccountID(ctx context.Context, groupID int64 key := buildSessionKey(groupID, sessionHash) return c.rdb.Del(ctx, key).Err() } + +// buildAffinityKey 构建客户端亲和 key +// 格式: client_affinity:{groupID}:{clientID} +func buildAffinityKey(groupID int64, clientID string) string { + return fmt.Sprintf("%s%d:%s", clientAffinityPrefix, groupID, clientID) +} + +// getAffinityScript: 清理过期成员后返回亲和账号列表(按最近使用降序) +// KEYS[1] = client_affinity:{groupID}:{clientID} +// ARGV[1] = 过期阈值时间戳 (now - ttl) +var getAffinityScript = redis.NewScript(` +redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]) +return redis.call('ZREVRANGE', KEYS[1], 0, -1) +`) + +// updateAffinityScript: 清理过期成员、添加/更新亲和关系、刷新 TTL +// KEYS[1] = client_affinity:{groupID}:{clientID} +// ARGV[1] = 当前时间戳 +// ARGV[2] = TTL 秒数 +// ARGV[3] = accountID +// ARGV[4] = 过期阈值时间戳 (now - ttl) +var updateAffinityScript = redis.NewScript(` +redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[4]) +redis.call('ZADD', KEYS[1], ARGV[1], ARGV[3]) +redis.call('EXPIRE', KEYS[1], ARGV[2]) +return 1 +`) + +func (c *gatewayCache) GetClientAffinityAccounts(ctx context.Context, groupID int64, clientID string, ttl time.Duration) ([]int64, error) { + key := buildAffinityKey(groupID, clientID) + now := time.Now().Unix() + expireThreshold := now - int64(ttl.Seconds()) + + result, err := getAffinityScript.Run(ctx, c.rdb, []string{key}, expireThreshold).StringSlice() + if err != nil { + if err == redis.Nil { + return nil, nil + } + return nil, err + } + + accountIDs := make([]int64, 0, len(result)) + for _, s := range result { + id, err := strconv.ParseInt(s, 10, 64) + if err != nil { + continue + } + accountIDs = append(accountIDs, id) + } + return accountIDs, nil +} + +func (c *gatewayCache) UpdateClientAffinity(ctx context.Context, groupID int64, clientID string, accountID int64, ttl time.Duration) error { + key := buildAffinityKey(groupID, clientID) + now := time.Now().Unix() + ttlSeconds := int64(ttl.Seconds()) + expireThreshold := now - ttlSeconds + + return updateAffinityScript.Run(ctx, c.rdb, []string{key}, + now, ttlSeconds, accountID, expireThreshold, + ).Err() +} + +func (c *gatewayCache) RemoveClientAffinity(ctx context.Context, groupID int64, clientID string, accountID int64) error { + key := buildAffinityKey(groupID, clientID) + return c.rdb.ZRem(ctx, key, accountID).Err() +} diff --git a/backend/internal/service/account.go b/backend/internal/service/account.go index 8eb3748c69..390790a325 100644 --- a/backend/internal/service/account.go +++ b/backend/internal/service/account.go @@ -1099,6 +1099,24 @@ func (a *Account) IsSessionIDMaskingEnabled() bool { return false } +// IsClientAffinityEnabled 检查是否启用客户端亲和调度 +// 仅适用于 Anthropic 平台账号(包括 anthropic 和 antigravity 平台) +// 启用后,新会话会优先调度到之前使用过的账号 +func (a *Account) IsClientAffinityEnabled() bool { + if a.Platform != PlatformAnthropic && a.Platform != PlatformAntigravity { + return false + } + if a.Extra == nil { + return false + } + if v, ok := a.Extra["client_affinity_enabled"]; ok { + if enabled, ok := v.(bool); ok { + return enabled + } + } + return false +} + // IsCacheTTLOverrideEnabled 检查是否启用缓存 TTL 强制替换 // 仅适用于 Anthropic OAuth/SetupToken 类型账号 // 启用后将所有 cache creation tokens 归入指定的 TTL 类型(5m 或 1h) diff --git a/backend/internal/service/antigravity_smart_retry_test.go b/backend/internal/service/antigravity_smart_retry_test.go index 432c80e59f..7f836bf47c 100644 --- a/backend/internal/service/antigravity_smart_retry_test.go +++ b/backend/internal/service/antigravity_smart_retry_test.go @@ -9,6 +9,7 @@ import ( "net/http" "strings" "testing" + "time" "github.com/stretchr/testify/require" ) @@ -29,6 +30,15 @@ func (c *stubSmartRetryCache) DeleteSessionAccountID(_ context.Context, groupID c.deleteCalls = append(c.deleteCalls, deleteSessionCall{groupID: groupID, sessionHash: sessionHash}) return nil } +func (c *stubSmartRetryCache) GetClientAffinityAccounts(_ context.Context, _ int64, _ string, _ time.Duration) ([]int64, error) { + return nil, nil +} +func (c *stubSmartRetryCache) UpdateClientAffinity(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error { + return nil +} +func (c *stubSmartRetryCache) RemoveClientAffinity(_ context.Context, _ int64, _ string, _ int64) error { + return nil +} // mockSmartRetryUpstream 用于 handleSmartRetry 测试的 mock upstream type mockSmartRetryUpstream struct { diff --git a/backend/internal/service/gateway_hotpath_optimization_test.go b/backend/internal/service/gateway_hotpath_optimization_test.go index 161c4ba4b1..137490e8db 100644 --- a/backend/internal/service/gateway_hotpath_optimization_test.go +++ b/backend/internal/service/gateway_hotpath_optimization_test.go @@ -143,6 +143,15 @@ func (s *stickyGatewayCacheHotpathStub) RefreshSessionTTL(ctx context.Context, g func (s *stickyGatewayCacheHotpathStub) DeleteSessionAccountID(ctx context.Context, groupID int64, sessionHash string) error { return nil } +func (s *stickyGatewayCacheHotpathStub) GetClientAffinityAccounts(_ context.Context, _ int64, _ string, _ time.Duration) ([]int64, error) { + return nil, nil +} +func (s *stickyGatewayCacheHotpathStub) UpdateClientAffinity(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error { + return nil +} +func (s *stickyGatewayCacheHotpathStub) RemoveClientAffinity(_ context.Context, _ int64, _ string, _ int64) error { + return nil +} func (s *modelsListAccountRepoStub) ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]Account, error) { s.listByGroupCalls.Add(1) diff --git a/backend/internal/service/gateway_multiplatform_test.go b/backend/internal/service/gateway_multiplatform_test.go index 320ceaa708..466ff294bb 100644 --- a/backend/internal/service/gateway_multiplatform_test.go +++ b/backend/internal/service/gateway_multiplatform_test.go @@ -235,6 +235,16 @@ func (m *mockGatewayCacheForPlatform) DeleteSessionAccountID(ctx context.Context return nil } +func (m *mockGatewayCacheForPlatform) GetClientAffinityAccounts(_ context.Context, _ int64, _ string, _ time.Duration) ([]int64, error) { + return nil, nil +} +func (m *mockGatewayCacheForPlatform) UpdateClientAffinity(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error { + return nil +} +func (m *mockGatewayCacheForPlatform) RemoveClientAffinity(_ context.Context, _ int64, _ string, _ int64) error { + return nil +} + type mockGroupRepoForGateway struct { groups map[int64]*Group getByIDCalls int diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index ccccdf4d10..0fe2545ec0 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -304,6 +304,10 @@ var ( sessionIDRegex = regexp.MustCompile(`session_([a-f0-9-]{36})`) claudeCliUserAgentRe = regexp.MustCompile(`^claude-cli/\d+\.\d+\.\d+`) + // clientIDFromMetadataRegex 从 metadata.user_id 中提取客户端 ID(64位 hex) + // 格式: user_{64位hex}_account_... + clientIDFromMetadataRegex = regexp.MustCompile(`^user_([a-f0-9]{64})_account_`) + // claudeCodePromptPrefixes 用于检测 Claude Code 系统提示词的前缀列表 // 支持多种变体:标准版、Agent SDK 版、Explore Agent 版、Compact 版等 // 注意:前缀之间不应存在包含关系,否则会导致冗余匹配 @@ -364,6 +368,13 @@ type GatewayCache interface { // DeleteSessionAccountID 删除粘性会话绑定,用于账号不可用时主动清理 // Delete sticky session binding, used to proactively clean up when account becomes unavailable DeleteSessionAccountID(ctx context.Context, groupID int64, sessionHash string) error + + // GetClientAffinityAccounts 获取客户端亲和账号列表(按最近使用降序),同时清理过期成员 + GetClientAffinityAccounts(ctx context.Context, groupID int64, clientID string, ttl time.Duration) ([]int64, error) + // UpdateClientAffinity 添加/更新客户端亲和关系(更新 score 为当前时间戳,刷新 key TTL) + UpdateClientAffinity(ctx context.Context, groupID int64, clientID string, accountID int64, ttl time.Duration) error + // RemoveClientAffinity 删除单个客户端亲和关系 + RemoveClientAffinity(ctx context.Context, groupID int64, clientID string, accountID int64) error } // derefGroupID safely dereferences *int64 to int64, returning 0 if nil @@ -434,6 +445,20 @@ func shouldClearStickySession(account *Account, requestedModel string) bool { return false } +// extractClientIDFromMetadata 从 metadata.user_id 中提取客户端 ID(64位 hex)。 +// 格式: user_{64位hex}_account_..._session_... +// 返回空字符串表示无法提取(非 Claude Code/Console 客户端)。 +func extractClientIDFromMetadata(metadataUserID string) string { + if metadataUserID == "" { + return "" + } + matches := clientIDFromMetadataRegex.FindStringSubmatch(metadataUserID) + if matches == nil { + return "" + } + return matches[1] +} + type AccountWaitPlan struct { AccountID int64 MaxConcurrency int @@ -1046,7 +1071,7 @@ func (s *GatewayService) SelectAccountForModelWithExclusions(ctx context.Context } // SelectAccountWithLoadAwareness selects account with load-awareness and wait plan. -// metadataUserID: 已废弃参数,会话限制现在统一使用 sessionHash +// metadataUserID: 用于客户端亲和调度,从中提取客户端 ID func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}, metadataUserID string) (*AccountSelectionResult, error) { // 调试日志:记录调度入口参数 excludedIDsList := make([]int64, 0, len(excludedIDs)) @@ -1077,6 +1102,9 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro } } + // 提取客户端 ID(用于客户端亲和调度) + affinityClientID := extractClientIDFromMetadata(metadataUserID) + if s.debugModelRoutingEnabled() && requestedModel != "" { groupPlatform := "" if group != nil { @@ -1374,6 +1402,9 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro if sessionHash != "" && s.cache != nil { _ = s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), sessionHash, item.account.ID, stickySessionTTL) } + if affinityClientID != "" && s.cache != nil && item.account.IsClientAffinityEnabled() { + _ = s.cache.UpdateClientAffinity(ctx, derefGroupID(groupID), affinityClientID, item.account.ID, stickySessionTTL) + } if s.debugModelRoutingEnabled() { logger.LegacyPrintf("service.gateway", "[ModelRoutingDebug] routed select: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), item.account.ID) } @@ -1470,6 +1501,66 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro } } + // ============ Layer 1.6: 客户端亲和(仅在粘性会话未命中时生效) ============ + if affinityClientID != "" && s.cache != nil && stickyAccountID <= 0 { + affinityAccountIDs, err := s.cache.GetClientAffinityAccounts(ctx, derefGroupID(groupID), affinityClientID, stickySessionTTL) + if err == nil && len(affinityAccountIDs) > 0 { + for _, affinityAccID := range affinityAccountIDs { + if isExcluded(affinityAccID) { + continue + } + account, ok := accountByID[affinityAccID] + if !ok || !s.isAccountSchedulableForSelection(account) { + continue + } + if !account.IsClientAffinityEnabled() { + continue + } + if !s.isAccountAllowedForPlatform(account, platform, useMixed) { + continue + } + if requestedModel != "" && !s.isModelSupportedByAccountWithContext(ctx, account, requestedModel) { + continue + } + if !s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) { + continue + } + if !s.isAccountSchedulableForQuota(account) { + continue + } + if !s.isAccountSchedulableForWindowCost(ctx, account, false) { + continue + } + if !s.isAccountSchedulableForRPM(ctx, account, false) { + continue + } + + result, err := s.tryAcquireAccountSlot(ctx, affinityAccID, account.Concurrency) + if err == nil && result.Acquired { + if !s.checkAndRegisterSession(ctx, account, sessionHash) { + result.ReleaseFunc() + continue + } + // 亲和命中:更新亲和 score + 绑定粘性会话 + _ = s.cache.UpdateClientAffinity(ctx, derefGroupID(groupID), affinityClientID, affinityAccID, stickySessionTTL) + if sessionHash != "" { + _ = s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), sessionHash, affinityAccID, stickySessionTTL) + } + slog.Debug("client_affinity_hit", + "group_id", derefGroupID(groupID), + "client_id", affinityClientID[:8]+"...", + "account_id", affinityAccID) + return &AccountSelectionResult{ + Account: account, + Acquired: true, + ReleaseFunc: result.ReleaseFunc, + }, nil + } + } + // 所有亲和账号不可用,继续到 Layer 2 + } + } + // ============ Layer 2: 负载感知选择 ============ candidates := make([]*Account, 0, len(accounts)) for i := range accounts { @@ -1522,6 +1613,9 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro loadMap, err := s.concurrencyService.GetAccountsLoadBatch(ctx, accountLoads) if err != nil { if result, ok := s.tryAcquireByLegacyOrder(ctx, candidates, groupID, sessionHash, preferOAuth); ok { + if affinityClientID != "" && s.cache != nil && result.Account != nil && result.Account.IsClientAffinityEnabled() { + _ = s.cache.UpdateClientAffinity(ctx, derefGroupID(groupID), affinityClientID, result.Account.ID, stickySessionTTL) + } return result, nil } } else { @@ -1560,6 +1654,10 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro if sessionHash != "" && s.cache != nil { _ = s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), sessionHash, selected.account.ID, stickySessionTTL) } + // 更新客户端亲和关系 + if affinityClientID != "" && s.cache != nil && selected.account.IsClientAffinityEnabled() { + _ = s.cache.UpdateClientAffinity(ctx, derefGroupID(groupID), affinityClientID, selected.account.ID, stickySessionTTL) + } return &AccountSelectionResult{ Account: selected.account, Acquired: true, diff --git a/backend/internal/service/gemini_multiplatform_test.go b/backend/internal/service/gemini_multiplatform_test.go index b0b804eb00..56f94d5f1f 100644 --- a/backend/internal/service/gemini_multiplatform_test.go +++ b/backend/internal/service/gemini_multiplatform_test.go @@ -288,6 +288,16 @@ func (m *mockGatewayCacheForGemini) DeleteSessionAccountID(ctx context.Context, return nil } +func (m *mockGatewayCacheForGemini) GetClientAffinityAccounts(_ context.Context, _ int64, _ string, _ time.Duration) ([]int64, error) { + return nil, nil +} +func (m *mockGatewayCacheForGemini) UpdateClientAffinity(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error { + return nil +} +func (m *mockGatewayCacheForGemini) RemoveClientAffinity(_ context.Context, _ int64, _ string, _ int64) error { + return nil +} + // TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform 测试 Gemini 单平台选择 func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform(t *testing.T) { ctx := context.Background() diff --git a/backend/internal/service/openai_gateway_service_test.go b/backend/internal/service/openai_gateway_service_test.go index 4f5f7f3c12..7e53c8996e 100644 --- a/backend/internal/service/openai_gateway_service_test.go +++ b/backend/internal/service/openai_gateway_service_test.go @@ -265,6 +265,16 @@ func (c *stubGatewayCache) DeleteSessionAccountID(ctx context.Context, groupID i return nil } +func (c *stubGatewayCache) GetClientAffinityAccounts(_ context.Context, _ int64, _ string, _ time.Duration) ([]int64, error) { + return nil, nil +} +func (c *stubGatewayCache) UpdateClientAffinity(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error { + return nil +} +func (c *stubGatewayCache) RemoveClientAffinity(_ context.Context, _ int64, _ string, _ int64) error { + return nil +} + func TestOpenAISelectAccountWithLoadAwareness_FiltersUnschedulable(t *testing.T) { now := time.Now() resetAt := now.Add(10 * time.Minute) diff --git a/backend/internal/service/openai_ws_state_store_test.go b/backend/internal/service/openai_ws_state_store_test.go index 235d42331d..6f77f64356 100644 --- a/backend/internal/service/openai_ws_state_store_test.go +++ b/backend/internal/service/openai_ws_state_store_test.go @@ -193,6 +193,16 @@ func (c *openAIWSStateStoreTimeoutProbeCache) DeleteSessionAccountID(ctx context return nil } +func (c *openAIWSStateStoreTimeoutProbeCache) GetClientAffinityAccounts(_ context.Context, _ int64, _ string, _ time.Duration) ([]int64, error) { + return nil, nil +} +func (c *openAIWSStateStoreTimeoutProbeCache) UpdateClientAffinity(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error { + return nil +} +func (c *openAIWSStateStoreTimeoutProbeCache) RemoveClientAffinity(_ context.Context, _ int64, _ string, _ int64) error { + return nil +} + func TestOpenAIWSStateStore_RedisOpsUseShortTimeout(t *testing.T) { probe := &openAIWSStateStoreTimeoutProbeCache{} store := NewOpenAIWSStateStore(probe) diff --git a/backend/internal/testutil/stubs.go b/backend/internal/testutil/stubs.go index 217a5f5699..669372ef51 100644 --- a/backend/internal/testutil/stubs.go +++ b/backend/internal/testutil/stubs.go @@ -97,6 +97,15 @@ func (c StubGatewayCache) RefreshSessionTTL(_ context.Context, _ int64, _ string func (c StubGatewayCache) DeleteSessionAccountID(_ context.Context, _ int64, _ string) error { return nil } +func (c StubGatewayCache) GetClientAffinityAccounts(_ context.Context, _ int64, _ string, _ time.Duration) ([]int64, error) { + return nil, nil +} +func (c StubGatewayCache) UpdateClientAffinity(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error { + return nil +} +func (c StubGatewayCache) RemoveClientAffinity(_ context.Context, _ int64, _ string, _ int64) error { + return nil +} // ============================================================ // StubSessionLimitCache — service.SessionLimitCache 的空实现 diff --git a/frontend/src/components/account/EditAccountModal.vue b/frontend/src/components/account/EditAccountModal.vue index be7d2d45ce..6026eb9570 100644 --- a/frontend/src/components/account/EditAccountModal.vue +++ b/frontend/src/components/account/EditAccountModal.vue @@ -1165,6 +1165,38 @@ + +
+ {{ t('admin.accounts.quotaControl.clientAffinity.hint') }} +
+