feat: add create workflow run error server action catch, redirect workflow parse error, add key revoked col

custom_routes.py

@@ -14,7 +14,7 @@ import io
 import time
 import execution
 import random
-
+import traceback
 import uuid
 import asyncio
 import atexit
@@ -90,7 +90,14 @@ async def comfy_deploy_run(request):
         "client_id": "fake_client" #api.client_id
     }
 
-    res = post_prompt(prompt)
+    try:
+        res = post_prompt(prompt)
+    except Exception as e:
+        error_type = type(e).__name__
+        stack_trace = traceback.format_exc().strip().split('\n')[-2]
+        print(f"error: {error_type}, {e}")
+        print(f"stack trace: {stack_trace}")
+        return web.Response(status=500, reason=f"{error_type}: {e}, {stack_trace}")
 
     prompt_metadata[res['prompt_id']] = {
         'status_endpoint': data.get('status_endpoint'),
@@ -252,5 +259,9 @@ async def update_run_with_output(prompt_id, data):
         }
         requests.post(status_endpoint, json=body)
 
+        await send('outputs_uploaded', {
+            "prompt_id": prompt_id
+        })
+
 prompt_server.send_json_original = prompt_server.send_json
 prompt_server.send_json = send_json_override.__get__(prompt_server, server.PromptServer)

web-plugin/index.js

