diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go index 3443bc9a6..f3ba1335c 100644 --- a/internal/runtime/executor/xai_executor.go +++ b/internal/runtime/executor/xai_executor.go @@ -1187,81 +1187,78 @@ func normalizeXAIImageRefs(body []byte) []byte { if !gjson.ValidBytes(body) { return body } - return normalizeXAIImageRefsAt(body, "") -} -func normalizeXAIImageRefsAt(body []byte, path string) []byte { - node := gjson.GetBytes(body, path) - if path == "" { - node = gjson.ParseBytes(body) - } - if !node.Exists() { + decoder := json.NewDecoder(bytes.NewReader(body)) + decoder.UseNumber() + var payload any + if errDecode := decoder.Decode(&payload); errDecode != nil { return body } - if node.IsObject() { - node.ForEach(func(key, value gjson.Result) bool { - childPath := key.String() - if path != "" { - childPath = path + "." + key.String() - } - switch key.String() { + if !normalizeXAIImageRefsValue(payload) { + return body + } + normalized, errMarshal := json.Marshal(payload) + if errMarshal != nil { + return body + } + return normalized +} + +func normalizeXAIImageRefsValue(value any) bool { + changed := false + switch node := value.(type) { + case map[string]any: + for key, child := range node { + switch key { case "image": - body = normalizeXAIImageRefAtPath(body, childPath) + changed = normalizeXAIImageRef(child) || changed case "images", "reference_images": - if value.IsArray() { - for i := range value.Array() { - body = normalizeXAIImageRefAtPath(body, fmt.Sprintf("%s.%d", childPath, i)) + if refs, ok := child.([]any); ok { + for _, ref := range refs { + changed = normalizeXAIImageRef(ref) || changed } } - default: - if value.IsObject() || value.IsArray() { - body = normalizeXAIImageRefsAt(body, childPath) - } } - return true - }) - return body - } - - if node.IsArray() { - for i := range node.Array() { - childPath := fmt.Sprintf("%d", i) - if path != "" { - childPath = fmt.Sprintf("%s.%d", path, i) - } - body = normalizeXAIImageRefsAt(body, childPath) + changed = normalizeXAIImageRefsValue(child) || changed + } + case []any: + for _, child := range node { + changed = normalizeXAIImageRefsValue(child) || changed } } - return body + return changed } -func normalizeXAIImageRefAtPath(body []byte, path string) []byte { - ref := gjson.GetBytes(body, path) - if !ref.Exists() || ref.Type != gjson.JSON || !ref.IsObject() { - return body +func normalizeXAIImageRef(value any) bool { + ref, ok := value.(map[string]any) + if !ok { + return false } - url := strings.TrimSpace(ref.Get("url").String()) + originalURL, _ := ref["url"].(string) + url := strings.TrimSpace(originalURL) + imageURL, hasImageURL := ref["image_url"] if url == "" { - if imageURL := ref.Get("image_url"); imageURL.Exists() { - if imageURL.Type == gjson.String { - url = strings.TrimSpace(imageURL.String()) - } else if imageURL.IsObject() { - url = strings.TrimSpace(imageURL.Get("url").String()) - } + switch imageURL := imageURL.(type) { + case string: + url = strings.TrimSpace(imageURL) + case map[string]any: + url, _ = imageURL["url"].(string) + url = strings.TrimSpace(url) } } if url == "" { - return body + return false + } + if url == originalURL && !hasImageURL { + return false } // Always emit the xAI field name and drop the OpenAI alias. - body, _ = sjson.SetBytes(body, path+".url", url) - if ref.Get("image_url").Exists() { - body, _ = sjson.DeleteBytes(body, path+".image_url") - } - return body + ref["url"] = url + delete(ref, "image_url") + return true } func xaiIsVideoRequest(opts cliproxyexecutor.Options) bool { diff --git a/internal/runtime/executor/xai_executor_test.go b/internal/runtime/executor/xai_executor_test.go index 892276fe4..9cd847025 100644 --- a/internal/runtime/executor/xai_executor_test.go +++ b/internal/runtime/executor/xai_executor_test.go @@ -5,6 +5,7 @@ import ( "context" "crypto/sha256" "encoding/base64" + "encoding/json" "errors" "fmt" "io" @@ -2302,6 +2303,42 @@ func TestNormalizeXAIImageRefsRewritesImageURLField(t *testing.T) { } } +func TestNormalizeXAIImageRefsSupportsSpecialJSONKeys(t *testing.T) { + t.Parallel() + + in := []byte(`{ + "metadata.with.dot":{"image":{"image_url":"https://example.com/dot.png"}}, + "back\\slash":{"image":{"image_url":"https://example.com/backslash.png"}}, + "":{"image":{"image_url":"https://example.com/empty-key.png"}} + }`) + out := normalizeXAIImageRefs(in) + + var payload map[string]any + if errUnmarshal := json.Unmarshal(out, &payload); errUnmarshal != nil { + t.Fatalf("unmarshal normalized payload: %v", errUnmarshal) + } + for key, wantURL := range map[string]string{ + "metadata.with.dot": "https://example.com/dot.png", + "back\\slash": "https://example.com/backslash.png", + "": "https://example.com/empty-key.png", + } { + nested, ok := payload[key].(map[string]any) + if !ok { + t.Fatalf("payload[%q] = %#v, want object", key, payload[key]) + } + image, ok := nested["image"].(map[string]any) + if !ok { + t.Fatalf("payload[%q].image = %#v, want object", key, nested["image"]) + } + if gotURL, _ := image["url"].(string); gotURL != wantURL { + t.Fatalf("payload[%q].image.url = %q, want %q", key, gotURL, wantURL) + } + if _, exists := image["image_url"]; exists { + t.Fatalf("payload[%q].image_url should be removed", key) + } + } +} + func TestXAIExecutorExecuteImagesRewritesImageURLToURL(t *testing.T) { var gotBody []byte server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {