chore: add pricing plan lock for machine types

web/src/app/(app)/machines/page.tsx

@@ -1,5 +1,6 @@
 import { AccessType } from "../../../lib/AccessType";
 import { MachineList } from "@/components/MachineList";
+import { SubscriptionProvider } from "@/components/useCurrentPlan";
 import { db } from "@/db/db";
 import { machinesTable } from "@/db/schema";
 import { getCurrentPlanWithAuth } from "@/server/getCurrentPlan";
@@ -32,12 +33,12 @@ async function MachineListServer() {
 
   return (
     <div className="w-full">
-      {/* <div>Machines</div> */}
-      <MachineList
-        sub={sub}
-        data={machines}
-        userMetadata={AccessType.parse(user.privateMetadata ?? {})}
-      />
+      <SubscriptionProvider sub={sub}>
+        <MachineList
+          data={machines}
+          userMetadata={AccessType.parse(user.privateMetadata ?? {})}
+        />
+      </SubscriptionProvider>
     </div>
   );
 }

web/src/components/CurrentPlanContextType.tsx

@@ -0,0 +1,24 @@
+"use client";
+import { getCurrentPlanWithAuth } from "@/server/getCurrentPlan";
+import * as React from "react";
+import { createContext } from "react";
+
+export type CurrentPlanContextType = Awaited<
+  ReturnType<typeof getCurrentPlanWithAuth>
+>;
+export const CurrentPlanContext = createContext<
+  CurrentPlanContextType | undefined
+>(undefined);
+export function SubscriptionProvider({
+  sub,
+  children,
+}: {
+  sub: CurrentPlanContextType;
+  children: React.ReactNode;
+}) {
+  return (
+    <CurrentPlanContext.Provider value={sub}>
+      {children}
+    </CurrentPlanContext.Provider>
+  );
+}

web/src/components/MachineList.tsx

@@ -183,6 +183,7 @@ export const columns: ColumnDef<Machine>[] = [
     cell: ({ row }) => {
       const machine = row.original;
       const [open, setOpen] = useState(false);
+      const sub = useCurrentPlan();
 
       return (
         <DropdownMenu>
@@ -291,7 +292,10 @@ export const columns: ColumnDef<Machine>[] = [
                   fieldType: "models",
                 },
                 gpu: {
-                  inputProps: {},
+                  fieldType: "gpuPicker",
+                  inputProps: {
+                    sub: sub,
+                  },
                 },
               }}
             />
@@ -319,14 +323,14 @@ export const columns: ColumnDef<Machine>[] = [
   },
 ];
 
