openteam/server/internal/proxy/jsonfast_test.go

package proxy

import (
	"encoding/json"
	"strings"
	"testing"
)

func TestRewriteModelFieldTopLevelOnly(t *testing.T) {
	in := `{"model":"a","messages":[{"role":"user","content":"{\"model\":\"nested\"}"}],"max_tokens":5}`
	out, ok := rewriteModelField([]byte(in), "b")
	if !ok {
		t.Fatal("应改写顶层 model")
	}
	var m map[string]any
	if err := json.Unmarshal(out, &m); err != nil {
		t.Fatalf("输出必须是合法 JSON: %v\n%s", err, out)
	}
	if m["model"] != "b" {
		t.Fatalf("顶层 model 未改写：%v", m["model"])
	}
	// 嵌套的 model 引用必须原样保留
	if !strings.Contains(string(out), `\"model\":\"nested\"`) {
		t.Fatalf("嵌套内容被改动了：%s", out)
	}
	if m["max_tokens"].(float64) != 5 {
		t.Fatalf("其它字段丢失：%v", out)
	}
}

func TestRewriteModelFieldKeepsByteShape(t *testing.T) {
	in := "{\n  \"model\" : \"old\",\n  \"stream\": true\n}"
	out, ok := rewriteModelField([]byte(in), "new")
	if !ok {
		t.Fatal("应改写")
	}
	want := "{\n  \"model\" : \"new\",\n  \"stream\": true\n}"
	if string(out) != want {
		t.Fatalf("应只替换值本身，实际：%q", out)
	}
}

func TestRewriteModelFieldMissing(t *testing.T) {
	if _, ok := rewriteModelField([]byte(`{"stream":true}`), "x"); ok {
		t.Fatal("没有 model 时不应改写")
	}
	if _, ok := rewriteModelField([]byte(`not json`), "x"); ok {
		t.Fatal("非 JSON 不应改写")
	}
}

func TestInjectStreamOptions(t *testing.T) {
	cases := []struct{ in, want string }{
		{`{"model":"m"}`, `{"model":"m","stream_options":{"include_usage":true}}`},
		{"{\n  \"model\": \"m\"\n}", "{\n  \"model\": \"m\"\n,\"stream_options\":{\"include_usage\":true}}"},
		{`{}`, `{"stream_options":{"include_usage":true}}`}, // 空对象也能安全追加
	}
	for _, c := range cases {
		out, ok := injectStreamOptions([]byte(c.in), `{"include_usage":true}`)
		if !ok {
			t.Fatalf("应注入：%s", c.in)
		}
		if string(out) != c.want {
			t.Fatalf("注入结果 = %q，期望 %q", out, c.want)
		}
		var m map[string]any
		if err := json.Unmarshal(out, &m); err != nil {
			t.Fatalf("注入后不是合法 JSON: %v (%s)", err, out)
		}
		if _, ok := m["stream_options"]; !ok {
			t.Fatalf("缺少 stream_options: %s", out)
		}
	}
	// 已有 stream_options 时不重复注入
	if _, ok := injectStreamOptions([]byte(`{"stream_options":{"include_usage":true}}`), `{}`); ok {
		t.Fatal("已有该键时不应再注入")
	}
}

// 快路径必须与 JSON 解析结果一致（含转义、非正文事件、三协议形态）。
func TestFastContentTextMatchesJSON(t *testing.T) {
	lines := []string{
		`data: {"choices":[{"delta":{"content":"hello"}}]}`,
		`data: {"choices":[{"delta":{"content":"中文内容"}}]}`,
		`data: {"choices":[{"delta":{"content":"带\"引号\"和\\反斜杠"}}]}`,
		`data: {"choices":[{"delta":{"role":"assistant"}}]}`,
		`data: {"choices":[{"delta":{},"finish_reason":"stop"}]}`,
		`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"块文本"}}`,
		`data: {"type":"response.output_text.delta","delta":"响应增量"}`,
		`data: {"type":"message_start","message":{"usage":{}}}`,
		`data: [DONE]`,
	}
	for _, line := range lines {
		got := fastContentText([]byte(line))
		want := sseContentText([]byte(line))
		if got != want {
			t.Fatalf("不一致\n行: %s\n快路径: %q\nJSON: %q", line, got, want)
		}
	}
}