openai_fast_policy_test.go 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520
  1. package service
  2. import (
  3. "context"
  4. "encoding/json"
  5. "errors"
  6. "testing"
  7. "github.com/Wei-Shaw/sub2api/internal/config"
  8. "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
  9. "github.com/stretchr/testify/require"
  10. "github.com/tidwall/gjson"
  11. )
  12. type openAIFastPolicyRepoStub struct {
  13. values map[string]string
  14. }
  15. func (s *openAIFastPolicyRepoStub) Get(ctx context.Context, key string) (*Setting, error) {
  16. panic("unexpected Get call")
  17. }
  18. func (s *openAIFastPolicyRepoStub) GetValue(ctx context.Context, key string) (string, error) {
  19. if v, ok := s.values[key]; ok {
  20. return v, nil
  21. }
  22. return "", ErrSettingNotFound
  23. }
  24. func (s *openAIFastPolicyRepoStub) Set(ctx context.Context, key, value string) error {
  25. if s.values == nil {
  26. s.values = map[string]string{}
  27. }
  28. s.values[key] = value
  29. return nil
  30. }
  31. func (s *openAIFastPolicyRepoStub) GetMultiple(ctx context.Context, keys []string) (map[string]string, error) {
  32. result := make(map[string]string, len(keys))
  33. for _, key := range keys {
  34. if value, ok := s.values[key]; ok {
  35. result[key] = value
  36. }
  37. }
  38. return result, nil
  39. }
  40. func (s *openAIFastPolicyRepoStub) SetMultiple(ctx context.Context, settings map[string]string) error {
  41. panic("unexpected SetMultiple call")
  42. }
  43. func (s *openAIFastPolicyRepoStub) GetAll(ctx context.Context) (map[string]string, error) {
  44. panic("unexpected GetAll call")
  45. }
  46. func (s *openAIFastPolicyRepoStub) Delete(ctx context.Context, key string) error {
  47. panic("unexpected Delete call")
  48. }
  49. func newOpenAIGatewayServiceWithSettings(t *testing.T, settings *OpenAIFastPolicySettings) *OpenAIGatewayService {
  50. t.Helper()
  51. repo := &openAIFastPolicyRepoStub{values: map[string]string{}}
  52. if settings != nil {
  53. raw, err := json.Marshal(settings)
  54. require.NoError(t, err)
  55. repo.values[SettingKeyOpenAIFastPolicySettings] = string(raw)
  56. }
  57. return &OpenAIGatewayService{
  58. settingService: NewSettingService(repo, &config.Config{}),
  59. }
  60. }
  61. func openAIFastFilterPriorityPolicy() *OpenAIFastPolicySettings {
  62. return &OpenAIFastPolicySettings{
  63. Rules: []OpenAIFastPolicyRule{{
  64. ServiceTier: OpenAIFastTierPriority,
  65. Action: BetaPolicyActionFilter,
  66. Scope: BetaPolicyScopeAll,
  67. ModelWhitelist: []string{},
  68. FallbackAction: BetaPolicyActionPass,
  69. }},
  70. }
  71. }
  72. func TestEvaluateOpenAIFastPolicy_DefaultPassesKnownTiers(t *testing.T) {
  73. require.Empty(t, DefaultOpenAIFastPolicySettings().Rules, "default policy must not rewrite service_tier unless admin configured rules")
  74. svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
  75. account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  76. action, _ := svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5", OpenAIFastTierPriority)
  77. require.Equal(t, BetaPolicyActionPass, action)
  78. action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5-turbo", OpenAIFastTierPriority)
  79. require.Equal(t, BetaPolicyActionPass, action)
  80. action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-4", OpenAIFastTierPriority)
  81. require.Equal(t, BetaPolicyActionPass, action)
  82. action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5", OpenAIFastTierFlex)
  83. require.Equal(t, BetaPolicyActionPass, action)
  84. action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.6-sol", OpenAIFastTierUltrafast)
  85. require.Equal(t, BetaPolicyActionPass, action)
  86. // empty tier → pass
  87. action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5", "")
  88. require.Equal(t, BetaPolicyActionPass, action)
  89. }
  90. func TestEvaluateOpenAIFastPolicy_BlockRuleCarriesMessage(t *testing.T) {
  91. settings := &OpenAIFastPolicySettings{
  92. Rules: []OpenAIFastPolicyRule{{
  93. ServiceTier: OpenAIFastTierPriority,
  94. Action: BetaPolicyActionBlock,
  95. Scope: BetaPolicyScopeAll,
  96. ErrorMessage: "fast mode is not allowed",
  97. ModelWhitelist: []string{"gpt-5.5"},
  98. FallbackAction: BetaPolicyActionPass,
  99. }},
  100. }
  101. svc := newOpenAIGatewayServiceWithSettings(t, settings)
  102. account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  103. action, msg := svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5", OpenAIFastTierPriority)
  104. require.Equal(t, BetaPolicyActionBlock, action)
  105. require.Equal(t, "fast mode is not allowed", msg)
  106. }
  107. func TestEvaluateOpenAIFastPolicy_ScopeFiltersOAuth(t *testing.T) {
  108. settings := &OpenAIFastPolicySettings{
  109. Rules: []OpenAIFastPolicyRule{{
  110. ServiceTier: OpenAIFastTierAny,
  111. Action: BetaPolicyActionFilter,
  112. Scope: BetaPolicyScopeOAuth,
  113. }},
  114. }
  115. svc := newOpenAIGatewayServiceWithSettings(t, settings)
  116. // OAuth account → rule matches
  117. oauthAccount := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  118. action, _ := svc.evaluateOpenAIFastPolicy(context.Background(), oauthAccount, "gpt-4", OpenAIFastTierPriority)
  119. require.Equal(t, BetaPolicyActionFilter, action)
  120. // API Key account → rule skipped → pass
  121. apiKeyAccount := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  122. action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), apiKeyAccount, "gpt-4", OpenAIFastTierPriority)
  123. require.Equal(t, BetaPolicyActionPass, action)
  124. }
  125. func TestEvaluateOpenAIFastPolicy_UserScopedRuleOverridesGlobalRule(t *testing.T) {
  126. settings := &OpenAIFastPolicySettings{
  127. Rules: []OpenAIFastPolicyRule{
  128. {
  129. ServiceTier: OpenAIFastTierPriority,
  130. Action: BetaPolicyActionFilter,
  131. Scope: BetaPolicyScopeAll,
  132. },
  133. {
  134. ServiceTier: OpenAIFastTierPriority,
  135. Action: BetaPolicyActionPass,
  136. Scope: BetaPolicyScopeAll,
  137. UserIDs: []int64{42},
  138. },
  139. },
  140. }
  141. svc := newOpenAIGatewayServiceWithSettings(t, settings)
  142. account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  143. allowedUserCtx := context.WithValue(context.Background(), ctxkey.UserID, int64(42))
  144. action, _ := svc.evaluateOpenAIFastPolicy(allowedUserCtx, account, "gpt-5.5", OpenAIFastTierPriority)
  145. require.Equal(t, BetaPolicyActionPass, action)
  146. otherUserCtx := context.WithValue(context.Background(), ctxkey.UserID, int64(43))
  147. action, _ = svc.evaluateOpenAIFastPolicy(otherUserCtx, account, "gpt-5.5", OpenAIFastTierPriority)
  148. require.Equal(t, BetaPolicyActionFilter, action)
  149. action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5", OpenAIFastTierPriority)
  150. require.Equal(t, BetaPolicyActionFilter, action)
  151. }
  152. func TestApplyOpenAIFastPolicyToBody_DefaultPassesPriorityAndFast(t *testing.T) {
  153. svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
  154. account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  155. body := []byte(`{"model":"gpt-5.5","service_tier":"priority","messages":[]}`)
  156. updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
  157. require.NoError(t, err)
  158. require.Equal(t, string(body), string(updated))
  159. body = []byte(`{"model":"gpt-5.5","service_tier":"fast"}`)
  160. updated, err = svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
  161. require.NoError(t, err)
  162. require.Equal(t, "priority", gjson.GetBytes(updated, "service_tier").String())
  163. body = []byte(`{"model":"gpt-4","service_tier":"priority"}`)
  164. updated, err = svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-4", body)
  165. require.NoError(t, err)
  166. require.Equal(t, string(body), string(updated))
  167. // No service_tier → no-op
  168. body = []byte(`{"model":"gpt-5.5"}`)
  169. updated, err = svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
  170. require.NoError(t, err)
  171. require.Equal(t, string(body), string(updated))
  172. }
  173. func TestApplyOpenAIFastPolicyToBody_ExplicitFilterRemovesField(t *testing.T) {
  174. svc := newOpenAIGatewayServiceWithSettings(t, openAIFastFilterPriorityPolicy())
  175. account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  176. body := []byte(`{"model":"gpt-5.5","service_tier":"priority","messages":[]}`)
  177. updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
  178. require.NoError(t, err)
  179. require.NotContains(t, string(updated), `"service_tier"`)
  180. body = []byte(`{"model":"gpt-5.5","service_tier":"fast"}`)
  181. updated, err = svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
  182. require.NoError(t, err)
  183. require.NotContains(t, string(updated), `"service_tier"`)
  184. }
  185. func TestApplyOpenAIFastPolicyToBody_UserScopedRuleOverridesGlobalRule(t *testing.T) {
  186. settings := &OpenAIFastPolicySettings{
  187. Rules: []OpenAIFastPolicyRule{
  188. {
  189. ServiceTier: OpenAIFastTierPriority,
  190. Action: BetaPolicyActionFilter,
  191. Scope: BetaPolicyScopeAll,
  192. },
  193. {
  194. ServiceTier: OpenAIFastTierPriority,
  195. Action: BetaPolicyActionPass,
  196. Scope: BetaPolicyScopeAll,
  197. UserIDs: []int64{42},
  198. },
  199. },
  200. }
  201. svc := newOpenAIGatewayServiceWithSettings(t, settings)
  202. account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  203. body := []byte(`{"model":"gpt-5.5","service_tier":"priority"}`)
  204. allowedUserCtx := context.WithValue(context.Background(), ctxkey.UserID, int64(42))
  205. updated, err := svc.applyOpenAIFastPolicyToBody(allowedUserCtx, account, "gpt-5.5", body)
  206. require.NoError(t, err)
  207. require.Equal(t, "priority", gjson.GetBytes(updated, "service_tier").String())
  208. otherUserCtx := context.WithValue(context.Background(), ctxkey.UserID, int64(43))
  209. updated, err = svc.applyOpenAIFastPolicyToBody(otherUserCtx, account, "gpt-5.5", body)
  210. require.NoError(t, err)
  211. require.NotContains(t, string(updated), `"service_tier"`)
  212. }
  213. func TestApplyOpenAIFastPolicyToBody_PriorityFilterLeavesUltrafast(t *testing.T) {
  214. svc := newOpenAIGatewayServiceWithSettings(t, openAIFastFilterPriorityPolicy())
  215. account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  216. body := []byte(`{"model":"gpt-5.6-sol","service_tier":"ultrafast"}`)
  217. updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.6-sol", body)
  218. require.NoError(t, err)
  219. require.Equal(t, OpenAIFastTierUltrafast, gjson.GetBytes(updated, "service_tier").String())
  220. }
  221. func TestApplyOpenAIFastPolicyToBody_ForcePriorityRewritesKnownTier(t *testing.T) {
  222. settings := &OpenAIFastPolicySettings{
  223. Rules: []OpenAIFastPolicyRule{{
  224. ServiceTier: OpenAIFastTierAny,
  225. Action: OpenAIFastPolicyActionForcePriority,
  226. Scope: BetaPolicyScopeAll,
  227. }},
  228. }
  229. svc := newOpenAIGatewayServiceWithSettings(t, settings)
  230. account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  231. for _, tier := range []string{"flex", "auto", "default", "scale", "fast", "priority", "ultrafast"} {
  232. body := []byte(`{"model":"gpt-5.5","service_tier":"` + tier + `"}`)
  233. updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
  234. require.NoError(t, err)
  235. require.Equal(t, OpenAIFastTierPriority, gjson.GetBytes(updated, "service_tier").String(),
  236. "tier %q should be forced to priority", tier)
  237. }
  238. }
  239. func TestApplyOpenAIFastPolicyToBody_GroupForceInjectsAndOverridesTier(t *testing.T) {
  240. svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
  241. account := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  242. ctx := context.WithValue(context.Background(), ctxkey.Group, &Group{
  243. ID: 7, Platform: PlatformOpenAI, Status: StatusActive, Hydrated: true, ForceOpenAIFast: true,
  244. })
  245. for _, body := range [][]byte{
  246. []byte(`{"model":"gpt-5.6-sol","input":"hi"}`),
  247. []byte(`{"model":"gpt-5.6-sol","service_tier":"default"}`),
  248. []byte(`{"model":"gpt-5.6-sol","service_tier":"flex"}`),
  249. []byte(`{"model":"gpt-5.6-sol","service_tier":"client-unknown"}`),
  250. } {
  251. updated, err := svc.applyOpenAIFastPolicyToBody(ctx, account, "gpt-5.6-sol", body)
  252. require.NoError(t, err)
  253. require.Equal(t, OpenAIFastTierPriority, gjson.GetBytes(updated, "service_tier").String())
  254. }
  255. }
  256. func TestApplyOpenAIFastPolicyToBody_GroupForceStillHonorsGlobalPolicy(t *testing.T) {
  257. svc := newOpenAIGatewayServiceWithSettings(t, openAIFastFilterPriorityPolicy())
  258. account := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  259. ctx := context.WithValue(context.Background(), ctxkey.Group, &Group{
  260. ID: 7, Platform: PlatformOpenAI, Status: StatusActive, Hydrated: true, ForceOpenAIFast: true,
  261. })
  262. updated, err := svc.applyOpenAIFastPolicyToBody(ctx, account, "gpt-5.6-sol", []byte(`{"model":"gpt-5.6-sol"}`))
  263. require.NoError(t, err)
  264. require.False(t, gjson.GetBytes(updated, "service_tier").Exists(),
  265. "the global filter remains authoritative after the group requests priority")
  266. }
  267. func TestApplyOpenAIFastPolicyToBody_GroupForceRequiresHydratedGroup(t *testing.T) {
  268. svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
  269. account := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  270. body := []byte(`{"model":"gpt-5.6-sol"}`)
  271. ctx := context.WithValue(context.Background(), ctxkey.Group, &Group{
  272. ID: 7, Platform: PlatformOpenAI, Status: StatusActive, ForceOpenAIFast: true,
  273. })
  274. updated, err := svc.applyOpenAIFastPolicyToBody(ctx, account, "gpt-5.6-sol", body)
  275. require.NoError(t, err)
  276. require.Equal(t, string(body), string(updated))
  277. }
  278. func TestApplyOpenAIFastPolicyToBody_GroupForceOnlyTargetsOpenAIAccounts(t *testing.T) {
  279. svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
  280. body := []byte(`{"model":"grok-4.1"}`)
  281. ctx := context.WithValue(context.Background(), ctxkey.Group, &Group{
  282. ID: 7, Platform: PlatformComposite, Status: StatusActive, Hydrated: true, ForceOpenAIFast: true,
  283. })
  284. updated, err := svc.applyOpenAIFastPolicyToBody(
  285. ctx,
  286. &Account{Platform: PlatformGrok, Type: AccountTypeOAuth},
  287. "grok-4.1",
  288. body,
  289. )
  290. require.NoError(t, err)
  291. require.Equal(t, string(body), string(updated))
  292. }
  293. func TestApplyOpenAIFastPolicyToBody_GroupForceRequiresSupportedGroupPlatform(t *testing.T) {
  294. svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
  295. body := []byte(`{"model":"gpt-5.6-sol"}`)
  296. ctx := context.WithValue(context.Background(), ctxkey.Group, &Group{
  297. ID: 7, Platform: PlatformAnthropic, Status: StatusActive, Hydrated: true, ForceOpenAIFast: true,
  298. })
  299. updated, err := svc.applyOpenAIFastPolicyToBody(
  300. ctx,
  301. &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth},
  302. "gpt-5.6-sol",
  303. body,
  304. )
  305. require.NoError(t, err)
  306. require.Equal(t, string(body), string(updated))
  307. }
  308. // TestApplyOpenAIFastPolicyToBody_OfficialTiersBypassDefaultRule 验证默认配置
  309. // 下客户端显式发送的 OpenAI 官方合法 tier 能透传到上游而不被静默剥离。
  310. func TestApplyOpenAIFastPolicyToBody_OfficialTiersBypassDefaultRule(t *testing.T) {
  311. svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
  312. account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  313. for _, tier := range []string{"auto", "default", "scale"} {
  314. body := []byte(`{"model":"gpt-5.5","service_tier":"` + tier + `"}`)
  315. updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
  316. require.NoError(t, err, "tier %q should pass without error", tier)
  317. require.Contains(t, string(updated), `"service_tier":"`+tier+`"`,
  318. "tier %q should be preserved in body under default policy", tier)
  319. }
  320. // evaluate 层也应判定为 pass(默认配置没有内置规则)。
  321. for _, tier := range []string{"auto", "default", "scale"} {
  322. action, _ := svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5", tier)
  323. require.Equal(t, BetaPolicyActionPass, action, "tier %q should evaluate to pass", tier)
  324. }
  325. }
  326. // TestApplyOpenAIFastPolicyToBody_AllRuleStripsOfficialTiers 验证管理员显式配置
  327. // ServiceTier=all + Action=filter 规则后,auto/default/scale 等官方 tier 也会
  328. // 被剥离。这是符合预期的——首条匹配 short-circuit,"all" 覆盖任意已识别 tier。
  329. func TestApplyOpenAIFastPolicyToBody_AllRuleStripsOfficialTiers(t *testing.T) {
  330. settings := &OpenAIFastPolicySettings{
  331. Rules: []OpenAIFastPolicyRule{{
  332. ServiceTier: OpenAIFastTierAny,
  333. Action: BetaPolicyActionFilter,
  334. Scope: BetaPolicyScopeAll,
  335. }},
  336. }
  337. svc := newOpenAIGatewayServiceWithSettings(t, settings)
  338. account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  339. for _, tier := range []string{"auto", "default", "scale", "priority", "flex"} {
  340. body := []byte(`{"model":"gpt-5.5","service_tier":"` + tier + `"}`)
  341. updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
  342. require.NoError(t, err)
  343. require.NotContains(t, string(updated), `"service_tier"`,
  344. "tier %q should be stripped under ServiceTier=all + filter rule", tier)
  345. }
  346. }
  347. // TestApplyOpenAIFastPolicyToBody_UnknownTierStripped 验证真未知 tier 仍被剥离
  348. // (normalize 返回 nil → normalizeResponsesBodyServiceTier 删除字段;
  349. // applyOpenAIFastPolicyToBody 在 normTier 为空时直接 no-op,因为字段已不可能存在
  350. // 于经过前置归一化的请求里。这里直接调 apply 验证它对未识别值不会异常)。
  351. func TestApplyOpenAIFastPolicyToBody_UnknownTierStripped(t *testing.T) {
  352. svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
  353. account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  354. // normalize 阶段会将未知值剥离
  355. require.Nil(t, normalizeOpenAIServiceTier("xxx"))
  356. // applyOpenAIFastPolicyToBody 收到未识别 tier 时不报错,body 透传不变
  357. // (不属于本函数职责——上层 normalizeResponsesBodyServiceTier 已剥离)
  358. body := []byte(`{"model":"gpt-5.5","service_tier":"xxx"}`)
  359. updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
  360. require.NoError(t, err)
  361. require.Equal(t, string(body), string(updated))
  362. }
  363. func TestApplyOpenAIFastPolicyToBody_BlockReturnsTypedError(t *testing.T) {
  364. settings := &OpenAIFastPolicySettings{
  365. Rules: []OpenAIFastPolicyRule{{
  366. ServiceTier: OpenAIFastTierPriority,
  367. Action: BetaPolicyActionBlock,
  368. Scope: BetaPolicyScopeAll,
  369. ErrorMessage: "fast mode is blocked for gpt-5.5",
  370. ModelWhitelist: []string{"gpt-5.5"},
  371. FallbackAction: BetaPolicyActionPass,
  372. }},
  373. }
  374. svc := newOpenAIGatewayServiceWithSettings(t, settings)
  375. account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  376. body := []byte(`{"model":"gpt-5.5","service_tier":"priority"}`)
  377. updated, err := svc.applyOpenAIFastPolicyToBody(context.Background(), account, "gpt-5.5", body)
  378. require.Error(t, err)
  379. var blocked *OpenAIFastBlockedError
  380. require.True(t, errors.As(err, &blocked))
  381. require.Contains(t, blocked.Message, "fast mode is blocked")
  382. require.Equal(t, string(body), string(updated)) // body not mutated on block
  383. }
  384. func TestSetOpenAIFastPolicySettings_Validation(t *testing.T) {
  385. repo := &openAIFastPolicyRepoStub{values: map[string]string{}}
  386. svc := NewSettingService(repo, &config.Config{})
  387. // Invalid action rejected
  388. err := svc.SetOpenAIFastPolicySettings(context.Background(), &OpenAIFastPolicySettings{
  389. Rules: []OpenAIFastPolicyRule{{
  390. ServiceTier: OpenAIFastTierPriority,
  391. Action: "bogus",
  392. Scope: BetaPolicyScopeAll,
  393. }},
  394. })
  395. require.Error(t, err)
  396. // Invalid service_tier rejected
  397. err = svc.SetOpenAIFastPolicySettings(context.Background(), &OpenAIFastPolicySettings{
  398. Rules: []OpenAIFastPolicyRule{{
  399. ServiceTier: "turbo",
  400. Action: BetaPolicyActionPass,
  401. Scope: BetaPolicyScopeAll,
  402. }},
  403. })
  404. require.Error(t, err)
  405. // Non-positive and duplicate user IDs are rejected.
  406. err = svc.SetOpenAIFastPolicySettings(context.Background(), &OpenAIFastPolicySettings{
  407. Rules: []OpenAIFastPolicyRule{{
  408. ServiceTier: OpenAIFastTierPriority,
  409. Action: BetaPolicyActionPass,
  410. Scope: BetaPolicyScopeAll,
  411. UserIDs: []int64{0},
  412. }},
  413. })
  414. require.Error(t, err)
  415. err = svc.SetOpenAIFastPolicySettings(context.Background(), &OpenAIFastPolicySettings{
  416. Rules: []OpenAIFastPolicyRule{{
  417. ServiceTier: OpenAIFastTierPriority,
  418. Action: BetaPolicyActionPass,
  419. Scope: BetaPolicyScopeAll,
  420. UserIDs: []int64{42, 42},
  421. }},
  422. })
  423. require.Error(t, err)
  424. // Valid settings persisted
  425. err = svc.SetOpenAIFastPolicySettings(context.Background(), &OpenAIFastPolicySettings{
  426. Rules: []OpenAIFastPolicyRule{
  427. {
  428. ServiceTier: OpenAIFastTierPriority,
  429. Action: OpenAIFastPolicyActionForcePriority,
  430. Scope: BetaPolicyScopeAll,
  431. UserIDs: []int64{42, 43},
  432. },
  433. {
  434. ServiceTier: OpenAIFastTierUltrafast,
  435. Action: BetaPolicyActionPass,
  436. Scope: BetaPolicyScopeAll,
  437. },
  438. },
  439. })
  440. require.NoError(t, err)
  441. got, err := svc.GetOpenAIFastPolicySettings(context.Background())
  442. require.NoError(t, err)
  443. require.Len(t, got.Rules, 2)
  444. require.Equal(t, OpenAIFastTierPriority, got.Rules[0].ServiceTier)
  445. require.Equal(t, OpenAIFastPolicyActionForcePriority, got.Rules[0].Action)
  446. require.Equal(t, []int64{42, 43}, got.Rules[0].UserIDs)
  447. require.Equal(t, OpenAIFastTierUltrafast, got.Rules[1].ServiceTier)
  448. require.Equal(t, BetaPolicyActionPass, got.Rules[1].Action)
  449. }