Austin/firecrawl
1
0
Fork 0
mirror of https://github.com/mendableai/firecrawl.git synced 2024-11-16 03:32:22 +08:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

Nick: retries for acuc

Browse Source
This commit is contained in:
Nicolas 2024-10-13 22:27:38 -03:00
parent 666082a7dd
commit 78b6127d88
2 changed files with 51 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

62
apps/api/src/controllers/auth.ts
View File

@ -37,12 +37,17 @@ function normalizedApiIsUuid(potentialUuid: string): boolean {
return validate(potentialUuid);
}
export async function setCachedACUC(api_key: string, acuc: AuthCreditUsageChunk | ((acuc: AuthCreditUsageChunk) => AuthCreditUsageChunk)) {
export async function setCachedACUC(
api_key: string,
acuc:
| AuthCreditUsageChunk
| ((acuc: AuthCreditUsageChunk) => AuthCreditUsageChunk)
) {
const cacheKeyACUC = `acuc_${api_key}`;
const redLockKey = `lock_${cacheKeyACUC}`;
try {
await redlock.using([redLockKey], 10000, {}, async signal => {
await redlock.using([redLockKey], 10000, {}, async (signal) => {
if (typeof acuc === "function") {
acuc = acuc(JSON.parse(await getValue(cacheKeyACUC)));
@ -68,7 +73,10 @@ export async function setCachedACUC(api_key: string, acuc: AuthCreditUsageChunk
}
}
export async function getACUC(api_key: string, cacheOnly = false): Promise<AuthCreditUsageChunk | null> {
export async function getACUC(
api_key: string,
cacheOnly = false
): Promise<AuthCreditUsageChunk | null> {
const cacheKeyACUC = `acuc_${api_key}`;
const cachedACUC = await getValue(cacheKeyACUC);
@ -76,18 +84,38 @@ export async function getACUC(api_key: string, cacheOnly = false): Promise<AuthC
if (cachedACUC !== null) {
return JSON.parse(cachedACUC);
} else if (!cacheOnly) {
const { data, error } =
await supabase_service.rpc("auth_credit_usage_chunk_test_3", { input_key: api_key });
if (error) {
throw new Error("Failed to retrieve authentication and credit usage data: " + JSON.stringify(error));
let data;
let error;
let retries = 0;
const maxRetries = 5;
while (retries < maxRetries) {
({ data, error } = await supabase_service.rpc(
"auth_credit_usage_chunk_test_3",
{ input_key: api_key }
));
if (!error) {
break;
}
Logger.warn(
`Failed to retrieve authentication and credit usage data after ${retries}, trying again...`
);
retries++;
if (retries === maxRetries) {
throw new Error(
"Failed to retrieve authentication and credit usage data after 3 attempts: " +
JSON.stringify(error)
);
}
// Wait for a short time before retrying
await new Promise((resolve) => setTimeout(resolve, 200));
}
const chunk: AuthCreditUsageChunk | null = data.length === 0
? null
: data[0].team_id === null
? null
: data[0];
const chunk: AuthCreditUsageChunk | null =
data.length === 0 ? null : data[0].team_id === null ? null : data[0];
// NOTE: Should we cache null chunks? - mogery
if (chunk !== null) {
@ -132,7 +160,11 @@ export async function supaAuthenticateUser(
plan?: PlanType;
chunk?: AuthCreditUsageChunk;
}> {
const authHeader = req.headers.authorization ?? (req.headers["sec-websocket-protocol"] ? `Bearer ${req.headers["sec-websocket-protocol"]}` : null);
const authHeader =
req.headers.authorization ??
(req.headers["sec-websocket-protocol"]
? `Bearer ${req.headers["sec-websocket-protocol"]}`
: null);
if (!authHeader) {
return { success: false, error: "Unauthorized", status: 401 };
}
@ -162,7 +194,7 @@ export async function supaAuthenticateUser(
rateLimiter = getRateLimiter(RateLimiterMode.CrawlStatus, token);
} else {
rateLimiter = getRateLimiter(RateLimiterMode.Preview, token);
}
}
teamId = "preview";
} else {
normalizedApi = parseApi(token);

6
apps/api/src/services/billing/credit_billing.ts
View File

@ -55,11 +55,13 @@ export async function supaCheckTeamCredits(chunk: AuthCreditUsageChunk, team_id:
const creditsWillBeUsed = chunk.adjusted_credits_used + credits;
// In case chunk.price_credits is undefined, set it to a large number to avoid division by zero
const totalPriceCredits = chunk.price_credits ?? 100000000;
// Removal of + credits
const creditUsagePercentage = creditsWillBeUsed / chunk.price_credits;
const creditUsagePercentage = creditsWillBeUsed / totalPriceCredits;
// Compare the adjusted total credits used with the credits allowed by the plan
if (creditsWillBeUsed > chunk.price_credits) {
if (creditsWillBeUsed > totalPriceCredits) {
sendNotification(
team_id,
NotificationType.LIMIT_REACHED,