fix(relay):修复转换层层数据断层导致连接中断没有正确计费的问题

This commit is contained in:
qian committed 2026-08-31 22:19:35 +08:00
1 parent bdfd1b6567
commit a2beb8230a
14 files changed
+295 -72

No files matched your search

+5 -1
View File
@@ -93,8 +93,12 @@ func (a *Adaptor) handleStreamToOpenAI(ctx context.Context, resp *http.Response,
return nil, constant.NewUpstreamError(resp.StatusCode, string(body), nil).WithRetryAfter(constant.RetryAfterFromHeader(resp.Header))
}
// relaykit 流式转换(特性开关控制,默认关闭)。未启用/无匹配回退旧路径。
// relaykit 流式转换(常开)。同格式/无匹配转换器时回退旧路径。
if usage, ok := relaykit_bridge.TryConvertStreamViaRelaykit(ctx, info, resp.Body, writer); ok {
// 流中断(客户端断开/写失败):桥接层已完成 usage 兜底,透传中断信号供上层按中断结算
if info.StreamStatus != nil && info.StreamStatus.IsPartialStreamEnd() {
return usage, common.ErrStreamInterrupted
}
return usage, nil
}
+7 -3
View File
@@ -133,8 +133,12 @@ func (a *Adaptor) DoResponse(ctx context.Context, resp *http.Response, info *com
// handleStreamResponse 流式模式:逐事件读取 Coze SSE,转换为 OpenAI SSE 格式输出
func (a *Adaptor) handleStreamResponse(ctx context.Context, resp *http.Response, info *common.RelayInfo, writer http.ResponseWriter) (*common.Usage, error) {
// relaykit 流式转换(特性开关控制,默认关闭)。未启用/无匹配回退旧路径。
// relaykit 流式转换(常开)。同格式/无匹配转换器时回退旧路径。
if usage, ok := relaykit_bridge.TryConvertStreamViaRelaykit(ctx, info, resp.Body, writer); ok {
// 流中断(客户端断开/写失败):桥接层已完成 usage 兜底,透传中断信号供上层按中断结算
if info.StreamStatus != nil && info.StreamStatus.IsPartialStreamEnd() {
return usage, common.ErrStreamInterrupted
}
return usage, nil
}
@@ -164,7 +168,7 @@ func (a *Adaptor) handleStreamResponse(ctx context.Context, resp *http.Response,
// 流中断计费兜底:输出按已转发文本 2 字符/token 估算,输入用请求侧估算值补齐
interruptedUsage := &common.Usage{}
helper.ApplyInterruptedUsageFallback(info, interruptedUsage, transferredTextLen)
return interruptedUsage, nil
return interruptedUsage, common.ErrStreamInterrupted
default:
}
@@ -207,7 +211,7 @@ func (a *Adaptor) handleStreamResponse(ctx context.Context, resp *http.Response,
// 流中断计费兜底:输出按已转发文本 2 字符/token 估算,输入用请求侧估算值补齐
interruptedUsage := &common.Usage{}
helper.ApplyInterruptedUsageFallback(info, interruptedUsage, transferredTextLen)
return interruptedUsage, nil
return interruptedUsage, common.ErrStreamInterrupted
}
case "conversation.message.completed":
+7 -3
View File
@@ -199,8 +199,12 @@ func (a *Adaptor) handleNonStreamResponse(ctx context.Context, resp *http.Respon
// handleStreamResponse 处理 Dify streaming 模式 SSE 响应
func (a *Adaptor) handleStreamResponse(ctx context.Context, resp *http.Response, info *common.RelayInfo, writer http.ResponseWriter) (*common.Usage, error) {
// relaykit 流式转换(特性开关控制,默认关闭)。未启用/无匹配回退旧路径。
// relaykit 流式转换(常开)。同格式/无匹配转换器时回退旧路径。
if usage, ok := relaykit_bridge.TryConvertStreamViaRelaykit(ctx, info, resp.Body, writer); ok {
// 流中断(客户端断开/写失败):桥接层已完成 usage 兜底,透传中断信号供上层按中断结算
if info.StreamStatus != nil && info.StreamStatus.IsPartialStreamEnd() {
return usage, common.ErrStreamInterrupted
}
return usage, nil
}
@@ -228,7 +232,7 @@ func (a *Adaptor) handleStreamResponse(ctx context.Context, resp *http.Response,
streamStatus.SetEndReason(common.StreamEndReasonClientGone, common.ErrStreamInterrupted)
// 流中断计费兜底:输出缺失按已转发文本 2 字符/token 估算,输入用请求侧估算值补齐
helper.ApplyInterruptedUsageFallback(info, &usage, transferredTextLen)
return &usage, nil
return &usage, common.ErrStreamInterrupted
default:
}
@@ -267,7 +271,7 @@ func (a *Adaptor) handleStreamResponse(ctx context.Context, resp *http.Response,
streamStatus.SetEndReason(common.StreamEndReasonClientGone, common.ErrStreamInterrupted)
// 流中断计费兜底:输出缺失按已转发文本 2 字符/token 估算,输入用请求侧估算值补齐
helper.ApplyInterruptedUsageFallback(info, &usage, transferredTextLen)
return &usage, nil
return &usage, common.ErrStreamInterrupted
}
case "message_end":
+5 -1
View File
@@ -240,8 +240,12 @@ func (a *Adaptor) handleStreamToOpenAI(ctx context.Context, resp *http.Response,
return nil, buildGeminiUpstreamError(body, resp.StatusCode)
}
// relaykit 流式转换(特性开关控制,默认关闭)。未启用/无匹配回退旧路径。
// relaykit 流式转换(常开)。同格式/无匹配转换器时回退旧路径。
if usage, ok := relaykit_bridge.TryConvertStreamViaRelaykit(ctx, info, resp.Body, writer); ok {
// 流中断(客户端断开/写失败):桥接层已完成 usage 兜底,透传中断信号供上层按中断结算
if info.StreamStatus != nil && info.StreamStatus.IsPartialStreamEnd() {
return usage, common.ErrStreamInterrupted
}
return usage, nil
}
+4
View File
@@ -240,6 +240,10 @@ func (a *Adaptor) handleChatStreamResponse(ctx context.Context, resp *http.Respo
// relaykit 流式转换(仅 chat;generate/embedding 不迁移,走旧路径)
if usage, ok := relaykit_bridge.TryConvertStreamViaRelaykit(ctx, info, resp.Body, writer); ok {
// 流中断(客户端断开/写失败):桥接层已完成 usage 兜底,透传中断信号供上层按中断结算
if info.StreamStatus != nil && info.StreamStatus.IsPartialStreamEnd() {
return usage, common.ErrStreamInterrupted
}
return usage, nil
}
+4
View File
@@ -1573,6 +1573,10 @@ func HandleResponsesStreamToChat(ctx context.Context, resp *http.Response, info
helper.WriteSSEData(writer, "[DONE]")
g.Log().Infof(ctx, "[HandleResponsesStreamToChat] completed: usage=%+v, endReason=%s", totalUsage, info.StreamStatus.GetEndReason())
// 流中断(客户端断开/写失败/上游 error 主动终止):透传中断信号,上层按已传输部分结算
if info.StreamStatus != nil && info.StreamStatus.IsPartialStreamEnd() {
return &totalUsage, common.ErrStreamInterrupted
}
return &totalUsage, nil
}
+10 -4
View File
@@ -13,7 +13,7 @@ const (
StreamEndReasonClientGone StreamEndReason = "client_gone"
StreamEndReasonError StreamEndReason = "error"
StreamEndReasonPanic StreamEndReason = "panic"
StreamEndReasonHandlerStop StreamEndReason = "handler_stop" // dataHandler 调用了 Stop()
StreamEndReasonHandlerStop StreamEndReason = "handler_stop" // dataHandler 调用了 Stop()(写客户端失败/上游 error 主动终止,按中断处理)
StreamEndReasonPingFail StreamEndReason = "ping_fail" // ping 写入失败
StreamEndReasonScannerErr StreamEndReason = "scanner_err" // bufio.Scanner 错误
StreamEndReasonEOF StreamEndReason = "eof" // scanner EOF 未收到 [DONE]
@@ -71,17 +71,23 @@ func (s *StreamStatus) ErrorCount() int {
}
// IsPartialStreamEnd 流已开始传输但未正常结束,应按已消费 usage 部分结算。
// handler_stop(dataHandler 调用 Stop)纳入中断判定:正常结束由 sr.Done()/scanner 的
// done 标记,实际走到 Stop 的只有「写客户端失败」(客户端不可达)与「上游 error 事件
// 主动终止」两种,均属已向客户端写出部分字节后的非正常结束,应按中断部分结算。
// 此前 handler_stop 不算中断,客户端断开时写失败若先于 ctx 取消被观察到(first-writer-wins
// 竞态),流会被误判为正常结束、跳过中断计费兜底,导致 0 token / 0 费用结算。
func (s *StreamStatus) IsPartialStreamEnd() bool {
r := s.GetEndReason()
return r == StreamEndReasonClientGone || r == StreamEndReasonScannerErr ||
r == StreamEndReasonTimeout || r == StreamEndReasonPingFail
r == StreamEndReasonTimeout || r == StreamEndReasonPingFail ||
r == StreamEndReasonHandlerStop
}
// IsNormalEnd 是否正常结束
// IsNormalEnd 是否正常结束(handler_stop 属异常终止,不算正常结束)
func (s *StreamStatus) IsNormalEnd() bool {
s.mu.RLock()
defer s.mu.RUnlock()
return s.reason == StreamEndReasonDone || s.reason == StreamEndReasonHandlerStop || s.reason == ""
return s.reason == StreamEndReasonDone || s.reason == ""
}
// HasErrors 是否有错误
+40
View File
@@ -0,0 +1,40 @@
package common
import (
"errors"
"testing"
)
// TestIsPartialStreamEnd_ReasonCoverage 流式结束原因的中断判定覆盖。
// handler_stop(写客户端失败/上游 error 主动终止)必须算部分流中断:
// 客户端断开时写失败常先于 ctx 取消被观察到(SetEndReason 为 first-writer-wins),
// 漏判会跳过中断计费兜底,按 0 token 成功结算。
func TestIsPartialStreamEnd_ReasonCoverage(t *testing.T) {
partialReasons := []StreamEndReason{
StreamEndReasonClientGone,
StreamEndReasonScannerErr,
StreamEndReasonTimeout,
StreamEndReasonPingFail,
StreamEndReasonHandlerStop,
}
for _, reason := range partialReasons {
s := NewStreamStatus()
s.SetEndReason(reason, errors.New("boom"))
if !s.IsPartialStreamEnd() {
t.Errorf("IsPartialStreamEnd(%s) = false, want true", reason)
}
if s.IsNormalEnd() {
t.Errorf("IsNormalEnd(%s) = true, want false", reason)
}
}
// 正常结束不算中断
s := NewStreamStatus()
s.SetEndReason(StreamEndReasonDone, nil)
if s.IsPartialStreamEnd() {
t.Error("IsPartialStreamEnd(done) = true, want false")
}
if !s.IsNormalEnd() {
t.Error("IsNormalEnd(done) = false, want true")
}
}
+15 -2
View File
@@ -718,6 +718,7 @@ func RelayHandler(ctx context.Context, body []byte, path string, headers http.He
if streamUsage == nil {
streamUsage = &common.Usage{}
}
var interruptSettle *common.SettlementResult
if billing != nil && preDeductAmount > 0 {
settleResult, err := billing.SettleStreamInterrupted(settleCtx, rc.TenantID, rc.UserID, rc.ApiKeyID, selection.ChannelID,
v.modelName, rc.RequestID, v.relayModeStr, streamUsage, preDeductAmount, rc.ProjectID, info.StartTime)
@@ -734,8 +735,9 @@ func RelayHandler(ctx context.Context, body []byte, path string, headers http.He
billing.IncrMemberQuotaUsed(settleCtx, rc.TenantID, rc.UserID, settleResult.ActualCost)
billing.IncrApiKeyQuotaUsed(settleCtx, rc.ApiKeyID, settleResult.ActualCost)
}
interruptSettle = settleResult
}
recordFailedUsageWithTokens(provider, rc, selection, v.modelName, v.relayMode, v.isStream, err, streamUsage)
recordFailedUsageWithTokens(provider, rc, selection, v.modelName, v.relayMode, v.isStream, err, streamUsage, interruptSettle)
finalizeTrace(trace, rc, hop, false, attempt, selection, err.Error(), info.LatencyMs())
dbgAttempt.MarkFinal(err)
return usage, v.billingResult, err
@@ -1056,7 +1058,7 @@ func recordFailedUsage(provider common.DataProvider, rc *RelayContext, selection
// recordFailedUsageWithTokens 记录失败用量(含 token 明细,用于流中断等已有部分 usage 的场景)。
// 与 recordFailedUsage 的区别:此函数会填充 token 字段,避免报表中流中断记录的 token 全为 0。
func recordFailedUsageWithTokens(provider common.DataProvider, rc *RelayContext, selection *common.ChannelSelection, modelName string, relayMode constant.RelayMode, isStream bool, err error, usage *common.Usage) {
func recordFailedUsageWithTokens(provider common.DataProvider, rc *RelayContext, selection *common.ChannelSelection, modelName string, relayMode constant.RelayMode, isStream bool, err error, usage *common.Usage, settleResult *common.SettlementResult) {
record := &common.UsageRecord{
TenantID: rc.TenantID,
UserID: rc.UserID,
@@ -1082,6 +1084,17 @@ func recordFailedUsageWithTokens(provider common.DataProvider, rc *RelayContext,
record.CompletionTokens = usage.CompletionTokens
record.TotalTokens = usage.TotalTokens
}
// 流中断结算成功时补充费用字段:中断已按已传输部分实际扣费,用量记录须与账本一致,
// 否则后台看到的中断请求永远是 0 费用(结算与用量记录是两条独立写入链路)
if settleResult != nil {
record.TotalCost = settleResult.BaseCost
record.ActualCost = settleResult.ActualCost
record.PreDeductAmount = settleResult.PreDeductAmount
record.RefundAmount = settleResult.RefundAmount
record.SupplementAmount = settleResult.SupplementAmount
record.BillingMode = settleResult.BillingMode
record.BillingSource = settleResult.BillingSource
}
provider.RecordUsage(context.Background(), record)
}
+18 -5
View File
@@ -69,7 +69,8 @@ func convertStreamViaRelaykit(ctx context.Context, info *common.RelayInfo, upstr
var (
gotFinish bool // 转换器是否已产出带 finish_reason 的结束 chunk
firstChunk bool
transferredTextLen int // 已转发的文本/思考内容长度,供流中断输出估算
writeFailed bool // 写客户端已失败(连接不可达),客户端断开的可靠信号
transferredTextLen int // 已转发的文本/思考内容长度,供流中断输出估算
)
// chunkWriter:将转换器产出的 *dto.ChatCompletionStreamResponse 序列化为 SSE 写出,
@@ -106,7 +107,11 @@ func convertStreamViaRelaykit(ctx context.Context, info *common.RelayInfo, upstr
if err != nil {
return err
}
return helper.WriteSSEData(safeWriter, string(data))
if err := helper.WriteSSEData(safeWriter, string(data)); err != nil {
writeFailed = true
return err
}
return nil
}
// writeTerminal 在转换器未产出结束 chunk 时补发一个终止 chunk,保证客户端正常收尾。
@@ -134,9 +139,17 @@ func convertStreamViaRelaykit(ctx context.Context, info *common.RelayInfo, upstr
monitor.TrackConverterCall(converterID, string(upstream), string(clientFormat), duration, err)
if err != nil {
if ctx.Err() != nil {
// 客户端断开 / 上下文取消:客户端已不可达,不写 [DONE]
setEndReason(common.StreamEndReasonClientGone, ctx.Err())
// 客户端断开判定:ctx 已取消,或已出现过写客户端失败。后者兜住「写失败先于
// ctx 取消被观察到」的竞态与断开信号经代理链路延迟传导的场景——此前仅凭
// ctx.Err() 区分,这类情况会被误判为转换失败:跳过中断计费兜底、向死连接
// 补写 [DONE],最终按 0 token 成功结算。
if writeFailed || ctx.Err() != nil {
endErr := ctx.Err()
if endErr == nil {
endErr = err
}
// 客户端已不可达,不写 terminal/[DONE]
setEndReason(common.StreamEndReasonClientGone, endErr)
// 流中断计费兜底:输出缺失按已转发文本 2 字符/token 估算,输入用请求侧估算值补齐
helper.ApplyInterruptedUsageFallback(info, capturedUsage, transferredTextLen)
return capturedUsage, true
+73
View File
@@ -2,6 +2,8 @@ package relaykit_bridge
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"strings"
"testing"
@@ -192,3 +194,74 @@ func tail(s string, n int) string {
}
return s[len(s)-n:]
}
// failOnSubstringWriter 写入载荷包含 marker 时返回错误(模拟客户端已断开后的写失败),
// 其余写入透传到内层 recorder。
type failOnSubstringWriter struct {
rec *httptest.ResponseRecorder
marker string
}
func (w *failOnSubstringWriter) Header() http.Header { return w.rec.Header() }
func (w *failOnSubstringWriter) WriteHeader(code int) { w.rec.WriteHeader(code) }
func (w *failOnSubstringWriter) Write(p []byte) (int, error) {
if strings.Contains(string(p), w.marker) {
return 0, errors.New("write failed: connection reset by peer")
}
return w.rec.Body.Write(p)
}
// TestConvertStreamViaRelaykit_ClientWriteFailure 复现中断误判竞态:客户端断开由「写失败」
// 首先暴露而 ctx 尚未取消(断开信号经代理传导延迟/first-writer-wins 场景)。
// 此时必须按流中断处理:end reason=client_gone、usage 走中断兜底(已转发文本 2 字符/token
// 估算),且不向死连接补写 [DONE]——而非误判为转换失败导致 0 token / 0 费用结算。
func TestConvertStreamViaRelaykit_ClientWriteFailure(t *testing.T) {
claudeStream := `data: {"type":"message_start","message":{"id":"msg_123","type":"message","role":"assistant","model":"claude-3-opus-20240229","usage":{"input_tokens":10,"output_tokens":0,"cache_read_input_tokens":4,"cache_creation_input_tokens":3}}}
data: {"type":"content_block_start","index":0,"content_block":{"type":"text"}}
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":", how can I help you?"}}
data: {"type":"content_block_stop","index":0}
data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":7}}
data: {"type":"message_stop"}
`
info := newStreamTestRelayInfo(constant.ProviderClaude, constant.RelayFormatOpenAI)
// 与生产一致(relay_handler buildRelayInfo 设置):流中断时上游 usage 缺失,
// 输入 token 按请求侧估算值补齐计费
info.SetEstimatePromptTokens(42)
// ctx 用 Background(永不取消),精确模拟「写失败先于 ctx 取消被观察到」的竞态
w := &failOnSubstringWriter{rec: httptest.NewRecorder(), marker: "how can I help you"}
usage, ok := convertStreamViaRelaykit(context.Background(), info, strings.NewReader(claudeStream), w)
if !ok {
t.Fatal("expected ok=true (handled), got false")
}
if usage == nil {
t.Fatal("expected non-nil usage")
}
if reason := info.StreamStatus.GetEndReason(); reason != common.StreamEndReasonClientGone {
t.Fatalf("end reason = %q, want %q (client write failure misclassified)", reason, common.StreamEndReasonClientGone)
}
if !info.StreamStatus.IsPartialStreamEnd() {
t.Fatal("IsPartialStreamEnd = false, want true")
}
// 转换器的 usage chunk 只在流末尾产出,中断时尚未捕获——prompt 按请求侧估算值(42)补齐;
// 两个 delta 文本("Hello"+", how can I help you?" = 5+21 = 26 字节)均在写失败前计入
// transferredTextLen(累计发生在写之前),中断估算 completion = (26+1)/2 = 13
if usage.PromptTokens != 42 {
t.Errorf("PromptTokens = %d, want 42 (request-side estimate fallback)", usage.PromptTokens)
}
if usage.CompletionTokens != 13 {
t.Errorf("CompletionTokens = %d, want 13 (interrupted estimate, got 0 => fallback skipped)", usage.CompletionTokens)
}
// 中断后不得向死连接补写 [DONE]
if strings.Contains(w.rec.Body.String(), "[DONE]") {
t.Error("output should not contain [DONE] after client disconnect")
}
}