BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add gateway billing estimate and rate limit details

Browse Source
This commit is contained in:
wangbo
2026-05-15 01:53:52 +08:00
parent bdc9be63d5
commit 37d0f919e5
14 changed files with 916 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/store/rate_limits.go
+50 -13
View File
@@ -31,9 +31,18 @@ func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, attemptID
}
if reservation.Metric == "" || reservation.Amount > reservation.Limit {
return RateLimitResult{}, &RateLimitExceededError{
Metric: reservation.Metric,
Message: fmt.Sprintf("rate limit exceeded: %s request amount %.0f is greater than limit %.0f", reservation.Metric, reservation.Amount, reservation.Limit),
Retryable: false,
ScopeType: reservation.ScopeType,
ScopeKey: reservation.ScopeKey,
ScopeName: reservation.ScopeName,
ScopeMetadata: reservation.ScopeMetadata,
Metric: reservation.Metric,
Limit: reservation.Limit,
Amount: reservation.Amount,
Projected: reservation.Amount,
WindowSeconds: reservation.WindowSeconds,
Policy: reservation.Policy,
Message: fmt.Sprintf("rate limit exceeded: %s request amount %.0f is greater than limit %.0f", reservation.Metric, reservation.Amount, reservation.Limit),
Retryable: false,
}
}
if reservation.WindowSeconds <= 0 {
@@ -78,10 +87,22 @@ WHERE scope_type = $1
}
if active+reservation.Amount > reservation.Limit {
return "", &RateLimitExceededError{
Metric: reservation.Metric,
Message: fmt.Sprintf("rate limit exceeded: concurrent active %.0f plus request %.0f is greater than limit %.0f", active, reservation.Amount, reservation.Limit),
RetryAfter: concurrencyRetryAfter(nextAvailableAt),
Retryable: true,
ScopeType: reservation.ScopeType,
ScopeKey: reservation.ScopeKey,
ScopeName: reservation.ScopeName,
ScopeMetadata: reservation.ScopeMetadata,
Metric: reservation.Metric,
Limit: reservation.Limit,
Amount: reservation.Amount,
Current: active,
Used: active,
Projected: active + reservation.Amount,
WindowSeconds: reservation.WindowSeconds,
ResetAt: nextAvailableAt,
Policy: reservation.Policy,
Message: fmt.Sprintf("rate limit exceeded: concurrent active %.0f plus request %.0f is greater than limit %.0f", active, reservation.Amount, reservation.Limit),
RetryAfter: concurrencyRetryAfter(nextAvailableAt),
Retryable: true,
}
}
var leaseID string
@@ -135,11 +156,13 @@ RETURNING window_start`,
if err != nil {
if errors.Is(err, pgx.ErrNoRows) {
resetAt := time.Now().Add(time.Duration(reservation.WindowSeconds) * time.Second)
currentUsed := 0.0
currentReserved := 0.0
_ = tx.QueryRow(ctx, `
WITH bounds AS (
SELECT to_timestamp(floor(extract(epoch FROM now()) / $4::int) * $4::int) AS window_start
)
SELECT counters.reset_at
SELECT counters.used_value::float8, counters.reserved_value::float8, counters.reset_at
FROM gateway_rate_limit_counters counters
JOIN bounds ON counters.window_start = bounds.window_start
WHERE scope_type = $1
@@ -149,12 +172,26 @@ WHERE scope_type = $1
reservation.ScopeKey,
reservation.Metric,
reservation.WindowSeconds,
).Scan(&resetAt)
).Scan(&currentUsed, &currentReserved, &resetAt)
current := currentUsed + currentReserved
return RateLimitReservation{}, &RateLimitExceededError{
Metric: reservation.Metric,
Message: fmt.Sprintf("rate limit exceeded: %s window has no remaining capacity", reservation.Metric),
RetryAfter: retryAfterUntil(resetAt),
Retryable: true,
ScopeType: reservation.ScopeType,
ScopeKey: reservation.ScopeKey,
ScopeName: reservation.ScopeName,
ScopeMetadata: reservation.ScopeMetadata,
Metric: reservation.Metric,
Limit: reservation.Limit,
Amount: reservation.Amount,
Current: current,
Used: currentUsed,
Reserved: currentReserved,
Projected: current + reservation.Amount,
WindowSeconds: reservation.WindowSeconds,
ResetAt: resetAt,
Policy: reservation.Policy,
Message: fmt.Sprintf("rate limit exceeded: %s window has no remaining capacity", reservation.Metric),
RetryAfter: retryAfterUntil(resetAt),
Retryable: true,
}
}
return RateLimitReservation{}, err
apps/api/internal/store/runtime_types.go
+20 -4
View File
@@ -33,10 +33,23 @@ func ModelCandidateErrorCode(err error) string {
}
type RateLimitExceededError struct {
Metric string
Message string
RetryAfter time.Duration
Retryable bool
ScopeType string
ScopeKey string
ScopeName string
ScopeMetadata map[string]any
Metric string
Limit float64
Amount float64
Current float64
Used float64
Reserved float64
Projected float64
WindowSeconds int
ResetAt time.Time
Policy map[string]any
Message string
RetryAfter time.Duration
Retryable bool
}
func (e *RateLimitExceededError) Error() string {
@@ -166,12 +179,15 @@ type RateLimitReservation struct {
ReservationID string
ScopeType string
ScopeKey string
ScopeName string
ScopeMetadata map[string]any
Metric string
Limit float64
Amount float64
WindowSeconds int
LeaseTTLSeconds int
WindowStart time.Time
Policy map[string]any
}
type RateLimitResult struct {
apps/api/internal/store/user_group_policy.go
+3 -2
View File
@@ -10,6 +10,7 @@ import (
type UserGroupPolicy struct {
ID string
GroupKey string
Name string
RateLimitPolicy map[string]any
BillingDiscountPolicy map[string]any
}
@@ -23,12 +24,12 @@ func (s *Store) ResolveUserGroupPolicy(ctx context.Context, user *auth.User) (Us
var rateLimit []byte
var billing []byte
err := s.pool.QueryRow(ctx, `
SELECT id::text, group_key, rate_limit_policy, billing_discount_policy
SELECT id::text, group_key, name, rate_limit_policy, billing_discount_policy
FROM gateway_user_groups
WHERE status = 'active'
AND (($1 <> '' AND id = NULLIF($1, '')::uuid) OR ($1 = '' AND group_key = 'default'))
ORDER BY CASE WHEN id::text = $1 THEN 0 ELSE 1 END, priority ASC
LIMIT 1`, userGroupID).Scan(&item.ID, &item.GroupKey, &rateLimit, &billing)
LIMIT 1`, userGroupID).Scan(&item.ID, &item.GroupKey, &item.Name, &rateLimit, &billing)
if err != nil {
if err == pgx.ErrNoRows {
return UserGroupPolicy{}, nil