chunzhi/comfyui-deploy
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
web/src/app/api/create-run/route.ts
-60
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,
}
);
}
}
web/src/app/api/run/route.ts
+21 -13
View File
@@ -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,
}
},
);
}
}
web/src/app/layout.tsx
+34 -15
View File
@@ -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}>
<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">
<a className="font-bold text-lg hover:underline" href="/">
Comfy Deploy
</a>
<NavbarRight />
{/* <div></div> */}
</div>
<div className="md:px-10 px-6 w-full flex items-start">
{children}
</div>
<Toaster richColors />
</main>
</body>
<ClerkProvider>
<body className={inter.className}>
<main className="flex min-h-screen flex-col items-center justify-start">
<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="/">
ComfyUI Deploy
</a>
<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 className="md:px-10 px-6 w-full flex items-start">
{children}
</div>
<Toaster richColors />
</main>
</body>
</ClerkProvider>
</html>
);
}
web/src/components/LiveStatus.tsx
+11 -1
View File
@@ -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
+2 -1
View File
@@ -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
+1 -1
View File
@@ -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
+2
View File
@@ -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
+3 -1
View File
@@ -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
+85 -86
View File
@@ -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<string, string>
) {
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<string, string>,
) => {
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",
};
},
);
web/src/server/curdApiKeys.ts
+21 -5
View File
@@ -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
+12
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));
}