ModelRouter/internal/gateway/messages_test.go

package gateway

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"testing"

	"llmsproxy/internal/config"
	"llmsproxy/internal/core"
)

// 本文件端到端验证 Phase 2 的核心主张：**对外协议由适配器决定，内核只调度**。
//
// 具体做法：起一个真实的 core + gateway，源挂 anthropic 适配器，上游是一个
// 只懂 Anthropic Messages 的假服务器。然后：
//
//	1. Anthropic 客户端形状的请求 → 拿到 Anthropic 形状的响应
//	2. 调度/配额/审计与 /v1/chat/completions 走同一条路径
//	3. 上游若不支持该模型，内核如实回 503，不伪造格式

// anthropicUpstream 是一个只实现 Anthropic Messages 的假上游。
func anthropicUpstream(t *testing.T) *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !strings.HasSuffix(r.URL.Path, "/messages") {
			// 网关把请求发到了非 Anthropic 路径 —— 显式暴露这个 bug
			t.Errorf("请求打到了非 Anthropic 路径: %s", r.URL.Path)
			w.WriteHeader(http.StatusNotFound)
			return
		}
		body, _ := io.ReadAll(r.Body)
		var probe map[string]interface{}
		_ = json.Unmarshal(body, &probe)
		if _, ok := probe["messages"]; !ok {
			t.Errorf("上游收到没有 messages 的请求体: %s", body)
		}
		if _, ok := probe["system"]; !ok && probe["system"] == nil {
			// system 是可选的，仅在客户端给了的时候才应出现；不强制
			_ = ok
		}
		if _, ok := probe["max_tokens"]; !ok {
			t.Errorf("Anthropic 要求 max_tokens，上游请求体却缺这个字段: %s", body)
		}

		// Streaming requests must get a real Anthropic SSE sequence; a JSON body
		// here would leave the adapter with zero frames and the gateway would
		// (correctly) reject the stream as empty.
		if stream, _ := probe["stream"].(bool); stream {
			w.Header().Set("Content-Type", "text/event-stream")
			fl, _ := w.(http.Flusher)
			frames := []string{
				`event: message_start` + "\n" + `data: {"type":"message_start","message":{"id":"msg_mock","type":"message","role":"assistant","model":"` + probeModel(probe) + `","content":[],"usage":{"input_tokens":3,"output_tokens":0}}}`,
				`event: content_block_start` + "\n" + `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
				`event: content_block_delta` + "\n" + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"pong"}}`,
				`event: content_block_stop` + "\n" + `data: {"type":"content_block_stop","index":0}`,
				`event: message_delta` + "\n" + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":5}}`,
				`event: message_stop` + "\n" + `data: {"type":"message_stop"}`,
			}
			for _, f := range frames {
				fmt.Fprint(w, f+"\n\n")
				if fl != nil {
					fl.Flush()
				}
			}
			return
		}
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, `{
			"id":"msg_mock","type":"message","role":"assistant","model":"`+probeModel(probe)+`",
			"content":[{"type":"thinking","thinking":"let me think"},{"type":"text","text":"pong"}],
			"stop_reason":"end_turn",
			"usage":{"input_tokens":3,"output_tokens":5}
		}`)
	}))
}

func probeModel(m map[string]interface{}) string {
	if s, ok := m["model"].(string); ok {
		return s
	}
	return ""
}

// messagesGateway 起一个真实网关，源挂 anthropic 适配器。
func messagesGateway(t *testing.T) (*httptest.Server, string) {
	t.Helper()
	up := anthropicUpstream(t)
	t.Cleanup(up.Close)

	dir := t.TempDir()
	// adapterDir 故意不预先创建：VM 首启会把内置适配器 seed 进去。
	// 预先创建会被当成"目录已存在即为权威"，于是没有任何内置适配器被加载，
	// 表现为 "adapter anthropic not loaded"。
	cfg := &config.Config{
		Path:         filepath.Join(dir, "config.yaml"),
		AdapterDir:   filepath.Join(dir, "adapters"),
		RuntimeFile:  filepath.Join(dir, "runtime.json"),
		Listen:       "127.0.0.1:0",
		DefaultModel: "claude-test",
		GatewayKeys:  []string{"sk-e2e-0001"},
		Sources: []config.Source{{
			Name: "anth", BaseURL: up.URL, Adapter: "anthropic",
			APIKey: "k", MaxConcurrent: 4,
			Models: []config.Model{{ID: "claude-test", Kind: "chat", Priority: 100}},
		}},
	}

	c, err := core.NewFromConfig(cfg)
	if err != nil {
		t.Fatalf("core: %v", err)
	}
	t.Cleanup(c.Close)
	gw, err := New(c)
	if err != nil {
		t.Fatalf("gateway: %v", err)
	}
	ts := httptest.NewServer(gw.Handler())
	t.Cleanup(ts.Close)
	return ts, "sk-e2e-0001"
}

