diff --git a/internal/gateway/chat.go b/internal/gateway/chat.go index de7631c..08432f6 100644 --- a/internal/gateway/chat.go +++ b/internal/gateway/chat.go @@ -817,29 +817,21 @@ func recordChatUsage(rec *Req, req *types.ChatRequest, resp *types.UnifiedRespon } } -// writeChatCompletion renders a unified response as an OpenAI -// chat.completion object. modelName is the id clients see as the serving -// model: the requested id for direct routes, the exact slot model for AUTO. -func writeChatCompletion(w http.ResponseWriter, resp *types.UnifiedResponse, modelName string) { - msg := RespMessage{Role: "assistant", Content: resp.Content} - if resp.ReasoningContent != "" { - msg.ReasoningContent = resp.ReasoningContent +// writeChatCompletion renders a unified response through the out-bound +// Serializer. The default serializer reproduces the OpenAI wire shape this +// function used to build inline (locked byte-for-byte by e2e golden); an +// adapter that declares serialize_response overrides it. modelName is the id +// clients see as the serving model: the requested id for direct routes, the +// exact slot model for AUTO. +func writeChatCompletion(w http.ResponseWriter, s Serializer, resp *types.UnifiedResponse, modelName string) { + body, err := s.NonStream(resp, modelName, newWireID(), newCreated()) + if err != nil { + writeError(w, http.StatusInternalServerError, "serialize_error", "failed to serialize response") + return } - if len(resp.ToolCalls) > 0 { - msg.ToolCalls = toolCallsWire(resp.ToolCalls) - } - out := ChatCompletion{ - ID: newID(), - Object: "chat.completion", - Created: time.Now().Unix(), - 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 - } - writeJSON(w, http.StatusOK, out) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(body) } // singleChat runs a direct (model-pinned) non-streaming request across the @@ -864,7 +856,7 @@ func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands [ // the total latency. rec.FirstByteMs = rec.LatMs g.writeRec(rec) - writeChatCompletion(w, resp, effective) + writeChatCompletion(w, g.serializerFor(usedSrc, cands), resp, effective) } // writeRec records a finished request (audit + aggregates) and fires the @@ -1116,24 +1108,29 @@ func mergeUsage(prev, cur *types.TokenUsage) *types.TokenUsage { return &out } -// pumpStream writes the full SSE sequence for a started stream: role -// preamble, one chunk per unified delta, the terminating finish_reason, the -// OpenAI-standard final usage chunk (empty choices) and [DONE]. modelName -// follows writeChatCompletion's rule (requested id for direct routes, exact -// slot model for AUTO). -func (g *Gateway) pumpStream(w http.ResponseWriter, rec *Req, chunks <-chan types.UnifiedChunk, modelName string, t0 time.Time) { +// pumpStream writes the full SSE sequence for a started stream. Every byte on +// the wire comes from the Serializer, so the OpenAI frame shapes live in +// serializer.go rather than here, and an adapter that declares serialize_chunk +// replaces them wholesale. +// +// This function owns only transport concerns — headers, the flush cadence, +// first-byte timing, and audit accounting — plus the protocol-agnostic rule +// that usage seen anywhere in the stream becomes the recorded usage. Whether a +// trailing usage frame is emitted, and in what shape, is the serializer's +// decision. +func (g *Gateway) pumpStream(w http.ResponseWriter, ser Serializer, rec *Req, chunks <-chan types.UnifiedChunk, modelName string, t0 time.Time) { w.Header().Set("Content-Type", "text/event-stream") w.Header().Set("Cache-Control", "no-cache") w.Header().Set("Connection", "keep-alive") w.WriteHeader(http.StatusOK) flusher, _ := w.(http.Flusher) - id := newID() - created := time.Now().Unix() - send := func(obj interface{}) bool { - b, err := json.Marshal(obj) - if err != nil { - return false + id := newWireID() + created := newCreated() + st := &StreamState{} + send := func(b []byte) bool { + if len(b) == 0 { + return true } if _, err := fmt.Fprintf(w, "data: %s\n\n", b); err != nil { return false @@ -1144,102 +1141,66 @@ func (g *Gateway) pumpStream(w http.ResponseWriter, rec *Req, chunks <-chan type return true } - if !send(ChatChunk{ - ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName, - Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{Role: "assistant"}}}, - }) { - return + if b, err := ser.OpenRolePreamble(modelName, id, created, st); err == nil { + if !send(b) { + return + } } // First SSE byte sent to the client: record time-to-first-byte for the // source's status-page latency average. rec.FirstByteMs = time.Since(t0).Milliseconds() - var lastUsage *types.TokenUsage - var lastCost string - lastFinish := "" + for ck := range chunks { if ck.Usage != nil { - lastUsage = mergeUsage(lastUsage, ck.Usage) + st.Usage = mergeUsage(st.Usage, ck.Usage) } if ck.Cost != "" { - lastCost = ck.Cost + st.Cost = ck.Cost } - chunk := ChatChunk{ - ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName, + b, err := ser.Chunk(&ck, modelName, id, created, st) + if err != nil { + return } - 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" - } - lastFinish = fin - choice.FinishReason = &fin - } - chunk.Choices = []ChunkChoice{choice} rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)+len(ck.ToolCalls)) / 3 - if !send(chunk) { + if !send(b) { return } } - finalFinish := lastFinish - if finalFinish == "" { - finalFinish = "stop" - } - send(ChatChunk{ - ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName, - Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &finalFinish}}, - }) - // Final usage chunk (OpenAI standard: empty choices + usage before [DONE]). - // Prefer the upstream's exact usage if the stream carried it; fall back to - // the gateway's estimate otherwise. - // Write the upstream's exact numbers back onto the audit record. Without - // this the streamed path kept only the per-chunk byte estimate, so the same + + // Write the upstream's exact usage back onto the audit record. Without this + // the streamed path kept only the per-chunk byte estimate, so the same // request recorded ~1.4x its real completion tokens (measured: upstream 100, // audit 145) while the non-streaming path recorded 100. Two paths, two // different numbers for one request is a reporting bug, not a rounding one. - if lastUsage != nil { - if lastUsage.Prompt > 0 { - rec.Prompt = int64(lastUsage.Prompt) + if st.Usage != nil { + if st.Usage.Prompt > 0 { + rec.Prompt = int64(st.Usage.Prompt) } - if lastUsage.Completion > 0 { - rec.Compl = int64(lastUsage.Completion) + if st.Usage.Completion > 0 { + rec.Compl = int64(st.Usage.Completion) } - } - var tut *types.TokenUsage - if lastUsage != nil { - tut = lastUsage - applyStreamUsageDetails(rec, lastUsage) + applyStreamUsageDetails(rec, st.Usage) } else if rec.Prompt+rec.Compl > 0 { - u := types.TokenUsage{ + // No upstream usage on the stream: the recorded per-chunk byte + // estimate is the best number available, so hand it to the serializer + // as the usage rather than dropping it. + st.Usage = &types.TokenUsage{ Prompt: int(rec.Prompt), Completion: int(rec.Compl), Total: int(rec.Prompt + rec.Compl), } - tut = &u } - if tut != nil { - send(ChatChunk{ - ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName, - Choices: []ChunkChoice{}, - Usage: tut, - Cost: lastCost, - }) - } else if lastCost != "" { - // Cost arrived without any usage (OpenCode sends it on its own frame). - send(ChatChunk{ - ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName, - Choices: []ChunkChoice{}, - Cost: lastCost, - }) + + if frames, err := ser.Close(modelName, id, created, st); err == nil { + for _, b := range frames { + if !send(b) { + return + } + } + } + if marker := ser.DoneMarker(); marker != "" { + fmt.Fprintf(w, "data: %s\n\n", marker) } - fmt.Fprintf(w, "data: [DONE]\n\n") if flusher != nil { flusher.Flush() } @@ -1270,7 +1231,7 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [ rec.Source = usedSrc g.fireRouted(ctx, "stream", usedSrc, rec.Model, -1, true) rec.Prompt = estimatePromptTokens(req) - g.pumpStream(w, rec, chunks, effective, t0) + g.pumpStream(w, g.serializerFor(usedSrc, cands), rec, chunks, effective, t0) } // singleChatAuto runs a non-streaming AUTO request down the chain (see @@ -1303,7 +1264,7 @@ func (g *Gateway) singleChatAuto(w http.ResponseWriter, ctx context.Context, cha g.fireRouted(ctx, "chat", usedSrc, usedModel, -2, false) rec.FirstByteMs = rec.LatMs g.writeRec(rec) - writeChatCompletion(w, resp, usedModel) + writeChatCompletion(w, g.serializerFor(usedSrc, nil), resp, usedModel) } // streamChatAuto streams an AUTO request down the chain. A slot is abandoned @@ -1333,7 +1294,7 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha rec.Source = usedSrc g.fireRouted(ctx, "stream", usedSrc, rec.Model, -2, true) rec.Prompt = estimatePromptTokens(req) - g.pumpStream(w, rec, chunks, usedModel, t0) + g.pumpStream(w, g.serializerFor(usedSrc, nil), rec, chunks, usedModel, t0) } func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) { diff --git a/internal/gateway/serializer.go b/internal/gateway/serializer.go new file mode 100644 index 0000000..c36212f --- /dev/null +++ b/internal/gateway/serializer.go @@ -0,0 +1,466 @@ +// 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) +} diff --git a/internal/gateway/serializer_test.go b/internal/gateway/serializer_test.go new file mode 100644 index 0000000..6b4403b --- /dev/null +++ b/internal/gateway/serializer_test.go @@ -0,0 +1,349 @@ +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) + } +} diff --git a/internal/lua/vm.go b/internal/lua/vm.go index d8c9ac4..fdc56d6 100644 --- a/internal/lua/vm.go +++ b/internal/lua/vm.go @@ -114,6 +114,12 @@ type adapterPool struct { lastGrow time.Time idleRounds int // consecutive janitor rounds that saw reclaimable slack peakInUse int // high-water mark of inUse, for observability + // revision counts how many times this adapter's SOURCE was replaced (or the + // adapter was removed and re-added). Callers that cache a decision derived + // from the script's declared surface read it via AdapterRevision, so a + // runtime rewrite (a plugin installing a new version of its own adapter) + // cannot leave them acting on a stale answer. + revision uint64 // pluginMode makes boot() store the returned table under pluginGlobal // instead of adapterGlobal. Everything else (elastic sizing, reclaim) is // identical, which is why plugins reuse this pool rather than getting a @@ -174,7 +180,7 @@ type PoolStats struct { } func newAdapterPool(name, script string, static staticInfo) *adapterPool { - p := &adapterPool{name: name, script: script, static: static, maxW: 1} + p := &adapterPool{name: name, script: script, static: static, maxW: 1, revision: 1} p.cond = sync.NewCond(&p.mu) return p } @@ -750,11 +756,28 @@ func (v *VM) LoadAdapterSource(name, code string) error { v.mu.RUnlock() v.mu.Lock() - if p, ok := v.pools[static.name]; ok { - p.shutdown() - } + old := v.pools[static.name] p := newAdapterPool(static.name, code, *static) p.setMax(maxW) + // Carry the revision forward so the new source keeps counting ABOVE the old + // one. A cache keyed on (vm, revision, name) must see a value different + // from the previous generation's; restarting at 1 could collide with a + // value an earlier generation already published, which would resurrect a + // stale cache entry — exactly the bug this revision exists to prevent. + // + // The old pool is read while v.mu is HELD (a nested RLock here would + // self-deadlock: sync.RWMutex is not reentrant) and while its own mu is + // free, because shutdown() — which takes p.mu — happens after v.mu is + // released below. No other path can mutate p.revision while we hold v.mu, + // because every writer reaches the pool through v.mu first. + if old != nil { + old.mu.Lock() + p.revision = old.revision + 1 + old.mu.Unlock() + } + if old != nil { + old.shutdown() + } v.pools[static.name] = p v.mu.Unlock() v.startJanitor() @@ -950,6 +973,46 @@ func (v *VM) StreamEndpoint(name string) string { return p.static.endpointStream } +// HasFn reports whether the named adapter declares a callable field fn, +// WITHOUT acquiring a worker and WITHOUT running the script. +// +// It exists for capability probing that runs on the request path: the +// out-bound serializer asks "does this adapter declare serialize_chunk?" once +// per adapter, not once per chunk. Probing by calling Transform would be +// wrong twice over — it costs a pool checkout, and a MISSING hook comes back +// as an error string rather than a boolean. +// +// The probe runs on a throwaway state built from the SAME script text, so the +// answer reflects what the script declares. A script whose table is built +// conditionally on its environment could differ per state; that is already +// true of every other staticInfo field (endpoint, headers), so the pool does +// not get a special exemption here. +func (v *VM) HasFn(name, fn string) bool { + p := v.pool(name) + if p == nil { + return false + } + p.mu.Lock() + script := p.script + p.mu.Unlock() + if script == "" { + return false + } + L := golua.NewState() + defer L.Close() + L.OpenLibs() + setupGlobals(L) + if err := L.DoString(script); err != nil { + return false + } + if L.Type(-1) != golua.LUA_TTABLE { + return false + } + L.GetField(-1, fn) + defer L.Pop(1) + return !L.IsNil(-1) && L.IsFunction(-1) +} + // ---- static inspection (compile-once at load time) ---- func inspectScript(name, code string) (*staticInfo, error) { @@ -1295,3 +1358,25 @@ func jsonDecode(s string) (interface{}, error) { } return v, nil } + +// AdapterRevision returns a monotonically increasing revision for the named +// adapter's CURRENT source. Every LoadAdapterSource / RemoveAdapter bumps it. +// +// Purpose: callers that cache a decision derived from an adapter's declared +// surface (e.g. "does this adapter declare serialize_chunk?") store the +// revision alongside the cached value and treat a mismatch as invalid. That +// removes the need for those callers to be told when to invalidate — which +// they previously were not, so a plugin rewriting an adapter at runtime could +// leave a stale decision in place forever. +// +// The revision is per-adapter rather than global so an unrelated adapter's +// reload does not discard a still-valid cache entry. +func (v *VM) AdapterRevision(name string) uint64 { + p := v.pool(name) + if p == nil { + return 0 + } + p.mu.Lock() + defer p.mu.Unlock() + return p.revision +}