From 1beb434a729949ce291673b0294954f28682810e Mon Sep 17 00:00:00 2001
From: emmpaz <59887297+emmpaz@users.noreply.github.com>
Date: Wed, 31 Jul 2024 07:57:27 -0500
Subject: [PATCH] fix: limits syncing issue (#1195)

Exposes `refreshLimits()` to be able to keep the limit in sync when
deleting/creating a document.
---
 .../documents/delete-document-dialog.tsx      |  3 +++
 .../(dashboard)/documents/upload-document.tsx |  4 ++-
 .../ee/server-only/limits/provider/client.tsx | 25 +++++++++++++------
 .../ee/server-only/limits/provider/server.tsx |  3 +--
 4 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/apps/web/src/app/(dashboard)/documents/delete-document-dialog.tsx b/apps/web/src/app/(dashboard)/documents/delete-document-dialog.tsx
index 558d39558..042dc5e93 100644
--- a/apps/web/src/app/(dashboard)/documents/delete-document-dialog.tsx
+++ b/apps/web/src/app/(dashboard)/documents/delete-document-dialog.tsx
@@ -4,6 +4,7 @@ import { useRouter } from 'next/navigation';
 
 import { match } from 'ts-pattern';
 
+import { useLimits } from '@documenso/ee/server-only/limits/provider/client';
 import { DocumentStatus } from '@documenso/prisma/client';
 import { trpc as trpcReact } from '@documenso/trpc/react';
 import { Alert, AlertDescription } from '@documenso/ui/primitives/alert';
@@ -41,6 +42,7 @@ export const DeleteDocumentDialog = ({
   const router = useRouter();
 
   const { toast } = useToast();
+  const { refreshLimits } = useLimits();
 
   const [inputValue, setInputValue] = useState('');
   const [isDeleteEnabled, setIsDeleteEnabled] = useState(status === DocumentStatus.DRAFT);
@@ -48,6 +50,7 @@ export const DeleteDocumentDialog = ({
   const { mutateAsync: deleteDocument, isLoading } = trpcReact.document.deleteDocument.useMutation({
     onSuccess: () => {
       router.refresh();
+      void refreshLimits();
 
       toast({
         title: 'Document deleted',
diff --git a/apps/web/src/app/(dashboard)/documents/upload-document.tsx b/apps/web/src/app/(dashboard)/documents/upload-document.tsx
index 0a0c029ab..f52aecb94 100644
--- a/apps/web/src/app/(dashboard)/documents/upload-document.tsx
+++ b/apps/web/src/app/(dashboard)/documents/upload-document.tsx
@@ -36,7 +36,7 @@ export const UploadDocument = ({ className, team }: UploadDocumentProps) => {
 
   const { toast } = useToast();
 
-  const { quota, remaining } = useLimits();
+  const { quota, remaining, refreshLimits } = useLimits();
 
   const [isLoading, setIsLoading] = useState(false);
 
@@ -71,6 +71,8 @@ export const UploadDocument = ({ className, team }: UploadDocumentProps) => {
         teamId: team?.id,
       });
 
+      void refreshLimits();
+
       toast({
         title: 'Document uploaded',
         description: 'Your document has been uploaded successfully.',
diff --git a/packages/ee/server-only/limits/provider/client.tsx b/packages/ee/server-only/limits/provider/client.tsx
index fdc00b439..624b61d98 100644
--- a/packages/ee/server-only/limits/provider/client.tsx
+++ b/packages/ee/server-only/limits/provider/client.tsx
@@ -1,6 +1,6 @@
 'use client';
 
-import { createContext, useContext, useEffect, useState } from 'react';
+import { createContext, useCallback, useContext, useEffect, useState } from 'react';
 
 import { equals } from 'remeda';
 
@@ -8,7 +8,7 @@ import { getLimits } from '../client';
 import { FREE_PLAN_LIMITS } from '../constants';
 import type { TLimitsResponseSchema } from '../schema';
 
-export type LimitsContextValue = TLimitsResponseSchema;
+export type LimitsContextValue = TLimitsResponseSchema & { refreshLimits: () => Promise<void> };
 
 const LimitsContext = createContext<LimitsContextValue | null>(null);
 
@@ -23,7 +23,7 @@ export const useLimits = () => {
 };
 
 export type LimitsProviderProps = {
-  initialValue?: LimitsContextValue;
+  initialValue?: TLimitsResponseSchema;
   teamId?: number;
   children?: React.ReactNode;
 };
@@ -38,7 +38,7 @@ export const LimitsProvider = ({
 }: LimitsProviderProps) => {
   const [limits, setLimits] = useState(() => initialValue);
 
-  const refreshLimits = async () => {
+  const refreshLimits = useCallback(async () => {
     const newLimits = await getLimits({ teamId });
 
     setLimits((oldLimits) => {
@@ -48,11 +48,11 @@ export const LimitsProvider = ({
 
       return newLimits;
     });
-  };
+  }, [teamId]);
 
   useEffect(() => {
     void refreshLimits();
-  }, []);
+  }, [refreshLimits]);
 
   useEffect(() => {
     const onFocus = () => {
@@ -64,7 +64,16 @@ export const LimitsProvider = ({
     return () => {
       window.removeEventListener('focus', onFocus);
     };
-  }, []);
+  }, [refreshLimits]);
 
-  return <LimitsContext.Provider value={limits}>{children}</LimitsContext.Provider>;
+  return (
+    <LimitsContext.Provider
+      value={{
+        ...limits,
+        refreshLimits,
+      }}
+    >
+      {children}
+    </LimitsContext.Provider>
+  );
 };
diff --git a/packages/ee/server-only/limits/provider/server.tsx b/packages/ee/server-only/limits/provider/server.tsx
index b7cde3573..969361060 100644
--- a/packages/ee/server-only/limits/provider/server.tsx
+++ b/packages/ee/server-only/limits/provider/server.tsx
@@ -3,7 +3,6 @@
 import { headers } from 'next/headers';
 
 import { getLimits } from '../client';
-import type { LimitsContextValue } from './client';
 import { LimitsProvider as ClientLimitsProvider } from './client';
 
 export type LimitsProviderProps = {
@@ -14,7 +13,7 @@ export type LimitsProviderProps = {
 export const LimitsProvider = async ({ children, teamId }: LimitsProviderProps) => {
   const requestHeaders = Object.fromEntries(headers().entries());
 
-  const limits: LimitsContextValue = await getLimits({ headers: requestHeaders, teamId });
+  const limits = await getLimits({ headers: requestHeaders, teamId });
 
   return (
     <ClientLimitsProvider initialValue={limits} teamId={teamId}>