okooooooo/team-api
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(model):优化响应头

Browse Source
This commit is contained in:
chengqian committed 2026-08-20 15:19:27 +08:00
1 parent 6a49559147
commit 7d82cb0cc1
5 files changed
+589 -132
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

internal/middleware/request_id.go
+23 -10
View File
@@ -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()
}
internal/middleware/request_id_security_test.go
+159
View File
@@ -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")
}
}
})
}
internal/middleware/request_id_test.go
+167 -40
View File
@@ -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)
})
}
relay/helper/stream.go
+22 -5
View File
@@ -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 {
relay/helper/stream_test.go
+218 -77
View File
@@ -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)
}
})
}
}