From 8b4fd28c95b9769ef1f61ddc5ca47b7fb9b77abb Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Mon, 20 Jul 2026 13:57:14 +0800 Subject: [PATCH] perf(executor): replace `sjson.SetBytes` with optimized helpers for conditional payload updates - Introduced `SetStringIfDifferent`, `SetBoolIfDifferent`, and similar helpers in `payload_helpers` to avoid unnecessary writes during JSON modifications. - Simplified message patching in `kimi_executor` for reasoning and tool call adjustments. - Updated all executors (Gemini, Codex, XAI, Antigravity, etc.) to use the new helpers, enhancing readability and potentially reducing overhead. - Removed obsolete `filterKimiEmptyAssistantMessages` function in `kimi_executor`. --- .../runtime/executor/antigravity_executor.go | 7 +- .../executor/antigravity_reasoning_replay.go | 20 +- internal/runtime/executor/claude_executor.go | 22 +- internal/runtime/executor/codex_executor.go | 41 +--- .../runtime/executor/codex_openai_images.go | 114 ++++++---- .../executor/codex_websockets_executor.go | 10 +- .../executor_payload_optimization_test.go | 194 ++++++++++++++++++ internal/runtime/executor/gemini_executor.go | 12 +- .../executor/gemini_vertex_executor.go | 8 +- .../runtime/executor/helps/payload_helpers.go | 71 +++++-- .../executor/helps/payload_mutations.go | 81 ++++++++ .../executor/helps/payload_mutations_test.go | 155 ++++++++++++++ .../executor/helps/vertex_payload_helpers.go | 70 +++++-- .../helps/vertex_payload_helpers_test.go | 45 ++++ internal/runtime/executor/kimi_executor.go | 153 ++++++++------ .../executor/openai_compat_executor.go | 9 +- internal/runtime/executor/xai_executor.go | 49 ++--- 17 files changed, 824 insertions(+), 237 deletions(-) create mode 100644 internal/runtime/executor/executor_payload_optimization_test.go create mode 100644 internal/runtime/executor/helps/payload_mutations.go create mode 100644 internal/runtime/executor/helps/payload_mutations_test.go create mode 100644 internal/runtime/executor/helps/vertex_payload_helpers_test.go diff --git a/internal/runtime/executor/antigravity_executor.go b/internal/runtime/executor/antigravity_executor.go index 0f0ca05e8..30b3b8d06 100644 --- a/internal/runtime/executor/antigravity_executor.go +++ b/internal/runtime/executor/antigravity_executor.go @@ -2223,7 +2223,6 @@ func (e *AntigravityExecutor) buildRequest(ctx context.Context, auth *cliproxyau return nil, errProject } payload = geminiToAntigravity(modelName, payload, projectID) - payload, _ = sjson.SetBytes(payload, "model", modelName) // Cap maxOutputTokens to model's max_completion_tokens from registry if maxOut := gjson.GetBytes(payload, "request.generationConfig.maxOutputTokens"); maxOut.Exists() && maxOut.Type == gjson.Number { @@ -2753,8 +2752,8 @@ func resolveCustomAntigravityBaseURL(auth *cliproxyauth.Auth) string { func geminiToAntigravity(modelName string, payload []byte, projectID string) []byte { template := payload - template, _ = sjson.SetBytes(template, "model", modelName) - template, _ = sjson.SetBytes(template, "userAgent", "antigravity") + template = helps.SetStringIfDifferent(template, "model", modelName) + template = helps.SetStringIfDifferent(template, "userAgent", "antigravity") isImageModel := strings.Contains(modelName, "image") reqType := strings.TrimSpace(gjson.GetBytes(template, "requestType").String()) @@ -2768,7 +2767,7 @@ func geminiToAntigravity(modelName string, payload []byte, projectID string) []b } if projectID != "" { - template, _ = sjson.SetBytes(template, "project", projectID) + template = helps.SetStringIfDifferent(template, "project", projectID) } else { template, _ = sjson.DeleteBytes(template, "project") } diff --git a/internal/runtime/executor/antigravity_reasoning_replay.go b/internal/runtime/executor/antigravity_reasoning_replay.go index 8276eadbd..4f0ff6214 100644 --- a/internal/runtime/executor/antigravity_reasoning_replay.go +++ b/internal/runtime/executor/antigravity_reasoning_replay.go @@ -525,12 +525,11 @@ func mergeAntigravityFunctionCallPartReplay(payload []byte, itemResult gjson.Res } type antigravityReasoningReplayAccumulator struct { - scope antigravityReasoningReplayScope - requestPayload []byte - items [][]byte - seenFC map[string]bool - contentIndex int - nextPartIndex int + scope antigravityReasoningReplayScope + items [][]byte + seenFC map[string]bool + contentIndex int + nextPartIndex int } func newAntigravityReasoningReplayAccumulator(scope antigravityReasoningReplayScope, requestPayload []byte) *antigravityReasoningReplayAccumulator { @@ -539,11 +538,10 @@ func newAntigravityReasoningReplayAccumulator(scope antigravityReasoningReplaySc } contentIndex, basePartIndex := antigravityReasoningReplayPendingModelContentIndex(requestPayload) return &antigravityReasoningReplayAccumulator{ - scope: scope, - requestPayload: append([]byte(nil), requestPayload...), - seenFC: make(map[string]bool), - contentIndex: contentIndex, - nextPartIndex: basePartIndex, + scope: scope, + seenFC: make(map[string]bool), + contentIndex: contentIndex, + nextPartIndex: basePartIndex, } } diff --git a/internal/runtime/executor/claude_executor.go b/internal/runtime/executor/claude_executor.go index 8c7bb207f..7edcf8e7e 100644 --- a/internal/runtime/executor/claude_executor.go +++ b/internal/runtime/executor/claude_executor.go @@ -289,7 +289,7 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r originalPayload := originalPayloadSource originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, stream) body := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, stream) - body, _ = sjson.SetBytes(body, "model", upstreamModel) + body = helps.SetStringIfDifferent(body, "model", upstreamModel) body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) if err != nil { @@ -484,7 +484,7 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A originalPayload := originalPayloadSource originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, true) body := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, true) - body, _ = sjson.SetBytes(body, "model", upstreamModel) + body = helps.SetStringIfDifferent(body, "model", upstreamModel) body, err = thinking.ApplyThinking(body, req.Model, from.String(), to.String(), e.Identifier()) if err != nil { @@ -779,7 +779,7 @@ func (e *ClaudeExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Aut // Use streaming translation to preserve function calling, except for claude. stream := from != to body := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, stream) - body, _ = sjson.SetBytes(body, "model", upstreamModel) + body = helps.SetStringIfDifferent(body, "model", upstreamModel) if rebuildMidSystemMessageEnabled(e.cfg, auth) { body = rebuildMidSystemMessagesToTopLevel(body) } @@ -1418,8 +1418,22 @@ func remapOAuthToolNames(body []byte) ([]byte, map[string]string) { // stale snapshot will preserve removals but overwrite renamed names back to their // original lowercase values. tools := gjson.GetBytes(body, "tools") + toolsNeedRewrite := false if tools.Exists() && tools.IsArray() { - + tools.ForEach(func(_, tool gjson.Result) bool { + if tool.Get("type").Exists() && tool.Get("type").String() != "" { + return true + } + name := tool.Get("name").String() + toolsNeedRewrite = oauthToolsToRemove[name] + if !toolsNeedRewrite { + newName, ok := oauthToolRenameMap[name] + toolsNeedRewrite = ok && newName != name + } + return !toolsNeedRewrite + }) + } + if toolsNeedRewrite { var toolsJSON strings.Builder toolsJSON.WriteByte('[') toolCount := 0 diff --git a/internal/runtime/executor/codex_executor.go b/internal/runtime/executor/codex_executor.go index 7abbbc63f..56b01069a 100644 --- a/internal/runtime/executor/codex_executor.go +++ b/internal/runtime/executor/codex_executor.go @@ -1136,8 +1136,8 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re requestedModel := helps.PayloadRequestedModel(opts, req.Model) requestPath := helps.PayloadRequestPath(opts) body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) - body, _ = sjson.SetBytes(body, "model", baseModel) - body, _ = sjson.SetBytes(body, "stream", true) + body = helps.SetStringIfDifferent(body, "model", baseModel) + body = helps.SetBoolIfDifferent(body, "stream", true) body, _ = sjson.DeleteBytes(body, "previous_response_id") body, _ = sjson.DeleteBytes(body, "prompt_cache_retention") body, _ = sjson.DeleteBytes(body, "safety_identifier") @@ -1250,28 +1250,7 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re } publishCodexImageToolUsage(ctx, reporter, body, eventData) - completedData := eventData - outputResult := gjson.GetBytes(completedData, "response.output") - shouldPatchOutput := (!outputResult.Exists() || !outputResult.IsArray() || len(outputResult.Array()) == 0) && (len(outputItemsByIndex) > 0 || len(outputItemsFallback) > 0) - if shouldPatchOutput { - completedDataPatched := completedData - completedDataPatched, _ = sjson.SetRawBytes(completedDataPatched, "response.output", []byte(`[]`)) - - indexes := make([]int64, 0, len(outputItemsByIndex)) - for idx := range outputItemsByIndex { - indexes = append(indexes, idx) - } - sort.Slice(indexes, func(i, j int) bool { - return indexes[i] < indexes[j] - }) - for _, idx := range indexes { - completedDataPatched, _ = sjson.SetRawBytes(completedDataPatched, "response.output.-1", outputItemsByIndex[idx]) - } - for _, item := range outputItemsFallback { - completedDataPatched, _ = sjson.SetRawBytes(completedDataPatched, "response.output.-1", item) - } - completedData = completedDataPatched - } + completedData := patchCodexCompletedOutput(eventData, outputItemsByIndex, outputItemsFallback) if eventType == "response.completed" { cacheCodexReasoningReplayFromCompleted(replayScope, completedData) } @@ -1323,7 +1302,7 @@ func (e *CodexExecutor) executeCompact(ctx context.Context, auth *cliproxyauth.A requestedModel := helps.PayloadRequestedModel(opts, req.Model) requestPath := helps.PayloadRequestPath(opts) body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) - body, _ = sjson.SetBytes(body, "model", baseModel) + body = helps.SetStringIfDifferent(body, "model", baseModel) body, _ = sjson.DeleteBytes(body, "stream") body = normalizeCodexInstructions(body) body = sanitizeOpenAIResponsesReasoningEncryptedContent(ctx, "codex executor", body) @@ -1432,7 +1411,7 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au body, _ = sjson.DeleteBytes(body, "prompt_cache_retention") body, _ = sjson.DeleteBytes(body, "safety_identifier") body, _ = sjson.DeleteBytes(body, "stream_options") - body, _ = sjson.SetBytes(body, "model", baseModel) + body = helps.SetStringIfDifferent(body, "model", baseModel) body = normalizeCodexInstructions(body) if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff { body = ensureImageGenerationTool(body, baseModel, auth, opts.Headers) @@ -1598,12 +1577,12 @@ func (e *CodexExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth return cliproxyexecutor.Response{}, err } - body, _ = sjson.SetBytes(body, "model", baseModel) + body = helps.SetStringIfDifferent(body, "model", baseModel) body, _ = sjson.DeleteBytes(body, "previous_response_id") body, _ = sjson.DeleteBytes(body, "prompt_cache_retention") body, _ = sjson.DeleteBytes(body, "safety_identifier") body, _ = sjson.DeleteBytes(body, "stream_options") - body, _ = sjson.SetBytes(body, "stream", false) + body = helps.SetBoolIfDifferent(body, "stream", false) body = normalizeCodexInstructions(body) enc, err := tokenizerForCodexModel(baseModel) @@ -1828,7 +1807,7 @@ func (e *CodexExecutor) cacheHelper(ctx context.Context, from sdktranslator.Form } if cache.ID != "" { - rawJSON, _ = sjson.SetBytes(rawJSON, "prompt_cache_key", cache.ID) + rawJSON = helps.SetStringIfDifferent(rawJSON, "prompt_cache_key", cache.ID) } rawJSON = helps.SanitizeCodexInputItemIDs(rawJSON) var identityState codexIdentityConfuseState @@ -1855,7 +1834,7 @@ func applyCodexIdentityConfuseBody(cfg *config.Config, auth *cliproxyauth.Auth, if promptCacheKey := strings.TrimSpace(gjson.GetBytes(userPayload, "prompt_cache_key").String()); promptCacheKey != "" { state.originalPromptCacheKey = promptCacheKey state.promptCacheKey = codexIdentityConfuseUUID(auth.ID, "prompt-cache", promptCacheKey) - rawJSON, _ = sjson.SetBytes(rawJSON, "prompt_cache_key", state.promptCacheKey) + rawJSON = helps.SetStringIfDifferent(rawJSON, "prompt_cache_key", state.promptCacheKey) } if installationID := strings.TrimSpace(gjson.GetBytes(userPayload, "client_metadata.x-codex-installation-id").String()); installationID != "" { rawJSON, _ = sjson.SetBytes(rawJSON, "client_metadata.x-codex-installation-id", codexIdentityConfuseUUID(auth.ID, "installation", installationID)) @@ -2193,7 +2172,7 @@ func ensureImageGenerationTool(body []byte, baseModel string, auth *cliproxyauth func normalizeCodexParallelToolCalls(body []byte, headers http.Header) []byte { if isCodexResponsesLiteRequest(body, headers) { - body, _ = sjson.SetBytes(body, "parallel_tool_calls", false) + body = helps.SetBoolIfDifferent(body, "parallel_tool_calls", false) return body } return normalizeCodexParallelToolCallsForTools(body) diff --git a/internal/runtime/executor/codex_openai_images.go b/internal/runtime/executor/codex_openai_images.go index 10019f0cd..6a514a6a8 100644 --- a/internal/runtime/executor/codex_openai_images.go +++ b/internal/runtime/executor/codex_openai_images.go @@ -550,13 +550,6 @@ func codexRewriteOpenAIImageEditMultipartToJSON(payload []byte, model string, bo out = codexSetOpenAIImageEditFormValues(out, key, values) } - for _, fileHeader := range codexMultipartImageFiles(form) { - dataURL, errData := codexMultipartFileToDataURL(fileHeader) - if errData != nil { - return nil, "", errData - } - out, _ = sjson.SetBytes(out, "images.-1.image_url", dataURL) - } if maskFiles := form.File["mask"]; len(maskFiles) > 0 && maskFiles[0] != nil { dataURL, errData := codexMultipartFileToDataURL(maskFiles[0]) if errData != nil { @@ -565,6 +558,35 @@ func codexRewriteOpenAIImageEditMultipartToJSON(payload []byte, model string, bo out, _ = sjson.SetBytes(out, "mask.image_url", dataURL) } + imageFiles := codexMultipartImageFiles(form) + if existingImages := gjson.GetBytes(out, "images"); !existingImages.Exists() || existingImages.IsArray() { + existingItems := existingImages.Array() + imageItems := make([][]byte, 0, len(existingItems)+len(imageFiles)) + for _, image := range existingItems { + imageItems = append(imageItems, []byte(image.Raw)) + } + for _, fileHeader := range imageFiles { + dataURL, errData := codexMultipartFileToDataURL(fileHeader) + if errData != nil { + return nil, "", errData + } + item := []byte(`{"image_url":""}`) + item, _ = sjson.SetBytes(item, "image_url", dataURL) + imageItems = append(imageItems, item) + } + if len(imageFiles) > 0 { + out, _ = sjson.SetRawBytes(out, "images", helps.JoinRawJSONArray(imageItems)) + } + } else { + for _, fileHeader := range imageFiles { + dataURL, errData := codexMultipartFileToDataURL(fileHeader) + if errData != nil { + return nil, "", errData + } + out, _ = sjson.SetBytes(out, "images.-1.image_url", dataURL) + } + } + return out, "application/json", nil } @@ -579,11 +601,11 @@ func codexSetOpenAIImageEditFormValues(out []byte, key string, values []string) if len(values) == 1 { return codexSetOpenAIImageEditFormValue(out, path, values[0]) } - out, _ = sjson.SetRawBytes(out, path, []byte(`[]`)) + items := make([][]byte, 0, len(values)) for _, value := range values { - item := codexOpenAIImageEditFormJSONValue(key, value) - out, _ = sjson.SetRawBytes(out, path+".-1", item) + items = append(items, codexOpenAIImageEditFormJSONValue(key, value)) } + out, _ = sjson.SetRawBytes(out, path, helps.JoinRawJSONArray(items)) return out } @@ -660,8 +682,8 @@ func (e *CodexExecutor) prepareCodexOpenAIImageBody(body []byte, req cliproxyexe requestedModel := helps.PayloadRequestedModel(opts, req.Model) requestPath := helps.PayloadRequestPath(opts) out = helps.ApplyPayloadConfigWithRequest(e.cfg, mainModel, "codex", codexOpenAIImageSourceFormat, "", out, body, requestedModel, requestPath, opts.Headers) - out, _ = sjson.SetBytes(out, "model", mainModel) - out, _ = sjson.SetBytes(out, "stream", true) + out = helps.SetStringIfDifferent(out, "model", mainModel) + out = helps.SetBoolIfDifferent(out, "stream", true) out, _ = sjson.DeleteBytes(out, "previous_response_id") out, _ = sjson.DeleteBytes(out, "prompt_cache_retention") out, _ = sjson.DeleteBytes(out, "safety_identifier") @@ -854,27 +876,38 @@ func codexBuildOpenAIImageTool(rawJSON []byte, routeModel string, action string, } func codexBuildImagesResponsesRequest(prompt string, images []string, toolJSON []byte) []byte { - req := []byte(`{"instructions":"","stream":true,"reasoning":{"effort":"medium","summary":"auto"},"parallel_tool_calls":true,"include":["reasoning.encrypted_content"],"model":"","store":false,"tool_choice":{"type":"image_generation"}}`) + req := []byte(`{"instructions":"","stream":true,"reasoning":{"effort":"medium","summary":"auto"},"parallel_tool_calls":true,"include":["reasoning.encrypted_content"],"model":"","store":false,"tool_choice":{"type":"image_generation"},"tools":[]}`) req, _ = sjson.SetBytes(req, "model", codexOpenAIImagesMainModel) + if len(toolJSON) > 0 && json.Valid(toolJSON) { + req, _ = sjson.SetRawBytes(req, "tools", helps.JoinRawJSONArray([][]byte{toolJSON})) + } - input := []byte(`[{"type":"message","role":"user","content":[{"type":"input_text","text":""}]}]`) - input, _ = sjson.SetBytes(input, "0.content.0.text", prompt) - contentIndex := 1 + textPart := []byte(`{"type":"input_text","text":""}`) + textPart, _ = sjson.SetBytes(textPart, "text", prompt) + contentItems := make([][]byte, 0, len(images)+1) + contentItems = append(contentItems, textPart) for _, img := range images { if strings.TrimSpace(img) == "" { continue } part := []byte(`{"type":"input_image","image_url":""}`) part, _ = sjson.SetBytes(part, "image_url", img) - input, _ = sjson.SetRawBytes(input, fmt.Sprintf("0.content.%d", contentIndex), part) - contentIndex++ + contentItems = append(contentItems, part) } + inputSize := len(`[{"type":"message","role":"user","content":[]}]`) + len(contentItems) + for _, item := range contentItems { + inputSize += len(item) + } + input := make([]byte, 0, inputSize) + input = append(input, `[{"type":"message","role":"user","content":[`...) + for index, item := range contentItems { + if index > 0 { + input = append(input, ',') + } + input = append(input, item...) + } + input = append(input, ']', '}', ']') req, _ = sjson.SetRawBytes(req, "input", input) - - req, _ = sjson.SetRawBytes(req, "tools", []byte(`[]`)) - if len(toolJSON) > 0 && json.Valid(toolJSON) { - req, _ = sjson.SetRawBytes(req, "tools.-1", toolJSON) - } return req } @@ -998,19 +1031,6 @@ func codexExtractImageResults(completed []byte, itemsByIndex map[int64][]byte, f func codexBuildImagesAPIResponse(results []codexImageCallResult, createdAt int64, usageRaw []byte, firstMeta codexImageCallResult, responseFormat string) ([]byte, error) { out := []byte(`{"created":0,"data":[]}`) out, _ = sjson.SetBytes(out, "created", createdAt) - responseFormat = codexNormalizeImageResponseFormat(responseFormat) - for _, img := range results { - item := []byte(`{}`) - if responseFormat == "url" { - item, _ = sjson.SetBytes(item, "url", "data:"+codexMimeTypeFromOutputFormat(img.OutputFormat)+";base64,"+img.Result) - } else { - item, _ = sjson.SetBytes(item, "b64_json", img.Result) - } - if img.RevisedPrompt != "" { - item, _ = sjson.SetBytes(item, "revised_prompt", img.RevisedPrompt) - } - out, _ = sjson.SetRawBytes(out, "data.-1", item) - } if firstMeta.Background != "" { out, _ = sjson.SetBytes(out, "background", firstMeta.Background) } @@ -1026,6 +1046,22 @@ func codexBuildImagesAPIResponse(results []codexImageCallResult, createdAt int64 if len(usageRaw) > 0 && json.Valid(usageRaw) { out, _ = sjson.SetRawBytes(out, "usage", usageRaw) } + + responseFormat = codexNormalizeImageResponseFormat(responseFormat) + items := make([][]byte, 0, len(results)) + for _, img := range results { + item := []byte(`{}`) + if img.RevisedPrompt != "" { + item, _ = sjson.SetBytes(item, "revised_prompt", img.RevisedPrompt) + } + if responseFormat == "url" { + item, _ = sjson.SetBytes(item, "url", "data:"+codexMimeTypeFromOutputFormat(img.OutputFormat)+";base64,"+img.Result) + } else { + item, _ = sjson.SetBytes(item, "b64_json", img.Result) + } + items = append(items, item) + } + out, _ = sjson.SetRawBytes(out, "data", helps.JoinRawJSONArray(items)) return out, nil } @@ -1051,14 +1087,14 @@ func codexBuildImageCompletedFrame(img codexImageCallResult, usageRaw []byte, re eventName := strings.TrimSpace(streamPrefix) + ".completed" data := []byte(`{"type":""}`) data, _ = sjson.SetBytes(data, "type", eventName) + if len(usageRaw) > 0 && json.Valid(usageRaw) { + data, _ = sjson.SetRawBytes(data, "usage", usageRaw) + } if codexNormalizeImageResponseFormat(responseFormat) == "url" { data, _ = sjson.SetBytes(data, "url", "data:"+codexMimeTypeFromOutputFormat(img.OutputFormat)+";base64,"+img.Result) } else { data, _ = sjson.SetBytes(data, "b64_json", img.Result) } - if len(usageRaw) > 0 && json.Valid(usageRaw) { - data, _ = sjson.SetRawBytes(data, "usage", usageRaw) - } return codexBuildSSEFrame(eventName, data) } diff --git a/internal/runtime/executor/codex_websockets_executor.go b/internal/runtime/executor/codex_websockets_executor.go index 5cbd4f01b..823726cb9 100644 --- a/internal/runtime/executor/codex_websockets_executor.go +++ b/internal/runtime/executor/codex_websockets_executor.go @@ -266,8 +266,8 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut requestedModel := helps.PayloadRequestedModel(opts, req.Model) requestPath := helps.PayloadRequestPath(opts) body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) - body, _ = sjson.SetBytes(body, "model", baseModel) - body, _ = sjson.SetBytes(body, "stream", true) + body = helps.SetStringIfDifferent(body, "model", baseModel) + body = helps.SetBoolIfDifferent(body, "stream", true) body, _ = sjson.DeleteBytes(body, "prompt_cache_retention") body, _ = sjson.DeleteBytes(body, "safety_identifier") body = normalizeCodexInstructions(body) @@ -515,7 +515,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr requestedModel := helps.PayloadRequestedModel(opts, req.Model) requestPath := helps.PayloadRequestPath(opts) body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) - body, _ = sjson.SetBytes(body, "model", baseModel) + body = helps.SetStringIfDifferent(body, "model", baseModel) body = normalizeCodexInstructions(body) if e.cfg == nil || e.cfg.DisableImageGeneration == config.DisableImageGenerationOff { body = ensureImageGenerationTool(body, baseModel, auth, opts.Headers) @@ -862,7 +862,7 @@ func normalizeCodexWebsocketParallelToolCalls(body []byte, headers http.Header) if !isCodexResponsesLiteRequest(body, headers) { return body } - body, _ = sjson.SetBytes(body, "parallel_tool_calls", false) + body = helps.SetBoolIfDifferent(body, "parallel_tool_calls", false) return body } @@ -1035,7 +1035,7 @@ func applyCodexPromptCacheHeadersWithContext(ctx context.Context, from sdktransl } if cache.ID != "" { - rawJSON, _ = sjson.SetBytes(rawJSON, "prompt_cache_key", cache.ID) + rawJSON = helps.SetStringIfDifferent(rawJSON, "prompt_cache_key", cache.ID) setHeaderCasePreserved(headers, "session_id", cache.ID) headers.Set("Conversation_id", cache.ID) } diff --git a/internal/runtime/executor/executor_payload_optimization_test.go b/internal/runtime/executor/executor_payload_optimization_test.go new file mode 100644 index 000000000..c60b934ff --- /dev/null +++ b/internal/runtime/executor/executor_payload_optimization_test.go @@ -0,0 +1,194 @@ +package executor + +import ( + "bytes" + "mime/multipart" + "strings" + "testing" + + "github.com/tidwall/gjson" +) + +func TestEnsureColonSpacedJSONLeavesInvalidPayloadUnchanged(t *testing.T) { + input := []byte(`{"text":"unterminated}`) + output := ensureColonSpacedJSON(input) + if &output[0] != &input[0] || string(output) != string(input) { + t.Fatal("invalid JSON payload changed") + } +} + +func TestNormalizeKimiToolMessageLinksReusesCanonicalPayload(t *testing.T) { + input := []byte(`{"messages":[{"role":"assistant","reasoning_content":"checking","tool_calls":[{"id":"call_1","type":"function","function":{"name":"lookup","arguments":"{}"}}]},{"role":"tool","tool_call_id":"call_1","content":"ok"}]}`) + output, errNormalize := normalizeKimiToolMessageLinks(input) + if errNormalize != nil { + t.Fatalf("normalizeKimiToolMessageLinks returned error: %v", errNormalize) + } + if &output[0] != &input[0] { + t.Fatal("canonical Kimi tool history was copied") + } +} + +func TestNormalizeKimiToolMessageLinksPreservesLargeArguments(t *testing.T) { + input := []byte(`{"messages":[{"role":"assistant","content":"lookup","tool_calls":[{"id":"call_1","type":"function","function":{"name":"lookup","arguments":{"id":9007199254740993}}}]},{"role":"tool","call_id":"call_1","content":"ok"}]}`) + output, errNormalize := normalizeKimiToolMessageLinks(input) + if errNormalize != nil { + t.Fatalf("normalizeKimiToolMessageLinks returned error: %v", errNormalize) + } + if got := gjson.GetBytes(output, "messages.0.tool_calls.0.function.arguments.id").Raw; got != "9007199254740993" { + t.Fatalf("argument id = %s, want exact large integer", got) + } + if got := gjson.GetBytes(output, "messages.1.tool_call_id").String(); got != "call_1" { + t.Fatalf("tool_call_id = %q, want call_1", got) + } + if got := gjson.GetBytes(output, "messages.0.reasoning_content").String(); got != "lookup" { + t.Fatalf("reasoning_content = %q, want lookup", got) + } +} + +func TestCodexMultipartImageEditAppendsExistingImages(t *testing.T) { + var body bytes.Buffer + writer := multipart.NewWriter(&body) + for _, value := range []string{"existing-1", "existing-2"} { + if errWrite := writer.WriteField("images", value); errWrite != nil { + t.Fatalf("write images field: %v", errWrite) + } + } + imagePart, errCreate := writer.CreateFormFile("image[]", "source.png") + if errCreate != nil { + t.Fatalf("create image field: %v", errCreate) + } + if _, errWrite := imagePart.Write([]byte("png-data")); errWrite != nil { + t.Fatalf("write image data: %v", errWrite) + } + if errClose := writer.Close(); errClose != nil { + t.Fatalf("close multipart writer: %v", errClose) + } + + output, _, errRewrite := codexRewriteOpenAIImageEditMultipartToJSON(body.Bytes(), "gpt-image-1.5", writer.Boundary(), false) + if errRewrite != nil { + t.Fatalf("rewrite multipart payload: %v", errRewrite) + } + if got := gjson.GetBytes(output, "images.0").String(); got != "existing-1" { + t.Fatalf("images.0 = %q", got) + } + if got := gjson.GetBytes(output, "images.1").String(); got != "existing-2" { + t.Fatalf("images.1 = %q", got) + } + if got := gjson.GetBytes(output, "images.2.image_url").String(); !strings.HasPrefix(got, "data:application/octet-stream;base64,") { + t.Fatalf("images.2.image_url = %q", got) + } +} + +func TestCodexImageBuildersPreservePayloads(t *testing.T) { + tool := []byte(`{"type":"image_generation","model":"gpt-image-2"}`) + request := codexBuildImagesResponsesRequest(`draw "this"`, []string{"data:image/png;base64,AA==", "", "data:image/jpeg;base64,BB=="}, tool) + if !gjson.ValidBytes(request) { + t.Fatalf("request is invalid JSON: %s", request) + } + if got := gjson.GetBytes(request, "input.0.content.0.text").String(); got != `draw "this"` { + t.Fatalf("prompt = %q", got) + } + if got := gjson.GetBytes(request, "input.0.content.#").Int(); got != 3 { + t.Fatalf("content count = %d, want 3", got) + } + if got := gjson.GetBytes(request, "tools.0.model").String(); got != "gpt-image-2" { + t.Fatalf("tool model = %q", got) + } + + result := codexImageCallResult{Result: "AA==", OutputFormat: "png", RevisedPrompt: `revised "prompt"`, Quality: "high", Size: "1024x1024"} + response, errBuild := codexBuildImagesAPIResponse([]codexImageCallResult{result}, 123, []byte(`{"images":1}`), result, "b64_json") + if errBuild != nil { + t.Fatalf("codexBuildImagesAPIResponse returned error: %v", errBuild) + } + if !gjson.ValidBytes(response) { + t.Fatalf("response is invalid JSON: %s", response) + } + if got := gjson.GetBytes(response, "data.0.b64_json").String(); got != "AA==" { + t.Fatalf("b64_json = %q", got) + } + if got := gjson.GetBytes(response, "data.0.revised_prompt").String(); got != `revised "prompt"` { + t.Fatalf("revised_prompt = %q", got) + } + if got := gjson.GetBytes(response, "usage.images").Int(); got != 1 { + t.Fatalf("usage.images = %d", got) + } +} + +var benchmarkExecutorPayloadOutput []byte + +func BenchmarkCodexBuildImagesAPIResponseLargePayload(b *testing.B) { + image := strings.Repeat("A", 2<<20) + results := []codexImageCallResult{ + {Result: image, OutputFormat: "png"}, + {Result: image, OutputFormat: "png"}, + {Result: image, OutputFormat: "png"}, + {Result: image, OutputFormat: "png"}, + } + b.ReportAllocs() + b.SetBytes(int64(len(image) * len(results))) + b.ResetTimer() + for b.Loop() { + benchmarkExecutorPayloadOutput, _ = codexBuildImagesAPIResponse(results, 1, []byte(`{"images":4}`), codexImageCallResult{}, "b64_json") + } +} + +func BenchmarkNormalizeKimiToolMessageLinksLargeSinglePatch(b *testing.B) { + content := strings.Repeat("x", 8<<20) + input := []byte(`{"messages":[{"role":"assistant","content":"` + content + `","tool_calls":[{"id":"call_1","type":"function","function":{"name":"lookup","arguments":"{}"}}]},{"role":"tool","tool_call_id":"call_1","content":"ok"}]}`) + b.ReportAllocs() + b.SetBytes(int64(len(input))) + b.ResetTimer() + for b.Loop() { + benchmarkExecutorPayloadOutput, _ = normalizeKimiToolMessageLinks(input) + } +} + +func BenchmarkNormalizeKimiToolMessageLinksLargeMultiplePatches(b *testing.B) { + content := strings.Repeat("x", (8<<20)/32) + var builder strings.Builder + builder.Grow(8 << 20) + builder.WriteString(`{"messages":[`) + for index := 0; index < 32; index++ { + if index > 0 { + builder.WriteByte(',') + } + builder.WriteString(`{"role":"assistant","content":"`) + builder.WriteString(content) + builder.WriteString(`","tool_calls":[{"id":"call_`) + builder.WriteString(strings.Repeat("x", index%3)) + builder.WriteString(`","type":"function","function":{"name":"lookup","arguments":"{}"}}]},{"role":"tool","call_id":"call_`) + builder.WriteString(strings.Repeat("x", index%3)) + builder.WriteString(`","content":"ok"}`) + } + builder.WriteString(`]}`) + input := []byte(builder.String()) + b.ReportAllocs() + b.SetBytes(int64(len(input))) + b.ResetTimer() + for b.Loop() { + benchmarkExecutorPayloadOutput, _ = normalizeKimiToolMessageLinks(input) + } +} + +func BenchmarkNormalizeKimiToolMessageLinksLargeCanonicalPayload(b *testing.B) { + content := strings.Repeat("x", (8<<20)/64) + var builder strings.Builder + builder.Grow(8 << 20) + builder.WriteString(`{"messages":[`) + for index := 0; index < 64; index++ { + if index > 0 { + builder.WriteByte(',') + } + builder.WriteString(`{"role":"user","content":"`) + builder.WriteString(content) + builder.WriteString(`"}`) + } + builder.WriteString(`]}`) + input := []byte(builder.String()) + b.ReportAllocs() + b.SetBytes(int64(len(input))) + b.ResetTimer() + for b.Loop() { + benchmarkExecutorPayloadOutput, _ = normalizeKimiToolMessageLinks(input) + } +} diff --git a/internal/runtime/executor/gemini_executor.go b/internal/runtime/executor/gemini_executor.go index 0607de863..0cd284dda 100644 --- a/internal/runtime/executor/gemini_executor.go +++ b/internal/runtime/executor/gemini_executor.go @@ -157,7 +157,7 @@ func (e *GeminiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r requestedModel := helps.PayloadRequestedModel(opts, req.Model) requestPath := helps.PayloadRequestPath(opts) body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) - body, _ = sjson.SetBytes(body, "model", baseModel) + body = helps.SetStringIfDifferent(body, "model", baseModel) body = capGeminiMaxOutputTokens(body, baseModel) action := "generateContent" @@ -270,7 +270,7 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A requestedModel := helps.PayloadRequestedModel(opts, req.Model) requestPath := helps.PayloadRequestPath(opts) body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) - body, _ = sjson.SetBytes(body, "model", baseModel) + body = helps.SetStringIfDifferent(body, "model", baseModel) body = capGeminiMaxOutputTokens(body, baseModel) baseURL := resolveGeminiBaseURL(auth) @@ -388,7 +388,7 @@ func (e *GeminiExecutor) executeInteractions(ctx context.Context, auth *cliproxy body := translateGeminiInteractionsRequestBody(targetName, req.Payload, opts, false) if gjson.GetBytes(body, "model").Exists() && targetName != "" { - body, _ = sjson.SetBytes(body, "model", targetName) + body = helps.SetStringIfDifferent(body, "model", targetName) } body, err = applyGeminiInteractionsThinking(body, req.Model) if err != nil { @@ -464,7 +464,7 @@ func (e *GeminiExecutor) executeInteractionsStream(ctx context.Context, auth *cl body := translateGeminiInteractionsRequestBody(targetName, req.Payload, opts, true) if gjson.GetBytes(body, "model").Exists() && targetName != "" { - body, _ = sjson.SetBytes(body, "model", targetName) + body = helps.SetStringIfDifferent(body, "model", targetName) } body, err = applyGeminiInteractionsThinking(body, req.Model) if err != nil { @@ -475,7 +475,7 @@ func (e *GeminiExecutor) executeInteractionsStream(ctx context.Context, auth *cl fromProtocol := opts.SourceFormat.String() originalTranslated := geminiInteractionsPayloadConfigSource(targetName, req.Payload, opts, true) body = helps.ApplyPayloadConfigWithRequest(e.cfg, targetName, "interactions", fromProtocol, "", body, originalTranslated, requestedModel, requestPath, opts.Headers) - body, _ = sjson.SetBytes(body, "stream", true) + body = helps.SetBoolIfDifferent(body, "stream", true) baseURL := resolveGeminiBaseURL(auth) url := fmt.Sprintf("%s/%s/interactions", baseURL, glAPIVersion) httpReq, errRequest := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body)) @@ -625,7 +625,7 @@ func (e *GeminiExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Aut translatedReq, _ = sjson.DeleteBytes(translatedReq, "tools") translatedReq, _ = sjson.DeleteBytes(translatedReq, "generationConfig") translatedReq, _ = sjson.DeleteBytes(translatedReq, "safetySettings") - translatedReq, _ = sjson.SetBytes(translatedReq, "model", baseModel) + translatedReq = helps.SetStringIfDifferent(translatedReq, "model", baseModel) baseURL := resolveGeminiBaseURL(auth) url := fmt.Sprintf("%s/%s/models/%s:%s", baseURL, glAPIVersion, baseModel, "countTokens") diff --git a/internal/runtime/executor/gemini_vertex_executor.go b/internal/runtime/executor/gemini_vertex_executor.go index b0677415a..f97eb84d1 100644 --- a/internal/runtime/executor/gemini_vertex_executor.go +++ b/internal/runtime/executor/gemini_vertex_executor.go @@ -340,7 +340,7 @@ func (e *GeminiVertexExecutor) executeWithServiceAccount(ctx context.Context, au requestedModel := helps.PayloadRequestedModel(opts, req.Model) requestPath := helps.PayloadRequestPath(opts) body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) - body, _ = sjson.SetBytes(body, "model", baseModel) + body = helps.SetStringIfDifferent(body, "model", baseModel) body = helps.StripVertexOpenAIResponsesToolCallIDs(body, from.String()) } @@ -465,7 +465,7 @@ func (e *GeminiVertexExecutor) executeWithAPIKey(ctx context.Context, auth *clip requestedModel := helps.PayloadRequestedModel(opts, req.Model) requestPath := helps.PayloadRequestPath(opts) body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) - body, _ = sjson.SetBytes(body, "model", baseModel) + body = helps.SetStringIfDifferent(body, "model", baseModel) body = helps.StripVertexOpenAIResponsesToolCallIDs(body, from.String()) action := getVertexAction(baseModel, false) @@ -580,7 +580,7 @@ func (e *GeminiVertexExecutor) executeStreamWithServiceAccount(ctx context.Conte requestedModel := helps.PayloadRequestedModel(opts, req.Model) requestPath := helps.PayloadRequestPath(opts) body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) - body, _ = sjson.SetBytes(body, "model", baseModel) + body = helps.SetStringIfDifferent(body, "model", baseModel) body = helps.StripVertexOpenAIResponsesToolCallIDs(body, from.String()) action := getVertexAction(baseModel, true) @@ -725,7 +725,7 @@ func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth requestedModel := helps.PayloadRequestedModel(opts, req.Model) requestPath := helps.PayloadRequestPath(opts) body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) - body, _ = sjson.SetBytes(body, "model", baseModel) + body = helps.SetStringIfDifferent(body, "model", baseModel) body = helps.StripVertexOpenAIResponsesToolCallIDs(body, from.String()) action := getVertexAction(baseModel, true) diff --git a/internal/runtime/executor/helps/payload_helpers.go b/internal/runtime/executor/helps/payload_helpers.go index 203589830..bb18f4151 100644 --- a/internal/runtime/executor/helps/payload_helpers.go +++ b/internal/runtime/executor/helps/payload_helpers.go @@ -120,11 +120,7 @@ func ApplyPayloadConfigWithRequest(cfg *config.Config, model, protocol, fromProt continue } for _, resolvedPath := range resolvePayloadRulePaths(out, fullPath) { - updated, errSet := sjson.SetBytes(out, resolvedPath, value) - if errSet != nil { - continue - } - out = updated + out = setPayloadValueIfDifferent(out, resolvedPath, value) } } } @@ -144,11 +140,7 @@ func ApplyPayloadConfigWithRequest(cfg *config.Config, model, protocol, fromProt continue } for _, resolvedPath := range resolvePayloadRulePaths(out, fullPath) { - updated, errSet := sjson.SetRawBytes(out, resolvedPath, rawValue) - if errSet != nil { - continue - } - out = updated + out = SetRawIfDifferent(out, resolvedPath, rawValue) } } } @@ -792,23 +784,64 @@ func removeToolTypeFromToolsArray(payload []byte, toolsPath string, toolType str if !tools.Exists() || !tools.IsArray() { return payload } + toolItems := tools.Array() removed := false - filtered := []byte(`[]`) - for _, tool := range tools.Array() { + for _, tool := range toolItems { if tool.Get("type").String() == toolType { removed = true - continue + break } - updated, errSet := sjson.SetRawBytes(filtered, "-1", []byte(tool.Raw)) - if errSet != nil { - continue - } - filtered = updated } if !removed { return payload } - updated, errSet := sjson.SetRawBytes(payload, toolsPath, filtered) + filtered := make([][]byte, 0, len(toolItems)) + for _, tool := range toolItems { + if tool.Get("type").String() != toolType { + filtered = append(filtered, []byte(tool.Raw)) + } + } + updated, errSet := sjson.SetRawBytes(payload, toolsPath, JoinRawJSONArray(filtered)) + if errSet != nil { + return payload + } + return updated +} + +func setPayloadValueIfDifferent(payload []byte, path string, value any) []byte { + current := gjson.GetBytes(payload, path) + switch typed := value.(type) { + case string: + if current.Type == gjson.String && current.String() == typed { + return payload + } + case bool: + if (typed && current.Type == gjson.True) || (!typed && current.Type == gjson.False) { + return payload + } + case nil: + if current.Raw == "null" { + return payload + } + default: + expectedJSON, errSet := sjson.SetBytes([]byte(`{}`), "value", value) + if errSet != nil { + return payload + } + expected := gjson.GetBytes(expectedJSON, "value") + if expected.Raw == "" { + return payload + } + if current.Raw == expected.Raw { + return payload + } + updated, errSet := sjson.SetRawBytes(payload, path, []byte(expected.Raw)) + if errSet != nil { + return payload + } + return updated + } + updated, errSet := sjson.SetBytes(payload, path, value) if errSet != nil { return payload } diff --git a/internal/runtime/executor/helps/payload_mutations.go b/internal/runtime/executor/helps/payload_mutations.go new file mode 100644 index 000000000..fc48f38cd --- /dev/null +++ b/internal/runtime/executor/helps/payload_mutations.go @@ -0,0 +1,81 @@ +package helps + +import ( + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +// SetStringIfDifferent updates path only when its value is not already the +// canonical JSON string. Values with another JSON type are still normalized. +func SetStringIfDifferent(payload []byte, path, value string) []byte { + current := gjson.GetBytes(payload, path) + if current.Type == gjson.String && current.String() == value { + return payload + } + updated, errSet := sjson.SetBytes(payload, path, value) + if errSet != nil { + return payload + } + return updated +} + +// SetBoolIfDifferent updates path only when its value is not already the +// canonical JSON boolean. Values with another JSON type are still normalized. +func SetBoolIfDifferent(payload []byte, path string, value bool) []byte { + current := gjson.GetBytes(payload, path) + if (value && current.Type == gjson.True) || (!value && current.Type == gjson.False) { + return payload + } + updated, errSet := sjson.SetBytes(payload, path, value) + if errSet != nil { + return payload + } + return updated +} + +// SetRawIfDifferent updates path only when the existing raw JSON is identical. +func SetRawIfDifferent(payload []byte, path string, value []byte) []byte { + current := gjson.GetBytes(payload, path) + if current.Exists() && current.Raw == string(value) { + return payload + } + updated, errSet := sjson.SetRawBytes(payload, path, value) + if errSet != nil { + return payload + } + return updated +} + +// JoinRawJSONArray joins validated raw JSON array items without re-encoding them. +func JoinRawJSONArray(items [][]byte) []byte { + size := len(items) + 1 + for _, item := range items { + size += len(item) + } + out := make([]byte, 0, size) + out = append(out, '[') + for index, item := range items { + if index > 0 { + out = append(out, ',') + } + out = append(out, item...) + } + return append(out, ']') +} + +// JoinRawJSONStrings joins raw JSON array items held as strings. +func JoinRawJSONStrings(items []string) []byte { + size := len(items) + 1 + for _, item := range items { + size += len(item) + } + out := make([]byte, 0, size) + out = append(out, '[') + for index, item := range items { + if index > 0 { + out = append(out, ',') + } + out = append(out, item...) + } + return append(out, ']') +} diff --git a/internal/runtime/executor/helps/payload_mutations_test.go b/internal/runtime/executor/helps/payload_mutations_test.go new file mode 100644 index 000000000..e96aa3900 --- /dev/null +++ b/internal/runtime/executor/helps/payload_mutations_test.go @@ -0,0 +1,155 @@ +package helps + +import ( + "bytes" + "encoding/json" + "strings" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/tidwall/gjson" +) + +type countingPayloadMarshaler struct { + calls *int + value string +} + +func (m countingPayloadMarshaler) MarshalJSON() ([]byte, error) { + *m.calls = *m.calls + 1 + return json.Marshal(m.value) +} + +func TestSetStringIfDifferentReusesCanonicalValue(t *testing.T) { + input := []byte(`{"model":"gpt-test","messages":[]}`) + output := SetStringIfDifferent(input, "model", "gpt-test") + if &output[0] != &input[0] { + t.Fatal("canonical string caused a payload copy") + } +} + +func TestSetStringIfDifferentNormalizesWrongType(t *testing.T) { + input := []byte(`{"model":123}`) + original := bytes.Clone(input) + output := SetStringIfDifferent(input, "model", "123") + model := gjson.GetBytes(output, "model") + if model.Type != gjson.String || model.String() != "123" { + t.Fatalf("model = %s, want string 123", model.Raw) + } + if !bytes.Equal(input, original) { + t.Fatal("input payload was modified in place") + } +} + +func TestSetBoolIfDifferentReusesCanonicalValue(t *testing.T) { + input := []byte(`{"stream":true,"input":[]}`) + output := SetBoolIfDifferent(input, "stream", true) + if &output[0] != &input[0] { + t.Fatal("canonical boolean caused a payload copy") + } +} + +func TestSetBoolIfDifferentNormalizesWrongType(t *testing.T) { + input := []byte(`{"stream":"true"}`) + output := SetBoolIfDifferent(input, "stream", true) + if stream := gjson.GetBytes(output, "stream"); stream.Type != gjson.True { + t.Fatalf("stream = %s, want boolean true", stream.Raw) + } +} + +func TestSetRawIfDifferentReusesIdenticalRawValue(t *testing.T) { + input := []byte(`{"metadata":{"source":"executor"},"input":[]}`) + output := SetRawIfDifferent(input, "metadata", []byte(`{"source":"executor"}`)) + if &output[0] != &input[0] { + t.Fatal("identical raw value caused a payload copy") + } +} + +func TestSetRawIfDifferentUpdatesDifferentRawValue(t *testing.T) { + input := []byte(`{"metadata":"executor"}`) + output := SetRawIfDifferent(input, "metadata", []byte(`{"source":"executor"}`)) + metadata := gjson.GetBytes(output, "metadata") + if !metadata.IsObject() || metadata.Get("source").String() != "executor" { + t.Fatalf("metadata = %s, want object", metadata.Raw) + } +} + +func TestApplyPayloadConfigReusesCanonicalOverrides(t *testing.T) { + cfg := &config.Config{Payload: config.PayloadConfig{ + Override: []config.PayloadRule{{ + Models: []config.PayloadModelRule{{Name: "gpt-test", Protocol: "openai"}}, + Params: map[string]any{"stream": true, "model": "gpt-test"}, + }}, + OverrideRaw: []config.PayloadRule{{ + Models: []config.PayloadModelRule{{Name: "gpt-test", Protocol: "openai"}}, + Params: map[string]any{"metadata": `{"source":"executor"}`}, + }}, + }} + input := []byte(`{"model":"gpt-test","stream":true,"metadata":{"source":"executor"},"messages":[]}`) + output := ApplyPayloadConfigWithRoot(cfg, "gpt-test", "openai", "", input, nil, "", "") + if &output[0] != &input[0] { + t.Fatal("canonical payload overrides caused a payload copy") + } +} + +func TestApplyPayloadConfigNormalizesByteSliceOverride(t *testing.T) { + cfg := &config.Config{Payload: config.PayloadConfig{ + Override: []config.PayloadRule{{ + Models: []config.PayloadModelRule{{Name: "gpt-test", Protocol: "openai"}}, + Params: map[string]any{"value": []byte("abc")}, + }}, + }} + input := []byte(`{"value":"YWJj"}`) + output := ApplyPayloadConfigWithRoot(cfg, "gpt-test", "openai", "", input, nil, "", "") + value := gjson.GetBytes(output, "value") + if value.Type != gjson.String || value.String() != "abc" { + t.Fatalf("value = %s, want string abc", value.Raw) + } +} + +func TestSetPayloadValueIfDifferentUsesSJSONNumberEncoding(t *testing.T) { + input := []byte(`{"value":1.2}`) + output := setPayloadValueIfDifferent(input, "value", float32(1.2)) + if got := gjson.GetBytes(output, "value").Raw; got != "1.2000000476837158" { + t.Fatalf("value = %s, want sjson float32 encoding", got) + } + canonical := []byte(`{"value":1.2000000476837158}`) + reused := setPayloadValueIfDifferent(canonical, "value", float32(1.2)) + if &reused[0] != &canonical[0] { + t.Fatal("canonical float32 encoding caused a payload copy") + } +} + +func TestSetPayloadValueIfDifferentCallsMarshalerOnce(t *testing.T) { + for _, input := range [][]byte{[]byte(`{"value":"old"}`), []byte(`{"value":"new"}`)} { + calls := 0 + value := countingPayloadMarshaler{calls: &calls, value: "new"} + output := setPayloadValueIfDifferent(input, "value", value) + if calls != 1 { + t.Fatalf("MarshalJSON calls = %d, want 1", calls) + } + if got := gjson.GetBytes(output, "value").String(); got != "new" { + t.Fatalf("value = %q, want new", got) + } + } +} + +func TestRemoveToolTypeReusesArrayWithoutMatch(t *testing.T) { + input := []byte(`{"tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}]}`) + output := removeToolTypeFromToolsArray(input, "tools", "image_generation") + if &output[0] != &input[0] { + t.Fatal("tool filtering without a match caused a payload copy") + } +} + +var benchmarkPayloadMutationOutput []byte + +func BenchmarkSetStringIfDifferentLargeCanonicalPayload(b *testing.B) { + input := []byte(`{"model":"gpt-test","messages":[{"role":"user","content":"` + strings.Repeat("x", 8<<20) + `"}]}`) + b.ReportAllocs() + b.SetBytes(int64(len(input))) + b.ResetTimer() + for b.Loop() { + benchmarkPayloadMutationOutput = SetStringIfDifferent(input, "model", "gpt-test") + } +} diff --git a/internal/runtime/executor/helps/vertex_payload_helpers.go b/internal/runtime/executor/helps/vertex_payload_helpers.go index 4c84fae45..2b2c46244 100644 --- a/internal/runtime/executor/helps/vertex_payload_helpers.go +++ b/internal/runtime/executor/helps/vertex_payload_helpers.go @@ -1,7 +1,6 @@ package helps import ( - "fmt" "strings" "github.com/tidwall/gjson" @@ -16,28 +15,71 @@ func StripVertexOpenAIResponsesToolCallIDs(payload []byte, sourceFormat string) } contents := gjson.GetBytes(payload, "contents") - if !contents.IsArray() { + if !contents.IsArray() || !vertexContentsHaveToolCallIDs(contents) { return payload } - out := payload - for contentIndex, content := range contents.Array() { + contentsChanged := false + contentItems := make([][]byte, 0, int(contents.Get("#").Int())) + contents.ForEach(func(_, content gjson.Result) bool { parts := content.Get("parts") if !parts.IsArray() { - continue + contentItems = append(contentItems, []byte(content.Raw)) + return true } - for partIndex, part := range parts.Array() { - if part.Get("functionCall.id").Exists() { - if updated, errDelete := sjson.DeleteBytes(out, fmt.Sprintf("contents.%d.parts.%d.functionCall.id", contentIndex, partIndex)); errDelete == nil { - out = updated + + partsChanged := false + partItems := make([][]byte, 0, int(parts.Get("#").Int())) + parts.ForEach(func(_, part gjson.Result) bool { + partJSON := []byte(part.Raw) + for _, path := range []string{"functionCall.id", "functionResponse.id"} { + if !part.Get(path).Exists() { + continue + } + updated, errDelete := sjson.DeleteBytes(partJSON, path) + if errDelete == nil { + partJSON = updated + partsChanged = true } } - if part.Get("functionResponse.id").Exists() { - if updated, errDelete := sjson.DeleteBytes(out, fmt.Sprintf("contents.%d.parts.%d.functionResponse.id", contentIndex, partIndex)); errDelete == nil { - out = updated - } + partItems = append(partItems, partJSON) + return true + }) + + contentJSON := []byte(content.Raw) + if partsChanged { + updated, errSet := sjson.SetRawBytes(contentJSON, "parts", JoinRawJSONArray(partItems)) + if errSet == nil { + contentJSON = updated + contentsChanged = true } } + contentItems = append(contentItems, contentJSON) + return true + }) + if !contentsChanged { + return payload } - return out + + updated, errSet := sjson.SetRawBytes(payload, "contents", JoinRawJSONArray(contentItems)) + if errSet != nil { + return payload + } + return updated +} + +func vertexContentsHaveToolCallIDs(contents gjson.Result) bool { + hasIDs := false + contents.ForEach(func(_, content gjson.Result) bool { + parts := content.Get("parts") + if !parts.IsArray() { + return true + } + parts.ForEach(func(_, part gjson.Result) bool { + hasIDs = part.Get("functionCall.id").Exists() || part.Get("functionResponse.id").Exists() + return !hasIDs + }) + return !hasIDs + }) + return hasIDs } diff --git a/internal/runtime/executor/helps/vertex_payload_helpers_test.go b/internal/runtime/executor/helps/vertex_payload_helpers_test.go new file mode 100644 index 000000000..f21217d36 --- /dev/null +++ b/internal/runtime/executor/helps/vertex_payload_helpers_test.go @@ -0,0 +1,45 @@ +package helps + +import ( + "strings" + "testing" + + "github.com/tidwall/gjson" +) + +func TestStripVertexToolCallIDsReusesPayloadWithoutIDs(t *testing.T) { + input := []byte(`{"contents":[{"role":"model","parts":[{"functionCall":{"name":"lookup","args":{"id":9007199254740993}}}]}]}`) + output := StripVertexOpenAIResponsesToolCallIDs(input, "openai-response") + if &output[0] != &input[0] { + t.Fatal("payload without tool call IDs was copied") + } +} + +func TestStripVertexToolCallIDsRebuildsContentsOnce(t *testing.T) { + input := []byte(`{"contents":[{"role":"model","parts":[{"functionCall":{"id":"call_1","name":"lookup","args":{"id":9007199254740993}}}]},{"role":"user","parts":[{"functionResponse":{"id":"call_1","name":"lookup","response":{"id":"keep"}}}]}]}`) + output := StripVertexOpenAIResponsesToolCallIDs(input, "openai-response") + if gjson.GetBytes(output, "contents.0.parts.0.functionCall.id").Exists() { + t.Fatal("functionCall.id was not removed") + } + if gjson.GetBytes(output, "contents.1.parts.0.functionResponse.id").Exists() { + t.Fatal("functionResponse.id was not removed") + } + if got := gjson.GetBytes(output, "contents.1.parts.0.functionResponse.response.id").String(); got != "keep" { + t.Fatalf("nested response id = %q, want keep", got) + } + if got := gjson.GetBytes(output, "contents.0.parts.0.functionCall.args.id").Raw; got != "9007199254740993" { + t.Fatalf("large integer = %s, want exact original value", got) + } +} + +var benchmarkVertexPayloadOutput []byte + +func BenchmarkStripVertexToolCallIDsLargeNoopPayload(b *testing.B) { + input := []byte(`{"contents":[{"role":"user","parts":[{"text":"` + strings.Repeat("x", 8<<20) + `"}]}]}`) + b.ReportAllocs() + b.SetBytes(int64(len(input))) + b.ResetTimer() + for b.Loop() { + benchmarkVertexPayloadOutput = StripVertexOpenAIResponsesToolCallIDs(input, "openai-response") + } +} diff --git a/internal/runtime/executor/kimi_executor.go b/internal/runtime/executor/kimi_executor.go index 53d19f81b..1799e6652 100644 --- a/internal/runtime/executor/kimi_executor.go +++ b/internal/runtime/executor/kimi_executor.go @@ -349,18 +349,18 @@ func normalizeKimiToolMessageLinks(body []byte) ([]byte, error) { return body, nil } - msgs := messages.Array() - out, dropped, err := filterKimiEmptyAssistantMessages(body, msgs) - if err != nil { - return body, err - } - if dropped > 0 { - log.WithField("dropped_assistant_messages", dropped).Debug("kimi executor: dropped empty assistant messages") + type messagePatch struct { + index int + path string + value string + errorContext string } - messages = gjson.GetBytes(out, "messages") - msgs = messages.Array() + msgs := messages.Array() + droppedMessages := make([]bool, len(msgs)) + patches := make([]messagePatch, 0) pending := make([]string, 0) + dropped := 0 patched := 0 patchedReasoning := 0 ambiguous := 0 @@ -377,8 +377,13 @@ func normalizeKimiToolMessageLinks(body []byte) ([]byte, error) { } } - for msgIdx := range msgs { - msg := msgs[msgIdx] + for msgIndex, msg := range msgs { + if shouldDropKimiAssistantMessage(msg) { + droppedMessages[msgIndex] = true + dropped++ + continue + } + role := strings.TrimSpace(msg.Get("role").String()) switch role { case "assistant": @@ -392,51 +397,39 @@ func normalizeKimiToolMessageLinks(body []byte) ([]byte, error) { } toolCalls := msg.Get("tool_calls") - if !toolCalls.Exists() || !toolCalls.IsArray() || len(toolCalls.Array()) == 0 { - continue - } - - if !reasoning.Exists() || strings.TrimSpace(reasoning.String()) == "" { - reasoningText := fallbackAssistantReasoning(msg, hasLatestReasoning, latestReasoning) - path := fmt.Sprintf("messages.%d.reasoning_content", msgIdx) - next, err := sjson.SetBytes(out, path, reasoningText) - if err != nil { - return body, fmt.Errorf("kimi executor: failed to set assistant reasoning_content: %w", err) + if toolCalls.Exists() && toolCalls.IsArray() { + toolCallItems := toolCalls.Array() + if len(toolCallItems) > 0 { + if !reasoning.Exists() || strings.TrimSpace(reasoning.String()) == "" { + patches = append(patches, messagePatch{ + index: msgIndex, + path: "reasoning_content", + value: fallbackAssistantReasoning(msg, hasLatestReasoning, latestReasoning), + errorContext: "failed to set assistant reasoning_content", + }) + patchedReasoning++ + } + for _, toolCall := range toolCallItems { + id := strings.TrimSpace(toolCall.Get("id").String()) + if id != "" { + pending = append(pending, id) + } + } } - out = next - patchedReasoning++ - } - - for _, tc := range toolCalls.Array() { - id := strings.TrimSpace(tc.Get("id").String()) - if id == "" { - continue - } - pending = append(pending, id) } case "tool": toolCallID := strings.TrimSpace(msg.Get("tool_call_id").String()) if toolCallID == "" { toolCallID = strings.TrimSpace(msg.Get("call_id").String()) if toolCallID != "" { - path := fmt.Sprintf("messages.%d.tool_call_id", msgIdx) - next, err := sjson.SetBytes(out, path, toolCallID) - if err != nil { - return body, fmt.Errorf("kimi executor: failed to set tool_call_id from call_id: %w", err) - } - out = next + patches = append(patches, messagePatch{index: msgIndex, path: "tool_call_id", value: toolCallID, errorContext: "failed to set tool_call_id from call_id"}) patched++ } } if toolCallID == "" { if len(pending) == 1 { toolCallID = pending[0] - path := fmt.Sprintf("messages.%d.tool_call_id", msgIdx) - next, err := sjson.SetBytes(out, path, toolCallID) - if err != nil { - return body, fmt.Errorf("kimi executor: failed to infer tool_call_id: %w", err) - } - out = next + patches = append(patches, messagePatch{index: msgIndex, path: "tool_call_id", value: toolCallID, errorContext: "failed to infer tool_call_id"}) patched++ } else if len(pending) > 1 { ambiguous++ @@ -448,6 +441,57 @@ func normalizeKimiToolMessageLinks(body []byte) ([]byte, error) { } } + if dropped > 0 { + log.WithField("dropped_assistant_messages", dropped).Debug("kimi executor: dropped empty assistant messages") + } + if dropped == 0 && len(patches) == 0 { + if ambiguous > 0 { + log.WithFields(log.Fields{ + "ambiguous_tool_messages": ambiguous, + "pending_tool_calls": len(pending), + }).Warn("kimi executor: tool messages missing tool_call_id with ambiguous candidates") + } + return body, nil + } + + var out []byte + if dropped == 0 && len(patches) == 1 { + patch := patches[0] + path := fmt.Sprintf("messages.%d.%s", patch.index, patch.path) + updated, errSet := sjson.SetBytes(body, path, patch.value) + if errSet != nil { + return body, fmt.Errorf("kimi executor: %s: %w", patch.errorContext, errSet) + } + out = updated + } else { + messageItems := make([]string, 0, len(msgs)-dropped) + patchIndex := 0 + for msgIndex, msg := range msgs { + if droppedMessages[msgIndex] { + continue + } + messageJSON := msg.Raw + for patchIndex < len(patches) && patches[patchIndex].index == msgIndex { + patch := patches[patchIndex] + next, errSet := sjson.SetBytes([]byte(messageJSON), patch.path, patch.value) + if errSet != nil { + return body, fmt.Errorf("kimi executor: %s: %w", patch.errorContext, errSet) + } + messageJSON = string(next) + patchIndex++ + } + messageItems = append(messageItems, messageJSON) + } + updated, errSet := sjson.SetRawBytes(body, "messages", helps.JoinRawJSONStrings(messageItems)) + if errSet != nil { + if dropped > 0 { + return body, fmt.Errorf("kimi executor: failed to drop empty assistant messages: %w", errSet) + } + return body, fmt.Errorf("kimi executor: %s: %w", patches[0].errorContext, errSet) + } + out = updated + } + if patched > 0 || patchedReasoning > 0 { log.WithFields(log.Fields{ "patched_tool_messages": patched, @@ -460,32 +504,9 @@ func normalizeKimiToolMessageLinks(body []byte) ([]byte, error) { "pending_tool_calls": len(pending), }).Warn("kimi executor: tool messages missing tool_call_id with ambiguous candidates") } - return out, nil } -func filterKimiEmptyAssistantMessages(body []byte, msgs []gjson.Result) ([]byte, int, error) { - kept := make([]string, 0, len(msgs)) - dropped := 0 - for _, msg := range msgs { - if shouldDropKimiAssistantMessage(msg) { - dropped++ - continue - } - kept = append(kept, msg.Raw) - } - if dropped == 0 { - return body, 0, nil - } - - rawMessages := []byte("[" + strings.Join(kept, ",") + "]") - out, err := sjson.SetRawBytes(body, "messages", rawMessages) - if err != nil { - return body, 0, fmt.Errorf("kimi executor: failed to drop empty assistant messages: %w", err) - } - return out, dropped, nil -} - func shouldDropKimiAssistantMessage(msg gjson.Result) bool { if strings.TrimSpace(msg.Get("role").String()) != "assistant" { return false diff --git a/internal/runtime/executor/openai_compat_executor.go b/internal/runtime/executor/openai_compat_executor.go index 758816143..d18ab9679 100644 --- a/internal/runtime/executor/openai_compat_executor.go +++ b/internal/runtime/executor/openai_compat_executor.go @@ -326,7 +326,7 @@ func (e *OpenAICompatExecutor) ExecuteStream(ctx context.Context, auth *cliproxy // Request usage data in the final streaming chunk so that token statistics // are captured even when the upstream is an OpenAI-compatible provider. - translated, _ = sjson.SetBytes(translated, "stream_options.include_usage", true) + translated = helps.SetBoolIfDifferent(translated, "stream_options.include_usage", true) reporter.SetTranslatedReasoningEffort(translated, to.String()) url := strings.TrimSuffix(baseURL, "/") + "/chat/completions" @@ -634,10 +634,10 @@ func prepareOpenAICompatImagesPayload(payload []byte, model string, contentType contentType = strings.TrimSpace(contentType) if json.Valid(payload) { if model != "" { - payload, _ = sjson.SetBytes(payload, "model", model) + payload = helps.SetStringIfDifferent(payload, "model", model) } if stream { - payload, _ = sjson.SetBytes(payload, "stream", true) + payload = helps.SetBoolIfDifferent(payload, "stream", true) } else { payload, _ = sjson.DeleteBytes(payload, "stream") } @@ -778,8 +778,7 @@ func (e *OpenAICompatExecutor) overrideModel(payload []byte, model string) []byt if len(payload) == 0 || model == "" { return payload } - payload, _ = sjson.SetBytes(payload, "model", model) - return payload + return helps.SetStringIfDifferent(payload, "model", model) } type statusErr struct { diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go index 5c2561a04..fa5713c33 100644 --- a/internal/runtime/executor/xai_executor.go +++ b/internal/runtime/executor/xai_executor.go @@ -904,8 +904,8 @@ func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliprox requestedModel := helps.PayloadRequestedModel(opts, req.Model) requestPath := helps.PayloadRequestPath(opts) body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) - body, _ = sjson.SetBytes(body, "model", baseModel) - body, _ = sjson.SetBytes(body, "stream", stream) + body = helps.SetStringIfDifferent(body, "model", baseModel) + body = helps.SetBoolIfDifferent(body, "stream", stream) body, _ = sjson.DeleteBytes(body, "previous_response_id") body, _ = sjson.DeleteBytes(body, "prompt_cache_retention") body, _ = sjson.DeleteBytes(body, "safety_identifier") @@ -941,7 +941,7 @@ func (e *XAIExecutor) prepareResponsesRequestTo(ctx context.Context, req cliprox return nil, errSession } if sessionID != "" { - body, _ = sjson.SetBytes(body, "prompt_cache_key", sessionID) + body = helps.SetStringIfDifferent(body, "prompt_cache_key", sessionID) } return &xaiPreparedRequest{ @@ -1415,27 +1415,24 @@ func pruneXAIAllowedToolsChoice(body []byte, available map[xaiToolChoiceKey]stru body, _ = sjson.DeleteBytes(body, "tool_choice") return body } - filtered := []byte(`[]`) + allowedItems := allowed.Array() + filtered := make([][]byte, 0, len(allowedItems)) changed := false - for _, tool := range allowed.Array() { + for _, tool := range allowedItems { if !xaiToolChoiceMatchesAvailable(tool, available) { changed = true continue } - updated, errSet := sjson.SetRawBytes(filtered, "-1", []byte(tool.Raw)) - if errSet != nil { - return body - } - filtered = updated + filtered = append(filtered, []byte(tool.Raw)) } if !changed { return body } - if len(gjson.ParseBytes(filtered).Array()) == 0 { + if len(filtered) == 0 { body, _ = sjson.DeleteBytes(body, "tool_choice") return body } - body, _ = sjson.SetRawBytes(body, "tool_choice.tools", filtered) + body, _ = sjson.SetRawBytes(body, "tool_choice.tools", helps.JoinRawJSONArray(filtered)) return body } @@ -1597,9 +1594,10 @@ func promoteXAIAdditionalTools(body []byte) []byte { } func normalizeXAIToolArray(tools gjson.Result) ([]byte, bool, bool) { + toolItems := tools.Array() + filtered := make([][]byte, 0, len(toolItems)) changed := false - filtered := []byte(`[]`) - for _, tool := range tools.Array() { + for _, tool := range toolItems { toolType := tool.Get("type").String() if toolType == xaiNamespaceToolType { changed = true @@ -1611,14 +1609,9 @@ func normalizeXAIToolArray(tools gjson.Result) ([]byte, bool, bool) { return nil, false, false } changed = changed || nestedChanged - if len(nestedRaw) == 0 { - continue + if len(nestedRaw) > 0 { + filtered = append(filtered, nestedRaw) } - updated, errSet := sjson.SetRawBytes(filtered, "-1", nestedRaw) - if errSet != nil { - return nil, false, false - } - filtered = updated } } continue @@ -1628,16 +1621,14 @@ func normalizeXAIToolArray(tools gjson.Result) ([]byte, bool, bool) { return nil, false, false } changed = changed || toolChanged - if len(raw) == 0 { - continue + if len(raw) > 0 { + filtered = append(filtered, raw) } - updated, errSet := sjson.SetRawBytes(filtered, "-1", raw) - if errSet != nil { - return nil, false, false - } - filtered = updated } - return filtered, changed, true + if !changed { + return nil, false, true + } + return helps.JoinRawJSONArray(filtered), true, true } // normalizeXAIToolChoiceForTools drops tool_choice and parallel_tool_calls