From 9bbf4301252d570efcd2b2c5d931078afbccadb2 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Tue, 16 Jul 2024 17:14:54 +1000
Subject: [PATCH] fix(ui): restore nodes output tracking

---
 .../socketio/socketInvocationComplete.ts      | 23 +++++++++++--------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketInvocationComplete.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketInvocationComplete.ts
index 0ba7fa1668..279f868d9f 100644
--- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketInvocationComplete.ts
+++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketInvocationComplete.ts
@@ -28,6 +28,19 @@ export const addInvocationCompleteEventListener = (startAppListening: AppStartLi
       );
 
       const { result, invocation_source_id } = data;
+
+      if (data.origin === 'workflows') {
+        const nes = deepClone($nodeExecutionStates.get()[invocation_source_id]);
+        if (nes) {
+          nes.status = zNodeStatus.enum.COMPLETED;
+          if (nes.progress !== null) {
+            nes.progress = 1;
+          }
+          nes.outputs.push(result);
+          upsertExecutionState(nes.nodeId, nes);
+        }
+      }
+
       // This complete event has an associated image output
       if (data.result.type === 'image_output' && !nodeTypeDenylist.includes(data.invocation.type)) {
         const { image_name } = data.result.image;
@@ -50,16 +63,6 @@ export const addInvocationCompleteEventListener = (startAppListening: AppStartLi
           } else if (!canvasV2.session.isActive) {
             $lastProgressEvent.set(null);
           }
-        } else if (data.origin === 'workflows') {
-          const nes = deepClone($nodeExecutionStates.get()[invocation_source_id]);
-          if (nes) {
-            nes.status = zNodeStatus.enum.COMPLETED;
-            if (nes.progress !== null) {
-              nes.progress = 1;
-            }
-            nes.outputs.push(result);
-            upsertExecutionState(nes.nodeId, nes);
-          }
         }
 
         if (!imageDTO.is_intermediate) {