Reapply "Merge pull request #773 from mendableai/nsc/retries-acuc-price-credits-fallback"

This reverts commit a6888ce17b.

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 } =
+    let data;
-      await supabase_service.rpc("auth_credit_usage_chunk_test_3", { input_key: api_key });
+    let error;
+    let retries = 0;
+    const maxRetries = 5;
 
-    if (error) {
+    while (retries < maxRetries) {
-      throw new Error("Failed to retrieve authentication and credit usage data: " + JSON.stringify(error));
+      ({ data, error } = await supabase_service.rpc(
+        "auth_credit_usage_chunk_test_3",
+        { input_key: api_key }
+      ));
+
+      if (!error) {
+        break;
       }
 
-    const chunk: AuthCreditUsageChunk | null = data.length === 0
+      Logger.warn(
-      ? null
+        `Failed to retrieve authentication and credit usage data after ${retries}, trying again...`
-      : data[0].team_id === null
+      );
-      ? null
+      retries++;
-      : data[0];
+      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];
 
     // 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 };
   }

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,