From 3022d2eba16fa0001046adc804ca191d16e20b5c Mon Sep 17 00:00:00 2001
From: Darcy Ye <darcyye@silverhand.io>
Date: Mon, 21 Oct 2024 18:28:05 +0800
Subject: [PATCH] fix(console): manually update tenant data when data get
 updated (#6697)

* fix(console): manually update tenant data when data get updated

* chore: update code
---
 .../use-new-subscription-data.ts                   | 14 +++++++++++---
 .../src/pages/CheckoutSuccessCallback/index.tsx    |  7 ++++++-
 .../pages/TenantSettings/BillingHistory/index.tsx  | 12 ++++++++++--
 .../pages/TenantSettings/Subscription/index.tsx    | 13 +++++++++++--
 4 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/packages/console/src/contexts/SubscriptionDataProvider/use-new-subscription-data.ts b/packages/console/src/contexts/SubscriptionDataProvider/use-new-subscription-data.ts
index 74da4339d..c355a1434 100644
--- a/packages/console/src/contexts/SubscriptionDataProvider/use-new-subscription-data.ts
+++ b/packages/console/src/contexts/SubscriptionDataProvider/use-new-subscription-data.ts
@@ -1,5 +1,5 @@
-import { cond, condString } from '@silverhand/essentials';
-import { useContext, useMemo } from 'react';
+import { cond, condString, pick } from '@silverhand/essentials';
+import { useContext, useEffect, useMemo } from 'react';
 import useSWR from 'swr';
 
 import { useCloudApi } from '@/cloud/hooks/use-cloud-api';
@@ -21,7 +21,7 @@ import { type NewSubscriptionContext } from './types';
 const useNewSubscriptionData: () => NewSubscriptionContext & { isLoading: boolean } = () => {
   const cloudApi = useCloudApi();
 
-  const { currentTenant } = useContext(TenantsContext);
+  const { currentTenant, updateTenant } = useContext(TenantsContext);
   const { isLoading: isLogtoSkusLoading, data: fetchedLogtoSkus } = useLogtoSkus();
   const tenantId = condString(currentTenant?.id);
 
@@ -50,6 +50,14 @@ const useNewSubscriptionData: () => NewSubscriptionContext & { isLoading: boolea
     [currentTenant?.planId, logtoSkus]
   );
 
+  useEffect(() => {
+    if (subscriptionUsageData?.quota) {
+      updateTenant(tenantId, {
+        quota: pick(subscriptionUsageData.quota, 'mauLimit', 'tokenLimit'),
+      });
+    }
+  }, [tenantId, subscriptionUsageData?.quota, updateTenant]);
+
   return useMemo(
     () => ({
       isLoading: isSubscriptionLoading || isLogtoSkusLoading || isSubscriptionUsageDataLoading,
diff --git a/packages/console/src/pages/CheckoutSuccessCallback/index.tsx b/packages/console/src/pages/CheckoutSuccessCallback/index.tsx
index f3f3c2c55..0c8a30bd4 100644
--- a/packages/console/src/pages/CheckoutSuccessCallback/index.tsx
+++ b/packages/console/src/pages/CheckoutSuccessCallback/index.tsx
@@ -26,7 +26,7 @@ function CheckoutSuccessCallback() {
   const { t } = useTranslation(undefined, { keyPrefix: 'admin_console.subscription' });
   const { navigate } = useTenantPathname();
   const cloudApi = useCloudApi({ hideErrorToast: true });
-  const { currentTenantId, navigateTenant } = useContext(TenantsContext);
+  const { currentTenantId, navigateTenant, updateTenant } = useContext(TenantsContext);
   const { onCurrentSubscriptionUpdated } = useContext(SubscriptionDataContext);
   const { search } = useLocation();
   const checkoutState = new URLSearchParams(search).get(checkoutStateQueryKey);
@@ -101,6 +101,10 @@ function CheckoutSuccessCallback() {
       }
 
       onCurrentSubscriptionUpdated(tenantSubscription);
+      updateTenant(checkoutTenantId, {
+        subscription: tenantSubscription,
+        ...conditional(tenantSubscription?.planId && { planId: tenantSubscription.planId }),
+      });
 
       // No need to check `isDowngrade` here, since a downgrade must occur in a tenant with a Pro
       // plan, and the purchase conversion has already been reported using the same tenant ID. We
@@ -131,6 +135,7 @@ function CheckoutSuccessCallback() {
     onCurrentSubscriptionUpdated,
     t,
     tenantSubscription,
+    updateTenant,
   ]);
 
   if (!isValidSession) {
diff --git a/packages/console/src/pages/TenantSettings/BillingHistory/index.tsx b/packages/console/src/pages/TenantSettings/BillingHistory/index.tsx
index 4bde6724d..dcaf56d78 100644
--- a/packages/console/src/pages/TenantSettings/BillingHistory/index.tsx
+++ b/packages/console/src/pages/TenantSettings/BillingHistory/index.tsx
@@ -1,7 +1,7 @@
 import { ReservedPlanId } from '@logto/schemas';
 import { conditional } from '@silverhand/essentials';
 import dayjs from 'dayjs';
-import { useCallback, useContext, useMemo } from 'react';
+import { useCallback, useContext, useEffect, useMemo } from 'react';
 
 import { useCloudApi } from '@/cloud/hooks/use-cloud-api';
 import EmptyDataPlaceholder from '@/components/EmptyDataPlaceholder';
@@ -18,7 +18,7 @@ import InvoiceStatusTag from './InvoiceStatusTag';
 
 function BillingHistory() {
   const cloudApi = useCloudApi();
-  const { currentTenantId } = useContext(TenantsContext);
+  const { currentTenantId, updateTenant } = useContext(TenantsContext);
   const { data: invoices, error } = useInvoices(currentTenantId);
   const isLoadingInvoices = !invoices && !error;
   const displayInvoices = useMemo(
@@ -41,6 +41,14 @@ function BillingHistory() {
     [cloudApi, currentTenantId]
   );
 
+  useEffect(() => {
+    if (invoices) {
+      updateTenant(currentTenantId, {
+        openInvoices: invoices,
+      });
+    }
+  }, [currentTenantId, invoices, updateTenant]);
+
   return (
     <div>
       <PageMeta titleKey={['tenants.tabs.billing_history', 'tenants.title']} />
diff --git a/packages/console/src/pages/TenantSettings/Subscription/index.tsx b/packages/console/src/pages/TenantSettings/Subscription/index.tsx
index 98708d19f..4c37bc2d2 100644
--- a/packages/console/src/pages/TenantSettings/Subscription/index.tsx
+++ b/packages/console/src/pages/TenantSettings/Subscription/index.tsx
@@ -19,7 +19,7 @@ function Subscription() {
   const cloudApi = useCloudApi();
   const { logtoSkus, currentSku, onCurrentSubscriptionUpdated } =
     useContext(SubscriptionDataContext);
-  const { currentTenantId } = useContext(TenantsContext);
+  const { currentTenantId, updateTenant } = useContext(TenantsContext);
 
   const reservedSkus = pickupFeaturedLogtoSkus(logtoSkus);
 
@@ -34,8 +34,17 @@ function Subscription() {
   useEffect(() => {
     if (isCloud) {
       onCurrentSubscriptionUpdated();
+
+      if (periodicUsage) {
+        updateTenant(currentTenantId, {
+          usage: {
+            activeUsers: periodicUsage.mauLimit,
+            tokenUsage: periodicUsage.tokenLimit,
+          },
+        });
+      }
     }
-  }, [onCurrentSubscriptionUpdated]);
+  }, [currentTenantId, onCurrentSubscriptionUpdated, periodicUsage, updateTenant]);
 
   if (isLoading) {
     return <Skeleton />;