From d3878edb013c78ed33be34f1e50e85cae7b8778c Mon Sep 17 00:00:00 2001 From: Robin Huang Date: Wed, 10 Jun 2026 14:17:46 -0700 Subject: [PATCH 1/3] Add Comfy-Usage-Source pass-through for API node requests Capture the Comfy-Usage-Source header (or extra_data.comfy_usage_source) on POST /prompt and forward it on API nodes' outbound requests to api.comfy.org, defaulting to comfyui-server when absent. --- comfy_api/latest/_io.py | 10 +++++++++- comfy_api_nodes/nodes_sonilo.py | 2 ++ comfy_api_nodes/util/_helpers.py | 5 +++++ comfy_api_nodes/util/client.py | 2 ++ comfy_api_nodes/util/download_helpers.py | 2 ++ execution.py | 4 ++++ server.py | 5 +++++ 7 files changed, 29 insertions(+), 1 deletion(-) diff --git a/comfy_api/latest/_io.py b/comfy_api/latest/_io.py index 37614a4c3325..012fae3aca8f 100644 --- a/comfy_api/latest/_io.py +++ b/comfy_api/latest/_io.py @@ -1400,7 +1400,8 @@ class V3Data(TypedDict): class HiddenHolder: 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, + comfy_usage_source: 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 @@ -1413,6 +1414,8 @@ def __init__(self, unique_id: str, prompt: Any, """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.comfy_usage_source = comfy_usage_source + """COMFY_USAGE_SOURCE identifies the client that submitted the prompt (e.g. comfyui-frontend, comfy-cli, comfyui-mcp); forwarded to API nodes' upstream requests via the Comfy-Usage-Source header.""" def __getattr__(self, key: str): '''If hidden variable not found, return None.''' @@ -1429,6 +1432,7 @@ def from_dict(cls, d: dict | None): 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), + comfy_usage_source=d.get(Hidden.comfy_usage_source, None), ) @classmethod @@ -1451,6 +1455,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.""" + comfy_usage_source = "COMFY_USAGE_SOURCE" + """COMFY_USAGE_SOURCE identifies the client that submitted the prompt (e.g. comfyui-frontend, comfy-cli, comfyui-mcp); forwarded to API nodes' upstream requests via the Comfy-Usage-Source header.""" @dataclass @@ -1654,6 +1660,8 @@ def finalize(self): self.hidden.append(Hidden.auth_token_comfy_org) if Hidden.api_key_comfy_org not in self.hidden: self.hidden.append(Hidden.api_key_comfy_org) + if Hidden.comfy_usage_source not in self.hidden: + self.hidden.append(Hidden.comfy_usage_source) # if is an output_node, will need prompt and extra_pnginfo if self.is_output_node: if Hidden.prompt not in self.hidden: diff --git a/comfy_api_nodes/nodes_sonilo.py b/comfy_api_nodes/nodes_sonilo.py index 9ce896ed0066..81bb423b7fcb 100644 --- a/comfy_api_nodes/nodes_sonilo.py +++ b/comfy_api_nodes/nodes_sonilo.py @@ -18,6 +18,7 @@ default_base_url, get_auth_header, get_node_id, + get_usage_source, is_processing_interrupted, ) from comfy_api_nodes.util.common_exceptions import ProcessingInterrupted @@ -176,6 +177,7 @@ async def _stream_sonilo_music( headers: dict[str, str] = {} headers.update(get_auth_header(cls)) + headers["Comfy-Usage-Source"] = get_usage_source(cls) headers.update(endpoint.headers) node_id = get_node_id(cls) diff --git a/comfy_api_nodes/util/_helpers.py b/comfy_api_nodes/util/_helpers.py index 648defe3deba..5d775e76a75e 100644 --- a/comfy_api_nodes/util/_helpers.py +++ b/comfy_api_nodes/util/_helpers.py @@ -35,6 +35,11 @@ def get_auth_header(node_cls: type[IO.ComfyNode]) -> dict[str, str]: return {} +def get_usage_source(node_cls: type[IO.ComfyNode]) -> str: + """Source of the prompt that triggered this API node, defaulting to this server itself.""" + return node_cls.hidden.comfy_usage_source or "comfyui-server" + + def default_base_url() -> str: return getattr(args, "comfy_api_base", "https://api.comfy.org") diff --git a/comfy_api_nodes/util/client.py b/comfy_api_nodes/util/client.py index 57c501724bec..1cb113ddad55 100644 --- a/comfy_api_nodes/util/client.py +++ b/comfy_api_nodes/util/client.py @@ -26,6 +26,7 @@ default_base_url, get_auth_header, get_node_id, + get_usage_source, is_processing_interrupted, sleep_with_interrupt, ) @@ -647,6 +648,7 @@ async def _monitor(stop_evt: asyncio.Event, start_ts: float): if not parsed_url.scheme and not parsed_url.netloc: # is URL relative? payload_headers.update(get_auth_header(cfg.node_cls)) payload_headers["Comfy-Env"] = get_deploy_environment() + payload_headers["Comfy-Usage-Source"] = get_usage_source(cfg.node_cls) if cfg.endpoint.headers: payload_headers.update(cfg.endpoint.headers) diff --git a/comfy_api_nodes/util/download_helpers.py b/comfy_api_nodes/util/download_helpers.py index aa588d038b4d..b2b2af8355a9 100644 --- a/comfy_api_nodes/util/download_helpers.py +++ b/comfy_api_nodes/util/download_helpers.py @@ -18,6 +18,7 @@ from ._helpers import ( default_base_url, get_auth_header, + get_usage_source, is_processing_interrupted, sleep_with_interrupt, to_aiohttp_url, @@ -65,6 +66,7 @@ async def download_url_to_bytesio( raise ValueError("For relative 'cloud' paths, the `cls` parameter is required.") url = urljoin(default_base_url().rstrip("/") + "/", url.lstrip("/")) headers = get_auth_header(cls) + headers["Comfy-Usage-Source"] = get_usage_source(cls) while True: attempt += 1 diff --git a/execution.py b/execution.py index 5246d651cfa5..32b8b2e01897 100644 --- a/execution.py +++ b/execution.py @@ -199,6 +199,8 @@ def mark_missing(): 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.comfy_usage_source.name in hidden: + hidden_inputs_v3[io.Hidden.comfy_usage_source] = extra_data.get("comfy_usage_source", None) else: if "hidden" in valid_inputs: h = valid_inputs["hidden"] @@ -215,6 +217,8 @@ def mark_missing(): 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] == "COMFY_USAGE_SOURCE": + input_data_all[x] = [extra_data.get("comfy_usage_source", None)] v3_data["hidden_inputs"] = hidden_inputs_v3 return input_data_all, missing_keys, v3_data diff --git a/server.py b/server.py index a85c1e59147c..ba00f99dd736 100644 --- a/server.py +++ b/server.py @@ -957,6 +957,11 @@ async def post_prompt(request): if "client_id" in json_data: extra_data["client_id"] = json_data["client_id"] + + if "comfy_usage_source" not in extra_data: + usage_source = request.headers.get("Comfy-Usage-Source") + if usage_source: + extra_data["comfy_usage_source"] = usage_source if valid[0]: outputs_to_execute = valid[2] sensitive = {} From b58af4279b16cbb0e3ff7cc6ac53cb4d12c82712 Mon Sep 17 00:00:00 2001 From: Robin Huang Date: Thu, 11 Jun 2026 12:10:09 -0700 Subject: [PATCH 2/3] Centralize shared Comfy API headers in get_comfy_api_headers helper --- comfy_api_nodes/nodes_sonilo.py | 7 ++----- comfy_api_nodes/util/_helpers.py | 16 ++++++++++++++++ comfy_api_nodes/util/client.py | 9 ++------- comfy_api_nodes/util/download_helpers.py | 6 ++---- 4 files changed, 22 insertions(+), 16 deletions(-) diff --git a/comfy_api_nodes/nodes_sonilo.py b/comfy_api_nodes/nodes_sonilo.py index 81bb423b7fcb..24a9a0b06fb8 100644 --- a/comfy_api_nodes/nodes_sonilo.py +++ b/comfy_api_nodes/nodes_sonilo.py @@ -16,9 +16,8 @@ ) from comfy_api_nodes.util._helpers import ( default_base_url, - get_auth_header, + get_comfy_api_headers, get_node_id, - get_usage_source, is_processing_interrupted, ) from comfy_api_nodes.util.common_exceptions import ProcessingInterrupted @@ -175,9 +174,7 @@ async def _stream_sonilo_music( """POST ``form`` to Sonilo, read the NDJSON stream, and return the first stream's audio bytes.""" url = urljoin(default_base_url().rstrip("/") + "/", endpoint.path.lstrip("/")) - headers: dict[str, str] = {} - headers.update(get_auth_header(cls)) - headers["Comfy-Usage-Source"] = get_usage_source(cls) + headers = get_comfy_api_headers(cls) headers.update(endpoint.headers) node_id = get_node_id(cls) diff --git a/comfy_api_nodes/util/_helpers.py b/comfy_api_nodes/util/_helpers.py index 5d775e76a75e..2f718e068720 100644 --- a/comfy_api_nodes/util/_helpers.py +++ b/comfy_api_nodes/util/_helpers.py @@ -9,6 +9,7 @@ from yarl import URL from comfy.cli_args import args +from comfy.deploy_environment import get_deploy_environment from comfy.model_management import processing_interrupted from comfy_api.latest import IO @@ -40,6 +41,21 @@ def get_usage_source(node_cls: type[IO.ComfyNode]) -> str: return node_cls.hidden.comfy_usage_source or "comfyui-server" +def get_comfy_api_headers(node_cls: type[IO.ComfyNode]) -> dict[str, str]: + """Common headers (auth, deploy environment, usage source) for Comfy API requests. + + Centralizes the shared header set so every Comfy API request sends a consistent + set and new shared headers only need to be added in one place. Intended for + relative/cloud URLs resolved against ``default_base_url()``; because the result + includes auth, callers must not attach it to arbitrary absolute/presigned URLs. + """ + return { + **get_auth_header(node_cls), + "Comfy-Env": get_deploy_environment(), + "Comfy-Usage-Source": get_usage_source(node_cls), + } + + def default_base_url() -> str: return getattr(args, "comfy_api_base", "https://api.comfy.org") diff --git a/comfy_api_nodes/util/client.py b/comfy_api_nodes/util/client.py index 1cb113ddad55..adcde7bcbd49 100644 --- a/comfy_api_nodes/util/client.py +++ b/comfy_api_nodes/util/client.py @@ -19,14 +19,11 @@ from comfy_api.latest import IO from server import PromptServer -from comfy.deploy_environment import get_deploy_environment - from . import request_logger from ._helpers import ( default_base_url, - get_auth_header, + get_comfy_api_headers, get_node_id, - get_usage_source, is_processing_interrupted, sleep_with_interrupt, ) @@ -646,9 +643,7 @@ async def _monitor(stop_evt: asyncio.Event, start_ts: float): payload_headers = {"Accept": "*/*"} if expect_binary else {"Accept": "application/json"} if not parsed_url.scheme and not parsed_url.netloc: # is URL relative? - payload_headers.update(get_auth_header(cfg.node_cls)) - payload_headers["Comfy-Env"] = get_deploy_environment() - payload_headers["Comfy-Usage-Source"] = get_usage_source(cfg.node_cls) + payload_headers.update(get_comfy_api_headers(cfg.node_cls)) if cfg.endpoint.headers: payload_headers.update(cfg.endpoint.headers) diff --git a/comfy_api_nodes/util/download_helpers.py b/comfy_api_nodes/util/download_helpers.py index b2b2af8355a9..0ec3c6e6648d 100644 --- a/comfy_api_nodes/util/download_helpers.py +++ b/comfy_api_nodes/util/download_helpers.py @@ -17,8 +17,7 @@ from . import request_logger from ._helpers import ( default_base_url, - get_auth_header, - get_usage_source, + get_comfy_api_headers, is_processing_interrupted, sleep_with_interrupt, to_aiohttp_url, @@ -65,8 +64,7 @@ async def download_url_to_bytesio( if cls is None: raise ValueError("For relative 'cloud' paths, the `cls` parameter is required.") url = urljoin(default_base_url().rstrip("/") + "/", url.lstrip("/")) - headers = get_auth_header(cls) - headers["Comfy-Usage-Source"] = get_usage_source(cls) + headers = get_comfy_api_headers(cls) while True: attempt += 1 From ee2cdaee5fe35cc9007a993c4672253636bcba3c Mon Sep 17 00:00:00 2001 From: Robin Huang Date: Thu, 11 Jun 2026 12:30:29 -0700 Subject: [PATCH 3/3] Default usage source to comfyui-api for direct API calls --- comfy_api_nodes/util/_helpers.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/comfy_api_nodes/util/_helpers.py b/comfy_api_nodes/util/_helpers.py index 2f718e068720..83cf7b0013c9 100644 --- a/comfy_api_nodes/util/_helpers.py +++ b/comfy_api_nodes/util/_helpers.py @@ -37,8 +37,12 @@ def get_auth_header(node_cls: type[IO.ComfyNode]) -> dict[str, str]: def get_usage_source(node_cls: type[IO.ComfyNode]) -> str: - """Source of the prompt that triggered this API node, defaulting to this server itself.""" - return node_cls.hidden.comfy_usage_source or "comfyui-server" + """Source of the prompt that triggered this API node. + + Defaults to "comfyui-api" when the submitting client didn't identify itself, + i.e. a direct API call to this server. + """ + return node_cls.hidden.comfy_usage_source or "comfyui-api" def get_comfy_api_headers(node_cls: type[IO.ComfyNode]) -> dict[str, str]: