ai-gateway-go/internal/gateway/token_quota.go

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) }