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.
This commit is contained in:
Robin Huang 2026-06-10 14:17:46 -07:00
parent 6d18f4adac
commit d3878edb01
7 changed files with 29 additions and 1 deletions

View File

@ -1400,7 +1400,8 @@ class V3Data(TypedDict):
class HiddenHolder: class HiddenHolder:
def __init__(self, unique_id: str, prompt: Any, def __init__(self, unique_id: str, prompt: Any,
extra_pnginfo: Any, dynprompt: 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 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).""" """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 self.prompt = prompt
@ -1413,6 +1414,8 @@ class HiddenHolder:
"""AUTH_TOKEN_COMFY_ORG is a token acquired from signing into a ComfyOrg account on frontend.""" """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 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.""" """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): def __getattr__(self, key: str):
'''If hidden variable not found, return None.''' '''If hidden variable not found, return None.'''
@ -1429,6 +1432,7 @@ class HiddenHolder:
dynprompt=d.get(Hidden.dynprompt, None), dynprompt=d.get(Hidden.dynprompt, None),
auth_token_comfy_org=d.get(Hidden.auth_token_comfy_org, 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), api_key_comfy_org=d.get(Hidden.api_key_comfy_org, None),
comfy_usage_source=d.get(Hidden.comfy_usage_source, None),
) )
@classmethod @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.""" """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 = "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.""" """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 @dataclass
@ -1654,6 +1660,8 @@ class Schema:
self.hidden.append(Hidden.auth_token_comfy_org) self.hidden.append(Hidden.auth_token_comfy_org)
if Hidden.api_key_comfy_org not in self.hidden: if Hidden.api_key_comfy_org not in self.hidden:
self.hidden.append(Hidden.api_key_comfy_org) 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 is an output_node, will need prompt and extra_pnginfo
if self.is_output_node: if self.is_output_node:
if Hidden.prompt not in self.hidden: if Hidden.prompt not in self.hidden:

View File

@ -18,6 +18,7 @@ from comfy_api_nodes.util._helpers import (
default_base_url, default_base_url,
get_auth_header, get_auth_header,
get_node_id, get_node_id,
get_usage_source,
is_processing_interrupted, is_processing_interrupted,
) )
from comfy_api_nodes.util.common_exceptions import ProcessingInterrupted from comfy_api_nodes.util.common_exceptions import ProcessingInterrupted
@ -176,6 +177,7 @@ async def _stream_sonilo_music(
headers: dict[str, str] = {} headers: dict[str, str] = {}
headers.update(get_auth_header(cls)) headers.update(get_auth_header(cls))
headers["Comfy-Usage-Source"] = get_usage_source(cls)
headers.update(endpoint.headers) headers.update(endpoint.headers)
node_id = get_node_id(cls) node_id = get_node_id(cls)

View File

@ -35,6 +35,11 @@ def get_auth_header(node_cls: type[IO.ComfyNode]) -> dict[str, str]:
return {} 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: def default_base_url() -> str:
return getattr(args, "comfy_api_base", "https://api.comfy.org") return getattr(args, "comfy_api_base", "https://api.comfy.org")

View File

@ -26,6 +26,7 @@ from ._helpers import (
default_base_url, default_base_url,
get_auth_header, get_auth_header,
get_node_id, get_node_id,
get_usage_source,
is_processing_interrupted, is_processing_interrupted,
sleep_with_interrupt, sleep_with_interrupt,
) )
@ -647,6 +648,7 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
if not parsed_url.scheme and not parsed_url.netloc: # is URL relative? if not parsed_url.scheme and not parsed_url.netloc: # is URL relative?
payload_headers.update(get_auth_header(cfg.node_cls)) payload_headers.update(get_auth_header(cfg.node_cls))
payload_headers["Comfy-Env"] = get_deploy_environment() payload_headers["Comfy-Env"] = get_deploy_environment()
payload_headers["Comfy-Usage-Source"] = get_usage_source(cfg.node_cls)
if cfg.endpoint.headers: if cfg.endpoint.headers:
payload_headers.update(cfg.endpoint.headers) payload_headers.update(cfg.endpoint.headers)

View File

@ -18,6 +18,7 @@ from . import request_logger
from ._helpers import ( from ._helpers import (
default_base_url, default_base_url,
get_auth_header, get_auth_header,
get_usage_source,
is_processing_interrupted, is_processing_interrupted,
sleep_with_interrupt, sleep_with_interrupt,
to_aiohttp_url, to_aiohttp_url,
@ -65,6 +66,7 @@ async def download_url_to_bytesio(
raise ValueError("For relative 'cloud' paths, the `cls` parameter is required.") raise ValueError("For relative 'cloud' paths, the `cls` parameter is required.")
url = urljoin(default_base_url().rstrip("/") + "/", url.lstrip("/")) url = urljoin(default_base_url().rstrip("/") + "/", url.lstrip("/"))
headers = get_auth_header(cls) headers = get_auth_header(cls)
headers["Comfy-Usage-Source"] = get_usage_source(cls)
while True: while True:
attempt += 1 attempt += 1

View File

@ -199,6 +199,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) 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: 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) 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: else:
if "hidden" in valid_inputs: if "hidden" in valid_inputs:
h = valid_inputs["hidden"] h = valid_inputs["hidden"]
@ -215,6 +217,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)] input_data_all[x] = [extra_data.get("auth_token_comfy_org", None)]
if h[x] == "API_KEY_COMFY_ORG": if h[x] == "API_KEY_COMFY_ORG":
input_data_all[x] = [extra_data.get("api_key_comfy_org", None)] 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 v3_data["hidden_inputs"] = hidden_inputs_v3
return input_data_all, missing_keys, v3_data return input_data_all, missing_keys, v3_data

View File

@ -957,6 +957,11 @@ class PromptServer():
if "client_id" in json_data: if "client_id" in json_data:
extra_data["client_id"] = json_data["client_id"] 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]: if valid[0]:
outputs_to_execute = valid[2] outputs_to_execute = valid[2]
sensitive = {} sensitive = {}