ModelRouter/internal/gateway/serializer.go

// Package gateway — 对外序列化（Phase 2）
//
// 职责边界（feature/unified-wire-v2）：
//
//	内核不硬编码任何对外协议。序列化由 Serializer 承载，默认实现是
//	openAISerializer —— 但它只是**最后容错**，不是默认业务路径。
//
// 优先级契约（用户明确要求：任何适配器与兜底 OpenAI 格式冲突时，
// 兜底 OpenAI 直接让步）：
//
//	适配器声明 serialize_response / serialize_chunk
//	  ↓ 未声明
//	openAISerializer（内核兜底，保证刚部署未配适配器时仍能服务）
//
// 「让步」的具体含义：一旦适配器声明了对外序列化钩子，内核在**该源**
// 上完全采用适配器的输出，绝不与兜底格式合并/兜底字段。A 请求经 Anthropic
// 序列化后，绝不会被再补上 OpenAI 的字段。这条规则由
// TestAdapterSerializerWinsOverFallback 锁定。
//
// 安全网：openAISerializer 的输出与改造前内核手写的 OpenAI wire 字节级
// 一致（e2e/.snapshots golden 锁定），所以「适配器未声明」这条路径上的
// 对外行为与本次改造之前完全相同。
package gateway

import (
	"encoding/json"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

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

// StreamState 是一个流会话在序列化侧的累积状态（角色前导是否已发、已见
// 的终止原因、累积 usage/cost）。跨 chunk 传递，避免每帧重新统计。
type StreamState struct {
	LastFinish string
	Usage      *types.TokenUsage
	Cost       string
}

// Serializer 把统一内部格式序列化成对外协议的响应体。
//
// 非流式与流式的终止语义对不同协议差异很大（OpenAI 在最后一个 chunk 上
// 带 finish_reason 后跟一个 usage-only 帧；Anthropic 用独立的
// message_delta 事件；Gemini 用空 candidates 帧），所以「流如何结束」
// 是接口的一部分，而不是调用方的 if。
type Serializer interface {
	// NonStream 把一次完成的统一响应序列化成对外 body。
	NonStream(resp *types.UnifiedResponse, modelName string, id string, created int64) ([]byte, error)

	// OpenRolePreamble 返回流开始时的角色前导帧（可能为空表示该协议不需要）。
	// 返回空切片时调用方不写帧。
	OpenRolePreamble(modelName, id string, created int64, st *StreamState) ([]byte, error)

	// Chunk 把一个统一分块序列化成对外 SSE 帧（不含 "data: " 前缀与空行）。
	Chunk(ck *types.UnifiedChunk, modelName, id string, created int64, st *StreamState) ([]byte, error)

	// Close 返回流结束时的收尾帧序列（可多帧），随后调用方写协议自己的
	// 流终止标记（OpenAI 是 [DONE]，Anthropic 无）。
	Close(modelName, id string, created int64, st *StreamState) ([][]byte, error)

	// DoneMarker 返回该协议的流终止标记；OpenAI 是 "[DONE]"，无此概念的
	// 协议返回空串（调用方因此不写那一行）。
	DoneMarker() string
}

// ---- 兜底实现：OpenAI wire ----

// openAISerializer 是内核的**容错兜底**，输出与改造前手写的 OpenAI wire
// 字节级一致（golden 锁定）。它不是默认业务路径——适配器一旦声明序列化
// 钩子就完全取代它。
type openAISerializer struct{}

// OpenAISerializer 返回兜底 OpenAI 序列化器。暴露给同包测试与
// 「未配置任何适配器」的启动兜底路径使用。
func OpenAISerializer() Serializer { return openAISerializer{} }

func (openAISerializer) NonStream(resp *types.UnifiedResponse, modelName, id string, created int64) ([]byte, error) {
	msg := RespMessage{Role: "assistant", Content: resp.Content}
	if resp.ReasoningContent != "" {
		msg.ReasoningContent = resp.ReasoningContent
	}
	if len(resp.ToolCalls) > 0 {
		msg.ToolCalls = toolCallsWire(resp.ToolCalls)
	}
	out := ChatCompletion{
		ID:      id,
		Object:  "chat.completion",
		Created: created,
		Model:   modelName,
		Choices: []ChatChoice{{Index: 0, Message: msg, FinishReason: resp.FinishReason}},
		Cost:    resp.Cost,
	}
	if resp.TokenUsage.Total > 0 || resp.TokenUsage.Prompt > 0 || resp.TokenUsage.Completion > 0 {
		out.Usage = &resp.TokenUsage
	}
	return json.Marshal(out)
}

func (openAISerializer) OpenRolePreamble(modelName, id string, created int64, st *StreamState) ([]byte, error) {
	return json.Marshal(ChatChunk{
		ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
		Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{Role: "assistant"}}},
	})
}

