diff --git a/backend/internal/repository/scheduler_cache.go b/backend/internal/repository/scheduler_cache.go index e9be8c7a..add0e501 100644 --- a/backend/internal/repository/scheduler_cache.go +++ b/backend/internal/repository/scheduler_cache.go @@ -426,6 +426,13 @@ func filterSchedulerExtra(extra map[string]any) map[string]any { "window_cost_sticky_reserve", "max_sessions", "session_idle_timeout_minutes", + "openai_oauth_responses_websockets_v2_enabled", + "openai_oauth_responses_websockets_v2_mode", + "openai_apikey_responses_websockets_v2_enabled", + "openai_apikey_responses_websockets_v2_mode", + "responses_websockets_v2_enabled", + "openai_ws_enabled", + "openai_ws_force_http", } filtered := make(map[string]any) for _, key := range keys { diff --git a/backend/internal/repository/scheduler_cache_unit_test.go b/backend/internal/repository/scheduler_cache_unit_test.go new file mode 100644 index 00000000..bcfd0e7a --- /dev/null +++ b/backend/internal/repository/scheduler_cache_unit_test.go @@ -0,0 +1,33 @@ +//go:build unit + +package repository + +import ( + "testing" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func TestBuildSchedulerMetadataAccount_KeepsOpenAIWSFlags(t *testing.T) { + account := service.Account{ + ID: 42, + Platform: service.PlatformOpenAI, + Type: service.AccountTypeOAuth, + Extra: map[string]any{ + "openai_oauth_responses_websockets_v2_enabled": true, + "openai_oauth_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough, + "openai_ws_force_http": true, + "mixed_scheduling": true, + "unused_large_field": "drop-me", + }, + } + + got := buildSchedulerMetadataAccount(account) + + require.Equal(t, true, got.Extra["openai_oauth_responses_websockets_v2_enabled"]) + require.Equal(t, service.OpenAIWSIngressModePassthrough, got.Extra["openai_oauth_responses_websockets_v2_mode"]) + require.Equal(t, true, got.Extra["openai_ws_force_http"]) + require.Equal(t, true, got.Extra["mixed_scheduling"]) + require.Nil(t, got.Extra["unused_large_field"]) +} diff --git a/backend/internal/service/openai_account_scheduler_ws_snapshot_test.go b/backend/internal/service/openai_account_scheduler_ws_snapshot_test.go new file mode 100644 index 00000000..c5de8203 --- /dev/null +++ b/backend/internal/service/openai_account_scheduler_ws_snapshot_test.go @@ -0,0 +1,62 @@ +//go:build unit + +package service + +import ( + "context" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +func TestOpenAIGatewayService_SelectAccountWithScheduler_UsesWSPassthroughSnapshotFlags(t *testing.T) { + ctx := context.Background() + groupID := int64(10105) + account := &Account{ + ID: 35001, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 10, + Extra: map[string]any{ + "openai_oauth_responses_websockets_v2_mode": OpenAIWSIngressModePassthrough, + }, + } + + snapshotCache := &openAISnapshotCacheStub{ + snapshotAccounts: []*Account{account}, + accountsByID: map[int64]*Account{account.ID: account}, + } + cfg := &config.Config{} + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.OAuthEnabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true + cfg.Gateway.OpenAIWS.IngressModeDefault = OpenAIWSIngressModeCtxPool + + svc := &OpenAIGatewayService{ + accountRepo: stubOpenAIAccountRepo{accounts: []Account{*account}}, + cache: &stubGatewayCache{}, + cfg: cfg, + schedulerSnapshot: &SchedulerSnapshotService{cache: snapshotCache}, + concurrencyService: NewConcurrencyService(stubConcurrencyCache{}), + } + + selection, decision, err := svc.SelectAccountWithScheduler( + ctx, + &groupID, + "", + "session_hash_ws_passthrough", + "gpt-5.1", + nil, + OpenAIUpstreamTransportResponsesWebsocketV2, + ) + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, account.ID, selection.Account.ID) + require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer) +}