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

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

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

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
 			}
 			// 扣余额（余额可为负：流式请求不中断；后续请求被拒）