From 84bff8add8ab93f28b4ef8f05b058448a929cd91 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gerg=C5=91=20M=C3=B3ricz?= <mo.geryy@gmail.com>
Date: Wed, 25 Sep 2024 21:37:01 +0200
Subject: [PATCH] fix(billTeam): update cached ACUC after billing

---
 apps/api/src/controllers/auth.ts              |  8 +++++---
 apps/api/src/controllers/v1/types.ts          |  2 +-
 .../src/services/billing/credit_billing.ts    | 19 ++++++++++++++++++-
 3 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/apps/api/src/controllers/auth.ts b/apps/api/src/controllers/auth.ts
index fe7baa6..96f94a9 100644
--- a/apps/api/src/controllers/auth.ts
+++ b/apps/api/src/controllers/auth.ts
@@ -37,7 +37,7 @@ function normalizedApiIsUuid(potentialUuid: string): boolean {
   return validate(potentialUuid);
 }
 
-async function setCachedACUC(api_key: string, acuc: AuthCreditUsageChunk) {
+export async function setCachedACUC(api_key: string, acuc: AuthCreditUsageChunk) {
   const cacheKeyACUC = `acuc_${api_key}`;
   const redLockKey = `lock_${cacheKeyACUC}`;
   const lockTTL = 10000; // 10 seconds
@@ -58,14 +58,14 @@ async function setCachedACUC(api_key: string, acuc: AuthCreditUsageChunk) {
   }
 }
 
-async function getACUC(api_key: string): 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);
 
   if (cachedACUC !== null) {
     return JSON.parse(cachedACUC);
-  } else {
+  } else if (!cacheOnly) {
     const { data, error } =
       await supabase_service.rpc("auth_credit_usage_chunk", { input_key: api_key });
     
@@ -85,6 +85,8 @@ async function getACUC(api_key: string): Promise<AuthCreditUsageChunk | null> {
     }
 
     return chunk;
+  } else {
+    return null;
   }
 }
 
diff --git a/apps/api/src/controllers/v1/types.ts b/apps/api/src/controllers/v1/types.ts
index 5e99fda..441980b 100644
--- a/apps/api/src/controllers/v1/types.ts
+++ b/apps/api/src/controllers/v1/types.ts
@@ -324,7 +324,7 @@ export type AuthCreditUsageChunk = {
   price_id: string | null;
   price_credits: number; // credit limit with assoicated price, or free_credits (500) if free plan
   credits_used: number;
-  coupon_credits: number;
+  coupon_credits: number; // do not rely on this number to be up to date after calling a billTeam
   coupons: any[];
   adjusted_credits_used: number; // credits this period minus coupons used
   remaining_credits: number;
diff --git a/apps/api/src/services/billing/credit_billing.ts b/apps/api/src/services/billing/credit_billing.ts
index 9d0d1f0..6fa3490 100644
--- a/apps/api/src/services/billing/credit_billing.ts
+++ b/apps/api/src/services/billing/credit_billing.ts
@@ -5,6 +5,7 @@ import { supabase_service } from "../supabase";
 import { Logger } from "../../lib/logger";
 import * as Sentry from "@sentry/node";
 import { AuthCreditUsageChunk } from "../../controllers/v1/types";
+import { getACUC, setCachedACUC } from "../../controllers/auth";
 
 const FREE_CREDITS = 500;
 
@@ -20,13 +21,29 @@ export async function supaBillTeam(team_id: string, subscription_id: string, cre
   }
   Logger.info(`Billing team ${team_id} for ${credits} credits`);
 
-  const { error } =
+  const { data, error } =
     await supabase_service.rpc("bill_team", { _team_id: team_id, sub_id: subscription_id ?? null, fetch_subscription: subscription_id === undefined, credits });
   
   if (error) {
     Sentry.captureException(error);
     Logger.error("Failed to bill team: " + JSON.stringify(error));
+    return;
   }
+
+  (async () => {
+    for (const apiKey of (data ?? []).map(x => x.api_key)) {
+      const acuc = await getACUC(apiKey, true);
+  
+      if (acuc !== null) {
+        await setCachedACUC(apiKey, {
+          ...acuc,
+          credits_used: acuc.credits_used + credits,
+          adjusted_credits_used: acuc.adjusted_credits_used + credits,
+          remaining_credits: acuc.remaining_credits - credits,
+        });
+      }
+    }
+  })();
 }
 
 export async function checkTeamCredits(chunk: AuthCreditUsageChunk, team_id: string, credits: number) {