From cb01c896a0c30a37371e35af7ab9f7c8c608111f Mon Sep 17 00:00:00 2001
From: bennykok <itechbenny@gmail.com>
Date: Sun, 28 Jan 2024 00:51:44 +0800
Subject: [PATCH] feat: adding add models lock according to prcing plan

---
 web/src/app/(app)/machines/page.tsx |  4 +++
 web/src/components/InsertModal.tsx  | 10 ++++----
 web/src/components/MachineList.tsx  | 38 +++++++++++++++++++++--------
 3 files changed, 37 insertions(+), 15 deletions(-)

diff --git a/web/src/app/(app)/machines/page.tsx b/web/src/app/(app)/machines/page.tsx
index 7366668..c4c3de2 100644
--- a/web/src/app/(app)/machines/page.tsx
+++ b/web/src/app/(app)/machines/page.tsx
@@ -2,6 +2,7 @@ import { AccessType } from "../../../lib/AccessType";
 import { MachineList } from "@/components/MachineList";
 import { db } from "@/db/db";
 import { machinesTable } from "@/db/schema";
+import { getCurrentPlanWithAuth } from "@/server/getCurrentPlan";
 import { auth } from "@clerk/nextjs";
 import { clerkClient } from "@clerk/nextjs/server";
 import { desc, eq, isNull, and } from "drizzle-orm";
@@ -19,6 +20,8 @@ async function MachineListServer() {
 
   const user = await clerkClient.users.getUser(userId);
 
+  const sub = await getCurrentPlanWithAuth();
+
   const machines = await db.query.machinesTable.findMany({
     orderBy: desc(machinesTable.updated_at),
     where:
@@ -31,6 +34,7 @@ async function MachineListServer() {
     <div className="w-full">
       {/* <div>Machines</div> */}
       <MachineList
+        sub={sub}
         data={machines}
         userMetadata={AccessType.parse(user.privateMetadata ?? {})}
       />
diff --git a/web/src/components/InsertModal.tsx b/web/src/components/InsertModal.tsx
index a6b3280..7df456a 100644
--- a/web/src/components/InsertModal.tsx
+++ b/web/src/components/InsertModal.tsx
@@ -26,11 +26,11 @@ import type { UnknownKeysParam, ZodObject, ZodRawShape, z } from "zod";
 export function InsertModal<
   K extends ZodRawShape,
   Y extends UnknownKeysParam,
-  Z extends ZodObject<K, Y>
+  Z extends ZodObject<K, Y>,
 >(props: {
   tooltip?: string;
   disabled?: boolean;
-  title: string;
+  title: React.ReactNode;
   description: string;
   dialogClassName?: string;
   serverAction: (data: z.infer<Z>) => Promise<unknown>;
@@ -105,7 +105,7 @@ export function InsertModal<
 export function UpdateModal<
   K extends ZodRawShape,
   Y extends UnknownKeysParam,
-  Z extends ZodObject<K, Y>
+  Z extends ZodObject<K, Y>,
 >(props: {
   open?: boolean;
   setOpen?: (open: boolean) => void;
@@ -118,7 +118,7 @@ export function UpdateModal<
   serverAction: (
     data: z.infer<Z> & {
       id: string;
-    }
+    },
   ) => Promise<unknown>;
   formSchema: Z;
   fieldConfig?: FieldConfig<z.infer<Z>>;
@@ -165,7 +165,7 @@ export function UpdateModal<
               props.serverAction({
                 ...data,
                 id: props.data.id,
-              })
+              }),
             );
             setIsLoading(false);
             setOpen(false);
diff --git a/web/src/components/MachineList.tsx b/web/src/components/MachineList.tsx
index 81db3f4..e6dde79 100644
--- a/web/src/components/MachineList.tsx
+++ b/web/src/components/MachineList.tsx
@@ -41,6 +41,7 @@ import {
   updateMachine,
 } from "@/server/curdMachine";
 import { editWorkflowOnMachine } from "@/server/editWorkflowOnMachine";
+import { getCurrentPlanWithAuth } from "@/server/getCurrentPlan";
 import type {
   ColumnDef,
   ColumnFiltersState,
@@ -55,7 +56,7 @@ import {
   getSortedRowModel,
   useReactTable,
 } from "@tanstack/react-table";
-import { ArrowUpDown, MoreHorizontal } from "lucide-react";
+import { ArrowUpDown, Lock, MoreHorizontal, Plus } from "lucide-react";
 import * as React from "react";
 import { useState } from "react";
 import { toast } from "sonner";
@@ -321,9 +322,11 @@ export const columns: ColumnDef<Machine>[] = [
 export function MachineList({
   data,
   userMetadata,
+  sub,
 }: {
   data: Machine[];
   userMetadata: z.infer<typeof AccessType>;
+  sub: Awaited<ReturnType<typeof getCurrentPlanWithAuth>>;
 }) {
   const [sorting, setSorting] = React.useState<SortingState>([]);
   const [columnFilters, setColumnFilters] = React.useState<ColumnFiltersState>(
@@ -352,6 +355,21 @@ export function MachineList({
     },
   });
 
+  let machineMaxCount = 2;
+
+  // Temp fixes for machine count
+  if (userMetadata.betaFeaturesAccess) machineMaxCount = 5;
+
+  if (sub?.plan == "pro") {
+    machineMaxCount = 10;
+  } else if (sub?.plan == "enterprise") {
+    machineMaxCount = 99;
+  }
+
+  const locked =
+    data.some((machine) => machine.type === "modal-serverless") &&
+    data.length >= machineMaxCount;
+
   return (
     <div className="w-full">
       <div className="flex items-center py-4">
@@ -366,17 +384,17 @@ export function MachineList({
         <div className="ml-auto flex gap-2">
           <InsertModal
             dialogClassName="sm:max-w-[600px]"
-            disabled={
-              data.some(
-                (machine) => machine.type === "comfy-deploy-serverless",
-              ) && !userMetadata.betaFeaturesAccess
-            }
+            disabled={locked}
             tooltip={
-              data.some((machine) => machine.type === "comfy-deploy-serverless")
-                ? "Only one hosted machine at preview stage"
-                : undefined
+              locked
+                ? `Max ${machineMaxCount} ComfyUI machine for your account, upgrade to unlock more cnfiguration.`
+                : `Max ${machineMaxCount} ComfyUI machine for your account`
+            }
+            title={
+              <>
+                New Machine {locked ? <Lock size={14} /> : <Plus size={14} />}
+              </>
             }
-            title="New Machine"
             description="Add custom ComfyUI machines to your account."
             serverAction={addCustomMachine}
             formSchema={addCustomMachineSchema}