+import { useCurrentPlan } from "./useCurrentPlan";
+
 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>(
@@ -336,6 +340,8 @@ export function MachineList({
     React.useState<VisibilityState>({});
   const [rowSelection, setRowSelection] = React.useState({});
 
+  const sub = useCurrentPlan();
+
   const table = useReactTable({
     data,
     columns,
@@ -420,11 +426,9 @@ export function MachineList({
                 },
               },
               gpu: {
-                fieldType: !userMetadata.betaFeaturesAccess
-                  ? "fallback"
-                  : "select",
+                fieldType: "gpuPicker",
                 inputProps: {
-                  disabled: !userMetadata.betaFeaturesAccess,
+                  sub: sub,
                 },
               },
             }}

web/src/components/custom-form/gpu-picker.tsx

@@ -0,0 +1,90 @@
+import { AutoFormInputComponentProps } from "@/components/ui/auto-form/types";
+import { getBaseSchema } from "@/components/ui/auto-form/utils";
+import {
+  FormItem,
+  FormLabel,
+  FormControl,
+  FormDescription,
+  FormMessage,
+} from "@/components/ui/form";
+import {
+  Select,
+  SelectContent,
+  SelectItem,
+  SelectTrigger,
+  SelectValue,
+} from "@/components/ui/select";
+import { Lock } from "lucide-react";
+import * as z from "zod";
+
+export default function AutoFormGPUPicker({
+  label,
+  isRequired,
+  field,
+  fieldConfigItem,
+  zodItem,
+}: AutoFormInputComponentProps) {
+  const baseValues = (getBaseSchema(zodItem) as unknown as z.ZodEnum<any>)._def
+    .values;
+
+  let values: [string, string][] = [];
+  if (!Array.isArray(baseValues)) {
+    values = Object.entries(baseValues);
+  } else {
+    values = baseValues.map((value) => [value, value]);
+  }
+
+  function findItem(value: any) {
+    return values.find((item) => item[0] === value);
+  }
+
+  const plan = fieldConfigItem.inputProps?.sub?.plan;
+  const enabledGPU = ["T4"];
+
+  if (plan == "pro") {
+    enabledGPU.push("A10G");
+  } else if (plan == "enterprise") {
+    enabledGPU.push("A10G");
+    enabledGPU.push("A100");
+  }
+
+  return (
+    <FormItem>
+      <FormLabel>
+        {label}
+        {isRequired && <span className="text-destructive"> *</span>}
+      </FormLabel>
+      <FormControl>
+        <Select onValueChange={field.onChange} defaultValue={field.value}>
+          <SelectTrigger>
+            <SelectValue
+              className="w-full"
+              placeholder={fieldConfigItem.inputProps?.placeholder}
+            >
+              {field.value ? findItem(field.value)?.[1] : "Select an option"}
+            </SelectValue>
+          </SelectTrigger>
+          <SelectContent>
+            {values.map(([value, label]) => {
+              const enabled = enabledGPU.includes(value);
+              return (
+                <SelectItem value={label} key={value} disabled={!enabled}>
+                  {label}
+                  {!enabled && (
+                    <span className="mx-2 inline-flex items-center justify-center gap-2">
+                      Upgrade to enabled this options <Lock size={14}></Lock>
+                    </span>
+                  )}
+                </SelectItem>
+              );
+            })}
+          </SelectContent>
+        </Select>
+      </FormControl>
+      {fieldConfigItem.description && (
+        <FormDescription>{fieldConfigItem.description}</FormDescription>
+      )}
+      <FormMessage />
+    </FormItem>
+  );
+}

web/src/components/ui/auto-form/config.ts

@@ -1,3 +1,4 @@
+import AutoFormGPUPicker from "@/components/custom-form/gpu-picker";
 import AutoFormCheckbox from "./fields/checkbox";
 import AutoFormDate from "./fields/date";
 import AutoFormEnum from "./fields/enum";
@@ -22,6 +23,7 @@ export const INPUT_COMPONENTS = {
   // Customs
   snapshot: AutoFormSnapshotPicker,
   models: AutoFormModelsPicker,
+  gpuPicker: AutoFormGPUPicker,
 };
 
 /**

web/src/components/ui/auto-form/types.ts

@@ -1,3 +1,4 @@
+import type { getCurrentPlanWithAuth } from "@/server/getCurrentPlan";
 import type { INPUT_COMPONENTS } from "./config";
 import type { ControllerRenderProps, FieldValues } from "react-hook-form";
 import type * as z from "zod";
@@ -6,6 +7,7 @@ export type FieldConfigItem = {
   description?: React.ReactNode;
   inputProps?: React.InputHTMLAttributes<HTMLInputElement> & {
     showLabel?: boolean;
+    sub?: Awaited<ReturnType<typeof getCurrentPlanWithAuth>>;
   };
   fieldType?:
     | keyof typeof INPUT_COMPONENTS

web/src/components/useCurrentPlan.tsx

@@ -0,0 +1,36 @@
+"use client";
+
+import { getCurrentPlanWithAuth } from "@/server/getCurrentPlan";
+import * as React from "react";
+import { createContext, useContext } from "react";
+
+type CurrentPlanContextType = Awaited<
+  ReturnType<typeof getCurrentPlanWithAuth>
+>;
+const CurrentPlanContext = createContext<CurrentPlanContextType | undefined>(
+  undefined,
+);
+
+export function SubscriptionProvider({
+  sub,
+  children,
+}: {
+  sub: CurrentPlanContextType;
+  children: React.ReactNode;
+}) {
+  return (
+    <CurrentPlanContext.Provider value={sub}>
+      {children}
+    </CurrentPlanContext.Provider>
+  );
+}
+
+export const useCurrentPlan = (): CurrentPlanContextType => {
+  const context = useContext(CurrentPlanContext);
+
+  // if (context === undefined) {
+  //   throw new Error("useCurrentPlan must be used within a CurrentPlanProvider");
+  // }
+
+  return context;
+};

web/src/server/getCurrentPlan.tsx

@@ -3,6 +3,7 @@ import { and, desc, eq, isNull, or } from "drizzle-orm";
 import { subscriptionStatusTable } from "@/db/schema";
 import { APIKeyUserType } from "@/server/APIKeyBodyRequest";
 import { auth } from "@clerk/nextjs";
+import "server-only";
 
 export async function getCurrentPlanWithAuth() {
   const { userId, orgId } = auth();
@@ -23,7 +24,10 @@ export async function getCurrentPlan({ user_id, org_id }: APIKeyUserType) {
       eq(subscriptionStatusTable.user_id, user_id),
       org_id
         ? eq(subscriptionStatusTable.org_id, org_id)
-        : or(isNull(subscriptionStatusTable.org_id), eq(subscriptionStatusTable.org_id, "")),
+        : or(
+            isNull(subscriptionStatusTable.org_id),
+            eq(subscriptionStatusTable.org_id, ""),
+          ),
     ),
     orderBy: desc(subscriptionStatusTable.created_at),
   });