diff --git a/backend/internal/repository/usage_billing_repo.go b/backend/internal/repository/usage_billing_repo.go index ce0adc2772..9732d52355 100644 --- a/backend/internal/repository/usage_billing_repo.go +++ b/backend/internal/repository/usage_billing_repo.go @@ -187,16 +187,18 @@ func (r *usageBillingRepository) applyUsageBillingEffects(ctx context.Context, t result.BalanceOverdrafted = !sufficient } + // API Key 在请求进行中被删除时，其自身的额度/限速计数已无意义，跳过即可； + // 但用户余额、订阅与账号额度必须照常结算，不能因此回滚整个计费事务。 if cmd.APIKeyQuotaCost > 0 { exhausted, err := incrementUsageBillingAPIKeyQuota(ctx, tx, cmd.APIKeyID, cmd.APIKeyQuotaCost) - if err != nil { + if err != nil && !errors.Is(err, service.ErrAPIKeyNotFound) { return err } result.APIKeyQuotaExhausted = exhausted } if cmd.APIKeyRateLimitCost > 0 { - if err := incrementUsageBillingAPIKeyRateLimit(ctx, tx, cmd.APIKeyID, cmd.APIKeyRateLimitCost); err != nil { + if err := incrementUsageBillingAPIKeyRateLimit(ctx, tx, cmd.APIKeyID, cmd.APIKeyRateLimitCost); err != nil && !errors.Is(err, service.ErrAPIKeyNotFound) { return err } } diff --git a/backend/internal/repository/usage_billing_repo_integration_test.go b/backend/internal/repository/usage_billing_repo_integration_test.go index e8d4d32707..9f2f0169dc 100644 --- a/backend/internal/repository/usage_billing_repo_integration_test.go +++ b/backend/internal/repository/usage_billing_repo_integration_test.go @@ -162,6 +162,61 @@ func TestUsageBillingRepositoryApply_RequestFingerprintConflict(t *testing.T) { require.ErrorIs(t, err, service.ErrUsageBillingRequestConflict) } +func TestUsageBillingRepositoryApply_DeletedAPIKeyStillBillsBalance(t *testing.T) { + ctx := context.Background() + client := testEntClient(t) + repo := NewUsageBillingRepository(client, integrationDB) + + user := mustCreateUser(t, client, &service.User{ + Email: fmt.Sprintf("usage-billing-deleted-key-user-%d@example.com", time.Now().UnixNano()), + PasswordHash: "hash", + Balance: 100, + }) + apiKey := mustCreateApiKey(t, client, &service.APIKey{ + UserID: user.ID, + Key: "sk-usage-billing-deleted-key-" + uuid.NewString(), + Name: "billing-deleted-key", + Quota: 50, + RateLimit5h: 50, + }) + account := mustCreateAccount(t, client, &service.Account{ + Name: "usage-billing-deleted-key-account-" + uuid.NewString(), + Type: service.AccountTypeAPIKey, + }) + + _, err := integrationDB.ExecContext(ctx, "UPDATE api_keys SET deleted_at = NOW() WHERE id = $1", apiKey.ID) + require.NoError(t, err) + + requestID := uuid.NewString() + result, err := repo.Apply(ctx, &service.UsageBillingCommand{ + RequestID: requestID, + APIKeyID: apiKey.ID, + UserID: user.ID, + AccountID: account.ID, + AccountType: service.AccountTypeAPIKey, + BalanceCost: 1.25, + APIKeyQuotaCost: 1.25, + APIKeyRateLimitCost: 1.25, + }) + require.NoError(t, err) + require.NotNil(t, result) + require.True(t, result.Applied) + require.False(t, result.APIKeyQuotaExhausted) + + var balance float64 + require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT balance FROM users WHERE id = $1", user.ID).Scan(&balance)) + require.InDelta(t, 98.75, balance, 0.000001) + + var quotaUsed, usage5h float64 + require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT quota_used, usage_5h FROM api_keys WHERE id = $1", apiKey.ID).Scan("aUsed, &usage5h)) + require.InDelta(t, 0, quotaUsed, 0.000001) + require.InDelta(t, 0, usage5h, 0.000001) + + var dedupCount int + require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM usage_billing_dedup WHERE request_id = $1 AND api_key_id = $2", requestID, apiKey.ID).Scan(&dedupCount)) + require.Equal(t, 1, dedupCount) +} + func TestUsageBillingRepositoryApply_UpdatesAccountQuota(t *testing.T) { ctx := context.Background() client := testEntClient(t) diff --git a/backend/internal/repository/usage_billing_repo_unit_test.go b/backend/internal/repository/usage_billing_repo_unit_test.go index 6fe4d4fc26..40394d819c 100644 --- a/backend/internal/repository/usage_billing_repo_unit_test.go +++ b/backend/internal/repository/usage_billing_repo_unit_test.go @@ -20,6 +20,8 @@ const ( captureBatchImageHoldSQL = `(?s)UPDATE users\s+SET balance = balance\s+\+ CASE WHEN \$1 > \$2 THEN \$1 - \$2 ELSE 0 END\s+- CASE WHEN \$2 > \$1 THEN \$2 - \$1 ELSE 0 END,\s+frozen_balance = COALESCE\(frozen_balance, 0\) - \$1,\s+updated_at = NOW\(\)\s+WHERE id = \$3 AND deleted_at IS NULL AND COALESCE\(frozen_balance, 0\) >= \$1\s+RETURNING balance, frozen_balance` releaseBatchImageHoldSQL = `(?s)UPDATE users\s+SET balance = balance \+ \$1,\s+frozen_balance = COALESCE\(frozen_balance, 0\) - \$1,\s+updated_at = NOW\(\)\s+WHERE id = \$2 AND deleted_at IS NULL AND COALESCE\(frozen_balance, 0\) >= \$1\s+RETURNING balance, frozen_balance` userExistsForBillingSQL = `(?s)SELECT 1\s+FROM users\s+WHERE id = \$1 AND deleted_at IS NULL` + apiKeyQuotaIncrementSQL = `(?s)UPDATE api_keys\s+SET quota_used = quota_used \+ \$1,.*WHERE id = \$2 AND deleted_at IS NULL\s+RETURNING` + apiKeyRateLimitIncrementSQL = `(?s)UPDATE api_keys SET\s+usage_5h = .*WHERE id = \$2 AND deleted_at IS NULL` ) func TestDeductUsageBillingBalance_UsesSufficientBalanceGuard(t *testing.T) { @@ -99,6 +101,66 @@ func TestApplyUsageBillingEffects_FlagsBalanceOverdraft(t *testing.T) { require.NoError(t, mock.ExpectationsWereMet()) } +func TestApplyUsageBillingEffects_DeletedAPIKeyStillBillsBalance(t *testing.T) { + ctx := context.Background() + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + mock.ExpectBegin() + tx, err := db.BeginTx(ctx, nil) + require.NoError(t, err) + mock.ExpectQuery(conditionalBalanceDeductSQL). + WithArgs(10.0, int64(42)). + WillReturnRows(sqlmock.NewRows([]string{"balance"}).AddRow(90.0)) + mock.ExpectQuery(apiKeyQuotaIncrementSQL). + WithArgs(10.0, int64(7), service.StatusAPIKeyActive, service.StatusAPIKeyQuotaExhausted). + WillReturnError(sql.ErrNoRows) + mock.ExpectExec(apiKeyRateLimitIncrementSQL). + WithArgs(10.0, int64(7)). + WillReturnResult(sqlmock.NewResult(0, 0)) + mock.ExpectCommit() + + result := &service.UsageBillingApplyResult{Applied: true} + err = (&usageBillingRepository{}).applyUsageBillingEffects(ctx, tx, &service.UsageBillingCommand{ + UserID: 42, + APIKeyID: 7, + BalanceCost: 10, + APIKeyQuotaCost: 10, + APIKeyRateLimitCost: 10, + }, result) + require.NoError(t, err) + require.NotNil(t, result.NewBalance) + require.InDelta(t, 90.0, *result.NewBalance, 0.000001) + require.False(t, result.APIKeyQuotaExhausted) + require.NoError(t, tx.Commit()) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestApplyUsageBillingEffects_APIKeyCounterErrorStillFails(t *testing.T) { + ctx := context.Background() + db, mock, err := sqlmock.New() + require.NoError(t, err) + defer func() { _ = db.Close() }() + + mock.ExpectBegin() + tx, err := db.BeginTx(ctx, nil) + require.NoError(t, err) + mock.ExpectExec(apiKeyRateLimitIncrementSQL). + WithArgs(10.0, int64(7)). + WillReturnError(sql.ErrConnDone) + mock.ExpectRollback() + + err = (&usageBillingRepository{}).applyUsageBillingEffects(ctx, tx, &service.UsageBillingCommand{ + UserID: 42, + APIKeyID: 7, + APIKeyRateLimitCost: 10, + }, &service.UsageBillingApplyResult{Applied: true}) + require.ErrorIs(t, err, sql.ErrConnDone) + require.NoError(t, tx.Rollback()) + require.NoError(t, mock.ExpectationsWereMet()) +} + func TestDeductUsageBillingBalance_ReturnsUserNotFoundWhenNoUserUpdated(t *testing.T) { ctx := context.Background() db, mock, err := sqlmock.New()