@@ -5,7 +5,7 @@ import { ComfyWidgets, LGraphNode } from "./widgets.js";
 /** @typedef {import('../../../web/types/comfy.js').ComfyExtension} ComfyExtension*/
 /** @type {ComfyExtension} */
 const ext = {
-  name: "BennyKok.ComfyDeploy",
+  name: "BennyKok.ComfyUIDeploy",
 
   init(app) {
     addButton();
@@ -65,7 +65,7 @@ const ext = {
     // Load default visibility
 
     LiteGraph.registerNodeType(
-      "Comfy Deploy",
+      "ComfyDeploy",
       Object.assign(ComfyDeploy, {
         title_mode: LiteGraph.NORMAL_TITLE,
         title: "Comfy Deploy",
@@ -112,7 +112,7 @@ function addButton() {
     /** @type {LGraph} */
     const graph = app.graph;
 
-    const deployMeta = graph.findNodesByType("Comfy Deploy");
+    const deployMeta = graph.findNodesByType("ComfyDeploy");
     const deployMetaNode = deployMeta[0];
 
     console.log(deployMetaNode);
@@ -137,6 +137,8 @@ function addButton() {
     deploy.textContent = "Deploying...";
     deploy.style.color = "orange";
 
+    console.log(prompt);
+
     const apiRoute = endpoint + "/api/upload"
     // const userId = apiKey
     try {

web/drizzle/0010_organic_marten_broadcloak.sql

@@ -0,0 +1 @@
+ALTER TABLE "comfyui_deploy"."api_keys" ADD COLUMN "revoked" boolean DEFAULT false NOT NULL;

web/drizzle/meta/0010_snapshot.json

@@ -0,0 +1,621 @@
+{
+  "id": "128aae04-697c-4253-b7de-b27e895e03f8",
+  "prevId": "a5d542ea-484f-431a-b31f-170e789a03a7",
+  "version": "5",
+  "dialect": "pg",
+  "tables": {
+    "api_keys": {
+      "name": "api_keys",
+      "schema": "comfyui_deploy",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": true,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "key": {
+          "name": "key",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "name": {
+          "name": "name",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "user_id": {
+          "name": "user_id",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "org_id": {
+          "name": "org_id",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": false
+        },
+        "revoked": {
+          "name": "revoked",
+          "type": "boolean",
+          "primaryKey": false,
+          "notNull": true,
+          "default": false
+        },
+        "created_at": {
+          "name": "created_at",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "now()"
+        },
+        "updated_at": {
+          "name": "updated_at",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "now()"
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "api_keys_user_id_users_id_fk": {
+          "name": "api_keys_user_id_users_id_fk",
+          "tableFrom": "api_keys",
+          "tableTo": "users",
+          "columnsFrom": [
+            "user_id"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "cascade",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {
+        "api_keys_key_unique": {
+          "name": "api_keys_key_unique",
+          "nullsNotDistinct": false,
+          "columns": [
+            "key"
+          ]
+        }
+      }
+    },
+    "deployments": {
+      "name": "deployments",
+      "schema": "comfyui_deploy",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": true,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "user_id": {
+          "name": "user_id",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "workflow_version_id": {
+          "name": "workflow_version_id",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "workflow_id": {
+          "name": "workflow_id",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "machine_id": {
+          "name": "machine_id",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "environment": {
+          "name": "environment",
+          "type": "deployment_environment",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "created_at": {
+          "name": "created_at",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "now()"
+        },
+        "updated_at": {
+          "name": "updated_at",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "now()"
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "deployments_user_id_users_id_fk": {
+          "name": "deployments_user_id_users_id_fk",
+          "tableFrom": "deployments",
+          "tableTo": "users",
+          "columnsFrom": [
+            "user_id"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "cascade",
+          "onUpdate": "no action"
+        },
+        "deployments_workflow_version_id_workflow_versions_id_fk": {
+          "name": "deployments_workflow_version_id_workflow_versions_id_fk",
+          "tableFrom": "deployments",
+          "tableTo": "workflow_versions",
+          "columnsFrom": [
+            "workflow_version_id"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        },
+        "deployments_workflow_id_workflows_id_fk": {
+          "name": "deployments_workflow_id_workflows_id_fk",
+          "tableFrom": "deployments",
+          "tableTo": "workflows",
+          "columnsFrom": [
+            "workflow_id"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        },
+        "deployments_machine_id_machines_id_fk": {
+          "name": "deployments_machine_id_machines_id_fk",
+          "tableFrom": "deployments",
+          "tableTo": "machines",
+          "columnsFrom": [
+            "machine_id"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {}
+    },
+    "machines": {
+      "name": "machines",
+      "schema": "comfyui_deploy",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": true,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "user_id": {
+          "name": "user_id",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "name": {
+          "name": "name",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "endpoint": {
+          "name": "endpoint",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "created_at": {
+          "name": "created_at",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "now()"
+        },
+        "updated_at": {
+          "name": "updated_at",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "now()"
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "machines_user_id_users_id_fk": {
+          "name": "machines_user_id_users_id_fk",
+          "tableFrom": "machines",
+          "tableTo": "users",
+          "columnsFrom": [
+            "user_id"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "cascade",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {}
+    },
+    "users": {
+      "name": "users",
+      "schema": "comfyui_deploy",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "text",
+          "primaryKey": true,
+          "notNull": true
+        },
+        "username": {
+          "name": "username",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "name": {
+          "name": "name",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "created_at": {
+          "name": "created_at",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": false,
+          "default": "now()"
+        },
+        "updated_at": {
+          "name": "updated_at",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": false,
+          "default": "now()"
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {},
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {}
+    },
+    "workflow_run_outputs": {
+      "name": "workflow_run_outputs",
+      "schema": "comfyui_deploy",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": true,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "run_id": {
+          "name": "run_id",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "data": {
+          "name": "data",
+          "type": "jsonb",
+          "primaryKey": false,
+          "notNull": false
+        },
+        "created_at": {
+          "name": "created_at",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "now()"
+        },
+        "updated_at": {
+          "name": "updated_at",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "now()"
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "workflow_run_outputs_run_id_workflow_runs_id_fk": {
+          "name": "workflow_run_outputs_run_id_workflow_runs_id_fk",
+          "tableFrom": "workflow_run_outputs",
+          "tableTo": "workflow_runs",
+          "columnsFrom": [
+            "run_id"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "cascade",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {}
+    },
+    "workflow_runs": {
+      "name": "workflow_runs",
+      "schema": "comfyui_deploy",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": true,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "workflow_version_id": {
+          "name": "workflow_version_id",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": false
+        },
+        "workflow_inputs": {
+          "name": "workflow_inputs",
+          "type": "jsonb",
+          "primaryKey": false,
+          "notNull": false
+        },
+        "workflow_id": {
+          "name": "workflow_id",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "machine_id": {
+          "name": "machine_id",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": false
+        },
+        "status": {
+          "name": "status",
+          "type": "workflow_run_status",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "'not-started'"
+        },
+        "ended_at": {
+          "name": "ended_at",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": false
+        },
+        "created_at": {
+          "name": "created_at",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "now()"
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "workflow_runs_workflow_version_id_workflow_versions_id_fk": {
+          "name": "workflow_runs_workflow_version_id_workflow_versions_id_fk",
+          "tableFrom": "workflow_runs",
+          "tableTo": "workflow_versions",
+          "columnsFrom": [
+            "workflow_version_id"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "set null",
+          "onUpdate": "no action"
+        },
+        "workflow_runs_workflow_id_workflows_id_fk": {
+          "name": "workflow_runs_workflow_id_workflows_id_fk",
+          "tableFrom": "workflow_runs",
+          "tableTo": "workflows",
+          "columnsFrom": [
+            "workflow_id"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "cascade",
+          "onUpdate": "no action"
+        },
+        "workflow_runs_machine_id_machines_id_fk": {
+          "name": "workflow_runs_machine_id_machines_id_fk",
+          "tableFrom": "workflow_runs",
+          "tableTo": "machines",
+          "columnsFrom": [
+            "machine_id"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "set null",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {}
+    },
+    "workflows": {
+      "name": "workflows",
+      "schema": "comfyui_deploy",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": true,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "user_id": {
+          "name": "user_id",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "name": {
+          "name": "name",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "created_at": {
+          "name": "created_at",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "now()"
+        },
+        "updated_at": {
+          "name": "updated_at",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "now()"
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "workflows_user_id_users_id_fk": {
+          "name": "workflows_user_id_users_id_fk",
+          "tableFrom": "workflows",
+          "tableTo": "users",
+          "columnsFrom": [
+            "user_id"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "cascade",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {}
+    },
+    "workflow_versions": {
+      "name": "workflow_versions",
+      "schema": "comfyui_deploy",
+      "columns": {
+        "workflow_id": {
+          "name": "workflow_id",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": true,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "workflow": {
+          "name": "workflow",
+          "type": "jsonb",
+          "primaryKey": false,
+          "notNull": false
+        },
+        "workflow_api": {
+          "name": "workflow_api",
+          "type": "jsonb",
+          "primaryKey": false,
+          "notNull": false
+        },
+        "version": {
+          "name": "version",
+          "type": "integer",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "created_at": {
+          "name": "created_at",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "now()"
+        },
+        "updated_at": {
+          "name": "updated_at",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "now()"
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "workflow_versions_workflow_id_workflows_id_fk": {
+          "name": "workflow_versions_workflow_id_workflows_id_fk",
+          "tableFrom": "workflow_versions",
+          "tableTo": "workflows",
+          "columnsFrom": [
+            "workflow_id"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "cascade",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {}
+    }
+  },
+  "enums": {
+    "deployment_environment": {
+      "name": "deployment_environment",
+      "values": {
+        "staging": "staging",
+        "production": "production"
+      }
+    },
+    "workflow_run_status": {
+      "name": "workflow_run_status",
+      "values": {
+        "not-started": "not-started",
+        "running": "running",
+        "success": "success",
+        "failed": "failed"
+      }
+    }
+  },
+  "schemas": {
+    "comfyui_deploy": "comfyui_deploy"
+  },
+  "_meta": {
+    "schemas": {},
+    "tables": {},
+    "columns": {}
+  }
+}

web/drizzle/meta/_journal.json

@@ -71,6 +71,13 @@
       "when": 1702632471725,
       "tag": "0009_easy_banshee",
       "breakpoints": true
+    },
+    {
+      "idx": 10,
+      "version": "5",
+      "when": 1702708886828,
+      "tag": "0010_organic_marten_broadcloak",
+      "breakpoints": true
     }
   ]
 }

web/src/app/api/create-run/route.ts

@@ -1,60 +0,0 @@
-import { parseDataSafe } from "../../../lib/parseDataSafe";
-import { createRun } from "../../../server/createRun";
-import { getRunsOutput } from "@/server/getRunsOutput";
-import { NextResponse } from "next/server";
-import { z } from "zod";
-
-const Request = z.object({
-  workflow_version_id: z.string(),
-  machine_id: z.string(),
-});
-
-const Request2 = z.object({
-  run_id: z.string(),
-});
-
-export async function GET(request: Request) {
-  const [data, error] = await parseDataSafe(Request2, request);
-  if (!data || error) return error;
-
-  const run = await getRunsOutput(data.run_id);
-
-  return NextResponse.json(run, {
-    status: 200,
-  });
-}
-
-export async function POST(request: Request) {
-  const [data, error] = await parseDataSafe(Request, request);
-  if (!data || error) return error;
-
-  const origin = new URL(request.url).origin;
-
-  const { workflow_version_id, machine_id } = data;
-
-  try {
-    const workflow_run_id = await createRun(
-      origin,
-      workflow_version_id,
-      machine_id
-    );
-
-    return NextResponse.json(
-      {
-        workflow_run_id: workflow_run_id.workflow_run_id,
-      },
-      {
-        status: 200,
-      }
-    );
-  } catch (error: any) {
-    return NextResponse.json(
-      {
-        error: error.message,
-      },
-      {
-        status: 500,
-      }
-    );
-  }
-}

web/src/app/api/run/route.ts

@@ -2,6 +2,7 @@ import { parseDataSafe } from "../../../lib/parseDataSafe";
 import { createRun } from "../../../server/createRun";
 import { db } from "@/db/db";
 import { deploymentsTable } from "@/db/schema";
+import { isKeyRevoked } from "@/server/curdApiKeys";
 import { getRunsData } from "@/server/getRunsOutput";
 import { parseJWT } from "@/server/parseJWT";
 import { replaceCDNUrl } from "@/server/resource";
@@ -18,14 +19,26 @@ const Request2 = z.object({
   run_id: z.string(),
 });
 
-export async function GET(request: Request) {
+async function checkToken(request: Request) {
   const token = request.headers.get("Authorization")?.split(" ")?.[1]; // Assuming token is sent as "Bearer your_token"
   const userData = token ? parseJWT(token) : undefined;
-  if (!userData) {
+  if (!userData || token === undefined) {
     return new NextResponse("Invalid or expired token", {
       status: 401,
     });
+  } else {
+    const revokedKey = await isKeyRevoked(token);
+    if (revokedKey)
+      return new NextResponse("Revoked token", {
+        status: 401,
+      });
   }
+}
+
+export async function GET(request: Request) {
+  const invalidRequest = await checkToken(request)
+  if (invalidRequest) return invalidRequest;
+
 
   const [data, error] = await parseDataSafe(Request2, request);
   if (!data || error) return error;
@@ -41,7 +54,7 @@ export async function GET(request: Request) {
       for (let j = 0; j < output.data?.images.length; j++) {
         const element = output.data?.images[j];
         element.url = replaceCDNUrl(
-          `${process.env.SPACES_ENDPOINT}/${process.env.SPACES_BUCKET}/outputs/runs/${run.id}/${element.filename}`
+          `${process.env.SPACES_ENDPOINT}/${process.env.SPACES_BUCKET}/outputs/runs/${run.id}/${element.filename}`,
         );
       }
     }
@@ -53,13 +66,8 @@ export async function GET(request: Request) {
 }
 
 export async function POST(request: Request) {
-  const token = request.headers.get("Authorization")?.split(" ")?.[1]; // Assuming token is sent as "Bearer your_token"
+  const invalidRequest = await checkToken(request)
-  const userData = token ? parseJWT(token) : undefined;
+  if (invalidRequest) return invalidRequest;
-  if (!userData) {
-    return new NextResponse("Invalid or expired token", {
-      status: 401,
-    });
-  }
 
   const [data, error] = await parseDataSafe(Request, request);
   if (!data || error) return error;
@@ -79,7 +87,7 @@ export async function POST(request: Request) {
       origin,
       deploymentData.workflow_version_id,
       deploymentData.machine_id,
-      inputs
+      inputs,
     );
 
     return NextResponse.json(
@@ -88,7 +96,7 @@ export async function POST(request: Request) {
       },
       {
         status: 200,
-      }
+      },
     );
   } catch (error: any) {
     return NextResponse.json(
@@ -97,7 +105,7 @@ export async function POST(request: Request) {
       },
       {
         status: 500,
-      }
+      },
     );
   }
 }

web/src/app/layout.tsx

@@ -1,5 +1,8 @@
 import "./globals.css";
 import { NavbarRight } from "@/components/NavbarRight";
+import { Button } from "@/components/ui/button";
+import { ClerkProvider, UserButton } from "@clerk/nextjs";
+import { Github } from "lucide-react";
 import type { Metadata } from "next";
 import { Inter } from "next/font/google";
 import { Toaster } from "sonner";
@@ -18,21 +21,37 @@ export default function RootLayout({
 }) {
   return (
     <html lang="en">
-      <body className={inter.className}>
+      <ClerkProvider>
-        <main className="flex min-h-screen flex-col items-center justify-start">
+        <body className={inter.className}>
-          <div className="w-full h-18 flex items-center gap-4 p-4 border-b border-gray-200">
+          <main className="flex min-h-screen flex-col items-center justify-start">
-            <a className="font-bold text-lg hover:underline" href="/">
+            <div className="w-full h-18 flex items-center justify-between gap-4 p-4 border-b border-gray-200">
-              Comfy Deploy
+              <div className="flex flex-row items-center gap-4">
-            </a>
+                <a className="font-bold text-lg hover:underline" href="/">
-            <NavbarRight />
+                  ComfyUI Deploy
-            {/* <div></div> */}
+                </a>
-          </div>
+                <NavbarRight />
-          <div className="md:px-10 px-6 w-full flex items-start">
+              </div>
-            {children}
+              <div className="flex flex-row items-center gap-2">
-          </div>
+                <UserButton />
-          <Toaster richColors />
+                <Button
-        </main>
+                  asChild
-      </body>
+                  variant={"outline"}
+                  className="rounded-full aspect-square p-2"
+                >
+                  <a target="_blank" href="https://github.com/BennyKok/comfyui-deploy">
+                    <Github />
+                  </a>
+                </Button>
+              </div>
+              {/* <div></div> */}
+            </div>
+            <div className="md:px-10 px-6 w-full flex items-start">
+              {children}
+            </div>
+            <Toaster richColors />
+          </main>
+        </body>
+      </ClerkProvider>
     </html>
   );
 }

web/src/components/LiveStatus.tsx

@@ -4,6 +4,8 @@ import { useStore } from "@/components/MachinesWS";
 import { StatusBadge } from "@/components/StatusBadge";
 import { TableCell } from "@/components/ui/table";
 import { type findAllRuns } from "@/server/findAllRuns";
+import { useRouter } from "next/navigation";
+import { useEffect, useState } from "react";
 
 export function LiveStatus({
   run,
@@ -14,7 +16,7 @@ export function LiveStatus({
     (state) =>
       state.data
         .filter((x) => x.id === run.id)
-        .sort((a, b) => b.timestamp - a.timestamp)?.[0]
+        .sort((a, b) => b.timestamp - a.timestamp)?.[0],
   );
 
   let status = run.status;
@@ -26,6 +28,14 @@ export function LiveStatus({
     status = "running";
   }
 
+  const router = useRouter();
+
+  useEffect(() => {
+    if (data?.json.event === "outputs_uploaded") {
+      router.refresh()
+    }
+  }, [data?.json.event]);
+
   return (
     <>
       <TableCell>

web/src/components/MachinesWS.tsx

@@ -55,7 +55,8 @@ function MachineWS({
   const { lastMessage, readyState } = useWebSocket(
     `${wsEndpoint}/comfyui-deploy/ws`,
     {
-      reconnectAttempts: 10,
+      shouldReconnect: ()=> true,
+      reconnectAttempts: 20,
       reconnectInterval: 1000,
     }
   );

web/src/components/RunsTable.tsx

@@ -38,7 +38,7 @@ export async function RunsTable(props: { workflow_id: string }) {
 export async function DeploymentsTable(props: { workflow_id: string }) {
   const allRuns = await findAllDeployments(props.workflow_id);
   return (
-    <div className="overflow-auto h-[400px] w-full">
+    <div className="overflow-auto h-fit lg:h-[400px] w-full">
       <Table className="">
         <TableCaption>A list of your deployments</TableCaption>
         <TableHeader className="bg-background top-0 sticky">

web/src/components/callServerPromise.tsx

@@ -7,6 +7,8 @@ export async function callServerPromise<T>(result: Promise<T>) {
     .then((x) => {
       if ((x as { message: string })?.message !== undefined) {
         toast.success((x as { message: string }).message);
+      } else if ((x as { error: string })?.error !== undefined) {
+        toast.error((x as { error: string }).error);
       }
       return x;
     })

web/src/db/schema.ts

@@ -7,6 +7,7 @@ import {
   timestamp,
   jsonb,
   pgEnum,
+  boolean,
 } from "drizzle-orm/pg-core";
 import { z } from "zod";
 
@@ -52,7 +53,7 @@ export const workflowType = z.any();
 export const workflowAPIType = z.record(
   z.object({
     inputs: z.record(z.any()),
-    class_type: z.string(),
+    class_type: z.string().optional(),
   })
 );
 
@@ -213,6 +214,7 @@ export const apiKeyTable = dbSchema.table("api_keys", {
     })
     .notNull(),
   org_id: text("org_id"),
+  revoked: boolean("revoked").default(false).notNull(),
   created_at: timestamp("created_at").defaultNow().notNull(),
   updated_at: timestamp("updated_at").defaultNow().notNull(),
 });

web/src/server/createRun.ts

@@ -6,102 +6,101 @@ import { ComfyAPI_Run } from "@/types/ComfyAPI_Run";
 import { eq } from "drizzle-orm";
 import { revalidatePath } from "next/cache";
 import "server-only";
+import { withServerPromise } from "./withServerPromise";
 
-export async function createRun(
+export const createRun = withServerPromise(
-  origin: string,
+  async (
-  workflow_version_id: string,
+    origin: string,
-  machine_id: string,
+    workflow_version_id: string,
-  inputs?: Record<string, string>
+    machine_id: string,
-) {
+    inputs?: Record<string, string>,
-  const machine = await db.query.machinesTable.findFirst({
+  ) => {
-    where: eq(workflowRunsTable.id, machine_id),
+    const machine = await db.query.machinesTable.findFirst({
-  });
+      where: eq(workflowRunsTable.id, machine_id),
-
-  if (!machine) {
-    throw new Error("Machine not found");
-    // return new Response("Machine not found", {
-    //   status: 404,
-    // });
-  }
-
-  const workflow_version_data =
-    // workflow_version_id
-    //   ?
-    await db.query.workflowVersionTable.findFirst({
-      where: eq(workflowRunsTable.id, workflow_version_id),
     });
-  // : workflow_version != undefined
-  // ? await db.query.workflowVersionTable.findFirst({
-  //     where: and(
-  //       eq(workflowVersionTable.version, workflow_version),
-  //       eq(workflowVersionTable.workflow_id)
-  //     ),
-  //   })
-  // : null;
-  if (!workflow_version_data) {
-    throw new Error("Workflow version not found");
-    // return new Response("Workflow version not found", {
-    //   status: 404,
-    // });
-  }
 
-  const comfyui_endpoint = `${machine.endpoint}/comfyui-deploy/run`;
+    if (!machine) {
-
+      throw new Error("Machine not found");
-  const workflow_api = workflow_version_data.workflow_api;
+      // return new Response("Machine not found", {
-
+      //   status: 404,
-  // Replace the inputs
+      // });
-  if (inputs && workflow_api) {
-    for (const key in inputs) {
-      Object.entries(workflow_api).forEach(([_, node]) => {
-        if (node.inputs["input_id"] === key) {
-          node.inputs["input_id"] = inputs[key];
-        }
-      });
     }
-  }
 
-  const body = {
+    const workflow_version_data =
-    workflow_api: workflow_api,
+      // workflow_version_id
-    status_endpoint: `${origin}/api/update-run`,
+      //   ?
-    file_upload_endpoint: `${origin}/api/file-upload`,
+      await db.query.workflowVersionTable.findFirst({
-  };
+        where: eq(workflowRunsTable.id, workflow_version_id),
-  // console.log(body);
+      });
-  const bodyJson = JSON.stringify(body);
+    // : workflow_version != undefined
-  // console.log(bodyJson);
+    // ? await db.query.workflowVersionTable.findFirst({
+    //     where: and(
+    //       eq(workflowVersionTable.version, workflow_version),
+    //       eq(workflowVersionTable.workflow_id)
+    //     ),
+    //   })
+    // : null;
+    if (!workflow_version_data) {
+      throw new Error("Workflow version not found");
+      // return new Response("Workflow version not found", {
+      //   status: 404,
+      // });
+    }
 
-  // Sending to comfyui
+    const comfyui_endpoint = `${machine.endpoint}/comfyui-deploy/run`;
-  const _result = await fetch(comfyui_endpoint, {
-    method: "POST",
-    body: bodyJson,
-    cache: "no-store",
-  });
 
-  if (!_result.ok) {
+    const workflow_api = workflow_version_data.workflow_api;
-    throw new Error(`Error creating run, ${_result.statusText}`);
-  }
 
-  console.log(_result);
+    // Replace the inputs
+    if (inputs && workflow_api) {
+      for (const key in inputs) {
+        Object.entries(workflow_api).forEach(([_, node]) => {
+          if (node.inputs["input_id"] === key) {
+            node.inputs["input_id"] = inputs[key];
+          }
+        });
+      }
+    }
 
-  const result = await ComfyAPI_Run.parseAsync(await _result.json());
+    const body = {
+      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);
 
-  console.log(result);
+    // Sending to comfyui
+    const _result = await fetch(comfyui_endpoint, {
+      method: "POST",
+      body: bodyJson,
+      cache: "no-store",
+    });
 
-  // Add to our db
+    if (!_result.ok) {
-  const workflow_run = await db
+      throw new Error(`Error creating run, ${_result.statusText}`);
-    .insert(workflowRunsTable)
+    }
-    .values({
-      id: result.prompt_id,
-      workflow_id: workflow_version_data.workflow_id,
-      workflow_version_id: workflow_version_data.id,
-      workflow_inputs: inputs,
-      machine_id,
-    })
-    .returning();
 
-  revalidatePath(`/${workflow_version_data.workflow_id}`);
+    const result = await ComfyAPI_Run.parseAsync(await _result.json());
 
-  return {
+    // Add to our db
-    workflow_run_id: workflow_run[0].id,
+    const workflow_run = await db
-    message: "Successfully workflow run",
+      .insert(workflowRunsTable)
-  };
+      .values({
-}
+        id: result.prompt_id,
+        workflow_id: workflow_version_data.workflow_id,
+        workflow_version_id: workflow_version_data.id,
+        workflow_inputs: inputs,
+        machine_id,
+      })
+      .returning();
+
+    revalidatePath(`/${workflow_version_data.workflow_id}`);
+
+    return {
+      workflow_run_id: workflow_run[0].id,
+      message: "Successfully workflow run",
+    };
+  },
+);

web/src/server/curdApiKeys.ts

@@ -29,7 +29,7 @@ export async function addNewAPIKey(name: string) {
   if (orgId) {
     token = jwt.sign(
       { user_id: userId, org_id: orgId },
-      process.env.JWT_SECRET!
+      process.env.JWT_SECRET!,
     );
   } else {
     token = jwt.sign({ user_id: userId }, process.env.JWT_SECRET!);
@@ -57,12 +57,20 @@ export async function deleteAPIKey(id: string) {
 
   if (orgId) {
     await db
-      .delete(apiKeyTable)
+      .update(apiKeyTable)
+      .set({
+        revoked: true,
+        updated_at: new Date(),
+      })
       .where(and(eq(apiKeyTable.id, id), eq(apiKeyTable.org_id, orgId)))
       .execute();
   } else {
     await db
-      .delete(apiKeyTable)
+      .update(apiKeyTable)
+      .set({
+        revoked: true,
+        updated_at: new Date(),
+      })
       .where(and(eq(apiKeyTable.id, id), eq(apiKeyTable.user_id, userId)))
       .execute();
   }
@@ -77,13 +85,21 @@ export async function getAPIKeys() {
 
   if (orgId) {
     return await db.query.apiKeyTable.findMany({
-      where: eq(apiKeyTable.org_id, orgId),
+      where: and(eq(apiKeyTable.org_id, orgId), eq(apiKeyTable.revoked, false)),
       orderBy: desc(apiKeyTable.created_at),
     });
   } else {
     return await db.query.apiKeyTable.findMany({
-      where: eq(apiKeyTable.user_id, userId),
+      where: and(eq(apiKeyTable.user_id, userId), eq(apiKeyTable.revoked, false)),
       orderBy: desc(apiKeyTable.created_at),
     });
   }
 }
+
+export async function isKeyRevoked(key: string) {
+  const revokedKey = await db.query.apiKeyTable.findFirst({
+    where: and(eq(apiKeyTable.key, key), eq(apiKeyTable.revoked, true)),
+  });
+
+  return revokedKey !== undefined;
+}

web/src/server/withServerPromise.ts

@@ -0,0 +1,12 @@
+export async function wrapServerPromise<T>(result: Promise<T>) {
+  return result.catch((error) => {
+    return {
+      error: error.message,
+    };
+  });
+}
+export function withServerPromise<T extends (...args: any[]) => Promise<any>>(
+  fn: T
+): (...args: Parameters<T>) => Promise<ReturnType<T> | { error: string; }> {
+  return (...args: Parameters<T>) => wrapServerPromise(fn(...args));
+}