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>

internal/gateway/token_quota.go

@@ -0,0 +1,129 @@
+package gateway
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strconv"
+	"time"
+
+	"aigateway.local/core/internal/apikey"
+	"github.com/redis/go-redis/v9"
+)
+
+var ErrTokenQuotaUnavailable = errors.New("token quota unavailable")
+
+type TokenReservation struct {
+	Allowed    bool
+	APIKeyID   string
+	CounterKey string
+	Reserved   int64
+	Limit      int64
+	Remaining  int64
+	ResetAt    time.Time
+}
+
+type TokenQuotaController interface {
+	Reserve(context.Context, apikey.Principal, int64, time.Time) (TokenReservation, error)
+	Commit(context.Context, TokenReservation, int64) error
+}
+
+type RedisTokenQuotaController struct {
+	client        *redis.Client
+	reserveScript *redis.Script
+	commitScript  *redis.Script
+}
+
+func NewRedisTokenQuotaController(client *redis.Client) *RedisTokenQuotaController {
+	return &RedisTokenQuotaController{
+		client: client, reserveScript: redis.NewScript(tokenReserveScript), commitScript: redis.NewScript(tokenCommitScript),
+	}
+}
+
+func (c *RedisTokenQuotaController) Reserve(ctx context.Context, principal apikey.Principal, estimate int64, now time.Time) (TokenReservation, error) {
+	if principal.APIKeyID == "" {
+		return TokenReservation{Allowed: true}, nil
+	}
+	if estimate < 0 {
+		estimate = 0
+	}
+	// Accounts without a monthly token quota never need a reservation and must
+	// not create a pointless monthly counter key in Redis.
+	if principal.MonthlyTokenQuota == 0 {
+		return TokenReservation{Allowed: true}, nil
+	}
+	if c == nil || c.client == nil {
+		return TokenReservation{}, ErrTokenQuotaUnavailable
+	}
+	now = now.UTC()
+	reset := time.Date(now.Year(), now.Month()+1, 1, 0, 0, 0, 0, time.UTC)
+	key := apikey.MonthlyTokenUsageKey(principal.APIKeyID, now)
+	result, err := c.reserveScript.Run(ctx, c.client, []string{key}, estimate, principal.MonthlyTokenQuota, int64(reset.Sub(now).Seconds())+86400).Slice()
+	if err != nil || len(result) != 2 {
+		if principal.MonthlyTokenQuota == 0 {
+			return TokenReservation{Allowed: true}, nil
+		}
+		return TokenReservation{}, fmt.Errorf("%w: %v", ErrTokenQuotaUnavailable, err)
+	}
+	allowed, err := redisInteger(result[0])
+	if err != nil {
+		return TokenReservation{}, ErrTokenQuotaUnavailable
+	}
+	current, err := redisInteger(result[1])
+	if err != nil {
+		return TokenReservation{}, ErrTokenQuotaUnavailable
+	}
+	return TokenReservation{
+		Allowed: allowed == 1, APIKeyID: principal.APIKeyID, CounterKey: key, Reserved: estimate,
+		Limit: principal.MonthlyTokenQuota, Remaining: max(principal.MonthlyTokenQuota-current, 0), ResetAt: reset,
+	}, nil
+}
+
+func (c *RedisTokenQuotaController) Commit(ctx context.Context, reservation TokenReservation, actual int64) error {
+	if c == nil || c.client == nil || reservation.CounterKey == "" || !reservation.Allowed {
+		return nil
+	}
+	if actual < 0 {
+		actual = 0
+	}
+	if _, err := c.commitScript.Run(ctx, c.client, []string{reservation.CounterKey}, actual-reservation.Reserved).Result(); err != nil {
+		return fmt.Errorf("%w: %v", ErrTokenQuotaUnavailable, err)
+	}
+	return nil
+}
+
+const tokenReserveScript = `
+local estimate = tonumber(ARGV[1])
+local quota = tonumber(ARGV[2])
+local current = tonumber(redis.call('GET', KEYS[1]) or '0')
+if quota > 0 and current + estimate > quota then
+  return {0, current}
+end
+if estimate > 0 then
+  current = redis.call('INCRBY', KEYS[1], estimate)
+  if current == estimate then redis.call('EXPIRE', KEYS[1], tonumber(ARGV[3])) end
+elseif redis.call('EXISTS', KEYS[1]) == 0 then
+  redis.call('SET', KEYS[1], 0, 'EX', tonumber(ARGV[3]))
+end
+return {1, current}
+`
+
+const tokenCommitScript = `
+local delta = tonumber(ARGV[1])
+local current = tonumber(redis.call('GET', KEYS[1]) or '0')
+local updated = current + delta
+if updated < 0 then updated = 0 end
+redis.call('SET', KEYS[1], updated, 'KEEPTTL')
+return updated
+`
+
+func writeTokenQuotaHeaders(header mapHeader, reservation TokenReservation) {
+	if reservation.Limit <= 0 || reservation.ResetAt.IsZero() {
+		return
+	}
+	header.Set("X-TokenLimit-Limit", strconv.FormatInt(reservation.Limit, 10))
+	header.Set("X-TokenLimit-Remaining", strconv.FormatInt(max(reservation.Remaining, 0), 10))
+	header.Set("X-TokenLimit-Reset", strconv.FormatInt(reservation.ResetAt.Unix(), 10))
+}
+
+type mapHeader interface{ Set(string, string) }