fix: make sure invocation_error arriving before execution-state initialization is processed correctly

Author: JPPhoto

invokeai/frontend/web/src/services/events/nodeExecutionState.ts

@@ -68,3 +68,20 @@ export const getUpdatedNodeExecutionStateOnInvocationComplete = (
 
   return _nodeExecutionState;
 };
+
+export const getUpdatedNodeExecutionStateOnInvocationError = (
+  nodeExecutionState: NodeExecutionState | undefined,
+  data: S['InvocationErrorEvent']
+) => {
+  const _nodeExecutionState = deepClone(nodeExecutionState ?? getInitialNodeExecutionState(data.invocation_source_id));
+  _nodeExecutionState.status = zNodeStatus.enum.FAILED;
+  _nodeExecutionState.progress = null;
+  _nodeExecutionState.progressImage = null;
+  _nodeExecutionState.error = {
+    error_type: data.error_type,
+    error_message: data.error_message,
+    error_traceback: data.error_traceback,
+  };
+
+  return _nodeExecutionState;
+};

invokeai/frontend/web/src/services/events/nodeExecutionStateHelpers.test.ts

@@ -5,6 +5,7 @@ import { describe, expect, it } from 'vitest';
 
 import {
   getUpdatedNodeExecutionStateOnInvocationComplete,
+  getUpdatedNodeExecutionStateOnInvocationError,
   getUpdatedNodeExecutionStateOnInvocationProgress,
   getUpdatedNodeExecutionStateOnInvocationStarted,
 } from './nodeExecutionState';
@@ -87,6 +88,26 @@ const buildInvocationCompleteEvent = (
     ...overrides,
   }) as S['InvocationCompleteEvent'];
 
+const buildInvocationErrorEvent = (overrides: Partial<S['InvocationErrorEvent']> = {}): S['InvocationErrorEvent'] =>
+  ({
+    queue_id: 'default',
+    item_id: 1,
+    batch_id: 'batch-1',
+    origin: 'workflows',
+    destination: 'gallery',
+    user_id: 'user-1',
+    session_id: 'session-1',
+    invocation_source_id: 'node-1',
+    invocation: {
+      id: 'prepared-node-1',
+      type: 'add',
+    },
+    error_type: 'TestError',
+    error_message: 'boom',
+    error_traceback: 'traceback',
+    ...overrides,
+  }) as S['InvocationErrorEvent'];
+
 describe(getUpdatedNodeExecutionStateOnInvocationStarted.name, () => {
   it('creates an execution state when started arrives before initialization', () => {
     const event = buildInvocationStartedEvent();
@@ -219,3 +240,41 @@ describe(getUpdatedNodeExecutionStateOnInvocationComplete.name, () => {
     expect(secondUpdate?.outputs).toEqual([firstEvent.result, secondEvent.result]);
   });
 });
+
+describe(getUpdatedNodeExecutionStateOnInvocationError.name, () => {
+  it('creates an execution state when error arrives before initialization', () => {
+    const event = buildInvocationErrorEvent();
+    const updated = getUpdatedNodeExecutionStateOnInvocationError(undefined, event);
+
+    expect(updated?.nodeId).toBe(event.invocation_source_id);
+    expect(updated?.status).toBe(zNodeStatus.enum.FAILED);
+    expect(updated?.progress).toBeNull();
+    expect(updated?.progressImage).toBeNull();
+    expect(updated?.error).toEqual({
+      error_type: event.error_type,
+      error_message: event.error_message,
+      error_traceback: event.error_traceback,
+    });
+  });
+
+  it('marks the node failed and records the error', () => {
+    const event = buildInvocationErrorEvent();
+    const updated = getUpdatedNodeExecutionStateOnInvocationError(
+      buildNodeExecutionState({
+        status: zNodeStatus.enum.IN_PROGRESS,
+        progress: 0.5,
+        progressImage: { dataURL: 'data:image/png;base64,abc', width: 64, height: 64 },
+      }),
+      event
+    );
+
+    expect(updated?.status).toBe(zNodeStatus.enum.FAILED);
+    expect(updated?.progress).toBeNull();
+    expect(updated?.progressImage).toBeNull();
+    expect(updated?.error).toEqual({
+      error_type: event.error_type,
+      error_message: event.error_message,
+      error_traceback: event.error_traceback,
+    });
+  });
+});

invokeai/frontend/web/src/services/events/setEventListeners.tsx

@@ -38,6 +38,7 @@ import { imagesApi } from 'services/api/endpoints/images';
 import { modelsApi } from 'services/api/endpoints/models';
 import { queueApi } from 'services/api/endpoints/queue';
 import {
+  getUpdatedNodeExecutionStateOnInvocationError,
   getUpdatedNodeExecutionStateOnInvocationProgress,
   getUpdatedNodeExecutionStateOnInvocationStarted,
 } from 'services/events/nodeExecutionState';
@@ -161,19 +162,12 @@ export const setEventListeners = ({ socket, store, setIsConnected }: SetEventLis
       log.trace({ data } as JsonObject, `Received event for already-finished queue item ${data.item_id}`);
       return;
     }
-    const { invocation_source_id, invocation, error_type, error_message, error_traceback } = data;
+    const { invocation_source_id, invocation } = data;
     log.error({ data } as JsonObject, `Invocation error (${invocation.type}, ${invocation_source_id})`);
-    const nes = deepClone($nodeExecutionStates.get()[invocation_source_id]);
-    if (nes) {
-      nes.status = zNodeStatus.enum.FAILED;
-      nes.progress = null;
-      nes.progressImage = null;
-      nes.error = {
-        error_type,
-        error_message,
-        error_traceback,
-      };
-      upsertExecutionState(nes.nodeId, nes);
+    const nes = $nodeExecutionStates.get()[invocation_source_id];
+    const updatedNodeExecutionState = getUpdatedNodeExecutionStateOnInvocationError(nes, data);
+    if (updatedNodeExecutionState) {
+      upsertExecutionState(updatedNodeExecutionState.nodeId, updatedNodeExecutionState);
     }
     // Clear canvas workflow integration processing state on error
     if (data.origin === 'canvas_workflow_integration') {