comfyui-node 1.4.4 → 1.6.0

package/dist/call-wrapper.js

@@ -1,5 +1,7 @@
 import { FailedCacheError, WentMissingError, EnqueueFailedError, DisconnectedError, CustomEventError, ExecutionFailedError, ExecutionInterruptedError, MissingNodeError } from "./types/error.js";
 import { buildEnqueueFailedError } from "./utils/response-error.js";
+const DISCONNECT_FAILURE_GRACE_MS = 5000;
+const CALL_WRAPPER_DEBUG = process.env.WORKFLOW_POOL_DEBUG === "1";
 /**
  * Represents a wrapper class for making API calls using the ComfyApi client.
  * Provides methods for setting callback functions and executing the job.
@@ -19,6 +21,15 @@ export class CallWrapper {
     onFinishedFn;
     onFailedFn;
     onProgressFn;
+    jobResolveFn;
+    jobDoneResolved = false;
+    pendingCompletion = null;
+    cancellationRequested = false;
+    promptLoadTrigger = null;
+    disconnectRecoveryActive = false;
+    disconnectFailureTimer = null;
+    onReconnectHandlerOffFn;
+    onReconnectFailedHandlerOffFn;
     onDisconnectedHandlerOffFn;
     checkExecutingOffFn;
     checkExecutedOffFn;
@@ -133,19 +144,39 @@ export class CallWrapper {
          * Start the job execution.
          */
         this.emitLog("CallWrapper.run", "enqueue start");
+        this.pendingCompletion = null;
+        this.jobResolveFn = undefined;
+        this.jobDoneResolved = false;
+        this.cancellationRequested = false;
+        this.promptLoadTrigger = null;
         const job = await this.enqueueJob();
         if (!job) {
             // enqueueJob already invoked onFailed with a rich error instance; just abort.
             this.emitLog("CallWrapper.run", "enqueue failed -> abort");
             return false;
         }
-        let promptLoadTrigger;
         const promptLoadCached = new Promise((resolve) => {
-            promptLoadTrigger = resolve;
+            this.promptLoadTrigger = (value) => {
+                if (this.promptLoadTrigger) {
+                    this.promptLoadTrigger = null;
+                }
+                resolve(value);
+            };
         });
