fix(gateway): address OpenAI OOM review lint

This commit is contained in:
name 2026-05-31 16:38:42 +08:00
parent 8ac2e23fb3
commit 1afae0a4cd
4 changed files with 109 additions and 51 deletions

View File

@ -229,16 +229,16 @@ func benchmarkBodySizes() []struct {
func buildSystemCacheableRequest(parts int) *ParsedRequest { func buildSystemCacheableRequest(parts int) *ParsedRequest {
var builder strings.Builder var builder strings.Builder
builder.WriteString(`{"system":[`) _, _ = builder.WriteString(`{"system":[`)
for i := 0; i < parts; i++ { for i := 0; i < parts; i++ {
if i > 0 { if i > 0 {
builder.WriteByte(',') _ = builder.WriteByte(',')
} }
builder.WriteString(`{"text":"system_part_`) _, _ = builder.WriteString(`{"text":"system_part_`)
builder.WriteString(strconv.Itoa(i)) _, _ = builder.WriteString(strconv.Itoa(i))
builder.WriteString(`","cache_control":{"type":"ephemeral"}}`) _, _ = builder.WriteString(`","cache_control":{"type":"ephemeral"}}`)
} }
builder.WriteString(`]}`) _, _ = builder.WriteString(`]}`)
parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(builder.String())), "") parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(builder.String())), "")
if err != nil { if err != nil {
panic(err) panic(err)
@ -249,93 +249,93 @@ func buildSystemCacheableRequest(parts int) *ParsedRequest {
func buildLargeAnthropicMessagesBody(targetBytes int, includeCacheControl bool) []byte { func buildLargeAnthropicMessagesBody(targetBytes int, includeCacheControl bool) []byte {
var builder strings.Builder var builder strings.Builder
builder.Grow(targetBytes + 1024) builder.Grow(targetBytes + 1024)
builder.WriteString(`{"model":"claude-sonnet-4-5","stream":true,"system":[{"type":"text","text":"system seed"}],"messages":[`) _, _ = builder.WriteString(`{"model":"claude-sonnet-4-5","stream":true,"system":[{"type":"text","text":"system seed"}],"messages":[`)
for i := 0; builder.Len() < targetBytes; i++ { for i := 0; builder.Len() < targetBytes; i++ {
if i > 0 { if i > 0 {
builder.WriteByte(',') _ = builder.WriteByte(',')
} }
builder.WriteString(`{"role":"user","content":[{"type":"text","text":"`) _, _ = builder.WriteString(`{"role":"user","content":[{"type":"text","text":"`)
builder.WriteString(strings.Repeat("anthropic payload ", 64)) _, _ = builder.WriteString(strings.Repeat("anthropic payload ", 64))
builder.WriteString(strconv.Itoa(i)) _, _ = builder.WriteString(strconv.Itoa(i))
builder.WriteByte('"') _ = builder.WriteByte('"')
if includeCacheControl && i%32 == 0 { if includeCacheControl && i%32 == 0 {
builder.WriteString(`,"cache_control":{"type":"ephemeral"}`) _, _ = builder.WriteString(`,"cache_control":{"type":"ephemeral"}`)
} }
builder.WriteString(`}]}`) _, _ = builder.WriteString(`}]}`)
} }
builder.WriteString(`]}`) _, _ = builder.WriteString(`]}`)
return []byte(builder.String()) return []byte(builder.String())
} }
func buildLargeGeminiContentsBody(targetBytes int) []byte { func buildLargeGeminiContentsBody(targetBytes int) []byte {
var builder strings.Builder var builder strings.Builder
builder.Grow(targetBytes + 1024) builder.Grow(targetBytes + 1024)
builder.WriteString(`{"model":"gemini-2.5-pro","systemInstruction":{"parts":[{"text":"system seed"}]},"contents":[`) _, _ = builder.WriteString(`{"model":"gemini-2.5-pro","systemInstruction":{"parts":[{"text":"system seed"}]},"contents":[`)
for i := 0; builder.Len() < targetBytes; i++ { for i := 0; builder.Len() < targetBytes; i++ {
if i > 0 { if i > 0 {
builder.WriteByte(',') _ = builder.WriteByte(',')
} }
builder.WriteString(`{"role":"user","parts":[{"text":"`) _, _ = builder.WriteString(`{"role":"user","parts":[{"text":"`)
builder.WriteString(strings.Repeat("gemini payload ", 64)) _, _ = builder.WriteString(strings.Repeat("gemini payload ", 64))
builder.WriteString(strconv.Itoa(i)) _, _ = builder.WriteString(strconv.Itoa(i))
builder.WriteString(`"}]}`) _, _ = builder.WriteString(`"}]}`)
} }
builder.WriteString(`]}`) _, _ = builder.WriteString(`]}`)
return []byte(builder.String()) return []byte(builder.String())
} }
func buildLargeOpenAIResponsesBody(targetBytes int) []byte { func buildLargeOpenAIResponsesBody(targetBytes int) []byte {
var builder strings.Builder var builder strings.Builder
builder.Grow(targetBytes + 1024) builder.Grow(targetBytes + 1024)
builder.WriteString(`{"model":"gpt-5.4","stream":true,"prompt_cache_key":"session-benchmark","input":[`) _, _ = builder.WriteString(`{"model":"gpt-5.4","stream":true,"prompt_cache_key":"session-benchmark","input":[`)
for i := 0; builder.Len() < targetBytes; i++ { for i := 0; builder.Len() < targetBytes; i++ {
if i > 0 { if i > 0 {
builder.WriteByte(',') _ = builder.WriteByte(',')
} }
builder.WriteString(`{"type":"message","role":"user","content":[{"type":"input_text","text":"`) _, _ = builder.WriteString(`{"type":"message","role":"user","content":[{"type":"input_text","text":"`)
builder.WriteString(strings.Repeat("openai responses payload ", 48)) _, _ = builder.WriteString(strings.Repeat("openai responses payload ", 48))
builder.WriteString(strconv.Itoa(i)) _, _ = builder.WriteString(strconv.Itoa(i))
builder.WriteString(`"}]}`) _, _ = builder.WriteString(`"}]}`)
} }
builder.WriteString(`],"tools":[{"type":"function","name":"lookup","parameters":{"type":"object","properties":{"query":{"type":"string"}}}}]}`) _, _ = builder.WriteString(`],"tools":[{"type":"function","name":"lookup","parameters":{"type":"object","properties":{"query":{"type":"string"}}}}]}`)
return []byte(builder.String()) return []byte(builder.String())
} }
func buildLargeOpenAIResponsesToolContinuationBody(targetBytes int) []byte { func buildLargeOpenAIResponsesToolContinuationBody(targetBytes int) []byte {
var builder strings.Builder var builder strings.Builder
builder.Grow(targetBytes + 1024) builder.Grow(targetBytes + 1024)
builder.WriteString(`{"model":"gpt-5.4","stream":true,"previous_response_id":"resp_benchmark","input":[`) _, _ = builder.WriteString(`{"model":"gpt-5.4","stream":true,"previous_response_id":"resp_benchmark","input":[`)
for i := 0; builder.Len() < targetBytes; i++ { for i := 0; builder.Len() < targetBytes; i++ {
if i > 0 { if i > 0 {
builder.WriteByte(',') _ = builder.WriteByte(',')
} }
callID := "call_" + strconv.Itoa(i) callID := "call_" + strconv.Itoa(i)
builder.WriteString(`{"type":"item_reference","id":"`) _, _ = builder.WriteString(`{"type":"item_reference","id":"`)
builder.WriteString(callID) _, _ = builder.WriteString(callID)
builder.WriteString(`"},{"type":"function_call_output","call_id":"`) _, _ = builder.WriteString(`"},{"type":"function_call_output","call_id":"`)
builder.WriteString(callID) _, _ = builder.WriteString(callID)
builder.WriteString(`","output":"`) _, _ = builder.WriteString(`","output":"`)
builder.WriteString(strings.Repeat("tool output payload ", 48)) _, _ = builder.WriteString(strings.Repeat("tool output payload ", 48))
builder.WriteString(strconv.Itoa(i)) _, _ = builder.WriteString(strconv.Itoa(i))
builder.WriteString(`"}`) _, _ = builder.WriteString(`"}`)
} }
builder.WriteString(`]}`) _, _ = builder.WriteString(`]}`)
return []byte(builder.String()) return []byte(builder.String())
} }
func buildLargeOpenAIResponsesImageToolBody(targetBytes int) []byte { func buildLargeOpenAIResponsesImageToolBody(targetBytes int) []byte {
var builder strings.Builder var builder strings.Builder
builder.Grow(targetBytes + 1024) builder.Grow(targetBytes + 1024)
builder.WriteString(`{"model":"gpt-5.4","stream":false,"tools":[{"type":"image_generation","model":"gpt-image-2","size":"2048x1152"}],"input":[`) _, _ = builder.WriteString(`{"model":"gpt-5.4","stream":false,"tools":[{"type":"image_generation","model":"gpt-image-2","size":"2048x1152"}],"input":[`)
for i := 0; builder.Len() < targetBytes; i++ { for i := 0; builder.Len() < targetBytes; i++ {
if i > 0 { if i > 0 {
builder.WriteByte(',') _ = builder.WriteByte(',')
} }
builder.WriteString(`{"type":"message","role":"user","content":[{"type":"input_text","text":"`) _, _ = builder.WriteString(`{"type":"message","role":"user","content":[{"type":"input_text","text":"`)
builder.WriteString(strings.Repeat("openai image billing payload ", 48)) _, _ = builder.WriteString(strings.Repeat("openai image billing payload ", 48))
builder.WriteString(strconv.Itoa(i)) _, _ = builder.WriteString(strconv.Itoa(i))
builder.WriteString(`"}]}`) _, _ = builder.WriteString(`"}]}`)
} }
builder.WriteString(`]}`) _, _ = builder.WriteString(`]}`)
return []byte(builder.String()) return []byte(builder.String())
} }

View File

@ -2544,6 +2544,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Added Codex image_generation bridge instructions") logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Added Codex image_generation bridge instructions")
} }
} else if imageGenerationAllowed && imageIntent && openAIRequestBodyHasImageGenerationTool(body) { } else if imageGenerationAllowed && imageIntent && openAIRequestBodyHasImageGenerationTool(body) {
// 完整 image_generation tool 只做 raw 计费读取,校验/桥接/旧字段迁移命中时才展开大 input map。
logger.LegacyPrintf("service.openai_gateway", "[OpenAI] /responses image_generation request inbound_model=%s mapped_model=%s account_type=%s", requestView.Model, upstreamModel, account.Type) logger.LegacyPrintf("service.openai_gateway", "[OpenAI] /responses image_generation request inbound_model=%s mapped_model=%s account_type=%s", requestView.Model, upstreamModel, account.Type)
} }
@ -6215,7 +6216,7 @@ func (v *openAIRequestView) MarkPatchSet(path string, value any) {
return return
} }
path = strings.TrimSpace(path) path = strings.TrimSpace(path)
if path == "" { if !isSimpleOpenAIRequestPatchPath(path) {
v.DisablePatches() v.DisablePatches()
return return
} }
@ -6227,13 +6228,25 @@ func (v *openAIRequestView) MarkPatchDelete(path string) {
return return
} }
path = strings.TrimSpace(path) path = strings.TrimSpace(path)
if path == "" { if !isSimpleOpenAIRequestPatchPath(path) {
v.DisablePatches() v.DisablePatches()
return return
} }
v.patches = append(v.patches, openAIRequestPatch{path: path, delete: true}) v.patches = append(v.patches, openAIRequestPatch{path: path, delete: true})
} }
func isSimpleOpenAIRequestPatchPath(path string) bool {
if path == "" || strings.ContainsRune(path, '\\') {
return false
}
for _, part := range strings.Split(path, ".") {
if strings.TrimSpace(part) == "" {
return false
}
}
return true
}
func (v *openAIRequestView) DisablePatches() { func (v *openAIRequestView) DisablePatches() {
if v == nil { if v == nil {
return return
@ -6772,7 +6785,7 @@ func openAIRequestBodyMayContainImageInput(body []byte) bool {
return false return false
} }
input := gjson.GetBytes(body, "input") input := gjson.GetBytes(body, "input")
messages := gjson.GetBytes(body, "messages") messages := gjson.GetBytes(body, "messages.#-1")
return openAIJSONValueMayContainImageInput(input) || openAIJSONValueMayContainImageInput(messages) return openAIJSONValueMayContainImageInput(input) || openAIJSONValueMayContainImageInput(messages)
} }

View File

@ -46,6 +46,15 @@ func TestOpenAIRequestView_ApplyPatches(t *testing.T) {
require.JSONEq(t, `{"model":"gpt-5.1","reasoning":{"effort":"none"},"input":[{"type":"message","content":"hi"}]}`, string(patched)) require.JSONEq(t, `{"model":"gpt-5.1","reasoning":{"effort":"none"},"input":[{"type":"message","content":"hi"}]}`, string(patched))
} }
func TestOpenAIRequestView_RejectsEscapedPatchPath(t *testing.T) {
view := newOpenAIRequestView([]byte(`{"metadata":{"user.id":"old"}}`))
view.MarkPatchSet(`metadata.user\.id`, "new")
require.False(t, view.HasPatches())
_, err := view.ApplyPatches()
require.Error(t, err)
}
func TestOpenAIRequestView_ApplyPatchesDisabled(t *testing.T) { func TestOpenAIRequestView_ApplyPatchesDisabled(t *testing.T) {
view := newOpenAIRequestView([]byte(`{"model":"gpt-5"}`)) view := newOpenAIRequestView([]byte(`{"model":"gpt-5"}`))
view.MarkPatchSet("model", "gpt-5.1") view.MarkPatchSet("model", "gpt-5.1")
@ -260,6 +269,42 @@ func TestOpenAIGatewayService_Forward_ImageToolBillingDoesNotForceFullDecode(t *
require.Equal(t, "gpt-image-2", result.BillingModel) require.Equal(t, "gpt-image-2", result.BillingModel)
} }
func TestOpenAIGatewayService_Forward_ImageToolWithImageOnlyModelIsNormalized(t *testing.T) {
gin.SetMode(gin.TestMode)
upstream := &httpUpstreamRecorder{
resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)),
},
}
cfg := &config.Config{}
cfg.Security.URLAllowlist.Enabled = false
svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream}
account := &Account{
ID: 11,
Name: "openai-apikey",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://example.com",
},
Extra: map[string]any{"use_responses_api": true},
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
SetOpenAIClientTransport(c, OpenAIClientTransportHTTP)
body := []byte(`{"model":"gpt-image-2","stream":false,"tools":[{"type":"image_generation","model":"gpt-image-2"}],"input":"draw"}`)
result, err := svc.Forward(context.Background(), c, account, body)
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, openAIImagesResponsesMainModel, gjson.GetBytes(upstream.lastBody, "model").String())
}
func TestOpenAIGatewayService_Forward_HTTPRetryRecoveryDoesNotDecodeBeforeError(t *testing.T) { func TestOpenAIGatewayService_Forward_HTTPRetryRecoveryDoesNotDecodeBeforeError(t *testing.T) {
gin.SetMode(gin.TestMode) gin.SetMode(gin.TestMode)
upstream := &httpUpstreamRecorder{ upstream := &httpUpstreamRecorder{

View File

@ -199,7 +199,7 @@ func ValidateFunctionCallOutputContextBytes(body []byte) FunctionCallOutputValid
} }
referenceIDs[idValue] = struct{}{} referenceIDs[idValue] = struct{}{}
} }
return !(result.HasFunctionCallOutput && result.HasToolCallContext) return !result.HasFunctionCallOutput || !result.HasToolCallContext
}) })
if !result.HasFunctionCallOutput || result.HasToolCallContext || len(callIDs) == 0 || len(referenceIDs) == 0 { if !result.HasFunctionCallOutput || result.HasToolCallContext || len(callIDs) == 0 || len(referenceIDs) == 0 {
return result return result