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

This commit is contained in:
BennyKok 2023-12-16 14:55:30 +08:00
parent 0cf6d97f2f
commit 21a17cb753
16 changed files with 839 additions and 188 deletions

View File

@ -14,7 +14,7 @@ import io
import time import time
import execution import execution
import random import random
import traceback
import uuid import uuid
import asyncio import asyncio
import atexit import atexit
@ -90,7 +90,14 @@ async def comfy_deploy_run(request):
"client_id": "fake_client" #api.client_id "client_id": "fake_client" #api.client_id
} }
try:
res = post_prompt(prompt) 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']] = { prompt_metadata[res['prompt_id']] = {
'status_endpoint': data.get('status_endpoint'), '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) 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_original = prompt_server.send_json
prompt_server.send_json = send_json_override.__get__(prompt_server, server.PromptServer) prompt_server.send_json = send_json_override.__get__(prompt_server, server.PromptServer)

View File

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

View File

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

View File

@ -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": {}
}
}

View File

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

View File

@ -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,
}
);
}
}

View File

@ -2,6 +2,7 @@ import { parseDataSafe } from "../../../lib/parseDataSafe";
import { createRun } from "../../../server/createRun"; import { createRun } from "../../../server/createRun";
import { db } from "@/db/db"; import { db } from "@/db/db";
import { deploymentsTable } from "@/db/schema"; import { deploymentsTable } from "@/db/schema";
import { isKeyRevoked } from "@/server/curdApiKeys";
import { getRunsData } from "@/server/getRunsOutput"; import { getRunsData } from "@/server/getRunsOutput";
import { parseJWT } from "@/server/parseJWT"; import { parseJWT } from "@/server/parseJWT";
import { replaceCDNUrl } from "@/server/resource"; import { replaceCDNUrl } from "@/server/resource";
@ -18,14 +19,26 @@ const Request2 = z.object({
run_id: z.string(), 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 token = request.headers.get("Authorization")?.split(" ")?.[1]; // Assuming token is sent as "Bearer your_token"
const userData = token ? parseJWT(token) : undefined; const userData = token ? parseJWT(token) : undefined;
if (!userData) { if (!userData || token === undefined) {
return new NextResponse("Invalid or expired token", { return new NextResponse("Invalid or expired token", {
status: 401, 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); const [data, error] = await parseDataSafe(Request2, request);
if (!data || error) return error; 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++) { for (let j = 0; j < output.data?.images.length; j++) {
const element = output.data?.images[j]; const element = output.data?.images[j];
element.url = replaceCDNUrl( 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) { 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); const [data, error] = await parseDataSafe(Request, request);
if (!data || error) return error; if (!data || error) return error;
@ -79,7 +87,7 @@ export async function POST(request: Request) {
origin, origin,
deploymentData.workflow_version_id, deploymentData.workflow_version_id,
deploymentData.machine_id, deploymentData.machine_id,
inputs inputs,
); );
return NextResponse.json( return NextResponse.json(
@ -88,7 +96,7 @@ export async function POST(request: Request) {
}, },
{ {
status: 200, status: 200,
} },
); );
} catch (error: any) { } catch (error: any) {
return NextResponse.json( return NextResponse.json(
@ -97,7 +105,7 @@ export async function POST(request: Request) {
}, },
{ {
status: 500, status: 500,
} },
); );
} }
} }

View File

@ -1,5 +1,8 @@
import "./globals.css"; import "./globals.css";
import { NavbarRight } from "@/components/NavbarRight"; 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 type { Metadata } from "next";
import { Inter } from "next/font/google"; import { Inter } from "next/font/google";
import { Toaster } from "sonner"; import { Toaster } from "sonner";
@ -18,13 +21,28 @@ export default function RootLayout({
}) { }) {
return ( return (
<html lang="en"> <html lang="en">
<ClerkProvider>
<body className={inter.className}> <body className={inter.className}>
<main className="flex min-h-screen flex-col items-center justify-start"> <main className="flex min-h-screen flex-col items-center justify-start">
<div className="w-full h-18 flex items-center gap-4 p-4 border-b border-gray-200"> <div className="w-full h-18 flex items-center justify-between gap-4 p-4 border-b border-gray-200">
<div className="flex flex-row items-center gap-4">
<a className="font-bold text-lg hover:underline" href="/"> <a className="font-bold text-lg hover:underline" href="/">
Comfy Deploy ComfyUI Deploy
</a> </a>
<NavbarRight /> <NavbarRight />
</div>
<div className="flex flex-row items-center gap-2">
<UserButton />
<Button
asChild
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> */}
</div> </div>
<div className="md:px-10 px-6 w-full flex items-start"> <div className="md:px-10 px-6 w-full flex items-start">
@ -33,6 +51,7 @@ export default function RootLayout({
<Toaster richColors /> <Toaster richColors />
</main> </main>
</body> </body>
</ClerkProvider>
</html> </html>
); );
} }

View File

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

View File

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

View File

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

View File

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

View File

@ -7,6 +7,7 @@ import {
timestamp, timestamp,
jsonb, jsonb,
pgEnum, pgEnum,
boolean,
} from "drizzle-orm/pg-core"; } from "drizzle-orm/pg-core";
import { z } from "zod"; import { z } from "zod";
@ -52,7 +53,7 @@ export const workflowType = z.any();
export const workflowAPIType = z.record( export const workflowAPIType = z.record(
z.object({ z.object({
inputs: z.record(z.any()), 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(), .notNull(),
org_id: text("org_id"), org_id: text("org_id"),
revoked: boolean("revoked").default(false).notNull(),
created_at: timestamp("created_at").defaultNow().notNull(), created_at: timestamp("created_at").defaultNow().notNull(),
updated_at: timestamp("updated_at").defaultNow().notNull(), updated_at: timestamp("updated_at").defaultNow().notNull(),
}); });

View File

@ -6,13 +6,15 @@ import { ComfyAPI_Run } from "@/types/ComfyAPI_Run";
import { eq } from "drizzle-orm"; import { eq } from "drizzle-orm";
import { revalidatePath } from "next/cache"; import { revalidatePath } from "next/cache";
import "server-only"; import "server-only";
import { withServerPromise } from "./withServerPromise";
export async function createRun( export const createRun = withServerPromise(
async (
origin: string, origin: string,
workflow_version_id: string, workflow_version_id: string,
machine_id: string, machine_id: string,
inputs?: Record<string, string> inputs?: Record<string, string>,
) { ) => {
const machine = await db.query.machinesTable.findFirst({ const machine = await db.query.machinesTable.findFirst({
where: eq(workflowRunsTable.id, machine_id), where: eq(workflowRunsTable.id, machine_id),
}); });
@ -80,12 +82,8 @@ export async function createRun(
throw new Error(`Error creating run, ${_result.statusText}`); throw new Error(`Error creating run, ${_result.statusText}`);
} }
console.log(_result);
const result = await ComfyAPI_Run.parseAsync(await _result.json()); const result = await ComfyAPI_Run.parseAsync(await _result.json());
console.log(result);
// Add to our db // Add to our db
const workflow_run = await db const workflow_run = await db
.insert(workflowRunsTable) .insert(workflowRunsTable)
@ -104,4 +102,5 @@ export async function createRun(
workflow_run_id: workflow_run[0].id, workflow_run_id: workflow_run[0].id,
message: "Successfully workflow run", message: "Successfully workflow run",
}; };
} },
);

View File

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

View File

@ -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));
}