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

Update auto_charge.ts

Browse Source
This commit is contained in:
Nicolas 2025-01-30 12:47:08 -03:00
parent a585340740
commit e2c3932dd2
1 changed files with 33 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -36,7 +36,7 @@ export async function autoCharge(
const cooldownKey = `auto-recharge-cooldown:${chunk.team_id}`; const cooldownKey = `auto-recharge-cooldown:${chunk.team_id}`;
const hourlyCounterKey = `auto-recharge-hourly:${chunk.team_id}`; const hourlyCounterKey = `auto-recharge-hourly:${chunk.team_id}`;
if (chunk.team_id === "285bb597-6eaf-4b96-801c-51461fc3c543") { if (chunk.team_id === "285bb597-6eaf-4b96-801c-51461fc3c543" || chunk.team_id === "dec639a0-98ca-4995-95b5-48ac1ffab5b7") {
return { return {
success: false, success: false,
message: "Auto-recharge failed: blocked team", message: "Auto-recharge failed: blocked team",
@ -46,14 +46,11 @@ export async function autoCharge(
} }
try { try {
// Check hourly rate limit // Check hourly rate limit first without incrementing
const hourlyCharges = await redisRateLimitClient.incr(hourlyCounterKey); const currentCharges = await redisRateLimitClient.get(hourlyCounterKey);
if (hourlyCharges === 1) { const hourlyCharges = currentCharges ? parseInt(currentCharges) : 0;
// Set expiry for the counter if it's new
await redisRateLimitClient.expire(hourlyCounterKey, HOURLY_COUNTER_EXPIRY);
}
if (hourlyCharges > MAX_CHARGES_PER_HOUR) { if (hourlyCharges >= MAX_CHARGES_PER_HOUR) {
logger.warn( logger.warn(
`Auto-recharge for team ${chunk.team_id} exceeded hourly limit of ${MAX_CHARGES_PER_HOUR}`, `Auto-recharge for team ${chunk.team_id} exceeded hourly limit of ${MAX_CHARGES_PER_HOUR}`,
); );
@ -91,9 +88,10 @@ export async function autoCharge(
remainingCredits: number; remainingCredits: number;
chunk: AuthCreditUsageChunk; chunk: AuthCreditUsageChunk;
}> => { }> => {
// Recheck the condition inside the lock to prevent race conditions // Recheck all conditions inside the lock to prevent race conditions
const updatedChunk = await getACUC(chunk.api_key, false, false); const updatedChunk = await getACUC(chunk.api_key, false, false);
// recheck cooldown
// Recheck cooldown
const cooldownValue = await getValue(cooldownKey); const cooldownValue = await getValue(cooldownKey);
if (cooldownValue) { if (cooldownValue) {
logger.info( logger.info(
@ -106,6 +104,19 @@ export async function autoCharge(
chunk, chunk,
}; };
} }
// Recheck hourly limit inside lock
const currentCharges = await redisRateLimitClient.get(hourlyCounterKey);
const hourlyCharges = currentCharges ? parseInt(currentCharges) : 0;
if (hourlyCharges >= MAX_CHARGES_PER_HOUR) {
return {
success: false,
message: "Auto-recharge hourly limit exceeded",
remainingCredits: chunk.remaining_credits,
chunk,
};
}
if ( if (
updatedChunk && updatedChunk &&
updatedChunk.remaining_credits < autoRechargeThreshold updatedChunk.remaining_credits < autoRechargeThreshold
@ -131,15 +142,16 @@ export async function autoCharge(
if (customer && customer.stripe_customer_id) { if (customer && customer.stripe_customer_id) {
let issueCreditsSuccess = false; let issueCreditsSuccess = false;
// Set cooldown BEFORE attempting payment
await setValue(cooldownKey, "true", AUTO_RECHARGE_COOLDOWN);
// Attempt to create a payment intent // Attempt to create a payment intent
const paymentStatus = await createPaymentIntent( const paymentStatus = await createPaymentIntent(
chunk.team_id, chunk.team_id,
customer.stripe_customer_id, customer.stripe_customer_id,
); );
// set cooldown
await setValue(cooldownKey, "true", AUTO_RECHARGE_COOLDOWN);
// If payment is successful or requires further action, issue credits // If payment is successful or requires further action, issue credits
if ( if (
paymentStatus.return_status === "succeeded" || paymentStatus.return_status === "succeeded" ||
@ -163,7 +175,10 @@ export async function autoCharge(
if (issueCreditsSuccess) { if (issueCreditsSuccess) {
// Increment hourly counter and set expiry if it doesn't exist // Increment hourly counter and set expiry if it doesn't exist
await redisRateLimitClient.incr(hourlyCounterKey); await redisRateLimitClient.incr(hourlyCounterKey);
await redisRateLimitClient.expire(hourlyCounterKey, HOURLY_COUNTER_EXPIRY); await redisRateLimitClient.expire(
hourlyCounterKey,
HOURLY_COUNTER_EXPIRY,
);
await sendNotification( await sendNotification(
chunk.team_id, chunk.team_id,
@ -188,7 +203,9 @@ export async function autoCharge(
false, false,
process.env.SLACK_ADMIN_WEBHOOK_URL, process.env.SLACK_ADMIN_WEBHOOK_URL,
).catch((error) => { ).catch((error) => {
logger.debug(`Error sending slack notification: ${error}`); logger.debug(
`Error sending slack notification: ${error}`,
);
}); });
// Set cooldown for 1 hour // Set cooldown for 1 hour