From 824ce16b4e7fd89d8ce309428cff32e08b581633 Mon Sep 17 00:00:00 2001
From: BennyKok <itechbenny@gmail.com>
Date: Fri, 29 Dec 2023 20:06:52 +0800
Subject: [PATCH] fix(plugin): pending upload status not updated

---
 custom_routes.py                   | 33 ++++++++++++++++++++++++------
 web/src/components/LiveStatus.tsx  |  6 +++++-
 web/src/components/StatusBadge.tsx |  1 +
 3 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/custom_routes.py b/custom_routes.py
index e2a2e26..a046248 100644
--- a/custom_routes.py
+++ b/custom_routes.py
@@ -206,15 +206,19 @@ async def send_json_override(self, event, data, sid=None):
         update_run(prompt_id, Status.RUNNING)
 
     # the last executing event is none, then the workflow is finished
-    if event == 'executing' and data.get('node') is None and not have_pending_upload(prompt_id):
-        update_run(prompt_id, Status.SUCCESS)
+    if event == 'executing' and data.get('node') is None:
+        mark_prompt_done(prompt_id=prompt_id)
+        if not have_pending_upload(prompt_id):
+            update_run(prompt_id, Status.SUCCESS)
 
     if event == 'execution_error':
         update_run(prompt_id, Status.FAILED)
         asyncio.create_task(update_run_with_output(prompt_id, data))
+        # await update_run_with_output(prompt_id, data)
 
     if event == 'executed' and 'node' in data and 'output' in data:
         asyncio.create_task(update_run_with_output(prompt_id, data.get('output'), node_id=data.get('node')))
+        # await update_run_with_output(prompt_id, data.get('output'), node_id=data.get('node'))
         # update_run_with_output(prompt_id, data.get('output'))
 
 
@@ -302,7 +306,22 @@ async def upload_file(prompt_id, filename, subfolder=None, content_type="image/p
 
 def have_pending_upload(prompt_id):
     if 'prompt_id' in prompt_metadata and 'uploading_nodes' in prompt_metadata[prompt_id] and len(prompt_metadata[prompt_id]['uploading_nodes']) > 0:
+        print("have pending upload ", len(prompt_metadata[prompt_id]['uploading_nodes']))
         return True
+
+    print("no pending upload")
+    return False
+
+def mark_prompt_done(prompt_id):
+    if prompt_id in prompt_metadata:
+        prompt_metadata[prompt_id]["done"] = True
+        print("Prompt done")
+
+def is_prompt_done(prompt_id):
+    if prompt_id in prompt_metadata and "done" in prompt_metadata[prompt_id]:
+        if prompt_metadata[prompt_id]["done"] == True:
+            return True
+    
     return False
 
 async def update_file_status(prompt_id, data, uploading, have_error=False, node_id=None):
@@ -315,7 +334,7 @@ async def update_file_status(prompt_id, data, uploading, have_error=False, node_
         else:
             prompt_metadata[prompt_id]['uploading_nodes'].discard(node_id)
 
-    # print(prompt_metadata[prompt_id])
+    print(prompt_metadata[prompt_id]['uploading_nodes'])
     # Update the remote status
 
     if have_error:
@@ -326,16 +345,18 @@ async def update_file_status(prompt_id, data, uploading, have_error=False, node_
         return
 
     # if there are still nodes that are uploading, then we set the status to uploading
-    if uploading and have_pending_upload(prompt_id):
+    if uploading:
         if prompt_metadata[prompt_id]['status'] != Status.UPLOADING:
             update_run(prompt_id, Status.UPLOADING)
+            print("Status: UPLOADING")
             await send("uploading", {
                 "prompt_id": prompt_id,
             })
     
     # if there are no nodes that are uploading, then we set the status to success
-    elif not uploading:
+    elif not uploading and not have_pending_upload(prompt_id) and is_prompt_done(prompt_id=prompt_id):
         update_run(prompt_id, Status.SUCCESS)
+        print("Status: SUCCUSS")
         await send("success", {
             "prompt_id": prompt_id,
         })
@@ -352,7 +373,7 @@ async def update_run_with_output(prompt_id, data, node_id=None):
         try:
             have_upload = 'images' in data or 'files' in data
 
-            print("have_upload", have_upload)
+            print("\nhave_upload", have_upload, node_id)
 
             if have_upload:
                 await update_file_status(prompt_id, data, True, node_id=node_id)
diff --git a/web/src/components/LiveStatus.tsx b/web/src/components/LiveStatus.tsx
index 8df690c..deeb5e4 100644
--- a/web/src/components/LiveStatus.tsx
+++ b/web/src/components/LiveStatus.tsx
@@ -26,7 +26,11 @@ export function LiveStatus({
   //   status = "success";
   // } else
   if (data?.json.event == "executing") {
-    status = "running";
+    if (data?.json?.data?.node == undefined) {
+      status = "success";
+    } else {
+      status = "running";
+    }
   } else if (data?.json.event == "uploading") {
     status = "uploading";
   } else if (data?.json.event == "success") {
diff --git a/web/src/components/StatusBadge.tsx b/web/src/components/StatusBadge.tsx
index 5775c93..b136b9d 100644
--- a/web/src/components/StatusBadge.tsx
+++ b/web/src/components/StatusBadge.tsx
@@ -8,6 +8,7 @@ export function StatusBadge({
   status: Awaited<ReturnType<typeof findAllRuns>>[0]["status"];
 }) {
   switch (status) {
+    case "uploading":
     case "running":
       return (
         <Badge variant="secondary" className="gap-2 w-fit">