fix: allow grok scheduler capability routing

This commit is contained in:
Heatherm Huang
2026-06-26 10:36:09 +08:00
parent e7a4f3f465
commit 1b9645ca32
3 changed files with 57 additions and 0 deletions
+3
View File
@@ -1258,6 +1258,9 @@ func (a *Account) openAIEndpointCapabilitySet() (map[string]bool, bool) {
}
func (a *Account) SupportsOpenAIImageCapability(capability OpenAIImagesCapability) bool {
if capability == "" {
return true
}
if !a.IsOpenAI() {
return false
}
@@ -475,6 +475,50 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_DefaultDisabled_Embeddi
require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
}
func TestOpenAIGatewayService_SelectAccountWithScheduler_DefaultDisabled_AllowsGrokChatAccount(t *testing.T) {
resetOpenAIAdvancedSchedulerSettingCacheForTest()
ctx := context.Background()
groupID := int64(10113)
accounts := []Account{
{
ID: 36041,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Priority: 0,
},
}
cfg := &config.Config{}
cfg.Gateway.Scheduling.LoadBatchEnabled = false
svc := &OpenAIGatewayService{
accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
cache: &schedulerTestGatewayCache{},
cfg: cfg,
concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
}
selection, decision, err := svc.SelectAccountWithSchedulerForCapability(
ctx,
&groupID,
"",
"",
"grok-4.3",
nil,
OpenAIUpstreamTransportAny,
OpenAIEndpointCapabilityChatCompletions,
false,
PlatformGrok,
)
require.NoError(t, err)
require.NotNil(t, selection)
require.NotNil(t, selection.Account)
require.Equal(t, int64(36041), selection.Account.ID)
require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
}
func TestOpenAIGatewayService_SelectAccountWithScheduler_EnabledUsesAdvancedPreviousResponseRouting(t *testing.T) {
resetOpenAIAdvancedSchedulerSettingCacheForTest()
@@ -442,6 +442,16 @@ func TestAccountSupportsOpenAIImageCapability_OAuthSupportsNative(t *testing.T)
require.True(t, account.SupportsOpenAIImageCapability(OpenAIImagesCapabilityNative))
}
func TestAccountSupportsOpenAIImageCapability_EmptyRequirementDoesNotRejectGrok(t *testing.T) {
account := &Account{
Platform: PlatformGrok,
Type: AccountTypeOAuth,
}
require.True(t, account.SupportsOpenAIImageCapability(""))
require.False(t, account.SupportsOpenAIImageCapability(OpenAIImagesCapabilityBasic))
}
func TestAccountSupportsOpenAIEndpointCapability(t *testing.T) {
t.Run("OpenAI APIKey 默认兼容 chat 和 embeddings", func(t *testing.T) {
account := &Account{