| package upstream
|
|
|
| import (
|
| "encoding/json"
|
| "errors"
|
| "net/http"
|
| "net/http/httptest"
|
| "strings"
|
| "testing"
|
| "time"
|
|
|
| "github.com/linguo2625469/workbuddy2api-panel/internal/auth"
|
| )
|
|
|
| func TestPrepareBodyForcesStream(t *testing.T) {
|
| out := PrepareBodyOpt([]byte(`{"model":"glm-5.2","messages":[]}`), true)
|
| var m map[string]any
|
| json.Unmarshal(out, &m)
|
| if m["stream"] != true {
|
| t.Errorf("stream=%v", m["stream"])
|
| }
|
| }
|
|
|
| func TestPrepareBodyToolChoiceFunctionObject(t *testing.T) {
|
| out := PrepareBodyOpt([]byte(`{"tool_choice":{"type":"function","function":{"name":"get_weather"}},"tools":[{"type":"function"}]}`), true)
|
| var m map[string]any
|
| json.Unmarshal(out, &m)
|
| if m["tool_choice"] != "get_weather" {
|
| t.Errorf("tool_choice=%v", m["tool_choice"])
|
| }
|
| if _, ok := m["tools"]; !ok {
|
| t.Error("tools should be kept for function choice")
|
| }
|
| }
|
|
|
| func TestPrepareBodyToolChoiceNone(t *testing.T) {
|
| for _, in := range []string{
|
| `{"tool_choice":"none","tools":[{}],"functions":[{}]}`,
|
| `{"tool_choice":{"type":"none"},"tools":[{}]}`,
|
| } {
|
| out := PrepareBodyOpt([]byte(in), true)
|
| var m map[string]any
|
| json.Unmarshal(out, &m)
|
| if _, ok := m["tool_choice"]; ok {
|
| t.Errorf("%s: tool_choice should be deleted", in)
|
| }
|
| if _, ok := m["tools"]; ok {
|
| t.Errorf("%s: tools should be deleted", in)
|
| }
|
| if _, ok := m["functions"]; ok {
|
| t.Errorf("%s: functions should be deleted", in)
|
| }
|
| }
|
| }
|
|
|
| func TestPrepareBodyToolChoiceAuto(t *testing.T) {
|
| out := PrepareBodyOpt([]byte(`{"tool_choice":{"type":"auto"}}`), true)
|
| var m map[string]any
|
| json.Unmarshal(out, &m)
|
| if m["tool_choice"] != "auto" {
|
| t.Errorf("tool_choice=%v", m["tool_choice"])
|
| }
|
| }
|
|
|
| func TestPrepareBodyInvalidJSON(t *testing.T) {
|
| in := []byte(`{broken`)
|
| out := PrepareBodyOpt(in, true)
|
| if string(out) != string(in) {
|
| t.Error("invalid json should pass through unchanged")
|
| }
|
| }
|
|
|
| const sseFixture = "data: {\"id\":\"chatcmpl-1\",\"object\":\"chat.completion.chunk\",\"created\":1753600000,\"model\":\"glm-5.2\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"你好\"}}]}\n\n" +
|
| "data: {\"id\":\"chatcmpl-1\",\"object\":\"chat.completion.chunk\",\"created\":1753600000,\"model\":\"glm-5.2\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"，世界\"}}]}\n\n" +
|
| "data: {\"id\":\"chatcmpl-1\",\"object\":\"chat.completion.chunk\",\"created\":1753600000,\"model\":\"glm-5.2\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":5,\"completion_tokens\":2,\"total_tokens\":7}}\n\n" +
|
| "data: [DONE]\n\n"
|
|
|
| func TestAggregate(t *testing.T) {
|
| resp, err := Aggregate(strings.NewReader(sseFixture))
|
| if err != nil {
|
| t.Fatal(err)
|
| }
|
| if resp["object"] != "chat.completion" {
|
| t.Errorf("object=%v", resp["object"])
|
| }
|
| if resp["model"] != "glm-5.2" {
|
| t.Errorf("model=%v", resp["model"])
|
| }
|
| choices := resp["choices"].([]any)
|
| msg := choices[0].(map[string]any)["message"].(map[string]any)
|
| if msg["content"] != "你好，世界" {
|
| t.Errorf("content=%q", msg["content"])
|
| }
|
| if msg["role"] != "assistant" {
|
| t.Errorf("role=%v", msg["role"])
|
| }
|
| if choices[0].(map[string]any)["finish_reason"] != "stop" {
|
| t.Errorf("finish_reason=%v", choices[0].(map[string]any)["finish_reason"])
|
| }
|
| usage := resp["usage"].(map[string]any)
|
| if usage["total_tokens"].(float64) != 7 {
|
| t.Errorf("usage=%v", usage)
|
| }
|
| }
|
|
|
| func TestAggregateSkipsNonDataLines(t *testing.T) {
|
| raw := ": comment\n\n" + sseFixture
|
| resp, err := Aggregate(strings.NewReader(raw))
|
| if err != nil {
|
| t.Fatal(err)
|
| }
|
| msg := resp["choices"].([]any)[0].(map[string]any)["message"].(map[string]any)
|
| if msg["content"] != "你好，世界" {
|
| t.Errorf("content=%q", msg["content"])
|
| }
|
| }
|
|
|
| func TestAggregateToolCalls(t *testing.T) {
|
|
|
| raw := `data: {"id":"x1","model":"deepseek-v4-pro","created":1,"choices":[{"index":0,"delta":{"role":"assistant","content":"","tool_calls":[{"id":"call_a","type":"function","function":{"name":"get_weather","arguments":""},"index":0}]}}],"usage":null}
|
|
|
| data: {"id":"x1","choices":[{"index":0,"delta":{"tool_calls":[{"function":{"arguments":"{\"city\":"},"index":0}]}}]}
|
|
|
| data: {"id":"x1","choices":[{"index":0,"delta":{"tool_calls":[{"function":{"arguments":"\"北京\"}"},"index":0}]}}]}
|
|
|
| data: {"id":"x1","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}],"usage":{"total_tokens":11}}
|
|
|
| data: [DONE]
|
|
|
| `
|
| resp, err := Aggregate(strings.NewReader(raw))
|
| if err != nil {
|
| t.Fatal(err)
|
| }
|
| choice := resp["choices"].([]any)[0].(map[string]any)
|
| if choice["finish_reason"] != "tool_calls" {
|
| t.Errorf("finish_reason=%v", choice["finish_reason"])
|
| }
|
| msg := choice["message"].(map[string]any)
|
| calls, ok := msg["tool_calls"].([]map[string]any)
|
| if !ok || len(calls) != 1 {
|
| t.Fatalf("tool_calls=%#v", msg["tool_calls"])
|
| }
|
| if calls[0]["id"] != "call_a" || calls[0]["type"] != "function" {
|
| t.Errorf("call meta=%v", calls[0])
|
| }
|
| fn := calls[0]["function"].(map[string]any)
|
| if fn["name"] != "get_weather" {
|
| t.Errorf("fn.name=%v", fn["name"])
|
| }
|
| if fn["arguments"] != `{"city":"北京"}` {
|
| t.Errorf("fn.arguments=%q", fn["arguments"])
|
| }
|
| }
|
|
|
|
|
|
|
| func TestStripToolCallNames(t *testing.T) {
|
| mkFrame := func(idx float64, name, args string) map[string]any {
|
| fn := map[string]any{}
|
| if name != "" {
|
| fn["name"] = name
|
| }
|
| if args != "" {
|
| fn["arguments"] = args
|
| }
|
| return map[string]any{"choices": []any{
|
| map[string]any{"delta": map[string]any{"tool_calls": []any{
|
| map[string]any{"index": idx, "function": fn},
|
| }}},
|
| }}
|
| }
|
| getFn := func(f map[string]any) map[string]any {
|
| return f["choices"].([]any)[0].(map[string]any)["delta"].(map[string]any)["tool_calls"].([]any)[0].(map[string]any)["function"].(map[string]any)
|
| }
|
| seen := map[int]bool{}
|
|
|
|
|
| f0 := mkFrame(0, "lookup", "")
|
| stripToolCallNames(f0, seen)
|
| if !seen[0] {
|
| t.Fatal("index 0 should be marked seen after first chunk")
|
| }
|
| if getFn(f0)["name"] != "lookup" {
|
| t.Errorf("first chunk name=%v want lookup", getFn(f0)["name"])
|
| }
|
|
|
|
|
| f1 := mkFrame(0, "", `{"term":"x"}`)
|
| stripToolCallNames(f1, seen)
|
| if _, ok := getFn(f1)["name"]; ok {
|
| t.Errorf("subsequent empty name should be stripped: %#v", getFn(f1))
|
| }
|
| if getFn(f1)["arguments"] != `{"term":"x"}` {
|
| t.Errorf("arguments altered: %#v", getFn(f1)["arguments"])
|
| }
|
|
|
|
|
| f2 := mkFrame(0, "lookup", "y")
|
| stripToolCallNames(f2, seen)
|
| if _, ok := getFn(f2)["name"]; ok {
|
| t.Errorf("subsequent duplicate non-empty name should be stripped: %#v", getFn(f2))
|
| }
|
| if getFn(f2)["arguments"] != "y" {
|
| t.Errorf("arguments altered: %#v", getFn(f2)["arguments"])
|
| }
|
|
|
|
|
| f3 := mkFrame(1, "other", "")
|
| stripToolCallNames(f3, seen)
|
| if getFn(f3)["name"] != "other" {
|
| t.Errorf("index 1 first name=%v want other", getFn(f3)["name"])
|
| }
|
| if !seen[1] {
|
| t.Error("index 1 should be marked seen")
|
| }
|
|
|
|
|
| f4 := map[string]any{"choices": []any{
|
| map[string]any{"delta": map[string]any{"content": "hi"}},
|
| }}
|
| stripToolCallNames(f4, seen)
|
| if got := f4["choices"].([]any)[0].(map[string]any)["delta"].(map[string]any); len(got) != 1 || got["content"] != "hi" {
|
| t.Errorf("content-only frame altered: %#v", got)
|
| }
|
| }
|
|
|
|
|
|
|
|
|
| func TestStreamToolCallNameOnce(t *testing.T) {
|
| const nFrames = 11
|
| var sb strings.Builder
|
| for i := 0; i < nFrames; i++ {
|
| sb.WriteString(`data: {"id":"x1","object":"chat.completion.chunk","created":1,"model":"m","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_a","type":"function","function":{"name":"Bash","arguments":"arg` + string(rune('0'+i)) + `"}}]}}]}`)
|
| sb.WriteString("\n\n")
|
| }
|
| sb.WriteString("data: [DONE]\n\n")
|
|
|
| frames, done := streamFrames(t, sb.String())
|
| if done != 1 {
|
| t.Fatalf("done=%d want 1", done)
|
| }
|
| if len(frames) != nFrames {
|
| t.Fatalf("frames=%d want %d", len(frames), nFrames)
|
| }
|
| gotName := 0
|
| for i, fr := range frames {
|
| chs, _ := fr["choices"].([]any)
|
| d, _ := chs[0].(map[string]any)["delta"].(map[string]any)
|
| tcs, _ := d["tool_calls"].([]any)
|
| if len(tcs) != 1 {
|
| t.Fatalf("frame %d tool_calls len=%d want 1", i, len(tcs))
|
| }
|
| fn, _ := tcs[0].(map[string]any)["function"].(map[string]any)
|
| if _, ok := fn["name"]; ok {
|
| gotName++
|
| if i != 0 || fn["name"] != "Bash" {
|
| t.Errorf("frame %d unexpected name=%v (name 只能出现在首帧且为 Bash)", i, fn["name"])
|
| }
|
| }
|
| if want := "arg" + string(rune('0'+i)); fn["arguments"] != want {
|
| t.Errorf("frame %d arguments=%v want %q", i, fn["arguments"], want)
|
| }
|
| }
|
| if gotName != 1 {
|
| t.Errorf("name 出现帧数=%d want 1", gotName)
|
| }
|
| }
|
|
|
|
|
|
|
| func TestStreamToolCallParallelFragments(t *testing.T) {
|
| raw := "data: {\"id\":\"x1\",\"object\":\"chat.completion.chunk\",\"created\":0,\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_0\",\"type\":\"function\",\"function\":{\"name\":\"Bash\",\"arguments\":\"\"}}]}}]}\n\n" +
|
| "data: {\"id\":\"x1\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":1,\"id\":\"call_1\",\"type\":\"function\",\"function\":{\"name\":\"Read\",\"arguments\":\"\"}}]}}]}\n\n" +
|
| "data: {\"id\":\"x1\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"name\":\"Bash\",\"arguments\":\"a0\"}}]}}]}\n\n" +
|
| "data: {\"id\":\"x1\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":1,\"function\":{\"name\":\"Read\",\"arguments\":\"a1\"}}]}}]}\n\n" +
|
| "data: {\"id\":\"x1\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"tool_calls\"}]}\n\n" +
|
| "data: [DONE]\n\n"
|
|
|
| frames, done := streamFrames(t, raw)
|
| if done != 1 {
|
| t.Fatalf("done=%d want 1", done)
|
| }
|
|
|
| argByIndex := map[int]string{}
|
| nameCount := map[int]int{}
|
| for _, fr := range frames {
|
| chs, _ := fr["choices"].([]any)
|
| d, _ := chs[0].(map[string]any)["delta"].(map[string]any)
|
| tcs, _ := d["tool_calls"].([]any)
|
| for _, tci := range tcs {
|
| tc, _ := tci.(map[string]any)
|
| idx := int(tc["index"].(float64))
|
| fn, _ := tc["function"].(map[string]any)
|
| if _, ok := fn["name"]; ok {
|
| nameCount[idx]++
|
| }
|
| if a, _ := fn["arguments"].(string); a != "" {
|
| argByIndex[idx] = a
|
| }
|
| }
|
| }
|
|
|
| if nameCount[0] != 1 || nameCount[1] != 1 {
|
| t.Errorf("name 出现次数 index0=%d index1=%d，各 want 1", nameCount[0], nameCount[1])
|
| }
|
| if argByIndex[0] != "a0" || argByIndex[1] != "a1" {
|
| t.Errorf("arguments index0=%q index1=%q want a0/a1", argByIndex[0], argByIndex[1])
|
| }
|
| }
|
|
|
|
|
|
|
| func TestStreamToolCallNoiseEmptyName(t *testing.T) {
|
| raw := "data: {\"id\":\"x1\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"m\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_a\",\"type\":\"function\",\"function\":{\"name\":\"lookup\",\"arguments\":\"\"}}]}}]}\n\n" +
|
| "data: {\"id\":\"x1\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"name\":\"\",\"arguments\":\"arg1\"}}]}}]}\n\n" +
|
| "data: {\"id\":\"x1\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"name\":\"\",\"arguments\":\"arg2\"}}]}}]}\n\n" +
|
| "data: {\"id\":\"x1\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"tool_calls\"}]}\n\n" +
|
| "data: [DONE]\n\n"
|
|
|
| frames, done := streamFrames(t, raw)
|
| if done != 1 {
|
| t.Fatalf("done=%d want 1", done)
|
| }
|
| for i, fr := range frames {
|
| chs, _ := fr["choices"].([]any)
|
| d, _ := chs[0].(map[string]any)["delta"].(map[string]any)
|
| tcs, _ := d["tool_calls"].([]any)
|
| for _, tci := range tcs {
|
| tc, _ := tci.(map[string]any)
|
| fn, _ := tc["function"].(map[string]any)
|
| if i == 0 {
|
| if fn["name"] != "lookup" {
|
| t.Errorf("frame 0 name=%v want lookup", fn["name"])
|
| }
|
| continue
|
| }
|
| if _, ok := fn["name"]; ok {
|
| t.Errorf("frame %d: 空串 name 应被剥离为键缺失, got %#v", i, fn)
|
| }
|
| }
|
| }
|
| }
|
|
|
|
|
|
|
|
|
| func TestStreamToolCallOverwriteClientSemantics(t *testing.T) {
|
| raw := "data: {\"id\":\"x1\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"m\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_a\",\"type\":\"function\",\"function\":{\"name\":\"Bash\",\"arguments\":\"\"}}]}}]}\n\n" +
|
| "data: {\"id\":\"x1\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"name\":\"\",\"arguments\":\"{\\\"cmd\\\":\\\"ls\\\"}\"}}]}}]}\n\n" +
|
| "data: {\"id\":\"x1\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"name\":\"Bash\",\"arguments\":\"\"}}]}}]}\n\n" +
|
| "data: [DONE]\n\n"
|
|
|
| frames, done := streamFrames(t, raw)
|
| if done != 1 {
|
| t.Fatalf("done=%d want 1", done)
|
| }
|
| state := map[int]string{}
|
| reconstructed := map[int]string{}
|
| for _, fr := range frames {
|
| chs, _ := fr["choices"].([]any)
|
| d, _ := chs[0].(map[string]any)["delta"].(map[string]any)
|
| tcs, _ := d["tool_calls"].([]any)
|
| for _, tci := range tcs {
|
| tc, _ := tci.(map[string]any)
|
| idx := int(tc["index"].(float64))
|
| fn, _ := tc["function"].(map[string]any)
|
|
|
| if name, ok := fn["name"]; ok {
|
| state[idx] = name.(string)
|
| }
|
| if state[idx] != "" {
|
| reconstructed[idx] = state[idx]
|
| }
|
| }
|
| }
|
|
|
| if len(reconstructed) != 1 || reconstructed[0] != "Bash" {
|
| t.Errorf("覆盖型重建 name=%v want map[0:Bash]", reconstructed)
|
| }
|
| }
|
|
|
|
|
| func streamFrames(t *testing.T, raw string) (frames []map[string]any, doneCount int) {
|
| t.Helper()
|
| rec := httptest.NewRecorder()
|
| if err := Stream(rec, strings.NewReader(raw)); err != nil {
|
| t.Fatal(err)
|
| }
|
| body := rec.Body.String()
|
| for _, ln := range strings.Split(body, "\n") {
|
| ln = strings.TrimSpace(ln)
|
| if strings.HasPrefix(ln, "data: [DONE]") {
|
| doneCount++
|
| continue
|
| }
|
| if strings.HasPrefix(ln, "data: ") {
|
| var obj map[string]any
|
| if err := json.Unmarshal([]byte(strings.TrimPrefix(ln, "data: ")), &obj); err != nil {
|
| t.Fatalf("bad frame %q: %v", ln, err)
|
| }
|
| frames = append(frames, obj)
|
| }
|
| }
|
| return frames, doneCount
|
| }
|
|
|
| func TestNormalizeFrame(t *testing.T) {
|
| cases := []struct {
|
| name string
|
| in map[string]any
|
| want string
|
| }{
|
| {"empty content/refusal and finish_reason empty string",
|
| map[string]any{"id": "x", "choices": []any{
|
| map[string]any{"index": 0, "delta": map[string]any{"content": "", "refusal": ""}, "finish_reason": ""},
|
| }},
|
| `{"choices":[{"delta":{},"finish_reason":null,"index":0}],"id":"x","object":"chat.completion.chunk","usage":null}`},
|
| {"non-empty tool_calls kept",
|
| map[string]any{"choices": []any{
|
| map[string]any{"index": 0, "delta": map[string]any{"tool_calls": []any{map[string]any{"id": "c1", "type": "function"}}}},
|
| }},
|
| `{"choices":[{"delta":{"tool_calls":[{"id":"c1","type":"function"}]},"finish_reason":null,"index":0}],"id":"chatcmpl-wb2api","object":"chat.completion.chunk","usage":null}`},
|
| {"empty tool_calls list dropped",
|
| map[string]any{"choices": []any{
|
| map[string]any{"index": 0, "delta": map[string]any{"tool_calls": []any{}, "content": "hi"}},
|
| }},
|
| `{"choices":[{"delta":{"content":"hi"},"finish_reason":null,"index":0}],"id":"chatcmpl-wb2api","object":"chat.completion.chunk","usage":null}`},
|
| {"empty placeholder function_call dropped",
|
| map[string]any{"choices": []any{
|
| map[string]any{"index": 0, "delta": map[string]any{"function_call": map[string]any{"name": "", "arguments": ""}}},
|
| }},
|
| `{"choices":[{"delta":{},"finish_reason":null,"index":0}],"id":"chatcmpl-wb2api","object":"chat.completion.chunk","usage":null}`},
|
| {"top-level unknown fields dropped, usage null when absent",
|
| map[string]any{"id": "x", "object": "chat.completion.chunk", "created": 1, "junk": "noise", "choices": []any{
|
| map[string]any{"index": 0, "delta": map[string]any{}, "finish_reason": "stop"},
|
| }},
|
| `{"choices":[{"delta":{},"finish_reason":"stop","index":0}],"created":1,"id":"x","object":"chat.completion.chunk","usage":null}`},
|
| }
|
| for _, c := range cases {
|
| t.Run(c.name, func(t *testing.T) {
|
| raw, err := json.Marshal(normalizeFrame(c.in))
|
| if err != nil {
|
| t.Fatal(err)
|
| }
|
| if string(raw) != c.want {
|
| t.Errorf("got %s\nwant %s", raw, c.want)
|
| }
|
| })
|
| }
|
| }
|
|
|
| func TestStreamNormalizesFrames(t *testing.T) {
|
|
|
|
|
| raw := "data: {\"id\":\"x1\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"m\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"\",\"reasoning_content\":\"\",\"refusal\":\"\",\"tool_calls\":[],\"function_call\":{\"name\":\"\",\"arguments\":\"\"}},\"finish_reason\":\"\"}],\"extra_field\":\"junk\"}\n\n" +
|
| "data: {\"id\":\"x1\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\",\"tool_calls\":[{\"id\":\"call_1\",\"type\":\"function\",\"function\":{\"name\":\"get_weather\",\"arguments\":\"{\\\"city\\\":\\\"北京\\\"}\"},\"index\":0}]},\"finish_reason\":\"\"}]}\n\n" +
|
| "data: {\"id\":\"x1\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"total_tokens\":7}}\n\n" +
|
| "data: [DONE]\n\n"
|
|
|
| frames, done := streamFrames(t, raw)
|
| if done != 1 {
|
| t.Fatalf("done frames=%d want 1", done)
|
| }
|
| if len(frames) != 3 {
|
| t.Fatalf("frames=%d want 3", len(frames))
|
| }
|
|
|
|
|
| f0 := frames[0]
|
| if _, ok := f0["extra_field"]; ok {
|
| t.Error("top-level extra_field should be dropped")
|
| }
|
| if f0["usage"] != nil {
|
| t.Errorf("usage should be null when absent, got %v", f0["usage"])
|
| }
|
| ch0 := f0["choices"].([]any)[0].(map[string]any)
|
| if ch0["finish_reason"] != nil {
|
| t.Errorf("frame1 finish_reason=%v want null", ch0["finish_reason"])
|
| }
|
| d := ch0["delta"].(map[string]any)
|
|
|
| if len(d) != 1 || d["role"] != "assistant" {
|
| t.Errorf("frame1 delta should only keep role, got %#v", d)
|
| }
|
| for _, noise := range []string{"content", "reasoning_content", "refusal", "tool_calls", "function_call"} {
|
| if _, ok := d[noise]; ok {
|
| t.Errorf("frame1 delta should drop %q, got %#v", noise, d)
|
| }
|
| }
|
|
|
|
|
| f1 := frames[1]
|
| ch1 := f1["choices"].([]any)[0].(map[string]any)
|
| d1 := ch1["delta"].(map[string]any)
|
| if d1["content"] != "hello" {
|
| t.Errorf("frame2 content=%v", d1["content"])
|
| }
|
| tcs, ok := d1["tool_calls"].([]any)
|
| if !ok || len(tcs) != 1 {
|
| t.Fatalf("frame2 tool_calls=%#v", d1["tool_calls"])
|
| }
|
| if ch1["finish_reason"] != nil {
|
| t.Errorf("frame2 finish_reason=%v want null (input empty string)", ch1["finish_reason"])
|
| }
|
|
|
|
|
| f2 := frames[2]
|
| ch2 := f2["choices"].([]any)[0].(map[string]any)
|
| if ch2["finish_reason"] != "stop" {
|
| t.Errorf("frame3 finish_reason=%v want stop", ch2["finish_reason"])
|
| }
|
| if f2["usage"].(map[string]any)["total_tokens"].(float64) != 7 {
|
| t.Errorf("frame3 usage=%v", f2["usage"])
|
| }
|
| }
|
|
|
|
|
|
|
|
|
| func TestStreamFirstIdPassthrough(t *testing.T) {
|
| raw := "data: {\"id\":\"chatcmpl-upstream-9\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"m\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"hello\"}}]}\n\n" +
|
|
|
| "data: {\"object\":\"chat.completion.chunk\",\"created\":1,\"choices\":[{\"index\":0,\"delta\":{\"content\":\" world\"}}]}\n\n" +
|
|
|
| "data: {\"id\":\"\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"!\"},\"finish_reason\":\"stop\"}]}\n\n" +
|
| "data: [DONE]\n\n"
|
|
|
| frames, done := streamFrames(t, raw)
|
| if done != 1 {
|
| t.Fatalf("done=%d want 1", done)
|
| }
|
| if len(frames) != 3 {
|
| t.Fatalf("frames=%d want 3", len(frames))
|
| }
|
| for i, fr := range frames {
|
| if got := fr["id"]; got != "chatcmpl-upstream-9" {
|
| t.Errorf("frame %d id=%v want chatcmpl-upstream-9 (首帧真实 id 续传)", i, got)
|
| }
|
| }
|
| }
|
|
|
|
|
|
|
| func TestStreamNoIdFallsBackToSentinel(t *testing.T) {
|
| raw := "data: {\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"}}]}\n\n" +
|
| "data: {\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{}," +
|
| "\"finish_reason\":\"stop\"}]}\n\n" +
|
| "data: [DONE]\n\n"
|
| frames, _ := streamFrames(t, raw)
|
| if len(frames) != 2 {
|
| t.Fatalf("frames=%d want 2", len(frames))
|
| }
|
| for i, fr := range frames {
|
| if got := fr["id"]; got != "chatcmpl-wb2api" {
|
| t.Errorf("frame %d id=%v want sentinel chatcmpl-wb2api (无真实 id)", i, got)
|
| }
|
| }
|
| }
|
|
|
| func TestStreamDoneFallback(t *testing.T) {
|
|
|
| rec := httptest.NewRecorder()
|
| err := Stream(rec, strings.NewReader("data: {\"id\":\"x1\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"}}]}\n\n"))
|
| if err != nil {
|
| t.Fatal(err)
|
| }
|
| body := rec.Body.String()
|
| if !strings.HasSuffix(strings.TrimRight(body, "\n"), "data: [DONE]") {
|
| t.Errorf("missing [DONE] fallback: %q", body)
|
| }
|
|
|
|
|
| rec2 := httptest.NewRecorder()
|
| if err := Stream(rec2, strings.NewReader("data: {\"id\":\"x1\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"}}]}\n\ndata: [DONE]\n\n")); err != nil {
|
| t.Fatal(err)
|
| }
|
| if n := strings.Count(rec2.Body.String(), "data: [DONE]"); n != 1 {
|
| t.Errorf("[DONE] count=%d want 1: %q", n, rec2.Body.String())
|
| }
|
| }
|
|
|
| func TestStreamPassthrough(t *testing.T) {
|
| rec := httptest.NewRecorder()
|
| err := Stream(rec, strings.NewReader(sseFixture))
|
| if err != nil {
|
| t.Fatal(err)
|
| }
|
| body := rec.Body.String()
|
| if !strings.Contains(body, "你好") || !strings.Contains(body, "data: [DONE]") {
|
| t.Errorf("body missing chunks: %q", body)
|
| }
|
|
|
| for _, ln := range strings.Split(strings.TrimRight(body, "\n"), "\n") {
|
| if ln != "" && !strings.HasPrefix(ln, "data: ") {
|
| t.Errorf("bad line: %q", ln)
|
| }
|
| }
|
| ct := rec.Header().Get("Content-Type")
|
| if !strings.Contains(ct, "text/event-stream") {
|
| t.Errorf("content-type=%q", ct)
|
| }
|
| }
|
|
|
|
|
|
|
| func TestAggregateEmptyStreamCases(t *testing.T) {
|
|
|
| valid := "data: {\"id\":\"c1\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"glm-5.2\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"hi\"}}]}\n\n" +
|
| "data: {\"id\":\"c1\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"total_tokens\":7}}\n\n" +
|
| "data: [DONE]\n\n"
|
|
|
| cases := []struct {
|
| name string
|
| raw string
|
| wantErr bool
|
|
|
| wantContent string
|
| wantUsage float64
|
| }{
|
| {
|
| name: "空流（EOF 即止）",
|
| raw: "",
|
| wantErr: true,
|
| },
|
| {
|
| name: "只有注释行和空行加 DONE",
|
| raw: ": comment\n\n: another comment\n\ndata: [DONE]\n\n",
|
| wantErr: true,
|
| },
|
| {
|
| name: "DONE 后跟垃圾帧不进聚合",
|
| raw: valid[:len(valid)-len("data: [DONE]\n\n")] + "data: [DONE]\n\ndata: {\"junk\":\"should not aggregate\"}\n\n",
|
| wantErr: false,
|
|
|
| wantContent: "hi",
|
| wantUsage: 7,
|
| },
|
| {
|
| name: "正常流回归",
|
| raw: valid,
|
| wantErr: false,
|
| wantContent: "hi",
|
| wantUsage: 7,
|
| },
|
| }
|
| for _, c := range cases {
|
| t.Run(c.name, func(t *testing.T) {
|
| resp, err := Aggregate(strings.NewReader(c.raw))
|
| if c.wantErr {
|
| if err == nil {
|
| t.Fatalf("expected error, got nil (resp=%v)", resp)
|
| }
|
| return
|
| }
|
| if err != nil {
|
| t.Fatal(err)
|
| }
|
| msg := resp["choices"].([]any)[0].(map[string]any)["message"].(map[string]any)
|
| if msg["content"] != c.wantContent {
|
| t.Errorf("content=%q want %q", msg["content"], c.wantContent)
|
| }
|
| if u, ok := resp["usage"].(map[string]any); ok {
|
| if u["total_tokens"].(float64) != c.wantUsage {
|
| t.Errorf("usage=%v want %v", u["total_tokens"], c.wantUsage)
|
| }
|
| } else {
|
| t.Errorf("usage missing")
|
| }
|
| })
|
| }
|
| }
|
|
|
|
|
| func TestAggregateEmptyStreamError(t *testing.T) {
|
| _, err := Aggregate(strings.NewReader(""))
|
| if err == nil || !strings.Contains(err.Error(), "no valid data events") {
|
| t.Fatalf("err=%v", err)
|
| }
|
| }
|
|
|
|
|
|
|
| func TestStreamEmptyFramesCase(t *testing.T) {
|
| cases := []struct {
|
| name string
|
| raw string
|
| }{
|
| {"空流", ""},
|
| {"只有注释行", ": comment\n\n"},
|
| {"只有 DONE", "data: [DONE]\n\n"},
|
| }
|
| for _, c := range cases {
|
| t.Run(c.name, func(t *testing.T) {
|
| rec := httptest.NewRecorder()
|
| err := Stream(rec, strings.NewReader(c.raw))
|
| if err == nil {
|
| t.Fatalf("expected error, got nil")
|
| }
|
| body := rec.Body.String()
|
| if n := strings.Count(body, "data: [DONE]"); n != 1 {
|
| t.Errorf("[DONE] count=%d want 1: %q", n, body)
|
| }
|
|
|
| var e map[string]any
|
| found := false
|
| for _, ln := range strings.Split(body, "\n") {
|
| ln = strings.TrimSpace(ln)
|
| if strings.HasPrefix(ln, "data: ") {
|
| payload := strings.TrimPrefix(ln, "data: ")
|
| if payload == "[DONE]" {
|
| continue
|
| }
|
| if json.Unmarshal([]byte(payload), &e) == nil {
|
| if em, ok := e["error"].(map[string]any); ok && em["message"] == "empty upstream stream" && em["type"] == "upstream_error" {
|
| found = true
|
| }
|
| }
|
| }
|
| }
|
| if !found {
|
| t.Errorf("error frame absent or error field stripped: %q", body)
|
| }
|
| })
|
| }
|
| }
|
|
|
|
|
| func TestStreamGarbageAfterDone(t *testing.T) {
|
| raw := "data: {\"id\":\"x1\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"}}]}\n\n" +
|
| "data: [DONE]\n\n" +
|
| "data: {\"should\":\"not appear\"}\n\n"
|
| rec := httptest.NewRecorder()
|
| if err := Stream(rec, strings.NewReader(raw)); err != nil {
|
| t.Fatal(err)
|
| }
|
| body := rec.Body.String()
|
| if strings.Contains(body, "should") {
|
| t.Errorf("garbage after DONE leaked into response: %q", body)
|
| }
|
| if n := strings.Count(body, "data: [DONE]"); n != 1 {
|
| t.Errorf("[DONE] count=%d want 1: %q", n, body)
|
| }
|
|
|
| if !strings.Contains(body, "hello") {
|
| t.Errorf("valid frame missing: %q", body)
|
| }
|
| }
|
|
|
|
|
|
|
| func TestStreamNormalPassthroughRegression(t *testing.T) {
|
| cases := []struct {
|
| name string
|
| raw string
|
| }{
|
| {"带 DONE 的正常流", sseFixture},
|
| {"漏发 DONE 自动补", "data: {\"id\":\"x1\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"}}]}\n\n"},
|
| }
|
| for _, c := range cases {
|
| t.Run(c.name, func(t *testing.T) {
|
| rec := httptest.NewRecorder()
|
| if err := Stream(rec, strings.NewReader(c.raw)); err != nil {
|
| t.Fatal(err)
|
| }
|
| body := rec.Body.String()
|
| if strings.Contains(body, `"error"`) {
|
| t.Errorf("unexpected error frame: %q", body)
|
| }
|
| if n := strings.Count(body, "data: [DONE]"); n != 1 {
|
| t.Errorf("[DONE] count=%d want 1: %q", n, body)
|
| }
|
|
|
| if !strings.Contains(body, `"object":"chat.completion.chunk"`) {
|
| t.Errorf("frame not normalized: %q", body)
|
| }
|
| })
|
| }
|
| }
|
|
|
|
|
|
|
|
|
| func TestNormalizeUsageCacheAliasesMirrorsNestedHit(t *testing.T) {
|
| usage := map[string]any{
|
| "prompt_tokens": 21041.0,
|
| "completion_tokens": 8.0,
|
| "total_tokens": 21049.0,
|
| "cache_read_input_tokens": 0.0,
|
| "cached_tokens": 0.0,
|
| "prompt_cache_hit_tokens": 0.0,
|
| "prompt_cache_miss_tokens": 177.0,
|
| "prompt_tokens_details": map[string]any{
|
| "cached_tokens": 20864.0,
|
| },
|
| }
|
|
|
| got := normalizeUsageCacheAliases(usage)
|
|
|
| for _, key := range []string{
|
| "cache_read_input_tokens",
|
| "cached_tokens",
|
| "prompt_cache_hit_tokens",
|
| } {
|
| if got[key] != 20864.0 {
|
| t.Fatalf("%s=%v want 20864", key, got[key])
|
| }
|
| }
|
| details := got["prompt_tokens_details"].(map[string]any)
|
| if details["cached_tokens"] != 20864.0 {
|
| t.Fatalf("prompt_tokens_details.cached_tokens=%v want 20864", details["cached_tokens"])
|
| }
|
| }
|
|
|
| func TestNormalizeUsageCacheAliasesPreservesZeroResult(t *testing.T) {
|
| usage := map[string]any{
|
| "prompt_tokens": 35.0,
|
| "completion_tokens": 2.0,
|
| "total_tokens": 37.0,
|
| "cache_read_input_tokens": 0.0,
|
| "prompt_cache_hit_tokens": 0.0,
|
| "prompt_tokens_details": map[string]any{
|
| "cached_tokens": 0.0,
|
| },
|
| }
|
|
|
| got := normalizeUsageCacheAliases(usage)
|
|
|
| details := got["prompt_tokens_details"].(map[string]any)
|
| if details["cached_tokens"] != 0.0 {
|
| t.Fatalf("prompt_tokens_details.cached_tokens=%v want 0", details["cached_tokens"])
|
| }
|
| }
|
|
|
|
|
|
|
|
|
| func TestStreamHintErrorFrameObserver(t *testing.T) {
|
| const errPayload = `{"error":{"code":6004,"message":"模型限流，将在 2026-09-27 01:00:00 重置"}}`
|
| raw := "data: {\"id\":\"x1\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"m\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"}}]}\n\n" +
|
| "data: " + errPayload + "\n\n" +
|
| "data: [DONE]\n\n"
|
|
|
| var observed []string
|
| rec := httptest.NewRecorder()
|
| if err := StreamHint(rec, strings.NewReader(raw), nil, WithErrorFrameObserver(func(payload string) {
|
| observed = append(observed, payload)
|
| })); err != nil {
|
| t.Fatal(err)
|
| }
|
|
|
| if len(observed) != 1 || observed[0] != errPayload {
|
| t.Fatalf("观察者回调 = %v want [%s]", observed, errPayload)
|
| }
|
|
|
| if !strings.Contains(rec.Body.String(), "6004") {
|
| t.Fatalf("error 帧原文未透传: %s", rec.Body.String())
|
| }
|
|
|
|
|
| var normalObserved int
|
| raw2 := "data: {\"id\":\"x2\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"m\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"}}]}\n\n" +
|
| "data: [DONE]\n\n"
|
| rec2 := httptest.NewRecorder()
|
| if err := StreamHint(rec2, strings.NewReader(raw2), nil, WithErrorFrameObserver(func(payload string) {
|
| normalObserved++
|
| })); err != nil {
|
| t.Fatal(err)
|
| }
|
| if normalObserved != 0 {
|
| t.Fatalf("正常流触发了观察者 %d 次，want 0", normalObserved)
|
| }
|
| }
|
|
|
| func TestUserResourceDetailedWithExpirySnapshot(t *testing.T) {
|
| now := time.Now().In(softRateResetLoc)
|
| soon := now.Add(24 * time.Hour).Truncate(time.Second)
|
| later := now.Add(10 * 24 * time.Hour).Truncate(time.Second)
|
| payload := `{"code":0,"data":{"Response":{"Data":{"Accounts":[` +
|
| `{"PackageName":"soon-a","CycleCapacitySize":10,"CycleCapacityRemain":10,"CycleCapacityUsed":0,"CycleEndTime":"` + soon.Format(packageEndLayout) + `"},` +
|
| `{"PackageName":"soon-b","CycleCapacitySize":15,"CycleCapacityRemain":15,"CycleCapacityUsed":0,"CycleEndTime":"` + soon.Format(packageEndLayout) + `"},` +
|
| `{"PackageName":"later","CycleCapacitySize":20,"CycleCapacityRemain":20,"CycleCapacityUsed":0,"CycleEndTime":"` + later.Format(packageEndLayout) + `"},` +
|
| `{"PackageName":"unknown","CycleCapacitySize":5,"CycleCapacityRemain":5,"CycleCapacityUsed":0}` +
|
| `]}}}}`
|
| c := testClient(func(r *http.Request) (*http.Response, error) {
|
| if !strings.HasSuffix(r.URL.Path, "/v2/billing/meter/get-user-resource") {
|
| return nil, errors.New("wrong path: " + r.URL.Path)
|
| }
|
| return jsonResp(200, payload), nil
|
| })
|
|
|
| remain, total, expiring, earliestAt, earliestRemaining, err := c.UserResourceDetailedWithExpiry(
|
| &auth.Auth{AccessToken: "at", UID: "u1"}, 48*time.Hour,
|
| )
|
| if err != nil {
|
| t.Fatalf("resource: %v", err)
|
| }
|
| if remain != 50 || total != 50 || expiring != 25 {
|
| t.Fatalf("remain/total/expiring=%d/%d/%d want 50/50/25", remain, total, expiring)
|
| }
|
| if earliestRemaining != 25 || !earliestAt.Equal(soon) {
|
| t.Fatalf("earliest=%v/%d want %v/25", earliestAt, earliestRemaining, soon)
|
| }
|
| }
|
| func TestCreditPackagesExpiryTimestamp(t *testing.T) {
|
| end := time.Now().In(softRateResetLoc).Add(7 * 24 * time.Hour).Truncate(time.Second)
|
| c := testClient(func(r *http.Request) (*http.Response, error) {
|
| if !strings.HasSuffix(r.URL.Path, "/v2/billing/meter/get-user-resource") {
|
| return nil, errors.New("wrong path: " + r.URL.Path)
|
| }
|
| return jsonResp(200, `{"code":0,"data":{"Response":{"Data":{"Accounts":[`+
|
| `{"PackageName":"gift","CycleCapacitySize":100,"CycleCapacityRemain":80,"CycleCapacityUsed":20,"CycleEndTime":"`+
|
| end.Format(packageEndLayout)+`"},`+
|
| `{"PackageName":"unknown","CycleCapacitySize":10,"CycleCapacityRemain":10,"CycleCapacityUsed":0}`+
|
| `]}}}}`), nil
|
| })
|
| packs, remain, size, err := c.CreditPackages(&auth.Auth{AccessToken: "at", UID: "u1"})
|
| if err != nil {
|
| t.Fatalf("packages: %v", err)
|
| }
|
| if remain != 90 || size != 110 {
|
| t.Fatalf("remain/size=%d/%d want 90/110", remain, size)
|
| }
|
| var found bool
|
| for _, p := range packs {
|
| if p.Name == "gift" {
|
| found = p.ExpiresAt == end.UnixMilli()
|
| }
|
| }
|
| if !found {
|
| t.Fatalf("gift pack missing Unix-ms expiry: %+v", packs)
|
| }
|
| }
|
|
|