func postJSON(t *testing.T, url, key, body string) (int, string) {
	t.Helper()
	req, _ := http.NewRequest(http.MethodPost, url, strings.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+key)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("request: %v", err)
	}
	defer resp.Body.Close()
	b, _ := io.ReadAll(resp.Body)
	return resp.StatusCode, string(b)
}

// TestMessagesEndpointSpeaksAnthropic 是本文件的主断言：Anthropic 客户端
// 打 /v1/messages，拿到 Anthropic 形状。
func TestMessagesEndpointSpeaksAnthropic(t *testing.T) {
	ts, key := messagesGateway(t)
	code, body := postJSON(t, ts.URL+"/v1/messages", key,
		`{"model":"claude-test","max_tokens":64,"messages":[{"role":"user","content":"hi"}]}`)
	if code != 200 {
		t.Fatalf("status = %d: %s", code, body)
	}
	var m map[string]interface{}
	if err := json.Unmarshal([]byte(body), &m); err != nil {
		t.Fatalf("响应不是 JSON: %v (%s)", err, body)
	}
	if m["type"] != "message" {
		t.Errorf("type = %v, want message（内核不该退回 OpenAI 的 chat.completion）", m["type"])
	}
	if m["stop_reason"] != "end_turn" {
		t.Errorf("stop_reason = %v, want end_turn", m["stop_reason"])
	}
	blocks, _ := m["content"].([]interface{})
	if len(blocks) != 2 {
		t.Fatalf("content blocks = %v, want [thinking, text]", m["content"])
	}
	if b0 := blocks[0].(map[string]interface{}); b0["type"] != "thinking" || b0["thinking"] != "let me think" {
		t.Errorf("首个块应为 thinking: %v", b0)
	}
	if b1 := blocks[1].(map[string]interface{}); b1["type"] != "text" || b1["text"] != "pong" {
		t.Errorf("第二个块应为 text/pong: %v", b1)
	}
	u := m["usage"].(map[string]interface{})
	if u["input_tokens"] != float64(3) || u["output_tokens"] != float64(5) {
		t.Errorf("usage 应为 Anthropic 的 input/output 命名: %v", u)
	}
	// 内核不得混入 OpenAI 形状
	for _, forbidden := range []string{`"object"`, `"choices"`, `"chat.completion"`, `"prompt_tokens"`} {
		if strings.Contains(body, forbidden) {
			t.Errorf("OpenAI 字段 %s 混入 Anthropic 响应: %s", forbidden, body)
		}
	}
}

// TestMessagesStreamSpeaksAnthropic 验证流式也是 Anthropic 事件序列，
// 而不是 OpenAI SSE。
func TestMessagesStreamSpeaksAnthropic(t *testing.T) {
	ts, key := messagesGateway(t)
	req, _ := http.NewRequest(http.MethodPost, ts.URL+"/v1/messages",
		strings.NewReader(`{"model":"claude-test","max_tokens":64,"stream":true,"messages":[{"role":"user","content":"hi"}]}`))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+key)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(resp.Body)
	s := string(body)

	if !strings.Contains(s, "message_start") {
		t.Errorf("Anthropic 流应以 message_start 开头: %s", s)
	}
	if !strings.Contains(s, "message_stop") {
		t.Errorf("Anthropic 流应以 message_stop 收尾: %s", s)
	}
	// 内核不得擅自给 Anthropic 流加 OpenAI 的 [DONE]
	if strings.Contains(s, "[DONE]") {
		t.Errorf("Anthropic 流里出现了 OpenAI 的 [DONE] 终止标记: %s", s)
	}
	if strings.Contains(s, "chat.completion.chunk") {
		t.Errorf("流里混入 OpenAI chunk 形状: %s", s)
	}
}

