From 1c7e656eb4e1d37c3f67af015f616f24acedc616 Mon Sep 17 00:00:00 2001 From: bymyself Date: Thu, 19 Feb 2026 23:35:24 -0800 Subject: [PATCH 1/5] Add prompt_id to progress_text binary WS messages Add supports_progress_text_metadata feature flag and extend send_progress_text() to accept optional prompt_id param. When prompt_id is provided and the client supports the new format, the binary wire format includes a length-prefixed prompt_id field: [4B event_type][4B prompt_id_len][prompt_id][4B node_id_len][node_id][text] Legacy format preserved for clients without the flag. Both callers (nodes_images.py, client.py) updated to pass prompt_id from get_executing_context(). Part of COM-12671: parallel workflow execution support. Amp-Thread-ID: https://ampcode.com/threads/T-019c79f7-f19b-70d9-b662-0687cc206282 --- comfy_api/feature_flags.py | 1 + comfy_api_nodes/util/client.py | 5 ++++- comfy_extras/nodes_images.py | 5 ++++- server.py | 23 ++++++++++++++++++++--- 4 files changed, 29 insertions(+), 5 deletions(-) diff --git a/comfy_api/feature_flags.py b/comfy_api/feature_flags.py index a90a5ca40..c3c804c96 100644 --- a/comfy_api/feature_flags.py +++ b/comfy_api/feature_flags.py @@ -12,6 +12,7 @@ from comfy.cli_args import args # Default server capabilities SERVER_FEATURE_FLAGS: dict[str, Any] = { "supports_preview_metadata": True, + "supports_progress_text_metadata": True, "max_upload_size": args.max_upload_size * 1024 * 1024, # Convert MB to bytes "extension": {"manager": {"supports_v4": True}}, "node_replacements": True, diff --git a/comfy_api_nodes/util/client.py b/comfy_api_nodes/util/client.py index 94886af7b..d0ec6e026 100644 --- a/comfy_api_nodes/util/client.py +++ b/comfy_api_nodes/util/client.py @@ -17,6 +17,7 @@ from pydantic import BaseModel from comfy import utils from comfy_api.latest import IO +from comfy_execution.utils import get_executing_context from server import PromptServer from . import request_logger @@ -446,7 +447,9 @@ def _display_text( if text is not None: display_lines.append(text) if display_lines: - PromptServer.instance.send_progress_text("\n".join(display_lines), get_node_id(node_cls)) + ctx = get_executing_context() + prompt_id = ctx.prompt_id if ctx is not None else None + PromptServer.instance.send_progress_text("\n".join(display_lines), get_node_id(node_cls), prompt_id=prompt_id) def _display_time_progress( diff --git a/comfy_extras/nodes_images.py b/comfy_extras/nodes_images.py index 727d7d09d..30aee6229 100644 --- a/comfy_extras/nodes_images.py +++ b/comfy_extras/nodes_images.py @@ -577,7 +577,10 @@ class GetImageSize(IO.ComfyNode): # Send progress text to display size on the node if cls.hidden.unique_id: - PromptServer.instance.send_progress_text(f"width: {width}, height: {height}\n batch size: {batch_size}", cls.hidden.unique_id) + from comfy_execution.utils import get_executing_context + ctx = get_executing_context() + prompt_id = ctx.prompt_id if ctx is not None else None + PromptServer.instance.send_progress_text(f"width: {width}, height: {height}\n batch size: {batch_size}", cls.hidden.unique_id, prompt_id=prompt_id) return IO.NodeOutput(width, height, batch_size) diff --git a/server.py b/server.py index 275bce5a7..2e9bacbf6 100644 --- a/server.py +++ b/server.py @@ -1233,13 +1233,30 @@ class PromptServer(): return json_data def send_progress_text( - self, text: Union[bytes, bytearray, str], node_id: str, sid=None + self, text: Union[bytes, bytearray, str], node_id: str, prompt_id: Optional[str] = None, sid=None ): if isinstance(text, str): text = text.encode("utf-8") node_id_bytes = str(node_id).encode("utf-8") - # Pack the node_id length as a 4-byte unsigned integer, followed by the node_id bytes - message = struct.pack(">I", len(node_id_bytes)) + node_id_bytes + text + # When prompt_id is provided and client supports the new format, + # prepend prompt_id as a length-prefixed field before node_id + target_sid = sid if sid is not None else self.client_id + if prompt_id and feature_flags.supports_feature( + self.sockets_metadata, target_sid, "supports_progress_text_metadata" + ): + prompt_id_bytes = prompt_id.encode("utf-8") + # Pack prompt_id length as a 4-byte unsigned integer, followed by prompt_id bytes, + # then node_id length as a 4-byte unsigned integer, followed by node_id bytes, then text + message = ( + struct.pack(">I", len(prompt_id_bytes)) + + prompt_id_bytes + + struct.pack(">I", len(node_id_bytes)) + + node_id_bytes + + text + ) + else: + # Pack the node_id length as a 4-byte unsigned integer, followed by the node_id bytes + message = struct.pack(">I", len(node_id_bytes)) + node_id_bytes + text self.send_sync(BinaryEventTypes.TEXT, message, sid) From 83df2a88bd00cd33256d67a04a573222eb4a4cf4 Mon Sep 17 00:00:00 2001 From: bymyself Date: Fri, 27 Feb 2026 17:21:14 -0800 Subject: [PATCH 2/5] refactor: add prompt_id as hidden type, fix imports, add docstrings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add PROMPT_ID as a new hidden type in the Hidden enum, HiddenHolder, HiddenInputTypeDict, and execution engine resolution (both V3 and legacy) - Refactor GetImageSize to use cls.hidden.prompt_id instead of manually calling get_executing_context() — addresses reviewer feedback - Remove lazy import of get_executing_context from nodes_images.py - Add docstrings to send_progress_text, _display_text, HiddenHolder, and HiddenHolder.from_dict Amp-Thread-ID: https://ampcode.com/threads/T-019ca1cb-0150-7549-8b1b-6713060d3408 --- comfy/comfy_types/node_typing.py | 2 ++ comfy_api/latest/_io.py | 22 +++++++++++++++++++++- comfy_api_nodes/util/client.py | 11 +++++++++++ comfy_extras/nodes_images.py | 7 ++----- execution.py | 8 ++++++-- server.py | 12 ++++++++++++ 6 files changed, 54 insertions(+), 8 deletions(-) diff --git a/comfy/comfy_types/node_typing.py b/comfy/comfy_types/node_typing.py index 92b1acbd5..fc5b895b1 100644 --- a/comfy/comfy_types/node_typing.py +++ b/comfy/comfy_types/node_typing.py @@ -193,6 +193,8 @@ class HiddenInputTypeDict(TypedDict): """EXTRA_PNGINFO is a dictionary that will be copied into the metadata of any .png files saved. Custom nodes can store additional information in this dictionary for saving (or as a way to communicate with a downstream node).""" dynprompt: NotRequired[Literal["DYNPROMPT"]] """DYNPROMPT is an instance of comfy_execution.graph.DynamicPrompt. It differs from PROMPT in that it may mutate during the course of execution in response to Node Expansion.""" + prompt_id: NotRequired[Literal["PROMPT_ID"]] + """PROMPT_ID is the unique identifier of the current prompt/job being executed. Useful for associating progress updates with specific jobs.""" class InputTypeDict(TypedDict): diff --git a/comfy_api/latest/_io.py b/comfy_api/latest/_io.py index 189d7d9bc..d8e6b356a 100644 --- a/comfy_api/latest/_io.py +++ b/comfy_api/latest/_io.py @@ -1269,9 +1269,16 @@ class V3Data(TypedDict): 'When True, the value of the dynamic input will be in the format (value, path_key).' class HiddenHolder: + """Holds hidden input values resolved during node execution. + + Hidden inputs are special values automatically provided by the execution + engine (e.g., node ID, prompt data, authentication tokens) rather than + being connected by the user in the graph. + """ def __init__(self, unique_id: str, prompt: Any, extra_pnginfo: Any, dynprompt: Any, - auth_token_comfy_org: str, api_key_comfy_org: str, **kwargs): + auth_token_comfy_org: str, api_key_comfy_org: str, + prompt_id: str = None, **kwargs): self.unique_id = unique_id """UNIQUE_ID is the unique identifier of the node, and matches the id property of the node on the client side. It is commonly used in client-server communications (see messages).""" self.prompt = prompt @@ -1284,6 +1291,8 @@ class HiddenHolder: """AUTH_TOKEN_COMFY_ORG is a token acquired from signing into a ComfyOrg account on frontend.""" self.api_key_comfy_org = api_key_comfy_org """API_KEY_COMFY_ORG is an API Key generated by ComfyOrg that allows skipping signing into a ComfyOrg account on frontend.""" + self.prompt_id = prompt_id + """PROMPT_ID is the unique identifier of the current prompt/job being executed.""" def __getattr__(self, key: str): '''If hidden variable not found, return None.''' @@ -1291,6 +1300,14 @@ class HiddenHolder: @classmethod def from_dict(cls, d: dict | None): + """Create a HiddenHolder from a dictionary of hidden input values. + + Args: + d: Dictionary mapping Hidden enum values to their resolved values. + + Returns: + A new HiddenHolder instance with values populated from the dict. + """ if d is None: d = {} return cls( @@ -1300,6 +1317,7 @@ class HiddenHolder: dynprompt=d.get(Hidden.dynprompt, None), auth_token_comfy_org=d.get(Hidden.auth_token_comfy_org, None), api_key_comfy_org=d.get(Hidden.api_key_comfy_org, None), + prompt_id=d.get(Hidden.prompt_id, None), ) @classmethod @@ -1322,6 +1340,8 @@ class Hidden(str, Enum): """AUTH_TOKEN_COMFY_ORG is a token acquired from signing into a ComfyOrg account on frontend.""" api_key_comfy_org = "API_KEY_COMFY_ORG" """API_KEY_COMFY_ORG is an API Key generated by ComfyOrg that allows skipping signing into a ComfyOrg account on frontend.""" + prompt_id = "PROMPT_ID" + """PROMPT_ID is the unique identifier of the current prompt/job being executed. Useful for associating progress updates with specific jobs.""" @dataclass diff --git a/comfy_api_nodes/util/client.py b/comfy_api_nodes/util/client.py index d0ec6e026..9e209fda9 100644 --- a/comfy_api_nodes/util/client.py +++ b/comfy_api_nodes/util/client.py @@ -437,6 +437,17 @@ def _display_text( status: str | int | None = None, price: float | None = None, ) -> None: + """Send a progress text message to the client for display on a node. + + Assembles status, price, and text lines, then sends them via WebSocket. + Automatically retrieves the current prompt_id from the execution context. + + Args: + node_cls: The ComfyNode class sending the progress text. + text: Optional text content to display. + status: Optional status string or code to display. + price: Optional price in dollars to display as credits. + """ display_lines: list[str] = [] if status: display_lines.append(f"Status: {status.capitalize() if isinstance(status, str) else status}") diff --git a/comfy_extras/nodes_images.py b/comfy_extras/nodes_images.py index 30aee6229..7b685d1cd 100644 --- a/comfy_extras/nodes_images.py +++ b/comfy_extras/nodes_images.py @@ -566,7 +566,7 @@ class GetImageSize(IO.ComfyNode): IO.Int.Output(display_name="height"), IO.Int.Output(display_name="batch_size"), ], - hidden=[IO.Hidden.unique_id], + hidden=[IO.Hidden.unique_id, IO.Hidden.prompt_id], ) @classmethod @@ -577,10 +577,7 @@ class GetImageSize(IO.ComfyNode): # Send progress text to display size on the node if cls.hidden.unique_id: - from comfy_execution.utils import get_executing_context - ctx = get_executing_context() - prompt_id = ctx.prompt_id if ctx is not None else None - PromptServer.instance.send_progress_text(f"width: {width}, height: {height}\n batch size: {batch_size}", cls.hidden.unique_id, prompt_id=prompt_id) + PromptServer.instance.send_progress_text(f"width: {width}, height: {height}\n batch size: {batch_size}", cls.hidden.unique_id, prompt_id=cls.hidden.prompt_id) return IO.NodeOutput(width, height, batch_size) diff --git a/execution.py b/execution.py index 75b021892..ba18569d6 100644 --- a/execution.py +++ b/execution.py @@ -149,7 +149,7 @@ class CacheSet: SENSITIVE_EXTRA_DATA_KEYS = ("auth_token_comfy_org", "api_key_comfy_org") -def get_input_data(inputs, class_def, unique_id, execution_list=None, dynprompt=None, extra_data={}): +def get_input_data(inputs, class_def, unique_id, execution_list=None, dynprompt=None, extra_data={}, prompt_id=None): is_v3 = issubclass(class_def, _ComfyNodeInternal) v3_data: io.V3Data = {} hidden_inputs_v3 = {} @@ -196,6 +196,8 @@ def get_input_data(inputs, class_def, unique_id, execution_list=None, dynprompt= hidden_inputs_v3[io.Hidden.auth_token_comfy_org] = extra_data.get("auth_token_comfy_org", None) if io.Hidden.api_key_comfy_org.name in hidden: hidden_inputs_v3[io.Hidden.api_key_comfy_org] = extra_data.get("api_key_comfy_org", None) + if io.Hidden.prompt_id.name in hidden: + hidden_inputs_v3[io.Hidden.prompt_id] = prompt_id else: if "hidden" in valid_inputs: h = valid_inputs["hidden"] @@ -212,6 +214,8 @@ def get_input_data(inputs, class_def, unique_id, execution_list=None, dynprompt= input_data_all[x] = [extra_data.get("auth_token_comfy_org", None)] if h[x] == "API_KEY_COMFY_ORG": input_data_all[x] = [extra_data.get("api_key_comfy_org", None)] + if h[x] == "PROMPT_ID": + input_data_all[x] = [prompt_id] v3_data["hidden_inputs"] = hidden_inputs_v3 return input_data_all, missing_keys, v3_data @@ -469,7 +473,7 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed, has_subgraph = False else: get_progress_state().start_progress(unique_id) - input_data_all, missing_keys, v3_data = get_input_data(inputs, class_def, unique_id, execution_list, dynprompt, extra_data) + input_data_all, missing_keys, v3_data = get_input_data(inputs, class_def, unique_id, execution_list, dynprompt, extra_data, prompt_id=prompt_id) if server.client_id is not None: server.last_node_id = display_node_id server.send_sync("executing", { "node": unique_id, "display_node": display_node_id, "prompt_id": prompt_id }, server.client_id) diff --git a/server.py b/server.py index 2e9bacbf6..3065bb7f7 100644 --- a/server.py +++ b/server.py @@ -1235,6 +1235,18 @@ class PromptServer(): def send_progress_text( self, text: Union[bytes, bytearray, str], node_id: str, prompt_id: Optional[str] = None, sid=None ): + """Send a progress text message to the client via WebSocket. + + Encodes the text as a binary message with length-prefixed node_id. When + prompt_id is provided and the client supports the ``supports_progress_text_metadata`` + feature flag, the prompt_id is prepended as an additional length-prefixed field. + + Args: + text: The progress text content to send. + node_id: The unique identifier of the node sending the progress. + prompt_id: Optional prompt/job identifier to associate the message with. + sid: Optional session ID to target a specific client. + """ if isinstance(text, str): text = text.encode("utf-8") node_id_bytes = str(node_id).encode("utf-8") From d74dfd2570e829f1787d6bb796fbeca596e54815 Mon Sep 17 00:00:00 2001 From: bymyself Date: Wed, 4 Mar 2026 20:40:46 +0000 Subject: [PATCH 3/5] fix: send_progress_text unicasts to client_id instead of broadcasting - Default sid to self.client_id when not explicitly provided, matching every other WS message dispatch (executing, executed, progress_state, etc.) - Previously sid=None caused broadcast to all connected clients - Format signature per ruff, remove redundant comments - Add unit tests for routing, legacy format, and new prompt_id format Amp-Thread-ID: https://ampcode.com/threads/T-019ca3ce-c530-75dd-8d68-349e745a022e --- server.py | 17 +- .../send_progress_text_test.py | 207 ++++++++++++++++++ 2 files changed, 217 insertions(+), 7 deletions(-) create mode 100644 tests-unit/prompt_server_test/send_progress_text_test.py diff --git a/server.py b/server.py index 3065bb7f7..608c98d45 100644 --- a/server.py +++ b/server.py @@ -1233,7 +1233,11 @@ class PromptServer(): return json_data def send_progress_text( - self, text: Union[bytes, bytearray, str], node_id: str, prompt_id: Optional[str] = None, sid=None + self, + text: Union[bytes, bytearray, str], + node_id: str, + prompt_id: Optional[str] = None, + sid=None, ): """Send a progress text message to the client via WebSocket. @@ -1251,15 +1255,15 @@ class PromptServer(): text = text.encode("utf-8") node_id_bytes = str(node_id).encode("utf-8") - # When prompt_id is provided and client supports the new format, - # prepend prompt_id as a length-prefixed field before node_id + # Auto-resolve sid to the currently executing client target_sid = sid if sid is not None else self.client_id + + # When prompt_id is available and client supports the new format, + # prepend prompt_id as a length-prefixed field before node_id if prompt_id and feature_flags.supports_feature( self.sockets_metadata, target_sid, "supports_progress_text_metadata" ): prompt_id_bytes = prompt_id.encode("utf-8") - # Pack prompt_id length as a 4-byte unsigned integer, followed by prompt_id bytes, - # then node_id length as a 4-byte unsigned integer, followed by node_id bytes, then text message = ( struct.pack(">I", len(prompt_id_bytes)) + prompt_id_bytes @@ -1268,7 +1272,6 @@ class PromptServer(): + text ) else: - # Pack the node_id length as a 4-byte unsigned integer, followed by the node_id bytes message = struct.pack(">I", len(node_id_bytes)) + node_id_bytes + text - self.send_sync(BinaryEventTypes.TEXT, message, sid) + self.send_sync(BinaryEventTypes.TEXT, message, target_sid) diff --git a/tests-unit/prompt_server_test/send_progress_text_test.py b/tests-unit/prompt_server_test/send_progress_text_test.py new file mode 100644 index 000000000..7631a4fb1 --- /dev/null +++ b/tests-unit/prompt_server_test/send_progress_text_test.py @@ -0,0 +1,207 @@ +"""Tests for send_progress_text routing and binary format logic. + +These tests verify: +1. sid defaults to client_id (unicast) instead of None (broadcast) +2. Legacy binary format when prompt_id absent or client unsupported +3. New binary format with prompt_id when client supports the feature flag +""" + +import struct + +from comfy_api import feature_flags + + +# --------------------------------------------------------------------------- +# Helpers – replicate the packing logic so we can assert on the wire format +# --------------------------------------------------------------------------- + + +def _unpack_legacy(message: bytes): + """Unpack a legacy progress_text binary message -> (node_id, text).""" + offset = 0 + node_id_len = struct.unpack_from(">I", message, offset)[0] + offset += 4 + node_id = message[offset : offset + node_id_len].decode("utf-8") + offset += node_id_len + text = message[offset:].decode("utf-8") + return node_id, text + + +def _unpack_with_prompt_id(message: bytes): + """Unpack new format -> (prompt_id, node_id, text).""" + offset = 0 + prompt_id_len = struct.unpack_from(">I", message, offset)[0] + offset += 4 + prompt_id = message[offset : offset + prompt_id_len].decode("utf-8") + offset += prompt_id_len + node_id_len = struct.unpack_from(">I", message, offset)[0] + offset += 4 + node_id = message[offset : offset + node_id_len].decode("utf-8") + offset += node_id_len + text = message[offset:].decode("utf-8") + return prompt_id, node_id, text + + +# --------------------------------------------------------------------------- +# Minimal stub that mirrors send_progress_text logic from server.py +# We can't import server.py directly (it pulls in torch via nodes.py), +# so we replicate the method body here. If the implementation changes, +# these tests should be updated in tandem. +# --------------------------------------------------------------------------- + + +class _StubServer: + """Stub that captures send_sync calls and runs the real packing logic.""" + + def __init__(self, client_id=None, sockets_metadata=None): + self.client_id = client_id + self.sockets_metadata = sockets_metadata or {} + self.sent = [] # list of (event, data, sid) + + def send_sync(self, event, data, sid=None): + self.sent.append((event, data, sid)) + + def send_progress_text(self, text, node_id, prompt_id=None, sid=None): + if isinstance(text, str): + text = text.encode("utf-8") + node_id_bytes = str(node_id).encode("utf-8") + + target_sid = sid if sid is not None else self.client_id + + if prompt_id and feature_flags.supports_feature( + self.sockets_metadata, target_sid, "supports_progress_text_metadata" + ): + prompt_id_bytes = prompt_id.encode("utf-8") + message = ( + struct.pack(">I", len(prompt_id_bytes)) + + prompt_id_bytes + + struct.pack(">I", len(node_id_bytes)) + + node_id_bytes + + text + ) + else: + message = struct.pack(">I", len(node_id_bytes)) + node_id_bytes + text + + self.send_sync(3, message, target_sid) # 3 == BinaryEventTypes.TEXT + + +# =========================================================================== +# Routing tests +# =========================================================================== + + +class TestSendProgressTextRouting: + """Verify sid resolution: defaults to client_id, overridable via sid param.""" + + def test_defaults_to_client_id_when_sid_not_provided(self): + server = _StubServer(client_id="active-client-123") + server.send_progress_text("hello", "node1") + + _, _, sid = server.sent[0] + assert sid == "active-client-123" + + def test_explicit_sid_overrides_client_id(self): + server = _StubServer(client_id="active-client-123") + server.send_progress_text("hello", "node1", sid="explicit-sid") + + _, _, sid = server.sent[0] + assert sid == "explicit-sid" + + def test_broadcasts_when_no_client_id_and_no_sid(self): + server = _StubServer(client_id=None) + server.send_progress_text("hello", "node1") + + _, _, sid = server.sent[0] + assert sid is None + + +# =========================================================================== +# Legacy format tests +# =========================================================================== + + +class TestSendProgressTextLegacyFormat: + """Verify legacy binary format: [4B node_id_len][node_id][text].""" + + def test_legacy_format_no_prompt_id(self): + server = _StubServer(client_id="c1") + server.send_progress_text("some text", "node-42") + + _, data, _ = server.sent[0] + node_id, text = _unpack_legacy(data) + assert node_id == "node-42" + assert text == "some text" + + def test_legacy_format_when_client_unsupported(self): + server = _StubServer( + client_id="c1", + sockets_metadata={"c1": {"feature_flags": {}}}, + ) + server.send_progress_text("text", "node1", prompt_id="prompt-abc") + + _, data, _ = server.sent[0] + node_id, text = _unpack_legacy(data) + assert node_id == "node1" + assert text == "text" + + def test_bytes_input_preserved(self): + server = _StubServer(client_id="c1") + server.send_progress_text(b"raw bytes", "node1") + + _, data, _ = server.sent[0] + node_id, text = _unpack_legacy(data) + assert text == "raw bytes" + + +# =========================================================================== +# New format tests +# =========================================================================== + + +class TestSendProgressTextNewFormat: + """Verify new format: [4B prompt_id_len][prompt_id][4B node_id_len][node_id][text].""" + + def test_includes_prompt_id_when_supported(self): + server = _StubServer( + client_id="c1", + sockets_metadata={ + "c1": {"feature_flags": {"supports_progress_text_metadata": True}} + }, + ) + server.send_progress_text("progress!", "node-7", prompt_id="prompt-xyz") + + _, data, _ = server.sent[0] + prompt_id, node_id, text = _unpack_with_prompt_id(data) + assert prompt_id == "prompt-xyz" + assert node_id == "node-7" + assert text == "progress!" + + def test_new_format_with_explicit_sid(self): + server = _StubServer( + client_id=None, + sockets_metadata={ + "my-sid": {"feature_flags": {"supports_progress_text_metadata": True}} + }, + ) + server.send_progress_text("txt", "n1", prompt_id="p1", sid="my-sid") + + _, data, sid = server.sent[0] + assert sid == "my-sid" + prompt_id, node_id, text = _unpack_with_prompt_id(data) + assert prompt_id == "p1" + assert node_id == "n1" + assert text == "txt" + + +# =========================================================================== +# Feature flag tests +# =========================================================================== + + +class TestProgressTextFeatureFlag: + """Verify the supports_progress_text_metadata flag exists in server features.""" + + def test_flag_in_server_features(self): + features = feature_flags.get_server_features() + assert "supports_progress_text_metadata" in features + assert features["supports_progress_text_metadata"] is True From 09e9bdbcad7d15307600b462126e4feb263eb2e9 Mon Sep 17 00:00:00 2001 From: bymyself Date: Wed, 4 Mar 2026 20:50:01 +0000 Subject: [PATCH 4/5] remove send_progress_text stub tests Copy-paste stub tests don't verify the real implementation and add maintenance burden without meaningful coverage. Amp-Thread-ID: https://ampcode.com/threads/T-019ca3ce-c530-75dd-8d68-349e745a022e --- .../send_progress_text_test.py | 207 ------------------ 1 file changed, 207 deletions(-) delete mode 100644 tests-unit/prompt_server_test/send_progress_text_test.py diff --git a/tests-unit/prompt_server_test/send_progress_text_test.py b/tests-unit/prompt_server_test/send_progress_text_test.py deleted file mode 100644 index 7631a4fb1..000000000 --- a/tests-unit/prompt_server_test/send_progress_text_test.py +++ /dev/null @@ -1,207 +0,0 @@ -"""Tests for send_progress_text routing and binary format logic. - -These tests verify: -1. sid defaults to client_id (unicast) instead of None (broadcast) -2. Legacy binary format when prompt_id absent or client unsupported -3. New binary format with prompt_id when client supports the feature flag -""" - -import struct - -from comfy_api import feature_flags - - -# --------------------------------------------------------------------------- -# Helpers – replicate the packing logic so we can assert on the wire format -# --------------------------------------------------------------------------- - - -def _unpack_legacy(message: bytes): - """Unpack a legacy progress_text binary message -> (node_id, text).""" - offset = 0 - node_id_len = struct.unpack_from(">I", message, offset)[0] - offset += 4 - node_id = message[offset : offset + node_id_len].decode("utf-8") - offset += node_id_len - text = message[offset:].decode("utf-8") - return node_id, text - - -def _unpack_with_prompt_id(message: bytes): - """Unpack new format -> (prompt_id, node_id, text).""" - offset = 0 - prompt_id_len = struct.unpack_from(">I", message, offset)[0] - offset += 4 - prompt_id = message[offset : offset + prompt_id_len].decode("utf-8") - offset += prompt_id_len - node_id_len = struct.unpack_from(">I", message, offset)[0] - offset += 4 - node_id = message[offset : offset + node_id_len].decode("utf-8") - offset += node_id_len - text = message[offset:].decode("utf-8") - return prompt_id, node_id, text - - -# --------------------------------------------------------------------------- -# Minimal stub that mirrors send_progress_text logic from server.py -# We can't import server.py directly (it pulls in torch via nodes.py), -# so we replicate the method body here. If the implementation changes, -# these tests should be updated in tandem. -# --------------------------------------------------------------------------- - - -class _StubServer: - """Stub that captures send_sync calls and runs the real packing logic.""" - - def __init__(self, client_id=None, sockets_metadata=None): - self.client_id = client_id - self.sockets_metadata = sockets_metadata or {} - self.sent = [] # list of (event, data, sid) - - def send_sync(self, event, data, sid=None): - self.sent.append((event, data, sid)) - - def send_progress_text(self, text, node_id, prompt_id=None, sid=None): - if isinstance(text, str): - text = text.encode("utf-8") - node_id_bytes = str(node_id).encode("utf-8") - - target_sid = sid if sid is not None else self.client_id - - if prompt_id and feature_flags.supports_feature( - self.sockets_metadata, target_sid, "supports_progress_text_metadata" - ): - prompt_id_bytes = prompt_id.encode("utf-8") - message = ( - struct.pack(">I", len(prompt_id_bytes)) - + prompt_id_bytes - + struct.pack(">I", len(node_id_bytes)) - + node_id_bytes - + text - ) - else: - message = struct.pack(">I", len(node_id_bytes)) + node_id_bytes + text - - self.send_sync(3, message, target_sid) # 3 == BinaryEventTypes.TEXT - - -# =========================================================================== -# Routing tests -# =========================================================================== - - -class TestSendProgressTextRouting: - """Verify sid resolution: defaults to client_id, overridable via sid param.""" - - def test_defaults_to_client_id_when_sid_not_provided(self): - server = _StubServer(client_id="active-client-123") - server.send_progress_text("hello", "node1") - - _, _, sid = server.sent[0] - assert sid == "active-client-123" - - def test_explicit_sid_overrides_client_id(self): - server = _StubServer(client_id="active-client-123") - server.send_progress_text("hello", "node1", sid="explicit-sid") - - _, _, sid = server.sent[0] - assert sid == "explicit-sid" - - def test_broadcasts_when_no_client_id_and_no_sid(self): - server = _StubServer(client_id=None) - server.send_progress_text("hello", "node1") - - _, _, sid = server.sent[0] - assert sid is None - - -# =========================================================================== -# Legacy format tests -# =========================================================================== - - -class TestSendProgressTextLegacyFormat: - """Verify legacy binary format: [4B node_id_len][node_id][text].""" - - def test_legacy_format_no_prompt_id(self): - server = _StubServer(client_id="c1") - server.send_progress_text("some text", "node-42") - - _, data, _ = server.sent[0] - node_id, text = _unpack_legacy(data) - assert node_id == "node-42" - assert text == "some text" - - def test_legacy_format_when_client_unsupported(self): - server = _StubServer( - client_id="c1", - sockets_metadata={"c1": {"feature_flags": {}}}, - ) - server.send_progress_text("text", "node1", prompt_id="prompt-abc") - - _, data, _ = server.sent[0] - node_id, text = _unpack_legacy(data) - assert node_id == "node1" - assert text == "text" - - def test_bytes_input_preserved(self): - server = _StubServer(client_id="c1") - server.send_progress_text(b"raw bytes", "node1") - - _, data, _ = server.sent[0] - node_id, text = _unpack_legacy(data) - assert text == "raw bytes" - - -# =========================================================================== -# New format tests -# =========================================================================== - - -class TestSendProgressTextNewFormat: - """Verify new format: [4B prompt_id_len][prompt_id][4B node_id_len][node_id][text].""" - - def test_includes_prompt_id_when_supported(self): - server = _StubServer( - client_id="c1", - sockets_metadata={ - "c1": {"feature_flags": {"supports_progress_text_metadata": True}} - }, - ) - server.send_progress_text("progress!", "node-7", prompt_id="prompt-xyz") - - _, data, _ = server.sent[0] - prompt_id, node_id, text = _unpack_with_prompt_id(data) - assert prompt_id == "prompt-xyz" - assert node_id == "node-7" - assert text == "progress!" - - def test_new_format_with_explicit_sid(self): - server = _StubServer( - client_id=None, - sockets_metadata={ - "my-sid": {"feature_flags": {"supports_progress_text_metadata": True}} - }, - ) - server.send_progress_text("txt", "n1", prompt_id="p1", sid="my-sid") - - _, data, sid = server.sent[0] - assert sid == "my-sid" - prompt_id, node_id, text = _unpack_with_prompt_id(data) - assert prompt_id == "p1" - assert node_id == "n1" - assert text == "txt" - - -# =========================================================================== -# Feature flag tests -# =========================================================================== - - -class TestProgressTextFeatureFlag: - """Verify the supports_progress_text_metadata flag exists in server features.""" - - def test_flag_in_server_features(self): - features = feature_flags.get_server_features() - assert "supports_progress_text_metadata" in features - assert features["supports_progress_text_metadata"] is True From 69d3bfa391b3710b48f7458d99d624d8c3ae76a6 Mon Sep 17 00:00:00 2001 From: bymyself Date: Thu, 12 Mar 2026 09:20:52 -0700 Subject: [PATCH 5/5] fix: always send new binary format when client supports feature flag When prompt_id is None, encode as zero-length string instead of falling back to old format. Prevents binary parse corruption on the frontend. Addresses review feedback: https://github.com/Comfy-Org/ComfyUI/pull/12540#discussion_r2923412491 --- server.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/server.py b/server.py index 608c98d45..97ebf476c 100644 --- a/server.py +++ b/server.py @@ -1242,8 +1242,9 @@ class PromptServer(): """Send a progress text message to the client via WebSocket. Encodes the text as a binary message with length-prefixed node_id. When - prompt_id is provided and the client supports the ``supports_progress_text_metadata`` - feature flag, the prompt_id is prepended as an additional length-prefixed field. + the client supports the ``supports_progress_text_metadata`` feature flag, + the prompt_id is always prepended as a length-prefixed field (empty string + when None) to ensure consistent binary framing. Args: text: The progress text content to send. @@ -1258,12 +1259,13 @@ class PromptServer(): # Auto-resolve sid to the currently executing client target_sid = sid if sid is not None else self.client_id - # When prompt_id is available and client supports the new format, - # prepend prompt_id as a length-prefixed field before node_id - if prompt_id and feature_flags.supports_feature( + # When client supports the new format, always send + # [prompt_id_len][prompt_id][node_id_len][node_id][text] + # even when prompt_id is None (encoded as zero-length string) + if feature_flags.supports_feature( self.sockets_metadata, target_sid, "supports_progress_text_metadata" ): - prompt_id_bytes = prompt_id.encode("utf-8") + prompt_id_bytes = (prompt_id or "").encode("utf-8") message = ( struct.pack(">I", len(prompt_id_bytes)) + prompt_id_bytes