diff --git a/custom_routes.py b/custom_routes.py index e86c957..d178bb2 100644 --- a/custom_routes.py +++ b/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) \ No newline at end of file diff --git a/web-plugin/index.js b/web-plugin/index.js index 41a568b..47a6f4e 100644 --- a/web-plugin/index.js +++ b/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 { diff --git a/web/drizzle/0010_organic_marten_broadcloak.sql b/web/drizzle/0010_organic_marten_broadcloak.sql new file mode 100644 index 0000000..8322260 --- /dev/null +++ b/web/drizzle/0010_organic_marten_broadcloak.sql @@ -0,0 +1 @@ +ALTER TABLE "comfyui_deploy"."api_keys" ADD COLUMN "revoked" boolean DEFAULT false NOT NULL; \ No newline at end of file diff --git a/web/drizzle/meta/0010_snapshot.json b/web/drizzle/meta/0010_snapshot.json new file mode 100644 index 0000000..4baeac7 --- /dev/null +++ b/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": {} + } +} \ No newline at end of file diff --git a/web/drizzle/meta/_journal.json b/web/drizzle/meta/_journal.json index 97faadb..a7da9f7 100644 --- a/web/drizzle/meta/_journal.json +++ b/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 } ] } \ No newline at end of file diff --git a/web/src/app/api/create-run/route.ts b/web/src/app/api/create-run/route.ts deleted file mode 100644 index 5499bf4..0000000 --- a/web/src/app/api/create-run/route.ts +++ /dev/null @@ -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, - } - ); - } -} diff --git a/web/src/app/api/run/route.ts b/web/src/app/api/run/route.ts index 477001e..d526f71 100644 --- a/web/src/app/api/run/route.ts +++ b/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 userData = token ? parseJWT(token) : undefined; - if (!userData) { - return new NextResponse("Invalid or expired token", { - status: 401, - }); - } + const invalidRequest = await checkToken(request) + if (invalidRequest) return invalidRequest; 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, - } + }, ); } } diff --git a/web/src/app/layout.tsx b/web/src/app/layout.tsx index 0c0184b..bb7cfd5 100644 --- a/web/src/app/layout.tsx +++ b/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 ( - -
-
- - Comfy Deploy - - - {/*
*/} -
-
- {children} -
- -
- + + +
+
+
+ + ComfyUI Deploy + + +
+
+ + +
+ {/*
*/} +
+
+ {children} +
+ +
+ + ); } diff --git a/web/src/components/LiveStatus.tsx b/web/src/components/LiveStatus.tsx index e20bc7e..e3f74ed 100644 --- a/web/src/components/LiveStatus.tsx +++ b/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 ( <> diff --git a/web/src/components/MachinesWS.tsx b/web/src/components/MachinesWS.tsx index e35f875..d151306 100644 --- a/web/src/components/MachinesWS.tsx +++ b/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, } ); diff --git a/web/src/components/RunsTable.tsx b/web/src/components/RunsTable.tsx index fd0effd..041dad4 100644 --- a/web/src/components/RunsTable.tsx +++ b/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 ( -
+
A list of your deployments diff --git a/web/src/components/callServerPromise.tsx b/web/src/components/callServerPromise.tsx index ac5d06a..17c0117 100644 --- a/web/src/components/callServerPromise.tsx +++ b/web/src/components/callServerPromise.tsx @@ -7,6 +7,8 @@ export async function callServerPromise(result: Promise) { .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; }) diff --git a/web/src/db/schema.ts b/web/src/db/schema.ts index e832462..c276b6f 100644 --- a/web/src/db/schema.ts +++ b/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(), }); diff --git a/web/src/server/createRun.ts b/web/src/server/createRun.ts index 538d9e1..5108d03 100644 --- a/web/src/server/createRun.ts +++ b/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( - origin: string, - workflow_version_id: string, - machine_id: string, - inputs?: Record -) { - 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), +export const createRun = withServerPromise( + async ( + origin: string, + workflow_version_id: string, + machine_id: string, + inputs?: Record, + ) => { + const machine = await db.query.machinesTable.findFirst({ + where: eq(workflowRunsTable.id, machine_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`; - - const workflow_api = workflow_version_data.workflow_api; - - // 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]; - } - }); + if (!machine) { + throw new Error("Machine not found"); + // return new Response("Machine not found", { + // status: 404, + // }); } - } - 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); + 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, + // }); + } - // Sending to comfyui - const _result = await fetch(comfyui_endpoint, { - method: "POST", - body: bodyJson, - cache: "no-store", - }); + const comfyui_endpoint = `${machine.endpoint}/comfyui-deploy/run`; - if (!_result.ok) { - throw new Error(`Error creating run, ${_result.statusText}`); - } + const workflow_api = workflow_version_data.workflow_api; - 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 - const workflow_run = await db - .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(); + if (!_result.ok) { + throw new Error(`Error creating run, ${_result.statusText}`); + } - revalidatePath(`/${workflow_version_data.workflow_id}`); + const result = await ComfyAPI_Run.parseAsync(await _result.json()); - return { - workflow_run_id: workflow_run[0].id, - message: "Successfully workflow run", - }; -} + // Add to our db + const workflow_run = await db + .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", + }; + }, +); diff --git a/web/src/server/curdApiKeys.ts b/web/src/server/curdApiKeys.ts index 0dbbf44..5db304c 100644 --- a/web/src/server/curdApiKeys.ts +++ b/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; +} \ No newline at end of file diff --git a/web/src/server/withServerPromise.ts b/web/src/server/withServerPromise.ts new file mode 100644 index 0000000..b2582f9 --- /dev/null +++ b/web/src/server/withServerPromise.ts @@ -0,0 +1,12 @@ +export async function wrapServerPromise(result: Promise) { + return result.catch((error) => { + return { + error: error.message, + }; + }); +} +export function withServerPromise Promise>( + fn: T +): (...args: Parameters) => Promise | { error: string; }> { + return (...args: Parameters) => wrapServerPromise(fn(...args)); +}