From 7e62731863c8f21ca32292e3923860e7f45221e9 Mon Sep 17 00:00:00 2001 From: Qilin You <12567356+asmoyou@users.noreply.github.com> Date: Tue, 8 Sep 2026 01:52:45 +0800 Subject: [PATCH] perf: reduce allocations for large Responses image payloads --- backend/internal/pkg/httputil/body.go | 55 ++++++++- .../internal/pkg/httputil/body_memory_test.go | 67 ++++++++++ .../service/openai_gateway_request_body.go | 82 ++++++++++++- .../openai_large_request_memory_test.go | 115 ++++++++++++++++++ ...openai_large_request_rewrite_bench_test.go | 22 ++++ .../service/openai_request_raw_input.go | 41 +++++++ .../openai_responses_ingress_compat.go | 17 +++ .../service/openai_responses_item_id.go | 33 ++--- 8 files changed, 406 insertions(+), 26 deletions(-) create mode 100644 backend/internal/pkg/httputil/body_memory_test.go create mode 100644 backend/internal/service/openai_large_request_memory_test.go create mode 100644 backend/internal/service/openai_large_request_rewrite_bench_test.go create mode 100644 backend/internal/service/openai_request_raw_input.go diff --git a/backend/internal/pkg/httputil/body.go b/backend/internal/pkg/httputil/body.go index e71e42775d..3a6fab4b78 100644 --- a/backend/internal/pkg/httputil/body.go +++ b/backend/internal/pkg/httputil/body.go @@ -83,11 +83,10 @@ func ReadRequestBodyWithPrealloc(req *http.Request) ([]byte, error) { } } - buf := bytes.NewBuffer(make([]byte, 0, capHint)) - if _, err := io.Copy(buf, req.Body); err != nil { + raw, err := readRequestBodyChunks(req.Body, capHint, req.ContentLength) + if err != nil { return nil, err } - raw := buf.Bytes() enc := strings.ToLower(strings.TrimSpace(req.Header.Get("Content-Encoding"))) if enc == "" || enc == "identity" { @@ -106,6 +105,56 @@ func ReadRequestBodyWithPrealloc(req *http.Request) ([]byte, error) { return decoded, nil } +// Read bounded chunks as bytes arrive, then assemble the exact-size result. +// This avoids doubling large buffers or eagerly allocating an untrusted +// Content-Length before the corresponding bytes have arrived. +func readRequestBodyChunks(reader io.Reader, initialCapacity int, contentLength int64) ([]byte, error) { + capacity := initialCapacity + var chunks [][]byte + total := 0 + for { + chunkCapacity := capacity + if remaining := contentLength - int64(total); remaining >= 0 && remaining < int64(chunkCapacity) { + chunkCapacity = int(remaining) + 1 + } + chunk := make([]byte, chunkCapacity) + n := 0 + var err error + for n < len(chunk) && err == nil { + var read int + read, err = reader.Read(chunk[n:]) + n += read + } + if err != nil && err != io.EOF { + return nil, err + } + if n > 0 { + chunks = append(chunks, chunk[:n]) + total += n + } + if err != nil { + if len(chunks) == 0 { + return chunk[:0], nil + } + if len(chunks) == 1 { + return chunks[0], nil + } + body := make([]byte, total) + offset := 0 + for _, part := range chunks { + offset += copy(body[offset:], part) + } + return body, nil + } + if capacity < requestBodyReadMaxInitCap { + capacity *= 2 + if capacity > requestBodyReadMaxInitCap { + capacity = requestBodyReadMaxInitCap + } + } + } +} + // ReadLenientJSONRequestBodyWithPrealloc reads a request body and normalizes // JSON string control bytes before strict validation. func ReadLenientJSONRequestBodyWithPrealloc(req *http.Request, maxNormalizedBytes int64) ([]byte, error) { diff --git a/backend/internal/pkg/httputil/body_memory_test.go b/backend/internal/pkg/httputil/body_memory_test.go new file mode 100644 index 0000000000..b5240998fa --- /dev/null +++ b/backend/internal/pkg/httputil/body_memory_test.go @@ -0,0 +1,67 @@ +package httputil + +import ( + "bytes" + "errors" + "fmt" + "io" + "net/http" + "testing" + "testing/iotest" +) + +type bodyReaderOnly struct{ io.Reader } + +func TestReadRequestBodyChunksPreservesContent(t *testing.T) { + for _, size := range []int{0, 1, 512, 513, 1 << 20, (2 << 20) + 17} { + body := bytes.Repeat([]byte{'x'}, size) + for _, declared := range []int64{-1, 0, 1, int64(size), 1 << 40} { + t.Run(fmt.Sprintf("size_%d_length_%d", size, declared), func(t *testing.T) { + req := &http.Request{Body: io.NopCloser(bodyReaderOnly{bytes.NewReader(body)}), ContentLength: declared, Header: make(http.Header)} + got, err := ReadRequestBodyWithPrealloc(req) + if err != nil || !bytes.Equal(got, body) { + t.Fatalf("body was lost or changed: %v", err) + } + }) + } + } +} + +func TestReadRequestBodyChunksPreservesReadErrors(t *testing.T) { + for _, readErr := range []error{io.ErrUnexpectedEOF, errors.New("connection reset")} { + reader := io.MultiReader(bytes.NewReader([]byte("partial")), iotest.ErrReader(readErr)) + req := &http.Request{Body: io.NopCloser(reader), ContentLength: 100, Header: make(http.Header)} + got, err := ReadRequestBodyWithPrealloc(req) + if !errors.Is(err, readErr) || got != nil { + t.Fatalf("truncated body accepted: body=%q err=%v", got, err) + } + } +} + +func TestReadRequestBodyChunksPreservesMaxBytesReader(t *testing.T) { + body := bytes.Repeat([]byte{'x'}, (2<<20)+1) + req := &http.Request{Body: http.MaxBytesReader(nil, io.NopCloser(bytes.NewReader(body)), 2<<20), ContentLength: int64(len(body)), Header: make(http.Header)} + got, err := ReadRequestBodyWithPrealloc(req) + var limitErr *http.MaxBytesError + if !errors.As(err, &limitErr) || limitErr.Limit != 2<<20 || got != nil { + t.Fatalf("body limit changed: len=%d err=%v", len(got), err) + } +} + +func BenchmarkReadLargeBody(b *testing.B) { + for _, size := range []int{512 << 10, 69 << 20} { + b.Run(fmt.Sprintf("bytes_%d", size), func(b *testing.B) { + body := bytes.Repeat([]byte{'x'}, size) + b.ReportAllocs() + b.SetBytes(int64(size)) + b.ResetTimer() + for i := 0; i < b.N; i++ { + req := &http.Request{Body: io.NopCloser(bodyReaderOnly{bytes.NewReader(body)}), ContentLength: int64(size), Header: make(http.Header)} + got, err := ReadRequestBodyWithPrealloc(req) + if err != nil || !bytes.Equal(got, body) { + b.Fatalf("body changed: %v", err) + } + } + }) + } +} diff --git a/backend/internal/service/openai_gateway_request_body.go b/backend/internal/service/openai_gateway_request_body.go index af4825762b..dc7b30a621 100644 --- a/backend/internal/service/openai_gateway_request_body.go +++ b/backend/internal/service/openai_gateway_request_body.go @@ -8,6 +8,7 @@ import ( "fmt" "net/http" "strings" + "unicode/utf8" "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" "github.com/Wei-Shaw/sub2api/internal/util/urlvalidator" @@ -474,7 +475,7 @@ func openAIRequestBodyHasTools(body []byte) bool { // ids, and opaque extensions). Callers scope this normalization to OpenAI // destinations; compatible providers may still consume their own content. func normalizeOpenAIResponsesReasoningContentReplay(body []byte) ([]byte, bool, error) { - input := gjson.GetBytes(body, "input") + input := parseRawJSONView(body).Get("input") if !input.IsArray() { return body, false, nil } @@ -527,6 +528,85 @@ func normalizeOpenAIResponsesReasoningContentReplay(body []byte) ([]byte, bool, } func normalizeOpenAIAPIKeyStoreFalseReasoningReplay(body []byte, knownStoreFalse bool) ([]byte, bool, error) { + if !knownStoreFalse && gjson.GetBytes(body, "store").Type != gjson.False { + return body, false, nil + } + root := parseRawJSONView(body) + input := root.Get("input") + if !input.IsArray() { + return body, false, nil + } + if !root.IsObject() || !gjson.ValidBytes(body) || !utf8.Valid(body) || hasDuplicateJSONObjectKeys(root) { + return normalizeOpenAIAPIKeyStoreFalseReasoningReplayDecoded(body, knownStoreFalse) + } + + // Only reasoning metadata needs decoding. Keep large image/tool results as + // slices of the original JSON and copy them once into the final request. + items := make([]string, 0) + changed := false + fallback := false + var itemErr error + input.ForEach(func(_, item gjson.Result) bool { + if !item.IsObject() { + items = append(items, item.Raw) + return true + } + if hasDuplicateJSONObjectKeys(item) { + fallback = true + return false + } + typ := strings.TrimSpace(item.Get("type").String()) + id := strings.TrimSpace(item.Get("id").String()) + encrypted := item.Get("encrypted_content") + if (typ == "reasoning" && (encrypted.Type != gjson.String || strings.TrimSpace(encrypted.Str) == "")) || + (typ == "item_reference" && strings.HasPrefix(id, "rs_")) { + changed = true + return true + } + stripID := typ == "reasoning" && strings.HasPrefix(id, "rs_") + addSummary := typ == "reasoning" && item.Get("summary").Type == gjson.Null + stripCallID := shouldStripOpenAIResponsesNonPairCallID(typ) && item.Get("call_id").Exists() + if !stripID && !addSummary && !stripCallID { + items = append(items, item.Raw) + return true + } + var decoded map[string]any + if err := decodeOpenAIJSONUseNumber([]byte(item.Raw), &decoded); err != nil { + itemErr = err + return false + } + if stripID { + delete(decoded, "id") + } + if addSummary { + decoded["summary"] = []any{} + } + if stripCallID { + delete(decoded, "call_id") + } + encoded, err := marshalOpenAIUpstreamJSON(decoded) + if err != nil { + itemErr = err + return false + } + items = append(items, string(encoded)) + changed = true + return true + }) + if fallback { + return normalizeOpenAIAPIKeyStoreFalseReasoningReplayDecoded(body, knownStoreFalse) + } + if itemErr != nil { + return body, false, fmt.Errorf("normalize API-key store=false reasoning replay: %w", itemErr) + } + if !changed { + return body, false, nil + } + return replaceOpenAIRawInput(body, input, items), true, nil +} + +// Preserve the decoder's handling of unusual or duplicate-key input objects. +func normalizeOpenAIAPIKeyStoreFalseReasoningReplayDecoded(body []byte, knownStoreFalse bool) ([]byte, bool, error) { if !knownStoreFalse && gjson.GetBytes(body, "store").Type != gjson.False { return body, false, nil } diff --git a/backend/internal/service/openai_large_request_memory_test.go b/backend/internal/service/openai_large_request_memory_test.go new file mode 100644 index 0000000000..673e8ca7b9 --- /dev/null +++ b/backend/internal/service/openai_large_request_memory_test.go @@ -0,0 +1,115 @@ +package service + +import ( + "bytes" + "runtime" + "testing" + + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func largeNativeResponsesBody(imageBytes int) []byte { + prefix := []byte(`{"model":"gpt-6-astra","store":false,"stream":true,"input":[{"type":"custom_tool_call_output","id":"fc_valid","call_id":"call_valid","output":[{"type":"input_image","image_url":"data:image/png;base64,`) + suffix := []byte(`"}]}]}`) + body := make([]byte, 0, len(prefix)+imageBytes+len(suffix)) + body = append(body, prefix...) + body = append(body, bytes.Repeat([]byte{'A'}, imageBytes)...) + return append(body, suffix...) +} + +func TestLargeNativeResponsesRemainsByteIdentical(t *testing.T) { + body := largeNativeResponsesBody(2 << 20) + for _, fn := range []func([]byte) ([]byte, bool, error){ + normalizeOpenAIResponsesLegacyIngress, + sanitizeOpenAIResponsesInputItemIDs, + normalizeOpenAIResponsesReasoningContentReplay, + func(body []byte) ([]byte, bool, error) { + return normalizeOpenAIAPIKeyStoreFalseReasoningReplay(body, false) + }, + } { + out, changed, err := fn(body) + require.NoError(t, err) + require.False(t, changed) + require.True(t, bytes.Equal(body, out)) + } +} + +func TestRawInputRewritePreservesImagesAndSurroundingFields(t *testing.T) { + body := []byte(" \n\t" + string(largeNativeResponsesBody(2<<20))) + body = bytes.Replace(body, []byte(`"id":"fc_valid"`), []byte(`"id":"bad"`), 1) + body = append(body[:len(body)-1], []byte(`,"opaque":9007199254740993}`)...) + image := gjson.GetBytes(body, "input.0.output.0.image_url").String() + out, changed, err := sanitizeOpenAIResponsesInputItemIDs(body) + require.NoError(t, err) + require.True(t, changed) + require.True(t, gjson.ValidBytes(out)) + require.False(t, gjson.GetBytes(out, "input.0.id").Exists()) + require.Equal(t, image, gjson.GetBytes(out, "input.0.output.0.image_url").String()) + require.Equal(t, "9007199254740993", gjson.GetBytes(out, "opaque").Raw) + require.Equal(t, "call_valid", gjson.GetBytes(out, "input.0.call_id").String()) +} + +func TestStoreFalseRawReplayMatchesDecodedBehavior(t *testing.T) { + fixtures := []string{ + ` {"store":false,"input":[{"type":"reasoning","id":"rs_a","encrypted_content":"opaque","summary":null,"opaque":9007199254740993},{"type":"custom_tool_call_output","id":"fc_a","output":[{"type":"input_image","image_url":"data:image/png;base64,AAAA"}]}],"tail":"keep"} `, + `{"store":false,"input":[null,42,"keep",{"type":"reasoning","encrypted_content":" "},{"type":"item_reference","id":"rs_drop"},{"type":"item_reference","id":"msg_keep"}]}`, + `{"store":false,"input":[{"type":"reasoning","id":"keep","encrypted_content":"opaque","summary":[]},{"type":"message","call_id":null,"content":"keep"}]}`, + `{"store":false,"input":[],"input":[{"type":"reasoning","id":"rs_drop"}]}`, + `{"store":false,"input":[{"type":"message","type":"reasoning","id":"rs_drop"}]}`, + `{"store":false,"input":[{"type":"reasoning","id":"rs_drop"}]} trailing`, + `{"store":false,"input":[{"type":"reasoning","id":"rs_drop"}]`, + `{"store":false,"input":[]}`, + } + for _, fixture := range fixtures { + body := []byte(fixture) + want, wantChanged, wantErr := normalizeOpenAIAPIKeyStoreFalseReasoningReplayDecoded(body, false) + got, changed, err := normalizeOpenAIAPIKeyStoreFalseReasoningReplay(body, false) + require.Equal(t, wantErr != nil, err != nil) + if wantErr != nil { + continue + } + require.Equal(t, wantChanged, changed) + require.JSONEq(t, string(want), string(got)) + } +} + +func TestLegacyIngressNativeFastPathKeepsValidation(t *testing.T) { + for _, body := range []string{`{"input":[]} trailing`, `{"input":[}`, `[{"input":[]}]`} { + _, _, err := normalizeOpenAIResponsesLegacyIngress([]byte(body)) + require.Error(t, err) + } + body := []byte(`{"model":"gpt-5.4","input":"keep","mess\u0061ges":[]}`) + out, changed, err := normalizeOpenAIResponsesLegacyIngress(body) + require.NoError(t, err) + require.True(t, changed) + require.False(t, gjson.GetBytes(out, "messages").Exists()) + require.Equal(t, "keep", gjson.GetBytes(out, "input").String()) +} + +func BenchmarkOpenAILargeNativeRequest(b *testing.B) { + body := largeNativeResponsesBody(69 << 20) + for _, tc := range []struct { + name string + fn func([]byte) ([]byte, bool, error) + }{ + {"legacy_ingress", normalizeOpenAIResponsesLegacyIngress}, + {"input_ids", sanitizeOpenAIResponsesInputItemIDs}, + {"reasoning_replay", normalizeOpenAIResponsesReasoningContentReplay}, + {"api_key_store_false", func(body []byte) ([]byte, bool, error) { + return normalizeOpenAIAPIKeyStoreFalseReasoningReplay(body, false) + }}, + } { + b.Run(tc.name, func(b *testing.B) { + b.ReportAllocs() + b.SetBytes(int64(len(body))) + for i := 0; i < b.N; i++ { + out, changed, err := tc.fn(body) + if err != nil || changed || !bytes.Equal(out, body) { + b.Fatalf("native image request changed: changed=%v err=%v", changed, err) + } + runtime.KeepAlive(out) + } + }) + } +} diff --git a/backend/internal/service/openai_large_request_rewrite_bench_test.go b/backend/internal/service/openai_large_request_rewrite_bench_test.go new file mode 100644 index 0000000000..9c6082bbd2 --- /dev/null +++ b/backend/internal/service/openai_large_request_rewrite_bench_test.go @@ -0,0 +1,22 @@ +package service + +import ( + "bytes" + "runtime" + "testing" +) + +func BenchmarkOpenAIStoreFalseMetadataRewrite(b *testing.B) { + body := largeNativeResponsesBody(69 << 20) + body = bytes.Replace(body, []byte(`"input":[`), []byte(`"input":[{"type":"reasoning","id":"rs_server_id","encrypted_content":"opaque","summary":[]},`), 1) + b.ReportAllocs() + b.SetBytes(int64(len(body))) + b.ResetTimer() + for i := 0; i < b.N; i++ { + out, changed, err := normalizeOpenAIAPIKeyStoreFalseReasoningReplay(body, false) + if err != nil || !changed || len(out) < 69<<20 { + b.Fatalf("metadata rewrite failed or image was lost: changed=%v err=%v", changed, err) + } + runtime.KeepAlive(out) + } +} diff --git a/backend/internal/service/openai_request_raw_input.go b/backend/internal/service/openai_request_raw_input.go new file mode 100644 index 0000000000..1ed88cb4eb --- /dev/null +++ b/backend/internal/service/openai_request_raw_input.go @@ -0,0 +1,41 @@ +package service + +import "github.com/tidwall/gjson" + +// replaceOpenAIRawInput copies unchanged image strings only into the final +// request, without allocating a second complete input array first. +func replaceOpenAIRawInput(body []byte, input gjson.Result, items []string) []byte { + size := len(body) - len(input.Raw) + 2 + for index, item := range items { + size += len(item) + if index > 0 { + size++ + } + } + result := make([]byte, 0, size) + result = append(result, body[:input.Index]...) + result = append(result, '[') + for index, item := range items { + if index > 0 { + result = append(result, ',') + } + result = append(result, item...) + } + result = append(result, ']') + return append(result, body[input.Index+len(input.Raw):]...) +} + +// The standard decoder keeps the last duplicate key; GJSON selects the first. +func hasDuplicateJSONObjectKeys(object gjson.Result) bool { + seen := make(map[string]struct{}) + duplicate := false + object.ForEach(func(key, _ gjson.Result) bool { + if _, exists := seen[key.Str]; exists { + duplicate = true + return false + } + seen[key.Str] = struct{}{} + return true + }) + return duplicate +} diff --git a/backend/internal/service/openai_responses_ingress_compat.go b/backend/internal/service/openai_responses_ingress_compat.go index f3162b3b02..b5a289dd3e 100644 --- a/backend/internal/service/openai_responses_ingress_compat.go +++ b/backend/internal/service/openai_responses_ingress_compat.go @@ -6,6 +6,7 @@ import ( "strings" "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat" + "github.com/tidwall/gjson" ) // normalizeOpenAIResponsesLegacyIngress accepts the Chat Completions-shaped @@ -15,6 +16,22 @@ func normalizeOpenAIResponsesLegacyIngress(body []byte) ([]byte, bool, error) { if len(body) == 0 { return body, false, nil } + // Native requests can contain large image data URLs. Decode the full object + // only when a legacy top-level field needs compatibility handling. + view := parseRawJSONView(body) + if view.IsObject() && gjson.ValidBytes(body) { + hasLegacyField := false + view.ForEach(func(key, _ gjson.Result) bool { + switch key.Str { + case "messages", "prompt", "commands": + hasLegacyField = true + } + return !hasLegacyField + }) + if !hasLegacyField { + return body, false, nil + } + } var request map[string]any if err := decodeOpenAIJSONUseNumber(body, &request); err != nil { diff --git a/backend/internal/service/openai_responses_item_id.go b/backend/internal/service/openai_responses_item_id.go index 3a10966e89..cb1522e4fb 100644 --- a/backend/internal/service/openai_responses_item_id.go +++ b/backend/internal/service/openai_responses_item_id.go @@ -63,20 +63,20 @@ func shouldStripOpenAIResponsesNonPairCallID(itemType string) bool { } func sanitizeOpenAIResponsesInputItemIDs(body []byte) ([]byte, bool, error) { - input := gjson.GetBytes(body, "input") + input := parseRawJSONView(body).Get("input") if !input.IsArray() { return body, false, nil } type inputItem struct { - body []byte + raw string stripID bool stripCallID bool } items := make([]inputItem, 0) input.ForEach(func(_, item gjson.Result) bool { - parsed := inputItem{body: []byte(item.Raw)} + parsed := inputItem{raw: item.Raw} if item.IsObject() { itemType := item.Get("type") id := item.Get("id") @@ -100,9 +100,13 @@ func sanitizeOpenAIResponsesInputItemIDs(body []byte) ([]byte, bool, error) { return body, false, nil } - rebuiltItems := make([][]byte, 0, len(items)) + rebuiltItems := make([]string, 0, len(items)) for index, item := range items { - itemBody := item.body + if !item.stripID && !item.stripCallID { + rebuiltItems = append(rebuiltItems, item.raw) + continue + } + itemBody := []byte(item.raw) if item.stripID { var err error itemBody, err = sjson.DeleteBytes(itemBody, "id") @@ -117,22 +121,7 @@ func sanitizeOpenAIResponsesInputItemIDs(body []byte) ([]byte, bool, error) { return nil, false, fmt.Errorf("delete input.%d.call_id: %w", index, err) } } - rebuiltItems = append(rebuiltItems, itemBody) + rebuiltItems = append(rebuiltItems, string(itemBody)) } - - rebuiltInput := make([]byte, 0, len(input.Raw)) - rebuiltInput = append(rebuiltInput, '[') - for i, item := range rebuiltItems { - if i > 0 { - rebuiltInput = append(rebuiltInput, ',') - } - rebuiltInput = append(rebuiltInput, item...) - } - rebuiltInput = append(rebuiltInput, ']') - - sanitized, err := sjson.SetRawBytes(body, "input", rebuiltInput) - if err != nil { - return nil, false, fmt.Errorf("replace sanitized input: %w", err) - } - return sanitized, true, nil + return replaceOpenAIRawInput(body, input, rebuiltItems), true, nil }