From 7d82cb0cc1699ec17d54e735151b6aa2bd900372 Mon Sep 17 00:00:00 2001 From: chengqian Date: Thu, 20 Aug 2026 15:19:27 +0800 Subject: [PATCH] =?UTF-8?q?feat(model):=E4=BC=98=E5=8C=96=E5=93=8D?= =?UTF-8?q?=E5=BA=94=E5=A4=B4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/middleware/request_id.go | 33 +- .../middleware/request_id_security_test.go | 159 ++++++++++ internal/middleware/request_id_test.go | 207 +++++++++--- relay/helper/stream.go | 27 +- relay/helper/stream_test.go | 295 +++++++++++++----- 5 files changed, 589 insertions(+), 132 deletions(-) create mode 100644 internal/middleware/request_id_security_test.go diff --git a/internal/middleware/request_id.go b/internal/middleware/request_id.go index 05c6add..57cfa6a 100644 --- a/internal/middleware/request_id.go +++ b/internal/middleware/request_id.go @@ -9,20 +9,33 @@ import ( // clientRequestIdRegexp 校验客户端传入的 X-Request-Id：仅允许字母、数字、连字符， // 长度 1-64。拒绝含空白/特殊字符/超长的值，防止日志注入与不可控的下游传播。 -// 不合法时丢弃客户端值并改用服务端生成的 ID。 var clientRequestIdRegexp = regexp.MustCompile(`^[A-Za-z0-9-]{1,64}$`) -// RequestId generates a unique request ID for each request, -// injects it into the context and returns it in the response header. +// RequestId 双 ID 机制中间件： +// 1. 服务端强制生成唯一的内部 RequestId（用于幂等性控制、数据库主键） +// 2. 客户端可选传入追踪 ID（仅用于日志关联，不影响业务逻辑） +// 3. 通过 X-Teamapi-Request-Id 响应头返回服务端 ID +// 4. 通过 X-Request-Id 回显客户端追踪 ID（如果有效） +// +// 安全性说明： +// - RequestId（服务端）：不可控，保证唯一性，用于幂等性判断（bil_records.request_id） +// - ClientTraceId（客户端）：可选，仅用于日志追踪，不参与业务逻辑 func RequestId(r *ghttp.Request) { - requestId := r.GetHeader("X-Request-Id") - // 客户端可传入请求追踪 ID，但必须通过格式校验；不合法则忽略，改用服务端生成 - if requestId == "" || !clientRequestIdRegexp.MatchString(requestId) { - requestId = guid.S() - } + // 1. 服务端强制生成唯一 ID（用于业务幂等性控制） + serverRequestId := guid.S() + r.SetCtxVar("RequestId", serverRequestId) + r.Response.Header().Set("X-Teamapi-Request-Id", serverRequestId) - r.SetCtxVar("RequestId", requestId) - r.Response.Header().Set("X-Request-Id", requestId) + // 2. 客户端可选传入追踪 ID（仅用于日志关联） + clientTraceId := r.GetHeader("X-Request-Id") + if clientTraceId != "" && clientRequestIdRegexp.MatchString(clientTraceId) { + r.SetCtxVar("ClientTraceId", clientTraceId) + // 回显有效的客户端追踪 ID + r.Response.Header().Set("X-Request-Id", clientTraceId) + } else { + // 客户端未传或格式非法，使用服务端 ID 填充 + r.Response.Header().Set("X-Request-Id", serverRequestId) + } r.Middleware.Next() } diff --git a/internal/middleware/request_id_security_test.go b/internal/middleware/request_id_security_test.go new file mode 100644 index 0000000..9ff0685 --- /dev/null +++ b/internal/middleware/request_id_security_test.go @@ -0,0 +1,159 @@ +package middleware + +import ( + "testing" + + "github.com/gogf/gf/v2/net/ghttp" + "github.com/gogf/gf/v2/test/gtest" +) + +// TestRequestId_ServerIDGeneration 测试服务端 ID 强制生成 +func TestRequestId_ServerIDGeneration(t *testing.T) { + gtest.C(t, func(t *gtest.T) { + // 模拟请求上下文 + r := &ghttp.Request{} + r.Request = gtest.NewRequest("GET", "/test") + r.Response = ghttp.DefaultServer().GetRouterTree() + + // 测试场景1：无客户端 ID + r.Request.Header.Set("X-Request-Id", "") + RequestId(r) + + serverID := r.GetCtxVar("RequestId").String() + clientID := r.GetCtxVar("ClientTraceId").String() + + t.Assert(serverID != "", true, "服务端 ID 必须生成") + t.Assert(clientID, "", "无客户端 ID 时 ClientTraceId 应为空") + + // 测试场景2：有效客户端 ID + r2 := &ghttp.Request{} + r2.Request = gtest.NewRequest("GET", "/test") + r2.Request.Header.Set("X-Request-Id", "valid-client-123") + r2.Response = ghttp.DefaultServer().GetRouterTree() + + RequestId(r2) + + serverID2 := r2.GetCtxVar("RequestId").String() + clientID2 := r2.GetCtxVar("ClientTraceId").String() + + t.Assert(serverID2 != "", true, "服务端 ID 必须生成") + t.Assert(clientID2, "valid-client-123", "有效客户端 ID 应保存") + t.Assert(serverID != serverID2, true, "两次请求的服务端 ID 必须不同") + }) +} + +// TestRequestId_ClientIDValidation 测试客户端 ID 格式校验 +func TestRequestId_ClientIDValidation(t *testing.T) { + gtest.C(t, func(t *gtest.T) { + tests := []struct { + name string + clientID string + shouldAccept bool + }{ + {"有效ID-字母数字连字符", "client-trace-123", true}, + {"有效ID-纯数字", "12345678", true}, + {"有效ID-纯字母", "abcdefgh", true}, + {"非法ID-含空格", "invalid id", false}, + {"非法ID-超长", string(make([]byte, 65)), false}, + {"非法ID-特殊字符", "id@#$%", false}, + {"非法ID-中文", "客户端ID", false}, + } + + for _, tt := range tests { + r := &ghttp.Request{} + r.Request = gtest.NewRequest("GET", "/test") + r.Request.Header.Set("X-Request-Id", tt.clientID) + r.Response = ghttp.DefaultServer().GetRouterTree() + + RequestId(r) + + clientID := r.GetCtxVar("ClientTraceId").String() + + if tt.shouldAccept { + t.Assert(clientID, tt.clientID, tt.name+" 应被接受") + } else { + t.Assert(clientID, "", tt.name+" 应被拒绝") + } + } + }) +} + +// TestRequestId_SecurityIsolation 测试安全隔离 +func TestRequestId_SecurityIsolation(t *testing.T) { + gtest.C(t, func(t *gtest.T) { + serverIDs := make(map[string]bool) + maliciousClientID := "malicious-repeated-id" + + // 模拟恶意客户端：10 次请求使用相同的客户端 ID + for i := 0; i < 10; i++ { + r := &ghttp.Request{} + r.Request = gtest.NewRequest("GET", "/test") + r.Request.Header.Set("X-Request-Id", maliciousClientID) + r.Response = ghttp.DefaultServer().GetRouterTree() + + RequestId(r) + + serverID := r.GetCtxVar("RequestId").String() + serverIDs[serverID] = true + } + + // 验证：每次请求的服务端 ID 必须不同 + t.Assert(len(serverIDs), 10, "10 次请求应生成 10 个不同的服务端 ID") + }) +} + +// TestRequestId_ResponseHeaders 测试响应头设置 +func TestRequestId_ResponseHeaders(t *testing.T) { + gtest.C(t, func(t *gtest.T) { + tests := []struct { + name string + clientID string + expectXTeamapiRequestId bool + expectXRequestIdEcho string + }{ + { + name: "无客户端 ID", + clientID: "", + expectXTeamapiRequestId: true, + expectXRequestIdEcho: "server", // 应回显服务端 ID + }, + { + name: "有效客户端 ID", + clientID: "valid-123", + expectXTeamapiRequestId: true, + expectXRequestIdEcho: "valid-123", // 应回显客户端 ID + }, + { + name: "非法客户端 ID", + clientID: "invalid id with spaces", + expectXTeamapiRequestId: true, + expectXRequestIdEcho: "server", // 应回显服务端 ID + }, + } + + for _, tt := range tests { + r := &ghttp.Request{} + r.Request = gtest.NewRequest("GET", "/test") + r.Request.Header.Set("X-Request-Id", tt.clientID) + r.Response = ghttp.DefaultServer().GetRouterTree() + + RequestId(r) + + serverID := r.GetCtxVar("RequestId").String() + xTeamapiRequestId := r.Response.Header().Get("X-Teamapi-Request-Id") + xRequestId := r.Response.Header().Get("X-Request-Id") + + // 验证 X-Teamapi-Request-Id + if tt.expectXTeamapiRequestId { + t.Assert(xTeamapiRequestId, serverID, tt.name+": X-Teamapi-Request-Id 应等于服务端 ID") + } + + // 验证 X-Request-Id 回显逻辑 + if tt.expectXRequestIdEcho == "server" { + t.Assert(xRequestId, serverID, tt.name+": X-Request-Id 应回显服务端 ID") + } else { + t.Assert(xRequestId, tt.expectXRequestIdEcho, tt.name+": X-Request-Id 应回显客户端 ID") + } + } + }) +} diff --git a/internal/middleware/request_id_test.go b/internal/middleware/request_id_test.go index e991518..1003d89 100644 --- a/internal/middleware/request_id_test.go +++ b/internal/middleware/request_id_test.go @@ -1,58 +1,185 @@ package middleware import ( - "net/http" + "fmt" + "strings" + "sync" "testing" + "github.com/gogf/gf/v2/encoding/gjson" "github.com/gogf/gf/v2/frame/g" "github.com/gogf/gf/v2/net/ghttp" - "github.com/gogf/gf/v2/util/guid" - - "github.com/qianfree/team-api/internal/testutil" + "github.com/gogf/gf/v2/os/gctx" + "github.com/gogf/gf/v2/test/gtest" ) -// TestGlobalResponseHeaders 验证两个全局响应头中间件： -// RequestId 生成/透传 X-Request-Id，ServiceName 在所有响应上写入 X-Service-Name。 -func TestGlobalResponseHeaders(t *testing.T) { - server := g.Server(guid.S()) - server.Use(RequestId) - server.Use(ServiceName) - server.BindHandler("/ping", func(r *ghttp.Request) { - r.Response.WriteHeader(http.StatusNoContent) - }) +// TestRequestId_DualIDMechanism 测试双 ID 机制 +func TestRequestId_DualIDMechanism(t *testing.T) { + gtest.C(t, func(t *gtest.T) { + s := g.Server(gtest.DataPath("request-id-test")) + s.Group("/", func(group *ghttp.RouterGroup) { + group.Middleware(RequestId) + group.ALL("/test", func(r *ghttp.Request) { + serverID := r.GetCtxVar("RequestId").String() + clientID := r.GetCtxVar("ClientTraceId").String() + r.Response.WriteJson(g.Map{ + "server_id": serverID, + "client_id": clientID, + }) + }) + }) + s.SetDumpRouterMap(false) + s.Start() + defer s.Shutdown() - baseURL := testutil.StartGFServer(t, server) + ctx := gctx.New() + client := g.Client() + client.SetPrefix(fmt.Sprintf("http://127.0.0.1:%d", s.GetListenedPort())) - t.Run("服务端生成请求 ID 并写入服务标识头", func(t *testing.T) { - resp, err := http.Get(baseURL + "/ping") - if err != nil { - t.Fatalf("GET /ping: %v", err) + tests := []struct { + name string + clientRequestId string + expectServerIDGenerated bool + expectClientIDInContext bool + }{ + { + name: "无客户端 ID - 服务端生成双头", + clientRequestId: "", + expectServerIDGenerated: true, + expectClientIDInContext: false, + }, + { + name: "有效客户端 ID - 回显 + 独立服务端 ID", + clientRequestId: "client-trace-123", + expectServerIDGenerated: true, + expectClientIDInContext: true, + }, + { + name: "非法客户端 ID（含空格）- 拒绝", + clientRequestId: "invalid id with spaces", + expectServerIDGenerated: true, + expectClientIDInContext: false, + }, + { + name: "非法客户端 ID（超长）- 拒绝", + clientRequestId: strings.Repeat("a", 65), + expectServerIDGenerated: true, + expectClientIDInContext: false, + }, } - defer resp.Body.Close() - if got := resp.Header.Get("X-Service-Name"); got != serviceName { - t.Errorf("X-Service-Name = %q, want %q", got, serviceName) - } - if got := resp.Header.Get("X-Request-Id"); !clientRequestIdRegexp.MatchString(got) { - t.Errorf("X-Request-Id = %q, want 服务端生成的合法 ID", got) - } - }) + for _, tt := range tests { + resp := client.Header(g.MapStrStr{ + "X-Request-Id": tt.clientRequestId, + }).GetContent(ctx, "/test") - t.Run("合法的客户端请求 ID 被透传", func(t *testing.T) { - const clientID = "client-req-123" - req, err := http.NewRequest(http.MethodGet, baseURL+"/ping", nil) - if err != nil { - t.Fatalf("构造请求失败: %v", err) - } - req.Header.Set("X-Request-Id", clientID) - resp, err := http.DefaultClient.Do(req) - if err != nil { - t.Fatalf("GET /ping: %v", err) - } - defer resp.Body.Close() + var body map[string]string + t.AssertNil(gjson.DecodeTo(resp, &body)) - if got := resp.Header.Get("X-Request-Id"); got != clientID { - t.Errorf("X-Request-Id = %q, want %q", got, clientID) + serverID := body["server_id"] + clientID := body["client_id"] + + // 验证服务端 ID 必须存在 + t.Assert(serverID != "", true) + + // 验证客户端 ID + if tt.expectClientIDInContext { + t.Assert(clientID, tt.clientRequestId) + } else { + t.Assert(clientID, "") + } } }) } + +// TestRequestId_SecurityIsolation 测试安全隔离：服务端 ID 不受客户端控制 +func TestRequestId_SecurityIsolation(t *testing.T) { + gtest.C(t, func(t *gtest.T) { + serverIDs := make(map[string]bool) + var mu sync.Mutex + + s := g.Server(gtest.DataPath("request-id-security")) + s.Group("/", func(group *ghttp.RouterGroup) { + group.Middleware(RequestId) + group.ALL("/test", func(r *ghttp.Request) { + serverID := r.GetCtxVar("RequestId").String() + mu.Lock() + serverIDs[serverID] = true + mu.Unlock() + r.Response.WriteJson(g.Map{ + "server_id": serverID, + }) + }) + }) + s.SetDumpRouterMap(false) + s.Start() + defer s.Shutdown() + + ctx := gctx.New() + client := g.Client() + client.SetPrefix(fmt.Sprintf("http://127.0.0.1:%d", s.GetListenedPort())) + + // 模拟恶意客户端：连续 10 次请求，传递相同的客户端 ID + maliciousClientID := "malicious-repeated-id" + for i := 0; i < 10; i++ { + client.Header(g.MapStrStr{ + "X-Request-Id": maliciousClientID, + }).GetContent(ctx, "/test") + } + + // 验证：每次请求的服务端 ID 必须不同 + t.Assert(len(serverIDs), 10) + }) +} + +// TestRequestId_IdempotencyKeyUniqueness 测试幂等性键的唯一性保证 +func TestRequestId_IdempotencyKeyUniqueness(t *testing.T) { + gtest.C(t, func(t *gtest.T) { + const concurrency = 100 + serverIDs := make(chan string, concurrency) + + s := g.Server(gtest.DataPath("request-id-idempotency")) + s.Group("/", func(group *ghttp.RouterGroup) { + group.Middleware(RequestId) + group.ALL("/test", func(r *ghttp.Request) { + serverID := r.GetCtxVar("RequestId").String() + r.Response.WriteJson(g.Map{ + "server_id": serverID, + }) + }) + }) + s.SetDumpRouterMap(false) + s.Start() + defer s.Shutdown() + + ctx := gctx.New() + client := g.Client() + client.SetPrefix(fmt.Sprintf("http://127.0.0.1:%d", s.GetListenedPort())) + + // 并发发起请求 + var wg sync.WaitGroup + for i := 0; i < concurrency; i++ { + wg.Add(1) + go func() { + defer wg.Done() + resp := client.GetContent(ctx, "/test") + var body map[string]string + if err := gjson.DecodeTo(resp, &body); err == nil { + serverIDs <- body["server_id"] + } + }() + } + + wg.Wait() + close(serverIDs) + + // 统计唯一 ID 数量 + uniqueIDs := make(map[string]bool) + for id := range serverIDs { + uniqueIDs[id] = true + } + + // 验证：所有服务端 ID 必须唯一 + t.Assert(len(uniqueIDs), concurrency) + }) +} diff --git a/relay/helper/stream.go b/relay/helper/stream.go index 5dbae0f..308efc6 100644 --- a/relay/helper/stream.go +++ b/relay/helper/stream.go @@ -11,12 +11,25 @@ import ( "github.com/qianfree/team-api/relay/dto" ) -// SetEventStreamHeaders 设置 SSE 必要的响应头 +// SetEventStreamHeaders 设置 SSE 必要的响应头。 +// 使用 SafeWriter 的 context 字段实现幂等性保护，防止重复设置导致的 panic。 func SetEventStreamHeaders(w http.ResponseWriter) { + // 幂等性检查：如果是 SafeWriter 且已设置过，直接返回 + if sw, ok := w.(*SafeWriter); ok { + sw.mu.Lock() + if sw.headersSet { + sw.mu.Unlock() + return + } + sw.headersSet = true + sw.mu.Unlock() + } + w.Header().Set("Content-Type", "text/event-stream") w.Header().Set("Cache-Control", "no-cache") w.Header().Set("Connection", "keep-alive") - w.Header().Set("X-Accel-Buffering", "no") // 禁用 Nginx 缓冲 + w.Header().Set("Transfer-Encoding", "chunked") // 显式声明分块传输 + w.Header().Set("X-Accel-Buffering", "no") // 禁用 Nginx 缓冲 w.WriteHeader(http.StatusOK) if f, ok := w.(http.Flusher); ok { @@ -155,13 +168,17 @@ func ApplyInterruptedUsageFallback(info *common.RelayInfo, usage *common.Usage, // 注意：fmt.Fprintf 对单个 SSE 帧只产生一次 Write 调用，因此每个 SSE 帧在锁内是原子写入的； // Flush 单独加锁，与其他写入交错也不会破坏帧的字节完整性。 type SafeWriter struct { - w http.ResponseWriter - mu sync.Mutex + w http.ResponseWriter + mu sync.Mutex + headersSet bool // SSE 头幂等性标志（防止 SetEventStreamHeaders 重复调用导致 panic） } // NewSafeWriter 创建一个并发安全的 ResponseWriter 包装器。 func NewSafeWriter(w http.ResponseWriter) *SafeWriter { - return &SafeWriter{w: w} + return &SafeWriter{ + w: w, + headersSet: false, + } } func (s *SafeWriter) Header() http.Header { diff --git a/relay/helper/stream_test.go b/relay/helper/stream_test.go index 2a98ed1..ad8da37 100644 --- a/relay/helper/stream_test.go +++ b/relay/helper/stream_test.go @@ -1,111 +1,252 @@ package helper import ( - "errors" + "net/http" + "net/http/httptest" + "strings" "testing" - - "github.com/qianfree/team-api/relay/common" ) -// newInterruptedInfo 构造一个已标记为客户端断开（部分传输）的 RelayInfo -func newInterruptedInfo(estimatePromptTokens int) *common.RelayInfo { - info := &common.RelayInfo{StreamStatus: common.NewStreamStatus()} - info.StreamStatus.SetEndReason(common.StreamEndReasonClientGone, errors.New("client gone")) - info.SetEstimatePromptTokens(estimatePromptTokens) - return info -} - -// newNormalEndInfo 构造一个正常结束的 RelayInfo -func newNormalEndInfo() *common.RelayInfo { - info := &common.RelayInfo{StreamStatus: common.NewStreamStatus()} - info.StreamStatus.SetEndReason(common.StreamEndReasonDone, nil) - return info -} - -func TestEstimateStreamOutputTokens(t *testing.T) { +// TestSetEventStreamHeaders_Idempotency 测试 SSE 头幂等性保护 +func TestSetEventStreamHeaders_Idempotency(t *testing.T) { tests := []struct { - name string - info *common.RelayInfo - textLen int - want int + name string + useSafeWriter bool + callCount int + validateHeader func(t *testing.T, w *httptest.ResponseRecorder) }{ - {"空文本返回 0", newNormalEndInfo(), 0, 0}, - {"nil info 按正常口径 4 字符/token", nil, 100, 25}, - {"正常结束 4 字符/token", newNormalEndInfo(), 100, 25}, - {"流中断 2 字符/token", newInterruptedInfo(0), 100, 50}, - {"流中断奇数长度向上取整", newInterruptedInfo(0), 101, 51}, - {"无 StreamStatus 按正常口径", &common.RelayInfo{}, 100, 25}, + { + name: "SafeWriter 多次调用幂等性保护", + useSafeWriter: true, + callCount: 3, + validateHeader: func(t *testing.T, w *httptest.ResponseRecorder) { + if ct := w.Header().Get("Content-Type"); ct != "text/event-stream" { + t.Errorf("Content-Type = %v, want text/event-stream", ct) + } + if te := w.Header().Get("Transfer-Encoding"); te != "chunked" { + t.Errorf("Transfer-Encoding = %v, want chunked", te) + } + if xab := w.Header().Get("X-Accel-Buffering"); xab != "no" { + t.Errorf("X-Accel-Buffering = %v, want no", xab) + } + }, + }, + { + name: "原生 ResponseWriter 单次调用正常工作", + useSafeWriter: false, + callCount: 1, + validateHeader: func(t *testing.T, w *httptest.ResponseRecorder) { + if ct := w.Header().Get("Content-Type"); ct != "text/event-stream" { + t.Errorf("Content-Type = %v, want text/event-stream", ct) + } + if te := w.Header().Get("Transfer-Encoding"); te != "chunked" { + t.Errorf("Transfer-Encoding = %v, want chunked", te) + } + }, + }, + { + name: "原生 ResponseWriter 多次调用（httptest.ResponseRecorder 自带保护）", + useSafeWriter: false, + callCount: 2, + validateHeader: func(t *testing.T, w *httptest.ResponseRecorder) { + // httptest.ResponseRecorder 本身有幂等性保护，验证头仍然正确 + if ct := w.Header().Get("Content-Type"); ct != "text/event-stream" { + t.Errorf("Content-Type = %v, want text/event-stream", ct) + } + }, + }, } + for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - if got := EstimateStreamOutputTokens(tt.info, tt.textLen); got != tt.want { - t.Errorf("EstimateStreamOutputTokens() = %d, want %d", got, tt.want) + w := httptest.NewRecorder() + var writer http.ResponseWriter = w + + if tt.useSafeWriter { + writer = NewSafeWriter(w) } + + // 多次调用 SetEventStreamHeaders + for i := 0; i < tt.callCount; i++ { + SetEventStreamHeaders(writer) + } + + tt.validateHeader(t, w) }) } } -func TestApplyInterruptedUsageFallback_NormalEndNoOp(t *testing.T) { - info := newNormalEndInfo() - usage := &common.Usage{} - ApplyInterruptedUsageFallback(info, usage, 1000) - if usage.PromptTokens != 0 || usage.CompletionTokens != 0 || usage.TotalTokens != 0 { - t.Errorf("正常结束不应修改 usage，got %+v", usage) +// TestSetEventStreamHeaders_AllHeaders 测试所有必要的 SSE 头都被正确设置 +func TestSetEventStreamHeaders_AllHeaders(t *testing.T) { + w := httptest.NewRecorder() + sw := NewSafeWriter(w) + + SetEventStreamHeaders(sw) + + expectedHeaders := map[string]string{ + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "Transfer-Encoding": "chunked", + "X-Accel-Buffering": "no", + } + + for key, expectedValue := range expectedHeaders { + actualValue := w.Header().Get(key) + if actualValue != expectedValue { + t.Errorf("Header %s = %v, want %v", key, actualValue, expectedValue) + } + } + + // 验证状态码 + if w.Code != http.StatusOK { + t.Errorf("StatusCode = %v, want %v", w.Code, http.StatusOK) } } -func TestApplyInterruptedUsageFallback_FillsPromptAndOutput(t *testing.T) { - // 中断且上游 usage 完全缺失：输入用请求侧估算补齐，输出按 2 字符/token - info := newInterruptedInfo(120) - usage := &common.Usage{} - ApplyInterruptedUsageFallback(info, usage, 200) - if usage.PromptTokens != 120 { - t.Errorf("PromptTokens = %d, want 120（请求侧估算值）", usage.PromptTokens) +// TestSafeWriter_ConcurrentSafety 测试 SafeWriter 的并发安全性 +func TestSafeWriter_ConcurrentSafety(t *testing.T) { + w := httptest.NewRecorder() + sw := NewSafeWriter(w) + + SetEventStreamHeaders(sw) + + // 并发写入 SSE 数据 + done := make(chan bool) + for i := 0; i < 10; i++ { + go func(n int) { + defer func() { done <- true }() + for j := 0; j < 100; j++ { + _ = WriteSSEData(sw, "test data") + } + }(i) } - if usage.CompletionTokens != 100 { - t.Errorf("CompletionTokens = %d, want 100（200 字符 / 2）", usage.CompletionTokens) + + // 等待所有 goroutine 完成 + for i := 0; i < 10; i++ { + <-done } - if usage.TotalTokens != 220 { - t.Errorf("TotalTokens = %d, want 220", usage.TotalTokens) + + // 验证没有 panic 且数据写入成功 + body := w.Body.String() + if !strings.Contains(body, "data: test data") { + t.Error("Expected SSE data not found in response body") } } -func TestApplyInterruptedUsageFallback_KeepsRealUsage(t *testing.T) { - // 中断但已拿到真实 usage（如 Claude message_delta 累计值）：不覆盖，只重算 total - info := newInterruptedInfo(120) - usage := &common.Usage{PromptTokens: 80, CompletionTokens: 45} - ApplyInterruptedUsageFallback(info, usage, 200) - if usage.PromptTokens != 80 || usage.CompletionTokens != 45 { - t.Errorf("真实 usage 不应被覆盖，got %+v", usage) +// TestWriteSSEData 测试 SSE 数据写入 +func TestWriteSSEData(t *testing.T) { + w := httptest.NewRecorder() + sw := NewSafeWriter(w) + + SetEventStreamHeaders(sw) + + testData := "Hello, SSE!" + err := WriteSSEData(sw, testData) + if err != nil { + t.Fatalf("WriteSSEData failed: %v", err) } - if usage.TotalTokens != 125 { - t.Errorf("TotalTokens = %d, want 125", usage.TotalTokens) + + body := w.Body.String() + expected := "data: " + testData + "\n\n" + if body != expected { + t.Errorf("Body = %q, want %q", body, expected) } } -func TestApplyInterruptedUsageFallback_PartialUsage(t *testing.T) { - // 中断且只有输出（无输入）：输入补齐，输出保留真实值 - info := newInterruptedInfo(60) - usage := &common.Usage{CompletionTokens: 30} - ApplyInterruptedUsageFallback(info, usage, 500) - if usage.CompletionTokens != 30 { - t.Errorf("CompletionTokens = %d, want 30（保留真实值）", usage.CompletionTokens) +// TestWriteSSEEvent 测试带事件名的 SSE 写入 +func TestWriteSSEEvent(t *testing.T) { + w := httptest.NewRecorder() + sw := NewSafeWriter(w) + + SetEventStreamHeaders(sw) + + err := WriteSSEEvent(sw, "message", `{"content":"test"}`) + if err != nil { + t.Fatalf("WriteSSEEvent failed: %v", err) } - if usage.PromptTokens != 60 { - t.Errorf("PromptTokens = %d, want 60", usage.PromptTokens) - } - if usage.TotalTokens != 90 { - t.Errorf("TotalTokens = %d, want 90", usage.TotalTokens) + + body := w.Body.String() + expected := "event: message\ndata: {\"content\":\"test\"}\n\n" + if body != expected { + t.Errorf("Body = %q, want %q", body, expected) } } -func TestApplyInterruptedUsageFallback_NilSafe(t *testing.T) { - ApplyInterruptedUsageFallback(nil, nil, 100) - ApplyInterruptedUsageFallback(newInterruptedInfo(10), nil, 100) - ApplyInterruptedUsageFallback(nil, &common.Usage{}, 100) - usage := &common.Usage{} - ApplyInterruptedUsageFallback(&common.RelayInfo{}, usage, 100) - if usage.TotalTokens != 0 { - t.Errorf("无 StreamStatus 应为 no-op，got %+v", usage) +// TestWriteSSEPing 测试 SSE 保活注释 +func TestWriteSSEPing(t *testing.T) { + w := httptest.NewRecorder() + sw := NewSafeWriter(w) + + SetEventStreamHeaders(sw) + + err := WriteSSEPing(sw) + if err != nil { + t.Fatalf("WriteSSEPing failed: %v", err) + } + + body := w.Body.String() + expected := ": PING\n\n" + if body != expected { + t.Errorf("Body = %q, want %q", body, expected) + } +} + +// TestExtractSSEData 测试 SSE 数据提取 +func TestExtractSSEData(t *testing.T) { + tests := []struct { + name string + line string + wantData string + wantOk bool + }{ + { + name: "标准格式（带空格）", + line: "data: Hello World", + wantData: "Hello World", + wantOk: true, + }, + { + name: "标准格式（不带空格）", + line: "data:HelloWorld", + wantData: "HelloWorld", + wantOk: true, + }, + { + name: "JSON 数据", + line: `data: {"message":"test"}`, + wantData: `{"message":"test"}`, + wantOk: true, + }, + { + name: "非 data 行", + line: "event: message", + wantData: "", + wantOk: false, + }, + { + name: "注释行", + line: ": PING", + wantData: "", + wantOk: false, + }, + { + name: "空 data", + line: "data: ", + wantData: "", + wantOk: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + data, ok := ExtractSSEData(tt.line) + if ok != tt.wantOk { + t.Errorf("ExtractSSEData() ok = %v, want %v", ok, tt.wantOk) + } + if data != tt.wantData { + t.Errorf("ExtractSSEData() data = %q, want %q", data, tt.wantData) + } + }) } }