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