fix quota reset reservation race

server/internal/tenant/app/ai_point_usage.go

@@ -185,7 +185,7 @@ func CompleteAIPoints(
 	if err != nil || !pending {
 		return err
 	}
-	if err := confirmAIPointReservation(ctx, tx, tenantID, reservation.ReservationID); err != nil {
+	if err := confirmAIPointReservationIfPending(ctx, tx, tenantID, reservation.ReservationID); err != nil {
 		return err
 	}
 	var modelPtr *string
@@ -226,29 +226,31 @@ func RefundAIPoints(
 		return err
 	}
 
-	status, err := repository.NewQuotaRepository(tx).GetAIQuotaStatus(ctx, tenantID, time.Now().UTC())
+	refunded, err := refundAIPointReservationIfPending(ctx, tx, tenantID, reservation.ReservationID)
 	if err != nil {
 		return err
 	}
-
-	quotaRepo := repository.NewQuotaRepository(tx)
-	reason := "ai_points_refund"
-	referenceType := "ai_point_usage"
-	usageID := reservation.UsageLogID
-	if _, err := quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{
-		TenantID:      tenantID,
-		OperatorID:    operatorID,
-		QuotaType:     aiPointsQuotaType,
-		Delta:         reservation.Points,
-		BalanceAfter:  status.Balance + reservation.Points,
-		Reason:        &reason,
-		ReferenceType: &referenceType,
-		ReferenceID:   &usageID,
-	}); err != nil {
-		return err
-	}
-	if err := refundAIPointReservation(ctx, tx, tenantID, reservation.ReservationID); err != nil {
-		return err
+	if refunded {
+		status, err := repository.NewQuotaRepository(tx).GetAIQuotaStatus(ctx, tenantID, time.Now().UTC())
+		if err != nil {
+			return err
+		}
+		quotaRepo := repository.NewQuotaRepository(tx)
+		reason := "ai_points_refund"
+		referenceType := "ai_point_usage"
+		usageID := reservation.UsageLogID
+		if _, err := quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{
+			TenantID:      tenantID,
+			OperatorID:    operatorID,
+			QuotaType:     aiPointsQuotaType,
+			Delta:         reservation.Points,
+			BalanceAfter:  status.Balance + reservation.Points,
+			Reason:        &reason,
+			ReferenceType: &referenceType,
+			ReferenceID:   &usageID,
+		}); err != nil {
+			return err
+		}
 	}
 	if err := repository.NewAIPointUsageRepository(tx).MarkRefunded(ctx, tenantID, reservation.UsageLogID, strings.TrimSpace(errorMessage)); err != nil {
 		return err
@@ -375,7 +377,7 @@ func lockPendingAIPointUsage(ctx context.Context, tx pgx.Tx, tenantID, usageLogI
 	return status == "pending", nil
 }
 
-func confirmAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error {
+func confirmAIPointReservationIfPending(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error {
 	tag, err := tx.Exec(ctx, `
 	UPDATE quota_reservations
 	SET status = 'confirmed',
@@ -389,13 +391,10 @@ func confirmAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reserva
 	if err != nil {
 		return err
 	}
-	if tag.RowsAffected() == 0 {
-		return errors.New("ai points reservation is not pending")
-	}
-	return nil
+	return ignoreAIPointReservationReset(ctx, tx, tag.RowsAffected() > 0, tenantID, reservationID)
 }
 
-func refundAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error {
+func refundAIPointReservationIfPending(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) (bool, error) {
 	tag, err := tx.Exec(ctx, `
 	UPDATE quota_reservations
 	SET status = 'refunded',
@@ -406,11 +405,32 @@ func refundAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservat
 	  AND quota_type = $3
 	  AND status = 'pending'
 	`, reservationID, tenantID, aiPointsQuotaType)
+	if err != nil {
+		return false, err
+	}
+	if tag.RowsAffected() > 0 {
+		return true, nil
+	}
+	return false, ignoreAIPointReservationReset(ctx, tx, false, tenantID, reservationID)
+}
+
+func ignoreAIPointReservationReset(ctx context.Context, tx pgx.Tx, updated bool, tenantID, reservationID int64) error {
+	if updated {
+		return nil
+	}
+	var status string
+	err := tx.QueryRow(ctx, `
+	SELECT status
+	FROM quota_reservations
+	WHERE id = $1
+	  AND tenant_id = $2
+	  AND quota_type = $3
+	`, reservationID, tenantID, aiPointsQuotaType).Scan(&status)
 	if err != nil {
 		return err
 	}
-	if tag.RowsAffected() == 0 {
-		return errors.New("ai points reservation is not pending")
+	if status == "reset" {
+		return nil
 	}
-	return nil
+	return errors.New("ai points reservation is not pending")
 }

server/internal/tenant/app/article_imitation_service.go

@@ -512,27 +512,29 @@ func (s *ArticleImitationService) failGeneration(ctx context.Context, job articl
 
 	if job.ReservationID > 0 {
 		balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation")
-		if balanceErr != nil {
-			cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
-		} else {
-			reason := "generation_refund"
-			referenceType := "generation_task"
-			taskReferenceID := job.TaskID
-			if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
-				TenantID:      job.TenantID,
-				OperatorID:    job.OperatorID,
-				QuotaType:     "article_generation",
-				Delta:         1,
-				BalanceAfter:  balance + 1,
-				Reason:        &reason,
-				ReferenceType: &referenceType,
-				ReferenceID:   &taskReferenceID,
-			}); err != nil {
-				cleanupErr = errors.Join(cleanupErr, fmt.Errorf("insert generation refund ledger: %w", err))
-			}
-		}
-		if err := quotaRepo.RefundReservation(cleanupCtx, job.ReservationID, job.TenantID); err != nil {
+		refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, job.ReservationID, job.TenantID)
+		if err != nil {
 			cleanupErr = errors.Join(cleanupErr, fmt.Errorf("refund generation reservation: %w", err))
+		} else if refunded {
+			if balanceErr != nil {
+				cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
+			} else {
+				reason := "generation_refund"
+				referenceType := "generation_task"
+				taskReferenceID := job.TaskID
+				if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
+					TenantID:      job.TenantID,
+					OperatorID:    job.OperatorID,
+					QuotaType:     "article_generation",
+					Delta:         1,
+					BalanceAfter:  balance + 1,
+					Reason:        &reason,
+					ReferenceType: &referenceType,
+					ReferenceID:   &taskReferenceID,
+				}); err != nil {
+					cleanupErr = errors.Join(cleanupErr, fmt.Errorf("insert generation refund ledger: %w", err))
+				}
+			}
 		}
 	}
 

server/internal/tenant/app/generation_task_runtime.go

@@ -79,29 +79,34 @@ func HandleClaimedGenerationTaskFailure(
 	}
 
 	if task.QuotaReservationID > 0 {
-		balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation")
-		if balanceErr != nil {
-			failureErr = errors.Join(failureErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
-		} else if task.OperatorID > 0 {
-			reason := "generation_refund"
-			referenceType := "generation_task"
-			taskReferenceID := task.ID
-			if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
-				TenantID:      task.TenantID,
-				OperatorID:    task.OperatorID,
-				QuotaType:     "article_generation",
-				Delta:         1,
-				BalanceAfter:  balance + 1,
-				Reason:        &reason,
-				ReferenceType: &referenceType,
-				ReferenceID:   &taskReferenceID,
-			}); err != nil {
-				failureErr = errors.Join(failureErr, fmt.Errorf("insert generation refund ledger: %w", err))
-			}
+		var balance int
+		balanceErr := error(nil)
+		if task.OperatorID > 0 {
+			balance, balanceErr = quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation")
 		}
-
-		if err := quotaRepo.RefundReservation(cleanupCtx, task.QuotaReservationID, task.TenantID); err != nil {
+		refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, task.QuotaReservationID, task.TenantID)
+		if err != nil {
 			failureErr = errors.Join(failureErr, fmt.Errorf("refund generation reservation: %w", err))
+		} else if refunded {
+			if balanceErr != nil {
+				failureErr = errors.Join(failureErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
+			} else if task.OperatorID > 0 {
+				reason := "generation_refund"
+				referenceType := "generation_task"
+				taskReferenceID := task.ID
+				if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
+					TenantID:      task.TenantID,
+					OperatorID:    task.OperatorID,
+					QuotaType:     "article_generation",
+					Delta:         1,
+					BalanceAfter:  balance + 1,
+					Reason:        &reason,
+					ReferenceType: &referenceType,
+					ReferenceID:   &taskReferenceID,
+				}); err != nil {
+					failureErr = errors.Join(failureErr, fmt.Errorf("insert generation refund ledger: %w", err))
+				}
+			}
 		}
 	}
 

server/internal/tenant/app/kol_generation_service.go

@@ -445,28 +445,34 @@ func HandleKolGenerationTaskFailure(
 	}
 
 	if task.QuotaReservationID > 0 {
-		balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation")
-		if balanceErr != nil {
-			failureErr = errors.Join(failureErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
-		} else if task.OperatorID > 0 {
-			reason := "generation_refund"
-			referenceType := "generation_task"
-			taskReferenceID := task.ID
-			if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
-				TenantID:      task.TenantID,
-				OperatorID:    task.OperatorID,
-				QuotaType:     "article_generation",
-				Delta:         1,
-				BalanceAfter:  balance + 1,
-				Reason:        &reason,
-				ReferenceType: &referenceType,
-				ReferenceID:   &taskReferenceID,
-			}); err != nil {
-				failureErr = errors.Join(failureErr, fmt.Errorf("insert generation refund ledger: %w", err))
-			}
+		var balance int
+		balanceErr := error(nil)
+		if task.OperatorID > 0 {
+			balance, balanceErr = quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation")
 		}
-		if err := quotaRepo.RefundReservation(cleanupCtx, task.QuotaReservationID, task.TenantID); err != nil {
+		refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, task.QuotaReservationID, task.TenantID)
+		if err != nil {
 			failureErr = errors.Join(failureErr, fmt.Errorf("refund generation reservation: %w", err))
+		} else if refunded {
+			if balanceErr != nil {
+				failureErr = errors.Join(failureErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
+			} else if task.OperatorID > 0 {
+				reason := "generation_refund"
+				referenceType := "generation_task"
+				taskReferenceID := task.ID
+				if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
+					TenantID:      task.TenantID,
+					OperatorID:    task.OperatorID,
+					QuotaType:     "article_generation",
+					Delta:         1,
+					BalanceAfter:  balance + 1,
+					Reason:        &reason,
+					ReferenceType: &referenceType,
+					ReferenceID:   &taskReferenceID,
+				}); err != nil {
+					failureErr = errors.Join(failureErr, fmt.Errorf("insert generation refund ledger: %w", err))
+				}
+			}
 		}
 	}
 

server/internal/tenant/app/prompt_generate_service.go

@@ -739,28 +739,29 @@ func (s *PromptRuleGenerationService) failGeneration(ctx context.Context, job pr
 	}
 
 	balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation")
-	if balanceErr != nil {
-		cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
-	} else {
-		reason := "generation_refund"
-		referenceType := "generation_task"
-		taskReferenceID := job.TaskID
-		if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
-			TenantID:      job.TenantID,
-			OperatorID:    job.OperatorID,
-			QuotaType:     "article_generation",
-			Delta:         1,
-			BalanceAfter:  balance + 1,
-			Reason:        &reason,
-			ReferenceType: &referenceType,
-			ReferenceID:   &taskReferenceID,
-		}); err != nil {
-			cleanupErr = errors.Join(cleanupErr, fmt.Errorf("insert generation refund ledger: %w", err))
-		}
-	}
-
-	if err := quotaRepo.RefundReservation(cleanupCtx, job.ReservationID, job.TenantID); err != nil {
+	refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, job.ReservationID, job.TenantID)
+	if err != nil {
 		cleanupErr = errors.Join(cleanupErr, fmt.Errorf("refund generation reservation: %w", err))
+	} else if refunded {
+		if balanceErr != nil {
+			cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
+		} else {
+			reason := "generation_refund"
+			referenceType := "generation_task"
+			taskReferenceID := job.TaskID
+			if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
+				TenantID:      job.TenantID,
+				OperatorID:    job.OperatorID,
+				QuotaType:     "article_generation",
+				Delta:         1,
+				BalanceAfter:  balance + 1,
+				Reason:        &reason,
+				ReferenceType: &referenceType,
+				ReferenceID:   &taskReferenceID,
+			}); err != nil {
+				cleanupErr = errors.Join(cleanupErr, fmt.Errorf("insert generation refund ledger: %w", err))
+			}
+		}
 	}
 	invalidateArticleCaches(cleanupCtx, s.cache, job.TenantID, &job.ArticleID)
 	if cleanupErr != nil {

server/internal/tenant/app/template_service.go

@@ -626,28 +626,29 @@ func (s *TemplateService) failGeneration(ctx context.Context, job generationJob,
 	}
 
 	balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation")
-	if balanceErr != nil {
-		cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
-	} else {
-		reason := "generation_refund"
-		referenceType := "generation_task"
-		taskReferenceID := job.TaskID
-		if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
-			TenantID:      job.TenantID,
-			OperatorID:    job.OperatorID,
-			QuotaType:     "article_generation",
-			Delta:         1,
-			BalanceAfter:  balance + 1,
-			Reason:        &reason,
-			ReferenceType: &referenceType,
-			ReferenceID:   &taskReferenceID,
-		}); err != nil {
-			cleanupErr = errors.Join(cleanupErr, fmt.Errorf("insert generation refund ledger: %w", err))
-		}
-	}
-
-	if err := quotaRepo.RefundReservation(cleanupCtx, job.ReservationID, job.TenantID); err != nil {
+	refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, job.ReservationID, job.TenantID)
+	if err != nil {
 		cleanupErr = errors.Join(cleanupErr, fmt.Errorf("refund generation reservation: %w", err))
+	} else if refunded {
+		if balanceErr != nil {
+			cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
+		} else {
+			reason := "generation_refund"
+			referenceType := "generation_task"
+			taskReferenceID := job.TaskID
+			if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
+				TenantID:      job.TenantID,
+				OperatorID:    job.OperatorID,
+				QuotaType:     "article_generation",
+				Delta:         1,
+				BalanceAfter:  balance + 1,
+				Reason:        &reason,
+				ReferenceType: &referenceType,
+				ReferenceID:   &taskReferenceID,
+			}); err != nil {
+				cleanupErr = errors.Join(cleanupErr, fmt.Errorf("insert generation refund ledger: %w", err))
+			}
+		}
 	}
 	invalidateArticleCaches(cleanupCtx, s.cache, job.TenantID, &job.ArticleID)
 	if cleanupErr != nil {