-        let jobDoneTrigger;
         const jobDonePromise = new Promise((resolve) => {
-            jobDoneTrigger = resolve;
+            this.jobDoneResolved = false;
+            this.jobResolveFn = (value) => {
+                if (this.jobDoneResolved) {
+                    return;
+                }
+                this.jobDoneResolved = true;
+                resolve(value);
+            };
+            if (this.pendingCompletion !== null) {
+                const pending = this.pendingCompletion;
+                this.pendingCompletion = null;
+                this.jobResolveFn?.(pending);
+            }
         });
         /**
          * Declare the function to check if the job is executing.
@@ -153,7 +184,7 @@ export class CallWrapper {
         const checkExecutingFn = (event) => {
             if (event.detail && event.detail.prompt_id === job.prompt_id) {
                 this.emitLog("CallWrapper.run", "executing observed", { node: event.detail.node });
-                promptLoadTrigger(false);
+                this.resolvePromptLoad(false);
             }
         };
         /**
@@ -171,7 +202,7 @@ export class CallWrapper {
                     nodes: event.detail.nodes,
                     expected: outputNodes
                 });
-                promptLoadTrigger(cached);
+                this.resolvePromptLoad(cached);
             }
         };
         /**
@@ -200,6 +231,15 @@ export class CallWrapper {
                 this.emitLog("CallWrapper.status", "cached output already handled");
                 return;
             }
+            if (this.cancellationRequested) {
+                this.emitLog("CallWrapper.status", "job missing after cancellation", {
+                    prompt_id: job.prompt_id
+                });
+                this.resolvePromptLoad(false);
+                this.resolveJob(false);
+                this.cleanupListeners("status handler cancellation");
+                return;
+            }
             wentMissing = true;
             const output = await this.handleCachedOutput(job.prompt_id);
             if (output) {
@@ -207,22 +247,31 @@ export class CallWrapper {
                 this.emitLog("CallWrapper.status", "output from history after missing", {
                     prompt_id: job.prompt_id
                 });
-                jobDoneTrigger(output);
+                this.resolvePromptLoad(false);
+                this.resolveJob(output);
                 this.cleanupListeners("status handler resolved from history");
                 return;
             }
+            if (this.disconnectRecoveryActive) {
+                this.emitLog("CallWrapper.status", "job missing but disconnect recovery active -> waiting", {
+                    prompt_id: job.prompt_id
+                });
+                this.resolvePromptLoad(false);
+                void this.attemptHistoryCompletion("status_missing");
+                return;
+            }
             cachedOutputDone = true;
             this.emitLog("CallWrapper.status", "job missing without cached output", {
                 prompt_id: job.prompt_id
             });
-            promptLoadTrigger(false);
-            jobDoneTrigger(false);
+            this.resolvePromptLoad(false);
+            this.resolveJob(false);
             this.cleanupListeners("status handler missing");
-            this.onFailedFn?.(new WentMissingError("The job went missing!"), job.prompt_id);
+            this.emitFailure(new WentMissingError("The job went missing!"), job.prompt_id);
         };
         this.statusHandlerOffFn = this.client.on("status", statusHandler);
         // Attach execution listeners immediately so fast jobs cannot finish before we subscribe
-        this.handleJobExecution(job.prompt_id, jobDoneTrigger);
+        this.handleJobExecution(job.prompt_id);
         await promptLoadCached;
         if (wentMissing) {
             return jobDonePromise;
@@ -232,14 +281,14 @@ export class CallWrapper {
         if (output) {
             cachedOutputDone = true;
             this.cleanupListeners("no cached output values returned");
-            jobDoneTrigger(output);
+            this.resolveJob(output);
             return output;
         }
         if (output === false) {
             cachedOutputDone = true;
             this.cleanupListeners("cached output ready before execution listeners");
-            this.onFailedFn?.(new FailedCacheError("Failed to get cached output"), this.promptId);
-            jobDoneTrigger(false);
+            this.emitFailure(new FailedCacheError("Failed to get cached output"), this.promptId);
+            this.resolveJob(false);
             return false;
         }
         this.emitLog("CallWrapper.run", "no cached output -> proceed with execution listeners");
@@ -305,10 +354,10 @@ export class CallWrapper {
             }
             catch (e) {
                 if (e instanceof Response) {
-                    this.onFailedFn?.(new MissingNodeError("Failed to get workflow node definitions", { cause: await e.json() }));
+                    this.emitFailure(new MissingNodeError("Failed to get workflow node definitions", { cause: await e.json() }), this.promptId);
                 }
                 else {
-                    this.onFailedFn?.(new MissingNodeError("There was a missing node in the workflow bypass.", { cause: e }));
+                    this.emitFailure(new MissingNodeError("There was a missing node in the workflow bypass.", { cause: e }), this.promptId);
                 }
                 return null;
             }
@@ -320,22 +369,22 @@ export class CallWrapper {
         catch (e) {
             try {
                 if (e instanceof EnqueueFailedError) {
-                    this.onFailedFn?.(e);
+                    this.emitFailure(e, this.promptId);
                 }
                 else if (e instanceof Response) {
                     const err = await buildEnqueueFailedError(e);
-                    this.onFailedFn?.(err);
+                    this.emitFailure(err, this.promptId);
                 }
                 else if (e && typeof e === "object" && "response" in e && e.response instanceof Response) {
                     const err = await buildEnqueueFailedError(e.response);
-                    this.onFailedFn?.(err);
+                    this.emitFailure(err, this.promptId);
                 }
                 else {
-                    this.onFailedFn?.(new EnqueueFailedError("Failed to queue prompt", { cause: e, reason: e?.message }));
+                    this.emitFailure(new EnqueueFailedError("Failed to queue prompt", { cause: e, reason: e?.message }), this.promptId);
                 }
             }
             catch (inner) {
-                this.onFailedFn?.(new EnqueueFailedError("Failed to queue prompt", { cause: inner }));
+                this.emitFailure(new EnqueueFailedError("Failed to queue prompt", { cause: inner }), this.promptId);
             }
             job = null;
         }
@@ -343,20 +392,162 @@ export class CallWrapper {
             return;
         }
         this.promptId = job.prompt_id;
+        console.log(`[CallWrapper] Enqueued with promptId=${this.promptId?.substring(0, 8)}...`);
+        console.log(`[CallWrapper] Full job object:`, JSON.stringify({ promptId: job.prompt_id }, null, 2));
         this.emitLog("CallWrapper.enqueueJob", "queued", { prompt_id: this.promptId });
         this.onPendingFn?.(this.promptId);
         this.onDisconnectedHandlerOffFn = this.client.on("disconnected", () => {
-            // Ignore disconnection if we are already successfully completing
-            // This prevents a race condition where outputs are collected successfully
-            // but the WebSocket disconnects before cleanupListeners() is called
             if (this.isCompletingSuccessfully) {
                 this.emitLog("CallWrapper.enqueueJob", "disconnected during success completion -> ignored");
                 return;
             }
-            this.onFailedFn?.(new DisconnectedError("Disconnected"), this.promptId);
+            this.emitLog("CallWrapper.enqueueJob", "socket disconnected -> enter recovery", { promptId: this.promptId });
+            this.startDisconnectRecovery();
+        });
+        this.onReconnectHandlerOffFn = this.client.on("reconnected", () => {
+            if (!this.disconnectRecoveryActive) {
+                return;
+            }
+            this.emitLog("CallWrapper.enqueueJob", "socket reconnected", { promptId: this.promptId });
+            this.stopDisconnectRecovery();
+            void this.attemptHistoryCompletion("reconnected");
+        });
+        this.onReconnectFailedHandlerOffFn = this.client.on("reconnection_failed", () => {
+            if (!this.disconnectRecoveryActive) {
+                return;
+            }
+            this.emitLog("CallWrapper.enqueueJob", "reconnection failed", { promptId: this.promptId });
+            this.failDisconnected("reconnection_failed");
         });
         return job;
     }
+    resolvePromptLoad(value) {
+        const trigger = this.promptLoadTrigger;
+        if (!trigger) {
+            return;
+        }
+        this.promptLoadTrigger = null;
+        try {
+            trigger(value);
+        }
+        catch (error) {
+            this.emitLog("CallWrapper.resolvePromptLoad", "prompt load trigger threw", {
+                error: error instanceof Error ? error.message : String(error),
+                promptId: this.promptId
+            });
+        }
+    }
+    resolveJob(value) {
+        if (CALL_WRAPPER_DEBUG) {
+            console.log("[debug] resolveJob", this.promptId, value, Boolean(this.jobResolveFn), this.jobDoneResolved);
+        }
+        if (this.jobResolveFn) {
+            if (this.jobDoneResolved) {
+                return;
+            }
+            this.jobDoneResolved = true;
+            this.jobResolveFn(value);
+            if (CALL_WRAPPER_DEBUG) {
+                console.log("[debug] jobResolveFn invoked", this.promptId);
+            }
+        }
+        else {
+            this.pendingCompletion = value;
+        }
+    }
+    emitFailure(error, promptId) {
+        const fn = this.onFailedFn;
+        if (!fn) {
+            return;
+        }
+        const targetPromptId = promptId ?? this.promptId;
+        try {
+            if (CALL_WRAPPER_DEBUG) {
+                console.log("[debug] emitFailure start", error.name);
+            }
+            fn(error, targetPromptId);
+            if (CALL_WRAPPER_DEBUG) {
+                console.log("[debug] emitFailure end", error.name);
+            }
+        }
+        catch (callbackError) {
+            this.emitLog("CallWrapper.emitFailure", "onFailed callback threw", {
+                prompt_id: targetPromptId,
+                error: callbackError instanceof Error ? callbackError.message : String(callbackError)
+            });
+        }
+    }
+    cancel(reason = "cancelled") {
+        if (this.cancellationRequested) {
+            this.emitLog("CallWrapper.cancel", "cancel already requested", {
+                promptId: this.promptId,
+                reason
+            });
+            return;
+        }
+        this.cancellationRequested = true;
+        this.emitLog("CallWrapper.cancel", "cancel requested", {
+            promptId: this.promptId,
+            reason
+        });
+        this.resolvePromptLoad(false);
+        this.emitFailure(new ExecutionInterruptedError("The execution was interrupted!", { cause: { reason } }), this.promptId);
+        this.cleanupListeners("cancel requested");
+        this.resolveJob(false);
+    }
+    startDisconnectRecovery() {
+        if (this.disconnectRecoveryActive || this.cancellationRequested) {
+            return;
+        }
+        this.disconnectRecoveryActive = true;
+        if (this.disconnectFailureTimer) {
+            clearTimeout(this.disconnectFailureTimer);
+        }
+        this.disconnectFailureTimer = setTimeout(() => this.failDisconnected("timeout"), DISCONNECT_FAILURE_GRACE_MS);
+        void this.attemptHistoryCompletion("disconnect_start");
+    }
+    stopDisconnectRecovery() {
+        if (!this.disconnectRecoveryActive) {
+            return;
+        }
+        this.disconnectRecoveryActive = false;
+        if (this.disconnectFailureTimer) {
+            clearTimeout(this.disconnectFailureTimer);
+            this.disconnectFailureTimer = null;
+        }
+    }
+    async attemptHistoryCompletion(reason) {
+        if (!this.promptId || this.cancellationRequested) {
+            return false;
+        }
+        try {
+            const output = await this.handleCachedOutput(this.promptId);
+            if (output && output !== false) {
+                this.emitLog("CallWrapper.historyRecovery", "completed from history", { reason, promptId: this.promptId });
+                this.stopDisconnectRecovery();
+                this.isCompletingSuccessfully = true;
+                this.resolvePromptLoad(false);
+                this.resolveJob(output);
+                this.cleanupListeners(`history recovery (${reason})`);
+                return true;
+            }
+        }
+        catch (error) {
+            this.emitLog("CallWrapper.historyRecovery", "history fetch failed", { reason, error: String(error) });
+        }
+        return false;
+    }
+    failDisconnected(reason) {
+        if (!this.disconnectRecoveryActive || this.isCompletingSuccessfully) {
+            return;
+        }
+        this.stopDisconnectRecovery();
+        this.emitLog("CallWrapper.enqueueJob", "disconnect recovery failed", { reason, promptId: this.promptId });
+        this.resolvePromptLoad(false);
+        this.resolveJob(false);
+        this.cleanupListeners("disconnect failure");
+        this.emitFailure(new DisconnectedError("Disconnected"), this.promptId);
+    }
     async handleCachedOutput(promptId) {
         const hisData = await this.client.ext.history.getHistory(promptId);
         this.emitLog("CallWrapper.handleCachedOutput", "history fetched", {
@@ -418,15 +609,31 @@ export class CallWrapper {
         }
         return output;
     }
-    handleJobExecution(promptId, jobDoneTrigger) {
+    handleJobExecution(promptId) {
         if (this.executionHandlerOffFn) {
             return;
         }
         const reverseOutputMapped = this.reverseMapOutputKeys();
+        const mapOutputKeys = this.prompt.mapOutputKeys;
+        console.log(`[CallWrapper] handleJobExecution for ${promptId.substring(0, 8)}... - mapOutputKeys:`, mapOutputKeys, "reverseOutputMapped:", reverseOutputMapped);
         this.progressHandlerOffFn = this.client.on("progress", (ev) => this.handleProgress(ev, promptId));
-        this.previewHandlerOffFn = this.client.on("b_preview", (ev) => this.onPreviewFn?.(ev.detail, this.promptId));
+        this.previewHandlerOffFn = this.client.on("b_preview", (ev) => {
+            // Note: b_preview events don't include prompt_id. They're scoped per connection.
+            // If multiple jobs use the same connection, they will all receive preview events.
+            // This is a limitation of the ComfyUI protocol - previews are not separated by prompt_id.
+            this.onPreviewFn?.(ev.detail, this.promptId);
+        });
         // Also forward preview with metadata if available
-        const offPreviewMeta = this.client.on("b_preview_meta", (ev) => this.onPreviewMetaFn?.(ev.detail, this.promptId));
+        const offPreviewMeta = this.client.on("b_preview_meta", (ev) => {
+            // Validate prompt_id from metadata if available to prevent cross-user preview leakage
+            const metadata = ev.detail.metadata;
+            const metaPromptId = metadata?.prompt_id;
+            if (metaPromptId && metaPromptId !== promptId) {
+                console.log(`[CallWrapper] Ignoring b_preview_meta for wrong prompt. Expected ${promptId.substring(0, 8)}..., got ${metaPromptId.substring(0, 8)}...`);
+                return;
+            }
+            this.onPreviewMetaFn?.(ev.detail, this.promptId);
+        });
         const prevCleanup = this.previewHandlerOffFn;
         this.previewHandlerOffFn = () => {
             prevCleanup?.();
@@ -434,16 +641,18 @@ export class CallWrapper {
         };
         const totalOutput = Object.keys(reverseOutputMapped).length;
         let remainingOutput = totalOutput;
+        console.log(`[CallWrapper] totalOutput=${totalOutput}, remainingOutput=${remainingOutput}`);
         const executionHandler = (ev) => {
-            if (ev.detail.prompt_id !== promptId)
+            console.log(`[CallWrapper.executionHandler] received executed event for promptId=${ev.detail.prompt_id?.substring(0, 8)}..., node=${ev.detail.node}, waitingFor=${promptId.substring(0, 8)}...`);
+            const eventPromptId = ev.detail.prompt_id;
+            const isCorrectPrompt = eventPromptId === promptId;
+            // STRICT: Only accept events where prompt_id matches our expected promptId
+            if (!isCorrectPrompt) {
+                console.log(`[CallWrapper.executionHandler] REJECTED - prompt_id mismatch (expected ${promptId.substring(0, 8)}..., got ${eventPromptId?.substring(0, 8)}...)`);
                 return;
+            }
             const outputKey = reverseOutputMapped[ev.detail.node];
-            this.emitLog("CallWrapper.executionHandler", "executed event received", {
-                node: ev.detail.node,
-                outputKey,
-                remainingBefore: remainingOutput,
-                isTrackedOutput: !!outputKey
-            });
+            console.log(`[CallWrapper] executionHandler - promptId: ${promptId.substring(0, 8)}... (event says: ${ev.detail.prompt_id?.substring(0, 8)}...), node: ${ev.detail.node}, outputKey: ${outputKey}, output:`, JSON.stringify(ev.detail.output));
             if (outputKey) {
                 this.output[outputKey] = ev.detail.output;
                 this.onOutputFn?.(outputKey, ev.detail.output, this.promptId);
@@ -454,54 +663,117 @@ export class CallWrapper {
                 this.output._raw[ev.detail.node] = ev.detail.output;
                 this.onOutputFn?.(ev.detail.node, ev.detail.output, this.promptId);
             }
-            this.emitLog("CallWrapper.executionHandler", "after processing executed event", {
-                remainingAfter: remainingOutput,
-                willTriggerCompletion: remainingOutput === 0
-            });
+            console.log(`[CallWrapper] afterProcessing - remainingAfter: ${remainingOutput}, willTriggerCompletion: ${remainingOutput === 0}`);
             if (remainingOutput === 0) {
-                this.emitLog("CallWrapper.handleJobExecution", "all outputs collected");
+                console.log(`[CallWrapper] all outputs collected for ${promptId.substring(0, 8)}...`);
                 // Mark as successfully completing BEFORE cleanup to prevent race condition with disconnection handler
                 this.isCompletingSuccessfully = true;
                 this.cleanupListeners("all outputs collected");
                 this.onFinishedFn?.(this.output, this.promptId);
-                jobDoneTrigger(this.output);
+                this.resolveJob(this.output);
             }
         };
         const executedEnd = async () => {
-            this.emitLog("CallWrapper.executedEnd", "execution_success fired", {
-                promptId,
-                remainingOutput,
-                totalOutput
-            });
+            console.log(`[CallWrapper] execution_success fired for ${promptId.substring(0, 8)}..., remainingOutput=${remainingOutput}, totalOutput=${totalOutput}`);
+            // If we've already marked this as successfully completing, don't fail it again
+            if (this.isCompletingSuccessfully) {
+                console.log(`[CallWrapper] Already marked as successfully completing, ignoring this execution_success`);
+                return;
+            }
+            if (remainingOutput === 0) {
+                console.log(`[CallWrapper] all outputs already collected, nothing to do`);
+                return;
+            }
+            // Wait briefly for outputs that might be arriving due to prompt ID mismatch
+            await new Promise(resolve => setTimeout(resolve, 100));
+            console.log(`[CallWrapper] After wait - remainingOutput=${remainingOutput}, this.output keys:`, Object.keys(this.output));
+            // Check if outputs arrived while we were waiting
             if (remainingOutput === 0) {
-                this.emitLog("CallWrapper.executedEnd", "all outputs already collected, nothing to do");
+                console.log(`[CallWrapper] Outputs arrived during wait - marking as complete`);
+                this.isCompletingSuccessfully = true;
+                this.cleanupListeners("executedEnd - outputs complete after wait");
+                this.onFinishedFn?.(this.output, this.promptId);
+                this.resolveJob(this.output);
+                return;
+            }
+            // Check if we have collected all outputs (even if prompt ID mismatch)
+            const hasAllOutputs = Object.keys(reverseOutputMapped).every(nodeId => this.output[reverseOutputMapped[nodeId]] !== undefined);
+            if (hasAllOutputs) {
+                console.log(`[CallWrapper] Have all required outputs despite promptId mismatch - marking as complete`);
+                this.isCompletingSuccessfully = true;
+                this.cleanupListeners("executedEnd - outputs complete despite promptId mismatch");
+                this.onFinishedFn?.(this.output, this.promptId);
+                this.resolveJob(this.output);
                 return;
             }
-            const hisData = await this.client.ext.history.getHistory(promptId);
-            if (hisData?.status?.completed) {
-                const outputCount = Object.keys(hisData.outputs ?? {}).length;
-                if (outputCount > 0 && outputCount - totalOutput === 0) {
-                    this.emitLog("CallWrapper.executedEnd", "outputs equal total after history check -> ignore false end");
+            // Try to fetch from history with retry logic
+            let hisData = null;
+            for (let retries = 0; retries < 5; retries++) {
+                hisData = await this.client.ext.history.getHistory(promptId);
+                console.log(`[CallWrapper] History query result for ${promptId.substring(0, 8)}... (attempt ${retries + 1}) - status:`, hisData?.status, 'outputs:', Object.keys(hisData?.outputs ?? {}).length);
+                if (hisData?.status?.completed && hisData.outputs) {
+                    console.log(`[CallWrapper] Found completed job in history with outputs - attempting to populate from history`);
+                    break;
+                }
+                if (retries < 4) {
+                    console.log(`[CallWrapper] History not ready yet, waiting 100ms before retry...`);
+                    await new Promise(resolve => setTimeout(resolve, 100));
+                }
+            }
+            if (hisData?.status?.completed && hisData.outputs) {
+                // Try to extract outputs from history data
+                let populatedCount = 0;
+                for (const [nodeIdStr, nodeOutput] of Object.entries(hisData.outputs)) {
+                    const nodeId = parseInt(nodeIdStr, 10);
+                    const outputKey = reverseOutputMapped[nodeId];
+                    if (outputKey && nodeOutput) {
+                        // nodeOutput is typically { images: [...] } or similar - take the first property
+                        const outputValue = Array.isArray(nodeOutput) ? nodeOutput[0] : Object.values(nodeOutput)[0];
+                        if (outputValue !== undefined) {
+                            this.output[outputKey] = outputValue;
+                            this.onOutputFn?.(outputKey, outputValue, this.promptId);
+                            populatedCount++;
+                            remainingOutput--;
+                            console.log(`[CallWrapper] Populated ${outputKey} from history`);
+                        }
+                    }
+                }
+                if (remainingOutput === 0) {
+                    console.log(`[CallWrapper] Successfully populated all outputs from history for ${promptId.substring(0, 8)}...`);
+                    this.isCompletingSuccessfully = true;
+                    this.cleanupListeners("executedEnd - populated from history");
+                    this.onFinishedFn?.(this.output, this.promptId);
+                    this.resolveJob(this.output);
                     return;
                 }
+                if (populatedCount > 0) {
+                    console.log(`[CallWrapper] Populated ${populatedCount} outputs from history (remainingOutput=${remainingOutput})`);
+                    if (remainingOutput === 0) {
+                        this.isCompletingSuccessfully = true;
+                        this.cleanupListeners("executedEnd - all outputs from history");
+                        this.onFinishedFn?.(this.output, this.promptId);
+                        this.resolveJob(this.output);
+                        return;
+                    }
+                }
             }
-            this.emitLog("CallWrapper.executedEnd", "execution failed due to missing outputs", {
-                remainingOutput,
-                totalOutput
-            });
-            this.onFailedFn?.(new ExecutionFailedError("Execution failed"), this.promptId);
+            console.log(`[CallWrapper] execution failed due to missing outputs - remainingOutput=${remainingOutput}, totalOutput=${totalOutput}`);
+            this.emitFailure(new ExecutionFailedError("Execution failed"), this.promptId);
+            this.resolvePromptLoad(false);
             this.cleanupListeners("executedEnd missing outputs");
-            jobDoneTrigger(false);
+            this.resolveJob(false);
         };
         this.executionEndSuccessOffFn = this.client.on("execution_success", executedEnd);
         this.executionHandlerOffFn = this.client.on("executed", executionHandler);
-        this.errorHandlerOffFn = this.client.on("execution_error", (ev) => this.handleError(ev, promptId, jobDoneTrigger));
+        console.log(`[CallWrapper] Registered listeners for ${promptId.substring(0, 8)}... - executionHandler and executedEnd`);
+        this.errorHandlerOffFn = this.client.on("execution_error", (ev) => this.handleError(ev, promptId));
         this.interruptionHandlerOffFn = this.client.on("execution_interrupted", (ev) => {
             if (ev.detail.prompt_id !== promptId)
                 return;
-            this.onFailedFn?.(new ExecutionInterruptedError("The execution was interrupted!", { cause: ev.detail }), ev.detail.prompt_id);
+            this.emitFailure(new ExecutionInterruptedError("The execution was interrupted!", { cause: ev.detail }), ev.detail.prompt_id);
+            this.resolvePromptLoad(false);
             this.cleanupListeners("execution interrupted");
-            jobDoneTrigger(false);
+            this.resolveJob(false);
         });
     }
     reverseMapOutputKeys() {
@@ -519,16 +791,26 @@ export class CallWrapper {
         }
         this.onProgressFn?.(ev.detail, this.promptId);
     }
-    handleError(ev, promptId, resolve) {
+    handleError(ev, promptId) {
         if (ev.detail.prompt_id !== promptId)
             return;
         this.emitLog("CallWrapper.handleError", ev.detail.exception_type, {
             prompt_id: ev.detail.prompt_id,
             node_id: ev.detail?.node_id
         });
-        this.onFailedFn?.(new CustomEventError(ev.detail.exception_type, { cause: ev.detail }), ev.detail.prompt_id);
+        this.emitFailure(new CustomEventError(ev.detail.exception_type, { cause: ev.detail }), ev.detail.prompt_id);
+        if (CALL_WRAPPER_DEBUG) {
+            console.log("[debug] handleError after emitFailure");
+        }
+        this.resolvePromptLoad(false);
+        if (CALL_WRAPPER_DEBUG) {
+            console.log("[debug] handleError before cleanup");
+        }
         this.cleanupListeners("execution_error received");
-        resolve(false);
+        if (CALL_WRAPPER_DEBUG) {
+            console.log("[debug] handleError after cleanup");
+        }
+        this.resolveJob(false);
     }
     emitLog(fnName, message, data) {
         const detail = { fnName, message, data };
@@ -543,6 +825,13 @@ export class CallWrapper {
     cleanupListeners(reason) {
         const debugPayload = { reason, promptId: this.promptId };
         this.emitLog("CallWrapper.cleanupListeners", "removing listeners", debugPayload);
+        this.resolvePromptLoad(false);
+        this.stopDisconnectRecovery();
+        this.onReconnectHandlerOffFn?.();
+        this.onReconnectHandlerOffFn = undefined;
+        this.onReconnectFailedHandlerOffFn?.();
+        this.onReconnectFailedHandlerOffFn = undefined;
+        this.disconnectFailureTimer = null;
         this.onDisconnectedHandlerOffFn?.();
         this.onDisconnectedHandlerOffFn = undefined;
         this.checkExecutingOffFn?.();