AI Gateway Go 0.10.0 源码快照 + 旗舰版需求规划报告
M0-M7 已完成:核心网关(身份/RBAC/TOTP/OIDC/SAML/Provider/配额/路由/内容策略/审计/定价)+ 资源市场(MCP/Skills/数字员工)。 含 22 个 PostgreSQL 迁移、管理端/门户端前端源码、OpenAPI 契约、部署 compose。 Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,261 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"log/slog"
|
||||
"mime"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// maxUsageDocumentBytes bounds the buffered tail of a non-streaming response.
|
||||
// The window is kept from the end of the body (where providers place the
|
||||
// "usage" object), so token accounting stays correct for responses far larger
|
||||
// than this bound while memory use stays bounded per in-flight response.
|
||||
const maxUsageDocumentBytes = 2 << 20
|
||||
|
||||
type usageSession struct {
|
||||
controller TokenQuotaController
|
||||
reservation TokenReservation
|
||||
fallback int64
|
||||
onFinish func(TokenUsage)
|
||||
once sync.Once
|
||||
log *slog.Logger
|
||||
}
|
||||
|
||||
type TokenUsage struct {
|
||||
Input int64
|
||||
Output int64
|
||||
Total int64
|
||||
}
|
||||
|
||||
type usageSessionContextKey struct{}
|
||||
|
||||
func withUsageSession(request *http.Request, session *usageSession) *http.Request {
|
||||
return request.WithContext(context.WithValue(request.Context(), usageSessionContextKey{}, session))
|
||||
}
|
||||
|
||||
func usageSessionFrom(request *http.Request) *usageSession {
|
||||
session, _ := request.Context().Value(usageSessionContextKey{}).(*usageSession)
|
||||
return session
|
||||
}
|
||||
|
||||
func (s *usageSession) finish(usage TokenUsage) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.once.Do(func() {
|
||||
usage.Total = max(usage.Total, usage.Input+usage.Output)
|
||||
if usage.Total <= 0 {
|
||||
usage.Total = s.fallback
|
||||
}
|
||||
if s.controller != nil {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
if err := s.controller.Commit(ctx, s.reservation, usage.Total); err != nil && s.log != nil {
|
||||
// The reservation already counted towards the quota, so a
|
||||
// failed reconciliation silently leaves the counter slightly
|
||||
// off. Surface it instead of dropping it.
|
||||
s.log.Warn("token quota commit failed", "error", err)
|
||||
}
|
||||
cancel()
|
||||
}
|
||||
if s.onFinish != nil {
|
||||
s.onFinish(usage)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
type usageReadCloser struct {
|
||||
io.ReadCloser
|
||||
collector *usageCollector
|
||||
session *usageSession
|
||||
}
|
||||
|
||||
func newUsageReadCloser(body io.ReadCloser, contentType string, session *usageSession) io.ReadCloser {
|
||||
mediaType, _, _ := mime.ParseMediaType(contentType)
|
||||
return &usageReadCloser{ReadCloser: body, collector: &usageCollector{sse: mediaType == "text/event-stream"}, session: session}
|
||||
}
|
||||
|
||||
func (r *usageReadCloser) Read(buffer []byte) (int, error) {
|
||||
n, err := r.ReadCloser.Read(buffer)
|
||||
if n > 0 {
|
||||
r.collector.feed(buffer[:n])
|
||||
}
|
||||
if err == io.EOF {
|
||||
r.session.finish(r.collector.usage())
|
||||
}
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (r *usageReadCloser) Close() error {
|
||||
r.session.finish(r.collector.usage())
|
||||
return r.ReadCloser.Close()
|
||||
}
|
||||
|
||||
type usageCollector struct {
|
||||
sse bool
|
||||
pending []byte
|
||||
doc []byte
|
||||
input int64
|
||||
output int64
|
||||
total int64
|
||||
}
|
||||
|
||||
func (c *usageCollector) feed(chunk []byte) {
|
||||
if !c.sse {
|
||||
c.doc = append(c.doc, chunk...)
|
||||
// Keep only a bounded tail window. The "usage" member lives at the end
|
||||
// of a non-streaming response, so dropping the head (never the tail)
|
||||
// preserves accounting for arbitrarily large bodies at a fixed memory
|
||||
// cost instead of truncating usage away.
|
||||
if len(c.doc) > maxUsageDocumentBytes {
|
||||
c.doc = append([]byte(nil), c.doc[len(c.doc)-maxUsageDocumentBytes:]...)
|
||||
}
|
||||
return
|
||||
}
|
||||
c.pending = append(c.pending, chunk...)
|
||||
for {
|
||||
index := bytes.IndexByte(c.pending, '\n')
|
||||
if index < 0 {
|
||||
if len(c.pending) > maxUsageDocumentBytes {
|
||||
c.pending = c.pending[:0]
|
||||
}
|
||||
return
|
||||
}
|
||||
line := strings.TrimSpace(string(c.pending[:index]))
|
||||
c.pending = c.pending[index+1:]
|
||||
if strings.HasPrefix(line, "data:") {
|
||||
payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
||||
if payload != "" && payload != "[DONE]" {
|
||||
c.consumeJSON([]byte(payload))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *usageCollector) tokens() int64 {
|
||||
return c.usage().Total
|
||||
}
|
||||
|
||||
func (c *usageCollector) usage() TokenUsage {
|
||||
if c.sse && len(c.pending) > 0 {
|
||||
line := strings.TrimSpace(string(c.pending))
|
||||
if strings.HasPrefix(line, "data:") {
|
||||
c.consumeJSON([]byte(strings.TrimSpace(strings.TrimPrefix(line, "data:"))))
|
||||
}
|
||||
c.pending = nil
|
||||
}
|
||||
if !c.sse && len(c.doc) > 0 {
|
||||
c.consumeJSON(c.doc) // fast path: whole valid JSON object
|
||||
c.consumeUsageObject(c.doc) // tail extraction: covers truncated bodies
|
||||
}
|
||||
return TokenUsage{Input: c.input, Output: c.output, Total: max(c.total, c.input+c.output)}
|
||||
}
|
||||
|
||||
// consumeUsageObject extracts the trailing `"usage"` object from a possibly
|
||||
// truncated response document. It scans for the final `"usage"` member and
|
||||
// decodes the JSON object that immediately follows — where OpenAI-compatible
|
||||
// providers place token accounting in non-streaming responses — so usage is
|
||||
// still counted when the buffered window starts mid-object.
|
||||
func (c *usageCollector) consumeUsageObject(doc []byte) {
|
||||
const key = `"usage"`
|
||||
index := bytes.LastIndex(doc, []byte(key))
|
||||
if index < 0 {
|
||||
return
|
||||
}
|
||||
rest := doc[index+len(key):]
|
||||
colon := bytes.IndexByte(rest, ':')
|
||||
if colon < 0 {
|
||||
return
|
||||
}
|
||||
rest = bytes.TrimSpace(rest[colon+1:])
|
||||
if len(rest) == 0 || rest[0] != '{' {
|
||||
return
|
||||
}
|
||||
depth := 0
|
||||
end := -1
|
||||
inString := false
|
||||
escaped := false
|
||||
for i := 0; i < len(rest); i++ {
|
||||
ch := rest[i]
|
||||
if inString {
|
||||
if escaped {
|
||||
escaped = false
|
||||
} else if ch == '\\' {
|
||||
escaped = true
|
||||
} else if ch == '"' {
|
||||
inString = false
|
||||
}
|
||||
continue
|
||||
}
|
||||
switch ch {
|
||||
case '"':
|
||||
inString = true
|
||||
case '{':
|
||||
depth++
|
||||
case '}':
|
||||
depth--
|
||||
if depth == 0 {
|
||||
end = i + 1
|
||||
}
|
||||
}
|
||||
if end > 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
if end > 0 {
|
||||
c.consumeJSON(rest[:end])
|
||||
}
|
||||
}
|
||||
|
||||
func (c *usageCollector) consumeJSON(payload []byte) {
|
||||
decoder := json.NewDecoder(bytes.NewReader(payload))
|
||||
decoder.UseNumber()
|
||||
var value any
|
||||
if decoder.Decode(&value) == nil {
|
||||
c.walk(value, false)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *usageCollector) walk(value any, inUsage bool) {
|
||||
switch typed := value.(type) {
|
||||
case map[string]any:
|
||||
for key, child := range typed {
|
||||
usage := inUsage || key == "usage"
|
||||
if usage {
|
||||
switch key {
|
||||
case "input_tokens", "prompt_tokens":
|
||||
c.input = max(c.input, jsonInt64(child))
|
||||
case "output_tokens", "completion_tokens":
|
||||
c.output = max(c.output, jsonInt64(child))
|
||||
case "total_tokens":
|
||||
c.total = max(c.total, jsonInt64(child))
|
||||
}
|
||||
}
|
||||
c.walk(child, usage)
|
||||
}
|
||||
case []any:
|
||||
for _, child := range typed {
|
||||
c.walk(child, inUsage)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func jsonInt64(value any) int64 {
|
||||
switch number := value.(type) {
|
||||
case json.Number:
|
||||
parsed, _ := number.Int64()
|
||||
return max(parsed, 0)
|
||||
case float64:
|
||||
return max(int64(number), 0)
|
||||
case int64:
|
||||
return max(number, 0)
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user