Commit a80ec5d8 authored by shaw's avatar shaw
Browse files

feat: apikey支持5h/1d/7d速率控制

parent b7df7ce5
...@@ -19,6 +19,11 @@ type APIKeyAuthSnapshot struct { ...@@ -19,6 +19,11 @@ type APIKeyAuthSnapshot struct {
// Expiration field for API Key expiration feature // Expiration field for API Key expiration feature
ExpiresAt *time.Time `json:"expires_at,omitempty"` // Expiration time (nil = never expires) ExpiresAt *time.Time `json:"expires_at,omitempty"` // Expiration time (nil = never expires)
// Rate limit configuration (only limits, not usage - usage read from Redis at check time)
RateLimit5h float64 `json:"rate_limit_5h"`
RateLimit1d float64 `json:"rate_limit_1d"`
RateLimit7d float64 `json:"rate_limit_7d"`
} }
// APIKeyAuthUserSnapshot 用户快照 // APIKeyAuthUserSnapshot 用户快照
......
...@@ -209,6 +209,9 @@ func (s *APIKeyService) snapshotFromAPIKey(apiKey *APIKey) *APIKeyAuthSnapshot { ...@@ -209,6 +209,9 @@ func (s *APIKeyService) snapshotFromAPIKey(apiKey *APIKey) *APIKeyAuthSnapshot {
Quota: apiKey.Quota, Quota: apiKey.Quota,
QuotaUsed: apiKey.QuotaUsed, QuotaUsed: apiKey.QuotaUsed,
ExpiresAt: apiKey.ExpiresAt, ExpiresAt: apiKey.ExpiresAt,
RateLimit5h: apiKey.RateLimit5h,
RateLimit1d: apiKey.RateLimit1d,
RateLimit7d: apiKey.RateLimit7d,
User: APIKeyAuthUserSnapshot{ User: APIKeyAuthUserSnapshot{
ID: apiKey.User.ID, ID: apiKey.User.ID,
Status: apiKey.User.Status, Status: apiKey.User.Status,
...@@ -262,6 +265,9 @@ func (s *APIKeyService) snapshotToAPIKey(key string, snapshot *APIKeyAuthSnapsho ...@@ -262,6 +265,9 @@ func (s *APIKeyService) snapshotToAPIKey(key string, snapshot *APIKeyAuthSnapsho
Quota: snapshot.Quota, Quota: snapshot.Quota,
QuotaUsed: snapshot.QuotaUsed, QuotaUsed: snapshot.QuotaUsed,
ExpiresAt: snapshot.ExpiresAt, ExpiresAt: snapshot.ExpiresAt,
RateLimit5h: snapshot.RateLimit5h,
RateLimit1d: snapshot.RateLimit1d,
RateLimit7d: snapshot.RateLimit7d,
User: &User{ User: &User{
ID: snapshot.User.ID, ID: snapshot.User.ID,
Status: snapshot.User.Status, Status: snapshot.User.Status,
......
...@@ -30,6 +30,11 @@ var ( ...@@ -30,6 +30,11 @@ var (
ErrAPIKeyExpired = infraerrors.Forbidden("API_KEY_EXPIRED", "api key 已过期") ErrAPIKeyExpired = infraerrors.Forbidden("API_KEY_EXPIRED", "api key 已过期")
// ErrAPIKeyQuotaExhausted = infraerrors.TooManyRequests("API_KEY_QUOTA_EXHAUSTED", "api key quota exhausted") // ErrAPIKeyQuotaExhausted = infraerrors.TooManyRequests("API_KEY_QUOTA_EXHAUSTED", "api key quota exhausted")
ErrAPIKeyQuotaExhausted = infraerrors.TooManyRequests("API_KEY_QUOTA_EXHAUSTED", "api key 额度已用完") ErrAPIKeyQuotaExhausted = infraerrors.TooManyRequests("API_KEY_QUOTA_EXHAUSTED", "api key 额度已用完")
// Rate limit errors
ErrAPIKeyRateLimit5hExceeded = infraerrors.TooManyRequests("API_KEY_RATE_5H_EXCEEDED", "api key 5小时限额已用完")
ErrAPIKeyRateLimit1dExceeded = infraerrors.TooManyRequests("API_KEY_RATE_1D_EXCEEDED", "api key 日限额已用完")
ErrAPIKeyRateLimit7dExceeded = infraerrors.TooManyRequests("API_KEY_RATE_7D_EXCEEDED", "api key 7天限额已用完")
) )
const ( const (
...@@ -64,6 +69,21 @@ type APIKeyRepository interface { ...@@ -64,6 +69,21 @@ type APIKeyRepository interface {
// Quota methods // Quota methods
IncrementQuotaUsed(ctx context.Context, id int64, amount float64) (float64, error) IncrementQuotaUsed(ctx context.Context, id int64, amount float64) (float64, error)
UpdateLastUsed(ctx context.Context, id int64, usedAt time.Time) error UpdateLastUsed(ctx context.Context, id int64, usedAt time.Time) error
// Rate limit methods
IncrementRateLimitUsage(ctx context.Context, id int64, cost float64) error
ResetRateLimitWindows(ctx context.Context, id int64) error
GetRateLimitData(ctx context.Context, id int64) (*APIKeyRateLimitData, error)
}
// APIKeyRateLimitData holds rate limit usage and window state for an API key.
type APIKeyRateLimitData struct {
Usage5h float64
Usage1d float64
Usage7d float64
Window5hStart *time.Time
Window1dStart *time.Time
Window7dStart *time.Time
} }
// APIKeyCache defines cache operations for API key service // APIKeyCache defines cache operations for API key service
...@@ -102,6 +122,11 @@ type CreateAPIKeyRequest struct { ...@@ -102,6 +122,11 @@ type CreateAPIKeyRequest struct {
// Quota fields // Quota fields
Quota float64 `json:"quota"` // Quota limit in USD (0 = unlimited) Quota float64 `json:"quota"` // Quota limit in USD (0 = unlimited)
ExpiresInDays *int `json:"expires_in_days"` // Days until expiry (nil = never expires) ExpiresInDays *int `json:"expires_in_days"` // Days until expiry (nil = never expires)
// Rate limit fields (0 = unlimited)
RateLimit5h float64 `json:"rate_limit_5h"`
RateLimit1d float64 `json:"rate_limit_1d"`
RateLimit7d float64 `json:"rate_limit_7d"`
} }
// UpdateAPIKeyRequest 更新API Key请求 // UpdateAPIKeyRequest 更新API Key请求
...@@ -117,9 +142,20 @@ type UpdateAPIKeyRequest struct { ...@@ -117,9 +142,20 @@ type UpdateAPIKeyRequest struct {
ExpiresAt *time.Time `json:"expires_at"` // Expiration time (nil = no change) ExpiresAt *time.Time `json:"expires_at"` // Expiration time (nil = no change)
ClearExpiration bool `json:"-"` // Clear expiration (internal use) ClearExpiration bool `json:"-"` // Clear expiration (internal use)
ResetQuota *bool `json:"reset_quota"` // Reset quota_used to 0 ResetQuota *bool `json:"reset_quota"` // Reset quota_used to 0
// Rate limit fields (nil = no change, 0 = unlimited)
RateLimit5h *float64 `json:"rate_limit_5h"`
RateLimit1d *float64 `json:"rate_limit_1d"`
RateLimit7d *float64 `json:"rate_limit_7d"`
ResetRateLimitUsage *bool `json:"reset_rate_limit_usage"` // Reset all usage counters to 0
} }
// APIKeyService API Key服务 // APIKeyService API Key服务
// RateLimitCacheInvalidator invalidates rate limit cache entries on manual reset.
type RateLimitCacheInvalidator interface {
InvalidateAPIKeyRateLimit(ctx context.Context, keyID int64) error
}
type APIKeyService struct { type APIKeyService struct {
apiKeyRepo APIKeyRepository apiKeyRepo APIKeyRepository
userRepo UserRepository userRepo UserRepository
...@@ -127,6 +163,7 @@ type APIKeyService struct { ...@@ -127,6 +163,7 @@ type APIKeyService struct {
userSubRepo UserSubscriptionRepository userSubRepo UserSubscriptionRepository
userGroupRateRepo UserGroupRateRepository userGroupRateRepo UserGroupRateRepository
cache APIKeyCache cache APIKeyCache
rateLimitCacheInvalid RateLimitCacheInvalidator // optional: invalidate Redis rate limit cache
cfg *config.Config cfg *config.Config
authCacheL1 *ristretto.Cache authCacheL1 *ristretto.Cache
authCfg apiKeyAuthCacheConfig authCfg apiKeyAuthCacheConfig
...@@ -158,6 +195,12 @@ func NewAPIKeyService( ...@@ -158,6 +195,12 @@ func NewAPIKeyService(
return svc return svc
} }
// SetRateLimitCacheInvalidator sets the optional rate limit cache invalidator.
// Called after construction (e.g. in wire) to avoid circular dependencies.
func (s *APIKeyService) SetRateLimitCacheInvalidator(inv RateLimitCacheInvalidator) {
s.rateLimitCacheInvalid = inv
}
func (s *APIKeyService) compileAPIKeyIPRules(apiKey *APIKey) { func (s *APIKeyService) compileAPIKeyIPRules(apiKey *APIKey) {
if apiKey == nil { if apiKey == nil {
return return
...@@ -327,6 +370,9 @@ func (s *APIKeyService) Create(ctx context.Context, userID int64, req CreateAPIK ...@@ -327,6 +370,9 @@ func (s *APIKeyService) Create(ctx context.Context, userID int64, req CreateAPIK
IPBlacklist: req.IPBlacklist, IPBlacklist: req.IPBlacklist,
Quota: req.Quota, Quota: req.Quota,
QuotaUsed: 0, QuotaUsed: 0,
RateLimit5h: req.RateLimit5h,
RateLimit1d: req.RateLimit1d,
RateLimit7d: req.RateLimit7d,
} }
// Set expiration time if specified // Set expiration time if specified
...@@ -519,6 +565,26 @@ func (s *APIKeyService) Update(ctx context.Context, id int64, userID int64, req ...@@ -519,6 +565,26 @@ func (s *APIKeyService) Update(ctx context.Context, id int64, userID int64, req
apiKey.IPWhitelist = req.IPWhitelist apiKey.IPWhitelist = req.IPWhitelist
apiKey.IPBlacklist = req.IPBlacklist apiKey.IPBlacklist = req.IPBlacklist
// Update rate limit configuration
if req.RateLimit5h != nil {
apiKey.RateLimit5h = *req.RateLimit5h
}
if req.RateLimit1d != nil {
apiKey.RateLimit1d = *req.RateLimit1d
}
if req.RateLimit7d != nil {
apiKey.RateLimit7d = *req.RateLimit7d
}
resetRateLimit := req.ResetRateLimitUsage != nil && *req.ResetRateLimitUsage
if resetRateLimit {
apiKey.Usage5h = 0
apiKey.Usage1d = 0
apiKey.Usage7d = 0
apiKey.Window5hStart = nil
apiKey.Window1dStart = nil
apiKey.Window7dStart = nil
}
if err := s.apiKeyRepo.Update(ctx, apiKey); err != nil { if err := s.apiKeyRepo.Update(ctx, apiKey); err != nil {
return nil, fmt.Errorf("update api key: %w", err) return nil, fmt.Errorf("update api key: %w", err)
} }
...@@ -526,6 +592,11 @@ func (s *APIKeyService) Update(ctx context.Context, id int64, userID int64, req ...@@ -526,6 +592,11 @@ func (s *APIKeyService) Update(ctx context.Context, id int64, userID int64, req
s.InvalidateAuthCacheByKey(ctx, apiKey.Key) s.InvalidateAuthCacheByKey(ctx, apiKey.Key)
s.compileAPIKeyIPRules(apiKey) s.compileAPIKeyIPRules(apiKey)
// Invalidate Redis rate limit cache so reset takes effect immediately
if resetRateLimit && s.rateLimitCacheInvalid != nil {
_ = s.rateLimitCacheInvalid.InvalidateAPIKeyRateLimit(ctx, apiKey.ID)
}
return apiKey, nil return apiKey, nil
} }
...@@ -746,3 +817,11 @@ func (s *APIKeyService) UpdateQuotaUsed(ctx context.Context, apiKeyID int64, cos ...@@ -746,3 +817,11 @@ func (s *APIKeyService) UpdateQuotaUsed(ctx context.Context, apiKeyID int64, cos
return nil return nil
} }
// UpdateRateLimitUsage atomically increments rate limit usage counters in the DB.
func (s *APIKeyService) UpdateRateLimitUsage(ctx context.Context, apiKeyID int64, cost float64) error {
if cost <= 0 {
return nil
}
return s.apiKeyRepo.IncrementRateLimitUsage(ctx, apiKeyID, cost)
}
...@@ -40,6 +40,7 @@ const ( ...@@ -40,6 +40,7 @@ const (
cacheWriteSetSubscription cacheWriteSetSubscription
cacheWriteUpdateSubscriptionUsage cacheWriteUpdateSubscriptionUsage
cacheWriteDeductBalance cacheWriteDeductBalance
cacheWriteUpdateRateLimitUsage
) )
// 异步缓存写入工作池配置 // 异步缓存写入工作池配置
...@@ -68,17 +69,24 @@ type cacheWriteTask struct { ...@@ -68,17 +69,24 @@ type cacheWriteTask struct {
kind cacheWriteKind kind cacheWriteKind
userID int64 userID int64
groupID int64 groupID int64
apiKeyID int64
balance float64 balance float64
amount float64 amount float64
subscriptionData *subscriptionCacheData subscriptionData *subscriptionCacheData
} }
// apiKeyRateLimitLoader defines the interface for loading rate limit data from DB.
type apiKeyRateLimitLoader interface {
GetRateLimitData(ctx context.Context, keyID int64) (*APIKeyRateLimitData, error)
}
// BillingCacheService 计费缓存服务 // BillingCacheService 计费缓存服务
// 负责余额和订阅数据的缓存管理,提供高性能的计费资格检查 // 负责余额和订阅数据的缓存管理,提供高性能的计费资格检查
type BillingCacheService struct { type BillingCacheService struct {
cache BillingCache cache BillingCache
userRepo UserRepository userRepo UserRepository
subRepo UserSubscriptionRepository subRepo UserSubscriptionRepository
apiKeyRateLimitLoader apiKeyRateLimitLoader
cfg *config.Config cfg *config.Config
circuitBreaker *billingCircuitBreaker circuitBreaker *billingCircuitBreaker
...@@ -96,11 +104,12 @@ type BillingCacheService struct { ...@@ -96,11 +104,12 @@ type BillingCacheService struct {
} }
// NewBillingCacheService 创建计费缓存服务 // NewBillingCacheService 创建计费缓存服务
func NewBillingCacheService(cache BillingCache, userRepo UserRepository, subRepo UserSubscriptionRepository, cfg *config.Config) *BillingCacheService { func NewBillingCacheService(cache BillingCache, userRepo UserRepository, subRepo UserSubscriptionRepository, apiKeyRepo APIKeyRepository, cfg *config.Config) *BillingCacheService {
svc := &BillingCacheService{ svc := &BillingCacheService{
cache: cache, cache: cache,
userRepo: userRepo, userRepo: userRepo,
subRepo: subRepo, subRepo: subRepo,
apiKeyRateLimitLoader: apiKeyRepo,
cfg: cfg, cfg: cfg,
} }
svc.circuitBreaker = newBillingCircuitBreaker(cfg.Billing.CircuitBreaker) svc.circuitBreaker = newBillingCircuitBreaker(cfg.Billing.CircuitBreaker)
...@@ -188,6 +197,12 @@ func (s *BillingCacheService) cacheWriteWorker(ch <-chan cacheWriteTask) { ...@@ -188,6 +197,12 @@ func (s *BillingCacheService) cacheWriteWorker(ch <-chan cacheWriteTask) {
logger.LegacyPrintf("service.billing_cache", "Warning: deduct balance cache failed for user %d: %v", task.userID, err) logger.LegacyPrintf("service.billing_cache", "Warning: deduct balance cache failed for user %d: %v", task.userID, err)
} }
} }
case cacheWriteUpdateRateLimitUsage:
if s.cache != nil {
if err := s.cache.UpdateAPIKeyRateLimitUsage(ctx, task.apiKeyID, task.amount); err != nil {
logger.LegacyPrintf("service.billing_cache", "Warning: update rate limit usage cache failed for api key %d: %v", task.apiKeyID, err)
}
}
} }
cancel() cancel()
} }
...@@ -204,6 +219,8 @@ func cacheWriteKindName(kind cacheWriteKind) string { ...@@ -204,6 +219,8 @@ func cacheWriteKindName(kind cacheWriteKind) string {
return "update_subscription_usage" return "update_subscription_usage"
case cacheWriteDeductBalance: case cacheWriteDeductBalance:
return "deduct_balance" return "deduct_balance"
case cacheWriteUpdateRateLimitUsage:
return "update_rate_limit_usage"
default: default:
return "unknown" return "unknown"
} }
...@@ -476,6 +493,137 @@ func (s *BillingCacheService) InvalidateSubscription(ctx context.Context, userID ...@@ -476,6 +493,137 @@ func (s *BillingCacheService) InvalidateSubscription(ctx context.Context, userID
return nil return nil
} }
// ============================================
// API Key 限速缓存方法
// ============================================
// checkAPIKeyRateLimits checks rate limit windows for an API key.
// It loads usage from Redis cache (falling back to DB on cache miss),
// resets expired windows in-memory and triggers async DB reset,
// and returns an error if any window limit is exceeded.
func (s *BillingCacheService) checkAPIKeyRateLimits(ctx context.Context, apiKey *APIKey) error {
if s.cache == nil {
// No cache: fall back to reading from DB directly
if s.apiKeyRateLimitLoader == nil {
return nil
}
data, err := s.apiKeyRateLimitLoader.GetRateLimitData(ctx, apiKey.ID)
if err != nil {
return nil // Don't block requests on DB errors
}
return s.evaluateRateLimits(ctx, apiKey, data.Usage5h, data.Usage1d, data.Usage7d,
data.Window5hStart, data.Window1dStart, data.Window7dStart)
}
cacheData, err := s.cache.GetAPIKeyRateLimit(ctx, apiKey.ID)
if err != nil {
// Cache miss: load from DB and populate cache
if s.apiKeyRateLimitLoader == nil {
return nil
}
dbData, dbErr := s.apiKeyRateLimitLoader.GetRateLimitData(ctx, apiKey.ID)
if dbErr != nil {
return nil // Don't block requests on DB errors
}
// Build cache entry from DB data
cacheEntry := &APIKeyRateLimitCacheData{
Usage5h: dbData.Usage5h,
Usage1d: dbData.Usage1d,
Usage7d: dbData.Usage7d,
}
if dbData.Window5hStart != nil {
cacheEntry.Window5h = dbData.Window5hStart.Unix()
}
if dbData.Window1dStart != nil {
cacheEntry.Window1d = dbData.Window1dStart.Unix()
}
if dbData.Window7dStart != nil {
cacheEntry.Window7d = dbData.Window7dStart.Unix()
}
_ = s.cache.SetAPIKeyRateLimit(ctx, apiKey.ID, cacheEntry)
cacheData = cacheEntry
}
var w5h, w1d, w7d *time.Time
if cacheData.Window5h > 0 {
t := time.Unix(cacheData.Window5h, 0)
w5h = &t
}
if cacheData.Window1d > 0 {
t := time.Unix(cacheData.Window1d, 0)
w1d = &t
}
if cacheData.Window7d > 0 {
t := time.Unix(cacheData.Window7d, 0)
w7d = &t
}
return s.evaluateRateLimits(ctx, apiKey, cacheData.Usage5h, cacheData.Usage1d, cacheData.Usage7d, w5h, w1d, w7d)
}
// evaluateRateLimits checks usage against limits, triggering async resets for expired windows.
func (s *BillingCacheService) evaluateRateLimits(ctx context.Context, apiKey *APIKey, usage5h, usage1d, usage7d float64, w5h, w1d, w7d *time.Time) error {
needsReset := false
// Reset expired windows in-memory for check purposes
if w5h != nil && time.Since(*w5h) >= 5*time.Hour {
usage5h = 0
needsReset = true
}
if w1d != nil && time.Since(*w1d) >= 24*time.Hour {
usage1d = 0
needsReset = true
}
if w7d != nil && time.Since(*w7d) >= 7*24*time.Hour {
usage7d = 0
needsReset = true
}
// Trigger async DB reset if any window expired
if needsReset {
keyID := apiKey.ID
go func() {
resetCtx, cancel := context.WithTimeout(context.Background(), cacheWriteTimeout)
defer cancel()
if s.apiKeyRateLimitLoader != nil {
// Use the repo directly - reset then reload cache
if loader, ok := s.apiKeyRateLimitLoader.(interface {
ResetRateLimitWindows(ctx context.Context, id int64) error
}); ok {
_ = loader.ResetRateLimitWindows(resetCtx, keyID)
}
}
// Invalidate cache so next request loads fresh data
if s.cache != nil {
_ = s.cache.InvalidateAPIKeyRateLimit(resetCtx, keyID)
}
}()
}
// Check limits
if apiKey.RateLimit5h > 0 && usage5h >= apiKey.RateLimit5h {
return ErrAPIKeyRateLimit5hExceeded
}
if apiKey.RateLimit1d > 0 && usage1d >= apiKey.RateLimit1d {
return ErrAPIKeyRateLimit1dExceeded
}
if apiKey.RateLimit7d > 0 && usage7d >= apiKey.RateLimit7d {
return ErrAPIKeyRateLimit7dExceeded
}
return nil
}
// QueueUpdateAPIKeyRateLimitUsage asynchronously updates rate limit usage in the cache.
func (s *BillingCacheService) QueueUpdateAPIKeyRateLimitUsage(apiKeyID int64, cost float64) {
if s.cache == nil {
return
}
s.enqueueCacheWrite(cacheWriteTask{
kind: cacheWriteUpdateRateLimitUsage,
apiKeyID: apiKeyID,
amount: cost,
})
}
// ============================================ // ============================================
// 统一检查方法 // 统一检查方法
// ============================================ // ============================================
...@@ -496,10 +644,23 @@ func (s *BillingCacheService) CheckBillingEligibility(ctx context.Context, user ...@@ -496,10 +644,23 @@ func (s *BillingCacheService) CheckBillingEligibility(ctx context.Context, user
isSubscriptionMode := group != nil && group.IsSubscriptionType() && subscription != nil isSubscriptionMode := group != nil && group.IsSubscriptionType() && subscription != nil
if isSubscriptionMode { if isSubscriptionMode {
return s.checkSubscriptionEligibility(ctx, user.ID, group, subscription) if err := s.checkSubscriptionEligibility(ctx, user.ID, group, subscription); err != nil {
return err
}
} else {
if err := s.checkBalanceEligibility(ctx, user.ID); err != nil {
return err
}
} }
return s.checkBalanceEligibility(ctx, user.ID) // Check API Key rate limits (applies to both billing modes)
if apiKey != nil && apiKey.HasRateLimits() {
if err := s.checkAPIKeyRateLimits(ctx, apiKey); err != nil {
return err
}
}
return nil
} }
// checkBalanceEligibility 检查余额模式资格 // checkBalanceEligibility 检查余额模式资格
......
...@@ -52,9 +52,25 @@ func (b *billingCacheWorkerStub) InvalidateSubscriptionCache(ctx context.Context ...@@ -52,9 +52,25 @@ func (b *billingCacheWorkerStub) InvalidateSubscriptionCache(ctx context.Context
return nil return nil
} }
func (b *billingCacheWorkerStub) GetAPIKeyRateLimit(ctx context.Context, keyID int64) (*APIKeyRateLimitCacheData, error) {
return nil, errors.New("not implemented")
}
func (b *billingCacheWorkerStub) SetAPIKeyRateLimit(ctx context.Context, keyID int64, data *APIKeyRateLimitCacheData) error {
return nil
}
func (b *billingCacheWorkerStub) UpdateAPIKeyRateLimitUsage(ctx context.Context, keyID int64, cost float64) error {
return nil
}
func (b *billingCacheWorkerStub) InvalidateAPIKeyRateLimit(ctx context.Context, keyID int64) error {
return nil
}
func TestBillingCacheServiceQueueHighLoad(t *testing.T) { func TestBillingCacheServiceQueueHighLoad(t *testing.T) {
cache := &billingCacheWorkerStub{} cache := &billingCacheWorkerStub{}
svc := NewBillingCacheService(cache, nil, nil, &config.Config{}) svc := NewBillingCacheService(cache, nil, nil, nil, &config.Config{})
t.Cleanup(svc.Stop) t.Cleanup(svc.Stop)
start := time.Now() start := time.Now()
...@@ -76,7 +92,7 @@ func TestBillingCacheServiceQueueHighLoad(t *testing.T) { ...@@ -76,7 +92,7 @@ func TestBillingCacheServiceQueueHighLoad(t *testing.T) {
func TestBillingCacheServiceEnqueueAfterStopReturnsFalse(t *testing.T) { func TestBillingCacheServiceEnqueueAfterStopReturnsFalse(t *testing.T) {
cache := &billingCacheWorkerStub{} cache := &billingCacheWorkerStub{}
svc := NewBillingCacheService(cache, nil, nil, &config.Config{}) svc := NewBillingCacheService(cache, nil, nil, nil, &config.Config{})
svc.Stop() svc.Stop()
enqueued := svc.enqueueCacheWrite(cacheWriteTask{ enqueued := svc.enqueueCacheWrite(cacheWriteTask{
......
...@@ -10,6 +10,16 @@ import ( ...@@ -10,6 +10,16 @@ import (
"github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/config"
) )
// APIKeyRateLimitCacheData holds rate limit usage data cached in Redis.
type APIKeyRateLimitCacheData struct {
Usage5h float64 `json:"usage_5h"`
Usage1d float64 `json:"usage_1d"`
Usage7d float64 `json:"usage_7d"`
Window5h int64 `json:"window_5h"` // unix timestamp, 0 = not started
Window1d int64 `json:"window_1d"`
Window7d int64 `json:"window_7d"`
}
// BillingCache defines cache operations for billing service // BillingCache defines cache operations for billing service
type BillingCache interface { type BillingCache interface {
// Balance operations // Balance operations
...@@ -23,6 +33,12 @@ type BillingCache interface { ...@@ -23,6 +33,12 @@ type BillingCache interface {
SetSubscriptionCache(ctx context.Context, userID, groupID int64, data *SubscriptionCacheData) error SetSubscriptionCache(ctx context.Context, userID, groupID int64, data *SubscriptionCacheData) error
UpdateSubscriptionUsage(ctx context.Context, userID, groupID int64, cost float64) error UpdateSubscriptionUsage(ctx context.Context, userID, groupID int64, cost float64) error
InvalidateSubscriptionCache(ctx context.Context, userID, groupID int64) error InvalidateSubscriptionCache(ctx context.Context, userID, groupID int64) error
// API Key rate limit operations
GetAPIKeyRateLimit(ctx context.Context, keyID int64) (*APIKeyRateLimitCacheData, error)
SetAPIKeyRateLimit(ctx context.Context, keyID int64, data *APIKeyRateLimitCacheData) error
UpdateAPIKeyRateLimitUsage(ctx context.Context, keyID int64, cost float64) error
InvalidateAPIKeyRateLimit(ctx context.Context, keyID int64) error
} }
// ModelPricing 模型价格配置(per-token价格,与LiteLLM格式一致) // ModelPricing 模型价格配置(per-token价格,与LiteLLM格式一致)
......
...@@ -6361,9 +6361,10 @@ type RecordUsageInput struct { ...@@ -6361,9 +6361,10 @@ type RecordUsageInput struct {
APIKeyService APIKeyQuotaUpdater // 可选:用于更新API Key配额 APIKeyService APIKeyQuotaUpdater // 可选:用于更新API Key配额
} }
// APIKeyQuotaUpdater defines the interface for updating API Key quota // APIKeyQuotaUpdater defines the interface for updating API Key quota and rate limit usage
type APIKeyQuotaUpdater interface { type APIKeyQuotaUpdater interface {
UpdateQuotaUsed(ctx context.Context, apiKeyID int64, cost float64) error UpdateQuotaUsed(ctx context.Context, apiKeyID int64, cost float64) error
UpdateRateLimitUsage(ctx context.Context, apiKeyID int64, cost float64) error
} }
// RecordUsage 记录使用量并扣费(或更新订阅用量) // RecordUsage 记录使用量并扣费(或更新订阅用量)
...@@ -6557,6 +6558,14 @@ func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInpu ...@@ -6557,6 +6558,14 @@ func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInpu
} }
} }
// Update API Key rate limit usage
if shouldBill && cost.ActualCost > 0 && apiKey.HasRateLimits() && input.APIKeyService != nil {
if err := input.APIKeyService.UpdateRateLimitUsage(ctx, apiKey.ID, cost.ActualCost); err != nil {
logger.LegacyPrintf("service.gateway", "Update API key rate limit usage failed: %v", err)
}
s.billingCacheService.QueueUpdateAPIKeyRateLimitUsage(apiKey.ID, cost.ActualCost)
}
// Schedule batch update for account last_used_at // Schedule batch update for account last_used_at
s.deferredService.ScheduleLastUsedUpdate(account.ID) s.deferredService.ScheduleLastUsedUpdate(account.ID)
...@@ -6746,6 +6755,14 @@ func (s *GatewayService) RecordUsageWithLongContext(ctx context.Context, input * ...@@ -6746,6 +6755,14 @@ func (s *GatewayService) RecordUsageWithLongContext(ctx context.Context, input *
} }
} }
// Update API Key rate limit usage
if shouldBill && cost.ActualCost > 0 && apiKey.HasRateLimits() && input.APIKeyService != nil {
if err := input.APIKeyService.UpdateRateLimitUsage(ctx, apiKey.ID, cost.ActualCost); err != nil {
logger.LegacyPrintf("service.gateway", "Update API key rate limit usage failed: %v", err)
}
s.billingCacheService.QueueUpdateAPIKeyRateLimitUsage(apiKey.ID, cost.ActualCost)
}
// Schedule batch update for account last_used_at // Schedule batch update for account last_used_at
s.deferredService.ScheduleLastUsedUpdate(account.ID) s.deferredService.ScheduleLastUsedUpdate(account.ID)
......
...@@ -3492,6 +3492,14 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec ...@@ -3492,6 +3492,14 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec
} }
} }
// Update API Key rate limit usage
if shouldBill && cost.ActualCost > 0 && apiKey.HasRateLimits() && input.APIKeyService != nil {
if err := input.APIKeyService.UpdateRateLimitUsage(ctx, apiKey.ID, cost.ActualCost); err != nil {
logger.LegacyPrintf("service.openai_gateway", "Update API key rate limit usage failed: %v", err)
}
s.billingCacheService.QueueUpdateAPIKeyRateLimitUsage(apiKey.ID, cost.ActualCost)
}
// Schedule batch update for account last_used_at // Schedule batch update for account last_used_at
s.deferredService.ScheduleLastUsedUpdate(account.ID) s.deferredService.ScheduleLastUsedUpdate(account.ID)
......
...@@ -46,6 +46,7 @@ export async function getById(id: number): Promise<ApiKey> { ...@@ -46,6 +46,7 @@ export async function getById(id: number): Promise<ApiKey> {
* @param ipBlacklist - Optional IP blacklist * @param ipBlacklist - Optional IP blacklist
* @param quota - Optional quota limit in USD (0 = unlimited) * @param quota - Optional quota limit in USD (0 = unlimited)
* @param expiresInDays - Optional days until expiry (undefined = never expires) * @param expiresInDays - Optional days until expiry (undefined = never expires)
* @param rateLimitData - Optional rate limit fields
* @returns Created API key * @returns Created API key
*/ */
export async function create( export async function create(
...@@ -55,7 +56,8 @@ export async function create( ...@@ -55,7 +56,8 @@ export async function create(
ipWhitelist?: string[], ipWhitelist?: string[],
ipBlacklist?: string[], ipBlacklist?: string[],
quota?: number, quota?: number,
expiresInDays?: number expiresInDays?: number,
rateLimitData?: { rate_limit_5h?: number; rate_limit_1d?: number; rate_limit_7d?: number }
): Promise<ApiKey> { ): Promise<ApiKey> {
const payload: CreateApiKeyRequest = { name } const payload: CreateApiKeyRequest = { name }
if (groupId !== undefined) { if (groupId !== undefined) {
...@@ -76,6 +78,15 @@ export async function create( ...@@ -76,6 +78,15 @@ export async function create(
if (expiresInDays !== undefined && expiresInDays > 0) { if (expiresInDays !== undefined && expiresInDays > 0) {
payload.expires_in_days = expiresInDays payload.expires_in_days = expiresInDays
} }
if (rateLimitData?.rate_limit_5h && rateLimitData.rate_limit_5h > 0) {
payload.rate_limit_5h = rateLimitData.rate_limit_5h
}
if (rateLimitData?.rate_limit_1d && rateLimitData.rate_limit_1d > 0) {
payload.rate_limit_1d = rateLimitData.rate_limit_1d
}
if (rateLimitData?.rate_limit_7d && rateLimitData.rate_limit_7d > 0) {
payload.rate_limit_7d = rateLimitData.rate_limit_7d
}
const { data } = await apiClient.post<ApiKey>('/keys', payload) const { data } = await apiClient.post<ApiKey>('/keys', payload)
return data return data
......
...@@ -560,6 +560,19 @@ export default { ...@@ -560,6 +560,19 @@ export default {
resetQuotaConfirmMessage: 'Are you sure you want to reset the used quota (${used}) for key "{name}" to 0? This action cannot be undone.', resetQuotaConfirmMessage: 'Are you sure you want to reset the used quota (${used}) for key "{name}" to 0? This action cannot be undone.',
quotaResetSuccess: 'Quota reset successfully', quotaResetSuccess: 'Quota reset successfully',
failedToResetQuota: 'Failed to reset quota', failedToResetQuota: 'Failed to reset quota',
rateLimitColumn: 'Rate Limit',
rateLimitSection: 'Rate Limit',
resetUsage: 'Reset',
rateLimit5h: '5-Hour Limit (USD)',
rateLimit1d: 'Daily Limit (USD)',
rateLimit7d: '7-Day Limit (USD)',
rateLimitHint: 'Set the maximum spending for this key within each time window. 0 = unlimited.',
rateLimitUsage: 'Rate Limit Usage',
resetRateLimitUsage: 'Reset Rate Limit Usage',
resetRateLimitTitle: 'Confirm Reset Rate Limit',
resetRateLimitConfirmMessage: 'Are you sure you want to reset the rate limit usage for key "{name}"? All time window usage will be reset to zero. This action cannot be undone.',
rateLimitResetSuccess: 'Rate limit usage reset successfully',
failedToResetRateLimit: 'Failed to reset rate limit usage',
expiration: 'Expiration', expiration: 'Expiration',
expiresInDays: '{days} days', expiresInDays: '{days} days',
extendDays: '+{days} days', extendDays: '+{days} days',
......
...@@ -566,6 +566,19 @@ export default { ...@@ -566,6 +566,19 @@ export default {
resetQuotaConfirmMessage: '确定要将密钥 "{name}" 的已用额度(${used})重置为 0 吗?此操作不可撤销。', resetQuotaConfirmMessage: '确定要将密钥 "{name}" 的已用额度(${used})重置为 0 吗?此操作不可撤销。',
quotaResetSuccess: '额度重置成功', quotaResetSuccess: '额度重置成功',
failedToResetQuota: '重置额度失败', failedToResetQuota: '重置额度失败',
rateLimitColumn: '速率限制',
rateLimitSection: '速率限制',
resetUsage: '重置',
rateLimit5h: '5小时限额 (USD)',
rateLimit1d: '日限额 (USD)',
rateLimit7d: '7天限额 (USD)',
rateLimitHint: '设置此密钥在指定时间窗口内的最大消费额。0 = 无限制。',
rateLimitUsage: '速率限制用量',
resetRateLimitUsage: '重置速率限制用量',
resetRateLimitTitle: '确认重置速率限制',
resetRateLimitConfirmMessage: '确定要重置密钥 "{name}" 的速率限制用量吗?所有时间窗口的已用额度将归零。此操作不可撤销。',
rateLimitResetSuccess: '速率限制已重置',
failedToResetRateLimit: '重置速率限制失败',
expiration: '密钥有效期', expiration: '密钥有效期',
expiresInDays: '{days} 天', expiresInDays: '{days} 天',
extendDays: '+{days} 天', extendDays: '+{days} 天',
......
...@@ -421,6 +421,15 @@ export interface ApiKey { ...@@ -421,6 +421,15 @@ export interface ApiKey {
created_at: string created_at: string
updated_at: string updated_at: string
group?: Group group?: Group
rate_limit_5h: number
rate_limit_1d: number
rate_limit_7d: number
usage_5h: number
usage_1d: number
usage_7d: number
window_5h_start: string | null
window_1d_start: string | null
window_7d_start: string | null
} }
export interface CreateApiKeyRequest { export interface CreateApiKeyRequest {
...@@ -431,6 +440,9 @@ export interface CreateApiKeyRequest { ...@@ -431,6 +440,9 @@ export interface CreateApiKeyRequest {
ip_blacklist?: string[] ip_blacklist?: string[]
quota?: number // Quota limit in USD (0 = unlimited) quota?: number // Quota limit in USD (0 = unlimited)
expires_in_days?: number // Days until expiry (null = never expires) expires_in_days?: number // Days until expiry (null = never expires)
rate_limit_5h?: number
rate_limit_1d?: number
rate_limit_7d?: number
} }
export interface UpdateApiKeyRequest { export interface UpdateApiKeyRequest {
...@@ -442,6 +454,10 @@ export interface UpdateApiKeyRequest { ...@@ -442,6 +454,10 @@ export interface UpdateApiKeyRequest {
quota?: number // Quota limit in USD (null = no change, 0 = unlimited) quota?: number // Quota limit in USD (null = no change, 0 = unlimited)
expires_at?: string | null // Expiration time (null = no change) expires_at?: string | null // Expiration time (null = no change)
reset_quota?: boolean // Reset quota_used to 0 reset_quota?: boolean // Reset quota_used to 0
rate_limit_5h?: number
rate_limit_1d?: number
rate_limit_7d?: number
reset_rate_limit_usage?: boolean
} }
export interface CreateGroupRequest { export interface CreateGroupRequest {
......
...@@ -137,6 +137,97 @@ ...@@ -137,6 +137,97 @@
</div> </div>
</template> </template>
<template #cell-rate_limit="{ row }">
<div v-if="row.rate_limit_5h > 0 || row.rate_limit_1d > 0 || row.rate_limit_7d > 0" class="space-y-1.5 min-w-[140px]">
<!-- 5h window -->
<div v-if="row.rate_limit_5h > 0">
<div class="flex items-center justify-between text-xs">
<span class="text-gray-500 dark:text-gray-400">5h</span>
<span :class="[
'font-medium tabular-nums',
row.usage_5h >= row.rate_limit_5h ? 'text-red-500' :
row.usage_5h >= row.rate_limit_5h * 0.8 ? 'text-yellow-500' :
'text-gray-700 dark:text-gray-300'
]">
${{ row.usage_5h?.toFixed(2) || '0.00' }}/${{ row.rate_limit_5h?.toFixed(2) }}
</span>
</div>
<div class="h-1 w-full overflow-hidden rounded-full bg-gray-200 dark:bg-dark-600">
<div
:class="[
'h-full rounded-full transition-all',
row.usage_5h >= row.rate_limit_5h ? 'bg-red-500' :
row.usage_5h >= row.rate_limit_5h * 0.8 ? 'bg-yellow-500' :
'bg-emerald-500'
]"
:style="{ width: Math.min((row.usage_5h / row.rate_limit_5h) * 100, 100) + '%' }"
/>
</div>
</div>
<!-- 1d window -->
<div v-if="row.rate_limit_1d > 0">
<div class="flex items-center justify-between text-xs">
<span class="text-gray-500 dark:text-gray-400">1d</span>
<span :class="[
'font-medium tabular-nums',
row.usage_1d >= row.rate_limit_1d ? 'text-red-500' :
row.usage_1d >= row.rate_limit_1d * 0.8 ? 'text-yellow-500' :
'text-gray-700 dark:text-gray-300'
]">
${{ row.usage_1d?.toFixed(2) || '0.00' }}/${{ row.rate_limit_1d?.toFixed(2) }}
</span>
</div>
<div class="h-1 w-full overflow-hidden rounded-full bg-gray-200 dark:bg-dark-600">
<div
:class="[
'h-full rounded-full transition-all',
row.usage_1d >= row.rate_limit_1d ? 'bg-red-500' :
row.usage_1d >= row.rate_limit_1d * 0.8 ? 'bg-yellow-500' :
'bg-emerald-500'
]"
:style="{ width: Math.min((row.usage_1d / row.rate_limit_1d) * 100, 100) + '%' }"
/>
</div>
</div>
<!-- 7d window -->
<div v-if="row.rate_limit_7d > 0">
<div class="flex items-center justify-between text-xs">
<span class="text-gray-500 dark:text-gray-400">7d</span>
<span :class="[
'font-medium tabular-nums',
row.usage_7d >= row.rate_limit_7d ? 'text-red-500' :
row.usage_7d >= row.rate_limit_7d * 0.8 ? 'text-yellow-500' :
'text-gray-700 dark:text-gray-300'
]">
${{ row.usage_7d?.toFixed(2) || '0.00' }}/${{ row.rate_limit_7d?.toFixed(2) }}
</span>
</div>
<div class="h-1 w-full overflow-hidden rounded-full bg-gray-200 dark:bg-dark-600">
<div
:class="[
'h-full rounded-full transition-all',
row.usage_7d >= row.rate_limit_7d ? 'bg-red-500' :
row.usage_7d >= row.rate_limit_7d * 0.8 ? 'bg-yellow-500' :
'bg-emerald-500'
]"
:style="{ width: Math.min((row.usage_7d / row.rate_limit_7d) * 100, 100) + '%' }"
/>
</div>
</div>
<!-- Reset button -->
<button
v-if="row.usage_5h > 0 || row.usage_1d > 0 || row.usage_7d > 0"
@click.stop="confirmResetRateLimitFromTable(row)"
class="mt-0.5 inline-flex items-center gap-1 rounded px-1.5 py-0.5 text-xs text-gray-500 transition-colors hover:bg-gray-100 hover:text-primary-600 dark:hover:bg-dark-700 dark:hover:text-primary-400"
:title="t('keys.resetRateLimitUsage')"
>
<Icon name="refresh" size="xs" />
{{ t('keys.resetUsage') }}
</button>
</div>
<span v-else class="text-sm text-gray-400 dark:text-dark-500">-</span>
</template>
<template #cell-expires_at="{ value }"> <template #cell-expires_at="{ value }">
<span v-if="value" :class="[ <span v-if="value" :class="[
'text-sm', 'text-sm',
...@@ -452,6 +543,180 @@ ...@@ -452,6 +543,180 @@
</div> </div>
</div> </div>
<!-- Rate Limit Section -->
<div class="space-y-3">
<div class="flex items-center justify-between">
<label class="input-label mb-0">{{ t('keys.rateLimitSection') }}</label>
<button
type="button"
@click="formData.enable_rate_limit = !formData.enable_rate_limit"
:class="[
'relative inline-flex h-5 w-9 flex-shrink-0 cursor-pointer rounded-full border-2 border-transparent transition-colors duration-200 ease-in-out focus:outline-none',
formData.enable_rate_limit ? 'bg-primary-600' : 'bg-gray-200 dark:bg-dark-600'
]"
>
<span
:class="[
'pointer-events-none inline-block h-4 w-4 transform rounded-full bg-white shadow ring-0 transition duration-200 ease-in-out',
formData.enable_rate_limit ? 'translate-x-4' : 'translate-x-0'
]"
/>
</button>
</div>
<div v-if="formData.enable_rate_limit" class="space-y-4 pt-2">
<p class="input-hint -mt-2">{{ t('keys.rateLimitHint') }}</p>
<!-- 5-Hour Limit -->
<div>
<label class="input-label">{{ t('keys.rateLimit5h') }}</label>
<div class="relative">
<span class="absolute left-3 top-1/2 -translate-y-1/2 text-gray-500">$</span>
<input
v-model.number="formData.rate_limit_5h"
type="number"
step="0.01"
min="0"
class="input pl-7"
:placeholder="'0'"
/>
</div>
<!-- Usage info (edit mode only) -->
<div v-if="showEditModal && selectedKey && selectedKey.rate_limit_5h > 0" class="mt-2">
<div class="flex items-center gap-2">
<div class="flex-1 rounded-lg bg-gray-100 px-3 py-2 dark:bg-dark-700 text-sm">
<span :class="[
'font-medium',
selectedKey.usage_5h >= selectedKey.rate_limit_5h ? 'text-red-500' :
selectedKey.usage_5h >= selectedKey.rate_limit_5h * 0.8 ? 'text-yellow-500' :
'text-gray-900 dark:text-white'
]">
${{ selectedKey.usage_5h?.toFixed(4) || '0.0000' }}
</span>
<span class="mx-2 text-gray-400">/</span>
<span class="text-gray-500 dark:text-gray-400">
${{ selectedKey.rate_limit_5h?.toFixed(2) || '0.00' }}
</span>
</div>
</div>
<div class="mt-1 h-1.5 w-full overflow-hidden rounded-full bg-gray-200 dark:bg-dark-600">
<div
:class="[
'h-full rounded-full transition-all',
selectedKey.usage_5h >= selectedKey.rate_limit_5h ? 'bg-red-500' :
selectedKey.usage_5h >= selectedKey.rate_limit_5h * 0.8 ? 'bg-yellow-500' :
'bg-green-500'
]"
:style="{ width: Math.min((selectedKey.usage_5h / selectedKey.rate_limit_5h) * 100, 100) + '%' }"
/>
</div>
</div>
</div>
<!-- Daily Limit -->
<div>
<label class="input-label">{{ t('keys.rateLimit1d') }}</label>
<div class="relative">
<span class="absolute left-3 top-1/2 -translate-y-1/2 text-gray-500">$</span>
<input
v-model.number="formData.rate_limit_1d"
type="number"
step="0.01"
min="0"
class="input pl-7"
:placeholder="'0'"
/>
</div>
<!-- Usage info (edit mode only) -->
<div v-if="showEditModal && selectedKey && selectedKey.rate_limit_1d > 0" class="mt-2">
<div class="flex items-center gap-2">
<div class="flex-1 rounded-lg bg-gray-100 px-3 py-2 dark:bg-dark-700 text-sm">
<span :class="[
'font-medium',
selectedKey.usage_1d >= selectedKey.rate_limit_1d ? 'text-red-500' :
selectedKey.usage_1d >= selectedKey.rate_limit_1d * 0.8 ? 'text-yellow-500' :
'text-gray-900 dark:text-white'
]">
${{ selectedKey.usage_1d?.toFixed(4) || '0.0000' }}
</span>
<span class="mx-2 text-gray-400">/</span>
<span class="text-gray-500 dark:text-gray-400">
${{ selectedKey.rate_limit_1d?.toFixed(2) || '0.00' }}
</span>
</div>
</div>
<div class="mt-1 h-1.5 w-full overflow-hidden rounded-full bg-gray-200 dark:bg-dark-600">
<div
:class="[
'h-full rounded-full transition-all',
selectedKey.usage_1d >= selectedKey.rate_limit_1d ? 'bg-red-500' :
selectedKey.usage_1d >= selectedKey.rate_limit_1d * 0.8 ? 'bg-yellow-500' :
'bg-green-500'
]"
:style="{ width: Math.min((selectedKey.usage_1d / selectedKey.rate_limit_1d) * 100, 100) + '%' }"
/>
</div>
</div>
</div>
<!-- 7-Day Limit -->
<div>
<label class="input-label">{{ t('keys.rateLimit7d') }}</label>
<div class="relative">
<span class="absolute left-3 top-1/2 -translate-y-1/2 text-gray-500">$</span>
<input
v-model.number="formData.rate_limit_7d"
type="number"
step="0.01"
min="0"
class="input pl-7"
:placeholder="'0'"
/>
</div>
<!-- Usage info (edit mode only) -->
<div v-if="showEditModal && selectedKey && selectedKey.rate_limit_7d > 0" class="mt-2">
<div class="flex items-center gap-2">
<div class="flex-1 rounded-lg bg-gray-100 px-3 py-2 dark:bg-dark-700 text-sm">
<span :class="[
'font-medium',
selectedKey.usage_7d >= selectedKey.rate_limit_7d ? 'text-red-500' :
selectedKey.usage_7d >= selectedKey.rate_limit_7d * 0.8 ? 'text-yellow-500' :
'text-gray-900 dark:text-white'
]">
${{ selectedKey.usage_7d?.toFixed(4) || '0.0000' }}
</span>
<span class="mx-2 text-gray-400">/</span>
<span class="text-gray-500 dark:text-gray-400">
${{ selectedKey.rate_limit_7d?.toFixed(2) || '0.00' }}
</span>
</div>
</div>
<div class="mt-1 h-1.5 w-full overflow-hidden rounded-full bg-gray-200 dark:bg-dark-600">
<div
:class="[
'h-full rounded-full transition-all',
selectedKey.usage_7d >= selectedKey.rate_limit_7d ? 'bg-red-500' :
selectedKey.usage_7d >= selectedKey.rate_limit_7d * 0.8 ? 'bg-yellow-500' :
'bg-green-500'
]"
:style="{ width: Math.min((selectedKey.usage_7d / selectedKey.rate_limit_7d) * 100, 100) + '%' }"
/>
</div>
</div>
</div>
<!-- Reset Rate Limit button (edit mode only) -->
<div v-if="showEditModal && selectedKey && (selectedKey.rate_limit_5h > 0 || selectedKey.rate_limit_1d > 0 || selectedKey.rate_limit_7d > 0)">
<button
type="button"
@click="confirmResetRateLimit"
class="btn btn-secondary text-sm"
>
{{ t('keys.resetRateLimitUsage') }}
</button>
</div>
</div>
</div>
<!-- Expiration Section --> <!-- Expiration Section -->
<div class="space-y-3"> <div class="space-y-3">
<div class="flex items-center justify-between"> <div class="flex items-center justify-between">
...@@ -593,6 +858,18 @@ ...@@ -593,6 +858,18 @@
@cancel="showResetQuotaDialog = false" @cancel="showResetQuotaDialog = false"
/> />
<!-- Reset Rate Limit Confirmation Dialog -->
<ConfirmDialog
:show="showResetRateLimitDialog"
:title="t('keys.resetRateLimitTitle')"
:message="t('keys.resetRateLimitConfirmMessage', { name: selectedKey?.name })"
:confirm-text="t('keys.reset')"
:cancel-text="t('common.cancel')"
:danger="true"
@confirm="resetRateLimitUsage"
@cancel="showResetRateLimitDialog = false"
/>
<!-- Use Key Modal --> <!-- Use Key Modal -->
<UseKeyModal <UseKeyModal
:show="showUseKeyModal" :show="showUseKeyModal"
...@@ -743,6 +1020,7 @@ const columns = computed<Column[]>(() => [ ...@@ -743,6 +1020,7 @@ const columns = computed<Column[]>(() => [
{ key: 'key', label: t('keys.apiKey'), sortable: false }, { key: 'key', label: t('keys.apiKey'), sortable: false },
{ key: 'group', label: t('keys.group'), sortable: false }, { key: 'group', label: t('keys.group'), sortable: false },
{ key: 'usage', label: t('keys.usage'), sortable: false }, { key: 'usage', label: t('keys.usage'), sortable: false },
{ key: 'rate_limit', label: t('keys.rateLimitColumn'), sortable: false },
{ key: 'expires_at', label: t('keys.expiresAt'), sortable: true }, { key: 'expires_at', label: t('keys.expiresAt'), sortable: true },
{ key: 'status', label: t('common.status'), sortable: true }, { key: 'status', label: t('common.status'), sortable: true },
{ key: 'last_used_at', label: t('keys.lastUsedAt'), sortable: true }, { key: 'last_used_at', label: t('keys.lastUsedAt'), sortable: true },
...@@ -768,6 +1046,7 @@ const showCreateModal = ref(false) ...@@ -768,6 +1046,7 @@ const showCreateModal = ref(false)
const showEditModal = ref(false) const showEditModal = ref(false)
const showDeleteDialog = ref(false) const showDeleteDialog = ref(false)
const showResetQuotaDialog = ref(false) const showResetQuotaDialog = ref(false)
const showResetRateLimitDialog = ref(false)
const showUseKeyModal = ref(false) const showUseKeyModal = ref(false)
const showCcsClientSelect = ref(false) const showCcsClientSelect = ref(false)
const pendingCcsRow = ref<ApiKey | null>(null) const pendingCcsRow = ref<ApiKey | null>(null)
...@@ -806,6 +1085,11 @@ const formData = ref({ ...@@ -806,6 +1085,11 @@ const formData = ref({
// Quota settings (empty = unlimited) // Quota settings (empty = unlimited)
enable_quota: false, enable_quota: false,
quota: null as number | null, quota: null as number | null,
// Rate limit settings
enable_rate_limit: false,
rate_limit_5h: null as number | null,
rate_limit_1d: null as number | null,
rate_limit_7d: null as number | null,
enable_expiration: false, enable_expiration: false,
expiration_preset: '30' as '7' | '30' | '90' | 'custom', expiration_preset: '30' as '7' | '30' | '90' | 'custom',
expiration_date: '' expiration_date: ''
...@@ -966,6 +1250,10 @@ const editKey = (key: ApiKey) => { ...@@ -966,6 +1250,10 @@ const editKey = (key: ApiKey) => {
ip_blacklist: (key.ip_blacklist || []).join('\n'), ip_blacklist: (key.ip_blacklist || []).join('\n'),
enable_quota: key.quota > 0, enable_quota: key.quota > 0,
quota: key.quota > 0 ? key.quota : null, quota: key.quota > 0 ? key.quota : null,
enable_rate_limit: (key.rate_limit_5h > 0) || (key.rate_limit_1d > 0) || (key.rate_limit_7d > 0),
rate_limit_5h: key.rate_limit_5h || null,
rate_limit_1d: key.rate_limit_1d || null,
rate_limit_7d: key.rate_limit_7d || null,
enable_expiration: hasExpiration, enable_expiration: hasExpiration,
expiration_preset: 'custom', expiration_preset: 'custom',
expiration_date: key.expires_at ? formatDateTimeLocal(key.expires_at) : '' expiration_date: key.expires_at ? formatDateTimeLocal(key.expires_at) : ''
...@@ -1078,6 +1366,13 @@ const handleSubmit = async () => { ...@@ -1078,6 +1366,13 @@ const handleSubmit = async () => {
expiresAt = '' expiresAt = ''
} }
// Calculate rate limit values (send 0 when toggle is off)
const rateLimitData = formData.value.enable_rate_limit ? {
rate_limit_5h: formData.value.rate_limit_5h && formData.value.rate_limit_5h > 0 ? formData.value.rate_limit_5h : 0,
rate_limit_1d: formData.value.rate_limit_1d && formData.value.rate_limit_1d > 0 ? formData.value.rate_limit_1d : 0,
rate_limit_7d: formData.value.rate_limit_7d && formData.value.rate_limit_7d > 0 ? formData.value.rate_limit_7d : 0,
} : { rate_limit_5h: 0, rate_limit_1d: 0, rate_limit_7d: 0 }
submitting.value = true submitting.value = true
try { try {
if (showEditModal.value && selectedKey.value) { if (showEditModal.value && selectedKey.value) {
...@@ -1088,7 +1383,10 @@ const handleSubmit = async () => { ...@@ -1088,7 +1383,10 @@ const handleSubmit = async () => {
ip_whitelist: ipWhitelist, ip_whitelist: ipWhitelist,
ip_blacklist: ipBlacklist, ip_blacklist: ipBlacklist,
quota: quota, quota: quota,
expires_at: expiresAt expires_at: expiresAt,
rate_limit_5h: rateLimitData.rate_limit_5h,
rate_limit_1d: rateLimitData.rate_limit_1d,
rate_limit_7d: rateLimitData.rate_limit_7d,
}) })
appStore.showSuccess(t('keys.keyUpdatedSuccess')) appStore.showSuccess(t('keys.keyUpdatedSuccess'))
} else { } else {
...@@ -1100,7 +1398,8 @@ const handleSubmit = async () => { ...@@ -1100,7 +1398,8 @@ const handleSubmit = async () => {
ipWhitelist, ipWhitelist,
ipBlacklist, ipBlacklist,
quota, quota,
expiresInDays expiresInDays,
rateLimitData
) )
appStore.showSuccess(t('keys.keyCreatedSuccess')) appStore.showSuccess(t('keys.keyCreatedSuccess'))
// Only advance tour if active, on submit step, and creation succeeded // Only advance tour if active, on submit step, and creation succeeded
...@@ -1154,6 +1453,10 @@ const closeModals = () => { ...@@ -1154,6 +1453,10 @@ const closeModals = () => {
ip_blacklist: '', ip_blacklist: '',
enable_quota: false, enable_quota: false,
quota: null, quota: null,
enable_rate_limit: false,
rate_limit_5h: null,
rate_limit_1d: null,
rate_limit_7d: null,
enable_expiration: false, enable_expiration: false,
expiration_preset: '30', expiration_preset: '30',
expiration_date: '' expiration_date: ''
...@@ -1190,6 +1493,37 @@ const resetQuotaUsed = async () => { ...@@ -1190,6 +1493,37 @@ const resetQuotaUsed = async () => {
} }
} }
// Show reset rate limit confirmation dialog (from edit modal)
const confirmResetRateLimit = () => {
showResetRateLimitDialog.value = true
}
// Show reset rate limit confirmation dialog (from table row)
const confirmResetRateLimitFromTable = (row: ApiKey) => {
selectedKey.value = row
showResetRateLimitDialog.value = true
}
// Reset rate limit usage for an API key
const resetRateLimitUsage = async () => {
if (!selectedKey.value) return
showResetRateLimitDialog.value = false
try {
await keysAPI.update(selectedKey.value.id, { reset_rate_limit_usage: true })
appStore.showSuccess(t('keys.rateLimitResetSuccess'))
// Refresh key data
await loadApiKeys()
// Update the editing key with fresh data
const refreshedKey = apiKeys.value.find(k => k.id === selectedKey.value!.id)
if (refreshedKey) {
selectedKey.value = refreshedKey
}
} catch (error: any) {
const errorMsg = error.response?.data?.detail || t('keys.failedToResetRateLimit')
appStore.showError(errorMsg)
}
}
const importToCcswitch = (row: ApiKey) => { const importToCcswitch = (row: ApiKey) => {
const platform = row.group?.platform || 'anthropic' const platform = row.group?.platform || 'anthropic'
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment