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 {
var builder strings.Builder
builder.WriteString(`{"system":[`)
_, _ = builder.WriteString(`{"system":[`)
for i := 0; i < parts; i++ {
if i > 0 {
builder.WriteByte(',')
_ = builder.WriteByte(',')
}
builder.WriteString(`{"text":"system_part_`)
builder.WriteString(strconv.Itoa(i))
builder.WriteString(`","cache_control":{"type":"ephemeral"}}`)
_, _ = builder.WriteString(`{"text":"system_part_`)
_, _ = builder.WriteString(strconv.Itoa(i))
_, _ = builder.WriteString(`","cache_control":{"type":"ephemeral"}}`)
}
builder.WriteString(`]}`)
_, _ = builder.WriteString(`]}`)
parsed, err := ParseGatewayRequest(NewRequestBodyRef([]byte(builder.String())), "")
if err != nil {
panic(err)
@ -249,93 +249,93 @@ func buildSystemCacheableRequest(parts int) *ParsedRequest {
func buildLargeAnthropicMessagesBody(targetBytes int, includeCacheControl bool) []byte {
var builder strings.Builder
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++ {
if i > 0 {
builder.WriteByte(',')
_ = builder.WriteByte(',')
}
builder.WriteString(`{"role":"user","content":[{"type":"text","text":"`)
builder.WriteString(strings.Repeat("anthropic payload ", 64))
builder.WriteString(strconv.Itoa(i))
builder.WriteByte('"')
_, _ = builder.WriteString(`{"role":"user","content":[{"type":"text","text":"`)
_, _ = builder.WriteString(strings.Repeat("anthropic payload ", 64))
_, _ = builder.WriteString(strconv.Itoa(i))
_ = builder.WriteByte('"')
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())
}
func buildLargeGeminiContentsBody(targetBytes int) []byte {
var builder strings.Builder
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++ {
if i > 0 {
builder.WriteByte(',')
_ = builder.WriteByte(',')
}
builder.WriteString(`{"role":"user","parts":[{"text":"`)
builder.WriteString(strings.Repeat("gemini payload ", 64))
builder.WriteString(strconv.Itoa(i))
builder.WriteString(`"}]}`)
_, _ = builder.WriteString(`{"role":"user","parts":[{"text":"`)
_, _ = builder.WriteString(strings.Repeat("gemini payload ", 64))
_, _ = builder.WriteString(strconv.Itoa(i))
_, _ = builder.WriteString(`"}]}`)
}
builder.WriteString(`]}`)
_, _ = builder.WriteString(`]}`)
return []byte(builder.String())
}
func buildLargeOpenAIResponsesBody(targetBytes int) []byte {
var builder strings.Builder
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++ {
if i > 0 {
builder.WriteByte(',')
_ = builder.WriteByte(',')
}
builder.WriteString(`{"type":"message","role":"user","content":[{"type":"input_text","text":"`)
builder.WriteString(strings.Repeat("openai responses payload ", 48))
builder.WriteString(strconv.Itoa(i))
builder.WriteString(`"}]}`)
_, _ = builder.WriteString(`{"type":"message","role":"user","content":[{"type":"input_text","text":"`)
_, _ = builder.WriteString(strings.Repeat("openai responses payload ", 48))
_, _ = builder.WriteString(strconv.Itoa(i))
_, _ = 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())
}
func buildLargeOpenAIResponsesToolContinuationBody(targetBytes int) []byte {
var builder strings.Builder
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++ {
if i > 0 {
builder.WriteByte(',')
_ = builder.WriteByte(',')
}
callID := "call_" + strconv.Itoa(i)
builder.WriteString(`{"type":"item_reference","id":"`)
builder.WriteString(callID)
builder.WriteString(`"},{"type":"function_call_output","call_id":"`)
builder.WriteString(callID)
builder.WriteString(`","output":"`)
builder.WriteString(strings.Repeat("tool output payload ", 48))
builder.WriteString(strconv.Itoa(i))
builder.WriteString(`"}`)
_, _ = builder.WriteString(`{"type":"item_reference","id":"`)
_, _ = builder.WriteString(callID)
_, _ = builder.WriteString(`"},{"type":"function_call_output","call_id":"`)
_, _ = builder.WriteString(callID)
_, _ = builder.WriteString(`","output":"`)
_, _ = builder.WriteString(strings.Repeat("tool output payload ", 48))
_, _ = builder.WriteString(strconv.Itoa(i))
_, _ = builder.WriteString(`"}`)
}
builder.WriteString(`]}`)
_, _ = builder.WriteString(`]}`)
return []byte(builder.String())
}
func buildLargeOpenAIResponsesImageToolBody(targetBytes int) []byte {
var builder strings.Builder
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++ {
if i > 0 {
builder.WriteByte(',')
_ = builder.WriteByte(',')
}
builder.WriteString(`{"type":"message","role":"user","content":[{"type":"input_text","text":"`)
builder.WriteString(strings.Repeat("openai image billing payload ", 48))
builder.WriteString(strconv.Itoa(i))
builder.WriteString(`"}]}`)
_, _ = builder.WriteString(`{"type":"message","role":"user","content":[{"type":"input_text","text":"`)
_, _ = builder.WriteString(strings.Repeat("openai image billing payload ", 48))
_, _ = builder.WriteString(strconv.Itoa(i))
_, _ = builder.WriteString(`"}]}`)
}
builder.WriteString(`]}`)
_, _ = builder.WriteString(`]}`)
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")
}
} 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)
}
@ -6215,7 +6216,7 @@ func (v *openAIRequestView) MarkPatchSet(path string, value any) {
return
}
path = strings.TrimSpace(path)
if path == "" {
if !isSimpleOpenAIRequestPatchPath(path) {
v.DisablePatches()
return
}
@ -6227,13 +6228,25 @@ func (v *openAIRequestView) MarkPatchDelete(path string) {
return
}
path = strings.TrimSpace(path)
if path == "" {
if !isSimpleOpenAIRequestPatchPath(path) {
v.DisablePatches()
return
}
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() {
if v == nil {
return
@ -6772,7 +6785,7 @@ func openAIRequestBodyMayContainImageInput(body []byte) bool {
return false
}
input := gjson.GetBytes(body, "input")
messages := gjson.GetBytes(body, "messages")
messages := gjson.GetBytes(body, "messages.#-1")
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))
}
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) {
view := newOpenAIRequestView([]byte(`{"model":"gpt-5"}`))
view.MarkPatchSet("model", "gpt-5.1")
@ -260,6 +269,42 @@ func TestOpenAIGatewayService_Forward_ImageToolBillingDoesNotForceFullDecode(t *
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) {
gin.SetMode(gin.TestMode)
upstream := &httpUpstreamRecorder{

View File

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