From 3b336b2eca577d995a16748916eddb9fe61ed9bf Mon Sep 17 00:00:00 2001 From: zhangcj <10606643+superzcj@users.noreply.github.com> Date: Tue, 22 Sep 2026 15:23:09 +0800 Subject: [PATCH] feat: optionally enforce API-key spending windows in simple mode --- README.md | 2 + README_CN.md | 2 + backend/internal/config/config.go | 4 + backend/internal/config/config_test.go | 9 ++ .../handler/openai_gateway_handler.go | 21 ++- .../handler/openai_gateway_handler_test.go | 62 +++++++- .../internal/service/billing_cache_service.go | 37 ++++- .../billing_cache_service_simple_mode_test.go | 96 ++++++++++++ .../gateway_simple_mode_record_usage_test.go | 62 ++++++++ .../internal/service/gateway_usage_billing.go | 66 ++++++-- .../gateway_usage_billing_simple_mode_test.go | 144 ++++++++++++++++++ .../internal/service/openai_gateway_usage.go | 31 ++-- deploy/.env.example | 2 + deploy/config.example.yaml | 3 + deploy/docker-compose.dev.yml | 1 + deploy/docker-compose.local.yml | 1 + deploy/docker-compose.standalone.yml | 1 + deploy/docker-compose.yml | 1 + 18 files changed, 509 insertions(+), 36 deletions(-) create mode 100644 backend/internal/service/billing_cache_service_simple_mode_test.go create mode 100644 backend/internal/service/gateway_simple_mode_record_usage_test.go create mode 100644 backend/internal/service/gateway_usage_billing_simple_mode_test.go diff --git a/README.md b/README.md index 82d0a4d7ae..f7e3f14f89 100644 --- a/README.md +++ b/README.md @@ -727,6 +727,8 @@ Simple Mode is designed for individual developers or internal teams who want qui - Enable: Set environment variable `RUN_MODE=simple` - Default groups are seeded on each startup. Set `SIMPLE_MODE_AUTO_CREATE_DEFAULT_GROUPS=false` (or YAML `simple_mode.auto_create_default_groups: false`) to manage groups yourself. The default is `true`; disabling it does not delete existing groups or change runtime auto-binding or admin concurrency setup. - Difference: Hides SaaS-related features and skips billing process +- Optional key windows: Set `SIMPLE_MODE_KEY_RATE_LIMIT_ENABLED=true` to enforce each API key's configured 5-hour, daily, and 7-day spending windows. The default is `false`; balance and subscription debit remain bypassed when enabled. +- Window enforcement uses the database as its source of truth and records only API-key window usage. It is a post-request soft cap, so concurrent in-flight requests can overshoot by their final costs. Historical simple-mode usage is not backfilled. - Security note: In production, you must also set `SIMPLE_MODE_CONFIRM=true` to allow startup --- diff --git a/README_CN.md b/README_CN.md index 6289d0c92b..00c6217218 100644 --- a/README_CN.md +++ b/README_CN.md @@ -737,6 +737,8 @@ OAuth / Setup Token 图片请求使用 Responses 主控模型调用 `image_gener - 启用方式：设置环境变量 `RUN_MODE=simple` - 默认每次启动会补齐默认分组。设置 `SIMPLE_MODE_AUTO_CREATE_DEFAULT_GROUPS=false`（或 YAML `simple_mode.auto_create_default_groups: false`）可自行管理分组。默认值为 `true`；关闭后不删除已有分组，不改变运行时自动绑定或管理员并发设置。 - 功能差异：隐藏 SaaS 相关功能，跳过计费流程 +- 可选密钥窗口：设置 `SIMPLE_MODE_KEY_RATE_LIMIT_ENABLED=true` 后，按每个 API Key 配置的 5 小时、1 天、7 天消费窗口进行限制，默认值为 `false`；启用后仍跳过余额和订阅扣费。 +- 窗口限制以数据库为准，只记录 API Key 窗口用量。它在请求完成后记账，并发中的请求可能以各自最终费用超过窗口上限。启用前的历史用量不会自动补算。 - 安全注意事项：生产环境需同时设置 `SIMPLE_MODE_CONFIRM=true` 才允许启动 --- diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index 2195c74ae5..11b86e6b22 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -105,6 +105,9 @@ type Config struct { BatchImage BatchImageConfig `mapstructure:"batch_image"` ImageStorage ImageStorageConfig `mapstructure:"image_storage"` Plugins PluginConfig `mapstructure:"plugins"` + + // Enforce only API-key spending windows in simple mode. + SimpleModeKeyRateLimitEnabled bool `mapstructure:"simple_mode_key_rate_limit_enabled" yaml:"simple_mode_key_rate_limit_enabled"` } // SimpleModeConfig controls startup behavior in simple mode. @@ -1999,6 +2002,7 @@ func configureConfigSource(setConfigFile, addConfigPath func(string)) { func setDefaults() { viper.SetDefault("run_mode", RunModeStandard) viper.SetDefault("simple_mode.auto_create_default_groups", true) + viper.SetDefault("simple_mode_key_rate_limit_enabled", false) // Server viper.SetDefault("server.host", "0.0.0.0") diff --git a/backend/internal/config/config_test.go b/backend/internal/config/config_test.go index 8783092b7e..b68a26d494 100644 --- a/backend/internal/config/config_test.go +++ b/backend/internal/config/config_test.go @@ -79,6 +79,15 @@ func TestLoadServerTimingConfig(t *testing.T) { }) } +func TestLoadSimpleModeKeyRateLimitEnabledFromEnvironment(t *testing.T) { + resetViperWithJWTSecret(t) + t.Setenv("SIMPLE_MODE_KEY_RATE_LIMIT_ENABLED", "true") + + cfg, err := Load() + require.NoError(t, err) + require.True(t, cfg.SimpleModeKeyRateLimitEnabled) +} + func TestLoadRedisUsernameFromEnvironment(t *testing.T) { resetViperWithJWTSecret(t) t.Setenv("REDIS_USERNAME", "app-user") diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index bf92e3a82e..bf4563f2f2 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -2526,6 +2526,20 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { return } + // A WebSocket may outlive a key's remaining spending window. Recheck + // after acquiring turn slots, including the first account-selection wait. + // Restrict this extra check to the opt-in mode so standard-mode RPM checks + // are not charged a second time for the same request. + checkSimpleModeTurnBilling := func() error { + if h.cfg == nil || h.cfg.RunMode != config.RunModeSimple || !h.cfg.SimpleModeKeyRateLimitEnabled { + return nil + } + if err := h.billingCacheService.CheckBillingEligibility(ctx, apiKey.User, apiKey, apiKey.Group, subscription, service.QuotaPlatform(ctx, apiKey)); err != nil { + return service.NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "billing check failed", err) + } + return nil + } + sessionHash := h.gatewayService.GenerateSessionHashWithFallback( c, firstMessage, @@ -2783,6 +2797,11 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { // turn 级定价：首轮回退到 TurnStarted 的所属 turn 时刻；后续 turn 由 // BeforeTurn 重新冻结 pricingAt 并按最新门复核当前账号。 var turnPricing openAIWSTurnPricing + // Passthrough ingress does not invoke BeforeTurn for the first frame. + if err := checkSimpleModeTurnBilling(); err != nil { + closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, "billing check failed") + return + } hooks := &service.OpenAIWSIngressHooks{ ClientLifecycleContext: clientLifecycleCtx, InitialRequestModel: reqModel, @@ -2893,7 +2912,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { } currentUserRelease = wrapReleaseOnDone(ctx, userReleaseFunc) currentAccountRelease = wrapReleaseOnDone(ctx, accountReleaseFunc) - return nil + return checkSimpleModeTurnBilling() }, AfterTurn: func(turn int, result *service.OpenAIForwardResult, turnErr error) { turnStart := getTurnStart(turn) diff --git a/backend/internal/handler/openai_gateway_handler_test.go b/backend/internal/handler/openai_gateway_handler_test.go index 29b17d1655..b960e90cd3 100644 --- a/backend/internal/handler/openai_gateway_handler_test.go +++ b/backend/internal/handler/openai_gateway_handler_test.go @@ -1916,7 +1916,9 @@ func newOpenAIWSHandlerTestServer(t *testing.T, h *OpenAIGatewayHandler, subject } type openAIResponsesWSUsageLogCase struct { - firstPayload string + simpleModeRejectAtRead int64 + closeReason string + firstPayload string // midPayload 在首个 turn 完成后发送（如 session.update），上游桩会为它 // 回一个 response.completed，客户端按普通事件读取。 midPayload string @@ -2963,7 +2965,13 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU }, nil, nil, nil, nil) } - billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil) + var keyRepo service.APIKeyRepository + if tc.simpleModeRejectAtRead > 0 { + cfg.SimpleModeKeyRateLimitEnabled = true + keyRepo = &simpleModeWSRateLimitRepo{rejectAt: tc.simpleModeRejectAtRead} + } + billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, keyRepo, nil, nil, cfg, nil) + t.Cleanup(billingCacheSvc.Stop) gatewaySvc := service.NewOpenAIGatewayService( accountRepo, usageRepo, @@ -2998,6 +3006,7 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU }, } h := &OpenAIGatewayHandler{ + cfg: cfg, gatewayService: gatewaySvc, billingCacheService: billingCacheSvc, apiKeyService: &service.APIKeyService{}, @@ -3009,6 +3018,9 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU GroupID: &groupID, User: &service.User{ID: 1701, Status: service.StatusActive}, } + if tc.simpleModeRejectAtRead > 0 { + apiKey.RateLimit5h = 1 + } if tc.group != nil { apiKey.Group = tc.group } @@ -3051,7 +3063,11 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU var closeErr coderws.CloseError require.ErrorAs(t, readErr, &closeErr) require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code) - require.Contains(t, closeErr.Reason, "not available for this group") + reason := tc.closeReason + if reason == "" { + reason = "not available for this group" + } + require.Contains(t, closeErr.Reason, reason) _ = clientConn.CloseNow() return openAIResponsesWSUsageLogResult{} } @@ -3086,7 +3102,11 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU var closeErr coderws.CloseError require.ErrorAs(t, readErr, &closeErr) require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code) - require.Contains(t, closeErr.Reason, "not available for this group") + reason := tc.closeReason + if reason == "" { + reason = "not available for this group" + } + require.Contains(t, closeErr.Reason, reason) _ = clientConn.CloseNow() return openAIResponsesWSUsageLogResult{} } @@ -3199,3 +3219,37 @@ data: {"type":"response.failed","error":{"message":"This content was flagged"}} require.False(t, openAIForwardErrorAlreadyCommunicated(c, c.Writer.Size(), errors.New("openai cyber_policy: blocked"))) }) } + +// The loader simulates window usage reaching its limit while a connection is +// waiting for its first account or between completed WebSocket turns. +type simpleModeWSRateLimitRepo struct { + service.APIKeyRepository + reads atomic.Int64 + rejectAt int64 +} + +func (r *simpleModeWSRateLimitRepo) GetRateLimitData(context.Context, int64) (*service.APIKeyRateLimitData, error) { + now := time.Now() + usage := 0.0 + if r.reads.Add(1) >= r.rejectAt { + usage = 1 + } + return &service.APIKeyRateLimitData{Usage5h: usage, Window5hStart: &now}, nil +} +func TestOpenAIResponsesWebSocketSimpleModeRechecksKeyWindows(t *testing.T) { + for _, mode := range []string{service.OpenAIWSIngressModePassthrough, service.OpenAIWSIngressModeCtxPool} { + t.Run(mode+"/after-account-selection", func(t *testing.T) { + runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{ + firstPayload: `{"type":"response.create","model":"gpt-5.1"}`, + ingressMode: mode, simpleModeRejectAtRead: 2, firstFrameCloseExpected: true, closeReason: "billing check failed", + }) + }) + t.Run(mode+"/followup-turn", func(t *testing.T) { + runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{ + firstPayload: `{"type":"response.create","model":"gpt-5.1"}`, + secondPayload: `{"type":"response.create","model":"gpt-5.1"}`, + ingressMode: mode, simpleModeRejectAtRead: 3, secondTurnCloseExpected: true, closeReason: "billing check failed", + }) + }) + } +} diff --git a/backend/internal/service/billing_cache_service.go b/backend/internal/service/billing_cache_service.go index 6ed70a5ba1..691d2558f7 100644 --- a/backend/internal/service/billing_cache_service.go +++ b/backend/internal/service/billing_cache_service.go @@ -733,8 +733,13 @@ func (s *BillingCacheService) IncrementUserPlatformQuotaUsage(userID int64, plat // 订阅模式：检查缓存用量未超过限额（Group限额从参数传入） // platform 为请求的目标平台（如 "anthropic"），传空串 "" 时跳过 user × platform quota 检查。 func (s *BillingCacheService) CheckBillingEligibility(ctx context.Context, user *User, apiKey *APIKey, group *Group, subscription *UserSubscription, platform string) error { - // 简易模式：跳过所有计费检查 - if s.cfg.RunMode == config.RunModeSimple { + // 简易模式默认跳过所有计费检查. An explicit key-window opt-in keeps + // balance/subscription/platform checks bypassed while enforcing the three + // API-key monetary windows from the database source of truth. + if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple { + if s.cfg.SimpleModeKeyRateLimitEnabled { + return s.checkSimpleModeAPIKeyRateLimits(ctx, apiKey) + } return nil } if s.circuitBreaker != nil && !s.circuitBreaker.Allow() { @@ -776,6 +781,34 @@ func (s *BillingCacheService) CheckBillingEligibility(ctx context.Context, user return nil } +// checkSimpleModeAPIKeyRateLimits is deliberately DB-authoritative. Redis +// updates are asynchronous and can be dropped or missed after a committed +// transaction, so using the cache here could let a limited simple-mode key +// continue past its configured window. A read failure fails closed because +// the operator explicitly opted into enforcement. +func (s *BillingCacheService) checkSimpleModeAPIKeyRateLimits(ctx context.Context, apiKey *APIKey) error { + if apiKey == nil || !apiKey.HasRateLimits() { + return nil + } + if s.apiKeyRateLimitLoader == nil { + return ErrBillingServiceUnavailable + } + data, err := s.apiKeyRateLimitLoader.GetRateLimitData(ctx, apiKey.ID) + if err != nil || data == nil { + return ErrBillingServiceUnavailable + } + if apiKey.RateLimit5h > 0 && data.EffectiveUsage5h() >= apiKey.RateLimit5h { + return ErrAPIKeyRateLimit5hExceeded + } + if apiKey.RateLimit1d > 0 && data.EffectiveUsage1d() >= apiKey.RateLimit1d { + return ErrAPIKeyRateLimit1dExceeded + } + if apiKey.RateLimit7d > 0 && data.EffectiveUsage7d() >= apiKey.RateLimit7d { + return ErrAPIKeyRateLimit7dExceeded + } + return nil +} + // checkRPM 执行并行 RPM 限流，所有适用的限制同时生效，任一超限即拒绝： // // 1. (用户, 分组) rpm_override — 最细粒度：管理员为特定用户在特定分组设定的专属限额。 diff --git a/backend/internal/service/billing_cache_service_simple_mode_test.go b/backend/internal/service/billing_cache_service_simple_mode_test.go new file mode 100644 index 0000000000..015563e685 --- /dev/null +++ b/backend/internal/service/billing_cache_service_simple_mode_test.go @@ -0,0 +1,96 @@ +package service + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +type simpleModeRateLimitLoaderStub struct { + data *APIKeyRateLimitData + err error +} + +func (s *simpleModeRateLimitLoaderStub) GetRateLimitData(context.Context, int64) (*APIKeyRateLimitData, error) { + return s.data, s.err +} + +func TestCheckBillingEligibilitySimpleModeKeyRateLimitsAreOptInAndDBAuthoritative(t *testing.T) { + now := time.Now() + tests := []struct { + name string + enabled bool + data *APIKeyRateLimitData + loaderErr error + wantErr error + }{ + { + name: "disabled preserves simple mode bypass", + enabled: false, + data: &APIKeyRateLimitData{Usage5h: 100}, + }, + { + name: "enabled allows below all windows", + enabled: true, + data: &APIKeyRateLimitData{Usage5h: 9, Usage1d: 19, Usage7d: 29, Window5hStart: &now, Window1dStart: &now, Window7dStart: &now}, + }, + { + name: "enabled rejects 5h", + enabled: true, + data: &APIKeyRateLimitData{Usage5h: 10, Window5hStart: &now}, + wantErr: ErrAPIKeyRateLimit5hExceeded, + }, + { + name: "enabled rejects 1d", + enabled: true, + data: &APIKeyRateLimitData{Usage1d: 20, Window1dStart: &now}, + wantErr: ErrAPIKeyRateLimit1dExceeded, + }, + { + name: "enabled rejects 7d", + enabled: true, + data: &APIKeyRateLimitData{Usage7d: 30, Window7dStart: &now}, + wantErr: ErrAPIKeyRateLimit7dExceeded, + }, + { + name: "expired window resets for eligibility", + enabled: true, + data: &APIKeyRateLimitData{Usage5h: 10, Window5hStart: simplePtrTime(now.Add(-RateLimitWindow5h - time.Minute))}, + }, + { + name: "database failure fails closed", + enabled: true, + loaderErr: errors.New("database unavailable"), + wantErr: ErrBillingServiceUnavailable, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := &config.Config{ + RunMode: config.RunModeSimple, + SimpleModeKeyRateLimitEnabled: tt.enabled, + } + svc := &BillingCacheService{ + cfg: cfg, + apiKeyRateLimitLoader: &simpleModeRateLimitLoaderStub{data: tt.data, err: tt.loaderErr}, + } + key := &APIKey{ID: 42, RateLimit5h: 10, RateLimit1d: 20, RateLimit7d: 30} + + err := svc.CheckBillingEligibility(context.Background(), &User{ID: 7, Balance: 0}, key, nil, nil, "") + if tt.wantErr == nil { + require.NoError(t, err) + return + } + require.ErrorIs(t, err, tt.wantErr) + }) + } +} + +func simplePtrTime(t time.Time) *time.Time { + return &t +} diff --git a/backend/internal/service/gateway_simple_mode_record_usage_test.go b/backend/internal/service/gateway_simple_mode_record_usage_test.go new file mode 100644 index 0000000000..fa388da886 --- /dev/null +++ b/backend/internal/service/gateway_simple_mode_record_usage_test.go @@ -0,0 +1,62 @@ +//go:build unit + +package service + +import ( + "context" + "fmt" + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" + "testing" +) + +func TestSimpleModeRecordUsageWindowOptIn(t *testing.T) { + for _, openAI := range []bool{false, true} { + for _, enabled := range []bool{false, true} { + for _, stream := range []bool{false, true} { + t.Run(fmt.Sprintf("openai=%v/enabled=%v/stream=%v", openAI, enabled, stream), func(t *testing.T) { + logs := &openAIRecordUsageLogRepoStub{inserted: true} + billing := &openAIRecordUsageBillingRepoStub{} + users := &openAIRecordUsageUserRepoStub{} + subs := &openAIRecordUsageSubRepoStub{} + key := &APIKey{ID: 1, Quota: 100, RateLimit5h: 30, Group: &Group{RateMultiplier: 1}} + user := &User{ID: 2, Balance: 0} + account := &Account{ID: 3, Type: AccountTypeAPIKey} + if openAI { + svc := newOpenAIRecordUsageServiceWithBillingRepoForTest(logs, billing, users, subs, nil) + svc.cfg.RunMode = config.RunModeSimple + svc.cfg.SimpleModeKeyRateLimitEnabled = enabled + err := svc.RecordUsage(context.Background(), &OpenAIRecordUsageInput{ + Result: &OpenAIForwardResult{RequestID: "simple-request", Model: "gpt-5.1", Usage: OpenAIUsage{InputTokens: 100, OutputTokens: 20}, Stream: stream}, + APIKey: key, User: user, Account: account, + }) + require.NoError(t, err) + } else { + svc := newGatewayRecordUsageServiceWithBillingRepoForTest(logs, billing, users, subs) + svc.cfg.RunMode = config.RunModeSimple + svc.cfg.SimpleModeKeyRateLimitEnabled = enabled + err := svc.RecordUsage(context.Background(), &RecordUsageInput{ + Result: &ForwardResult{RequestID: "simple-request", Model: "claude-sonnet-4", Usage: ClaudeUsage{InputTokens: 100, OutputTokens: 20}, Stream: stream}, + APIKey: key, User: user, Account: account, + }) + require.NoError(t, err) + } + require.Equal(t, 1, logs.calls) + require.Positive(t, logs.lastLog.ActualCost) + require.Zero(t, users.deductCalls) + require.Zero(t, subs.incrementCalls) + if !enabled { + require.Zero(t, billing.calls) + return + } + require.Equal(t, 1, billing.calls) + require.InDelta(t, logs.lastLog.ActualCost, billing.lastCmd.APIKeyRateLimitCost, 1e-12) + require.Zero(t, billing.lastCmd.BalanceCost) + require.Zero(t, billing.lastCmd.SubscriptionCost) + require.Zero(t, billing.lastCmd.APIKeyQuotaCost) + require.Zero(t, billing.lastCmd.AccountQuotaCost) + }) + } + } + } +} diff --git a/backend/internal/service/gateway_usage_billing.go b/backend/internal/service/gateway_usage_billing.go index 5a9a00aea7..53bd3d8ccd 100644 --- a/backend/internal/service/gateway_usage_billing.go +++ b/backend/internal/service/gateway_usage_billing.go @@ -2,6 +2,7 @@ package service import ( "context" + "errors" "log/slog" "strings" "time" @@ -82,6 +83,16 @@ type postUsageBillingParams struct { AccountRateMultiplier float64 APIKeyService APIKeyQuotaUpdater Platform string // 来自 APIKey 关联 Group 的平台标识 + // SimpleModeKeyRateLimitOnly opts the request into the simple-mode billing + // path that records only API-key 5h/1d/7d window usage. It must not trigger + // balance, subscription, account, platform, or lifetime-key-quota effects. + SimpleModeKeyRateLimitOnly bool +} + +var ErrSimpleModeKeyRateLimitBillingUnavailable = errors.New("simple mode api key rate-limit billing unavailable") + +func simpleModeKeyRateLimitBillingEnabled(cfg *config.Config, apiKey *APIKey) bool { + return cfg != nil && cfg.RunMode == config.RunModeSimple && cfg.SimpleModeKeyRateLimitEnabled && apiKey != nil && apiKey.HasRateLimits() } // PlatformFromAPIKey 从 APIKey 关联的 Group 推导 platform 名称。 @@ -306,6 +317,14 @@ func buildUsageBillingCommand(requestID string, usageLog *UsageLog, p *postUsage } } + if p.SimpleModeKeyRateLimitOnly { + if p.Cost.ActualCost > 0 && p.APIKey.HasRateLimits() { + cmd.APIKeyRateLimitCost = p.Cost.ActualCost + } + cmd.Normalize() + return cmd + } + // Record subscription / balance cost using ActualCost so the group (and any // user-specific) rate multiplier consumes subscription quota at the expected // speed. TotalCost remains the raw (pre-multiplier) value; downstream guards @@ -338,6 +357,11 @@ func applyUsageBilling(ctx context.Context, requestID string, usageLog *UsageLog cmd := buildUsageBillingCommand(requestID, usageLog, p) if cmd == nil || cmd.RequestID == "" || repo == nil { + if p.SimpleModeKeyRateLimitOnly { + return false, ErrSimpleModeKeyRateLimitBillingUnavailable + } + // The legacy path is only a fallback for standard billing. Simple mode + // must retain request-id deduplication and never bill other balances. postUsageBilling(ctx, p, deps) return true, nil } @@ -370,6 +394,20 @@ func finalizePostUsageBilling(ctx context.Context, p *postUsageBillingParams, de return } + if p.SimpleModeKeyRateLimitOnly { + // Simple-mode key windows are DB-authoritative during preflight. Clear + // any Redis snapshot after the committed increment as a best-effort aid + // when switching back to standard mode. Enforcement here never relies on + // successful cache invalidation. + if p.APIKey != nil && deps.billingCacheService != nil { + if err := deps.billingCacheService.InvalidateAPIKeyRateLimit(ctx, p.APIKey.ID); err != nil { + logger.LegacyPrintf("service.gateway", "Warning: invalidate simple-mode api key rate-limit cache failed for key %d: %v", p.APIKey.ID, err) + } + } + deps.deferredService.ScheduleLastUsedUpdate(p.Account.ID) + return + } + if p.IsSubscriptionBill { if p.Cost.ActualCost > 0 && p.User != nil && p.APIKey != nil && p.APIKey.GroupID != nil { deps.billingCacheService.QueueUpdateSubscriptionUsage(p.User.ID, *p.APIKey.GroupID, p.Cost.ActualCost) @@ -378,7 +416,7 @@ func finalizePostUsageBilling(ctx context.Context, p *postUsageBillingParams, de syncBalanceCacheAfterDeduction(ctx, p, deps, result) } - if p.Cost.ActualCost > 0 && p.APIKey != nil && p.APIKey.HasRateLimits() { + if p.Cost.ActualCost > 0 && p.APIKey != nil && p.APIKey.HasRateLimits() && deps.billingCacheService != nil { deps.billingCacheService.QueueUpdateAPIKeyRateLimitUsage(p.APIKey.ID, p.Cost.ActualCost) } @@ -391,7 +429,7 @@ func finalizePostUsageBilling(ctx context.Context, p *postUsageBillingParams, de // 限制在并发 in-flight 请求数量内（旧实现的异步入队会让超支无限累积直到 worker 处理） // - DB 异步(flusher_enabled=false):在独立 goroutine 中走 detached context,失败用 ALERT log 触发 oncall 对账 // - flusher_enabled=true:不直写 DB,由 flusher 异步批量刷（markDirty 已在 IncrementUserPlatformQuotaUsage 内部完成） - if !p.IsSubscriptionBill && p.Platform != "" && p.Cost.ActualCost > 0 && p.User != nil && deps.userPlatformQuotaRepo != nil { + if !p.IsSubscriptionBill && p.Platform != "" && p.Cost.ActualCost > 0 && p.User != nil && deps.userPlatformQuotaRepo != nil && deps.billingCacheService != nil { if deps.billingCacheService.HasUserPlatformQuotaLimit(ctx, p.User.ID, p.Platform) { deps.billingCacheService.IncrementUserPlatformQuotaUsage(p.User.ID, p.Platform, p.Cost.ActualCost) if deps.cfg == nil || !deps.cfg.Database.UserPlatformQuotaFlusherEnabled { @@ -835,7 +873,8 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage ) } - if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple { + simpleModeKeyRateLimitOnly := simpleModeKeyRateLimitBillingEnabled(s.cfg, apiKey) + if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple && !simpleModeKeyRateLimitOnly { writeUsageLogBestEffort(ctx, s.usageLogRepo, usageLog, "service.gateway") logger.LegacyPrintf("service.gateway", "[SIMPLE MODE] Usage recorded (not billed): user=%d, tokens=%d", usageLog.UserID, usageLog.TotalTokens()) s.deferredService.ScheduleLastUsedUpdate(account.ID) @@ -854,16 +893,17 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage } requestID := usageLog.RequestID _, billingErr := applyUsageBilling(ctx, requestID, usageLog, &postUsageBillingParams{ - Cost: cost, - User: user, - APIKey: apiKey, - Account: account, - Subscription: subscription, - RequestPayloadHash: resolveUsageBillingPayloadFingerprint(ctx, input.RequestPayloadHash), - IsSubscriptionBill: isSubscriptionBilling, - AccountRateMultiplier: accountRateMultiplier, - APIKeyService: input.APIKeyService, - Platform: quotaPlatform, + Cost: cost, + User: user, + APIKey: apiKey, + Account: account, + Subscription: subscription, + RequestPayloadHash: resolveUsageBillingPayloadFingerprint(ctx, input.RequestPayloadHash), + IsSubscriptionBill: isSubscriptionBilling && !simpleModeKeyRateLimitOnly, + AccountRateMultiplier: accountRateMultiplier, + APIKeyService: input.APIKeyService, + Platform: quotaPlatform, + SimpleModeKeyRateLimitOnly: simpleModeKeyRateLimitOnly, }, s.billingDeps(), s.usageBillingRepo) if billingErr != nil { diff --git a/backend/internal/service/gateway_usage_billing_simple_mode_test.go b/backend/internal/service/gateway_usage_billing_simple_mode_test.go new file mode 100644 index 0000000000..cd65253398 --- /dev/null +++ b/backend/internal/service/gateway_usage_billing_simple_mode_test.go @@ -0,0 +1,144 @@ +package service + +import ( + "context" + "errors" + "strconv" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestBuildUsageBillingCommandSimpleModeOnlyChargesAPIKeyWindows(t *testing.T) { + apiKey := &APIKey{ID: 13, Quota: 100, RateLimit5h: 10} + user := &User{ID: 7} + account := &Account{ID: 9, Type: AccountTypeAPIKey} + p := &postUsageBillingParams{ + Cost: &CostBreakdown{ActualCost: 3.25, TotalCost: 2.5}, + User: user, + APIKey: apiKey, + Account: account, + APIKeyService: &apiKeyQuotaUpdaterStub{}, + SimpleModeKeyRateLimitOnly: true, + } + + cmd := buildUsageBillingCommand("simple-req", nil, p) + require.NotNil(t, cmd) + require.Zero(t, cmd.BalanceCost) + require.Zero(t, cmd.SubscriptionCost) + require.Zero(t, cmd.APIKeyQuotaCost) + require.Zero(t, cmd.AccountQuotaCost) + require.Equal(t, 3.25, cmd.APIKeyRateLimitCost) +} + +type apiKeyQuotaUpdaterStub struct{} + +func (apiKeyQuotaUpdaterStub) UpdateQuotaUsed(context.Context, int64, float64) error { return nil } +func (apiKeyQuotaUpdaterStub) UpdateRateLimitUsage(context.Context, int64, float64) error { + return nil +} + +type simpleModeUsageBillingRepoStub struct { + UsageBillingRepository + seen map[string]struct{} + cmds []*UsageBillingCommand +} + +func (s *simpleModeUsageBillingRepoStub) Apply(_ context.Context, cmd *UsageBillingCommand) (*UsageBillingApplyResult, error) { + s.cmds = append(s.cmds, cmd) + if s.seen == nil { + s.seen = make(map[string]struct{}) + } + key := cmd.RequestID + ":" + strconv.FormatInt(cmd.APIKeyID, 10) + if _, ok := s.seen[key]; ok { + return &UsageBillingApplyResult{Applied: false}, nil + } + s.seen[key] = struct{}{} + return &UsageBillingApplyResult{Applied: true}, nil +} + +func TestApplyUsageBillingSimpleModeDeduplicatesWithoutBalanceEffects(t *testing.T) { + repo := &simpleModeUsageBillingRepoStub{} + p := &postUsageBillingParams{ + Cost: &CostBreakdown{ActualCost: 3.25, TotalCost: 3.25}, + User: &User{ID: 7, Balance: 0}, + APIKey: &APIKey{ID: 13, Quota: 100, RateLimit5h: 10}, + Account: &Account{ID: 9, Type: AccountTypeAPIKey}, + APIKeyService: &apiKeyQuotaUpdaterStub{}, + SimpleModeKeyRateLimitOnly: true, + } + + first, err := applyUsageBilling(context.Background(), "simple-req", nil, p, &billingDeps{deferredService: &DeferredService{}}, repo) + require.NoError(t, err) + require.True(t, first) + second, err := applyUsageBilling(context.Background(), "simple-req", nil, p, &billingDeps{deferredService: &DeferredService{}}, repo) + require.NoError(t, err) + require.False(t, second) + require.Len(t, repo.cmds, 2) + for _, cmd := range repo.cmds { + require.Zero(t, cmd.BalanceCost) + require.Zero(t, cmd.SubscriptionCost) + require.Zero(t, cmd.APIKeyQuotaCost) + require.Zero(t, cmd.AccountQuotaCost) + require.Equal(t, 3.25, cmd.APIKeyRateLimitCost) + } +} + +func TestApplyUsageBillingSimpleModeRejectsLegacyFallback(t *testing.T) { + for _, missing := range []string{"repository", "request_id", "api_key"} { + t.Run(missing, func(t *testing.T) { + p := &postUsageBillingParams{ + Cost: &CostBreakdown{ActualCost: 1}, User: &User{ID: 7}, + APIKey: &APIKey{ID: 13, RateLimit5h: 10}, Account: &Account{ID: 9}, + SimpleModeKeyRateLimitOnly: true, + } + var repo UsageBillingRepository = &simpleModeUsageBillingRepoStub{} + requestID := "simple-req" + switch missing { + case "repository": + repo = nil + case "request_id": + requestID = "" + case "api_key": + p.APIKey = nil + } + _, err := applyUsageBilling(context.Background(), requestID, nil, p, &billingDeps{deferredService: &DeferredService{}}, repo) + require.ErrorIs(t, err, ErrSimpleModeKeyRateLimitBillingUnavailable) + }) + } +} + +// A failed cache eviction must not fall through into standard-mode charges. +type simpleModeInvalidationCache struct { + BillingCache + invalidated []int64 +} + +func (c *simpleModeInvalidationCache) InvalidateAPIKeyRateLimit(_ context.Context, keyID int64) error { + c.invalidated = append(c.invalidated, keyID) + return errors.New("redis unavailable") +} +func TestFinalizeSimpleModePreservesLastUsedWithoutFinancialCacheWrites(t *testing.T) { + for _, subscription := range []bool{false, true} { + t.Run(strconv.FormatBool(subscription), func(t *testing.T) { + cache := &simpleModeInvalidationCache{} + writes := make(chan cacheWriteTask, 4) + deferred := &DeferredService{} + groupID := int64(5) + p := &postUsageBillingParams{ + Cost: &CostBreakdown{ActualCost: 3}, User: &User{ID: 7}, + APIKey: &APIKey{ID: 13, GroupID: &groupID, RateLimit5h: 10}, + Account: &Account{ID: 9}, IsSubscriptionBill: subscription, + SimpleModeKeyRateLimitOnly: true, + } + finalizePostUsageBilling(context.Background(), p, &billingDeps{ + billingCacheService: &BillingCacheService{cache: cache, cacheWriteChan: writes}, + deferredService: deferred, + }, &UsageBillingApplyResult{Applied: true}) + require.Equal(t, []int64{13}, cache.invalidated) + require.Empty(t, writes, "simple mode must not enqueue balance, subscription or window increments") + _, scheduled := deferred.lastUsedUpdates.Load(int64(9)) + require.True(t, scheduled, "early return must preserve the account activity update") + }) + } +} diff --git a/backend/internal/service/openai_gateway_usage.go b/backend/internal/service/openai_gateway_usage.go index 86ebf21ebd..cd2c7aadea 100644 --- a/backend/internal/service/openai_gateway_usage.go +++ b/backend/internal/service/openai_gateway_usage.go @@ -485,7 +485,8 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec ) } - if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple { + simpleModeKeyRateLimitOnly := simpleModeKeyRateLimitBillingEnabled(s.cfg, apiKey) + if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple && !simpleModeKeyRateLimitOnly { writeUsageLogBestEffort(ctx, s.usageLogRepo, usageLog, "service.openai_gateway") logger.LegacyPrintf("service.openai_gateway", "[SIMPLE MODE] Usage recorded (not billed): user=%d, tokens=%d", usageLog.UserID, usageLog.TotalTokens()) s.deferredService.ScheduleLastUsedUpdate(account.ID) @@ -499,21 +500,19 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec quotaPlatform = PlatformFromAPIKey(apiKey) } - billingErr := func() error { - _, err := applyUsageBilling(ctx, requestID, usageLog, &postUsageBillingParams{ - Cost: cost, - User: user, - APIKey: apiKey, - Account: account, - Subscription: subscription, - RequestPayloadHash: resolveUsageBillingPayloadFingerprint(ctx, input.RequestPayloadHash), - IsSubscriptionBill: isSubscriptionBilling, - AccountRateMultiplier: accountRateMultiplier, - APIKeyService: input.APIKeyService, - Platform: quotaPlatform, - }, s.billingDeps(), s.usageBillingRepo) - return err - }() + _, billingErr := applyUsageBilling(ctx, requestID, usageLog, &postUsageBillingParams{ + Cost: cost, + User: user, + APIKey: apiKey, + Account: account, + Subscription: subscription, + RequestPayloadHash: resolveUsageBillingPayloadFingerprint(ctx, input.RequestPayloadHash), + IsSubscriptionBill: isSubscriptionBilling && !simpleModeKeyRateLimitOnly, + AccountRateMultiplier: accountRateMultiplier, + APIKeyService: input.APIKeyService, + Platform: quotaPlatform, + SimpleModeKeyRateLimitOnly: simpleModeKeyRateLimitOnly, + }, s.billingDeps(), s.usageBillingRepo) if billingErr != nil { usageLog.ActualCost = 0 diff --git a/deploy/.env.example b/deploy/.env.example index 31b913e1cd..1299e45c56 100644 --- a/deploy/.env.example +++ b/deploy/.env.example @@ -122,6 +122,8 @@ RUN_MODE=standard # 简易模式启动时自动补齐默认分组；留空或未设置时使用 YAML 配置，未配置时应用默认 true # 显式 true/false 会覆盖 YAML；false 由用户自行管理，已有分组和运行时行为不变 SIMPLE_MODE_AUTO_CREATE_DEFAULT_GROUPS= +# Opt into API-key 5h/1d/7d spending windows while RUN_MODE=simple. Default false. +# SIMPLE_MODE_KEY_RATE_LIMIT_ENABLED=false # Timezone TZ=Asia/Shanghai diff --git a/deploy/config.example.yaml b/deploy/config.example.yaml index 3507cbc03d..e1e1c6df7e 100644 --- a/deploy/config.example.yaml +++ b/deploy/config.example.yaml @@ -105,6 +105,9 @@ simple_mode: # 环境变量：SIMPLE_MODE_AUTO_CREATE_DEFAULT_GROUPS auto_create_default_groups: true +# Enforce API-key 5h/1d/7d spending windows in simple mode; other billing stays disabled. +simple_mode_key_rate_limit_enabled: false + # ============================================================================= # CORS Configuration # 跨域资源共享 (CORS) 配置 diff --git a/deploy/docker-compose.dev.yml b/deploy/docker-compose.dev.yml index 92a075ceb7..f2acd1ff36 100644 --- a/deploy/docker-compose.dev.yml +++ b/deploy/docker-compose.dev.yml @@ -31,6 +31,7 @@ services: - ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false} - RUN_MODE=${RUN_MODE:-standard} - SIMPLE_MODE_AUTO_CREATE_DEFAULT_GROUPS=${SIMPLE_MODE_AUTO_CREATE_DEFAULT_GROUPS:-} + - SIMPLE_MODE_KEY_RATE_LIMIT_ENABLED=${SIMPLE_MODE_KEY_RATE_LIMIT_ENABLED:-false} - DATABASE_HOST=postgres - DATABASE_PORT=5432 - DATABASE_USER=${POSTGRES_USER:-sub2api} diff --git a/deploy/docker-compose.local.yml b/deploy/docker-compose.local.yml index 8e8956cd69..5b6c8d9eb5 100644 --- a/deploy/docker-compose.local.yml +++ b/deploy/docker-compose.local.yml @@ -56,6 +56,7 @@ services: - ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false} - RUN_MODE=${RUN_MODE:-standard} - SIMPLE_MODE_AUTO_CREATE_DEFAULT_GROUPS=${SIMPLE_MODE_AUTO_CREATE_DEFAULT_GROUPS:-} + - SIMPLE_MODE_KEY_RATE_LIMIT_ENABLED=${SIMPLE_MODE_KEY_RATE_LIMIT_ENABLED:-false} - UPDATE_GITHUB_TOKEN=${UPDATE_GITHUB_TOKEN:-} # ======================================================================= diff --git a/deploy/docker-compose.standalone.yml b/deploy/docker-compose.standalone.yml index af98c34a60..abbd8c23f7 100644 --- a/deploy/docker-compose.standalone.yml +++ b/deploy/docker-compose.standalone.yml @@ -42,6 +42,7 @@ services: - ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false} - RUN_MODE=${RUN_MODE:-standard} - SIMPLE_MODE_AUTO_CREATE_DEFAULT_GROUPS=${SIMPLE_MODE_AUTO_CREATE_DEFAULT_GROUPS:-} + - SIMPLE_MODE_KEY_RATE_LIMIT_ENABLED=${SIMPLE_MODE_KEY_RATE_LIMIT_ENABLED:-false} - UPDATE_GITHUB_TOKEN=${UPDATE_GITHUB_TOKEN:-} # ======================================================================= diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index f9086dda4a..1a09032156 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -52,6 +52,7 @@ services: - ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false} - RUN_MODE=${RUN_MODE:-standard} - SIMPLE_MODE_AUTO_CREATE_DEFAULT_GROUPS=${SIMPLE_MODE_AUTO_CREATE_DEFAULT_GROUPS:-} + - SIMPLE_MODE_KEY_RATE_LIMIT_ENABLED=${SIMPLE_MODE_KEY_RATE_LIMIT_ENABLED:-false} - UPDATE_GITHUB_TOKEN=${UPDATE_GITHUB_TOKEN:-} - ALIPAY_MOBILE_PRECREATE_DEEP_LINK=${ALIPAY_MOBILE_PRECREATE_DEEP_LINK:-}