From 75d2c4a4b4e29a4ea0fc3279dee3b097f96380ef Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Wed, 12 Aug 2026 20:32:10 +0800 Subject: [PATCH] fix(openai): avoid JSON copies in websocket responses tool-call repair path Closes: #4925 --- .../openai/openai_responses_websocket_test.go | 119 ++++++++++++- ...nai_responses_websocket_toolcall_repair.go | 168 ++++++++++++++---- 2 files changed, 249 insertions(+), 38 deletions(-) diff --git a/sdk/api/handlers/openai/openai_responses_websocket_test.go b/sdk/api/handlers/openai/openai_responses_websocket_test.go index 0f561e687..029f6f51b 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket_test.go +++ b/sdk/api/handlers/openai/openai_responses_websocket_test.go @@ -1270,6 +1270,120 @@ func TestResponsesWebsocketFallbackTurnBoundsTranscriptAllocations(t *testing.T) } } +func TestResponsesWebsocketToolCacheScansDoNotCopyLargePayloads(t *testing.T) { + const maxAllocatedBytes = 256 << 10 + padding := strings.Repeat("x", 4<<20) + + requestPayload := []byte(fmt.Sprintf( + `{"input":[{"type":"message","id":"message-1","call_id":"not-a-tool","content":%q},{"type":"function_call","id":"fc-1","call_id":"call-1","name":"lookup","arguments":"{}"},{"type":"function_call_output","id":"fco-1","call_id":"call-1","output":"ok"}]}`, + padding, + )) + t.Run("request", func(t *testing.T) { + result := testing.Benchmark(func(b *testing.B) { + b.ReportAllocs() + for index := 0; index < b.N; index++ { + payload, turn := prepareResponsesWebsocketFallbackTurn("large-request-session", requestPayload) + runtime.KeepAlive(payload) + runtime.KeepAlive(turn) + } + }) + t.Logf("request tool-cache scan allocated %d bytes per operation", result.AllocedBytesPerOp()) + if allocatedBytes := result.AllocedBytesPerOp(); allocatedBytes > maxAllocatedBytes { + t.Fatalf("request tool-cache scan allocated %d bytes per operation, want at most %d", allocatedBytes, maxAllocatedBytes) + } + }) + + responsePayload := []byte(fmt.Sprintf( + `{"type":"response.completed","response":{"output":[{"type":"message","id":"message-1","content":%q},{"type":"function_call","id":"fc-1","call_id":"call-1","name":"lookup","arguments":"{}"}]}}`, + padding, + )) + t.Run("response", func(t *testing.T) { + turn := newResponsesWebsocketToolCacheTurn("large-response-session") + result := testing.Benchmark(func(b *testing.B) { + b.ReportAllocs() + for index := 0; index < b.N; index++ { + turn.recordResponse(responsePayload) + runtime.KeepAlive(turn) + } + }) + t.Logf("response tool-cache scan allocated %d bytes per operation", result.AllocedBytesPerOp()) + if allocatedBytes := result.AllocedBytesPerOp(); allocatedBytes > maxAllocatedBytes { + t.Fatalf("response tool-cache scan allocated %d bytes per operation, want at most %d", allocatedBytes, maxAllocatedBytes) + } + }) +} + +func TestResponsesWebsocketToolCacheScanPreservesJSONRequestSemantics(t *testing.T) { + t.Run("rejects trailing data", func(t *testing.T) { + payload := []byte(`{"input":[{"type":"function_call","id":"fc-1","call_id":"call-1","name":"lookup","arguments":"{}"}]} trailing`) + repaired, turn := prepareResponsesWebsocketFallbackTurn("trailing-data-session", payload) + if !bytes.Equal(repaired, payload) { + t.Fatalf("repaired payload = %s, want original malformed payload", repaired) + } + if len(turn.calls) != 0 || len(turn.outputs) != 0 { + t.Fatalf("malformed payload recorded calls=%d outputs=%d, want none", len(turn.calls), len(turn.outputs)) + } + }) + + t.Run("uses last duplicate input", func(t *testing.T) { + payload := []byte(`{"input":[{"type":"function_call","id":"fc-1","call_id":"call-1","name":"lookup","arguments":"{}"}],"input":[{"type":"message","id":"message-1","role":"user","content":"hello"}]}`) + repaired, turn := prepareResponsesWebsocketFallbackTurn("duplicate-input-session", payload) + if !bytes.Equal(repaired, payload) { + t.Fatalf("repaired payload = %s, want original payload", repaired) + } + if len(turn.calls) != 0 || len(turn.outputs) != 0 { + t.Fatalf("duplicate input recorded calls=%d outputs=%d from the shadowed value, want none", len(turn.calls), len(turn.outputs)) + } + }) + + t.Run("repairs last case-insensitive duplicate input", func(t *testing.T) { + payload := []byte(`{"input":[{"type":"message","id":"shadowed","role":"user","content":"ignore"}],"INPUT":[{"type":"function_call_output","id":"fco-1","call_id":"missing-call","output":"orphan"}]}`) + repaired, _ := prepareResponsesWebsocketFallbackTurn("duplicate-case-input-session", payload) + var request struct { + Input []json.RawMessage `json:"input"` + } + if errUnmarshal := json.Unmarshal(repaired, &request); errUnmarshal != nil { + t.Fatalf("unmarshal repaired payload: %v", errUnmarshal) + } + if len(request.Input) != 0 { + t.Fatalf("repaired effective input count = %d, want 0", len(request.Input)) + } + }) + + t.Run("repairs last exact duplicate input", func(t *testing.T) { + payload := []byte(`{"input":[{"type":"message","id":"shadowed","role":"user","content":"ignore"}],"input":[{"type":"function_call_output","id":"fco-1","call_id":"missing-call","output":"orphan"}]}`) + repaired, _ := prepareResponsesWebsocketFallbackTurn("duplicate-exact-input-session", payload) + var request struct { + Input []json.RawMessage `json:"input"` + } + if errUnmarshal := json.Unmarshal(repaired, &request); errUnmarshal != nil { + t.Fatalf("unmarshal repaired payload: %v", errUnmarshal) + } + if len(request.Input) != 0 { + t.Fatalf("repaired effective input count = %d, want 0", len(request.Input)) + } + }) + + t.Run("rejects invalid earlier duplicate input", func(t *testing.T) { + payload := []byte(`{"input":{},"input":[{"type":"function_call","id":"fc-1","call_id":"call-1","name":"lookup","arguments":"{}"},{"type":"function_call_output","id":"fco-1","call_id":"call-1","output":"ok"}]}`) + repaired, turn := prepareResponsesWebsocketFallbackTurn("invalid-duplicate-input-session", payload) + if !bytes.Equal(repaired, payload) { + t.Fatalf("repaired payload = %s, want original payload with invalid duplicate input", repaired) + } + if len(turn.calls) != 0 || len(turn.outputs) != 0 { + t.Fatalf("invalid duplicate input recorded calls=%d outputs=%d, want none", len(turn.calls), len(turn.outputs)) + } + }) + + t.Run("uses last duplicate previous response id", func(t *testing.T) { + payload := []byte(`{"previous_response_id":"resp-first","previous_response_id":null,"input":[{"type":"function_call_output","id":"fco-1","call_id":"missing-call","output":"orphan"}]}`) + repaired, _ := prepareResponsesWebsocketFallbackTurn("duplicate-previous-response-session", payload) + if inputCount := gjson.GetBytes(repaired, "input.#").Int(); inputCount != 0 { + t.Fatalf("repaired input count = %d, want 0 when the last previous_response_id is null", inputCount) + } + }) +} + func TestNormalizeResponsesWebsocketRequestCreateWithHistory(t *testing.T) { lastRequest := []byte(`{"model":"test-model","stream":true,"input":[{"type":"message","id":"msg-1"}]}`) lastResponseOutput := []byte(`[ @@ -1981,13 +2095,16 @@ func TestResponsesWebsocketToolCacheTurnDoesNotRetainRequestBackingStorage(t *te } copy(payload[len(prefix)+paddingSize:], suffix) - _, turn := prepareResponsesWebsocketFallbackTurn("backing-storage-session", payload) + repaired, turn := prepareResponsesWebsocketFallbackTurn("backing-storage-session", payload) payload = nil + repaired = nil + runtime.GC() runtime.GC() var after runtime.MemStats runtime.ReadMemStats(&after) runtime.KeepAlive(turn) + runtime.KeepAlive(repaired) retainedHeap := int64(after.HeapAlloc) - int64(before.HeapAlloc) if retainedHeap > maxRetainedHeap { t.Fatalf("tool cache turn retained %d bytes after request release, want at most %d", retainedHeap, maxRetainedHeap) diff --git a/sdk/api/handlers/openai/openai_responses_websocket_toolcall_repair.go b/sdk/api/handlers/openai/openai_responses_websocket_toolcall_repair.go index 73d189825..ce503dac2 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket_toolcall_repair.go +++ b/sdk/api/handlers/openai/openai_responses_websocket_toolcall_repair.go @@ -8,8 +8,8 @@ import ( "sync" "time" + "github.com/router-for-me/CLIProxyAPI/v7/internal/util" "github.com/tidwall/gjson" - "github.com/tidwall/sjson" ) const ( @@ -248,30 +248,35 @@ func (t *responsesWebsocketToolCacheTurn) recordResponse(payload []byte) { if t == nil || len(payload) == 0 { return } - switch strings.TrimSpace(gjson.GetBytes(payload, "type").String()) { + switch strings.TrimSpace(util.GetGJSONBytesNoCopy(payload, "type").String()) { case "response.completed": - output := gjson.GetBytes(payload, "response.output") + output := util.GetGJSONBytesNoCopy(payload, "response.output") if !output.Exists() || !output.IsArray() { return } - for _, item := range output.Array() { + output.ForEach(func(_, item gjson.Result) bool { if isCompleteResponsesWebsocketToolCall(item) { - t.recordItem(item) + t.recordItem(payload, item) } - } + return true + }) case "response.output_item.added", "response.output_item.done": - item := gjson.GetBytes(payload, "item") + item := util.GetGJSONBytesNoCopy(payload, "item") if isCompleteResponsesWebsocketToolCall(item) { - t.recordItem(item) + t.recordItem(payload, item) } } } -func (t *responsesWebsocketToolCacheTurn) recordItem(item gjson.Result) { +func (t *responsesWebsocketToolCacheTurn) recordItem(payload []byte, item gjson.Result) { if t == nil || !item.Exists() { return } - t.recordRawItem(item.Get("type").String(), item.Get("call_id").String(), []byte(item.Raw)) + rawItem, ok := responsesWebsocketRawMessageForResult(payload, item) + if !ok { + return + } + t.recordRawItem(item.Get("type").String(), item.Get("call_id").String(), rawItem) } func (t *responsesWebsocketToolCacheTurn) recordInputItem(item responsesWebsocketInputItem) { @@ -282,23 +287,25 @@ func (t *responsesWebsocketToolCacheTurn) recordInputItem(item responsesWebsocke } func (t *responsesWebsocketToolCacheTurn) recordRawItem(itemType string, callID string, rawItem []byte) { + if t == nil || (!isResponsesToolCallOutputType(itemType) && !isResponsesToolCallType(itemType)) { + return + } callID = strings.Clone(strings.TrimSpace(callID)) - if t == nil || callID == "" || len(bytes.TrimSpace(rawItem)) == 0 { + if callID == "" || len(bytes.TrimSpace(rawItem)) == 0 { return } raw := append(json.RawMessage(nil), rawItem...) - switch { - case isResponsesToolCallOutputType(itemType): + if isResponsesToolCallOutputType(itemType) { if _, exists := t.outputs[callID]; !exists { t.outputOrder = append(t.outputOrder, callID) } t.outputs[callID] = raw - case isResponsesToolCallType(itemType): - if _, exists := t.calls[callID]; !exists { - t.callOrder = append(t.callOrder, callID) - } - t.calls[callID] = raw + return } + if _, exists := t.calls[callID]; !exists { + t.callOrder = append(t.callOrder, callID) + } + t.calls[callID] = raw } func (t *responsesWebsocketToolCacheTurn) commit() { @@ -363,15 +370,12 @@ func repairResponsesWebsocketToolCallsWithCachesMode( return payload } - var request struct { - Input []json.RawMessage `json:"input"` - PreviousResponseID json.RawMessage `json:"previous_response_id"` - } - if errUnmarshal := json.Unmarshal(payload, &request); errUnmarshal != nil || request.Input == nil { + input, previousResponseID, ok := parseResponsesWebsocketRepairRequest(payload) + if !ok { return payload } - items, errParse := appendResponsesWebsocketRawInputItems(nil, request.Input) - if errParse != nil { + items, rawItems, ok := parseResponsesWebsocketInputItemsNoCopy(payload, input) + if !ok { return payload } @@ -382,12 +386,12 @@ func repairResponsesWebsocketToolCallsWithCachesMode( callCache, sessionKey, items, - repairEnabled && responsesWebsocketMetadataString(request.PreviousResponseID) != "", + repairEnabled && responsesWebsocketMetadataString(previousResponseID) != "", record && repairEnabled, turn, repairEnabled, ) - if errRepair != nil || responsesWebsocketInputItemsEqualRaw(updatedItems, request.Input) { + if errRepair != nil || responsesWebsocketInputItemsEqualRaw(updatedItems, rawItems) { return payload } @@ -395,13 +399,94 @@ func repairResponsesWebsocketToolCallsWithCachesMode( if errMarshal != nil { return payload } - updated, errSet := sjson.SetRawBytes(payload, "input", []byte(updatedRaw)) - if errSet != nil { + updated, ok := replaceResponsesWebsocketRawResult(payload, input, []byte(updatedRaw)) + if !ok { return payload } return updated } +func parseResponsesWebsocketRepairRequest(payload []byte) (gjson.Result, json.RawMessage, bool) { + if !json.Valid(payload) { + return gjson.Result{}, nil, false + } + root := util.ParseGJSONBytesNoCopy(payload) + if !root.IsObject() { + return gjson.Result{}, nil, false + } + + var input gjson.Result + var previousResponseID json.RawMessage + inputFound := false + valid := true + root.ForEach(func(key, value gjson.Result) bool { + switch { + case strings.EqualFold(key.String(), "input"): + if !value.IsArray() && strings.TrimSpace(value.Raw) != "null" { + valid = false + return false + } + input = value + inputFound = true + case strings.EqualFold(key.String(), "previous_response_id"): + var ok bool + previousResponseID, ok = responsesWebsocketRawMessageForResult(payload, value) + if !ok { + valid = false + return false + } + } + return true + }) + if !valid || !inputFound || !input.IsArray() { + return gjson.Result{}, nil, false + } + return input, previousResponseID, true +} + +func replaceResponsesWebsocketRawResult(payload []byte, result gjson.Result, replacement []byte) ([]byte, bool) { + if result.Index < 0 || result.Index > len(payload) || len(result.Raw) > len(payload)-result.Index { + return nil, false + } + updated := make([]byte, 0, len(payload)-len(result.Raw)+len(replacement)) + updated = append(updated, payload[:result.Index]...) + updated = append(updated, replacement...) + updated = append(updated, payload[result.Index+len(result.Raw):]...) + return updated, true +} + +func parseResponsesWebsocketInputItemsNoCopy(payload []byte, input gjson.Result) ([]responsesWebsocketInputItem, []json.RawMessage, bool) { + var items []responsesWebsocketInputItem + var rawItems []json.RawMessage + valid := true + input.ForEach(func(_, itemResult gjson.Result) bool { + rawItem, ok := responsesWebsocketRawMessageForResult(payload, itemResult) + if !ok { + valid = false + return false + } + item, errItem := parseResponsesWebsocketInputItem(rawItem) + if errItem != nil { + valid = false + return false + } + items = append(items, item) + rawItems = append(rawItems, rawItem) + return true + }) + if !valid { + return nil, nil, false + } + return items, rawItems, true +} + +func responsesWebsocketRawMessageForResult(payload []byte, result gjson.Result) (json.RawMessage, bool) { + if result.Index < 0 || result.Index > len(payload) || len(result.Raw) > len(payload)-result.Index { + return nil, false + } + return payload[result.Index : result.Index+len(result.Raw)], true +} + func repairResponsesToolCallItems( outputCache, callCache *websocketToolOutputCache, sessionKey string, @@ -538,27 +623,36 @@ func recordResponsesWebsocketToolCallsFromPayloadWithCache(cache *websocketToolO return } - eventType := strings.TrimSpace(gjson.GetBytes(payload, "type").String()) + eventType := strings.TrimSpace(util.GetGJSONBytesNoCopy(payload, "type").String()) switch eventType { case "response.completed": - output := gjson.GetBytes(payload, "response.output") + output := util.GetGJSONBytesNoCopy(payload, "response.output") if !output.Exists() || !output.IsArray() { return } - for _, item := range output.Array() { + output.ForEach(func(_, item gjson.Result) bool { if !isCompleteResponsesWebsocketToolCall(item) { - continue + return true + } + rawItem, ok := responsesWebsocketRawMessageForResult(payload, item) + if !ok { + return false } callID := strings.TrimSpace(item.Get("call_id").String()) - cache.record(sessionKey, callID, json.RawMessage(item.Raw)) - } + cache.record(sessionKey, callID, rawItem) + return true + }) case "response.output_item.added", "response.output_item.done": - item := gjson.GetBytes(payload, "item") + item := util.GetGJSONBytesNoCopy(payload, "item") if !isCompleteResponsesWebsocketToolCall(item) { return } + rawItem, ok := responsesWebsocketRawMessageForResult(payload, item) + if !ok { + return + } callID := strings.TrimSpace(item.Get("call_id").String()) - cache.record(sessionKey, callID, json.RawMessage(item.Raw)) + cache.record(sessionKey, callID, rawItem) } }