Merge pull request #622 from mfts/fix/limit-links

Author: mfts

Diff for: components/Sidebar.tsx

@@ -17,7 +17,8 @@ import MenuIcon from "@/components/shared/icons/menu";
 import { Sheet, SheetContent, SheetTrigger } from "@/components/ui/sheet";
 
 import { usePlan } from "@/lib/swr/use-billing";
-import { cn } from "@/lib/utils";
+import useLimits from "@/lib/swr/use-limits";
+import { cn, nFormatter } from "@/lib/utils";
 
 import Banner from "./banner";
 import ProBanner from "./billing/pro-banner";
@@ -26,6 +27,7 @@ import ProfileMenu from "./profile-menu";
 import SiderbarFolders from "./sidebar-folders";
 import { AddTeamModal } from "./teams/add-team-modal";
 import SelectTeam from "./teams/select-team";
+import { Progress } from "./ui/progress";
 import { ScrollArea } from "./ui/scroll-area";
 
 export default function Sidebar() {
@@ -66,6 +68,9 @@ export const SidebarComponent = ({ className }: { className?: string }) => {
   const { data: session, status } = useSession();
   const { plan: userPlan, trial: userTrial, loading } = usePlan();
   const isTrial = !!userTrial;
+  const { limits } = useLimits();
+  const linksLimit = limits?.links;
+  const documentsLimit = limits?.documents;
 
   const router = useRouter();
   const { currentTeam, teams, isLoading }: TeamContextType =
@@ -281,6 +286,25 @@ export const SidebarComponent = ({ className }: { className?: string }) => {
             <ProBanner setShowProBanner={setShowProBanner} />
           ) : null}
 
+          <div className="mb-2">
+            {linksLimit ? (
+              <UsageProgress
+                title="Links"
+                unit="links"
+                usage={limits?.usage?.links}
+                usageLimit={linksLimit}
+              />
+            ) : null}
+            {documentsLimit ? (
+              <UsageProgress
+                title="Documents"
+                unit="documents"
+                usage={limits?.usage?.documents}
+                usageLimit={documentsLimit}
+              />
+            ) : null}
+          </div>
+
           <div className="hidden w-full lg:block">
             <ProfileMenu size="large" />
           </div>
@@ -289,3 +313,50 @@ export const SidebarComponent = ({ className }: { className?: string }) => {
     </div>
   );
 };
+
+function UsageProgress(data: {
+  title: string;
+  unit: string;
+  usage?: number;
+  usageLimit?: number;
+}) {
+  let { title, unit, usage, usageLimit } = data;
+  let usagePercentage = 0;
+  if (usage !== undefined && usageLimit !== undefined) {
+    usagePercentage = (usage / usageLimit) * 100;
+  }
+
+  return (
+    <div className="p-2">
+      {/* <div className="flex items-center space-x-2">
+        <h3 className="font-medium">{title}</h3>
+      </div> */}
+
+      <div className="mt-2 flex flex-col space-y-2">
+        {usage !== undefined && usageLimit !== undefined ? (
+          <p className="text-sm text-foreground">
+            <span
+              className={cn(
+                usagePercentage > 100 && "text-destructive dark:text-red-600",
+              )}
+            >
+              {nFormatter(usage)}
+            </span>{" "}
+            / {nFormatter(usageLimit)} {unit}{" "}
+            <span
+              className={cn(
+                usagePercentage > 100 && "text-destructive dark:text-red-600",
+              )}
+            >
+              ({usagePercentage.toFixed(0)}
+              %)
+            </span>
+          </p>
+        ) : (
+          <div className="h-5 w-32 animate-pulse rounded-md bg-muted" />
+        )}
+        <Progress value={usagePercentage} className="h-2 bg-muted" max={100} />
+      </div>
+    </div>
+  );
+}

Diff for: components/documents/add-document-modal.tsx

@@ -28,9 +28,12 @@ import {
   createNewDocumentVersion,
 } from "@/lib/documents/create-document";
 import { putFile } from "@/lib/files/put-file";
+import useLimits from "@/lib/swr/use-limits";
 import { copyToClipboard } from "@/lib/utils";
 import { getSupportedContentType } from "@/lib/utils/get-content-type";
 
