ModelRouter/internal/gateway/serializer_test.go

package gateway

import (
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"

	"llmsproxy/internal/lua"
	"llmsproxy/internal/types"
)

// 本文件锁定对外序列化的**优先级契约**：
//
//	适配器声明 serialize_* → 适配器完全接管
//	适配器未声明          → 兜底 OpenAI（保证刚部署/未配适配器仍能服务）
//
// 以及「让步」的定义：适配器接管后，内核绝不再往它的输出上补任何 OpenAI
// 字段。两个序列化器混在一起产出的东西比任何一个都糟——客户端无法解析，
// 而排查时会指向错误的层。

// luaVMFor 建一个会 seed 内置适配器的 VM（目录名必须是 "adapters"）。
func luaVMFor(t *testing.T) *lua.VM {
	t.Helper()
	vm := lua.NewVM(filepath.Join(t.TempDir(), "adapters"))
	if err := vm.Start(); err != nil {
		t.Fatalf("vm start: %v", err)
	}
	t.Cleanup(vm.Stop)
	return vm
}

// TestFallbackUsedWhenAdapterDeclaresNothing 是「兜底」本身的测试：
// 没有任何内置适配器声明 serialize_*，所以默认路径必须是 OpenAI。
func TestFallbackUsedWhenAdapterDeclaresNothing(t *testing.T) {
	vm := luaVMFor(t)
	// 先证明前提：内置适配器确实没声明这些钩子（否则下面测的是别的东西）
	for _, name := range []string{"openai", "anthropic", "gemini", "ollama", "deepseek"} {
		for _, hook := range []string{"serialize_response", "serialize_chunk"} {
			if vm.HasFn(name, hook) {
				t.Fatalf("内置适配器 %s 声明了 %s —— 本测试的前提（默认走兜底）已不成立", name, hook)
			}
		}
	}
	ser := serializerFor(vm, "openai")
	if _, isFallback := ser.(openAISerializer); !isFallback {
		t.Fatalf("未声明钩子的适配器应得到兜底序列化器，得到 %T", ser)
	}
}

// TestAdapterSerializerWinsOverFallback 锁定核心规则：适配器声明钩子即接管，
// 且内核不与兜底混合任何字段。
func TestAdapterSerializerWinsOverFallback(t *testing.T) {
	dir := t.TempDir()
	if err := os.MkdirAll(filepath.Join(dir, "adapters"), 0o755); err != nil {
		t.Fatal(err)
	}
	// 一个刻意用**非 OpenAI 形状**的适配器：如果内核混入了兜底字段
	// （object/choices/usage 等），下面断言会失败。
	adapter := `
local adapter = {}
adapter.name = "anthropicish"
adapter.version = "1.0.0"
adapter.endpoint = "/chat/completions"

-- 非流式：故意用 content[] 分块 + 没有 choices/usage
function adapter.serialize_response(payload)
  local p = json.decode(payload)
  return json.encode({
    type = "message",
    role = "assistant",
    model = p.model,
    content = { { type = "text", text = p.content } },
    stop_reason = (p.finish_reason ~= "" and p.finish_reason or "end_turn"),
  })
end

-- 流式：event: style 的帧，且不带 OpenAI 的 object 字段
function adapter.serialize_chunk(payload)
  local p = json.decode(payload)
  return json.encode({
    type = "content_block_delta",
    delta = { type = "text_delta", text = p.content },
  })
end

return adapter
`
	if err := os.WriteFile(filepath.Join(dir, "adapters", "anthropicish.lua"), []byte(adapter), 0o644); err != nil {
		t.Fatal(err)
	}
	vm := lua.NewVM(filepath.Join(dir, "adapters"))
	if err := vm.Start(); err != nil {
		t.Fatalf("vm: %v", err)
	}
	defer vm.Stop()

	ser := serializerFor(vm, "anthropicish")
	if _, isFallback := ser.(openAISerializer); isFallback {
		t.Fatal("声明了 serialize_* 的适配器必须接管序列化，而不是拿到兜底")
	}

	// 非流式输出必须是适配器形状，且**不含**任何 OpenAI 字段
	resp := &types.UnifiedResponse{
		Model:        "m",
		Content:      "hello",
		FinishReason: "stop",
		TokenUsage:   types.TokenUsage{Prompt: 3, Completion: 1, Total: 4},
	}
	body, err := ser.NonStream(resp, "m", "id-1", 1700000000)
	if err != nil {
		t.Fatalf("NonStream: %v", err)
	}
	var got map[string]interface{}
	if err := json.Unmarshal(body, &got); err != nil {
		t.Fatalf("adapter output not JSON: %v (%s)", err, body)
	}
	if got["type"] != "message" {
		t.Errorf("适配器输出未被采用：type = %v, want message", got["type"])
	}
	// 「让步」的硬断言：OpenAI 字段一个都不许出现
	for _, k := range []string{"object", "choices", "usage", "created"} {
		if _, leaked := got[k]; leaked {
			t.Errorf("兜底 OpenAI 字段 %q 混进了适配器的输出（让步规则被违反）: %s", k, body)
		}
	}
	// 适配器没有读 usage，但兜底也不许替它补上
	if _, has := got["token_usage"]; has {
		t.Errorf("兜底补了 token_usage 字段，适配器输出被污染: %s", body)
	}
}

