package app import ( "context" "encoding/json" "io" "net/http" "net/http/httptest" "strings" "testing" ) func captureAIRequest(t *testing.T, response string, status int) (*httptest.Server, *map[string]any, *http.Header) { t.Helper() body := map[string]any{} header := http.Header{} server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { raw, _ := io.ReadAll(io.LimitReader(r.Body, 1<<20)) _ = json.Unmarshal(raw, &body) for key, values := range r.Header { header[key] = values } body["__path"] = r.URL.Path w.Header().Set("Content-Type", "application/json") w.WriteHeader(status) _, _ = w.Write([]byte(response)) })) t.Cleanup(server.Close) return server, &body, &header } func testModel(protocol, baseURL string) aiModel { return aiModel{ ID: 1, Name: "test", Protocol: protocol, BaseURL: baseURL, ModelName: "demo-model", APIKey: "secret-key", Temperature: 0.5, MaxTokens: 128, TimeoutMS: 5000, Status: 1, } } func chatOnce(t *testing.T, protocol, baseURL string) (aiChatResult, error) { t.Helper() provider, err := newAIProvider(testModel(protocol, baseURL), false) if err != nil { t.Fatalf("newAIProvider(%s) error: %v", protocol, err) } return provider.Chat(context.Background(), aiChatRequest{ System: "你是测试人设", Messages: []aiChatMessage{{Role: "user", Text: "在吗"}, {Role: "assistant", Text: "在的"}, {Role: "user", Text: "吃了吗"}}, }) } func TestOpenAIProviderRequestAndResponse(t *testing.T) { server, body, header := captureAIRequest(t, `{"id":"cmpl-1","choices":[{"message":{"content":" 吃过啦 "},"finish_reason":"stop"}],"usage":{"prompt_tokens":12,"completion_tokens":4}}`, http.StatusOK) result, err := chatOnce(t, "openai", server.URL) if err != nil { t.Fatalf("openai chat failed: %v", err) } if result.Text != "吃过啦" || result.PromptTokens != 12 || result.CompletionTokens != 4 || result.FinishReason != "stop" { t.Fatalf("unexpected openai result: %+v", result) } if (*body)["__path"] != "/chat/completions" { t.Fatalf("unexpected openai path: %v", (*body)["__path"]) } if header.Get("Authorization") != "Bearer secret-key" { t.Fatalf("unexpected openai authorization: %q", header.Get("Authorization")) } // The system prompt is the first message for this protocol. messages, _ := (*body)["messages"].([]any) if len(messages) != 4 { t.Fatalf("expected system + 3 messages, got %d", len(messages)) } first, _ := messages[0].(map[string]any) if first["role"] != "system" || first["content"] != "你是测试人设" { t.Fatalf("unexpected first message: %v", first) } if (*body)["max_tokens"] != float64(128) || (*body)["temperature"] != 0.5 { t.Fatalf("model defaults were not applied: %v %v", (*body)["max_tokens"], (*body)["temperature"]) } } func TestAnthropicProviderUsesTopLevelSystem(t *testing.T) { server, body, header := captureAIRequest(t, `{"id":"msg-1","content":[{"type":"text","text":"吃过啦"}],"stop_reason":"end_turn","usage":{"input_tokens":9,"output_tokens":3}}`, http.StatusOK) result, err := chatOnce(t, "anthropic", server.URL) if err != nil { t.Fatalf("anthropic chat failed: %v", err) } if result.Text != "吃过啦" || result.PromptTokens != 9 || result.CompletionTokens != 3 { t.Fatalf("unexpected anthropic result: %+v", result) } if (*body)["__path"] != "/v1/messages" { t.Fatalf("unexpected anthropic path: %v", (*body)["__path"]) } if header.Get("X-Api-Key") != "secret-key" || header.Get("Anthropic-Version") == "" { t.Fatalf("unexpected anthropic headers: %v", *header) } if (*body)["system"] != "你是测试人设" { t.Fatalf("system prompt must stay outside the message list: %v", (*body)["system"]) } if messages, _ := (*body)["messages"].([]any); len(messages) != 3 { t.Fatalf("expected 3 messages, got %d", len(messages)) } } func TestGeminiProviderMapsAssistantToModel(t *testing.T) { server, body, header := captureAIRequest(t, `{"candidates":[{"content":{"parts":[{"text":"吃过啦"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":7,"candidatesTokenCount":2}}`, http.StatusOK) result, err := chatOnce(t, "gemini", server.URL) if err != nil { t.Fatalf("gemini chat failed: %v", err) } if result.Text != "吃过啦" || result.PromptTokens != 7 || result.CompletionTokens != 2 { t.Fatalf("unexpected gemini result: %+v", result) } if (*body)["__path"] != "/v1beta/models/demo-model:generateContent" { t.Fatalf("unexpected gemini path: %v", (*body)["__path"]) } // The key must not travel in the query string where proxies log it. if header.Get("X-Goog-Api-Key") != "secret-key" { t.Fatalf("gemini key must be sent as a header: %v", *header) } contents, _ := (*body)["contents"].([]any) if len(contents) != 3 { t.Fatalf("expected 3 contents, got %d", len(contents)) } second, _ := contents[1].(map[string]any) if second["role"] != "model" { t.Fatalf("assistant turns must be mapped to role model: %v", second["role"]) } } func TestWebhookProviderAcceptsAlternateReplyFields(t *testing.T) { server, _, _ := captureAIRequest(t, `{"reply":"吃过啦","usage":{"promptTokens":5,"completionTokens":2}}`, http.StatusOK) result, err := chatOnce(t, "webhook", server.URL) if err != nil { t.Fatalf("webhook chat failed: %v", err) } if result.Text != "吃过啦" || result.PromptTokens != 5 { t.Fatalf("unexpected webhook result: %+v", result) } } func TestAIProviderSurfacesUpstreamErrors(t *testing.T) { server, _, _ := captureAIRequest(t, `{"error":{"message":"insufficient quota"}}`, http.StatusTooManyRequests) if _, err := chatOnce(t, "openai", server.URL); err == nil || !strings.Contains(err.Error(), "HTTP 429") { t.Fatalf("upstream status must be reported: %v", err) } okServer, _, _ := captureAIRequest(t, `{"error":{"message":"model not found"}}`, http.StatusOK) _, err := chatOnce(t, "openai", okServer.URL) if err == nil || !strings.Contains(err.Error(), "model not found") { t.Fatalf("an error body with HTTP 200 must still fail: %v", err) } emptyServer, _, _ := captureAIRequest(t, `{"choices":[]}`, http.StatusOK) if _, err = chatOnce(t, "openai", emptyServer.URL); err == nil { t.Fatal("an empty choice list must fail rather than send an empty reply") } brokenServer, _, _ := captureAIRequest(t, `not json`, http.StatusOK) if _, err = chatOnce(t, "openai", brokenServer.URL); err == nil { t.Fatal("a non-JSON body must fail") } } func TestParseAIBaseURL(t *testing.T) { tests := []struct { protocol string endpoint string production bool valid bool }{ {protocol: "openai", endpoint: "https://api.openai.com/v1", production: true, valid: true}, {protocol: "openai", endpoint: "http://127.0.0.1:11434/v1", production: false, valid: true}, {protocol: "openai", endpoint: "http://127.0.0.1:11434/v1", production: true, valid: false}, {protocol: "openai", endpoint: "https://user:pass@api.example.com/v1", production: true, valid: false}, {protocol: "openai", endpoint: "https://api.example.com/v1?key=secret", production: true, valid: false}, {protocol: "openai", endpoint: "", production: false, valid: false}, {protocol: "debug", endpoint: "", production: false, valid: true}, } for _, test := range tests { _, err := parseAIBaseURL(test.protocol, test.endpoint, test.production) if (err == nil) != test.valid { t.Errorf("parseAIBaseURL(%q, %q, production=%v) error = %v", test.protocol, test.endpoint, test.production, err) } } } func TestAIProviderFactoryGuards(t *testing.T) { if _, err := newAIProvider(testModel("debug", ""), true); err == nil { t.Fatal("production must refuse the debug protocol") } if _, err := newAIProvider(testModel("unknown", "https://example.com"), false); err == nil { t.Fatal("unknown protocols must be refused") } missingKey := testModel("openai", "https://example.com") missingKey.APIKey = "" if _, err := newAIProvider(missingKey, false); err == nil { t.Fatal("a missing key must be refused before the request is made") } // A webhook may be protected by network rules instead of a key. if _, err := newAIProvider(func() aiModel { m := testModel("webhook", "https://example.com/hook"); m.APIKey = ""; return m }(), false); err != nil { t.Fatalf("webhook without a key must be allowed: %v", err) } } func TestDebugProviderIsDeterministicAndOffline(t *testing.T) { provider, err := newAIProvider(testModel("debug", ""), false) if err != nil { t.Fatalf("debug provider unavailable: %v", err) } request := aiChatRequest{Messages: []aiChatMessage{{Role: "user", Text: "今天天气怎么样呀在不在"}}} first, err := provider.Chat(context.Background(), request) if err != nil { t.Fatalf("debug chat failed: %v", err) } second, _ := provider.Chat(context.Background(), request) if first.Text != second.Text || first.Text == "" { t.Fatalf("debug replies must be deterministic: %q vs %q", first.Text, second.Text) } if empty, _ := provider.Chat(context.Background(), aiChatRequest{}); empty.Text == "" { t.Fatal("debug provider must answer even without history") } } // A reasoning model answers in two parts. When the token ceiling is reached // before it stops thinking, the reply field is empty and the whole budget shows // up as completion tokens — exactly what the admin log showed as "成功" with a // blank reply. func TestReasoningModelExhaustingItsBudgetIsRecognised(t *testing.T) { server, _, _ := captureAIRequest(t, `{"id":"call-1","choices":[{"message":{"content":"","reasoning_content":"先想想怎么回答"},"finish_reason":"length"}],"usage":{"prompt_tokens":240,"completion_tokens":256}}`, 200) result, err := chatOnce(t, "openai", server.URL) if err != nil { t.Fatalf("空正文本身不是传输错误: %v", err) } if result.Text != "" || result.Reasoning != "先想想怎么回答" { t.Fatalf("result = %+v", result) } if !result.truncated() { t.Fatal("必须能识别出被截断") } // The ordinary case must not be mistaken for it. plain, _, _ := captureAIRequest(t, `{"id":"call-2","choices":[{"message":{"content":"在的,刚吃过"},"finish_reason":"stop"}],"usage":{"prompt_tokens":240,"completion_tokens":62}}`, 200) answered, err := chatOnce(t, "openai", plain.URL) if err != nil { t.Fatal(err) } if answered.Text != "在的,刚吃过" || answered.truncated() { t.Fatalf("result = %+v", answered) } // An empty reply with no reasoning and a normal stop is still just empty. blank, _, _ := captureAIRequest(t, `{"id":"call-3","choices":[{"message":{"content":""},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":0}}`, 200) empty, err := chatOnce(t, "openai", blank.URL) if err != nil { t.Fatal(err) } if empty.truncated() { t.Fatal("普通空回复不该被当成截断") } } func TestAIRetryMaxTokensStaysWithinBounds(t *testing.T) { for _, test := range []struct{ current, want int }{ {256, 1024}, // the configured value that produced blank replies {64, 1024}, // never retry with less than a usable budget {512, 2048}, // four times when that is already sensible {2048, 4096}, // capped, so a misconfiguration cannot bill without limit {8192, 4096}, } { if got := aiRetryMaxTokens(test.current); got != test.want { t.Fatalf("aiRetryMaxTokens(%d) = %d, want %d", test.current, got, test.want) } } } // The request the retry sends has to actually carry the raised ceiling. func TestChatRequestMaxTokensOverridesTheModelConfiguration(t *testing.T) { server, body, _ := captureAIRequest(t, `{"id":"c","choices":[{"message":{"content":"好"},"finish_reason":"stop"}],"usage":{}}`, 200) provider, err := newAIProvider(testModel("openai", server.URL), false) if err != nil { t.Fatal(err) } if _, err = provider.Chat(context.Background(), aiChatRequest{ Messages: []aiChatMessage{{Role: "user", Text: "在吗"}}, MaxTokens: aiRetryMaxTokens(128), }); err != nil { t.Fatal(err) } if got := (*body)["max_tokens"]; got != float64(1024) { t.Fatalf("max_tokens = %v, want 1024", got) } } // A reply that opens with a number used to arrive with the number missing: the // list-marker trimming removed the digits themselves. func TestSanitizeAIReplyKeepsLeadingNumbers(t *testing.T) { for _, test := range []struct{ raw, want string }{ {"30岁啦,你呢", "30岁啦,你呢"}, {"1. 先自我介绍", "先自我介绍"}, {"2、我在深圳", "我在深圳"}, {"- 喜欢跑步", "喜欢跑步"}, {"• 也喜欢音乐", "也喜欢音乐"}, {"**在的**", "在的"}, {"18612725546 是我的号码", "18612725546 是我的号码"}, {"3 个月前搬来的", "3 个月前搬来的"}, } { if got := sanitizeAIReply(test.raw); got != test.want { t.Fatalf("sanitizeAIReply(%q) = %q, want %q", test.raw, got, test.want) } } } // A cut-off answer is worth one more attempt too, not just an empty one. func TestTruncatedAnswerIsRetriedAndKeptIfItSurvives(t *testing.T) { cut, _, _ := captureAIRequest(t, `{"id":"c","choices":[{"message":{"content":"30岁啦,"},"finish_reason":"length"}],"usage":{"prompt_tokens":240,"completion_tokens":1024}}`, 200) result, err := chatOnce(t, "openai", cut.URL) if err != nil { t.Fatal(err) } if !result.truncated() { t.Fatal("被截断的正文也要触发重试") } if sanitizeAIReply(result.Text) != "30岁啦," { t.Fatalf("text = %q", result.Text) } } // A model that reasons by default is told not to through its stored protocol // fields, without a code change and without a second model row. func TestModelExtraFieldsReachTheProviderWithoutOverridingTheCall(t *testing.T) { server, body, _ := captureAIRequest(t, `{"id":"c","choices":[{"message":{"content":"在的"},"finish_reason":"stop"}],"usage":{}}`, 200) model := testModel("openai", server.URL) model.Extra = map[string]any{ "thinking": map[string]any{"type": "disabled"}, // Structural fields and the budget stay under this code's control. "model": "someone-elses-model", "max_tokens": 999999, "stream": true, } provider, err := newAIProvider(model, false) if err != nil { t.Fatal(err) } if _, err = provider.Chat(context.Background(), aiChatRequest{Messages: []aiChatMessage{{Role: "user", Text: "在吗"}}}); err != nil { t.Fatal(err) } thinking, ok := (*body)["thinking"].(map[string]any) if !ok || thinking["type"] != "disabled" { t.Fatalf("thinking = %v", (*body)["thinking"]) } if (*body)["model"] != "demo-model" || (*body)["max_tokens"] != float64(128) || (*body)["stream"] != false { t.Fatalf("扩展参数不得覆盖调用本身: %v", *body) } } func TestAIModelExtraIsValidatedBeforeItIsStored(t *testing.T) { keep := "{\"thinking\":{\"type\":\"disabled\"}}" value, err := aiModelExtra(&keep, nil) if err != nil || value != keep { t.Fatalf("value = %v err = %v", value, err) } blank := " " if value, err = aiModelExtra(&blank, "old"); err != nil || value != nil { t.Fatalf("留空应清除扩展参数: %v %v", value, err) } broken := "{not json" if _, err = aiModelExtra(&broken, nil); err == nil { t.Fatal("非法 JSON 必须在写库前被拒绝") } // Omitted entirely: the stored value survives an edit that does not mention it. if value, err = aiModelExtra(nil, "old"); err != nil || value != "old" { t.Fatalf("value = %v err = %v", value, err) } }