AI/firecrawl
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/mendableai/firecrawl synced 2025-08-14 04:05:58 +08:00
Code Issues Packages Projects Releases Wiki Activity

Added more safety guards to auto-rech

Browse Source
This commit is contained in:
Nicolas 2025-01-30 12:24:37 -03:00
parent c88176a596
commit ec2c0f671c
1 changed files with 82 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

63
apps/api/src/services/billing/auto_charge.ts
View File

@ -8,12 +8,15 @@ import { issueCredits } from "./issue_credits";
import { sendNotification } from "../notification/email_notification";
import { NotificationType } from "../../types";
import { deleteKey, getValue, setValue } from "../redis";
import { redisRateLimitClient } from "../rate-limiter";
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 = 300; // 5 minutes in seconds
const AUTO_RECHARGE_COOLDOWN = 600; // 10 minutes in seconds
const MAX_CHARGES_PER_HOUR = 5; // Maximum number of auto-charges per hour
const HOURLY_COUNTER_EXPIRY = 3600; // 1 hour in seconds
/**
* Attempt to automatically charge a user's account when their credit balance falls below a threshold
@ -31,19 +34,38 @@ export async function autoCharge(
}> {
const resource = `auto-recharge:${chunk.team_id}`;
const cooldownKey = `auto-recharge-cooldown:${chunk.team_id}`;
const hourlyCounterKey = `auto-recharge-hourly:${chunk.team_id}`;
if (chunk.team_id === "285bb597-6eaf-4b96-801c-51461fc3c543") {
return {
success: false,
message: "Auto-recharge failed",
message: "Auto-recharge failed: blocked team",
remainingCredits: chunk.remaining_credits,
chunk,
};
}
try {
// Check if the team is in the cooldown period
// Another check to prevent race conditions, double charging - cool down of 5 minutes
// Check hourly rate limit
const hourlyCharges = await redisRateLimitClient.incr(hourlyCounterKey);
if (hourlyCharges === 1) {
// Set expiry for the counter if it's new
await redisRateLimitClient.expire(hourlyCounterKey, HOURLY_COUNTER_EXPIRY);
}
if (hourlyCharges > MAX_CHARGES_PER_HOUR) {
logger.warn(
`Auto-recharge for team ${chunk.team_id} exceeded hourly limit of ${MAX_CHARGES_PER_HOUR}`,
);
return {
success: false,
message: "Auto-recharge hourly limit exceeded",
remainingCredits: chunk.remaining_credits,
chunk,
};
}
// Check cooldown period
const cooldownValue = await getValue(cooldownKey);
if (cooldownValue) {
logger.info(
@ -71,6 +93,19 @@ export async function autoCharge(
}> => {
// Recheck the condition inside the lock to prevent race conditions
const updatedChunk = await getACUC(chunk.api_key, false, false);
// recheck cooldown
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,
};
}
if (
updatedChunk &&
updatedChunk.remaining_credits < autoRechargeThreshold
@ -102,6 +137,9 @@ export async function autoCharge(
customer.stripe_customer_id,
);
// set cooldown
await setValue(cooldownKey, "true", AUTO_RECHARGE_COOLDOWN);
// If payment is successful or requires further action, issue credits
if (
paymentStatus.return_status === "succeeded" ||
@ -123,6 +161,10 @@ export async function autoCharge(
// Send a notification if credits were successfully issued
if (issueCreditsSuccess) {
// Increment hourly counter and set expiry if it doesn't exist
await redisRateLimitClient.incr(hourlyCounterKey);
await redisRateLimitClient.expire(hourlyCounterKey, HOURLY_COUNTER_EXPIRY);
await sendNotification(
chunk.team_id,
NotificationType.AUTO_RECHARGE_SUCCESS,
@ -132,10 +174,6 @@ export async function autoCharge(
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);
@ -177,6 +215,15 @@ export async function autoCharge(
chunk,
};
}
} 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 {