// TestAdapterSerializerStreamWinsOverFallback 同上，流式路径。
func TestAdapterSerializerStreamWinsOverFallback(t *testing.T) {
	dir := t.TempDir()
	if err := os.MkdirAll(filepath.Join(dir, "adapters"), 0o755); err != nil {
		t.Fatal(err)
	}
	adapter := `
local adapter = {}
adapter.name = "streamsh"
adapter.version = "1.0.0"
adapter.endpoint = "/chat/completions"

function adapter.serialize_chunk(payload)
  local p = json.decode(payload)
  return json.encode({ type = "content_block_delta", text = p.content })
end

-- 显式声明收尾与终止标记，证明协议语义完全由适配器决定
function adapter.serialize_close(payload)
  local p = json.decode(payload)
  return json.encode({ type = "message_stop", stop_reason = p.finish_reason })
end

return adapter
`
	if err := os.WriteFile(filepath.Join(dir, "adapters", "streamsh.lua"), []byte(adapter), 0o644); err != nil {
		t.Fatal(err)
	}
	vm := lua.NewVM(filepath.Join(dir, "adapters"))
	if err := vm.Start(); err != nil {
		t.Fatalf("vm: %v", err)
	}
	defer vm.Stop()

	ser := serializerFor(vm, "streamsh")
	if _, isFallback := ser.(openAISerializer); isFallback {
		t.Fatal("声明 serialize_chunk 的适配器必须接管流式序列化")
	}
	st := &StreamState{}
	b, err := ser.Chunk(&types.UnifiedChunk{Content: "hi"}, "m", "id", 1, st)
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(string(b), `"content_block_delta"`) {
		t.Errorf("chunk 未采用适配器形状: %s", b)
	}
	if strings.Contains(string(b), "chat.completion.chunk") || strings.Contains(string(b), `"choices"`) {
		t.Errorf("兜底 OpenAI 帧混进了适配器输出: %s", b)
	}

	// 终止标记：适配器接管流式后，内核不得自作主张发 [DONE]
	if m := ser.DoneMarker(); m != "" {
		t.Errorf("适配器接管流式时 DoneMarker 应为空（不擅自发 [DONE]），实得 %q", m)
	}
	// 但它没声明 serialize_close 时，收尾仍借用兜底形状 —— 这是「没实现」
	// 不是「冲突」，允许兜底。
	if _, isFallbackClose := ser.(interface {
		Close(model, id string, created int64, st *StreamState) ([][]byte, error)
	}); !isFallbackClose {
		t.Fatal("Close 不应在接口上")
	}
}

// TestFallbackIsUsedWhenAdapterHookFails 验证钩子**执行失败**时的行为。
//
// 契约：退回兜底的完整输出，而不是产出「一半适配器字段 + 一半 OpenAI
// 字段」的混合体。混合体比任何一个单独结果都更难排查。
func TestFallbackIsUsedWhenAdapterHookFails(t *testing.T) {
	dir := t.TempDir()
	if err := os.MkdirAll(filepath.Join(dir, "adapters"), 0o755); err != nil {
		t.Fatal(err)
	}
	adapter := `
local adapter = {}
adapter.name = "broken"
adapter.version = "1.0.0"
adapter.endpoint = "/chat/completions"
function adapter.serialize_response(payload)
  error("deliberate failure")
end
return adapter
`
	if err := os.WriteFile(filepath.Join(dir, "adapters", "broken.lua"), []byte(adapter), 0o644); err != nil {
		t.Fatal(err)
	}
	vm := lua.NewVM(filepath.Join(dir, "adapters"))
	if err := vm.Start(); err != nil {
		t.Fatalf("vm: %v", err)
	}
	defer vm.Stop()

	ser := serializerFor(vm, "broken")
	resp := &types.UnifiedResponse{Content: "fallback still works", FinishReason: "stop"}
	body, err := ser.NonStream(resp, "m", "id", 1)
	if err != nil {
		t.Fatalf("钩子失败后必须仍能产出兜底输出，而非报错终止: %v", err)
	}
	if !strings.Contains(string(body), "chat.completion") {
		t.Errorf("钩子失败后未退回兜底 OpenAI 形状: %s", body)
	}
}

