diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index f7cfcead2e..1b07ae3a6a 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -941,6 +941,9 @@ type InflightReservationConfig struct { MaxInputTokens int `mapstructure:"max_input_tokens"` // MaxReservationUSD 单请求预留金额上限;0 表示不设上限。 MaxReservationUSD float64 `mapstructure:"max_reservation_usd"` + // FailClosedOnUnpriced 无法为请求估算费用(模型/分组/渠道均无定价)时是否拒绝请求。 + // 默认 false:放行且不预留(fail-open,节流告警日志)。 + FailClosedOnUnpriced bool `mapstructure:"fail_closed_on_unpriced"` } type CircuitBreakerConfig struct { @@ -2119,6 +2122,7 @@ func setDefaults() { viper.SetDefault("billing.inflight_reservation.max_output_tokens", 128000) viper.SetDefault("billing.inflight_reservation.max_input_tokens", 200000) viper.SetDefault("billing.inflight_reservation.max_reservation_usd", 0) + viper.SetDefault("billing.inflight_reservation.fail_closed_on_unpriced", false) // Turnstile viper.SetDefault("turnstile.required", false) diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index acf4e21acb..f83494ef9a 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -264,7 +264,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) { } // 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。 - inflightRelease, err := reserveInflightBalance(c.Request.Context(), h.billingCacheService, h.gatewayService, apiKey, subscription, reqModel, body) + inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body)) if err != nil { reqLog.Info("gateway.inflight_reservation_rejected", zap.Error(err)) status, code, message, retryAfter := billingErrorDetails(err) @@ -2527,9 +2527,12 @@ func (h *GatewayHandler) submitUsageRecordTask(parent context.Context, task serv if task == nil { return } - task = wrapUsageRecordTaskContext(parent, task) + task, abandon := wrapUsageRecordTaskContext(parent, task) if h.usageRecordWorkerPool != nil { if mode := h.usageRecordWorkerPool.Submit(task); mode != service.UsageRecordSubmitModeDroppedStopped { + if mode.Dropped() { + abandon() + } return } // 池已停止(进程关停窗口):计费任务不能静默丢失,降级为内联同步执行。 @@ -2557,7 +2560,7 @@ func (h *GatewayHandler) submitMandatoryUsageRecordTask(parent context.Context, if task == nil { return } - task = wrapUsageRecordTaskContext(parent, task) + task, _ = wrapUsageRecordTaskContext(parent, task) if h.usageRecordWorkerPool != nil { if mode := h.usageRecordWorkerPool.Submit(task); !mode.Dropped() { return diff --git a/backend/internal/handler/gateway_handler_chat_completions.go b/backend/internal/handler/gateway_handler_chat_completions.go index b844d27d1f..79ca025303 100644 --- a/backend/internal/handler/gateway_handler_chat_completions.go +++ b/backend/internal/handler/gateway_handler_chat_completions.go @@ -145,7 +145,7 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) { } // 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。 - inflightRelease, err := reserveInflightBalance(c.Request.Context(), h.billingCacheService, h.gatewayService, apiKey, subscription, reqModel, body) + inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body)) if err != nil { reqLog.Info("gateway.cc.inflight_reservation_rejected", zap.Error(err)) status, code, message, retryAfter := billingErrorDetails(err) diff --git a/backend/internal/handler/gateway_handler_responses.go b/backend/internal/handler/gateway_handler_responses.go index a10ad8d4ec..a13e438775 100644 --- a/backend/internal/handler/gateway_handler_responses.go +++ b/backend/internal/handler/gateway_handler_responses.go @@ -154,7 +154,7 @@ func (h *GatewayHandler) Responses(c *gin.Context) { } // 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。 - inflightRelease, err := reserveInflightBalance(requestCtx, h.billingCacheService, h.gatewayService, apiKey, subscription, reqModel, body) + inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body)) if err != nil { reqLog.Info("gateway.responses.inflight_reservation_rejected", zap.Error(err)) status, code, message, retryAfter := billingErrorDetails(err) diff --git a/backend/internal/handler/gateway_inflight_reservation.go b/backend/internal/handler/gateway_inflight_reservation.go index 08ebaafb9f..5928220dc0 100644 --- a/backend/internal/handler/gateway_inflight_reservation.go +++ b/backend/internal/handler/gateway_inflight_reservation.go @@ -4,12 +4,13 @@ import ( "context" "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" "github.com/tidwall/gjson" ) -// inflightReservationEstimator 估算单请求在途预留金额(USD)。 +// inflightReservationEstimator 估算单请求在途预留金额(USD);false 表示无法定价。 type inflightReservationEstimator interface { - EstimateInflightReservation(ctx context.Context, apiKey *service.APIKey, model string, bodyBytes, maxTokens int) float64 + EstimateInflightReservation(ctx context.Context, apiKey *service.APIKey, req service.InflightEstimateRequest) (float64, bool) } // requestMaxOutputTokens 从请求体中提取输出 token 上限(兼容 Anthropic / OpenAI Chat / Responses / Gemini)。 @@ -22,28 +23,110 @@ func requestMaxOutputTokens(body []byte) int { return 0 } +// tokenInflightEstimate 文本类请求的估算输入。 +func tokenInflightEstimate(model string, body []byte) service.InflightEstimateRequest { + return service.InflightEstimateRequest{ + Model: model, + BodyBytes: len(body), + MaxTokens: requestMaxOutputTokens(body), + Kind: service.InflightEstimateToken, + } +} + +func inflightNoop() {} + // reserveInflightBalance 在 CheckBillingEligibility 之后为余额模式请求登记在途预留。 -// 返回的 release 必须 defer 调用(覆盖成功/错误/客户端断开/failover 等所有路径); -// 开关关闭、订阅模式、估算失败或 Redis 故障时返回 no-op(fail-open)。 +// +// 成功时把预留句柄挂到 c.Request 的 context 上:之后通过 submit*UsageRecordTask 提交的 +// 计费任务会接管一个引用,直到余额缓存实际扣减后才释放。返回的 done 必须 defer 调用 +// (停止续期并归还 handler 引用;没有提交计费任务时即刻释放)。 +// 开关关闭、订阅模式、Redis 故障时返回 no-op(fail-open);无法定价时默认 fail-open, +// 配置 fail_closed_on_unpriced=true 时返回 ErrInsufficientBalance。 func reserveInflightBalance( + c *gin.Context, + billing *service.BillingCacheService, + estimator inflightReservationEstimator, + apiKey *service.APIKey, + subscription *service.UserSubscription, + req service.InflightEstimateRequest, +) (func(), error) { + if c == nil || c.Request == nil { + return inflightNoop, nil + } + ctx, done, err := reserveInflightBalanceCtx(c.Request.Context(), billing, estimator, apiKey, subscription, req) + if err != nil { + return inflightNoop, err + } + c.Request = c.Request.WithContext(ctx) + return done, nil +} + +// reserveInflightBalanceCtx 同 reserveInflightBalance,但返回携带预留句柄的新 context +// (供 WebSocket 等自管 context 的路径)。 +func reserveInflightBalanceCtx( ctx context.Context, billing *service.BillingCacheService, estimator inflightReservationEstimator, apiKey *service.APIKey, subscription *service.UserSubscription, - model string, - body []byte, -) (func(), error) { - noop := func() {} + req service.InflightEstimateRequest, +) (context.Context, func(), error) { if billing == nil || estimator == nil || apiKey == nil || apiKey.User == nil || !billing.InflightReservationEnabled() { - return noop, nil + return ctx, inflightNoop, nil } if apiKey.Group != nil && apiKey.Group.IsSubscriptionType() && subscription != nil { - return noop, nil + return ctx, inflightNoop, nil + } + estimate, priced := estimator.EstimateInflightReservation(ctx, apiKey, req) + if !priced && billing.InflightReservationFailClosedOnUnpriced() { + return ctx, inflightNoop, service.ErrInsufficientBalance } - estimate := estimator.EstimateInflightReservation(ctx, apiKey, model, len(body), requestMaxOutputTokens(body)) if estimate <= 0 { - return noop, nil + return ctx, inflightNoop, nil } - return billing.ReserveInflightBalance(ctx, apiKey.User, apiKey.Group, subscription, estimate) + res, err := billing.ReserveInflight(ctx, apiKey.User, apiKey.Group, subscription, estimate) + if err != nil { + return ctx, inflightNoop, err + } + if res == nil { + return ctx, inflightNoop, nil + } + return service.WithInflightReservation(ctx, res), res.HandlerDone, nil +} + +// grokMediaInflightEstimate 媒体生成请求的估算输入;状态/内容查询返回空模型(不预留: +// 查询会为已生成的媒体计费,不能因余额预留而拦截用户取回已付费结果)。 +func grokMediaInflightEstimate(endpoint service.GrokMediaEndpoint, model string, info service.GrokMediaRequestInfo, body []byte) service.InflightEstimateRequest { + if !endpoint.IsGenerationRequest() { + return service.InflightEstimateRequest{} + } + switch endpoint { + case service.GrokMediaEndpointImagesGenerations, service.GrokMediaEndpointImagesEdits: + return service.InflightEstimateRequest{Model: model, BodyBytes: len(body), Kind: service.InflightEstimateImage, Units: info.N} + default: + return service.InflightEstimateRequest{ + Model: model, + BodyBytes: len(body), + Kind: service.InflightEstimateVideo, + Units: 1, + VideoResolution: info.Resolution, + VideoDurationSeconds: info.DurationSeconds, + } + } +} + +// grokVoiceSTTBytesPerSecond STT 时长粗估(~128kbps 压缩音频)。 +const grokVoiceSTTBytesPerSecond = 16000 + +// grokVoiceInflightEstimate 语音 HTTP 接口估算:TTS 按输入字符数(百万字符),STT 按音频字节粗估时长(小时)。 +// 其他接口(custom-voices)无音频计量,返回的估算为 0(不预留)。 +func grokVoiceInflightEstimate(endpoint string, body []byte) service.InflightEstimateRequest { + req := service.InflightEstimateRequest{Model: endpoint, Kind: service.InflightEstimateAudio, AudioMode: endpoint} + switch endpoint { + case "tts": + req.AudioUnits = float64(len([]rune(extractGrokTTSInputText(body)))) / 1e6 + case "stt": + req.AudioUnits = float64(len(body)) / grokVoiceSTTBytesPerSecond / 3600 + } + return req } diff --git a/backend/internal/handler/gateway_inflight_reservation_test.go b/backend/internal/handler/gateway_inflight_reservation_test.go index a76f5eb6cb..fa99884c6a 100644 --- a/backend/internal/handler/gateway_inflight_reservation_test.go +++ b/backend/internal/handler/gateway_inflight_reservation_test.go @@ -1,11 +1,18 @@ package handler import ( + "bytes" "context" + "net/http" + "net/http/httptest" + "sync" "testing" + "time" "github.com/Wei-Shaw/sub2api/internal/config" + middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) @@ -18,29 +25,180 @@ func TestRequestMaxOutputTokens(t *testing.T) { require.Equal(t, 0, requestMaxOutputTokens([]byte(`{}`))) } -type countingEstimator struct{ calls int } +type countingEstimator struct { + calls int + cost float64 + priced bool +} -func (e *countingEstimator) EstimateInflightReservation(context.Context, *service.APIKey, string, int, int) float64 { +func (e *countingEstimator) EstimateInflightReservation(context.Context, *service.APIKey, service.InflightEstimateRequest) (float64, bool) { e.calls++ - return 1 + return e.cost, e.priced +} + +func newInflightTestGinContext() *gin.Context { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/", nil) + return c } func TestReserveInflightBalance_SkipsWhenDisabledOrSubscription(t *testing.T) { cfg := &config.Config{} billing := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil) t.Cleanup(billing.Stop) - est := &countingEstimator{} + est := &countingEstimator{cost: 1, priced: true} apiKey := &service.APIKey{User: &service.User{ID: 1}} - release, err := reserveInflightBalance(context.Background(), billing, est, apiKey, nil, "m", []byte(`{}`)) + done, err := reserveInflightBalance(newInflightTestGinContext(), billing, est, apiKey, nil, tokenInflightEstimate("m", []byte(`{}`))) require.NoError(t, err) - release() + done() require.Equal(t, 0, est.calls, "disabled switch must not even estimate") cfg.Billing.InflightReservation.Enabled = true apiKey.Group = &service.Group{SubscriptionType: service.SubscriptionTypeSubscription} - release, err = reserveInflightBalance(context.Background(), billing, est, apiKey, &service.UserSubscription{}, "m", []byte(`{}`)) + done, err = reserveInflightBalance(newInflightTestGinContext(), billing, est, apiKey, &service.UserSubscription{}, tokenInflightEstimate("m", []byte(`{}`))) require.NoError(t, err) - release() + done() require.Equal(t, 0, est.calls, "subscription mode must be unaffected") } + +func TestReserveInflightBalance_UnpricedFailOpenByDefaultFailClosedOptIn(t *testing.T) { + cache := newHandlerInflightCache(10) + cfg := &config.Config{} + cfg.Billing.InflightReservation = config.InflightReservationConfig{Enabled: true, TTLSeconds: 60} + billing := service.NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(billing.Stop) + apiKey := &service.APIKey{User: &service.User{ID: 1}} + est := &countingEstimator{priced: false} + + done, err := reserveInflightBalance(newInflightTestGinContext(), billing, est, apiKey, nil, tokenInflightEstimate("unknown", nil)) + require.NoError(t, err) + done() + + cfg.Billing.InflightReservation.FailClosedOnUnpriced = true + _, err = reserveInflightBalance(newInflightTestGinContext(), billing, est, apiKey, nil, tokenInflightEstimate("unknown", nil)) + require.ErrorIs(t, err, service.ErrInsufficientBalance) +} + +// handlerInflightCache 内存版余额缓存 + 在途预留(语义同 Redis Lua)。 +type handlerInflightCache struct { + service.BillingCache + mu sync.Mutex + balance float64 + res map[string]float64 +} + +func newHandlerInflightCache(balance float64) *handlerInflightCache { + return &handlerInflightCache{balance: balance, res: map[string]float64{}} +} + +func (m *handlerInflightCache) GetUserBalance(context.Context, int64) (float64, error) { + m.mu.Lock() + defer m.mu.Unlock() + return m.balance, nil +} + +func (m *handlerInflightCache) GetUserPlatformQuotaCache(context.Context, int64, string) (*service.UserPlatformQuotaCacheEntry, bool, error) { + return nil, false, nil +} + +func (m *handlerInflightCache) ReserveInflightBalance(_ context.Context, _ int64, id string, amount, balance float64, _ time.Duration) (bool, float64, error) { + m.mu.Lock() + defer m.mu.Unlock() + sum := 0.0 + for _, v := range m.res { + sum += v + } + if len(m.res) > 0 && balance-sum < amount { + return false, sum, nil + } + m.res[id] = amount + return true, sum, nil +} + +func (m *handlerInflightCache) ReleaseInflightBalance(_ context.Context, _ int64, id string) error { + m.mu.Lock() + defer m.mu.Unlock() + delete(m.res, id) + return nil +} + +func (m *handlerInflightCache) count() int { + m.mu.Lock() + defer m.mu.Unlock() + return len(m.res) +} + +func TestWrapUsageRecordTaskContext_HandsReservationToBillingTask(t *testing.T) { + cache := newHandlerInflightCache(1) + cfg := &config.Config{} + cfg.Billing.InflightReservation = config.InflightReservationConfig{Enabled: true, TTLSeconds: 60} + billing := service.NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(billing.Stop) + apiKey := &service.APIKey{User: &service.User{ID: 5}} + + c := newInflightTestGinContext() + done, err := reserveInflightBalance(c, billing, &countingEstimator{cost: 0.9, priced: true}, apiKey, nil, tokenInflightEstimate("m", nil)) + require.NoError(t, err) + require.Equal(t, 1, cache.count()) + + ran := false + task, abandon := wrapUsageRecordTaskContext(c.Request.Context(), func(context.Context) { ran = true }) + done() // handler returns; billing still pending + require.Equal(t, 1, cache.count(), "reservation held until the billing task finishes") + task(context.Background()) + require.True(t, ran) + require.Equal(t, 0, cache.count()) + abandon() // idempotent with the task's own done + + // Dropped task: the submitter abandons it and the reservation is released. + c2 := newInflightTestGinContext() + done2, err := reserveInflightBalance(c2, billing, &countingEstimator{cost: 0.9, priced: true}, apiKey, nil, tokenInflightEstimate("m", nil)) + require.NoError(t, err) + _, abandon2 := wrapUsageRecordTaskContext(c2.Request.Context(), func(context.Context) {}) + done2() + require.Equal(t, 1, cache.count()) + abandon2() + require.Equal(t, 0, cache.count()) +} + +// 新接入的端点(独立 web_search):在途预留超过余额时拒绝,且不残留预留。 +func TestWebSearch_RejectsWhenInflightExceedsBalance(t *testing.T) { + cache := newHandlerInflightCache(1.5) + cfg := &config.Config{} + cfg.Billing.InflightReservation = config.InflightReservationConfig{Enabled: true, TTLSeconds: 60} + billingCache := service.NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(billingCache.Stop) + billing := service.NewBillingService(cfg, nil) + gw := service.NewGatewayService( + nil, nil, nil, nil, nil, nil, nil, nil, cfg, nil, nil, billing, nil, nil, + nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, service.NewModelPricingResolver(nil, billing), nil, nil, nil, + ) + h := &GatewayHandler{gatewayService: gw, billingCacheService: billingCache} + + groupID := int64(3) + perK := 1000.0 // $1 per search + apiKey := &service.APIKey{ + ID: 9, User: &service.User{ID: 42, Balance: 1.5}, GroupID: &groupID, + Group: &service.Group{ID: groupID, Platform: service.PlatformGrok, RateMultiplier: 1, SearchPricePer1k: &perK}, + } + + // Another in-flight request of this user already holds $1. + held, err := billingCache.ReserveInflight(context.Background(), apiKey.User, apiKey.Group, nil, 1.0) + require.NoError(t, err) + defer held.HandlerDone() + + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/web_search", bytes.NewBufferString(`{"query":"sub2api"}`)) + c.Request.Header.Set("Content-Type", "application/json") + c.Set(string(middleware2.ContextKeyAPIKey), apiKey) + + h.WebSearch(c) + + require.NotEqual(t, http.StatusOK, w.Code) + require.Contains(t, w.Body.String(), "balance") + require.Equal(t, 1, cache.count(), "rejected request must not leave a reservation behind") +} diff --git a/backend/internal/handler/gateway_web_search.go b/backend/internal/handler/gateway_web_search.go index 741cc322f1..70fbc1b7e3 100644 --- a/backend/internal/handler/gateway_web_search.go +++ b/backend/internal/handler/gateway_web_search.go @@ -88,6 +88,18 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) { return } + // 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。 + inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, service.InflightEstimateRequest{Model: "grok-" + strings.ReplaceAll(searchLabel, "_", "-"), Kind: service.InflightEstimatePerRequest, SearchCalls: 1}) + if inflightErr != nil { + status, code, message, retryAfter := billingErrorDetails(inflightErr) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + c.JSON(status, gin.H{"error": gin.H{"type": code, "message": message}}) + return + } + defer inflightDone() + subject, _ := middleware2.GetAuthSubjectFromContext(c) reqLog := requestLogger(c, "handler.gateway.web_search") // Audit user search query before upstream Grok web_search traffic. diff --git a/backend/internal/handler/gemini_v1beta_handler.go b/backend/internal/handler/gemini_v1beta_handler.go index 0fc9e5eba5..4e56c37232 100644 --- a/backend/internal/handler/gemini_v1beta_handler.go +++ b/backend/internal/handler/gemini_v1beta_handler.go @@ -391,7 +391,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { } // 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。 - inflightRelease, err := reserveInflightBalance(c.Request.Context(), h.billingCacheService, h.gatewayService, apiKey, subscription, modelName, body) + inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(modelName, body)) if err != nil { reqLog.Info("gemini.inflight_reservation_rejected", zap.Error(err)) status, _, message, retryAfter := billingErrorDetails(err) diff --git a/backend/internal/handler/grok_audio.go b/backend/internal/handler/grok_audio.go index 28ae40a39b..7455485d7b 100644 --- a/backend/internal/handler/grok_audio.go +++ b/backend/internal/handler/grok_audio.go @@ -49,6 +49,18 @@ func (h *OpenAIGatewayHandler) GrokRealtime(c *gin.Context) { if strings.TrimSpace(model) == "" { model = "grok-voice-latest" } + // 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。 + inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, service.InflightEstimateRequest{Model: model, Kind: service.InflightEstimateAudio, AudioMode: "realtime", AudioUnits: 1}) + if inflightErr != nil { + status, code, message, retryAfter := billingErrorDetails(inflightErr) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + h.errorResponse(c, status, code, message) + return + } + defer inflightDone() + // Keep the HTTP response uncommitted while selecting and probing an account. // Realtime is not an HTTP streaming response; using reqStream=true here would // let the wait queue flush an SSE ping before the WebSocket handshake succeeds. @@ -212,6 +224,18 @@ func (h *OpenAIGatewayHandler) GrokVoice(c *gin.Context, endpoint string) { return } } + // 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。 + inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, grokVoiceInflightEstimate(endpoint, body)) + if inflightErr != nil { + status, code, message, retryAfter := billingErrorDetails(inflightErr) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + h.errorResponse(c, status, code, message) + return + } + defer inflightDone() + contentType := c.GetHeader("Content-Type") if strings.TrimSpace(contentType) == "" { contentType = "application/json" diff --git a/backend/internal/handler/grok_media.go b/backend/internal/handler/grok_media.go index b61cd9515b..03dabf7520 100644 --- a/backend/internal/handler/grok_media.go +++ b/backend/internal/handler/grok_media.go @@ -173,6 +173,18 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. return } + // 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。 + inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, grokMediaInflightEstimate(endpoint, routingModel, requestInfo, body)) + if inflightErr != nil { + status, code, message, retryAfter := billingErrorDetails(inflightErr) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + h.errorResponse(c, status, code, message) + return + } + defer inflightDone() + sessionSeed := body if len(sessionSeed) == 0 && strings.TrimSpace(requestID) != "" { sessionSeed = []byte(requestID) diff --git a/backend/internal/handler/openai_alpha_search.go b/backend/internal/handler/openai_alpha_search.go index 66c50be59a..3c25b66a08 100644 --- a/backend/internal/handler/openai_alpha_search.go +++ b/backend/internal/handler/openai_alpha_search.go @@ -109,6 +109,18 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) { return } + // 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。 + inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, service.InflightEstimateRequest{Model: requestedModel, BodyBytes: len(body), MaxTokens: requestMaxOutputTokens(body), Kind: service.InflightEstimateToken, SearchCalls: 1}) + if inflightErr != nil { + status, code, message, retryAfter := billingErrorDetails(inflightErr) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + h.errorResponse(c, status, code, message) + return + } + defer inflightDone() + searchID := strings.TrimSpace(gjson.GetBytes(body, "id").String()) sessionHash := h.gatewayService.GenerateSessionHashWithFallback(c, nil, searchID) profitVetoCount := 0 diff --git a/backend/internal/handler/openai_chat_completions.go b/backend/internal/handler/openai_chat_completions.go index afecc1b06a..0a371f12eb 100644 --- a/backend/internal/handler/openai_chat_completions.go +++ b/backend/internal/handler/openai_chat_completions.go @@ -147,7 +147,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { } // 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。 - inflightRelease, err := reserveInflightBalance(c.Request.Context(), h.billingCacheService, h.gatewayService, apiKey, subscription, reqModel, body) + inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body)) if err != nil { reqLog.Info("openai_chat_completions.inflight_reservation_rejected", zap.Error(err)) status, code, message, retryAfter := billingErrorDetails(err) diff --git a/backend/internal/handler/openai_embeddings.go b/backend/internal/handler/openai_embeddings.go index 8dda9faf3f..6c900c482e 100644 --- a/backend/internal/handler/openai_embeddings.go +++ b/backend/internal/handler/openai_embeddings.go @@ -108,6 +108,18 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) { return } + // 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。 + inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, service.InflightEstimateRequest{Model: reqModel, BodyBytes: len(body), MaxTokens: 1, Kind: service.InflightEstimateToken}) + if inflightErr != nil { + status, code, message, retryAfter := billingErrorDetails(inflightErr) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + h.errorResponse(c, status, code, message) + return + } + defer inflightDone() + profitVetoCount := 0 failedAccountIDs := make(map[int64]struct{}) var lastFailoverErr *service.UpstreamFailoverError diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 5278219eaa..dcc8a0952d 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -257,13 +257,21 @@ func usageRecordContext(parent context.Context, base context.Context) context.Co return base } -func wrapUsageRecordTaskContext(parent context.Context, task service.UsageRecordTask) service.UsageRecordTask { +// wrapUsageRecordTaskContext 包装计费任务:复制请求级 context 值,并接管请求的在途余额预留引用。 +// 返回的 abandon 在任务未被执行(被丢弃)时必须调用以归还预留引用;任务执行结束时(含 panic) +// 自动归还,此时余额缓存已在计费路径中同步扣减。 +func wrapUsageRecordTaskContext(parent context.Context, task service.UsageRecordTask) (service.UsageRecordTask, func()) { if task == nil { - return nil + return nil, func() {} + } + done := func() {} + if parent != nil { + done = service.InflightReservationFromContext(parent).Acquire() } return func(ctx context.Context) { + defer done() task(usageRecordContext(parent, ctx)) - } + }, done } func openAICompatibleRequestPlatform(ctx context.Context, apiKey *service.APIKey) string { @@ -607,7 +615,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { } // 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。 - inflightRelease, err := reserveInflightBalance(c.Request.Context(), h.billingCacheService, h.gatewayService, apiKey, subscription, reqModel, body) + inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body)) if err != nil { reqLog.Info("openai.inflight_reservation_rejected", zap.Error(err)) status, code, message, retryAfter := billingErrorDetails(err) @@ -1255,6 +1263,18 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { return } + // 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。 + inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body)) + if inflightErr != nil { + status, code, message, retryAfter := billingErrorDetails(inflightErr) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + h.anthropicStreamingAwareError(c, status, code, message, streamStarted) + return + } + defer inflightDone() + sessionHash := h.gatewayService.GenerateSessionHash(c, body) promptCacheKey := h.gatewayService.ExtractSessionID(c, body) sessionHash, promptCacheKey = resolveOpenAIMessagesMetadataSession(c, sessionHash, promptCacheKey, reqModel, body) @@ -2539,6 +2559,17 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { return } + // 余额模式在途预留(会话级):按首帧估算一次,会话期间续期;每轮计费任务接管引用, + // 会话结束且所有轮次扣减落地后释放。 + inflightCtx, inflightDone, inflightErr := reserveInflightBalanceCtx(ctx, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, firstMessage)) + if inflightErr != nil { + reqLog.Info("openai.websocket_inflight_reservation_rejected", zap.Error(inflightErr)) + closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, "billing check failed") + return + } + defer inflightDone() + ctx = inflightCtx + // 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 @@ -3279,9 +3310,12 @@ func (h *OpenAIGatewayHandler) submitUsageRecordTask(parent context.Context, tas if task == nil { return } - task = wrapUsageRecordTaskContext(parent, task) + task, abandon := wrapUsageRecordTaskContext(parent, task) if h.usageRecordWorkerPool != nil { if mode := h.usageRecordWorkerPool.Submit(task); mode != service.UsageRecordSubmitModeDroppedStopped { + if mode.Dropped() { + abandon() + } return } // 池已停止(进程关停窗口):计费任务不能静默丢失,降级为内联同步执行。 @@ -3318,7 +3352,7 @@ func (h *OpenAIGatewayHandler) submitMandatoryUsageRecordTask(parent context.Con if task == nil { return } - task = wrapUsageRecordTaskContext(parent, task) + task, _ = wrapUsageRecordTaskContext(parent, task) if h.usageRecordWorkerPool != nil { if mode := h.usageRecordWorkerPool.Submit(task); !mode.Dropped() { return diff --git a/backend/internal/handler/openai_images.go b/backend/internal/handler/openai_images.go index 9900861715..beda748cff 100644 --- a/backend/internal/handler/openai_images.go +++ b/backend/internal/handler/openai_images.go @@ -143,6 +143,18 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { return } + // 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。 + inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, service.InflightEstimateRequest{Model: routingModel, BodyBytes: len(body), Kind: service.InflightEstimateImage, Units: parsed.N}) + if inflightErr != nil { + status, code, message, retryAfter := billingErrorDetails(inflightErr) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + h.handleStreamingAwareError(c, status, code, message, streamStarted) + return + } + defer inflightDone() + sessionHash := h.gatewayService.GenerateExplicitSessionHash(c, body) requestCtx := service.WithOpenAIImagesEndpoint(service.WithOpenAIImageGenerationIntent(c.Request.Context())) diff --git a/backend/internal/repository/billing_inflight_cache.go b/backend/internal/repository/billing_inflight_cache.go index 47692fcf74..aba021d87c 100644 --- a/backend/internal/repository/billing_inflight_cache.go +++ b/backend/internal/repository/billing_inflight_cache.go @@ -56,6 +56,23 @@ var ( return {1, tostring(sum), count} `) + // KEYS: [1]=zset [2]=hash ARGV: [1]=member [2]=expire_at_ms [3]=key_ttl_ms + // 仅当成员仍存在时续期(XX),并把两个 key 的 TTL 至少延长到 key_ttl_ms。 + renewInflightBalanceScript = redis.NewScript(` + if not redis.call('ZSCORE', KEYS[1], ARGV[1]) then + return 0 + end + redis.call('ZADD', KEYS[1], 'XX', tonumber(ARGV[2]), ARGV[1]) + local ttl = tonumber(ARGV[3]) + if redis.call('PTTL', KEYS[1]) < ttl then + redis.call('PEXPIRE', KEYS[1], ttl) + end + if redis.call('PTTL', KEYS[2]) < ttl then + redis.call('PEXPIRE', KEYS[2], ttl) + end + return 1 + `) + releaseInflightBalanceScript = redis.NewScript(` redis.call('ZREM', KEYS[1], ARGV[1]) redis.call('HDEL', KEYS[2], ARGV[1]) @@ -99,4 +116,21 @@ func (c *billingCache) ReleaseInflightBalance(ctx context.Context, userID int64, return releaseInflightBalanceScript.Run(ctx, c.rdb, []string{zkey, hkey}, requestID).Err() } -var _ service.InflightBalanceReservationCache = (*billingCache)(nil) +// RenewInflightBalance 实现 service.InflightBalanceReservationRenewer。 +func (c *billingCache) RenewInflightBalance(ctx context.Context, userID int64, requestID string, ttl time.Duration) (bool, error) { + zkey, hkey := billingInflightKeys(userID) + ttlMs := ttl.Milliseconds() + if ttlMs <= 0 { + ttlMs = 1 + } + n, err := renewInflightBalanceScript.Run(ctx, c.rdb, []string{zkey, hkey}, requestID, time.Now().UnixMilli()+ttlMs, ttlMs).Int64() + if err != nil { + return false, err + } + return n == 1, nil +} + +var ( + _ service.InflightBalanceReservationCache = (*billingCache)(nil) + _ service.InflightBalanceReservationRenewer = (*billingCache)(nil) +) diff --git a/backend/internal/repository/billing_inflight_cache_test.go b/backend/internal/repository/billing_inflight_cache_test.go index 6540c83bdf..d9a338a943 100644 --- a/backend/internal/repository/billing_inflight_cache_test.go +++ b/backend/internal/repository/billing_inflight_cache_test.go @@ -209,3 +209,61 @@ func TestInflightReservation_ScriptErrorSurfaces(t *testing.T) { require.False(t, errors.Is(err, context.Canceled)) _ = rdb.Close() } + +// 回归:原实现在 handler 返回时即释放预留,而计费(RecordUsage → 余额缓存扣减)是异步的, +// 窗口内的顺序请求看到「在途=0 且余额未扣」全部放行(余额 $1、5×$0.9 → -4.40)。 +// 现在预留由计费任务在余额缓存扣减之后才释放。 +func TestInflightReservation_SequentialInBillingWindowAdmitsOnlyWhatBalanceCovers(t *testing.T) { + _, cache, svc := newInflightTestEnv(t, true, 60) + ctx := context.Background() + user := &service.User{ID: 77} + require.NoError(t, cache.SetUserBalance(ctx, user.ID, 1.0)) + + const cost = 0.9 + var pendingBilling []func() + admitted := 0 + for i := 0; i < 5; i++ { + res, err := svc.ReserveInflight(ctx, user, nil, nil, cost) + if err != nil { + require.ErrorIs(t, err, service.ErrInsufficientBalance) + continue + } + admitted++ + taskDone := res.Acquire() // handler submits the async billing task + res.HandlerDone() // handler returns before billing lands + pendingBilling = append(pendingBilling, func() { + require.NoError(t, cache.DeductUserBalance(ctx, user.ID, cost)) + taskDone() + }) + } + require.Equal(t, 1, admitted, "requests arriving before billing lands must not be admitted") + + for _, bill := range pendingBilling { + bill() + } + bal, err := cache.GetUserBalance(ctx, user.ID) + require.NoError(t, err) + require.InDelta(t, 0.1, bal, 1e-9) + require.Equal(t, int64(0), inflightCount(t, cache, user.ID)) +} + +func TestInflightReservation_RenewalKeepsStreamingReservationAlive(t *testing.T) { + _, cache, svc := newInflightTestEnv(t, true, 1) + ctx := context.Background() + user := &service.User{ID: 78} + require.NoError(t, cache.SetUserBalance(ctx, user.ID, 1.0)) + + res, err := svc.ReserveInflight(ctx, user, nil, nil, 0.8) + require.NoError(t, err) + time.Sleep(2500 * time.Millisecond) // streaming well past the 1s TTL + _, err = svc.ReserveInflight(ctx, user, nil, nil, 0.8) + require.ErrorIs(t, err, service.ErrInsufficientBalance, "renewed reservation must still count") + require.Equal(t, int64(1), inflightCount(t, cache, user.ID)) + + res.HandlerDone() + require.Equal(t, int64(0), inflightCount(t, cache, user.ID)) + + ok, err := cache.RenewInflightBalance(ctx, user.ID, "missing", time.Second) + require.NoError(t, err) + require.False(t, ok, "renew must not resurrect a released reservation") +} diff --git a/backend/internal/service/billing_inflight_reservation.go b/backend/internal/service/billing_inflight_reservation.go index 86b7e69a86..971bb472bd 100644 --- a/backend/internal/service/billing_inflight_reservation.go +++ b/backend/internal/service/billing_inflight_reservation.go @@ -4,10 +4,12 @@ import ( "context" "math" "sync" + "sync/atomic" "time" "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/Wei-Shaw/sub2api/internal/pkg/timezone" "github.com/google/uuid" ) @@ -22,14 +24,167 @@ type InflightBalanceReservationCache interface { ReleaseInflightBalance(ctx context.Context, userID int64, requestID string) error } +// InflightBalanceReservationRenewer 可选:续期仍存活的预留(流式长请求期间防止 TTL 过期)。 +type InflightBalanceReservationRenewer interface { + // RenewInflightBalance 若 requestID 仍存在则把其过期时间推迟到 now+ttl;返回是否仍存在。 + RenewInflightBalance(ctx context.Context, userID int64, requestID string, ttl time.Duration) (bool, error) +} + const ( defaultInflightReservationTTL = 15 * time.Minute defaultInflightDefaultMaxTokens = 8192 inflightReservationReleaseTimeout = 2 * time.Second inflightReservationReserveTimeout = 2 * time.Second + inflightReservationRenewTimeout = 2 * time.Second inflightInputBytesPerTokenEstimate = 4 + inflightUnpricedLogInterval = time.Minute ) +// InflightReservation 一次请求的在途预留句柄。 +// +// 生命周期(引用计数,归零时释放一次): +// - 创建时持有 1 个「handler」引用,并启动续期协程(每 ttl/3 续期一次)。 +// - handler 提交计费任务时通过 Acquire 再取一个引用,由计费任务在余额缓存 +// 实际扣减之后归还(任务被丢弃时由提交方立即归还)。 +// - handler 返回时调用 HandlerDone:停止续期并归还 handler 引用。 +// +// 因此预留会一直保持到「handler 结束 且 所有计费任务已完成扣减」, +// 从而消除「handler 已返回、异步计费尚未落地」窗口内的透支。handler 结束后 +// 不再续期,所以最长持有时间受 TTL 约束(计费任务卡死时预留自动过期)。 +// 所有方法对 nil 接收者安全。 +type InflightReservation struct { + cache InflightBalanceReservationCache + userID int64 + requestID string + amount float64 + ttl time.Duration + + refs atomic.Int64 + releaseOnce sync.Once + stopOnce sync.Once + closeOnce sync.Once + stopRenew chan struct{} + renewDone chan struct{} +} + +// Amount 预留金额。 +func (r *InflightReservation) Amount() float64 { + if r == nil { + return 0 + } + return r.amount +} + +// Acquire 为一个异步计费任务增加引用;返回的 done 幂等,必须在任务结束(或被丢弃)时调用。 +// 预留已释放时返回 no-op。 +func (r *InflightReservation) Acquire() func() { + if r == nil { + return noopRelease + } + for { + cur := r.refs.Load() + if cur <= 0 { + return noopRelease + } + if r.refs.CompareAndSwap(cur, cur+1) { + break + } + } + var once sync.Once + return func() { once.Do(r.decRef) } +} + +// HandlerDone handler 结束:停止续期并归还 handler 引用(幂等)。 +func (r *InflightReservation) HandlerDone() { + if r == nil { + return + } + r.stopOnce.Do(func() { + r.stopRenewal() + r.decRef() + }) +} + +// Release 立即释放预留(幂等),不论引用计数。 +func (r *InflightReservation) Release() { + if r == nil { + return + } + r.stopRenewal() + r.releaseOnce.Do(func() { + r.refs.Store(0) + relCtx, relCancel := context.WithTimeout(context.Background(), inflightReservationReleaseTimeout) + defer relCancel() + if err := r.cache.ReleaseInflightBalance(relCtx, r.userID, r.requestID); err != nil { + logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation release failed for user %d (expires by ttl): %v", r.userID, err) + } + }) +} + +func (r *InflightReservation) decRef() { + if r.refs.Add(-1) <= 0 { + r.Release() + } +} + +func (r *InflightReservation) stopRenewal() { + if r.stopRenew == nil { + return + } + r.closeOnce.Do(func() { close(r.stopRenew) }) + <-r.renewDone +} + +func (r *InflightReservation) startRenewal(renewer InflightBalanceReservationRenewer) { + interval := r.ttl / 3 + if interval < 10*time.Millisecond { + interval = 10 * time.Millisecond + } + r.stopRenew = make(chan struct{}) + r.renewDone = make(chan struct{}) + go func() { + defer close(r.renewDone) + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + select { + case <-r.stopRenew: + return + case <-ticker.C: + ctx, cancel := context.WithTimeout(context.Background(), inflightReservationRenewTimeout) + alive, err := renewer.RenewInflightBalance(ctx, r.userID, r.requestID, r.ttl) + cancel() + if err != nil { + logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation renew failed for user %d: %v", r.userID, err) + continue + } + if !alive { + return + } + } + } + }() +} + +type inflightReservationCtxKey struct{} + +// WithInflightReservation 把预留句柄挂到 context 上,供计费任务提交时交接。 +func WithInflightReservation(ctx context.Context, r *InflightReservation) context.Context { + if r == nil { + return ctx + } + return context.WithValue(ctx, inflightReservationCtxKey{}, r) +} + +// InflightReservationFromContext 读取 context 中的预留句柄(可能为 nil)。 +func InflightReservationFromContext(ctx context.Context) *InflightReservation { + if ctx == nil { + return nil + } + r, _ := ctx.Value(inflightReservationCtxKey{}).(*InflightReservation) + return r +} + func noopRelease() {} func (s *BillingCacheService) inflightReservationConfig() (config.InflightReservationConfig, bool) { @@ -49,37 +204,59 @@ func (s *BillingCacheService) InflightReservationEnabled() bool { return ok } -// ReserveInflightBalance 在余额模式下为本次请求登记在途预留。 +// InflightReservationFailClosedOnUnpriced 无法估算费用时是否拒绝请求(默认 false = fail-open)。 +func (s *BillingCacheService) InflightReservationFailClosedOnUnpriced() bool { + cfg, ok := s.inflightReservationConfig() + return ok && cfg.FailClosedOnUnpriced +} + +// ReserveInflightBalance 简化封装:返回释放函数,不续期、不做计费交接(预留最长存活 TTL)。 +func (s *BillingCacheService) ReserveInflightBalance(ctx context.Context, user *User, group *Group, subscription *UserSubscription, estimate float64) (func(), error) { + r, err := s.reserveInflight(ctx, user, group, subscription, estimate, false) + if err != nil { + return noopRelease, err + } + if r == nil { + return noopRelease, nil + } + return r.Release, nil +} + +// ReserveInflight 在余额模式下为本次请求登记在途预留。 // -// 必须在 CheckBillingEligibility 通过后调用。返回的 release 必须在请求结束时调用 -// (所有路径,建议 defer),多次调用安全。 +// 必须在 CheckBillingEligibility 通过后调用。返回的句柄可能为 nil(未预留); +// 调用方须在 handler 结束时调用 HandlerDone(nil 安全)。 // // 以下情况直接放行且不登记预留(fail-open,保持旧行为): // 开关关闭 / 简易模式 / 订阅模式 / estimate <= 0 / 缓存不支持 / 余额读取失败 / Redis 执行失败。 // 仅当 Redis 明确判定 缓存余额 - 在途合计 < estimate(且已有在途请求)时返回 ErrInsufficientBalance。 -func (s *BillingCacheService) ReserveInflightBalance(ctx context.Context, user *User, group *Group, subscription *UserSubscription, estimate float64) (func(), error) { +func (s *BillingCacheService) ReserveInflight(ctx context.Context, user *User, group *Group, subscription *UserSubscription, estimate float64) (*InflightReservation, error) { + return s.reserveInflight(ctx, user, group, subscription, estimate, true) +} + +func (s *BillingCacheService) reserveInflight(ctx context.Context, user *User, group *Group, subscription *UserSubscription, estimate float64, renew bool) (*InflightReservation, error) { cfg, ok := s.inflightReservationConfig() if !ok || user == nil { - return noopRelease, nil + return nil, nil } if group != nil && group.IsSubscriptionType() && subscription != nil { - return noopRelease, nil + return nil, nil } if cfg.MaxReservationUSD > 0 && estimate > cfg.MaxReservationUSD { estimate = cfg.MaxReservationUSD } if estimate <= 0 || math.IsNaN(estimate) || math.IsInf(estimate, 0) { - return noopRelease, nil + return nil, nil } rc, ok := s.cache.(InflightBalanceReservationCache) if !ok || rc == nil { - return noopRelease, nil + return nil, nil } balance, err := s.GetUserBalance(ctx, user.ID) if err != nil { logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation balance read failed for user %d (fail-open): %v", user.ID, err) - return noopRelease, nil + return nil, nil } ttl := defaultInflightReservationTTL @@ -92,41 +269,121 @@ func (s *BillingCacheService) ReserveInflightBalance(ctx context.Context, user * cancel() if err != nil { logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation failed for user %d (fail-open): %v", user.ID, err) - return noopRelease, nil + return nil, nil } if !allowed { logger.LegacyPrintf("service.billing_cache", "inflight reservation rejected: user=%d balance=%.6f inflight=%.6f estimate=%.6f", user.ID, balance, inflight, estimate) - return noopRelease, ErrInsufficientBalance + return nil, ErrInsufficientBalance } - var once sync.Once - userID := user.ID - return func() { - once.Do(func() { - relCtx, relCancel := context.WithTimeout(context.Background(), inflightReservationReleaseTimeout) - defer relCancel() - if err := rc.ReleaseInflightBalance(relCtx, userID, requestID); err != nil { - logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation release failed for user %d (expires by ttl): %v", userID, err) - } - }) - }, nil + r := &InflightReservation{cache: rc, userID: user.ID, requestID: requestID, amount: estimate, ttl: ttl} + r.refs.Store(1) + if renewer, ok := s.cache.(InflightBalanceReservationRenewer); renew && ok && renewer != nil { + r.startRenewal(renewer) + } + return r, nil } -// EstimateInflightReservationCost 估算单请求的保守费用(USD,已乘倍率)。 -// -// input_tokens = min(bodyBytes / 4, max_input_tokens) -// output_tokens = min(max_tokens 或 default_max_tokens, max_output_tokens) -// cost = (input_tokens × 输入单价 + output_tokens × 输出单价) × rateMultiplier -// -// 无法取得定价时返回 (0, false),调用方应 fail-open。 -func EstimateInflightReservationCost(billing *BillingService, cfg config.InflightReservationConfig, model string, bodyBytes, maxTokens int, rateMultiplier float64) (float64, bool) { - if billing == nil || model == "" || rateMultiplier <= 0 { - return 0, false +// ============================================ +// 费用估算 +// ============================================ + +// InflightEstimateKind 估算口径。 +type InflightEstimateKind int + +const ( + // InflightEstimateToken 文本/对话类:输入估算 + 输出上限(按次渠道定价时按次计)。 + InflightEstimateToken InflightEstimateKind = iota + // InflightEstimateImage 图片生成:按张计(取所有尺寸档最高单价)。 + InflightEstimateImage + // InflightEstimatePerRequest 按次类(独立搜索等):仅按渠道/分组按次价与搜索附加费,不做 token 估算。 + InflightEstimatePerRequest + // InflightEstimateVideo 视频生成:按秒 × 条数(分组/模型视频价)。 + InflightEstimateVideo + // InflightEstimateAudio 语音(tts/stt/realtime):按 AudioMode/AudioUnits 计。 + InflightEstimateAudio +) + +// InflightEstimateRequest 单请求估算输入。 +type InflightEstimateRequest struct { + Model string + BodyBytes int + MaxTokens int + Kind InflightEstimateKind + // Units 按次/按张数量(<=0 视为 1)。 + Units int + // SearchCalls 叠加的搜索次数(按分组 search_price_per_1k 计)。 + SearchCalls int + // 视频:分辨率与时长(秒)。 + VideoResolution string + VideoDurationSeconds int + // 音频:模式(tts/stt/realtime)与计量单位(百万字符/小时/分钟)。 + AudioMode string + AudioUnits float64 +} + +// inflightEstimateDeps 两种网关 service 共用的估算依赖。 +type inflightEstimateDeps struct { + cfg *config.Config + billing *BillingService + resolver *ModelPricingResolver + resolveMapping func(ctx context.Context, groupID int64, model string) ChannelMappingResult + userGroupRate func(ctx context.Context, userID, groupID int64, groupDefault float64) float64 +} + +var inflightUnpricedLastLog atomic.Int64 + +func logInflightUnpriced(model string, groupID *int64) { + now := time.Now().UnixNano() + last := inflightUnpricedLastLog.Load() + if now-last < int64(inflightUnpricedLogInterval) || !inflightUnpricedLastLog.CompareAndSwap(last, now) { + return } - pricing, err := billing.GetModelPricing(model) - if err != nil || pricing == nil { - return 0, false + var gid int64 + if groupID != nil { + gid = *groupID } + logger.LegacyPrintf("service.billing_cache", "Warning: inflight reservation cannot price model=%q group=%d; request admitted without reservation (fail-open, throttled log)", model, gid) +} + +// inflightBillingModelCandidates 与计费路径一致地挑选计费模型: +// channel_mapped(默认)→ 映射后模型;requested → 请求模型; +// upstream / response_model 在准入时未知 → 取请求模型与映射模型两者较高估算。 +func inflightBillingModelCandidates(ctx context.Context, deps inflightEstimateDeps, apiKey *APIKey, model string) []string { + if apiKey == nil || apiKey.GroupID == nil || deps.resolveMapping == nil { + return []string{model} + } + m := deps.resolveMapping(ctx, *apiKey.GroupID, model) + mapped := m.MappedModel + if mapped == "" || mapped == model { + return []string{model} + } + switch m.BillingModelSource { + case BillingModelSourceRequested: + return []string{model} + case BillingModelSourceChannelMapped, "": + // 计费侧对无价映射模型会回落到具体转发模型;此处同时估算请求模型,取较高者。 + return []string{mapped, model} + default: + return []string{mapped, model} + } +} + +func (d inflightEstimateDeps) rates(ctx context.Context, apiKey *APIKey) (text, image float64) { + rate := 1.0 + if d.cfg != nil && d.cfg.Default.RateMultiplier > 0 { + rate = d.cfg.Default.RateMultiplier + } + if apiKey != nil && apiKey.GroupID != nil && apiKey.Group != nil { + rate = apiKey.Group.RateMultiplier + if d.userGroupRate != nil && apiKey.User != nil { + rate = d.userGroupRate(ctx, apiKey.User.ID, *apiKey.GroupID, rate) + } + } + return computePeakAwareMultipliers(apiKey, rate, timezone.Now()) +} + +func tokenCounts(cfg config.InflightReservationConfig, bodyBytes, maxTokens int) (int, int) { inputTokens := 0 if bodyBytes > 0 { inputTokens = bodyBytes / inflightInputBytesPerTokenEstimate @@ -144,8 +401,165 @@ func EstimateInflightReservationCost(billing *BillingService, cfg config.Infligh if cfg.MaxOutputTokens > 0 && outputTokens > cfg.MaxOutputTokens { outputTokens = cfg.MaxOutputTokens } + return inputTokens, outputTokens +} + +func maxPerRequestPrice(resolved *ResolvedPricing) float64 { + if resolved == nil { + return 0 + } + p := resolved.DefaultPerRequestPrice + for _, tier := range resolved.RequestTiers { + if tier.PerRequestPrice != nil && *tier.PerRequestPrice > p { + p = *tier.PerRequestPrice + } + } + return p +} + +func validCost(v float64) bool { return v > 0 && !math.IsNaN(v) && !math.IsInf(v, 0) } + +// estimateOne 估算单个候选计费模型(未乘倍率的 token 部分与按次部分分开返回,便于套用不同倍率)。 +func (d inflightEstimateDeps) estimateOne(ctx context.Context, apiKey *APIKey, model string, req InflightEstimateRequest, textRate, imageRate float64) float64 { + cfg := inflightReservationCfg(d.cfg) + units := req.Units + if units <= 0 { + units = 1 + } + var resolved *ResolvedPricing + if d.resolver != nil { + in := PricingInput{Model: model} + if apiKey != nil { + in.GroupID = apiKey.GroupID + in.Group = apiKey.Group + } + resolved = d.resolver.Resolve(ctx, in) + } + + inputTokens, outputTokens := tokenCounts(cfg, req.BodyBytes, req.MaxTokens) + tokenCost := func() float64 { + var pricing *ModelPricing + if resolved != nil && (resolved.Mode == BillingModeToken || resolved.Mode == "") && d.resolver != nil { + pricing = d.resolver.GetIntervalPricing(resolved, inputTokens) + } + if pricing == nil && d.billing != nil { + pricing, _ = d.billing.GetModelPricing(model) + } + if pricing == nil { + return 0 + } + return (float64(inputTokens)*pricing.InputPricePerToken + float64(outputTokens)*pricing.OutputPricePerToken) * textRate + } + + var cost float64 + perRequestMode := resolved != nil && (resolved.Mode == BillingModePerRequest || resolved.Mode == BillingModeImage || resolved.Mode == BillingModeVideo) + switch req.Kind { + case InflightEstimateImage: + if perRequestMode { + cost = maxPerRequestPrice(resolved) * float64(units) * imageRate + } + if d.billing != nil { + cfgImg := imagePriceConfigFromAPIKey(apiKey) + for _, tier := range []string{ImageBillingSize1K, ImageBillingSize2K, ImageBillingSize4K} { + if b := d.billing.CalculateImageCost(model, tier, units, cfgImg, imageRate); b != nil && b.ActualCost > cost { + cost = b.ActualCost + } + } + } + if cost <= 0 { + cost = tokenCost() + } + case InflightEstimateVideo: + if perRequestMode { + cost = maxPerRequestPrice(resolved) * float64(units) * math.Max(textRate, imageRate) + } + if d.billing != nil { + if b := d.billing.CalculateVideoCost(model, req.VideoResolution, units, req.VideoDurationSeconds, videoPriceConfigFromAPIKey(apiKey), math.Max(textRate, imageRate)); b != nil && b.ActualCost > cost { + cost = b.ActualCost + } + } + case InflightEstimateAudio: + if perRequestMode && resolved.Mode == BillingModePerRequest { + u := req.AudioUnits + if u <= 0 { + u = 1 + } + cost = maxPerRequestPrice(resolved) * u * textRate + } + if d.billing != nil && req.AudioUnits > 0 { + if b := d.billing.CalculateAudioCost(req.AudioMode, req.AudioUnits, groupAudioPriceConfigFromAPIKey(apiKey), textRate); b != nil && b.ActualCost > cost { + cost = b.ActualCost + } + } + default: + if perRequestMode { + rate := textRate + if resolved.Mode == BillingModeImage { + rate = imageRate + } + cost = maxPerRequestPrice(resolved) * float64(units) * rate + } else if req.Kind != InflightEstimatePerRequest { + cost = tokenCost() + } + } + if req.SearchCalls > 0 { + if d.billing != nil { + if b := d.billing.CalculateSearchCost(req.SearchCalls, groupSearchPricePer1kFromAPIKey(apiKey), textRate); b != nil && b.ActualCost > 0 { + cost += b.ActualCost + } + } + } + if !validCost(cost) { + return 0 + } + return cost +} + +// estimate 返回保守的单请求费用(USD,已乘倍率);无法定价返回 (0,false)。 +func (d inflightEstimateDeps) estimate(ctx context.Context, apiKey *APIKey, req InflightEstimateRequest) (float64, bool) { + if apiKey == nil || apiKey.User == nil { + return 0, false + } + if req.Model == "" || (req.Kind == InflightEstimateAudio && req.AudioUnits <= 0) { + // 非计量请求(媒体状态查询、custom-voices 等):无需预留,也不算「无法定价」。 + return 0, true + } + textRate, imageRate := d.rates(ctx, apiKey) + if textRate <= 0 && imageRate <= 0 { + // 免费分组:不计费,也无需预留。 + return 0, true + } + best := 0.0 + for _, m := range inflightBillingModelCandidates(ctx, d, apiKey, req.Model) { + if c := d.estimateOne(ctx, apiKey, m, req, textRate, imageRate); c > best { + best = c + } + } + if best <= 0 { + logInflightUnpriced(req.Model, apiKey.GroupID) + return 0, false + } + return best, true +} + +// EstimateInflightReservationCost 仅用基础定价的简化估算(保留给无 resolver 的调用方/测试)。 +// +// input_tokens = min(bodyBytes / 4, max_input_tokens) +// output_tokens = min(max_tokens 或 default_max_tokens, max_output_tokens) +// cost = (input_tokens × 输入单价 + output_tokens × 输出单价) × rateMultiplier +// +// 无法取得定价时返回 (0, false),调用方应 fail-open。 +func EstimateInflightReservationCost(billing *BillingService, cfg config.InflightReservationConfig, model string, bodyBytes, maxTokens int, rateMultiplier float64) (float64, bool) { + if billing == nil || model == "" || rateMultiplier <= 0 { + return 0, false + } + pricing, err := billing.GetModelPricing(model) + if err != nil || pricing == nil { + return 0, false + } + inputTokens, outputTokens := tokenCounts(cfg, bodyBytes, maxTokens) cost := (float64(inputTokens)*pricing.InputPricePerToken + float64(outputTokens)*pricing.OutputPricePerToken) * rateMultiplier - if cost <= 0 || math.IsNaN(cost) || math.IsInf(cost, 0) { + if !validCost(cost) { return 0, false } return cost, true @@ -158,41 +572,37 @@ func inflightReservationCfg(cfg *config.Config) config.InflightReservationConfig return cfg.Billing.InflightReservation } -func groupRateMultiplierOrOne(group *Group) float64 { - if group == nil { - return 1 +func (s *GatewayService) inflightEstimateDeps() inflightEstimateDeps { + d := inflightEstimateDeps{cfg: s.cfg, billing: s.billingService, resolver: s.resolver} + if s.channelService != nil { + d.resolveMapping = s.channelService.ResolveChannelMapping } - return group.RateMultiplier + d.userGroupRate = s.getUserGroupRateMultiplier + return d } -// EstimateInflightReservation 使用网关计费倍率估算在途预留金额;0 表示不预留。 -func (s *GatewayService) EstimateInflightReservation(ctx context.Context, apiKey *APIKey, model string, bodyBytes, maxTokens int) float64 { - if s == nil || apiKey == nil || apiKey.User == nil { - return 0 +// EstimateInflightReservation 与计费路径同口径(计费模型 / ModelPricingResolver / 倍率)估算在途预留金额。 +// 第二个返回值为 false 表示无法定价(调用方 fail-open 或按配置 fail-closed)。 +func (s *GatewayService) EstimateInflightReservation(ctx context.Context, apiKey *APIKey, req InflightEstimateRequest) (float64, bool) { + if s == nil { + return 0, false } - rate := groupRateMultiplierOrOne(apiKey.Group) - if apiKey.GroupID != nil && apiKey.Group != nil { - rate = s.getUserGroupRateMultiplier(ctx, apiKey.User.ID, *apiKey.GroupID, rate) - } - cost, ok := EstimateInflightReservationCost(s.billingService, inflightReservationCfg(s.cfg), model, bodyBytes, maxTokens, rate) - if !ok { - return 0 - } - return cost + return s.inflightEstimateDeps().estimate(ctx, apiKey, req) } -// EstimateInflightReservation 使用 OpenAI 网关计费倍率估算在途预留金额;0 表示不预留。 -func (s *OpenAIGatewayService) EstimateInflightReservation(ctx context.Context, apiKey *APIKey, model string, bodyBytes, maxTokens int) float64 { - if s == nil || apiKey == nil || apiKey.User == nil { - return 0 +func (s *OpenAIGatewayService) inflightEstimateDeps() inflightEstimateDeps { + d := inflightEstimateDeps{cfg: s.cfg, billing: s.billingService, resolver: s.resolver} + if s.channelService != nil { + d.resolveMapping = s.channelService.ResolveChannelMapping } - rate := groupRateMultiplierOrOne(apiKey.Group) - if apiKey.GroupID != nil && apiKey.Group != nil && s.userGroupRateResolver != nil { - rate = s.userGroupRateResolver.Resolve(ctx, apiKey.User.ID, *apiKey.GroupID, rate) + d.userGroupRate = s.ResolveUserGroupRateMultiplier + return d +} + +// EstimateInflightReservation 同 GatewayService.EstimateInflightReservation(OpenAI 网关倍率口径)。 +func (s *OpenAIGatewayService) EstimateInflightReservation(ctx context.Context, apiKey *APIKey, req InflightEstimateRequest) (float64, bool) { + if s == nil { + return 0, false } - cost, ok := EstimateInflightReservationCost(s.billingService, inflightReservationCfg(s.cfg), model, bodyBytes, maxTokens, rate) - if !ok { - return 0 - } - return cost + return s.inflightEstimateDeps().estimate(ctx, apiKey, req) } diff --git a/backend/internal/service/billing_inflight_reservation_test.go b/backend/internal/service/billing_inflight_reservation_test.go index f1dff41ec8..bef02bc18a 100644 --- a/backend/internal/service/billing_inflight_reservation_test.go +++ b/backend/internal/service/billing_inflight_reservation_test.go @@ -4,7 +4,10 @@ package service import ( "context" + "sync" + "sync/atomic" "testing" + "time" "github.com/Wei-Shaw/sub2api/internal/config" "github.com/stretchr/testify/require" @@ -58,3 +61,263 @@ func TestReserveInflightBalance_SimpleModeDisabled(t *testing.T) { t.Cleanup(svc.Stop) require.False(t, svc.InflightReservationEnabled()) } + +// --------------------------------------------------------------------------- +// estimate: same model / resolver as billing +// --------------------------------------------------------------------------- + +func newInflightEstimateGateway(t *testing.T, channelService *ChannelService) *GatewayService { + t.Helper() + cfg := &config.Config{} + cfg.Billing.InflightReservation = config.InflightReservationConfig{Enabled: true, DefaultMaxTokens: 1000, MaxInputTokens: 200000, MaxOutputTokens: 128000} + billing := NewBillingService(cfg, nil) + return &GatewayService{ + cfg: cfg, + billingService: billing, + resolver: NewModelPricingResolver(channelService, billing), + channelService: channelService, + } +} + +func TestInflightEstimate_ChannelAliasUsesMappedModel(t *testing.T) { + groupID := int64(10) + ch := Channel{ + ID: 1, + Status: StatusActive, + GroupIDs: []int64{groupID}, + ModelMapping: map[string]map[string]string{ + "anthropic": {"my-alias": "claude-sonnet-4-5"}, + }, + } + cs := newTestChannelService(makeStandardRepo(ch, map[int64]string{groupID: "anthropic"})) + svc := newInflightEstimateGateway(t, cs) + apiKey := &APIKey{User: &User{ID: 1}, GroupID: &groupID, Group: &Group{ID: groupID, Platform: PlatformAnthropic, RateMultiplier: 1}} + + // Old estimate (client model, base pricing only) → 0 → reservation bypassed. + _, ok := EstimateInflightReservationCost(svc.billingService, svc.cfg.Billing.InflightReservation, "my-alias", 4000, 1000, 1) + require.False(t, ok, "precondition: alias has no base pricing") + + est, priced := svc.EstimateInflightReservation(context.Background(), apiKey, InflightEstimateRequest{Model: "my-alias", BodyBytes: 4000, MaxTokens: 1000}) + require.True(t, priced) + require.Greater(t, est, 0.0) + + direct, _ := svc.EstimateInflightReservation(context.Background(), apiKey, InflightEstimateRequest{Model: "claude-sonnet-4-5", BodyBytes: 4000, MaxTokens: 1000}) + require.InDelta(t, direct, est, 1e-12, "alias must be estimated as the channel-mapped billing model") +} + +func TestInflightEstimate_GroupPerRequestPricing(t *testing.T) { + groupID := int64(20) + price := 0.5 + group := &Group{ID: groupID, Platform: PlatformAnthropic, RateMultiplier: 2, ModelPricing: []ChannelModelPricing{ + {Models: []string{"custom-per-request"}, BillingMode: BillingModePerRequest, PerRequestPrice: &price}, + }} + svc := newInflightEstimateGateway(t, nil) + apiKey := &APIKey{User: &User{ID: 1}, GroupID: &groupID, Group: group} + + est, priced := svc.EstimateInflightReservation(context.Background(), apiKey, InflightEstimateRequest{Model: "custom-per-request", BodyBytes: 100}) + require.True(t, priced) + require.InDelta(t, 1.0, est, 1e-12, "per_request 0.5 × group rate 2") + + // Wildcard group token pricing is honored as well. + in, out := 1e-6, 2e-6 + group.ModelPricing = []ChannelModelPricing{{Models: []string{"wild-*"}, BillingMode: BillingModeToken, InputPrice: &in, OutputPrice: &out}} + est, priced = svc.EstimateInflightReservation(context.Background(), apiKey, InflightEstimateRequest{Model: "wild-x", BodyBytes: 4000, MaxTokens: 500}) + require.True(t, priced) + require.InDelta(t, (1000*in+500*out)*2, est, 1e-12) +} + +func TestInflightEstimate_UnpricedIsReportedAndNonMeteredIsNot(t *testing.T) { + svc := newInflightEstimateGateway(t, nil) + apiKey := &APIKey{User: &User{ID: 1}} + est, priced := svc.EstimateInflightReservation(context.Background(), apiKey, InflightEstimateRequest{Model: "totally-unknown-model-xyz"}) + require.False(t, priced) + require.Zero(t, est) + + est, priced = svc.EstimateInflightReservation(context.Background(), apiKey, InflightEstimateRequest{}) + require.True(t, priced, "non-metered requests (e.g. media status lookups) are not 'unpriced'") + require.Zero(t, est) +} + +func TestInflightEstimate_MediaKinds(t *testing.T) { + groupID := int64(30) + p4k := 0.3 + search := 1000.0 + group := &Group{ID: groupID, Platform: PlatformOpenAI, RateMultiplier: 1, ImagePrice4K: &p4k, SearchPricePer1k: &search} + svc := newInflightEstimateGateway(t, nil) + apiKey := &APIKey{User: &User{ID: 1}, GroupID: &groupID, Group: group} + + est, priced := svc.EstimateInflightReservation(context.Background(), apiKey, InflightEstimateRequest{Model: "gpt-image-1", Kind: InflightEstimateImage, Units: 2}) + require.True(t, priced) + require.GreaterOrEqual(t, est, 0.6, "image estimate uses the highest size tier × n") + + est, priced = svc.EstimateInflightReservation(context.Background(), apiKey, InflightEstimateRequest{Model: "grok-web-search", Kind: InflightEstimatePerRequest, SearchCalls: 1}) + require.True(t, priced) + require.InDelta(t, 1.0, est, 1e-12) + + est, priced = svc.EstimateInflightReservation(context.Background(), apiKey, InflightEstimateRequest{Model: "realtime", Kind: InflightEstimateAudio, AudioMode: "realtime", AudioUnits: 1}) + require.True(t, priced) + require.Greater(t, est, 0.0) +} + +// --------------------------------------------------------------------------- +// reservation handle: hand-off to billing task, renewal +// --------------------------------------------------------------------------- + +type memInflightCache struct { + BillingCache + mu sync.Mutex + balance float64 + res map[string]float64 + exp map[string]time.Time + renews atomic.Int32 + deducts atomic.Int32 + releaseCt atomic.Int32 +} + +func newMemInflightCache(balance float64) *memInflightCache { + return &memInflightCache{balance: balance, res: map[string]float64{}, exp: map[string]time.Time{}} +} + +func (m *memInflightCache) GetUserBalance(context.Context, int64) (float64, error) { + m.mu.Lock() + defer m.mu.Unlock() + return m.balance, nil +} + +func (m *memInflightCache) DeductUserBalance(_ context.Context, _ int64, amount float64) error { + m.mu.Lock() + defer m.mu.Unlock() + m.balance -= amount + m.deducts.Add(1) + return nil +} + +func (m *memInflightCache) gc() { + now := time.Now() + for id, e := range m.exp { + if !e.After(now) { + delete(m.exp, id) + delete(m.res, id) + } + } +} + +func (m *memInflightCache) ReserveInflightBalance(_ context.Context, _ int64, id string, amount, balance float64, ttl time.Duration) (bool, float64, error) { + m.mu.Lock() + defer m.mu.Unlock() + m.gc() + sum := 0.0 + for _, v := range m.res { + sum += v + } + if len(m.res) > 0 && balance-sum < amount { + return false, sum, nil + } + m.res[id] = amount + m.exp[id] = time.Now().Add(ttl) + return true, sum, nil +} + +func (m *memInflightCache) ReleaseInflightBalance(_ context.Context, _ int64, id string) error { + m.mu.Lock() + defer m.mu.Unlock() + delete(m.res, id) + delete(m.exp, id) + m.releaseCt.Add(1) + return nil +} + +func (m *memInflightCache) RenewInflightBalance(_ context.Context, _ int64, id string, ttl time.Duration) (bool, error) { + m.mu.Lock() + defer m.mu.Unlock() + m.renews.Add(1) + if _, ok := m.res[id]; !ok { + return false, nil + } + m.exp[id] = time.Now().Add(ttl) + return true, nil +} + +func (m *memInflightCache) count() int { + m.mu.Lock() + defer m.mu.Unlock() + m.gc() + return len(m.res) +} + +func newInflightSvc(t *testing.T, cache BillingCache, ttlSeconds int) *BillingCacheService { + t.Helper() + cfg := &config.Config{} + cfg.Billing.InflightReservation = config.InflightReservationConfig{Enabled: true, TTLSeconds: ttlSeconds} + svc := NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(svc.Stop) + return svc +} + +func TestInflightReservation_HeldUntilBillingTaskDone(t *testing.T) { + cache := newMemInflightCache(1) + svc := newInflightSvc(t, cache, 60) + user := &User{ID: 1} + + res, err := svc.ReserveInflight(context.Background(), user, nil, nil, 0.9) + require.NoError(t, err) + require.NotNil(t, res) + taskDone := res.Acquire() // billing task submitted + res.HandlerDone() // handler returned + require.Equal(t, 1, cache.count(), "reservation must survive handler return while billing is pending") + + _, err = svc.ReserveInflight(context.Background(), user, nil, nil, 0.9) + require.ErrorIs(t, err, ErrInsufficientBalance) + + taskDone() + taskDone() // idempotent + require.Equal(t, 0, cache.count()) + require.Equal(t, int32(1), cache.releaseCt.Load(), "released exactly once") + + // No billing task → HandlerDone releases immediately; Acquire after release is a no-op. + res2, err := svc.ReserveInflight(context.Background(), user, nil, nil, 0.1) + require.NoError(t, err) + res2.HandlerDone() + res2.HandlerDone() + require.Equal(t, 0, cache.count()) + res2.Acquire()() + require.Equal(t, int32(2), cache.releaseCt.Load()) + + // Nil handle is safe. + var nilRes *InflightReservation + nilRes.Acquire()() + nilRes.HandlerDone() + nilRes.Release() +} + +func TestInflightReservation_RenewedWhileHandlerActive(t *testing.T) { + cache := newMemInflightCache(1) + svc := newInflightSvc(t, cache, 1) + user := &User{ID: 2} + + res, err := svc.ReserveInflight(context.Background(), user, nil, nil, 0.9) + require.NoError(t, err) + time.Sleep(2500 * time.Millisecond) // > 2 × TTL + require.Equal(t, 1, cache.count(), "renewal keeps the reservation alive past its TTL") + require.Greater(t, cache.renews.Load(), int32(3)) + _, err = svc.ReserveInflight(context.Background(), user, nil, nil, 0.9) + require.ErrorIs(t, err, ErrInsufficientBalance) + + done := res.Acquire() + res.HandlerDone() + renewsAtStop := cache.renews.Load() + time.Sleep(1200 * time.Millisecond) + require.Equal(t, renewsAtStop, cache.renews.Load(), "no renewal after handler end") + require.Equal(t, 0, cache.count(), "post-handler hold is bounded by TTL") + done() +} + +func TestSyncBalanceCacheAfterDeduction_SynchronousWhenInflightEnabled(t *testing.T) { + cache := newMemInflightCache(1) + svc := newInflightSvc(t, cache, 60) + p := &postUsageBillingParams{Cost: &CostBreakdown{ActualCost: 0.25}, User: &User{ID: 3}} + syncBalanceCacheAfterDeduction(context.Background(), p, &billingDeps{billingCacheService: svc}, nil) + require.Equal(t, int32(1), cache.deducts.Load(), "cache deduction must land before the billing task returns") + bal, _ := cache.GetUserBalance(context.Background(), 3) + require.InDelta(t, 0.75, bal, 1e-12) +} diff --git a/backend/internal/service/gateway_usage_billing.go b/backend/internal/service/gateway_usage_billing.go index 64b4e05fbe..3a612d17d0 100644 --- a/backend/internal/service/gateway_usage_billing.go +++ b/backend/internal/service/gateway_usage_billing.go @@ -477,6 +477,16 @@ func syncBalanceCacheAfterDeduction(ctx context.Context, p *postUsageBillingPara } return } + if deps.billingCacheService.InflightReservationEnabled() { + // 在途预留开启时同步扣减余额缓存:计费任务结束后才会释放预留, + // 必须保证此时准入读取的缓存余额已反映本次扣费,否则释放与扣减之间 + // 仍存在「在途=0 且余额未扣」的窗口。本函数运行在计费 worker 中,不在请求热路径。 + err := deps.billingCacheService.DeductBalanceCache(ctx, p.User.ID, p.Cost.ActualCost) + if err == nil { + return + } + logger.LegacyPrintf("service.gateway", "Warning: sync deduct balance cache failed for user %d, falling back to queue: %v", p.User.ID, err) + } deps.billingCacheService.QueueDeductBalance(p.User.ID, p.Cost.ActualCost) } diff --git a/deploy/config.example.yaml b/deploy/config.example.yaml index de77ce6ad8..a49c400f73 100644 --- a/deploy/config.example.yaml +++ b/deploy/config.example.yaml @@ -1135,8 +1135,11 @@ billing: # 用户无在途请求时始终放行(与旧行为一致);估算失败/Redis 不可用时 fail-open。 inflight_reservation: enabled: true - # Max lifetime of a single reservation (self-heals leaked reservations) - # 单条预留最长存活时间(秒),泄漏的预留到期自动释放 + # Reservation TTL. Renewed every ttl/3 while the request is running; after the + # handler returns it is held (not renewed) until the async billing task has + # deducted the cached balance, so a stuck billing task frees it after at most one TTL. + # 预留 TTL(秒)。请求进行中每 ttl/3 续期;请求结束后不再续期,保持到异步计费 + # 扣减余额缓存后释放,计费任务卡死时最多一个 TTL 后自动失效 ttl_seconds: 900 # Output tokens assumed when the request has no max_tokens # 请求未携带 max_tokens 时假定的输出 token 数 @@ -1148,6 +1151,10 @@ billing: # Cap for a single reservation in USD (0 = no cap) # 单请求预留金额上限(USD,0 为不设上限) max_reservation_usd: 0 + # When a request cannot be priced (no base / group / channel pricing for the + # billing model), reject it instead of admitting it without reservation. + # 无法为请求定价(计费模型无基础/分组/渠道定价)时拒绝请求;默认 false 为放行不预留(节流告警日志) + fail_closed_on_unpriced: false # ============================================================================= # Turnstile Configuration