comfyanonymous/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-02-19 06:30:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: replay execution context on WS reconnect

Browse Source 
Store current execution state (prompt_id, cached nodes, executed nodes,
outputs to execute) on the server instance during execution. On WS
reconnect, replay execution_start, execution_cached, progress_state,
and executing events so the frontend can restore progress tracking.

Refactor ProgressRegistry to expose get_serialized_state() and reuse
it in WebUIProgressHandler._send_progress_state().
This commit is contained in:
Johnpaul
2026-02-16 23:53:50 +01:00
parent e5ae670a40
commit 3278ff75fd
3 changed files with 64 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
comfy_execution/progress.py
View File

@@ -164,21 +164,7 @@ class WebUIProgressHandler(ProgressHandler):
if self.server_instance is None:
return
# Only send info for non-pending nodes
active_nodes = {
node_id: {
"value": state["value"],
"max": state["max"],
"state": state["state"].value,
"node_id": node_id,
"prompt_id": prompt_id,
"display_node_id": self.registry.dynprompt.get_display_node_id(node_id),
"parent_node_id": self.registry.dynprompt.get_parent_node_id(node_id),
"real_node_id": self.registry.dynprompt.get_real_node_id(node_id),
}
for node_id, state in nodes.items()
if state["state"] != NodeState.Pending
}
active_nodes = self.registry.get_serialized_state()
# Send a combined progress_state message with all node states
# Include client_id to ensure message is only sent to the initiating client
@@ -314,6 +300,24 @@ class ProgressRegistry:
if handler.enabled:
handler.finish_handler(node_id, entry, self.prompt_id)
def get_serialized_state(self) -> Dict[str, dict]:
"""Return current node progress as a dict suitable for WS progress_state."""
active: Dict[str, dict] = {}
for nid, state in self.nodes.items():
if state["state"] == NodeState.Pending:
continue
active[nid] = {
"value": state["value"],
"max": state["max"],
"state": state["state"].value,
"node_id": nid,
"prompt_id": self.prompt_id,
"display_node_id": self.dynprompt.get_display_node_id(nid),
"parent_node_id": self.dynprompt.get_parent_node_id(nid),
"real_node_id": self.dynprompt.get_real_node_id(nid),
}
return active
def reset_handlers(self) -> None:
"""Reset all handlers"""
for handler in self.handlers.values():