From c3a9630e330e4f09d77cb2b095b36cdb62043d42 Mon Sep 17 00:00:00 2001
From: Nicolas <nicolascamara29@gmail.com>
Date: Mon, 14 Oct 2024 12:24:34 -0300
Subject: [PATCH] Reapply "Merge pull request #773 from
 mendableai/nsc/retries-acuc-price-credits-fallback"

This reverts commit a6888ce17b98e4a47fbfc9796911446a6513959c.
---
 apps/api/src/controllers/auth.ts              | 62 ++++++++++++++-----
 .../src/services/billing/credit_billing.ts    |  6 +-
 2 files changed, 51 insertions(+), 17 deletions(-)

diff --git a/apps/api/src/controllers/auth.ts b/apps/api/src/controllers/auth.ts
index 0367358f..5546bc17 100644
--- a/apps/api/src/controllers/auth.ts
+++ b/apps/api/src/controllers/auth.ts
@@ -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);
diff --git a/apps/api/src/services/billing/credit_billing.ts b/apps/api/src/services/billing/credit_billing.ts
index 3346e291..694d0e5c 100644
--- a/apps/api/src/services/billing/credit_billing.ts
+++ b/apps/api/src/services/billing/credit_billing.ts
@@ -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 mistakes
+  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,