mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-11-02 00:30:32 +00:00
Accept service_tier=ultrafast on the OpenAI/Codex gateway, advertise it on GPT-5.6 Sol in the Codex model catalog, and include it in ChatGPT OAuth routing hints so Codex clients can select Ultrafast.
521 lines
21 KiB
Go
521 lines
21 KiB
Go
package service
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"testing"
|
||
|
||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||
"github.com/stretchr/testify/require"
|
||
"github.com/tidwall/gjson"
|
||
)
|
||
|
||
type openAIFastPolicyRepoStub struct {
|
||
values map[string]string
|
||
}
|
||
|
||
func (s *openAIFastPolicyRepoStub) Get(ctx context.Context, key string) (*Setting, error) {
|
||
panic("unexpected Get call")
|
||
}
|
||
|
||
func (s *openAIFastPolicyRepoStub) GetValue(ctx context.Context, key string) (string, error) {
|
||
if v, ok := s.values[key]; ok {
|
||
return v, nil
|
||
}
|
||
return "", ErrSettingNotFound
|
||
}
|
||
|
||
func (s *openAIFastPolicyRepoStub) Set(ctx context.Context, key, value string) error {
|
||
if s.values == nil {
|
||
s.values = map[string]string{}
|
||
}
|
||
s.values[key] = value
|
||
return nil
|
||
}
|
||
|
||
func (s *openAIFastPolicyRepoStub) GetMultiple(ctx context.Context, keys []string) (map[string]string, error) {
|
||
result := make(map[string]string, len(keys))
|
||
for _, key := range keys {
|
||
if value, ok := s.values[key]; ok {
|
||
result[key] = value
|
||
}
|
||
}
|
||
return result, nil
|
||
}
|
||
|
||
func (s *openAIFastPolicyRepoStub) SetMultiple(ctx context.Context, settings map[string]string) error {
|
||
panic("unexpected SetMultiple call")
|
||
}
|
||
|
||
func (s *openAIFastPolicyRepoStub) GetAll(ctx context.Context) (map[string]string, error) {
|
||
panic("unexpected GetAll call")
|
||
}
|
||
|
||
func (s *openAIFastPolicyRepoStub) Delete(ctx context.Context, key string) error {
|
||
panic("unexpected Delete call")
|
||
}
|
||
|
||
func newOpenAIGatewayServiceWithSettings(t *testing.T, settings *OpenAIFastPolicySettings) *OpenAIGatewayService {
|
||
t.Helper()
|
||
repo := &openAIFastPolicyRepoStub{values: map[string]string{}}
|
||
if settings != nil {
|
||
raw, err := json.Marshal(settings)
|
||
require.NoError(t, err)
|
||
repo.values[SettingKeyOpenAIFastPolicySettings] = string(raw)
|
||
}
|
||
return &OpenAIGatewayService{
|
||
settingService: NewSettingService(repo, &config.Config{}),
|
||
}
|
||
}
|
||
|
||
func openAIFastFilterPriorityPolicy() *OpenAIFastPolicySettings {
|
||
return &OpenAIFastPolicySettings{
|
||
Rules: []OpenAIFastPolicyRule{{
|
||
ServiceTier: OpenAIFastTierPriority,
|
||
Action: BetaPolicyActionFilter,
|
||
Scope: BetaPolicyScopeAll,
|
||
ModelWhitelist: []string{},
|
||
FallbackAction: BetaPolicyActionPass,
|
||
}},
|
||
}
|
||
}
|
||
|
||
func TestEvaluateOpenAIFastPolicy_DefaultPassesKnownTiers(t *testing.T) {
|
||
require.Empty(t, DefaultOpenAIFastPolicySettings().Rules, "default policy must not rewrite service_tier unless admin configured rules")
|
||
|
||
svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
|
||
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
||
|
||
action, _ := svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5", OpenAIFastTierPriority)
|
||
require.Equal(t, BetaPolicyActionPass, action)
|
||
|
||
action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5-turbo", OpenAIFastTierPriority)
|
||
require.Equal(t, BetaPolicyActionPass, action)
|
||
|
||
action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-4", OpenAIFastTierPriority)
|
||
require.Equal(t, BetaPolicyActionPass, action)
|
||
|
||
action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5", OpenAIFastTierFlex)
|
||
require.Equal(t, BetaPolicyActionPass, action)
|
||
|
||
action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.6-sol", OpenAIFastTierUltrafast)
|
||
require.Equal(t, BetaPolicyActionPass, action)
|
||
|
||
// empty tier → pass
|
||
action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5", "")
|
||
require.Equal(t, BetaPolicyActionPass, action)
|
||
}
|
||
|
||
func TestEvaluateOpenAIFastPolicy_BlockRuleCarriesMessage(t *testing.T) {
|
||
settings := &OpenAIFastPolicySettings{
|
||
Rules: []OpenAIFastPolicyRule{{
|
||
ServiceTier: OpenAIFastTierPriority,
|
||
Action: BetaPolicyActionBlock,
|
||
Scope: BetaPolicyScopeAll,
|
||
ErrorMessage: "fast mode is not allowed",
|
||
ModelWhitelist: []string{"gpt-5.5"},
|
||
FallbackAction: BetaPolicyActionPass,
|
||
}},
|
||
}
|
||
svc := newOpenAIGatewayServiceWithSettings(t, settings)
|
||
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
||
|
||
action, msg := svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5", OpenAIFastTierPriority)
|
||
require.Equal(t, BetaPolicyActionBlock, action)
|
||
require.Equal(t, "fast mode is not allowed", msg)
|
||
}
|
||
|
||
func TestEvaluateOpenAIFastPolicy_ScopeFiltersOAuth(t *testing.T) {
|
||
settings := &OpenAIFastPolicySettings{
|
||
Rules: []OpenAIFastPolicyRule{{
|
||
ServiceTier: OpenAIFastTierAny,
|
||
Action: BetaPolicyActionFilter,
|
||
Scope: BetaPolicyScopeOAuth,
|
||
}},
|
||
}
|
||
svc := newOpenAIGatewayServiceWithSettings(t, settings)
|
||
|
||
// OAuth account → rule matches
|
||
oauthAccount := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
||
action, _ := svc.evaluateOpenAIFastPolicy(context.Background(), oauthAccount, "gpt-4", OpenAIFastTierPriority)
|
||
require.Equal(t, BetaPolicyActionFilter, action)
|
||
|
||
// API Key account → rule skipped → pass
|
||
apiKeyAccount := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
||
action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), apiKeyAccount, "gpt-4", OpenAIFastTierPriority)
|
||
require.Equal(t, BetaPolicyActionPass, action)
|
||
}
|
||
|
||
func TestEvaluateOpenAIFastPolicy_UserScopedRuleOverridesGlobalRule(t *testing.T) {
|
||
settings := &OpenAIFastPolicySettings{
|
||
Rules: []OpenAIFastPolicyRule{
|
||
{
|
||
ServiceTier: OpenAIFastTierPriority,
|
||
Action: BetaPolicyActionFilter,
|
||
Scope: BetaPolicyScopeAll,
|
||
},
|
||
{
|
||
ServiceTier: OpenAIFastTierPriority,
|
||
Action: BetaPolicyActionPass,
|
||
Scope: BetaPolicyScopeAll,
|
||
UserIDs: []int64{42},
|
||
},
|
||
},
|
||
}
|
||
svc := newOpenAIGatewayServiceWithSettings(t, settings)
|
||
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
||
|
||
allowedUserCtx := context.WithValue(context.Background(), ctxkey.UserID, int64(42))
|
||
action, _ := svc.evaluateOpenAIFastPolicy(allowedUserCtx, account, "gpt-5.5", OpenAIFastTierPriority)
|
||
require.Equal(t, BetaPolicyActionPass, action)
|
||
|
||
otherUserCtx := context.WithValue(context.Background(), ctxkey.UserID, int64(43))
|
||
action, _ = svc.evaluateOpenAIFastPolicy(otherUserCtx, account, "gpt-5.5", OpenAIFastTierPriority)
|
||
require.Equal(t, BetaPolicyActionFilter, action)
|
||
|
||
action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5", OpenAIFastTierPriority)
|
||
require.Equal(t, BetaPolicyActionFilter, action)
|
||
}
|
||
|
||
func TestApplyOpenAIFastPolicyToBody_DefaultPassesPriorityAndFast(t *testing.T) {
|
||
svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
|
||
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
||
|
||
body := []byte(`{"model":"gpt-5.5","service_tier":"priority","messages":[]}`)
|
||
updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
|
||
require.NoError(t, err)
|
||
require.Equal(t, string(body), string(updated))
|
||
|
||
body = []byte(`{"model":"gpt-5.5","service_tier":"fast"}`)
|
||
updated, err = svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
|
||
require.NoError(t, err)
|
||
require.Equal(t, "priority", gjson.GetBytes(updated, "service_tier").String())
|
||
|
||
body = []byte(`{"model":"gpt-4","service_tier":"priority"}`)
|
||
updated, err = svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-4", body)
|
||
require.NoError(t, err)
|
||
require.Equal(t, string(body), string(updated))
|
||
|
||
// No service_tier → no-op
|
||
body = []byte(`{"model":"gpt-5.5"}`)
|
||
updated, err = svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
|
||
require.NoError(t, err)
|
||
require.Equal(t, string(body), string(updated))
|
||
}
|
||
|
||
func TestApplyOpenAIFastPolicyToBody_ExplicitFilterRemovesField(t *testing.T) {
|
||
svc := newOpenAIGatewayServiceWithSettings(t, openAIFastFilterPriorityPolicy())
|
||
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
||
|
||
body := []byte(`{"model":"gpt-5.5","service_tier":"priority","messages":[]}`)
|
||
updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
|
||
require.NoError(t, err)
|
||
require.NotContains(t, string(updated), `"service_tier"`)
|
||
|
||
body = []byte(`{"model":"gpt-5.5","service_tier":"fast"}`)
|
||
updated, err = svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
|
||
require.NoError(t, err)
|
||
require.NotContains(t, string(updated), `"service_tier"`)
|
||
}
|
||
|
||
func TestApplyOpenAIFastPolicyToBody_UserScopedRuleOverridesGlobalRule(t *testing.T) {
|
||
settings := &OpenAIFastPolicySettings{
|
||
Rules: []OpenAIFastPolicyRule{
|
||
{
|
||
ServiceTier: OpenAIFastTierPriority,
|
||
Action: BetaPolicyActionFilter,
|
||
Scope: BetaPolicyScopeAll,
|
||
},
|
||
{
|
||
ServiceTier: OpenAIFastTierPriority,
|
||
Action: BetaPolicyActionPass,
|
||
Scope: BetaPolicyScopeAll,
|
||
UserIDs: []int64{42},
|
||
},
|
||
},
|
||
}
|
||
svc := newOpenAIGatewayServiceWithSettings(t, settings)
|
||
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
||
body := []byte(`{"model":"gpt-5.5","service_tier":"priority"}`)
|
||
|
||
allowedUserCtx := context.WithValue(context.Background(), ctxkey.UserID, int64(42))
|
||
updated, err := svc.applyOpenAIFastPolicyToBody(allowedUserCtx, account, "gpt-5.5", body)
|
||
require.NoError(t, err)
|
||
require.Equal(t, "priority", gjson.GetBytes(updated, "service_tier").String())
|
||
|
||
otherUserCtx := context.WithValue(context.Background(), ctxkey.UserID, int64(43))
|
||
updated, err = svc.applyOpenAIFastPolicyToBody(otherUserCtx, account, "gpt-5.5", body)
|
||
require.NoError(t, err)
|
||
require.NotContains(t, string(updated), `"service_tier"`)
|
||
}
|
||
|
||
func TestApplyOpenAIFastPolicyToBody_PriorityFilterLeavesUltrafast(t *testing.T) {
|
||
svc := newOpenAIGatewayServiceWithSettings(t, openAIFastFilterPriorityPolicy())
|
||
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
||
body := []byte(`{"model":"gpt-5.6-sol","service_tier":"ultrafast"}`)
|
||
|
||
updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.6-sol", body)
|
||
require.NoError(t, err)
|
||
require.Equal(t, OpenAIFastTierUltrafast, gjson.GetBytes(updated, "service_tier").String())
|
||
}
|
||
|
||
func TestApplyOpenAIFastPolicyToBody_ForcePriorityRewritesKnownTier(t *testing.T) {
|
||
settings := &OpenAIFastPolicySettings{
|
||
Rules: []OpenAIFastPolicyRule{{
|
||
ServiceTier: OpenAIFastTierAny,
|
||
Action: OpenAIFastPolicyActionForcePriority,
|
||
Scope: BetaPolicyScopeAll,
|
||
}},
|
||
}
|
||
svc := newOpenAIGatewayServiceWithSettings(t, settings)
|
||
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
||
|
||
for _, tier := range []string{"flex", "auto", "default", "scale", "fast", "priority", "ultrafast"} {
|
||
body := []byte(`{"model":"gpt-5.5","service_tier":"` + tier + `"}`)
|
||
updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
|
||
require.NoError(t, err)
|
||
require.Equal(t, OpenAIFastTierPriority, gjson.GetBytes(updated, "service_tier").String(),
|
||
"tier %q should be forced to priority", tier)
|
||
}
|
||
}
|
||
|
||
func TestApplyOpenAIFastPolicyToBody_GroupForceInjectsAndOverridesTier(t *testing.T) {
|
||
svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
|
||
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
||
ctx := context.WithValue(context.Background(), ctxkey.Group, &Group{
|
||
ID: 7, Platform: PlatformOpenAI, Status: StatusActive, Hydrated: true, ForceOpenAIFast: true,
|
||
})
|
||
|
||
for _, body := range [][]byte{
|
||
[]byte(`{"model":"gpt-5.6-sol","input":"hi"}`),
|
||
[]byte(`{"model":"gpt-5.6-sol","service_tier":"default"}`),
|
||
[]byte(`{"model":"gpt-5.6-sol","service_tier":"flex"}`),
|
||
[]byte(`{"model":"gpt-5.6-sol","service_tier":"client-unknown"}`),
|
||
} {
|
||
updated, err := svc.applyOpenAIFastPolicyToBody(ctx, account, "gpt-5.6-sol", body)
|
||
require.NoError(t, err)
|
||
require.Equal(t, OpenAIFastTierPriority, gjson.GetBytes(updated, "service_tier").String())
|
||
}
|
||
}
|
||
|
||
func TestApplyOpenAIFastPolicyToBody_GroupForceStillHonorsGlobalPolicy(t *testing.T) {
|
||
svc := newOpenAIGatewayServiceWithSettings(t, openAIFastFilterPriorityPolicy())
|
||
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
||
ctx := context.WithValue(context.Background(), ctxkey.Group, &Group{
|
||
ID: 7, Platform: PlatformOpenAI, Status: StatusActive, Hydrated: true, ForceOpenAIFast: true,
|
||
})
|
||
|
||
updated, err := svc.applyOpenAIFastPolicyToBody(ctx, account, "gpt-5.6-sol", []byte(`{"model":"gpt-5.6-sol"}`))
|
||
require.NoError(t, err)
|
||
require.False(t, gjson.GetBytes(updated, "service_tier").Exists(),
|
||
"the global filter remains authoritative after the group requests priority")
|
||
}
|
||
|
||
func TestApplyOpenAIFastPolicyToBody_GroupForceRequiresHydratedGroup(t *testing.T) {
|
||
svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
|
||
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
||
body := []byte(`{"model":"gpt-5.6-sol"}`)
|
||
ctx := context.WithValue(context.Background(), ctxkey.Group, &Group{
|
||
ID: 7, Platform: PlatformOpenAI, Status: StatusActive, ForceOpenAIFast: true,
|
||
})
|
||
|
||
updated, err := svc.applyOpenAIFastPolicyToBody(ctx, account, "gpt-5.6-sol", body)
|
||
require.NoError(t, err)
|
||
require.Equal(t, string(body), string(updated))
|
||
}
|
||
|
||
func TestApplyOpenAIFastPolicyToBody_GroupForceOnlyTargetsOpenAIAccounts(t *testing.T) {
|
||
svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
|
||
body := []byte(`{"model":"grok-4.1"}`)
|
||
ctx := context.WithValue(context.Background(), ctxkey.Group, &Group{
|
||
ID: 7, Platform: PlatformComposite, Status: StatusActive, Hydrated: true, ForceOpenAIFast: true,
|
||
})
|
||
|
||
updated, err := svc.applyOpenAIFastPolicyToBody(
|
||
ctx,
|
||
&Account{Platform: PlatformGrok, Type: AccountTypeOAuth},
|
||
"grok-4.1",
|
||
body,
|
||
)
|
||
require.NoError(t, err)
|
||
require.Equal(t, string(body), string(updated))
|
||
}
|
||
|
||
func TestApplyOpenAIFastPolicyToBody_GroupForceRequiresSupportedGroupPlatform(t *testing.T) {
|
||
svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
|
||
body := []byte(`{"model":"gpt-5.6-sol"}`)
|
||
ctx := context.WithValue(context.Background(), ctxkey.Group, &Group{
|
||
ID: 7, Platform: PlatformAnthropic, Status: StatusActive, Hydrated: true, ForceOpenAIFast: true,
|
||
})
|
||
|
||
updated, err := svc.applyOpenAIFastPolicyToBody(
|
||
ctx,
|
||
&Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth},
|
||
"gpt-5.6-sol",
|
||
body,
|
||
)
|
||
require.NoError(t, err)
|
||
require.Equal(t, string(body), string(updated))
|
||
}
|
||
|
||
// TestApplyOpenAIFastPolicyToBody_OfficialTiersBypassDefaultRule 验证默认配置
|
||
// 下客户端显式发送的 OpenAI 官方合法 tier 能透传到上游而不被静默剥离。
|
||
func TestApplyOpenAIFastPolicyToBody_OfficialTiersBypassDefaultRule(t *testing.T) {
|
||
svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
|
||
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
||
|
||
for _, tier := range []string{"auto", "default", "scale"} {
|
||
body := []byte(`{"model":"gpt-5.5","service_tier":"` + tier + `"}`)
|
||
updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
|
||
require.NoError(t, err, "tier %q should pass without error", tier)
|
||
require.Contains(t, string(updated), `"service_tier":"`+tier+`"`,
|
||
"tier %q should be preserved in body under default policy", tier)
|
||
}
|
||
|
||
// evaluate 层也应判定为 pass(默认配置没有内置规则)。
|
||
for _, tier := range []string{"auto", "default", "scale"} {
|
||
action, _ := svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5", tier)
|
||
require.Equal(t, BetaPolicyActionPass, action, "tier %q should evaluate to pass", tier)
|
||
}
|
||
}
|
||
|
||
// TestApplyOpenAIFastPolicyToBody_AllRuleStripsOfficialTiers 验证管理员显式配置
|
||
// ServiceTier=all + Action=filter 规则后,auto/default/scale 等官方 tier 也会
|
||
// 被剥离。这是符合预期的——首条匹配 short-circuit,"all" 覆盖任意已识别 tier。
|
||
func TestApplyOpenAIFastPolicyToBody_AllRuleStripsOfficialTiers(t *testing.T) {
|
||
settings := &OpenAIFastPolicySettings{
|
||
Rules: []OpenAIFastPolicyRule{{
|
||
ServiceTier: OpenAIFastTierAny,
|
||
Action: BetaPolicyActionFilter,
|
||
Scope: BetaPolicyScopeAll,
|
||
}},
|
||
}
|
||
svc := newOpenAIGatewayServiceWithSettings(t, settings)
|
||
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
||
|
||
for _, tier := range []string{"auto", "default", "scale", "priority", "flex"} {
|
||
body := []byte(`{"model":"gpt-5.5","service_tier":"` + tier + `"}`)
|
||
updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
|
||
require.NoError(t, err)
|
||
require.NotContains(t, string(updated), `"service_tier"`,
|
||
"tier %q should be stripped under ServiceTier=all + filter rule", tier)
|
||
}
|
||
}
|
||
|
||
// TestApplyOpenAIFastPolicyToBody_UnknownTierStripped 验证真未知 tier 仍被剥离
|
||
// (normalize 返回 nil → normalizeResponsesBodyServiceTier 删除字段;
|
||
// applyOpenAIFastPolicyToBody 在 normTier 为空时直接 no-op,因为字段已不可能存在
|
||
// 于经过前置归一化的请求里。这里直接调 apply 验证它对未识别值不会异常)。
|
||
func TestApplyOpenAIFastPolicyToBody_UnknownTierStripped(t *testing.T) {
|
||
svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
|
||
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
||
|
||
// normalize 阶段会将未知值剥离
|
||
require.Nil(t, normalizeOpenAIServiceTier("xxx"))
|
||
|
||
// applyOpenAIFastPolicyToBody 收到未识别 tier 时不报错,body 透传不变
|
||
// (不属于本函数职责——上层 normalizeResponsesBodyServiceTier 已剥离)
|
||
body := []byte(`{"model":"gpt-5.5","service_tier":"xxx"}`)
|
||
updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
|
||
require.NoError(t, err)
|
||
require.Equal(t, string(body), string(updated))
|
||
}
|
||
|
||
func TestApplyOpenAIFastPolicyToBody_BlockReturnsTypedError(t *testing.T) {
|
||
settings := &OpenAIFastPolicySettings{
|
||
Rules: []OpenAIFastPolicyRule{{
|
||
ServiceTier: OpenAIFastTierPriority,
|
||
Action: BetaPolicyActionBlock,
|
||
Scope: BetaPolicyScopeAll,
|
||
ErrorMessage: "fast mode is blocked for gpt-5.5",
|
||
ModelWhitelist: []string{"gpt-5.5"},
|
||
FallbackAction: BetaPolicyActionPass,
|
||
}},
|
||
}
|
||
svc := newOpenAIGatewayServiceWithSettings(t, settings)
|
||
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
||
|
||
body := []byte(`{"model":"gpt-5.5","service_tier":"priority"}`)
|
||
updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
|
||
require.Error(t, err)
|
||
var blocked *OpenAIFastBlockedError
|
||
require.True(t, errors.As(err, &blocked))
|
||
require.Contains(t, blocked.Message, "fast mode is blocked")
|
||
require.Equal(t, string(body), string(updated)) // body not mutated on block
|
||
}
|
||
|
||
func TestSetOpenAIFastPolicySettings_Validation(t *testing.T) {
|
||
repo := &openAIFastPolicyRepoStub{values: map[string]string{}}
|
||
svc := NewSettingService(repo, &config.Config{})
|
||
|
||
// Invalid action rejected
|
||
err := svc.SetOpenAIFastPolicySettings(context.Background(), &OpenAIFastPolicySettings{
|
||
Rules: []OpenAIFastPolicyRule{{
|
||
ServiceTier: OpenAIFastTierPriority,
|
||
Action: "bogus",
|
||
Scope: BetaPolicyScopeAll,
|
||
}},
|
||
})
|
||
require.Error(t, err)
|
||
|
||
// Invalid service_tier rejected
|
||
err = svc.SetOpenAIFastPolicySettings(context.Background(), &OpenAIFastPolicySettings{
|
||
Rules: []OpenAIFastPolicyRule{{
|
||
ServiceTier: "turbo",
|
||
Action: BetaPolicyActionPass,
|
||
Scope: BetaPolicyScopeAll,
|
||
}},
|
||
})
|
||
require.Error(t, err)
|
||
|
||
// Non-positive and duplicate user IDs are rejected.
|
||
err = svc.SetOpenAIFastPolicySettings(context.Background(), &OpenAIFastPolicySettings{
|
||
Rules: []OpenAIFastPolicyRule{{
|
||
ServiceTier: OpenAIFastTierPriority,
|
||
Action: BetaPolicyActionPass,
|
||
Scope: BetaPolicyScopeAll,
|
||
UserIDs: []int64{0},
|
||
}},
|
||
})
|
||
require.Error(t, err)
|
||
|
||
err = svc.SetOpenAIFastPolicySettings(context.Background(), &OpenAIFastPolicySettings{
|
||
Rules: []OpenAIFastPolicyRule{{
|
||
ServiceTier: OpenAIFastTierPriority,
|
||
Action: BetaPolicyActionPass,
|
||
Scope: BetaPolicyScopeAll,
|
||
UserIDs: []int64{42, 42},
|
||
}},
|
||
})
|
||
require.Error(t, err)
|
||
|
||
// Valid settings persisted
|
||
err = svc.SetOpenAIFastPolicySettings(context.Background(), &OpenAIFastPolicySettings{
|
||
Rules: []OpenAIFastPolicyRule{
|
||
{
|
||
ServiceTier: OpenAIFastTierPriority,
|
||
Action: OpenAIFastPolicyActionForcePriority,
|
||
Scope: BetaPolicyScopeAll,
|
||
UserIDs: []int64{42, 43},
|
||
},
|
||
{
|
||
ServiceTier: OpenAIFastTierUltrafast,
|
||
Action: BetaPolicyActionPass,
|
||
Scope: BetaPolicyScopeAll,
|
||
},
|
||
},
|
||
})
|
||
require.NoError(t, err)
|
||
|
||
got, err := svc.GetOpenAIFastPolicySettings(context.Background())
|
||
require.NoError(t, err)
|
||
require.Len(t, got.Rules, 2)
|
||
require.Equal(t, OpenAIFastTierPriority, got.Rules[0].ServiceTier)
|
||
require.Equal(t, OpenAIFastPolicyActionForcePriority, got.Rules[0].Action)
|
||
require.Equal(t, []int64{42, 43}, got.Rules[0].UserIDs)
|
||
require.Equal(t, OpenAIFastTierUltrafast, got.Rules[1].ServiceTier)
|
||
require.Equal(t, BetaPolicyActionPass, got.Rules[1].Action)
|
||
}
|