// TestMessagesRejectsBadRequest 确认入向校验（Anthropic 要求 max_tokens）
// 与 OpenAI 路径一样在入口拒绝，不占用上游槽位。
func TestMessagesRejectsBadRequest(t *testing.T) {
	ts, key := messagesGateway(t)
	// Anthropic 允许省略 max_tokens（上游会给默认），所以「有 messages 无
	// max_tokens」必须被接受；真正该拒的是没有 messages。
	code, body := postJSON(t, ts.URL+"/v1/messages", key,
		`{"model":"claude-test","messages":[{"role":"user","content":"hi"}]}`)
	if code != 200 {
		t.Fatalf("有 messages 的请求应被接受（max_tokens 可省），却得 %d: %s", code, body)
	}
	code, body = postJSON(t, ts.URL+"/v1/messages", key, `{"model":"claude-test"}`)
	if code != 400 {
		t.Errorf("无 messages 的请求 status = %d, want 400: %s", code, body)
	}
	// 坏 JSON
	code, _ = postJSON(t, ts.URL+"/v1/messages", key, `{"model":`)
	if code != 400 {
		t.Errorf("坏 JSON status = %d, want 400", code)
	}
}

// TestMessagesUnconfiguredFallsBackToOpenai 验证 R5 的兜底在新增端点上同样
// 生效：源没有声明序列化钩子时（此例用 openai 适配器），回 OpenAI 形状。
func TestMessagesUnconfiguredFallsBackToOpenai(t *testing.T) {
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, `{"choices":[{"message":{"content":"pong"},"finish_reason":"stop"}],
			"usage":{"prompt_tokens":3,"completion_tokens":1,"total_tokens":4}}`)
	}))
	defer up.Close()

	dir := t.TempDir()
	// adapterDir 故意不预先创建：VM 首启会把内置适配器 seed 进去。
	// 预先创建会被当成"目录已存在即为权威"，于是没有任何内置适配器被加载，
	// 表现为 "adapter anthropic not loaded"。
	cfg := &config.Config{
		Path:         filepath.Join(dir, "config.yaml"),
		AdapterDir:   filepath.Join(dir, "adapters"),
		RuntimeFile:  filepath.Join(dir, "runtime.json"),
		Listen:       "127.0.0.1:0",
		DefaultModel: "gpt-test",
		GatewayKeys:  []string{"sk-e2e-0001"},
		Sources: []config.Source{{
			Name: "oai", BaseURL: up.URL, Adapter: "openai", APIKey: "k", MaxConcurrent: 4,
			Models: []config.Model{{ID: "gpt-test", Kind: "chat"}},
		}},
	}
	c, err := core.NewFromConfig(cfg)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	gw, err := New(c)
	if err != nil {
		t.Fatal(err)
	}
	ts := httptest.NewServer(gw.Handler())
	defer ts.Close()

	code, body := postJSON(t, ts.URL+"/v1/messages", "sk-e2e-0001",
		`{"model":"gpt-test","max_tokens":64,"messages":[{"role":"user","content":"hi"}]}`)
	if code != 200 {
		t.Fatalf("status = %d: %s", code, body)
	}
	if !strings.Contains(body, "chat.completion") {
		t.Errorf("无序列化钩子的源应回落 OpenAI 兜底形状: %s", body)
	}
}

// TestMessagesRequiresAuth 确认新端点走同一套鉴权（不因为是新协议而放行）。
func TestMessagesRequiresAuth(t *testing.T) {
	ts, _ := messagesGateway(t)
	code, _ := postJSON(t, ts.URL+"/v1/messages", "wrong-key",
		`{"model":"claude-test","max_tokens":8,"messages":[{"role":"user","content":"hi"}]}`)
	if code != 401 {
		t.Errorf("错误 key status = %d, want 401", code)
	}
	// 无 auth
	req, _ := http.NewRequest(http.MethodPost, ts.URL+"/v1/messages", strings.NewReader("{}"))
	req.Header.Set("Content-Type", "application/json")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != 401 {
		t.Errorf("无 key status = %d, want 401", resp.StatusCode)
	}
	_ = context.Background()
}