// TestFallbackSavesUnconfiguredDeployment 直接覆盖用户说的场景：
// 「刚部署没有配置任何适配器」。
func TestFallbackSavesUnconfiguredDeployment(t *testing.T) {
	// 完全不建 VM（模拟适配器目录不存在/为空）
	ser := serializerFor(nil, "whatever")
	body, err := ser.NonStream(&types.UnifiedResponse{Content: "no adapters configured"}, "m", "id", 1)
	if err != nil {
		t.Fatalf("无适配器时兜底必须可用: %v", err)
	}
	if !strings.Contains(string(body), "chat.completion") {
		t.Errorf("无适配器兜底输出形状不对: %s", body)
	}
}

// TestSerializerCacheFollowsAdapterRewrite 验证插件运行时改写适配器后，
// 缓存不会继续持有旧的序列化器判定（否则新声明的钩子永远不生效，而因为
// 判定本身不报错，表现只是「适配器说好的格式没出现」，极难定位）。
func TestSerializerCacheFollowsAdapterRewrite(t *testing.T) {
	vm := luaVMFor(t)
	if _, isFallback := serializerForCached(vm, "openai").(openAISerializer); !isFallback {
		t.Fatal("前提不成立：改写前 openai 应走兜底")
	}
	// 插件改写该适配器，加一个 serialize_response
	src := `local a = {}
a.name = "openai"
a.version = "9"
a.endpoint = "/chat/completions"
function a.serialize_response(p) return '{"rewritten":true}' end
return a`
	if err := vm.LoadAdapterSource("openai", src); err != nil {
		t.Fatalf("rewrite: %v", err)
	}
	// 无需任何显式清缓存动作：修订号变化即让旧条目失效。
	after := serializerForCached(vm, "openai")
	if _, stillFallback := after.(openAISerializer); stillFallback {
		t.Error("缓存未随适配器改写失效：已声明 serialize_response 却仍走兜底")
	}
	body, err := after.NonStream(&types.UnifiedResponse{Content: "x"}, "m", "id", 1)
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(string(body), "rewritten") {
		t.Errorf("改写后的序列化器未生效: %s", body)
	}
}

// TestSerializerCacheIsPerVM 是那个 bug 的回归测试：缓存曾经只按适配器名做键，
// 于是同一进程内两个 VM 共享同一条判定——一个适配器被重写后，另一个 VM 的
// 缓存会说「没声明钩子」，新钩子永远不生效。
func TestSerializerCacheIsPerVM(t *testing.T) {
	vm1 := luaVMFor(t)
	vm2 := luaVMFor(t)
	// 两个 VM 各自独立：改写 vm1 不得影响 vm2 的判定
	if err := vm1.LoadAdapterSource("openai", `local a={} a.name="openai" a.version="9"
a.endpoint="/x" function a.serialize_response(p) return '{"only_vm1":true}' end return a`); err != nil {
		t.Fatal(err)
	}
	if _, f := serializerForCached(vm1, "openai").(openAISerializer); f {
		t.Error("vm1 改写后应改走适配器序列化")
	}
	if _, f := serializerForCached(vm2, "openai").(openAISerializer); !f {
		t.Error("vm2 未被改写，却因共享缓存而丢失了兜底判定（跨 VM 串扰）")
	}
}

// ---- 并发安全 ----

// TestSerializerForConcurrent 是压力测试：多 goroutine 并发取序列化器
// （缓存是 sync.Map，HasFn 会编译脚本 —— 这里验证并发下不 panic、不返回
// 半初始化对象，且每次结果自洽）。
func TestSerializerForConcurrent(t *testing.T) {
	vm := luaVMFor(t)
	const workers, iters = 16, 40
	var wg sync.WaitGroup
	errs := make(chan string, workers*iters)
	for w := 0; w < workers; w++ {
		wg.Add(1)
		go func(w int) {
			defer wg.Done()
			for i := 0; i < iters; i++ {
				name := "openai"
				if w%2 == 1 {
					name = "anthropic"
				}
				ser := serializerForCached(vm, name)
				if ser == nil {
					errs <- "nil serializer"
					return
				}
				// 结果必须自洽：同一适配器每次拿到同一类型
				_, isFallback := ser.(openAISerializer)
				if w%2 == 0 && !isFallback {
					errs <- "openai 未走兜底"
					return
				}
				body, err := ser.NonStream(&types.UnifiedResponse{Content: "x"}, "m", "id", 1)
				if err != nil {
					errs <- "serialize: " + err.Error()
					return
				}
				if !strings.Contains(string(body), "chat.completion") {
					errs <- "并发下输出形状不一致"
					return
				}
			}
		}(w)
	}
	wg.Wait()
	close(errs)
	for e := range errs {
		t.Error(e)
	}
}