func (openAISerializer) Chunk(ck *types.UnifiedChunk, modelName, id string, created int64, st *StreamState) ([]byte, error) {
	delta := RespMessage{Role: "assistant", Content: ck.Content}
	if ck.ReasoningContent != "" {
		delta.ReasoningContent = ck.ReasoningContent
	}
	if len(ck.ToolCalls) > 0 {
		delta.ToolCalls = ck.ToolCalls
	}
	choice := ChunkChoice{Index: 0, Delta: delta}
	if ck.Done {
		fin := ck.FinishReason
		if fin == "" {
			fin = "stop"
		}
		st.LastFinish = fin
		choice.FinishReason = &fin
	}
	return json.Marshal(ChatChunk{
		ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
		Choices: []ChunkChoice{choice},
	})
}

func (openAISerializer) Close(modelName, id string, created int64, st *StreamState) ([][]byte, error) {
	fin := st.LastFinish
	if fin == "" {
		fin = "stop"
	}
	term, err := json.Marshal(ChatChunk{
		ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
		Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &fin}},
	})
	if err != nil {
		return nil, err
	}
	out := [][]byte{term}
	// OpenAI 的 usage 帧：空 choices + usage + 可选 cost。
	if st.Usage != nil || st.Cost != "" {
		var tut *types.TokenUsage
		if st.Usage != nil {
			tut = st.Usage
		} else {
			tut = &types.TokenUsage{}
		}
		u, err := json.Marshal(ChatChunk{
			ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
			Choices: []ChunkChoice{}, Usage: tut, Cost: st.Cost,
		})
		if err != nil {
			return nil, err
		}
		out = append(out, u)
	}
	return out, nil
}

func (openAISerializer) DoneMarker() string { return "[DONE]" }

// ---- 适配器钩子通道 ----

// luaSerializer 把对外序列化委托给适配器的 serialize_response /
// serialize_chunk 钩子。适配器声明了钩子就**完全取代**兜底实现。
type luaSerializer struct {
	vm      *lua.VM
	adapter string
	// responseOK / chunkOK 在构造时一次性探测（HasFn 不占用 worker）。
	responseOK bool
	chunkOK    bool
	// openPreambleOK / closeOK 是可选钩子；未声明时借用兜底实现的对应行为，
	// 这样「只实现 serialize_chunk」的适配器也能立刻可用。
	openPreambleOK bool
	closeOK        bool
}

// serializerFor 为某个适配器名挑选序列化器。
//
// 这是「适配器优先、兜底让步」的唯一决策点：适配器声明了对外序列化
// 钩子就返回 luaSerializer（其输出不再与 OpenAI 兜底混合），否则返回
// 兜底 OpenAI。
func serializerFor(vm *lua.VM, adapter string) Serializer {
	if vm == nil {
		return OpenAISerializer()
	}
	resp := vm.HasFn(adapter, "serialize_response")
	chunk := vm.HasFn(adapter, "serialize_chunk")
	if !resp && !chunk {
		return OpenAISerializer()
	}
	return &luaSerializer{
		vm:             vm,
		adapter:        adapter,
		responseOK:     resp,
		chunkOK:        chunk,
		openPreambleOK: vm.HasFn(adapter, "serialize_open_preamble"),
		closeOK:        vm.HasFn(adapter, "serialize_close"),
	}
}

func (s *luaSerializer) NonStream(resp *types.UnifiedResponse, modelName, id string, created int64) ([]byte, error) {
	if !s.responseOK {
		// 该适配器只实现了流式钩子：非流式沿用兜底。这是唯一允许兜底
		// 参与的场景——适配器**没有**就该方向的实现，而不是「冲突」。
		return openAISerializer{}.NonStream(resp, modelName, id, created)
	}
	payload := map[string]interface{}{
		"model":         modelName,
		"id":            id,
		"created":       created,
		"content":       resp.Content,
		"reasoning":     resp.ReasoningContent,
		"finish_reason": resp.FinishReason,
		"token_usage":   resp.TokenUsage,
		"cost":          resp.Cost,
	}
	if len(resp.ToolCalls) > 0 {
		payload["tool_calls"] = resp.ToolCalls
	}
	out, err := s.vm.Transform(s.adapter, "serialize_response", wireJSON(payload))
	if err != nil || out == "" {
		// 钩子失败：明确退回兜底，而不是产出一半适配器字段 + 一半 OpenAI
		// 字段的混合体（那正是「让步」规则要避免的）。
		return openAISerializer{}.NonStream(resp, modelName, id, created)
	}
	return []byte(out), nil
}

