From b8e478c8ee3a7747a918eb7f389a8ef489eb52ab Mon Sep 17 00:00:00 2001
From: BennyKok <itechbenny@gmail.com>
Date: Fri, 12 Jan 2024 18:56:02 +0800
Subject: [PATCH] feat: add rebuild machine option

---
 web/src/components/MachineList.tsx | 29 ++++++++++++++++++++-------
 web/src/server/curdMachine.ts      | 32 +++++++++++++++++++++++++++---
 2 files changed, 51 insertions(+), 10 deletions(-)

diff --git a/web/src/components/MachineList.tsx b/web/src/components/MachineList.tsx
index af21e41..c6c7148 100644
--- a/web/src/components/MachineList.tsx
+++ b/web/src/components/MachineList.tsx
@@ -32,6 +32,7 @@ import {
 import {
   addCustomMachine,
   addMachine,
+  buildMachine,
   deleteMachine,
   disableMachine,
   enableMachine,
@@ -202,14 +203,28 @@ export const columns: ColumnDef<Machine>[] = [
               </DropdownMenuItem>
             )}
             {machine.type === "comfy-deploy-serverless" && (
-              <DropdownMenuItem asChild>
-                <a
-                  target="_blank"
-                  href={machine.endpoint.replace("comfyui-api", "comfyui-app")}
+              <>
+                <DropdownMenuItem asChild>
+                  <a
+                    target="_blank"
+                    href={machine.endpoint.replace(
+                      "comfyui-api",
+                      "comfyui-app"
+                    )}
+                  >
+                    Open ComfyUI
+                  </a>
+                </DropdownMenuItem>
+                <DropdownMenuItem
+                  onClick={() => {
+                    buildMachine({
+                      id: machine.id,
+                    });
+                  }}
                 >
-                  Open ComfyUI
-                </a>
-              </DropdownMenuItem>
+                  Rebuild
+                </DropdownMenuItem>
+              </>
             )}
             <DropdownMenuItem onClick={() => setOpen(true)}>
               Edit
diff --git a/web/src/server/curdMachine.ts b/web/src/server/curdMachine.ts
index 4a31685..b59415f 100644
--- a/web/src/server/curdMachine.ts
+++ b/web/src/server/curdMachine.ts
@@ -113,7 +113,7 @@ export const updateCustomMachine = withServerPromise(
         .where(eq(machinesTable.id, id));
 
       // Perform custom build if there are changes
-      await buildMachine(data, currentMachine);
+      await _buildMachine(data, currentMachine);
       redirect(`/machines/${id}`);
     } else {
       revalidatePath("/machines");
@@ -123,6 +123,32 @@ export const updateCustomMachine = withServerPromise(
   }
 );
 
+export const buildMachine = withServerPromise(
+  async ({ id }: { id: string }) => {
+    const { userId } = auth();
+    if (!userId) return { error: "No user id" };
+
+    const currentMachine = await db.query.machinesTable.findFirst({
+      where: eq(machinesTable.id, id),
+    });
+
+    if (!currentMachine) return { error: "Machine not found" };
+
+    const datas = await db
+      .update(machinesTable)
+      .set({
+        status: "building",
+        endpoint: "not-ready",
+      })
+      .where(eq(machinesTable.id, id))
+      .returning();
+
+    // Perform custom build if there are changes
+    await _buildMachine(datas[0], currentMachine);
+    redirect(`/machines/${id}`);
+  }
+);
+
 export const addCustomMachine = withServerPromise(
   async (data: z.infer<typeof addCustomMachineSchema>) => {
     const { userId, orgId } = auth();
@@ -143,14 +169,14 @@ export const addCustomMachine = withServerPromise(
 
     const b = a[0];
 
-    await buildMachine(data, b);
+    await _buildMachine(data, b);
     redirect(`/machines/${b.id}`);
     // revalidatePath("/machines");
     return { message: "Machine Building" };
   }
 );
 
-async function buildMachine(
+async function _buildMachine(
   data: z.infer<typeof addCustomMachineSchema>,
   b: MachineType
 ) {