// sse.go 处理上游 SSE 流:聚合成单个 OpenAI 响应,或透传给客户端。 package upstream import ( "bufio" "encoding/json" "errors" "fmt" "io" "net/http" "sort" "strings" "time" ) // errEmptyStream 上游返回 200 但没有有效 SSE 数据帧(空流/只有注释/[DONE])。 // 用哨兵错误替代裸 fmt.Errorf:StreamHint 的调用方(handler 流式路径)需要区分 // 「上游空流」与「客户端断连写失败」——空流是上游缺陷,应记 502 观测;写失败是 // 客户端已走,日志口径不同。Aggregate 与 StreamHint 共用同一哨兵(errors.Is 判定)。 var errEmptyStream = errors.New("upstream stream contained no valid data events") // IsEmptyStreamError 报告错误是否为「上游空流」(无有效 SSE 帧)——供 handler // 在流式路径把空流记为失败观测(HTTP 头已发出只能 200,但日志/状态应收敛到 // upstream_parse 同语义),与客户端断连类错误区分。 func IsEmptyStreamError(err error) bool { return errors.Is(err, errEmptyStream) } // Aggregate 读取完整 SSE 流,聚合 delta.content 为单个 OpenAI chat.completion 响应。 // 分片/半行由 bufio.Reader.ReadString 处理;遇到 "data: [DONE]" 结束。 // tool_calls 以流式 delta 到达(按 index 合并:首片带 id/type/name,后续只带 arguments 片段)。 func Aggregate(r io.Reader) (map[string]any, error) { br := bufio.NewReaderSize(r, 64*1024) var ( id, model string created float64 content strings.Builder reasoning strings.Builder role = "assistant" finishReason = "stop" usage map[string]any gotAnyContent bool validEvents int sawDone bool // 上游显式发过 data: [DONE](正常收尾) toolCalls = map[int]map[string]any{} toolOrder []int // toolSeq 缺 index 的 tool_call 的分配序号源:跨帧延续「最近分配」槽位, // 同帧内递增(见 mergeToolCallsChunk 注释)。 toolSeq int // idIndex id → 已分配的 index:跨帧持续(延续碎片的 id 常出现在后续帧), // 供缺 index 时按 id 归位既有调用(见 mergeToolCallsChunk 注释)。 idIndex = map[string]int{} ) // appendContent 是「已取到正文」(gotAnyContent latch)的唯一写入点:delta 与 // message 两路 content 都必须经此并入,规约只有一份(issue #142)—— // S1 空串不算「已取到正文」、不占 latch 名额:OpenAI 风格 role-only 首帧 // (delta.content="")和整条空 message 帧是常态帧,若空串置位 latch, // 后续真正文会被 message 回退分支的 !gotAnyContent 守卫静默拒绝; // S2 空串本身也无可追加字节,跳过 WriteString 与追加语义自洽。 appendContent := func(txt string) { if txt == "" { return } content.WriteString(txt) gotAnyContent = true } // mergeToolCallsChunk 把一段 tool_calls 数组按 index 合并进累计表。 // delta(流式分片,按 index 累积)与 message(非 delta 整条)共用同一合并逻辑, // 保证「上游给的身份/函数名不丢、arguments 拼接语义一致」。 // // index 缺失兼容:OpenAI 规范要求 delta 帧的 tool_call 带 index(标记分片归属), // 但部分上游省略它。此前缺 index 一律归 0——多调用场景下不同 call 被合并进同一 // index 槽,arguments 串联、name 互相覆盖(数据污染)。修法按「id 优先、lastIdx 兜底」: // - 带 index → 按 index 累积(合规形态,零改动); // - 缺 index 带 id 且 id 已见过 → 延续该 id 所在 index; // - 缺 index 带 id 且 id 是新的 → 开新序号(多调用不合并); // - 缺 index 无 id → 追加到最近收到碎片的 index(单个调用的延续分片无 id // 是标准形态),无既往则开新号。 // 带 index 的碎片照常按 index 累积,不受影响。 // nextToolIndex 分配下一个不冲突的缺 index 序号:从 toolSeq 起递增跳过既有 // index(合规流的 index 是 0..N-1,缺 index 的补位不能覆盖它们)。 nextToolIndex := func() int { for { idx := toolSeq toolSeq++ if _, used := toolCalls[idx]; !used { return idx } } } mergeToolCallsChunk := func(tcs []any) { for _, tc := range tcs { call, ok := tc.(map[string]any) if !ok { continue } idx := -1 if v, ok := call["index"].(float64); ok { idx = int(v) } else if cid, _ := call["id"].(string); cid != "" { if mid, seen := idIndex[cid]; seen { idx = mid // 该 id 已归位过:延续既有调用(跨帧有效) } else { idx = nextToolIndex() } } else if len(toolOrder) > 0 { idx = toolOrder[len(toolOrder)-1] // 无 id 碎片:延续最近调用 } else { idx = nextToolIndex() } merged, seen := toolCalls[idx] if !seen { merged = map[string]any{"index": idx} toolCalls[idx] = merged toolOrder = append(toolOrder, idx) } if cid, _ := call["id"].(string); cid != "" { idIndex[cid] = idx } if cid, _ := merged["id"].(string); cid != "" { idIndex[cid] = idx } mergeToolCallDelta(merged, call) } } // mergeMessageFields 把非 delta 的完整 message 内容并入聚合(message 是整条下发, // 非流式拼接,content 只取一次)。role/reasoning_content/tool_calls 与 delta 分支 // 同构透出;content 同样置 gotAnyContent,与 delta 路径的 latch 语义一致 // (一帧整条 message 之后,后续 delta 帧不重复追加)。 mergeMessageFields := func(msg map[string]any) { if r2, ok := msg["role"].(string); ok && r2 != "" { role = r2 } if txt, ok := msg["content"].(string); ok { appendContent(txt) } if rc, ok := msg["reasoning_content"].(string); ok { reasoning.WriteString(rc) } if tcs, ok := msg["tool_calls"].([]any); ok { mergeToolCallsChunk(tcs) } } for { line, err := br.ReadString('\n') if err != nil && err != io.EOF { return nil, err } line = strings.TrimRight(line, "\r\n") if strings.HasPrefix(line, "data: ") { payload := strings.TrimPrefix(line, "data: ") if payload == "[DONE]" { // 上游显式结束:停止读取,DONE 之后的任何数据一律忽略。 sawDone = true break } else { var chunk map[string]any if json.Unmarshal([]byte(payload), &chunk) == nil { // 有效事件计数:仅 JSON 解析成功的数据帧计入(解析失败沿用静默 continue)。 validEvents++ if v, ok := chunk["id"].(string); ok && id == "" { id = v } if v, ok := chunk["model"].(string); ok && model == "" { model = v } if v, ok := chunk["created"].(float64); ok && created == 0 { created = v } if u, ok := chunk["usage"].(map[string]any); ok { usage = u } if ch, ok := chunk["choices"].([]any); ok { for _, ci := range ch { c, _ := ci.(map[string]any) if c == nil { continue } if fr, ok := c["finish_reason"].(string); ok && fr != "" { finishReason = fr } if delta, ok := c["delta"].(map[string]any); ok { if r2, ok := delta["role"].(string); ok && r2 != "" { role = r2 } if txt, ok := delta["content"].(string); ok { appendContent(txt) } if rc, ok := delta["reasoning_content"].(string); ok { reasoning.WriteString(rc) } if tcs, ok := delta["tool_calls"].([]any); ok { mergeToolCallsChunk(tcs) } } // 有的上游把完整消息放在 message 里(非 delta): // 整条并入(role/content/reasoning_content/tool_calls), // 与 delta 分支同构。delta 已取过正文(gotAnyContent)则跳过 // (避免与 delta 路径重复拼接——PR #134 的 latch 语义)。 if msg, ok := c["message"].(map[string]any); ok && !gotAnyContent { mergeMessageFields(msg) } } } } } } if err == io.EOF { break } } if validEvents == 0 { // 上游返回 200 但没有任何有效数据事件(空流/只有 [DONE]/只有注释行): // 不再合成空 content 的假成功响应,直接报错,由 handler 映射为 502 upstream_parse。 return nil, errEmptyStream } if id == "" { id = fmt.Sprintf("chatcmpl-%d", time.Now().UnixNano()) } if created == 0 { created = float64(time.Now().Unix()) } message := map[string]any{ "role": role, "content": content.String(), } if reasoning.Len() > 0 { message["reasoning_content"] = reasoning.String() } if len(toolOrder) > 0 { sort.Ints(toolOrder) calls := make([]map[string]any, 0, len(toolOrder)) for _, idx := range toolOrder { calls = append(calls, toolCalls[idx]) } // P1b:流被截断时 tool_call 的 arguments 是残缺 JSON(解析失败),不把脏参数 // 交给客户端——残留分片会被客户端解析成非法 JSON 卡死会话。截断的两个来源: // - finish_reason=="length"(模型因 max_tokens 提前中止); // - 上游连接中断(EOF 收尾但未发 data: [DONE],sawDone=false)。 // 完整参数原样保留(正例零改动);空参数(无参工具)不是截断,同样保留。 // 此前只认 finish_reason=="length",EOF 截断的 tool_calls 残缺参数被原样下发。 if finishReason == "length" || !sawDone { calls = dropTruncatedToolCalls(calls) } if len(calls) > 0 { message["tool_calls"] = calls } } resp := map[string]any{ "id": id, "object": "chat.completion", "created": int64(created), "model": model, "choices": []any{ map[string]any{ "index": 0, "message": message, "finish_reason": finishReason, }, }, } if usage != nil { // OpenAI 非流式 usage 必含 total_tokens。上游若只发 prompt_tokens + // completion_tokens(部分上游末帧缺 total),网关合成补齐——否则严格按 // schema 校验的客户端收不到 total_tokens。已有 total 或二者缺一不补 // (不臆造:单边有值无法合成可信的 total)。 resp["usage"] = normalizeUsageCacheAliases(ensureUsageTotal(usage)) } return resp, nil } // ensureUsageTotal 在 usage 缺 total_tokens 但 prompt_tokens/completion_tokens 都在时 // 补齐 total = prompt + completion(通过新 map 合并,不修改原上游 map)。 // 任一缺失或已有 total 时原样返回。 func ensureUsageTotal(u map[string]any) map[string]any { if _, ok := u["total_tokens"]; ok { return u } pt, pok := num64(u["prompt_tokens"]) ct, cok := num64(u["completion_tokens"]) if !pok || !cok { return u } out := make(map[string]any, len(u)+1) for k, v := range u { out[k] = v } out["total_tokens"] = pt + ct return out } // num64 把 JSON number(float64/int64 均可)归一为 float64;非数字返回 ok=false。 // // 注意与 payload.go 的类型 switch 语义不同:那边是翻译请求别名字段(非数字拒绝 // 整个请求),这边是聚合响应求和(非数字只跳过合成)。防后人合并两处。 func num64(v any) (float64, bool) { switch n := v.(type) { case float64: return n, true case int64: return float64(n), true case int: return float64(n), true default: return 0, false } } // mergeToolCallDelta 把流式 tool_call 片段合并到累计对象: // id/type/function.name 直覆盖(后续分片通常缺省),function.arguments 拼接。 func mergeToolCallDelta(merged, delta map[string]any) { if v, ok := delta["id"].(string); ok && v != "" { merged["id"] = v } if v, ok := delta["type"].(string); ok && v != "" { merged["type"] = v } df, _ := delta["function"].(map[string]any) if df == nil { return } mf, _ := merged["function"].(map[string]any) if mf == nil { mf = map[string]any{} merged["function"] = mf } if v, ok := df["name"].(string); ok && v != "" { mf["name"] = v } if v, ok := df["arguments"].(string); ok && v != "" { if prev, _ := mf["arguments"].(string); prev != "" { mf["arguments"] = prev + v } else { mf["arguments"] = v } } } // stripToolCallNames 收敛流式 tool_calls 的 name 语义为「每个 index 只出现一次」: // 首片保留 function.name,同一 index 后续分片里的 name 键一律删除(无论上游是 // 空串还是重复非空串)。这是 OpenAI 官方流的真实形态——首帧带 name,后续帧只带 // arguments 片段、不再出现 name 键——因此是累加型与覆盖型客户端的共同祖先行为。 // // 两类消费模型在该形态下同时正确: // - 累加型(官方 WorkBuddy/CodeBuddy `name += tc_function?.name || ""`): // 后续分片 name 键缺失 → 追加空串,累积 name 保持唯一,不再拼成 Bash×帧数(issue #82)。 // - 覆盖型(hawklithm#2 / Grok Build `name ?? state.name` 或 `if (name) state.name = name`): // 后续分片 name 键缺失 → 保留已建好的首帧 name,不被空串意外清空。 // 键缺失是比空串更安全的形态:`??` 与 truthy 守卫对缺失键必然保留旧值, // 而对空串,`??` 会误判为重设并清空工具名。 // // seen 记录每个 index 是否已发过首片(与 name 是否非空无关);删除是幂等的。 // 只动 function.name 键,id/type/arguments 原样透传。 func stripToolCallNames(obj map[string]any, seen map[int]bool) { choices, _ := obj["choices"].([]any) for _, ci := range choices { c, _ := ci.(map[string]any) if c == nil { continue } delta, _ := c["delta"].(map[string]any) if delta == nil { continue } tcs, _ := delta["tool_calls"].([]any) for _, tci := range tcs { tc, _ := tci.(map[string]any) if tc == nil { continue } idx := 0 if v, ok := tc["index"].(float64); ok { idx = int(v) } if seen[idx] { // 已发过首片:删除本分片的 name 键(存在即删,幂等)。 if fn, _ := tc["function"].(map[string]any); fn != nil { delete(fn, "name") } continue } // 首现:保留 name 键原样(上游首片通常带非空 name;空 name 也照发, // 与 OpenAI 对「首帧无 name」的容忍一致),随后分片统一删除。 seen[idx] = true } } } // normalizeFrame 以 OpenAI 流式规范白名单重建帧:仅保留标准字段, // 剔除上游噪声(finish_reason:"" → null、空 content/refusal、空 tool_calls 列表、 // 空占位 function_call、顶层未知字段),空 delta 键一律省略, // usage 缺失 → null,保证任意标准客户端按规范解析。 func normalizeFrame(obj map[string]any) map[string]any { out := map[string]any{} for _, k := range []string{"id", "object", "created", "model", "system_fingerprint", "service_tier"} { if v, ok := obj[k]; ok && v != nil { out[k] = v } } if _, ok := out["object"]; !ok { out["object"] = "chat.completion.chunk" } if _, ok := out["id"]; !ok { out["id"] = "chatcmpl-wb2api" } if chs, ok := obj["choices"].([]any); ok { nchs := make([]any, 0, len(chs)) for _, ci := range chs { c, ok := ci.(map[string]any) if !ok { continue } nc := map[string]any{} if idx, ok := c["index"]; ok { nc["index"] = idx } delta := map[string]any{} if d, ok := c["delta"].(map[string]any); ok { if v, ok := d["role"].(string); ok && v != "" { delta["role"] = v } if v, ok := d["content"].(string); ok && v != "" { delta["content"] = v } if v, ok := d["reasoning_content"].(string); ok && v != "" { delta["reasoning_content"] = v } if v, ok := d["refusal"].(string); ok && v != "" { delta["refusal"] = v } if tcs, ok := d["tool_calls"].([]any); ok && len(tcs) > 0 { delta["tool_calls"] = tcs } if fc, ok := d["function_call"]; ok && fc != nil { // 空占位 function_call(name/arguments 全空)视为噪声剔除 keep := false if fcm, ok2 := fc.(map[string]any); ok2 { n, _ := fcm["name"].(string) a, _ := fcm["arguments"].(string) keep = n != "" || a != "" } else { keep = true } if keep { delta["function_call"] = fc } } } nc["delta"] = delta if fr, ok := c["finish_reason"].(string); ok && fr != "" { nc["finish_reason"] = fr } else { nc["finish_reason"] = nil } nchs = append(nchs, nc) } out["choices"] = nchs } if rawUsage, ok := obj["usage"]; ok { if u, ok := rawUsage.(map[string]any); ok { out["usage"] = normalizeUsageCacheAliases(u) } else { out["usage"] = rawUsage } } else { out["usage"] = nil } return out } // Stream 透传上游 SSE 到 w(逐帧规范化后 flush),保证至少写一个 [DONE]。 // 调用方必须先设置过 status 200;本函数自设 SSE headers。 // 流式策略:逐帧透传(规范化已剥空 content 噪声),恢复与上游一致的平滑流式。 // // StreamHint 变体(gateway_hint 任务):上游 error 帧透传时附加 // error.gateway_hint 字段——message 原文不动,hint 并列补充;hintFn 返回空串 // 或 nil 时与 Stream 行为逐字节一致。 func Stream(w http.ResponseWriter, r io.Reader) error { return StreamHint(w, r, nil) } // StreamOption StreamHint 的可选行为开关(均不影响透传字节,只做旁路观测)。 type StreamOption func(*streamOptions) type streamOptions struct { onErrorFrame func(payload string) } // WithErrorFrameObserver 注册"上游以 error 帧报错"的观察者:透传该帧之前先用 // payload(原始 data 内容)回调一次。给调用方一个**账号处置**的挂载点—— // 上游「200 + error 帧」是真实形态(6004 限流、内容拦截、审核),此前网关在读 // 第一帧之前就把账号记成功,限流号被当成健康号;有观察者后可在流尾按帧分类处置。 func WithErrorFrameObserver(fn func(payload string)) StreamOption { return func(o *streamOptions) { o.onErrorFrame = fn } } // StreamHint 同 Stream,但上游 error 帧透出前把 hintFn(payload) 的返回值写入 // error.gateway_hint。hintFn 为 nil 或返回空串 → 原样透传(零改写)。 // 空流兜底 error 帧("empty upstream stream")不带 hint(网关本地故障形态 // 未覆盖,不编造)。 // opts:旁路观测开关(见 WithErrorFrameObserver),不改变任何透传字节。 func StreamHint(w http.ResponseWriter, r io.Reader, hintFn func(string) string, opts ...StreamOption) error { var o streamOptions for _, f := range opts { if f != nil { f(&o) } } h := w.Header() h.Set("Content-Type", "text/event-stream") h.Set("Cache-Control", "no-cache") h.Set("Connection", "keep-alive") h.Set("X-Accel-Buffering", "no") fl, _ := w.(http.Flusher) // toolCallSeen 跨帧记录 delta.tool_calls 里已发过首片的 index, // 供逐 chunk 透传时收敛 name 为「每 index 一次」(对齐 OpenAI 官方流)。 toolCallSeen := map[int]bool{} // firstID 透传流的消息级 id 基准:缓存首个非空上游 id,后续帧缺失/空串时复用 // (issue #35:同一条 SSE 消息所有帧共用一个真实 id,后台按 id 归并;此前中间帧 // 一律补 chatcmpl-wb2api 哨兵,造成同流 id 分裂)。全流无真实 id → 才出现哨兵。 firstID := "" // writeRaw 原样写出一帧(绕过 normalizeFrame)并 flush。上游 error 帧(error-passthrough) // 与空流错误帧需保留 error 字段,不能被白名单剥掉,故经此写出。 // gateway_hint:上游 error 帧透出前按 hintFn 附加 error.gateway_hint 字段 // (error 对象上加一个键,message/code/requestId 等原文不动;hintFn 为 // nil / 空串 / 非 JSON 帧 → 原样写出,零改写)。 writeRaw := func(payload string) error { if hint := frameGatewayHint(hintFn, payload); hint != "" { payload = attachHintToErrorFrame(payload, hint) } if _, werr := io.WriteString(w, "data: "+payload+"\n\n"); werr != nil { return werr } if fl != nil { fl.Flush() } return nil } // writeFrame 把 payload 按规范白名单重建后以 data: 帧写出并 flush。 // 仅 JSON 解析成功时计数记为一次有效转发(JSON 解析失败照常降级原样写出,但不计数)。 writeFrame := func(payload string) (int, error) { var obj map[string]any valid := 0 if json.Unmarshal([]byte(payload), &obj) == nil { // 上游错误帧透传(error-passthrough):带 error 键的帧**原样写出**,不走 // normalizeFrame 白名单——白名单会剥掉 error 字段,客户端就看不到上游 // code/msg/requestId。error.message 即上游原文(如 6004 限流、审核拦截), // 计入有效帧(避免误判空流补写 "empty upstream stream")。 if _, hasErr := obj["error"]; hasErr { // 旁路通知(账号处置用):原始 payload 交给观察者,透传字节不变。 if o.onErrorFrame != nil { o.onErrorFrame(payload) } if werr := writeRaw(payload); werr != nil { return 0, werr } return 1, nil } // 先按 index 收敛 tool_calls name(每 index 仅首片保留,后续分片删 name 键),再规范化透传。 stripToolCallNames(obj, toolCallSeen) // id 续传:首帧非空真实 id 缓存;后续帧缺 id / 空 id 一律用缓存值, // 有自己 id 的帧保持原样(不同流分裂的帧允许各自 id)。 if firstID == "" { if v, ok := obj["id"].(string); ok && v != "" { firstID = v } } else { if v, ok := obj["id"].(string); !ok || v == "" { obj["id"] = firstID } } if raw, err := json.Marshal(normalizeFrame(obj)); err == nil { payload = string(raw) } valid = 1 } if _, werr := io.WriteString(w, "data: "+payload+"\n\n"); werr != nil { return 0, werr } if fl != nil { fl.Flush() } return valid, nil } br := bufio.NewReaderSize(r, 64*1024) validFrames := 0 readLoop: for { line, err := br.ReadString('\n') trimmed := strings.TrimRight(line, "\r\n") switch { case strings.HasPrefix(trimmed, "data: [DONE]"): // 上游显式结束:停止读取,DONE 之后的任何数据(含垃圾帧)一律不再透传。 // [DONE] 统一在循环结束后写出,保证恰好一个。 break readLoop case strings.HasPrefix(trimmed, "data: "): n, werr := writeFrame(strings.TrimPrefix(trimmed, "data: ")) validFrames += n if werr != nil { return werr } case trimmed != "": // 注释/其他行:原样透传 if _, werr := io.WriteString(w, line); werr != nil { return werr } if fl != nil { fl.Flush() } } // 空行(帧分隔)吞掉:本函数自产 "\n\n" if err != nil { if err == io.EOF { break } return err } } // 空流(0 有效帧):先写一帧 error(绕过 normalizeFrame 原样保留 error 字段), // 再补 [DONE] 保证客户端能正常收尾,并返回非 nil error 供调用方记录。 // 网关本地空流兜底帧走 hintFn=nil 的直写路径:该形态未覆盖(不编造 hint), // 且 writeRaw 的 hintFn 闭包在空流路径下可能携带上一帧的上下文造成误配。 if validFrames == 0 { _ = writeRaw(`{"error":{"message":"empty upstream stream","type":"upstream_error","code":"upstream_parse"}}`) } // 保证恰好写一个 [DONE](上游漏发时兜底补上)。 if _, err := io.WriteString(w, "data: [DONE]\n\n"); err != nil { return err } if fl != nil { fl.Flush() } if validFrames == 0 { return errEmptyStream } return nil } // frameGatewayHint 取 error 帧的 gateway_hint(hintFn 缺失/异常返回空串 → 不附加)。 func frameGatewayHint(hintFn func(string) string, payload string) string { if hintFn == nil { return "" } // panic 隔离:hint 判定是补充功能,任何实现缺陷不得击穿流透传主路径。 defer func() { _ = recover() }() return strings.TrimSpace(hintFn(payload)) } // attachHintToErrorFrame 在 error 帧的 error 对象上附加 gateway_hint 字段。 // message/code/requestId 等既有键原样保留(只加不改);非 JSON / 无 error 对象 → // payload 原样返回(宁可不加 hint 也不破坏原文透传)。 func attachHintToErrorFrame(payload, hint string) string { var obj map[string]any if json.Unmarshal([]byte(payload), &obj) != nil { return payload } e, ok := obj["error"].(map[string]any) if !ok { return payload } e["gateway_hint"] = hint out, err := json.Marshal(obj) if err != nil { return payload } return string(out) }