diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index 1cd33c4f6305..04412de85138 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -373,7 +373,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { upstreamBillingProbeService := service.ProvideUpstreamBillingProbeService(accountRepository, accountTestService, settingService, leaderLockCache, db, configConfig, newAPIAuthorizationRepository, secretEncryptor) openCodeGoUsageService := service.ProvideOpenCodeGoUsageService(accountRepository, httpUpstream, settingService, leaderLockCache, db, configConfig) adminHandlers := handler.ProvideAdminHandlers(requestCaptureHandler, dashboardHandler, adminUserHandler, groupHandler, accountHandler, openAIOAuthReauthService, adminAnnouncementHandler, dataManagementHandler, backupHandler, oAuthHandler, openAIOAuthHandler, openAIOAuthReauthHandler, geminiOAuthHandler, antigravityOAuthHandler, grokOAuthHandler, cnProviderHandler, proxyHandler, adminRedeemHandler, promoHandler, settingHandler, opsHandler, systemHandler, adminSubscriptionHandler, adminUsageHandler, userAttributeHandler, errorPassthroughHandler, tlsFingerprintProfileHandler, pluginHandler, adminAPIKeyHandler, scheduledTestHandler, pelicanGroupTestHandler, controlledExperimentHandler, accountOpsHandler, accountTokenGuardHandler, accountTokenGuardV2Handler, channelHandler, channelMonitorHandler, channelMonitorRequestTemplateHandler, contentModerationHandler, promptAdminHandler, paymentHandler, affiliateHandler, complianceHandler, auditLogHandler, upstreamBillingProbeService, ollamaCloudUsageService, openCodeGoUsageService, settingService, codexHarvestService, openAIGatewayService) - usageRecordWorkerPool := service.NewUsageRecordWorkerPool(configConfig) + usageRecordWorkerPool := service.ProvideUsageRecordWorkerPool(configConfig, usageBillingRepository, billingCacheService, apiKeyAuthCacheInvalidator) userMsgQueueCache := repository.NewUserMsgQueueCache(redisClient) userMessageQueueService := service.ProvideUserMessageQueueService(userMsgQueueCache, rpmCache, configConfig) legacyEngine := securityaudit.NewLegacyModerationAdapter(contentModerationService) diff --git a/backend/internal/handler/gateway_billing_delivery_http_test.go b/backend/internal/handler/gateway_billing_delivery_http_test.go new file mode 100644 index 000000000000..0744ba36b231 --- /dev/null +++ b/backend/internal/handler/gateway_billing_delivery_http_test.go @@ -0,0 +1,82 @@ +//go:build unit + +package handler + +import ( + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type httpDeliveryBilling struct { + service.UsageBillingRepository + event *service.UsageBillingEvent + failure error + stageFailure error + applied bool +} + +func (r *httpDeliveryBilling) StageUsageBilling(_ context.Context, c *service.UsageBillingCommand, u *service.UsageLog) error { + if r.stageFailure != nil { + return r.stageFailure + } + r.event = &service.UsageBillingEvent{Version: 1, Command: *c, Usage: *u} + return nil +} +func (r *httpDeliveryBilling) ApplyStagedUsageBilling(context.Context, service.UsageBillingEventKey) (*service.UsageBillingRecovery, error) { + if r.failure != nil { + return nil, r.failure + } + fresh := !r.applied + r.applied = true + return &service.UsageBillingRecovery{Event: *r.event, Result: service.UsageBillingApplyResult{Applied: fresh}, NeedsCache: true}, nil +} +func (*httpDeliveryBilling) PendingUsageBilling(context.Context, int) ([]service.UsageBillingEventKey, error) { + return nil, nil +} +func (*httpDeliveryBilling) AcknowledgeUsageBillingCache(context.Context, service.UsageBillingEventKey) error { + return nil +} + +// Real HTTP transport + production protocol adapter + GatewayHandler. The SQL +// crash/atomicity half is exercised separately with PostgreSQL integration tests. +func TestGatewayHTTPAdapterStagesUsageBeforeExecutionAndNeverRecallsModel(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, endpoint := range []string{"responses", "chat/completions"} { + t.Run(endpoint, func(t *testing.T) { + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, compatStreamUsagePartialStream+compatStreamUsageCompletion) + })) + defer upstream.Close() + billing := &httpDeliveryBilling{failure: errors.New("injected execution failure after durable stage")} + run := runCompatStreamUsageHandlerWithDelivery(t, endpoint, "", true, billing, upstream.URL) + require.Equal(t, int32(1), calls.Load()) + require.Contains(t, run.rec.Body.String(), "partial") + require.NotNil(t, billing.event, "handler must stage before returning, not queue an in-memory closure") + require.Equal(t, 12, billing.event.Usage.InputTokens) + require.Equal(t, 7, billing.event.Usage.OutputTokens) + require.Greater(t, billing.event.Usage.ActualCost, 0.0) + require.Empty(t, run.usageRepo.created, "failure is not a fake zero-cost success") + billing.failure = nil + key := service.UsageBillingEventKey{RequestID: billing.event.Command.RequestID, APIKeyID: billing.event.Command.APIKeyID} + first, err := billing.ApplyStagedUsageBilling(context.Background(), key) + require.NoError(t, err) + require.True(t, first.Result.Applied) + second, err := billing.ApplyStagedUsageBilling(context.Background(), key) + require.NoError(t, err) + require.False(t, second.Result.Applied) + require.Equal(t, int32(1), calls.Load(), "billing recovery cannot invoke a model transport") + }) + } +} diff --git a/backend/internal/handler/gateway_compat_stream_usage_test.go b/backend/internal/handler/gateway_compat_stream_usage_test.go index 73ae8ef53419..31e9367e8f32 100644 --- a/backend/internal/handler/gateway_compat_stream_usage_test.go +++ b/backend/internal/handler/gateway_compat_stream_usage_test.go @@ -14,6 +14,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" + "github.com/Wei-Shaw/sub2api/internal/repository" middleware "github.com/Wei-Shaw/sub2api/internal/server/middleware" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/gin-gonic/gin" @@ -68,8 +69,14 @@ type compatStreamUsageRun struct { // runCompatStreamUsageHandler 经 GatewayHandler 转发一次流式请求;上游写完 upstreamStream 后, // closeUpstream 为 false 时保持连接不再发送数据。 func runCompatStreamUsageHandler(t *testing.T, endpoint, upstreamStream string, closeUpstream bool) compatStreamUsageRun { + return runCompatStreamUsageHandlerWithDelivery(t, endpoint, upstreamStream, closeUpstream, nil, "") +} + +func runCompatStreamUsageHandlerWithDelivery(t *testing.T, endpoint, upstreamStream string, closeUpstream bool, durable service.UsageBillingRepository, baseURL string) compatStreamUsageRun { t.Helper() cfg := &config.Config{Gateway: config.GatewayConfig{StreamDataIntervalTimeout: 1}} + cfg.Security.URLAllowlist.AllowPrivateHosts = true + cfg.Security.URLAllowlist.AllowInsecureHTTP = true groupID := int64(9300) group := &service.Group{ID: groupID, Hydrated: true, Platform: service.PlatformAnthropic, Status: service.StatusActive, RateMultiplier: 1} account := &service.Account{ @@ -99,9 +106,18 @@ func runCompatStreamUsageHandler(t *testing.T, endpoint, upstreamStream string, StatusCode: http.StatusOK, Body: pr, Header: http.Header{"X-Request-Id": []string{"compat-stream-usage"}}, }} + var billing service.UsageBillingRepository = run.billingRepo + var transport service.HTTPUpstream = run.upstream + if durable != nil { + billing = durable + } + if baseURL != "" { + account.Credentials["base_url"] = baseURL + transport = repository.NewHTTPUpstream(cfg) + } gw := service.NewGatewayService( - nil, &fakeGroupRepo{group: group}, run.usageRepo, run.billingRepo, nil, nil, nil, nil, cfg, - snapshot, nil, service.NewBillingService(cfg, nil), nil, nil, nil, run.upstream, + nil, &fakeGroupRepo{group: group}, run.usageRepo, billing, nil, nil, nil, nil, cfg, + snapshot, nil, service.NewBillingService(cfg, nil), nil, nil, nil, transport, &service.DeferredService{}, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, ) h := &GatewayHandler{ diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index ad0ee0071bdd..6d0f2554c31a 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -2493,7 +2493,7 @@ func (h *GatewayHandler) submitUsageRecordTask(parent context.Context, task serv return } task, abandon := wrapUsageRecordTaskContext(parent, task) - if h.usageRecordWorkerPool != nil { + if h.usageRecordWorkerPool != nil && !h.gatewayService.DurableUsageBillingEnabled() { if mode := h.usageRecordWorkerPool.Submit(task); mode != service.UsageRecordSubmitModeDroppedStopped { if mode.Dropped() { abandon() @@ -2526,7 +2526,7 @@ func (h *GatewayHandler) submitMandatoryUsageRecordTask(parent context.Context, return } task, _ = wrapUsageRecordTaskContext(parent, task) - if h.usageRecordWorkerPool != nil { + if h.usageRecordWorkerPool != nil && !h.gatewayService.DurableUsageBillingEnabled() { if mode := h.usageRecordWorkerPool.Submit(task); !mode.Dropped() { return } diff --git a/backend/internal/handler/openai_billing_delivery_http_test.go b/backend/internal/handler/openai_billing_delivery_http_test.go new file mode 100644 index 000000000000..1a5602648ee3 --- /dev/null +++ b/backend/internal/handler/openai_billing_delivery_http_test.go @@ -0,0 +1,111 @@ +//go:build unit + +package handler + +import ( + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/repository" + "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" +) + +func TestOpenAIGatewayHandlerRealHTTPDeliveryFailureDoesNotInventUsageOrRecallModel(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, phase := range []string{"stage", "execute"} { + t.Run(phase, func(t *testing.T) { + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"resp_billing_boundary","object":"response","status":"completed","model":"gpt-5.1","output":[],"usage":{"input_tokens":12,"output_tokens":7,"total_tokens":19}}`) + })) + defer upstream.Close() + cfg := &config.Config{} + cfg.Default.RateMultiplier = 1 + cfg.Security.URLAllowlist.AllowPrivateHosts = true + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + groupID := int64(9201) + group := &service.Group{ID: groupID, Platform: service.PlatformOpenAI, Status: service.StatusActive, Hydrated: true, RateMultiplier: 1} + account := service.Account{ID: 9202, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, + Status: service.StatusActive, Schedulable: true, Concurrency: 1, GroupIDs: []int64{groupID}, + Credentials: map[string]any{"api_key": "synthetic", "base_url": upstream.URL}} + accounts := &openAIWSUsageHandlerAccountRepoStub{account: account} + logs := &openAIWSUsageHandlerUsageLogRepoStub{created: make(chan *service.UsageLog, 2)} + billing := &httpDeliveryBilling{} + if phase == "stage" { + billing.stageFailure = errors.New("injected durable stage failure") + } else { + billing.failure = errors.New("injected execution failure") + } + billingCache := service.NewBillingCacheService(newHandlerInflightCache(100), nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(billingCache.Stop) + concurrency := service.NewConcurrencyService(&concurrencyCacheMock{ + acquireUserSlotFn: func(context.Context, int64, int, string) (bool, error) { return true, nil }, + acquireAccountSlotFn: func(context.Context, int64, int, string) (bool, error) { return true, nil }, + }) + gateway := service.NewOpenAIGatewayService(accounts, nil, logs, billing, nil, nil, nil, nil, cfg, nil, concurrency, + service.NewBillingService(cfg, nil), nil, nil, repository.NewHTTPUpstream(cfg), &service.DeferredService{}, nil, nil, nil, nil, nil, nil, nil) + t.Cleanup(gateway.CloseOpenAIWSPool) + h := &OpenAIGatewayHandler{gatewayService: gateway, cfg: cfg, billingCacheService: billingCache, + apiKeyService: &service.APIKeyService{}, concurrencyHelper: NewConcurrencyHelper(concurrency, SSEPingFormatNone, 0)} + key := &service.APIKey{ID: 9203, UserID: 9204, GroupID: &groupID, Group: group, + User: &service.User{ID: 9204, Status: service.StatusActive, Balance: 100}} + router := gin.New() + router.Use(middleware.ClientRequestID(), func(c *gin.Context) { + c.Set(string(middleware.ContextKeyAPIKey), key) + c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: key.UserID, Concurrency: 1}) + c.Next() + }) + finished := make(chan struct{}) + router.POST("/openai/v1/responses", func(c *gin.Context) { + defer close(finished) + h.Responses(c) + }) + ingress := httptest.NewServer(router) + defer ingress.Close() + client := &http.Client{Timeout: 5 * time.Second} + res, err := client.Post(ingress.URL+"/openai/v1/responses", "application/json", strings.NewReader(`{"model":"gpt-5.1","input":"synthetic","stream":false}`)) + require.NoError(t, err) + body, err := io.ReadAll(res.Body) + _ = res.Body.Close() + require.NoError(t, err) + select { + case <-finished: + case <-time.After(5 * time.Second): + t.Fatal("actual OpenAI handler did not finish its delivery boundary") + } + require.Equal(t, http.StatusOK, res.StatusCode) + require.Contains(t, string(body), "resp_billing_boundary") + require.Equal(t, int32(1), calls.Load()) + require.Empty(t, logs.created, "neither a stage failure nor an execution failure is a successful zero-cost usage row") + if phase == "stage" { + require.Nil(t, billing.event, "pre-storage failure is an explicit unrecoverable observation boundary") + return + } + require.NotNil(t, billing.event, "the actual OpenAI handler must stage synchronously before it returns") + require.Equal(t, 12, billing.event.Usage.InputTokens) + require.Equal(t, 7, billing.event.Usage.OutputTokens) + require.Greater(t, billing.event.Usage.ActualCost, 0.0) + billing.failure = nil + eventKey := service.UsageBillingEventKey{RequestID: billing.event.Command.RequestID, APIKeyID: key.ID} + for repeat := 0; repeat < 2; repeat++ { + result, err := billing.ApplyStagedUsageBilling(context.Background(), eventKey) + require.NoError(t, err) + require.Equal(t, repeat == 0, result.Result.Applied) + } + require.Equal(t, int32(1), calls.Load()) + }) + } +} diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index ad518c9d68fb..dbc5616a0968 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -4022,7 +4022,7 @@ func (h *OpenAIGatewayHandler) submitUsageRecordTask(parent context.Context, tas return } task, abandon := wrapUsageRecordTaskContext(parent, task) - if h.usageRecordWorkerPool != nil { + if h.usageRecordWorkerPool != nil && !h.gatewayService.DurableUsageBillingEnabled() { if mode := h.usageRecordWorkerPool.Submit(task); mode != service.UsageRecordSubmitModeDroppedStopped { if mode.Dropped() { abandon() @@ -4064,7 +4064,7 @@ func (h *OpenAIGatewayHandler) submitMandatoryUsageRecordTask(parent context.Con return } task, _ = wrapUsageRecordTaskContext(parent, task) - if h.usageRecordWorkerPool != nil { + if h.usageRecordWorkerPool != nil && !h.gatewayService.DurableUsageBillingEnabled() { if mode := h.usageRecordWorkerPool.Submit(task); !mode.Dropped() { return } diff --git a/backend/internal/repository/batch_image_money_integration_test.go b/backend/internal/repository/batch_image_money_integration_test.go new file mode 100644 index 000000000000..6720ab292d6c --- /dev/null +++ b/backend/internal/repository/batch_image_money_integration_test.go @@ -0,0 +1,157 @@ +//go:build integration + +package repository + +import ( + "context" + "fmt" + "strings" + "sync" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +// Real NUMERIC columns and simultaneous jobs, rather than a float-only mock. +func TestBatchImageMoneyConcurrentJobsConserveBalance(t *testing.T) { + ctx := context.Background() + db, client := newBillingAuditDatabase(t) + user := mustCreateUser(t, client, &service.User{Email: uuid.NewString() + "@example.com", PasswordHash: "hash", Balance: 100}) + repo := NewUsageBillingRepository(nil, db) + amounts := []float64{0.07619951, 0.07621186} + jobs := make([]string, len(amounts)) + var wg sync.WaitGroup + errs := make(chan error, len(amounts)) + for i, amount := range amounts { + jobs[i] = uuid.NewString() + wg.Add(1) + go func(id string, hold float64) { + defer wg.Done() + _, err := repo.ReserveBatchImageBalance(ctx, &service.BatchImageBalanceHoldCommand{ + RequestID: service.BatchImageHoldRequestID(id), APIKeyID: 17, UserID: user.ID, BatchID: id, HoldAmount: hold, MonetaryVersion: 2, + }) + errs <- err + }(jobs[i], amount) + } + wg.Wait() + close(errs) + for err := range errs { + require.NoError(t, err) + } + // Capture one job, refund the other, and replay both operations. + for attempt := 0; attempt < 2; attempt++ { + _, err := repo.CaptureBatchImageBalance(ctx, &service.BatchImageBalanceHoldCommand{ + RequestID: service.BatchImageCaptureRequestID(jobs[0]), APIKeyID: 17, UserID: user.ID, BatchID: jobs[0], HoldAmount: amounts[0], ActualAmount: 0.05, MonetaryVersion: 2, + }) + require.NoError(t, err) + _, err = repo.ReleaseBatchImageBalance(ctx, &service.BatchImageBalanceHoldCommand{ + RequestID: service.BatchImageReleaseRequestID(jobs[0]), APIKeyID: 17, UserID: user.ID, BatchID: jobs[0], HoldAmount: amounts[0], MonetaryVersion: 2, + }) + require.ErrorIs(t, err, service.ErrUsageBillingRequestConflict, "capture must not be refunded out of another job's frozen funds") + _, err = repo.ReleaseBatchImageBalance(ctx, &service.BatchImageBalanceHoldCommand{ + RequestID: service.BatchImageReleaseRequestID(jobs[1]), APIKeyID: 17, UserID: user.ID, BatchID: jobs[1], HoldAmount: amounts[1], MonetaryVersion: 2, + }) + require.NoError(t, err) + } + var balance, frozen string + require.NoError(t, db.QueryRowContext(ctx, "SELECT balance::text,frozen_balance::text FROM users WHERE id=$1", user.ID).Scan(&balance, &frozen)) + require.Equal(t, "99.95000000", balance) + require.Equal(t, "0.00000000", frozen) + // A legacy ambiguous hold must not consume money belonging to another job. + _, err := repo.ReleaseBatchImageBalance(ctx, &service.BatchImageBalanceHoldCommand{ + RequestID: fmt.Sprintf("legacy:%s", jobs[0]), APIKeyID: 17, UserID: user.ID, BatchID: jobs[0], HoldAmount: 0.0761995125, + }) + require.ErrorIs(t, err, service.ErrBatchImageLegacyMoneyPrecision) +} + +func TestBatchImageOwnedMoneyCaptureFailureRollsBackAndFreeJobsStayFree(t *testing.T) { + ctx := context.Background() + db, client := newBillingAuditDatabase(t) + user := mustCreateUser(t, client, &service.User{Email: uuid.NewString() + "@example.com", PasswordHash: "hash", Balance: 100}) + repo := NewUsageBillingRepository(nil, db) + id := uuid.NewString() + _, err := repo.ReserveBatchImageBalance(ctx, &service.BatchImageBalanceHoldCommand{RequestID: service.BatchImageHoldRequestID(id), + APIKeyID: 17, UserID: user.ID, BatchID: id, HoldAmount: 1, MonetaryVersion: 2}) + require.NoError(t, err) + name := "hold_fail_" + strings.ReplaceAll(uuid.NewString(), "-", "") + _, err = db.ExecContext(ctx, fmt.Sprintf(`CREATE FUNCTION %s() RETURNS trigger LANGUAGE plpgsql AS $$ + BEGIN IF NEW.batch_id='%s' THEN RAISE EXCEPTION 'injected receipt failure'; END IF; RETURN NEW; END; $$; + CREATE TRIGGER %s BEFORE UPDATE ON batch_image_money_holds FOR EACH ROW EXECUTE FUNCTION %s()`, name, id, name, name)) + require.NoError(t, err) + cleanup := func() { + _, _ = db.ExecContext(ctx, fmt.Sprintf("DROP TRIGGER IF EXISTS %s ON batch_image_money_holds; DROP FUNCTION IF EXISTS %s()", name, name)) + } + t.Cleanup(cleanup) + capture := &service.BatchImageBalanceHoldCommand{RequestID: service.BatchImageCaptureRequestID(id), APIKeyID: 17, + UserID: user.ID, BatchID: id, HoldAmount: 1, ActualAmount: 0.5, MonetaryVersion: 2} + _, err = repo.CaptureBatchImageBalance(ctx, capture) + require.Error(t, err) + var balance, frozen string + require.NoError(t, db.QueryRowContext(ctx, "SELECT balance::text,frozen_balance::text FROM users WHERE id=$1", user.ID).Scan(&balance, &frozen)) + require.Equal(t, "99.00000000", balance) + require.Equal(t, "1.00000000", frozen) + cleanup() + _, err = repo.CaptureBatchImageBalance(ctx, capture) + require.NoError(t, err) + free := uuid.NewString() + _, err = repo.CaptureBatchImageBalance(ctx, &service.BatchImageBalanceHoldCommand{RequestID: service.BatchImageCaptureRequestID(free), + APIKeyID: 17, UserID: user.ID, BatchID: free, MonetaryVersion: 2}) + require.NoError(t, err) + require.NoError(t, db.QueryRowContext(ctx, "SELECT balance::text,frozen_balance::text FROM users WHERE id=$1", user.ID).Scan(&balance, &frozen)) + require.Equal(t, "99.50000000", balance) + require.Equal(t, "0.00000000", frozen) +} + +func TestBatchImageSameTaskCaptureReleaseCompeteWithoutTakingOtherHold(t *testing.T) { + ctx := context.Background() + db, client := newBillingAuditDatabase(t) + user := mustCreateUser(t, client, &service.User{Email: uuid.NewString() + "@example.com", PasswordHash: "hash", Balance: 100}) + repo := NewUsageBillingRepository(nil, db) + batch, other := uuid.NewString(), uuid.NewString() + for id, amount := range map[string]float64{batch: 1, other: 2} { + _, err := repo.ReserveBatchImageBalance(ctx, &service.BatchImageBalanceHoldCommand{RequestID: service.BatchImageHoldRequestID(id), + BatchID: id, UserID: user.ID, APIKeyID: 17, HoldAmount: amount, MonetaryVersion: 2}) + require.NoError(t, err) + } + start := make(chan struct{}) + errs := make(chan error, 2) + go func() { + <-start + _, err := repo.CaptureBatchImageBalance(ctx, &service.BatchImageBalanceHoldCommand{RequestID: service.BatchImageCaptureRequestID(batch), + BatchID: batch, UserID: user.ID, APIKeyID: 17, HoldAmount: 1, ActualAmount: 0.4, MonetaryVersion: 2}) + errs <- err + }() + go func() { + <-start + _, err := repo.ReleaseBatchImageBalance(ctx, &service.BatchImageBalanceHoldCommand{RequestID: service.BatchImageReleaseRequestID(batch), + BatchID: batch, UserID: user.ID, APIKeyID: 17, HoldAmount: 1, MonetaryVersion: 2}) + errs <- err + }() + close(start) + success, conflict := 0, 0 + for i := 0; i < 2; i++ { + if err := <-errs; err != nil { + require.ErrorIs(t, err, service.ErrUsageBillingRequestConflict) + conflict++ + } else { + success++ + } + } + require.Equal(t, 1, success) + require.Equal(t, 1, conflict) + var state, balance, frozen string + var receipts int + require.NoError(t, db.QueryRowContext(ctx, "SELECT state FROM batch_image_money_holds WHERE batch_id=$1 AND api_key_id=17", batch).Scan(&state)) + require.NoError(t, db.QueryRowContext(ctx, "SELECT balance::text,frozen_balance::text FROM users WHERE id=$1", user.ID).Scan(&balance, &frozen)) + require.NoError(t, db.QueryRowContext(ctx, "SELECT count(*) FROM usage_billing_dedup WHERE api_key_id=17 AND request_id IN ($1,$2)", service.BatchImageCaptureRequestID(batch), service.BatchImageReleaseRequestID(batch)).Scan(&receipts)) + require.Equal(t, 1, receipts) + require.Equal(t, "2.00000000", frozen, "the other task's hold remains intact") + if state == "captured" { + require.Equal(t, "97.60000000", balance) + } else { + require.Equal(t, "released", state) + require.Equal(t, "98.00000000", balance) + } +} diff --git a/backend/internal/repository/batch_image_owned_money.go b/backend/internal/repository/batch_image_owned_money.go new file mode 100644 index 000000000000..2d726016eec5 --- /dev/null +++ b/backend/internal/repository/batch_image_owned_money.go @@ -0,0 +1,71 @@ +package repository + +import ( + "context" + "database/sql" + "errors" + + "github.com/Wei-Shaw/sub2api/internal/service" +) + +// Lock order is dedup -> user -> owned hold. All changes commit with the claim. +// A completed hold cannot spend another job's share of pooled frozen_balance. +func applyOwnedBatchImageMoney(ctx context.Context, tx *sql.Tx, cmd *service.BatchImageBalanceHoldCommand, operation string, + apply func(context.Context, *sql.Tx, *service.BatchImageBalanceHoldCommand) (*service.BatchImageBalanceHoldResult, error), +) (*service.BatchImageBalanceHoldResult, error) { + wantID := map[string]string{"hold": service.BatchImageHoldRequestID(cmd.BatchID), + "capture": service.BatchImageCaptureRequestID(cmd.BatchID), "release": service.BatchImageReleaseRequestID(cmd.BatchID)}[operation] + if cmd.RequestID != wantID || cmd.BatchID == "" { + return nil, service.ErrUsageBillingRequestConflict + } + var user int64 + if err := tx.QueryRowContext(ctx, "SELECT id FROM users WHERE id=$1 AND deleted_at IS NULL FOR NO KEY UPDATE", cmd.UserID).Scan(&user); err != nil { + return nil, err + } + var owner int64 + var hold float64 + var state string + err := tx.QueryRowContext(ctx, "SELECT user_id,hold_amount,state FROM batch_image_money_holds WHERE batch_id=$1 AND api_key_id=$2 FOR UPDATE", cmd.BatchID, cmd.APIKeyID).Scan(&owner, &hold, &state) + if errors.Is(err, sql.ErrNoRows) { + if operation == "capture" && cmd.HoldAmount == 0 && cmd.ActualAmount == 0 { + // Free jobs deliberately skip ReserveBatchImageBalance. Record a + // zero-money terminal receipt without touching pooled frozen funds. + _, err = tx.ExecContext(ctx, `INSERT INTO batch_image_money_holds(batch_id,api_key_id,user_id,hold_amount,actual_amount,state) + VALUES($1,$2,$3,0,0,'captured')`, cmd.BatchID, cmd.APIKeyID, cmd.UserID) + return &service.BatchImageBalanceHoldResult{}, err + } + if operation == "release" { + claimed, err := batchImageHoldClaimExists(ctx, tx, service.BatchImageHoldRequestID(cmd.BatchID), cmd.APIKeyID) + if err != nil { + return nil, err + } + if !claimed { + return &service.BatchImageBalanceHoldResult{}, nil + } + } + if operation != "hold" { + return nil, service.ErrBatchImageLegacyMoneyPrecision + } + if _, err = tx.ExecContext(ctx, `INSERT INTO batch_image_money_holds(batch_id,api_key_id,user_id,hold_amount,state) + VALUES($1,$2,$3,$4,'held')`, cmd.BatchID, cmd.APIKeyID, cmd.UserID, cmd.HoldAmount); err != nil { + return nil, err + } + return apply(ctx, tx, cmd) + } + if err != nil { + return nil, err + } + if owner != cmd.UserID || hold != cmd.HoldAmount || state != "held" || operation == "hold" { + return nil, service.ErrUsageBillingRequestConflict + } + result, err := apply(ctx, tx, cmd) + if err != nil { + return nil, err + } + next := "captured" + if operation == "release" { + next = "released" + } + _, err = tx.ExecContext(ctx, `UPDATE batch_image_money_holds SET state=$3,actual_amount=$4 WHERE batch_id=$1 AND api_key_id=$2`, cmd.BatchID, cmd.APIKeyID, next, cmd.ActualAmount) + return result, err +} diff --git a/backend/internal/repository/batch_image_receipt.go b/backend/internal/repository/batch_image_receipt.go new file mode 100644 index 000000000000..edefdc592bf7 --- /dev/null +++ b/backend/internal/repository/batch_image_receipt.go @@ -0,0 +1,39 @@ +package repository + +import ( + "context" + "database/sql" + "math" + + "github.com/Wei-Shaw/sub2api/internal/service" +) + +// One read snapshot spans hot/archive movement. This method neither claims an +// ID nor changes money, owned holds or task state. User/batch/amount/payload are +// bound by the original fingerprint, recomputed rather than trusting a caller. +func (r *usageBillingRepository) ConfirmBatchImageBillingReceipt(ctx context.Context, cmd *service.BatchImageBalanceHoldCommand) (bool, error) { + if cmd == nil || cmd.UserID <= 0 || cmd.APIKeyID <= 0 || cmd.BatchID == "" || + math.IsNaN(cmd.HoldAmount) || math.IsInf(cmd.HoldAmount, 0) || cmd.HoldAmount < 0 || + math.IsNaN(cmd.ActualAmount) || math.IsInf(cmd.ActualAmount, 0) || cmd.ActualAmount < 0 { + return false, service.ErrBatchImageLegacyMoneyPrecision + } + copy := *cmd + copy.RequestFingerprint = "" + copy.Normalize() + if copy.RequestID != service.BatchImageHoldRequestID(copy.BatchID) && + copy.RequestID != service.BatchImageCaptureRequestID(copy.BatchID) && copy.RequestID != service.BatchImageReleaseRequestID(copy.BatchID) { + return false, service.ErrUsageBillingRequestConflict + } + var matches sql.NullBool + err := r.db.QueryRowContext(ctx, `SELECT bool_and(request_fingerprint=$3) FROM ( + SELECT request_fingerprint FROM usage_billing_dedup WHERE request_id=$1 AND api_key_id=$2 + UNION ALL SELECT request_fingerprint FROM usage_billing_dedup_archive WHERE request_id=$1 AND api_key_id=$2 + ) receipts`, copy.RequestID, copy.APIKeyID, copy.RequestFingerprint).Scan(&matches) + if err != nil || !matches.Valid { + return false, err + } + if !matches.Bool { + return false, service.ErrUsageBillingRequestConflict + } + return true, nil +} diff --git a/backend/internal/repository/batch_image_receipt_integration_test.go b/backend/internal/repository/batch_image_receipt_integration_test.go new file mode 100644 index 000000000000..e52599ddc778 --- /dev/null +++ b/backend/internal/repository/batch_image_receipt_integration_test.go @@ -0,0 +1,90 @@ +//go:build integration + +package repository + +import ( + "context" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +type batchSnapshotOnlyPrice struct{} + +func (*batchSnapshotOnlyPrice) BatchImageUnitPrice(context.Context, *service.BatchImageJob) (float64, error) { + panic("submitted snapshot must be used instead of repricing") +} + +func TestLegacyBatchReceiptsReadOnlyMatchHotArchiveAndOwnership(t *testing.T) { + ctx := context.Background() + db, _ := newBillingAuditDatabase(t) + repo := &usageBillingRepository{db: db} + for _, archive := range []bool{false, true} { + for _, prefix := range []func(string) string{service.BatchImageHoldRequestID, service.BatchImageCaptureRequestID, service.BatchImageReleaseRequestID} { + batch := uuid.NewString() + cmd := &service.BatchImageBalanceHoldCommand{BatchID: batch, RequestID: prefix(batch), UserID: 31, APIKeyID: 41, + HoldAmount: 0.0761995125, RequestPayloadHash: "original"} + cmd.Normalize() + confirmed, err := repo.ConfirmBatchImageBillingReceipt(ctx, cmd) + require.NoError(t, err) + require.False(t, confirmed) + table := "usage_billing_dedup" + if archive { + table = "usage_billing_dedup_archive" + } + _, err = db.ExecContext(ctx, "INSERT INTO "+table+"(request_id,api_key_id,request_fingerprint,created_at) VALUES($1,$2,$3,NOW())", cmd.RequestID, cmd.APIKeyID, cmd.RequestFingerprint) + require.NoError(t, err) + confirmed, err = repo.ConfirmBatchImageBillingReceipt(ctx, cmd) + require.NoError(t, err) + require.True(t, confirmed) + wrong := *cmd + wrong.UserID++ + confirmed, err = repo.ConfirmBatchImageBillingReceipt(ctx, &wrong) + require.ErrorIs(t, err, service.ErrUsageBillingRequestConflict) + require.False(t, confirmed) + wrong = *cmd + wrong.RequestPayloadHash = "changed" + _, err = repo.ConfirmBatchImageBillingReceipt(ctx, &wrong) + require.ErrorIs(t, err, service.ErrUsageBillingRequestConflict) + var ownedRows int + require.NoError(t, db.QueryRowContext(ctx, "SELECT count(*) FROM batch_image_money_holds WHERE batch_id=$1", batch).Scan(&ownedRows)) + require.Zero(t, ownedRows, "receipt confirmation must not install historical owned holds") + } + } +} + +func TestLegacyCaptureCommittedBeforeTaskStatusUpgradeRetryDoesNotMoveMoney(t *testing.T) { + ctx := context.Background() + db, client := newBillingAuditDatabase(t) + user := mustCreateUser(t, client, &service.User{Email: uuid.NewString() + "@example.com", PasswordHash: "hash", Balance: 99.93876543}) + key, account := int64(17), int64(23) + hold, actual := 0.0761995125, 0.0612345678 + jobs := NewBatchImageRepository(db) + job, err := jobs.CreateBatchImageJob(ctx, service.CreateBatchImageJobParams{ + BatchID: uuid.NewString(), UserID: user.ID, APIKeyID: &key, AccountID: &account, + Provider: service.BatchImageProviderGeminiAPI, Model: "legacy-image", Status: service.BatchImageJobStatusSettling, + ItemCount: 1, SuccessCount: 1, HoldAmount: &hold, EstimatedCost: hold, BillableUnitPrice: actual, PricingSnapshotVersion: 1, + }) + require.NoError(t, err) + receipt := &service.BatchImageBalanceHoldCommand{RequestID: service.BatchImageCaptureRequestID(job.BatchID), BatchID: job.BatchID, + UserID: user.ID, APIKeyID: key, HoldAmount: hold, ActualAmount: actual, RequestPayloadHash: service.BuildBatchImageSettlementManifestHash(job)} + receipt.Normalize() + _, err = db.ExecContext(ctx, "INSERT INTO usage_billing_dedup(request_id,api_key_id,request_fingerprint) VALUES($1,$2,$3)", receipt.RequestID, key, receipt.RequestFingerprint) + require.NoError(t, err) + s := &service.BatchImageSettlementService{Repo: jobs, BillingRepo: &usageBillingRepository{db: db}, Pricing: &batchSnapshotOnlyPrice{}} + result, err := s.Settle(ctx, job.BatchID) + require.NoError(t, err) + require.False(t, result.AlreadySettled) + result, err = s.Settle(ctx, job.BatchID) + require.NoError(t, err) + require.True(t, result.AlreadySettled) + saved, err := jobs.GetBatchImageJobByBatchID(ctx, job.BatchID) + require.NoError(t, err) + require.Equal(t, service.BatchImageJobStatusCompleted, saved.Status) + var balance, frozen string + require.NoError(t, db.QueryRowContext(ctx, "SELECT balance::text,frozen_balance::text FROM users WHERE id=$1", user.ID).Scan(&balance, &frozen)) + require.Equal(t, "99.93876543", balance) + require.Equal(t, "0.00000000", frozen) +} diff --git a/backend/internal/repository/batch_image_snapshot_roundtrip_integration_test.go b/backend/internal/repository/batch_image_snapshot_roundtrip_integration_test.go new file mode 100644 index 000000000000..3cfa318b9252 --- /dev/null +++ b/backend/internal/repository/batch_image_snapshot_roundtrip_integration_test.go @@ -0,0 +1,107 @@ +//go:build integration + +package repository + +import ( + "context" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +type snapshotSubmitAccounts struct{ account *service.Account } + +func (r *snapshotSubmitAccounts) GetByID(context.Context, int64) (*service.Account, error) { + return r.account, nil +} +func (r *snapshotSubmitAccounts) ListSchedulableByPlatform(context.Context, string) ([]service.Account, error) { + return []service.Account{*r.account}, nil +} +func (r *snapshotSubmitAccounts) ListSchedulableByGroupIDAndPlatform(context.Context, int64, string) ([]service.Account, error) { + return []service.Account{*r.account}, nil +} + +type snapshotSubmitGroup struct{ group *service.Group } + +func (r *snapshotSubmitGroup) GetByIDLite(context.Context, int64) (*service.Group, error) { + return r.group, nil +} + +type snapshotSubmitProvider struct { + service.BatchImageProvider + calls int +} + +func (*snapshotSubmitProvider) Name() string { return service.BatchImageProviderGeminiAPI } +func (*snapshotSubmitProvider) SupportsAccount(*service.Account) bool { return true } +func (p *snapshotSubmitProvider) Submit(context.Context, *service.BatchImageJob, *service.Account, service.BatchImageInput) (*service.BatchProviderJob, error) { + p.calls++ + return &service.BatchProviderJob{ProviderJobName: "synthetic/provider-job"}, nil +} + +func TestBatchSubmitSnapshotDatabaseRoundTripThenSettlement(t *testing.T) { + for _, tc := range []struct { + rate, unit, hold float64 + balance string + }{ + {1.2345, 0.0761995125, 0.07619951, "99.92380049"}, + {1.2346, 0.0762118580, 0.07621186, "99.92378814"}, + } { + ctx := context.Background() + db, client := newBillingAuditDatabase(t) + user := mustCreateUser(t, client, &service.User{Email: uuid.NewString() + "@example.com", PasswordHash: "hash", Balance: 100}) + price := 0.05 + group := mustCreateGroup(t, client, &service.Group{Name: "snapshot-" + uuid.NewString(), Platform: service.PlatformGemini, + RateMultiplier: tc.rate, ImagePrice1K: &price, AllowImageGeneration: true, AllowBatchImageGeneration: true, + BatchImageDiscountMultiplier: 1, BatchImageHoldMultiplier: 1}) + key := mustCreateApiKey(t, client, &service.APIKey{UserID: user.ID, GroupID: &group.ID, Key: "sk-mock-" + uuid.NewString(), Name: "snapshot"}) + account := mustCreateAccount(t, client, &service.Account{Name: "snapshot-" + uuid.NewString(), Platform: service.PlatformGemini, + Type: service.AccountTypeAPIKey, Status: service.StatusActive, Schedulable: true, RateMultiplier: &tc.rate}) + account.GroupIDs = []int64{group.ID} + jobs := NewBatchImageRepository(db) + billing := NewUsageBillingRepository(nil, db) + provider := &snapshotSubmitProvider{} + public := &service.BatchImagePublicService{Repo: jobs, BillingRepo: billing, AccountRepo: &snapshotSubmitAccounts{account}, + GroupRepo: &snapshotSubmitGroup{group}, ProviderRegistry: service.NewBatchImageProviderRegistry(provider), Pricing: &batchSnapshotOnlyPrice{}, + Config: &config.Config{BatchImage: config.BatchImageConfig{Enabled: true, MaxItemsPerJobDefault: 2, MaxPromptCharsPerItem: 64, + DefaultResponseMimeType: "image/png", DefaultImageSize: "1K"}}} + owner := service.BatchImageOwner{UserID: user.ID, APIKeyID: key.ID, GroupID: &group.ID} + // Keep the request stable across wall-clock seconds: the service + // generates a time-based default TaskName when the field is empty. + req := service.BatchImageSubmitRequest{TaskName: "snapshot-roundtrip", Model: "gemini-2.5-flash-image", Provider: service.BatchImageProviderGeminiAPI, + ImageSize: "1K", Items: []service.BatchImageSubmitItem{{CustomID: "one", Prompt: "synthetic"}}} + idempotency := uuid.NewString() + submitted, err := public.Submit(ctx, owner, req, idempotency) + require.NoError(t, err) + _, err = public.Submit(ctx, owner, req, idempotency) + require.NoError(t, err) + require.Equal(t, 1, provider.calls) + // A fresh repository reads NUMERIC(20,10) prices and the exact eight-place + // hold, rather than keeping the Submit calculation's in-memory floats. + saved, err := NewBatchImageRepository(db).GetBatchImageJobByBatchID(ctx, submitted.ID) + require.NoError(t, err) + require.Equal(t, 2, saved.PricingSnapshotVersion) + require.Equal(t, tc.hold, *saved.HoldAmount) + require.Equal(t, tc.unit, saved.BillableUnitPrice) + require.NoError(t, jobs.TransitionBatchImageJobStatus(ctx, saved.BatchID, service.BatchImageJobStatusIndexing, service.BatchImageTransitionOptions{})) + _, err = db.ExecContext(ctx, "UPDATE batch_image_jobs SET success_count=1 WHERE batch_id=$1", saved.BatchID) + require.NoError(t, err) // synthetic provider result, not a real image/model call + require.NoError(t, jobs.TransitionBatchImageJobStatus(ctx, saved.BatchID, service.BatchImageJobStatusSettling, service.BatchImageTransitionOptions{})) + settlement := &service.BatchImageSettlementService{Repo: NewBatchImageRepository(db), BillingRepo: billing, Pricing: &batchSnapshotOnlyPrice{}} + for repeat := 0; repeat < 2; repeat++ { + result, err := settlement.Settle(ctx, saved.BatchID) + require.NoError(t, err) + require.Equal(t, tc.hold, result.ActualCost) + require.Equal(t, repeat > 0, result.AlreadySettled) + } + var balance, frozen, ownedState string + require.NoError(t, db.QueryRowContext(ctx, "SELECT balance::text,frozen_balance::text FROM users WHERE id=$1", user.ID).Scan(&balance, &frozen)) + require.NoError(t, db.QueryRowContext(ctx, "SELECT state FROM batch_image_money_holds WHERE batch_id=$1 AND api_key_id=$2", saved.BatchID, key.ID).Scan(&ownedState)) + require.Equal(t, tc.balance, balance) + require.Equal(t, "0.00000000", frozen) + require.Equal(t, "captured", ownedState) + } +} diff --git a/backend/internal/repository/billing_audit_fixture_integration_test.go b/backend/internal/repository/billing_audit_fixture_integration_test.go new file mode 100644 index 000000000000..159c84819d70 --- /dev/null +++ b/backend/internal/repository/billing_audit_fixture_integration_test.go @@ -0,0 +1,62 @@ +//go:build integration + +package repository + +import ( + "context" + "database/sql" + "testing" + "time" + + dbent "github.com/Wei-Shaw/sub2api/ent" + "github.com/stretchr/testify/require" +) + +// Billing repositories commit their own transactions and recovery workers scan +// every pending event. An outer rollback cannot isolate these fixtures. Reuse +// the existing migrated, container-only database fixture instead of the shared +// testEntClient, including for independently constructed repositories/workers. +func newBillingAuditDatabase(t *testing.T) (*sql.DB, *dbent.Client) { + t.Helper() + before := sharedBillingAuditRows(t) + var name string + // LIFO: worker/trigger cleanups registered by the caller run first, then the + // fixture closes its pool and drops only its own database, then we verify. + t.Cleanup(func() { + require.Equal(t, before, sharedBillingAuditRows(t), "billing fixture must not change shared routing/statistics or recovery rows") + if name != "" { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + var exists bool + require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT EXISTS(SELECT 1 FROM pg_database WHERE datname=$1)", name).Scan(&exists)) + require.False(t, exists, "the complete private billing fixture must be removed, including task/usage/owned-money references") + } + }) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() + db, client := newImageUpgradeDatabase(t, ctx) + require.NoError(t, db.QueryRowContext(ctx, "SELECT current_database()").Scan(&name)) + return db, client +} + +func sharedBillingAuditRows(t *testing.T) []int64 { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + // These are observations only. Never truncate/reset shared tables or delete + // another test's rows to make global routing/dashboard assertions pass. + var users, keys, accounts, groups, usage, events, claims, archive, quotes int64 + var jobs, items, jobEvents, holds, quotas, totals, requests, allocations int64 + require.NoError(t, integrationDB.QueryRowContext(ctx, `SELECT + (SELECT count(*) FROM users), (SELECT count(*) FROM api_keys), + (SELECT count(*) FROM accounts), (SELECT count(*) FROM groups), + (SELECT count(*) FROM usage_logs), (SELECT count(*) FROM usage_billing_events), + (SELECT count(*) FROM usage_billing_dedup), (SELECT count(*) FROM usage_billing_dedup_archive), + (SELECT count(*) FROM usage_billing_quotes), (SELECT count(*) FROM batch_image_jobs), + (SELECT count(*) FROM batch_image_items), (SELECT count(*) FROM batch_image_events), + (SELECT count(*) FROM batch_image_money_holds), (SELECT count(*) FROM user_platform_quotas), + (SELECT count(*) FROM usage_platform_quota_totals), (SELECT count(*) FROM regional_primary_requests), + (SELECT count(*) FROM regional_budget_allocations)`).Scan(&users, &keys, &accounts, &groups, &usage, &events, + &claims, &archive, "es, &jobs, &items, &jobEvents, &holds, "as, &totals, &requests, &allocations)) + return []int64{users, keys, accounts, groups, usage, events, claims, archive, quotes, jobs, items, jobEvents, holds, quotas, totals, requests, allocations} +} diff --git a/backend/internal/repository/billing_cache.go b/backend/internal/repository/billing_cache.go index af8ec49c08e1..cb3979288507 100644 --- a/backend/internal/repository/billing_cache.go +++ b/backend/internal/repository/billing_cache.go @@ -438,7 +438,10 @@ func parseUserPlatformQuotaHash(m map[string]string) *service.UserPlatformQuotaC } func (c *billingCache) GetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string) (*service.UserPlatformQuotaCacheEntry, bool, error) { - key := userPlatformQuotaCacheKey(userID, platform) + return c.getPlatformQuotaCache(ctx, userPlatformQuotaCacheKey(userID, platform)) +} + +func (c *billingCache) getPlatformQuotaCache(ctx context.Context, key string) (*service.UserPlatformQuotaCacheEntry, bool, error) { m, err := c.rdb.HGetAll(ctx, key).Result() if err != nil { return nil, false, err @@ -452,10 +455,13 @@ func (c *billingCache) GetUserPlatformQuotaCache(ctx context.Context, userID int } func (c *billingCache) SetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string, entry *service.UserPlatformQuotaCacheEntry, ttl time.Duration) error { + return c.setPlatformQuotaCache(ctx, userPlatformQuotaCacheKey(userID, platform), entry, ttl) +} + +func (c *billingCache) setPlatformQuotaCache(ctx context.Context, key string, entry *service.UserPlatformQuotaCacheEntry, ttl time.Duration) error { if entry == nil { return nil } - key := userPlatformQuotaCacheKey(userID, platform) pipe := c.rdb.TxPipeline() // 浮点可空字段:nil → 空字符串(读取时 parseFloatPtr 返回 nil,表示无限额) @@ -492,7 +498,7 @@ func (c *billingCache) SetUserPlatformQuotaCache(ctx context.Context, userID int } func (c *billingCache) DeleteUserPlatformQuotaCache(ctx context.Context, userID int64, platform string) error { - return c.rdb.Del(ctx, userPlatformQuotaCacheKey(userID, platform)).Err() + return c.rdb.Del(ctx, userPlatformQuotaCacheKey(userID, platform), durablePlatformQuotaCacheKey(userID, platform)).Err() } // updateUserPlatformQuotaUsageScript 缓存累加:EXISTS + schema_version 双重守卫。 @@ -544,17 +550,20 @@ func (c *billingCache) IncrUserPlatformQuotaUsageCache(ctx context.Context, user if markDirty { member = userPlatformQuotaDirtyMember(userID, platform) } - _, err := c.rdb.Eval(ctx, updateUserPlatformQuotaUsageScript, + applied, err := c.rdb.Eval(ctx, updateUserPlatformQuotaUsageScript, []string{userPlatformQuotaCacheKey(userID, platform), userPlatformQuotaDirtySetKey()}, strconv.FormatFloat(cost, 'f', -1, 64), int(ttl.Seconds()), service.UserPlatformQuotaCacheSchemaV1, member, userPlatformQuotaDirtyTTLSeconds, - ).Result() + ).Int64() if err != nil && !errors.Is(err, redis.Nil) { return err } + if applied != 1 { + return service.ErrUserPlatformQuotaCacheNotUpdated + } return nil } diff --git a/backend/internal/repository/billing_cache_user_platform_quota_test.go b/backend/internal/repository/billing_cache_user_platform_quota_test.go index 15b185e77693..539d86d3a31a 100644 --- a/backend/internal/repository/billing_cache_user_platform_quota_test.go +++ b/backend/internal/repository/billing_cache_user_platform_quota_test.go @@ -4,6 +4,7 @@ package repository import ( "context" + "errors" "testing" "time" @@ -86,10 +87,10 @@ func TestUserPlatformQuotaCache_NilLimitSetThenGet(t *testing.T) { } } -func TestUserPlatformQuotaCache_IncrMissIsNoop(t *testing.T) { +func TestUserPlatformQuotaCache_IncrMissReportsNotApplied(t *testing.T) { c, _ := newMiniRedisCache(t) - if err := c.IncrUserPlatformQuotaUsageCache(context.Background(), 1, "openai", 0.5, time.Minute, false); err != nil { - t.Fatal(err) + if err := c.IncrUserPlatformQuotaUsageCache(context.Background(), 1, "openai", 0.5, time.Minute, false); !errors.Is(err, service.ErrUserPlatformQuotaCacheNotUpdated) { + t.Fatalf("missing cache must report not applied, got %v", err) } _, ok, _ := c.GetUserPlatformQuotaCache(context.Background(), 1, "openai") if ok { diff --git a/backend/internal/repository/durable_platform_quota_cache.go b/backend/internal/repository/durable_platform_quota_cache.go new file mode 100644 index 000000000000..fd79138ea0e8 --- /dev/null +++ b/backend/internal/repository/durable_platform_quota_cache.go @@ -0,0 +1,25 @@ +package repository + +import ( + "context" + "fmt" + "time" + + "github.com/Wei-Shaw/sub2api/internal/service" +) + +func durablePlatformQuotaCacheKey(user int64, platform string) string { + return fmt.Sprintf("billing:durable_platform_quota:%d:%s", user, platform) +} + +func (c *billingCache) GetDurablePlatformQuotaCache(ctx context.Context, user int64, platform string) (*service.UserPlatformQuotaCacheEntry, bool, error) { + return c.getPlatformQuotaCache(ctx, durablePlatformQuotaCacheKey(user, platform)) +} + +func (c *billingCache) SetDurablePlatformQuotaCache(ctx context.Context, user int64, platform string, entry *service.UserPlatformQuotaCacheEntry, ttl time.Duration) error { + return c.setPlatformQuotaCache(ctx, durablePlatformQuotaCacheKey(user, platform), entry, ttl) +} + +func (c *billingCache) DeleteDurablePlatformQuotaCache(ctx context.Context, user int64, platform string) error { + return c.rdb.Del(ctx, durablePlatformQuotaCacheKey(user, platform)).Err() +} diff --git a/backend/internal/repository/durable_platform_quota_cache_test.go b/backend/internal/repository/durable_platform_quota_cache_test.go new file mode 100644 index 000000000000..3f39d9f60531 --- /dev/null +++ b/backend/internal/repository/durable_platform_quota_cache_test.go @@ -0,0 +1,39 @@ +//go:build unit + +package repository + +import ( + "context" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func TestDurablePlatformQuotaCacheMissSchemaAndLegacyDirtyIsolation(t *testing.T) { + cache, _ := newMiniRedisCache(t) + ctx := context.Background() + _, ok, err := cache.GetDurablePlatformQuotaCache(ctx, 1, "openai") + require.NoError(t, err) + require.False(t, ok) + // A legacy schema miss is distinguishable from an applied increment. + require.NoError(t, cache.SetUserPlatformQuotaCache(ctx, 1, "openai", &service.UserPlatformQuotaCacheEntry{}, time.Minute)) + require.ErrorIs(t, cache.IncrUserPlatformQuotaUsageCache(ctx, 1, "openai", 1, time.Minute, true), service.ErrUserPlatformQuotaCacheNotUpdated) + legacy := &service.UserPlatformQuotaCacheEntry{SchemaVersion: service.UserPlatformQuotaCacheSchemaV1, DailyUsageUSD: 2} + require.NoError(t, cache.SetUserPlatformQuotaCache(ctx, 1, "openai", legacy, time.Minute)) + require.NoError(t, cache.IncrUserPlatformQuotaUsageCache(ctx, 1, "openai", 0.5, time.Minute, true)) + combined := &service.UserPlatformQuotaCacheEntry{SchemaVersion: service.UserPlatformQuotaCacheSchemaV1, DailyUsageUSD: 3.5} + require.NoError(t, cache.SetDurablePlatformQuotaCache(ctx, 1, "openai", combined, time.Minute)) + // Retrying the new event can invalidate repeatedly without erasing old + // unflushed increments, or placing a combined total into the legacy dirty set. + for i := 0; i < 2; i++ { + require.NoError(t, cache.DeleteDurablePlatformQuotaCache(ctx, 1, "openai")) + } + keys, err := cache.PopDirtyUserPlatformQuotaKeys(ctx, 10) + require.NoError(t, err) + require.Len(t, keys, 1) + entries, err := cache.BatchGetUserPlatformQuotaCache(ctx, keys) + require.NoError(t, err) + require.Equal(t, 2.5, entries[0].DailyUsageUSD, "legacy flusher must never receive the combined 3.5 total") +} diff --git a/backend/internal/repository/platform_quota_writer_handoff_integration_test.go b/backend/internal/repository/platform_quota_writer_handoff_integration_test.go new file mode 100644 index 000000000000..7d8b182543a4 --- /dev/null +++ b/backend/internal/repository/platform_quota_writer_handoff_integration_test.go @@ -0,0 +1,96 @@ +//go:build integration + +package repository + +import ( + "context" + "database/sql" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/regional" + "github.com/Wei-Shaw/sub2api/internal/regional/ledger" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +func TestPlatformQuotaOrdinaryRegionalHandoffPreservesRollingMonthAndMoney(t *testing.T) { + ctx := context.Background() + repo, cmd, usage, client := durableBillingFixture(t) + db := repo.db + group := mustCreateGroup(t, client, &service.Group{Name: "handoff-" + uuid.NewString(), Platform: service.PlatformOpenAI, RateMultiplier: 1}) + _, err := db.ExecContext(ctx, "UPDATE api_keys SET group_id=$2 WHERE id=$1", cmd.APIKeyID, group.ID) + require.NoError(t, err) + _, err = db.ExecContext(ctx, "UPDATE accounts SET platform='openai',status='active',schedulable=true WHERE id=$1", cmd.AccountID) + require.NoError(t, err) + _, err = db.ExecContext(ctx, "INSERT INTO account_groups(account_id,group_id)VALUES($1,$2)", cmd.AccountID, group.ID) + require.NoError(t, err) + now := time.Now().UTC().Truncate(time.Microsecond) + ordinaryAt := now.Add(-10 * 24 * time.Hour) + _, err = db.ExecContext(ctx, `INSERT INTO user_platform_quotas + (user_id,platform,monthly_limit_usd,monthly_usage_usd,monthly_window_start) + VALUES($1,'openai',6,4,$2)`, cmd.UserID, now.Add(-45*24*time.Hour)) + require.NoError(t, err) + ordinary := func(at time.Time, amount float64) { + c, u := cmd, usage + c.RequestID, c.RequestFingerprint = uuid.NewString(), "" + c.BalanceCost, c.APIKeyQuotaCost = amount, amount + c.PlatformQuotaPlatform, c.PlatformQuotaCost, c.PlatformUsageAt = "openai", amount, at + u.RequestID, u.ActualCost, u.TotalCost, u.CreatedAt = c.RequestID, amount, amount, at + require.NoError(t, repo.StageUsageBilling(ctx, &c, &u)) + _, err := repo.ApplyStagedUsageBilling(ctx, service.UsageBillingEventKey{RequestID: c.RequestID, APIKeyID: c.APIKeyID}) + require.NoError(t, err) + } + ordinary(ordinaryAt, 5) + quotaRepo := NewUserPlatformQuotaRepository(client) + assertMonth := func(amount float64, start time.Time) { + rec, err := quotaRepo.GetByUserPlatform(ctx, cmd.UserID, "openai") + require.NoError(t, err) + require.Equal(t, amount, rec.MonthlyUsageUSD) + require.NotNil(t, rec.MonthlyWindowStart) + require.True(t, rec.MonthlyWindowStart.Equal(start)) + } + assertMonth(5, ordinaryAt) + _, holds, _ := NewRegionalFundsAuthorities(db, func(context.Context, *sql.Tx) error { return nil }) + price := ledger.Price{Model: "handoff-" + uuid.NewString(), Input: 500_000_000_000_000, MaxInput: 1, MaxOutput: 1} + require.NoError(t, ledger.InstallPrice(ctx, db, price)) + request := func(at time.Time, maximum regional.Amount) ledger.Reservation { + return ledger.Reservation{RequestID: uuid.NewString(), UserID: cmd.UserID, APIKeyID: cmd.APIKeyID, AccountID: cmd.AccountID, + RegionID: "primary", AdmittedAt: at, Maximum: maximum, Model: price.Model, + PricingVersion: price.Version(), PayloadSHA256: strings.Repeat("a", 64)} + } + // Actual primary admission, not just a helper invocation: 5 already spent + // plus a maximum of 2 must exceed the monthly limit of 6. + _, err = holds.Reserve(ctx, request(now, 2_000_000_000)) + require.ErrorIs(t, err, ledger.ErrExhausted) + assertMonth(5, ordinaryAt) + complete := func(at time.Time) { + r := request(at, 1_000_000_000) + fresh, err := holds.Reserve(ctx, r) + require.NoError(t, err) + require.True(t, fresh) + require.NoError(t, holds.Start(ctx, r.RequestID)) + e := regionalLedgerEvent(r, 500_000_000) + e.InputTokens, e.OutputTokens = 1, 0 + e.Usage.CompletedAt = at.Add(time.Millisecond) + require.NoError(t, holds.Complete(ctx, e)) + require.NoError(t, holds.Complete(ctx, e)) + } + complete(now) + assertMonth(5.5, ordinaryAt) + boundary := ordinaryAt.Add(30 * 24 * time.Hour) + _, err = holds.Reserve(ctx, request(boundary.Add(-time.Microsecond), 1_000_000_000)) + require.ErrorIs(t, err, ledger.ErrExhausted) + assertMonth(5.5, ordinaryAt) + complete(boundary) + assertMonth(0.5, boundary) + // Switching back to ordinary delivery must share the newly advanced + // regional anchor; an hour-later anchor must not hide the regional 0.5. + ordinary(boundary.Add(time.Hour), 0.25) + assertMonth(0.75, boundary) + var balance string + require.NoError(t, db.QueryRowContext(ctx, "SELECT balance::text FROM users WHERE id=$1", cmd.UserID).Scan(&balance)) + require.Equal(t, "93.75000000", balance) +} diff --git a/backend/internal/repository/regional_budget_scope.go b/backend/internal/repository/regional_budget_scope.go index 5590950755d5..d616bd1e0459 100644 --- a/backend/internal/repository/regional_budget_scope.go +++ b/backend/internal/repository/regional_budget_scope.go @@ -134,7 +134,7 @@ func (s RegionalBudgetScope) Reserve(ctx context.Context, tx *sql.Tx, id string, } if err == nil { for _, window := range []string{"daily", "weekly", "monthly"} { - query := "SELECT " + window + "_limit_usd IS NULL OR " + window + "_usage_usd+$2::numeric+COALESCE((SELECT SUM(maximum-spent)::numeric/1000000000 FROM regional_budget_allocations WHERE user_id=$3 AND platform=$4 AND state='held'),0)<=" + window + "_limit_usd FROM user_platform_quotas WHERE id=$1" + query := "SELECT " + window + "_limit_usd IS NULL OR " + window + "_usage_usd+$2::numeric+COALESCE((SELECT SUM(maximum-spent)::numeric/1000000000 FROM regional_budget_allocations WHERE user_id=$3 AND platform=$4 AND state='held'),0)<=" + window + "_limit_usd FROM usage_platform_quota_effective WHERE id=$1" if err = tx.QueryRowContext(ctx, query, quotaID, amount, user, platform).Scan(&enough); err != nil { return err } diff --git a/backend/internal/repository/regional_budget_window_rollover.go b/backend/internal/repository/regional_budget_window_rollover.go index 449d9addf132..de0852ec8e6d 100644 --- a/backend/internal/repository/regional_budget_window_rollover.go +++ b/backend/internal/repository/regional_budget_window_rollover.go @@ -3,7 +3,6 @@ package repository import ( "context" "database/sql" - "github.com/Wei-Shaw/sub2api/internal/pkg/timezone" "github.com/Wei-Shaw/sub2api/internal/service" "strings" "time" @@ -59,17 +58,20 @@ func regionalSubscriptionWindows(ctx context.Context, tx *sql.Tx, id int64, admi } func regionalPlatformWindows(ctx context.Context, tx *sql.Tx, user int64, platform string, admitted time.Time) error { var id int64 - var starts [3]sql.NullTime - err := tx.QueryRowContext(ctx, "SELECT id,daily_window_start,weekly_window_start,monthly_window_start FROM user_platform_quotas WHERE user_id=$1 AND platform=$2 AND deleted_at IS NULL FOR UPDATE", user, platform).Scan(&id, &starts[0], &starts[1], &starts[2]) + err := tx.QueryRowContext(ctx, "SELECT id FROM user_platform_quotas WHERE user_id=$1 AND platform=$2 AND deleted_at IS NULL FOR UPDATE", user, platform).Scan(&id) if err == sql.ErrNoRows { return nil } if err != nil { return err } - desired := [3]time.Time{timezone.StartOfDay(admitted), timezone.StartOfWeek(admitted), admitted} - if starts[2].Valid && admitted.Before(starts[2].Time.Add(30*24*time.Hour)) { - desired[2] = starts[2].Time + // Ordinary delivery may already have advanced its durable projection while + // the legacy baseline still has an expired anchor. Join that effective + // window before admission checks; inventing a newer baseline would hide + // current ordinary usage. The shared helper also retains exact reset cutoffs. + desired, err := effectivePlatformQuotaWindows(ctx, tx, id, admitted) + if err != nil { + return err } for i, w := range []string{"daily", "weekly", "monthly"} { if err := advanceRegionalWindow(ctx, tx, "user_platform_quotas", w+"_usage_usd", w+"_window_start", id, desired[i]); err != nil { diff --git a/backend/internal/repository/usage_billing_auth_recovery_integration_test.go b/backend/internal/repository/usage_billing_auth_recovery_integration_test.go new file mode 100644 index 000000000000..9eefd12382d0 --- /dev/null +++ b/backend/internal/repository/usage_billing_auth_recovery_integration_test.go @@ -0,0 +1,71 @@ +//go:build integration + +package repository + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +type failingCheckedBillingAuth struct{ fail atomic.Bool } + +func (*failingCheckedBillingAuth) InvalidateAuthCacheByKey(context.Context, string) { + panic("unchecked invalidation cannot confirm durable billing") +} +func (*failingCheckedBillingAuth) InvalidateAuthCacheByUserID(context.Context, int64) { + panic("unchecked invalidation cannot confirm durable billing") +} +func (*failingCheckedBillingAuth) InvalidateAuthCacheByGroupID(context.Context, int64) { + panic("unchecked invalidation cannot confirm durable billing") +} +func (a *failingCheckedBillingAuth) InvalidateAuthCacheByUserIDChecked(context.Context, int64) error { + if a.fail.Load() { + return errors.New("injected checked broadcast failure") + } + return nil +} + +func TestUsageBillingQuotaExhaustionAuthFailureRetriesWithoutSecondDebit(t *testing.T) { + ctx := context.Background() + repo, cmd, usage, _ := durableBillingFixture(t) + db := repo.db + _, err := db.ExecContext(ctx, "UPDATE api_keys SET quota=1 WHERE id=$1", cmd.APIKeyID) + require.NoError(t, err) + require.NoError(t, repo.StageUsageBilling(ctx, &cmd, &usage)) + result, err := repo.ApplyStagedUsageBilling(ctx, service.UsageBillingEventKey{RequestID: cmd.RequestID, APIKeyID: cmd.APIKeyID}) + require.NoError(t, err) + require.True(t, result.Result.APIKeyQuotaExhausted) + auth := &failingCheckedBillingAuth{} + auth.fail.Store(true) + worker := service.ProvideUsageRecordWorkerPool(&config.Config{}, repo, nil, auth) + t.Cleanup(worker.Stop) + var state, failure string + require.Eventually(t, func() bool { + err := db.QueryRowContext(ctx, `SELECT state,COALESCE(failure_code,'') FROM usage_billing_events WHERE request_id=$1 AND api_key_id=$2`, cmd.RequestID, cmd.APIKeyID).Scan(&state, &failure) + return err == nil && state == "applied" && failure == "cache_invalidation_failed" + }, 10*time.Second, 10*time.Millisecond) + auth.fail.Store(false) + // Advance the synthetic retry clock after the failure receipt is durable. + _, err = db.ExecContext(ctx, "UPDATE usage_billing_events SET next_attempt_at=NOW() WHERE request_id=$1 AND api_key_id=$2", cmd.RequestID, cmd.APIKeyID) + require.NoError(t, err) + require.Eventually(t, func() bool { + err := db.QueryRowContext(ctx, "SELECT state FROM usage_billing_events WHERE request_id=$1 AND api_key_id=$2", cmd.RequestID, cmd.APIKeyID).Scan(&state) + return err == nil && state == "complete" + }, 10*time.Second, 10*time.Millisecond) + var balance, quota, status string + var canonicalRows int + require.NoError(t, db.QueryRowContext(ctx, "SELECT balance::text FROM users WHERE id=$1", cmd.UserID).Scan(&balance)) + require.NoError(t, db.QueryRowContext(ctx, "SELECT quota_used::text,status FROM api_keys WHERE id=$1", cmd.APIKeyID).Scan("a, &status)) + require.NoError(t, db.QueryRowContext(ctx, "SELECT count(*) FROM usage_logs WHERE request_id=$1 AND api_key_id=$2", cmd.RequestID, cmd.APIKeyID).Scan(&canonicalRows)) + require.Equal(t, "99.00000000", balance) + require.Equal(t, "1.00000000", quota) + require.Equal(t, service.StatusAPIKeyQuotaExhausted, status) + require.Equal(t, 1, canonicalRows) +} diff --git a/backend/internal/repository/usage_billing_delivery.go b/backend/internal/repository/usage_billing_delivery.go new file mode 100644 index 000000000000..a4a6f58890f8 --- /dev/null +++ b/backend/internal/repository/usage_billing_delivery.go @@ -0,0 +1,235 @@ +package repository + +import ( + "context" + "encoding/json" + "errors" + "time" + + "github.com/Wei-Shaw/sub2api/internal/service" +) + +var _ service.DurableUsageBillingRepository = (*usageBillingRepository)(nil) + +func (r *usageBillingRepository) StageUsageBilling(ctx context.Context, cmd *service.UsageBillingCommand, usage *service.UsageLog) error { + if cmd == nil || usage == nil || cmd.RequestID == "" || cmd.APIKeyID <= 0 || cmd.UserID <= 0 || + usage.RequestID != cmd.RequestID || usage.APIKeyID != cmd.APIKeyID || usage.UserID != cmd.UserID || usage.AccountID != cmd.AccountID { + return service.ErrUsageBillingRequestConflict + } + cmd.Normalize() + e := service.UsageBillingEvent{Version: 1, Command: *cmd, Usage: *usage} + e.Usage.User, e.Usage.APIKey, e.Usage.Account, e.Usage.Group, e.Usage.Subscription = nil, nil, nil, nil, nil + if e.Usage.CreatedAt.IsZero() { + e.Usage.CreatedAt = time.Now().UTC() + } + raw, err := json.Marshal(e) + if err != nil { + return err + } + command, err := json.Marshal(e.Command) + if err != nil { + return err + } + tx, err := r.db.BeginTx(ctx, nil) + if err != nil { + return err + } + defer func() { _ = tx.Rollback() }() + _, err = tx.ExecContext(ctx, `INSERT INTO usage_billing_events + (request_id,api_key_id,request_fingerprint,schema_version,payload,observed_at) + VALUES($1,$2,$3,1,$4::jsonb,$5) ON CONFLICT(request_id,api_key_id) DO NOTHING`, + cmd.RequestID, cmd.APIKeyID, cmd.RequestFingerprint, string(raw), e.Usage.CreatedAt) + if err != nil { + return err + } + // Keep the first immutable usage snapshot. Repeated delivery may observe a + // later wall clock, but must never change the command, owner or monetary basis. + var same bool + err = tx.QueryRowContext(ctx, `SELECT schema_version=1 AND request_fingerprint=$3 AND ((payload->'command')-'PlatformUsageAt')=($4::jsonb-'PlatformUsageAt') + FROM usage_billing_events WHERE request_id=$1 AND api_key_id=$2`, cmd.RequestID, cmd.APIKeyID, cmd.RequestFingerprint, string(command)).Scan(&same) + if err != nil { + return err + } + if !same { + return service.ErrUsageBillingRequestConflict + } + if err = tx.Commit(); err == nil { + return nil + } + // A lost COMMIT response is not proof of rollback. Check a fresh connection + // before returning; do not re-run either the model or a monetary operation. + check, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + var stored bool + checkErr := r.db.QueryRowContext(check, `SELECT schema_version=1 AND request_fingerprint=$3 AND ((payload->'command')-'PlatformUsageAt')=($4::jsonb-'PlatformUsageAt') + FROM usage_billing_events WHERE request_id=$1 AND api_key_id=$2`, cmd.RequestID, cmd.APIKeyID, cmd.RequestFingerprint, string(command)).Scan(&stored) + if checkErr == nil && stored { + return nil + } + return errors.Join(service.ErrUsageBillingOutcomeUnknown, err, checkErr) +} + +func validateUsageBillingEvent(e *service.UsageBillingEvent, key service.UsageBillingEventKey, fingerprint string) error { + if e.Version != 1 { + return service.ErrUsageBillingEventVersion + } + c, u := &e.Command, &e.Usage + if c.RequestID != key.RequestID || c.APIKeyID != key.APIKeyID || c.RequestFingerprint != fingerprint || + u.RequestID != c.RequestID || u.APIKeyID != c.APIKeyID || u.UserID != c.UserID || u.AccountID != c.AccountID || + u.CreatedAt.IsZero() || u.User != nil || u.APIKey != nil || u.Account != nil || u.Group != nil || u.Subscription != nil { + return service.ErrUsageBillingRequestConflict + } + return nil +} + +func (r *usageBillingRepository) ApplyStagedUsageBilling(ctx context.Context, key service.UsageBillingEventKey) (_ *service.UsageBillingRecovery, err error) { + defer func() { + if err == nil { + return + } + code, state := "apply_failed", "pending" + if errors.Is(err, service.ErrUsageBillingRequestConflict) || errors.Is(err, service.ErrUsageBillingEventVersion) { + code, state = "identity_or_version_conflict", "blocked" + } + retry, cancel := context.WithTimeout(context.WithoutCancel(ctx), time.Second) + defer cancel() + // No monetary mutation here. Only schedule a still-pending event; never + // overwrite an applied/complete result after an ambiguous commit. + _, _ = r.db.ExecContext(retry, `UPDATE usage_billing_events SET state=$3,failure_code=$4, + attempts=attempts+1,next_attempt_at=NOW()+INTERVAL '30 seconds' + WHERE request_id=$1 AND api_key_id=$2 AND state='pending'`, key.RequestID, key.APIKeyID, state, code) + }() + tx, err := r.db.BeginTx(ctx, nil) + if err != nil { + return nil, err + } + defer func() { _ = tx.Rollback() }() + var raw, resultRaw []byte + var state, fingerprint string + var version int + if err = tx.QueryRowContext(ctx, `SELECT schema_version,payload,state,result,request_fingerprint + FROM usage_billing_events WHERE request_id=$1 AND api_key_id=$2 FOR UPDATE`, key.RequestID, key.APIKeyID). + Scan(&version, &raw, &state, &resultRaw, &fingerprint); err != nil { + return nil, err + } + recovered := &service.UsageBillingRecovery{} + if err = json.Unmarshal(raw, &recovered.Event); err != nil { + return nil, err + } + if version != recovered.Event.Version { + return nil, service.ErrUsageBillingEventVersion + } + if err = validateUsageBillingEvent(&recovered.Event, key, fingerprint); err != nil { + return nil, err + } + if state == "applied" || state == "complete" || state == "legacy" { + if len(resultRaw) > 0 { + if err = json.Unmarshal(resultRaw, &recovered.Result); err != nil { + return nil, err + } + } + recovered.Result.Applied = false + recovered.NeedsCache = state == "applied" + return recovered, nil + } + if state != "pending" { + return nil, service.ErrUsageBillingRequestConflict + } + cmd := &recovered.Event.Command + applied, err := r.claimUsageBillingKey(ctx, tx, cmd) + if err != nil { + return nil, err + } + if !applied { + // An old hot/archive claim may predate this journal. Do not reconstruct + // or alter historical usage or reapply its post-commit monetary effects. + _, err = tx.ExecContext(ctx, `UPDATE usage_billing_events SET state='legacy',completed_at=NOW(),result='{}' + WHERE request_id=$1 AND api_key_id=$2`, key.RequestID, key.APIKeyID) + if err != nil { + return nil, err + } + return recovered, tx.Commit() + } + var existing bool + if err = tx.QueryRowContext(ctx, `SELECT EXISTS(SELECT 1 FROM usage_logs WHERE request_id=$1 AND api_key_id=$2)`, key.RequestID, key.APIKeyID).Scan(&existing); err != nil { + return nil, err + } + if existing { + // A pre-existing usage row without our committed claim is not authority + // to back-charge an old failed/zero-cost request automatically. + return nil, service.ErrUsageBillingRequestConflict + } + recovered.Result.Applied = true + if err = r.applyUsageBillingEffects(ctx, tx, cmd, &recovered.Result); err != nil { + return nil, err + } + if err = insertUsageBillingQuote(ctx, tx, cmd); err != nil { + return nil, err + } + if err = execUsageLogInsertNoResult(ctx, tx, prepareUsageLogInsert(&recovered.Event.Usage)); err != nil { + return nil, err + } + encodedResult, err := json.Marshal(recovered.Result) + if err != nil { + return nil, err + } + _, err = tx.ExecContext(ctx, `UPDATE usage_billing_events SET state='applied',result=$3::jsonb, + applied_at=NOW(),failure_code=NULL,next_attempt_at=NOW() WHERE request_id=$1 AND api_key_id=$2`, + key.RequestID, key.APIKeyID, string(encodedResult)) + if err != nil { + return nil, err + } + if err = tx.Commit(); err == nil { + recovered.NeedsCache = true + return recovered, nil + } + // Inspect durable evidence before any replay. A later retry takes the row + // lock again, so even a transaction still resolving here cannot double debit. + check, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + var committedState string + var committedResult []byte + checkErr := r.db.QueryRowContext(check, `SELECT state,result FROM usage_billing_events + WHERE request_id=$1 AND api_key_id=$2`, key.RequestID, key.APIKeyID).Scan(&committedState, &committedResult) + if checkErr == nil && (committedState == "applied" || committedState == "complete") { + if checkErr = json.Unmarshal(committedResult, &recovered.Result); checkErr == nil { + recovered.NeedsCache = committedState == "applied" + return recovered, nil + } + } + return nil, errors.Join(service.ErrUsageBillingOutcomeUnknown, err, checkErr) +} + +func (r *usageBillingRepository) PendingUsageBilling(ctx context.Context, limit int) ([]service.UsageBillingEventKey, error) { + if limit < 1 || limit > 256 { + return nil, errors.New("invalid usage recovery batch size") + } + rows, err := r.db.QueryContext(ctx, `SELECT request_id,api_key_id FROM usage_billing_events + WHERE state IN ('pending','applied') AND next_attempt_at<=NOW() + ORDER BY CASE WHEN state='pending' THEN 0 ELSE 1 END,next_attempt_at,staged_at LIMIT $1`, limit) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + var keys []service.UsageBillingEventKey + for rows.Next() { + var key service.UsageBillingEventKey + if err := rows.Scan(&key.RequestID, &key.APIKeyID); err != nil { + return nil, err + } + keys = append(keys, key) + } + return keys, rows.Err() +} + +func (r *usageBillingRepository) AcknowledgeUsageBillingCache(ctx context.Context, key service.UsageBillingEventKey) error { + _, err := r.db.ExecContext(ctx, `UPDATE usage_billing_events SET state='complete',completed_at=NOW(),failure_code=NULL + WHERE request_id=$1 AND api_key_id=$2 AND state='applied'`, key.RequestID, key.APIKeyID) + return err +} + +func (r *usageBillingRepository) DeferUsageBillingCache(ctx context.Context, key service.UsageBillingEventKey) error { + _, err := r.db.ExecContext(ctx, `UPDATE usage_billing_events SET next_attempt_at=NOW()+INTERVAL '30 seconds', + attempts=attempts+1,failure_code='cache_invalidation_failed' WHERE request_id=$1 AND api_key_id=$2 AND state='applied'`, key.RequestID, key.APIKeyID) + return err +} diff --git a/backend/internal/repository/usage_billing_delivery_integration_test.go b/backend/internal/repository/usage_billing_delivery_integration_test.go new file mode 100644 index 000000000000..5d68ea8304fa --- /dev/null +++ b/backend/internal/repository/usage_billing_delivery_integration_test.go @@ -0,0 +1,130 @@ +//go:build integration + +package repository + +import ( + "context" + "fmt" + "strings" + "sync" + "testing" + "time" + + dbent "github.com/Wei-Shaw/sub2api/ent" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +func durableBillingFixture(t *testing.T) (*usageBillingRepository, service.UsageBillingCommand, service.UsageLog, *dbent.Client) { + t.Helper() + db, client := newBillingAuditDatabase(t) + user := mustCreateUser(t, client, &service.User{Email: uuid.NewString() + "@example.com", PasswordHash: "hash", Balance: 100}) + key := mustCreateApiKey(t, client, &service.APIKey{UserID: user.ID, Key: "sk-mock-" + uuid.NewString(), Name: "delivery", Quota: 100}) + account := mustCreateAccount(t, client, &service.Account{Name: "delivery-" + uuid.NewString(), Type: service.AccountTypeAPIKey}) + cmd := service.UsageBillingCommand{RequestID: uuid.NewString(), UserID: user.ID, APIKeyID: key.ID, AccountID: account.ID, BalanceCost: 1, APIKeyQuotaCost: 1} + usage := service.UsageLog{RequestID: cmd.RequestID, UserID: cmd.UserID, APIKeyID: cmd.APIKeyID, AccountID: cmd.AccountID, + Model: "mock-model", InputTokens: 10, ActualCost: 1, TotalCost: 1, CreatedAt: time.Now().UTC(), + // Relations must never enter the persistent billing envelope. + User: &service.User{ID: user.ID, PasswordHash: "must-not-be-persisted"}, + APIKey: &service.APIKey{Key: "must-not-be-persisted"}, + } + return &usageBillingRepository{db: db}, cmd, usage, client +} + +func TestUsageBillingDeliveryCrashBeforeExecutionAndTwoWriters(t *testing.T) { + ctx := context.Background() + repo, cmd, usage, _ := durableBillingFixture(t) + db := repo.db + require.NoError(t, repo.StageUsageBilling(ctx, &cmd, &usage)) + // No execution in the first process. Only the staged event survives. + var payload string + require.NoError(t, db.QueryRowContext(ctx, "SELECT payload::text FROM usage_billing_events WHERE request_id=$1 AND api_key_id=$2", cmd.RequestID, cmd.APIKeyID).Scan(&payload)) + require.NotContains(t, payload, "must-not-be-persisted") + var wg sync.WaitGroup + errs := make(chan error, 2) + for i := 0; i < 2; i++ { + wg.Add(1) + go func() { + defer wg.Done() + freshProcess := &usageBillingRepository{db: db} + _, err := freshProcess.ApplyStagedUsageBilling(ctx, service.UsageBillingEventKey{RequestID: cmd.RequestID, APIKeyID: cmd.APIKeyID}) + errs <- err + }() + } + wg.Wait() + close(errs) + for err := range errs { + require.NoError(t, err) + } + var balance, quota string + var logs int + require.NoError(t, db.QueryRowContext(ctx, "SELECT balance::text FROM users WHERE id=$1", cmd.UserID).Scan(&balance)) + require.NoError(t, db.QueryRowContext(ctx, "SELECT quota_used::text FROM api_keys WHERE id=$1", cmd.APIKeyID).Scan("a)) + require.NoError(t, db.QueryRowContext(ctx, "SELECT count(*) FROM usage_logs WHERE request_id=$1 AND api_key_id=$2", cmd.RequestID, cmd.APIKeyID).Scan(&logs)) + require.Equal(t, "99.00000000", balance) + require.Equal(t, "1.00000000", quota) + require.Equal(t, 1, logs) + // A restart between cache invalidation and acknowledgement repeats only + // invalidation, not the amount. Different command ownership is rejected. + wrong := cmd + wrong.UserID++ + require.ErrorIs(t, repo.StageUsageBilling(ctx, &wrong, &usage), service.ErrUsageBillingRequestConflict) +} + +func TestUsageBillingDeliveryLogFailureRollsBackMoneyAndRecovers(t *testing.T) { + ctx := context.Background() + repo, cmd, usage, _ := durableBillingFixture(t) + db := repo.db + require.NoError(t, repo.StageUsageBilling(ctx, &cmd, &usage)) + name := "billing_fail_" + strings.ReplaceAll(uuid.NewString(), "-", "") + // A real PostgreSQL failure between monetary UPDATEs and canonical INSERT. + _, err := db.ExecContext(ctx, fmt.Sprintf(`CREATE FUNCTION %s() RETURNS trigger LANGUAGE plpgsql AS $$ + BEGIN IF NEW.request_id='%s' THEN RAISE EXCEPTION 'injected canonical failure'; END IF; RETURN NEW; END; $$; + CREATE TRIGGER %s BEFORE INSERT ON usage_logs FOR EACH ROW EXECUTE FUNCTION %s()`, name, cmd.RequestID, name, name)) + require.NoError(t, err) + cleanup := func() { + _, _ = db.ExecContext(ctx, fmt.Sprintf("DROP TRIGGER IF EXISTS %s ON usage_logs; DROP FUNCTION IF EXISTS %s()", name, name)) + } + t.Cleanup(cleanup) + key := service.UsageBillingEventKey{RequestID: cmd.RequestID, APIKeyID: cmd.APIKeyID} + _, err = repo.ApplyStagedUsageBilling(ctx, key) + require.Error(t, err) + var balance, state string + var claims int + require.NoError(t, db.QueryRowContext(ctx, "SELECT balance::text FROM users WHERE id=$1", cmd.UserID).Scan(&balance)) + require.NoError(t, db.QueryRowContext(ctx, "SELECT state FROM usage_billing_events WHERE request_id=$1 AND api_key_id=$2", cmd.RequestID, cmd.APIKeyID).Scan(&state)) + require.NoError(t, db.QueryRowContext(ctx, "SELECT count(*) FROM usage_billing_dedup WHERE request_id=$1 AND api_key_id=$2", cmd.RequestID, cmd.APIKeyID).Scan(&claims)) + require.Equal(t, "100.00000000", balance) + require.Equal(t, "pending", state) + require.Zero(t, claims) + cleanup() + result, err := (&usageBillingRepository{db: db}).ApplyStagedUsageBilling(ctx, key) + require.NoError(t, err) + require.True(t, result.Result.Applied) + require.NoError(t, repo.AcknowledgeUsageBillingCache(ctx, key)) + result, err = repo.ApplyStagedUsageBilling(ctx, key) + require.NoError(t, err) + require.False(t, result.Result.Applied) + require.False(t, result.NeedsCache) +} + +func TestUsageBillingDeliveryArchiveClaimNeverBackChargesHistory(t *testing.T) { + ctx := context.Background() + repo, cmd, usage, _ := durableBillingFixture(t) + db := repo.db + cmd.Normalize() + _, err := db.ExecContext(ctx, `INSERT INTO usage_billing_dedup_archive + (request_id,api_key_id,request_fingerprint,created_at) VALUES($1,$2,$3,NOW())`, cmd.RequestID, cmd.APIKeyID, cmd.RequestFingerprint) + require.NoError(t, err) + require.NoError(t, repo.StageUsageBilling(ctx, &cmd, &usage)) + r, err := repo.ApplyStagedUsageBilling(ctx, service.UsageBillingEventKey{RequestID: cmd.RequestID, APIKeyID: cmd.APIKeyID}) + require.NoError(t, err) + require.False(t, r.Result.Applied) + var balance string + var logs int + require.NoError(t, db.QueryRowContext(ctx, "SELECT balance::text FROM users WHERE id=$1", cmd.UserID).Scan(&balance)) + require.NoError(t, db.QueryRowContext(ctx, "SELECT count(*) FROM usage_logs WHERE request_id=$1 AND api_key_id=$2", cmd.RequestID, cmd.APIKeyID).Scan(&logs)) + require.Equal(t, "100.00000000", balance) + require.Zero(t, logs, "no fabricated historical canonical usage") +} diff --git a/backend/internal/repository/usage_billing_delivery_unit_test.go b/backend/internal/repository/usage_billing_delivery_unit_test.go new file mode 100644 index 000000000000..bc20d4f9bebe --- /dev/null +++ b/backend/internal/repository/usage_billing_delivery_unit_test.go @@ -0,0 +1,61 @@ +//go:build unit + +package repository + +import ( + "context" + "encoding/json" + "errors" + "testing" + "time" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func TestUsageBillingDeliveryCommitUnknownChecksReceiptBeforeReplay(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + cmd := service.UsageBillingCommand{RequestID: "unknown", APIKeyID: 2, UserID: 1, AccountID: 3, BalanceCost: 1} + cmd.Normalize() + e := service.UsageBillingEvent{Version: 1, Command: cmd, Usage: service.UsageLog{ + RequestID: "unknown", APIKeyID: 2, UserID: 1, AccountID: 3, ActualCost: 1, Model: "mock", CreatedAt: time.Now(), + }} + payload, err := json.Marshal(e) + require.NoError(t, err) + result := service.UsageBillingApplyResult{Applied: true} + storedResult, err := json.Marshal(result) + require.NoError(t, err) + mock.ExpectBegin() + mock.ExpectQuery("SELECT schema_version,payload,state,result,request_fingerprint").WillReturnRows( + sqlmock.NewRows([]string{"schema_version", "payload", "state", "result", "request_fingerprint"}).AddRow(1, payload, "pending", nil, cmd.RequestFingerprint)) + mock.ExpectQuery("INSERT INTO usage_billing_dedup").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1)) + mock.ExpectQuery("SELECT request_fingerprint.*FROM usage_billing_dedup_archive").WillReturnRows(sqlmock.NewRows([]string{"request_fingerprint"})) + mock.ExpectQuery("SELECT EXISTS.*FROM usage_logs").WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(false)) + mock.ExpectQuery(conditionalBalanceDeductSQL).WillReturnRows(sqlmock.NewRows([]string{"balance"}).AddRow(99)) + mock.ExpectExec("INSERT INTO usage_logs").WillReturnResult(sqlmock.NewResult(1, 1)) + mock.ExpectExec("UPDATE usage_billing_events SET state='applied'").WillReturnResult(sqlmock.NewResult(0, 1)) + mock.ExpectCommit().WillReturnError(errors.New("lost commit response")) + mock.ExpectQuery("SELECT state,result FROM usage_billing_events").WillReturnRows(sqlmock.NewRows([]string{"state", "result"}).AddRow("applied", storedResult)) + recovered, err := (&usageBillingRepository{db: db}).ApplyStagedUsageBilling(context.Background(), service.UsageBillingEventKey{RequestID: "unknown", APIKeyID: 2}) + require.NoError(t, err) + require.True(t, recovered.NeedsCache) + require.NoError(t, mock.ExpectationsWereMet(), "only one money UPDATE is allowed") +} + +func TestUsageBillingStageCommitUnknownChecksStoredCommand(t *testing.T) { + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + cmd := service.UsageBillingCommand{RequestID: "stage-unknown", APIKeyID: 2, UserID: 1, AccountID: 3} + usage := service.UsageLog{RequestID: cmd.RequestID, APIKeyID: 2, UserID: 1, AccountID: 3, CreatedAt: time.Now()} + mock.ExpectBegin() + mock.ExpectExec("INSERT INTO usage_billing_events").WillReturnResult(sqlmock.NewResult(0, 1)) + mock.ExpectQuery("SELECT schema_version=1 AND request_fingerprint").WillReturnRows(sqlmock.NewRows([]string{"same"}).AddRow(true)) + mock.ExpectCommit().WillReturnError(errors.New("stage commit response lost")) + mock.ExpectQuery("SELECT schema_version=1 AND request_fingerprint").WillReturnRows(sqlmock.NewRows([]string{"same"}).AddRow(true)) + require.NoError(t, (&usageBillingRepository{db: db}).StageUsageBilling(context.Background(), &cmd, &usage)) + require.NoError(t, mock.ExpectationsWereMet(), "stage must not execute money effects") +} diff --git a/backend/internal/repository/usage_billing_platform_integration_test.go b/backend/internal/repository/usage_billing_platform_integration_test.go new file mode 100644 index 000000000000..254438157374 --- /dev/null +++ b/backend/internal/repository/usage_billing_platform_integration_test.go @@ -0,0 +1,91 @@ +//go:build integration + +package repository + +import ( + "context" + "fmt" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/timezone" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +func TestUsageBillingPlatformDBFailureRecoveryAndLegacySnapshotIsolation(t *testing.T) { + ctx := context.Background() + repo, cmd, usage, client := durableBillingFixture(t) + db := repo.db + cmd.PlatformQuotaPlatform, cmd.PlatformQuotaCost, cmd.PlatformUsageAt = "openai", 1, usage.CreatedAt + var quotaID int64 + day, week := timezone.StartOfDay(usage.CreatedAt), timezone.StartOfWeek(usage.CreatedAt) + require.NoError(t, db.QueryRowContext(ctx, `INSERT INTO user_platform_quotas + (user_id,platform,daily_limit_usd,daily_usage_usd,weekly_usage_usd,monthly_usage_usd,daily_window_start,weekly_window_start,monthly_window_start) + VALUES($1,'openai',100,2,2,2,$2,$3,$4) RETURNING id`, cmd.UserID, day, week, usage.CreatedAt).Scan("aID)) + require.NoError(t, repo.StageUsageBilling(ctx, &cmd, &usage)) + name := "quota_fail_" + strings.ReplaceAll(uuid.NewString(), "-", "") + _, err := db.ExecContext(ctx, fmt.Sprintf(`CREATE FUNCTION %s() RETURNS trigger LANGUAGE plpgsql AS $$ + BEGIN IF NEW.quota_id=%d THEN RAISE EXCEPTION 'injected quota DB failure'; END IF; RETURN NEW; END; $$; + CREATE TRIGGER %s BEFORE INSERT ON usage_platform_quota_totals FOR EACH ROW EXECUTE FUNCTION %s()`, name, quotaID, name, name)) + require.NoError(t, err) + cleanup := func() { + _, _ = db.ExecContext(ctx, fmt.Sprintf("DROP TRIGGER IF EXISTS %s ON usage_platform_quota_totals; DROP FUNCTION IF EXISTS %s()", name, name)) + } + t.Cleanup(cleanup) + key := service.UsageBillingEventKey{RequestID: cmd.RequestID, APIKeyID: cmd.APIKeyID} + _, err = repo.ApplyStagedUsageBilling(ctx, key) + require.Error(t, err) + var balance string + require.NoError(t, db.QueryRowContext(ctx, "SELECT balance::text FROM users WHERE id=$1", cmd.UserID).Scan(&balance)) + require.Equal(t, "100.00000000", balance, "quota failure rolls back the debit, not an untracked post-commit goroutine") + cleanup() + for attempt := 0; attempt < 2; attempt++ { + _, err = repo.ApplyStagedUsageBilling(ctx, key) + require.NoError(t, err) + } + quotaRepo := NewUserPlatformQuotaRepository(client) + rec, err := quotaRepo.GetByUserPlatform(ctx, cmd.UserID, "openai") + require.NoError(t, err) + require.Equal(t, 3.0, rec.DailyUsageUSD) + // A late legacy flusher snapshot changes only the old baseline (2 -> 2.5). + // It cannot replace the new durable delta of 1, under either flusher setting. + require.NoError(t, quotaRepo.BatchSnapshotUsage(ctx, []UserPlatformQuotaSnapshot{{UserID: cmd.UserID, Platform: "openai", + DailyUsageUSD: 2.5, WeeklyUsageUSD: 2.5, MonthlyUsageUSD: 2.5, + DailyWindowStart: day, WeeklyWindowStart: week, MonthlyWindowStart: usage.CreatedAt}}, time.Now())) + rec, err = quotaRepo.GetByUserPlatform(ctx, cmd.UserID, "openai") + require.NoError(t, err) + require.Equal(t, 3.5, rec.DailyUsageUSD) + require.NoError(t, quotaRepo.ResetExpiredWindow(ctx, cmd.UserID, "openai", "daily", day)) + rec, err = quotaRepo.GetByUserPlatform(ctx, cmd.UserID, "openai") + require.NoError(t, err) + require.Zero(t, rec.DailyUsageUSD) + require.Equal(t, 3.5, rec.WeeklyUsageUSD, "explicit daily reset must not reset another window") +} + +func TestUsageBillingPlatformLateRetryDoesNotRewindNewWindow(t *testing.T) { + ctx := context.Background() + repo, first, usage, client := durableBillingFixture(t) + db := repo.db + first.PlatformQuotaPlatform, first.PlatformQuotaCost = "openai", 1 + first.PlatformUsageAt = usage.CreatedAt.Add(-48 * time.Hour) + _, err := db.ExecContext(ctx, "INSERT INTO user_platform_quotas(user_id,platform,daily_limit_usd) VALUES($1,'openai',100)", first.UserID) + require.NoError(t, err) + second := first + second.RequestID = uuid.NewString() + second.PlatformUsageAt = usage.CreatedAt + // Apply the current event first, then a recovered older observation. + for _, cmd := range []*service.UsageBillingCommand{&second, &first} { + log := usage + log.RequestID = cmd.RequestID + require.NoError(t, repo.StageUsageBilling(ctx, cmd, &log)) + _, err = repo.ApplyStagedUsageBilling(ctx, service.UsageBillingEventKey{RequestID: cmd.RequestID, APIKeyID: cmd.APIKeyID}) + require.NoError(t, err) + } + rec, err := NewUserPlatformQuotaRepository(client).GetByUserPlatform(ctx, first.UserID, "openai") + require.NoError(t, err) + require.Equal(t, 1.0, rec.DailyUsageUSD) + require.True(t, rec.DailyWindowStart.Equal(timezone.StartOfDay(usage.CreatedAt))) +} diff --git a/backend/internal/repository/usage_billing_platform_quota.go b/backend/internal/repository/usage_billing_platform_quota.go new file mode 100644 index 000000000000..825320884dbe --- /dev/null +++ b/backend/internal/repository/usage_billing_platform_quota.go @@ -0,0 +1,78 @@ +package repository + +import ( + "context" + "database/sql" + "errors" + "time" + + "github.com/Wei-Shaw/sub2api/internal/domain" + "github.com/Wei-Shaw/sub2api/internal/pkg/timezone" + "github.com/Wei-Shaw/sub2api/internal/service" +) + +func incrementUsageBillingPlatformQuota(ctx context.Context, tx *sql.Tx, cmd *service.UsageBillingCommand) error { + if !domain.IsConcretePlatform(cmd.PlatformQuotaPlatform) || cmd.PlatformUsageAt.IsZero() { + return service.ErrUsageBillingRequestConflict + } + var id int64 + err := tx.QueryRowContext(ctx, `SELECT id FROM user_platform_quotas + WHERE user_id=$1 AND platform=$2 AND deleted_at IS NULL + AND (daily_limit_usd IS NOT NULL OR weekly_limit_usd IS NOT NULL OR monthly_limit_usd IS NOT NULL) FOR UPDATE`, + cmd.UserID, cmd.PlatformQuotaPlatform).Scan(&id) + if errors.Is(err, sql.ErrNoRows) { + return nil // absent quota is unlimited; never create limits implicitly + } + if err != nil { + return err + } + at := cmd.PlatformUsageAt + starts, err := effectivePlatformQuotaWindows(ctx, tx, id, at) + if err != nil { + return err + } + _, err = tx.ExecContext(ctx, `INSERT INTO usage_platform_quota_totals + (quota_id,daily_usage_usd,weekly_usage_usd,monthly_usage_usd,daily_window_start,weekly_window_start,monthly_window_start,legacy_daily_anchor,legacy_weekly_anchor,legacy_monthly_anchor) + VALUES($1,CASE WHEN $3::timestamptz>$6::timestamptz THEN 0 ELSE $2::numeric END,CASE WHEN $4::timestamptz>$6::timestamptz THEN 0 ELSE $2::numeric END, + CASE WHEN $5::timestamptz>$6::timestamptz THEN 0 ELSE $2::numeric END,$3,$4,$5,$3,$4,$5) + ON CONFLICT(quota_id) DO UPDATE SET + daily_usage_usd=CASE WHEN usage_platform_quota_totals.daily_window_start>EXCLUDED.daily_window_start THEN usage_platform_quota_totals.daily_usage_usd + WHEN usage_platform_quota_totals.daily_window_startEXCLUDED.weekly_window_start THEN usage_platform_quota_totals.weekly_usage_usd + WHEN usage_platform_quota_totals.weekly_window_startEXCLUDED.monthly_window_start THEN usage_platform_quota_totals.monthly_usage_usd + WHEN usage_platform_quota_totals.monthly_window_start= 2 { + result, err = applyOwnedBatchImageMoney(ctx, tx, cmd, operation, apply) + } else { + var owned bool + if err = tx.QueryRowContext(ctx, `SELECT EXISTS(SELECT 1 FROM batch_image_money_holds WHERE batch_id=$1 AND api_key_id=$2)`, cmd.BatchID, cmd.APIKeyID).Scan(&owned); err != nil { + return nil, err + } + if owned { + return nil, service.ErrBatchImageLegacyMoneyPrecision + } + result, err = apply(ctx, tx, cmd) + } if err != nil { return nil, err } @@ -210,6 +256,12 @@ func (r *usageBillingRepository) applyUsageBillingEffects(ctx context.Context, t } result.QuotaState = quotaState } + if cmd.PlatformQuotaCost > 0 && cmd.PlatformQuotaPlatform != "" { + if err := incrementUsageBillingPlatformQuota(ctx, tx, cmd); err != nil { + return err + } + result.PlatformQuotaApplied = true + } return nil } diff --git a/backend/internal/repository/usage_billing_reset_cutoff_integration_test.go b/backend/internal/repository/usage_billing_reset_cutoff_integration_test.go new file mode 100644 index 000000000000..b1b0cecac18a --- /dev/null +++ b/backend/internal/repository/usage_billing_reset_cutoff_integration_test.go @@ -0,0 +1,81 @@ +//go:build integration + +package repository + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/timezone" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +func TestUsageBillingNonBoundaryResetPreservesNewUsageAndRejectsLateWindowUsage(t *testing.T) { + for _, window := range []string{"daily", "weekly"} { + for _, existingTotals := range []bool{false, true} { + for _, lateFirst := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/totals=%t/late_first=%t", window, existingTotals, lateFirst), func(t *testing.T) { + ctx := context.Background() + repo, template, usage, client := durableBillingFixture(t) + db := repo.db + cutoff := time.Date(2026, 10, 21, 12, 34, 56, 123000000, timezone.Location()) + _, err := db.ExecContext(ctx, `INSERT INTO user_platform_quotas + (user_id,platform,daily_limit_usd,weekly_limit_usd,daily_usage_usd,weekly_usage_usd,daily_window_start,weekly_window_start) + VALUES($1,'openai',100,100,2,2,$2,$3)`, template.UserID, timezone.StartOfDay(cutoff), timezone.StartOfWeek(cutoff)) + require.NoError(t, err) + stage := func(at time.Time, amount float64) service.UsageBillingEventKey { + cmd, log := template, usage + cmd.RequestID, log.RequestID = uuid.NewString(), "" + cmd.RequestFingerprint = "" + cmd.BalanceCost, cmd.APIKeyQuotaCost = amount, amount + cmd.PlatformQuotaPlatform, cmd.PlatformQuotaCost, cmd.PlatformUsageAt = "openai", amount, at + log.RequestID, log.ActualCost, log.TotalCost, log.CreatedAt = cmd.RequestID, amount, amount, at + require.NoError(t, repo.StageUsageBilling(ctx, &cmd, &log)) + return service.UsageBillingEventKey{RequestID: cmd.RequestID, APIKeyID: cmd.APIKeyID} + } + if existingTotals { + _, err = repo.ApplyStagedUsageBilling(ctx, stage(cutoff.Add(-time.Hour), 0.5)) + require.NoError(t, err) + } + // Persist an observed request before the exact administrative reset, + // but leave its monetary execution pending across that reset. + late := stage(cutoff.Add(-time.Minute), 0.25) + quotaRepo := NewUserPlatformQuotaRepository(client) + require.NoError(t, quotaRepo.ResetExpiredWindow(ctx, template.UserID, "openai", window, cutoff)) + fresh := stage(cutoff.Add(time.Minute), 0.4) + order := []service.UsageBillingEventKey{fresh, late} + if lateFirst { + order = []service.UsageBillingEventKey{late, fresh} + } + for repeat := 0; repeat < 2; repeat++ { + for _, key := range order { + result, err := repo.ApplyStagedUsageBilling(ctx, key) + require.NoError(t, err) + require.Equal(t, repeat == 0, result.Result.Applied) + } + } + rec, err := quotaRepo.GetByUserPlatform(ctx, template.UserID, "openai") + require.NoError(t, err) + used, start := rec.DailyUsageUSD, rec.DailyWindowStart + if window == "weekly" { + used, start = rec.WeeklyUsageUSD, rec.WeeklyWindowStart + } + require.Equal(t, 0.4, used, "new post-reset usage counts; pre-reset late usage and replays do not") + require.NotNil(t, start) + require.True(t, start.Equal(cutoff), "retain the precise cutoff rather than rounding it to a natural boundary") + var balance string + require.NoError(t, db.QueryRowContext(ctx, "SELECT balance::text FROM users WHERE id=$1", template.UserID).Scan(&balance)) + want := "99.35000000" + if existingTotals { + want = "98.85000000" + } + require.Equal(t, want, balance, "the quota reset never refunds or repeats the monetary charge") + }) + } + } + } +} diff --git a/backend/internal/repository/user_platform_quota_repo.go b/backend/internal/repository/user_platform_quota_repo.go index 26f6e55062a0..3b67544a1016 100644 --- a/backend/internal/repository/user_platform_quota_repo.go +++ b/backend/internal/repository/user_platform_quota_repo.go @@ -134,42 +134,52 @@ func (r *userPlatformQuotaRepository) BulkInsertInitial(ctx context.Context, rec return err } -// GetByUserPlatform 通过 ent 查询单条配额(排除软删除)。未找到返回 (nil, nil)。 -func (r *userPlatformQuotaRepository) GetByUserPlatform(ctx context.Context, userID int64, platform string) (*UserPlatformQuotaRecord, error) { +// Read the legacy baseline plus the independent, transactionally applied new +// request deltas. Old flusher snapshots cannot overwrite the new projection. +const effectivePlatformQuotaColumns = `user_id,platform,daily_limit_usd,weekly_limit_usd,monthly_limit_usd, + daily_usage_usd,weekly_usage_usd,monthly_usage_usd,daily_window_start,weekly_window_start,monthly_window_start` + +func (r *userPlatformQuotaRepository) readEffective(ctx context.Context, userID int64, platform string) ([]UserPlatformQuotaRecord, error) { client := clientFromContext(ctx, r.client) - entity, err := client.UserPlatformQuota.Query(). - Where( - userplatformquota.UserIDEQ(userID), - userplatformquota.PlatformEQ(platform), - userplatformquota.DeletedAtIsNil(), - ). - Only(ctx) - if dbent.IsNotFound(err) { - return nil, nil + query := "SELECT " + effectivePlatformQuotaColumns + " FROM usage_platform_quota_effective WHERE user_id=$1 AND deleted_at IS NULL" + args := []any{userID} + if platform != "" { + query += " AND platform=$2" + args = append(args, platform) } + rows, err := client.QueryContext(ctx, query, args...) if err != nil { return nil, err } - return entQuotaToRecord(entity), nil + defer func() { _ = rows.Close() }() + out := make([]UserPlatformQuotaRecord, 0) + for rows.Next() { + var rec UserPlatformQuotaRecord + if err := rows.Scan(&rec.UserID, &rec.Platform, &rec.DailyLimitUSD, &rec.WeeklyLimitUSD, &rec.MonthlyLimitUSD, + &rec.DailyUsageUSD, &rec.WeeklyUsageUSD, &rec.MonthlyUsageUSD, &rec.DailyWindowStart, &rec.WeeklyWindowStart, &rec.MonthlyWindowStart); err != nil { + return nil, err + } + out = append(out, rec) + } + return out, rows.Err() } -// ListByUser 查询用户的所有平台配额记录(排除软删除)。 -func (r *userPlatformQuotaRepository) ListByUser(ctx context.Context, userID int64) ([]UserPlatformQuotaRecord, error) { - client := clientFromContext(ctx, r.client) - rows, err := client.UserPlatformQuota.Query(). - Where( - userplatformquota.UserIDEQ(userID), - userplatformquota.DeletedAtIsNil(), - ). - All(ctx) +func (r *userPlatformQuotaRepository) GetByUserPlatform(ctx context.Context, userID int64, platform string) (*UserPlatformQuotaRecord, error) { + if platform == "" { + return nil, nil + } + rows, err := r.readEffective(ctx, userID, platform) if err != nil { return nil, err } - out := make([]UserPlatformQuotaRecord, 0, len(rows)) - for _, e := range rows { - out = append(out, *entQuotaToRecord(e)) + if len(rows) == 0 { + return nil, nil } - return out, nil + return &rows[0], nil +} + +func (r *userPlatformQuotaRepository) ListByUser(ctx context.Context, userID int64) ([]UserPlatformQuotaRecord, error) { + return r.readEffective(ctx, userID, "") } // IncrementUsageWithReset 原子累加 cost 到 (user, platform) 三个窗口的 *_usage_usd。 @@ -226,31 +236,33 @@ func (r *userPlatformQuotaRepository) IncrementUsageWithReset(ctx context.Contex // // 未命中活跃记录时返回 ErrUserPlatformQuotaNotFound。 func (r *userPlatformQuotaRepository) ResetExpiredWindow(ctx context.Context, userID int64, platform string, window string, newStart time.Time) error { - client := clientFromContext(ctx, r.client) - upd := client.UserPlatformQuota.Update(). - Where( - userplatformquota.UserIDEQ(userID), - userplatformquota.PlatformEQ(platform), - userplatformquota.DeletedAtIsNil(), - ) switch window { - case "daily": - upd = upd.SetDailyUsageUsd(0).SetDailyWindowStart(newStart) - case "weekly": - upd = upd.SetWeeklyUsageUsd(0).SetWeeklyWindowStart(newStart) - case "monthly": - upd = upd.SetMonthlyUsageUsd(0).SetMonthlyWindowStart(newStart) + case "daily", "weekly", "monthly": default: return fmt.Errorf("unknown window %q", window) } - n, err := upd.Save(ctx) - if err != nil { + return r.withTx(ctx, func(txCtx context.Context, client *dbent.Client) error { + rows, err := client.QueryContext(txCtx, "UPDATE user_platform_quotas SET "+window+"_usage_usd=0,"+window+"_window_start=$3 WHERE user_id=$1 AND platform=$2 AND deleted_at IS NULL RETURNING id", userID, platform, newStart) + if err != nil { + return err + } + var id int64 + if !rows.Next() { + rowErr := rows.Err() + _ = rows.Close() + if rowErr != nil { + return rowErr + } + return ErrUserPlatformQuotaNotFound + } + err = rows.Scan(&id) + _ = rows.Close() + if err != nil { + return err + } + _, err = client.ExecContext(txCtx, "UPDATE usage_platform_quota_totals SET "+window+"_usage_usd=0,"+window+"_window_start=$2 WHERE quota_id=$1", id, newStart) return err - } - if n == 0 { - return ErrUserPlatformQuotaNotFound - } - return nil + }) } // withTx 在事务中执行 fn,若 ctx 中已有事务则复用。 @@ -276,24 +288,6 @@ func (r *userPlatformQuotaRepository) withTx(ctx context.Context, fn func(txCtx return nil } -// entQuotaToRecord 将 ent entity 映射为 repository record。 -// 注意 ent 生成字段名为 DailyLimitUsd(非 DailyLimitUSD)。 -func entQuotaToRecord(e *dbent.UserPlatformQuota) *UserPlatformQuotaRecord { - return &UserPlatformQuotaRecord{ - UserID: e.UserID, - Platform: e.Platform, - DailyLimitUSD: e.DailyLimitUsd, - WeeklyLimitUSD: e.WeeklyLimitUsd, - MonthlyLimitUSD: e.MonthlyLimitUsd, - DailyUsageUSD: e.DailyUsageUsd, - WeeklyUsageUSD: e.WeeklyUsageUsd, - MonthlyUsageUSD: e.MonthlyUsageUsd, - DailyWindowStart: e.DailyWindowStart, - WeeklyWindowStart: e.WeeklyWindowStart, - MonthlyWindowStart: e.MonthlyWindowStart, - } -} - // maybeReset 判断是否需要重置窗口用量: // - 若 prevStart 为 nil 或与 currStart 不同,表示窗口已过期,返回 cost(重置) // - 否则返回 prevUsage + cost(累加) diff --git a/backend/internal/service/batch_image.go b/backend/internal/service/batch_image.go index 130428f18e35..23fa6224cccd 100644 --- a/backend/internal/service/batch_image.go +++ b/backend/internal/service/batch_image.go @@ -79,6 +79,7 @@ var ( ErrBatchImageQueueFailed = infraerrors.New(http.StatusBadGateway, "BATCH_IMAGE_QUEUE_FAILED", "batch image queue failed") ErrBatchImageIdempotencyConflict = infraerrors.New(http.StatusConflict, "BATCH_IMAGE_IDEMPOTENCY_CONFLICT", "idempotency key reused with different batch image request") ErrBatchImageCancelFailed = infraerrors.New(http.StatusBadGateway, "BATCH_IMAGE_CANCEL_FAILED", "batch image cancel failed") + ErrBatchImageLegacyMoneyPrecision = infraerrors.New(http.StatusConflict, "BATCH_IMAGE_MONEY_REVIEW_REQUIRED", "legacy batch hold requires explicit monetary reconciliation") ErrBatchImageVertexGCSBucketMissing = infraerrors.New(http.StatusBadGateway, "BATCH_IMAGE_VERTEX_GCS_BUCKET_MISSING", "Vertex managed GCS bucket is not configured") ErrBatchImageNotReady = infraerrors.New(http.StatusConflict, "BATCH_IMAGE_NOT_READY", "batch image job is not completed") diff --git a/backend/internal/service/batch_image_billing_hold.go b/backend/internal/service/batch_image_billing_hold.go index 80f77208baa5..49491b7557bf 100644 --- a/backend/internal/service/batch_image_billing_hold.go +++ b/backend/internal/service/batch_image_billing_hold.go @@ -3,10 +3,8 @@ package service import ( "context" "errors" + "math" "strings" - - "github.com/Wei-Shaw/sub2api/internal/pkg/logger" - "go.uber.org/zap" ) const ( @@ -38,6 +36,8 @@ func buildBatchImageHoldCommand(job *BatchImageJob, requestID string, actualAmou if job.HoldAmount != nil { holdAmount = *job.HoldAmount } + // Preserve the original amounts for receipt verification. A legacy operation + // may already have committed before the task's status write failed. if holdAmount < 0 { holdAmount = 0 } @@ -51,10 +51,36 @@ func buildBatchImageHoldCommand(job *BatchImageJob, requestID string, actualAmou BatchID: job.BatchID, HoldAmount: holdAmount, ActualAmount: actualAmount, + MonetaryVersion: job.PricingSnapshotVersion, RequestPayloadHash: strings.TrimSpace(payloadHash), }, nil } +// This is a read-only compatibility check, never permission to execute an +// ambiguous amount. Only the exact committed operation fingerprint can succeed. +func confirmAmbiguousBatchImageRetry(ctx context.Context, repo UsageBillingRepository, cmd *BatchImageBalanceHoldCommand) (bool, error) { + if math.IsNaN(cmd.HoldAmount) || math.IsInf(cmd.HoldAmount, 0) || math.IsNaN(cmd.ActualAmount) || math.IsInf(cmd.ActualAmount, 0) { + return true, ErrBatchImageLegacyMoneyPrecision + } + if cmd.HoldAmount == QuantizeUsageBillingAmount(cmd.HoldAmount) && cmd.ActualAmount == QuantizeUsageBillingAmount(cmd.ActualAmount) { + return false, nil + } + reader, ok := repo.(interface { + ConfirmBatchImageBillingReceipt(context.Context, *BatchImageBalanceHoldCommand) (bool, error) + }) + if !ok { + return true, ErrBatchImageLegacyMoneyPrecision + } + confirmed, err := reader.ConfirmBatchImageBillingReceipt(ctx, cmd) + if err != nil { + return true, err + } + if !confirmed { + return true, ErrBatchImageLegacyMoneyPrecision + } + return true, nil +} + func reserveBatchImageBalanceHold(ctx context.Context, repo UsageBillingRepository, job *BatchImageJob, payloadHash string) error { if repo == nil { return ErrBatchImageBillingHoldFailed.WithCause(errors.New("batch image billing repository is not configured")) @@ -63,6 +89,9 @@ func reserveBatchImageBalanceHold(ctx context.Context, repo UsageBillingReposito if err != nil { return err } + if handled, err := confirmAmbiguousBatchImageRetry(ctx, repo, cmd); handled { + return err + } if cmd.HoldAmount <= 0 { return nil } @@ -83,6 +112,9 @@ func captureBatchImageBalanceHold(ctx context.Context, repo UsageBillingReposito if err != nil { return err } + if handled, err := confirmAmbiguousBatchImageRetry(ctx, repo, cmd); handled { + return err + } if _, err := repo.CaptureBatchImageBalance(ctx, cmd); err != nil { return ErrBatchImageSettlementBillingFailed.WithCause(err) } @@ -97,19 +129,13 @@ func releaseBatchImageBalanceHold(ctx context.Context, repo UsageBillingReposito if err != nil { return err } + if handled, err := confirmAmbiguousBatchImageRetry(ctx, repo, cmd); handled { + return err + } if cmd.HoldAmount <= 0 { return nil } if _, err := repo.ReleaseBatchImageBalance(ctx, cmd); err != nil { - // 同一 release request id 出现指纹冲突,说明此前已有一次携带不同 - // payloadHash 的释放成功提交(资金已归还)。视为幂等成功, - // 避免历史指纹不一致的 job 永远卡在释放失败的毒消息循环里。 - if errors.Is(err, ErrUsageBillingRequestConflict) { - logger.L().Warn("batch_image.release_fingerprint_conflict_treated_as_released", - zap.String("batch_id", job.BatchID), - ) - return nil - } return ErrBatchImageBillingHoldFailed.WithCause(err) } return nil diff --git a/backend/internal/service/batch_image_legacy_receipt_test.go b/backend/internal/service/batch_image_legacy_receipt_test.go new file mode 100644 index 000000000000..56ab5696ae98 --- /dev/null +++ b/backend/internal/service/batch_image_legacy_receipt_test.go @@ -0,0 +1,57 @@ +//go:build unit + +package service + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" +) + +type legacyReceiptOnlyBilling struct { + UsageBillingRepository // money methods panic if the compatibility path calls them + confirmed bool + conflict error + checks []BatchImageBalanceHoldCommand +} + +func (r *legacyReceiptOnlyBilling) ConfirmBatchImageBillingReceipt(_ context.Context, cmd *BatchImageBalanceHoldCommand) (bool, error) { + r.checks = append(r.checks, *cmd) + return r.confirmed, r.conflict +} + +func TestLegacyHighPrecisionCommittedRetryOnlyConfirmsExactReceipt(t *testing.T) { + for _, operation := range []string{"hold", "capture", "release"} { + for _, committed := range []bool{false, true} { + t.Run(operation+map[bool]string{false: "/uncommitted", true: "/committed"}[committed], func(t *testing.T) { + job := testSettlingBatchImageJob("legacy-committed-retry") + hold := 0.0761995125 + job.HoldAmount = &hold + reader := &legacyReceiptOnlyBilling{confirmed: committed} + var err error + switch operation { + case "hold": + err = reserveBatchImageBalanceHold(context.Background(), reader, job, "original-payload") + case "capture": + err = captureBatchImageBalanceHold(context.Background(), reader, job, 0.0612345678, "original-payload") + case "release": + err = releaseBatchImageBalanceHold(context.Background(), reader, job, "original-payload") + } + if committed { + require.NoError(t, err) + } else { + require.ErrorIs(t, err, ErrBatchImageLegacyMoneyPrecision) + } + require.Len(t, reader.checks, 1) + require.Equal(t, hold, reader.checks[0].HoldAmount) + require.Equal(t, job.UserID, reader.checks[0].UserID) + }) + } + } + job := testSettlingBatchImageJob("legacy-conflict") + hold := 0.0761995125 + job.HoldAmount = &hold + r := &legacyReceiptOnlyBilling{conflict: ErrUsageBillingRequestConflict} + require.ErrorIs(t, releaseBatchImageBalanceHold(context.Background(), r, job, "other-payload"), ErrUsageBillingRequestConflict) +} diff --git a/backend/internal/service/batch_image_money_precision_test.go b/backend/internal/service/batch_image_money_precision_test.go new file mode 100644 index 000000000000..f3e0115fda9f --- /dev/null +++ b/backend/internal/service/batch_image_money_precision_test.go @@ -0,0 +1,74 @@ +//go:build unit + +package service + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestBatchImageNewSnapshotUsesOneMoneyQuantum(t *testing.T) { + for _, tc := range []struct { + unit, rate, want float64 + }{ + {0.05, 1.2345, 0.07619951}, // down + {0.05, 1.2346, 0.07621186}, // up + } { + svc, _, _, _, _ := newTestBatchImagePublicService(true) + groupID := int64(7) + svc.GroupRepo = &publicBatchImageGroupRepo{groups: map[int64]*Group{ + groupID: {ID: groupID, Platform: PlatformGemini, RateMultiplier: tc.rate, + AllowBatchImageGeneration: true, ImagePrice1K: &tc.unit, + BatchImageDiscountMultiplier: 1, BatchImageHoldMultiplier: 1}, + }} + account := testBatchImageAccount(101, AccountTypeAPIKey) + account.RateMultiplier = &tc.rate + req := validBatchImageSubmitRequest() + req.ImageSize = "1K" + req.Items = req.Items[:1] + price, err := svc.resolvePricingSnapshot(context.Background(), BatchImageOwner{GroupID: &groupID}, req, BatchImageProviderGeminiAPI, &account) + require.NoError(t, err) + require.Equal(t, tc.want, price.HoldAmount) + require.Equal(t, price.HoldAmount, price.EstimatedCost) + } +} + +func TestBatchImageLegacyAmbiguousHoldRequiresReview(t *testing.T) { + for _, hold := range []float64{0.0761995125, 0.076211858} { + job := testSettlingBatchImageJob("legacy-precision") + job.HoldAmount = &hold + for _, id := range []string{BatchImageHoldRequestID(job.BatchID), BatchImageCaptureRequestID(job.BatchID), BatchImageReleaseRequestID(job.BatchID)} { + cmd, err := buildBatchImageHoldCommand(job, id, 0, "payload") + require.NoError(t, err) + require.Equal(t, hold, cmd.HoldAmount, "retain the exact input for read-only receipt verification") + handled, err := confirmAmbiguousBatchImageRetry(context.Background(), &fakeBatchImageBillingRepo{}, cmd) + require.True(t, handled) + require.ErrorIs(t, err, ErrBatchImageLegacyMoneyPrecision) + } + } + job := testSettlingBatchImageJob("legacy-actual-precision") + cmd, err := buildBatchImageHoldCommand(job, BatchImageCaptureRequestID(job.BatchID), 0.000078125, "payload") + require.NoError(t, err) + handled, err := confirmAmbiguousBatchImageRetry(context.Background(), &fakeBatchImageBillingRepo{}, cmd) + require.True(t, handled) + require.ErrorIs(t, err, ErrBatchImageLegacyMoneyPrecision) +} + +func TestBatchImageVersionTwoSettlementRoundsCompleteAmountOnce(t *testing.T) { + repo := newFakeBatchImageRepository() + job := testSettlingBatchImageJob("money-v2") + job.PricingSnapshotVersion = 2 + job.BillableUnitPrice = 0.0761995125 + job.SuccessCount, job.FailCount, job.ItemCount = 1, 0, 1 + hold := 0.07619951 + job.HoldAmount = &hold + repo.jobs[job.BatchID] = job + billing := &fakeBatchImageBillingRepo{} + svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{}} + result, err := svc.Settle(context.Background(), job.BatchID) + require.NoError(t, err) + require.Equal(t, hold, result.ActualCost) + require.Equal(t, hold, billing.captures[0].ActualAmount) +} diff --git a/backend/internal/service/batch_image_public.go b/backend/internal/service/batch_image_public.go index 409bc590eba4..cb10a285be7d 100644 --- a/backend/internal/service/batch_image_public.go +++ b/backend/internal/service/batch_image_public.go @@ -277,7 +277,7 @@ func (s *BatchImagePublicService) Submit(ctx context.Context, owner BatchImageOw HoldMultiplier: pricingSnapshot.HoldMultiplier, BillableUnitPrice: pricingSnapshot.BillableUnitPrice, HoldUnitPrice: pricingSnapshot.HoldUnitPrice, - PricingSnapshotVersion: 1, + PricingSnapshotVersion: 2, Currency: "USD", HoldID: &holdID, IdempotencyKey: batchImageOptionalStringPtr(idempotencyKey), @@ -1083,8 +1083,10 @@ func (s *BatchImagePublicService) resolvePricingSnapshot(ctx context.Context, ow HoldMultiplier: holdMultiplier, BillableUnitPrice: billableUnitPrice, HoldUnitPrice: holdUnitPrice, - EstimatedCost: billableUnitPrice * float64(len(req.Items)), - HoldAmount: holdUnitPrice * float64(len(req.Items)), + // Round the complete amount once, using the existing monetary policy. + // Version 2 persists exactly the amount moved between the two money columns. + EstimatedCost: QuantizeUsageBillingAmount(billableUnitPrice * float64(len(req.Items))), + HoldAmount: QuantizeUsageBillingAmount(holdUnitPrice * float64(len(req.Items))), }, nil } diff --git a/backend/internal/service/batch_image_public_test.go b/backend/internal/service/batch_image_public_test.go index 67ec00e13ae4..81196182fa16 100644 --- a/backend/internal/service/batch_image_public_test.go +++ b/backend/internal/service/batch_image_public_test.go @@ -84,7 +84,7 @@ func TestBatchImagePublicService_Submit(t *testing.T) { require.Equal(t, "files/gemini_api/output", batchImageDerefString(job.ProviderOutputRef)) require.NotNil(t, job.AccountID) require.Equal(t, int64(101), *job.AccountID) - require.Equal(t, 1, job.PricingSnapshotVersion) + require.Equal(t, 2, job.PricingSnapshotVersion) require.InDelta(t, 0.25, job.BaseUnitPrice, 1e-12) require.InDelta(t, 1.0, job.GroupRateMultiplier, 1e-12) require.InDelta(t, 1.0, job.AccountRateMultiplier, 1e-12) diff --git a/backend/internal/service/batch_image_settlement.go b/backend/internal/service/batch_image_settlement.go index e9b8a3685e97..cf36812d670f 100644 --- a/backend/internal/service/batch_image_settlement.go +++ b/backend/internal/service/batch_image_settlement.go @@ -136,6 +136,9 @@ func (s *BatchImageSettlementService) Settle(ctx context.Context, batchID string return nil, err } actualCost := float64(job.SuccessCount) * unitPrice + if job.PricingSnapshotVersion >= 2 { + actualCost = QuantizeUsageBillingAmount(actualCost) + } result.ActualCost = actualCost holdAmount := job.EstimatedCost if job.HoldAmount != nil { diff --git a/backend/internal/service/batch_image_settlement_test.go b/backend/internal/service/batch_image_settlement_test.go index f1fc3a15eb2a..f1012870393c 100644 --- a/backend/internal/service/batch_image_settlement_test.go +++ b/backend/internal/service/batch_image_settlement_test.go @@ -341,13 +341,11 @@ func TestBatchImageSettlementService_InvalidCountsExhaustsAndReleases(t *testing require.Equal(t, requestHash, billing.releases[0].RequestPayloadHash) } -func TestReleaseBatchImageBalanceHold_TreatsFingerprintConflictAsReleased(t *testing.T) { +func TestReleaseBatchImageBalanceHold_ConflictRequiresMatchingReceipt(t *testing.T) { job := testSettlingBatchImageJob("imgbatch_release_conflict") - // 历史版本用 manifestHash 释放过一次:同一 request id 再以 RequestHash - // 释放会命中指纹冲突。资金已归还,必须视为幂等成功而非毒消息。 billing := &fakeBatchImageBillingRepo{releaseErr: ErrUsageBillingRequestConflict} err := releaseBatchImageBalanceHold(context.Background(), billing, job, "request-hash") - require.NoError(t, err) + require.ErrorIs(t, err, ErrUsageBillingRequestConflict) require.Len(t, billing.releases, 1) } diff --git a/backend/internal/service/billing_cache_service.go b/backend/internal/service/billing_cache_service.go index cef54d266ae6..26bce21eee36 100644 --- a/backend/internal/service/billing_cache_service.go +++ b/backend/internal/service/billing_cache_service.go @@ -1247,7 +1247,7 @@ func (s *BillingCacheService) checkUserPlatformQuotaEligibility( cacheErr error ) if s.cache != nil { - entry, ok, cacheErr = s.cache.GetUserPlatformQuotaCache(ctx, userID, platform) + entry, ok, cacheErr = s.getPlatformQuotaCache(ctx, userID, platform) } else { // 标记为"cache 故障"分支:跳过 HIT 路径、不回填、走 DB 一次性检查 cacheErr = errBillingCacheUnavailable @@ -1311,7 +1311,7 @@ func (s *BillingCacheService) checkUserPlatformQuotaEligibility( } ttl := time.Duration(s.cfg.Billing.UserPlatformQuotaCacheTTLSeconds) * time.Second setCtx, setCancel := context.WithTimeout(context.Background(), 50*time.Millisecond) - if setErr := s.cache.SetUserPlatformQuotaCache(setCtx, userID, platform, refreshed, ttl); setErr != nil { + if setErr := s.setPlatformQuotaCache(setCtx, userID, platform, refreshed, ttl); setErr != nil { logger.LegacyPrintf("service.billing_cache", "Warning: refresh expired user platform quota cache failed user=%d platform=%s: %v", userID, platform, setErr) @@ -1380,7 +1380,7 @@ func (s *BillingCacheService) checkUserPlatformQuotaEligibility( sentinelTTL = time.Hour } setCtx, setCancel := context.WithTimeout(context.Background(), 50*time.Millisecond) - if setErr := s.cache.SetUserPlatformQuotaCache(setCtx, userID, platform, sentinel, sentinelTTL); setErr != nil { + if setErr := s.setPlatformQuotaCache(setCtx, userID, platform, sentinel, sentinelTTL); setErr != nil { userPlatformQuotaSentinelSetCacheErrorTotal.Add(1) logger.LegacyPrintf("service.billing_cache", "Warning: set sentinel quota cache failed user=%d platform=%s: %v", userID, platform, setErr) } @@ -1436,7 +1436,7 @@ func (s *BillingCacheService) checkUserPlatformQuotaEligibility( // 避免请求 ctx 提前取消(客户端断连/上游超时)导致 cache 回填失败, // 让下一次 preflight 仍然 MISS 并击穿到 DB(高并发下增大 DB 压力)。 setCtx, setCancel := context.WithTimeout(context.Background(), 50*time.Millisecond) - if setErr := s.cache.SetUserPlatformQuotaCache(setCtx, userID, platform, newEntry, ttl); setErr != nil { + if setErr := s.setPlatformQuotaCache(setCtx, userID, platform, newEntry, ttl); setErr != nil { logger.LegacyPrintf("service.billing_cache", "Warning: set user platform quota cache failed user=%d platform=%s: %v", userID, platform, setErr) } setCancel() @@ -1516,7 +1516,7 @@ func (s *BillingCacheService) HasUserPlatformQuotaLimit(ctx context.Context, use if s.cache == nil { return true } - entry, ok, err := s.cache.GetUserPlatformQuotaCache(ctx, userID, platform) + entry, ok, err := s.getPlatformQuotaCache(ctx, userID, platform) if err != nil || !ok || entry == nil { return true } diff --git a/backend/internal/service/billing_service.go b/backend/internal/service/billing_service.go index 593e824d2036..b90fd2b26b8f 100644 --- a/backend/internal/service/billing_service.go +++ b/backend/internal/service/billing_service.go @@ -84,7 +84,8 @@ type BillingCache interface { GetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string) (*UserPlatformQuotaCacheEntry, bool, error) SetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string, entry *UserPlatformQuotaCacheEntry, ttl time.Duration) error DeleteUserPlatformQuotaCache(ctx context.Context, userID int64, platform string) error - // IncrUserPlatformQuotaUsageCache 在缓存命中时累加用量;缓存未命中(key 不存在)静默返回 nil。 + // IncrUserPlatformQuotaUsageCache returns ErrUserPlatformQuotaCacheNotUpdated + // when the key/schema is missing; only a real increment returns nil. // markDirty=true 时将该 key 的 member 写入 Redis 脏集,供 flusher 批量回写 DB。 IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration, markDirty bool) error diff --git a/backend/internal/service/durable_platform_quota_cache.go b/backend/internal/service/durable_platform_quota_cache.go new file mode 100644 index 000000000000..cff09c95d8f4 --- /dev/null +++ b/backend/internal/service/durable_platform_quota_cache.go @@ -0,0 +1,41 @@ +package service + +import ( + "context" + "errors" + "time" +) + +var ErrUserPlatformQuotaCacheNotUpdated = errors.New("platform quota cache missing or schema mismatch; increment was not applied") + +// Combined DB totals must never be flushed back as a legacy baseline. A separate +// cache namespace leaves old dirty entries available to the legacy flusher. +type durablePlatformQuotaCache interface { + GetDurablePlatformQuotaCache(context.Context, int64, string) (*UserPlatformQuotaCacheEntry, bool, error) + SetDurablePlatformQuotaCache(context.Context, int64, string, *UserPlatformQuotaCacheEntry, time.Duration) error + DeleteDurablePlatformQuotaCache(context.Context, int64, string) error +} + +func (s *BillingCacheService) getPlatformQuotaCache(ctx context.Context, user int64, platform string) (*UserPlatformQuotaCacheEntry, bool, error) { + if cache, ok := s.cache.(durablePlatformQuotaCache); ok { + return cache.GetDurablePlatformQuotaCache(ctx, user, platform) + } + return s.cache.GetUserPlatformQuotaCache(ctx, user, platform) +} + +func (s *BillingCacheService) setPlatformQuotaCache(ctx context.Context, user int64, platform string, entry *UserPlatformQuotaCacheEntry, ttl time.Duration) error { + if cache, ok := s.cache.(durablePlatformQuotaCache); ok { + return cache.SetDurablePlatformQuotaCache(ctx, user, platform, entry, ttl) + } + return s.cache.SetUserPlatformQuotaCache(ctx, user, platform, entry, ttl) +} + +func (s *BillingCacheService) InvalidateDurablePlatformQuota(ctx context.Context, user int64, platform string) error { + if s == nil || s.cache == nil { + return nil + } + if cache, ok := s.cache.(durablePlatformQuotaCache); ok { + return cache.DeleteDurablePlatformQuotaCache(ctx, user, platform) + } + return s.cache.DeleteUserPlatformQuotaCache(ctx, user, platform) +} diff --git a/backend/internal/service/gateway_record_usage_test.go b/backend/internal/service/gateway_record_usage_test.go index 9dfa55629882..6b17970869ed 100644 --- a/backend/internal/service/gateway_record_usage_test.go +++ b/backend/internal/service/gateway_record_usage_test.go @@ -224,8 +224,8 @@ func TestGatewayServiceRecordUsage_GeminiFlashThinkingTierUsesCatalogPrice(t *te require.NotNil(t, usageRepo.lastLog) require.Equal(t, model, usageRepo.lastLog.Model) require.InDelta(t, 0.02007585, usageRepo.lastLog.TotalCost, 1e-12) - require.InDelta(t, 0.0030113775, usageRepo.lastLog.ActualCost, 1e-12) - require.InDelta(t, 0.0030113775, userRepo.lastAmount, 1e-12) + require.InDelta(t, 0.00301138, usageRepo.lastLog.ActualCost, 1e-12) + require.InDelta(t, 0.00301138, userRepo.lastAmount, 1e-12) }) } } @@ -377,8 +377,8 @@ func TestGatewayServiceRecordUsage_PeakRateAffectsTokenModeImageOutputTokens(t * require.InDelta(t, textInput+textOutput+imageOutput, usageRepo.lastLog.TotalCost, 1e-12) require.InDelta(t, imageOutput, usageRepo.lastLog.ImageOutputCost, 1e-12) - require.InDelta(t, expectedActual, usageRepo.lastLog.ActualCost, 1e-12) - require.InDelta(t, expectedActual, userRepo.lastAmount, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(expectedActual), usageRepo.lastLog.ActualCost, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(expectedActual), userRepo.lastAmount, 1e-12) } func TestGatewayServiceRecordUsage_TimePricingUsesPricingAt(t *testing.T) { @@ -410,7 +410,7 @@ func TestGatewayServiceRecordUsage_TimePricingUsesPricingAt(t *testing.T) { require.NotNil(t, usageRepo.lastLog) baseCost := 1000*3e-6 + 500*15e-6 require.InDelta(t, baseCost*2, usageRepo.lastLog.TotalCost, 1e-12) - require.InDelta(t, baseCost*2*0.8, usageRepo.lastLog.ActualCost, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(baseCost*2*0.8), usageRepo.lastLog.ActualCost, 1e-12) require.InDelta(t, 0.8, usageRepo.lastLog.RateMultiplier, 1e-12) } func TestGatewayServiceRecordUsage_UsesExplicitPricingAtForPeakRate(t *testing.T) { @@ -512,9 +512,9 @@ func TestGatewayServiceRecordUsage_DeepSeekAccountStatsUsesRequestPricingAtAndUp require.False(t, log.CreatedAt.Equal(slot.pricingAt), "request pricing time must differ from record creation") customerTotal := 1000*inputPrice + 500*outputPrice + 1000*cachePrice require.InDelta(t, customerTotal, log.TotalCost, 1e-12) - require.InDelta(t, customerTotal*0.8, log.ActualCost, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(customerTotal*0.8), log.ActualCost, 1e-12) require.Equal(t, 1, userRepo.deductCalls) - require.InDelta(t, customerTotal*0.8, userRepo.lastAmount, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(customerTotal*0.8), userRepo.lastAmount, 1e-12) require.NotNil(t, log.AccountStatsCost) require.InDelta(t, model.offPeakCost*slot.multiplier, *log.AccountStatsCost, 1e-12, "account cost must use the upstream model and historical PricingAt") @@ -669,7 +669,7 @@ func TestGatewayServiceRecordUsage_DroppedUsageLogFallsBackToSyncCreate(t *testi require.NoError(t, usageRepo.lastCtxErr) } -func TestGatewayServiceRecordUsage_BillingErrorWritesUnsettledUsageLog(t *testing.T) { +func TestGatewayServiceRecordUsage_BillingErrorDoesNotWriteFalseZeroCostLog(t *testing.T) { usageRepo := &openAIRecordUsageLogRepoStub{} billingErr := errors.New("billing tx failed") billingRepo := &openAIRecordUsageBillingRepoStub{err: billingErr} @@ -694,14 +694,8 @@ func TestGatewayServiceRecordUsage_BillingErrorWritesUnsettledUsageLog(t *testin require.ErrorIs(t, err, billingErr) require.Equal(t, 1, billingRepo.calls) - require.Equal(t, 1, usageRepo.calls) - require.NotNil(t, usageRepo.lastLog) - require.Equal(t, 10, usageRepo.lastLog.InputTokens) - require.Equal(t, 6, usageRepo.lastLog.OutputTokens) - require.Greater(t, usageRepo.lastLog.InputCost, 0.0) - require.Greater(t, usageRepo.lastLog.OutputCost, 0.0) - require.Greater(t, usageRepo.lastLog.TotalCost, 0.0) - require.Zero(t, usageRepo.lastLog.ActualCost) + require.Zero(t, usageRepo.calls) + require.Nil(t, usageRepo.lastLog, "unknown settlement is not a successful zero-cost usage row") } func TestGatewayServiceRecordUsage_ReasoningEffortPersisted(t *testing.T) { diff --git a/backend/internal/service/gateway_usage_billing.go b/backend/internal/service/gateway_usage_billing.go index 4b01e5a888de..f9bca5d9792c 100644 --- a/backend/internal/service/gateway_usage_billing.go +++ b/backend/internal/service/gateway_usage_billing.go @@ -299,6 +299,8 @@ func buildUsageBillingCommand(requestID string, usageLog *UsageLog, p *postUsage AccountType: p.Account.Type, RequestPayloadHash: strings.TrimSpace(p.RequestPayloadHash), } + quoted := p.Cost.ActualCost + cmd.QuotedActualCost = "ed if usageLog != nil { cmd.Model = usageLog.Model cmd.BillingType = usageLog.BillingType @@ -348,6 +350,14 @@ func buildUsageBillingCommand(requestID string, usageLog *UsageLog, p *postUsage cmd.AccountQuotaHasDailyLimit = p.Account.GetQuotaDailyLimit() > 0 cmd.AccountQuotaHasWeeklyLimit = p.Account.GetQuotaWeeklyLimit() > 0 } + if !p.IsSubscriptionBill && p.Platform != "" && p.Cost.ActualCost > 0 { + cmd.PlatformQuotaPlatform = p.Platform + cmd.PlatformQuotaCost = p.Cost.ActualCost + cmd.PlatformUsageAt = time.Now().UTC() + if usageLog != nil && !usageLog.CreatedAt.IsZero() { + cmd.PlatformUsageAt = usageLog.CreatedAt + } + } cmd.Normalize() return cmd @@ -359,6 +369,19 @@ func applyUsageBilling(ctx context.Context, requestID string, usageLog *UsageLog } cmd := buildUsageBillingCommand(requestID, usageLog, p) + // The fingerprint above retains the original monetary inputs. From here on, + // every customer-facing amount uses the exact same quantum as the SQL charge. + // Keep TotalCost and the independent account-cost multiplier untouched. + if p.Cost != nil { + canonicalParams := *p + canonicalCost := *p.Cost + canonicalCost.ActualCost = QuantizeUsageBillingAmount(p.Cost.ActualCost) + canonicalParams.Cost = &canonicalCost + p = &canonicalParams + if usageLog != nil { + usageLog.ActualCost = canonicalCost.ActualCost + } + } if cmd == nil || cmd.RequestID == "" || repo == nil { if p.SimpleModeKeyRateLimitOnly { return false, ErrSimpleModeKeyRateLimitBillingUnavailable @@ -372,6 +395,9 @@ func applyUsageBilling(ctx context.Context, requestID string, usageLog *UsageLog billingCtx, cancel := detachedBillingContext(ctx) defer cancel() + if durable, ok := repo.(DurableUsageBillingRepository); ok && usageLog != nil { + return applyDurableUsageBilling(billingCtx, durable, cmd, usageLog, p, deps) + } result, err := repo.Apply(billingCtx, cmd) if err != nil { return false, err @@ -432,7 +458,11 @@ 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 && deps.billingCacheService != nil { + if result != nil && result.PlatformQuotaApplied && deps.billingCacheService != nil && p.User != nil { + if err := deps.billingCacheService.InvalidateDurablePlatformQuota(ctx, p.User.ID, p.Platform); err != nil { + logger.LegacyPrintf("service.gateway", "durable platform quota cache invalidation failed: %v", err) + } + } else 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 { @@ -659,6 +689,14 @@ func writeUsageLogBestEffort(ctx context.Context, repo UsageLogRepository, usage } } +func recordDurableUsageTiming(ctx context.Context, repo UsageLogRepository, usage *UsageLog) { + if recorder, ok := repo.(interface { + RecordRequestTiming(context.Context, string, int64) + }); ok && usage != nil { + recorder.RecordRequestTiming(ctx, usage.RequestID, usage.APIKeyID) + } +} + // RecordUsage 记录使用量并扣费(或更新订阅用量) func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInput) error { return s.recordUsageCore(ctx, &recordUsageCoreInput{ @@ -941,11 +979,13 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage }, s.billingDeps(), s.usageBillingRepo) if billingErr != nil { - usageLog.ActualCost = 0 - writeUsageLogBestEffort(ctx, s.usageLogRepo, usageLog, "service.gateway") return billingErr } - writeUsageLogBestEffort(ctx, s.usageLogRepo, usageLog, "service.gateway") + if !s.DurableUsageBillingEnabled() { + writeUsageLogBestEffort(ctx, s.usageLogRepo, usageLog, "service.gateway") + } else { + recordDurableUsageTiming(ctx, s.usageLogRepo, usageLog) + } return nil } diff --git a/backend/internal/service/openai_gateway_record_usage_test.go b/backend/internal/service/openai_gateway_record_usage_test.go index f7ca6f22b17f..ec1f2b1b9fb4 100644 --- a/backend/internal/service/openai_gateway_record_usage_test.go +++ b/backend/internal/service/openai_gateway_record_usage_test.go @@ -98,9 +98,9 @@ func TestRecordCyberPolicyUsageLog_BillsRealUpstreamTokens(t *testing.T) { expected := expectedOpenAICost(t, svc, "gpt-5.1", usage, 1.1) require.Greater(t, usageRepo.lastLog.ActualCost, 0.0, "流式 cyber 有真实 token,须计费") - require.InDelta(t, expected.ActualCost, usageRepo.lastLog.ActualCost, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(expected.ActualCost), usageRepo.lastLog.ActualCost, 1e-12) require.Equal(t, 1, userRepo.deductCalls, "按真实 token 扣费,与 WS/正常请求一致") - require.InDelta(t, expected.ActualCost, userRepo.lastAmount, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(expected.ActualCost), userRepo.lastAmount, 1e-12) } func TestRecordCyberPolicyUsageLog_NonStreamZeroTokensZeroCost(t *testing.T) { @@ -434,8 +434,8 @@ func TestOpenAIGatewayServiceRecordUsage_UsesUserSpecificGroupRate(t *testing.T) require.Equal(t, 3, usageRepo.lastLog.CacheReadTokens) expected := expectedOpenAICost(t, svc, "gpt-5.1", usage, userRate) - require.InDelta(t, expected.ActualCost, usageRepo.lastLog.ActualCost, 1e-12) - require.InDelta(t, expected.ActualCost, userRepo.lastAmount, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(expected.ActualCost), usageRepo.lastLog.ActualCost, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(expected.ActualCost), userRepo.lastAmount, 1e-12) require.Equal(t, 1, userRepo.deductCalls) } @@ -522,8 +522,8 @@ func TestOpenAIGatewayServiceRecordUsage_PeakRateAffectsTokenModeImageOutputToke require.InDelta(t, expected.TotalCost, usageRepo.lastLog.TotalCost, 1e-12) require.InDelta(t, expected.ImageOutputCost, usageRepo.lastLog.ImageOutputCost, 1e-12) - require.InDelta(t, expectedActual, usageRepo.lastLog.ActualCost, 1e-12) - require.InDelta(t, expectedActual, userRepo.lastAmount, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(expectedActual), usageRepo.lastLog.ActualCost, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(expectedActual), userRepo.lastAmount, 1e-12) } func TestOpenAIGatewayServiceRecordUsage_TimePricingUsesPricingAt(t *testing.T) { @@ -555,7 +555,7 @@ func TestOpenAIGatewayServiceRecordUsage_TimePricingUsesPricingAt(t *testing.T) require.NotNil(t, usageRepo.lastLog) baseCost := 1000*3e-6 + 500*15e-6 require.InDelta(t, baseCost*2, usageRepo.lastLog.TotalCost, 1e-12) - require.InDelta(t, baseCost*2*0.8, usageRepo.lastLog.ActualCost, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(baseCost*2*0.8), usageRepo.lastLog.ActualCost, 1e-12) require.InDelta(t, 0.8, usageRepo.lastLog.RateMultiplier, 1e-12) } @@ -588,7 +588,7 @@ func TestOpenAIGatewayServiceRecordUsage_TimePricingUsesExplicitPricingAt(t *tes require.NotNil(t, usageRepo.lastLog) baseCost := 1000*3e-6 + 500*15e-6 require.InDelta(t, baseCost, usageRepo.lastLog.TotalCost, 1e-12) - require.InDelta(t, baseCost*0.8, usageRepo.lastLog.ActualCost, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(baseCost*0.8), usageRepo.lastLog.ActualCost, 1e-12) require.InDelta(t, 0.8, usageRepo.lastLog.RateMultiplier, 1e-12) } @@ -649,9 +649,9 @@ func TestOpenAIGatewayServiceRecordUsage_DeepSeekAccountStatsUsesRequestPricingA require.False(t, log.CreatedAt.Equal(slot.pricingAt), "request pricing time must differ from record creation") customerTotal := 1000*inputPrice + 500*outputPrice + 1000*cachePrice require.InDelta(t, customerTotal, log.TotalCost, 1e-12) - require.InDelta(t, customerTotal*0.8, log.ActualCost, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(customerTotal*0.8), log.ActualCost, 1e-12) require.Equal(t, 1, userRepo.deductCalls) - require.InDelta(t, customerTotal*0.8, userRepo.lastAmount, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(customerTotal*0.8), userRepo.lastAmount, 1e-12) require.NotNil(t, log.AccountStatsCost) require.InDelta(t, model.offPeakCost*slot.multiplier, *log.AccountStatsCost, 1e-12, "account cost must use the upstream model and historical PricingAt") @@ -731,7 +731,7 @@ func TestOpenAIGatewayServiceRecordUsage_FallsBackToGroupDefaultRateOnResolverEr require.Equal(t, groupRate, usageRepo.lastLog.RateMultiplier) expected := expectedOpenAICost(t, svc, "gpt-5.1", usage, groupRate) - require.InDelta(t, expected.ActualCost, userRepo.lastAmount, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(expected.ActualCost), userRepo.lastAmount, 1e-12) } func TestOpenAIGatewayServiceRecordUsage_FallsBackToGroupDefaultRateWhenResolverMissing(t *testing.T) { @@ -1103,7 +1103,7 @@ func TestOpenAIGatewayServiceRecordUsage_GeneratesRequestIDWhenAllSourcesMissing require.Equal(t, billingRepo.lastCmd.RequestID, usageRepo.lastLog.RequestID) } -func TestOpenAIGatewayServiceRecordUsage_BillingErrorWritesUnsettledUsageLog(t *testing.T) { +func TestOpenAIGatewayServiceRecordUsage_BillingErrorDoesNotWriteFalseZeroCostLog(t *testing.T) { usageRepo := &openAIRecordUsageLogRepoStub{} billingErr := errors.New("billing tx failed") billingRepo := &openAIRecordUsageBillingRepoStub{err: billingErr} @@ -1128,14 +1128,8 @@ func TestOpenAIGatewayServiceRecordUsage_BillingErrorWritesUnsettledUsageLog(t * require.ErrorIs(t, err, billingErr) require.Equal(t, 1, billingRepo.calls) - require.Equal(t, 1, usageRepo.calls) - require.NotNil(t, usageRepo.lastLog) - require.Equal(t, 8, usageRepo.lastLog.InputTokens) - require.Equal(t, 4, usageRepo.lastLog.OutputTokens) - require.Greater(t, usageRepo.lastLog.InputCost, 0.0) - require.Greater(t, usageRepo.lastLog.OutputCost, 0.0) - require.Greater(t, usageRepo.lastLog.TotalCost, 0.0) - require.Zero(t, usageRepo.lastLog.ActualCost) + require.Zero(t, usageRepo.calls) + require.Nil(t, usageRepo.lastLog, "unknown settlement is not a successful zero-cost usage row") } func TestOpenAIGatewayServiceRecordUsage_UpdatesAPIKeyQuotaWhenConfigured(t *testing.T) { @@ -1166,7 +1160,7 @@ func TestOpenAIGatewayServiceRecordUsage_UpdatesAPIKeyQuotaWhenConfigured(t *tes require.Equal(t, 1, quotaSvc.quotaCalls) require.Equal(t, 0, quotaSvc.rateLimitCalls) expected := expectedOpenAICost(t, svc, "gpt-5.1", usage, 1.1) - require.InDelta(t, expected.ActualCost, quotaSvc.lastAmount, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(expected.ActualCost), quotaSvc.lastAmount, 1e-12) } func TestOpenAIGatewayServiceRecordUsage_ClampsActualInputTokensToZero(t *testing.T) { @@ -1236,7 +1230,7 @@ func TestOpenAIGatewayServiceRecordUsage_GPT56SeparatesCacheWriteForBillingAndSt require.InDelta(t, 200*6.25e-6, usageRepo.lastLog.CacheCreationCost, 1e-12) require.InDelta(t, 100*0.5e-6, usageRepo.lastLog.CacheReadCost, 1e-12) require.InDelta(t, 50*30e-6, usageRepo.lastLog.OutputCost, 1e-12) - require.InDelta(t, usageRepo.lastLog.TotalCost*1.1, usageRepo.lastLog.ActualCost, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(usageRepo.lastLog.TotalCost*1.1), usageRepo.lastLog.ActualCost, 1e-12) } func TestOpenAIGatewayServiceRecordUsage_Gpt54LongContextBillingDisabledWhenGroupAndAccountOff(t *testing.T) { @@ -1268,7 +1262,7 @@ func TestOpenAIGatewayServiceRecordUsage_Gpt54LongContextBillingDisabledWhenGrou require.InDelta(t, expectedInput, usageRepo.lastLog.InputCost, 1e-10) require.InDelta(t, expectedOutput, usageRepo.lastLog.OutputCost, 1e-10) require.InDelta(t, expectedInput+expectedOutput, usageRepo.lastLog.TotalCost, 1e-10) - require.InDelta(t, (expectedInput+expectedOutput)*1.1, usageRepo.lastLog.ActualCost, 1e-10) + require.InDelta(t, QuantizeUsageBillingAmount((expectedInput+expectedOutput)*1.1), usageRepo.lastLog.ActualCost, 1e-10) require.False(t, usageRepo.lastLog.LongContextBillingApplied) require.Equal(t, 1, userRepo.deductCalls) } @@ -1316,7 +1310,7 @@ func TestOpenAIGatewayServiceRecordUsage_Gpt54LongContextBillingEnabledPerAccoun require.InDelta(t, expectedInput, usageRepo.lastLog.InputCost, 1e-10) require.InDelta(t, expectedOutput, usageRepo.lastLog.OutputCost, 1e-10) require.InDelta(t, expectedInput+expectedOutput, usageRepo.lastLog.TotalCost, 1e-10) - require.InDelta(t, (expectedInput+expectedOutput)*1.1, usageRepo.lastLog.ActualCost, 1e-10) + require.InDelta(t, QuantizeUsageBillingAmount((expectedInput+expectedOutput)*1.1), usageRepo.lastLog.ActualCost, 1e-10) require.True(t, usageRepo.lastLog.LongContextBillingApplied) } @@ -1773,9 +1767,9 @@ func TestOpenAIGatewayServiceRecordUsage_BillsMappedRequestsUsingRequestedModel( require.NoError(t, err) require.NotNil(t, usageRepo.lastLog) require.Equal(t, "gpt-5.1", usageRepo.lastLog.Model) - require.Equal(t, expectedCost.ActualCost, usageRepo.lastLog.ActualCost) + require.Equal(t, QuantizeUsageBillingAmount(expectedCost.ActualCost), usageRepo.lastLog.ActualCost) require.Equal(t, expectedCost.TotalCost, usageRepo.lastLog.TotalCost) - require.Equal(t, expectedCost.ActualCost, userRepo.lastAmount) + require.Equal(t, QuantizeUsageBillingAmount(expectedCost.ActualCost), userRepo.lastAmount) } func TestOpenAIGatewayServiceRecordUsage_ChannelMappedDoesNotOverrideBillingModelWhenUnmapped(t *testing.T) { @@ -1815,7 +1809,7 @@ func TestOpenAIGatewayServiceRecordUsage_ChannelMappedDoesNotOverrideBillingMode require.NoError(t, err) require.NotNil(t, usageRepo.lastLog) - require.Equal(t, expectedCost.ActualCost, usageRepo.lastLog.ActualCost) + require.Equal(t, QuantizeUsageBillingAmount(expectedCost.ActualCost), usageRepo.lastLog.ActualCost) require.True(t, usageRepo.lastLog.ActualCost > 0, "cost must not be zero") } @@ -1856,7 +1850,7 @@ func TestOpenAIGatewayServiceRecordUsage_ChannelMappedOverridesBillingModelWhenM require.NoError(t, err) require.NotNil(t, usageRepo.lastLog) - require.Equal(t, expectedCost.ActualCost, usageRepo.lastLog.ActualCost) + require.Equal(t, QuantizeUsageBillingAmount(expectedCost.ActualCost), usageRepo.lastLog.ActualCost) require.True(t, usageRepo.lastLog.ActualCost > 0, "cost must not be zero") } @@ -1913,8 +1907,8 @@ func TestOpenAIGatewayServiceRecordUsage_ResponsesMappedBillingModelHonorsBillin require.NoError(t, err) require.NotNil(t, usageRepo.lastLog) require.Equal(t, "gpt-5.4", usageRepo.lastLog.Model) - require.InDelta(t, expectedCost.ActualCost, usageRepo.lastLog.ActualCost, 1e-12) - require.InDelta(t, expectedCost.ActualCost, userRepo.lastAmount, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(expectedCost.ActualCost), usageRepo.lastLog.ActualCost, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(expectedCost.ActualCost), userRepo.lastAmount, 1e-12) require.True(t, usageRepo.lastLog.ActualCost > 0, "cost must not be zero") }) } @@ -1951,9 +1945,9 @@ func TestOpenAIGatewayServiceRecordUsage_BillsCompactOpenAIModelAlias(t *testing require.Equal(t, "gpt5.5", usageRepo.lastLog.Model) require.NotNil(t, usageRepo.lastLog.UpstreamModel) require.Equal(t, "gpt-5.4", *usageRepo.lastLog.UpstreamModel) - require.InDelta(t, expectedCost.ActualCost, usageRepo.lastLog.ActualCost, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(expectedCost.ActualCost), usageRepo.lastLog.ActualCost, 1e-12) require.True(t, usageRepo.lastLog.ActualCost > 0, "cost must not be zero") - require.InDelta(t, expectedCost.ActualCost, userRepo.lastAmount, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(expectedCost.ActualCost), userRepo.lastAmount, 1e-12) } func TestOpenAIGatewayServiceRecordUsage_FallsBackToUpstreamModelWhenPrimaryUnpriceable(t *testing.T) { @@ -1985,9 +1979,9 @@ func TestOpenAIGatewayServiceRecordUsage_FallsBackToUpstreamModelWhenPrimaryUnpr require.NoError(t, err) require.NotNil(t, usageRepo.lastLog) - require.InDelta(t, expectedCost.ActualCost, usageRepo.lastLog.ActualCost, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(expectedCost.ActualCost), usageRepo.lastLog.ActualCost, 1e-12) require.True(t, usageRepo.lastLog.ActualCost > 0, "cost must not be zero") - require.InDelta(t, expectedCost.ActualCost, userRepo.lastAmount, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(expectedCost.ActualCost), userRepo.lastAmount, 1e-12) } func TestOpenAIGatewayServiceRecordUsage_UnpricedTokenModelFallsBackToZeroCostUsageLog(t *testing.T) { @@ -2451,7 +2445,7 @@ func TestOpenAIGatewayServiceRecordUsage_GrokVideoUsesDefaultRateCard(t *testing require.Nil(t, usageRepo.lastLog.ImageSize) // 结果未携带 duration 时按上游默认 8 秒计费:0.14 USD/s × 8s。 require.InDelta(t, 0.14*8, usageRepo.lastLog.TotalCost, 1e-12) - require.InDelta(t, 0.14*8, usageRepo.lastLog.ActualCost, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(0.14*8), usageRepo.lastLog.ActualCost, 1e-12) require.Equal(t, 0, usageRepo.lastLog.ImageCount) require.NotNil(t, usageRepo.lastLog.BillingMode) require.Equal(t, string(BillingModeVideo), *usageRepo.lastLog.BillingMode) @@ -2590,7 +2584,7 @@ func TestOpenAIGatewayServiceRecordUsage_GroupVideoModelPriceOverridesFlatAndCha require.NoError(t, err) require.NotNil(t, usageRepo.lastLog) require.InDelta(t, modelVideoPrice720P*2, usageRepo.lastLog.TotalCost, 1e-12) - require.InDelta(t, modelVideoPrice720P*2, usageRepo.lastLog.ActualCost, 1e-12) + require.InDelta(t, QuantizeUsageBillingAmount(modelVideoPrice720P*2), usageRepo.lastLog.ActualCost, 1e-12) require.NotNil(t, usageRepo.lastLog.BillingMode) require.Equal(t, string(BillingModeVideo), *usageRepo.lastLog.BillingMode) } @@ -3229,7 +3223,7 @@ func TestOpenAIGatewayServiceRecordUsage_FreeOpenAIFastChargesStandard(t *testin standardTotal := float64(tokens.InputTokens)*inputPrice + float64(tokens.OutputTokens)*outputPrice require.InDelta(t, standardTotal*fastMultiplier, usageRepo.lastLog.TotalCost, 1e-10) - require.InDelta(t, standardTotal*0.5, usageRepo.lastLog.ActualCost, 1e-10) + require.InDelta(t, QuantizeUsageBillingAmount(standardTotal*0.5), usageRepo.lastLog.ActualCost, 1e-10) } func TestOpenAIGatewayServiceRecordUsage_FreeOpenAIFastMissingPricingRecordsZeroCostUsageLog(t *testing.T) { diff --git a/backend/internal/service/openai_gateway_usage.go b/backend/internal/service/openai_gateway_usage.go index 1b7d43f92191..f634e4449735 100644 --- a/backend/internal/service/openai_gateway_usage.go +++ b/backend/internal/service/openai_gateway_usage.go @@ -551,11 +551,13 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec }, s.billingDeps(), s.usageBillingRepo) if billingErr != nil { - usageLog.ActualCost = 0 - writeUsageLogBestEffort(ctx, s.usageLogRepo, usageLog, "service.openai_gateway") return billingErr } - writeUsageLogBestEffort(ctx, s.usageLogRepo, usageLog, "service.openai_gateway") + if !s.DurableUsageBillingEnabled() { + writeUsageLogBestEffort(ctx, s.usageLogRepo, usageLog, "service.openai_gateway") + } else { + recordDurableUsageTiming(ctx, s.usageLogRepo, usageLog) + } return nil } diff --git a/backend/internal/service/usage_billing.go b/backend/internal/service/usage_billing.go index 447260ecb451..40c27f21625c 100644 --- a/backend/internal/service/usage_billing.go +++ b/backend/internal/service/usage_billing.go @@ -8,6 +8,7 @@ import ( "fmt" "math" "strings" + "time" "github.com/shopspring/decimal" ) @@ -42,6 +43,14 @@ type UsageBillingCommand struct { APIKeyQuotaCost float64 APIKeyRateLimitCost float64 AccountQuotaCost float64 + // QuotedActualCost preserves the original quote before the existing eight-place + // monetary rounding. It does not replace TotalCost or account-cost accounting. + QuotedActualCost *float64 + // Platform usage belongs to this same deduplicated request, not a Redis + // snapshot. UsageAt is frozen at observation so retries cannot move windows. + PlatformQuotaPlatform string + PlatformQuotaCost float64 + PlatformUsageAt time.Time // AccountQuotaHasDailyLimit / AccountQuotaHasWeeklyLimit 记录构造命令时账号是否设置了日/周上限, // 只用于选择账号配额累加语句,不参与幂等指纹;实际累加哪些维度以数据库中的账号配置为准。 @@ -91,6 +100,7 @@ func (c *UsageBillingCommand) quantizeMonetaryFields() { c.APIKeyQuotaCost = QuantizeUsageBillingAmount(c.APIKeyQuotaCost) c.APIKeyRateLimitCost = QuantizeUsageBillingAmount(c.APIKeyRateLimitCost) c.AccountQuotaCost = QuantizeUsageBillingAmount(c.AccountQuotaCost) + c.PlatformQuotaCost = QuantizeUsageBillingAmount(c.PlatformQuotaCost) } // QuantizeUsageBillingAmount 把金额舍入到 UsageBillingMonetaryScale 位小数, @@ -170,6 +180,7 @@ type AccountQuotaState struct { type UsageBillingApplyResult struct { Applied bool APIKeyQuotaExhausted bool + PlatformQuotaApplied bool // durable SQL effect; do not replay legacy cache/DB increments NewBalance *float64 // post-deduction balance (nil = no balance deduction) BalanceOverdrafted bool // true when the sufficient-balance guard missed and debt was still recorded QuotaState *AccountQuotaState // post-increment quota state (nil = no quota increment) @@ -185,6 +196,7 @@ type BatchImageBalanceHoldCommand struct { BatchID string HoldAmount float64 ActualAmount float64 + MonetaryVersion int // >=2 requires an individually owned durable hold } func (c *BatchImageBalanceHoldCommand) Normalize() { diff --git a/backend/internal/service/usage_billing_auth_recovery_test.go b/backend/internal/service/usage_billing_auth_recovery_test.go new file mode 100644 index 000000000000..c3b3992fed40 --- /dev/null +++ b/backend/internal/service/usage_billing_auth_recovery_test.go @@ -0,0 +1,150 @@ +//go:build unit + +package service + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" + "github.com/stretchr/testify/require" +) + +type checkedBillingAuthRedis struct { + APIKeyCache + client *redis.Client + fault *atomic.Int32 + ready chan struct{} +} + +func (c *checkedBillingAuthRedis) DeleteAuthCache(ctx context.Context, key string) error { + if c.fault.Load() == 2 { + return errors.New("injected Redis delete failure") + } + return c.client.Del(ctx, key).Err() +} +func (c *checkedBillingAuthRedis) PublishAuthCacheInvalidation(ctx context.Context, key string) error { + if c.fault.Load() == 3 { + return errors.New("injected Redis broadcast failure") + } + return c.client.Publish(ctx, "billing-auth-test", key).Err() +} +func (c *checkedBillingAuthRedis) SubscribeAuthCacheInvalidation(ctx context.Context, handler func(string)) error { + sub := c.client.Subscribe(ctx, "billing-auth-test") + defer func() { _ = sub.Close() }() + if _, err := sub.Receive(ctx); err != nil { + return err + } + NotifyAuthCacheSubscriptionReady(ctx) + c.ready <- struct{}{} + for { + select { + case <-ctx.Done(): + return ctx.Err() + case msg, ok := <-sub.Channel(): + if !ok { + return errors.New("test subscriber closed") + } + handler(msg.Payload) + } + } +} + +type exhaustedBillingReceipt struct { + UsageBillingRepository + debits int + applied bool + complete bool +} + +func (*exhaustedBillingReceipt) StageUsageBilling(context.Context, *UsageBillingCommand, *UsageLog) error { + return nil +} +func (r *exhaustedBillingReceipt) PendingUsageBilling(context.Context, int) ([]UsageBillingEventKey, error) { + if r.complete { + return nil, nil + } + return []UsageBillingEventKey{{RequestID: "exhausted", APIKeyID: 7}}, nil +} +func (r *exhaustedBillingReceipt) ApplyStagedUsageBilling(context.Context, UsageBillingEventKey) (*UsageBillingRecovery, error) { + fresh := !r.applied + if fresh { + r.debits++ + r.applied = true + } + return &UsageBillingRecovery{Event: UsageBillingEvent{Command: UsageBillingCommand{UserID: 9}}, + Result: UsageBillingApplyResult{Applied: fresh, APIKeyQuotaExhausted: true}, NeedsCache: true}, nil +} +func (r *exhaustedBillingReceipt) AcknowledgeUsageBillingCache(context.Context, UsageBillingEventKey) error { + r.complete = true + return nil +} + +func TestExhaustedBillingCheckedInvalidationFailuresRetryAcrossTwoNodes(t *testing.T) { + for _, failure := range []int32{1, 2, 3} { + t.Run(map[int32]string{1: "lookup", 2: "delete", 3: "broadcast"}[failure], func(t *testing.T) { + ctx := context.Background() + server := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: server.Addr()}) + var fault atomic.Int32 + cache := &checkedBillingAuthRedis{client: client, fault: &fault, ready: make(chan struct{}, 2)} + repo := &authRepoStub{listKeysByUserID: func(context.Context, int64) ([]string, error) { + if fault.Load() == 1 { + return nil, errors.New("injected key lookup failure") + } + return []string{"synthetic-key"}, nil + }} + cfg := &config.Config{APIKeyAuth: config.APIKeyAuthCacheConfig{L1Size: 128, L1TTLSeconds: 60, L2TTLSeconds: 60}} + a := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg) + b := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg) + t.Cleanup(func() { + a.StopAuthCacheInvalidationSubscriber() + b.StopAuthCacheInvalidationSubscriber() + _ = client.Close() + }) + a.StartAuthCacheInvalidationSubscriber(ctx) + b.StartAuthCacheInvalidationSubscriber(ctx) + for i := 0; i < 2; i++ { + select { + case <-cache.ready: + case <-time.After(3 * time.Second): + t.Fatal("both auth nodes must subscribe before invalidation") + } + } + key := a.authCacheKey("synthetic-key") + for _, node := range []*APIKeyService{a, b} { + require.NotNil(t, node.authCacheL1) + node.setAuthCacheL1(key, &APIKeyAuthCacheEntry{}) + node.authCacheL1.Wait() + _, ok := node.authCacheL1.Get(key) + require.True(t, ok) + } + require.NoError(t, client.Set(ctx, key, "cached", time.Minute).Err()) + fault.Store(failure) + receipt := &exhaustedBillingReceipt{} + require.Error(t, recoverUsageBillingBatch(ctx, receipt, nil, a)) + require.False(t, receipt.complete, "lookup/delete/broadcast errors must not acknowledge completion") + require.Equal(t, 1, receipt.debits) + _, stillCached := b.authCacheL1.Get(key) + require.True(t, stillCached, "failed broadcast must not masquerade as remote invalidation") + fault.Store(0) + require.NoError(t, recoverUsageBillingBatch(ctx, receipt, nil, a)) + require.True(t, receipt.complete) + require.Equal(t, 1, receipt.debits, "cache recovery must not debit again") + require.Eventually(t, func() bool { _, ok := b.authCacheL1.Get(key); return !ok }, 3*time.Second, time.Millisecond) + require.Equal(t, int64(0), client.Exists(ctx, key).Val()) + }) + } +} + +func TestExhaustedBillingMissingCheckedInvalidatorRemainsPending(t *testing.T) { + err := invalidateDurableUsageBilling(context.Background(), &UsageBillingRecovery{ + Result: UsageBillingApplyResult{APIKeyQuotaExhausted: true}, NeedsCache: true, + }, nil, nil) + require.ErrorContains(t, err, "checked auth cache invalidation unavailable") +} diff --git a/backend/internal/service/usage_billing_canonical_test.go b/backend/internal/service/usage_billing_canonical_test.go new file mode 100644 index 000000000000..4cccd769a50a --- /dev/null +++ b/backend/internal/service/usage_billing_canonical_test.go @@ -0,0 +1,28 @@ +//go:build unit + +package service + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestUsageBillingCanonicalLogKeepsOriginalFingerprintAndAccountCost(t *testing.T) { + raw := 0.000078125 + cost := &CostBreakdown{ActualCost: raw, TotalCost: 0.0000625} + p := &postUsageBillingParams{Cost: cost, User: &User{ID: 1}, APIKey: &APIKey{ID: 2}, + Account: &Account{ID: 3, Type: AccountTypeAPIKey}, AccountRateMultiplier: 2} + usage := &UsageLog{RequestID: "precision", UserID: 1, APIKeyID: 2, AccountID: 3, ActualCost: raw} + wantFingerprint := buildUsageBillingCommand("precision", usage, p).RequestFingerprint + repo := &openAIRecordUsageBillingRepoStub{result: &UsageBillingApplyResult{Applied: false}} + _, err := applyUsageBilling(context.Background(), "precision", usage, p, &billingDeps{deferredService: &DeferredService{}}, repo) + require.NoError(t, err) + require.Equal(t, 0.00007813, usage.ActualCost) + require.Equal(t, usage.ActualCost, repo.lastCmd.BalanceCost) + require.Equal(t, raw, *repo.lastCmd.QuotedActualCost) + require.Equal(t, wantFingerprint, repo.lastCmd.RequestFingerprint) + require.Equal(t, raw, cost.ActualCost, "the original pricing breakdown is immutable") + require.Equal(t, 0.0000625, cost.TotalCost, "account cost keeps its separate basis") +} diff --git a/backend/internal/service/usage_billing_delivery.go b/backend/internal/service/usage_billing_delivery.go new file mode 100644 index 000000000000..3517fb2e763f --- /dev/null +++ b/backend/internal/service/usage_billing_delivery.go @@ -0,0 +1,189 @@ +package service + +import ( + "context" + "errors" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" +) + +var ErrUsageBillingOutcomeUnknown = errors.New("usage billing outcome is unknown; durable state must be checked") +var ErrUsageBillingEventVersion = errors.New("unsupported usage billing event version") + +type UsageBillingEventKey struct { + RequestID string + APIKeyID int64 +} + +// The delivery snapshot contains only billing metadata, never relation objects +// (which may contain API keys, account credentials or customer profile data). +type UsageBillingEvent struct { + Version int `json:"version"` + Command UsageBillingCommand `json:"command"` + Usage UsageLog `json:"usage"` +} + +type UsageBillingRecovery struct { + Event UsageBillingEvent + Result UsageBillingApplyResult + NeedsCache bool +} + +// Stage is a separate durable boundary before monetary execution. A crash after +// it commits can be recovered without the handler closure or another model call. +type DurableUsageBillingRepository interface { + StageUsageBilling(context.Context, *UsageBillingCommand, *UsageLog) error + ApplyStagedUsageBilling(context.Context, UsageBillingEventKey) (*UsageBillingRecovery, error) + PendingUsageBilling(context.Context, int) ([]UsageBillingEventKey, error) + AcknowledgeUsageBillingCache(context.Context, UsageBillingEventKey) error +} + +func (s *GatewayService) DurableUsageBillingEnabled() bool { + if s == nil { + return false + } + _, ok := s.usageBillingRepo.(DurableUsageBillingRepository) + return ok +} + +func (s *OpenAIGatewayService) DurableUsageBillingEnabled() bool { + if s == nil || s.regionalMetering != nil { + return false + } + _, ok := s.usageBillingRepo.(DurableUsageBillingRepository) + return ok +} + +// Invalidations can safely be repeated after a crash. Never replay a Redis +// monetary delta when the caller cannot know whether it was applied previously. +func invalidateDurableUsageBilling(ctx context.Context, recovered *UsageBillingRecovery, cache *BillingCacheService, auth APIKeyAuthCacheInvalidator) error { + if recovered == nil || !recovered.NeedsCache { + return nil + } + e := &recovered.Event + var err error + if cache != nil { + if e.Command.BalanceCost > 0 { + err = errors.Join(err, cache.InvalidateUserBalance(ctx, e.Command.UserID)) + } + if e.Command.SubscriptionCost > 0 && e.Usage.GroupID != nil { + err = errors.Join(err, cache.InvalidateSubscription(ctx, e.Command.UserID, *e.Usage.GroupID)) + } + if e.Command.APIKeyRateLimitCost > 0 { + err = errors.Join(err, cache.InvalidateAPIKeyRateLimit(ctx, e.Command.APIKeyID)) + } + if e.Command.PlatformQuotaCost > 0 { + err = errors.Join(err, cache.InvalidateDurablePlatformQuota(ctx, e.Command.UserID, e.Command.PlatformQuotaPlatform)) + } + } + if recovered.Result.APIKeyQuotaExhausted { + checked, ok := auth.(interface { + InvalidateAuthCacheByUserIDChecked(context.Context, int64) error + }) + if !ok { + err = errors.Join(err, errors.New("checked auth cache invalidation unavailable; receipt remains pending")) + } else { + err = errors.Join(err, checked.InvalidateAuthCacheByUserIDChecked(ctx, e.Command.UserID)) + } + } + return err +} + +func applyDurableUsageBilling(ctx context.Context, repo DurableUsageBillingRepository, cmd *UsageBillingCommand, usage *UsageLog, p *postUsageBillingParams, deps *billingDeps) (bool, error) { + if err := repo.StageUsageBilling(ctx, cmd, usage); err != nil { + return false, err + } + key := UsageBillingEventKey{RequestID: cmd.RequestID, APIKeyID: cmd.APIKeyID} + recovered, err := repo.ApplyStagedUsageBilling(ctx, key) + if err != nil { + return false, err + } + auth, _ := p.APIKeyService.(APIKeyAuthCacheInvalidator) + if err := invalidateDurableUsageBilling(ctx, recovered, deps.billingCacheService, auth); err != nil { + // The committed canonical usage is already authoritative. Cache work + // stays pending for the recovery worker; never call legacy billing. + deferUsageBillingCache(repo, ctx, key) + return false, err + } + if recovered != nil && recovered.NeedsCache { + if err := repo.AcknowledgeUsageBillingCache(ctx, key); err != nil { + return false, err + } + } + if deps.deferredService != nil && p.Account != nil { + deps.deferredService.ScheduleLastUsedUpdate(p.Account.ID) + } + if recovered != nil && recovered.Result.Applied && !p.SimpleModeKeyRateLimitOnly { + go notifyBalanceLow(p, deps, &recovered.Result) + go notifyAccountQuota(p, deps, &recovered.Result) + } + return recovered != nil && recovered.Result.Applied, nil +} + +// This provider deliberately shares the pool's shutdown/join boundary. It does +// not start financial work merely by constructing a repository in a test/tool. +func ProvideUsageRecordWorkerPool(cfg *config.Config, billing UsageBillingRepository, cache *BillingCacheService, auth APIKeyAuthCacheInvalidator) *UsageRecordWorkerPool { + p := NewUsageRecordWorkerPool(cfg) + repo, ok := billing.(DurableUsageBillingRepository) + if !ok { + return p + } + ctx, cancel := context.WithCancel(context.Background()) + p.billingRecoveryCancel = cancel + p.lifecycleWg.Add(1) + go func() { + defer p.lifecycleWg.Done() + ticker := time.NewTicker(time.Second) + defer ticker.Stop() + for { + cycle, stop := context.WithTimeout(ctx, 5*time.Second) + if err := recoverUsageBillingBatch(cycle, repo, cache, auth); err != nil && ctx.Err() == nil { + logger.LegacyPrintf("usage_billing.recovery", "durable billing cycle incomplete: %v", err) + } + stop() + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + } + }() + return p +} + +func recoverUsageBillingBatch(ctx context.Context, repo DurableUsageBillingRepository, cache *BillingCacheService, auth APIKeyAuthCacheInvalidator) error { + keys, err := repo.PendingUsageBilling(ctx, 64) + if err != nil { + return err + } + var failures error + for _, key := range keys { + if ctx.Err() != nil { + return errors.Join(failures, ctx.Err()) + } + r, err := repo.ApplyStagedUsageBilling(ctx, key) + if err == nil { + err = invalidateDurableUsageBilling(ctx, r, cache, auth) + } + if err == nil && r != nil && r.NeedsCache { + err = repo.AcknowledgeUsageBillingCache(ctx, key) + } + if err != nil && r != nil && r.NeedsCache { + deferUsageBillingCache(repo, ctx, key) + } + failures = errors.Join(failures, err) + } + return failures +} + +func deferUsageBillingCache(repo DurableUsageBillingRepository, parent context.Context, key UsageBillingEventKey) { + if retry, ok := repo.(interface { + DeferUsageBillingCache(context.Context, UsageBillingEventKey) error + }); ok { + ctx, cancel := context.WithTimeout(context.WithoutCancel(parent), time.Second) + defer cancel() + _ = retry.DeferUsageBillingCache(ctx, key) + } +} diff --git a/backend/internal/service/usage_billing_recovery_cache_test.go b/backend/internal/service/usage_billing_recovery_cache_test.go new file mode 100644 index 000000000000..2b3d7d550498 --- /dev/null +++ b/backend/internal/service/usage_billing_recovery_cache_test.go @@ -0,0 +1,66 @@ +//go:build unit + +package service + +import ( + "context" + "errors" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +type recoveryCacheFailure struct{ BillingCache } + +func (*recoveryCacheFailure) InvalidateUserBalance(_ context.Context, user int64) error { + if user == 1 { + return errors.New("injected cache failure") + } + return nil +} + +type recoveryCacheRepo struct { + UsageBillingRepository + seen map[string]bool + acked []string + deferred []string +} + +func (*recoveryCacheRepo) StageUsageBilling(context.Context, *UsageBillingCommand, *UsageLog) error { + return nil +} +func (*recoveryCacheRepo) PendingUsageBilling(context.Context, int) ([]UsageBillingEventKey, error) { + return []UsageBillingEventKey{{RequestID: "first", APIKeyID: 1}, {RequestID: "second", APIKeyID: 2}}, nil +} +func (r *recoveryCacheRepo) ApplyStagedUsageBilling(_ context.Context, key UsageBillingEventKey) (*UsageBillingRecovery, error) { + fresh := !r.seen[key.RequestID] + r.seen[key.RequestID] = true + return &UsageBillingRecovery{Event: UsageBillingEvent{Command: UsageBillingCommand{UserID: key.APIKeyID, BalanceCost: 1}}, + Result: UsageBillingApplyResult{Applied: fresh}, NeedsCache: true}, nil +} +func (r *recoveryCacheRepo) AcknowledgeUsageBillingCache(_ context.Context, key UsageBillingEventKey) error { + r.acked = append(r.acked, key.RequestID) + return nil +} +func (r *recoveryCacheRepo) DeferUsageBillingCache(_ context.Context, key UsageBillingEventKey) error { + r.deferred = append(r.deferred, key.RequestID) + return nil +} + +func TestUsageRecoveryCacheFailureDoesNotStopAnotherEventOrReplayDeltas(t *testing.T) { + for _, flusher := range []bool{false, true} { + r := &recoveryCacheRepo{seen: map[string]bool{}} + cfg := &config.Config{} + cfg.Database.UserPlatformQuotaFlusherEnabled = flusher + cache := &BillingCacheService{cache: &recoveryCacheFailure{}, cfg: cfg} + for attempt := 0; attempt < 2; attempt++ { + require.Error(t, recoverUsageBillingBatch(context.Background(), r, cache, nil)) + } + require.Equal(t, []string{"second", "second"}, r.acked) + require.Equal(t, []string{"first", "first"}, r.deferred) + require.Len(t, r.seen, 2) + // The embedded nil BillingCache panics on a monetary delta method. This + // path only invalidates, under both legacy flusher configurations. + } +} diff --git a/backend/internal/service/usage_record_worker_pool.go b/backend/internal/service/usage_record_worker_pool.go index babcc2599e02..86213904b812 100644 --- a/backend/internal/service/usage_record_worker_pool.go +++ b/backend/internal/service/usage_record_worker_pool.go @@ -111,6 +111,7 @@ type UsageRecordWorkerPool struct { autoScaleCooldown time.Duration lastScaleNanos atomic.Int64 autoScaleCancel context.CancelFunc + billingRecoveryCancel context.CancelFunc lifecycleWg sync.WaitGroup stopOnce sync.Once } @@ -222,6 +223,9 @@ func (p *UsageRecordWorkerPool) Stop() { if p.autoScaleCancel != nil { p.autoScaleCancel() } + if p.billingRecoveryCancel != nil { + p.billingRecoveryCancel() + } p.lifecycleWg.Wait() p.pool.StopAndWait() }) diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index 9c6bfe4e58ee..ef199957730b 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -1099,7 +1099,7 @@ var ProviderSet = wire.NewSet( wire.Bind(new(DefaultSubscriptionAssigner), new(*SubscriptionService)), ProvideConcurrencyService, ProvideUserMessageQueueService, - NewUsageRecordWorkerPool, + ProvideUsageRecordWorkerPool, ProvideSchedulerSnapshotService, NewIdentityService, ProvideCRSSyncService, diff --git a/backend/migrations/302_usage_billing_quotes.sql b/backend/migrations/302_usage_billing_quotes.sql new file mode 100644 index 000000000000..917702c80c10 --- /dev/null +++ b/backend/migrations/302_usage_billing_quotes.sql @@ -0,0 +1,10 @@ +-- New charges only: no historical quote backfill or monetary adjustment. +CREATE TABLE usage_billing_quotes ( + request_id VARCHAR(255) NOT NULL, + api_key_id BIGINT NOT NULL, + request_fingerprint VARCHAR(64) NOT NULL, + quoted_actual_cost NUMERIC(20,10) NOT NULL, + canonical_actual_cost NUMERIC(20,8) NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + PRIMARY KEY (request_id, api_key_id) +); diff --git a/backend/migrations/303_usage_billing_delivery.sql b/backend/migrations/303_usage_billing_delivery.sql new file mode 100644 index 000000000000..34674ccce30c --- /dev/null +++ b/backend/migrations/303_usage_billing_delivery.sql @@ -0,0 +1,38 @@ +-- Only newly observed, explicitly staged events are recovered. This migration +-- does not infer or charge historical usage, or change any monetary row. +CREATE TABLE usage_billing_events ( + request_id VARCHAR(255) NOT NULL, + api_key_id BIGINT NOT NULL, + request_fingerprint VARCHAR(64) NOT NULL, + schema_version INTEGER NOT NULL CHECK (schema_version > 0), + payload JSONB NOT NULL CHECK (jsonb_typeof(payload) = 'object'), + state TEXT NOT NULL DEFAULT 'pending' + CHECK (state IN ('pending','applied','complete','legacy','blocked')), + result JSONB, + attempts INTEGER NOT NULL DEFAULT 0, + failure_code TEXT, + next_attempt_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + observed_at TIMESTAMPTZ NOT NULL, + staged_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + applied_at TIMESTAMPTZ, + completed_at TIMESTAMPTZ, + PRIMARY KEY (request_id, api_key_id) +); +CREATE INDEX usage_billing_events_pending ON usage_billing_events(next_attempt_at,staged_at) + WHERE state IN ('pending','applied'); + +CREATE FUNCTION reject_usage_billing_event_identity_mutation() RETURNS trigger LANGUAGE plpgsql AS $$ +BEGIN + IF OLD.request_id IS DISTINCT FROM NEW.request_id + OR OLD.api_key_id IS DISTINCT FROM NEW.api_key_id + OR OLD.request_fingerprint IS DISTINCT FROM NEW.request_fingerprint + OR OLD.schema_version IS DISTINCT FROM NEW.schema_version + OR OLD.payload IS DISTINCT FROM NEW.payload + OR OLD.observed_at IS DISTINCT FROM NEW.observed_at THEN + RAISE EXCEPTION 'usage billing event identity is immutable'; + END IF; + RETURN NEW; +END; +$$; +CREATE TRIGGER usage_billing_event_identity_immutable BEFORE UPDATE ON usage_billing_events + FOR EACH ROW EXECUTE FUNCTION reject_usage_billing_event_identity_mutation(); diff --git a/backend/migrations/304_usage_platform_quota_deltas.sql b/backend/migrations/304_usage_platform_quota_deltas.sql new file mode 100644 index 000000000000..fbe14781f45c --- /dev/null +++ b/backend/migrations/304_usage_platform_quota_deltas.sql @@ -0,0 +1,31 @@ +-- New request charges have an independent durable projection. The legacy +-- Redis flusher may still finish old snapshots into user_platform_quotas; +-- it cannot overwrite these new increments. No historical backfill/import. +CREATE TABLE usage_platform_quota_totals ( + quota_id BIGINT PRIMARY KEY REFERENCES user_platform_quotas(id) ON DELETE CASCADE, + daily_usage_usd NUMERIC(20,8) NOT NULL DEFAULT 0 CHECK(daily_usage_usd>=0), + weekly_usage_usd NUMERIC(20,8) NOT NULL DEFAULT 0 CHECK(weekly_usage_usd>=0), + monthly_usage_usd NUMERIC(20,8) NOT NULL DEFAULT 0 CHECK(monthly_usage_usd>=0), + legacy_daily_anchor TIMESTAMPTZ NOT NULL, + legacy_weekly_anchor TIMESTAMPTZ NOT NULL, + legacy_monthly_anchor TIMESTAMPTZ NOT NULL, + daily_window_start TIMESTAMPTZ NOT NULL, + weekly_window_start TIMESTAMPTZ NOT NULL, + monthly_window_start TIMESTAMPTZ NOT NULL +); + +CREATE VIEW usage_platform_quota_effective AS +SELECT q.id,q.user_id,q.platform,q.daily_limit_usd,q.weekly_limit_usd,q.monthly_limit_usd,q.deleted_at, + CASE WHEN p.quota_id IS NULL THEN q.daily_usage_usd + ELSE CASE WHEN COALESCE(q.daily_window_start,p.legacy_daily_anchor) IS NULL OR COALESCE(q.daily_window_start,p.legacy_daily_anchor)>=p.daily_window_start THEN q.daily_usage_usd ELSE 0 END + + CASE WHEN COALESCE(q.daily_window_start,p.legacy_daily_anchor) IS NULL OR p.daily_window_start>=COALESCE(q.daily_window_start,p.legacy_daily_anchor) THEN p.daily_usage_usd ELSE 0 END END AS daily_usage_usd, + CASE WHEN p.quota_id IS NULL THEN q.weekly_usage_usd + ELSE CASE WHEN COALESCE(q.weekly_window_start,p.legacy_weekly_anchor) IS NULL OR COALESCE(q.weekly_window_start,p.legacy_weekly_anchor)>=p.weekly_window_start THEN q.weekly_usage_usd ELSE 0 END + + CASE WHEN COALESCE(q.weekly_window_start,p.legacy_weekly_anchor) IS NULL OR p.weekly_window_start>=COALESCE(q.weekly_window_start,p.legacy_weekly_anchor) THEN p.weekly_usage_usd ELSE 0 END END AS weekly_usage_usd, + CASE WHEN p.quota_id IS NULL THEN q.monthly_usage_usd + ELSE CASE WHEN COALESCE(q.monthly_window_start,p.legacy_monthly_anchor) IS NULL OR COALESCE(q.monthly_window_start,p.legacy_monthly_anchor)>=p.monthly_window_start THEN q.monthly_usage_usd ELSE 0 END + + CASE WHEN COALESCE(q.monthly_window_start,p.legacy_monthly_anchor) IS NULL OR p.monthly_window_start>=COALESCE(q.monthly_window_start,p.legacy_monthly_anchor) THEN p.monthly_usage_usd ELSE 0 END END AS monthly_usage_usd, + GREATEST(COALESCE(q.daily_window_start,p.legacy_daily_anchor),p.daily_window_start) AS daily_window_start, + GREATEST(COALESCE(q.weekly_window_start,p.legacy_weekly_anchor),p.weekly_window_start) AS weekly_window_start, + GREATEST(COALESCE(q.monthly_window_start,p.legacy_monthly_anchor),p.monthly_window_start) AS monthly_window_start +FROM user_platform_quotas q LEFT JOIN usage_platform_quota_totals p ON p.quota_id=q.id; diff --git a/backend/migrations/305_batch_image_owned_money.sql b/backend/migrations/305_batch_image_owned_money.sql new file mode 100644 index 000000000000..ddd916fe2d4f --- /dev/null +++ b/backend/migrations/305_batch_image_owned_money.sql @@ -0,0 +1,10 @@ +-- Version-2 jobs only. Never infer historical holds from pooled frozen money. +CREATE TABLE batch_image_money_holds ( + batch_id VARCHAR(64) NOT NULL, + api_key_id BIGINT NOT NULL, + user_id BIGINT NOT NULL, + hold_amount NUMERIC(20,8) NOT NULL CHECK (hold_amount>=0), + actual_amount NUMERIC(20,8), + state TEXT NOT NULL CHECK (state IN ('held','captured','released')), + PRIMARY KEY (batch_id,api_key_id) +); diff --git a/backend/migrations/306_usage_billing_quote_precision.sql b/backend/migrations/306_usage_billing_quote_precision.sql new file mode 100644 index 000000000000..fb080150b612 --- /dev/null +++ b/backend/migrations/306_usage_billing_quote_precision.sql @@ -0,0 +1,9 @@ +-- Removing the typmod preserves all decimal digits of new quote snapshots. +-- Existing ten-place values cannot be made lossless retrospectively: version 0 +-- explicitly marks their provenance as unknown/rounded, without backfilling it. +ALTER TABLE usage_billing_quotes ALTER COLUMN quoted_actual_cost TYPE NUMERIC; +ALTER TABLE usage_billing_quotes ADD COLUMN quote_encoding_version SMALLINT NOT NULL DEFAULT 0; +COMMENT ON COLUMN usage_billing_quotes.quoted_actual_cost IS + 'Version 1: exact shortest decimal round-trip of the original float64 quotation; version 0 may already be rounded'; +COMMENT ON COLUMN usage_billing_quotes.canonical_actual_cost IS + 'Committed normalized customer amount from the billing command, never recomputed from a stored quote'; diff --git a/docs/billing-audit-recovery.md b/docs/billing-audit-recovery.md new file mode 100644 index 000000000000..d530d296dfc4 --- /dev/null +++ b/docs/billing-audit-recovery.md @@ -0,0 +1,91 @@ +# Billing audit recovery boundary + +New batch image jobs use pricing snapshot version 2. Their complete hold and +actual charge are rounded once to eight USD decimal places with the existing +half-away-from-zero monetary rule. Prices, discounts and account cost rates are +unchanged. The persisted hold is the exact amount moved between spendable and +frozen balance. Version 0/1 jobs with an exact eight-place hold retain their +existing behavior. + +Version-2 holds also carry individual ownership and a held/captured/released +receipt, committed with the money operation. Capture and release are mutually +exclusive even when another job has enough money in the user's frozen pool. +Free jobs retain their zero-money behavior. Existing holds are not backfilled. + +A legacy amount with more than eight places cannot execute a new money change. +A retry first makes a read-only hot/archive receipt check using the original +operation ID, key, owner, batch, amounts and payload fingerprint. An exact match +confirms that operation already committed, so the task status can catch up +without moving money. A missing/conflicting receipt fails; a fingerprint conflict +is never itself evidence that a refund succeeded. No pooled balance inference, +historical rounding, backfill, refund or back-charge is supplied. + +Ordinary usage is staged synchronously as a versioned immutable event before +monetary execution. The money effects, dedup claim, raw quote, canonical usage +and applied receipt commit together. A worker retries only pending events and +idempotent cache invalidations. A lost commit response is checked against the +durable receipt before any replay. Hot and archived legacy claims remain +authoritative; they are not permission to fabricate history or charge again. + +There remains an explicit observation boundary: a process can die after the +upstream produces usage but before the complete event has committed to the +journal. A database unavailable at that boundary also prevents durable delivery. +The application cannot reconstruct final usage it never persisted and does not +call the model again or invent a zero-cost success. This change removes the +in-memory queue window and recovers events committed before execution; it does +not promise that this pre-storage window or unknown upstream usage is solved. + +Explicit recovery would first need an authorized reconciliation of the original +hold transaction, all other holds for that user, captures/releases and current +balances. An individually reviewed correction would require separate approval, +an immutable correction identity and a transaction proving both ownership and +conservation. This PR supplies no operator action to perform that correction. + +New ordinary platform-quota increments commit with the charge and dedup claim, +independently of the legacy flusher flag. `usage_platform_quota_totals` keeps the +new projection separate from the existing baseline. The effective view combines +only matching current windows; old events cannot rewind a newer window. Explicit +administrator resets update both components in one transaction. Quota removal +and recreation use different quota row identities. + +The combined read cache has its own namespace. It is invalidated idempotently; +new billing never replays Redis monetary deltas. Legacy dirty entries retain +their original namespace so the old flusher can finish its baseline snapshots +without overwriting or double-counting new deltas. Missing/old-schema legacy +increments return an explicit not-applied error. No historical cache values are +imported, no missing old charges are inferred, and no balance adjustments occur +at migration time. Existing legacy backlog requires separate evidence and must +not be described as repaired by this PR. + +The customer quote is stored in scale-free NUMERIC as the shortest exact decimal +round-trip of the original float64 quotation. `quote_encoding_version=1` marks +that representation; version 0 may have already lost digits under the old +ten-place schema and is not retrospectively repaired. `canonical_actual_cost` +comes from the already-normalized customer command, never a re-rounded stored +quote. Canonical +usage and cache reloads use that same customer amount. `total_cost` and the +account cost multiplier retain their independent meaning. A subscription or +simple-mode key-window charge is not new cash revenue. + +Mixed old financial writers/binaries are a release blocker, including a rolling +period in which an old batch worker can claim a version-2 job. Approval requires +an explicitly reviewed admission stop and drain of old in-flight requests, +usage workers, batch workers and flusher cycles (including already-popped +snapshots), then all relevant nodes on the same verified financial protocol and +migrations before new admission. Unknown started work remains occupied; expiry +is not authority to refund or repeat a model request. This document executes no +drain and supplies no automatic migration of old monetary values. + +Rollback to an old reader/writer is not approved just because all journal events +are complete: it still cannot account for new projections and owned holds. +A rollback build must understand those states, or a separate, explicitly +authorized reconciliation/drain plan must be reviewed. No command in this PR +merges historical money or projections into a legacy ledger. + +Quota window starts retain precise administrative reset cutoffs. Day/week +delivery uses the later of that cutoff and the natural boundary; an event before +the cutoff contributes zero to that window, even if it is the first projection +row. Ordinary and regional writers read the same effective monthly anchor under +the baseline row lock, and advance together only at the 30-day boundary. Checked +auth invalidation errors (including key lookup and broadcast) leave cache work +pending; an unchecked/absent invalidator cannot acknowledge key-quota exhaustion.