fix(openai): avoid JSON copies in websocket responses tool-call repair path

Closes: #4925
2026-09-07 00:24:17 +08:00 · 2026-08-12 20:32:10 +08:00
parent f43aad7637
commit 75d2c4a4b4
2 changed files with 249 additions and 38 deletions
--- 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)
--- 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)
 	}
 }