fix(openai-ws): 对齐 OpenAI 新版上下文池容量计算

This commit is contained in:
xlplbo
2026-09-12 16:13:16 +08:00
parent 4726bdd08b
commit 956f4672e4
6 changed files with 95 additions and 23 deletions
+3 -1
View File
@@ -1274,7 +1274,9 @@ type GatewayOpenAIWSConfig struct {
MaxConnsPerAccount int `mapstructure:"max_conns_per_account"`
MinIdlePerAccount int `mapstructure:"min_idle_per_account"`
MaxIdlePerAccount int `mapstructure:"max_idle_per_account"`
// DynamicMaxConnsByAccountConcurrencyEnabled: 是否按账号并发动态计算连接池上限
// DynamicMaxConnsByAccountConcurrencyEnabled: 是否按账号并发动态计算连接池上限。
// 旧版及 mode_router_v2 的 ctx_pool 共用此开关和类型系数;关闭后使用 max_conns_per_account。
// mode_router_v2 下并发数 <= 0 的账号仍不可调度。
DynamicMaxConnsByAccountConcurrencyEnabled bool `mapstructure:"dynamic_max_conns_by_account_concurrency_enabled"`
// OAuthMaxConnsFactor: OAuth 账号连接池系数(effective=ceil(concurrency*factor))
OAuthMaxConnsFactor float64 `mapstructure:"oauth_max_conns_factor"`
+2 -8
View File
@@ -2156,14 +2156,8 @@ func (p *openAIWSConnPool) effectiveMaxConnsByAccount(account *Account) int {
if hardCap <= 0 {
return 0
}
if p.modeRouterV2Enabled() {
if account == nil {
return hardCap
}
if account.Concurrency <= 0 {
return 0
}
return min(account.Concurrency, hardCap)
if p.modeRouterV2Enabled() && account != nil && account.Concurrency <= 0 {
return 0
}
if account == nil || !p.dynamicMaxConnsEnabled() {
return hardCap
+86 -12
View File
@@ -1164,21 +1164,95 @@ func TestOpenAIWSConnPool_EffectiveMaxConnsDisabledFallbackHardCap(t *testing.T)
require.Equal(t, 8, pool.effectiveMaxConnsByAccount(account), "关闭动态模式后应保持旧行为")
}
func TestOpenAIWSConnPool_EffectiveMaxConnsByAccount_ModeRouterV2RespectsHardCap(t *testing.T) {
cfg := &config.Config{}
cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 8
cfg.Gateway.OpenAIWS.DynamicMaxConnsByAccountConcurrencyEnabled = true
cfg.Gateway.OpenAIWS.OAuthMaxConnsFactor = 0.3
cfg.Gateway.OpenAIWS.APIKeyMaxConnsFactor = 0.6
func TestOpenAIWSConnPool_EffectiveMaxConnsByAccount_ModeRouterV2(t *testing.T) {
tests := []struct {
name string
accountType string
concurrency int
factor float64
dynamic bool
want int
}{
{name: "oauth expansion", accountType: AccountTypeOAuth, concurrency: 1, factor: 5, dynamic: true, want: 5},
{name: "apikey expansion", accountType: AccountTypeAPIKey, concurrency: 1, factor: 5, dynamic: true, want: 5},
{name: "oauth fraction", accountType: AccountTypeOAuth, concurrency: 20, factor: 0.3, dynamic: true, want: 6},
{name: "apikey rounding", accountType: AccountTypeAPIKey, concurrency: 3, factor: 0.6, dynamic: true, want: 2},
{name: "minimum one", accountType: AccountTypeOAuth, concurrency: 1, factor: 0.3, dynamic: true, want: 1},
{name: "hard cap", accountType: AccountTypeOAuth, concurrency: 3, factor: 5, dynamic: true, want: 8},
{name: "dynamic disabled", accountType: AccountTypeAPIKey, concurrency: 2, factor: 0.6, want: 8},
{name: "zero concurrency", accountType: AccountTypeOAuth, factor: 5, dynamic: true, want: 0},
{name: "negative concurrency", accountType: AccountTypeAPIKey, concurrency: -1, factor: 5, dynamic: true, want: 0},
{name: "zero concurrency dynamic disabled", accountType: AccountTypeOAuth, factor: 5, want: 0},
{name: "negative concurrency dynamic disabled", accountType: AccountTypeAPIKey, concurrency: -1, factor: 5, want: 0},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
cfg := &config.Config{}
cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 8
cfg.Gateway.OpenAIWS.DynamicMaxConnsByAccountConcurrencyEnabled = tt.dynamic
cfg.Gateway.OpenAIWS.OAuthMaxConnsFactor = 1
cfg.Gateway.OpenAIWS.APIKeyMaxConnsFactor = 1
if tt.accountType == AccountTypeOAuth {
cfg.Gateway.OpenAIWS.OAuthMaxConnsFactor = tt.factor
} else {
cfg.Gateway.OpenAIWS.APIKeyMaxConnsFactor = tt.factor
}
pool := newOpenAIWSConnPool(cfg)
t.Cleanup(pool.Close)
account := &Account{Platform: PlatformOpenAI, Type: tt.accountType, Concurrency: tt.concurrency}
require.Equal(t, tt.want, pool.effectiveMaxConnsByAccount(account))
require.Equal(t, 8, pool.effectiveMaxConnsByAccount(nil))
})
}
}
pool := newOpenAIWSConnPool(cfg)
func TestOpenAIWSConnPool_AcquireRetainedSessionsUsesScaledCapacity(t *testing.T) {
for _, mode := range []struct {
name string
v2 bool
}{
{name: "legacy"},
{name: "mode router v2", v2: true},
} {
t.Run(mode.name, func(t *testing.T) {
cfg := &config.Config{}
cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = mode.v2
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 3
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 3
cfg.Gateway.OpenAIWS.PoolTargetUtilization = 1
cfg.Gateway.OpenAIWS.DynamicMaxConnsByAccountConcurrencyEnabled = true
cfg.Gateway.OpenAIWS.OAuthMaxConnsFactor = 2
pool := newOpenAIWSConnPool(cfg)
pool.setClientDialerForTest(&openAIWSFakeDialer{})
t.Cleanup(pool.Close)
req := openAIWSAcquireRequest{
Account: &Account{ID: 902, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 2},
WSURL: "wss://example.com/v1/responses",
}
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
leases := make([]*openAIWSConnLease, 0, 3)
for range 3 {
lease, err := pool.Acquire(ctx, req)
require.NoError(t, err, "轮次之间仍持有连接的会话应可使用系数扩出的容量")
t.Cleanup(lease.Release)
leases = append(leases, lease)
}
high := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 20}
require.Equal(t, 8, pool.effectiveMaxConnsByAccount(high), "v2 路径也必须受连接池硬上限约束")
fullCtx, fullCancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
defer fullCancel()
_, err := pool.Acquire(fullCtx, req)
require.ErrorIs(t, err, context.DeadlineExceeded, "扩容后仍必须受全局硬上限约束")
nonPositive := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 0}
require.Equal(t, 0, pool.effectiveMaxConnsByAccount(nonPositive), "并发数<=0 时应不可调度")
leases[0].Release()
reused, err := pool.Acquire(ctx, req)
require.NoError(t, err)
defer reused.Release()
require.Equal(t, leases[0].ConnID(), reused.ConnID())
require.True(t, reused.Reused())
})
}
}
func TestOpenAIWSConnPool_AcquireRejectsWhenEffectiveMaxConnsIsZero(t *testing.T) {
+2
View File
@@ -357,6 +357,8 @@ gateway:
max_idle_per_account: 12
# 是否按账号并发动态计算连接池上限:
# effective_max_conns = min(max_conns_per_account, ceil(account.concurrency * factor))
# 旧版及 mode_router_v2 的 ctx_pool 共用此开关和下方类型系数;关闭后使用 max_conns_per_account。
# mode_router_v2 下并发数 <= 0 的账号仍不可调度。
dynamic_max_conns_by_account_concurrency_enabled: true
# 按账号类型分别设置系数(OAuth / API Key)
oauth_max_conns_factor: 1.0
@@ -613,7 +613,7 @@ export default {
wsModeShared: 'Shared (shared)',
wsModeDedicated: 'Dedicated (dedicated)',
wsModeConcurrencyHint:
'When WS mode is enabled, account concurrency becomes the WS connection pool limit for this account.',
'Account concurrency limits requests in flight; the context pool connection limit is controlled by the gateway connection pool settings.',
wsModePassthroughHint: 'Passthrough mode does not use the WS connection pool.',
oauthResponsesWebsocketsV2: 'OAuth WebSocket Mode',
oauthResponsesWebsocketsV2Desc:
@@ -698,7 +698,7 @@ export default {
wsModeHttpBridge: 'HTTP 桥接(http_bridge)',
wsModeShared: '共享(shared)',
wsModeDedicated: '独享(dedicated)',
wsModeConcurrencyHint: '启用 WS mode 后,该账号并发数将作为该账号 WS 连接池上限。',
wsModeConcurrencyHint: '账号并发数限制同时执行的请求数;上下文池的连接上限由网关连接池配置决定。',
wsModePassthroughHint: 'passthrough 模式不使用 WS 连接池。',
oauthResponsesWebsocketsV2: 'OAuth WebSocket Mode',
oauthResponsesWebsocketsV2Desc: