feat: add affinity scheduling controls

This commit is contained in:
erio
2026-03-18 04:08:45 +08:00
parent 338e3df8d9
commit 44ec2626cf
44 changed files with 3711 additions and 496 deletions
+1
View File
@@ -128,6 +128,7 @@ deploy/docker-compose.override.yml
.gocache/
vite.config.js
docs/*
!docs/ACCOUNT_SCHEDULING_FLOW.md
.serena/
# ===================
+1 -1
View File
@@ -1 +1 @@
0.1.101.3
0.1.101.4
@@ -0,0 +1,254 @@
package admin
import (
"context"
"log/slog"
"strconv"
"strings"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// 账号亲和调度 API 与契约说明见 docs/ACCOUNT_SCHEDULING_FLOW.md 。
// GetAffinityClients returns the list of affinity clients for an account with last active timestamps.
// GET /api/v1/admin/accounts/:id/affinity-clients
func (h *AccountHandler) GetAffinityClients(c *gin.Context) {
accountID, err := strconv.ParseInt(c.Param("id"), 10, 64)
if err != nil {
response.BadRequest(c, "Invalid account ID")
return
}
account, err := h.adminService.GetAccount(c.Request.Context(), accountID)
if err != nil {
response.ErrorFrom(c, err)
return
}
if !account.IsAffinityEnabled() {
response.Success(c, []service.AffinityClient{})
return
}
if h.gatewayCache == nil || len(account.GroupIDs) == 0 {
response.Success(c, []service.AffinityClient{})
return
}
clients, err := h.gatewayCache.GetAccountAffinityClientsWithScores(
c.Request.Context(), accountID, account.GroupIDs, service.ClientAffinityTTL,
)
if err != nil {
response.Success(c, []service.AffinityClient{})
return
}
response.Success(c, clients)
}
// clearAccountAffinity 清除指定账号在所有分组的亲和记录。
func (h *AccountHandler) clearAccountAffinity(ctx context.Context, accountID int64, groupIDs []int64) {
if h.gatewayCache == nil || len(groupIDs) == 0 {
return
}
if err := h.gatewayCache.ClearAccountAffinity(ctx, accountID, groupIDs); err != nil {
// 清理失败不影响主流程,记录日志即可
slog.Warn("clear account affinity failed",
"account_id", accountID,
"error", err,
)
}
}
// countUniqueUsersFromAffinityMembers 从 "{userID}/{clientID}" 格式的成员列表中计算唯一用户数。
func countUniqueUsersFromAffinityMembers(members []string) int64 {
users := make(map[string]struct{}, len(members))
for _, m := range members {
if idx := strings.Index(m, "/"); idx > 0 {
users[m[:idx]] = struct{}{}
}
}
return int64(len(users))
}
// AffinityDetailClient 亲和详情中的单个客户端信息
type AffinityDetailClient struct {
ClientID string `json:"client_id"`
LastActive time.Time `json:"last_active"`
IsPinned bool `json:"is_pinned"`
}
// AffinityDetailUser 亲和详情中按用户分组的信息
type AffinityDetailUser struct {
UserID int64 `json:"user_id"`
UserEmail string `json:"user_email"`
ClientCount int `json:"client_count"`
IsPinned bool `json:"is_pinned"`
Clients []AffinityDetailClient `json:"clients"`
}
// AffinityDetailsResponse 亲和详情响应
type AffinityDetailsResponse struct {
Users []AffinityDetailUser `json:"users"`
TotalUsers int `json:"total_users"`
TotalClients int `json:"total_clients"`
PinnedUsers []int64 `json:"pinned_users"`
}
type affinityStateSnapshot struct {
enabled bool
groupIDs []int64
}
func (h *AccountHandler) captureAffinityStates(ctx context.Context, accountIDs []int64) map[int64]affinityStateSnapshot {
states := make(map[int64]affinityStateSnapshot, len(accountIDs))
if h.gatewayCache == nil || len(accountIDs) == 0 {
return states
}
for _, accountID := range accountIDs {
account, err := h.adminService.GetAccount(ctx, accountID)
if err != nil || account == nil {
continue
}
states[accountID] = affinityStateSnapshot{
enabled: account.IsAffinityEnabled(),
groupIDs: append([]int64(nil), account.GroupIDs...),
}
}
return states
}
func (h *AccountHandler) clearAffinityCacheIfDisabled(ctx context.Context, accountID int64, oldState affinityStateSnapshot) {
if h.gatewayCache == nil || !oldState.enabled {
return
}
account, err := h.adminService.GetAccount(ctx, accountID)
if err != nil || account == nil || account.IsAffinityEnabled() {
return
}
groupIDs := oldState.groupIDs
if len(account.GroupIDs) > 0 {
groupIDs = mergeGroupIDs(groupIDs, account.GroupIDs)
}
h.clearAccountAffinity(ctx, accountID, groupIDs)
}
func (h *AccountHandler) clearAffinityCacheForBulkIfDisabled(ctx context.Context, accountIDs []int64, oldStates map[int64]affinityStateSnapshot) {
for _, accountID := range accountIDs {
oldState, ok := oldStates[accountID]
if !ok {
continue
}
h.clearAffinityCacheIfDisabled(ctx, accountID, oldState)
}
}
// GetAffinityDetails returns the affinity details grouped by user for an account.
// GET /api/v1/admin/accounts/:id/affinity-details
func (h *AccountHandler) GetAffinityDetails(c *gin.Context) {
accountID, err := strconv.ParseInt(c.Param("id"), 10, 64)
if err != nil {
response.BadRequest(c, "Invalid account ID")
return
}
emptyResp := AffinityDetailsResponse{
Users: []AffinityDetailUser{},
TotalUsers: 0,
TotalClients: 0,
PinnedUsers: []int64{},
}
account, err := h.adminService.GetAccount(c.Request.Context(), accountID)
if err != nil {
response.ErrorFrom(c, err)
return
}
if !account.IsAffinityEnabled() {
response.Success(c, emptyResp)
return
}
pinnedUsers := account.GetPinnedUsers()
if pinnedUsers == nil {
pinnedUsers = []int64{}
}
if h.gatewayCache == nil || len(account.GroupIDs) == 0 {
emptyResp.PinnedUsers = pinnedUsers
response.Success(c, emptyResp)
return
}
clients, err := h.gatewayCache.GetAccountAffinityClientsWithScores(
c.Request.Context(), accountID, account.GroupIDs, service.ClientAffinityTTL,
)
if err != nil {
emptyResp.PinnedUsers = pinnedUsers
response.Success(c, emptyResp)
return
}
pinnedSet := make(map[int64]struct{}, len(pinnedUsers))
for _, uid := range pinnedUsers {
pinnedSet[uid] = struct{}{}
}
// 按 UserID 分组
userMap := make(map[int64]*AffinityDetailUser)
var userOrder []int64
for _, cl := range clients {
u, ok := userMap[cl.UserID]
if !ok {
_, pinned := pinnedSet[cl.UserID]
u = &AffinityDetailUser{
UserID: cl.UserID,
UserEmail: "",
ClientCount: 0,
IsPinned: pinned,
Clients: []AffinityDetailClient{},
}
userMap[cl.UserID] = u
userOrder = append(userOrder, cl.UserID)
}
u.Clients = append(u.Clients, AffinityDetailClient{
ClientID: cl.ClientID,
LastActive: cl.LastActive,
IsPinned: false, // 客户端级别暂无 pinned 概念
})
}
// 关联用户邮箱(查询失败时保持空字符串,不影响主流程)
for _, uid := range userOrder {
user, uErr := h.adminService.GetUser(c.Request.Context(), uid)
if uErr != nil || user == nil {
continue
}
if u, ok := userMap[uid]; ok {
if user.Email != "" {
u.UserEmail = user.Email
} else {
u.UserEmail = user.Username
}
}
}
users := make([]AffinityDetailUser, 0, len(userOrder))
for _, uid := range userOrder {
u := userMap[uid]
u.ClientCount = len(u.Clients)
users = append(users, *u)
}
response.Success(c, AffinityDetailsResponse{
Users: users,
TotalUsers: len(users),
TotalClients: len(clients),
PinnedUsers: pinnedUsers,
})
}
@@ -9,7 +9,6 @@ import (
"errors"
"fmt"
"log"
"log/slog"
"net/http"
"strconv"
"strings"
@@ -211,7 +210,7 @@ func (h *AccountHandler) buildAccountResponseWithRuntime(ctx context.Context, ac
}
// 亲和客户端数据(启用亲和的账号始终返回 count,即使为 0)
if account.IsClientAffinityEnabled() {
if account.IsAffinityEnabled() {
if h.gatewayCache != nil && len(account.GroupIDs) > 0 {
accountGroups := map[int64][]int64{account.ID: account.GroupIDs}
if clients, err := h.gatewayCache.GetAccountAffinityClientsBatch(ctx, accountGroups, service.ClientAffinityTTL); err == nil {
@@ -219,17 +218,22 @@ func (h *AccountHandler) buildAccountResponseWithRuntime(ctx context.Context, ac
count := int64(len(cl))
item.AffinityClientCount = &count
item.AffinityClients = cl
userCount := countUniqueUsersFromAffinityMembers(cl)
item.AffinityUserCount = &userCount
} else {
zero := int64(0)
item.AffinityClientCount = &zero
item.AffinityUserCount = &zero
}
} else {
zero := int64(0)
item.AffinityClientCount = &zero
item.AffinityUserCount = &zero
}
} else {
zero := int64(0)
item.AffinityClientCount = &zero
item.AffinityUserCount = &zero
}
}
@@ -351,7 +355,7 @@ func (h *AccountHandler) List(c *gin.Context) {
accountGroups := make(map[int64][]int64)
for i := range accounts {
acc := &accounts[i]
if acc.IsClientAffinityEnabled() && len(acc.GroupIDs) > 0 {
if acc.IsAffinityEnabled() && len(acc.GroupIDs) > 0 {
accountGroups[acc.ID] = acc.GroupIDs
}
}
@@ -391,14 +395,18 @@ func (h *AccountHandler) List(c *gin.Context) {
}
// 注入亲和客户端数据到 DTO(启用亲和的账号始终返回 count,即使为 0)
if acc.IsClientAffinityEnabled() {
if acc.IsAffinityEnabled() {
if clients, ok := affinityClients[acc.ID]; ok && len(clients) > 0 {
count := int64(len(clients))
item.AffinityClientCount = &count
item.AffinityClients = clients
// 从成员列表中解析唯一用户数
userCount := countUniqueUsersFromAffinityMembers(clients)
item.AffinityUserCount = &userCount
} else {
zero := int64(0)
item.AffinityClientCount = &zero
item.AffinityUserCount = &zero
}
}
@@ -622,15 +630,7 @@ func (h *AccountHandler) Update(c *gin.Context) {
// base_rpm 输入校验:负值归零,超过 10000 截断
sanitizeExtraBaseRPM(req.Extra)
// 记录更新前的亲和状态,用于检测亲和关闭时清理 Redis 记录
oldAffinityEnabled := false
var oldGroupIDs []int64
if len(req.Extra) > 0 && h.gatewayCache != nil {
if oldAccount, err := h.adminService.GetAccount(c.Request.Context(), accountID); err == nil {
oldAffinityEnabled = oldAccount.IsClientAffinityEnabled()
oldGroupIDs = oldAccount.GroupIDs
}
}
oldStates := h.captureAffinityStates(c.Request.Context(), []int64{accountID})
// 确定是否跳过混合渠道检查
skipCheck := req.ConfirmMixedChannelRisk != nil && *req.ConfirmMixedChannelRisk
@@ -668,14 +668,7 @@ func (h *AccountHandler) Update(c *gin.Context) {
return
}
// 亲和关闭时清理 Redis 中的亲和记录
if oldAffinityEnabled && !account.IsClientAffinityEnabled() {
groupIDs := oldGroupIDs
if len(account.GroupIDs) > 0 {
groupIDs = mergeGroupIDs(oldGroupIDs, account.GroupIDs)
}
h.clearAccountAffinity(c.Request.Context(), accountID, groupIDs)
}
h.clearAffinityCacheForBulkIfDisabled(c.Request.Context(), []int64{accountID}, oldStates)
response.Success(c, h.buildAccountResponseWithRuntime(c.Request.Context(), account))
}
@@ -1393,6 +1386,8 @@ func (h *AccountHandler) BulkUpdate(c *gin.Context) {
return
}
oldStates := h.captureAffinityStates(c.Request.Context(), req.AccountIDs)
result, err := h.adminService.BulkUpdateAccounts(c.Request.Context(), &service.BulkUpdateAccountsInput{
AccountIDs: req.AccountIDs,
Name: req.Name,
@@ -1427,6 +1422,8 @@ func (h *AccountHandler) BulkUpdate(c *gin.Context) {
return
}
h.clearAffinityCacheForBulkIfDisabled(c.Request.Context(), req.AccountIDs, oldStates)
response.Success(c, result)
}
@@ -1639,56 +1636,6 @@ func (h *AccountHandler) ResetQuota(c *gin.Context) {
response.Success(c, h.buildAccountResponseWithRuntime(c.Request.Context(), account))
}
// GetAffinityClients returns the list of affinity clients for an account with last active timestamps.
// GET /api/v1/admin/accounts/:id/affinity-clients
func (h *AccountHandler) GetAffinityClients(c *gin.Context) {
accountID, err := strconv.ParseInt(c.Param("id"), 10, 64)
if err != nil {
response.BadRequest(c, "Invalid account ID")
return
}
account, err := h.adminService.GetAccount(c.Request.Context(), accountID)
if err != nil {
response.ErrorFrom(c, err)
return
}
if !account.IsClientAffinityEnabled() {
response.Success(c, []service.AffinityClient{})
return
}
if h.gatewayCache == nil || len(account.GroupIDs) == 0 {
response.Success(c, []service.AffinityClient{})
return
}
clients, err := h.gatewayCache.GetAccountAffinityClientsWithScores(
c.Request.Context(), accountID, account.GroupIDs, service.ClientAffinityTTL,
)
if err != nil {
response.Success(c, []service.AffinityClient{})
return
}
response.Success(c, clients)
}
// clearAccountAffinity 清除指定账号在所有分组的亲和记录。
func (h *AccountHandler) clearAccountAffinity(ctx context.Context, accountID int64, groupIDs []int64) {
if h.gatewayCache == nil || len(groupIDs) == 0 {
return
}
if err := h.gatewayCache.ClearAccountAffinity(ctx, accountID, groupIDs); err != nil {
// 清理失败不影响主流程,记录日志即可
slog.Warn("clear account affinity failed",
"account_id", accountID,
"error", err,
)
}
}
// mergeGroupIDs 合并两个 groupID 切片并去重。
func mergeGroupIDs(a, b []int64) []int64 {
seen := make(map[int64]struct{}, len(a)+len(b))
@@ -0,0 +1,298 @@
package admin
import (
"bytes"
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
type affinityCacheAdminService struct {
*stubAdminService
accounts map[int64]*service.Account
bulkResultIDs []int64
bulkFailedCount int
}
func (s *affinityCacheAdminService) GetAccount(_ context.Context, id int64) (*service.Account, error) {
if acc, ok := s.accounts[id]; ok {
cloned := *acc
return &cloned, nil
}
return s.stubAdminService.GetAccount(context.Background(), id)
}
func (s *affinityCacheAdminService) UpdateAccount(_ context.Context, id int64, input *service.UpdateAccountInput) (*service.Account, error) {
acc := s.accounts[id]
if acc == nil {
return s.stubAdminService.UpdateAccount(context.Background(), id, input)
}
if input.Extra != nil {
acc.Extra = input.Extra
}
if input.GroupIDs != nil {
acc.GroupIDs = *input.GroupIDs
}
cloned := *acc
return &cloned, nil
}
func (s *affinityCacheAdminService) BulkUpdateAccounts(_ context.Context, input *service.BulkUpdateAccountsInput) (*service.BulkUpdateAccountsResult, error) {
for _, accountID := range input.AccountIDs {
acc := s.accounts[accountID]
if acc == nil {
continue
}
if input.Extra != nil {
acc.Extra = input.Extra
}
if input.GroupIDs != nil {
acc.GroupIDs = *input.GroupIDs
}
}
successIDs := append([]int64(nil), input.AccountIDs...)
failed := 0
if len(s.bulkResultIDs) > 0 {
successIDs = append([]int64(nil), s.bulkResultIDs...)
failed = s.bulkFailedCount
}
return &service.BulkUpdateAccountsResult{
Success: len(successIDs),
Failed: failed,
SuccessIDs: successIDs,
Results: []service.BulkUpdateAccountResult{},
}, nil
}
type affinityCacheClearRecorder struct {
clearCalls []struct {
accountID int64
groupIDs []int64
}
}
func (r *affinityCacheClearRecorder) GetSessionAccountID(_ context.Context, _ int64, _ string) (int64, error) {
return 0, nil
}
func (r *affinityCacheClearRecorder) SetSessionAccountID(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error {
return nil
}
func (r *affinityCacheClearRecorder) RefreshSessionTTL(_ context.Context, _ int64, _ string, _ time.Duration) error {
return nil
}
func (r *affinityCacheClearRecorder) DeleteSessionAccountID(_ context.Context, _ int64, _ string) error {
return nil
}
func (r *affinityCacheClearRecorder) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
return nil, nil
}
func (r *affinityCacheClearRecorder) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error {
return nil
}
func (r *affinityCacheClearRecorder) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) {
return map[int64]int64{}, nil
}
func (r *affinityCacheClearRecorder) GetAccountAffinityClientsBatch(_ context.Context, _ map[int64][]int64, _ time.Duration) (map[int64][]string, error) {
return map[int64][]string{}, nil
}
func (r *affinityCacheClearRecorder) GetAccountAffinityClientsWithScores(_ context.Context, _ int64, _ []int64, _ time.Duration) ([]service.AffinityClient, error) {
return nil, nil
}
func (r *affinityCacheClearRecorder) ClearAccountAffinity(_ context.Context, accountID int64, groupIDs []int64) error {
r.clearCalls = append(r.clearCalls, struct {
accountID int64
groupIDs []int64
}{
accountID: accountID,
groupIDs: append([]int64(nil), groupIDs...),
})
return nil
}
func (r *affinityCacheClearRecorder) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (users, clients, perUser int64, err error) {
return 0, 0, 0, nil
}
func setupAffinityCacheRouter(adminSvc service.AdminService, gatewayCache service.GatewayCache) *gin.Engine {
gin.SetMode(gin.TestMode)
router := gin.New()
handler := NewAccountHandler(adminSvc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, gatewayCache)
router.PUT("/api/v1/admin/accounts/:id", handler.Update)
router.POST("/api/v1/admin/accounts/bulk-update", handler.BulkUpdate)
return router
}
func TestAccountHandlerUpdate_DisablingAffinityClearsCache(t *testing.T) {
adminSvc := &affinityCacheAdminService{
stubAdminService: newStubAdminService(),
accounts: map[int64]*service.Account{
3: {
ID: 3,
Platform: service.PlatformAnthropic,
Type: service.AccountTypeOAuth,
Status: service.StatusActive,
GroupIDs: []int64{1, 2},
Extra: map[string]any{
"client_affinity_enabled": true,
},
},
},
}
cache := &affinityCacheClearRecorder{}
router := setupAffinityCacheRouter(adminSvc, cache)
body, _ := json.Marshal(map[string]any{
"extra": map[string]any{},
})
rec := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPut, "/api/v1/admin/accounts/3", bytes.NewReader(body))
req.Header.Set("Content-Type", "application/json")
router.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
require.Len(t, cache.clearCalls, 1)
require.Equal(t, int64(3), cache.clearCalls[0].accountID)
require.Equal(t, []int64{1, 2}, cache.clearCalls[0].groupIDs)
}
func TestAccountHandlerBulkUpdate_DisablingAffinityClearsCache(t *testing.T) {
adminSvc := &affinityCacheAdminService{
stubAdminService: newStubAdminService(),
accounts: map[int64]*service.Account{
1: {
ID: 1,
Platform: service.PlatformAnthropic,
Type: service.AccountTypeOAuth,
Status: service.StatusActive,
GroupIDs: []int64{11},
Extra: map[string]any{
"client_affinity_enabled": true,
},
},
2: {
ID: 2,
Platform: service.PlatformAnthropic,
Type: service.AccountTypeSetupToken,
Status: service.StatusActive,
GroupIDs: []int64{22},
Extra: map[string]any{
"client_affinity_enabled": true,
},
},
},
}
cache := &affinityCacheClearRecorder{}
router := setupAffinityCacheRouter(adminSvc, cache)
body, _ := json.Marshal(map[string]any{
"account_ids": []int64{1, 2},
"extra": map[string]any{
"client_affinity_enabled": false,
},
})
rec := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/bulk-update", bytes.NewReader(body))
req.Header.Set("Content-Type", "application/json")
router.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
require.Len(t, cache.clearCalls, 2)
require.Equal(t, int64(1), cache.clearCalls[0].accountID)
require.Equal(t, []int64{11}, cache.clearCalls[0].groupIDs)
require.Equal(t, int64(2), cache.clearCalls[1].accountID)
require.Equal(t, []int64{22}, cache.clearCalls[1].groupIDs)
}
func TestAccountHandlerBulkUpdate_DisablingAffinityWithNewFlagClearsCache(t *testing.T) {
adminSvc := &affinityCacheAdminService{
stubAdminService: newStubAdminService(),
accounts: map[int64]*service.Account{
1: {
ID: 1,
Platform: service.PlatformAnthropic,
Type: service.AccountTypeOAuth,
Status: service.StatusActive,
GroupIDs: []int64{33},
Extra: map[string]any{
"affinity_enabled": true,
},
},
},
}
cache := &affinityCacheClearRecorder{}
router := setupAffinityCacheRouter(adminSvc, cache)
body, _ := json.Marshal(map[string]any{
"account_ids": []int64{1},
"extra": map[string]any{
"affinity_enabled": false,
},
})
rec := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/bulk-update", bytes.NewReader(body))
req.Header.Set("Content-Type", "application/json")
router.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
require.Len(t, cache.clearCalls, 1)
require.Equal(t, int64(1), cache.clearCalls[0].accountID)
require.Equal(t, []int64{33}, cache.clearCalls[0].groupIDs)
}
func TestAccountHandlerBulkUpdate_PartialFailureStillClearsDisabledAffinityCache(t *testing.T) {
adminSvc := &affinityCacheAdminService{
stubAdminService: newStubAdminService(),
accounts: map[int64]*service.Account{
1: {
ID: 1,
Platform: service.PlatformAnthropic,
Type: service.AccountTypeOAuth,
Status: service.StatusActive,
GroupIDs: []int64{44},
Extra: map[string]any{
"affinity_enabled": true,
},
},
2: {
ID: 2,
Platform: service.PlatformAnthropic,
Type: service.AccountTypeSetupToken,
Status: service.StatusActive,
GroupIDs: []int64{55},
Extra: map[string]any{
"client_affinity_enabled": true,
},
},
},
bulkResultIDs: []int64{1},
bulkFailedCount: 1,
}
cache := &affinityCacheClearRecorder{}
router := setupAffinityCacheRouter(adminSvc, cache)
body, _ := json.Marshal(map[string]any{
"account_ids": []int64{1, 2},
"group_ids": []int64{101},
"extra": map[string]any{
"affinity_enabled": false,
"client_affinity_enabled": false,
},
})
rec := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/bulk-update", bytes.NewReader(body))
req.Header.Set("Content-Type", "application/json")
router.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
require.Len(t, cache.clearCalls, 2)
require.Equal(t, int64(1), cache.clearCalls[0].accountID)
require.Equal(t, []int64{44, 101}, cache.clearCalls[0].groupIDs)
require.Equal(t, int64(2), cache.clearCalls[1].accountID)
require.Equal(t, []int64{55, 101}, cache.clearCalls[1].groupIDs)
}
@@ -0,0 +1,205 @@
package admin
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
type affinityDetailsAdminService struct {
*stubAdminService
account service.Account
usersByID map[int64]service.User
}
func (s *affinityDetailsAdminService) GetAccount(_ context.Context, id int64) (*service.Account, error) {
if s.account.ID == id {
acc := s.account
return &acc, nil
}
return s.stubAdminService.GetAccount(context.Background(), id)
}
func (s *affinityDetailsAdminService) GetUser(_ context.Context, id int64) (*service.User, error) {
if u, ok := s.usersByID[id]; ok {
user := u
return &user, nil
}
return s.stubAdminService.GetUser(context.Background(), id)
}
type affinityDetailsGatewayCacheStub struct {
clients []service.AffinityClient
}
func (s *affinityDetailsGatewayCacheStub) GetSessionAccountID(_ context.Context, _ int64, _ string) (int64, error) {
return 0, nil
}
func (s *affinityDetailsGatewayCacheStub) SetSessionAccountID(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error {
return nil
}
func (s *affinityDetailsGatewayCacheStub) RefreshSessionTTL(_ context.Context, _ int64, _ string, _ time.Duration) error {
return nil
}
func (s *affinityDetailsGatewayCacheStub) DeleteSessionAccountID(_ context.Context, _ int64, _ string) error {
return nil
}
func (s *affinityDetailsGatewayCacheStub) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
return nil, nil
}
func (s *affinityDetailsGatewayCacheStub) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error {
return nil
}
func (s *affinityDetailsGatewayCacheStub) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) {
return map[int64]int64{}, nil
}
func (s *affinityDetailsGatewayCacheStub) GetAccountAffinityClientsBatch(_ context.Context, _ map[int64][]int64, _ time.Duration) (map[int64][]string, error) {
return map[int64][]string{}, nil
}
func (s *affinityDetailsGatewayCacheStub) GetAccountAffinityClientsWithScores(_ context.Context, _ int64, _ []int64, _ time.Duration) ([]service.AffinityClient, error) {
return s.clients, nil
}
func (s *affinityDetailsGatewayCacheStub) ClearAccountAffinity(_ context.Context, _ int64, _ []int64) error {
return nil
}
func (s *affinityDetailsGatewayCacheStub) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (users, clients, perUser int64, err error) {
return 0, 0, 0, nil
}
func setupAffinityDetailsRouter(adminSvc service.AdminService, gatewayCache service.GatewayCache) *gin.Engine {
gin.SetMode(gin.TestMode)
router := gin.New()
handler := NewAccountHandler(adminSvc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, gatewayCache)
router.GET("/api/v1/admin/accounts/:id/affinity-details", handler.GetAffinityDetails)
return router
}
func TestAccountHandlerGetAffinityDetails_ContractFields(t *testing.T) {
adminSvc := &affinityDetailsAdminService{
stubAdminService: newStubAdminService(),
account: service.Account{
ID: 88,
Platform: service.PlatformAnthropic,
Type: service.AccountTypeOAuth,
Status: service.StatusActive,
GroupIDs: []int64{1},
Extra: map[string]any{
"client_affinity_enabled": true,
"pinned_users": []any{float64(42), float64(99)},
},
},
usersByID: map[int64]service.User{
42: {ID: 42, Email: "user42@example.com"},
7: {ID: 7, Username: "user7"},
},
}
cache := &affinityDetailsGatewayCacheStub{
clients: []service.AffinityClient{
{UserID: 42, ClientID: "client-a", LastActive: time.Now().Add(-time.Minute)},
{UserID: 7, ClientID: "client-b", LastActive: time.Now().Add(-2 * time.Minute)},
{UserID: 42, ClientID: "client-c", LastActive: time.Now().Add(-3 * time.Minute)},
},
}
router := setupAffinityDetailsRouter(adminSvc, cache)
rec := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts/88/affinity-details", nil)
router.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
var resp struct {
Code int `json:"code"`
Data struct {
Users []struct {
UserID int64 `json:"user_id"`
UserEmail string `json:"user_email"`
ClientCount int `json:"client_count"`
IsPinned bool `json:"is_pinned"`
Clients []struct {
ClientID string `json:"client_id"`
} `json:"clients"`
} `json:"users"`
TotalUsers int `json:"total_users"`
TotalClients int `json:"total_clients"`
PinnedUsers []int64 `json:"pinned_users"`
} `json:"data"`
}
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
require.Equal(t, 0, resp.Code)
require.Equal(t, 2, resp.Data.TotalUsers)
require.Equal(t, 3, resp.Data.TotalClients)
require.Equal(t, []int64{42, 99}, resp.Data.PinnedUsers)
require.Len(t, resp.Data.Users, 2)
usersByID := make(map[int64]struct {
UserEmail string
ClientCount int
IsPinned bool
ClientLen int
}, len(resp.Data.Users))
for _, u := range resp.Data.Users {
usersByID[u.UserID] = struct {
UserEmail string
ClientCount int
IsPinned bool
ClientLen int
}{
UserEmail: u.UserEmail,
ClientCount: u.ClientCount,
IsPinned: u.IsPinned,
ClientLen: len(u.Clients),
}
}
require.Equal(t, "user42@example.com", usersByID[42].UserEmail)
require.Equal(t, 2, usersByID[42].ClientCount)
require.True(t, usersByID[42].IsPinned)
require.Equal(t, 2, usersByID[42].ClientLen)
require.Equal(t, "user7", usersByID[7].UserEmail)
require.Equal(t, 1, usersByID[7].ClientCount)
require.False(t, usersByID[7].IsPinned)
require.Equal(t, 1, usersByID[7].ClientLen)
}
func TestAccountHandlerGetAffinityDetails_DisabledReturnsEmptyContract(t *testing.T) {
adminSvc := &affinityDetailsAdminService{
stubAdminService: newStubAdminService(),
account: service.Account{
ID: 89,
Platform: service.PlatformAnthropic,
Type: service.AccountTypeOAuth,
Status: service.StatusActive,
GroupIDs: []int64{1},
},
usersByID: map[int64]service.User{},
}
router := setupAffinityDetailsRouter(adminSvc, &affinityDetailsGatewayCacheStub{})
rec := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts/89/affinity-details", nil)
router.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
var resp struct {
Code int `json:"code"`
Data struct {
Users []any `json:"users"`
TotalUsers int `json:"total_users"`
TotalClients int `json:"total_clients"`
PinnedUsers []int64 `json:"pinned_users"`
} `json:"data"`
}
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
require.Equal(t, 0, resp.Code)
require.Empty(t, resp.Data.Users)
require.Equal(t, 0, resp.Data.TotalUsers)
require.Equal(t, 0, resp.Data.TotalClients)
require.Empty(t, resp.Data.PinnedUsers)
}
+6 -1
View File
@@ -266,9 +266,14 @@ func AccountFromServiceShallow(a *service.Account) *Account {
}
// 客户端亲和调度(Anthropic 和 Antigravity 账号)
if a.IsClientAffinityEnabled() {
if a.IsAffinityEnabled() {
enabled := true
out.ClientAffinityEnabled = &enabled
allow := a.IsAffinityAllowSwitch()
out.AffinityAllowSwitch = &allow
if pinnedUsers := a.GetPinnedUsers(); len(pinnedUsers) > 0 {
out.PinnedUserIDs = pinnedUsers
}
}
// 提取账号配额限制(apikey / bedrock 类型有效)
+9
View File
@@ -201,6 +201,15 @@ type Account struct {
// 启用后新会话会优先调度到客户端之前使用过的账号
ClientAffinityEnabled *bool `json:"client_affinity_enabled,omitempty"`
// 亲和允许切换(默认 true)
AffinityAllowSwitch *bool `json:"affinity_allow_switch,omitempty"`
// 亲和用户数量(admin 列表端点注入)
AffinityUserCount *int64 `json:"affinity_user_count,omitempty"`
// 指定亲和用户 ID 列表
PinnedUserIDs []int64 `json:"pinned_user_ids,omitempty"`
// 亲和客户端数据(仅 admin 列表端点注入,不由 mapper 填充)
AffinityClientCount *int64 `json:"affinity_client_count,omitempty"`
AffinityClients []string `json:"affinity_clients,omitempty"`
+10 -2
View File
@@ -291,7 +291,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
}
for {
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, reqModel, fs.FailedAccountIDs, "") // Gemini 不使用会话限制
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, reqModel, fs.FailedAccountIDs, "", int64(0)) // Gemini 不使用会话限制
if err != nil {
if len(fs.FailedAccountIDs) == 0 {
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts: "+err.Error(), streamStarted)
@@ -501,8 +501,16 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
for {
// 选择支持该模型的账号
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), currentAPIKey.GroupID, sessionKey, reqModel, fs.FailedAccountIDs, parsedReq.MetadataUserID)
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), currentAPIKey.GroupID, sessionKey, reqModel, fs.FailedAccountIDs, parsedReq.MetadataUserID, subject.UserID)
if err != nil {
if errors.Is(err, service.ErrAffinityNoSwitch) {
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "Affinity account unavailable and switching is disabled", streamStarted)
return
}
if errors.Is(err, service.ErrAffinityLimitExceeded) {
h.handleStreamingAwareError(c, http.StatusTooManyRequests, "api_error", "Affinity client limit exceeded", streamStarted)
return
}
if len(fs.FailedAccountIDs) == 0 {
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts: "+err.Error(), streamStarted)
return
@@ -352,7 +352,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
}
for {
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, modelName, fs.FailedAccountIDs, "") // Gemini 不使用会话限制
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, modelName, fs.FailedAccountIDs, "", int64(0)) // Gemini 不使用会话限制
if err != nil {
if len(fs.FailedAccountIDs) == 0 {
googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts: "+err.Error())
@@ -225,7 +225,7 @@ func (h *SoraGatewayHandler) ChatCompletions(c *gin.Context) {
var lastFailoverHeaders http.Header
for {
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionHash, reqModel, failedAccountIDs, "")
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionHash, reqModel, failedAccountIDs, "", int64(0))
if err != nil {
reqLog.Warn("sora.account_select_failed",
zap.Error(err),
@@ -323,7 +323,7 @@ func (h *SoraVideosHandler) selectAccount(c *gin.Context, model string) (*servic
}
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(
c.Request.Context(), apiKey.GroupID, "", model, nil, "",
c.Request.Context(), apiKey.GroupID, "", model, nil, "", int64(0),
)
if err != nil {
soraErrorResponse(c, http.StatusServiceUnavailable, "server_error", "No available accounts")
+102 -35
View File
@@ -5,6 +5,7 @@ import (
_ "embed"
"fmt"
"strconv"
"strings"
"time"
"github.com/Wei-Shaw/sub2api/internal/service"
@@ -12,9 +13,9 @@ import (
)
const (
stickySessionPrefix = "sticky_session:"
clientAffinityPrefix = "client_affinity:"
clientAffinityReversePrefix = "client_affinity_rev:"
stickySessionPrefix = "sticky_session:"
affinityKeyPrefix = "affinity:"
affinityRevKeyPrefix = "affinity_rev:"
)
var (
@@ -30,6 +31,8 @@ var (
getAffinityClientsWithScoresLua string
//go:embed lua/clear_account_affinity.lua
clearAccountAffinityLua string
//go:embed lua/get_affinity_multi_count.lua
getAffinityMultiCountLua string
getAffinityScript = redis.NewScript(getAffinityLua)
updateAffinityScript = redis.NewScript(updateAffinityLua)
@@ -37,6 +40,7 @@ var (
getAffinityClientsScript = redis.NewScript(getAffinityClientsLua)
getAffinityClientsWithScoresScript = redis.NewScript(getAffinityClientsWithScoresLua)
clearAccountAffinityScript = redis.NewScript(clearAccountAffinityLua)
getAffinityMultiCountScript = redis.NewScript(getAffinityMultiCountLua)
)
type gatewayCache struct {
@@ -84,20 +88,39 @@ func (c *gatewayCache) DeleteSessionAccountID(ctx context.Context, groupID int64
return c.rdb.Del(ctx, key).Err()
}
// buildAffinityKey 构建正向亲和 key(client → accounts)
// 格式: client_affinity:{groupID}:{clientID}
func buildAffinityKey(groupID int64, clientID string) string {
return fmt.Sprintf("%s%d:%s", clientAffinityPrefix, groupID, clientID)
// buildAffinityKey 构建正向亲和 key(member → accounts)
// 格式: affinity:{groupID}:{userID}/{clientID}
func buildAffinityKey(groupID int64, userID int64, clientID string) string {
return fmt.Sprintf("%s%d:%s", affinityKeyPrefix, groupID, buildAffinityMember(userID, clientID))
}
// buildAffinityReverseKey 构建反向亲和 key(account → clients)
// 格式: client_affinity_rev:{groupID}:{accountID}
// buildAffinityReverseKey 构建反向亲和 key(account → members)
// 格式: affinity_rev:{groupID}:{accountID}
func buildAffinityReverseKey(groupID int64, accountID int64) string {
return fmt.Sprintf("%s%d:%d", clientAffinityReversePrefix, groupID, accountID)
return fmt.Sprintf("%s%d:%d", affinityRevKeyPrefix, groupID, accountID)
}
func (c *gatewayCache) GetClientAffinityAccounts(ctx context.Context, groupID int64, clientID string, ttl time.Duration) ([]int64, error) {
key := buildAffinityKey(groupID, clientID)
// buildAffinityMember 构建亲和成员标识
// 格式: {userID}/{clientID}
func buildAffinityMember(userID int64, clientID string) string {
return fmt.Sprintf("%d/%s", userID, clientID)
}
// parseAffinityMember 解析亲和成员标识为 userID 和 clientID
func parseAffinityMember(member string) (userID int64, clientID string) {
idx := strings.IndexByte(member, '/')
if idx < 0 {
// 兼容旧格式(纯 clientID,无 userID 前缀)
return 0, member
}
userID, _ = strconv.ParseInt(member[:idx], 10, 64)
clientID = member[idx+1:]
return userID, clientID
}
// GetAffinityAccounts 获取亲和账号列表(按最近使用降序),同时清理过期成员
func (c *gatewayCache) GetAffinityAccounts(ctx context.Context, groupID int64, userID int64, clientID string, ttl time.Duration) ([]int64, error) {
key := buildAffinityKey(groupID, userID, clientID)
now := time.Now().Unix()
expireThreshold := now - int64(ttl.Seconds())
@@ -120,19 +143,21 @@ func (c *gatewayCache) GetClientAffinityAccounts(ctx context.Context, groupID in
return accountIDs, nil
}
func (c *gatewayCache) UpdateClientAffinity(ctx context.Context, groupID int64, clientID string, accountID int64, ttl time.Duration) error {
fwdKey := buildAffinityKey(groupID, clientID)
// UpdateAffinity 添加/更新亲和关系(更新 score 为当前时间戳,刷新 key TTL)
func (c *gatewayCache) UpdateAffinity(ctx context.Context, groupID int64, userID int64, clientID string, accountID int64, ttl time.Duration) error {
fwdKey := buildAffinityKey(groupID, userID, clientID)
revKey := buildAffinityReverseKey(groupID, accountID)
now := time.Now().Unix()
ttlSeconds := int64(ttl.Seconds())
expireThreshold := now - ttlSeconds
member := buildAffinityMember(userID, clientID)
return updateAffinityScript.Run(ctx, c.rdb, []string{fwdKey, revKey},
now, ttlSeconds, accountID, expireThreshold, clientID,
now, ttlSeconds, accountID, expireThreshold, member,
).Err()
}
// GetAccountAffinityCountBatch 批量获取账号的亲和客户端数量(惰性清理过期成员)
// GetAccountAffinityCountBatch 批量获取账号的亲和成员数量(惰性清理过期成员)
func (c *gatewayCache) GetAccountAffinityCountBatch(ctx context.Context, groupID int64, accountIDs []int64, ttl time.Duration) (map[int64]int64, error) {
if len(accountIDs) == 0 {
return map[int64]int64{}, nil
@@ -162,8 +187,9 @@ func (c *gatewayCache) GetAccountAffinityCountBatch(ctx context.Context, groupID
return result, nil
}
// GetAccountAffinityClientsBatch 批量获取每个账号跨所有分组的亲和客户端列表(去重)。
// GetAccountAffinityClientsBatch 批量获取每个账号跨所有分组的亲和成员列表(去重)。
// accountGroups: map[accountID][]groupID,对每个 (groupID, accountID) 组合查询反向索引。
// 返回值成员格式为 {userID}/{clientID}。
func (c *gatewayCache) GetAccountAffinityClientsBatch(ctx context.Context, accountGroups map[int64][]int64, ttl time.Duration) (map[int64][]string, error) {
if len(accountGroups) == 0 {
return map[int64][]string{}, nil
@@ -197,21 +223,21 @@ func (c *gatewayCache) GetAccountAffinityClientsBatch(ctx context.Context, accou
return nil, err
}
// 合并结果:同一个 accountID 跨多个 group 的 clientID 去重
// 合并结果:同一个 accountID 跨多个 group 的成员去重
result := make(map[int64][]string, len(accountGroups))
seen := make(map[int64]map[string]struct{}, len(accountGroups))
for i, q := range queries {
clients, _ := cmds[i].StringSlice()
if len(clients) == 0 {
members, _ := cmds[i].StringSlice()
if len(members) == 0 {
continue
}
if seen[q.accountID] == nil {
seen[q.accountID] = make(map[string]struct{})
}
for _, clientID := range clients {
if _, exists := seen[q.accountID][clientID]; !exists {
seen[q.accountID][clientID] = struct{}{}
result[q.accountID] = append(result[q.accountID], clientID)
for _, member := range members {
if _, exists := seen[q.accountID][member]; !exists {
seen[q.accountID][member] = struct{}{}
result[q.accountID] = append(result[q.accountID], member)
}
}
}
@@ -245,25 +271,32 @@ func (c *gatewayCache) GetAccountAffinityClientsWithScores(
return nil, err
}
// 合并跨组结果,同一 clientID 取最近的 lastActive
seen := make(map[string]int64) // clientID → max timestamp
// 合并跨组结果,同一 member 取最近的 lastActive
type memberInfo struct {
userID int64
clientID string
ts int64
}
seen := make(map[string]*memberInfo) // member string → info
for _, cmd := range cmds {
vals, _ := cmd.StringSlice()
// vals 格式: [clientID1, score1, clientID2, score2, ...]
// vals 格式: [member1, score1, member2, score2, ...]
for j := 0; j+1 < len(vals); j += 2 {
clientID := vals[j]
member := vals[j]
ts, _ := strconv.ParseInt(vals[j+1], 10, 64)
if existing, ok := seen[clientID]; !ok || ts > existing {
seen[clientID] = ts
if existing, ok := seen[member]; !ok || ts > existing.ts {
uid, cid := parseAffinityMember(member)
seen[member] = &memberInfo{userID: uid, clientID: cid, ts: ts}
}
}
}
result := make([]service.AffinityClient, 0, len(seen))
for clientID, ts := range seen {
for _, info := range seen {
result = append(result, service.AffinityClient{
ClientID: clientID,
LastActive: time.Unix(ts, 0),
UserID: info.userID,
ClientID: info.clientID,
LastActive: time.Unix(info.ts, 0),
})
}
@@ -274,8 +307,8 @@ func (c *gatewayCache) GetAccountAffinityClientsWithScores(
}
// ClearAccountAffinity 清除指定账号在所有分组的亲和记录(正向+反向索引)。
// 对每个 groupID 执行 Lua 脚本:读取反向索引获取所有客户端,
// 从每个客户端的正向索引中移除该账号,然后删除反向索引。
// 对每个 groupID 执行 Lua 脚本:读取反向索引获取所有成员,
// 从每个成员的正向索引中移除该账号,然后删除反向索引。
func (c *gatewayCache) ClearAccountAffinity(ctx context.Context, accountID int64, groupIDs []int64) error {
if len(groupIDs) == 0 {
return nil
@@ -294,3 +327,37 @@ func (c *gatewayCache) ClearAccountAffinity(ctx context.Context, accountID int64
}
return nil
}
// GetAffinityMultiCount 获取账号的多维度亲和计数。
// 返回: uniqueUsers(独立用户数), uniqueClients(独立客户端数), perUserClients(目标用户的客户端数)
func (c *gatewayCache) GetAffinityMultiCount(
ctx context.Context,
groupID int64,
accountID int64,
targetUserID int64,
ttl time.Duration,
) (users, clients, perUser int64, err error) {
key := buildAffinityReverseKey(groupID, accountID)
now := time.Now().Unix()
expireThreshold := now - int64(ttl.Seconds())
targetUserStr := ""
if targetUserID > 0 {
targetUserStr = strconv.FormatInt(targetUserID, 10)
}
result, err := getAffinityMultiCountScript.Run(ctx, c.rdb, []string{key}, expireThreshold, targetUserStr).Int64Slice()
if err != nil {
if err == redis.Nil {
return 0, 0, 0, nil
}
return 0, 0, 0, err
}
if len(result) < 4 {
return 0, 0, 0, nil
}
// result: {totalMembers, uniqueUsers, uniqueClients, perUserClients}
return result[1], result[2], result[3], nil
}
@@ -1,21 +1,21 @@
-- 清除单个账号在指定分组的所有亲和记录(正向+反向)
-- KEYS[1] = client_affinity_rev:{groupID}:{accountID} (反向索引)
-- KEYS[1] = affinity_rev:{groupID}:{accountID} (反向索引)
-- ARGV[1] = groupID (用于构建正向 key)
-- ARGV[2] = accountID (正向索引中要移除的成员)
-- 返回: 清理的客户端数量
-- 返回: 清理的成员数量
local rev_key = KEYS[1]
local group_id = ARGV[1]
local account_id = ARGV[2]
-- 获取反向索引中所有客户端 ID
local clients = redis.call('ZRANGE', rev_key, 0, -1)
if #clients == 0 then
-- 获取反向索引中所有成员 ({userID}/{clientID})
local members = redis.call('ZRANGE', rev_key, 0, -1)
if #members == 0 then
return 0
end
-- 从每个客户端的正向索引中移除该账号
for _, client_id in ipairs(clients) do
local fwd_key = 'client_affinity:' .. group_id .. ':' .. client_id
-- 从每个成员的正向索引中移除该账号
for _, member in ipairs(members) do
local fwd_key = 'affinity:' .. group_id .. ':' .. member
redis.call('ZREM', fwd_key, account_id)
-- 如果正向索引为空,删除 key
if redis.call('ZCARD', fwd_key) == 0 then
@@ -26,4 +26,4 @@ end
-- 删除反向索引
redis.call('DEL', rev_key)
return #clients
return #members
@@ -1,5 +1,5 @@
-- 清理过期成员后返回亲和账号列表(按最近使用降序)
-- KEYS[1] = client_affinity:{groupID}:{clientID}
-- KEYS[1] = affinity:{groupID}:{userID}/{clientID}
-- ARGV[1] = 过期阈值时间戳 (now - ttl)
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1])
return redis.call('ZREVRANGE', KEYS[1], 0, -1)
@@ -1,5 +1,6 @@
-- 清理过期成员后返回反向索引的 clientID 列表(按最近使用降序)
-- KEYS[1] = client_affinity_rev:{groupID}:{accountID}
-- 清理过期成员后返回反向索引的成员列表(按最近使用降序)
-- 成员格式: {userID}/{clientID}
-- KEYS[1] = affinity_rev:{groupID}:{accountID}
-- ARGV[1] = 过期阈值时间戳 (now - ttl)
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1])
return redis.call('ZREVRANGE', KEYS[1], 0, -1)
@@ -1,6 +1,7 @@
-- 清理过期成员后返回反向索引的 clientID 列表及其 score(最后活跃时间戳)
-- KEYS[1] = client_affinity_rev:{groupID}:{accountID}
-- 清理过期成员后返回反向索引的成员列表及其 score(最后活跃时间戳)
-- 成员格式: {userID}/{clientID}
-- KEYS[1] = affinity_rev:{groupID}:{accountID}
-- ARGV[1] = 过期阈值时间戳 (now - ttl)
-- 返回: {clientID1, score1, clientID2, score2, ...}(按最近使用降序)
-- 返回: {member1, score1, member2, score2, ...}(按最近使用降序)
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1])
return redis.call('ZREVRANGEBYSCORE', KEYS[1], '+inf', '-inf', 'WITHSCORES')
@@ -1,5 +1,5 @@
-- 清理过期成员后返回反向索引的成员数量
-- KEYS[1] = client_affinity_rev:{groupID}:{accountID}
-- KEYS[1] = affinity_rev:{groupID}:{accountID}
-- ARGV[1] = 过期阈值时间戳 (now - ttl)
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1])
return redis.call('ZCARD', KEYS[1])
@@ -0,0 +1,40 @@
-- 从反向索引解析多维度计数(用户/客户端/每用户客户端)
-- KEYS[1] = affinity_rev:{groupID}:{accountID}
-- ARGV[1] = 过期阈值时间戳 (now - ttl)
-- ARGV[2] = 目标 userID(传 "" 则不计算 perUserClients)
-- 返回: {totalMembers, uniqueUsers, uniqueClients, perUserClients}
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1])
local members = redis.call('ZRANGE', KEYS[1], 0, -1)
local total = #members
if total == 0 then
return {0, 0, 0, 0}
end
local target_user = ARGV[2]
local users = {}
local clients = {}
local user_count = 0
local client_count = 0
local per_user_count = 0
for _, member in ipairs(members) do
local sep = string.find(member, '/', 1, true)
if sep then
local uid = string.sub(member, 1, sep - 1)
local cid = string.sub(member, sep + 1)
if not users[uid] then
users[uid] = true
user_count = user_count + 1
end
if not clients[cid] then
clients[cid] = true
client_count = client_count + 1
end
if target_user ~= '' and uid == target_user then
per_user_count = per_user_count + 1
end
end
end
return {total, user_count, client_count, per_user_count}
@@ -1,11 +1,11 @@
-- 原子双写正向+反向索引
-- KEYS[1] = client_affinity:{groupID}:{clientID} (正向: client → accounts)
-- KEYS[2] = client_affinity_rev:{groupID}:{accountID} (反向: account → clients)
-- KEYS[1] = affinity:{groupID}:{userID}/{clientID} (正向: member → accounts)
-- KEYS[2] = affinity_rev:{groupID}:{accountID} (反向: account → members)
-- ARGV[1] = 当前时间戳 (score)
-- ARGV[2] = TTL 秒数
-- ARGV[3] = accountID (正向索引的成员)
-- ARGV[4] = 过期阈值时间戳 (now - ttl)
-- ARGV[5] = clientID (反向索引的成员)
-- ARGV[5] = {userID}/{clientID} (反向索引的成员)
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[4])
redis.call('ZADD', KEYS[1], ARGV[1], ARGV[3])
redis.call('EXPIRE', KEYS[1], ARGV[2])
+1
View File
@@ -263,6 +263,7 @@ func registerAccountRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
accounts.POST("/:id/clear-rate-limit", h.Admin.Account.ClearRateLimit)
accounts.POST("/:id/reset-quota", h.Admin.Account.ResetQuota)
accounts.GET("/:id/affinity-clients", h.Admin.Account.GetAffinityClients)
accounts.GET("/:id/affinity-details", h.Admin.Account.GetAffinityDetails)
accounts.GET("/:id/temp-unschedulable", h.Admin.Account.GetTempUnschedulable)
accounts.DELETE("/:id/temp-unschedulable", h.Admin.Account.ClearTempUnschedulable)
accounts.POST("/:id/schedulable", h.Admin.Account.SetSchedulable)
+158 -4
View File
@@ -1204,16 +1204,21 @@ func (a *Account) IsSessionIDMaskingEnabled() bool {
return false
}
// IsClientAffinityEnabled 检查是否启用客户端亲和调度
// 仅适用于 Anthropic 账号(OAuth/SetupToken/APIKey)
// 启用后,新会话会优先调度到之前使用过的账号
func (a *Account) IsClientAffinityEnabled() bool {
// IsAffinityEnabled 检查是否启用亲和调度(统一入口)
// 仅适用于 Anthropic 账号,同时检查新字段 affinity_enabled 和旧字段 client_affinity_enabled(向后兼容)
func (a *Account) IsAffinityEnabled() bool {
if a.Platform != PlatformAnthropic {
return false
}
if a.Extra == nil {
return false
}
if v, ok := a.Extra["affinity_enabled"]; ok {
if enabled, ok := v.(bool); ok {
return enabled
}
return false
}
if v, ok := a.Extra["client_affinity_enabled"]; ok {
if enabled, ok := v.(bool); ok {
return enabled
@@ -1222,6 +1227,11 @@ func (a *Account) IsClientAffinityEnabled() bool {
return false
}
// IsClientAffinityEnabled 向后兼容别名,内部调用 IsAffinityEnabled
func (a *Account) IsClientAffinityEnabled() bool {
return a.IsAffinityEnabled()
}
// AffinityZone 表示账号的客户端亲和分区
type AffinityZone int
@@ -1288,6 +1298,150 @@ func (a *Account) GetAffinityZone(clientCount int64) AffinityZone {
return AffinityZoneRed
}
// IsAffinityAllowSwitch 检查亲和账号是否允许在无绿区时切换到其他账号
// 默认 true(允许切换),设为 false 时即使全部红区也不切换
func (a *Account) IsAffinityAllowSwitch() bool {
if a.Extra == nil {
return true
}
if v, ok := a.Extra["affinity_allow_switch"]; ok {
if allow, ok := v.(bool); ok {
return allow
}
}
return true
}
// GetPinnedUsers 获取指定亲和用户列表
func (a *Account) GetPinnedUsers() []int64 {
if a.Extra == nil {
return nil
}
v, ok := a.Extra["pinned_users"]
if !ok || v == nil {
return nil
}
arr, ok := v.([]any)
if !ok || len(arr) == 0 {
return nil
}
result := make([]int64, 0, len(arr))
for _, item := range arr {
switch id := item.(type) {
case float64:
result = append(result, int64(id))
case int64:
result = append(result, id)
case int:
result = append(result, int64(id))
case json.Number:
if i, err := id.Int64(); err == nil {
result = append(result, i)
}
}
}
return result
}
// IsPinnedUser 检查 userID 是否在指定亲和用户列表中
func (a *Account) IsPinnedUser(userID int64) bool {
for _, id := range a.GetPinnedUsers() {
if id == userID {
return true
}
}
return false
}
// GetAffinityUserBase 获取用户维度亲和基础限制(绿区上限),0 表示未配置
func (a *Account) GetAffinityUserBase() int {
if a.Extra == nil {
return 0
}
if v, ok := a.Extra["affinity_user_base"]; ok {
return parseExtraInt(v)
}
return 0
}
// GetAffinityUserBuffer 获取用户维度亲和缓冲区大小(黄区范围)
// 返回 (value, configured),语义与 GetAffinityBuffer 相同
func (a *Account) GetAffinityUserBuffer() (int, bool) {
if a.Extra == nil {
return 0, false
}
v, ok := a.Extra["affinity_user_buffer"]
if !ok {
return 0, false
}
if v == nil {
return 0, false
}
return parseExtraInt(v), true
}
// GetPerUserClientLimit 获取每用户客户端限制,0 表示不限制
func (a *Account) GetPerUserClientLimit() int {
if a.Extra == nil {
return 0
}
if v, ok := a.Extra["per_user_client_limit"]; ok {
return parseExtraInt(v)
}
return 0
}
// getAffinityZoneForDim 根据单一维度的计数和三区参数计算亲和分区
func getAffinityZoneForDim(count int64, base int, buffer int, bufferConfigured bool) AffinityZone {
if base <= 0 {
return AffinityZoneGreen
}
if count <= int64(base) {
return AffinityZoneGreen
}
if !bufferConfigured {
return AffinityZoneYellow
}
if buffer == 0 {
return AffinityZoneRed
}
if count <= int64(base+buffer) {
return AffinityZoneYellow
}
return AffinityZoneRed
}
// GetMultiDimAffinityZone 根据多维度计数计算亲和分区,取所有维度中最严格的区域。
// 未开启亲和的账号永远返回绿区。
func (a *Account) GetMultiDimAffinityZone(userCount, clientCount, perUserCount int64) AffinityZone {
if !a.IsAffinityEnabled() {
return AffinityZoneGreen
}
worst := AffinityZoneGreen
// 客户端维度
clientBase := a.GetAffinityBase()
clientBuffer, clientBufCfg := a.GetAffinityBuffer()
if z := getAffinityZoneForDim(clientCount, clientBase, clientBuffer, clientBufCfg); z > worst {
worst = z
}
// 用户维度
userBase := a.GetAffinityUserBase()
userBuffer, userBufCfg := a.GetAffinityUserBuffer()
if z := getAffinityZoneForDim(userCount, userBase, userBuffer, userBufCfg); z > worst {
worst = z
}
// 每用户客户端维度
perUserLimit := a.GetPerUserClientLimit()
if perUserLimit > 0 && perUserCount > int64(perUserLimit) {
worst = AffinityZoneRed
}
return worst
}
// IsCacheTTLOverrideEnabled 检查是否启用缓存 TTL 强制替换
// 仅适用于 Anthropic OAuth/SetupToken 类型账号
// 启用后将所有 cache creation tokens 归入指定的 TTL 类型(5m 或 1h)
@@ -30,12 +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) {
func (c *stubSmartRetryCache) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
return nil, nil
}
func (c *stubSmartRetryCache) UpdateClientAffinity(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error {
func (c *stubSmartRetryCache) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error {
return nil
}
func (c *stubSmartRetryCache) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (int64, int64, int64, error) {
return 0, 0, 0, nil
}
func (c *stubSmartRetryCache) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) {
return map[int64]int64{}, nil
}
@@ -0,0 +1,239 @@
package service
import "context"
// gatewayAffinityFlow encapsulates affinity-specific scheduling steps so the
// main account selection flow can stay focused on generic scheduling.
type gatewayAffinityFlow struct {
svc *GatewayService
ctx context.Context
groupID *int64
sessionHash string
requestedModel string
affinityClientID string
affinityUserID int64
platform string
useMixed bool
accountByID map[int64]*Account
isExcluded func(int64) bool
}
type affinityWaitCandidate struct {
account *Account
}
func newGatewayAffinityFlow(
svc *GatewayService,
ctx context.Context,
groupID *int64,
sessionHash string,
requestedModel string,
affinityClientID string,
affinityUserID int64,
platform string,
useMixed bool,
accountByID map[int64]*Account,
isExcluded func(int64) bool,
) *gatewayAffinityFlow {
return &gatewayAffinityFlow{
svc: svc,
ctx: ctx,
groupID: groupID,
sessionHash: sessionHash,
requestedModel: requestedModel,
affinityClientID: affinityClientID,
affinityUserID: affinityUserID,
platform: platform,
useMixed: useMixed,
accountByID: accountByID,
isExcluded: isExcluded,
}
}
// shouldFilterAccountWithoutClientID excludes affinity-enabled Anthropic accounts
// when metadata.user_id does not provide a usable client_id.
func shouldFilterAccountWithoutClientID(account *Account, affinityClientID string) bool {
if account == nil || affinityClientID != "" {
return false
}
if account.Platform != PlatformAnthropic {
return false
}
return account.IsAffinityEnabled()
}
func filterAccountsWithoutClientID(accounts []Account, affinityClientID string) []Account {
if affinityClientID != "" {
return accounts
}
filtered := make([]Account, 0, len(accounts))
for _, acc := range accounts {
if shouldFilterAccountWithoutClientID(&acc, affinityClientID) {
continue
}
filtered = append(filtered, acc)
}
return filtered
}
func (f *gatewayAffinityFlow) preprocessPinnedUsers(accounts []Account) {
if f.affinityUserID <= 0 || f.affinityClientID == "" || f.svc.cache == nil {
return
}
for i := range accounts {
if accounts[i].IsPinnedUser(f.affinityUserID) && accounts[i].IsAffinityEnabled() {
_ = f.svc.cache.UpdateAffinity(
f.ctx,
derefGroupID(f.groupID),
f.affinityUserID,
f.affinityClientID,
accounts[i].ID,
ClientAffinityTTL,
)
}
}
}
// trySelectAffinityAccount runs Layer 1.4 and returns:
// - result != nil: affinity path selected an account or wait plan
// - affinityHit == true: an effective affinity-enabled record was considered and should suppress sticky fallback
func (f *gatewayAffinityFlow) trySelectAffinityAccount() (*AccountSelectionResult, bool, error) {
if f.affinityClientID == "" || f.affinityUserID <= 0 || f.svc.cache == nil {
return nil, false, nil
}
gid := derefGroupID(f.groupID)
affinityAccountIDs, err := f.svc.cache.GetAffinityAccounts(f.ctx, gid, f.affinityUserID, f.affinityClientID, ClientAffinityTTL)
if err != nil || len(affinityAccountIDs) == 0 {
return nil, false, nil
}
noSwitchBlocked := false
anyAllowSwitch := false
effectiveAffinityHit := false
waitCandidates := make(map[int64]*affinityWaitCandidate)
for _, affinityAccID := range affinityAccountIDs {
account, ok := f.accountByID[affinityAccID]
if f.isExcluded != nil && f.isExcluded(affinityAccID) {
checkAcc := account
if !ok && f.svc.accountRepo != nil {
if acc, repoErr := f.svc.accountRepo.GetByID(f.ctx, affinityAccID); repoErr == nil && acc != nil {
checkAcc = acc
}
}
if checkAcc != nil && checkAcc.IsAffinityEnabled() {
effectiveAffinityHit = true
if !checkAcc.IsAffinityAllowSwitch() {
noSwitchBlocked = true
} else {
anyAllowSwitch = true
}
}
continue
}
if !ok || !f.svc.isAccountSchedulableForSelection(account) {
checkAcc := account
if !ok && f.svc.accountRepo != nil {
if acc, repoErr := f.svc.accountRepo.GetByID(f.ctx, affinityAccID); repoErr == nil && acc != nil {
checkAcc = acc
}
}
if checkAcc != nil && checkAcc.IsAffinityEnabled() {
effectiveAffinityHit = true
if !checkAcc.IsAffinityAllowSwitch() {
noSwitchBlocked = true
} else {
anyAllowSwitch = true
}
}
continue
}
if !account.IsAffinityEnabled() {
continue
}
effectiveAffinityHit = true
if account.IsAffinityAllowSwitch() {
anyAllowSwitch = true
}
if !f.svc.isAccountAllowedForPlatform(account, f.platform, f.useMixed) ||
(f.requestedModel != "" && !f.svc.isModelSupportedByAccountWithContext(f.ctx, account, f.requestedModel)) ||
!f.svc.isAccountSchedulableForModelSelection(f.ctx, account, f.requestedModel) ||
!f.svc.isAccountSchedulableForQuota(account) ||
!f.svc.isAccountSchedulableForWindowCost(f.ctx, account, false) ||
!f.svc.isAccountSchedulableForRPM(f.ctx, account, false) {
if !account.IsAffinityAllowSwitch() {
noSwitchBlocked = true
}
continue
}
userCount, clientCount, perUserCount, multiErr := f.svc.cache.GetAffinityMultiCount(
f.ctx, gid, affinityAccID, f.affinityUserID, ClientAffinityTTL,
)
if multiErr == nil {
zone := account.GetMultiDimAffinityZone(userCount, clientCount, perUserCount)
if zone == AffinityZoneRed {
if !account.IsAffinityAllowSwitch() {
noSwitchBlocked = true
}
continue
}
}
result, acquireErr := f.svc.tryAcquireAccountSlot(f.ctx, affinityAccID, account.Concurrency)
if acquireErr == nil && result.Acquired {
if !f.svc.checkAndRegisterSession(f.ctx, account, f.sessionHash) {
result.ReleaseFunc()
continue
}
_ = f.svc.cache.UpdateAffinity(f.ctx, gid, f.affinityUserID, f.affinityClientID, affinityAccID, ClientAffinityTTL)
if f.sessionHash != "" {
_ = f.svc.cache.SetSessionAccountID(f.ctx, gid, f.sessionHash, affinityAccID, stickySessionTTL)
}
return &AccountSelectionResult{
Account: account,
Acquired: true,
ReleaseFunc: result.ReleaseFunc,
}, true, nil
}
if acquireErr == nil && !result.Acquired && !account.IsAffinityAllowSwitch() {
noSwitchBlocked = true
waitCandidates[affinityAccID] = &affinityWaitCandidate{account: account}
}
}
if noSwitchBlocked && !anyAllowSwitch && f.svc.concurrencyService != nil {
for _, waitAccID := range affinityAccountIDs {
candidate, ok := waitCandidates[waitAccID]
if !ok || candidate == nil || candidate.account == nil {
continue
}
acc := candidate.account
waitingCount, _ := f.svc.concurrencyService.GetAccountWaitingCount(f.ctx, waitAccID)
if waitingCount >= f.svc.schedulingConfig().StickySessionMaxWaiting {
continue
}
if !f.svc.checkAndRegisterSession(f.ctx, acc, f.sessionHash) {
continue
}
cfg := f.svc.schedulingConfig()
return &AccountSelectionResult{
Account: acc,
WaitPlan: &AccountWaitPlan{
AccountID: waitAccID,
MaxConcurrency: acc.Concurrency,
Timeout: cfg.StickySessionWaitTimeout,
MaxWaiting: cfg.StickySessionMaxWaiting,
},
}, true, nil
}
return nil, effectiveAffinityHit, ErrAffinityNoSwitch
}
return nil, effectiveAffinityHit, nil
}
@@ -9,6 +9,7 @@ import (
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
@@ -36,12 +37,15 @@ func (m *mockAffinityCache) RefreshSessionTTL(_ context.Context, _ int64, _ stri
func (m *mockAffinityCache) DeleteSessionAccountID(_ context.Context, _ int64, _ string) error {
return nil
}
func (m *mockAffinityCache) GetClientAffinityAccounts(_ context.Context, _ int64, _ string, _ time.Duration) ([]int64, error) {
func (m *mockAffinityCache) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
return nil, nil
}
func (m *mockAffinityCache) UpdateClientAffinity(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error {
func (m *mockAffinityCache) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error {
return nil
}
func (m *mockAffinityCache) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (int64, int64, int64, error) {
return 0, 0, 0, nil
}
func (m *mockAffinityCache) GetAccountAffinityCountBatch(ctx context.Context, groupID int64, accountIDs []int64, ttl time.Duration) (map[int64]int64, error) {
m.getCountBatchCalls++
if m.getCountBatchFunc != nil {
@@ -613,6 +617,28 @@ func TestAffinityIsClientAffinityEnabled(t *testing.T) {
}
assert.False(t, acc.IsClientAffinityEnabled(), "string 'true' should not enable affinity")
})
t.Run("new flag false overrides legacy flag true", func(t *testing.T) {
acc := &Account{
Platform: PlatformAnthropic,
Extra: map[string]any{
"affinity_enabled": false,
"client_affinity_enabled": true,
},
}
assert.False(t, acc.IsClientAffinityEnabled())
})
t.Run("new flag true overrides legacy flag false", func(t *testing.T) {
acc := &Account{
Platform: PlatformAnthropic,
Extra: map[string]any{
"affinity_enabled": true,
"client_affinity_enabled": false,
},
}
assert.True(t, acc.IsClientAffinityEnabled())
})
}
// ===========================================================================
@@ -798,3 +824,936 @@ func TestClassifyByAffinityZone(t *testing.T) {
require.Len(t, result, 2)
})
}
// ===========================================================================
// GetMultiDimAffinityZone 测试
// ===========================================================================
func TestGetMultiDimAffinityZone(t *testing.T) {
makeAccount := func(opts map[string]any) *Account {
extra := map[string]any{"affinity_enabled": true}
for k, v := range opts {
extra[k] = v
}
return &Account{
Platform: PlatformAnthropic,
Extra: extra,
}
}
t.Run("user green + client green = green", func(t *testing.T) {
acc := makeAccount(map[string]any{
"affinity_base": 10,
"affinity_buffer": 5,
"affinity_user_base": 3,
"affinity_user_buffer": 2,
})
assert.Equal(t, AffinityZoneGreen, acc.GetMultiDimAffinityZone(2, 5, 0))
})
t.Run("user green + client yellow = yellow", func(t *testing.T) {
acc := makeAccount(map[string]any{
"affinity_base": 10,
"affinity_buffer": 5,
"affinity_user_base": 3,
"affinity_user_buffer": 2,
})
assert.Equal(t, AffinityZoneYellow, acc.GetMultiDimAffinityZone(2, 12, 0))
})
t.Run("user yellow + client green = yellow", func(t *testing.T) {
acc := makeAccount(map[string]any{
"affinity_base": 10,
"affinity_buffer": 5,
"affinity_user_base": 3,
"affinity_user_buffer": 2,
})
assert.Equal(t, AffinityZoneYellow, acc.GetMultiDimAffinityZone(4, 5, 0))
})
t.Run("user red + client green = red", func(t *testing.T) {
acc := makeAccount(map[string]any{
"affinity_base": 10,
"affinity_buffer": 5,
"affinity_user_base": 3,
"affinity_user_buffer": 0, // no yellow, direct red
})
assert.Equal(t, AffinityZoneRed, acc.GetMultiDimAffinityZone(4, 5, 0))
})
t.Run("user green + client red = red", func(t *testing.T) {
acc := makeAccount(map[string]any{
"affinity_base": 10,
"affinity_buffer": 0, // no yellow, direct red
"affinity_user_base": 3,
"affinity_user_buffer": 2,
})
assert.Equal(t, AffinityZoneRed, acc.GetMultiDimAffinityZone(2, 11, 0))
})
t.Run("user_base=0 only checks client dimension", func(t *testing.T) {
acc := makeAccount(map[string]any{
"affinity_base": 10,
"affinity_buffer": 5,
// no affinity_user_base → 0 → user dimension always green
})
// client in yellow range
assert.Equal(t, AffinityZoneYellow, acc.GetMultiDimAffinityZone(100, 12, 0))
})
t.Run("perUserCount exceeds limit = red", func(t *testing.T) {
acc := makeAccount(map[string]any{
"affinity_base": 10,
"affinity_buffer": 5,
"per_user_client_limit": 3,
})
// client green, user green, but perUser > limit
assert.Equal(t, AffinityZoneRed, acc.GetMultiDimAffinityZone(1, 5, 4))
})
t.Run("perUserCount within limit = green", func(t *testing.T) {
acc := makeAccount(map[string]any{
"affinity_base": 10,
"affinity_buffer": 5,
"per_user_client_limit": 3,
})
assert.Equal(t, AffinityZoneGreen, acc.GetMultiDimAffinityZone(1, 5, 3))
})
t.Run("affinity disabled always green", func(t *testing.T) {
acc := &Account{
Platform: PlatformAnthropic,
Extra: map[string]any{
"affinity_enabled": false,
"affinity_base": 1,
"affinity_buffer": 0,
},
}
assert.Equal(t, AffinityZoneGreen, acc.GetMultiDimAffinityZone(100, 100, 100))
})
}
// ===========================================================================
// IsAffinityAllowSwitch 测试
// ===========================================================================
func TestIsAffinityAllowSwitch(t *testing.T) {
t.Run("default true when no field in Extra", func(t *testing.T) {
acc := &Account{
Platform: PlatformAnthropic,
Extra: map[string]any{"affinity_enabled": true},
}
assert.True(t, acc.IsAffinityAllowSwitch())
})
t.Run("explicit false", func(t *testing.T) {
acc := &Account{
Platform: PlatformAnthropic,
Extra: map[string]any{"affinity_allow_switch": false},
}
assert.False(t, acc.IsAffinityAllowSwitch())
})
t.Run("Extra nil defaults to true", func(t *testing.T) {
acc := &Account{
Platform: PlatformAnthropic,
Extra: nil,
}
assert.True(t, acc.IsAffinityAllowSwitch())
})
t.Run("explicit true", func(t *testing.T) {
acc := &Account{
Platform: PlatformAnthropic,
Extra: map[string]any{"affinity_allow_switch": true},
}
assert.True(t, acc.IsAffinityAllowSwitch())
})
}
// ===========================================================================
// GetPinnedUsers / IsPinnedUser 测试
// ===========================================================================
func TestGetPinnedUsers(t *testing.T) {
t.Run("normal list", func(t *testing.T) {
acc := &Account{
Platform: PlatformAnthropic,
Extra: map[string]any{
"pinned_users": []any{float64(10), float64(20), float64(30)},
},
}
result := acc.GetPinnedUsers()
assert.Equal(t, []int64{10, 20, 30}, result)
})
t.Run("empty list", func(t *testing.T) {
acc := &Account{
Platform: PlatformAnthropic,
Extra: map[string]any{
"pinned_users": []any{},
},
}
result := acc.GetPinnedUsers()
assert.Nil(t, result)
})
t.Run("Extra nil", func(t *testing.T) {
acc := &Account{
Platform: PlatformAnthropic,
Extra: nil,
}
result := acc.GetPinnedUsers()
assert.Nil(t, result)
})
t.Run("pinned_users not set", func(t *testing.T) {
acc := &Account{
Platform: PlatformAnthropic,
Extra: map[string]any{},
}
result := acc.GetPinnedUsers()
assert.Nil(t, result)
})
}
func TestIsPinnedUser(t *testing.T) {
acc := &Account{
Platform: PlatformAnthropic,
Extra: map[string]any{
"pinned_users": []any{float64(10), float64(20), float64(30)},
},
}
t.Run("hit", func(t *testing.T) {
assert.True(t, acc.IsPinnedUser(20))
})
t.Run("miss", func(t *testing.T) {
assert.False(t, acc.IsPinnedUser(99))
})
t.Run("nil Extra", func(t *testing.T) {
nilAcc := &Account{Platform: PlatformAnthropic, Extra: nil}
assert.False(t, nilAcc.IsPinnedUser(10))
})
}
// ===========================================================================
// Enhanced mock for integration tests (supports affinity lookups + tracking)
// ===========================================================================
type affinityIntegrationCache struct {
mockAffinityCache
getAffinityAccountsFunc func(ctx context.Context, groupID int64, userID int64, clientID string, ttl time.Duration) ([]int64, error)
updateAffinityCalls []affinityUpdateCall
getMultiCountFunc func(ctx context.Context, groupID int64, accountID int64, targetUserID int64, ttl time.Duration) (int64, int64, int64, error)
sessionBindings map[string]int64
getAffinityAccountsCalled bool // 标记 GetAffinityAccounts 是否已被调用,用于区分预处理/Layer 2 阶段
}
type affinityUpdateCall struct {
groupID int64
userID int64
clientID string
accountID int64
beforeAffinityLookup bool // true 表示该调用发生在 GetAffinityAccounts 之前(预处理阶段)
}
func (c *affinityIntegrationCache) GetSessionAccountID(_ context.Context, _ int64, hash string) (int64, error) {
if c.sessionBindings != nil {
if id, ok := c.sessionBindings[hash]; ok {
return id, nil
}
}
return 0, errors.New("not found")
}
func (c *affinityIntegrationCache) GetAffinityAccounts(ctx context.Context, groupID int64, userID int64, clientID string, ttl time.Duration) ([]int64, error) {
c.getAffinityAccountsCalled = true
if c.getAffinityAccountsFunc != nil {
return c.getAffinityAccountsFunc(ctx, groupID, userID, clientID, ttl)
}
return nil, nil
}
func (c *affinityIntegrationCache) UpdateAffinity(_ context.Context, groupID int64, userID int64, clientID string, accountID int64, _ time.Duration) error {
c.updateAffinityCalls = append(c.updateAffinityCalls, affinityUpdateCall{
groupID: groupID, userID: userID, clientID: clientID, accountID: accountID,
beforeAffinityLookup: !c.getAffinityAccountsCalled,
})
return nil
}
func (c *affinityIntegrationCache) GetAffinityMultiCount(ctx context.Context, groupID int64, accountID int64, targetUserID int64, ttl time.Duration) (int64, int64, int64, error) {
if c.getMultiCountFunc != nil {
return c.getMultiCountFunc(ctx, groupID, accountID, targetUserID, ttl)
}
return 0, 0, 0, nil
}
// ===========================================================================
// TestAffinityPreprocessPinnedUsers — 验证 pinned_users 预处理逻辑
// ===========================================================================
func TestAffinityPreprocessPinnedUsers(t *testing.T) {
ctx := context.Background()
testUserID := int64(42)
testClientID := "user_" + "a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2" + "_account_test"
t.Run("pinned user triggers UpdateAffinity", func(t *testing.T) {
cache := &affinityIntegrationCache{}
repo := &mockAccountRepoForPlatform{
accounts: []Account{
{
ID: 1, Platform: PlatformAnthropic, Priority: 1,
Status: StatusActive, Schedulable: true, Concurrency: 5,
Extra: map[string]any{
"affinity_enabled": true,
"pinned_users": []any{float64(42)},
},
},
{
ID: 2, Platform: PlatformAnthropic, Priority: 1,
Status: StatusActive, Schedulable: true, Concurrency: 5,
Extra: map[string]any{"affinity_enabled": true},
},
},
accountsByID: map[int64]*Account{},
}
for i := range repo.accounts {
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
}
concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{1: true, 2: true}}
concSvc := NewConcurrencyService(concCache)
cfg := &config.Config{RunMode: config.RunModeStandard}
cfg.Gateway.Scheduling.LoadBatchEnabled = true
svc := &GatewayService{
accountRepo: repo,
cache: cache,
cfg: cfg,
concurrencyService: concSvc,
}
_, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, testClientID, testUserID)
require.NoError(t, err)
// 验证 UpdateAffinity 至少被调用了一次(pinned preprocess for account 1)
found := false
for _, call := range cache.updateAffinityCalls {
if call.accountID == 1 && call.userID == testUserID {
found = true
break
}
}
assert.True(t, found, "UpdateAffinity should be called for pinned account 1 with userID %d", testUserID)
})
t.Run("non-pinned user does not trigger UpdateAffinity preprocess", func(t *testing.T) {
cache := &affinityIntegrationCache{
getAffinityAccountsFunc: func(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
return nil, nil // 无亲和记录,直接进入 Layer 2
},
}
repo := &mockAccountRepoForPlatform{
accounts: []Account{
{
ID: 1, Platform: PlatformAnthropic, Priority: 1,
Status: StatusActive, Schedulable: true, Concurrency: 5,
Extra: map[string]any{
"affinity_enabled": true,
"pinned_users": []any{float64(99)}, // different user
},
},
},
accountsByID: map[int64]*Account{},
}
for i := range repo.accounts {
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
}
concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{1: true}}
concSvc := NewConcurrencyService(concCache)
cfg := &config.Config{RunMode: config.RunModeStandard}
cfg.Gateway.Scheduling.LoadBatchEnabled = true
svc := &GatewayService{
accountRepo: repo,
cache: cache,
cfg: cfg,
concurrencyService: concSvc,
}
_, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, testClientID, testUserID)
require.NoError(t, err)
// 验证预处理阶段(beforeAffinityLookup=true)没有为 userID=42 在 account 1 上调用 UpdateAffinity
for _, call := range cache.updateAffinityCalls {
if call.beforeAffinityLookup && call.accountID == 1 && call.userID == testUserID {
t.Errorf("UpdateAffinity should NOT be called in preprocess for non-pinned user %d on account %d", testUserID, call.accountID)
}
}
})
}
// ===========================================================================
// TestAffinityNoSwitchError — 验证 ErrAffinityNoSwitch 行为
// ===========================================================================
func TestAffinityNoSwitchError(t *testing.T) {
ctx := context.Background()
testUserID := int64(42)
testClientID := "user_" + "a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2" + "_account_test"
futureTime := time.Now().Add(1 * time.Hour)
t.Run("affinity hit but unschedulable + allow_switch=false returns ErrAffinityNoSwitch", func(t *testing.T) {
cache := &affinityIntegrationCache{
getAffinityAccountsFunc: func(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
return []int64{1}, nil
},
}
repo := &mockAccountRepoForPlatform{
accounts: []Account{
{
ID: 1, Platform: PlatformAnthropic, Priority: 1,
Status: StatusActive, Schedulable: true,
RateLimitResetAt: &futureTime, // rate limited = unschedulable for selection
Concurrency: 5,
Extra: map[string]any{
"affinity_enabled": true,
"affinity_allow_switch": false,
},
},
{
ID: 2, Platform: PlatformAnthropic, Priority: 1,
Status: StatusActive, Schedulable: true, Concurrency: 5,
},
},
accountsByID: map[int64]*Account{},
}
for i := range repo.accounts {
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
}
concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{2: true}}
concSvc := NewConcurrencyService(concCache)
cfg := &config.Config{RunMode: config.RunModeStandard}
cfg.Gateway.Scheduling.LoadBatchEnabled = true
svc := &GatewayService{
accountRepo: repo,
cache: cache,
cfg: cfg,
concurrencyService: concSvc,
}
_, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, testClientID, testUserID)
require.Error(t, err)
assert.ErrorIs(t, err, ErrAffinityNoSwitch)
})
t.Run("affinity hit but unschedulable + allow_switch=true continues to Layer 2", func(t *testing.T) {
cache := &affinityIntegrationCache{
getAffinityAccountsFunc: func(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
return []int64{1}, nil
},
}
repo := &mockAccountRepoForPlatform{
accounts: []Account{
{
ID: 1, Platform: PlatformAnthropic, Priority: 1,
Status: StatusActive, Schedulable: true,
RateLimitResetAt: &futureTime, // rate limited = unschedulable
Concurrency: 5,
Extra: map[string]any{
"affinity_enabled": true,
"affinity_allow_switch": true,
},
},
{
ID: 2, Platform: PlatformAnthropic, Priority: 1,
Status: StatusActive, Schedulable: true, Concurrency: 5,
},
},
accountsByID: map[int64]*Account{},
}
for i := range repo.accounts {
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
}
concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{2: true}}
concSvc := NewConcurrencyService(concCache)
cfg := &config.Config{RunMode: config.RunModeStandard}
cfg.Gateway.Scheduling.LoadBatchEnabled = true
svc := &GatewayService{
accountRepo: repo,
cache: cache,
cfg: cfg,
concurrencyService: concSvc,
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, testClientID, testUserID)
require.NoError(t, err)
require.NotNil(t, result)
assert.Equal(t, int64(2), result.Account.ID, "should fall through to Layer 2 and select account 2")
})
t.Run("no affinity records - not affected by allow_switch", func(t *testing.T) {
cache := &affinityIntegrationCache{} // returns nil from GetAffinityAccounts
repo := &mockAccountRepoForPlatform{
accounts: []Account{
{
ID: 1, Platform: PlatformAnthropic, Priority: 1,
Status: StatusActive, Schedulable: true, Concurrency: 5,
Extra: map[string]any{
"affinity_enabled": true,
"affinity_allow_switch": false,
},
},
},
accountsByID: map[int64]*Account{},
}
for i := range repo.accounts {
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
}
concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{1: true}}
concSvc := NewConcurrencyService(concCache)
cfg := &config.Config{RunMode: config.RunModeStandard}
cfg.Gateway.Scheduling.LoadBatchEnabled = true
svc := &GatewayService{
accountRepo: repo,
cache: cache,
cfg: cfg,
concurrencyService: concSvc,
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, testClientID, testUserID)
require.NoError(t, err)
require.NotNil(t, result)
assert.Equal(t, int64(1), result.Account.ID, "should select normally via Layer 2")
})
t.Run("affinity disabled account ignores allow_switch=false", func(t *testing.T) {
cache := &affinityIntegrationCache{
getAffinityAccountsFunc: func(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
return []int64{1}, nil
},
}
repo := &mockAccountRepoForPlatform{
accounts: []Account{
{
ID: 1, Platform: PlatformAnthropic, Priority: 1,
Status: StatusActive, Schedulable: true, Concurrency: 5,
Extra: map[string]any{
"affinity_enabled": false,
"affinity_allow_switch": false,
},
},
{
ID: 2, Platform: PlatformAnthropic, Priority: 2,
Status: StatusActive, Schedulable: true, Concurrency: 5,
},
},
accountsByID: map[int64]*Account{},
}
for i := range repo.accounts {
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
}
concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{2: true}}
concSvc := NewConcurrencyService(concCache)
cfg := &config.Config{RunMode: config.RunModeStandard}
cfg.Gateway.Scheduling.LoadBatchEnabled = true
svc := &GatewayService{
accountRepo: repo,
cache: cache,
cfg: cfg,
concurrencyService: concSvc,
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, testClientID, testUserID)
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
assert.NotEqual(t, int64(0), result.Account.ID)
})
t.Run("allow_switch=false + model_unsupported returns ErrAffinityNoSwitch", func(t *testing.T) {
cache := &affinityIntegrationCache{
getAffinityAccountsFunc: func(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
return []int64{1}, nil
},
}
repo := &mockAccountRepoForPlatform{
accounts: []Account{
{
ID: 1, Platform: PlatformAnthropic, Priority: 1,
Status: StatusActive, Schedulable: true, Concurrency: 5,
Extra: map[string]any{
"affinity_enabled": true,
"affinity_allow_switch": false,
},
// model_mapping 仅包含 claude-3-haiku → 请求 claude-3-5-sonnet 将被拒绝
Credentials: map[string]any{
"model_mapping": map[string]any{
"claude-3-haiku-20240307": "claude-3-haiku-20240307",
},
},
},
{
ID: 2, Platform: PlatformAnthropic, Priority: 1,
Status: StatusActive, Schedulable: true, Concurrency: 5,
},
},
accountsByID: map[int64]*Account{},
}
for i := range repo.accounts {
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
}
concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{2: true}}
concSvc := NewConcurrencyService(concCache)
cfg := &config.Config{RunMode: config.RunModeStandard}
cfg.Gateway.Scheduling.LoadBatchEnabled = true
svc := &GatewayService{
accountRepo: repo,
cache: cache,
cfg: cfg,
concurrencyService: concSvc,
}
_, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, testClientID, testUserID)
require.Error(t, err)
assert.ErrorIs(t, err, ErrAffinityNoSwitch)
})
t.Run("allow_switch=false + red_zone returns ErrAffinityNoSwitch", func(t *testing.T) {
cache := &affinityIntegrationCache{
getAffinityAccountsFunc: func(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
return []int64{1}, nil
},
getMultiCountFunc: func(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (int64, int64, int64, error) {
// userCount=5, clientCount=5, perUserCount=0
// 账号 affinity_base=1, affinity_buffer=0 → clientCount(5) > base(1) + buffer(0) → 红区
return 5, 5, 0, nil
},
}
repo := &mockAccountRepoForPlatform{
accounts: []Account{
{
ID: 1, Platform: PlatformAnthropic, Priority: 1,
Status: StatusActive, Schedulable: true, Concurrency: 5,
Extra: map[string]any{
"affinity_enabled": true,
"affinity_allow_switch": false,
"affinity_base": 1,
"affinity_buffer": 0, // buffer=0 → 超过 base 直接红区
},
},
{
ID: 2, Platform: PlatformAnthropic, Priority: 1,
Status: StatusActive, Schedulable: true, Concurrency: 5,
},
},
accountsByID: map[int64]*Account{},
}
for i := range repo.accounts {
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
}
concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{2: true}}
concSvc := NewConcurrencyService(concCache)
cfg := &config.Config{RunMode: config.RunModeStandard}
cfg.Gateway.Scheduling.LoadBatchEnabled = true
svc := &GatewayService{
accountRepo: repo,
cache: cache,
cfg: cfg,
concurrencyService: concSvc,
}
_, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, testClientID, testUserID)
require.Error(t, err)
assert.ErrorIs(t, err, ErrAffinityNoSwitch)
})
t.Run("one_allow_switch=true overrides another allow_switch=false (一票放行)", func(t *testing.T) {
// 账号 1: allow_switch=false + rate limited(不可调度)
// 账号 3: allow_switch=true + rate limited(不可调度)
// 一票放行:账号 3 允许切换 → 不阻断 → 降级到 Layer 2 选中账号 2
cache := &affinityIntegrationCache{
getAffinityAccountsFunc: func(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
return []int64{1, 3}, nil
},
}
repo := &mockAccountRepoForPlatform{
accounts: []Account{
{
ID: 1, Platform: PlatformAnthropic, Priority: 1,
Status: StatusActive, Schedulable: true,
RateLimitResetAt: &futureTime, // rate limited
Concurrency: 5,
Extra: map[string]any{
"affinity_enabled": true,
"affinity_allow_switch": false,
},
},
{
ID: 2, Platform: PlatformAnthropic, Priority: 1,
Status: StatusActive, Schedulable: true, Concurrency: 5,
},
{
ID: 3, Platform: PlatformAnthropic, Priority: 1,
Status: StatusActive, Schedulable: true,
RateLimitResetAt: &futureTime, // rate limited
Concurrency: 5,
Extra: map[string]any{
"affinity_enabled": true,
"affinity_allow_switch": true,
},
},
},
accountsByID: map[int64]*Account{},
}
for i := range repo.accounts {
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
}
concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{2: true}}
concSvc := NewConcurrencyService(concCache)
cfg := &config.Config{RunMode: config.RunModeStandard}
cfg.Gateway.Scheduling.LoadBatchEnabled = true
svc := &GatewayService{
accountRepo: repo,
cache: cache,
cfg: cfg,
concurrencyService: concSvc,
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, testClientID, testUserID)
require.NoError(t, err)
require.NotNil(t, result)
assert.Equal(t, int64(2), result.Account.ID, "一票放行: should fall through to Layer 2")
})
t.Run("excluded affinity account with allow_switch=false returns ErrAffinityNoSwitch", func(t *testing.T) {
cache := &affinityIntegrationCache{
getAffinityAccountsFunc: func(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
return []int64{1}, nil
},
}
repo := &mockAccountRepoForPlatform{
accounts: []Account{
{
ID: 1, Platform: PlatformAnthropic, Priority: 1,
Status: StatusActive, Schedulable: true, Concurrency: 5,
Extra: map[string]any{
"affinity_enabled": true,
"affinity_allow_switch": false,
},
},
{
ID: 2, Platform: PlatformAnthropic, Priority: 1,
Status: StatusActive, Schedulable: true, Concurrency: 5,
},
},
accountsByID: map[int64]*Account{},
}
for i := range repo.accounts {
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
}
concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{2: true}}
concSvc := NewConcurrencyService(concCache)
cfg := &config.Config{RunMode: config.RunModeStandard}
cfg.Gateway.Scheduling.LoadBatchEnabled = true
svc := &GatewayService{
accountRepo: repo,
cache: cache,
cfg: cfg,
concurrencyService: concSvc,
}
excluded := map[int64]struct{}{1: {}}
_, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", excluded, testClientID, testUserID)
require.Error(t, err)
assert.ErrorIs(t, err, ErrAffinityNoSwitch)
})
t.Run("disabled affinity record does not suppress sticky fallback", func(t *testing.T) {
sessionHash := "sticky-disabled-affinity"
cache := &affinityIntegrationCache{
getAffinityAccountsFunc: func(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
return []int64{1}, nil
},
sessionBindings: map[string]int64{sessionHash: 2},
}
repo := &mockAccountRepoForPlatform{
accounts: []Account{
{
ID: 1, Platform: PlatformAnthropic, Priority: 1,
Status: StatusActive, Schedulable: true, Concurrency: 5,
Extra: map[string]any{
"affinity_enabled": false,
},
},
{
ID: 2, Platform: PlatformAnthropic, Priority: 1,
Status: StatusActive, Schedulable: true, Concurrency: 5,
},
},
accountsByID: map[int64]*Account{},
}
for i := range repo.accounts {
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
}
concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{2: true}}
concSvc := NewConcurrencyService(concCache)
cfg := &config.Config{RunMode: config.RunModeStandard}
cfg.Gateway.Scheduling.LoadBatchEnabled = true
svc := &GatewayService{
accountRepo: repo,
cache: cache,
cfg: cfg,
concurrencyService: concSvc,
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, sessionHash, "claude-3-5-sonnet-20241022", nil, testClientID, testUserID)
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
assert.Equal(t, int64(2), result.Account.ID)
})
t.Run("historical affinity record with disabled account still falls back to sticky account", func(t *testing.T) {
sessionHash := "sticky-disabled-status-affinity"
cache := &affinityIntegrationCache{
getAffinityAccountsFunc: func(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
return []int64{1}, nil
},
sessionBindings: map[string]int64{sessionHash: 2},
}
repo := &mockAccountRepoForPlatform{
accounts: []Account{
{
ID: 1, Platform: PlatformAnthropic, Priority: 1,
Status: StatusDisabled, Schedulable: true, Concurrency: 5,
Extra: map[string]any{
"affinity_enabled": true,
},
},
{
ID: 2, Platform: PlatformAnthropic, Priority: 1,
Status: StatusActive, Schedulable: true, Concurrency: 5,
},
},
accountsByID: map[int64]*Account{},
}
for i := range repo.accounts {
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
}
concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{2: true}}
concSvc := NewConcurrencyService(concCache)
cfg := &config.Config{RunMode: config.RunModeStandard}
cfg.Gateway.Scheduling.LoadBatchEnabled = true
svc := &GatewayService{
accountRepo: repo,
cache: cache,
cfg: cfg,
concurrencyService: concSvc,
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, sessionHash, "claude-3-5-sonnet-20241022", nil, testClientID, testUserID)
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
assert.Equal(t, int64(2), result.Account.ID)
})
}
// ===========================================================================
// TestAffinityMultiDimClassify — 验证 classifyByAffinityZone 与多维度计数
// ===========================================================================
func TestAffinityMultiDimClassify(t *testing.T) {
t.Run("multi-dim zone classification with user and client dimensions", func(t *testing.T) {
makeAWL := func(id int64, extra map[string]any, count int64) accountWithLoad {
return accountWithLoad{
account: &Account{ID: id, Platform: PlatformAnthropic, Extra: extra},
loadInfo: &AccountLoadInfo{AccountID: id},
affinityCount: count,
}
}
// Account 1: client green (3 <= 5), classified as green by old classifyByAffinityZone
// Account 2: client red (10 > 8), classified as red
accs := []accountWithLoad{
makeAWL(1, map[string]any{
"affinity_enabled": true,
"affinity_base": 5,
"affinity_buffer": 3,
}, 3),
makeAWL(2, map[string]any{
"affinity_enabled": true,
"affinity_base": 5,
"affinity_buffer": 3,
}, 10),
}
result := classifyByAffinityZone(accs)
require.Len(t, result, 1)
assert.Equal(t, int64(1), result[0].account.ID, "only green account should remain")
})
t.Run("all green returns all", func(t *testing.T) {
makeAWL := func(id int64, count int64) accountWithLoad {
return accountWithLoad{
account: &Account{
ID: id,
Platform: PlatformAnthropic,
Extra: map[string]any{
"affinity_enabled": true,
"affinity_base": 10,
"affinity_buffer": 5,
},
},
loadInfo: &AccountLoadInfo{AccountID: id},
affinityCount: count,
}
}
accs := []accountWithLoad{
makeAWL(1, 3),
makeAWL(2, 5),
makeAWL(3, 10), // exactly at base
}
result := classifyByAffinityZone(accs)
require.Len(t, result, 3)
})
}
@@ -143,12 +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) {
func (s *stickyGatewayCacheHotpathStub) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
return nil, nil
}
func (s *stickyGatewayCacheHotpathStub) UpdateClientAffinity(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error {
func (s *stickyGatewayCacheHotpathStub) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error {
return nil
}
func (s *stickyGatewayCacheHotpathStub) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (int64, int64, int64, error) {
return 0, 0, 0, nil
}
func (s *stickyGatewayCacheHotpathStub) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) {
return map[int64]int64{}, nil
}
@@ -750,7 +753,7 @@ func TestSelectAccountWithLoadAwareness_StickyReadReuse(t *testing.T) {
modelsListCacheTTL: time.Minute,
}
result, err := svc.SelectAccountWithLoadAwareness(baseCtx, nil, "sess-hash", "", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(baseCtx, nil, "sess-hash", "", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
@@ -772,7 +775,7 @@ func TestSelectAccountWithLoadAwareness_StickyReadReuse(t *testing.T) {
ctx := context.WithValue(baseCtx, ctxkey.PrefetchedStickyAccountID, account.ID)
ctx = context.WithValue(ctx, ctxkey.PrefetchedStickyGroupID, int64(0))
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "sess-hash", "", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "sess-hash", "", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
@@ -794,7 +797,7 @@ func TestSelectAccountWithLoadAwareness_StickyReadReuse(t *testing.T) {
ctx := context.WithValue(baseCtx, ctxkey.PrefetchedStickyAccountID, int64(999))
ctx = context.WithValue(ctx, ctxkey.PrefetchedStickyGroupID, int64(77))
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "sess-hash", "", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "sess-hash", "", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
@@ -235,12 +235,15 @@ func (m *mockGatewayCacheForPlatform) DeleteSessionAccountID(ctx context.Context
return nil
}
func (m *mockGatewayCacheForPlatform) GetClientAffinityAccounts(_ context.Context, _ int64, _ string, _ time.Duration) ([]int64, error) {
func (m *mockGatewayCacheForPlatform) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
return nil, nil
}
func (m *mockGatewayCacheForPlatform) UpdateClientAffinity(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error {
func (m *mockGatewayCacheForPlatform) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error {
return nil
}
func (m *mockGatewayCacheForPlatform) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (int64, int64, int64, error) {
return 0, 0, 0, nil
}
func (m *mockGatewayCacheForPlatform) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) {
return map[int64]int64{}, nil
}
@@ -2050,7 +2053,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: nil, // No concurrency service
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
@@ -2103,7 +2106,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: nil, // legacy path
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, sessionHash, "claude-b", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, sessionHash, "claude-b", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
@@ -2135,7 +2138,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: nil,
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
@@ -2167,13 +2170,314 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
}
excludedIDs := map[int64]struct{}{1: {}}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", excludedIDs, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", excludedIDs, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
require.Equal(t, int64(2), result.Account.ID, "不应选择被排除的账号")
})
t.Run("无客户端ID时过滤Anthropic OAuth亲和账号-load-aware路径", func(t *testing.T) {
repo := &mockAccountRepoForPlatform{
accounts: []Account{
{
ID: 1,
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Priority: 1,
Status: StatusActive,
Schedulable: true,
Concurrency: 5,
Extra: map[string]any{
"client_affinity_enabled": true,
},
},
{
ID: 2,
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Priority: 2,
Status: StatusActive,
Schedulable: true,
Concurrency: 5,
},
},
accountsByID: map[int64]*Account{},
}
for i := range repo.accounts {
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
}
cfg := testConfig()
cfg.Gateway.Scheduling.LoadBatchEnabled = true
svc := &GatewayService{
accountRepo: repo,
cache: &mockGatewayCacheForPlatform{},
cfg: cfg,
concurrencyService: NewConcurrencyService(&mockConcurrencyCache{}),
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, "", 0)
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
require.Equal(t, int64(2), result.Account.ID, "无客户端ID时应过滤亲和OAuth账号")
})
t.Run("无客户端ID且亲和关闭时-load-aware路径恢复旧逻辑", func(t *testing.T) {
repo := &mockAccountRepoForPlatform{
accounts: []Account{
{
ID: 1,
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Priority: 1,
Status: StatusActive,
Schedulable: true,
Concurrency: 5,
Extra: map[string]any{
"affinity_enabled": false,
},
},
{
ID: 2,
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Priority: 2,
Status: StatusActive,
Schedulable: true,
Concurrency: 5,
},
},
accountsByID: map[int64]*Account{},
}
for i := range repo.accounts {
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
}
cfg := testConfig()
cfg.Gateway.Scheduling.LoadBatchEnabled = true
svc := &GatewayService{
accountRepo: repo,
cache: &mockGatewayCacheForPlatform{},
cfg: cfg,
concurrencyService: NewConcurrencyService(&mockConcurrencyCache{}),
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, "", 0)
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
require.Equal(t, int64(1), result.Account.ID, "关闭亲和后不应再被无clientID过滤")
})
t.Run("有客户端ID时不过滤Anthropic OAuth亲和账号", func(t *testing.T) {
repo := &mockAccountRepoForPlatform{
accounts: []Account{
{
ID: 1,
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Priority: 1,
Status: StatusActive,
Schedulable: true,
Concurrency: 5,
Extra: map[string]any{
"client_affinity_enabled": true,
},
},
{
ID: 2,
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Priority: 2,
Status: StatusActive,
Schedulable: true,
Concurrency: 5,
},
},
accountsByID: map[int64]*Account{},
}
for i := range repo.accounts {
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
}
cfg := testConfig()
cfg.Gateway.Scheduling.LoadBatchEnabled = true
svc := &GatewayService{
accountRepo: repo,
cache: &mockGatewayCacheForPlatform{},
cfg: cfg,
concurrencyService: NewConcurrencyService(&mockConcurrencyCache{}),
}
metadataUserID := "user_0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef_account_test_session_test"
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, metadataUserID, 123)
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
require.Equal(t, int64(1), result.Account.ID, "有客户端ID时应允许亲和OAuth账号参与调度")
})
t.Run("过滤对Anthropic全类型生效-SetupToken也受影响", func(t *testing.T) {
repo := &mockAccountRepoForPlatform{
accounts: []Account{
{
ID: 1,
Platform: PlatformAnthropic,
Type: AccountTypeSetupToken,
Priority: 1,
Status: StatusActive,
Schedulable: true,
Concurrency: 5,
Extra: map[string]any{
"client_affinity_enabled": true,
},
},
{
ID: 2,
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Priority: 2,
Status: StatusActive,
Schedulable: true,
Concurrency: 5,
Extra: map[string]any{
"client_affinity_enabled": true,
},
},
{
ID: 3,
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Priority: 3,
Status: StatusActive,
Schedulable: true,
Concurrency: 5,
},
},
accountsByID: map[int64]*Account{},
}
for i := range repo.accounts {
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
}
cfg := testConfig()
cfg.Gateway.Scheduling.LoadBatchEnabled = true
svc := &GatewayService{
accountRepo: repo,
cache: &mockGatewayCacheForPlatform{},
cfg: cfg,
concurrencyService: NewConcurrencyService(&mockConcurrencyCache{}),
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, "", 0)
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
require.Equal(t, int64(3), result.Account.ID, "无客户端ID时,Anthropic SetupToken/OAuth(开启亲和)都应被过滤")
})
t.Run("无客户端ID时过滤Anthropic OAuth亲和账号-legacy路径", func(t *testing.T) {
repo := &mockAccountRepoForPlatform{
accounts: []Account{
{
ID: 1,
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Priority: 1,
Status: StatusActive,
Schedulable: true,
Concurrency: 5,
Extra: map[string]any{
"client_affinity_enabled": true,
},
},
{
ID: 2,
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Priority: 2,
Status: StatusActive,
Schedulable: true,
Concurrency: 5,
},
},
accountsByID: map[int64]*Account{},
}
for i := range repo.accounts {
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
}
cfg := testConfig()
cfg.Gateway.Scheduling.LoadBatchEnabled = false
svc := &GatewayService{
accountRepo: repo,
cache: &mockGatewayCacheForPlatform{},
cfg: cfg,
concurrencyService: nil,
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, "", 0)
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
require.Equal(t, int64(2), result.Account.ID, "legacy路径也应过滤亲和OAuth账号")
})
t.Run("无客户端ID且亲和关闭时-legacy路径恢复旧逻辑", func(t *testing.T) {
repo := &mockAccountRepoForPlatform{
accounts: []Account{
{
ID: 1,
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Priority: 1,
Status: StatusActive,
Schedulable: true,
Concurrency: 5,
Extra: map[string]any{
"client_affinity_enabled": false,
},
},
{
ID: 2,
Platform: PlatformAnthropic,
Type: AccountTypeOAuth,
Priority: 2,
Status: StatusActive,
Schedulable: true,
Concurrency: 5,
},
},
accountsByID: map[int64]*Account{},
}
for i := range repo.accounts {
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
}
cfg := testConfig()
cfg.Gateway.Scheduling.LoadBatchEnabled = false
svc := &GatewayService{
accountRepo: repo,
cache: &mockGatewayCacheForPlatform{},
cfg: cfg,
concurrencyService: nil,
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, "", 0)
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
require.Equal(t, int64(1), result.Account.ID, "关闭亲和后 legacy 路径不应再被无clientID过滤")
})
t.Run("粘性命中-不调用GetByID", func(t *testing.T) {
repo := &mockAccountRepoForPlatform{
accounts: []Account{
@@ -2201,7 +2505,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: NewConcurrencyService(concurrencyCache),
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "sticky", "claude-3-5-sonnet-20241022", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "sticky", "claude-3-5-sonnet-20241022", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
@@ -2237,7 +2541,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: NewConcurrencyService(concurrencyCache),
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "sticky", "claude-3-5-sonnet-20241022", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "sticky", "claude-3-5-sonnet-20241022", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
@@ -2278,7 +2582,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: NewConcurrencyService(concurrencyCache),
}
result, err := svc.SelectAccountWithLoadAwareness(testCtx, nil, "sticky", "claude-3-5-sonnet-20241022", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(testCtx, nil, "sticky", "claude-3-5-sonnet-20241022", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
@@ -2306,7 +2610,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: nil,
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, "", int64(0))
require.Error(t, err)
require.Nil(t, result)
require.ErrorIs(t, err, ErrNoAvailableAccounts)
@@ -2338,7 +2642,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: nil,
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
@@ -2371,7 +2675,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: nil,
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
@@ -2409,7 +2713,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: NewConcurrencyService(concurrencyCache),
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "sticky", "claude-3-5-sonnet-20241022", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "sticky", "claude-3-5-sonnet-20241022", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.WaitPlan)
@@ -2445,7 +2749,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: NewConcurrencyService(concurrencyCache),
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "legacy", "claude-3-5-sonnet-20241022", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "legacy", "claude-3-5-sonnet-20241022", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
@@ -2504,7 +2808,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: NewConcurrencyService(concurrencyCache),
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, sessionHash, "claude-3-5-sonnet-20241022", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, sessionHash, "claude-3-5-sonnet-20241022", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.WaitPlan)
@@ -2558,7 +2862,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: NewConcurrencyService(concurrencyCache),
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, sessionHash, "claude-3-5-sonnet-20241022", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, sessionHash, "claude-3-5-sonnet-20241022", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
@@ -2612,7 +2916,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: NewConcurrencyService(concurrencyCache),
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, sessionHash, "claude-3-5-sonnet-20241022", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, sessionHash, "claude-3-5-sonnet-20241022", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
@@ -2670,7 +2974,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: NewConcurrencyService(concurrencyCache),
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, "route", "claude-3-5-sonnet-20241022", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, "route", "claude-3-5-sonnet-20241022", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
@@ -2728,7 +3032,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: NewConcurrencyService(concurrencyCache),
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, "route-full", "claude-3-5-sonnet-20241022", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, "route-full", "claude-3-5-sonnet-20241022", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.WaitPlan)
@@ -2786,7 +3090,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: NewConcurrencyService(concurrencyCache),
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, "fallback", "claude-3-5-sonnet-20241022", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, "fallback", "claude-3-5-sonnet-20241022", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
@@ -2823,7 +3127,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: NewConcurrencyService(concurrencyCache),
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.WaitPlan)
@@ -2875,7 +3179,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: NewConcurrencyService(concurrencyCache),
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, "gemini", "gemini-2.5-pro", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, "gemini", "gemini-2.5-pro", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
@@ -2953,7 +3257,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
}
excluded := map[int64]struct{}{1: {}}
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, "", "claude-3-5-sonnet-20241022", excluded, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, "", "claude-3-5-sonnet-20241022", excluded, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
@@ -3007,7 +3311,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: nil,
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, "", "gemini-2.5-pro", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, "", "gemini-2.5-pro", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
@@ -3040,7 +3344,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: nil,
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, "", "claude-3-5-sonnet-20241022", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, "", "claude-3-5-sonnet-20241022", nil, "", int64(0))
require.Error(t, err)
require.Nil(t, result)
require.ErrorIs(t, err, ErrClaudeCodeOnly)
@@ -3078,7 +3382,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: NewConcurrencyService(concurrencyCache),
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "wait", "claude-3-5-sonnet-20241022", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "wait", "claude-3-5-sonnet-20241022", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.WaitPlan)
@@ -3116,7 +3420,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
concurrencyService: NewConcurrencyService(concurrencyCache),
}
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "missing-load", "claude-3-5-sonnet-20241022", nil, "")
result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "missing-load", "claude-3-5-sonnet-20241022", nil, "", int64(0))
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, result.Account)
+71 -100
View File
@@ -364,6 +364,12 @@ var ErrNoAvailableAccounts = errors.New("no available accounts")
// ErrClaudeCodeOnly 表示分组仅允许 Claude Code 客户端访问
var ErrClaudeCodeOnly = errors.New("this group only allows Claude Code clients")
// ErrAffinityNoSwitch 表示亲和账号不可用且不允许切换到其他账号
var ErrAffinityNoSwitch = errors.New("affinity account unavailable and switching is disabled")
// ErrAffinityLimitExceeded 表示亲和客户端限制已达上限
var ErrAffinityLimitExceeded = errors.New("affinity client limit exceeded")
// allowedHeaders 白名单headers(参考CRS项目)
var allowedHeaders = map[string]bool{
"accept": true,
@@ -405,24 +411,29 @@ type GatewayCache interface {
// 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
// GetAccountAffinityCountBatch 批量获取账号的亲和客户端数量(惰性清理过期成员)
// GetAffinityAccounts 获取亲和账号列表(按最近使用降序),同时清理过期成员
GetAffinityAccounts(ctx context.Context, groupID int64, userID int64, clientID string, ttl time.Duration) ([]int64, error)
// UpdateAffinity 添加/更新亲和关系(更新 score 为当前时间戳,刷新 key TTL)
UpdateAffinity(ctx context.Context, groupID int64, userID int64, clientID string, accountID int64, ttl time.Duration) error
// GetAccountAffinityCountBatch 批量获取账号的亲和成员数量(惰性清理过期成员)
GetAccountAffinityCountBatch(ctx context.Context, groupID int64, accountIDs []int64, ttl time.Duration) (map[int64]int64, error)
// GetAccountAffinityClientsBatch 批量获取每个账号跨所有分组的亲和客户端列表(去重)
// GetAccountAffinityClientsBatch 批量获取每个账号跨所有分组的亲和成员列表(去重)
// accountGroups: map[accountID][]groupID
// 返回值成员格式为 {userID}/{clientID}
GetAccountAffinityClientsBatch(ctx context.Context, accountGroups map[int64][]int64, ttl time.Duration) (map[int64][]string, error)
// GetAccountAffinityClientsWithScores 获取单个账号跨所有分组的亲和客户端列表(含最后活跃时间)
GetAccountAffinityClientsWithScores(ctx context.Context, accountID int64, groupIDs []int64, ttl time.Duration) ([]AffinityClient, error)
// ClearAccountAffinity 清除指定账号在所有分组的亲和记录(正向+反向索引)
// 用于账号关闭客户端亲和时立即清理旧绑定
// 用于账号关闭亲和时立即清理旧绑定
ClearAccountAffinity(ctx context.Context, accountID int64, groupIDs []int64) error
// GetAffinityMultiCount 获取账号的多维度亲和计数
// 返回: uniqueUsers, uniqueClients, perUserClients
GetAffinityMultiCount(ctx context.Context, groupID int64, accountID int64, targetUserID int64, ttl time.Duration) (users, clients, perUser int64, err error)
}
// AffinityClient 亲和客户端信息(含最后活跃时间)
// AffinityClient 亲和客户端信息(含用户 ID 和最后活跃时间)
type AffinityClient struct {
UserID int64 `json:"user_id"`
ClientID string `json:"client_id"`
LastActive time.Time `json:"last_active"`
}
@@ -1148,8 +1159,10 @@ func (s *GatewayService) SelectAccountForModelWithExclusions(ctx context.Context
}
// SelectAccountWithLoadAwareness selects account with load-awareness and wait plan.
// 调度流程文档见 docs/ACCOUNT_SCHEDULING_FLOW.md 。
// metadataUserID: 用于客户端亲和调度,从中提取客户端 ID
func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}, metadataUserID string) (*AccountSelectionResult, error) {
// sub2apiUserID: 系统用户 ID,用于二维亲和调度
func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}, metadataUserID string, sub2apiUserID int64) (*AccountSelectionResult, error) {
// 调试日志:记录调度入口参数
excludedIDsList := make([]int64, 0, len(excludedIDs))
for id := range excludedIDs {
@@ -1181,6 +1194,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
// 提取客户端 ID(用于客户端亲和调度)
affinityClientID := extractClientIDFromMetadata(metadataUserID)
affinityUserID := sub2apiUserID
if s.debugModelRoutingEnabled() && requestedModel != "" {
groupPlatform := ""
@@ -1203,6 +1217,10 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
if err != nil {
return nil, err
}
if shouldFilterAccountWithoutClientID(account, affinityClientID) {
localExcluded[account.ID] = struct{}{}
continue
}
result, err := s.tryAcquireAccountSlot(ctx, account.ID, account.Concurrency)
if err == nil && result.Acquired {
@@ -1264,12 +1282,18 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
if err != nil {
return nil, err
}
accounts = filterAccountsWithoutClientID(accounts, affinityClientID)
if len(accounts) == 0 {
return nil, ErrNoAvailableAccounts
}
ctx = s.withWindowCostPrefetch(ctx, accounts)
ctx = s.withRPMPrefetch(ctx, accounts)
// 提前构建 accountByID(供 Layer 1 和 Layer 1.5 使用)
accountByID := make(map[int64]*Account, len(accounts))
for i := range accounts {
accountByID[accounts[i].ID] = &accounts[i]
}
isExcluded := func(accountID int64) bool {
if excludedIDs == nil {
return false
@@ -1277,12 +1301,19 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
_, excluded := excludedIDs[accountID]
return excluded
}
// 提前构建 accountByID(供 Layer 1 和 Layer 1.5 使用)
accountByID := make(map[int64]*Account, len(accounts))
for i := range accounts {
accountByID[accounts[i].ID] = &accounts[i]
}
affinityFlow := newGatewayAffinityFlow(
s,
ctx,
groupID,
sessionHash,
requestedModel,
affinityClientID,
affinityUserID,
platform,
useMixed,
accountByID,
isExcluded,
)
// 获取模型路由配置(仅 anthropic 平台)
var routingAccountIDs []int64
@@ -1485,8 +1516,8 @@ 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, ClientAffinityTTL)
if affinityClientID != "" && affinityUserID > 0 && s.cache != nil && item.account.IsAffinityEnabled() {
_ = s.cache.UpdateAffinity(ctx, derefGroupID(groupID), affinityUserID, affinityClientID, item.account.ID, ClientAffinityTTL)
}
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)
@@ -1525,14 +1556,27 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
}
}
// ============ Layer 1.5: 粘性会话(仅在无模型路由配置时生效) ============
if len(routingAccountIDs) == 0 && sessionHash != "" && stickyAccountID > 0 && !isExcluded(stickyAccountID) {
// ============ Layer 1.3: 用户亲和预处理(pinned_users 自动注入) ============
affinityFlow.preprocessPinnedUsers(accounts)
// ============ Layer 1.4: 客户端亲和调度(优先于粘性会话) ============
affinityHit := false
if affinityResult, hit, err := affinityFlow.trySelectAffinityAccount(); err != nil {
return nil, err
} else {
affinityHit = hit
if affinityResult != nil {
return affinityResult, nil
}
}
// ============ Layer 1.5: 粘性会话(仅在无模型路由配置 且 亲和未命中时生效) ============
if !affinityHit && len(routingAccountIDs) == 0 && sessionHash != "" && stickyAccountID > 0 && !isExcluded(stickyAccountID) {
accountID := stickyAccountID
if accountID > 0 && !isExcluded(accountID) {
account, ok := accountByID[accountID]
if ok {
// 检查账户是否需要清理粘性会话绑定
// Check if the account needs sticky session cleanup
clearSticky := shouldClearStickySession(account, requestedModel)
if clearSticky {
_ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash)
@@ -1548,7 +1592,6 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
result, err := s.tryAcquireAccountSlot(ctx, accountID, account.Concurrency)
if err == nil && result.Acquired {
// 会话数量限制检查
// Session count limit check
if !s.checkAndRegisterSession(ctx, account, sessionHash) {
result.ReleaseFunc() // 释放槽位,继续到 Layer 2
} else {
@@ -1563,10 +1606,8 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
waitingCount, _ := s.concurrencyService.GetAccountWaitingCount(ctx, accountID)
if waitingCount < cfg.StickySessionMaxWaiting {
// 会话数量限制检查(等待计划也需要占用会话配额)
// Session count limit check (wait plan also requires session quota)
if !s.checkAndRegisterSession(ctx, account, sessionHash) {
// 会话限制已满,继续到 Layer 2
// Session limit full, continue to Layer 2
} else {
return &AccountSelectionResult{
Account: account,
@@ -1584,76 +1625,6 @@ 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, ClientAffinityTTL)
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
}
// 亲和三区检查:红区账号不可通过亲和命中调度
if account.GetAffinityBase() > 0 && s.cache != nil {
countMap, err := s.cache.GetAccountAffinityCountBatch(ctx, derefGroupID(groupID), []int64{affinityAccID}, ClientAffinityTTL)
if err == nil {
zone := account.GetAffinityZone(countMap[affinityAccID])
if zone == AffinityZoneRed {
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, ClientAffinityTTL)
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 {
@@ -1706,8 +1677,8 @@ 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, ClientAffinityTTL)
if affinityClientID != "" && affinityUserID > 0 && s.cache != nil && result.Account != nil && result.Account.IsAffinityEnabled() {
_ = s.cache.UpdateAffinity(ctx, derefGroupID(groupID), affinityUserID, affinityClientID, result.Account.ID, ClientAffinityTTL)
}
return result, nil
}
@@ -1771,9 +1742,9 @@ 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, ClientAffinityTTL)
// 更新亲和关系
if affinityClientID != "" && affinityUserID > 0 && s.cache != nil && selected.account.IsAffinityEnabled() {
_ = s.cache.UpdateAffinity(ctx, derefGroupID(groupID), affinityUserID, affinityClientID, selected.account.ID, ClientAffinityTTL)
}
return &AccountSelectionResult{
Account: selected.account,
@@ -2541,7 +2512,7 @@ func (s *GatewayService) populateAffinityCounts(ctx context.Context, accounts []
// 快速检查:是否有任何账号开启了亲和
hasAffinity := false
for _, acc := range accounts {
if acc.account.IsClientAffinityEnabled() {
if acc.account.IsAffinityEnabled() {
hasAffinity = true
break
}
@@ -2632,7 +2603,7 @@ func classifyByAffinityZone(accounts []accountWithLoad) []accountWithLoad {
// 快速检查:是否有任何账号配置了 affinity_base
hasZoneConfig := false
for _, acc := range accounts {
if acc.account.IsClientAffinityEnabled() && acc.account.GetAffinityBase() > 0 {
if acc.account.IsAffinityEnabled() && acc.account.GetAffinityBase() > 0 {
hasZoneConfig = true
break
}
@@ -288,12 +288,15 @@ func (m *mockGatewayCacheForGemini) DeleteSessionAccountID(ctx context.Context,
return nil
}
func (m *mockGatewayCacheForGemini) GetClientAffinityAccounts(_ context.Context, _ int64, _ string, _ time.Duration) ([]int64, error) {
func (m *mockGatewayCacheForGemini) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
return nil, nil
}
func (m *mockGatewayCacheForGemini) UpdateClientAffinity(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error {
func (m *mockGatewayCacheForGemini) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error {
return nil
}
func (m *mockGatewayCacheForGemini) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (int64, int64, int64, error) {
return 0, 0, 0, nil
}
func (m *mockGatewayCacheForGemini) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) {
return map[int64]int64{}, nil
}
@@ -282,12 +282,15 @@ func (c *stubGatewayCache) DeleteSessionAccountID(ctx context.Context, groupID i
return nil
}
func (c *stubGatewayCache) GetClientAffinityAccounts(_ context.Context, _ int64, _ string, _ time.Duration) ([]int64, error) {
func (c *stubGatewayCache) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
return nil, nil
}
func (c *stubGatewayCache) UpdateClientAffinity(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error {
func (c *stubGatewayCache) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error {
return nil
}
func (c *stubGatewayCache) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (int64, int64, int64, error) {
return 0, 0, 0, nil
}
func (c *stubGatewayCache) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) {
return map[int64]int64{}, nil
}
@@ -193,12 +193,15 @@ func (c *openAIWSStateStoreTimeoutProbeCache) DeleteSessionAccountID(ctx context
return nil
}
func (c *openAIWSStateStoreTimeoutProbeCache) GetClientAffinityAccounts(_ context.Context, _ int64, _ string, _ time.Duration) ([]int64, error) {
func (c *openAIWSStateStoreTimeoutProbeCache) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
return nil, nil
}
func (c *openAIWSStateStoreTimeoutProbeCache) UpdateClientAffinity(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error {
func (c *openAIWSStateStoreTimeoutProbeCache) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error {
return nil
}
func (c *openAIWSStateStoreTimeoutProbeCache) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (int64, int64, int64, error) {
return 0, 0, 0, nil
}
func (c *openAIWSStateStoreTimeoutProbeCache) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) {
return map[int64]int64{}, nil
}
+1 -1
View File
@@ -519,7 +519,7 @@ func (s *OpsService) selectAccountForRetry(ctx context.Context, reqType opsRetry
if s.gatewayService == nil {
return nil, fmt.Errorf("gateway service not available")
}
return s.gatewayService.SelectAccountWithLoadAwareness(ctx, groupID, "", model, excludedIDs, "") // 重试不使用会话限制
return s.gatewayService.SelectAccountWithLoadAwareness(ctx, groupID, "", model, excludedIDs, "", int64(0)) // 重试不使用会话限制
default:
return nil, fmt.Errorf("unsupported retry type: %s", reqType)
}
+5 -2
View File
@@ -100,12 +100,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) {
func (c StubGatewayCache) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) {
return nil, nil
}
func (c StubGatewayCache) UpdateClientAffinity(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error {
func (c StubGatewayCache) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error {
return nil
}
func (c StubGatewayCache) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (int64, int64, int64, error) {
return 0, 0, 0, nil
}
func (c StubGatewayCache) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) {
return map[int64]int64{}, nil
}
+61
View File
@@ -0,0 +1,61 @@
# Account Scheduling Flow(SelectAccountWithLoadAwareness)
本文档对应后端主调度入口 `SelectAccountWithLoadAwareness`,用于说明 Anthropic 网关的账号选择流程、亲和策略与关键过滤规则。
## 代码入口与引用
- 调度主入口:`backend/internal/service/gateway_service.go` 中的 `SelectAccountWithLoadAwareness`
- 亲和详情 API:`backend/internal/handler/admin/account_affinity_handler.go` 中的 `GetAffinityDetails`
- Admin 路由:`backend/internal/server/routes/admin.go`(`/api/v1/admin/accounts/:id/affinity-details`)
## 流程概览(修正版)
1. 解析输入参数
`sessionHash`、`stickyAccountID`、`affinityClientID(metadata.user_id 提取)`、`affinityUserID(sub2api user id)`。
2. Claude Code 限制与分组降级
先执行 `checkClaudeCodeRestriction`,必要时替换 `groupID`。
3. Legacy / Load-aware 分支
若 `concurrencyService == nil || !LoadBatchEnabled` 走传统路径;否则走负载感知路径。
4. Layer 1(模型路由优先)
- 路由候选过滤(排除、平台、模型、配额、窗口费用、RPM)
- 路由范围内 sticky 优先
- 路由内按(优先级 > 负载 > 亲和数 > LRU)尝试获取槽位
5. Layer 1.3(pinned_users 预处理)
仅做 `UpdateAffinity` 预热,不直接做调度决策。
6. Layer 1.4(客户端亲和调度)
按亲和记录尝试命中;`allow_switch=false` 且无一票放行时尝试等待计划,否则返回 `ErrAffinityNoSwitch`。
7. Layer 1.5(粘性会话)
仅在 `!affinityHit && routingAccountIDs==0` 时生效;先过 `shouldClearStickySession`。
8. Layer 2(负载感知选择)
分层过滤:优先级 -> 亲和区(单维客户端) -> 最低负载 -> 最少亲和客户端 -> LRU。
9. Layer 3(兜底排队)
按 `FallbackSelectionMode`(`last_used` 或 `random`)排序后返回等待计划。
注意:此层不是“按最低负载”。
## 关键业务规则(2026-03)
### 无客户端 ID 时过滤亲和 Anthropic 账号
当 `metadata.user_id` 无法提取 `client_id`(即 `affinityClientID == ""`)时:
- 会直接过滤 **Anthropic 平台且开启客户端亲和** 的账号(覆盖 OAuth / SetupToken / API Key / Bedrock 等类型)
- 该规则只看平台与亲和开关,不再限定账号类型
设计目标:避免无客户端标识的请求误用客户端亲和账号。
## 亲和详情 API 契约
`GET /api/v1/admin/accounts/:id/affinity-details` 返回结构(后端已对齐前端):
- `users[]`:包含 `user_id`、`user_email`、`client_count`、`is_pinned`、`clients[]`
- `total_users`
- `total_clients`
- `pinned_users`
+16 -2
View File
@@ -17,7 +17,8 @@ import type {
AdminDataPayload,
AdminDataImportResult,
CheckMixedChannelRequest,
CheckMixedChannelResponse
CheckMixedChannelResponse,
AffinityDetailsResponse
} from '@/types'
/**
@@ -630,6 +631,18 @@ export async function getAffinityClients(id: number): Promise<{ client_id: strin
return data
}
/**
* Get affinity details for an account with user-level grouping
* @param id - Account ID
* @returns Affinity details with user groups
*/
export async function getAffinityDetails(id: number): Promise<AffinityDetailsResponse> {
const { data } = await apiClient.get<AffinityDetailsResponse>(
`/admin/accounts/${id}/affinity-details`
)
return data
}
export const accountsAPI = {
list,
listWithEtag,
@@ -667,7 +680,8 @@ export const accountsAPI = {
getAntigravityDefaultModelMapping,
batchClearError,
batchRefresh,
getAffinityClients
getAffinityClients,
getAffinityDetails
}
export default accountsAPI
@@ -29,7 +29,7 @@
</CapacityBadge>
<!-- 客户端亲和 -->
<AffinityBadge v-if="showAffinity" :account-id="account.id" :count="affinityClientCount" :base="affinityBase" :buffer="affinityBuffer" />
<AffinityBadge v-if="showAffinity" :account-id="account.id" :client-count="affinityClientCount" :user-count="affinityUserCount" :base="affinityBase" :buffer="affinityBuffer" :user-base="affinityUserBase" :user-buffer="affinityUserBuffer" />
<!-- API Key 账号配额限制 -->
<QuotaBadge v-if="showDailyQuota" :used="account.quota_daily_used ?? 0" :limit="account.quota_daily_limit!" label="D" />
@@ -177,10 +177,11 @@ const rpmTooltip = computed(() => {
const showAffinity = computed(() =>
props.account.platform === 'anthropic' &&
props.account.client_affinity_enabled === true &&
props.account.affinity_client_count != null
(props.account.affinity_client_count != null || props.account.affinity_user_count != null)
)
const affinityClientCount = computed(() => props.account.affinity_client_count ?? 0)
const affinityUserCount = computed(() => props.account.affinity_user_count ?? 0)
const affinityBase = computed(() => {
const extra = props.account.extra as Record<string, unknown> | undefined
@@ -197,6 +198,21 @@ const affinityBuffer = computed((): number | null => {
return null
})
const affinityUserBase = computed(() => {
const extra = props.account.extra as Record<string, unknown> | undefined
if (!extra) return 0
const v = extra.affinity_user_base
return (typeof v === 'number' && v > 0) ? v : 0
})
const affinityUserBuffer = computed((): number | null => {
const extra = props.account.extra as Record<string, unknown> | undefined
if (!extra) return null
const v = extra.affinity_user_buffer
if (typeof v === 'number') return v
return null
})
// 格式化费用显示
const formatCost = (value: number | null | undefined) => {
if (value === null || value === undefined) return '0'
+106 -67
View File
@@ -11,9 +11,24 @@
<svg class="h-2.5 w-2.5" fill="none" viewBox="0 0 24 24" stroke="currentColor" stroke-width="2">
<path stroke-linecap="round" stroke-linejoin="round" d="M13.19 8.688a4.5 4.5 0 011.242 7.244l-4.5 4.5a4.5 4.5 0 01-6.364-6.364l1.757-1.757m13.35-.622l1.757-1.757a4.5 4.5 0 00-6.364-6.364l-4.5 4.5a4.5 4.5 0 001.242 7.244" />
</svg>
<span class="font-mono">{{ count }}</span>
<span class="text-gray-400 dark:text-gray-500">/</span>
<span class="font-mono">{{ limitDisplay }}</span>
<!-- Dual dimension display -->
<template v-if="showUserDimension && showClientDimension">
<span class="font-mono">U{{ userCount }}/{{ userLimitDisplay }}</span>
<span class="text-gray-400 dark:text-gray-500 mx-px">|</span>
<span class="font-mono">C{{ clientCount }}/{{ clientLimitDisplay }}</span>
</template>
<!-- Single dimension: user only -->
<template v-else-if="showUserDimension">
<span class="font-mono">{{ userCount }}</span>
<span class="text-gray-400 dark:text-gray-500">/</span>
<span class="font-mono">{{ userLimitDisplay }}</span>
</template>
<!-- Single dimension: client only (original) -->
<template v-else>
<span class="font-mono">{{ clientCount }}</span>
<span class="text-gray-400 dark:text-gray-500">/</span>
<span class="font-mono">{{ clientLimitDisplay }}</span>
</template>
</span>
<!-- Popover -->
@@ -28,7 +43,7 @@
>
<div
v-if="showPopover"
class="fixed z-50 w-72 rounded-lg border border-gray-200 bg-white shadow-lg dark:border-dark-600 dark:bg-dark-800"
class="fixed z-50 w-80 rounded-lg border border-gray-200 bg-white shadow-lg dark:border-dark-600 dark:bg-dark-800"
:style="popoverStyle"
@mouseenter="handlePopoverEnter"
@mouseleave="handlePopoverLeave"
@@ -36,7 +51,7 @@
<!-- Header -->
<div class="flex items-center justify-between border-b border-gray-100 px-3 py-2 dark:border-dark-700">
<span class="text-xs font-medium text-gray-500 dark:text-gray-400">
{{ t('admin.accounts.affinityClients', { count }) }}
{{ t('admin.accounts.affinityDetailTitle') }}
</span>
<span v-if="loading" class="text-xs text-gray-400">
<svg class="h-3.5 w-3.5 animate-spin" viewBox="0 0 24 24" fill="none">
@@ -46,26 +61,43 @@
</span>
</div>
<!-- Client list -->
<!-- Summary -->
<div v-if="details" class="flex items-center gap-3 border-b border-gray-100 px-3 py-1.5 dark:border-dark-700 text-[10px] text-gray-500 dark:text-gray-400">
<span v-if="showUserDimension">{{ t('admin.accounts.affinityUsers', { count: details.total_users }) }}</span>
<span>{{ t('admin.accounts.affinityClients', { count: details.total_clients }) }}</span>
</div>
<!-- User groups tree -->
<div class="max-h-60 overflow-y-auto">
<div v-if="loading && clients.length === 0" class="px-3 py-4 text-center text-xs text-gray-400">
<div v-if="loading && !details" class="px-3 py-4 text-center text-xs text-gray-400">
{{ t('common.loading') }}...
</div>
<div v-else-if="clients.length === 0" class="px-3 py-4 text-center text-xs text-gray-400">
{{ t('admin.accounts.affinityNoClients') }}
<div v-else-if="!details || details.users.length === 0" class="px-3 py-4 text-center text-xs text-gray-400">
{{ t('admin.accounts.affinityNoUsers') }}
</div>
<div v-else class="divide-y divide-gray-50 dark:divide-dark-700">
<div
v-for="(client, index) in clients"
:key="index"
class="flex items-center justify-between px-3 py-1.5"
>
<span class="font-mono text-xs text-gray-700 dark:text-gray-300 truncate mr-2" :title="client.client_id">
{{ client.client_id }}
</span>
<span class="text-[10px] text-gray-400 dark:text-gray-500 whitespace-nowrap shrink-0">
{{ formatRelativeTime(client.last_active) }}
</span>
<div v-for="userGroup in details.users" :key="userGroup.user_id" class="py-1">
<!-- User row -->
<div class="flex items-center justify-between px-3 py-1">
<div class="flex items-center gap-1.5">
<span v-if="userGroup.is_pinned" class="text-[10px]" :title="t('admin.accounts.affinityPinnedLabel')">P</span>
<span class="text-xs font-medium text-gray-700 dark:text-gray-300 truncate" :title="userGroup.user_email">
{{ userGroup.user_email || `User #${userGroup.user_id}` }}
</span>
</div>
<span class="text-[10px] text-gray-400 dark:text-gray-500 shrink-0">
{{ t('admin.accounts.affinityClientCountLabel', { count: userGroup.client_count }) }}
</span>
</div>
<!-- Client rows under user -->
<div v-for="(client, idx) in userGroup.clients" :key="idx" class="flex items-center justify-between px-3 pl-7 py-0.5">
<span class="font-mono text-[10px] text-gray-500 dark:text-gray-400 truncate mr-2" :title="client.client_id">
{{ client.client_id }}
</span>
<span class="text-[10px] text-gray-400 dark:text-gray-500 whitespace-nowrap shrink-0">
{{ formatRelativeTime(client.last_active) }}
</span>
</div>
</div>
</div>
</div>
@@ -78,66 +110,76 @@
<script setup lang="ts">
import { ref, computed } from 'vue'
import { useI18n } from 'vue-i18n'
import { getAffinityClients } from '@/api/admin/accounts'
import { getAffinityDetails } from '@/api/admin/accounts'
import type { AffinityDetailsResponse } from '@/types'
interface Props {
accountId: number
count: number
clientCount: number
userCount: number
base: number // 0 = not configured
buffer: number | null // null = infinite yellow
userBase: number // 0 = not configured
userBuffer: number | null
}
const props = defineProps<Props>()
const props = withDefaults(defineProps<Props>(), {
userCount: 0,
userBase: 0,
userBuffer: null
})
const { t } = useI18n()
const containerRef = ref<HTMLElement | null>(null)
const showPopover = ref(false)
const loading = ref(false)
const clients = ref<{ client_id: string; last_active: string }[]>([])
const details = ref<AffinityDetailsResponse | null>(null)
let loaded = false
let hideTimer: ReturnType<typeof setTimeout> | null = null
let showTimer: ReturnType<typeof setTimeout> | null = null
// 显示的上限文本
const limitDisplay = computed(() => {
if (props.base <= 0) return '\u221E' // ∞
if (props.buffer === null) return `${props.base}+` // base + infinite yellow
// Dimension visibility
const showClientDimension = computed(() => props.base > 0 || props.clientCount > 0)
const showUserDimension = computed(() => props.userBase > 0 || props.userCount > 0)
// Client limit display
const clientLimitDisplay = computed(() => {
if (props.base <= 0) return '\u221E' // infinity
if (props.buffer === null) return `${props.base}+`
if (props.buffer === 0) return `${props.base}`
return `${props.base + props.buffer}`
})
// 根据三区模型着色
// User limit display
const userLimitDisplay = computed(() => {
if (props.userBase <= 0) return '\u221E'
if (props.userBuffer === null) return `${props.userBase}+`
if (props.userBuffer === 0) return `${props.userBase}`
return `${props.userBase + props.userBuffer}`
})
// Zone calculation for a dimension
function calcZone(count: number, base: number, buffer: number | null): number {
if (base <= 0) return 0 // no limit = green
if (count <= base) return 0 // green
if (buffer === null) return 1 // infinite yellow
if (buffer === 0) return 2 // no yellow, red
if (count <= base + buffer) return 1 // yellow
return 2 // red
}
const clientZone = computed(() => calcZone(props.clientCount, props.base, props.buffer))
const userZone = computed(() => calcZone(props.userCount, props.userBase, props.userBuffer))
const overallZone = computed(() => Math.max(userZone.value, clientZone.value))
const badgeClass = computed(() => {
const c = props.count
const base = props.base
// 未配置 base → 总是 emerald
if (base <= 0) {
if (c > 0) return 'bg-emerald-100 text-emerald-700 dark:bg-emerald-900/30 dark:text-emerald-400'
return 'bg-gray-100 text-gray-600 dark:bg-gray-800 dark:text-gray-400'
const maxCount = Math.max(props.clientCount, props.userCount)
if (maxCount <= 0) return 'bg-gray-100 text-gray-600 dark:bg-gray-800 dark:text-gray-400'
switch (overallZone.value) {
case 2: return 'bg-red-100 text-red-700 dark:bg-red-900/30 dark:text-red-400'
case 1: return 'bg-yellow-100 text-yellow-700 dark:bg-yellow-900/30 dark:text-yellow-400'
default: return 'bg-emerald-100 text-emerald-700 dark:bg-emerald-900/30 dark:text-emerald-400'
}
// 绿区
if (c <= base) {
if (c > 0) return 'bg-emerald-100 text-emerald-700 dark:bg-emerald-900/30 dark:text-emerald-400'
return 'bg-gray-100 text-gray-600 dark:bg-gray-800 dark:text-gray-400'
}
// 黄区
if (props.buffer === null) {
// 无限黄区
return 'bg-yellow-100 text-yellow-700 dark:bg-yellow-900/30 dark:text-yellow-400'
}
if (props.buffer === 0) {
// 无黄区,直接红区
return 'bg-red-100 text-red-700 dark:bg-red-900/30 dark:text-red-400'
}
if (c <= base + props.buffer) {
return 'bg-yellow-100 text-yellow-700 dark:bg-yellow-900/30 dark:text-yellow-400'
}
// 红区
return 'bg-red-100 text-red-700 dark:bg-red-900/30 dark:text-red-400'
})
const popoverStyle = computed(() => {
@@ -149,15 +191,12 @@ const popoverStyle = computed(() => {
let top = rect.bottom + 6
let left = rect.left - 40
// 下方空间不足时显示在上方
if (top + 280 > viewportHeight) {
top = Math.max(8, rect.top - 280)
}
// 右侧空间不足时向左偏移
if (left + 288 > viewportWidth) {
left = Math.max(8, viewportWidth - 296)
if (left + 320 > viewportWidth) {
left = Math.max(8, viewportWidth - 328)
}
// 不超出左边界
if (left < 8) left = 8
return { top: `${top}px`, left: `${left}px` }
@@ -172,7 +211,7 @@ function handleMouseEnter() {
clearTimers()
showTimer = setTimeout(() => {
showPopover.value = true
if (!loaded) fetchClients()
if (!loaded) fetchDetails()
}, 200)
}
@@ -190,13 +229,13 @@ function handlePopoverLeave() {
hideTimer = setTimeout(() => { showPopover.value = false }, 150)
}
async function fetchClients() {
async function fetchDetails() {
loading.value = true
try {
clients.value = await getAffinityClients(props.accountId)
details.value = await getAffinityDetails(props.accountId)
loaded = true
} catch {
clients.value = []
details.value = null
} finally {
loading.value = false
}
@@ -1,6 +1,8 @@
<script setup lang="ts">
import { ref, watch, computed } from 'vue'
import { useI18n } from 'vue-i18n'
import { list as listUsers } from '@/api/admin/users'
import type { AdminUser } from '@/types'
const { t } = useI18n()
@@ -8,12 +10,23 @@ const props = defineProps<{
enabled: boolean
base: number | null
buffer: number | null // null = infinite yellow zone
// New v2 props
allowSwitch: boolean
userBase: number | null
userBuffer: number | null
perUserLimit: number | null
pinnedUsers: number[]
}>()
const emit = defineEmits<{
'update:enabled': [value: boolean]
'update:base': [value: number | null]
'update:buffer': [value: number | null]
'update:allowSwitch': [value: boolean]
'update:userBase': [value: number | null]
'update:userBuffer': [value: number | null]
'update:perUserLimit': [value: number | null]
'update:pinnedUsers': [value: number[]]
}>()
const localEnabled = ref(props.enabled)
@@ -27,20 +40,61 @@ watch(localEnabled, (val) => {
if (!val) {
emit('update:base', null)
emit('update:buffer', null)
emit('update:userBase', null)
emit('update:userBuffer', null)
emit('update:perUserLimit', null)
emit('update:pinnedUsers', [])
}
})
// Green zone: toggle + input
const baseLimitEnabled = ref(props.base != null && props.base > 0)
// ====== Allow Switch ======
const localAllowSwitch = ref(props.allowSwitch)
watch(() => props.allowSwitch, (val) => { localAllowSwitch.value = val })
watch(localAllowSwitch, (val) => { emit('update:allowSwitch', val) })
watch(() => props.base, (val) => {
baseLimitEnabled.value = val != null && val > 0
})
// ====== User Affinity ======
const userLimitEnabled = ref(props.userBase != null && props.userBase > 0)
watch(() => props.userBase, (val) => { userLimitEnabled.value = val != null && val > 0 })
const toggleUserLimit = () => {
userLimitEnabled.value = !userLimitEnabled.value
if (userLimitEnabled.value) {
emit('update:userBase', 5)
} else {
emit('update:userBase', null)
emit('update:userBuffer', null)
emit('update:perUserLimit', null)
}
}
const onUserBaseInput = (e: Event) => {
const raw = (e.target as HTMLInputElement).valueAsNumber
emit('update:userBase', Number.isNaN(raw) ? null : Math.max(1, Math.floor(raw)))
}
const userBufferIsInfinite = ref(props.userBuffer === null || props.userBuffer === undefined)
watch(() => props.userBuffer, (val) => { userBufferIsInfinite.value = val === null || val === undefined })
const toggleUserBufferInfinite = () => {
userBufferIsInfinite.value = !userBufferIsInfinite.value
emit('update:userBuffer', userBufferIsInfinite.value ? null : 3)
}
const onUserBufferInput = (e: Event) => {
const raw = (e.target as HTMLInputElement).valueAsNumber
emit('update:userBuffer', Number.isNaN(raw) ? null : Math.max(0, Math.floor(raw)))
}
const userZonePreview = computed(() => buildZonePreview(props.userBase, props.userBuffer))
// ====== Client Affinity (existing) ======
const baseLimitEnabled = ref(props.base != null && props.base > 0)
watch(() => props.base, (val) => { baseLimitEnabled.value = val != null && val > 0 })
const toggleBaseLimit = () => {
baseLimitEnabled.value = !baseLimitEnabled.value
if (baseLimitEnabled.value) {
emit('update:base', 5) // default base
emit('update:base', 5)
} else {
emit('update:base', null)
emit('update:buffer', null)
@@ -52,20 +106,12 @@ const onBaseInput = (e: Event) => {
emit('update:base', Number.isNaN(raw) ? null : Math.max(1, Math.floor(raw)))
}
// Yellow zone: "unlimited" checkbox + input
const bufferIsInfinite = ref(props.buffer === null || props.buffer === undefined)
watch(() => props.buffer, (val) => {
bufferIsInfinite.value = val === null || val === undefined
})
watch(() => props.buffer, (val) => { bufferIsInfinite.value = val === null || val === undefined })
const toggleBufferInfinite = () => {
bufferIsInfinite.value = !bufferIsInfinite.value
if (bufferIsInfinite.value) {
emit('update:buffer', null)
} else {
emit('update:buffer', 3) // default buffer
}
emit('update:buffer', bufferIsInfinite.value ? null : 3)
}
const onBufferInput = (e: Event) => {
@@ -73,126 +119,410 @@ const onBufferInput = (e: Event) => {
emit('update:buffer', Number.isNaN(raw) ? null : Math.max(0, Math.floor(raw)))
}
// Zone preview
const zonePreview = computed(() => {
const base = props.base ?? 0
if (base <= 0) return null
const buf = props.buffer
const greenMax = base
if (buf === null || buf === undefined) {
return { green: `1~${greenMax}`, yellow: `${greenMax + 1}+`, red: null }
}
if (buf === 0) {
return { green: `1~${greenMax}`, yellow: null, red: `${greenMax + 1}+` }
}
const yellowMax = base + buf
return {
green: `1~${greenMax}`,
yellow: `${greenMax + 1}~${yellowMax}`,
red: `${yellowMax + 1}+`
}
const clientZonePreview = computed(() => buildZonePreview(props.base, props.buffer))
// ====== Per-user Client Limit ======
const perUserEnabled = ref(props.perUserLimit != null && props.perUserLimit > 0)
watch(() => props.perUserLimit, (val) => { perUserEnabled.value = val != null && val > 0 })
const togglePerUserLimit = () => {
perUserEnabled.value = !perUserEnabled.value
emit('update:perUserLimit', perUserEnabled.value ? 3 : null)
}
const onPerUserLimitInput = (e: Event) => {
const raw = (e.target as HTMLInputElement).valueAsNumber
emit('update:perUserLimit', Number.isNaN(raw) ? null : Math.max(1, Math.floor(raw)))
}
// ====== Pinned Users ======
const searchQuery = ref('')
const searchResults = ref<AdminUser[]>([])
const searching = ref(false)
let searchTimer: ReturnType<typeof setTimeout> | null = null
const pinnedUserDisplay = computed(() => {
return props.pinnedUsers.map(id => ({ id }))
})
const searchUsers = () => {
if (searchTimer) clearTimeout(searchTimer)
if (!searchQuery.value.trim()) {
searchResults.value = []
return
}
searchTimer = setTimeout(async () => {
searching.value = true
try {
const resp = await listUsers(1, 10, { search: searchQuery.value.trim() })
searchResults.value = resp.items.filter(
u => !props.pinnedUsers.includes(u.id)
)
} catch {
searchResults.value = []
} finally {
searching.value = false
}
}, 300)
}
const addPinnedUser = (user: AdminUser) => {
if (!props.pinnedUsers.includes(user.id)) {
emit('update:pinnedUsers', [...props.pinnedUsers, user.id])
}
searchQuery.value = ''
searchResults.value = []
}
const removePinnedUser = (userId: number) => {
emit('update:pinnedUsers', props.pinnedUsers.filter(id => id !== userId))
}
// ====== Shared helper ======
function buildZonePreview(base: number | null, buffer: number | null) {
const b = base ?? 0
if (b <= 0) return null
if (buffer === null || buffer === undefined) {
return { green: `1~${b}`, yellow: `${b + 1}+`, red: null }
}
if (buffer === 0) {
return { green: `1~${b}`, yellow: null, red: `${b + 1}+` }
}
const yellowMax = b + buffer
return { green: `1~${b}`, yellow: `${b + 1}~${yellowMax}`, red: `${yellowMax + 1}+` }
}
</script>
<template>
<div class="rounded-lg border border-gray-200 p-4 dark:border-dark-600">
<div class="flex items-center justify-between" :class="{ 'mb-3': localEnabled }">
<div class="mb-4">
<label class="input-label mb-0 text-base font-semibold">{{ t('admin.accounts.affinityConfigTitle') }}</label>
<p class="mt-1 text-xs text-gray-500 dark:text-gray-400">
{{ t('admin.accounts.affinityConfigHint') }}
</p>
</div>
<!-- Section 1: Basic Config -->
<div class="flex items-center justify-between" :class="{ 'mb-3': localEnabled }">
<div>
<label class="input-label mb-0">{{ t('admin.accounts.affinityToggle') }}</label>
<p class="mt-1 text-xs text-gray-500 dark:text-gray-400">
{{ t('admin.accounts.affinityToggleHint') }}
</p>
</div>
<button
type="button"
@click="localEnabled = !localEnabled"
:class="[
'relative inline-flex h-6 w-11 flex-shrink-0 cursor-pointer rounded-full border-2 border-transparent transition-colors duration-200 ease-in-out focus:outline-none focus:ring-2 focus:ring-primary-500 focus:ring-offset-2',
localEnabled ? 'bg-primary-600' : 'bg-gray-200 dark:bg-dark-600'
]"
>
<span
:class="[
'pointer-events-none inline-block h-5 w-5 transform rounded-full bg-white shadow ring-0 transition duration-200 ease-in-out',
localEnabled ? 'translate-x-5' : 'translate-x-0'
]"
/>
</button>
</div>
<div v-if="localEnabled" class="space-y-4">
<!-- Allow Switch toggle -->
<div class="flex items-center justify-between">
<div>
<label class="input-label mb-0">{{ t('admin.accounts.affinityToggle') }}</label>
<label class="input-label mb-0">{{ t('admin.accounts.affinityAllowSwitch') }}</label>
<p class="mt-1 text-xs text-gray-500 dark:text-gray-400">
{{ t('admin.accounts.affinityToggleHint') }}
{{ t('admin.accounts.affinityAllowSwitchHint') }}
</p>
</div>
<button
type="button"
@click="localEnabled = !localEnabled"
@click="localAllowSwitch = !localAllowSwitch"
:class="[
'relative inline-flex h-6 w-11 flex-shrink-0 cursor-pointer rounded-full border-2 border-transparent transition-colors duration-200 ease-in-out focus:outline-none focus:ring-2 focus:ring-primary-500 focus:ring-offset-2',
localEnabled ? 'bg-primary-600' : 'bg-gray-200 dark:bg-dark-600'
'relative inline-flex h-5 w-9 flex-shrink-0 cursor-pointer rounded-full border-2 border-transparent transition-colors duration-200 ease-in-out focus:outline-none',
localAllowSwitch ? 'bg-primary-600' : 'bg-gray-200 dark:bg-dark-600'
]"
>
<span
:class="[
'pointer-events-none inline-block h-5 w-5 transform rounded-full bg-white shadow ring-0 transition duration-200 ease-in-out',
localEnabled ? 'translate-x-5' : 'translate-x-0'
'pointer-events-none inline-block h-4 w-4 transform rounded-full bg-white shadow ring-0 transition duration-200 ease-in-out',
localAllowSwitch ? 'translate-x-4' : 'translate-x-0'
]"
/>
</button>
</div>
<p v-if="!localAllowSwitch" class="text-xs text-amber-600 dark:text-amber-400">
{{ t('admin.accounts.affinityAllowSwitchWarning') }}
</p>
<div v-if="localEnabled" class="space-y-3">
<!-- Green zone toggle + input -->
<div>
<div class="flex items-center justify-between mb-1">
<label class="input-label mb-0">{{ t('admin.accounts.affinityBase') }}</label>
<button
type="button"
@click="toggleBaseLimit"
<!-- Section 2: User Affinity -->
<div class="rounded-lg border border-gray-100 p-3 dark:border-dark-700">
<div class="flex items-center justify-between mb-2">
<div>
<label class="input-label mb-0 text-sm font-medium">{{ t('admin.accounts.affinityUserSection') }}</label>
<p class="mt-0.5 text-xs text-gray-500 dark:text-gray-400">
{{ t('admin.accounts.affinityUserSectionHint') }}
</p>
</div>
<button
type="button"
@click="toggleUserLimit"
:class="[
'relative inline-flex h-5 w-9 flex-shrink-0 cursor-pointer rounded-full border-2 border-transparent transition-colors duration-200 ease-in-out focus:outline-none',
userLimitEnabled ? 'bg-primary-600' : 'bg-gray-200 dark:bg-dark-600'
]"
>
<span
:class="[
'relative inline-flex h-5 w-9 flex-shrink-0 cursor-pointer rounded-full border-2 border-transparent transition-colors duration-200 ease-in-out focus:outline-none',
baseLimitEnabled ? 'bg-primary-600' : 'bg-gray-200 dark:bg-dark-600'
'pointer-events-none inline-block h-4 w-4 transform rounded-full bg-white shadow ring-0 transition duration-200 ease-in-out',
userLimitEnabled ? 'translate-x-4' : 'translate-x-0'
]"
>
<span
/>
</button>
</div>
<div v-if="userLimitEnabled" class="space-y-2">
<!-- User green zone -->
<div>
<label class="input-label mb-0">{{ t('admin.accounts.affinityUserBase') }}</label>
<input
:value="userBase"
@input="onUserBaseInput"
type="number"
min="1"
step="1"
class="input"
:placeholder="t('admin.accounts.affinityBasePlaceholder')"
/>
<p class="input-hint">{{ t('admin.accounts.affinityUserBaseHint') }}</p>
</div>
<!-- User yellow zone -->
<div>
<div class="flex items-center justify-between mb-1">
<label class="input-label mb-0">{{ t('admin.accounts.affinityUserBuffer') }}</label>
<label class="flex items-center gap-1.5 text-xs text-gray-500 dark:text-gray-400 cursor-pointer">
<input
type="checkbox"
:checked="userBufferIsInfinite"
@change="toggleUserBufferInfinite"
class="h-3.5 w-3.5 rounded border-gray-300 text-primary-600"
/>
{{ t('admin.accounts.affinityBufferInfinite') }}
</label>
</div>
<input
v-if="!userBufferIsInfinite"
:value="userBuffer"
@input="onUserBufferInput"
type="number"
min="0"
step="1"
class="input"
:placeholder="t('admin.accounts.affinityBufferPlaceholder')"
/>
<p class="input-hint">{{ t('admin.accounts.affinityUserBufferHint') }}</p>
</div>
<!-- User zone preview -->
<div v-if="userZonePreview" class="flex items-center gap-2 text-xs">
<span class="inline-flex items-center gap-1 rounded-full bg-emerald-100 px-2 py-0.5 text-emerald-700 dark:bg-emerald-900/30 dark:text-emerald-400">
{{ userZonePreview.green }}
</span>
<span v-if="userZonePreview.yellow" class="inline-flex items-center gap-1 rounded-full bg-yellow-100 px-2 py-0.5 text-yellow-700 dark:bg-yellow-900/30 dark:text-yellow-400">
{{ userZonePreview.yellow }}
</span>
<span v-if="userZonePreview.red" class="inline-flex items-center gap-1 rounded-full bg-red-100 px-2 py-0.5 text-red-700 dark:bg-red-900/30 dark:text-red-400">
{{ userZonePreview.red }}
</span>
</div>
<div class="border-t border-gray-100 pt-3 dark:border-dark-700">
<div class="flex items-center justify-between mb-2">
<div>
<label class="input-label mb-0 text-sm font-medium">{{ t('admin.accounts.affinityPerUserLimit') }}</label>
<p class="mt-0.5 text-xs text-gray-500 dark:text-gray-400">
{{ t('admin.accounts.affinityPerUserLimitHint') }}
</p>
</div>
<button
type="button"
@click="togglePerUserLimit"
:class="[
'pointer-events-none inline-block h-4 w-4 transform rounded-full bg-white shadow ring-0 transition duration-200 ease-in-out',
baseLimitEnabled ? 'translate-x-4' : 'translate-x-0'
'relative inline-flex h-5 w-9 flex-shrink-0 cursor-pointer rounded-full border-2 border-transparent transition-colors duration-200 ease-in-out focus:outline-none',
perUserEnabled ? 'bg-primary-600' : 'bg-gray-200 dark:bg-dark-600'
]"
/>
</button>
</div>
<input
v-if="baseLimitEnabled"
:value="base"
@input="onBaseInput"
type="number"
min="1"
step="1"
class="input"
:placeholder="t('admin.accounts.affinityBasePlaceholder')"
/>
<p class="input-hint">{{ baseLimitEnabled ? t('admin.accounts.affinityBaseHint') : t('admin.accounts.affinityBaseOffHint') }}</p>
</div>
<!-- Buffer (yellow zone) - only shown when base is set -->
<div v-if="baseLimitEnabled">
<div class="flex items-center justify-between mb-1">
<label class="input-label mb-0">{{ t('admin.accounts.affinityBuffer') }}</label>
<label class="flex items-center gap-1.5 text-xs text-gray-500 dark:text-gray-400 cursor-pointer">
>
<span
:class="[
'pointer-events-none inline-block h-4 w-4 transform rounded-full bg-white shadow ring-0 transition duration-200 ease-in-out',
perUserEnabled ? 'translate-x-4' : 'translate-x-0'
]"
/>
</button>
</div>
<div v-if="perUserEnabled">
<label class="input-label mb-0">{{ t('admin.accounts.affinityPerUserMax') }}</label>
<input
type="checkbox"
:checked="bufferIsInfinite"
@change="toggleBufferInfinite"
class="h-3.5 w-3.5 rounded border-gray-300 text-primary-600"
:value="perUserLimit"
@input="onPerUserLimitInput"
type="number"
min="1"
step="1"
class="input"
/>
{{ t('admin.accounts.affinityBufferInfinite') }}
</label>
</div>
</div>
<input
v-if="!bufferIsInfinite"
:value="buffer"
@input="onBufferInput"
type="number"
min="0"
step="1"
class="input"
:placeholder="t('admin.accounts.affinityBufferPlaceholder')"
/>
<p class="input-hint">{{ t('admin.accounts.affinityBufferHint') }}</p>
</div>
<!-- Zone preview -->
<div v-if="zonePreview" class="flex items-center gap-2 text-xs">
<span class="inline-flex items-center gap-1 rounded-full bg-emerald-100 px-2 py-0.5 text-emerald-700 dark:bg-emerald-900/30 dark:text-emerald-400">
{{ zonePreview.green }}
</span>
<span v-if="zonePreview.yellow" class="inline-flex items-center gap-1 rounded-full bg-yellow-100 px-2 py-0.5 text-yellow-700 dark:bg-yellow-900/30 dark:text-yellow-400">
{{ zonePreview.yellow }}
</span>
<span v-if="zonePreview.red" class="inline-flex items-center gap-1 rounded-full bg-red-100 px-2 py-0.5 text-red-700 dark:bg-red-900/30 dark:text-red-400">
{{ zonePreview.red }}
</span>
<div class="border-t border-gray-100 pt-3 dark:border-dark-700">
<div class="mb-2">
<label class="input-label mb-0 text-sm font-medium">{{ t('admin.accounts.affinityPinnedUsers') }}</label>
<p class="mt-0.5 text-xs text-gray-500 dark:text-gray-400">
{{ t('admin.accounts.affinityPinnedUsersHint') }}
</p>
</div>
<div class="relative">
<input
v-model="searchQuery"
@input="searchUsers"
type="text"
class="input"
:placeholder="t('admin.accounts.affinityPinnedUsersSearch')"
/>
<div
v-if="searchResults.length > 0"
class="absolute z-10 mt-1 w-full rounded-md border border-gray-200 bg-white shadow-lg dark:border-dark-600 dark:bg-dark-800 max-h-40 overflow-y-auto"
>
<div
v-for="user in searchResults"
:key="user.id"
@click="addPinnedUser(user)"
class="cursor-pointer px-3 py-1.5 text-xs hover:bg-gray-50 dark:hover:bg-dark-700 flex items-center justify-between"
>
<span class="text-gray-700 dark:text-gray-300 truncate">{{ user.email || user.username }}</span>
<span class="text-gray-400 dark:text-gray-500 ml-2 shrink-0">#{{ user.id }}</span>
</div>
</div>
<div v-if="searching" class="absolute right-2 top-1/2 -translate-y-1/2">
<svg class="h-4 w-4 animate-spin text-gray-400" viewBox="0 0 24 24" fill="none">
<circle class="opacity-25" cx="12" cy="12" r="10" stroke="currentColor" stroke-width="4" />
<path class="opacity-75" fill="currentColor" d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4z" />
</svg>
</div>
</div>
<div v-if="pinnedUserDisplay.length > 0" class="mt-2 space-y-1">
<div
v-for="pu in pinnedUserDisplay"
:key="pu.id"
class="flex items-center justify-between rounded bg-gray-50 px-2 py-1 dark:bg-dark-700"
>
<span class="font-mono text-xs text-gray-700 dark:text-gray-300">User #{{ pu.id }}</span>
<button
type="button"
@click="removePinnedUser(pu.id)"
class="text-gray-400 hover:text-red-500 dark:text-gray-500 dark:hover:text-red-400"
>
<svg class="h-3.5 w-3.5" fill="none" viewBox="0 0 24 24" stroke="currentColor" stroke-width="2">
<path stroke-linecap="round" stroke-linejoin="round" d="M6 18L18 6M6 6l12 12" />
</svg>
</button>
</div>
</div>
<p v-else class="mt-2 text-xs text-gray-400 dark:text-gray-500">
{{ t('admin.accounts.affinityPinnedUsersEmpty') }}
</p>
</div>
</div>
</div>
<!-- Section 3: Client Affinity -->
<div class="rounded-lg border border-gray-100 p-3 dark:border-dark-700">
<div class="mb-2">
<label class="input-label mb-0 text-sm font-medium">{{ t('admin.accounts.affinityClientSection') }}</label>
<p class="mt-0.5 text-xs text-gray-500 dark:text-gray-400">
{{ t('admin.accounts.affinityClientSectionHint') }}
</p>
</div>
<div class="space-y-2">
<!-- Client green zone toggle + input -->
<div>
<div class="flex items-center justify-between mb-1">
<label class="input-label mb-0">{{ t('admin.accounts.affinityBase') }}</label>
<button
type="button"
@click="toggleBaseLimit"
:class="[
'relative inline-flex h-5 w-9 flex-shrink-0 cursor-pointer rounded-full border-2 border-transparent transition-colors duration-200 ease-in-out focus:outline-none',
baseLimitEnabled ? 'bg-primary-600' : 'bg-gray-200 dark:bg-dark-600'
]"
>
<span
:class="[
'pointer-events-none inline-block h-4 w-4 transform rounded-full bg-white shadow ring-0 transition duration-200 ease-in-out',
baseLimitEnabled ? 'translate-x-4' : 'translate-x-0'
]"
/>
</button>
</div>
<input
v-if="baseLimitEnabled"
:value="base"
@input="onBaseInput"
type="number"
min="1"
step="1"
class="input"
:placeholder="t('admin.accounts.affinityBasePlaceholder')"
/>
<p class="input-hint">{{ baseLimitEnabled ? t('admin.accounts.affinityBaseHint') : t('admin.accounts.affinityBaseOffHint') }}</p>
</div>
<!-- Client buffer (yellow zone) -->
<div v-if="baseLimitEnabled">
<div class="flex items-center justify-between mb-1">
<label class="input-label mb-0">{{ t('admin.accounts.affinityBuffer') }}</label>
<label class="flex items-center gap-1.5 text-xs text-gray-500 dark:text-gray-400 cursor-pointer">
<input
type="checkbox"
:checked="bufferIsInfinite"
@change="toggleBufferInfinite"
class="h-3.5 w-3.5 rounded border-gray-300 text-primary-600"
/>
{{ t('admin.accounts.affinityBufferInfinite') }}
</label>
</div>
<input
v-if="!bufferIsInfinite"
:value="buffer"
@input="onBufferInput"
type="number"
min="0"
step="1"
class="input"
:placeholder="t('admin.accounts.affinityBufferPlaceholder')"
/>
<p class="input-hint">{{ t('admin.accounts.affinityBufferHint') }}</p>
</div>
<!-- Client zone preview -->
<div v-if="clientZonePreview" class="flex items-center gap-2 text-xs">
<span class="inline-flex items-center gap-1 rounded-full bg-emerald-100 px-2 py-0.5 text-emerald-700 dark:bg-emerald-900/30 dark:text-emerald-400">
{{ clientZonePreview.green }}
</span>
<span v-if="clientZonePreview.yellow" class="inline-flex items-center gap-1 rounded-full bg-yellow-100 px-2 py-0.5 text-yellow-700 dark:bg-yellow-900/30 dark:text-yellow-400">
{{ clientZonePreview.yellow }}
</span>
<span v-if="clientZonePreview.red" class="inline-flex items-center gap-1 rounded-full bg-red-100 px-2 py-0.5 text-red-700 dark:bg-red-900/30 dark:text-red-400">
{{ clientZonePreview.red }}
</span>
</div>
</div>
</div>
</div>
</div>
</template>
@@ -1564,7 +1564,7 @@
<div class="mb-3">
<h3 class="input-label mb-0 text-base font-semibold">{{ t('admin.accounts.quotaControl.title') }}</h3>
<p class="mt-1 text-xs text-gray-500 dark:text-gray-400">
{{ t('admin.accounts.quotaLimitHint') }}
{{ t('admin.accounts.quotaControl.hint') }}
</p>
</div>
<QuotaLimitCard
@@ -1591,9 +1591,19 @@
:enabled="clientAffinityEnabled"
:base="affinityBase"
:buffer="affinityBuffer"
:allow-switch="affinityAllowSwitch"
:user-base="affinityUserBase"
:user-buffer="affinityUserBuffer"
:per-user-limit="perUserClientLimit"
:pinned-users="pinnedUsers"
@update:enabled="clientAffinityEnabled = $event"
@update:base="affinityBase = $event"
@update:buffer="affinityBuffer = $event"
@update:allow-switch="affinityAllowSwitch = $event"
@update:user-base="affinityUserBase = $event"
@update:user-buffer="affinityUserBuffer = $event"
@update:per-user-limit="perUserClientLimit = $event"
@update:pinned-users="pinnedUsers = $event"
/>
</div>
@@ -1962,9 +1972,19 @@
:enabled="clientAffinityEnabled"
:base="affinityBase"
:buffer="affinityBuffer"
:allow-switch="affinityAllowSwitch"
:user-base="affinityUserBase"
:user-buffer="affinityUserBuffer"
:per-user-limit="perUserClientLimit"
:pinned-users="pinnedUsers"
@update:enabled="clientAffinityEnabled = $event"
@update:base="affinityBase = $event"
@update:buffer="affinityBuffer = $event"
@update:allow-switch="affinityAllowSwitch = $event"
@update:user-base="affinityUserBase = $event"
@update:user-buffer="affinityUserBuffer = $event"
@update:per-user-limit="perUserClientLimit = $event"
@update:pinned-users="pinnedUsers = $event"
/>
<!-- Window Cost Limit -->
@@ -3120,6 +3140,11 @@ const showGeminiHelpDialog = ref(false)
const clientAffinityEnabled = ref(false)
const affinityBase = ref<number | null>(null)
const affinityBuffer = ref<number | null>(null)
const affinityAllowSwitch = ref(true)
const affinityUserBase = ref<number | null>(null)
const affinityUserBuffer = ref<number | null>(null)
const perUserClientLimit = ref<number | null>(null)
const pinnedUsers = ref<number[]>([])
// Quota control state (Anthropic OAuth/SetupToken only)
const windowCostEnabled = ref(false)
@@ -3773,6 +3798,11 @@ const resetForm = () => {
clientAffinityEnabled.value = false
affinityBase.value = null
affinityBuffer.value = null
affinityAllowSwitch.value = true
affinityUserBase.value = null
affinityUserBuffer.value = null
perUserClientLimit.value = null
pinnedUsers.value = []
modelMappings.value = []
modelRestrictionMode.value = 'whitelist'
allowedModels.value = [...claudeModels] // Default fill related models
@@ -3887,6 +3917,7 @@ const buildAnthropicExtra = (base?: Record<string, unknown>): Record<string, unk
/** 将客户端亲和设置写入 extra(Anthropic 全类型通用) */
const applyClientAffinity = (extra: Record<string, unknown>) => {
if (clientAffinityEnabled.value) {
extra.affinity_enabled = true
extra.client_affinity_enabled = true
if (affinityBase.value != null && affinityBase.value > 0) {
extra.affinity_base = affinityBase.value
@@ -3898,10 +3929,38 @@ const applyClientAffinity = (extra: Record<string, unknown>) => {
} else {
delete extra.affinity_buffer
}
// v2 fields
extra.affinity_allow_switch = affinityAllowSwitch.value
if (affinityUserBase.value != null && affinityUserBase.value > 0) {
extra.affinity_user_base = affinityUserBase.value
} else {
delete extra.affinity_user_base
}
if (affinityUserBase.value != null && affinityUserBase.value > 0 && affinityUserBuffer.value != null) {
extra.affinity_user_buffer = affinityUserBuffer.value
} else {
delete extra.affinity_user_buffer
}
if (perUserClientLimit.value != null && perUserClientLimit.value > 0) {
extra.per_user_client_limit = perUserClientLimit.value
} else {
delete extra.per_user_client_limit
}
if (pinnedUsers.value.length > 0) {
extra.pinned_users = pinnedUsers.value
} else {
delete extra.pinned_users
}
} else {
delete extra.affinity_enabled
delete extra.client_affinity_enabled
delete extra.affinity_base
delete extra.affinity_buffer
delete extra.affinity_allow_switch
delete extra.affinity_user_base
delete extra.affinity_user_buffer
delete extra.per_user_client_limit
delete extra.pinned_users
}
}
@@ -1158,7 +1158,7 @@
<div class="mb-3">
<h3 class="input-label mb-0 text-base font-semibold">{{ t('admin.accounts.quotaControl.title') }}</h3>
<p class="mt-1 text-xs text-gray-500 dark:text-gray-400">
{{ t('admin.accounts.quotaLimitHint') }}
{{ t('admin.accounts.quotaControl.hint') }}
</p>
</div>
<QuotaLimitCard
@@ -1185,12 +1185,21 @@
:enabled="clientAffinityEnabled"
:base="affinityBase"
:buffer="affinityBuffer"
:allow-switch="affinityAllowSwitch"
:user-base="affinityUserBase"
:user-buffer="affinityUserBuffer"
:per-user-limit="perUserClientLimit"
:pinned-users="pinnedUsers"
@update:enabled="clientAffinityEnabled = $event"
@update:base="affinityBase = $event"
@update:buffer="affinityBuffer = $event"
@update:allow-switch="affinityAllowSwitch = $event"
@update:user-base="affinityUserBase = $event"
@update:user-buffer="affinityUserBuffer = $event"
@update:per-user-limit="perUserClientLimit = $event"
@update:pinned-users="pinnedUsers = $event"
/>
</div>
<!-- 配额控制 (非 Anthropic apikey/bedrock) -->
<div
v-else-if="account?.type === 'apikey' || account?.type === 'bedrock'"
@@ -1298,9 +1307,19 @@
:enabled="clientAffinityEnabled"
:base="affinityBase"
:buffer="affinityBuffer"
:allow-switch="affinityAllowSwitch"
:user-base="affinityUserBase"
:user-buffer="affinityUserBuffer"
:per-user-limit="perUserClientLimit"
:pinned-users="pinnedUsers"
@update:enabled="clientAffinityEnabled = $event"
@update:base="affinityBase = $event"
@update:buffer="affinityBuffer = $event"
@update:allow-switch="affinityAllowSwitch = $event"
@update:user-base="affinityUserBase = $event"
@update:user-buffer="affinityUserBuffer = $event"
@update:per-user-limit="perUserClientLimit = $event"
@update:pinned-users="pinnedUsers = $event"
/>
<!-- Window Cost Limit -->
@@ -1902,6 +1921,11 @@ const cacheTTLOverrideTarget = ref<string>('5m')
const clientAffinityEnabled = ref(false)
const affinityBase = ref<number | null>(null)
const affinityBuffer = ref<number | null>(null)
const affinityAllowSwitch = ref(true)
const affinityUserBase = ref<number | null>(null)
const affinityUserBuffer = ref<number | null>(null)
const perUserClientLimit = ref<number | null>(null)
const pinnedUsers = ref<number[]>([])
// OpenAI 自动透传开关(OAuth/API Key)
const openaiPassthroughEnabled = ref(false)
@@ -2512,6 +2536,11 @@ function loadQuotaControlSettings(account: Account) {
clientAffinityEnabled.value = false
affinityBase.value = null
affinityBuffer.value = null
affinityAllowSwitch.value = true
affinityUserBase.value = null
affinityUserBuffer.value = null
perUserClientLimit.value = null
pinnedUsers.value = []
// Remaining quota control settings only apply to Anthropic accounts
if (account.platform !== 'anthropic') {
@@ -2530,6 +2559,17 @@ function loadQuotaControlSettings(account: Account) {
// buffer: null = infinite yellow, 0 = no yellow, >0 = yellow range
const buf = extra.affinity_buffer
affinityBuffer.value = (typeof buf === 'number') ? buf : null
// New v2 fields
affinityAllowSwitch.value = extra.affinity_allow_switch !== false
const ub = extra.affinity_user_base
affinityUserBase.value = (typeof ub === 'number' && ub > 0) ? ub : null
const ubuf = extra.affinity_user_buffer
affinityUserBuffer.value = (typeof ubuf === 'number') ? ubuf : null
const pul = extra.per_user_client_limit
perUserClientLimit.value = (typeof pul === 'number' && pul > 0) ? pul : null
const pu = extra.pinned_users
pinnedUsers.value = Array.isArray(pu) ? pu : []
}
// Window cost / session limit only apply to Anthropic OAuth/SetupToken accounts
@@ -2951,11 +2991,12 @@ const handleSubmit = async () => {
}
// For all Anthropic accounts, handle client_affinity in extra
if (props.account.platform === 'anthropic') {
const currentExtra = (props.account.extra as Record<string, unknown>) || {}
const newExtra: Record<string, unknown> = { ...currentExtra }
if (clientAffinityEnabled.value) {
newExtra.client_affinity_enabled = true
if (props.account.platform === 'anthropic') {
const currentExtra = (props.account.extra as Record<string, unknown>) || {}
const newExtra: Record<string, unknown> = { ...currentExtra }
if (clientAffinityEnabled.value) {
newExtra.affinity_enabled = true
newExtra.client_affinity_enabled = true
if (affinityBase.value != null && affinityBase.value > 0) {
newExtra.affinity_base = affinityBase.value
} else {
@@ -2967,10 +3008,38 @@ const handleSubmit = async () => {
} else {
delete newExtra.affinity_buffer
}
} else {
delete newExtra.client_affinity_enabled
delete newExtra.affinity_base
delete newExtra.affinity_buffer
// v2 fields
newExtra.affinity_allow_switch = affinityAllowSwitch.value
if (affinityUserBase.value != null && affinityUserBase.value > 0) {
newExtra.affinity_user_base = affinityUserBase.value
} else {
delete newExtra.affinity_user_base
}
if (affinityUserBase.value != null && affinityUserBase.value > 0 && affinityUserBuffer.value != null) {
newExtra.affinity_user_buffer = affinityUserBuffer.value
} else {
delete newExtra.affinity_user_buffer
}
if (perUserClientLimit.value != null && perUserClientLimit.value > 0) {
newExtra.per_user_client_limit = perUserClientLimit.value
} else {
delete newExtra.per_user_client_limit
}
if (pinnedUsers.value.length > 0) {
newExtra.pinned_users = pinnedUsers.value
} else {
delete newExtra.pinned_users
}
} else {
newExtra.affinity_enabled = false
newExtra.client_affinity_enabled = false
delete newExtra.affinity_base
delete newExtra.affinity_buffer
delete newExtra.affinity_allow_switch
delete newExtra.affinity_user_base
delete newExtra.affinity_user_buffer
delete newExtra.per_user_client_limit
delete newExtra.pinned_users
}
updatePayload.extra = newExtra
}
@@ -3,7 +3,8 @@ import { mount } from '@vue/test-utils'
import AffinityBadge from '../AffinityBadge.vue'
vi.mock('@/api/admin/accounts', () => ({
getAffinityClients: vi.fn()
getAffinityClients: vi.fn(),
getAffinityDetails: vi.fn()
}))
vi.mock('vue-i18n', async () => {
@@ -16,79 +17,147 @@ vi.mock('vue-i18n', async () => {
}
})
function mountBadge(count: number, base = 5, buffer: number | null = 10) {
function mountBadge(opts: {
clientCount: number
base?: number
buffer?: number | null
userCount?: number
userBase?: number
userBuffer?: number | null
}) {
return mount(AffinityBadge, {
props: {
accountId: 42,
count,
base,
buffer
clientCount: opts.clientCount,
base: opts.base ?? 5,
buffer: opts.buffer ?? 10,
userCount: opts.userCount ?? 0,
userBase: opts.userBase ?? 0,
userBuffer: opts.userBuffer ?? null
}
})
}
describe('AffinityBadge', () => {
it('renders the correct count number', () => {
const wrapper = mountBadge(5)
// ====== Original tests (client dimension only) ======
it('renders the correct client count number', () => {
const wrapper = mountBadge({ clientCount: 5 })
expect(wrapper.text()).toContain('5')
})
it('renders configured limit text', () => {
const wrapper = mountBadge(5, 5, 10)
const wrapper = mountBadge({ clientCount: 5, base: 5, buffer: 10 })
expect(wrapper.text()).toContain('15')
})
it('renders infinity limit text when base is not configured', () => {
const wrapper = mountBadge(5, 0, null)
expect(wrapper.text()).toContain('∞')
const wrapper = mountBadge({ clientCount: 5, base: 0, buffer: null })
expect(wrapper.text()).toContain('\u221E')
})
it('applies red badge class when count exceeds base plus buffer', () => {
const wrapper = mountBadge(16, 5, 10)
const wrapper = mountBadge({ clientCount: 16, base: 5, buffer: 10 })
const badge = wrapper.find('span')
expect(badge.classes()).toContain('bg-red-100')
expect(badge.classes()).toContain('text-red-700')
})
it('applies yellow badge class when count is in buffer range', () => {
const wrapper = mountBadge(6, 5, 10)
const wrapper = mountBadge({ clientCount: 6, base: 5, buffer: 10 })
const badge = wrapper.find('span')
expect(badge.classes()).toContain('bg-yellow-100')
expect(badge.classes()).toContain('text-yellow-700')
})
it('applies yellow badge class when buffer is infinite', () => {
const wrapper = mountBadge(6, 5, null)
const wrapper = mountBadge({ clientCount: 6, base: 5, buffer: null })
const badge = wrapper.find('span')
expect(badge.classes()).toContain('bg-yellow-100')
expect(badge.classes()).toContain('text-yellow-700')
})
it('applies green badge class when count is within base', () => {
const wrapper = mountBadge(5, 5, 10)
const wrapper = mountBadge({ clientCount: 5, base: 5, buffer: 10 })
const badge = wrapper.find('span')
expect(badge.classes()).toContain('bg-emerald-100')
expect(badge.classes()).toContain('text-emerald-700')
})
it('applies gray badge class when count is 0', () => {
const wrapper = mountBadge(0, 5, 10)
const wrapper = mountBadge({ clientCount: 0, base: 5, buffer: 10 })
const badge = wrapper.find('span')
expect(badge.classes()).toContain('bg-gray-100')
expect(badge.classes()).toContain('text-gray-600')
})
it('does NOT show popover initially', () => {
const wrapper = mountBadge(3)
const wrapper = mountBadge({ clientCount: 3 })
expect(wrapper.html()).not.toContain('divide-y')
expect(wrapper.html()).not.toContain('affinityClients')
})
it('has mouseenter and mouseleave handlers on the badge', () => {
const wrapper = mountBadge(3)
const wrapper = mountBadge({ clientCount: 3 })
const badge = wrapper.find('span')
expect(badge.exists()).toBe(true)
badge.trigger('mouseenter')
badge.trigger('mouseleave')
})
// ====== Dual dimension tests ======
it('dual dimension: user green + client yellow = yellow', () => {
const wrapper = mountBadge({
clientCount: 6, base: 5, buffer: 10,
userCount: 3, userBase: 5, userBuffer: 5
})
const badge = wrapper.find('span')
expect(badge.classes()).toContain('bg-yellow-100')
expect(badge.classes()).toContain('text-yellow-700')
})
it('dual dimension: user green + client red = red', () => {
const wrapper = mountBadge({
clientCount: 16, base: 5, buffer: 10,
userCount: 3, userBase: 5, userBuffer: 5
})
const badge = wrapper.find('span')
expect(badge.classes()).toContain('bg-red-100')
expect(badge.classes()).toContain('text-red-700')
})
it('dual dimension: user red + client green = red', () => {
const wrapper = mountBadge({
clientCount: 3, base: 5, buffer: 10,
userCount: 11, userBase: 5, userBuffer: 5
})
const badge = wrapper.find('span')
expect(badge.classes()).toContain('bg-red-100')
expect(badge.classes()).toContain('text-red-700')
})
it('user only dimension shows user count text', () => {
const wrapper = mountBadge({
clientCount: 0, base: 0, buffer: null,
userCount: 3, userBase: 5, userBuffer: 5
})
expect(wrapper.text()).toContain('3')
expect(wrapper.text()).toContain('10')
})
it('client only dimension shows client count text', () => {
const wrapper = mountBadge({
clientCount: 4, base: 5, buffer: 10,
userCount: 0, userBase: 0, userBuffer: null
})
expect(wrapper.text()).toContain('4')
expect(wrapper.text()).toContain('15')
})
it('dual dimension limit text shows both U and C prefixes', () => {
const wrapper = mountBadge({
clientCount: 3, base: 5, buffer: 10,
userCount: 2, userBase: 4, userBuffer: 3
})
expect(wrapper.text()).toContain('U2/7')
expect(wrapper.text()).toContain('C3/15')
})
})
+29 -4
View File
@@ -2177,7 +2177,7 @@ export default {
// Quota control (Anthropic OAuth/SetupToken only)
quotaControl: {
title: 'Quota Control',
hint: 'Configure cost window, session limits, client affinity and other scheduling controls.',
hint: 'Configure affinity settings (user affinity, client affinity), cost windows, session limits and other scheduling controls.',
windowCost: {
label: '5h Window Cost Limit',
hint: 'Limit account cost usage within the 5-hour window',
@@ -2238,12 +2238,15 @@ export default {
hint: 'When enabled, new sessions prefer accounts previously used by this client to reduce account switching'
}
},
affinityConfigTitle: 'Affinity Settings',
affinityConfigHint: 'Configure user affinity and client affinity scheduling rules under quota control.',
affinityNoClients: 'No affinity clients',
affinityClients: '{count} affinity clients:',
affinityClients: '{count} client affinity',
affinityClientCountLabel: '{count} clients',
affinitySection: 'Client Affinity',
affinitySectionHint: 'Control how clients are distributed across accounts. Configure zone thresholds to balance load.',
affinityToggle: 'Enable Client Affinity',
affinityToggleHint: 'New sessions prefer accounts previously used by this client',
affinityToggle: 'Enable Affinity Scheduling',
affinityToggleHint: 'When enabled, new sessions prefer user affinity and client affinity matches',
affinityBase: 'Base Limit (Green Zone)',
affinityBasePlaceholder: 'Empty = no limit',
affinityBaseHint: 'Max clients in green zone (full priority scheduling)',
@@ -2252,6 +2255,28 @@ export default {
affinityBufferPlaceholder: 'e.g. 3',
affinityBufferHint: 'Additional clients allowed in the yellow zone (degraded priority)',
affinityBufferInfinite: 'Unlimited',
affinityAllowSwitch: 'Allow Switch',
affinityAllowSwitchHint: 'Allow scheduling to other accounts when affinity account is unavailable',
affinityAllowSwitchWarning: 'When disabled, unavailable affinity account will return error directly',
affinityUserSection: 'User Affinity',
affinityUserSectionHint: 'Configure affinity capacity and scheduling rules for users',
affinityUserBase: 'Base Limit (Green Zone)',
affinityUserBaseHint: 'Max users in green zone (full priority scheduling)',
affinityUserBuffer: 'Buffer (Yellow Zone)',
affinityUserBufferHint: 'Additional users allowed in the yellow zone (degraded priority)',
affinityClientSection: 'Client Affinity',
affinityClientSectionHint: 'Configure affinity capacity and scheduling rules for clients',
affinityPerUserLimit: 'Per-User Client Limit',
affinityPerUserLimitHint: 'Limit the number of distinct clients each user can use',
affinityPerUserMax: 'Max Clients',
affinityPinnedUsers: 'Pinned Affinity Users',
affinityPinnedUsersHint: 'Pre-bound users will automatically create affinity cache and occupy user affinity slots',
affinityPinnedUsersSearch: 'Search user email or username...',
affinityPinnedUsersEmpty: 'No pinned affinity users',
affinityUsers: '{count} user affinity',
affinityNoUsers: 'No affinity users',
affinityDetailTitle: 'Affinity Details',
affinityPinnedLabel: 'Pinned',
expired: 'Expired',
proxy: 'Proxy',
noProxy: 'No Proxy',
+29 -4
View File
@@ -2319,7 +2319,7 @@ export default {
// Quota control (Anthropic OAuth/SetupToken only)
quotaControl: {
title: '配额控制',
hint: '配置费用窗口、会话限制、客户端亲和等调度控制。',
hint: '配置亲和配置(用户亲和、客户端亲和)、费用窗口、会话限制等调度控制。',
windowCost: {
label: '5h窗口费用控制',
hint: '限制账号在5小时窗口内的费用使用',
@@ -2380,12 +2380,15 @@ export default {
hint: '启用后,新会话会优先调度到该客户端之前使用过的账号,避免频繁切换账号'
}
},
affinityConfigTitle: '亲和配置',
affinityConfigHint: '在配额控制下配置用户亲和与客户端亲和调度规则。',
affinityNoClients: '无亲和客户端',
affinityClients: '{count} 个亲和客户端:',
affinityClients: '{count} 个客户端亲和',
affinityClientCountLabel: '{count} 个客户端',
affinitySection: '客户端亲和',
affinitySectionHint: '控制客户端在账号间的分布。通过配置区域阈值来平衡负载。',
affinityToggle: '启用客户端亲和',
affinityToggleHint: '新会话优先调度到该客户端之前使用过的账号',
affinityToggle: '启用亲和调度',
affinityToggleHint: '启用后,新会话会优先命中用户亲和与客户端亲和规则',
affinityBase: '基础限额(绿区)',
affinityBasePlaceholder: '留空表示不限制',
affinityBaseHint: '绿区最大客户端数量(完整优先级调度)',
@@ -2394,6 +2397,28 @@ export default {
affinityBufferPlaceholder: '例如 3',
affinityBufferHint: '黄区允许的额外客户端数量(降级优先级调度)',
affinityBufferInfinite: '不限制',
affinityAllowSwitch: '允许切换',
affinityAllowSwitchHint: '亲和账号不可用时允许调度到其他账号',
affinityAllowSwitchWarning: '关闭后,亲和账号不可用将直接返回错误',
affinityUserSection: '用户亲和',
affinityUserSectionHint: '配置用户维度的亲和容量与调度规则',
affinityUserBase: '基础限额(绿区)',
affinityUserBaseHint: '绿区最大用户数量(完整优先级调度)',
affinityUserBuffer: '缓冲区(黄区)',
affinityUserBufferHint: '黄区允许的额外用户数量(降级优先级调度)',
affinityClientSection: '客户端亲和',
affinityClientSectionHint: '配置客户端维度的亲和容量与调度规则',
affinityPerUserLimit: '每用户客户端限制',
affinityPerUserLimitHint: '限制每个用户可使用的不同客户端数量',
affinityPerUserMax: '最大客户端数',
affinityPinnedUsers: '指定亲和用户',
affinityPinnedUsersHint: '预先绑定的用户将自动创建亲和缓存,占用用户亲和名额',
affinityPinnedUsersSearch: '搜索用户邮箱或用户名...',
affinityPinnedUsersEmpty: '未指定亲和用户',
affinityUsers: '{count} 个用户亲和',
affinityNoUsers: '无亲和用户',
affinityDetailTitle: '亲和详情',
affinityPinnedLabel: '指定',
expired: '已过期',
proxy: '代理',
noProxy: '无代理',
+26
View File
@@ -731,6 +731,11 @@ export interface Account {
affinity_client_count?: number | null
affinity_clients?: string[] | null
// 二维亲和扩展
affinity_allow_switch?: boolean | null // 允许切换
affinity_user_count?: number | null // 关联用户数
pinned_user_ids?: number[] | null // 指定亲和用户 ID 列表
// API Key 账号配额限制
quota_limit?: number | null
quota_used?: number | null
@@ -755,6 +760,27 @@ export interface Account {
current_rpm?: number | null // 当前分钟 RPM 计数
}
// Affinity Details types
export interface AffinityClientInfo {
client_id: string
last_active: string
}
export interface AffinityUserGroup {
user_id: number
user_email: string
client_count: number
is_pinned: boolean
clients: AffinityClientInfo[]
}
export interface AffinityDetailsResponse {
users: AffinityUserGroup[]
total_users: number
total_clients: number
pinned_users: number[]
}
// Account Usage types
export interface WindowStats {
requests: number