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