func (s *luaSerializer) OpenRolePreamble(modelName, id string, created int64, st *StreamState) ([]byte, error) {
	if !s.openPreambleOK {
		return openAISerializer{}.OpenRolePreamble(modelName, id, created, st)
	}
	out, err := s.vm.Transform(s.adapter, "serialize_open_preamble",
		wireJSON(map[string]interface{}{"model": modelName, "id": id, "created": created}))
	if err != nil {
		return openAISerializer{}.OpenRolePreamble(modelName, id, created, st)
	}
	if out == "" {
		return nil, nil // 该协议不需要角色前导
	}
	return []byte(out), nil
}

func (s *luaSerializer) Chunk(ck *types.UnifiedChunk, modelName, id string, created int64, st *StreamState) ([]byte, error) {
	if !s.chunkOK {
		return openAISerializer{}.Chunk(ck, modelName, id, created, st)
	}
	payload := map[string]interface{}{
		"model":         modelName,
		"id":            id,
		"created":       created,
		"content":       ck.Content,
		"reasoning":     ck.ReasoningContent,
		"done":          ck.Done,
		"finish_reason": ck.FinishReason,
	}
	if len(ck.ToolCalls) > 0 {
		payload["tool_calls"] = ck.ToolCalls
	}
	if ck.Usage != nil {
		payload["usage"] = ck.Usage
	}
	if ck.Cost != "" {
		payload["cost"] = ck.Cost
	}
	out, err := s.vm.Transform(s.adapter, "serialize_chunk", wireJSON(payload))
	if err != nil {
		return openAISerializer{}.Chunk(ck, modelName, id, created, st)
	}
	if out == "" {
		return nil, nil
	}
	// 适配器的终止帧要驱动审计记录与收尾逻辑，所以把终止原因抄回状态。
	if ck.Done && st.LastFinish == "" {
		if fin := ck.FinishReason; fin != "" {
			st.LastFinish = fin
		}
	}
	return []byte(out), nil
}

func (s *luaSerializer) Close(modelName, id string, created int64, st *StreamState) ([][]byte, error) {
	if !s.closeOK {
		return openAISerializer{}.Close(modelName, id, created, st)
	}
	out, err := s.vm.Transform(s.adapter, "serialize_close",
		wireJSON(map[string]interface{}{
			"model": modelName, "id": id, "created": created,
			"finish_reason": st.LastFinish, "cost": st.Cost,
			"has_usage": st.Usage != nil,
		}))
	if err != nil {
		return openAISerializer{}.Close(modelName, id, created, st)
	}
	if out == "" {
		return nil, nil
	}
	// 适配器可以返回多帧（换行分隔），因为 Anthropic 的收尾本身就是
	// message_stop + 若干 delta；单帧限制会逼适配器把多事件塞进一个 JSON。
	var frames [][]byte
	for _, line := range splitLines(out) {
		if line == "" {
			continue
		}
		frames = append(frames, []byte(line))
	}
	return frames, nil
}

// DoneMarker 返回协议自己的流终止标记。适配器可用 serialize_done_marker
// 声明；未声明时若有自定义流式序列化就返回空串（表示「没有 [DONE] 这种
// 概念」），否则沿用 OpenAI 的 [DONE]。
func (s *luaSerializer) DoneMarker() string {
	if s.vm.HasFn(s.adapter, "serialize_done_marker") {
		out, err := s.vm.Transform(s.adapter, "serialize_done_marker", "{}")
		if err == nil {
			return out
		}
	}
	if s.chunkOK {
		// 适配器接管了流式序列化：只有它没显式声明标记时才不发 [DONE]，
		// 因为擅自给 Anthropic 流加一行 "[DONE]" 会让严格客户端报错。
		return ""
	}
	return "[DONE]"
}

// splitLines 切分多行帧输出。
func splitLines(s string) []string {
	var out []string
	start := 0
	for i := 0; i < len(s); i++ {
		if s[i] == '\n' {
			out = append(out, trimCR(s[start:i]))
			start = i + 1
		}
	}
	if start < len(s) {
		out = append(out, trimCR(s[start:]))
	}
	return out
}

func trimCR(s string) string {
	if n := len(s); n > 0 && s[n-1] == '\r' {
		return s[:n-1]
	}
	return s
}

