opensource/sub2api
3
0
Fork 0
mirror of https://github.com/Wei-Shaw/sub2api.git synced 2026-10-16 14:57:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: optionally enforce API-key spending windows in simple mode

Browse Source
This commit is contained in:
zhangcj
2026-09-23 08:17:48 +08:00
parent 20a94fbb56
commit 3b336b2eca
18 changed files with 509 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+2
View File
@@ -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
---
README_CN.md
+2
View File
@@ -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` 才允许启动
---
backend/internal/config/config.go
+4
View File
@@ -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")
backend/internal/config/config_test.go
+9
View File
@@ -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")
backend/internal/handler/openai_gateway_handler.go
+20 -1
View File
@@ -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)
backend/internal/handler/openai_gateway_handler_test.go
+58 -4
View File
@@ -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",
})
})
}
}
backend/internal/service/billing_cache_service.go
+35 -2
View File
@@ -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 — 最细粒度：管理员为特定用户在特定分组设定的专属限额。
backend/internal/service/billing_cache_service_simple_mode_test.go
+96
View File
@@ -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
}
backend/internal/service/gateway_simple_mode_record_usage_test.go
+62
View File
@@ -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)
})
}
}
}
}
backend/internal/service/gateway_usage_billing.go
+53 -13
View File
@@ -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 {
backend/internal/service/gateway_usage_billing_simple_mode_test.go
+144
View File
@@ -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")
})
}
}
backend/internal/service/openai_gateway_usage.go
+15 -16
View File
@@ -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
deploy/.env.example
+2
View File
@@ -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
deploy/config.example.yaml
+3
View File
@@ -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) 配置
deploy/docker-compose.dev.yml
+1
View File
@@ -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}
deploy/docker-compose.local.yml
+1
View File
@@ -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:-}
# =======================================================================
deploy/docker-compose.standalone.yml
+1
View File
@@ -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:-}
# =======================================================================
deploy/docker-compose.yml
+1
View File
@@ -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:-}