Force OpenAI Fast across gateway paths

This commit is contained in:
Ritel-T committed 2026-09-01 04:34:06 +09:00
1 parent a44a1019f8
commit f8eae45048
6 files changed
+170 -1

No files matched your search

@@ -262,6 +262,85 @@ func TestApplyOpenAIFastPolicyToBody_ForcePriorityRewritesKnownTier(t *testing.T
}
}
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) {
@@ -122,6 +122,34 @@ func TestWSResponseCreate_ForcePriorityRewritesKnownTier(t *testing.T) {
}
}
func TestWSResponseCreate_GroupForceInjectsMissingTier(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,
})
frame := []byte(`{"type":"response.create","model":"gpt-5.6-sol","input":[]}`)
updated, blocked, err := svc.applyOpenAIFastPolicyToWSResponseCreate(ctx, account, "gpt-5.6-sol", frame)
require.NoError(t, err)
require.Nil(t, blocked)
require.Equal(t, OpenAIFastTierPriority, gjson.GetBytes(updated, "service_tier").String())
}
func TestWSResponseCreate_GroupForceDoesNotTouchOtherFrames(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,
})
frame := []byte(`{"type":"response.cancel","model":"gpt-5.6-sol"}`)
updated, blocked, err := svc.applyOpenAIFastPolicyToWSResponseCreate(ctx, account, "gpt-5.6-sol", frame)
require.NoError(t, err)
require.Nil(t, blocked)
require.Equal(t, string(frame), string(updated))
}
func TestWSResponseCreate_FlexPassThrough(t *testing.T) {
svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
@@ -305,6 +305,7 @@ func (s *OpenAIGatewayService) ForwardAsChatCompletions(
return nil, policyErr
}
responsesBody = updatedBody
responsesReq.ServiceTier = normalizedOpenAIServiceTierValue(gjson.GetBytes(responsesBody, "service_tier").String())
// 5. Get access token
token, _, err := s.GetAccessToken(ctx, account)
@@ -291,6 +291,7 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic(
return nil, policyErr
}
responsesBody = updatedBody
responsesReq.ServiceTier = normalizedOpenAIServiceTierValue(gjson.GetBytes(responsesBody, "service_tier").String())
grokCacheIdentity := ""
if account.Platform == PlatformGrok {
grokIntentBody := responsesBody
@@ -1638,11 +1638,23 @@ func openAIFastPolicySettingsFromContext(ctx context.Context) *OpenAIFastPolicyS
return nil
}
func openAIGroupForcesFast(ctx context.Context, account *Account) bool {
if ctx == nil || account == nil || account.Platform != PlatformOpenAI {
return false
}
group, _ := ctx.Value(ctxkey.Group).(*Group)
return IsGroupContextValid(group) && groupSupportsOpenAIFast(group.Platform) && group.ForceOpenAIFast
}
// applyOpenAIFastPolicyToBody applies the OpenAI fast policy to a raw request
// body. When action=filter it removes the service_tier field; when
// action=block it returns (body, *OpenAIFastBlockedError). On pass it
// normalizes the service_tier value (e.g. client alias "fast" → "priority").
// action=force_priority rewrites any matched known tier to "priority".
// action=force_priority rewrites any matched known tier to "priority". Before
// the global policy is evaluated, a trusted request Group with
// ForceOpenAIFast enabled unconditionally sets service_tier to "priority".
// The global policy remains authoritative and may still pass, filter, or block
// that final value.
//
// Rationale for normalize-on-pass: chat-completions / messages 入口在调用本
// 函数之前已经通过 normalizeResponsesBodyServiceTier 把 service_tier 归一化
@@ -1653,6 +1665,13 @@ func (s *OpenAIGatewayService) applyOpenAIFastPolicyToBody(ctx context.Context,
if len(body) == 0 {
return body, nil
}
if openAIGroupForcesFast(ctx, account) {
updated, err := sjson.SetBytes(body, "service_tier", OpenAIFastTierPriority)
if err != nil {
return body, fmt.Errorf("force group service_tier priority on body: %w", err)
}
body = updated
}
rawTier := gjson.GetBytes(body, "service_tier").String()
if rawTier == "" {
return body, nil
@@ -1726,6 +1745,7 @@ func writeOpenAIFastPolicyBlockedResponse(c *gin.Context, err *OpenAIFastBlocked
// - filter: returns a copy with top-level service_tier removed
// - force_priority: keeps service_tier and rewrites it to "priority"
// - block: returns (frame, *OpenAIFastBlockedError)
// - Group ForceOpenAIFast: sets priority first, then applies the global rule
//
// Only frames whose "type" field strictly equals "response.create" are
// inspected/mutated. Any other frame type — including the empty string —
@@ -1764,6 +1784,13 @@ func (s *OpenAIGatewayService) applyOpenAIFastPolicyToWSResponseCreate(
if frameType != "response.create" {
return frame, nil, nil
}
if openAIGroupForcesFast(ctx, account) {
updated, err := sjson.SetBytes(frame, "service_tier", OpenAIFastTierPriority)
if err != nil {
return frame, nil, fmt.Errorf("force group service_tier priority in ws frame: %w", err)
}
frame = updated
}
rawTier := gjson.GetBytes(frame, "service_tier").String()
if rawTier == "" {
return frame, nil, nil
@@ -15,6 +15,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
"github.com/Wei-Shaw/sub2api/internal/pkg/openai_compat"
@@ -473,6 +474,38 @@ func TestOpenAIGatewayService_OAuthPassthrough_StreamKeepsToolNameAndBodyNormali
require.NotContains(t, body, "\"name\":\"edit\"")
}
func TestOpenAIGatewayService_OAuthPassthrough_GroupForceOpenAIFastInjectsMissingTier(t *testing.T) {
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(nil))
c.Request.Header.Set("User-Agent", "codex_cli_rs/0.144.1")
body := []byte(`{"model":"gpt-5.6-sol","stream":true,"instructions":"test","input":"hi"}`)
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader("data: [DONE]\n\n")),
}}
svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
account := &Account{
ID: 123, Name: "oauth", Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 1,
Credentials: map[string]any{"access_token": "oauth-token", "chatgpt_account_id": "chatgpt-acc"},
Extra: map[string]any{"openai_passthrough": true}, Status: StatusActive, Schedulable: true,
RateMultiplier: f64p(1),
}
ctx := context.WithValue(context.Background(), ctxkey.Group, &Group{
ID: 7, Platform: PlatformOpenAI, Status: StatusActive, Hydrated: true, ForceOpenAIFast: true,
})
result, err := svc.Forward(ctx, c, account, body)
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, OpenAIFastTierPriority, gjson.GetBytes(upstream.lastBody, "service_tier").String())
require.NotNil(t, result.ServiceTier)
require.Equal(t, OpenAIFastTierPriority, *result.ServiceTier)
}
// 「自动透传(仅替换认证)」的默认行为必须真的只替换认证:namespace 声明、
// namespace 形态的 tool_choice、历史调用项上的 namespace 都原样转发,只清掉
// 非调用项上的残留 namespace(Codex 协议里只有调用项会带该字段)。