Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Sign in / Register
Toggle navigation
Menu
Open sidebar
陈曦
sub2api
Commits
f6ed3d14
"backend/vscode:/vscode.git/clone" did not exist on "e1a9c1ecd96e5d9ea22c0c28d749da81a23e0d2b"
Commit
f6ed3d14
authored
Jan 20, 2026
by
yangjianbo
Browse files
Merge branch 'test' into dev
parents
bdc426a7
84686753
Changes
19
Hide whitespace changes
Inline
Side-by-side
backend/cmd/server/wire_gen.go
View file @
f6ed3d14
...
...
@@ -84,7 +84,8 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
}
dashboardAggregationService
:=
service
.
ProvideDashboardAggregationService
(
dashboardAggregationRepository
,
timingWheelService
,
configConfig
)
dashboardHandler
:=
admin
.
NewDashboardHandler
(
dashboardService
,
dashboardAggregationService
)
accountRepository
:=
repository
.
NewAccountRepository
(
client
,
db
)
schedulerCache
:=
repository
.
NewSchedulerCache
(
redisClient
)
accountRepository
:=
repository
.
NewAccountRepository
(
client
,
db
,
schedulerCache
)
proxyRepository
:=
repository
.
NewProxyRepository
(
client
,
db
)
proxyExitInfoProber
:=
repository
.
NewProxyExitInfoProber
(
configConfig
)
proxyLatencyCache
:=
repository
.
NewProxyLatencyCache
(
redisClient
)
...
...
@@ -128,7 +129,6 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
adminRedeemHandler
:=
admin
.
NewRedeemHandler
(
adminService
)
promoHandler
:=
admin
.
NewPromoHandler
(
promoService
)
opsRepository
:=
repository
.
NewOpsRepository
(
db
)
schedulerCache
:=
repository
.
NewSchedulerCache
(
redisClient
)
schedulerOutboxRepository
:=
repository
.
NewSchedulerOutboxRepository
(
db
)
schedulerSnapshotService
:=
service
.
ProvideSchedulerSnapshotService
(
schedulerCache
,
schedulerOutboxRepository
,
accountRepository
,
groupRepository
,
configConfig
)
pricingRemoteClient
:=
repository
.
ProvidePricingRemoteClient
(
configConfig
)
...
...
backend/internal/repository/account_repo.go
View file @
f6ed3d14
...
...
@@ -39,9 +39,15 @@ import (
// 设计说明:
// - client: Ent 客户端,用于类型安全的 ORM 操作
// - sql: 原生 SQL 执行器,用于复杂查询和批量操作
// - schedulerCache: 调度器缓存,用于在账号状态变更时同步快照
type
accountRepository
struct
{
client
*
dbent
.
Client
// Ent ORM 客户端
sql
sqlExecutor
// 原生 SQL 执行接口
// schedulerCache 用于在账号状态变更时主动同步快照到缓存,
// 确保粘性会话能及时感知账号不可用状态。
// Used to proactively sync account snapshot to cache when status changes,
// ensuring sticky sessions can promptly detect unavailable accounts.
schedulerCache
service
.
SchedulerCache
}
type
tempUnschedSnapshot
struct
{
...
...
@@ -51,14 +57,14 @@ type tempUnschedSnapshot struct {
// NewAccountRepository 创建账户仓储实例。
// 这是对外暴露的构造函数,返回接口类型以便于依赖注入。
func
NewAccountRepository
(
client
*
dbent
.
Client
,
sqlDB
*
sql
.
DB
)
service
.
AccountRepository
{
return
newAccountRepositoryWithSQL
(
client
,
sqlDB
)
func
NewAccountRepository
(
client
*
dbent
.
Client
,
sqlDB
*
sql
.
DB
,
schedulerCache
service
.
SchedulerCache
)
service
.
AccountRepository
{
return
newAccountRepositoryWithSQL
(
client
,
sqlDB
,
schedulerCache
)
}
// newAccountRepositoryWithSQL 是内部构造函数,支持依赖注入 SQL 执行器。
// 这种设计便于单元测试时注入 mock 对象。
func
newAccountRepositoryWithSQL
(
client
*
dbent
.
Client
,
sqlq
sqlExecutor
)
*
accountRepository
{
return
&
accountRepository
{
client
:
client
,
sql
:
sqlq
}
func
newAccountRepositoryWithSQL
(
client
*
dbent
.
Client
,
sqlq
sqlExecutor
,
schedulerCache
service
.
SchedulerCache
)
*
accountRepository
{
return
&
accountRepository
{
client
:
client
,
sql
:
sqlq
,
schedulerCache
:
schedulerCache
}
}
func
(
r
*
accountRepository
)
Create
(
ctx
context
.
Context
,
account
*
service
.
Account
)
error
{
...
...
@@ -356,6 +362,9 @@ func (r *accountRepository) Update(ctx context.Context, account *service.Account
if
err
:=
enqueueSchedulerOutbox
(
ctx
,
r
.
sql
,
service
.
SchedulerOutboxEventAccountChanged
,
&
account
.
ID
,
nil
,
buildSchedulerGroupPayload
(
account
.
GroupIDs
));
err
!=
nil
{
log
.
Printf
(
"[SchedulerOutbox] enqueue account update failed: account=%d err=%v"
,
account
.
ID
,
err
)
}
if
account
.
Status
==
service
.
StatusError
||
account
.
Status
==
service
.
StatusDisabled
||
!
account
.
Schedulable
{
r
.
syncSchedulerAccountSnapshot
(
ctx
,
account
.
ID
)
}
return
nil
}
...
...
@@ -540,9 +549,32 @@ func (r *accountRepository) SetError(ctx context.Context, id int64, errorMsg str
if
err
:=
enqueueSchedulerOutbox
(
ctx
,
r
.
sql
,
service
.
SchedulerOutboxEventAccountChanged
,
&
id
,
nil
,
nil
);
err
!=
nil
{
log
.
Printf
(
"[SchedulerOutbox] enqueue set error failed: account=%d err=%v"
,
id
,
err
)
}
r
.
syncSchedulerAccountSnapshot
(
ctx
,
id
)
return
nil
}
// syncSchedulerAccountSnapshot 在账号状态变更时主动同步快照到调度器缓存。
// 当账号被设置为错误、禁用、不可调度或临时不可调度时调用,
// 确保调度器和粘性会话逻辑能及时感知账号的最新状态,避免继续使用不可用账号。
//
// syncSchedulerAccountSnapshot proactively syncs account snapshot to scheduler cache
// when account status changes. Called when account is set to error, disabled,
// unschedulable, or temporarily unschedulable, ensuring scheduler and sticky session
// logic can promptly detect the latest account state and avoid using unavailable accounts.
func
(
r
*
accountRepository
)
syncSchedulerAccountSnapshot
(
ctx
context
.
Context
,
accountID
int64
)
{
if
r
==
nil
||
r
.
schedulerCache
==
nil
||
accountID
<=
0
{
return
}
account
,
err
:=
r
.
GetByID
(
ctx
,
accountID
)
if
err
!=
nil
{
log
.
Printf
(
"[Scheduler] sync account snapshot read failed: id=%d err=%v"
,
accountID
,
err
)
return
}
if
err
:=
r
.
schedulerCache
.
SetAccount
(
ctx
,
account
);
err
!=
nil
{
log
.
Printf
(
"[Scheduler] sync account snapshot write failed: id=%d err=%v"
,
accountID
,
err
)
}
}
func
(
r
*
accountRepository
)
AddToGroup
(
ctx
context
.
Context
,
accountID
,
groupID
int64
,
priority
int
)
error
{
_
,
err
:=
r
.
client
.
AccountGroup
.
Create
()
.
SetAccountID
(
accountID
)
.
...
...
@@ -864,6 +896,7 @@ func (r *accountRepository) SetTempUnschedulable(ctx context.Context, id int64,
if
err
:=
enqueueSchedulerOutbox
(
ctx
,
r
.
sql
,
service
.
SchedulerOutboxEventAccountChanged
,
&
id
,
nil
,
nil
);
err
!=
nil
{
log
.
Printf
(
"[SchedulerOutbox] enqueue temp unschedulable failed: account=%d err=%v"
,
id
,
err
)
}
r
.
syncSchedulerAccountSnapshot
(
ctx
,
id
)
return
nil
}
...
...
@@ -974,6 +1007,9 @@ func (r *accountRepository) SetSchedulable(ctx context.Context, id int64, schedu
if
err
:=
enqueueSchedulerOutbox
(
ctx
,
r
.
sql
,
service
.
SchedulerOutboxEventAccountChanged
,
&
id
,
nil
,
nil
);
err
!=
nil
{
log
.
Printf
(
"[SchedulerOutbox] enqueue schedulable change failed: account=%d err=%v"
,
id
,
err
)
}
if
!
schedulable
{
r
.
syncSchedulerAccountSnapshot
(
ctx
,
id
)
}
return
nil
}
...
...
@@ -1128,6 +1164,18 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates
if
err
:=
enqueueSchedulerOutbox
(
ctx
,
r
.
sql
,
service
.
SchedulerOutboxEventAccountBulkChanged
,
nil
,
nil
,
payload
);
err
!=
nil
{
log
.
Printf
(
"[SchedulerOutbox] enqueue bulk update failed: err=%v"
,
err
)
}
shouldSync
:=
false
if
updates
.
Status
!=
nil
&&
(
*
updates
.
Status
==
service
.
StatusError
||
*
updates
.
Status
==
service
.
StatusDisabled
)
{
shouldSync
=
true
}
if
updates
.
Schedulable
!=
nil
&&
!*
updates
.
Schedulable
{
shouldSync
=
true
}
if
shouldSync
{
for
_
,
id
:=
range
ids
{
r
.
syncSchedulerAccountSnapshot
(
ctx
,
id
)
}
}
}
return
rows
,
nil
}
...
...
backend/internal/repository/account_repo_integration_test.go
View file @
f6ed3d14
...
...
@@ -21,11 +21,56 @@ type AccountRepoSuite struct {
repo
*
accountRepository
}
type
schedulerCacheRecorder
struct
{
setAccounts
[]
*
service
.
Account
}
func
(
s
*
schedulerCacheRecorder
)
GetSnapshot
(
ctx
context
.
Context
,
bucket
service
.
SchedulerBucket
)
([]
*
service
.
Account
,
bool
,
error
)
{
return
nil
,
false
,
nil
}
func
(
s
*
schedulerCacheRecorder
)
SetSnapshot
(
ctx
context
.
Context
,
bucket
service
.
SchedulerBucket
,
accounts
[]
service
.
Account
)
error
{
return
nil
}
func
(
s
*
schedulerCacheRecorder
)
GetAccount
(
ctx
context
.
Context
,
accountID
int64
)
(
*
service
.
Account
,
error
)
{
return
nil
,
nil
}
func
(
s
*
schedulerCacheRecorder
)
SetAccount
(
ctx
context
.
Context
,
account
*
service
.
Account
)
error
{
s
.
setAccounts
=
append
(
s
.
setAccounts
,
account
)
return
nil
}
func
(
s
*
schedulerCacheRecorder
)
DeleteAccount
(
ctx
context
.
Context
,
accountID
int64
)
error
{
return
nil
}
func
(
s
*
schedulerCacheRecorder
)
UpdateLastUsed
(
ctx
context
.
Context
,
updates
map
[
int64
]
time
.
Time
)
error
{
return
nil
}
func
(
s
*
schedulerCacheRecorder
)
TryLockBucket
(
ctx
context
.
Context
,
bucket
service
.
SchedulerBucket
,
ttl
time
.
Duration
)
(
bool
,
error
)
{
return
true
,
nil
}
func
(
s
*
schedulerCacheRecorder
)
ListBuckets
(
ctx
context
.
Context
)
([]
service
.
SchedulerBucket
,
error
)
{
return
nil
,
nil
}
func
(
s
*
schedulerCacheRecorder
)
GetOutboxWatermark
(
ctx
context
.
Context
)
(
int64
,
error
)
{
return
0
,
nil
}
func
(
s
*
schedulerCacheRecorder
)
SetOutboxWatermark
(
ctx
context
.
Context
,
id
int64
)
error
{
return
nil
}
func
(
s
*
AccountRepoSuite
)
SetupTest
()
{
s
.
ctx
=
context
.
Background
()
tx
:=
testEntTx
(
s
.
T
())
s
.
client
=
tx
.
Client
()
s
.
repo
=
newAccountRepositoryWithSQL
(
s
.
client
,
tx
)
s
.
repo
=
newAccountRepositoryWithSQL
(
s
.
client
,
tx
,
nil
)
}
func
TestAccountRepoSuite
(
t
*
testing
.
T
)
{
...
...
@@ -73,6 +118,20 @@ func (s *AccountRepoSuite) TestUpdate() {
s
.
Require
()
.
Equal
(
"updated"
,
got
.
Name
)
}
func
(
s
*
AccountRepoSuite
)
TestUpdate_SyncSchedulerSnapshotOnDisabled
()
{
account
:=
mustCreateAccount
(
s
.
T
(),
s
.
client
,
&
service
.
Account
{
Name
:
"sync-update"
,
Status
:
service
.
StatusActive
,
Schedulable
:
true
})
cacheRecorder
:=
&
schedulerCacheRecorder
{}
s
.
repo
.
schedulerCache
=
cacheRecorder
account
.
Status
=
service
.
StatusDisabled
err
:=
s
.
repo
.
Update
(
s
.
ctx
,
account
)
s
.
Require
()
.
NoError
(
err
,
"Update"
)
s
.
Require
()
.
Len
(
cacheRecorder
.
setAccounts
,
1
)
s
.
Require
()
.
Equal
(
account
.
ID
,
cacheRecorder
.
setAccounts
[
0
]
.
ID
)
s
.
Require
()
.
Equal
(
service
.
StatusDisabled
,
cacheRecorder
.
setAccounts
[
0
]
.
Status
)
}
func
(
s
*
AccountRepoSuite
)
TestDelete
()
{
account
:=
mustCreateAccount
(
s
.
T
(),
s
.
client
,
&
service
.
Account
{
Name
:
"to-delete"
})
...
...
@@ -174,7 +233,7 @@ func (s *AccountRepoSuite) TestListWithFilters() {
// 每个 case 重新获取隔离资源
tx
:=
testEntTx
(
s
.
T
())
client
:=
tx
.
Client
()
repo
:=
newAccountRepositoryWithSQL
(
client
,
tx
)
repo
:=
newAccountRepositoryWithSQL
(
client
,
tx
,
nil
)
ctx
:=
context
.
Background
()
tt
.
setup
(
client
)
...
...
@@ -365,12 +424,38 @@ func (s *AccountRepoSuite) TestListSchedulableByGroupIDAndPlatform() {
func
(
s
*
AccountRepoSuite
)
TestSetSchedulable
()
{
account
:=
mustCreateAccount
(
s
.
T
(),
s
.
client
,
&
service
.
Account
{
Name
:
"acc-sched"
,
Schedulable
:
true
})
cacheRecorder
:=
&
schedulerCacheRecorder
{}
s
.
repo
.
schedulerCache
=
cacheRecorder
s
.
Require
()
.
NoError
(
s
.
repo
.
SetSchedulable
(
s
.
ctx
,
account
.
ID
,
false
))
got
,
err
:=
s
.
repo
.
GetByID
(
s
.
ctx
,
account
.
ID
)
s
.
Require
()
.
NoError
(
err
)
s
.
Require
()
.
False
(
got
.
Schedulable
)
s
.
Require
()
.
Len
(
cacheRecorder
.
setAccounts
,
1
)
s
.
Require
()
.
Equal
(
account
.
ID
,
cacheRecorder
.
setAccounts
[
0
]
.
ID
)
}
func
(
s
*
AccountRepoSuite
)
TestBulkUpdate_SyncSchedulerSnapshotOnDisabled
()
{
account1
:=
mustCreateAccount
(
s
.
T
(),
s
.
client
,
&
service
.
Account
{
Name
:
"bulk-1"
,
Status
:
service
.
StatusActive
,
Schedulable
:
true
})
account2
:=
mustCreateAccount
(
s
.
T
(),
s
.
client
,
&
service
.
Account
{
Name
:
"bulk-2"
,
Status
:
service
.
StatusActive
,
Schedulable
:
true
})
cacheRecorder
:=
&
schedulerCacheRecorder
{}
s
.
repo
.
schedulerCache
=
cacheRecorder
disabled
:=
service
.
StatusDisabled
rows
,
err
:=
s
.
repo
.
BulkUpdate
(
s
.
ctx
,
[]
int64
{
account1
.
ID
,
account2
.
ID
},
service
.
AccountBulkUpdate
{
Status
:
&
disabled
,
})
s
.
Require
()
.
NoError
(
err
)
s
.
Require
()
.
Equal
(
int64
(
2
),
rows
)
s
.
Require
()
.
Len
(
cacheRecorder
.
setAccounts
,
2
)
ids
:=
map
[
int64
]
struct
{}{}
for
_
,
acc
:=
range
cacheRecorder
.
setAccounts
{
ids
[
acc
.
ID
]
=
struct
{}{}
}
s
.
Require
()
.
Contains
(
ids
,
account1
.
ID
)
s
.
Require
()
.
Contains
(
ids
,
account2
.
ID
)
}
// --- SetOverloaded / SetRateLimited / ClearRateLimit ---
...
...
backend/internal/repository/gateway_cache.go
View file @
f6ed3d14
...
...
@@ -39,3 +39,15 @@ func (c *gatewayCache) RefreshSessionTTL(ctx context.Context, groupID int64, ses
key
:=
buildSessionKey
(
groupID
,
sessionHash
)
return
c
.
rdb
.
Expire
(
ctx
,
key
,
ttl
)
.
Err
()
}
// DeleteSessionAccountID 删除粘性会话与账号的绑定关系。
// 当检测到绑定的账号不可用(如状态错误、禁用、不可调度等)时调用,
// 以便下次请求能够重新选择可用账号。
//
// DeleteSessionAccountID removes the sticky session binding for the given session.
// Called when the bound account becomes unavailable (e.g., error status, disabled,
// or unschedulable), allowing subsequent requests to select a new available account.
func
(
c
*
gatewayCache
)
DeleteSessionAccountID
(
ctx
context
.
Context
,
groupID
int64
,
sessionHash
string
)
error
{
key
:=
buildSessionKey
(
groupID
,
sessionHash
)
return
c
.
rdb
.
Del
(
ctx
,
key
)
.
Err
()
}
backend/internal/repository/gateway_cache_integration_test.go
View file @
f6ed3d14
...
...
@@ -78,6 +78,19 @@ func (s *GatewayCacheSuite) TestRefreshSessionTTL_MissingKey() {
require
.
NoError
(
s
.
T
(),
err
,
"RefreshSessionTTL on missing key should not error"
)
}
func
(
s
*
GatewayCacheSuite
)
TestDeleteSessionAccountID
()
{
sessionID
:=
"openai:s4"
accountID
:=
int64
(
102
)
groupID
:=
int64
(
1
)
sessionTTL
:=
1
*
time
.
Minute
require
.
NoError
(
s
.
T
(),
s
.
cache
.
SetSessionAccountID
(
s
.
ctx
,
groupID
,
sessionID
,
accountID
,
sessionTTL
),
"SetSessionAccountID"
)
require
.
NoError
(
s
.
T
(),
s
.
cache
.
DeleteSessionAccountID
(
s
.
ctx
,
groupID
,
sessionID
),
"DeleteSessionAccountID"
)
_
,
err
:=
s
.
cache
.
GetSessionAccountID
(
s
.
ctx
,
groupID
,
sessionID
)
require
.
True
(
s
.
T
(),
errors
.
Is
(
err
,
redis
.
Nil
),
"expected redis.Nil after delete"
)
}
func
(
s
*
GatewayCacheSuite
)
TestGetSessionAccountID_CorruptedValue
()
{
sessionID
:=
"corrupted"
groupID
:=
int64
(
1
)
...
...
backend/internal/repository/gateway_routing_integration_test.go
View file @
f6ed3d14
...
...
@@ -24,7 +24,7 @@ func (s *GatewayRoutingSuite) SetupTest() {
s
.
ctx
=
context
.
Background
()
tx
:=
testEntTx
(
s
.
T
())
s
.
client
=
tx
.
Client
()
s
.
accountRepo
=
newAccountRepositoryWithSQL
(
s
.
client
,
tx
)
s
.
accountRepo
=
newAccountRepositoryWithSQL
(
s
.
client
,
tx
,
nil
)
}
func
TestGatewayRoutingSuite
(
t
*
testing
.
T
)
{
...
...
backend/internal/repository/openai_oauth_service.go
View file @
f6ed3d14
...
...
@@ -4,6 +4,7 @@ import (
"context"
"fmt"
"net/url"
"strings"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
...
...
@@ -21,7 +22,7 @@ type openaiOAuthService struct {
}
func
(
s
*
openaiOAuthService
)
ExchangeCode
(
ctx
context
.
Context
,
code
,
codeVerifier
,
redirectURI
,
proxyURL
string
)
(
*
openai
.
TokenResponse
,
error
)
{
client
:=
createOpenAIReqClient
(
proxyURL
)
client
:=
createOpenAIReqClient
(
s
.
tokenURL
,
proxyURL
)
if
redirectURI
==
""
{
redirectURI
=
openai
.
DefaultRedirectURI
...
...
@@ -54,7 +55,7 @@ func (s *openaiOAuthService) ExchangeCode(ctx context.Context, code, codeVerifie
}
func
(
s
*
openaiOAuthService
)
RefreshToken
(
ctx
context
.
Context
,
refreshToken
,
proxyURL
string
)
(
*
openai
.
TokenResponse
,
error
)
{
client
:=
createOpenAIReqClient
(
proxyURL
)
client
:=
createOpenAIReqClient
(
s
.
tokenURL
,
proxyURL
)
formData
:=
url
.
Values
{}
formData
.
Set
(
"grant_type"
,
"refresh_token"
)
...
...
@@ -81,9 +82,14 @@ func (s *openaiOAuthService) RefreshToken(ctx context.Context, refreshToken, pro
return
&
tokenResp
,
nil
}
func
createOpenAIReqClient
(
proxyURL
string
)
*
req
.
Client
{
func
createOpenAIReqClient
(
tokenURL
,
proxyURL
string
)
*
req
.
Client
{
forceHTTP2
:=
false
if
parsedURL
,
err
:=
url
.
Parse
(
tokenURL
);
err
==
nil
{
forceHTTP2
=
strings
.
EqualFold
(
parsedURL
.
Scheme
,
"https"
)
}
return
getSharedReqClient
(
reqClientOptions
{
ProxyURL
:
proxyURL
,
Timeout
:
60
*
time
.
Second
,
ProxyURL
:
proxyURL
,
Timeout
:
120
*
time
.
Second
,
ForceHTTP2
:
forceHTTP2
,
})
}
backend/internal/repository/openai_oauth_service_test.go
View file @
f6ed3d14
...
...
@@ -244,6 +244,13 @@ func (s *OpenAIOAuthServiceSuite) TestRefreshToken_NonSuccessStatus() {
require
.
ErrorContains
(
s
.
T
(),
err
,
"status 401"
)
}
func
TestNewOpenAIOAuthClient_DefaultTokenURL
(
t
*
testing
.
T
)
{
client
:=
NewOpenAIOAuthClient
()
svc
,
ok
:=
client
.
(
*
openaiOAuthService
)
require
.
True
(
t
,
ok
)
require
.
Equal
(
t
,
openai
.
TokenURL
,
svc
.
tokenURL
)
}
func
TestOpenAIOAuthServiceSuite
(
t
*
testing
.
T
)
{
suite
.
Run
(
t
,
new
(
OpenAIOAuthServiceSuite
))
}
backend/internal/repository/req_client_pool.go
View file @
f6ed3d14
...
...
@@ -14,6 +14,7 @@ type reqClientOptions struct {
ProxyURL
string
// 代理 URL(支持 http/https/socks5)
Timeout
time
.
Duration
// 请求超时时间
Impersonate
bool
// 是否模拟 Chrome 浏览器指纹
ForceHTTP2
bool
// 是否强制使用 HTTP/2
}
// sharedReqClients 存储按配置参数缓存的 req 客户端实例
...
...
@@ -41,6 +42,9 @@ func getSharedReqClient(opts reqClientOptions) *req.Client {
}
client
:=
req
.
C
()
.
SetTimeout
(
opts
.
Timeout
)
if
opts
.
ForceHTTP2
{
client
=
client
.
EnableForceHTTP2
()
}
if
opts
.
Impersonate
{
client
=
client
.
ImpersonateChrome
()
}
...
...
@@ -56,9 +60,10 @@ func getSharedReqClient(opts reqClientOptions) *req.Client {
}
func
buildReqClientKey
(
opts
reqClientOptions
)
string
{
return
fmt
.
Sprintf
(
"%s|%s|%t"
,
return
fmt
.
Sprintf
(
"%s|%s|%t
|%t
"
,
strings
.
TrimSpace
(
opts
.
ProxyURL
),
opts
.
Timeout
.
String
(),
opts
.
Impersonate
,
opts
.
ForceHTTP2
,
)
}
backend/internal/repository/req_client_pool_test.go
0 → 100644
View file @
f6ed3d14
package
repository
import
(
"reflect"
"sync"
"testing"
"time"
"unsafe"
"github.com/imroc/req/v3"
"github.com/stretchr/testify/require"
)
func
forceHTTPVersion
(
t
*
testing
.
T
,
client
*
req
.
Client
)
string
{
t
.
Helper
()
transport
:=
client
.
GetTransport
()
field
:=
reflect
.
ValueOf
(
transport
)
.
Elem
()
.
FieldByName
(
"forceHttpVersion"
)
require
.
True
(
t
,
field
.
IsValid
(),
"forceHttpVersion field not found"
)
require
.
True
(
t
,
field
.
CanAddr
(),
"forceHttpVersion field not addressable"
)
return
reflect
.
NewAt
(
field
.
Type
(),
unsafe
.
Pointer
(
field
.
UnsafeAddr
()))
.
Elem
()
.
String
()
}
func
TestGetSharedReqClient_ForceHTTP2SeparatesCache
(
t
*
testing
.
T
)
{
sharedReqClients
=
sync
.
Map
{}
base
:=
reqClientOptions
{
ProxyURL
:
"http://proxy.local:8080"
,
Timeout
:
time
.
Second
,
}
clientDefault
:=
getSharedReqClient
(
base
)
force
:=
base
force
.
ForceHTTP2
=
true
clientForce
:=
getSharedReqClient
(
force
)
require
.
NotSame
(
t
,
clientDefault
,
clientForce
)
require
.
NotEqual
(
t
,
buildReqClientKey
(
base
),
buildReqClientKey
(
force
))
}
func
TestGetSharedReqClient_ReuseCachedClient
(
t
*
testing
.
T
)
{
sharedReqClients
=
sync
.
Map
{}
opts
:=
reqClientOptions
{
ProxyURL
:
"http://proxy.local:8080"
,
Timeout
:
2
*
time
.
Second
,
}
first
:=
getSharedReqClient
(
opts
)
second
:=
getSharedReqClient
(
opts
)
require
.
Same
(
t
,
first
,
second
)
}
func
TestGetSharedReqClient_IgnoresNonClientCache
(
t
*
testing
.
T
)
{
sharedReqClients
=
sync
.
Map
{}
opts
:=
reqClientOptions
{
ProxyURL
:
" http://proxy.local:8080 "
,
Timeout
:
3
*
time
.
Second
,
}
key
:=
buildReqClientKey
(
opts
)
sharedReqClients
.
Store
(
key
,
"invalid"
)
client
:=
getSharedReqClient
(
opts
)
require
.
NotNil
(
t
,
client
)
loaded
,
ok
:=
sharedReqClients
.
Load
(
key
)
require
.
True
(
t
,
ok
)
require
.
IsType
(
t
,
"invalid"
,
loaded
)
}
func
TestGetSharedReqClient_ImpersonateAndProxy
(
t
*
testing
.
T
)
{
sharedReqClients
=
sync
.
Map
{}
opts
:=
reqClientOptions
{
ProxyURL
:
" http://proxy.local:8080 "
,
Timeout
:
4
*
time
.
Second
,
Impersonate
:
true
,
}
client
:=
getSharedReqClient
(
opts
)
require
.
NotNil
(
t
,
client
)
require
.
Equal
(
t
,
"http://proxy.local:8080|4s|true|false"
,
buildReqClientKey
(
opts
))
}
func
TestCreateOpenAIReqClient_ForceHTTP2Enabled
(
t
*
testing
.
T
)
{
sharedReqClients
=
sync
.
Map
{}
client
:=
createOpenAIReqClient
(
"https://auth.openai.com/oauth/token"
,
"http://proxy.local:8080"
)
require
.
Equal
(
t
,
"2"
,
forceHTTPVersion
(
t
,
client
))
}
func
TestCreateOpenAIReqClient_ForceHTTP2DisabledForHTTP
(
t
*
testing
.
T
)
{
sharedReqClients
=
sync
.
Map
{}
client
:=
createOpenAIReqClient
(
"http://localhost/oauth/token"
,
"http://proxy.local:8080"
)
require
.
Equal
(
t
,
""
,
forceHTTPVersion
(
t
,
client
))
}
func
TestCreateOpenAIReqClient_Timeout120Seconds
(
t
*
testing
.
T
)
{
sharedReqClients
=
sync
.
Map
{}
client
:=
createOpenAIReqClient
(
"https://auth.openai.com/oauth/token"
,
"http://proxy.local:8080"
)
require
.
Equal
(
t
,
120
*
time
.
Second
,
client
.
GetClient
()
.
Timeout
)
}
func
TestCreateGeminiReqClient_ForceHTTP2Disabled
(
t
*
testing
.
T
)
{
sharedReqClients
=
sync
.
Map
{}
client
:=
createGeminiReqClient
(
"http://proxy.local:8080"
)
require
.
Equal
(
t
,
""
,
forceHTTPVersion
(
t
,
client
))
}
backend/internal/repository/scheduler_snapshot_outbox_integration_test.go
View file @
f6ed3d14
...
...
@@ -19,7 +19,7 @@ func TestSchedulerSnapshotOutboxReplay(t *testing.T) {
_
,
_
=
integrationDB
.
ExecContext
(
ctx
,
"TRUNCATE scheduler_outbox"
)
accountRepo
:=
newAccountRepositoryWithSQL
(
client
,
integrationDB
)
accountRepo
:=
newAccountRepositoryWithSQL
(
client
,
integrationDB
,
nil
)
outboxRepo
:=
NewSchedulerOutboxRepository
(
integrationDB
)
cache
:=
NewSchedulerCache
(
rdb
)
...
...
backend/internal/service/gateway_multiplatform_test.go
View file @
f6ed3d14
...
...
@@ -179,6 +179,7 @@ var _ AccountRepository = (*mockAccountRepoForPlatform)(nil)
// mockGatewayCacheForPlatform 单平台测试用的 cache mock
type
mockGatewayCacheForPlatform
struct
{
sessionBindings
map
[
string
]
int64
deletedSessions
map
[
string
]
int
}
func
(
m
*
mockGatewayCacheForPlatform
)
GetSessionAccountID
(
ctx
context
.
Context
,
groupID
int64
,
sessionHash
string
)
(
int64
,
error
)
{
...
...
@@ -200,6 +201,18 @@ func (m *mockGatewayCacheForPlatform) RefreshSessionTTL(ctx context.Context, gro
return
nil
}
func
(
m
*
mockGatewayCacheForPlatform
)
DeleteSessionAccountID
(
ctx
context
.
Context
,
groupID
int64
,
sessionHash
string
)
error
{
if
m
.
sessionBindings
==
nil
{
return
nil
}
if
m
.
deletedSessions
==
nil
{
m
.
deletedSessions
=
make
(
map
[
string
]
int
)
}
m
.
deletedSessions
[
sessionHash
]
++
delete
(
m
.
sessionBindings
,
sessionHash
)
return
nil
}
type
mockGroupRepoForGateway
struct
{
groups
map
[
int64
]
*
Group
getByIDCalls
int
...
...
@@ -623,183 +636,433 @@ func TestGatewayService_SelectAccountForModelWithPlatform_StickySession(t *testi
})
}
func
TestGatewayService_isModelSupportedByAccount
(
t
*
testing
.
T
)
{
svc
:=
&
GatewayService
{}
func
TestGatewayService_SelectAccountForModelWithExclusions_ForcePlatform
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
ctx
=
context
.
WithValue
(
ctx
,
ctxkey
.
ForcePlatform
,
PlatformAntigravity
)
tests
:=
[]
struct
{
name
string
account
*
Account
model
string
expected
bool
}{
{
name
:
"Antigravity平台-支持claude模型"
,
account
:
&
Account
{
Platform
:
PlatformAntigravity
},
model
:
"claude-3-5-sonnet-20241022"
,
expected
:
true
,
},
{
name
:
"Antigravity平台-支持gemini模型"
,
account
:
&
Account
{
Platform
:
PlatformAntigravity
},
model
:
"gemini-2.5-flash"
,
expected
:
true
,
},
{
name
:
"Antigravity平台-不支持gpt模型"
,
account
:
&
Account
{
Platform
:
PlatformAntigravity
},
model
:
"gpt-4"
,
expected
:
false
,
},
{
name
:
"Anthropic平台-无映射配置-支持所有模型"
,
account
:
&
Account
{
Platform
:
PlatformAnthropic
},
model
:
"claude-3-5-sonnet-20241022"
,
expected
:
true
,
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
{
ID
:
2
,
Platform
:
PlatformAntigravity
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
},
},
{
name
:
"Anthropic平台-有映射配置-只支持配置的模型"
,
account
:
&
Account
{
Platform
:
PlatformAnthropic
,
Credentials
:
map
[
string
]
any
{
"model_mapping"
:
map
[
string
]
any
{
"claude-opus-4"
:
"x"
}},
},
model
:
"claude-3-5-sonnet-20241022"
,
expected
:
false
,
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
}
acc
,
err
:=
svc
.
SelectAccountForModelWithExclusions
(
ctx
,
nil
,
""
,
"claude-3-5-sonnet-20241022"
,
nil
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
)
require
.
Equal
(
t
,
PlatformAntigravity
,
acc
.
Platform
)
}
func
TestGatewayService_SelectAccountForModelWithPlatform_RoutedStickySessionClears
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
groupID
:=
int64
(
10
)
requestedModel
:=
"claude-3-5-sonnet-20241022"
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
2
,
Status
:
StatusDisabled
,
Schedulable
:
true
},
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
},
{
name
:
"Anthropic平台-有映射配置-支持配置的模型"
,
account
:
&
Account
{
Platform
:
PlatformAnthropic
,
Credentials
:
map
[
string
]
any
{
"model_mapping"
:
map
[
string
]
any
{
"claude-3-5-sonnet-20241022"
:
"x"
}},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
"session-123"
:
1
},
}
groupRepo
:=
&
mockGroupRepoForGateway
{
groups
:
map
[
int64
]
*
Group
{
groupID
:
{
ID
:
groupID
,
Name
:
"route-group"
,
Platform
:
PlatformAnthropic
,
Status
:
StatusActive
,
Hydrated
:
true
,
ModelRoutingEnabled
:
true
,
ModelRouting
:
map
[
string
][]
int64
{
requestedModel
:
{
1
,
2
},
},
},
model
:
"claude-3-5-sonnet-20241022"
,
expected
:
true
,
},
}
for
_
,
tt
:=
range
tests
{
t
.
Run
(
tt
.
name
,
func
(
t
*
testing
.
T
)
{
got
:=
svc
.
isModelSupportedByAccount
(
tt
.
account
,
tt
.
model
)
require
.
Equal
(
t
,
tt
.
expected
,
got
)
})
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
groupRepo
:
groupRepo
,
}
acc
,
err
:=
svc
.
selectAccountForModelWithPlatform
(
ctx
,
&
groupID
,
"session-123"
,
requestedModel
,
nil
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
)
require
.
Equal
(
t
,
1
,
cache
.
deletedSessions
[
"session-123"
])
require
.
Equal
(
t
,
int64
(
2
),
cache
.
sessionBindings
[
"session-123"
])
}
// TestGatewayService_selectAccountWithMixedScheduling 测试混合调度
func
TestGatewayService_selectAccountWithMixedScheduling
(
t
*
testing
.
T
)
{
func
TestGatewayService_SelectAccountForModelWithPlatform_RoutedStickySessionHit
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
groupID
:=
int64
(
11
)
requestedModel
:=
"claude-3-5-sonnet-20241022"
t
.
Run
(
"混合调度-Gemini优先选择OAuth账户"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformGemini
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Type
:
AccountTypeAPIKey
},
{
ID
:
2
,
Platform
:
PlatformGemini
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Type
:
AccountTypeOAuth
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
"session-456"
:
1
},
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
}
groupRepo
:=
&
mockGroupRepoForGateway
{
groups
:
map
[
int64
]
*
Group
{
groupID
:
{
ID
:
groupID
,
Name
:
"route-group-hit"
,
Platform
:
PlatformAnthropic
,
Status
:
StatusActive
,
Hydrated
:
true
,
ModelRoutingEnabled
:
true
,
ModelRouting
:
map
[
string
][]
int64
{
requestedModel
:
{
1
,
2
},
},
},
},
}
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
nil
,
""
,
"gemini-2.5-pro"
,
nil
,
PlatformGemini
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
,
"同优先级且未使用时应优先选择OAuth账户"
)
})
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
groupRepo
:
groupRepo
,
}
t
.
Run
(
"混合调度-包含启用mixed_scheduling的antigravity账户"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
},
{
ID
:
2
,
Platform
:
PlatformAntigravity
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Extra
:
map
[
string
]
any
{
"mixed_scheduling"
:
true
}},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
acc
,
err
:=
svc
.
selectAccountForModelWithPlatform
(
ctx
,
&
groupID
,
"session-456"
,
requestedModel
,
nil
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
1
),
acc
.
ID
)
}
cache
:=
&
mockGatewayCacheForPlatform
{}
func
TestGatewayService_SelectAccountForModelWithPlatform_RoutedFallbackToNormal
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
groupID
:=
int64
(
12
)
requestedModel
:=
"claude-3-5-sonnet-20241022"
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
}
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
nil
,
""
,
"claude-3-5-sonnet-20241022"
,
nil
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
,
"应选择优先级最高的账户(包含启用混合调度的antigravity)"
)
})
cache
:=
&
mockGatewayCacheForPlatform
{}
t
.
Run
(
"混合调度-过滤未启用mixed_scheduling的antigravity账户"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
},
{
ID
:
2
,
Platform
:
PlatformAntigravity
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
// 未启用 mixed_scheduling
groupRepo
:=
&
mockGroupRepoForGateway
{
groups
:
map
[
int64
]
*
Group
{
groupID
:
{
ID
:
groupID
,
Name
:
"route-fallback"
,
Platform
:
PlatformAnthropic
,
Status
:
StatusActive
,
Hydrated
:
true
,
ModelRoutingEnabled
:
true
,
ModelRouting
:
map
[
string
][]
int64
{
requestedModel
:
{
99
},
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
},
}
cache
:=
&
mockGatewayCacheForPlatform
{}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
groupRepo
:
groupRepo
,
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
}
acc
,
err
:=
svc
.
selectAccountForModelWithPlatform
(
ctx
,
&
groupID
,
""
,
requestedModel
,
nil
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
1
),
acc
.
ID
)
}
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
nil
,
""
,
"claude-3-5-sonnet-20241022"
,
nil
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
1
),
acc
.
ID
,
"未启用mixed_scheduling的antigravity账户应被过滤"
)
require
.
Equal
(
t
,
PlatformAnthropic
,
acc
.
Platform
)
})
func
TestGatewayService_SelectAccountForModelWithPlatform_NoModelSupport
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
t
.
Run
(
"混合调度-粘性会话命中启用mixed_scheduling的antigravity账户"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
{
ID
:
2
,
Platform
:
PlatformAntigravity
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
,
Extra
:
map
[
string
]
any
{
"mixed_scheduling"
:
true
}},
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Credentials
:
map
[
string
]
any
{
"model_mapping"
:
map
[
string
]
any
{
"claude-3-5-haiku-20241022"
:
"claude-3-5-haiku-20241022"
}},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
"session-123"
:
2
},
}
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
}
cache
:=
&
mockGatewayCacheForPlatform
{}
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
nil
,
"session-123"
,
"claude-3-5-sonnet-20241022"
,
nil
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
,
"应返回粘性会话绑定的启用mixed_scheduling的antigravity账户"
)
}
)
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
}
t
.
Run
(
"混合调度-粘性会话命中未启用mixed_scheduling的antigravity账户-降级选择"
,
func
(
t
*
testing
.
T
)
{
acc
,
err
:=
svc
.
selectAccountForModelWithPlatform
(
ctx
,
nil
,
""
,
"claude-3-5-sonnet-20241022"
,
nil
,
PlatformAnthropic
)
require
.
Error
(
t
,
err
)
require
.
Nil
(
t
,
acc
)
require
.
Contains
(
t
,
err
.
Error
(),
"supporting model"
)
}
func
TestGatewayService_SelectAccountForModelWithPlatform_GeminiPreferOAuth
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformGemini
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Type
:
AccountTypeAPIKey
},
{
ID
:
2
,
Platform
:
PlatformGemini
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Type
:
AccountTypeOAuth
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
}
acc
,
err
:=
svc
.
selectAccountForModelWithPlatform
(
ctx
,
nil
,
""
,
"gemini-2.5-pro"
,
nil
,
PlatformGemini
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
)
}
func
TestGatewayService_SelectAccountForModelWithPlatform_StickyInGroup
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
groupID
:=
int64
(
50
)
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
AccountGroups
:
[]
AccountGroup
{{
GroupID
:
groupID
}}},
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
,
AccountGroups
:
[]
AccountGroup
{{
GroupID
:
groupID
}}},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
"session-group"
:
1
},
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
}
acc
,
err
:=
svc
.
selectAccountForModelWithPlatform
(
ctx
,
&
groupID
,
"session-group"
,
""
,
nil
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
1
),
acc
.
ID
)
}
func
TestGatewayService_SelectAccountForModelWithPlatform_StickyModelMismatchFallback
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Credentials
:
map
[
string
]
any
{
"model_mapping"
:
map
[
string
]
any
{
"claude-3-5-haiku-20241022"
:
"claude-3-5-haiku-20241022"
}},
},
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
"session-miss"
:
1
},
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
}
acc
,
err
:=
svc
.
selectAccountForModelWithPlatform
(
ctx
,
nil
,
"session-miss"
,
"claude-3-5-sonnet-20241022"
,
nil
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
)
}
func
TestGatewayService_SelectAccountForModelWithPlatform_PreferNeverUsed
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
lastUsed
:=
time
.
Now
()
.
Add
(
-
1
*
time
.
Hour
)
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
LastUsedAt
:
&
lastUsed
},
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
}
acc
,
err
:=
svc
.
selectAccountForModelWithPlatform
(
ctx
,
nil
,
""
,
"claude-3-5-sonnet-20241022"
,
nil
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
)
}
func
TestGatewayService_SelectAccountForModelWithPlatform_NoAccounts
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{},
accountsByID
:
map
[
int64
]
*
Account
{},
}
cache
:=
&
mockGatewayCacheForPlatform
{}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
}
acc
,
err
:=
svc
.
selectAccountForModelWithPlatform
(
ctx
,
nil
,
""
,
""
,
nil
,
PlatformAnthropic
)
require
.
Error
(
t
,
err
)
require
.
Nil
(
t
,
acc
)
require
.
Contains
(
t
,
err
.
Error
(),
"no available accounts"
)
}
func
TestGatewayService_isModelSupportedByAccount
(
t
*
testing
.
T
)
{
svc
:=
&
GatewayService
{}
tests
:=
[]
struct
{
name
string
account
*
Account
model
string
expected
bool
}{
{
name
:
"Antigravity平台-支持claude模型"
,
account
:
&
Account
{
Platform
:
PlatformAntigravity
},
model
:
"claude-3-5-sonnet-20241022"
,
expected
:
true
,
},
{
name
:
"Antigravity平台-支持gemini模型"
,
account
:
&
Account
{
Platform
:
PlatformAntigravity
},
model
:
"gemini-2.5-flash"
,
expected
:
true
,
},
{
name
:
"Antigravity平台-不支持gpt模型"
,
account
:
&
Account
{
Platform
:
PlatformAntigravity
},
model
:
"gpt-4"
,
expected
:
false
,
},
{
name
:
"Anthropic平台-无映射配置-支持所有模型"
,
account
:
&
Account
{
Platform
:
PlatformAnthropic
},
model
:
"claude-3-5-sonnet-20241022"
,
expected
:
true
,
},
{
name
:
"Anthropic平台-有映射配置-只支持配置的模型"
,
account
:
&
Account
{
Platform
:
PlatformAnthropic
,
Credentials
:
map
[
string
]
any
{
"model_mapping"
:
map
[
string
]
any
{
"claude-opus-4"
:
"x"
}},
},
model
:
"claude-3-5-sonnet-20241022"
,
expected
:
false
,
},
{
name
:
"Anthropic平台-有映射配置-支持配置的模型"
,
account
:
&
Account
{
Platform
:
PlatformAnthropic
,
Credentials
:
map
[
string
]
any
{
"model_mapping"
:
map
[
string
]
any
{
"claude-3-5-sonnet-20241022"
:
"x"
}},
},
model
:
"claude-3-5-sonnet-20241022"
,
expected
:
true
,
},
}
for
_
,
tt
:=
range
tests
{
t
.
Run
(
tt
.
name
,
func
(
t
*
testing
.
T
)
{
got
:=
svc
.
isModelSupportedByAccount
(
tt
.
account
,
tt
.
model
)
require
.
Equal
(
t
,
tt
.
expected
,
got
)
})
}
}
// TestGatewayService_selectAccountWithMixedScheduling 测试混合调度
func
TestGatewayService_selectAccountWithMixedScheduling
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
t
.
Run
(
"混合调度-Gemini优先选择OAuth账户"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
Platform
Anthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
{
ID
:
2
,
Platform
:
Platform
Antigravity
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
}
,
// 未启用 mixed_scheduling
{
ID
:
1
,
Platform
:
Platform
Gemini
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Type
:
AccountTypeAPIKey
},
{
ID
:
2
,
Platform
:
Platform
Gemini
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Type
:
AccountTypeOAuth
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
...
...
@@ -807,9 +1070,7 @@ func TestGatewayService_selectAccountWithMixedScheduling(t *testing.T) {
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
"session-123"
:
2
},
}
cache
:=
&
mockGatewayCacheForPlatform
{}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
...
...
@@ -817,16 +1078,17 @@ func TestGatewayService_selectAccountWithMixedScheduling(t *testing.T) {
cfg
:
testConfig
(),
}
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
nil
,
"
session-123"
,
"claude-3-5-sonnet-20241022
"
,
nil
,
Platform
Anthropic
)
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
nil
,
"
"
,
"gemini-2.5-pro
"
,
nil
,
Platform
Gemini
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
1
),
acc
.
ID
,
"
粘性会话绑定的账户未启用mixed_scheduling,应降级选择anthropic
账户"
)
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
,
"
同优先级且未使用时应优先选择OAuth
账户"
)
})
t
.
Run
(
"混合调度-
仅有
启用mixed_scheduling的antigravity账户"
,
func
(
t
*
testing
.
T
)
{
t
.
Run
(
"混合调度-
包含
启用mixed_scheduling的antigravity账户"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAntigravity
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Extra
:
map
[
string
]
any
{
"mixed_scheduling"
:
true
}},
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
},
{
ID
:
2
,
Platform
:
PlatformAntigravity
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Extra
:
map
[
string
]
any
{
"mixed_scheduling"
:
true
}},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
...
...
@@ -845,14 +1107,16 @@ func TestGatewayService_selectAccountWithMixedScheduling(t *testing.T) {
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
nil
,
""
,
"claude-3-5-sonnet-20241022"
,
nil
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
1
),
acc
.
ID
)
require
.
Equal
(
t
,
PlatformAntigravity
,
acc
.
Platform
)
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
,
"应选择优先级最高的账户(包含启用混合调度的antigravity)"
)
})
t
.
Run
(
"混合调度-无可用账户"
,
func
(
t
*
testing
.
T
)
{
t
.
Run
(
"混合调度-路由优先选择路由账号"
,
func
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
30
)
requestedModel
:=
"claude-3-5-sonnet-20241022"
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAntigravity
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
// 未启用 mixed_scheduling
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
{
ID
:
2
,
Platform
:
PlatformAntigravity
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
,
Extra
:
map
[
string
]
any
{
"mixed_scheduling"
:
true
}},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
...
...
@@ -862,26 +1126,565 @@ func TestGatewayService_selectAccountWithMixedScheduling(t *testing.T) {
cache
:=
&
mockGatewayCacheForPlatform
{}
groupRepo
:=
&
mockGroupRepoForGateway
{
groups
:
map
[
int64
]
*
Group
{
groupID
:
{
ID
:
groupID
,
Name
:
"route-mixed-select"
,
Platform
:
PlatformAnthropic
,
Status
:
StatusActive
,
Hydrated
:
true
,
ModelRoutingEnabled
:
true
,
ModelRouting
:
map
[
string
][]
int64
{
requestedModel
:
{
2
},
},
},
},
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
groupRepo
:
groupRepo
,
}
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
nil
,
""
,
"claude-3-5-sonnet-20241022"
,
nil
,
PlatformAnthropic
)
require
.
Error
(
t
,
err
)
require
.
Nil
(
t
,
acc
)
require
.
Contains
(
t
,
err
.
Error
(),
"no available accounts"
)
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
&
groupID
,
""
,
requestedModel
,
nil
,
PlatformAnthropic
)
require
.
No
Error
(
t
,
err
)
require
.
Not
Nil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
)
})
}
// TestAccount_IsMixedSchedulingEnabled 测试混合调度开关检查
func
TestAccount_IsMixedSchedulingEnabled
(
t
*
testing
.
T
)
{
tests
:=
[]
struct
{
name
string
account
Account
expected
bool
}{
t
.
Run
(
"混合调度-路由粘性命中"
,
func
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
31
)
requestedModel
:=
"claude-3-5-sonnet-20241022"
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
{
ID
:
2
,
Platform
:
PlatformAntigravity
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
,
Extra
:
map
[
string
]
any
{
"mixed_scheduling"
:
true
},
AccountGroups
:
[]
AccountGroup
{{
GroupID
:
groupID
}}},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
"session-777"
:
2
},
}
groupRepo
:=
&
mockGroupRepoForGateway
{
groups
:
map
[
int64
]
*
Group
{
groupID
:
{
ID
:
groupID
,
Name
:
"route-mixed-sticky"
,
Platform
:
PlatformAnthropic
,
Status
:
StatusActive
,
Hydrated
:
true
,
ModelRoutingEnabled
:
true
,
ModelRouting
:
map
[
string
][]
int64
{
requestedModel
:
{
2
},
},
},
},
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
groupRepo
:
groupRepo
,
}
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
&
groupID
,
"session-777"
,
requestedModel
,
nil
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
)
})
t
.
Run
(
"混合调度-路由账号缺失回退"
,
func
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
32
)
requestedModel
:=
"claude-3-5-sonnet-20241022"
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
{
ID
:
2
,
Platform
:
PlatformAntigravity
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
,
Extra
:
map
[
string
]
any
{
"mixed_scheduling"
:
true
}},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{}
groupRepo
:=
&
mockGroupRepoForGateway
{
groups
:
map
[
int64
]
*
Group
{
groupID
:
{
ID
:
groupID
,
Name
:
"route-mixed-miss"
,
Platform
:
PlatformAnthropic
,
Status
:
StatusActive
,
Hydrated
:
true
,
ModelRoutingEnabled
:
true
,
ModelRouting
:
map
[
string
][]
int64
{
requestedModel
:
{
99
},
},
},
},
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
groupRepo
:
groupRepo
,
}
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
&
groupID
,
""
,
requestedModel
,
nil
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
1
),
acc
.
ID
)
})
t
.
Run
(
"混合调度-路由账号未启用mixed_scheduling回退"
,
func
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
33
)
requestedModel
:=
"claude-3-5-sonnet-20241022"
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
{
ID
:
2
,
Platform
:
PlatformAntigravity
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
},
// 未启用 mixed_scheduling
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{}
groupRepo
:=
&
mockGroupRepoForGateway
{
groups
:
map
[
int64
]
*
Group
{
groupID
:
{
ID
:
groupID
,
Name
:
"route-mixed-disabled"
,
Platform
:
PlatformAnthropic
,
Status
:
StatusActive
,
Hydrated
:
true
,
ModelRoutingEnabled
:
true
,
ModelRouting
:
map
[
string
][]
int64
{
requestedModel
:
{
2
},
},
},
},
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
groupRepo
:
groupRepo
,
}
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
&
groupID
,
""
,
requestedModel
,
nil
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
1
),
acc
.
ID
)
})
t
.
Run
(
"混合调度-路由过滤覆盖"
,
func
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
35
)
requestedModel
:=
"claude-3-5-sonnet-20241022"
resetAt
:=
time
.
Now
()
.
Add
(
10
*
time
.
Minute
)
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
false
},
{
ID
:
3
,
Platform
:
PlatformAntigravity
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
{
ID
:
4
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Extra
:
map
[
string
]
any
{
"model_rate_limits"
:
map
[
string
]
any
{
"claude_sonnet"
:
map
[
string
]
any
{
"rate_limit_reset_at"
:
resetAt
.
Format
(
time
.
RFC3339
),
},
},
},
},
{
ID
:
5
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Credentials
:
map
[
string
]
any
{
"model_mapping"
:
map
[
string
]
any
{
"claude-3-5-haiku-20241022"
:
"claude-3-5-haiku-20241022"
}},
},
{
ID
:
6
,
Platform
:
PlatformAnthropic
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
},
{
ID
:
7
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{}
groupRepo
:=
&
mockGroupRepoForGateway
{
groups
:
map
[
int64
]
*
Group
{
groupID
:
{
ID
:
groupID
,
Name
:
"route-mixed-filter"
,
Platform
:
PlatformAnthropic
,
Status
:
StatusActive
,
Hydrated
:
true
,
ModelRoutingEnabled
:
true
,
ModelRouting
:
map
[
string
][]
int64
{
requestedModel
:
{
1
,
2
,
3
,
4
,
5
,
6
,
7
},
},
},
},
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
groupRepo
:
groupRepo
,
}
excluded
:=
map
[
int64
]
struct
{}{
1
:
{}}
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
&
groupID
,
""
,
requestedModel
,
excluded
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
7
),
acc
.
ID
)
})
t
.
Run
(
"混合调度-粘性命中分组账号"
,
func
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
34
)
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
AccountGroups
:
[]
AccountGroup
{{
GroupID
:
groupID
}}},
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
,
AccountGroups
:
[]
AccountGroup
{{
GroupID
:
groupID
}}},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
"session-group"
:
1
},
}
groupRepo
:=
&
mockGroupRepoForGateway
{
groups
:
map
[
int64
]
*
Group
{
groupID
:
{
ID
:
groupID
,
Platform
:
PlatformAnthropic
,
Status
:
StatusActive
,
Hydrated
:
true
,
},
},
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
groupRepo
:
groupRepo
,
}
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
&
groupID
,
"session-group"
,
"claude-3-5-sonnet-20241022"
,
nil
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
1
),
acc
.
ID
)
})
t
.
Run
(
"混合调度-过滤未启用mixed_scheduling的antigravity账户"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
},
{
ID
:
2
,
Platform
:
PlatformAntigravity
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
// 未启用 mixed_scheduling
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
}
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
nil
,
""
,
"claude-3-5-sonnet-20241022"
,
nil
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
1
),
acc
.
ID
,
"未启用mixed_scheduling的antigravity账户应被过滤"
)
require
.
Equal
(
t
,
PlatformAnthropic
,
acc
.
Platform
)
})
t
.
Run
(
"混合调度-粘性会话命中启用mixed_scheduling的antigravity账户"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
{
ID
:
2
,
Platform
:
PlatformAntigravity
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
,
Extra
:
map
[
string
]
any
{
"mixed_scheduling"
:
true
}},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
"session-123"
:
2
},
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
}
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
nil
,
"session-123"
,
"claude-3-5-sonnet-20241022"
,
nil
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
,
"应返回粘性会话绑定的启用mixed_scheduling的antigravity账户"
)
})
t
.
Run
(
"混合调度-粘性会话命中未启用mixed_scheduling的antigravity账户-降级选择"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
{
ID
:
2
,
Platform
:
PlatformAntigravity
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
},
// 未启用 mixed_scheduling
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
"session-123"
:
2
},
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
}
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
nil
,
"session-123"
,
"claude-3-5-sonnet-20241022"
,
nil
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
1
),
acc
.
ID
,
"粘性会话绑定的账户未启用mixed_scheduling,应降级选择anthropic账户"
)
})
t
.
Run
(
"混合调度-粘性会话不可调度-清理并回退"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAntigravity
,
Priority
:
1
,
Status
:
StatusDisabled
,
Schedulable
:
true
,
Extra
:
map
[
string
]
any
{
"mixed_scheduling"
:
true
}},
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
"session-123"
:
1
},
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
}
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
nil
,
"session-123"
,
"claude-3-5-sonnet-20241022"
,
nil
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
)
require
.
Equal
(
t
,
1
,
cache
.
deletedSessions
[
"session-123"
])
require
.
Equal
(
t
,
int64
(
2
),
cache
.
sessionBindings
[
"session-123"
])
})
t
.
Run
(
"混合调度-路由粘性不可调度-清理并回退"
,
func
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
12
)
requestedModel
:=
"claude-3-5-sonnet-20241022"
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAntigravity
,
Priority
:
1
,
Status
:
StatusDisabled
,
Schedulable
:
true
,
Extra
:
map
[
string
]
any
{
"mixed_scheduling"
:
true
}},
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
"session-123"
:
1
},
}
groupRepo
:=
&
mockGroupRepoForGateway
{
groups
:
map
[
int64
]
*
Group
{
groupID
:
{
ID
:
groupID
,
Name
:
"route-mixed"
,
Platform
:
PlatformAnthropic
,
Status
:
StatusActive
,
Hydrated
:
true
,
ModelRoutingEnabled
:
true
,
ModelRouting
:
map
[
string
][]
int64
{
requestedModel
:
{
1
,
2
},
},
},
},
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
groupRepo
:
groupRepo
,
}
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
&
groupID
,
"session-123"
,
requestedModel
,
nil
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
)
require
.
Equal
(
t
,
1
,
cache
.
deletedSessions
[
"session-123"
])
require
.
Equal
(
t
,
int64
(
2
),
cache
.
sessionBindings
[
"session-123"
])
})
t
.
Run
(
"混合调度-仅有启用mixed_scheduling的antigravity账户"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAntigravity
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Extra
:
map
[
string
]
any
{
"mixed_scheduling"
:
true
}},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
}
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
nil
,
""
,
"claude-3-5-sonnet-20241022"
,
nil
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
1
),
acc
.
ID
)
require
.
Equal
(
t
,
PlatformAntigravity
,
acc
.
Platform
)
})
t
.
Run
(
"混合调度-无可用账户"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAntigravity
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
// 未启用 mixed_scheduling
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
}
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
nil
,
""
,
"claude-3-5-sonnet-20241022"
,
nil
,
PlatformAnthropic
)
require
.
Error
(
t
,
err
)
require
.
Nil
(
t
,
acc
)
require
.
Contains
(
t
,
err
.
Error
(),
"no available accounts"
)
})
t
.
Run
(
"混合调度-不支持模型返回错误"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Credentials
:
map
[
string
]
any
{
"model_mapping"
:
map
[
string
]
any
{
"claude-3-5-haiku-20241022"
:
"claude-3-5-haiku-20241022"
}},
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
}
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
nil
,
""
,
"claude-3-5-sonnet-20241022"
,
nil
,
PlatformAnthropic
)
require
.
Error
(
t
,
err
)
require
.
Nil
(
t
,
acc
)
require
.
Contains
(
t
,
err
.
Error
(),
"supporting model"
)
})
t
.
Run
(
"混合调度-优先未使用账号"
,
func
(
t
*
testing
.
T
)
{
lastUsed
:=
time
.
Now
()
.
Add
(
-
2
*
time
.
Hour
)
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
LastUsedAt
:
&
lastUsed
},
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
testConfig
(),
}
acc
,
err
:=
svc
.
selectAccountWithMixedScheduling
(
ctx
,
nil
,
""
,
"claude-3-5-sonnet-20241022"
,
nil
,
PlatformAnthropic
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
)
})
}
// TestAccount_IsMixedSchedulingEnabled 测试混合调度开关检查
func
TestAccount_IsMixedSchedulingEnabled
(
t
*
testing
.
T
)
{
tests
:=
[]
struct
{
name
string
account
Account
expected
bool
}{
{
name
:
"非antigravity平台-返回false"
,
account
:
Account
{
Platform
:
PlatformAnthropic
},
...
...
@@ -914,125 +1717,754 @@ func TestAccount_IsMixedSchedulingEnabled(t *testing.T) {
},
}
for
_
,
tt
:=
range
tests
{
t
.
Run
(
tt
.
name
,
func
(
t
*
testing
.
T
)
{
got
:=
tt
.
account
.
IsMixedSchedulingEnabled
()
require
.
Equal
(
t
,
tt
.
expected
,
got
)
})
}
}
for
_
,
tt
:=
range
tests
{
t
.
Run
(
tt
.
name
,
func
(
t
*
testing
.
T
)
{
got
:=
tt
.
account
.
IsMixedSchedulingEnabled
()
require
.
Equal
(
t
,
tt
.
expected
,
got
)
})
}
}
// mockConcurrencyService for testing
type
mockConcurrencyService
struct
{
accountLoads
map
[
int64
]
*
AccountLoadInfo
accountWaitCounts
map
[
int64
]
int
acquireResults
map
[
int64
]
bool
}
func
(
m
*
mockConcurrencyService
)
GetAccountsLoadBatch
(
ctx
context
.
Context
,
accounts
[]
AccountWithConcurrency
)
(
map
[
int64
]
*
AccountLoadInfo
,
error
)
{
if
m
.
accountLoads
==
nil
{
return
map
[
int64
]
*
AccountLoadInfo
{},
nil
}
result
:=
make
(
map
[
int64
]
*
AccountLoadInfo
)
for
_
,
acc
:=
range
accounts
{
if
load
,
ok
:=
m
.
accountLoads
[
acc
.
ID
];
ok
{
result
[
acc
.
ID
]
=
load
}
else
{
result
[
acc
.
ID
]
=
&
AccountLoadInfo
{
AccountID
:
acc
.
ID
,
CurrentConcurrency
:
0
,
WaitingCount
:
0
,
LoadRate
:
0
,
}
}
}
return
result
,
nil
}
func
(
m
*
mockConcurrencyService
)
GetAccountWaitingCount
(
ctx
context
.
Context
,
accountID
int64
)
(
int
,
error
)
{
if
m
.
accountWaitCounts
==
nil
{
return
0
,
nil
}
return
m
.
accountWaitCounts
[
accountID
],
nil
}
type
mockConcurrencyCache
struct
{
acquireAccountCalls
int
loadBatchCalls
int
acquireResults
map
[
int64
]
bool
loadBatchErr
error
loadMap
map
[
int64
]
*
AccountLoadInfo
waitCounts
map
[
int64
]
int
skipDefaultLoad
bool
}
func
(
m
*
mockConcurrencyCache
)
AcquireAccountSlot
(
ctx
context
.
Context
,
accountID
int64
,
maxConcurrency
int
,
requestID
string
)
(
bool
,
error
)
{
m
.
acquireAccountCalls
++
if
m
.
acquireResults
!=
nil
{
if
result
,
ok
:=
m
.
acquireResults
[
accountID
];
ok
{
return
result
,
nil
}
}
return
true
,
nil
}
func
(
m
*
mockConcurrencyCache
)
ReleaseAccountSlot
(
ctx
context
.
Context
,
accountID
int64
,
requestID
string
)
error
{
return
nil
}
func
(
m
*
mockConcurrencyCache
)
GetAccountConcurrency
(
ctx
context
.
Context
,
accountID
int64
)
(
int
,
error
)
{
return
0
,
nil
}
func
(
m
*
mockConcurrencyCache
)
IncrementAccountWaitCount
(
ctx
context
.
Context
,
accountID
int64
,
maxWait
int
)
(
bool
,
error
)
{
return
true
,
nil
}
func
(
m
*
mockConcurrencyCache
)
DecrementAccountWaitCount
(
ctx
context
.
Context
,
accountID
int64
)
error
{
return
nil
}
func
(
m
*
mockConcurrencyCache
)
GetAccountWaitingCount
(
ctx
context
.
Context
,
accountID
int64
)
(
int
,
error
)
{
if
m
.
waitCounts
!=
nil
{
if
count
,
ok
:=
m
.
waitCounts
[
accountID
];
ok
{
return
count
,
nil
}
}
return
0
,
nil
}
func
(
m
*
mockConcurrencyCache
)
AcquireUserSlot
(
ctx
context
.
Context
,
userID
int64
,
maxConcurrency
int
,
requestID
string
)
(
bool
,
error
)
{
return
true
,
nil
}
func
(
m
*
mockConcurrencyCache
)
ReleaseUserSlot
(
ctx
context
.
Context
,
userID
int64
,
requestID
string
)
error
{
return
nil
}
func
(
m
*
mockConcurrencyCache
)
GetUserConcurrency
(
ctx
context
.
Context
,
userID
int64
)
(
int
,
error
)
{
return
0
,
nil
}
func
(
m
*
mockConcurrencyCache
)
IncrementWaitCount
(
ctx
context
.
Context
,
userID
int64
,
maxWait
int
)
(
bool
,
error
)
{
return
true
,
nil
}
func
(
m
*
mockConcurrencyCache
)
DecrementWaitCount
(
ctx
context
.
Context
,
userID
int64
)
error
{
return
nil
}
func
(
m
*
mockConcurrencyCache
)
GetAccountsLoadBatch
(
ctx
context
.
Context
,
accounts
[]
AccountWithConcurrency
)
(
map
[
int64
]
*
AccountLoadInfo
,
error
)
{
m
.
loadBatchCalls
++
if
m
.
loadBatchErr
!=
nil
{
return
nil
,
m
.
loadBatchErr
}
result
:=
make
(
map
[
int64
]
*
AccountLoadInfo
,
len
(
accounts
))
if
m
.
skipDefaultLoad
&&
m
.
loadMap
!=
nil
{
for
_
,
acc
:=
range
accounts
{
if
load
,
ok
:=
m
.
loadMap
[
acc
.
ID
];
ok
{
result
[
acc
.
ID
]
=
load
}
}
return
result
,
nil
}
for
_
,
acc
:=
range
accounts
{
if
m
.
loadMap
!=
nil
{
if
load
,
ok
:=
m
.
loadMap
[
acc
.
ID
];
ok
{
result
[
acc
.
ID
]
=
load
continue
}
}
result
[
acc
.
ID
]
=
&
AccountLoadInfo
{
AccountID
:
acc
.
ID
,
CurrentConcurrency
:
0
,
WaitingCount
:
0
,
LoadRate
:
0
,
}
}
return
result
,
nil
}
func
(
m
*
mockConcurrencyCache
)
CleanupExpiredAccountSlots
(
ctx
context
.
Context
,
accountID
int64
)
error
{
return
nil
}
// TestGatewayService_SelectAccountWithLoadAwareness tests load-aware account selection
func
TestGatewayService_SelectAccountWithLoadAwareness
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
t
.
Run
(
"禁用负载批量查询-降级到传统选择"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
{
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
]
}
cache
:=
&
mockGatewayCacheForPlatform
{}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
false
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
nil
,
// No concurrency service
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
nil
,
""
,
"claude-3-5-sonnet-20241022"
,
nil
,
""
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
Account
)
require
.
Equal
(
t
,
int64
(
1
),
result
.
Account
.
ID
,
"应选择优先级最高的账号"
)
})
t
.
Run
(
"模型路由-无ConcurrencyService也生效"
,
func
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
1
)
sessionHash
:=
"sticky"
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
,
AccountGroups
:
[]
AccountGroup
{{
GroupID
:
groupID
}}},
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
,
AccountGroups
:
[]
AccountGroup
{{
GroupID
:
groupID
}}},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
sessionHash
:
1
},
}
groupRepo
:=
&
mockGroupRepoForGateway
{
groups
:
map
[
int64
]
*
Group
{
groupID
:
{
ID
:
groupID
,
Platform
:
PlatformAnthropic
,
Status
:
StatusActive
,
Hydrated
:
true
,
ModelRoutingEnabled
:
true
,
ModelRouting
:
map
[
string
][]
int64
{
"claude-a"
:
{
1
},
"claude-b"
:
{
2
},
},
},
},
}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
true
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
groupRepo
:
groupRepo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
nil
,
// legacy path
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
&
groupID
,
sessionHash
,
"claude-b"
,
nil
,
""
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
Account
)
require
.
Equal
(
t
,
int64
(
2
),
result
.
Account
.
ID
,
"切换到 claude-b 时应按模型路由切换账号"
)
require
.
Equal
(
t
,
int64
(
2
),
cache
.
sessionBindings
[
sessionHash
],
"粘性绑定应更新为路由选择的账号"
)
})
t
.
Run
(
"无ConcurrencyService-降级到传统选择"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
{
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
]
}
cache
:=
&
mockGatewayCacheForPlatform
{}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
true
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
nil
,
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
nil
,
""
,
"claude-3-5-sonnet-20241022"
,
nil
,
""
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
Account
)
require
.
Equal
(
t
,
int64
(
2
),
result
.
Account
.
ID
,
"应选择优先级最高的账号"
)
})
t
.
Run
(
"排除账号-不选择被排除的账号"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
{
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
]
}
cache
:=
&
mockGatewayCacheForPlatform
{}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
false
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
nil
,
}
excludedIDs
:=
map
[
int64
]
struct
{}{
1
:
{}}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
nil
,
""
,
"claude-3-5-sonnet-20241022"
,
excludedIDs
,
""
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
Account
)
require
.
Equal
(
t
,
int64
(
2
),
result
.
Account
.
ID
,
"不应选择被排除的账号"
)
})
t
.
Run
(
"粘性命中-不调用GetByID"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
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
]
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
"sticky"
:
1
},
}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
true
concurrencyCache
:=
&
mockConcurrencyCache
{}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
),
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
nil
,
"sticky"
,
"claude-3-5-sonnet-20241022"
,
nil
,
""
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
Account
)
require
.
Equal
(
t
,
int64
(
1
),
result
.
Account
.
ID
)
require
.
Equal
(
t
,
0
,
repo
.
getByIDCalls
,
"粘性命中不应调用GetByID"
)
require
.
Equal
(
t
,
0
,
concurrencyCache
.
loadBatchCalls
,
"粘性命中应在负载批量查询前返回"
)
})
t
.
Run
(
"粘性账号不在候选集-回退负载感知选择"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
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
]
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
"sticky"
:
1
},
}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
true
concurrencyCache
:=
&
mockConcurrencyCache
{}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
),
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
nil
,
"sticky"
,
"claude-3-5-sonnet-20241022"
,
nil
,
""
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
Account
)
require
.
Equal
(
t
,
int64
(
2
),
result
.
Account
.
ID
,
"粘性账号不在候选集时应回退到可用账号"
)
require
.
Equal
(
t
,
0
,
repo
.
getByIDCalls
,
"粘性账号缺失不应回退到GetByID"
)
require
.
Equal
(
t
,
1
,
concurrencyCache
.
loadBatchCalls
,
"应继续进行负载批量查询"
)
})
t
.
Run
(
"粘性账号禁用-清理会话并回退选择"
,
func
(
t
*
testing
.
T
)
{
testCtx
:=
context
.
WithValue
(
ctx
,
ctxkey
.
ForcePlatform
,
PlatformAnthropic
)
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
false
,
Concurrency
:
5
},
{
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
]
}
repo
.
listPlatformFunc
=
func
(
ctx
context
.
Context
,
platform
string
)
([]
Account
,
error
)
{
return
repo
.
accounts
,
nil
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
"sticky"
:
1
},
}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
true
concurrencyCache
:=
&
mockConcurrencyCache
{}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
),
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
testCtx
,
nil
,
"sticky"
,
"claude-3-5-sonnet-20241022"
,
nil
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
Account
)
require
.
Equal
(
t
,
int64
(
2
),
result
.
Account
.
ID
,
"粘性账号禁用时应回退到可用账号"
)
updatedID
,
ok
:=
cache
.
sessionBindings
[
"sticky"
]
require
.
True
(
t
,
ok
,
"粘性会话应更新绑定"
)
require
.
Equal
(
t
,
int64
(
2
),
updatedID
,
"粘性会话应绑定到新账号"
)
})
t
.
Run
(
"无可用账号-返回错误"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{},
accountsByID
:
map
[
int64
]
*
Account
{},
}
cache
:=
&
mockGatewayCacheForPlatform
{}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
false
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
nil
,
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
nil
,
""
,
"claude-3-5-sonnet-20241022"
,
nil
,
""
)
require
.
Error
(
t
,
err
)
require
.
Nil
(
t
,
result
)
require
.
Contains
(
t
,
err
.
Error
(),
"no available accounts"
)
})
t
.
Run
(
"过滤不可调度账号-限流账号被跳过"
,
func
(
t
*
testing
.
T
)
{
now
:=
time
.
Now
()
resetAt
:=
now
.
Add
(
10
*
time
.
Minute
)
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
,
RateLimitResetAt
:
&
resetAt
},
{
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
]
}
cache
:=
&
mockGatewayCacheForPlatform
{}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
false
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
nil
,
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
nil
,
""
,
"claude-3-5-sonnet-20241022"
,
nil
,
""
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
Account
)
require
.
Equal
(
t
,
int64
(
2
),
result
.
Account
.
ID
,
"应跳过限流账号,选择可用账号"
)
})
t
.
Run
(
"过滤不可调度账号-过载账号被跳过"
,
func
(
t
*
testing
.
T
)
{
now
:=
time
.
Now
()
overloadUntil
:=
now
.
Add
(
10
*
time
.
Minute
)
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
,
OverloadUntil
:
&
overloadUntil
},
{
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
]
}
cache
:=
&
mockGatewayCacheForPlatform
{}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
false
// mockConcurrencyService for testing
type
mockConcurrencyService
struct
{
accountLoads
map
[
int64
]
*
AccountLoadInfo
accountWaitCounts
map
[
int64
]
int
acquireResults
map
[
int64
]
bool
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
nil
,
}
func
(
m
*
mockConcurrencyService
)
GetAccountsLoadBatch
(
ctx
context
.
Context
,
accounts
[]
AccountWithConcurrency
)
(
map
[
int64
]
*
AccountLoadInfo
,
error
)
{
if
m
.
accountLoads
==
nil
{
return
map
[
int64
]
*
AccountLoadInfo
{},
nil
}
result
:=
make
(
map
[
int64
]
*
AccountLoadInfo
)
for
_
,
acc
:=
range
accounts
{
if
load
,
ok
:=
m
.
accountLoads
[
acc
.
ID
];
ok
{
result
[
acc
.
ID
]
=
load
}
else
{
result
[
acc
.
ID
]
=
&
AccountLoadInfo
{
AccountID
:
acc
.
ID
,
CurrentConcurrency
:
0
,
WaitingCount
:
0
,
LoadRate
:
0
,
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
nil
,
""
,
"claude-3-5-sonnet-20241022"
,
nil
,
""
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
Account
)
require
.
Equal
(
t
,
int64
(
2
),
result
.
Account
.
ID
,
"应跳过过载账号,选择可用账号"
)
})
t
.
Run
(
"粘性账号槽位满-返回粘性等待计划"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
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
]
}
}
return
result
,
nil
}
func
(
m
*
mockConcurrencyService
)
GetAccountWaitingCount
(
ctx
context
.
Context
,
accountID
int64
)
(
int
,
error
)
{
if
m
.
accountWaitCounts
==
nil
{
return
0
,
nil
}
return
m
.
accountWaitCounts
[
accountID
],
nil
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
"sticky"
:
1
},
}
type
mockConcurrencyCache
struct
{
acquireAccountCalls
int
loadBatchCalls
int
}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
true
cfg
.
Gateway
.
Scheduling
.
StickySessionMaxWaiting
=
1
func
(
m
*
mockC
oncurrencyCache
)
AcquireAccountSlot
(
ctx
context
.
Context
,
accountID
int64
,
maxConcurrency
int
,
requestID
string
)
(
bool
,
error
)
{
m
.
acquire
AccountCalls
++
return
true
,
nil
}
c
oncurrencyCache
:=
&
mockConcurrencyCache
{
acquire
Results
:
map
[
int64
]
bool
{
1
:
false
},
waitCounts
:
map
[
int64
]
int
{
1
:
0
},
}
func
(
m
*
mockConcurrencyCache
)
ReleaseAccountSlot
(
ctx
context
.
Context
,
accountID
int64
,
requestID
string
)
error
{
return
nil
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
),
}
func
(
m
*
mockConcurrencyCache
)
GetAccountConcurrency
(
ctx
context
.
Context
,
accountID
int64
)
(
int
,
error
)
{
return
0
,
nil
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
nil
,
"sticky"
,
"claude-3-5-sonnet-20241022"
,
nil
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
WaitPlan
)
require
.
Equal
(
t
,
int64
(
1
),
result
.
Account
.
ID
)
require
.
Equal
(
t
,
0
,
concurrencyCache
.
loadBatchCalls
)
})
func
(
m
*
mockConcurrencyCache
)
IncrementAccountWaitCount
(
ctx
context
.
Context
,
accountID
int64
,
maxWait
int
)
(
bool
,
error
)
{
return
true
,
nil
}
t
.
Run
(
"负载批量查询失败-降级旧顺序选择"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
{
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
]
}
func
(
m
*
mockConcurrencyCache
)
DecrementAccountWaitCount
(
ctx
context
.
Context
,
accountID
int64
)
error
{
return
nil
}
cache
:=
&
mockGatewayCacheForPlatform
{}
func
(
m
*
mockConcurrencyCache
)
GetAccountWaitingCount
(
ctx
context
.
Context
,
accountID
int64
)
(
int
,
error
)
{
return
0
,
nil
}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
true
func
(
m
*
mockC
oncurrencyCache
)
AcquireUserSlot
(
ctx
context
.
Context
,
userID
int64
,
maxConcurrency
int
,
requestID
string
)
(
bool
,
error
)
{
return
true
,
nil
}
c
oncurrencyCache
:=
&
mockConcurrencyCache
{
loadBatchErr
:
errors
.
New
(
"load batch failed"
),
}
func
(
m
*
mockConcurrencyCache
)
ReleaseUserSlot
(
ctx
context
.
Context
,
userID
int64
,
requestID
string
)
error
{
return
nil
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
),
}
func
(
m
*
mockConcurrencyCache
)
GetUserConcurrency
(
ctx
context
.
Context
,
userID
int64
)
(
int
,
error
)
{
return
0
,
nil
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
nil
,
"legacy"
,
"claude-3-5-sonnet-20241022"
,
nil
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
Account
)
require
.
Equal
(
t
,
int64
(
2
),
result
.
Account
.
ID
)
require
.
Equal
(
t
,
int64
(
2
),
cache
.
sessionBindings
[
"legacy"
])
})
t
.
Run
(
"模型路由-粘性账号等待计划"
,
func
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
20
)
sessionHash
:=
"route-sticky"
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
{
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
]
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
sessionHash
:
1
},
}
groupRepo
:=
&
mockGroupRepoForGateway
{
groups
:
map
[
int64
]
*
Group
{
groupID
:
{
ID
:
groupID
,
Platform
:
PlatformAnthropic
,
Status
:
StatusActive
,
Hydrated
:
true
,
ModelRoutingEnabled
:
true
,
ModelRouting
:
map
[
string
][]
int64
{
"claude-3-5-sonnet-20241022"
:
{
1
,
2
},
},
},
},
}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
true
cfg
.
Gateway
.
Scheduling
.
StickySessionMaxWaiting
=
1
concurrencyCache
:=
&
mockConcurrencyCache
{
acquireResults
:
map
[
int64
]
bool
{
1
:
false
},
waitCounts
:
map
[
int64
]
int
{
1
:
0
},
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
groupRepo
:
groupRepo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
),
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
&
groupID
,
sessionHash
,
"claude-3-5-sonnet-20241022"
,
nil
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
WaitPlan
)
require
.
Equal
(
t
,
int64
(
1
),
result
.
Account
.
ID
)
})
t
.
Run
(
"模型路由-粘性账号命中"
,
func
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
20
)
sessionHash
:=
"route-hit"
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
{
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
]
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
sessionHash
:
1
},
}
groupRepo
:=
&
mockGroupRepoForGateway
{
groups
:
map
[
int64
]
*
Group
{
groupID
:
{
ID
:
groupID
,
Platform
:
PlatformAnthropic
,
Status
:
StatusActive
,
Hydrated
:
true
,
ModelRoutingEnabled
:
true
,
ModelRouting
:
map
[
string
][]
int64
{
"claude-3-5-sonnet-20241022"
:
{
1
,
2
},
},
},
},
}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
true
concurrencyCache
:=
&
mockConcurrencyCache
{}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
groupRepo
:
groupRepo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
),
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
&
groupID
,
sessionHash
,
"claude-3-5-sonnet-20241022"
,
nil
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
Account
)
require
.
Equal
(
t
,
int64
(
1
),
result
.
Account
.
ID
)
require
.
Equal
(
t
,
0
,
concurrencyCache
.
loadBatchCalls
)
})
t
.
Run
(
"模型路由-粘性账号缺失-清理并回退"
,
func
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
22
)
sessionHash
:=
"route-missing"
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
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
]
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
sessionHash
:
1
},
}
groupRepo
:=
&
mockGroupRepoForGateway
{
groups
:
map
[
int64
]
*
Group
{
groupID
:
{
ID
:
groupID
,
Platform
:
PlatformAnthropic
,
Status
:
StatusActive
,
Hydrated
:
true
,
ModelRoutingEnabled
:
true
,
ModelRouting
:
map
[
string
][]
int64
{
"claude-3-5-sonnet-20241022"
:
{
1
,
2
},
},
},
},
}
func
(
m
*
mockConcurrencyCache
)
IncrementWaitCount
(
ctx
context
.
Context
,
userID
int64
,
maxWait
int
)
(
bool
,
error
)
{
return
true
,
nil
}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
true
func
(
m
*
mockConcurrencyCache
)
DecrementWaitCount
(
ctx
context
.
Context
,
userID
int64
)
error
{
return
nil
}
concurrencyCache
:=
&
mockConcurrencyCache
{}
func
(
m
*
mockConcurrencyCache
)
GetAccountsLoadBatch
(
ctx
context
.
Context
,
accounts
[]
AccountWithConcurrency
)
(
map
[
int64
]
*
AccountLoadInfo
,
error
)
{
m
.
loadBatchCalls
++
result
:=
make
(
map
[
int64
]
*
AccountLoadInfo
,
len
(
accounts
))
for
_
,
acc
:=
range
accounts
{
result
[
acc
.
ID
]
=
&
AccountLoadInfo
{
AccountID
:
acc
.
ID
,
CurrentConcurrency
:
0
,
WaitingCount
:
0
,
LoadRate
:
0
,
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
groupRepo
:
groupRepo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
),
}
}
return
result
,
nil
}
func
(
m
*
mockConcurrencyCache
)
CleanupExpiredAccountSlots
(
ctx
context
.
Context
,
accountID
int64
)
error
{
return
nil
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
&
groupID
,
sessionHash
,
"claude-3-5-sonnet-20241022"
,
nil
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
Account
)
require
.
Equal
(
t
,
int64
(
2
),
result
.
Account
.
ID
)
require
.
Equal
(
t
,
1
,
cache
.
deletedSessions
[
sessionHash
])
require
.
Equal
(
t
,
int64
(
2
),
cache
.
sessionBindings
[
sessionHash
])
})
// TestGatewayService_SelectAccountWithLoadAwareness tests load-aware account selection
func
TestGatewayService_SelectAccountWithLoadAwareness
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
t
.
Run
(
"模型路由-按负载选择账号"
,
func
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
21
)
t
.
Run
(
"禁用负载批量查询-降级到传统选择"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
...
...
@@ -1042,31 +2474,54 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
cache
:=
&
mockGatewayCacheForPlatform
{}
groupRepo
:=
&
mockGroupRepoForGateway
{
groups
:
map
[
int64
]
*
Group
{
groupID
:
{
ID
:
groupID
,
Platform
:
PlatformAnthropic
,
Status
:
StatusActive
,
Hydrated
:
true
,
ModelRoutingEnabled
:
true
,
ModelRouting
:
map
[
string
][]
int64
{
"claude-3-5-sonnet-20241022"
:
{
1
,
2
},
},
},
},
}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
false
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
true
concurrencyCache
:=
&
mockConcurrencyCache
{
loadMap
:
map
[
int64
]
*
AccountLoadInfo
{
1
:
{
AccountID
:
1
,
LoadRate
:
80
},
2
:
{
AccountID
:
2
,
LoadRate
:
20
},
},
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
groupRepo
:
groupRepo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
nil
,
// No c
oncurrency
s
ervice
concurrencyService
:
NewC
oncurrency
S
ervice
(
concurrencyCache
),
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
nil
,
"
"
,
"claude-3-5-sonnet-20241022"
,
nil
,
""
)
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
&
groupID
,
"route
"
,
"claude-3-5-sonnet-20241022"
,
nil
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
Account
)
require
.
Equal
(
t
,
int64
(
1
),
result
.
Account
.
ID
,
"应选择优先级最高的账号"
)
require
.
Equal
(
t
,
int64
(
2
),
result
.
Account
.
ID
)
require
.
Equal
(
t
,
int64
(
2
),
cache
.
sessionBindings
[
"route"
])
})
t
.
Run
(
"模型路由-无ConcurrencyService也生效"
,
func
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
1
)
sessionHash
:=
"sticky"
t
.
Run
(
"模型路由-路由账号全满返回等待计划"
,
func
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
23
)
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
,
AccountGroups
:
[]
AccountGroup
{{
GroupID
:
groupID
}}
},
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
,
AccountGroups
:
[]
AccountGroup
{{
GroupID
:
groupID
}}
},
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
...
...
@@ -1074,9 +2529,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
sessionHash
:
1
},
}
cache
:=
&
mockGatewayCacheForPlatform
{}
groupRepo
:=
&
mockGroupRepoForGateway
{
groups
:
map
[
int64
]
*
Group
{
...
...
@@ -1087,8 +2540,7 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
Hydrated
:
true
,
ModelRoutingEnabled
:
true
,
ModelRouting
:
map
[
string
][]
int64
{
"claude-a"
:
{
1
},
"claude-b"
:
{
2
},
"claude-3-5-sonnet-20241022"
:
{
1
,
2
},
},
},
},
...
...
@@ -1097,27 +2549,37 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
true
concurrencyCache
:=
&
mockConcurrencyCache
{
acquireResults
:
map
[
int64
]
bool
{
1
:
false
,
2
:
false
},
loadMap
:
map
[
int64
]
*
AccountLoadInfo
{
1
:
{
AccountID
:
1
,
LoadRate
:
10
},
2
:
{
AccountID
:
2
,
LoadRate
:
20
},
},
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
groupRepo
:
groupRepo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
nil
,
// legacy path
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
),
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
&
groupID
,
sessionHash
,
"claude-b
"
,
nil
,
""
)
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
&
groupID
,
"route-full"
,
"claude-3-5-sonnet-20241022
"
,
nil
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
Account
)
require
.
Equal
(
t
,
int64
(
2
),
result
.
Account
.
ID
,
"切换到 claude-b 时应按模型路由切换账号"
)
require
.
Equal
(
t
,
int64
(
2
),
cache
.
sessionBindings
[
sessionHash
],
"粘性绑定应更新为路由选择的账号"
)
require
.
NotNil
(
t
,
result
.
WaitPlan
)
require
.
Equal
(
t
,
int64
(
1
),
result
.
Account
.
ID
)
})
t
.
Run
(
"无ConcurrencyService-降级到传统选择"
,
func
(
t
*
testing
.
T
)
{
t
.
Run
(
"模型路由-路由账号全满-回退普通选择"
,
func
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
22
)
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
{
ID
:
3
,
Platform
:
PlatformAnthropic
,
Priority
:
0
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
...
...
@@ -1127,24 +2589,49 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
cache
:=
&
mockGatewayCacheForPlatform
{}
groupRepo
:=
&
mockGroupRepoForGateway
{
groups
:
map
[
int64
]
*
Group
{
groupID
:
{
ID
:
groupID
,
Platform
:
PlatformAnthropic
,
Status
:
StatusActive
,
Hydrated
:
true
,
ModelRoutingEnabled
:
true
,
ModelRouting
:
map
[
string
][]
int64
{
"claude-3-5-sonnet-20241022"
:
{
1
,
2
},
},
},
},
}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
true
concurrencyCache
:=
&
mockConcurrencyCache
{
loadMap
:
map
[
int64
]
*
AccountLoadInfo
{
1
:
{
AccountID
:
1
,
LoadRate
:
100
},
2
:
{
AccountID
:
2
,
LoadRate
:
100
},
3
:
{
AccountID
:
3
,
LoadRate
:
0
},
},
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
groupRepo
:
groupRepo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
nil
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
)
,
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
nil
,
"
"
,
"claude-3-5-sonnet-20241022"
,
nil
,
""
)
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
&
groupID
,
"fallback
"
,
"claude-3-5-sonnet-20241022"
,
nil
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
Account
)
require
.
Equal
(
t
,
int64
(
2
),
result
.
Account
.
ID
,
"应选择优先级最高的账号"
)
require
.
Equal
(
t
,
int64
(
3
),
result
.
Account
.
ID
)
require
.
Equal
(
t
,
int64
(
3
),
cache
.
sessionBindings
[
"fallback"
])
})
t
.
Run
(
"
排除账号-不选择被排除的账号
"
,
func
(
t
*
testing
.
T
)
{
t
.
Run
(
"
负载批量失败且无法获取-兜底等待
"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
...
...
@@ -1159,27 +2646,34 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
cache
:=
&
mockGatewayCacheForPlatform
{}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
false
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
true
concurrencyCache
:=
&
mockConcurrencyCache
{
loadBatchErr
:
errors
.
New
(
"load batch failed"
),
acquireResults
:
map
[
int64
]
bool
{
1
:
false
,
2
:
false
},
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
nil
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
)
,
}
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"
,
nil
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
Account
)
require
.
Equal
(
t
,
int64
(
2
),
result
.
Account
.
ID
,
"不应选择被排除的账号"
)
require
.
NotNil
(
t
,
result
.
WaitPlan
)
require
.
Equal
(
t
,
int64
(
1
),
result
.
Account
.
ID
)
})
t
.
Run
(
"粘性命中-不调用GetByID"
,
func
(
t
*
testing
.
T
)
{
t
.
Run
(
"Gemini负载排序-优先OAuth"
,
func
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
24
)
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
{
ID
:
1
,
Platform
:
PlatformGemini
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
,
Type
:
AccountTypeAPIKey
},
{
ID
:
2
,
Platform
:
PlatformGemini
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
,
Type
:
AccountTypeOAuth
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
...
...
@@ -1187,35 +2681,77 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
"sticky"
:
1
},
cache
:=
&
mockGatewayCacheForPlatform
{}
groupRepo
:=
&
mockGroupRepoForGateway
{
groups
:
map
[
int64
]
*
Group
{
groupID
:
{
ID
:
groupID
,
Platform
:
PlatformGemini
,
Status
:
StatusActive
,
Hydrated
:
true
,
},
},
}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
true
concurrencyCache
:=
&
mockConcurrencyCache
{}
concurrencyCache
:=
&
mockConcurrencyCache
{
loadMap
:
map
[
int64
]
*
AccountLoadInfo
{
1
:
{
AccountID
:
1
,
LoadRate
:
10
},
2
:
{
AccountID
:
2
,
LoadRate
:
10
},
},
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
groupRepo
:
groupRepo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
),
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
nil
,
"sticky"
,
"claude-3-5-sonnet-20241022
"
,
nil
,
""
)
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
&
groupID
,
"gemini"
,
"gemini-2.5-pro
"
,
nil
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
Account
)
require
.
Equal
(
t
,
int64
(
1
),
result
.
Account
.
ID
)
require
.
Equal
(
t
,
0
,
repo
.
getByIDCalls
,
"粘性命中不应调用GetByID"
)
require
.
Equal
(
t
,
0
,
concurrencyCache
.
loadBatchCalls
,
"粘性命中应在负载批量查询前返回"
)
require
.
Equal
(
t
,
int64
(
2
),
result
.
Account
.
ID
)
})
t
.
Run
(
"粘性账号不在候选集-回退负载感知选择"
,
func
(
t
*
testing
.
T
)
{
t
.
Run
(
"模型路由-过滤路径覆盖"
,
func
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
70
)
now
:=
time
.
Now
()
.
Add
(
10
*
time
.
Minute
)
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
{
ID
:
3
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
false
,
Concurrency
:
5
},
{
ID
:
4
,
Platform
:
PlatformAntigravity
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
{
ID
:
5
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
,
Extra
:
map
[
string
]
any
{
"model_rate_limits"
:
map
[
string
]
any
{
"claude_sonnet"
:
map
[
string
]
any
{
"rate_limit_reset_at"
:
now
.
Format
(
time
.
RFC3339
),
},
},
},
},
{
ID
:
6
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
,
Credentials
:
map
[
string
]
any
{
"model_mapping"
:
map
[
string
]
any
{
"claude-3-5-haiku-20241022"
:
"claude-3-5-haiku-20241022"
}},
},
{
ID
:
7
,
Platform
:
PlatformAnthropic
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
...
...
@@ -1223,8 +2759,21 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{
sessionBindings
:
map
[
string
]
int64
{
"sticky"
:
1
},
cache
:=
&
mockGatewayCacheForPlatform
{}
groupRepo
:=
&
mockGroupRepoForGateway
{
groups
:
map
[
int64
]
*
Group
{
groupID
:
{
ID
:
groupID
,
Platform
:
PlatformAnthropic
,
Status
:
StatusActive
,
Hydrated
:
true
,
ModelRoutingEnabled
:
true
,
ModelRouting
:
map
[
string
][]
int64
{
"claude-3-5-sonnet-20241022"
:
{
1
,
2
,
3
,
4
,
5
,
6
},
},
},
},
}
cfg
:=
testConfig
()
...
...
@@ -1234,51 +2783,110 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
groupRepo
:
groupRepo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
),
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
nil
,
"sticky"
,
"claude-3-5-sonnet-20241022"
,
nil
,
""
)
excluded
:=
map
[
int64
]
struct
{}{
1
:
{}}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
&
groupID
,
""
,
"claude-3-5-sonnet-20241022"
,
excluded
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
Account
)
require
.
Equal
(
t
,
int64
(
2
),
result
.
Account
.
ID
,
"粘性账号不在候选集时应回退到可用账号"
)
require
.
Equal
(
t
,
0
,
repo
.
getByIDCalls
,
"粘性账号缺失不应回退到GetByID"
)
require
.
Equal
(
t
,
1
,
concurrencyCache
.
loadBatchCalls
,
"应继续进行负载批量查询"
)
require
.
Equal
(
t
,
int64
(
7
),
result
.
Account
.
ID
)
})
t
.
Run
(
"无可用账号-返回错误"
,
func
(
t
*
testing
.
T
)
{
t
.
Run
(
"ClaudeCode限制-回退分组"
,
func
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
60
)
fallbackID
:=
int64
(
61
)
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{},
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformGemini
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForPlatform
{}
groupRepo
:=
&
mockGroupRepoForGateway
{
groups
:
map
[
int64
]
*
Group
{
groupID
:
{
ID
:
groupID
,
Platform
:
PlatformAnthropic
,
Status
:
StatusActive
,
Hydrated
:
true
,
ClaudeCodeOnly
:
true
,
FallbackGroupID
:
func
()
*
int64
{
v
:=
fallbackID
return
&
v
}(),
},
fallbackID
:
{
ID
:
fallbackID
,
Platform
:
PlatformGemini
,
Status
:
StatusActive
,
Hydrated
:
true
,
},
},
}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
false
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
groupRepo
:
groupRepo
,
cache
:
&
mockGatewayCacheForPlatform
{},
cfg
:
cfg
,
concurrencyService
:
nil
,
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
nil
,
""
,
"claude-3-5-sonnet-20241022"
,
nil
,
""
)
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
&
groupID
,
""
,
"gemini-2.5-pro"
,
nil
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
Account
)
require
.
Equal
(
t
,
int64
(
1
),
result
.
Account
.
ID
)
})
t
.
Run
(
"ClaudeCode限制-无降级返回错误"
,
func
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
62
)
groupRepo
:=
&
mockGroupRepoForGateway
{
groups
:
map
[
int64
]
*
Group
{
groupID
:
{
ID
:
groupID
,
Platform
:
PlatformAnthropic
,
Status
:
StatusActive
,
Hydrated
:
true
,
ClaudeCodeOnly
:
true
,
},
},
}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
false
svc
:=
&
GatewayService
{
accountRepo
:
&
mockAccountRepoForPlatform
{},
groupRepo
:
groupRepo
,
cache
:
&
mockGatewayCacheForPlatform
{},
cfg
:
cfg
,
concurrencyService
:
nil
,
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
&
groupID
,
""
,
"claude-3-5-sonnet-20241022"
,
nil
)
require
.
Error
(
t
,
err
)
require
.
Nil
(
t
,
result
)
require
.
Contain
s
(
t
,
err
.
Err
or
(),
"no available accounts"
)
require
.
ErrorI
s
(
t
,
err
,
Err
ClaudeCodeOnly
)
})
t
.
Run
(
"过滤不可调度账号-限流账号被跳过"
,
func
(
t
*
testing
.
T
)
{
now
:=
time
.
Now
()
resetAt
:=
now
.
Add
(
10
*
time
.
Minute
)
t
.
Run
(
"负载可用但无法获取槽位-兜底等待"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
,
RateLimitResetAt
:
&
resetAt
},
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
...
...
@@ -1288,31 +2896,37 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
}
cache
:=
&
mockGatewayCacheForPlatform
{}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
false
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
true
concurrencyCache
:=
&
mockConcurrencyCache
{
acquireResults
:
map
[
int64
]
bool
{
1
:
false
,
2
:
false
},
loadMap
:
map
[
int64
]
*
AccountLoadInfo
{
1
:
{
AccountID
:
1
,
LoadRate
:
10
},
2
:
{
AccountID
:
2
,
LoadRate
:
20
},
},
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
nil
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
)
,
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
nil
,
""
,
"claude-3-5-sonnet-20241022"
,
nil
,
""
)
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
nil
,
"
wait
"
,
"claude-3-5-sonnet-20241022"
,
nil
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
Account
)
require
.
Equal
(
t
,
int64
(
2
),
result
.
Account
.
ID
,
"应跳过限流账号,选择可用账号"
)
require
.
NotNil
(
t
,
result
.
WaitPlan
)
require
.
Equal
(
t
,
int64
(
1
),
result
.
Account
.
ID
)
})
t
.
Run
(
"过滤不可调度账号-过载账号被跳过"
,
func
(
t
*
testing
.
T
)
{
now
:=
time
.
Now
()
overloadUntil
:=
now
.
Add
(
10
*
time
.
Minute
)
t
.
Run
(
"负载信息缺失-使用默认负载"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForPlatform
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
,
OverloadUntil
:
&
overloadUntil
},
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
{
ID
:
2
,
Platform
:
PlatformAnthropic
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
5
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
...
...
@@ -1321,21 +2935,29 @@ func TestGatewayService_SelectAccountWithLoadAwareness(t *testing.T) {
}
cache
:=
&
mockGatewayCacheForPlatform
{}
cfg
:=
testConfig
()
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
false
cfg
.
Gateway
.
Scheduling
.
LoadBatchEnabled
=
true
concurrencyCache
:=
&
mockConcurrencyCache
{
loadMap
:
map
[
int64
]
*
AccountLoadInfo
{
1
:
{
AccountID
:
1
,
LoadRate
:
50
},
},
skipDefaultLoad
:
true
,
}
svc
:=
&
GatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
cfg
:
cfg
,
concurrencyService
:
nil
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
)
,
}
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
nil
,
""
,
"claude-3-5-sonnet-20241022"
,
nil
,
""
)
result
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
ctx
,
nil
,
"
missing-load
"
,
"claude-3-5-sonnet-20241022"
,
nil
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
result
)
require
.
NotNil
(
t
,
result
.
Account
)
require
.
Equal
(
t
,
int64
(
2
),
result
.
Account
.
ID
,
"应跳过过载账号,选择可用账号"
)
require
.
Equal
(
t
,
int64
(
2
),
result
.
Account
.
ID
)
})
}
...
...
backend/internal/service/gateway_service.go
View file @
f6ed3d14
...
...
@@ -97,11 +97,24 @@ var allowedHeaders = map[string]bool{
"content-type"
:
true
,
}
// GatewayCache defines cache operations for gateway service
// GatewayCache 定义网关服务的缓存操作接口。
// 提供粘性会话(Sticky Session)的存储、查询、刷新和删除功能。
//
// GatewayCache defines cache operations for gateway service.
// Provides sticky session storage, retrieval, refresh and deletion capabilities.
type
GatewayCache
interface
{
// GetSessionAccountID 获取粘性会话绑定的账号 ID
// Get the account ID bound to a sticky session
GetSessionAccountID
(
ctx
context
.
Context
,
groupID
int64
,
sessionHash
string
)
(
int64
,
error
)
// SetSessionAccountID 设置粘性会话与账号的绑定关系
// Set the binding between sticky session and account
SetSessionAccountID
(
ctx
context
.
Context
,
groupID
int64
,
sessionHash
string
,
accountID
int64
,
ttl
time
.
Duration
)
error
// RefreshSessionTTL 刷新粘性会话的过期时间
// Refresh the expiration time of a sticky session
RefreshSessionTTL
(
ctx
context
.
Context
,
groupID
int64
,
sessionHash
string
,
ttl
time
.
Duration
)
error
// DeleteSessionAccountID 删除粘性会话绑定,用于账号不可用时主动清理
// Delete sticky session binding, used to proactively clean up when account becomes unavailable
DeleteSessionAccountID
(
ctx
context
.
Context
,
groupID
int64
,
sessionHash
string
)
error
}
// derefGroupID safely dereferences *int64 to int64, returning 0 if nil
...
...
@@ -112,6 +125,28 @@ func derefGroupID(groupID *int64) int64 {
return
*
groupID
}
// shouldClearStickySession 检查账号是否处于不可调度状态,需要清理粘性会话绑定。
// 当账号状态为错误、禁用、不可调度,或处于临时不可调度期间时,返回 true。
// 这确保后续请求不会继续使用不可用的账号。
//
// shouldClearStickySession checks if an account is in an unschedulable state
// and the sticky session binding should be cleared.
// Returns true when account status is error/disabled, schedulable is false,
// or within temporary unschedulable period.
// This ensures subsequent requests won't continue using unavailable accounts.
func
shouldClearStickySession
(
account
*
Account
)
bool
{
if
account
==
nil
{
return
false
}
if
account
.
Status
==
StatusError
||
account
.
Status
==
StatusDisabled
||
!
account
.
Schedulable
{
return
true
}
if
account
.
TempUnschedulableUntil
!=
nil
&&
time
.
Now
()
.
Before
(
*
account
.
TempUnschedulableUntil
)
{
return
true
}
return
false
}
type
AccountWaitPlan
struct
{
AccountID
int64
MaxConcurrency
int
...
...
@@ -620,6 +655,8 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
}
// 粘性账号槽位满且等待队列已满,继续使用负载感知选择
}
}
else
{
_
=
s
.
cache
.
DeleteSessionAccountID
(
ctx
,
derefGroupID
(
groupID
),
sessionHash
)
}
}
}
...
...
@@ -720,38 +757,47 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
accountID
,
err
:=
s
.
cache
.
GetSessionAccountID
(
ctx
,
derefGroupID
(
groupID
),
sessionHash
)
if
err
==
nil
&&
accountID
>
0
&&
!
isExcluded
(
accountID
)
{
account
,
ok
:=
accountByID
[
accountID
]
if
ok
&&
s
.
isAccountInGroup
(
account
,
groupID
)
&&
s
.
isAccountAllowedForPlatform
(
account
,
platform
,
useMixed
)
&&
account
.
IsSchedulableForModel
(
requestedModel
)
&&
(
requestedModel
==
""
||
s
.
isModelSupportedByAccount
(
account
,
requestedModel
))
&&
s
.
isAccountSchedulableForWindowCost
(
ctx
,
account
,
true
)
{
// 粘性会话窗口费用检查
result
,
err
:=
s
.
tryAcquireAccountSlot
(
ctx
,
accountID
,
account
.
Concurrency
)
if
err
==
nil
&&
result
.
Acquired
{
// 会话数量限制检查
if
!
s
.
checkAndRegisterSession
(
ctx
,
account
,
sessionUUID
)
{
result
.
ReleaseFunc
()
// 释放槽位,继续到 Layer 2
}
else
{
_
=
s
.
cache
.
RefreshSessionTTL
(
ctx
,
derefGroupID
(
groupID
),
sessionHash
,
stickySessionTTL
)
if
ok
{
// 检查账户是否需要清理粘性会话绑定
// Check if the account needs sticky session cleanup
clearSticky
:=
shouldClearStickySession
(
account
)
if
clearSticky
{
_
=
s
.
cache
.
DeleteSessionAccountID
(
ctx
,
derefGroupID
(
groupID
),
sessionHash
)
}
if
!
clearSticky
&&
s
.
isAccountInGroup
(
account
,
groupID
)
&&
s
.
isAccountAllowedForPlatform
(
account
,
platform
,
useMixed
)
&&
account
.
IsSchedulableForModel
(
requestedModel
)
&&
(
requestedModel
==
""
||
s
.
isModelSupportedByAccount
(
account
,
requestedModel
))
&&
s
.
isAccountSchedulableForWindowCost
(
ctx
,
account
,
true
)
{
// 粘性会话窗口费用检查
result
,
err
:=
s
.
tryAcquireAccountSlot
(
ctx
,
accountID
,
account
.
Concurrency
)
if
err
==
nil
&&
result
.
Acquired
{
// 会话数量限制检查
// Session count limit check
if
!
s
.
checkAndRegisterSession
(
ctx
,
account
,
sessionUUID
)
{
result
.
ReleaseFunc
()
// 释放槽位,继续到 Layer 2
}
else
{
_
=
s
.
cache
.
RefreshSessionTTL
(
ctx
,
derefGroupID
(
groupID
),
sessionHash
,
stickySessionTTL
)
return
&
AccountSelectionResult
{
Account
:
account
,
Acquired
:
true
,
ReleaseFunc
:
result
.
ReleaseFunc
,
},
nil
}
}
waitingCount
,
_
:=
s
.
concurrencyService
.
GetAccountWaitingCount
(
ctx
,
accountID
)
if
waitingCount
<
cfg
.
StickySessionMaxWaiting
{
return
&
AccountSelectionResult
{
Account
:
account
,
Acquired
:
true
,
ReleaseFunc
:
result
.
ReleaseFunc
,
Account
:
account
,
WaitPlan
:
&
AccountWaitPlan
{
AccountID
:
accountID
,
MaxConcurrency
:
account
.
Concurrency
,
Timeout
:
cfg
.
StickySessionWaitTimeout
,
MaxWaiting
:
cfg
.
StickySessionMaxWaiting
,
},
},
nil
}
}
waitingCount
,
_
:=
s
.
concurrencyService
.
GetAccountWaitingCount
(
ctx
,
accountID
)
if
waitingCount
<
cfg
.
StickySessionMaxWaiting
{
return
&
AccountSelectionResult
{
Account
:
account
,
WaitPlan
:
&
AccountWaitPlan
{
AccountID
:
accountID
,
MaxConcurrency
:
account
.
Concurrency
,
Timeout
:
cfg
.
StickySessionWaitTimeout
,
MaxWaiting
:
cfg
.
StickySessionMaxWaiting
,
},
},
nil
}
}
}
}
...
...
@@ -1278,14 +1324,20 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
if
_
,
excluded
:=
excludedIDs
[
accountID
];
!
excluded
{
account
,
err
:=
s
.
getSchedulableAccount
(
ctx
,
accountID
)
// 检查账号分组归属和平台匹配(确保粘性会话不会跨分组或跨平台)
if
err
==
nil
&&
s
.
isAccountInGroup
(
account
,
groupID
)
&&
account
.
Platform
==
platform
&&
account
.
IsSchedulableForModel
(
requestedModel
)
&&
(
requestedModel
==
""
||
s
.
isModelSupportedByAccount
(
account
,
requestedModel
))
{
if
err
:=
s
.
cache
.
RefreshSessionTTL
(
ctx
,
derefGroupID
(
groupID
),
sessionHash
,
stickySessionTTL
);
err
!=
nil
{
log
.
Printf
(
"refresh session ttl failed: session=%s err=%v"
,
sessionHash
,
err
)
if
err
==
nil
{
clearSticky
:=
shouldClearStickySession
(
account
)
if
clearSticky
{
_
=
s
.
cache
.
DeleteSessionAccountID
(
ctx
,
derefGroupID
(
groupID
),
sessionHash
)
}
if
s
.
debugModelRoutingEnabled
()
{
log
.
Printf
(
"[ModelRoutingDebug] legacy routed sticky hit: group_id=%v model=%s session=%s account=%d"
,
derefGroupID
(
groupID
),
requestedModel
,
shortSessionHash
(
sessionHash
),
accountID
)
if
!
clearSticky
&&
s
.
isAccountInGroup
(
account
,
groupID
)
&&
account
.
Platform
==
platform
&&
account
.
IsSchedulableForModel
(
requestedModel
)
&&
(
requestedModel
==
""
||
s
.
isModelSupportedByAccount
(
account
,
requestedModel
))
{
if
err
:=
s
.
cache
.
RefreshSessionTTL
(
ctx
,
derefGroupID
(
groupID
),
sessionHash
,
stickySessionTTL
);
err
!=
nil
{
log
.
Printf
(
"refresh session ttl failed: session=%s err=%v"
,
sessionHash
,
err
)
}
if
s
.
debugModelRoutingEnabled
()
{
log
.
Printf
(
"[ModelRoutingDebug] legacy routed sticky hit: group_id=%v model=%s session=%s account=%d"
,
derefGroupID
(
groupID
),
requestedModel
,
shortSessionHash
(
sessionHash
),
accountID
)
}
return
account
,
nil
}
return
account
,
nil
}
}
}
...
...
@@ -1375,11 +1427,17 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
if
_
,
excluded
:=
excludedIDs
[
accountID
];
!
excluded
{
account
,
err
:=
s
.
getSchedulableAccount
(
ctx
,
accountID
)
// 检查账号分组归属和平台匹配(确保粘性会话不会跨分组或跨平台)
if
err
==
nil
&&
s
.
isAccountInGroup
(
account
,
groupID
)
&&
account
.
Platform
==
platform
&&
account
.
IsSchedulableForModel
(
requestedModel
)
&&
(
requestedModel
==
""
||
s
.
isModelSupportedByAccount
(
account
,
requestedModel
))
{
if
err
:=
s
.
cache
.
RefreshSessionTTL
(
ctx
,
derefGroupID
(
groupID
),
sessionHash
,
stickySessionTTL
);
err
!=
nil
{
log
.
Printf
(
"refresh session ttl failed: session=%s err=%v"
,
sessionHash
,
err
)
if
err
==
nil
{
clearSticky
:=
shouldClearStickySession
(
account
)
if
clearSticky
{
_
=
s
.
cache
.
DeleteSessionAccountID
(
ctx
,
derefGroupID
(
groupID
),
sessionHash
)
}
if
!
clearSticky
&&
s
.
isAccountInGroup
(
account
,
groupID
)
&&
account
.
Platform
==
platform
&&
account
.
IsSchedulableForModel
(
requestedModel
)
&&
(
requestedModel
==
""
||
s
.
isModelSupportedByAccount
(
account
,
requestedModel
))
{
if
err
:=
s
.
cache
.
RefreshSessionTTL
(
ctx
,
derefGroupID
(
groupID
),
sessionHash
,
stickySessionTTL
);
err
!=
nil
{
log
.
Printf
(
"refresh session ttl failed: session=%s err=%v"
,
sessionHash
,
err
)
}
return
account
,
nil
}
return
account
,
nil
}
}
}
...
...
@@ -1479,15 +1537,21 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
if
_
,
excluded
:=
excludedIDs
[
accountID
];
!
excluded
{
account
,
err
:=
s
.
getSchedulableAccount
(
ctx
,
accountID
)
// 检查账号分组归属和有效性:原生平台直接匹配,antigravity 需要启用混合调度
if
err
==
nil
&&
s
.
isAccountInGroup
(
account
,
groupID
)
&&
account
.
IsSchedulableForModel
(
requestedModel
)
&&
(
requestedModel
==
""
||
s
.
isModelSupportedByAccount
(
account
,
requestedModel
))
{
if
account
.
Platform
==
nativePlatform
||
(
account
.
Platform
==
PlatformAntigravity
&&
account
.
IsMixedSchedulingEnabled
())
{
if
err
:=
s
.
cache
.
RefreshSessionTTL
(
ctx
,
derefGroupID
(
groupID
),
sessionHash
,
stickySessionTTL
);
err
!=
nil
{
log
.
Printf
(
"refresh session ttl failed: session=%s err=%v"
,
sessionHash
,
err
)
}
if
s
.
debugModelRoutingEnabled
()
{
log
.
Printf
(
"[ModelRoutingDebug] legacy mixed routed sticky hit: group_id=%v model=%s session=%s account=%d"
,
derefGroupID
(
groupID
),
requestedModel
,
shortSessionHash
(
sessionHash
),
accountID
)
if
err
==
nil
{
clearSticky
:=
shouldClearStickySession
(
account
)
if
clearSticky
{
_
=
s
.
cache
.
DeleteSessionAccountID
(
ctx
,
derefGroupID
(
groupID
),
sessionHash
)
}
if
!
clearSticky
&&
s
.
isAccountInGroup
(
account
,
groupID
)
&&
account
.
IsSchedulableForModel
(
requestedModel
)
&&
(
requestedModel
==
""
||
s
.
isModelSupportedByAccount
(
account
,
requestedModel
))
{
if
account
.
Platform
==
nativePlatform
||
(
account
.
Platform
==
PlatformAntigravity
&&
account
.
IsMixedSchedulingEnabled
())
{
if
err
:=
s
.
cache
.
RefreshSessionTTL
(
ctx
,
derefGroupID
(
groupID
),
sessionHash
,
stickySessionTTL
);
err
!=
nil
{
log
.
Printf
(
"refresh session ttl failed: session=%s err=%v"
,
sessionHash
,
err
)
}
if
s
.
debugModelRoutingEnabled
()
{
log
.
Printf
(
"[ModelRoutingDebug] legacy mixed routed sticky hit: group_id=%v model=%s session=%s account=%d"
,
derefGroupID
(
groupID
),
requestedModel
,
shortSessionHash
(
sessionHash
),
accountID
)
}
return
account
,
nil
}
return
account
,
nil
}
}
}
...
...
@@ -1578,12 +1642,18 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
if
_
,
excluded
:=
excludedIDs
[
accountID
];
!
excluded
{
account
,
err
:=
s
.
getSchedulableAccount
(
ctx
,
accountID
)
// 检查账号分组归属和有效性:原生平台直接匹配,antigravity 需要启用混合调度
if
err
==
nil
&&
s
.
isAccountInGroup
(
account
,
groupID
)
&&
account
.
IsSchedulableForModel
(
requestedModel
)
&&
(
requestedModel
==
""
||
s
.
isModelSupportedByAccount
(
account
,
requestedModel
))
{
if
account
.
Platform
==
nativePlatform
||
(
account
.
Platform
==
PlatformAntigravity
&&
account
.
IsMixedSchedulingEnabled
())
{
if
err
:=
s
.
cache
.
RefreshSessionTTL
(
ctx
,
derefGroupID
(
groupID
),
sessionHash
,
stickySessionTTL
);
err
!=
nil
{
log
.
Printf
(
"refresh session ttl failed: session=%s err=%v"
,
sessionHash
,
err
)
if
err
==
nil
{
clearSticky
:=
shouldClearStickySession
(
account
)
if
clearSticky
{
_
=
s
.
cache
.
DeleteSessionAccountID
(
ctx
,
derefGroupID
(
groupID
),
sessionHash
)
}
if
!
clearSticky
&&
s
.
isAccountInGroup
(
account
,
groupID
)
&&
account
.
IsSchedulableForModel
(
requestedModel
)
&&
(
requestedModel
==
""
||
s
.
isModelSupportedByAccount
(
account
,
requestedModel
))
{
if
account
.
Platform
==
nativePlatform
||
(
account
.
Platform
==
PlatformAntigravity
&&
account
.
IsMixedSchedulingEnabled
())
{
if
err
:=
s
.
cache
.
RefreshSessionTTL
(
ctx
,
derefGroupID
(
groupID
),
sessionHash
,
stickySessionTTL
);
err
!=
nil
{
log
.
Printf
(
"refresh session ttl failed: session=%s err=%v"
,
sessionHash
,
err
)
}
return
account
,
nil
}
return
account
,
nil
}
}
}
...
...
backend/internal/service/gemini_messages_compat_service.go
View file @
f6ed3d14
...
...
@@ -82,145 +82,276 @@ func (s *GeminiMessagesCompatService) SelectAccountForModel(ctx context.Context,
}
func
(
s
*
GeminiMessagesCompatService
)
SelectAccountForModelWithExclusions
(
ctx
context
.
Context
,
groupID
*
int64
,
sessionHash
string
,
requestedModel
string
,
excludedIDs
map
[
int64
]
struct
{})
(
*
Account
,
error
)
{
// 1. 确定目标平台和调度模式
// Determine target platform and scheduling mode
platform
,
useMixedScheduling
,
hasForcePlatform
,
err
:=
s
.
resolvePlatformAndSchedulingMode
(
ctx
,
groupID
)
if
err
!=
nil
{
return
nil
,
err
}
cacheKey
:=
"gemini:"
+
sessionHash
// 2. 尝试粘性会话命中
// Try sticky session hit
if
account
:=
s
.
tryStickySessionHit
(
ctx
,
groupID
,
sessionHash
,
cacheKey
,
requestedModel
,
excludedIDs
,
platform
,
useMixedScheduling
);
account
!=
nil
{
return
account
,
nil
}
// 3. 查询可调度账户(强制平台模式:优先按分组查找,找不到再查全部)
// Query schedulable accounts (force platform mode: try group first, fallback to all)
accounts
,
err
:=
s
.
listSchedulableAccountsOnce
(
ctx
,
groupID
,
platform
,
hasForcePlatform
)
if
err
!=
nil
{
return
nil
,
fmt
.
Errorf
(
"query accounts failed: %w"
,
err
)
}
// 强制平台模式下,分组中找不到账户时回退查询全部
if
len
(
accounts
)
==
0
&&
groupID
!=
nil
&&
hasForcePlatform
{
accounts
,
err
=
s
.
listSchedulableAccountsOnce
(
ctx
,
nil
,
platform
,
hasForcePlatform
)
if
err
!=
nil
{
return
nil
,
fmt
.
Errorf
(
"query accounts failed: %w"
,
err
)
}
}
// 4. 按优先级 + LRU 选择最佳账号
// Select best account by priority + LRU
selected
:=
s
.
selectBestGeminiAccount
(
ctx
,
accounts
,
requestedModel
,
excludedIDs
,
platform
,
useMixedScheduling
)
if
selected
==
nil
{
if
requestedModel
!=
""
{
return
nil
,
fmt
.
Errorf
(
"no available Gemini accounts supporting model: %s"
,
requestedModel
)
}
return
nil
,
errors
.
New
(
"no available Gemini accounts"
)
}
// 5. 设置粘性会话绑定
// Set sticky session binding
if
sessionHash
!=
""
{
_
=
s
.
cache
.
SetSessionAccountID
(
ctx
,
derefGroupID
(
groupID
),
cacheKey
,
selected
.
ID
,
geminiStickySessionTTL
)
}
return
selected
,
nil
}
// resolvePlatformAndSchedulingMode 解析目标平台和调度模式。
// 返回:平台名称、是否使用混合调度、是否强制平台、错误。
//
// resolvePlatformAndSchedulingMode resolves target platform and scheduling mode.
// Returns: platform name, whether to use mixed scheduling, whether force platform, error.
func
(
s
*
GeminiMessagesCompatService
)
resolvePlatformAndSchedulingMode
(
ctx
context
.
Context
,
groupID
*
int64
)
(
platform
string
,
useMixedScheduling
bool
,
hasForcePlatform
bool
,
err
error
)
{
// 优先检查 context 中的强制平台(/antigravity 路由)
var
platform
string
forcePlatform
,
hasForcePlatform
:=
ctx
.
Value
(
ctxkey
.
ForcePlatform
)
.
(
string
)
if
hasForcePlatform
&&
forcePlatform
!=
""
{
platform
=
forcePlatform
}
else
if
groupID
!=
nil
{
return
forcePlatform
,
false
,
true
,
nil
}
if
groupID
!=
nil
{
// 根据分组 platform 决定查询哪种账号
var
group
*
Group
if
ctxGroup
,
ok
:=
ctx
.
Value
(
ctxkey
.
Group
)
.
(
*
Group
);
ok
&&
IsGroupContextValid
(
ctxGroup
)
&&
ctxGroup
.
ID
==
*
groupID
{
group
=
ctxGroup
}
else
{
var
err
error
group
,
err
=
s
.
groupRepo
.
GetByIDLite
(
ctx
,
*
groupID
)
if
err
!=
nil
{
return
nil
,
fmt
.
Errorf
(
"get group failed: %w"
,
err
)
return
""
,
false
,
false
,
fmt
.
Errorf
(
"get group failed: %w"
,
err
)
}
}
platform
=
group
.
Platform
}
else
{
// 无分组时只使用原生 gemini 平台
platform
=
PlatformGemini
// gemini 分组支持混合调度(包含启用了 mixed_scheduling 的 antigravity 账户)
return
group
.
Platform
,
group
.
Platform
==
PlatformGemini
,
false
,
nil
}
//
gemini 分组支持混合调度(包含启用了 mixed_scheduling 的 antigravity 账户)
// 注意:强制平台模式不走混合调度
useMixedScheduling
:=
platform
==
PlatformGemini
&&
!
hasForcePlatform
//
无分组时只使用原生 gemini 平台
return
PlatformGemini
,
true
,
false
,
nil
}
cacheKey
:=
"gemini:"
+
sessionHash
// tryStickySessionHit 尝试从粘性会话获取账号。
// 如果命中且账号可用则返回账号;如果账号不可用则清理会话并返回 nil。
//
// tryStickySessionHit attempts to get account from sticky session.
// Returns account if hit and usable; clears session and returns nil if account unavailable.
func
(
s
*
GeminiMessagesCompatService
)
tryStickySessionHit
(
ctx
context
.
Context
,
groupID
*
int64
,
sessionHash
,
cacheKey
,
requestedModel
string
,
excludedIDs
map
[
int64
]
struct
{},
platform
string
,
useMixedScheduling
bool
,
)
*
Account
{
if
sessionHash
==
""
{
return
nil
}
if
sessionHash
!=
""
{
accountID
,
err
:=
s
.
cache
.
GetSessionAccountID
(
ctx
,
derefGroupID
(
groupID
),
cacheKey
)
if
err
==
nil
&&
accountID
>
0
{
if
_
,
excluded
:=
excludedIDs
[
accountID
];
!
excluded
{
account
,
err
:=
s
.
getSchedulableAccount
(
ctx
,
accountID
)
// 检查账号是否有效:原生平台直接匹配,antigravity 需要启用混合调度
if
err
==
nil
&&
account
.
IsSchedulableForModel
(
requestedModel
)
&&
(
requestedModel
==
""
||
s
.
isModelSupportedByAccount
(
account
,
requestedModel
))
{
valid
:=
false
if
account
.
Platform
==
platform
{
valid
=
true
}
else
if
useMixedScheduling
&&
account
.
Platform
==
PlatformAntigravity
&&
account
.
IsMixedSchedulingEnabled
()
{
valid
=
true
}
if
valid
{
usable
:=
true
if
s
.
rateLimitService
!=
nil
&&
requestedModel
!=
""
{
ok
,
err
:=
s
.
rateLimitService
.
PreCheckUsage
(
ctx
,
account
,
requestedModel
)
if
err
!=
nil
{
log
.
Printf
(
"[Gemini PreCheck] Account %d precheck error: %v"
,
account
.
ID
,
err
)
}
if
!
ok
{
usable
=
false
}
}
if
usable
{
_
=
s
.
cache
.
RefreshSessionTTL
(
ctx
,
derefGroupID
(
groupID
),
cacheKey
,
geminiStickySessionTTL
)
return
account
,
nil
}
}
}
}
}
accountID
,
err
:=
s
.
cache
.
GetSessionAccountID
(
ctx
,
derefGroupID
(
groupID
),
cacheKey
)
if
err
!=
nil
||
accountID
<=
0
{
return
nil
}
// 查询可调度账户(强制平台模式:优先按分组查找,找不到再查全部)
accounts
,
err
:=
s
.
listSchedulableAccountsOnce
(
ctx
,
groupID
,
platform
,
hasForcePlatform
)
if
_
,
excluded
:=
excludedIDs
[
accountID
];
excluded
{
return
nil
}
account
,
err
:=
s
.
getSchedulableAccount
(
ctx
,
accountID
)
if
err
!=
nil
{
return
nil
,
fmt
.
Errorf
(
"query accounts failed: %w"
,
err
)
return
nil
}
// 强制平台模式下,分组中找不到账户时回退查询全部
if
len
(
accounts
)
==
0
&&
groupID
!=
nil
&&
hasForcePlatform
{
accounts
,
err
=
s
.
listSchedulableAccountsOnce
(
ctx
,
nil
,
platform
,
hasForcePlatform
)
if
err
!=
nil
{
return
nil
,
fmt
.
Errorf
(
"query accounts failed: %w"
,
err
)
}
// 检查账号是否需要清理粘性会话
// Check if sticky session should be cleared
if
shouldClearStickySession
(
account
)
{
_
=
s
.
cache
.
DeleteSessionAccountID
(
ctx
,
derefGroupID
(
groupID
),
cacheKey
)
return
nil
}
// 验证账号是否可用于当前请求
// Verify account is usable for current request
if
!
s
.
isAccountUsableForRequest
(
ctx
,
account
,
requestedModel
,
platform
,
useMixedScheduling
)
{
return
nil
}
// 刷新会话 TTL 并返回账号
// Refresh session TTL and return account
_
=
s
.
cache
.
RefreshSessionTTL
(
ctx
,
derefGroupID
(
groupID
),
cacheKey
,
geminiStickySessionTTL
)
return
account
}
// isAccountUsableForRequest 检查账号是否可用于当前请求。
// 验证:模型调度、模型支持、平台匹配、速率限制预检。
//
// isAccountUsableForRequest checks if account is usable for current request.
// Validates: model scheduling, model support, platform matching, rate limit precheck.
func
(
s
*
GeminiMessagesCompatService
)
isAccountUsableForRequest
(
ctx
context
.
Context
,
account
*
Account
,
requestedModel
,
platform
string
,
useMixedScheduling
bool
,
)
bool
{
// 检查模型调度能力
// Check model scheduling capability
if
!
account
.
IsSchedulableForModel
(
requestedModel
)
{
return
false
}
// 检查模型支持
// Check model support
if
requestedModel
!=
""
&&
!
s
.
isModelSupportedByAccount
(
account
,
requestedModel
)
{
return
false
}
// 检查平台匹配
// Check platform matching
if
!
s
.
isAccountValidForPlatform
(
account
,
platform
,
useMixedScheduling
)
{
return
false
}
// 速率限制预检
// Rate limit precheck
if
!
s
.
passesRateLimitPreCheck
(
ctx
,
account
,
requestedModel
)
{
return
false
}
return
true
}
// isAccountValidForPlatform 检查账号是否匹配目标平台。
// 原生平台直接匹配;混合调度模式下 antigravity 需要启用 mixed_scheduling。
//
// isAccountValidForPlatform checks if account matches target platform.
// Native platform matches directly; mixed scheduling mode requires antigravity to enable mixed_scheduling.
func
(
s
*
GeminiMessagesCompatService
)
isAccountValidForPlatform
(
account
*
Account
,
platform
string
,
useMixedScheduling
bool
)
bool
{
if
account
.
Platform
==
platform
{
return
true
}
if
useMixedScheduling
&&
account
.
Platform
==
PlatformAntigravity
&&
account
.
IsMixedSchedulingEnabled
()
{
return
true
}
return
false
}
// passesRateLimitPreCheck 执行速率限制预检。
// 返回 true 表示通过预检或无需预检。
//
// passesRateLimitPreCheck performs rate limit precheck.
// Returns true if passed or precheck not required.
func
(
s
*
GeminiMessagesCompatService
)
passesRateLimitPreCheck
(
ctx
context
.
Context
,
account
*
Account
,
requestedModel
string
)
bool
{
if
s
.
rateLimitService
==
nil
||
requestedModel
==
""
{
return
true
}
ok
,
err
:=
s
.
rateLimitService
.
PreCheckUsage
(
ctx
,
account
,
requestedModel
)
if
err
!=
nil
{
log
.
Printf
(
"[Gemini PreCheck] Account %d precheck error: %v"
,
account
.
ID
,
err
)
}
return
ok
}
// selectBestGeminiAccount 从候选账号中选择最佳账号(优先级 + LRU + OAuth 优先)。
// 返回 nil 表示无可用账号。
//
// selectBestGeminiAccount selects best account from candidates (priority + LRU + OAuth preferred).
// Returns nil if no available account.
func
(
s
*
GeminiMessagesCompatService
)
selectBestGeminiAccount
(
ctx
context
.
Context
,
accounts
[]
Account
,
requestedModel
string
,
excludedIDs
map
[
int64
]
struct
{},
platform
string
,
useMixedScheduling
bool
,
)
*
Account
{
var
selected
*
Account
for
i
:=
range
accounts
{
acc
:=
&
accounts
[
i
]
// 跳过被排除的账号
if
_
,
excluded
:=
excludedIDs
[
acc
.
ID
];
excluded
{
continue
}
// 混合调度模式下:原生平台直接通过,antigravity 需要启用 mixed_scheduling
// 非混合调度模式(antigravity 分组):不需要过滤
if
useMixedScheduling
&&
acc
.
Platform
==
PlatformAntigravity
&&
!
acc
.
IsMixedSchedulingEnabled
()
{
continue
}
if
!
acc
.
IsSchedulableForModel
(
requestedModel
)
{
continue
}
if
requestedModel
!=
""
&&
!
s
.
isModelSupportedByAccount
(
acc
,
requestedModel
)
{
// 检查账号是否可用于当前请求
if
!
s
.
isAccountUsableForRequest
(
ctx
,
acc
,
requestedModel
,
platform
,
useMixedScheduling
)
{
continue
}
if
s
.
rateLimitService
!=
nil
&&
requestedModel
!=
""
{
ok
,
err
:=
s
.
rateLimitService
.
PreCheckUsage
(
ctx
,
acc
,
requestedModel
)
if
err
!=
nil
{
log
.
Printf
(
"[Gemini PreCheck] Account %d precheck error: %v"
,
acc
.
ID
,
err
)
}
if
!
ok
{
continue
}
}
// 选择最佳账号
if
selected
==
nil
{
selected
=
acc
continue
}
if
acc
.
Priority
<
selected
.
Priority
{
if
s
.
isBetterGeminiAccount
(
acc
,
selected
)
{
selected
=
acc
}
else
if
acc
.
Priority
==
selected
.
Priority
{
switch
{
case
acc
.
LastUsedAt
==
nil
&&
selected
.
LastUsedAt
!=
nil
:
selected
=
acc
case
acc
.
LastUsedAt
!=
nil
&&
selected
.
LastUsedAt
==
nil
:
// keep selected (never used is preferred)
case
acc
.
LastUsedAt
==
nil
&&
selected
.
LastUsedAt
==
nil
:
// Prefer OAuth accounts when both are unused (more compatible for Code Assist flows).
if
acc
.
Type
==
AccountTypeOAuth
&&
selected
.
Type
!=
AccountTypeOAuth
{
selected
=
acc
}
default
:
if
acc
.
LastUsedAt
.
Before
(
*
selected
.
LastUsedAt
)
{
selected
=
acc
}
}
}
}
if
selected
==
nil
{
if
requestedModel
!=
""
{
return
nil
,
fmt
.
Errorf
(
"no available Gemini accounts supporting model: %s"
,
requestedModel
)
}
return
nil
,
errors
.
New
(
"no available Gemini accounts"
)
}
return
selected
}
if
sessionHash
!=
""
{
_
=
s
.
cache
.
SetSessionAccountID
(
ctx
,
derefGroupID
(
groupID
),
cacheKey
,
selected
.
ID
,
geminiStickySessionTTL
)
// isBetterGeminiAccount 判断 candidate 是否比 current 更优。
// 规则:优先级更高(数值更小)优先;同优先级时,未使用过的优先(OAuth > 非 OAuth),其次是最久未使用的。
//
// isBetterGeminiAccount checks if candidate is better than current.
// Rules: higher priority (lower value) wins; same priority: never used (OAuth > non-OAuth) > least recently used.
func
(
s
*
GeminiMessagesCompatService
)
isBetterGeminiAccount
(
candidate
,
current
*
Account
)
bool
{
// 优先级更高(数值更小)
if
candidate
.
Priority
<
current
.
Priority
{
return
true
}
if
candidate
.
Priority
>
current
.
Priority
{
return
false
}
return
selected
,
nil
// 同优先级,比较最后使用时间
switch
{
case
candidate
.
LastUsedAt
==
nil
&&
current
.
LastUsedAt
!=
nil
:
// candidate 从未使用,优先
return
true
case
candidate
.
LastUsedAt
!=
nil
&&
current
.
LastUsedAt
==
nil
:
// current 从未使用,保持
return
false
case
candidate
.
LastUsedAt
==
nil
&&
current
.
LastUsedAt
==
nil
:
// 都未使用,优先选择 OAuth 账号(更兼容 Code Assist 流程)
return
candidate
.
Type
==
AccountTypeOAuth
&&
current
.
Type
!=
AccountTypeOAuth
default
:
// 都使用过,选择最久未使用的
return
candidate
.
LastUsedAt
.
Before
(
*
current
.
LastUsedAt
)
}
}
// isModelSupportedByAccount 根据账户平台检查模型支持
...
...
backend/internal/service/gemini_multiplatform_test.go
View file @
f6ed3d14
...
...
@@ -15,8 +15,10 @@ import (
// mockAccountRepoForGemini Gemini 测试用的 mock
type
mockAccountRepoForGemini
struct
{
accounts
[]
Account
accountsByID
map
[
int64
]
*
Account
accounts
[]
Account
accountsByID
map
[
int64
]
*
Account
listByGroupFunc
func
(
ctx
context
.
Context
,
groupID
int64
,
platforms
[]
string
)
([]
Account
,
error
)
listByPlatformFunc
func
(
ctx
context
.
Context
,
platforms
[]
string
)
([]
Account
,
error
)
}
func
(
m
*
mockAccountRepoForGemini
)
GetByID
(
ctx
context
.
Context
,
id
int64
)
(
*
Account
,
error
)
{
...
...
@@ -104,6 +106,9 @@ func (m *mockAccountRepoForGemini) ListSchedulableByGroupID(ctx context.Context,
return
nil
,
nil
}
func
(
m
*
mockAccountRepoForGemini
)
ListSchedulableByPlatforms
(
ctx
context
.
Context
,
platforms
[]
string
)
([]
Account
,
error
)
{
if
m
.
listByPlatformFunc
!=
nil
{
return
m
.
listByPlatformFunc
(
ctx
,
platforms
)
}
var
result
[]
Account
platformSet
:=
make
(
map
[
string
]
bool
)
for
_
,
p
:=
range
platforms
{
...
...
@@ -117,6 +122,9 @@ func (m *mockAccountRepoForGemini) ListSchedulableByPlatforms(ctx context.Contex
return
result
,
nil
}
func
(
m
*
mockAccountRepoForGemini
)
ListSchedulableByGroupIDAndPlatforms
(
ctx
context
.
Context
,
groupID
int64
,
platforms
[]
string
)
([]
Account
,
error
)
{
if
m
.
listByGroupFunc
!=
nil
{
return
m
.
listByGroupFunc
(
ctx
,
groupID
,
platforms
)
}
return
m
.
ListSchedulableByPlatforms
(
ctx
,
platforms
)
}
func
(
m
*
mockAccountRepoForGemini
)
SetRateLimited
(
ctx
context
.
Context
,
id
int64
,
resetAt
time
.
Time
)
error
{
...
...
@@ -212,6 +220,7 @@ var _ GroupRepository = (*mockGroupRepoForGemini)(nil)
// mockGatewayCacheForGemini Gemini 测试用的 cache mock
type
mockGatewayCacheForGemini
struct
{
sessionBindings
map
[
string
]
int64
deletedSessions
map
[
string
]
int
}
func
(
m
*
mockGatewayCacheForGemini
)
GetSessionAccountID
(
ctx
context
.
Context
,
groupID
int64
,
sessionHash
string
)
(
int64
,
error
)
{
...
...
@@ -233,6 +242,18 @@ func (m *mockGatewayCacheForGemini) RefreshSessionTTL(ctx context.Context, group
return
nil
}
func
(
m
*
mockGatewayCacheForGemini
)
DeleteSessionAccountID
(
ctx
context
.
Context
,
groupID
int64
,
sessionHash
string
)
error
{
if
m
.
sessionBindings
==
nil
{
return
nil
}
if
m
.
deletedSessions
==
nil
{
m
.
deletedSessions
=
make
(
map
[
string
]
int
)
}
m
.
deletedSessions
[
sessionHash
]
++
delete
(
m
.
sessionBindings
,
sessionHash
)
return
nil
}
// TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform 测试 Gemini 单平台选择
func
TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
...
...
@@ -523,6 +544,274 @@ func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_StickyS
// 粘性会话未命中,按优先级选择
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
,
"粘性会话未命中,应按优先级选择"
)
})
t
.
Run
(
"粘性会话不可调度-清理并回退选择"
,
func
(
t
*
testing
.
T
)
{
repo
:=
&
mockAccountRepoForGemini
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformGemini
,
Priority
:
2
,
Status
:
StatusDisabled
,
Schedulable
:
true
},
{
ID
:
2
,
Platform
:
PlatformGemini
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForGemini
{
sessionBindings
:
map
[
string
]
int64
{
"gemini:session-123"
:
1
},
}
groupRepo
:=
&
mockGroupRepoForGemini
{
groups
:
map
[
int64
]
*
Group
{}}
svc
:=
&
GeminiMessagesCompatService
{
accountRepo
:
repo
,
groupRepo
:
groupRepo
,
cache
:
cache
,
}
acc
,
err
:=
svc
.
SelectAccountForModelWithExclusions
(
ctx
,
nil
,
"session-123"
,
"gemini-2.5-flash"
,
nil
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
)
require
.
Equal
(
t
,
1
,
cache
.
deletedSessions
[
"gemini:session-123"
])
require
.
Equal
(
t
,
int64
(
2
),
cache
.
sessionBindings
[
"gemini:session-123"
])
})
}
func
TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_ForcePlatformFallback
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
groupID
:=
int64
(
9
)
ctx
=
context
.
WithValue
(
ctx
,
ctxkey
.
ForcePlatform
,
PlatformAntigravity
)
repo
:=
&
mockAccountRepoForGemini
{
listByGroupFunc
:
func
(
ctx
context
.
Context
,
groupID
int64
,
platforms
[]
string
)
([]
Account
,
error
)
{
return
nil
,
nil
},
listByPlatformFunc
:
func
(
ctx
context
.
Context
,
platforms
[]
string
)
([]
Account
,
error
)
{
return
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAntigravity
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
},
nil
},
accountsByID
:
map
[
int64
]
*
Account
{
1
:
{
ID
:
1
,
Platform
:
PlatformAntigravity
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
},
}
cache
:=
&
mockGatewayCacheForGemini
{}
groupRepo
:=
&
mockGroupRepoForGemini
{
groups
:
map
[
int64
]
*
Group
{}}
svc
:=
&
GeminiMessagesCompatService
{
accountRepo
:
repo
,
groupRepo
:
groupRepo
,
cache
:
cache
,
}
acc
,
err
:=
svc
.
SelectAccountForModelWithExclusions
(
ctx
,
&
groupID
,
""
,
"gemini-2.5-flash"
,
nil
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
1
),
acc
.
ID
)
}
func
TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_NoModelSupport
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
repo
:=
&
mockAccountRepoForGemini
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformGemini
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Credentials
:
map
[
string
]
any
{
"model_mapping"
:
map
[
string
]
any
{
"gemini-1.0-pro"
:
"gemini-1.0-pro"
}},
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForGemini
{}
groupRepo
:=
&
mockGroupRepoForGemini
{
groups
:
map
[
int64
]
*
Group
{}}
svc
:=
&
GeminiMessagesCompatService
{
accountRepo
:
repo
,
groupRepo
:
groupRepo
,
cache
:
cache
,
}
acc
,
err
:=
svc
.
SelectAccountForModelWithExclusions
(
ctx
,
nil
,
""
,
"gemini-2.5-flash"
,
nil
)
require
.
Error
(
t
,
err
)
require
.
Nil
(
t
,
acc
)
require
.
Contains
(
t
,
err
.
Error
(),
"supporting model"
)
}
func
TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_StickyMixedScheduling
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
repo
:=
&
mockAccountRepoForGemini
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAntigravity
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Extra
:
map
[
string
]
any
{
"mixed_scheduling"
:
true
}},
{
ID
:
2
,
Platform
:
PlatformGemini
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForGemini
{
sessionBindings
:
map
[
string
]
int64
{
"gemini:session-999"
:
1
},
}
groupRepo
:=
&
mockGroupRepoForGemini
{
groups
:
map
[
int64
]
*
Group
{}}
svc
:=
&
GeminiMessagesCompatService
{
accountRepo
:
repo
,
groupRepo
:
groupRepo
,
cache
:
cache
,
}
acc
,
err
:=
svc
.
SelectAccountForModelWithExclusions
(
ctx
,
nil
,
"session-999"
,
"gemini-2.5-flash"
,
nil
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
1
),
acc
.
ID
)
}
func
TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_SkipDisabledMixedScheduling
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
repo
:=
&
mockAccountRepoForGemini
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAntigravity
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
{
ID
:
2
,
Platform
:
PlatformGemini
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForGemini
{}
groupRepo
:=
&
mockGroupRepoForGemini
{
groups
:
map
[
int64
]
*
Group
{}}
svc
:=
&
GeminiMessagesCompatService
{
accountRepo
:
repo
,
groupRepo
:
groupRepo
,
cache
:
cache
,
}
acc
,
err
:=
svc
.
SelectAccountForModelWithExclusions
(
ctx
,
nil
,
""
,
"gemini-2.5-flash"
,
nil
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
)
}
func
TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_ExcludedAccount
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
repo
:=
&
mockAccountRepoForGemini
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformGemini
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
},
{
ID
:
2
,
Platform
:
PlatformGemini
,
Priority
:
2
,
Status
:
StatusActive
,
Schedulable
:
true
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForGemini
{}
groupRepo
:=
&
mockGroupRepoForGemini
{
groups
:
map
[
int64
]
*
Group
{}}
svc
:=
&
GeminiMessagesCompatService
{
accountRepo
:
repo
,
groupRepo
:
groupRepo
,
cache
:
cache
,
}
excluded
:=
map
[
int64
]
struct
{}{
1
:
{}}
acc
,
err
:=
svc
.
SelectAccountForModelWithExclusions
(
ctx
,
nil
,
""
,
"gemini-2.5-flash"
,
excluded
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
)
}
func
TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_ListError
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
repo
:=
&
mockAccountRepoForGemini
{
listByPlatformFunc
:
func
(
ctx
context
.
Context
,
platforms
[]
string
)
([]
Account
,
error
)
{
return
nil
,
errors
.
New
(
"query failed"
)
},
}
cache
:=
&
mockGatewayCacheForGemini
{}
groupRepo
:=
&
mockGroupRepoForGemini
{
groups
:
map
[
int64
]
*
Group
{}}
svc
:=
&
GeminiMessagesCompatService
{
accountRepo
:
repo
,
groupRepo
:
groupRepo
,
cache
:
cache
,
}
acc
,
err
:=
svc
.
SelectAccountForModelWithExclusions
(
ctx
,
nil
,
""
,
"gemini-2.5-flash"
,
nil
)
require
.
Error
(
t
,
err
)
require
.
Nil
(
t
,
acc
)
require
.
Contains
(
t
,
err
.
Error
(),
"query accounts failed"
)
}
func
TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_PreferOAuth
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
repo
:=
&
mockAccountRepoForGemini
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformGemini
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Type
:
AccountTypeAPIKey
},
{
ID
:
2
,
Platform
:
PlatformGemini
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
Type
:
AccountTypeOAuth
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForGemini
{}
groupRepo
:=
&
mockGroupRepoForGemini
{
groups
:
map
[
int64
]
*
Group
{}}
svc
:=
&
GeminiMessagesCompatService
{
accountRepo
:
repo
,
groupRepo
:
groupRepo
,
cache
:
cache
,
}
acc
,
err
:=
svc
.
SelectAccountForModelWithExclusions
(
ctx
,
nil
,
""
,
"gemini-2.5-pro"
,
nil
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
)
}
func
TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_PreferLeastRecentlyUsed
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
oldTime
:=
time
.
Now
()
.
Add
(
-
2
*
time
.
Hour
)
newTime
:=
time
.
Now
()
.
Add
(
-
1
*
time
.
Hour
)
repo
:=
&
mockAccountRepoForGemini
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformGemini
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
LastUsedAt
:
&
newTime
},
{
ID
:
2
,
Platform
:
PlatformGemini
,
Priority
:
1
,
Status
:
StatusActive
,
Schedulable
:
true
,
LastUsedAt
:
&
oldTime
},
},
accountsByID
:
map
[
int64
]
*
Account
{},
}
for
i
:=
range
repo
.
accounts
{
repo
.
accountsByID
[
repo
.
accounts
[
i
]
.
ID
]
=
&
repo
.
accounts
[
i
]
}
cache
:=
&
mockGatewayCacheForGemini
{}
groupRepo
:=
&
mockGroupRepoForGemini
{
groups
:
map
[
int64
]
*
Group
{}}
svc
:=
&
GeminiMessagesCompatService
{
accountRepo
:
repo
,
groupRepo
:
groupRepo
,
cache
:
cache
,
}
acc
,
err
:=
svc
.
SelectAccountForModelWithExclusions
(
ctx
,
nil
,
""
,
"gemini-2.5-pro"
,
nil
)
require
.
NoError
(
t
,
err
)
require
.
NotNil
(
t
,
acc
)
require
.
Equal
(
t
,
int64
(
2
),
acc
.
ID
)
}
// TestGeminiPlatformRouting_DocumentRouteDecision 测试平台路由决策逻辑
...
...
backend/internal/service/openai_gateway_service.go
View file @
f6ed3d14
...
...
@@ -180,81 +180,164 @@ func (s *OpenAIGatewayService) SelectAccountForModel(ctx context.Context, groupI
}
// SelectAccountForModelWithExclusions selects an account supporting the requested model while excluding specified accounts.
// SelectAccountForModelWithExclusions 选择支持指定模型的账号,同时排除指定的账号。
func
(
s
*
OpenAIGatewayService
)
SelectAccountForModelWithExclusions
(
ctx
context
.
Context
,
groupID
*
int64
,
sessionHash
string
,
requestedModel
string
,
excludedIDs
map
[
int64
]
struct
{})
(
*
Account
,
error
)
{
// 1. Check sticky session
if
sessionHash
!=
""
{
accountID
,
err
:=
s
.
cache
.
GetSessionAccountID
(
ctx
,
derefGroupID
(
groupID
),
"openai:"
+
sessionHash
)
if
err
==
nil
&&
accountID
>
0
{
if
_
,
excluded
:=
excludedIDs
[
accountID
];
!
excluded
{
account
,
err
:=
s
.
getSchedulableAccount
(
ctx
,
accountID
)
if
err
==
nil
&&
account
.
IsSchedulable
()
&&
account
.
IsOpenAI
()
&&
(
requestedModel
==
""
||
account
.
IsModelSupported
(
requestedModel
))
{
// Refresh sticky session TTL
_
=
s
.
cache
.
RefreshSessionTTL
(
ctx
,
derefGroupID
(
groupID
),
"openai:"
+
sessionHash
,
openaiStickySessionTTL
)
return
account
,
nil
}
}
}
cacheKey
:=
"openai:"
+
sessionHash
// 1. 尝试粘性会话命中
// Try sticky session hit
if
account
:=
s
.
tryStickySessionHit
(
ctx
,
groupID
,
sessionHash
,
cacheKey
,
requestedModel
,
excludedIDs
);
account
!=
nil
{
return
account
,
nil
}
// 2. Get schedulable OpenAI accounts
// 2. 获取可调度的 OpenAI 账号
// Get schedulable OpenAI accounts
accounts
,
err
:=
s
.
listSchedulableAccounts
(
ctx
,
groupID
)
if
err
!=
nil
{
return
nil
,
fmt
.
Errorf
(
"query accounts failed: %w"
,
err
)
}
// 3. Select by priority + LRU
// 3. 按优先级 + LRU 选择最佳账号
// Select by priority + LRU
selected
:=
s
.
selectBestAccount
(
accounts
,
requestedModel
,
excludedIDs
)
if
selected
==
nil
{
if
requestedModel
!=
""
{
return
nil
,
fmt
.
Errorf
(
"no available OpenAI accounts supporting model: %s"
,
requestedModel
)
}
return
nil
,
errors
.
New
(
"no available OpenAI accounts"
)
}
// 4. 设置粘性会话绑定
// Set sticky session binding
if
sessionHash
!=
""
{
_
=
s
.
cache
.
SetSessionAccountID
(
ctx
,
derefGroupID
(
groupID
),
cacheKey
,
selected
.
ID
,
openaiStickySessionTTL
)
}
return
selected
,
nil
}
// tryStickySessionHit 尝试从粘性会话获取账号。
// 如果命中且账号可用则返回账号;如果账号不可用则清理会话并返回 nil。
//
// tryStickySessionHit attempts to get account from sticky session.
// Returns account if hit and usable; clears session and returns nil if account is unavailable.
func
(
s
*
OpenAIGatewayService
)
tryStickySessionHit
(
ctx
context
.
Context
,
groupID
*
int64
,
sessionHash
,
cacheKey
,
requestedModel
string
,
excludedIDs
map
[
int64
]
struct
{})
*
Account
{
if
sessionHash
==
""
{
return
nil
}
accountID
,
err
:=
s
.
cache
.
GetSessionAccountID
(
ctx
,
derefGroupID
(
groupID
),
cacheKey
)
if
err
!=
nil
||
accountID
<=
0
{
return
nil
}
if
_
,
excluded
:=
excludedIDs
[
accountID
];
excluded
{
return
nil
}
account
,
err
:=
s
.
getSchedulableAccount
(
ctx
,
accountID
)
if
err
!=
nil
{
return
nil
}
// 检查账号是否需要清理粘性会话
// Check if sticky session should be cleared
if
shouldClearStickySession
(
account
)
{
_
=
s
.
cache
.
DeleteSessionAccountID
(
ctx
,
derefGroupID
(
groupID
),
cacheKey
)
return
nil
}
// 验证账号是否可用于当前请求
// Verify account is usable for current request
if
!
account
.
IsSchedulable
()
||
!
account
.
IsOpenAI
()
{
return
nil
}
if
requestedModel
!=
""
&&
!
account
.
IsModelSupported
(
requestedModel
)
{
return
nil
}
// 刷新会话 TTL 并返回账号
// Refresh session TTL and return account
_
=
s
.
cache
.
RefreshSessionTTL
(
ctx
,
derefGroupID
(
groupID
),
cacheKey
,
openaiStickySessionTTL
)
return
account
}
// selectBestAccount 从候选账号中选择最佳账号(优先级 + LRU)。
// 返回 nil 表示无可用账号。
//
// selectBestAccount selects the best account from candidates (priority + LRU).
// Returns nil if no available account.
func
(
s
*
OpenAIGatewayService
)
selectBestAccount
(
accounts
[]
Account
,
requestedModel
string
,
excludedIDs
map
[
int64
]
struct
{})
*
Account
{
var
selected
*
Account
for
i
:=
range
accounts
{
acc
:=
&
accounts
[
i
]
// 跳过被排除的账号
// Skip excluded accounts
if
_
,
excluded
:=
excludedIDs
[
acc
.
ID
];
excluded
{
continue
}
// Scheduler snapshots can be temporarily stale; re-check schedulability here to
// avoid selecting accounts that were recently rate-limited/overloaded.
if
!
acc
.
IsSchedulable
()
{
// 调度器快照可能暂时过时,这里重新检查可调度性和平台
// Scheduler snapshots can be temporarily stale; re-check schedulability and platform
if
!
acc
.
IsSchedulable
()
||
!
acc
.
IsOpenAI
()
{
continue
}
// 检查模型支持
// Check model support
if
requestedModel
!=
""
&&
!
acc
.
IsModelSupported
(
requestedModel
)
{
continue
}
// 选择优先级最高且最久未使用的账号
// Select highest priority and least recently used
if
selected
==
nil
{
selected
=
acc
continue
}
// Lower priority value means higher priority
if
acc
.
Priority
<
selected
.
Priority
{
if
s
.
isBetterAccount
(
acc
,
selected
)
{
selected
=
acc
}
else
if
acc
.
Priority
==
selected
.
Priority
{
switch
{
case
acc
.
LastUsedAt
==
nil
&&
selected
.
LastUsedAt
!=
nil
:
selected
=
acc
case
acc
.
LastUsedAt
!=
nil
&&
selected
.
LastUsedAt
==
nil
:
// keep selected (never used is preferred)
case
acc
.
LastUsedAt
==
nil
&&
selected
.
LastUsedAt
==
nil
:
// keep selected (both never used)
default
:
// Same priority, select least recently used
if
acc
.
LastUsedAt
.
Before
(
*
selected
.
LastUsedAt
)
{
selected
=
acc
}
}
}
}
if
selected
==
nil
{
if
requestedModel
!=
""
{
return
nil
,
fmt
.
Errorf
(
"no available OpenAI accounts supporting model: %s"
,
requestedModel
)
}
return
nil
,
errors
.
New
(
"no available OpenAI accounts"
)
}
return
selected
}
// 4. Set sticky session
if
sessionHash
!=
""
{
_
=
s
.
cache
.
SetSessionAccountID
(
ctx
,
derefGroupID
(
groupID
),
"openai:"
+
sessionHash
,
selected
.
ID
,
openaiStickySessionTTL
)
// isBetterAccount 判断 candidate 是否比 current 更优。
// 规则:优先级更高(数值更小)优先;同优先级时,未使用过的优先,其次是最久未使用的。
//
// isBetterAccount checks if candidate is better than current.
// Rules: higher priority (lower value) wins; same priority: never used > least recently used.
func
(
s
*
OpenAIGatewayService
)
isBetterAccount
(
candidate
,
current
*
Account
)
bool
{
// 优先级更高(数值更小)
// Higher priority (lower value)
if
candidate
.
Priority
<
current
.
Priority
{
return
true
}
if
candidate
.
Priority
>
current
.
Priority
{
return
false
}
return
selected
,
nil
// 同优先级,比较最后使用时间
// Same priority, compare last used time
switch
{
case
candidate
.
LastUsedAt
==
nil
&&
current
.
LastUsedAt
!=
nil
:
// candidate 从未使用,优先
return
true
case
candidate
.
LastUsedAt
!=
nil
&&
current
.
LastUsedAt
==
nil
:
// current 从未使用,保持
return
false
case
candidate
.
LastUsedAt
==
nil
&&
current
.
LastUsedAt
==
nil
:
// 都未使用,保持
return
false
default
:
// 都使用过,选择最久未使用的
return
candidate
.
LastUsedAt
.
Before
(
*
current
.
LastUsedAt
)
}
}
// SelectAccountWithLoadAwareness selects an account with load-awareness and wait plan.
...
...
@@ -325,29 +408,35 @@ func (s *OpenAIGatewayService) SelectAccountWithLoadAwareness(ctx context.Contex
accountID
,
err
:=
s
.
cache
.
GetSessionAccountID
(
ctx
,
derefGroupID
(
groupID
),
"openai:"
+
sessionHash
)
if
err
==
nil
&&
accountID
>
0
&&
!
isExcluded
(
accountID
)
{
account
,
err
:=
s
.
getSchedulableAccount
(
ctx
,
accountID
)
if
err
==
nil
&&
account
.
IsSchedulable
()
&&
account
.
IsOpenAI
()
&&
(
requestedModel
==
""
||
account
.
IsModelSupported
(
requestedModel
))
{
result
,
err
:=
s
.
tryAcquireAccountSlot
(
ctx
,
accountID
,
account
.
Concurrency
)
if
err
==
nil
&&
result
.
Acquired
{
_
=
s
.
cache
.
RefreshSessionTTL
(
ctx
,
derefGroupID
(
groupID
),
"openai:"
+
sessionHash
,
openaiStickySessionTTL
)
return
&
AccountSelectionResult
{
Account
:
account
,
Acquired
:
true
,
ReleaseFunc
:
result
.
ReleaseFunc
,
},
nil
if
err
==
nil
{
clearSticky
:=
shouldClearStickySession
(
account
)
if
clearSticky
{
_
=
s
.
cache
.
DeleteSessionAccountID
(
ctx
,
derefGroupID
(
groupID
),
"openai:"
+
sessionHash
)
}
if
!
clearSticky
&&
account
.
IsSchedulable
()
&&
account
.
IsOpenAI
()
&&
(
requestedModel
==
""
||
account
.
IsModelSupported
(
requestedModel
))
{
result
,
err
:=
s
.
tryAcquireAccountSlot
(
ctx
,
accountID
,
account
.
Concurrency
)
if
err
==
nil
&&
result
.
Acquired
{
_
=
s
.
cache
.
RefreshSessionTTL
(
ctx
,
derefGroupID
(
groupID
),
"openai:"
+
sessionHash
,
openaiStickySessionTTL
)
return
&
AccountSelectionResult
{
Account
:
account
,
Acquired
:
true
,
ReleaseFunc
:
result
.
ReleaseFunc
,
},
nil
}
waitingCount
,
_
:=
s
.
concurrencyService
.
GetAccountWaitingCount
(
ctx
,
accountID
)
if
waitingCount
<
cfg
.
StickySessionMaxWaiting
{
return
&
AccountSelectionResult
{
Account
:
account
,
WaitPlan
:
&
AccountWaitPlan
{
AccountID
:
accountID
,
MaxConcurrency
:
account
.
Concurrency
,
Timeout
:
cfg
.
StickySessionWaitTimeout
,
MaxWaiting
:
cfg
.
StickySessionMaxWaiting
,
},
},
nil
waitingCount
,
_
:=
s
.
concurrencyService
.
GetAccountWaitingCount
(
ctx
,
accountID
)
if
waitingCount
<
cfg
.
StickySessionMaxWaiting
{
return
&
AccountSelectionResult
{
Account
:
account
,
WaitPlan
:
&
AccountWaitPlan
{
AccountID
:
accountID
,
MaxConcurrency
:
account
.
Concurrency
,
Timeout
:
cfg
.
StickySessionWaitTimeout
,
MaxWaiting
:
cfg
.
StickySessionMaxWaiting
,
},
},
nil
}
}
}
}
...
...
backend/internal/service/openai_gateway_service_test.go
View file @
f6ed3d14
...
...
@@ -21,19 +21,50 @@ type stubOpenAIAccountRepo struct {
accounts
[]
Account
}
func
(
r
stubOpenAIAccountRepo
)
GetByID
(
ctx
context
.
Context
,
id
int64
)
(
*
Account
,
error
)
{
for
i
:=
range
r
.
accounts
{
if
r
.
accounts
[
i
]
.
ID
==
id
{
return
&
r
.
accounts
[
i
],
nil
}
}
return
nil
,
errors
.
New
(
"account not found"
)
}
func
(
r
stubOpenAIAccountRepo
)
ListSchedulableByGroupIDAndPlatform
(
ctx
context
.
Context
,
groupID
int64
,
platform
string
)
([]
Account
,
error
)
{
return
append
([]
Account
(
nil
),
r
.
accounts
...
),
nil
var
result
[]
Account
for
_
,
acc
:=
range
r
.
accounts
{
if
acc
.
Platform
==
platform
{
result
=
append
(
result
,
acc
)
}
}
return
result
,
nil
}
func
(
r
stubOpenAIAccountRepo
)
ListSchedulableByPlatform
(
ctx
context
.
Context
,
platform
string
)
([]
Account
,
error
)
{
return
append
([]
Account
(
nil
),
r
.
accounts
...
),
nil
var
result
[]
Account
for
_
,
acc
:=
range
r
.
accounts
{
if
acc
.
Platform
==
platform
{
result
=
append
(
result
,
acc
)
}
}
return
result
,
nil
}
type
stubConcurrencyCache
struct
{
ConcurrencyCache
loadBatchErr
error
loadMap
map
[
int64
]
*
AccountLoadInfo
acquireResults
map
[
int64
]
bool
waitCounts
map
[
int64
]
int
skipDefaultLoad
bool
}
func
(
c
stubConcurrencyCache
)
AcquireAccountSlot
(
ctx
context
.
Context
,
accountID
int64
,
maxConcurrency
int
,
requestID
string
)
(
bool
,
error
)
{
if
c
.
acquireResults
!=
nil
{
if
result
,
ok
:=
c
.
acquireResults
[
accountID
];
ok
{
return
result
,
nil
}
}
return
true
,
nil
}
...
...
@@ -42,8 +73,25 @@ func (c stubConcurrencyCache) ReleaseAccountSlot(ctx context.Context, accountID
}
func
(
c
stubConcurrencyCache
)
GetAccountsLoadBatch
(
ctx
context
.
Context
,
accounts
[]
AccountWithConcurrency
)
(
map
[
int64
]
*
AccountLoadInfo
,
error
)
{
if
c
.
loadBatchErr
!=
nil
{
return
nil
,
c
.
loadBatchErr
}
out
:=
make
(
map
[
int64
]
*
AccountLoadInfo
,
len
(
accounts
))
if
c
.
skipDefaultLoad
&&
c
.
loadMap
!=
nil
{
for
_
,
acc
:=
range
accounts
{
if
load
,
ok
:=
c
.
loadMap
[
acc
.
ID
];
ok
{
out
[
acc
.
ID
]
=
load
}
}
return
out
,
nil
}
for
_
,
acc
:=
range
accounts
{
if
c
.
loadMap
!=
nil
{
if
load
,
ok
:=
c
.
loadMap
[
acc
.
ID
];
ok
{
out
[
acc
.
ID
]
=
load
continue
}
}
out
[
acc
.
ID
]
=
&
AccountLoadInfo
{
AccountID
:
acc
.
ID
,
LoadRate
:
0
}
}
return
out
,
nil
...
...
@@ -92,6 +140,51 @@ func TestOpenAIGatewayService_GenerateSessionHash_Priority(t *testing.T) {
}
}
func
(
c
stubConcurrencyCache
)
GetAccountWaitingCount
(
ctx
context
.
Context
,
accountID
int64
)
(
int
,
error
)
{
if
c
.
waitCounts
!=
nil
{
if
count
,
ok
:=
c
.
waitCounts
[
accountID
];
ok
{
return
count
,
nil
}
}
return
0
,
nil
}
type
stubGatewayCache
struct
{
sessionBindings
map
[
string
]
int64
deletedSessions
map
[
string
]
int
}
func
(
c
*
stubGatewayCache
)
GetSessionAccountID
(
ctx
context
.
Context
,
groupID
int64
,
sessionHash
string
)
(
int64
,
error
)
{
if
id
,
ok
:=
c
.
sessionBindings
[
sessionHash
];
ok
{
return
id
,
nil
}
return
0
,
errors
.
New
(
"not found"
)
}
func
(
c
*
stubGatewayCache
)
SetSessionAccountID
(
ctx
context
.
Context
,
groupID
int64
,
sessionHash
string
,
accountID
int64
,
ttl
time
.
Duration
)
error
{
if
c
.
sessionBindings
==
nil
{
c
.
sessionBindings
=
make
(
map
[
string
]
int64
)
}
c
.
sessionBindings
[
sessionHash
]
=
accountID
return
nil
}
func
(
c
*
stubGatewayCache
)
RefreshSessionTTL
(
ctx
context
.
Context
,
groupID
int64
,
sessionHash
string
,
ttl
time
.
Duration
)
error
{
return
nil
}
func
(
c
*
stubGatewayCache
)
DeleteSessionAccountID
(
ctx
context
.
Context
,
groupID
int64
,
sessionHash
string
)
error
{
if
c
.
sessionBindings
==
nil
{
return
nil
}
if
c
.
deletedSessions
==
nil
{
c
.
deletedSessions
=
make
(
map
[
string
]
int
)
}
c
.
deletedSessions
[
sessionHash
]
++
delete
(
c
.
sessionBindings
,
sessionHash
)
return
nil
}
func
TestOpenAISelectAccountWithLoadAwareness_FiltersUnschedulable
(
t
*
testing
.
T
)
{
now
:=
time
.
Now
()
resetAt
:=
now
.
Add
(
10
*
time
.
Minute
)
...
...
@@ -182,6 +275,515 @@ func TestOpenAISelectAccountWithLoadAwareness_FiltersUnschedulableWhenNoConcurre
}
}
func
TestOpenAISelectAccountForModelWithExclusions_StickyUnschedulableClearsSession
(
t
*
testing
.
T
)
{
sessionHash
:=
"session-1"
repo
:=
stubOpenAIAccountRepo
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformOpenAI
,
Status
:
StatusDisabled
,
Schedulable
:
true
,
Concurrency
:
1
},
{
ID
:
2
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
1
},
},
}
cache
:=
&
stubGatewayCache
{
sessionBindings
:
map
[
string
]
int64
{
"openai:"
+
sessionHash
:
1
},
}
svc
:=
&
OpenAIGatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
}
acc
,
err
:=
svc
.
SelectAccountForModelWithExclusions
(
context
.
Background
(),
nil
,
sessionHash
,
"gpt-4"
,
nil
)
if
err
!=
nil
{
t
.
Fatalf
(
"SelectAccountForModelWithExclusions error: %v"
,
err
)
}
if
acc
==
nil
||
acc
.
ID
!=
2
{
t
.
Fatalf
(
"expected account 2, got %+v"
,
acc
)
}
if
cache
.
deletedSessions
[
"openai:"
+
sessionHash
]
!=
1
{
t
.
Fatalf
(
"expected sticky session to be deleted"
)
}
if
cache
.
sessionBindings
[
"openai:"
+
sessionHash
]
!=
2
{
t
.
Fatalf
(
"expected sticky session to bind to account 2"
)
}
}
func
TestOpenAISelectAccountWithLoadAwareness_StickyUnschedulableClearsSession
(
t
*
testing
.
T
)
{
sessionHash
:=
"session-2"
groupID
:=
int64
(
1
)
repo
:=
stubOpenAIAccountRepo
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformOpenAI
,
Status
:
StatusDisabled
,
Schedulable
:
true
,
Concurrency
:
1
},
{
ID
:
2
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
1
},
},
}
cache
:=
&
stubGatewayCache
{
sessionBindings
:
map
[
string
]
int64
{
"openai:"
+
sessionHash
:
1
},
}
svc
:=
&
OpenAIGatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
concurrencyService
:
NewConcurrencyService
(
stubConcurrencyCache
{}),
}
selection
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
context
.
Background
(),
&
groupID
,
sessionHash
,
"gpt-4"
,
nil
)
if
err
!=
nil
{
t
.
Fatalf
(
"SelectAccountWithLoadAwareness error: %v"
,
err
)
}
if
selection
==
nil
||
selection
.
Account
==
nil
||
selection
.
Account
.
ID
!=
2
{
t
.
Fatalf
(
"expected account 2, got %+v"
,
selection
)
}
if
cache
.
deletedSessions
[
"openai:"
+
sessionHash
]
!=
1
{
t
.
Fatalf
(
"expected sticky session to be deleted"
)
}
if
cache
.
sessionBindings
[
"openai:"
+
sessionHash
]
!=
2
{
t
.
Fatalf
(
"expected sticky session to bind to account 2"
)
}
if
selection
.
ReleaseFunc
!=
nil
{
selection
.
ReleaseFunc
()
}
}
func
TestOpenAISelectAccountForModelWithExclusions_NoModelSupport
(
t
*
testing
.
T
)
{
repo
:=
stubOpenAIAccountRepo
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Credentials
:
map
[
string
]
any
{
"model_mapping"
:
map
[
string
]
any
{
"gpt-3.5-turbo"
:
"gpt-3.5-turbo"
}},
},
},
}
cache
:=
&
stubGatewayCache
{}
svc
:=
&
OpenAIGatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
}
acc
,
err
:=
svc
.
SelectAccountForModelWithExclusions
(
context
.
Background
(),
nil
,
""
,
"gpt-4"
,
nil
)
if
err
==
nil
{
t
.
Fatalf
(
"expected error for unsupported model"
)
}
if
acc
!=
nil
{
t
.
Fatalf
(
"expected nil account for unsupported model"
)
}
if
!
strings
.
Contains
(
err
.
Error
(),
"supporting model"
)
{
t
.
Fatalf
(
"unexpected error: %v"
,
err
)
}
}
func
TestOpenAISelectAccountWithLoadAwareness_LoadBatchErrorFallback
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
1
)
repo
:=
stubOpenAIAccountRepo
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
1
,
Priority
:
2
},
{
ID
:
2
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
1
,
Priority
:
1
},
},
}
cache
:=
&
stubGatewayCache
{}
concurrencyCache
:=
stubConcurrencyCache
{
loadBatchErr
:
errors
.
New
(
"load batch failed"
),
}
svc
:=
&
OpenAIGatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
),
}
selection
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
context
.
Background
(),
&
groupID
,
"fallback"
,
"gpt-4"
,
nil
)
if
err
!=
nil
{
t
.
Fatalf
(
"SelectAccountWithLoadAwareness error: %v"
,
err
)
}
if
selection
==
nil
||
selection
.
Account
==
nil
{
t
.
Fatalf
(
"expected selection"
)
}
if
selection
.
Account
.
ID
!=
2
{
t
.
Fatalf
(
"expected account 2, got %d"
,
selection
.
Account
.
ID
)
}
if
cache
.
sessionBindings
[
"openai:fallback"
]
!=
2
{
t
.
Fatalf
(
"expected sticky session updated"
)
}
if
selection
.
ReleaseFunc
!=
nil
{
selection
.
ReleaseFunc
()
}
}
func
TestOpenAISelectAccountWithLoadAwareness_NoSlotFallbackWait
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
1
)
repo
:=
stubOpenAIAccountRepo
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
1
,
Priority
:
1
},
},
}
cache
:=
&
stubGatewayCache
{}
concurrencyCache
:=
stubConcurrencyCache
{
acquireResults
:
map
[
int64
]
bool
{
1
:
false
},
loadMap
:
map
[
int64
]
*
AccountLoadInfo
{
1
:
{
AccountID
:
1
,
LoadRate
:
10
},
},
}
svc
:=
&
OpenAIGatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
),
}
selection
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
context
.
Background
(),
&
groupID
,
""
,
"gpt-4"
,
nil
)
if
err
!=
nil
{
t
.
Fatalf
(
"SelectAccountWithLoadAwareness error: %v"
,
err
)
}
if
selection
==
nil
||
selection
.
WaitPlan
==
nil
{
t
.
Fatalf
(
"expected wait plan fallback"
)
}
if
selection
.
Account
==
nil
||
selection
.
Account
.
ID
!=
1
{
t
.
Fatalf
(
"expected account 1"
)
}
}
func
TestOpenAISelectAccountForModelWithExclusions_SetsStickyBinding
(
t
*
testing
.
T
)
{
sessionHash
:=
"bind"
repo
:=
stubOpenAIAccountRepo
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
1
,
Priority
:
1
},
},
}
cache
:=
&
stubGatewayCache
{}
svc
:=
&
OpenAIGatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
}
acc
,
err
:=
svc
.
SelectAccountForModelWithExclusions
(
context
.
Background
(),
nil
,
sessionHash
,
"gpt-4"
,
nil
)
if
err
!=
nil
{
t
.
Fatalf
(
"SelectAccountForModelWithExclusions error: %v"
,
err
)
}
if
acc
==
nil
||
acc
.
ID
!=
1
{
t
.
Fatalf
(
"expected account 1"
)
}
if
cache
.
sessionBindings
[
"openai:"
+
sessionHash
]
!=
1
{
t
.
Fatalf
(
"expected sticky session binding"
)
}
}
func
TestOpenAISelectAccountWithLoadAwareness_StickyWaitPlan
(
t
*
testing
.
T
)
{
sessionHash
:=
"sticky-wait"
groupID
:=
int64
(
1
)
repo
:=
stubOpenAIAccountRepo
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
1
,
Priority
:
1
},
},
}
cache
:=
&
stubGatewayCache
{
sessionBindings
:
map
[
string
]
int64
{
"openai:"
+
sessionHash
:
1
},
}
concurrencyCache
:=
stubConcurrencyCache
{
acquireResults
:
map
[
int64
]
bool
{
1
:
false
},
waitCounts
:
map
[
int64
]
int
{
1
:
0
},
}
svc
:=
&
OpenAIGatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
),
}
selection
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
context
.
Background
(),
&
groupID
,
sessionHash
,
"gpt-4"
,
nil
)
if
err
!=
nil
{
t
.
Fatalf
(
"SelectAccountWithLoadAwareness error: %v"
,
err
)
}
if
selection
==
nil
||
selection
.
WaitPlan
==
nil
{
t
.
Fatalf
(
"expected sticky wait plan"
)
}
if
selection
.
Account
==
nil
||
selection
.
Account
.
ID
!=
1
{
t
.
Fatalf
(
"expected account 1"
)
}
}
func
TestOpenAISelectAccountWithLoadAwareness_PrefersLowerLoad
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
1
)
repo
:=
stubOpenAIAccountRepo
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
1
,
Priority
:
1
},
{
ID
:
2
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
1
,
Priority
:
1
},
},
}
cache
:=
&
stubGatewayCache
{}
concurrencyCache
:=
stubConcurrencyCache
{
loadMap
:
map
[
int64
]
*
AccountLoadInfo
{
1
:
{
AccountID
:
1
,
LoadRate
:
80
},
2
:
{
AccountID
:
2
,
LoadRate
:
10
},
},
}
svc
:=
&
OpenAIGatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
),
}
selection
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
context
.
Background
(),
&
groupID
,
"load"
,
"gpt-4"
,
nil
)
if
err
!=
nil
{
t
.
Fatalf
(
"SelectAccountWithLoadAwareness error: %v"
,
err
)
}
if
selection
==
nil
||
selection
.
Account
==
nil
||
selection
.
Account
.
ID
!=
2
{
t
.
Fatalf
(
"expected account 2"
)
}
if
cache
.
sessionBindings
[
"openai:load"
]
!=
2
{
t
.
Fatalf
(
"expected sticky session updated"
)
}
}
func
TestOpenAISelectAccountForModelWithExclusions_StickyExcludedFallback
(
t
*
testing
.
T
)
{
sessionHash
:=
"excluded"
repo
:=
stubOpenAIAccountRepo
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
1
,
Priority
:
1
},
{
ID
:
2
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
1
,
Priority
:
2
},
},
}
cache
:=
&
stubGatewayCache
{
sessionBindings
:
map
[
string
]
int64
{
"openai:"
+
sessionHash
:
1
},
}
svc
:=
&
OpenAIGatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
}
excluded
:=
map
[
int64
]
struct
{}{
1
:
{}}
acc
,
err
:=
svc
.
SelectAccountForModelWithExclusions
(
context
.
Background
(),
nil
,
sessionHash
,
"gpt-4"
,
excluded
)
if
err
!=
nil
{
t
.
Fatalf
(
"SelectAccountForModelWithExclusions error: %v"
,
err
)
}
if
acc
==
nil
||
acc
.
ID
!=
2
{
t
.
Fatalf
(
"expected account 2"
)
}
}
func
TestOpenAISelectAccountForModelWithExclusions_StickyNonOpenAI
(
t
*
testing
.
T
)
{
sessionHash
:=
"non-openai"
repo
:=
stubOpenAIAccountRepo
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformAnthropic
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
1
,
Priority
:
1
},
{
ID
:
2
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
1
,
Priority
:
2
},
},
}
cache
:=
&
stubGatewayCache
{
sessionBindings
:
map
[
string
]
int64
{
"openai:"
+
sessionHash
:
1
},
}
svc
:=
&
OpenAIGatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
}
acc
,
err
:=
svc
.
SelectAccountForModelWithExclusions
(
context
.
Background
(),
nil
,
sessionHash
,
"gpt-4"
,
nil
)
if
err
!=
nil
{
t
.
Fatalf
(
"SelectAccountForModelWithExclusions error: %v"
,
err
)
}
if
acc
==
nil
||
acc
.
ID
!=
2
{
t
.
Fatalf
(
"expected account 2"
)
}
}
func
TestOpenAISelectAccountForModelWithExclusions_NoAccounts
(
t
*
testing
.
T
)
{
repo
:=
stubOpenAIAccountRepo
{
accounts
:
[]
Account
{}}
cache
:=
&
stubGatewayCache
{}
svc
:=
&
OpenAIGatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
}
acc
,
err
:=
svc
.
SelectAccountForModelWithExclusions
(
context
.
Background
(),
nil
,
""
,
""
,
nil
)
if
err
==
nil
{
t
.
Fatalf
(
"expected error for no accounts"
)
}
if
acc
!=
nil
{
t
.
Fatalf
(
"expected nil account"
)
}
if
!
strings
.
Contains
(
err
.
Error
(),
"no available OpenAI accounts"
)
{
t
.
Fatalf
(
"unexpected error: %v"
,
err
)
}
}
func
TestOpenAISelectAccountWithLoadAwareness_NoCandidates
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
1
)
resetAt
:=
time
.
Now
()
.
Add
(
1
*
time
.
Hour
)
repo
:=
stubOpenAIAccountRepo
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
1
,
Priority
:
1
,
RateLimitResetAt
:
&
resetAt
},
},
}
cache
:=
&
stubGatewayCache
{}
concurrencyCache
:=
stubConcurrencyCache
{}
svc
:=
&
OpenAIGatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
),
}
selection
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
context
.
Background
(),
&
groupID
,
""
,
"gpt-4"
,
nil
)
if
err
==
nil
{
t
.
Fatalf
(
"expected error for no candidates"
)
}
if
selection
!=
nil
{
t
.
Fatalf
(
"expected nil selection"
)
}
}
func
TestOpenAISelectAccountWithLoadAwareness_AllFullWaitPlan
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
1
)
repo
:=
stubOpenAIAccountRepo
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
1
,
Priority
:
1
},
},
}
cache
:=
&
stubGatewayCache
{}
concurrencyCache
:=
stubConcurrencyCache
{
loadMap
:
map
[
int64
]
*
AccountLoadInfo
{
1
:
{
AccountID
:
1
,
LoadRate
:
100
},
},
}
svc
:=
&
OpenAIGatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
),
}
selection
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
context
.
Background
(),
&
groupID
,
""
,
"gpt-4"
,
nil
)
if
err
!=
nil
{
t
.
Fatalf
(
"SelectAccountWithLoadAwareness error: %v"
,
err
)
}
if
selection
==
nil
||
selection
.
WaitPlan
==
nil
{
t
.
Fatalf
(
"expected wait plan"
)
}
}
func
TestOpenAISelectAccountWithLoadAwareness_LoadBatchErrorNoAcquire
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
1
)
repo
:=
stubOpenAIAccountRepo
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
1
,
Priority
:
1
},
},
}
cache
:=
&
stubGatewayCache
{}
concurrencyCache
:=
stubConcurrencyCache
{
loadBatchErr
:
errors
.
New
(
"load batch failed"
),
acquireResults
:
map
[
int64
]
bool
{
1
:
false
},
}
svc
:=
&
OpenAIGatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
),
}
selection
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
context
.
Background
(),
&
groupID
,
""
,
"gpt-4"
,
nil
)
if
err
!=
nil
{
t
.
Fatalf
(
"SelectAccountWithLoadAwareness error: %v"
,
err
)
}
if
selection
==
nil
||
selection
.
WaitPlan
==
nil
{
t
.
Fatalf
(
"expected wait plan"
)
}
}
func
TestOpenAISelectAccountWithLoadAwareness_MissingLoadInfo
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
1
)
repo
:=
stubOpenAIAccountRepo
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
1
,
Priority
:
1
},
{
ID
:
2
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
1
,
Priority
:
1
},
},
}
cache
:=
&
stubGatewayCache
{}
concurrencyCache
:=
stubConcurrencyCache
{
loadMap
:
map
[
int64
]
*
AccountLoadInfo
{
1
:
{
AccountID
:
1
,
LoadRate
:
50
},
},
skipDefaultLoad
:
true
,
}
svc
:=
&
OpenAIGatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
),
}
selection
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
context
.
Background
(),
&
groupID
,
""
,
"gpt-4"
,
nil
)
if
err
!=
nil
{
t
.
Fatalf
(
"SelectAccountWithLoadAwareness error: %v"
,
err
)
}
if
selection
==
nil
||
selection
.
Account
==
nil
||
selection
.
Account
.
ID
!=
2
{
t
.
Fatalf
(
"expected account 2"
)
}
}
func
TestOpenAISelectAccountForModelWithExclusions_LeastRecentlyUsed
(
t
*
testing
.
T
)
{
oldTime
:=
time
.
Now
()
.
Add
(
-
2
*
time
.
Hour
)
newTime
:=
time
.
Now
()
.
Add
(
-
1
*
time
.
Hour
)
repo
:=
stubOpenAIAccountRepo
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Priority
:
1
,
LastUsedAt
:
&
newTime
},
{
ID
:
2
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Priority
:
1
,
LastUsedAt
:
&
oldTime
},
},
}
cache
:=
&
stubGatewayCache
{}
svc
:=
&
OpenAIGatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
}
acc
,
err
:=
svc
.
SelectAccountForModelWithExclusions
(
context
.
Background
(),
nil
,
""
,
"gpt-4"
,
nil
)
if
err
!=
nil
{
t
.
Fatalf
(
"SelectAccountForModelWithExclusions error: %v"
,
err
)
}
if
acc
==
nil
||
acc
.
ID
!=
2
{
t
.
Fatalf
(
"expected account 2"
)
}
}
func
TestOpenAISelectAccountWithLoadAwareness_PreferNeverUsed
(
t
*
testing
.
T
)
{
groupID
:=
int64
(
1
)
lastUsed
:=
time
.
Now
()
.
Add
(
-
1
*
time
.
Hour
)
repo
:=
stubOpenAIAccountRepo
{
accounts
:
[]
Account
{
{
ID
:
1
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
1
,
Priority
:
1
,
LastUsedAt
:
&
lastUsed
},
{
ID
:
2
,
Platform
:
PlatformOpenAI
,
Status
:
StatusActive
,
Schedulable
:
true
,
Concurrency
:
1
,
Priority
:
1
},
},
}
cache
:=
&
stubGatewayCache
{}
concurrencyCache
:=
stubConcurrencyCache
{
loadMap
:
map
[
int64
]
*
AccountLoadInfo
{
1
:
{
AccountID
:
1
,
LoadRate
:
10
},
2
:
{
AccountID
:
2
,
LoadRate
:
10
},
},
}
svc
:=
&
OpenAIGatewayService
{
accountRepo
:
repo
,
cache
:
cache
,
concurrencyService
:
NewConcurrencyService
(
concurrencyCache
),
}
selection
,
err
:=
svc
.
SelectAccountWithLoadAwareness
(
context
.
Background
(),
&
groupID
,
""
,
"gpt-4"
,
nil
)
if
err
!=
nil
{
t
.
Fatalf
(
"SelectAccountWithLoadAwareness error: %v"
,
err
)
}
if
selection
==
nil
||
selection
.
Account
==
nil
||
selection
.
Account
.
ID
!=
2
{
t
.
Fatalf
(
"expected account 2"
)
}
}
func
TestOpenAIStreamingTimeout
(
t
*
testing
.
T
)
{
gin
.
SetMode
(
gin
.
TestMode
)
cfg
:=
&
config
.
Config
{
...
...
backend/internal/service/sticky_session_test.go
0 → 100644
View file @
f6ed3d14
//go:build unit
// Package service 提供 API 网关核心服务。
// 本文件包含 shouldClearStickySession 函数的单元测试,
// 验证粘性会话清理逻辑在各种账号状态下的正确行为。
//
// This file contains unit tests for the shouldClearStickySession function,
// verifying correct sticky session clearing behavior under various account states.
package
service
import
(
"testing"
"time"
"github.com/stretchr/testify/require"
)
// TestShouldClearStickySession 测试粘性会话清理判断逻辑。
// 验证在以下情况下是否正确判断需要清理粘性会话:
// - nil 账号:不清理(返回 false)
// - 状态为错误或禁用:清理
// - 不可调度:清理
// - 临时不可调度且未过期:清理
// - 临时不可调度已过期:不清理
// - 正常可调度状态:不清理
//
// TestShouldClearStickySession tests the sticky session clearing logic.
// Verifies correct behavior for various account states including:
// nil account, error/disabled status, unschedulable, temporary unschedulable.
func
TestShouldClearStickySession
(
t
*
testing
.
T
)
{
now
:=
time
.
Now
()
future
:=
now
.
Add
(
1
*
time
.
Hour
)
past
:=
now
.
Add
(
-
1
*
time
.
Hour
)
tests
:=
[]
struct
{
name
string
account
*
Account
want
bool
}{
{
name
:
"nil account"
,
account
:
nil
,
want
:
false
},
{
name
:
"status error"
,
account
:
&
Account
{
Status
:
StatusError
,
Schedulable
:
true
},
want
:
true
},
{
name
:
"status disabled"
,
account
:
&
Account
{
Status
:
StatusDisabled
,
Schedulable
:
true
},
want
:
true
},
{
name
:
"schedulable false"
,
account
:
&
Account
{
Status
:
StatusActive
,
Schedulable
:
false
},
want
:
true
},
{
name
:
"temp unschedulable"
,
account
:
&
Account
{
Status
:
StatusActive
,
Schedulable
:
true
,
TempUnschedulableUntil
:
&
future
},
want
:
true
},
{
name
:
"temp unschedulable expired"
,
account
:
&
Account
{
Status
:
StatusActive
,
Schedulable
:
true
,
TempUnschedulableUntil
:
&
past
},
want
:
false
},
{
name
:
"active schedulable"
,
account
:
&
Account
{
Status
:
StatusActive
,
Schedulable
:
true
},
want
:
false
},
}
for
_
,
tt
:=
range
tests
{
t
.
Run
(
tt
.
name
,
func
(
t
*
testing
.
T
)
{
require
.
Equal
(
t
,
tt
.
want
,
shouldClearStickySession
(
tt
.
account
))
})
}
}
backend/internal/service/usage_cleanup_service_test.go
View file @
f6ed3d14
...
...
@@ -345,6 +345,9 @@ func TestUsageCleanupServiceRunOnceSuccess(t *testing.T) {
repo
.
mu
.
Lock
()
defer
repo
.
mu
.
Unlock
()
require
.
Len
(
t
,
repo
.
deleteCalls
,
3
)
require
.
Equal
(
t
,
2
,
repo
.
deleteCalls
[
0
]
.
limit
)
require
.
True
(
t
,
repo
.
deleteCalls
[
0
]
.
filters
.
StartTime
.
Equal
(
start
))
require
.
True
(
t
,
repo
.
deleteCalls
[
0
]
.
filters
.
EndTime
.
Equal
(
end
))
require
.
Len
(
t
,
repo
.
markSucceeded
,
1
)
require
.
Empty
(
t
,
repo
.
markFailed
)
require
.
Equal
(
t
,
int64
(
5
),
repo
.
markSucceeded
[
0
]
.
taskID
)
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment