代理: 流式中断按已生成部分计费

- 新增 tokenizer 包(tiktoken-go)按模型估算 token, 未知模型回退 cl100k_base
- passthrough 提取请求输入文本 + SSE 已生成内容, 客户端断开时记 canceled
- usage 计费范围扩展: canceled(流式中断)按已生成部分收费

Co-Authored-By: Claude <noreply@anthropic.com>

server/go.mod

@@ -1,6 +1,6 @@
 module github.com/openteam/server
 
-go 1.25.0
+go 1.26
 
 require (
 	github.com/gin-gonic/gin v1.12.0
@@ -8,6 +8,7 @@ require (
 	github.com/go-webauthn/webauthn v0.17.4
 	github.com/golang-jwt/jwt/v5 v5.3.1
 	github.com/spf13/viper v1.21.0
+	github.com/tiktoken-go/tokenizer v0.8.1
 	golang.org/x/crypto v0.55.0
 	gorm.io/driver/postgres v1.6.2
 	gorm.io/gorm v1.31.2
@@ -18,6 +19,7 @@ require (
 	github.com/bytedance/sonic v1.15.0 // indirect
 	github.com/bytedance/sonic/loader v0.5.0 // indirect
 	github.com/cloudwego/base64x v0.1.6 // indirect
+	github.com/dlclark/regexp2/v2 v2.5.1 // indirect
 	github.com/dustin/go-humanize v1.0.1 // indirect
 	github.com/fsnotify/fsnotify v1.9.0 // indirect
 	github.com/fxamacker/cbor/v2 v2.9.2 // indirect

server/go.sum

@@ -9,6 +9,8 @@ github.com/cloudwego/base64x v0.1.6/go.mod h1:OFcloc187FXDaYHvrNIjxSe8ncn0OOM8gE
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/dlclark/regexp2/v2 v2.5.1 h1:E5Ug7Dh264W1ymdySmiHNcDG7fmsR307APCE5R07a20=
+github.com/dlclark/regexp2/v2 v2.5.1/go.mod h1:avUrQvPaLz2DrFNHJF0taWAFFX2C1GMSSoeiqFjcBmU=
 github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
 github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
 github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8=
@@ -130,6 +132,8 @@ github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu
 github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
 github.com/subosito/gotenv v1.6.0 h1:9NlTDc1FTs4qu0DDq7AEtTPNw6SVm7uBMsUCUjABIf8=
 github.com/subosito/gotenv v1.6.0/go.mod h1:Dk4QP5c2W3ibzajGcXpNraDfq2IrhjMIvMSWPKKo0FU=
+github.com/tiktoken-go/tokenizer v0.8.1 h1:4obDoB6/dhdBt9xMweX4nww5cjdOq/nYF4ecwPq2+mg=
+github.com/tiktoken-go/tokenizer v0.8.1/go.mod h1:eLA0t6nGvn9mDc7gt90qt7pMat+gE9ViqwQ6l9B+tA4=
 github.com/tinylib/msgp v1.6.4 h1:mOwYbyYDLPj35mkA2BjjYejgJk9BuHxDdvRnb6v2ZcQ=
 github.com/tinylib/msgp v1.6.4/go.mod h1:RSp0LW9oSxFut3KzESt5Voq4GVWyS+PSulT77roAqEA=
 github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI=

server/internal/pkg/tokenizer/tokenizer.go

@@ -0,0 +1,24 @@
+// Package tokenizer 按模型估算 token 数（tiktoken-go）。
+// 用于流式中断时对已生成内容做近似计费；未知模型回退 cl100k_base。
+package tokenizer
+
+import "github.com/tiktoken-go/tokenizer"
+
+// Count 估算文本 token 数；空文本返回 0。
+func Count(text, model string) int {
+	if text == "" {
+		return 0
+	}
+	enc, err := tokenizer.ForModel(tokenizer.Model(model))
+	if err != nil {
+		enc, err = tokenizer.Get(tokenizer.Cl100kBase)
+		if err != nil {
+			return 0
+		}
+	}
+	toks, _, err := enc.Encode(text)
+	if err != nil {
+		return 0
+	}
+	return len(toks)
+}

server/internal/proxy/passthrough.go

@@ -16,6 +16,7 @@ import (
 
 	"github.com/gin-gonic/gin"
 	"github.com/openteam/server/internal/channel"
+	"github.com/openteam/server/internal/pkg/tokenizer"
 	"github.com/openteam/server/internal/store"
 )
 
@@ -39,11 +40,107 @@ func parseBody(c *gin.Context) (*bodyReq, []byte, error) {
 		return nil, nil, err
 	}
 	c.Request.Body = io.NopCloser(bytes.NewReader(body))
+	// 中断请求输入文本（token 估算用，仅存文本不立即分词）
+	c.Set("est_input_text", requestText(body))
 	br := &bodyReq{}
 	_ = json.Unmarshal(body, br) // 解析失败按空处理，直通仍可转发
 	return br, body, nil
 }
 
+// requestText 提取请求体中的用户输入文本（chat/messages 的 content、responses 的 input/instructions），
+// 用于中断时估算输入 token。
+func requestText(body []byte) string {
+	var m map[string]any
+	if json.Unmarshal(body, &m) != nil {
+		return ""
+	}
+	var parts []string
+	add := func(s string) {
+		if s = strings.TrimSpace(s); s != "" {
+			parts = append(parts, s)
+		}
+	}
+	if s, ok := m["instructions"].(string); ok {
+		add(s)
+	}
+	if s, ok := m["system"].(string); ok {
+		add(s)
+	}
+	switch input := m["input"].(type) {
+	case string:
+		add(input)
+	case []any:
+		for _, it := range input {
+			if im, ok := it.(map[string]any); ok {
+				if s, ok := im["content"].(string); ok {
+					add(s)
+				}
+			}
+		}
+	}
+	if msgs, ok := m["messages"].([]any); ok {
+		for _, msg := range msgs {
+			mm, ok := msg.(map[string]any)
+			if !ok {
+				continue
+			}
+			switch c := mm["content"].(type) {
+			case string:
+				add(c)
+			case []any:
+				for _, b := range c {
+					if bm, ok := b.(map[string]any); ok {
+						if s, ok := bm["text"].(string); ok {
+							add(s)
+						}
+					}
+				}
+			}
+		}
+	}
+	return strings.Join(parts, "\n")
+}
+
+// sseContentText 提取一条 SSE data 行中的内容文本（chat delta.content / responses delta / messages delta.text）。
+func sseContentText(line []byte) string {
+	s := string(line)
+	if strings.HasPrefix(s, "data: ") {
+		s = strings.TrimPrefix(s, "data: ")
+	}
+	s = strings.TrimSpace(s)
+	if s == "" || s == "[DONE]" {
+		return ""
+	}
+	var m map[string]any
+	if json.Unmarshal([]byte(s), &m) != nil {
+		return ""
+	}
+	// responses output_text.delta: {"delta":"..."}
+	if d, ok := m["delta"].(string); ok {
+		return d
+	}
+	// messages content_block_delta: {"delta":{"text":"..."}}
+	if dm, ok := m["delta"].(map[string]any); ok {
+		if t, ok := dm["text"].(string); ok {
+			return t
+		}
+	}
+	// chat chunk: {"choices":[{"delta":{"content":"..."}}]}（GLM 思考模型另有 reasoning_content）
+	if choices, ok := m["choices"].([]any); ok && len(choices) > 0 {
+		if c0, ok := choices[0].(map[string]any); ok {
+			if delta, ok := c0["delta"].(map[string]any); ok {
+				if t, ok := delta["content"].(string); ok {
+					return t
+				}
+				if t, ok := delta["reasoning_content"].(string); ok {
+					return t
+				}
+			}
+		}
+	}
+	return ""
+}
+
 // upstreamURL 组装上游地址：按协议选 base_url 再拼资源路径（见 store.Channel.UpstreamURL）。
 func upstreamURL(ch *store.Channel, proto, path string) string {
 	return ch.UpstreamURL(proto, path)
@@ -192,10 +289,14 @@ func (g *Gateway) streamCopy(c *gin.Context, ch *store.Channel, r io.Reader, sta
 			}
 			if out != nil {
 				if _, werr := w.Write(out); werr != nil {
-					g.recordError(c, ch, nil, start, "client_disconnect")
+					// 客户端意外断开：按已生成部分收费（canceled）
+					g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect")
 					return
 				}
 				flusher.Flush()
+				if sink != nil {
+					sink.outputText += sseContentText(out)
+				}
 			}
 			if usageRaw := scanUsage(line); usageRaw != nil && sink != nil {
 				sink.push(usageRaw)
@@ -205,7 +306,8 @@ func (g *Gateway) streamCopy(c *gin.Context, ch *store.Channel, r io.Reader, sta
 			if err == io.EOF {
 				g.finishUsage(c, ch, start, store.UsageStatusSuccess, "")
 			} else if c.Request.Context().Err() != nil {
-				g.recordError(c, ch, nil, start, "client_disconnect")
+				// 客户端意外断开：按已生成部分收费（canceled）
+				g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect")
 			} else {
 				g.recordError(c, ch, nil, start, "stream_read_error")
 			}
@@ -336,8 +438,10 @@ func (s *sseScanner) Next() ([]byte, error) {
 // 记账
 
 // usageSink 累积多次 usage：合并各事件字段（message_start 给 input，message_delta 给 output）。
+// outputText 累积已转发内容文本，用于流式中断时按 tiktoken 估算输出 token。
 type usageSink struct {
-	us usageShape
+	us         usageShape
+	outputText string
 }
 
 func (u *usageSink) push(raw json.RawMessage) {
@@ -379,9 +483,11 @@ func (g *Gateway) finishUsage(c *gin.Context, ch *store.Channel, start time.Time
 	trace, _ := c.Get(CtxTrace)
 
 	var us usageShape
+	var sink *usageSink
 	if h, ok := c.Get("usage_raw"); ok {
 		if holder, ok := h.(*sinkHolder); ok && holder.sink != nil {
 			us = holder.sink.Shape()
+			sink = holder.sink
 		}
 	}
 
@@ -393,6 +499,20 @@ func (g *Gateway) finishUsage(c *gin.Context, ch *store.Channel, start time.Time
 	modelName, _ := c.Get("model_name")
 	mn, _ := modelName.(string)
 
+	// 流式中断（canceled）：上游最终 usage 可能未返回，按已生成内容用 tiktoken 估算
+	if status == store.UsageStatusCanceled {
+		if in == 0 {
+			if est, ok := c.Get("est_input_text"); ok {
+				if v, ok2 := est.(string); ok2 {
+					in = int64(tokenizer.Count(v, mn))
+				}
+			}
+		}
+		if out == 0 && sink != nil && sink.outputText != "" {
+			out = int64(tokenizer.Count(sink.outputText, mn))
+		}
+	}
+
 	var model store.Model
 	var cost float64
 	var modelID uint64

server/internal/usage/usage.go

@@ -97,7 +97,8 @@ func (r *Recorder) flush(logs []*store.UsageLog) error {
 			return err
 		}
 		for _, l := range logs {
-			if l.Status != store.UsageStatusSuccess || l.Cost <= 0 {
+			// 计费范围：success（正常完成）与 canceled（流式中断，按已生成部分收费）
+			if (l.Status != store.UsageStatusSuccess && l.Status != store.UsageStatusCanceled) || l.Cost <= 0 {
 				continue
 			}
 			// 扣余额（余额可为负：流式请求不中断；后续请求被拒）