Skip to content

Commit e6e73b4

Browse files
Wei-ShawWei-Shaw
authored
Merge pull request #1690 from KnowSky404/fix/ws-codex-scheduler-cache-1662
fix: preserve openai ws flags in scheduler cache
2 parents 7ea8e7e + 836092a commit e6e73b4

6 files changed

Lines changed: 108 additions & 6 deletions

File tree

‎backend/internal/repository/scheduler_cache.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -426,6 +426,13 @@ func filterSchedulerExtra(extra map[string]any) map[string]any {
426426
"window_cost_sticky_reserve",
427427
"max_sessions",
428428
"session_idle_timeout_minutes",
429+
"openai_oauth_responses_websockets_v2_enabled",
430+
"openai_oauth_responses_websockets_v2_mode",
431+
"openai_apikey_responses_websockets_v2_enabled",
432+
"openai_apikey_responses_websockets_v2_mode",
433+
"responses_websockets_v2_enabled",
434+
"openai_ws_enabled",
435+
"openai_ws_force_http",
429436
}
430437
filtered := make(map[string]any)
431438
for _, key := range keys {

‎backend/internal/repository/scheduler_cache_unit_test.go‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
//go:build unit
2+
3+
package repository
4+
5+
import (
6+
"testing"
7+
8+
"github.com/Wei-Shaw/sub2api/internal/service"
9+
"github.com/stretchr/testify/require"
10+
)
11+
12+
func TestBuildSchedulerMetadataAccount_KeepsOpenAIWSFlags(t *testing.T) {
13+
account := service.Account{
14+
ID: 42,
15+
Platform: service.PlatformOpenAI,
16+
Type: service.AccountTypeOAuth,
17+
Extra: map[string]any{
18+
"openai_oauth_responses_websockets_v2_enabled": true,
19+
"openai_oauth_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough,
20+
"openai_ws_force_http": true,
21+
"mixed_scheduling": true,
22+
"unused_large_field": "drop-me",
23+
},
24+
}
25+
26+
got := buildSchedulerMetadataAccount(account)
27+
28+
require.Equal(t, true, got.Extra["openai_oauth_responses_websockets_v2_enabled"])
29+
require.Equal(t, service.OpenAIWSIngressModePassthrough, got.Extra["openai_oauth_responses_websockets_v2_mode"])
30+
require.Equal(t, true, got.Extra["openai_ws_force_http"])
31+
require.Equal(t, true, got.Extra["mixed_scheduling"])
32+
require.Nil(t, got.Extra["unused_large_field"])
33+
}

‎backend/internal/service/openai_account_scheduler_ws_snapshot_test.go‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
//go:build unit
2+
3+
package service
4+
5+
import (
6+
"context"
7+
"testing"
8+
9+
"github.com/Wei-Shaw/sub2api/internal/config"
10+
"github.com/stretchr/testify/require"
11+
)
12+
13+
func TestOpenAIGatewayService_SelectAccountWithScheduler_UsesWSPassthroughSnapshotFlags(t *testing.T) {
14+
ctx := context.Background()
15+
groupID := int64(10105)
16+
account := &Account{
17+
ID: 35001,
18+
Platform: PlatformOpenAI,
19+
Type: AccountTypeOAuth,
20+
Status: StatusActive,
21+
Schedulable: true,
22+
Concurrency: 10,
23+
Extra: map[string]any{
24+
"openai_oauth_responses_websockets_v2_mode": OpenAIWSIngressModePassthrough,
25+
},
26+
}
27+
28+
snapshotCache := &openAISnapshotCacheStub{
29+
snapshotAccounts: []*Account{account},
30+
accountsByID: map[int64]*Account{account.ID: account},
31+
}
32+
cfg := &config.Config{}
33+
cfg.Gateway.OpenAIWS.Enabled = true
34+
cfg.Gateway.OpenAIWS.OAuthEnabled = true
35+
cfg.Gateway.OpenAIWS.APIKeyEnabled = true
36+
cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
37+
cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
38+
cfg.Gateway.OpenAIWS.IngressModeDefault = OpenAIWSIngressModeCtxPool
39+
40+
svc := &OpenAIGatewayService{
41+
accountRepo: stubOpenAIAccountRepo{accounts: []Account{*account}},
42+
cache: &stubGatewayCache{},
43+
cfg: cfg,
44+
schedulerSnapshot: &SchedulerSnapshotService{cache: snapshotCache},
45+
concurrencyService: NewConcurrencyService(stubConcurrencyCache{}),
46+
}
47+
48+
selection, decision, err := svc.SelectAccountWithScheduler(
49+
ctx,
50+
&groupID,
51+
"",
52+
"session_hash_ws_passthrough",
53+
"gpt-5.1",
54+
nil,
55+
OpenAIUpstreamTransportResponsesWebsocketV2,
56+
)
57+
require.NoError(t, err)
58+
require.NotNil(t, selection)
59+
require.NotNil(t, selection.Account)
60+
require.Equal(t, account.ID, selection.Account.ID)
61+
require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
62+
}

‎frontend/src/components/account/BulkEditAccountModal.vue‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -921,6 +921,7 @@ import {
921921
getPresetMappingsByPlatform
922922
} from '@/composables/useModelWhitelist'
923923
import {
924+
OPENAI_WS_MODE_CTX_POOL,
924925
OPENAI_WS_MODE_OFF,
925926
OPENAI_WS_MODE_PASSTHROUGH,
926927
isOpenAIWSModeEnabled,
@@ -1069,6 +1070,7 @@ const isOpenAIModelRestrictionDisabled = computed(
10691070
10701071
const openAIWSModeOptions = computed(() => [
10711072
{ value: OPENAI_WS_MODE_OFF, label: t('admin.accounts.openai.wsModeOff') },
1073+
{ value: OPENAI_WS_MODE_CTX_POOL, label: t('admin.accounts.openai.wsModeCtxPool') },
10721074
{ value: OPENAI_WS_MODE_PASSTHROUGH, label: t('admin.accounts.openai.wsModePassthrough') }
10731075
])
10741076
const openAIWSModeConcurrencyHintKey = computed(() =>

‎frontend/src/components/account/CreateAccountModal.vue‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2932,7 +2932,7 @@ import { applyInterceptWarmup } from '@/components/account/credentialsBuilder'
29322932
import { formatDateTimeLocalInput, parseDateTimeLocalInput } from '@/utils/format'
29332933
import { createStableObjectKeyResolver } from '@/utils/stableObjectKey'
29342934
import {
2935-
// OPENAI_WS_MODE_CTX_POOL,
2935+
OPENAI_WS_MODE_CTX_POOL,
29362936
OPENAI_WS_MODE_OFF,
29372937
OPENAI_WS_MODE_PASSTHROUGH,
29382938
isOpenAIWSModeEnabled,
@@ -3180,8 +3180,7 @@ const geminiSelectedTier = computed(() => {
31803180
31813181
const openAIWSModeOptions = computed(() => [
31823182
{ value: OPENAI_WS_MODE_OFF, label: t('admin.accounts.openai.wsModeOff') },
3183-
// TODO: ctx_pool 选项暂时隐藏，待测试完成后恢复
3184-
// { value: OPENAI_WS_MODE_CTX_POOL, label: t('admin.accounts.openai.wsModeCtxPool') },
3183+
{ value: OPENAI_WS_MODE_CTX_POOL, label: t('admin.accounts.openai.wsModeCtxPool') },
31853184
{ value: OPENAI_WS_MODE_PASSTHROUGH, label: t('admin.accounts.openai.wsModePassthrough') }
31863185
])
31873186

‎frontend/src/components/account/EditAccountModal.vue‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1858,7 +1858,7 @@ import { applyInterceptWarmup } from '@/components/account/credentialsBuilder'
18581858
import { formatDateTimeLocalInput, parseDateTimeLocalInput } from '@/utils/format'
18591859
import { createStableObjectKeyResolver } from '@/utils/stableObjectKey'
18601860
import {
1861-
// OPENAI_WS_MODE_CTX_POOL,
1861+
OPENAI_WS_MODE_CTX_POOL,
18621862
OPENAI_WS_MODE_OFF,
18631863
OPENAI_WS_MODE_PASSTHROUGH,
18641864
isOpenAIWSModeEnabled,
@@ -2020,8 +2020,7 @@ const editWeeklyResetHour = ref<number | null>(null)
20202020
const editResetTimezone = ref<string | null>(null)
20212021
const openAIWSModeOptions = computed(() => [
20222022
{ value: OPENAI_WS_MODE_OFF, label: t('admin.accounts.openai.wsModeOff') },
2023-
// TODO: ctx_pool 选项暂时隐藏，待测试完成后恢复
2024-
// { value: OPENAI_WS_MODE_CTX_POOL, label: t('admin.accounts.openai.wsModeCtxPool') },
2023+
{ value: OPENAI_WS_MODE_CTX_POOL, label: t('admin.accounts.openai.wsModeCtxPool') },
20252024
{ value: OPENAI_WS_MODE_PASSTHROUGH, label: t('admin.accounts.openai.wsModePassthrough') }
20262025
])
20272026
const openaiResponsesWebSocketV2Mode = computed({

0 commit comments

Comments
 (0)