fix: handle file upload status correctly, add serverless machine type

web/src/server/addMachineSchema.ts

@@ -0,0 +1,8 @@
+import { insertMachineSchema } from "@/db/schema";
+
+export const addMachineSchema = insertMachineSchema.pick({
+  name: true,
+  endpoint: true,
+  type: true,
+  auth_token: true,
+});

web/src/server/createRun.ts

@@ -7,6 +7,7 @@ import { ComfyAPI_Run } from "@/types/ComfyAPI_Run";
 import { and, eq } from "drizzle-orm";
 import { revalidatePath } from "next/cache";
 import "server-only";
+import { v4 } from "uuid";
 
 export const createRun = withServerPromise(
   async (
@@ -37,8 +38,6 @@ export const createRun = withServerPromise(
       throw new Error("Workflow version not found");
     }
 
-    const comfyui_endpoint = `${machine.endpoint}/comfyui-deploy/run`;
-
     const workflow_api = workflow_version_data.workflow_api;
 
     // Replace the inputs
@@ -52,33 +51,57 @@ export const createRun = withServerPromise(
       }
     }
 
-    const body = {
+    let prompt_id: string | undefined = undefined;
+    const shareData = {
       workflow_api: workflow_api,
       status_endpoint: `${origin}/api/update-run`,
       file_upload_endpoint: `${origin}/api/file-upload`,
     };
-    // console.log(body);
-    const bodyJson = JSON.stringify(body);
-    // console.log(bodyJson);
 
-    // Sending to comfyui
-    const _result = await fetch(comfyui_endpoint, {
-      method: "POST",
-      body: bodyJson,
-      cache: "no-store",
-    });
-
-    if (!_result.ok) {
-      throw new Error(`Error creating run, ${_result.statusText}`);
+    switch (machine.type) {
+      case "runpod-serverless":
+        prompt_id = v4();
+        const data = {
+          input: {
+            ...shareData,
+            prompt_id: prompt_id,
+          },
+        };
+        console.log(data);
+        const __result = await fetch(`${machine.endpoint}/run`, {
+          method: "POST",
+          headers: {
+            "Content-Type": "application/json",
+            Authorization: `Bearer ${machine.auth_token}`,
+          },
+          body: JSON.stringify(data),
+          cache: "no-store",
+        });
+        console.log(__result);
+        if (!__result.ok)
+          throw new Error(`Error creating run, ${__result.statusText}`);
+        console.log(data, __result);
+        break;
+      case "classic":
+        const body = shareData;
+        const comfyui_endpoint = `${machine.endpoint}/comfyui-deploy/run`;
+        const _result = await fetch(comfyui_endpoint, {
+          method: "POST",
+          body: JSON.stringify(body),
+          cache: "no-store",
+        });
+        if (!_result.ok)
+          throw new Error(`Error creating run, ${_result.statusText}`);
+        const result = await ComfyAPI_Run.parseAsync(await _result.json());
+        prompt_id = result.prompt_id;
+        break;
     }
 
-    const result = await ComfyAPI_Run.parseAsync(await _result.json());
-
     // Add to our db
     const workflow_run = await db
       .insert(workflowRunsTable)
       .values({
-        id: result.prompt_id,
+        id: prompt_id,
         workflow_id: workflow_version_data.workflow_id,
         workflow_version_id: workflow_version_data.id,
         workflow_inputs: inputs,

web/src/server/curdMachine.ts

@@ -1,5 +1,6 @@
 "use server";
 
+import type { addMachineSchema } from "./addMachineSchema";
 import { withServerPromise } from "./withServerPromise";
 import { db } from "@/db/db";
 import { machinesTable } from "@/db/schema";
@@ -7,6 +8,7 @@ import { auth } from "@clerk/nextjs";
 import { and, eq } from "drizzle-orm";
 import { revalidatePath } from "next/cache";
 import "server-only";
+import type { z } from "zod";
 
 export async function getMachines() {
   const { userId } = auth();
@@ -20,17 +22,36 @@ export async function getMachines() {
   return machines;
 }
 
-export async function addMachine(name: string, endpoint: string) {
-  const { userId } = auth();
-  if (!userId) throw new Error("No user id");
-  console.log(name, endpoint);
-  await db.insert(machinesTable).values({
-    user_id: userId,
-    name,
-    endpoint,
-  });
-  revalidatePath("/machines");
-}
+export const addMachine = withServerPromise(
+  async ({ name, endpoint, type }: z.infer<typeof addMachineSchema>) => {
+    const { userId } = auth();
+    if (!userId) return { error: "No user id" };
+    console.log(name, endpoint);
+    await db.insert(machinesTable).values({
+      user_id: userId,
+      name,
+      endpoint,
+      type,
+    });
+    revalidatePath("/machines");
+    return { message: "Machine Added" };
+  }
+);
+
+export const updateMachine = withServerPromise(
+  async ({
+    id,
+    ...data
+  }: z.infer<typeof addMachineSchema> & {
+    id: string;
+  }) => {
+    const { userId } = auth();
+    if (!userId) return { error: "No user id" };
+    await db.update(machinesTable).set(data).where(eq(machinesTable.id, id));
+    revalidatePath("/machines");
+    return { message: "Machine Updated" };
+  }
+);
 
 export const deleteMachine = withServerPromise(
   async (machine_id: string): Promise<{ message: string }> => {