From 8f7f06470c1e924807abe5a1df24f219ac850284 Mon Sep 17 00:00:00 2001
From: bennykok <itechbenny@gmail.com>
Date: Sat, 27 Jan 2024 11:38:10 +0800
Subject: [PATCH] feat: ensure open in comfy ui will carry over the auth token

---
 web/src/components/MachineList.tsx       | 23 ++++++++++++++++-------
 web/src/components/VersionSelect.tsx     |  2 +-
 web/src/server/editWorkflowOnMachine.tsx | 20 ++++++++++++++------
 3 files changed, 31 insertions(+), 14 deletions(-)

diff --git a/web/src/components/MachineList.tsx b/web/src/components/MachineList.tsx
index bc8b26a..eec5713 100644
--- a/web/src/components/MachineList.tsx
+++ b/web/src/components/MachineList.tsx
@@ -40,6 +40,7 @@ import {
   updateCustomMachine,
   updateMachine,
 } from "@/server/curdMachine";
+import { editWorkflowOnMachine } from "@/server/editWorkflowOnMachine";
 import type {
   ColumnDef,
   ColumnFiltersState,
@@ -57,6 +58,7 @@ import {
 import { ArrowUpDown, MoreHorizontal } from "lucide-react";
 import * as React from "react";
 import { useState } from "react";
+import { toast } from "sonner";
 import type { z } from "zod";
 
 export type Machine = MachineType;
@@ -218,15 +220,22 @@ export const columns: ColumnDef<Machine>[] = [
             {machine.type === "comfy-deploy-serverless" && (
               <>
                 <DropdownMenuItem asChild>
-                  <a
-                    target="_blank"
-                    href={machine.endpoint.replace(
-                      "comfyui-api",
-                      "comfyui-app"
-                    )} rel="noreferrer"
+                  <button
+                    onClick={async () => {
+                      const id = toast.loading("Getting machine url...")
+                      const url = await callServerPromise(
+                        editWorkflowOnMachine(machine.id),
+                      );
+                      if (url && typeof url !== "object") {
+                        window.open(url, "_blank");
+                      } else if (url && typeof url === "object" && url.error) {
+                        console.error(url.error);
+                      }
+                      toast.dismiss(id)
+                    }}
                   >
                     Open ComfyUI
-                  </a>
+                  </button>
                 </DropdownMenuItem>
                 <DropdownMenuItem
                   onClick={() => {
diff --git a/web/src/components/VersionSelect.tsx b/web/src/components/VersionSelect.tsx
index 1ac739b..aea9e7c 100644
--- a/web/src/components/VersionSelect.tsx
+++ b/web/src/components/VersionSelect.tsx
@@ -402,7 +402,7 @@ export function OpenEditButton({
         onClick={async () => {
           setIsLoading(true);
           const url = await callServerPromise(
-            editWorkflowOnMachine(workflow_version_id, machine),
+            editWorkflowOnMachine(machine, workflow_version_id),
           );
           if (url && typeof url !== "object") {
             window.open(url, "_blank");
diff --git a/web/src/server/editWorkflowOnMachine.tsx b/web/src/server/editWorkflowOnMachine.tsx
index 0f7c581..6387293 100644
--- a/web/src/server/editWorkflowOnMachine.tsx
+++ b/web/src/server/editWorkflowOnMachine.tsx
@@ -9,7 +9,7 @@ import "server-only";
 import { getUrlServerSide } from "./getUrlServerSide";
 
 export const editWorkflowOnMachine = withServerPromise(
-  async (workflow_version_id: string, machine_id: string) => {
+  async (machine_id: string, workflow_version_id?: string) => {
     const { userId, orgId } = auth();
 
     const domain = getUrlServerSide();
@@ -37,10 +37,18 @@ export const editWorkflowOnMachine = withServerPromise(
       endpoint = machine.endpoint.replace("comfyui-api", "comfyui-app");
     }
 
-    return `${endpoint}?workflow_version_id=${encodeURIComponent(
-      workflow_version_id,
-    )}&auth_token=${encodeURIComponent(token)}&org_display=${encodeURIComponent(
-      userName,
-    )}&origin=${encodeURIComponent(domain)}`;
+    const params = {
+      workflow_version_id: workflow_version_id,
+      auth_token: token,
+      org_display: userName,
+      origin: domain,
+    };
+
+    const queryString = Object.entries(params)
+      .filter(([key, value]) => value !== undefined)
+      .map(([key, value]) => `${encodeURIComponent(key)}=${encodeURIComponent(value as string)}`)
+      .join('&');
+
+    return `${endpoint}?${queryString}`;
   },
 );