diff --git a/relay/channel/openai/converter.go b/relay/channel/openai/converter.go index 645959c..c42bcab 100644 --- a/relay/channel/openai/converter.go +++ b/relay/channel/openai/converter.go @@ -1369,11 +1369,9 @@ func HandleResponsesStreamToChat(ctx context.Context, resp *http.Response, info if chunk == nil { return true } - data, err := json.Marshal(chunk) - if err != nil { - return false - } - return helper.WriteSSEData(writer, string(data)) == nil + // 池化缓冲 + 绑定 encoder:序列化与拼帧零分配,整帧单次写出。 + // 序列化失败与写失败都返回 false(同旧行为),由调用方终止流。 + return helper.WriteSSEDataJSON(writer, chunk) == nil } sendStartIfNeeded := func() bool { diff --git a/relay/channel/openai/gemini_response.go b/relay/channel/openai/gemini_response.go index 64729a1..f5b064a 100644 --- a/relay/channel/openai/gemini_response.go +++ b/relay/channel/openai/gemini_response.go @@ -157,8 +157,8 @@ func handleGeminiInboundStream(ctx context.Context, resp *http.Response, info *c } if len(geminiChunk.Candidates) > 0 { - chunkData, _ := json.Marshal(geminiChunk) - _ = helper.WriteSSEData(writer, string(chunkData)) + // 传指针而非结构体值:值传进 any 会额外堆分配一份拷贝用于装箱 + _ = helper.WriteSSEDataJSON(writer, &geminiChunk) } } @@ -194,10 +194,9 @@ func handleGeminiInboundStream(ctx context.Context, resp *http.Response, info *c finalChunk.UsageMetadata.ThoughtsTokenCount = usage.CompletionTokenDetails.ReasoningTokens } } - chunkData, _ := json.Marshal(finalChunk) - _ = helper.WriteSSEData(writer, string(chunkData)) + _ = helper.WriteSSEDataJSON(writer, &finalChunk) - helper.WriteSSEData(writer, "[DONE]") + _ = helper.WriteSSEData(writer, "[DONE]") info.StreamStatus.SetEndReason(common.StreamEndReasonDone, nil) if err := scanner.Err(); err != nil && err != io.EOF && ctx.Err() == nil { diff --git a/relay/channel/openai/pooled_sse_stream_test.go b/relay/channel/openai/pooled_sse_stream_test.go new file mode 100644 index 0000000..36d02bc --- /dev/null +++ b/relay/channel/openai/pooled_sse_stream_test.go @@ -0,0 +1,211 @@ +package openai + +import ( + "context" + "encoding/json" + "io" + "net/http" + "strings" + "testing" + + "github.com/qianfree/team-api/relay/common" + "github.com/qianfree/team-api/relay/constant" + "github.com/qianfree/team-api/relay/dto" +) + +// 本文件覆盖两条改用池化帧写入器(helper.WriteSSEDataJSON)的转换路径: +// - handleGeminiInboundStream:OpenAI 上游 → Gemini 客户端 +// - HandleResponsesStreamToChat:Responses 上游 → Chat Completions 客户端 +// +// 两条路径此前无测试覆盖。断言两件事: +// 1. 输出仍是合法且内容正确的 SSE(池化不能改变对外协议); +// 2. 每次 Write 都落在帧边界(整帧单次写出,并发保活 ping 插不进帧中间)。 + +// assertFramesAreSingleWrites 断言每次 Write 都是一个以空行收尾的完整帧 +func assertFramesAreSingleWrites(t *testing.T, chunks []string) { + t.Helper() + if len(chunks) == 0 { + t.Fatal("没有任何写出") + } + for i, c := range chunks { + if !strings.HasSuffix(c, "\n\n") { + t.Errorf("第 %d 次 Write 未落在帧边界,保活 ping 可插入其中:\n%q", i, c) + } + } +} + +// collectSSEDataFrames 从 SSE 文本中取出所有 data 行的内容 +func collectSSEDataFrames(raw string) []string { + var out []string + for _, line := range strings.Split(raw, "\n") { + line = strings.TrimSuffix(line, "\r") + if strings.HasPrefix(line, "data:") { + out = append(out, strings.TrimSpace(strings.TrimPrefix(line, "data:"))) + } + } + return out +} + +func geminiInboundInfo() *common.RelayInfo { + return &common.RelayInfo{ + IsStream: true, + InboundFormat: constant.RelayFormatGemini, + ClientFormat: constant.RelayFormatGemini, + OriginModelName: "gemini-3-pro", + StreamStatus: common.NewStreamStatus(), + ChannelMeta: &common.ChannelMeta{ + ChannelType: int(constant.ProviderOpenAI), + BaseURL: "https://upstream.example.com", + UpstreamModelName: "gpt-4o", + }, + } +} + +// TestHandleGeminiInboundStream_PooledFramesValid OpenAI SSE → Gemini SSE: +// 每个 chunk 都是完整可解析的 Gemini JSON,且整帧单次写出。 +func TestHandleGeminiInboundStream_PooledFramesValid(t *testing.T) { + upstream := strings.Join([]string{ + `data: {"id":"c1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant","content":"你好"}}]}`, + ``, + `data: {"id":"c1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"世界"}}]}`, + ``, + `data: {"id":"c1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":11,"completion_tokens":7,"total_tokens":18}}`, + ``, + `data: [DONE]`, + ``, + }, "\n") + + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": {"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(upstream)), + } + + info := geminiInboundInfo() + rec := &sseFrameRecorder{} + + usage, err := handleGeminiInboundStream(context.Background(), resp, info, rec) + if err != nil { + t.Fatalf("handleGeminiInboundStream error: %v", err) + } + + assertFramesAreSingleWrites(t, rec.chunks) + + frames := collectSSEDataFrames(rec.body.String()) + if len(frames) < 2 { + t.Fatalf("data 帧数 = %d, 期望至少 2(内容 chunk + 终止 chunk)\n%s", len(frames), rec.body.String()) + } + + var text strings.Builder + sawDone := false + for i, f := range frames { + if f == "[DONE]" { + sawDone = true + if i != len(frames)-1 { + t.Errorf("[DONE] 不在最后一帧(位置 %d/%d)", i, len(frames)-1) + } + continue + } + var chunk dto.GeminiChatResponse + if err := json.Unmarshal([]byte(f), &chunk); err != nil { + t.Fatalf("第 %d 帧不是合法 Gemini JSON(池化若出错会产生这种残片): %v\ndata=%q", i, err, f) + } + for _, cand := range chunk.Candidates { + if cand.Content == nil { + continue + } + for _, p := range cand.Content.Parts { + text.WriteString(p.Text) + } + } + } + if !sawDone { + t.Error("缺少 [DONE] 终止帧") + } + if got := text.String(); got != "你好世界" { + t.Errorf("转发文本 = %q, want %q", got, "你好世界") + } + if usage == nil || usage.PromptTokens != 11 || usage.CompletionTokens != 7 { + t.Errorf("usage = %+v, want prompt=11 completion=7", usage) + } +} + +func responsesToChatInfo() *common.RelayInfo { + return &common.RelayInfo{ + IsStream: true, + InboundFormat: constant.RelayFormatOpenAI, + ClientFormat: constant.RelayFormatOpenAI, + OriginModelName: "gpt-4o", + StreamStatus: common.NewStreamStatus(), + ChannelMeta: &common.ChannelMeta{ + ChannelType: int(constant.ProviderOpenAI), + BaseURL: "https://upstream.example.com", + UpstreamModelName: "gpt-4o", + SupportsResponses: true, + }, + } +} + +// TestHandleResponsesStreamToChat_PooledFramesValid Responses SSE → Chat SSE: +// 每个 chunk 都是完整可解析的 chat.completion.chunk,且整帧单次写出。 +func TestHandleResponsesStreamToChat_PooledFramesValid(t *testing.T) { + upstream := strings.Join([]string{ + "event: response.created", + `data: {"type":"response.created","response":{"id":"resp_1","object":"response","status":"in_progress"}}`, + ``, + "event: response.output_text.delta", + `data: {"type":"response.output_text.delta","delta":"你好"}`, + ``, + "event: response.output_text.delta", + `data: {"type":"response.output_text.delta","delta":"世界"}`, + ``, + "event: response.completed", + `data: {"type":"response.completed","response":{"id":"resp_1","object":"response","status":"completed","usage":{"input_tokens":11,"output_tokens":7,"total_tokens":18}}}`, + ``, + `data: [DONE]`, + ``, + }, "\n") + + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": {"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(upstream)), + } + + info := responsesToChatInfo() + rec := &sseFrameRecorder{} + + if _, err := HandleResponsesStreamToChat(context.Background(), resp, info, rec); err != nil { + t.Fatalf("HandleResponsesStreamToChat error: %v", err) + } + + assertFramesAreSingleWrites(t, rec.chunks) + + frames := collectSSEDataFrames(rec.body.String()) + var text strings.Builder + sawDone := false + for i, f := range frames { + if f == "[DONE]" { + sawDone = true + continue + } + var chunk dto.ChatCompletionStreamResponse + if err := json.Unmarshal([]byte(f), &chunk); err != nil { + t.Fatalf("第 %d 帧不是合法 chat.completion.chunk: %v\ndata=%q", i, err, f) + } + if chunk.Object != "chat.completion.chunk" { + t.Errorf("第 %d 帧 object = %q, want chat.completion.chunk", i, chunk.Object) + } + for _, c := range chunk.Choices { + if s, ok := c.Delta.Content.(string); ok { + text.WriteString(s) + } + } + } + if !sawDone { + t.Error("缺少 [DONE] 终止帧") + } + if got := text.String(); got != "你好世界" { + t.Errorf("转发文本 = %q, want %q", got, "你好世界") + } +} diff --git a/relay/helper/stream_test.go b/relay/helper/stream_test.go index 7b36b13..de87706 100644 --- a/relay/helper/stream_test.go +++ b/relay/helper/stream_test.go @@ -832,6 +832,33 @@ func TestWriteSSEEventJSON_RealWorldFrameSizes(t *testing.T) { } } +// BenchmarkSSEData_PooledJSON_Pointer / _Value 验证「传指针而非结构体值」的理由: +// 结构体值传进 any 形参需要装箱,会额外堆分配一份拷贝;传指针则不会。 +// 迁移调用点时据此统一传 &chunk。 +func BenchmarkSSEData_PooledJSON_Pointer(b *testing.B) { + w := &discardWriter{} + payload := benchClaudeDelta() // 已是指针 + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if err := WriteSSEDataJSON(w, payload); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkSSEData_PooledJSON_Value(b *testing.B) { + w := &discardWriter{} + payload := *benchClaudeDelta() // 解引用成值 + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if err := WriteSSEDataJSON(w, payload); err != nil { + b.Fatal(err) + } + } +} + // TestSSEFrameWriter_MaxBufDegradeSuppressesPing 超大帧降级为分次写出时,写出的这批字节 // 不含帧尾 —— 此刻放行 ping 依然会把帧劈成两个事件。降级期间必须抑制 ping,直到分隔空行到达。 func TestSSEFrameWriter_MaxBufDegradeSuppressesPing(t *testing.T) {