Unverified Commit c0c9c984 authored by Wesley Liddick's avatar Wesley Liddick Committed by GitHub
Browse files

Merge pull request #471 from bayma888/feature/api-key-quota-expiration

feat(api-key): 添加API密钥独立配额和过期时间功能
parents ba5a0d47 3fed478e
...@@ -75,6 +75,9 @@ func (f fakeAPIKeyRepo) ListKeysByUserID(ctx context.Context, userID int64) ([]s ...@@ -75,6 +75,9 @@ func (f fakeAPIKeyRepo) ListKeysByUserID(ctx context.Context, userID int64) ([]s
func (f fakeAPIKeyRepo) ListKeysByGroupID(ctx context.Context, groupID int64) ([]string, error) { func (f fakeAPIKeyRepo) ListKeysByGroupID(ctx context.Context, groupID int64) ([]string, error) {
return nil, errors.New("not implemented") return nil, errors.New("not implemented")
} }
func (f fakeAPIKeyRepo) IncrementQuotaUsed(ctx context.Context, id int64, amount float64) (float64, error) {
return 0, errors.New("not implemented")
}
type googleErrorResponse struct { type googleErrorResponse struct {
Error struct { Error struct {
......
...@@ -319,6 +319,10 @@ func (r *stubApiKeyRepo) ListKeysByGroupID(ctx context.Context, groupID int64) ( ...@@ -319,6 +319,10 @@ func (r *stubApiKeyRepo) ListKeysByGroupID(ctx context.Context, groupID int64) (
return nil, errors.New("not implemented") return nil, errors.New("not implemented")
} }
func (r *stubApiKeyRepo) IncrementQuotaUsed(ctx context.Context, id int64, amount float64) (float64, error) {
return 0, errors.New("not implemented")
}
type stubUserSubscriptionRepo struct { type stubUserSubscriptionRepo struct {
getActive func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) getActive func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error)
updateStatus func(ctx context.Context, subscriptionID int64, status string) error updateStatus func(ctx context.Context, subscriptionID int64, status string) error
......
...@@ -2,6 +2,14 @@ package service ...@@ -2,6 +2,14 @@ package service
import "time" import "time"
// API Key status constants
const (
StatusAPIKeyActive = "active"
StatusAPIKeyDisabled = "disabled"
StatusAPIKeyQuotaExhausted = "quota_exhausted"
StatusAPIKeyExpired = "expired"
)
type APIKey struct { type APIKey struct {
ID int64 ID int64
UserID int64 UserID int64
...@@ -15,8 +23,53 @@ type APIKey struct { ...@@ -15,8 +23,53 @@ type APIKey struct {
UpdatedAt time.Time UpdatedAt time.Time
User *User User *User
Group *Group Group *Group
// Quota fields
Quota float64 // Quota limit in USD (0 = unlimited)
QuotaUsed float64 // Used quota amount
ExpiresAt *time.Time // Expiration time (nil = never expires)
} }
func (k *APIKey) IsActive() bool { func (k *APIKey) IsActive() bool {
return k.Status == StatusActive return k.Status == StatusActive
} }
// IsExpired checks if the API key has expired
func (k *APIKey) IsExpired() bool {
if k.ExpiresAt == nil {
return false
}
return time.Now().After(*k.ExpiresAt)
}
// IsQuotaExhausted checks if the API key quota is exhausted
func (k *APIKey) IsQuotaExhausted() bool {
if k.Quota <= 0 {
return false // unlimited
}
return k.QuotaUsed >= k.Quota
}
// GetQuotaRemaining returns remaining quota (-1 for unlimited)
func (k *APIKey) GetQuotaRemaining() float64 {
if k.Quota <= 0 {
return -1 // unlimited
}
remaining := k.Quota - k.QuotaUsed
if remaining < 0 {
return 0
}
return remaining
}
// GetDaysUntilExpiry returns days until expiry (-1 for never expires)
func (k *APIKey) GetDaysUntilExpiry() int {
if k.ExpiresAt == nil {
return -1 // never expires
}
duration := time.Until(*k.ExpiresAt)
if duration < 0 {
return 0
}
return int(duration.Hours() / 24)
}
package service package service
import "time"
// APIKeyAuthSnapshot API Key 认证缓存快照(仅包含认证所需字段) // APIKeyAuthSnapshot API Key 认证缓存快照(仅包含认证所需字段)
type APIKeyAuthSnapshot struct { type APIKeyAuthSnapshot struct {
APIKeyID int64 `json:"api_key_id"` APIKeyID int64 `json:"api_key_id"`
...@@ -10,6 +12,13 @@ type APIKeyAuthSnapshot struct { ...@@ -10,6 +12,13 @@ type APIKeyAuthSnapshot struct {
IPBlacklist []string `json:"ip_blacklist,omitempty"` IPBlacklist []string `json:"ip_blacklist,omitempty"`
User APIKeyAuthUserSnapshot `json:"user"` User APIKeyAuthUserSnapshot `json:"user"`
Group *APIKeyAuthGroupSnapshot `json:"group,omitempty"` Group *APIKeyAuthGroupSnapshot `json:"group,omitempty"`
// Quota fields for API Key independent quota feature
Quota float64 `json:"quota"` // Quota limit in USD (0 = unlimited)
QuotaUsed float64 `json:"quota_used"` // Used quota amount
// Expiration field for API Key expiration feature
ExpiresAt *time.Time `json:"expires_at,omitempty"` // Expiration time (nil = never expires)
} }
// APIKeyAuthUserSnapshot 用户快照 // APIKeyAuthUserSnapshot 用户快照
......
...@@ -213,6 +213,9 @@ func (s *APIKeyService) snapshotFromAPIKey(apiKey *APIKey) *APIKeyAuthSnapshot { ...@@ -213,6 +213,9 @@ func (s *APIKeyService) snapshotFromAPIKey(apiKey *APIKey) *APIKeyAuthSnapshot {
Status: apiKey.Status, Status: apiKey.Status,
IPWhitelist: apiKey.IPWhitelist, IPWhitelist: apiKey.IPWhitelist,
IPBlacklist: apiKey.IPBlacklist, IPBlacklist: apiKey.IPBlacklist,
Quota: apiKey.Quota,
QuotaUsed: apiKey.QuotaUsed,
ExpiresAt: apiKey.ExpiresAt,
User: APIKeyAuthUserSnapshot{ User: APIKeyAuthUserSnapshot{
ID: apiKey.User.ID, ID: apiKey.User.ID,
Status: apiKey.User.Status, Status: apiKey.User.Status,
...@@ -256,6 +259,9 @@ func (s *APIKeyService) snapshotToAPIKey(key string, snapshot *APIKeyAuthSnapsho ...@@ -256,6 +259,9 @@ func (s *APIKeyService) snapshotToAPIKey(key string, snapshot *APIKeyAuthSnapsho
Status: snapshot.Status, Status: snapshot.Status,
IPWhitelist: snapshot.IPWhitelist, IPWhitelist: snapshot.IPWhitelist,
IPBlacklist: snapshot.IPBlacklist, IPBlacklist: snapshot.IPBlacklist,
Quota: snapshot.Quota,
QuotaUsed: snapshot.QuotaUsed,
ExpiresAt: snapshot.ExpiresAt,
User: &User{ User: &User{
ID: snapshot.User.ID, ID: snapshot.User.ID,
Status: snapshot.User.Status, Status: snapshot.User.Status,
......
...@@ -24,6 +24,10 @@ var ( ...@@ -24,6 +24,10 @@ var (
ErrAPIKeyInvalidChars = infraerrors.BadRequest("API_KEY_INVALID_CHARS", "api key can only contain letters, numbers, underscores, and hyphens") ErrAPIKeyInvalidChars = infraerrors.BadRequest("API_KEY_INVALID_CHARS", "api key can only contain letters, numbers, underscores, and hyphens")
ErrAPIKeyRateLimited = infraerrors.TooManyRequests("API_KEY_RATE_LIMITED", "too many failed attempts, please try again later") ErrAPIKeyRateLimited = infraerrors.TooManyRequests("API_KEY_RATE_LIMITED", "too many failed attempts, please try again later")
ErrInvalidIPPattern = infraerrors.BadRequest("INVALID_IP_PATTERN", "invalid IP or CIDR pattern") ErrInvalidIPPattern = infraerrors.BadRequest("INVALID_IP_PATTERN", "invalid IP or CIDR pattern")
// ErrAPIKeyExpired = infraerrors.Forbidden("API_KEY_EXPIRED", "api key has expired")
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 额度已用完")
) )
const ( const (
...@@ -51,6 +55,9 @@ type APIKeyRepository interface { ...@@ -51,6 +55,9 @@ type APIKeyRepository interface {
CountByGroupID(ctx context.Context, groupID int64) (int64, error) CountByGroupID(ctx context.Context, groupID int64) (int64, error)
ListKeysByUserID(ctx context.Context, userID int64) ([]string, error) ListKeysByUserID(ctx context.Context, userID int64) ([]string, error)
ListKeysByGroupID(ctx context.Context, groupID int64) ([]string, error) ListKeysByGroupID(ctx context.Context, groupID int64) ([]string, error)
// Quota methods
IncrementQuotaUsed(ctx context.Context, id int64, amount float64) (float64, error)
} }
// APIKeyCache defines cache operations for API key service // APIKeyCache defines cache operations for API key service
...@@ -85,6 +92,10 @@ type CreateAPIKeyRequest struct { ...@@ -85,6 +92,10 @@ type CreateAPIKeyRequest struct {
CustomKey *string `json:"custom_key"` // 可选的自定义key CustomKey *string `json:"custom_key"` // 可选的自定义key
IPWhitelist []string `json:"ip_whitelist"` // IP 白名单 IPWhitelist []string `json:"ip_whitelist"` // IP 白名单
IPBlacklist []string `json:"ip_blacklist"` // IP 黑名单 IPBlacklist []string `json:"ip_blacklist"` // IP 黑名单
// Quota fields
Quota float64 `json:"quota"` // Quota limit in USD (0 = unlimited)
ExpiresInDays *int `json:"expires_in_days"` // Days until expiry (nil = never expires)
} }
// UpdateAPIKeyRequest 更新API Key请求 // UpdateAPIKeyRequest 更新API Key请求
...@@ -94,6 +105,12 @@ type UpdateAPIKeyRequest struct { ...@@ -94,6 +105,12 @@ type UpdateAPIKeyRequest struct {
Status *string `json:"status"` Status *string `json:"status"`
IPWhitelist []string `json:"ip_whitelist"` // IP 白名单(空数组清空) IPWhitelist []string `json:"ip_whitelist"` // IP 白名单(空数组清空)
IPBlacklist []string `json:"ip_blacklist"` // IP 黑名单(空数组清空) IPBlacklist []string `json:"ip_blacklist"` // IP 黑名单(空数组清空)
// Quota fields
Quota *float64 `json:"quota"` // Quota limit in USD (nil = no change, 0 = unlimited)
ExpiresAt *time.Time `json:"expires_at"` // Expiration time (nil = no change)
ClearExpiration bool `json:"-"` // Clear expiration (internal use)
ResetQuota *bool `json:"reset_quota"` // Reset quota_used to 0
} }
// APIKeyService API Key服务 // APIKeyService API Key服务
...@@ -289,6 +306,14 @@ func (s *APIKeyService) Create(ctx context.Context, userID int64, req CreateAPIK ...@@ -289,6 +306,14 @@ func (s *APIKeyService) Create(ctx context.Context, userID int64, req CreateAPIK
Status: StatusActive, Status: StatusActive,
IPWhitelist: req.IPWhitelist, IPWhitelist: req.IPWhitelist,
IPBlacklist: req.IPBlacklist, IPBlacklist: req.IPBlacklist,
Quota: req.Quota,
QuotaUsed: 0,
}
// Set expiration time if specified
if req.ExpiresInDays != nil && *req.ExpiresInDays > 0 {
expiresAt := time.Now().AddDate(0, 0, *req.ExpiresInDays)
apiKey.ExpiresAt = &expiresAt
} }
if err := s.apiKeyRepo.Create(ctx, apiKey); err != nil { if err := s.apiKeyRepo.Create(ctx, apiKey); err != nil {
...@@ -436,6 +461,35 @@ func (s *APIKeyService) Update(ctx context.Context, id int64, userID int64, req ...@@ -436,6 +461,35 @@ func (s *APIKeyService) Update(ctx context.Context, id int64, userID int64, req
} }
} }
// Update quota fields
if req.Quota != nil {
apiKey.Quota = *req.Quota
// If quota is increased and status was quota_exhausted, reactivate
if apiKey.Status == StatusAPIKeyQuotaExhausted && *req.Quota > apiKey.QuotaUsed {
apiKey.Status = StatusActive
}
}
if req.ResetQuota != nil && *req.ResetQuota {
apiKey.QuotaUsed = 0
// If resetting quota and status was quota_exhausted, reactivate
if apiKey.Status == StatusAPIKeyQuotaExhausted {
apiKey.Status = StatusActive
}
}
if req.ClearExpiration {
apiKey.ExpiresAt = nil
// If clearing expiry and status was expired, reactivate
if apiKey.Status == StatusAPIKeyExpired {
apiKey.Status = StatusActive
}
} else if req.ExpiresAt != nil {
apiKey.ExpiresAt = req.ExpiresAt
// If extending expiry and status was expired, reactivate
if apiKey.Status == StatusAPIKeyExpired && time.Now().Before(*req.ExpiresAt) {
apiKey.Status = StatusActive
}
}
// 更新 IP 限制(空数组会清空设置) // 更新 IP 限制(空数组会清空设置)
apiKey.IPWhitelist = req.IPWhitelist apiKey.IPWhitelist = req.IPWhitelist
apiKey.IPBlacklist = req.IPBlacklist apiKey.IPBlacklist = req.IPBlacklist
...@@ -572,3 +626,51 @@ func (s *APIKeyService) SearchAPIKeys(ctx context.Context, userID int64, keyword ...@@ -572,3 +626,51 @@ func (s *APIKeyService) SearchAPIKeys(ctx context.Context, userID int64, keyword
} }
return keys, nil return keys, nil
} }
// CheckAPIKeyQuotaAndExpiry checks if the API key is valid for use (not expired, quota not exhausted)
// Returns nil if valid, error if invalid
func (s *APIKeyService) CheckAPIKeyQuotaAndExpiry(apiKey *APIKey) error {
// Check expiration
if apiKey.IsExpired() {
return ErrAPIKeyExpired
}
// Check quota
if apiKey.IsQuotaExhausted() {
return ErrAPIKeyQuotaExhausted
}
return nil
}
// UpdateQuotaUsed updates the quota_used field after a request
// Also checks if quota is exhausted and updates status accordingly
func (s *APIKeyService) UpdateQuotaUsed(ctx context.Context, apiKeyID int64, cost float64) error {
if cost <= 0 {
return nil
}
// Use repository to atomically increment quota_used
newQuotaUsed, err := s.apiKeyRepo.IncrementQuotaUsed(ctx, apiKeyID, cost)
if err != nil {
return fmt.Errorf("increment quota used: %w", err)
}
// Check if quota is now exhausted and update status if needed
apiKey, err := s.apiKeyRepo.GetByID(ctx, apiKeyID)
if err != nil {
return nil // Don't fail the request, just log
}
// If quota is set and now exhausted, update status
if apiKey.Quota > 0 && newQuotaUsed >= apiKey.Quota {
apiKey.Status = StatusAPIKeyQuotaExhausted
if err := s.apiKeyRepo.Update(ctx, apiKey); err != nil {
return nil // Don't fail the request
}
// Invalidate cache so next request sees the new status
s.InvalidateAuthCacheByKey(ctx, apiKey.Key)
}
return nil
}
...@@ -99,6 +99,10 @@ func (s *authRepoStub) ListKeysByGroupID(ctx context.Context, groupID int64) ([] ...@@ -99,6 +99,10 @@ func (s *authRepoStub) ListKeysByGroupID(ctx context.Context, groupID int64) ([]
return s.listKeysByGroupID(ctx, groupID) return s.listKeysByGroupID(ctx, groupID)
} }
func (s *authRepoStub) IncrementQuotaUsed(ctx context.Context, id int64, amount float64) (float64, error) {
panic("unexpected IncrementQuotaUsed call")
}
type authCacheStub struct { type authCacheStub struct {
getAuthCache func(ctx context.Context, key string) (*APIKeyAuthCacheEntry, error) getAuthCache func(ctx context.Context, key string) (*APIKeyAuthCacheEntry, error)
setAuthKeys []string setAuthKeys []string
......
...@@ -118,6 +118,10 @@ func (s *apiKeyRepoStub) ListKeysByGroupID(ctx context.Context, groupID int64) ( ...@@ -118,6 +118,10 @@ func (s *apiKeyRepoStub) ListKeysByGroupID(ctx context.Context, groupID int64) (
panic("unexpected ListKeysByGroupID call") panic("unexpected ListKeysByGroupID call")
} }
func (s *apiKeyRepoStub) IncrementQuotaUsed(ctx context.Context, id int64, amount float64) (float64, error) {
panic("unexpected IncrementQuotaUsed call")
}
// apiKeyCacheStub 是 APIKeyCache 接口的测试桩实现。 // apiKeyCacheStub 是 APIKeyCache 接口的测试桩实现。
// 用于验证删除操作时缓存清理逻辑是否被正确调用。 // 用于验证删除操作时缓存清理逻辑是否被正确调用。
// //
......
...@@ -4489,13 +4489,19 @@ func (s *GatewayService) replaceToolNamesInResponseBody(body []byte, toolNameMap ...@@ -4489,13 +4489,19 @@ func (s *GatewayService) replaceToolNamesInResponseBody(body []byte, toolNameMap
// RecordUsageInput 记录使用量的输入参数 // RecordUsageInput 记录使用量的输入参数
type RecordUsageInput struct { type RecordUsageInput struct {
Result *ForwardResult Result *ForwardResult
APIKey *APIKey APIKey *APIKey
User *User User *User
Account *Account Account *Account
Subscription *UserSubscription // 可选:订阅信息 Subscription *UserSubscription // 可选:订阅信息
UserAgent string // 请求的 User-Agent UserAgent string // 请求的 User-Agent
IPAddress string // 请求的客户端 IP 地址 IPAddress string // 请求的客户端 IP 地址
APIKeyService APIKeyQuotaUpdater // 可选:用于更新API Key配额
}
// APIKeyQuotaUpdater defines the interface for updating API Key quota
type APIKeyQuotaUpdater interface {
UpdateQuotaUsed(ctx context.Context, apiKeyID int64, cost float64) error
} }
// RecordUsage 记录使用量并扣费(或更新订阅用量) // RecordUsage 记录使用量并扣费(或更新订阅用量)
...@@ -4635,6 +4641,13 @@ func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInpu ...@@ -4635,6 +4641,13 @@ func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInpu
} }
} }
// 更新 API Key 配额(如果设置了配额限制)
if shouldBill && cost.ActualCost > 0 && apiKey.Quota > 0 && input.APIKeyService != nil {
if err := input.APIKeyService.UpdateQuotaUsed(ctx, apiKey.ID, cost.ActualCost); err != nil {
log.Printf("Update API key quota failed: %v", err)
}
}
// 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)
...@@ -4652,6 +4665,7 @@ type RecordUsageLongContextInput struct { ...@@ -4652,6 +4665,7 @@ type RecordUsageLongContextInput struct {
IPAddress string // 请求的客户端 IP 地址 IPAddress string // 请求的客户端 IP 地址
LongContextThreshold int // 长上下文阈值(如 200000) LongContextThreshold int // 长上下文阈值(如 200000)
LongContextMultiplier float64 // 超出阈值部分的倍率(如 2.0) LongContextMultiplier float64 // 超出阈值部分的倍率(如 2.0)
APIKeyService *APIKeyService // API Key 配额服务(可选)
} }
// RecordUsageWithLongContext 记录使用量并扣费,支持长上下文双倍计费(用于 Gemini) // RecordUsageWithLongContext 记录使用量并扣费,支持长上下文双倍计费(用于 Gemini)
...@@ -4788,6 +4802,12 @@ func (s *GatewayService) RecordUsageWithLongContext(ctx context.Context, input * ...@@ -4788,6 +4802,12 @@ func (s *GatewayService) RecordUsageWithLongContext(ctx context.Context, input *
} }
// 异步更新余额缓存 // 异步更新余额缓存
s.billingCacheService.QueueDeductBalance(user.ID, cost.ActualCost) s.billingCacheService.QueueDeductBalance(user.ID, cost.ActualCost)
// API Key 独立配额扣费
if input.APIKeyService != nil && apiKey.Quota > 0 {
if err := input.APIKeyService.UpdateQuotaUsed(ctx, apiKey.ID, cost.ActualCost); err != nil {
log.Printf("Add API key quota used failed: %v", err)
}
}
} }
} }
......
...@@ -1681,13 +1681,14 @@ func (s *OpenAIGatewayService) replaceModelInResponseBody(body []byte, fromModel ...@@ -1681,13 +1681,14 @@ func (s *OpenAIGatewayService) replaceModelInResponseBody(body []byte, fromModel
// OpenAIRecordUsageInput input for recording usage // OpenAIRecordUsageInput input for recording usage
type OpenAIRecordUsageInput struct { type OpenAIRecordUsageInput struct {
Result *OpenAIForwardResult Result *OpenAIForwardResult
APIKey *APIKey APIKey *APIKey
User *User User *User
Account *Account Account *Account
Subscription *UserSubscription Subscription *UserSubscription
UserAgent string // 请求的 User-Agent UserAgent string // 请求的 User-Agent
IPAddress string // 请求的客户端 IP 地址 IPAddress string // 请求的客户端 IP 地址
APIKeyService APIKeyQuotaUpdater
} }
// RecordUsage records usage and deducts balance // RecordUsage records usage and deducts balance
...@@ -1799,6 +1800,13 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec ...@@ -1799,6 +1800,13 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec
} }
} }
// Update API key quota if applicable (only for balance mode with quota set)
if shouldBill && cost.ActualCost > 0 && apiKey.Quota > 0 && input.APIKeyService != nil {
if err := input.APIKeyService.UpdateQuotaUsed(ctx, apiKey.ID, cost.ActualCost); err != nil {
log.Printf("Update API key quota failed: %v", err)
}
}
// 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)
......
-- Migration: Add quota fields to api_keys table
-- This migration adds independent quota and expiration support for API keys
-- Add quota limit field (0 = unlimited)
ALTER TABLE api_keys ADD COLUMN IF NOT EXISTS quota DECIMAL(20, 8) NOT NULL DEFAULT 0;
-- Add used quota amount field
ALTER TABLE api_keys ADD COLUMN IF NOT EXISTS quota_used DECIMAL(20, 8) NOT NULL DEFAULT 0;
-- Add expiration time field (NULL = never expires)
ALTER TABLE api_keys ADD COLUMN IF NOT EXISTS expires_at TIMESTAMPTZ;
-- Add indexes for efficient quota queries
CREATE INDEX IF NOT EXISTS idx_api_keys_quota_quota_used ON api_keys(quota, quota_used) WHERE deleted_at IS NULL;
CREATE INDEX IF NOT EXISTS idx_api_keys_expires_at ON api_keys(expires_at) WHERE deleted_at IS NULL;
-- Comment on columns for documentation
COMMENT ON COLUMN api_keys.quota IS 'Quota limit in USD for this API key (0 = unlimited)';
COMMENT ON COLUMN api_keys.quota_used IS 'Used quota amount in USD';
COMMENT ON COLUMN api_keys.expires_at IS 'Expiration time for this API key (null = never expires)';
//go:build tools
// +build tools
package tools
import (
_ "entgo.io/ent/cmd/ent"
_ "github.com/google/wire/cmd/wire"
)
...@@ -44,6 +44,8 @@ export async function getById(id: number): Promise<ApiKey> { ...@@ -44,6 +44,8 @@ export async function getById(id: number): Promise<ApiKey> {
* @param customKey - Optional custom key value * @param customKey - Optional custom key value
* @param ipWhitelist - Optional IP whitelist * @param ipWhitelist - Optional IP whitelist
* @param ipBlacklist - Optional IP blacklist * @param ipBlacklist - Optional IP blacklist
* @param quota - Optional quota limit in USD (0 = unlimited)
* @param expiresInDays - Optional days until expiry (undefined = never expires)
* @returns Created API key * @returns Created API key
*/ */
export async function create( export async function create(
...@@ -51,7 +53,9 @@ export async function create( ...@@ -51,7 +53,9 @@ export async function create(
groupId?: number | null, groupId?: number | null,
customKey?: string, customKey?: string,
ipWhitelist?: string[], ipWhitelist?: string[],
ipBlacklist?: string[] ipBlacklist?: string[],
quota?: number,
expiresInDays?: number
): Promise<ApiKey> { ): Promise<ApiKey> {
const payload: CreateApiKeyRequest = { name } const payload: CreateApiKeyRequest = { name }
if (groupId !== undefined) { if (groupId !== undefined) {
...@@ -66,6 +70,12 @@ export async function create( ...@@ -66,6 +70,12 @@ export async function create(
if (ipBlacklist && ipBlacklist.length > 0) { if (ipBlacklist && ipBlacklist.length > 0) {
payload.ip_blacklist = ipBlacklist payload.ip_blacklist = ipBlacklist
} }
if (quota !== undefined && quota > 0) {
payload.quota = quota
}
if (expiresInDays !== undefined && expiresInDays > 0) {
payload.expires_in_days = expiresInDays
}
const { data } = await apiClient.post<ApiKey>('/keys', payload) const { data } = await apiClient.post<ApiKey>('/keys', payload)
return data return data
......
...@@ -407,6 +407,7 @@ export default { ...@@ -407,6 +407,7 @@ export default {
usage: 'Usage', usage: 'Usage',
today: 'Today', today: 'Today',
total: 'Total', total: 'Total',
quota: 'Quota',
useKey: 'Use Key', useKey: 'Use Key',
useKeyModal: { useKeyModal: {
title: 'Use API Key', title: 'Use API Key',
...@@ -470,6 +471,33 @@ export default { ...@@ -470,6 +471,33 @@ export default {
geminiCli: 'Gemini CLI', geminiCli: 'Gemini CLI',
geminiCliDesc: 'Import as Gemini CLI configuration', geminiCliDesc: 'Import as Gemini CLI configuration',
}, },
// Quota and expiration
quotaLimit: 'Quota Limit',
quotaAmount: 'Quota Amount (USD)',
quotaAmountPlaceholder: 'Enter quota limit in USD',
quotaAmountHint: 'Set the maximum amount this key can spend. 0 = unlimited.',
quotaUsed: 'Quota Used',
reset: 'Reset',
resetQuotaUsed: 'Reset used quota to 0',
resetQuotaTitle: 'Confirm Reset Quota',
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',
failedToResetQuota: 'Failed to reset quota',
expiration: 'Expiration',
expiresInDays: '{days} days',
extendDays: '+{days} days',
customDate: 'Custom',
expirationDate: 'Expiration Date',
expirationDateHint: 'Select when this API key should expire.',
currentExpiration: 'Current expiration',
expiresAt: 'Expires',
noExpiration: 'Never',
status: {
active: 'Active',
inactive: 'Inactive',
quota_exhausted: 'Quota Exhausted',
expired: 'Expired',
},
}, },
// Usage // Usage
......
This diff is collapsed.
...@@ -374,9 +374,12 @@ export interface ApiKey { ...@@ -374,9 +374,12 @@ export interface ApiKey {
key: string key: string
name: string name: string
group_id: number | null group_id: number | null
status: 'active' | 'inactive' status: 'active' | 'inactive' | 'quota_exhausted' | 'expired'
ip_whitelist: string[] ip_whitelist: string[]
ip_blacklist: string[] ip_blacklist: string[]
quota: number // Quota limit in USD (0 = unlimited)
quota_used: number // Used quota amount in USD
expires_at: string | null // Expiration time (null = never expires)
created_at: string created_at: string
updated_at: string updated_at: string
group?: Group group?: Group
...@@ -388,6 +391,8 @@ export interface CreateApiKeyRequest { ...@@ -388,6 +391,8 @@ export interface CreateApiKeyRequest {
custom_key?: string // Optional custom API Key custom_key?: string // Optional custom API Key
ip_whitelist?: string[] ip_whitelist?: string[]
ip_blacklist?: string[] ip_blacklist?: string[]
quota?: number // Quota limit in USD (0 = unlimited)
expires_in_days?: number // Days until expiry (null = never expires)
} }
export interface UpdateApiKeyRequest { export interface UpdateApiKeyRequest {
...@@ -396,6 +401,9 @@ export interface UpdateApiKeyRequest { ...@@ -396,6 +401,9 @@ export interface UpdateApiKeyRequest {
status?: 'active' | 'inactive' status?: 'active' | 'inactive'
ip_whitelist?: string[] ip_whitelist?: string[]
ip_blacklist?: string[] ip_blacklist?: string[]
quota?: number // Quota limit in USD (null = no change, 0 = unlimited)
expires_at?: string | null // Expiration time (null = no change)
reset_quota?: boolean // Reset quota_used to 0
} }
export interface CreateGroupRequest { export interface CreateGroupRequest {
......
...@@ -108,12 +108,53 @@ ...@@ -108,12 +108,53 @@
${{ (usageStats[row.id]?.total_actual_cost ?? 0).toFixed(4) }} ${{ (usageStats[row.id]?.total_actual_cost ?? 0).toFixed(4) }}
</span> </span>
</div> </div>
<!-- Quota progress (if quota is set) -->
<div v-if="row.quota > 0" class="mt-1.5">
<div class="flex items-center gap-1.5">
<span class="text-gray-500 dark:text-gray-400">{{ t('keys.quota') }}:</span>
<span :class="[
'font-medium',
row.quota_used >= row.quota ? 'text-red-500' :
row.quota_used >= row.quota * 0.8 ? 'text-yellow-500' :
'text-gray-900 dark:text-white'
]">
${{ row.quota_used?.toFixed(2) || '0.00' }} / ${{ row.quota?.toFixed(2) }}
</span>
</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',
row.quota_used >= row.quota ? 'bg-red-500' :
row.quota_used >= row.quota * 0.8 ? 'bg-yellow-500' :
'bg-primary-500'
]"
:style="{ width: Math.min((row.quota_used / row.quota) * 100, 100) + '%' }"
/>
</div>
</div>
</div> </div>
</template> </template>
<template #cell-expires_at="{ value }">
<span v-if="value" :class="[
'text-sm',
new Date(value) < new Date() ? 'text-red-500 dark:text-red-400' : 'text-gray-500 dark:text-dark-400'
]">
{{ formatDateTime(value) }}
</span>
<span v-else class="text-sm text-gray-400 dark:text-dark-500">{{ t('keys.noExpiration') }}</span>
</template>
<template #cell-status="{ value }"> <template #cell-status="{ value }">
<span :class="['badge', value === 'active' ? 'badge-success' : 'badge-gray']"> <span :class="[
{{ t('admin.accounts.status.' + value) }} 'badge',
value === 'active' ? 'badge-success' :
value === 'quota_exhausted' ? 'badge-warning' :
value === 'expired' ? 'badge-danger' :
'badge-gray'
]">
{{ t('keys.status.' + value) }}
</span> </span>
</template> </template>
...@@ -334,6 +375,145 @@ ...@@ -334,6 +375,145 @@
</div> </div>
</div> </div>
</div> </div>
<!-- Quota Limit Section -->
<div class="space-y-3">
<label class="input-label">{{ t('keys.quotaLimit') }}</label>
<!-- Switch commented out - always show input, 0 = unlimited
<div class="flex items-center justify-between">
<label class="input-label mb-0">{{ t('keys.quotaLimit') }}</label>
<button
type="button"
@click="formData.enable_quota = !formData.enable_quota"
: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_quota ? '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_quota ? 'translate-x-4' : 'translate-x-0'
]"
/>
</button>
</div>
-->
<div class="space-y-4">
<div>
<div class="relative">
<span class="absolute left-3 top-1/2 -translate-y-1/2 text-gray-500">$</span>
<input
v-model.number="formData.quota"
type="number"
step="0.01"
min="0"
class="input pl-7"
:placeholder="t('keys.quotaAmountPlaceholder')"
/>
</div>
<p class="input-hint">{{ t('keys.quotaAmountHint') }}</p>
</div>
<!-- Quota used display (only in edit mode) -->
<div v-if="showEditModal && selectedKey && selectedKey.quota > 0">
<label class="input-label">{{ t('keys.quotaUsed') }}</label>
<div class="flex items-center gap-2">
<div class="flex-1 rounded-lg bg-gray-100 px-3 py-2 dark:bg-dark-700">
<span class="font-medium text-gray-900 dark:text-white">
${{ selectedKey.quota_used?.toFixed(4) || '0.0000' }}
</span>
<span class="mx-2 text-gray-400">/</span>
<span class="text-gray-500 dark:text-gray-400">
${{ selectedKey.quota?.toFixed(2) || '0.00' }}
</span>
</div>
<button
type="button"
@click="confirmResetQuota"
class="btn btn-secondary text-sm"
:title="t('keys.resetQuotaUsed')"
>
{{ t('keys.reset') }}
</button>
</div>
</div>
</div>
</div>
<!-- Expiration Section -->
<div class="space-y-3">
<div class="flex items-center justify-between">
<label class="input-label mb-0">{{ t('keys.expiration') }}</label>
<button
type="button"
@click="formData.enable_expiration = !formData.enable_expiration"
: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_expiration ? '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_expiration ? 'translate-x-4' : 'translate-x-0'
]"
/>
</button>
</div>
<div v-if="formData.enable_expiration" class="space-y-4 pt-2">
<!-- Quick select buttons (for both create and edit mode) -->
<div class="flex flex-wrap gap-2">
<button
v-for="days in ['7', '30', '90']"
:key="days"
type="button"
@click="setExpirationDays(parseInt(days))"
:class="[
'rounded-lg px-3 py-1.5 text-sm transition-colors',
formData.expiration_preset === days
? 'bg-primary-100 text-primary-700 dark:bg-primary-900/30 dark:text-primary-400'
: 'bg-gray-100 text-gray-600 hover:bg-gray-200 dark:bg-dark-700 dark:text-gray-400 dark:hover:bg-dark-600'
]"
>
{{ showEditModal ? t('keys.extendDays', { days }) : t('keys.expiresInDays', { days }) }}
</button>
<button
type="button"
@click="formData.expiration_preset = 'custom'"
:class="[
'rounded-lg px-3 py-1.5 text-sm transition-colors',
formData.expiration_preset === 'custom'
? 'bg-primary-100 text-primary-700 dark:bg-primary-900/30 dark:text-primary-400'
: 'bg-gray-100 text-gray-600 hover:bg-gray-200 dark:bg-dark-700 dark:text-gray-400 dark:hover:bg-dark-600'
]"
>
{{ t('keys.customDate') }}
</button>
</div>
<!-- Date picker (always show for precise adjustment) -->
<div>
<label class="input-label">{{ t('keys.expirationDate') }}</label>
<input
v-model="formData.expiration_date"
type="datetime-local"
class="input"
/>
<p class="input-hint">{{ t('keys.expirationDateHint') }}</p>
</div>
<!-- Current expiration display (only in edit mode) -->
<div v-if="showEditModal && selectedKey?.expires_at" class="text-sm">
<span class="text-gray-500 dark:text-gray-400">{{ t('keys.currentExpiration') }}: </span>
<span class="font-medium text-gray-900 dark:text-white">
{{ formatDateTime(selectedKey.expires_at) }}
</span>
</div>
</div>
</div>
</form> </form>
<template #footer> <template #footer>
<div class="flex justify-end gap-3"> <div class="flex justify-end gap-3">
...@@ -391,6 +571,18 @@ ...@@ -391,6 +571,18 @@
@cancel="showDeleteDialog = false" @cancel="showDeleteDialog = false"
/> />
<!-- Reset Quota Confirmation Dialog -->
<ConfirmDialog
:show="showResetQuotaDialog"
:title="t('keys.resetQuotaTitle')"
:message="t('keys.resetQuotaConfirmMessage', { name: selectedKey?.name, used: selectedKey?.quota_used?.toFixed(4) })"
:confirm-text="t('keys.reset')"
:cancel-text="t('common.cancel')"
:danger="true"
@confirm="resetQuotaUsed"
@cancel="showResetQuotaDialog = false"
/>
<!-- Use Key Modal --> <!-- Use Key Modal -->
<UseKeyModal <UseKeyModal
:show="showUseKeyModal" :show="showUseKeyModal"
...@@ -514,6 +706,13 @@ import type { Column } from '@/components/common/types' ...@@ -514,6 +706,13 @@ import type { Column } from '@/components/common/types'
import type { BatchApiKeyUsageStats } from '@/api/usage' import type { BatchApiKeyUsageStats } from '@/api/usage'
import { formatDateTime } from '@/utils/format' import { formatDateTime } from '@/utils/format'
// Helper to format date for datetime-local input
const formatDateTimeLocal = (isoDate: string): string => {
const date = new Date(isoDate)
const pad = (n: number) => n.toString().padStart(2, '0')
return `${date.getFullYear()}-${pad(date.getMonth() + 1)}-${pad(date.getDate())}T${pad(date.getHours())}:${pad(date.getMinutes())}`
}
interface GroupOption { interface GroupOption {
value: number value: number
label: string label: string
...@@ -532,6 +731,7 @@ const columns = computed<Column[]>(() => [ ...@@ -532,6 +731,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: '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: 'created_at', label: t('keys.created'), sortable: true }, { key: 'created_at', label: t('keys.created'), sortable: true },
{ key: 'actions', label: t('common.actions'), sortable: false } { key: 'actions', label: t('common.actions'), sortable: false }
...@@ -553,6 +753,7 @@ const pagination = ref({ ...@@ -553,6 +753,7 @@ const pagination = ref({
const showCreateModal = ref(false) 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 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)
...@@ -587,7 +788,13 @@ const formData = ref({ ...@@ -587,7 +788,13 @@ const formData = ref({
custom_key: '', custom_key: '',
enable_ip_restriction: false, enable_ip_restriction: false,
ip_whitelist: '', ip_whitelist: '',
ip_blacklist: '' ip_blacklist: '',
// Quota settings (empty = unlimited)
enable_quota: false,
quota: null as number | null,
enable_expiration: false,
expiration_preset: '30' as '7' | '30' | '90' | 'custom',
expiration_date: ''
}) })
// 自定义Key验证 // 自定义Key验证
...@@ -724,15 +931,21 @@ const handlePageSizeChange = (pageSize: number) => { ...@@ -724,15 +931,21 @@ const handlePageSizeChange = (pageSize: number) => {
const editKey = (key: ApiKey) => { const editKey = (key: ApiKey) => {
selectedKey.value = key selectedKey.value = key
const hasIPRestriction = (key.ip_whitelist?.length > 0) || (key.ip_blacklist?.length > 0) const hasIPRestriction = (key.ip_whitelist?.length > 0) || (key.ip_blacklist?.length > 0)
const hasExpiration = !!key.expires_at
formData.value = { formData.value = {
name: key.name, name: key.name,
group_id: key.group_id, group_id: key.group_id,
status: key.status, status: key.status === 'quota_exhausted' || key.status === 'expired' ? 'inactive' : key.status,
use_custom_key: false, use_custom_key: false,
custom_key: '', custom_key: '',
enable_ip_restriction: hasIPRestriction, enable_ip_restriction: hasIPRestriction,
ip_whitelist: (key.ip_whitelist || []).join('\n'), ip_whitelist: (key.ip_whitelist || []).join('\n'),
ip_blacklist: (key.ip_blacklist || []).join('\n') ip_blacklist: (key.ip_blacklist || []).join('\n'),
enable_quota: key.quota > 0,
quota: key.quota > 0 ? key.quota : null,
enable_expiration: hasExpiration,
expiration_preset: 'custom',
expiration_date: key.expires_at ? formatDateTimeLocal(key.expires_at) : ''
} }
showEditModal.value = true showEditModal.value = true
} }
...@@ -820,6 +1033,28 @@ const handleSubmit = async () => { ...@@ -820,6 +1033,28 @@ const handleSubmit = async () => {
const ipWhitelist = formData.value.enable_ip_restriction ? parseIPList(formData.value.ip_whitelist) : [] const ipWhitelist = formData.value.enable_ip_restriction ? parseIPList(formData.value.ip_whitelist) : []
const ipBlacklist = formData.value.enable_ip_restriction ? parseIPList(formData.value.ip_blacklist) : [] const ipBlacklist = formData.value.enable_ip_restriction ? parseIPList(formData.value.ip_blacklist) : []
// Calculate quota value (null/empty/0 = unlimited, stored as 0)
const quota = formData.value.quota && formData.value.quota > 0 ? formData.value.quota : 0
// Calculate expiration
let expiresInDays: number | undefined
let expiresAt: string | null | undefined
if (formData.value.enable_expiration && formData.value.expiration_date) {
if (!showEditModal.value) {
// Create mode: calculate days from date
const expDate = new Date(formData.value.expiration_date)
const now = new Date()
const diffDays = Math.ceil((expDate.getTime() - now.getTime()) / (1000 * 60 * 60 * 24))
expiresInDays = diffDays > 0 ? diffDays : 1
} else {
// Edit mode: use custom date directly
expiresAt = new Date(formData.value.expiration_date).toISOString()
}
} else if (showEditModal.value) {
// Edit mode: if expiration disabled or date cleared, send empty string to clear
expiresAt = ''
}
submitting.value = true submitting.value = true
try { try {
if (showEditModal.value && selectedKey.value) { if (showEditModal.value && selectedKey.value) {
...@@ -828,12 +1063,22 @@ const handleSubmit = async () => { ...@@ -828,12 +1063,22 @@ const handleSubmit = async () => {
group_id: formData.value.group_id, group_id: formData.value.group_id,
status: formData.value.status, status: formData.value.status,
ip_whitelist: ipWhitelist, ip_whitelist: ipWhitelist,
ip_blacklist: ipBlacklist ip_blacklist: ipBlacklist,
quota: quota,
expires_at: expiresAt
}) })
appStore.showSuccess(t('keys.keyUpdatedSuccess')) appStore.showSuccess(t('keys.keyUpdatedSuccess'))
} else { } else {
const customKey = formData.value.use_custom_key ? formData.value.custom_key : undefined const customKey = formData.value.use_custom_key ? formData.value.custom_key : undefined
await keysAPI.create(formData.value.name, formData.value.group_id, customKey, ipWhitelist, ipBlacklist) await keysAPI.create(
formData.value.name,
formData.value.group_id,
customKey,
ipWhitelist,
ipBlacklist,
quota,
expiresInDays
)
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
if (onboardingStore.isCurrentStep('[data-tour="key-form-submit"]')) { if (onboardingStore.isCurrentStep('[data-tour="key-form-submit"]')) {
...@@ -883,7 +1128,42 @@ const closeModals = () => { ...@@ -883,7 +1128,42 @@ const closeModals = () => {
custom_key: '', custom_key: '',
enable_ip_restriction: false, enable_ip_restriction: false,
ip_whitelist: '', ip_whitelist: '',
ip_blacklist: '' ip_blacklist: '',
enable_quota: false,
quota: null,
enable_expiration: false,
expiration_preset: '30',
expiration_date: ''
}
}
// Show reset quota confirmation dialog
const confirmResetQuota = () => {
showResetQuotaDialog.value = true
}
// Set expiration date based on quick select days
const setExpirationDays = (days: number) => {
formData.value.expiration_preset = days.toString() as '7' | '30' | '90'
const expDate = new Date()
expDate.setDate(expDate.getDate() + days)
formData.value.expiration_date = formatDateTimeLocal(expDate.toISOString())
}
// Reset quota used for an API key
const resetQuotaUsed = async () => {
if (!selectedKey.value) return
showResetQuotaDialog.value = false
try {
await keysAPI.update(selectedKey.value.id, { reset_quota: true })
appStore.showSuccess(t('keys.quotaResetSuccess'))
// Update local state
if (selectedKey.value) {
selectedKey.value.quota_used = 0
}
} catch (error: any) {
const errorMsg = error.response?.data?.detail || t('keys.failedToResetQuota')
appStore.showError(errorMsg)
} }
} }
......
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