firecrawl/apps/api/src/services/billing/auto_charge.ts

// Import necessary dependencies and types
import { AuthCreditUsageChunk } from "../../controllers/v1/types";
import { getACUC, setCachedACUC } from "../../controllers/auth";
import { redlock } from "../redlock";
import { supabase_service } from "../supabase";
import { createPaymentIntent } from "./stripe";
import { issueCredits } from "./issue_credits";
import { sendNotification } from "../notification/email_notification";
import { NotificationType } from "../../types";
import { deleteKey, getValue, setValue } from "../redis";
import { sendSlackWebhook } from "../alerts/slack";
import { Logger } from "../../lib/logger";

// Define the number of credits to be added during auto-recharge
const AUTO_RECHARGE_CREDITS = 1000;
const AUTO_RECHARGE_COOLDOWN = 600; // 10 minutes in seconds

/**
 * Attempt to automatically charge a user's account when their credit balance falls below a threshold
 * @param chunk The user's current usage data
 * @param autoRechargeThreshold The credit threshold that triggers auto-recharge
 */
export async function autoCharge(
  chunk: AuthCreditUsageChunk,
  autoRechargeThreshold: number
): Promise<{ success: boolean; message: string; remainingCredits: number; chunk: AuthCreditUsageChunk }> {
  const resource = `auto-recharge:${chunk.team_id}`;
  const cooldownKey = `auto-recharge-cooldown:${chunk.team_id}`;

  try {
    // Check if the team is in the cooldown period
    // Another check to prevent race conditions, double charging - cool down of 10 minutes
    const cooldownValue = await getValue(cooldownKey);
    if (cooldownValue) {
      Logger.info(`Auto-recharge for team ${chunk.team_id} is in cooldown period`);
      return {
        success: false,
        message: "Auto-recharge is in cooldown period",
        remainingCredits: chunk.remaining_credits,
        chunk,
      };
    }

    // Use a distributed lock to prevent concurrent auto-charge attempts
    return await redlock.using([resource], 5000, async (signal) : Promise<{ success: boolean; message: string; remainingCredits: number; chunk: AuthCreditUsageChunk }> => {
      // Recheck the condition inside the lock to prevent race conditions
      const updatedChunk = await getACUC(chunk.api_key, false, false);
      if (
        updatedChunk &&
        updatedChunk.remaining_credits < autoRechargeThreshold
      ) {
        if (chunk.sub_user_id) {
          // Fetch the customer's Stripe information
          const { data: customer, error: customersError } =
            await supabase_service
              .from("customers")
              .select("id, stripe_customer_id")
              .eq("id", chunk.sub_user_id)
              .single();
          
          if (customersError) {
            Logger.error(`Error fetching customer data: ${customersError}`);
            return {
              success: false,
              message: "Error fetching customer data",
              remainingCredits: chunk.remaining_credits,
              chunk,
            };
          }

          if (customer && customer.stripe_customer_id) {
            let issueCreditsSuccess = false;
            // Attempt to create a payment intent
            const paymentStatus = await createPaymentIntent(
              chunk.team_id,
              customer.stripe_customer_id
            );
            
            // If payment is successful or requires further action, issue credits
            if (
              paymentStatus.return_status === "succeeded" ||
              paymentStatus.return_status === "requires_action"
            ) {
              issueCreditsSuccess = await issueCredits(
                chunk.team_id,
                AUTO_RECHARGE_CREDITS
              );
            }

            // Record the auto-recharge transaction
            await supabase_service.from("auto_recharge_transactions").insert({
              team_id: chunk.team_id,
              initial_payment_status: paymentStatus.return_status,
              credits_issued: issueCreditsSuccess ? AUTO_RECHARGE_CREDITS : 0,
              stripe_charge_id: paymentStatus.charge_id,
            });

            // Send a notification if credits were successfully issued
            if (issueCreditsSuccess) {
              await sendNotification(
                chunk.team_id, 
                NotificationType.AUTO_RECHARGE_SUCCESS, 
                chunk.sub_current_period_start, 
                chunk.sub_current_period_end, 
                chunk, 
                true
              );

              // Set cooldown period
              await setValue(cooldownKey, 'true', AUTO_RECHARGE_COOLDOWN);
            } 

            // Reset ACUC cache to reflect the new credit balance
            const cacheKeyACUC = `acuc_${chunk.api_key}`;
            await deleteKey(cacheKeyACUC);
            if (process.env.SLACK_ADMIN_WEBHOOK_URL ) {
                sendSlackWebhook(
                  `Auto-recharge successful: Team ${chunk.team_id}. ${AUTO_RECHARGE_CREDITS} credits added. Payment status: ${paymentStatus.return_status}. User was notified via email.`,
                  false,
                  process.env.SLACK_ADMIN_WEBHOOK_URL
                ).catch((error) => {
                  Logger.debug(`Error sending slack notification: ${error}`);
                });
              }
            return {
              success: true,
              message: "Auto-recharge successful",
              remainingCredits: chunk.remaining_credits + AUTO_RECHARGE_CREDITS,
              chunk: {...chunk, remaining_credits: chunk.remaining_credits + AUTO_RECHARGE_CREDITS},
            };
          } else {
            Logger.error("No Stripe customer ID found for user");
            return {
              success: false,
              message: "No Stripe customer ID found for user",
              remainingCredits: chunk.remaining_credits,
              chunk,
            };
          }
        } else {
            Logger.error("No sub_user_id found in chunk");
          return {
            success: false,
            message: "No sub_user_id found in chunk",
            remainingCredits: chunk.remaining_credits,
            chunk,
          };
        }
      }
      return {
        success: false,
        message: "No need to auto-recharge",
        remainingCredits: chunk.remaining_credits,
        chunk,
      };

    });
  } catch (error) {
    Logger.error(`Failed to acquire lock for auto-recharge: ${error}`);
    return {
      success: false,
      message: "Failed to acquire lock for auto-recharge",
      remainingCredits: chunk.remaining_credits,
      chunk,
    };
  }
}