+import { UpgradePlanModal } from "../billing/upgrade-plan-modal";
+
 export function AddDocumentModal({
   newVersion,
   children,
@@ -52,6 +55,7 @@ export function AddDocumentModal({
   const [currentFile, setCurrentFile] = useState<File | null>(null);
   const [notionLink, setNotionLink] = useState<string | null>(null);
   const teamInfo = useTeam();
+  const { canAddDocuments } = useLimits();
 
   const teamId = teamInfo?.currentTeam?.id as string;
 
@@ -69,6 +73,11 @@ export function AddDocumentModal({
       return; // prevent form from submitting
     }
 
+    if (!canAddDocuments) {
+      toast.error("You have reached the maximum number of documents.");
+      return;
+    }
+
     try {
       setUploading(true);
 
@@ -251,6 +260,12 @@ export function AddDocumentModal({
     event: FormEvent<HTMLFormElement>,
   ): Promise<void> => {
     event.preventDefault();
+
+    if (!canAddDocuments) {
+      toast.error("You have reached the maximum number of documents.");
+      return;
+    }
+
     const validateNotionPageURL = parsePageId(notionLink);
     // Check if it's a valid URL or not by Regx
     const isValidURL =
@@ -353,6 +368,24 @@ export function AddDocumentModal({
     setAddDocumentModalOpen && setAddDocumentModalOpen(!isOpen);
   };
 
+  if (!canAddDocuments && children) {
+    if (newVersion) {
+      return (
+        <UpgradePlanModal
+          clickedPlan="Pro"
+          trigger={"limit_upload_document_version"}
+        >
+          {children}
+        </UpgradePlanModal>
+      );
+    }
+    return (
+      <UpgradePlanModal clickedPlan="Pro" trigger={"limit_upload_documents"}>
+        <Button>Upgrade to Add Documents</Button>
+      </UpgradePlanModal>
+    );
+  }
+
   return (
     <Dialog open={isOpen} onOpenChange={clearModelStates}>
       <DialogTrigger asChild>{children}</DialogTrigger>

Diff for: components/ui/progress.tsx

@@ -27,7 +27,10 @@ const Progress = React.forwardRef<
       {...props}
     >
       <ProgressPrimitive.Indicator
-        className="h-full w-full flex-1 bg-primary transition-all"
+        className={cn(
+          "h-full w-full flex-1 bg-primary transition-all",
+          value && value > 100 ? "bg-destructive dark:bg-red-600" : "",
+        )}
         style={{ transform: `translateX(-${100 - (value || 0)}%)` }}
       />
       {text && !error ? (

Diff for: components/upload-zone.tsx

@@ -6,12 +6,14 @@ import { useTeam } from "@/context/team-context";
 import { DocumentStorageType } from "@prisma/client";
 import { useSession } from "next-auth/react";
 import { FileRejection, useDropzone } from "react-dropzone";
+import { toast } from "sonner";
 import { mutate } from "swr";
 
 import { useAnalytics } from "@/lib/analytics";
 import { DocumentData, createDocument } from "@/lib/documents/create-document";
 import { resumableUpload } from "@/lib/files/tus-upload";
 import { usePlan } from "@/lib/swr/use-billing";
+import useLimits from "@/lib/swr/use-limits";
 import { CustomUser } from "@/lib/types";
 import { cn } from "@/lib/utils";
 import { getSupportedContentType } from "@/lib/utils/get-content-type";
@@ -70,13 +72,21 @@ export default function UploadZone({
   const isTrial = !!trial;
   const maxSize = plan === "business" || plan === "datarooms" ? 250 : 30;
   const maxNumPages = plan === "business" || plan === "datarooms" ? 500 : 100;
+  const { limits, canAddDocuments } = useLimits();
+  const remainingDocuments = limits?.documents
+    ? limits?.documents - limits?.usage?.documents
+    : 0;
 
   const [progress, setProgress] = useState<number>(0);
   const [showProgress, setShowProgress] = useState(false);
   const uploadProgress = useRef<number[]>([]);
 
   const onDrop = useCallback(
     (acceptedFiles: FileWithPath[]) => {
+      if (!canAddDocuments && acceptedFiles.length > remainingDocuments) {
+        toast.error("You have reached the maximum number of documents.");
+        return;
+      }
       const newUploads = acceptedFiles.map((file) => ({
         fileName: file.name,
         progress: 0,

Diff for: ee/limits/constants.ts

@@ -1,5 +1,20 @@
+// INFO: for numeric values,`null` means unlimited
+
+export type TPlanLimits = {
+  users: number;
+  links: number | null;
+  documents: number | null;
+  domains: number;
+  datarooms: number;
+  customDomainOnPro: boolean;
+  customDomainInDataroom: boolean;
+  advancedLinkControlsOnPro: boolean | null;
+};
+
 export const FREE_PLAN_LIMITS = {
   users: 1,
+  links: 10,
+  documents: 10,
   domains: 0,
   datarooms: 0,
   customDomainOnPro: false,
@@ -9,7 +24,9 @@ export const FREE_PLAN_LIMITS = {
 
 export const PRO_PLAN_LIMITS = {
   users: 2,
-  domains: 5,
+  links: null,
+  documents: 100,
+  domains: 2,
   datarooms: 0,
   customDomainOnPro: false,
   customDomainInDataroom: false,
@@ -18,6 +35,8 @@ export const PRO_PLAN_LIMITS = {
 
 export const BUSINESS_PLAN_LIMITS = {
   users: 3,
+  links: null,
+  documents: null,
   domains: 5,
   datarooms: 1,
   customDomainOnPro: true,
@@ -27,6 +46,8 @@ export const BUSINESS_PLAN_LIMITS = {
 
 export const DATAROOMS_PLAN_LIMITS = {
   users: 3,
+  links: null,
+  documents: null,
   domains: 10,
   datarooms: 100,
   customDomainOnPro: true,

Diff for: ee/limits/server.ts

@@ -2,6 +2,44 @@ import { z } from "zod";
 
 import prisma from "@/lib/prisma";
 
+import {
+  BUSINESS_PLAN_LIMITS,
+  DATAROOMS_PLAN_LIMITS,
+  FREE_PLAN_LIMITS,
+  PRO_PLAN_LIMITS,
+  TPlanLimits,
+} from "./constants";
+
+// Function to determine if a plan is free or free+drtrial
+const isFreePlan = (plan: string) => plan === "free" || plan === "free+drtrial";
+
+// Function to get the base plan from a plan string
+const getBasePlan = (plan: string) => plan.split("+")[0];
+
+const planLimitsMap: Record<string, TPlanLimits> = {
+  free: FREE_PLAN_LIMITS,
+  pro: PRO_PLAN_LIMITS,
+  business: BUSINESS_PLAN_LIMITS,
+  datarooms: DATAROOMS_PLAN_LIMITS,
+};
+
+const configSchema = z.object({
+  datarooms: z.number(),
+  links: z
+    .preprocess((v) => (v === null ? Infinity : Number(v)), z.number())
+    .optional()
+    .default(10),
+  documents: z
+    .preprocess((v) => (v === null ? Infinity : Number(v)), z.number())
+    .optional()
+    .default(10),
+  users: z.number(),
+  domains: z.number(),
+  customDomainOnPro: z.boolean(),
+  customDomainInDataroom: z.boolean(),
+  advancedLinkControlsOnPro: z.boolean().nullish(),
+});
+
 export async function getLimits({
   teamId,
   userId,
@@ -21,31 +59,63 @@ export async function getLimits({
     select: {
       plan: true,
       limits: true,
+      _count: {
+        select: {
+          documents: true,
+        },
+      },
+      documents: {
+        select: {
+          _count: {
+            select: {
+              links: true,
+            },
+          },
+        },
+      },
     },
   });
 
   if (!team) {
     throw new Error("Team not found");
   }
 
+  const documentCount = team._count.documents;
+  const linkCount = team.documents.reduce(
+    (sum, doc) => sum + doc._count.links,
+    0,
+  );
+
   // parse the limits json with zod and return the limits
   // {datarooms: 1, users: 1, domains: 1, customDomainOnPro: boolean, customDomainInDataroom: boolean}
 
-  const configSchema = z.object({
-    datarooms: z.number(),
-    users: z.number(),
-    domains: z.number(),
-    customDomainOnPro: z.boolean(),
-    customDomainInDataroom: z.boolean(),
-    advancedLinkControlsOnPro: z.boolean().nullish(),
-  });
-
   try {
-    const parsedData = configSchema.parse(team.limits);
+    let parsedData = configSchema.parse(team.limits);
 
-    return parsedData;
+    // Adjust limits based on the plan if they're at the default value
+    if (isFreePlan(team.plan)) {
+      return {
+        ...parsedData,
+        usage: { documents: documentCount, links: linkCount },
+      };
+    } else {
+      return {
+        ...parsedData,
+        // if account is paid, but link and document limits are not set, then set them to Infinity
+        links: parsedData.links === 10 ? Infinity : parsedData.links,
+        documents:
+          parsedData.documents === 10 ? Infinity : parsedData.documents,
+        usage: { documents: documentCount, links: linkCount },
+      };
+    }
   } catch (error) {
-    // if no limits set, then return null and don't block the team
-    return null;
+    // if no limits set or parsing fails, return default limits based on the plan
+    const basePlan = getBasePlan(team.plan);
+    const defaultLimits = planLimitsMap[basePlan] || FREE_PLAN_LIMITS;
+
+    return {
+      ...defaultLimits,
+      usage: { documents: documentCount, links: linkCount },
+    };
   }
 }