fix: preserve openai ws flags in scheduler cache
This commit is contained in:
@@ -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 {
|
||||
|
||||
33
backend/internal/repository/scheduler_cache_unit_test.go
Normal file
33
backend/internal/repository/scheduler_cache_unit_test.go
Normal file
@@ -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"])
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
Reference in New Issue
Block a user