// ---- 共享辅助 ----

var wireSeq int64

func newWireID() string {
	n := atomic.AddInt64(&wireSeq, 1)
	return fmt.Sprintf("chatcmpl-%d", n)
}

func newCreated() int64 { return time.Now().Unix() }

// wireJSON 是本文件私有的 JSON 编码助手（包内已有 discover.go 的 mustJSON，
// 那个返回 []byte 且忽略错误，语义不同，故不共用）。
func wireJSON(v interface{}) string {
	b, err := json.Marshal(v)
	if err != nil {
		return "{}"
	}
	return string(b)
}

// serializerCache 按 (VM, 适配器名) 缓存 Serializer。
//
// VM 必须进键里：缓存曾经只按适配器名做键，结果一个进程内两个 VM（或
// 热重载后的新旧 VM）会共享同一条判定。插件在运行时改写适配器后声明了
// serialize_response，另一条缓存记录仍会说「没声明」，于是新钩子永远不
// 生效——而且因为判定本身不会报错，表现只是「适配器说好的对外格式没
// 出现」，极难定位。
//
// HasFn 的代价决定了它必须被缓存：它会**编译一次适配器脚本**来回答
// 「这个适配器声明了哪些函数」，放在每个请求上等于把加载成本翻倍。
//
// 失效由 vmGeneration 负责：适配器重写时 VM 的代号会变（见 adapterEpoch），
// 于是旧条目再也匹配不上，无需显式清扫。
var serializerCache sync.Map // cacheKey -> Serializer

// cacheKey 是 (VM 指针, VM 代号, 适配器名) 的复合键。
type cacheKey struct {
	vm   *lua.VM
	gen  uint64
	name string
}

func serializerForCached(vm *lua.VM, adapter string) Serializer {
	if vm == nil {
		return OpenAISerializer()
	}
	k := cacheKey{vm: vm, gen: adapterEpoch(vm, adapter), name: adapter}
	if v, ok := serializerCache.Load(k); ok {
		return v.(Serializer)
	}
	s := serializerFor(vm, adapter)
	actual, _ := serializerCache.LoadOrStore(k, s)
	return actual.(Serializer)
}

// adapterEpoch 返回该适配器「当前这份代码」的代号。适配器被重写后代号变化，
// 使旧缓存条目自动失效（见 serializerCache 的注释）。
func adapterEpoch(vm *lua.VM, adapter string) uint64 {
	return vm.AdapterRevision(adapter)
}

// serializerForSource 解析「实际服务该请求的那个源」该用哪个序列化器。
//
// 必须按**服务者**选，而不是按调用方请求的首选源选：direct 请求在候选源
// 之间故障转移时，usedSrc 可能与请求 model 所属的源完全不同；按首选源选
// 会让故障转移后的响应带着另一个源的对外格式。
//
// 源名查不到时（配置在选择之后被改、或插件装的源尚未注册）退回兜底
// OpenAI：那是「刚部署 / 异常态下仍能服务」的保证，而不是静默产出错误格式。
func serializerForSource(vm *lua.VM, sources []SourceRef, sourceName string) Serializer {
	for _, s := range sources {
		if s.Name == sourceName {
			return serializerForCached(vm, s.Adapter)
		}
	}
	return OpenAISerializer()
}

// SourceRef 是解析序列化器所需的最小源信息（名字 + 适配器）。
type SourceRef struct {
	Name    string
	Adapter string
}

// serializerFor 是 Gateway 侧的解析入口：按**实际服务该请求的源**挑序列化器。
//
// cands 传候选列表时用它做故障转移后的解析（服务源可能不是首选）；传 nil
// 表示 AUTO 路径——AUTO 的服务源由链决定，不在调用方的候选集里，退回按
// 当前配置里全部源查找。
func (g *Gateway) serializerFor(sourceName string, cands []*provider.Provider) Serializer {
	vm := g.core.VM()
	if vm == nil {
		return OpenAISerializer()
	}
	if cands != nil {
		refs := make([]SourceRef, 0, len(cands))
		for _, p := range cands {
			refs = append(refs, SourceRef{Name: p.Name(), Adapter: p.Adapter()})
		}
		return serializerForSource(vm, refs, sourceName)
	}
	refs := make([]SourceRef, 0, 16)
	for _, s := range g.core.Sources() {
		refs = append(refs, SourceRef{Name: s.Name, Adapter: s.Adapter})
	}
	return serializerForSource(vm, refs, sourceName)
}