diff --git a/relay/channel/claude/native_stream_test.go b/relay/channel/claude/native_stream_test.go new file mode 100644 index 0000000..c9f11ed --- /dev/null +++ b/relay/channel/claude/native_stream_test.go @@ -0,0 +1,175 @@ +package claude + +import ( + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/qianfree/team-api/relay/common" + "github.com/qianfree/team-api/relay/constant" +) + +// nativeStreamInfo 构造 Claude 原生透传（/v1/messages 直通）的 RelayInfo +func nativeStreamInfo() *common.RelayInfo { + return &common.RelayInfo{ + RelayMode: int(constant.RelayModeClaudeMessages), + IsStream: true, + InboundFormat: constant.RelayFormatClaude, + ClientFormat: constant.RelayFormatClaude, + OriginModelName: "claude-opus-4-8", + StartTime: time.Now(), + StreamStatus: common.NewStreamStatus(), + ChannelMeta: &common.ChannelMeta{ + BaseURL: "https://upstream.example.com", + UpstreamModelName: "claude-opus-4-8", + IsModelMapped: false, + }, + } +} + +func sseResponse(body io.Reader) *http.Response { + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": {"text/event-stream"}}, + Body: io.NopCloser(body), + } +} + +// TestHandleClaudeNativeStream_FramesForwardedIntact 正常路径：上游 SSE 逐字节原样转发， +// 每个事件都带完整 data（客户端解码后不会出现空 data）。 +func TestHandleClaudeNativeStream_FramesForwardedIntact(t *testing.T) { + upstream := strings.Join([]string{ + `event: message_start`, + `data: {"type":"message_start","message":{"id":"msg_01","model":"claude-opus-4-8","usage":{"input_tokens":2,"output_tokens":0,"cache_creation_input_tokens":5161}}}`, + ``, + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"你好"}}`, + ``, + `event: message_delta`, + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":7}}`, + ``, + `event: message_stop`, + `data: {"type":"message_stop"}`, + ``, + ``, + }, "\n") + + a := &Adaptor{} + info := nativeStreamInfo() + rec := httptest.NewRecorder() + + usage, err := a.handleClaudeNativeStream(context.Background(), sseResponse(strings.NewReader(upstream)), info, rec) + if err != nil { + t.Fatalf("handleClaudeNativeStream error: %v", err) + } + + if got := rec.Body.String(); got != upstream { + t.Errorf("转发内容与上游不一致:\ngot: %q\nwant: %q", got, upstream) + } + if reason := info.StreamStatus.GetEndReason(); reason != common.StreamEndReasonDone { + t.Errorf("EndReason = %q, want %q", reason, common.StreamEndReasonDone) + } + if usage == nil || usage.PromptTokens != 2 || usage.CompletionTokens != 7 || usage.CacheCreationTokens != 5161 { + t.Errorf("usage = %+v, want prompt=2 completion=7 cacheCreation=5161", usage) + } +} + +// TestHandleClaudeNativeStream_LastLineWithoutNewline 上游最后一行没有换行结尾时 +// 不得丢弃 —— ReadString 会同时返回「残留数据 + io.EOF」，旧实现直接 break 会吞掉这一行。 +func TestHandleClaudeNativeStream_LastLineWithoutNewline(t *testing.T) { + upstream := "event: message_stop\ndata: {\"type\":\"message_stop\"}" + + a := &Adaptor{} + info := nativeStreamInfo() + rec := httptest.NewRecorder() + + if _, err := a.handleClaudeNativeStream(context.Background(), sseResponse(strings.NewReader(upstream)), info, rec); err != nil { + t.Fatalf("handleClaudeNativeStream error: %v", err) + } + + if got := rec.Body.String(); got != upstream { + t.Errorf("无换行结尾的最后一行被丢弃:\ngot: %q\nwant: %q", got, upstream) + } +} + +// TestHandleClaudeNativeStream_ClientGoneDiscardsPartialFrame 流中断时， +// 尚未攒齐的半帧必须丢弃，客户端只应看到完整的 error 帧 —— 否则半帧与 error 帧 +// 拼在一起会变成非法 SSE，客户端反而先在 JSON 解析上炸掉。 +func TestHandleClaudeNativeStream_ClientGoneDiscardsPartialFrame(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + // 吐出一个完整帧后取消 ctx，再交出半帧（只有 event 行）：主循环把半帧收进缓冲、 + // 回到循环顶部时 select 命中 Done，模拟客户端在帧中途断开 + halfSent := false + body := io.MultiReader( + strings.NewReader("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":2,\"cache_creation_input_tokens\":5161}}}\n\n"), + readerFunc(func(p []byte) (int, error) { + if halfSent { + return 0, io.EOF + } + halfSent = true + cancel() + return copy(p, "event: content_block_delta\n"), nil + }), + ) + + a := &Adaptor{} + info := nativeStreamInfo() + rec := httptest.NewRecorder() + + usage, err := a.handleClaudeNativeStream(ctx, sseResponse(body), info, rec) + if !errors.Is(err, common.ErrStreamInterrupted) { + t.Fatalf("err = %v, want ErrStreamInterrupted", err) + } + if reason := info.StreamStatus.GetEndReason(); reason != common.StreamEndReasonClientGone { + t.Errorf("EndReason = %q, want %q", reason, common.StreamEndReasonClientGone) + } + + got := rec.Body.String() + if strings.Contains(got, "content_block_delta") { + t.Errorf("半帧未被丢弃，泄漏到客户端:\n%q", got) + } + if !strings.Contains(got, `"type":"error"`) { + t.Errorf("未写出 error 帧，客户端会挂起:\n%q", got) + } + // 中断兜底仍要保留已拿到的 usage，避免按 0 token 结算 + if usage == nil || usage.PromptTokens != 2 || usage.CacheCreationTokens != 5161 { + t.Errorf("usage = %+v, want prompt=2 cacheCreation=5161", usage) + } +} + +// TestHandleClaudeNativeStream_UpstreamErrorDiscardsPartialFrame 上游读取出错时， +// 缓冲里的残留是不完整的帧，不得写给客户端。 +func TestHandleClaudeNativeStream_UpstreamErrorDiscardsPartialFrame(t *testing.T) { + upstreamErr := errors.New("unexpected EOF from upstream") + body := io.MultiReader( + strings.NewReader("event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: content_block_delta\n"), + readerFunc(func(p []byte) (int, error) { return 0, upstreamErr }), + ) + + a := &Adaptor{} + info := nativeStreamInfo() + rec := httptest.NewRecorder() + + _, err := a.handleClaudeNativeStream(context.Background(), sseResponse(body), info, rec) + if err == nil || !strings.Contains(err.Error(), "upstream stream interrupted") { + t.Fatalf("err = %v, want upstream stream interrupted", err) + } + if reason := info.StreamStatus.GetEndReason(); reason != common.StreamEndReasonError { + t.Errorf("EndReason = %q, want %q", reason, common.StreamEndReasonError) + } + if got := rec.Body.String(); strings.Contains(got, "content_block_delta") { + t.Errorf("上游出错时半帧泄漏到客户端:\n%q", got) + } +} + +// readerFunc 把函数适配成 io.Reader +type readerFunc func(p []byte) (int, error) + +func (f readerFunc) Read(p []byte) (int, error) { return f(p) } diff --git a/relay/channel/claude/response.go b/relay/channel/claude/response.go index ebf4b24..ff82f49 100644 --- a/relay/channel/claude/response.go +++ b/relay/channel/claude/response.go @@ -462,9 +462,25 @@ func (a *Adaptor) handleClaudeNativeStream(ctx context.Context, resp *http.Respo defer stopPing() reader := bufio.NewReaderSize(resp.Body, 64*1024) + // 按帧攒齐再写：逐行写会让并发的保活 ping 把自带空行插进帧中间，客户端据此派发出 + // 一个 data 为空的事件（JSON.parse("") 报错并中止请求），网关侧只看到 ctx 取消被记成 + // client_gone，真实成因被掩盖。SafeWriter 只保证单次 Write 原子，保证不了整帧原子。 + frame := helper.NewSSEFrameWriter(writer) var usage dto.ClaudeUsage var transferredTextLen int // 已转发的文本/思考内容长度，供流中断输出估算 + // onWriteFail 写客户端失败：立即关闭上游连接，停止 token 生成，按流中断结算 + onWriteFail := func(writeErr error) (*common.Usage, error) { + g.Log().Warningf(context.Background(), + "[ClaudeNativeStream] 写入客户端失败 request_id=%s writeErr=%v ctx.Err=%v elapsed=%v", + info.RequestID, writeErr, ctx.Err(), time.Since(info.StartTime)) + cleanup() + info.StreamStatus.SetEndReason(common.StreamEndReasonClientGone, writeErr) + interruptedUsage := buildUsageFromClaude(&usage) + helper.ApplyInterruptedUsageFallback(info, interruptedUsage, transferredTextLen) + return interruptedUsage, common.ErrStreamInterrupted + } + for { select { case <-streamCtx.Done(): @@ -472,6 +488,8 @@ func (a *Adaptor) handleClaudeNativeStream(ctx context.Context, resp *http.Respo // 让 SDK 收到"200 + SSE头 + 无事件 + EOF"，Anthropic SDK 进入等待状态， // 后续请求的响应会被误判为当前流的 SSE 数据 → "Failed to parse JSON"。 // 发送一个 Claude 格式的 error event，让 SDK 以正常 API Error 退出，而非挂起。 + // 未完成的半帧尚未出网，先丢弃，避免与 error 帧拼成非法 SSE 输出。 + frame.Discard() _, _ = fmt.Fprintf(writer, "event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"api_error\",\"message\":\"upstream disconnected\"}}\n\n") if f, ok := writer.(http.Flusher); ok { @@ -485,11 +503,13 @@ func (a *Adaptor) handleClaudeNativeStream(ctx context.Context, resp *http.Respo default: } + // ReadString 在上游 EOF 时可能同时返回「最后一段无换行的数据 + io.EOF」， + // 必须先处理 line 再判 err，否则会丢掉上游未以换行结尾的最后一行。 line, err := reader.ReadString('\n') - if err != nil { - if err == io.EOF { - break - } + + if err != nil && err != io.EOF { + // 上游读取出错：缓冲里的残留是不完整的帧，写给客户端只会造成非法 SSE，丢弃 + frame.Discard() info.StreamStatus.SetEndReason(common.StreamEndReasonError, err) // 已有部分输出时按部分成功处理（避免标记为完全失败） interruptedUsage := buildUsageFromClaude(&usage) @@ -499,75 +519,71 @@ func (a *Adaptor) handleClaudeNativeStream(ctx context.Context, resp *http.Respo return interruptedUsage, fmt.Errorf("upstream stream interrupted: %w", err) } - if strings.HasPrefix(line, "data:") { - data, _ := helper.ExtractSSEData(line) + if line != "" { + if strings.HasPrefix(line, "data:") { + data, _ := helper.ExtractSSEData(line) - if data != "" && data != "[DONE]" { - info.SetFirstResponseTime() - } + if data != "" && data != "[DONE]" { + info.SetFirstResponseTime() + } - var event dto.ClaudeResponse - if json.Unmarshal([]byte(data), &event) != nil { - // JSON 解析失败：静默跳过 - } else { - switch event.Type { - case "message_start": - if event.Message != nil && event.Message.Usage != nil { - usage = *event.Message.Usage + var event dto.ClaudeResponse + if json.Unmarshal([]byte(data), &event) != nil { + // JSON 解析失败：静默跳过 + } else { + switch event.Type { + case "message_start": + if event.Message != nil && event.Message.Usage != nil { + usage = *event.Message.Usage + } + case "content_block_delta": + // 累计已转发文本长度，供流中断（message_delta 未到达时）输出估算 + if event.Delta != nil { + if event.Delta.Text != nil { + transferredTextLen += len(*event.Delta.Text) + } + if event.Delta.Thinking != nil { + transferredTextLen += len(*event.Delta.Thinking) + } + } + case "message_delta": + if event.Usage != nil { + if event.Usage.InputTokens > 0 { + usage.InputTokens = event.Usage.InputTokens + } + usage.OutputTokens = event.Usage.OutputTokens + if event.Usage.CacheReadInputTokens > 0 { + usage.CacheReadInputTokens = event.Usage.CacheReadInputTokens + } + if event.Usage.CacheCreationInputTokens > 0 { + usage.CacheCreationInputTokens = event.Usage.CacheCreationInputTokens + } + if event.Usage.CacheCreation != nil { + usage.CacheCreation = event.Usage.CacheCreation + } + } + case "error": + info.StreamStatus.SetEndReason(common.StreamEndReasonError, fmt.Errorf("claude upstream stream error")) } - case "content_block_delta": - // 累计已转发文本长度，供流中断（message_delta 未到达时）输出估算 - if event.Delta != nil { - if event.Delta.Text != nil { - transferredTextLen += len(*event.Delta.Text) - } - if event.Delta.Thinking != nil { - transferredTextLen += len(*event.Delta.Thinking) - } - } - case "message_delta": - if event.Usage != nil { - if event.Usage.InputTokens > 0 { - usage.InputTokens = event.Usage.InputTokens - } - usage.OutputTokens = event.Usage.OutputTokens - if event.Usage.CacheReadInputTokens > 0 { - usage.CacheReadInputTokens = event.Usage.CacheReadInputTokens - } - if event.Usage.CacheCreationInputTokens > 0 { - usage.CacheCreationInputTokens = event.Usage.CacheCreationInputTokens - } - if event.Usage.CacheCreation != nil { - usage.CacheCreation = event.Usage.CacheCreation - } - } - case "error": - info.StreamStatus.SetEndReason(common.StreamEndReasonError, fmt.Errorf("claude upstream stream error")) + } + + if info.ChannelMeta.IsModelMapped { + replaced := string(helper.ReplaceModelName([]byte(data), info.OriginModelName)) + line = fmt.Sprintf("data: %s\n", replaced) } } - if info.ChannelMeta.IsModelMapped { - replaced := string(helper.ReplaceModelName([]byte(data), info.OriginModelName)) - line = fmt.Sprintf("data: %s\n", replaced) + if werr := frame.WriteLine(line); werr != nil { + return onWriteFail(werr) } } - if _, err := writer.Write([]byte(line)); err != nil { - // 写入客户端失败：立即关闭上游连接，停止生成 - g.Log().Warningf(context.Background(), - "[ClaudeNativeStream] 写入客户端失败 request_id=%s writeErr=%v ctx.Err=%v elapsed=%v", - info.RequestID, err, ctx.Err(), time.Since(info.StartTime)) - cleanup() - info.StreamStatus.SetEndReason(common.StreamEndReasonClientGone, err) - interruptedUsage := buildUsageFromClaude(&usage) - helper.ApplyInterruptedUsageFallback(info, interruptedUsage, transferredTextLen) - return interruptedUsage, common.ErrStreamInterrupted - } - - if len(line) == 1 && line[0] == '\n' { - if f, ok := writer.(http.Flusher); ok { - f.Flush() + if err == io.EOF { + // 上游未以空行结尾时冲刷残留，避免丢掉最后一帧 + if ferr := frame.FlushPartial(); ferr != nil { + return onWriteFail(ferr) } + break } } diff --git a/relay/constant/errors.go b/relay/constant/errors.go index 8fc237b..a02e2d2 100644 --- a/relay/constant/errors.go +++ b/relay/constant/errors.go @@ -15,7 +15,7 @@ import ( ) // Relay 层哨兵错误 -var ErrAllChannelsFailed = gerror.New("all channels failed") +var ErrAllChannelsFailed = gerror.New("no available channels") // ErrStatefulResponsesUnsupported responses 有状态请求（previous_response_id）落在 // 不支持 Responses 协议的 chat-only 渠道上，降级转换会静默丢失会话上下文，必须快速失败。 diff --git a/relay/helper/stream.go b/relay/helper/stream.go index 308efc6..4829b63 100644 --- a/relay/helper/stream.go +++ b/relay/helper/stream.go @@ -165,8 +165,10 @@ func ApplyInterruptedUsageFallback(info *common.RelayInfo, usage *common.Usage, // SafeWriter 包装 http.ResponseWriter，用互斥锁串行化 Write/WriteHeader/Flush， // 使保活 ping goroutine 与主循环可以安全地并发写同一个 ResponseWriter。 // -// 注意：fmt.Fprintf 对单个 SSE 帧只产生一次 Write 调用，因此每个 SSE 帧在锁内是原子写入的； -// Flush 单独加锁，与其他写入交错也不会破坏帧的字节完整性。 +// 注意：本类型只保证**单次 Write 调用**原子，不保证 SSE 帧原子。 +// 调用方一次 Write 写完整帧时（如 WriteSSEData/WriteSSEEvent 的单次 Fprintf），帧即原子； +// 而按行转发上游 SSE 的调用方必须改用 SSEFrameWriter 攒齐整帧再写，否则 ping 的空行 +// 会插进帧中间，让客户端派发出 data 为空的事件（详见 SSEFrameWriter 的说明）。 type SafeWriter struct { w http.ResponseWriter mu sync.Mutex @@ -211,6 +213,84 @@ var ( _ http.Flusher = (*SafeWriter)(nil) ) +const ( + // sseFrameMaxBuf 单帧缓冲上限：超过则先把已缓冲内容写出（降级为逐行写）， + // 防止上游异常（始终不发空行）把整条流攒进内存。正常 SSE 帧远小于此。 + sseFrameMaxBuf = 1 << 20 // 1 MiB + // sseFrameKeepBuf 帧写出后保留复用的缓冲上限，超过则释放，避免个别大帧长期占用内存 + sseFrameKeepBuf = 64 * 1024 +) + +// IsSSEFrameBoundary 判断一行是否为 SSE 帧分隔空行（兼容 LF / CRLF 两种行尾）。 +func IsSSEFrameBoundary(line string) bool { + return line == "\n" || line == "\r\n" +} + +// SSEFrameWriter 为「逐行转发上游 SSE」的场景提供帧级原子写入。 +// +// SSE 以空行分帧，客户端读到空行即派发事件。若转发方按行写，而保活 ping goroutine +// 又并发写同一个 writer，ping 自带的空行就可能插进一帧中间，把只写了 `event: X` 的 +// 半帧提前派发出去 —— 客户端拿到 data 为空的事件，对其做 JSON.parse 直接抛错 +// （如 Claude Code 的 "JSON Parse error: Unexpected EOF"）并中止请求，网关侧随后 +// 只观察到 ctx 取消，被记成 client_gone，真实成因被掩盖。 +// SafeWriter 只保证单次 Write 原子，保证不了整帧原子，故需本类型按帧攒齐再写。 +// +// 用法：逐行 WriteLine（遇空行自动整帧写出并 Flush）；上游 EOF 时 FlushPartial +// 冲刷无空行结尾的残留；中断或上游读错时 Discard 丢弃半帧（半帧尚未出网，丢弃即可）。 +// 本类型非并发安全，只应由转发主循环单协程使用，底层 w 需为 SafeWriter。 +type SSEFrameWriter struct { + w http.ResponseWriter + buf []byte +} + +// NewSSEFrameWriter 创建帧级原子写入器，w 应为并发安全的 writer（如 SafeWriter）。 +func NewSSEFrameWriter(w http.ResponseWriter) *SSEFrameWriter { + return &SSEFrameWriter{w: w} +} + +// WriteLine 追加一行（含行尾换行符）。该行为帧分隔空行、或缓冲超过上限时整帧写出。 +func (f *SSEFrameWriter) WriteLine(line string) error { + f.buf = append(f.buf, line...) + if IsSSEFrameBoundary(line) || len(f.buf) >= sseFrameMaxBuf { + return f.flush() + } + return nil +} + +// FlushPartial 冲刷未以空行结尾的残留内容（上游 EOF 时调用，避免丢掉最后一帧）。 +func (f *SSEFrameWriter) FlushPartial() error { + return f.flush() +} + +// Discard 丢弃尚未写出的半帧内容。 +func (f *SSEFrameWriter) Discard() { + f.buf = f.buf[:0] +} + +// Buffered 返回当前缓冲的字节数（尚未写给客户端的半帧）。 +func (f *SSEFrameWriter) Buffered() int { + return len(f.buf) +} + +func (f *SSEFrameWriter) flush() error { + if len(f.buf) == 0 { + return nil + } + _, err := f.w.Write(f.buf) + if cap(f.buf) > sseFrameKeepBuf { + f.buf = nil + } else { + f.buf = f.buf[:0] + } + if err != nil { + return err + } + if fl, ok := f.w.(http.Flusher); ok { + fl.Flush() + } + return nil +} + // PingTicker 在后台定期发送 SSE 保活注释 // 调用方必须传入并发安全的 writer（如 SafeWriter），以避免与主循环并发写 ResponseWriter // 返回一个 stop 函数用于停止 goroutine diff --git a/relay/helper/stream_test.go b/relay/helper/stream_test.go index ad8da37..ad22d00 100644 --- a/relay/helper/stream_test.go +++ b/relay/helper/stream_test.go @@ -1,8 +1,10 @@ package helper import ( + "errors" "net/http" "net/http/httptest" + "strconv" "strings" "testing" ) @@ -192,7 +194,280 @@ func TestWriteSSEPing(t *testing.T) { } } -// TestExtractSSEData 测试 SSE 数据提取 +// sseEvent 一个被客户端派发出去的 SSE 事件 +type sseEvent struct { + event string + data string +} + +// decodeSSE 按 SSE 规范模拟客户端解码：`:` 开头为注释忽略，空行即派发当前事件。 +// 用于断言网关吐出的字节在客户端眼里是什么，而不是只比对原始字符串。 +func decodeSSE(raw string) []sseEvent { + var events []sseEvent + cur := sseEvent{} + var dataLines []string + dispatch := func() { + if cur.event == "" && len(dataLines) == 0 { + return + } + cur.data = strings.Join(dataLines, "\n") + events = append(events, cur) + cur = sseEvent{} + dataLines = nil + } + for _, line := range strings.Split(raw, "\n") { + line = strings.TrimSuffix(line, "\r") + switch { + case line == "": + dispatch() + case strings.HasPrefix(line, ":"): + // 注释（保活 ping），忽略 + case strings.HasPrefix(line, "event:"): + cur.event = strings.TrimSpace(strings.TrimPrefix(line, "event:")) + case strings.HasPrefix(line, "data:"): + dataLines = append(dataLines, strings.TrimSpace(strings.TrimPrefix(line, "data:"))) + } + } + dispatch() + return events +} + +// assertNoEmptyDataEvent 断言没有「有事件名但 data 为空」的事件 —— 这正是客户端 +// 对空串做 JSON.parse 报 "JSON Parse error: Unexpected EOF" 并中止请求的成因。 +func assertNoEmptyDataEvent(t *testing.T, raw string) { + t.Helper() + for i, ev := range decodeSSE(raw) { + if ev.event != "" && ev.data == "" { + t.Errorf("事件 #%d (event=%q) 的 data 为空，客户端 JSON.parse 会失败；原始输出:\n%s", + i, ev.event, raw) + } + } +} + +// TestSafeWriter_LineByLineWriteIsSplitByPing 对照用例：逐行写 SSE 时， +// SafeWriter 挡不住保活 ping 把自带空行插进帧中间 —— 半帧被提前派发， +// 客户端拿到 data 为空的事件。这是 SSEFrameWriter 存在的理由，也证明 +// assertNoEmptyDataEvent 确实能抓到该缺陷。 +func TestSafeWriter_LineByLineWriteIsSplitByPing(t *testing.T) { + rec := httptest.NewRecorder() + sw := NewSafeWriter(rec) + + _, _ = sw.Write([]byte("event: content_block_delta\n")) + // 主循环阻塞在读上游下一行时，ping goroutine 抢到了写入 + _ = WriteSSEPing(sw) + _, _ = sw.Write([]byte(`data: {"type":"content_block_delta"}` + "\n")) + _, _ = sw.Write([]byte("\n")) + + events := decodeSSE(rec.Body.String()) + if len(events) == 0 || events[0].data != "" { + t.Fatalf("对照用例失效：期望逐行写会派发出空 data 事件，实际 events=%+v", events) + } +} + +// TestSSEFrameWriter_PingCannotSplitFrame 回归：改用 SSEFrameWriter 后， +// 帧攒齐才出网，ping 只能落在帧边界，客户端不会再收到空 data 事件。 +func TestSSEFrameWriter_PingCannotSplitFrame(t *testing.T) { + rec := httptest.NewRecorder() + sw := NewSafeWriter(rec) + frame := NewSSEFrameWriter(sw) + + if err := frame.WriteLine("event: content_block_delta\n"); err != nil { + t.Fatalf("WriteLine failed: %v", err) + } + // 半帧必须还留在缓冲里，一个字节都不能出网 + if rec.Body.Len() != 0 { + t.Errorf("半帧提前出网: %q", rec.Body.String()) + } + if frame.Buffered() == 0 { + t.Error("Buffered() = 0, 期望半帧仍在缓冲中") + } + + // ping goroutine 此刻抢到写入 + if err := WriteSSEPing(sw); err != nil { + t.Fatalf("WriteSSEPing failed: %v", err) + } + + if err := frame.WriteLine(`data: {"type":"content_block_delta"}` + "\n"); err != nil { + t.Fatalf("WriteLine failed: %v", err) + } + if err := frame.WriteLine("\n"); err != nil { + t.Fatalf("WriteLine failed: %v", err) + } + + body := rec.Body.String() + assertNoEmptyDataEvent(t, body) + + events := decodeSSE(body) + if len(events) != 1 { + t.Fatalf("派发事件数 = %d, want 1; body=%q", len(events), body) + } + if events[0].event != "content_block_delta" || events[0].data != `{"type":"content_block_delta"}` { + t.Errorf("事件 = %+v, 与写入内容不符", events[0]) + } + // ping 必须整体落在帧之前，不能夹在帧内部 + if !strings.HasPrefix(body, ": PING\n\n") { + t.Errorf("ping 未落在帧边界: %q", body) + } +} + +// TestSSEFrameWriter_ConcurrentPingDuringFrames 并发压力下的帧完整性： +// 主循环持续逐行写帧，ping goroutine 同时高频写入，客户端解码后不得出现空 data 事件。 +func TestSSEFrameWriter_ConcurrentPingDuringFrames(t *testing.T) { + rec := httptest.NewRecorder() + sw := NewSafeWriter(rec) + frame := NewSSEFrameWriter(sw) + + stopPing := make(chan struct{}) + pingDone := make(chan struct{}) + go func() { + defer close(pingDone) + for { + select { + case <-stopPing: + return + default: + _ = WriteSSEPing(sw) + } + } + }() + + for i := 0; i < 500; i++ { + if err := frame.WriteLine("event: content_block_delta\n"); err != nil { + t.Fatalf("WriteLine failed: %v", err) + } + if err := frame.WriteLine(`data: {"i":` + strconv.Itoa(i) + "}\n"); err != nil { + t.Fatalf("WriteLine failed: %v", err) + } + if err := frame.WriteLine("\n"); err != nil { + t.Fatalf("WriteLine failed: %v", err) + } + } + + close(stopPing) + <-pingDone + + body := rec.Body.String() + assertNoEmptyDataEvent(t, body) + + events := decodeSSE(body) + if len(events) != 500 { + t.Fatalf("派发事件数 = %d, want 500", len(events)) + } + for i, ev := range events { + if want := `{"i":` + strconv.Itoa(i) + "}"; ev.data != want { + t.Fatalf("事件 #%d data = %q, want %q（帧内容错位）", i, ev.data, want) + } + } +} + +// TestSSEFrameWriter_CRLFBoundary 上游用 CRLF 行尾时同样按帧冲刷。 +// 旧实现只认 len(line)==1 && line[0]=='\n'，CRLF 上游会导致整条流从不 Flush。 +func TestSSEFrameWriter_CRLFBoundary(t *testing.T) { + rec := httptest.NewRecorder() + frame := NewSSEFrameWriter(NewSafeWriter(rec)) + + _ = frame.WriteLine("event: ping\r\n") + if rec.Body.Len() != 0 { + t.Errorf("半帧提前出网: %q", rec.Body.String()) + } + _ = frame.WriteLine("data: {}\r\n") + _ = frame.WriteLine("\r\n") + + if frame.Buffered() != 0 { + t.Errorf("CRLF 空行未被识别为帧边界，仍有 %d 字节滞留", frame.Buffered()) + } + if got := rec.Body.String(); got != "event: ping\r\ndata: {}\r\n\r\n" { + t.Errorf("Body = %q", got) + } +} + +// TestSSEFrameWriter_FlushPartial 上游未以空行结尾时，残留必须冲刷，不能丢帧。 +func TestSSEFrameWriter_FlushPartial(t *testing.T) { + rec := httptest.NewRecorder() + frame := NewSSEFrameWriter(NewSafeWriter(rec)) + + _ = frame.WriteLine("data: {\"last\":true}\n") + if rec.Body.Len() != 0 { + t.Fatalf("无空行结尾的帧不应提前出网: %q", rec.Body.String()) + } + if err := frame.FlushPartial(); err != nil { + t.Fatalf("FlushPartial failed: %v", err) + } + if got := rec.Body.String(); got != "data: {\"last\":true}\n" { + t.Errorf("Body = %q, 最后一帧丢失", got) + } + // 空缓冲重复调用应为 no-op + if err := frame.FlushPartial(); err != nil { + t.Errorf("空缓冲 FlushPartial 应为 no-op, got %v", err) + } +} + +// TestSSEFrameWriter_Discard 中断时丢弃半帧，不把不完整的帧写给客户端。 +func TestSSEFrameWriter_Discard(t *testing.T) { + rec := httptest.NewRecorder() + frame := NewSSEFrameWriter(NewSafeWriter(rec)) + + _ = frame.WriteLine("event: content_block_delta\n") + frame.Discard() + if frame.Buffered() != 0 { + t.Errorf("Discard 后 Buffered() = %d, want 0", frame.Buffered()) + } + if err := frame.FlushPartial(); err != nil { + t.Fatalf("FlushPartial failed: %v", err) + } + if rec.Body.Len() != 0 { + t.Errorf("被丢弃的半帧仍然出网: %q", rec.Body.String()) + } +} + +// TestSSEFrameWriter_MaxBufDegrades 上游异常（始终不发空行）时按上限降级冲刷， +// 不把整条流攒进内存。 +func TestSSEFrameWriter_MaxBufDegrades(t *testing.T) { + rec := httptest.NewRecorder() + frame := NewSSEFrameWriter(NewSafeWriter(rec)) + + line := "data: " + strings.Repeat("x", 64*1024) + "\n" + for i := 0; i < sseFrameMaxBuf/len(line)+2; i++ { + if err := frame.WriteLine(line); err != nil { + t.Fatalf("WriteLine failed: %v", err) + } + } + if frame.Buffered() >= sseFrameMaxBuf { + t.Errorf("Buffered() = %d, 超过上限仍未降级冲刷", frame.Buffered()) + } + if rec.Body.Len() == 0 { + t.Error("超过上限后仍未向客户端写出任何数据") + } +} + +// TestSSEFrameWriter_WriteErrorPropagates 写客户端失败时错误必须上报， +// 供调用方走流中断结算（而非静默丢数据按成功结算）。 +func TestSSEFrameWriter_WriteErrorPropagates(t *testing.T) { + fw := &failingWriter{err: errors.New("connection reset by peer")} + frame := NewSSEFrameWriter(fw) + + if err := frame.WriteLine("event: x\n"); err != nil { + t.Fatalf("半帧不触发写入，不应报错: %v", err) + } + err := frame.WriteLine("\n") + if err == nil || !strings.Contains(err.Error(), "connection reset") { + t.Fatalf("WriteLine 帧结束时错误未上报: %v", err) + } + // 失败后缓冲已清空，不会重复写出 + if frame.Buffered() != 0 { + t.Errorf("写失败后 Buffered() = %d, want 0", frame.Buffered()) + } +} + +// failingWriter 写入必定失败的 ResponseWriter，用于验证错误传播 +type failingWriter struct { + err error +} + +func (f *failingWriter) Header() http.Header { return http.Header{} } +func (f *failingWriter) WriteHeader(int) {} +func (f *failingWriter) Write([]byte) (int, error) { return 0, f.err } + func TestExtractSSEData(t *testing.T) { tests := []struct { name string