Add Comfy-Usage-Source header to all CLI API requests

comfy_cli/cloud/oauth.py

@@ -54,7 +54,7 @@
     CLIENT_NAME,
     get_base_url,
 )
-from comfy_cli.http import NoRedirectHandler, build_http_only_opener
+from comfy_cli.http import USAGE_SOURCE_HEADERS, NoRedirectHandler, build_http_only_opener
 
 # ---------------------------------------------------------------------------
 # Error types — caller maps these to renderer.error(code=...) codes.
@@ -976,6 +976,8 @@ def _assert_https_or_loopback(url: str) -> None:
 
 def _send_and_parse(req: urllib.request.Request, *, timeout: float = _HTTP_TIMEOUT_S) -> dict:
     _assert_https_or_loopback(req.full_url)
+    for k, v in USAGE_SOURCE_HEADERS.items():
+        req.add_header(k, v)
     try:
         with _OAUTH_OPENER.open(req, timeout=timeout) as resp:
             raw = resp.read().decode("utf-8", errors="replace") or "{}"

comfy_cli/comfy_client.py

@@ -26,6 +26,7 @@
 
 from comfy_cli.caller import usage_source
 from comfy_cli.http import (
+    USAGE_SOURCE_HEADERS,
     NoRedirectHandler,
     ResponseTooLarge,
     build_http_only_opener,
@@ -262,9 +263,8 @@ def _request(
         data = json.dumps(body).encode("utf-8") if body is not None else None
         req = urllib.request.Request(url, data=data, method=method)
         req.add_header("Accept", "application/json")
-        # Usage-source attribution on every ComfyUI/cloud API request so the
-        # server can tell CLI-originated traffic apart from the web UI (#468).
-        req.add_header("Comfy-Usage-Source", "comfy-cli")
+        for k, v in USAGE_SOURCE_HEADERS.items():
+            req.add_header(k, v)
         if data is not None:
             req.add_header("Content-Type", "application/json")
         # Cloud auth: `target_auth_headers` owns the header selection for

comfy_cli/command/build.py

@@ -100,6 +100,7 @@
 )
 from comfy_cli.command.pack_scan import read_pyproject
 from comfy_cli.constants import SUPPORTED_PT_EXTENSIONS
+from comfy_cli.http import USAGE_SOURCE_HEADERS
 from comfy_cli.interaction import confirm, require_option
 from comfy_cli.output import get_renderer
 from comfy_cli.registry.api import sanitize_error_body
@@ -539,7 +540,7 @@ def detect_comfy_version_from_server(base_url: str) -> str | None:
     at the data root. Returns None if no server answers (fast-fail) or the field
     is absent."""
     try:
-        resp = requests.get(base_url.rstrip("/") + "/system_stats", timeout=2)
+        resp = requests.get(base_url.rstrip("/") + "/system_stats", headers=USAGE_SOURCE_HEADERS, timeout=2)
         resp.raise_for_status()
         version = resp.json().get("system", {}).get("comfyui_version")
     except (requests.RequestException, ValueError):

comfy_cli/command/cloud_http.py

@@ -36,7 +36,9 @@ def _authed_request(
     loosely to keep urllib out of the module's top-level imports."""
     import urllib.request
 
-    req = urllib.request.Request(url, data=data, method=method)
+    from comfy_cli.http import USAGE_SOURCE_HEADERS
+
+    req = urllib.request.Request(url, data=data, method=method, headers=USAGE_SOURCE_HEADERS)
     if target.api_key:
         req.add_header("X-API-Key", target.api_key)
     elif target.auth_token:

comfy_cli/command/code_search.py

@@ -15,6 +15,7 @@
 
 from comfy_cli import tracking
 from comfy_cli.caller import detect_caller
+from comfy_cli.http import USAGE_SOURCE_HEADERS
 
 app = typer.Typer()
 console = Console()
@@ -196,6 +197,7 @@ def _fetch_results(query: str, *, agent_origin: bool = False) -> dict:
         # otherwise expand far past the cap before the check fires (a
         # decompression bomb). Refuse compression altogether instead.
         headers={
+            **USAGE_SOURCE_HEADERS,
             "Accept-Encoding": "identity",
             **({"X-Comfy-Code-Search-Origin": "agent"} if agent_origin else {}),
         },

comfy_cli/command/generate/app.py

@@ -51,6 +51,7 @@
 from comfy_cli import constants, knowledge, tracking, ui
 from comfy_cli.command.generate import adapters, client, emit, output, poll, schema, spec, upload
 from comfy_cli.config_manager import ConfigManager
+from comfy_cli.http import USAGE_SOURCE_HEADERS
 from comfy_cli.output.renderer import Renderer, get_renderer
 from comfy_cli.output.sanitize import sanitize_markup
 
@@ -1069,7 +1070,7 @@ def _schema(extra_args: list[str]) -> None:
 
 def _fetch_spec(url: str) -> httpx.Response:
     with httpx.Client(timeout=30.0, follow_redirects=True) as cli:
-        r = cli.get(url, headers={"Comfy-Env": "comfy-cli", "User-Agent": "comfy-cli/api"})
+        r = cli.get(url, headers={"Comfy-Env": "comfy-cli", "User-Agent": "comfy-cli/api", **USAGE_SOURCE_HEADERS})
         r.raise_for_status()
         return r
 

comfy_cli/command/generate/client.py

@@ -16,6 +16,7 @@
 
 from comfy_cli.command.generate import spec
 from comfy_cli.command.generate.schema import FlagDef
+from comfy_cli.http import USAGE_SOURCE_HEADERS
 
 
 class ApiError(RuntimeError):
@@ -141,7 +142,7 @@ def _auth_headers(api_key: str, extra: dict[str, str] | None = None) -> dict[str
     #   - "comfyui-..." API keys → X-API-Key (validated by sha256 lookup)
     #   - Firebase ID tokens     → Authorization: Bearer (validated as a JWT)
     # See comfy-api server/middleware/authentication/comfy_firebase_auth.go.
-    headers = {"User-Agent": "comfy-cli/api", "Comfy-Env": "comfy-cli", "Comfy-Usage-Source": "comfy-cli"}
+    headers = {"User-Agent": "comfy-cli/api", "Comfy-Env": "comfy-cli", **USAGE_SOURCE_HEADERS}
     if api_key.startswith("comfyui-"):
         headers["X-API-Key"] = api_key
     else:

comfy_cli/command/job_watcher.py

@@ -22,6 +22,7 @@
 import typer
 
 from comfy_cli import execution_errors, jobs_state
+from comfy_cli.http import USAGE_SOURCE_HEADERS
 
 app = typer.Typer(hidden=True)
 
@@ -279,7 +280,9 @@ def _probe_local_server(host: str, port: int) -> str:
 
     try:
         # ``host`` arrives already bracketed from _resolve_watch_target.
-        resp = requests.get(f"http://{host}:{port}/history?max_items=1", timeout=_PROBE_TIMEOUT_S)
+        resp = requests.get(
+            f"http://{host}:{port}/history?max_items=1", headers=USAGE_SOURCE_HEADERS, timeout=_PROBE_TIMEOUT_S
+        )
     except requests.exceptions.Timeout:
         # Listed before ConnectionError: ConnectTimeout subclasses both, and a
         # timeout is never evidence that nothing is listening.

comfy_cli/command/jobs.py

@@ -45,7 +45,7 @@
 from comfy_cli.env_checker import check_comfy_server_running
 from comfy_cli.host_port import report_usage_error
 from comfy_cli.host_port import resolve_host_port as _resolve_host_port
-from comfy_cli.http import ResponseTooLarge, authed_urlopen, plain_urlopen, read_capped
+from comfy_cli.http import USAGE_SOURCE_HEADERS, ResponseTooLarge, authed_urlopen, plain_urlopen, read_capped
 from comfy_cli.output import get_renderer
 from comfy_cli.output.sanitize import sanitize, sanitize_markup
 from comfy_cli.utils import parse_rfc3339
@@ -167,7 +167,7 @@ def _http_get_json(url: str, *, timeout: float = 10.0) -> Any:
     covers its siblings that are not ``OSError`` at all (``IncompleteRead``,
     ``BadStatusLine``).
     """
-    req = urllib.request.Request(url)
+    req = urllib.request.Request(url, headers=USAGE_SOURCE_HEADERS)
     try:
         with plain_urlopen(req, timeout=timeout) as resp:
             return json.loads(read_capped(resp, url))
@@ -1728,7 +1728,10 @@ def _local_cancel(prompt_id: str, host: str, port: int) -> None:
     # 2. Remove from the pending queue (no-op if not pending).
     queue_body = json.dumps({"delete": [prompt_id]}).encode("utf-8")
     queue_req = urllib.request.Request(
-        f"{base}/queue", data=queue_body, method="POST", headers={"Content-Type": "application/json"}
+        f"{base}/queue",
+        data=queue_body,
+        method="POST",
+        headers={"Content-Type": "application/json", **USAGE_SOURCE_HEADERS},
     )
     queue_ok = True
     try:
@@ -1762,7 +1765,7 @@ def _local_cancel(prompt_id: str, host: str, port: int) -> None:
 
     interrupt_ok = True
     if is_running:
-        interrupt_req = urllib.request.Request(f"{base}/interrupt", method="POST")
+        interrupt_req = urllib.request.Request(f"{base}/interrupt", method="POST", headers=USAGE_SOURCE_HEADERS)
         try:
             with plain_urlopen(interrupt_req, timeout=10) as resp:
                 _ = resp.read()

comfy_cli/command/run/execution.py

@@ -41,7 +41,7 @@
 from comfy_cli import execution_errors, jobs_state
 from comfy_cli.caller import usage_source
 from comfy_cli.command.run.loader import _MAX_BODY_PREVIEW, _node_errors_to_list
-from comfy_cli.http import no_redirect_urlopen
+from comfy_cli.http import USAGE_SOURCE_HEADERS, no_redirect_urlopen
 from comfy_cli.output import get_renderer
 from comfy_cli.output import rprint as pprint
 from comfy_cli.output.sanitize import sanitize_markup
@@ -242,8 +242,9 @@ def queue(self):
             data["extra_data"][jobs_state.BORROWED_CLIENT_ID_KEY] = True
         else:
             data["extra_data"].pop(jobs_state.BORROWED_CLIENT_ID_KEY, None)
-        req = request.Request(f"http://{self.host}:{self.port}/prompt", json.dumps(data).encode("utf-8"))
-        req.add_header("Comfy-Usage-Source", "comfy-cli")
+        req = request.Request(
+            f"http://{self.host}:{self.port}/prompt", json.dumps(data).encode("utf-8"), headers=USAGE_SOURCE_HEADERS
+        )
         try:
             # No-redirect, not ``plain_urlopen``: ``extra_data`` can carry a
             # Comfy Org credential, so this submit gets the same refuse-a-30x

comfy_cli/command/run/preflight.py

@@ -11,11 +11,12 @@
 
 import json
 import urllib.error
+import urllib.request
 
 import typer
 
 from comfy_cli.command.run.loader import _MAX_BODY_PREVIEW
-from comfy_cli.http import plain_urlopen
+from comfy_cli.http import USAGE_SOURCE_HEADERS, plain_urlopen
 from comfy_cli.output import get_renderer
 from comfy_cli.output import rprint as pprint
 from comfy_cli.output.sanitize import sanitize_markup
@@ -43,7 +44,7 @@ def fetch_object_info(host, port, timeout):
     renderer = get_renderer()
     url = f"http://{host}:{port}/object_info"
     try:
-        with plain_urlopen(url, timeout=timeout) as resp:
+        with plain_urlopen(urllib.request.Request(url, headers=USAGE_SOURCE_HEADERS), timeout=timeout) as resp:
             body = resp.read(_MAX_OBJECT_INFO_BYTES)
     except urllib.error.HTTPError as e:
         body_text = e.read(_MAX_OBJECT_INFO_BYTES).decode("utf-8", errors="replace").strip()

comfy_cli/command/transfer.py

@@ -30,7 +30,7 @@
 from comfy_cli.comfy_client import Client, Unauthenticated, extract_output_entries
 from comfy_cli.credentials import DataPlaneToken
 from comfy_cli.host_port import report_usage_error, validate_host
-from comfy_cli.http import NoRedirectHandler, build_http_only_opener
+from comfy_cli.http import USAGE_SOURCE_HEADERS, NoRedirectHandler, build_http_only_opener
 from comfy_cli.http import target_auth_headers as _auth_headers
 from comfy_cli.output import get_renderer
 from comfy_cli.output import rprint as pprint
@@ -393,7 +393,7 @@ def _upload_file(path: Path, target: Any, *, overwrite: bool) -> dict:
     url = target.url("upload/image")
     req = urllib.request.Request(url, data=body, method="POST")
     req.add_header("Content-Type", f"multipart/form-data; boundary={boundary}")
-    for hdr, val in _auth_headers(target).items():
+    for hdr, val in {**USAGE_SOURCE_HEADERS, **_auth_headers(target)}.items():
         req.add_header(hdr, val)
 
     with _TRANSFER_OPENER.open(req, timeout=_UPLOAD_TIMEOUT_S) as resp:
@@ -946,7 +946,7 @@ def execute_download(
     dest = Path(out_dir or _default_out_dir())
     dest.mkdir(parents=True, exist_ok=True)
 
-    auth_hdrs = _auth_headers(target)
+    auth_hdrs = {**USAGE_SOURCE_HEADERS, **_auth_headers(target)}
     saved_files: list[dict[str, Any]] = []
     saved_paths: list[str] = []
     short_id = prompt_id[:8]

comfy_cli/cql/engine.py

@@ -27,7 +27,7 @@
 from typing import Any
 
 from comfy_cli.cql._net import is_loopback_host
-from comfy_cli.http import NoRedirectHandler, build_http_only_opener
+from comfy_cli.http import USAGE_SOURCE_HEADERS, NoRedirectHandler, build_http_only_opener
 
 # ---------------------------------------------------------------------------
 # Types — mirrors nodegraph/types.go
@@ -3664,7 +3664,7 @@ def _load_from_target(*, mode: str = "local", host: str | None = None, port: int
                 f"in local mode (potential SSRF). Use --where cloud for remote targets."
             )
 
-    req = urllib.request.Request(url)
+    req = urllib.request.Request(url, headers=USAGE_SOURCE_HEADERS)
     req.add_header("Accept", "application/json")
 
     # Auth headers (cloud only — local has no auth)

comfy_cli/deploy_assets.py

@@ -19,6 +19,7 @@
 from comfy_cli.deploy_api_errors import DeployAPIError, assert_safe_deploy_url
 from comfy_cli.hashing import blake3_file
 from comfy_cli.http import (
+    USAGE_SOURCE_HEADERS,
     ResponseTooLarge,
     no_redirect_urlopen,
     read_capped,
@@ -132,7 +133,7 @@ def _multipart_request(
     upload = urllib.request.Request(target.url("assets"), data=body, method="POST")
     upload.add_header("Content-Type", f"multipart/form-data; boundary={boundary}")
     upload.add_header("Content-Length", str(body.content_length))
-    for name, value in target_auth_headers(target).items():
+    for name, value in {**USAGE_SOURCE_HEADERS, **target_auth_headers(target)}.items():
         upload.add_header(name, value)
     return upload, file_size
 

comfy_cli/env_checker.py

@@ -108,8 +108,12 @@ def check_comfy_server_running(port=8188, host="localhost", timeout: float = 5.0
     # actually probing the server, not for every CLI invocation.
     import requests
 
+    from comfy_cli.http import USAGE_SOURCE_HEADERS
+
     try:
-        response = requests.get(f"http://{_bracket_host(host)}:{port}/history", timeout=timeout)
+        response = requests.get(
+            f"http://{_bracket_host(host)}:{port}/history", headers=USAGE_SOURCE_HEADERS, timeout=timeout
+        )
         return response.status_code == 200
     except requests.exceptions.RequestException:
         return False

comfy_cli/http.py

@@ -22,6 +22,12 @@
 # capping a legitimately long transfer.
 DOWNLOAD_TIMEOUT = (10.0, 60.0)
 
+# Usage-source attribution, sent on every request to a Comfy-owned API (local
+# or cloud ComfyUI, the builder, deploy, the registry, sign-in) so the server
+# can tell CLI-originated traffic apart from the web UI (#468). Third-party
+# hosts (GitHub, PyPI, model hosts, pre-signed storage URLs) never get it.
+USAGE_SOURCE_HEADERS: Final[dict[str, str]] = {"Comfy-Usage-Source": "comfy-cli"}
+
 _LOOPBACK_HOSTS = {"localhost", "127.0.0.1", "::1", "[::1]"}
 
 
@@ -353,7 +359,7 @@ def build_authed_request(
     Target can't carry a stray credential into this path either.
     """
     req = urllib.request.Request(url, data=data, method=method)
-    for k, v in target_auth_headers(target).items():
+    for k, v in {**USAGE_SOURCE_HEADERS, **target_auth_headers(target)}.items():
         req.add_header(k, v)
     if content_type:
         req.add_header("Content-Type", content_type)
@@ -468,7 +474,7 @@ def request_json(
         assert_safe_url(url)
     data = encode_json_body(body) if body is not None else None
     req = urllib.request.Request(url, data=data, method=method)
-    for k, v in (headers or {}).items():
+    for k, v in {**USAGE_SOURCE_HEADERS, **(headers or {})}.items():
         req.add_header(k, v)
     for k, v in auth_headers.items():
         req.add_header(k, v)

comfy_cli/registry/api.py

@@ -2,7 +2,7 @@
 import logging
 import os
 
-from comfy_cli.http import DEFAULT_HTTP_TIMEOUT
+from comfy_cli.http import DEFAULT_HTTP_TIMEOUT, USAGE_SOURCE_HEADERS
 
 # Reduced global imports from comfy_cli.registry
 from comfy_cli.registry.types import (
@@ -130,7 +130,7 @@ def publish_node_version(
         if changelog:
             request_body["node_version"]["changelog"] = changelog
         url = f"{self.base_url}/publishers/{node_config.tool_comfy.publisher_id}/nodes/{node_config.project.name}/versions"
-        headers = {"Content-Type": "application/json"}
+        headers = {"Content-Type": "application/json", **USAGE_SOURCE_HEADERS}
         body = request_body
 
         # Imported lazily: requests costs ~30ms to import and this module is
@@ -165,7 +165,7 @@ def list_all_nodes(self):
         import requests  # deferred; see publish_node_version
 
         url = f"{self.base_url}/nodes"
-        response = requests.get(url, timeout=DEFAULT_HTTP_TIMEOUT)
+        response = requests.get(url, headers=USAGE_SOURCE_HEADERS, timeout=DEFAULT_HTTP_TIMEOUT)
         if response.status_code == 200:
             raw_nodes = response.json()["nodes"]
             return [map_node_to_node_class(node) for node in raw_nodes]
@@ -197,7 +197,7 @@ def install_node(self, node_id, version=None):
 
         # A stalled/blackholed registry must not hang callers indefinitely.
         # A Timeout surfaces as a RequestException for callers to catch.
-        response = requests.get(url, timeout=DEFAULT_HTTP_TIMEOUT)
+        response = requests.get(url, headers=USAGE_SOURCE_HEADERS, timeout=DEFAULT_HTTP_TIMEOUT)
         if response.status_code == 200:
             # Convert the API response to a NodeVersion object
             logging.debug(f"RegistryAPI install_node response: {response.json()}")
@@ -228,7 +228,7 @@ def get_node(self, node_id):
 
         url = f"{self.base_url}/nodes/{node_id}"
         # Same rationale as install_node: a stalled registry must not hang callers.
-        response = requests.get(url, timeout=DEFAULT_HTTP_TIMEOUT)
+        response = requests.get(url, headers=USAGE_SOURCE_HEADERS, timeout=DEFAULT_HTTP_TIMEOUT)
         if response.status_code == 200:
             logging.debug(f"RegistryAPI get_node response: {response.json()}")
             return map_node_to_node_class(response.json())

tests/comfy_cli/command/test_build.py

@@ -121,12 +121,19 @@ def raise_for_status(self):
         def json(self):
             return {"system": {"comfyui_version": "0.3.55"}}
 
-    monkeypatch.setattr("comfy_cli.command.build.requests.get", lambda url, timeout: _Resp())
+    seen = {}
+
+    def fake_get(url, headers, timeout):
+        seen["headers"] = headers
+        return _Resp()
+
+    monkeypatch.setattr("comfy_cli.command.build.requests.get", fake_get)
     assert build.detect_comfy_version_from_server("http://127.0.0.1:8188") == "0.3.55"
+    assert seen["headers"] == {"Comfy-Usage-Source": "comfy-cli"}
 
 
 def test_detect_comfy_version_from_server_none_when_down(monkeypatch):
-    def boom(url, timeout):
+    def boom(url, headers, timeout):
         raise requests.ConnectionError("refused")
 
     monkeypatch.setattr("comfy_cli.command.build.requests.get", boom)