kefu/im/backend/internal/app/ai_providers_test.go

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)
	}
}