Skip to content

Add Comfy-Usage-Source header to ComfyUI and cloud API requests #468

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bigcat88 merged 2 commits into from
Jun 11, 2026
+55 −7
Merged

Add Comfy-Usage-Source header to ComfyUI and cloud API requests #468

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ff2d9e7
Add Comfy-Usage-Source header to ComfyUI and cloud API requests
robinjhuang Jun 10, 2026
9a68f18
ci: bump codecov-action to v6.0.2 to fix GPG verification
robinjhuang Jun 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/workflows/pytest.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ jobs:
pytest --cov=comfy_cli --cov-report=xml .

- name: Upload coverage to Codecov
uses: codecov/codecov-action@57e3a136b779b570ffcdbf80b3bdc90e7fab3de2 # v6.0.0
uses: codecov/codecov-action@fb8b3582c8e4def4969c97caa2f19720cb33a72f # v6.0.2
Comment thread
bigcat88 marked this conversation as resolved.
with:
token: ${{ secrets.CODECOV_TOKEN }}
fail_ci_if_error: true
Expand Down
2 changes: 1 addition & 1 deletion comfy_cli/command/generate/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,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"}
headers = {"User-Agent": "comfy-cli/api", "Comfy-Env": "comfy-cli", "Comfy-Usage-Source": "comfy-cli"}
if api_key.startswith("comfyui-"):
headers["X-API-Key"] = api_key
else:
Expand Down
4 changes: 3 additions & 1 deletion comfy_cli/command/run.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -579,12 +579,14 @@ def connect(self):

def queue(self):
data: dict = {"prompt": self.workflow, "client_id": self.client_id}
data["extra_data"] = {"comfy_usage_source": "comfy-cli"}
if self.api_key:
data["extra_data"] = {"api_key_comfy_org": self.api_key}
data["extra_data"]["api_key_comfy_org"] = self.api_key
req = request.Request(
f"http://{self.host}:{self.port}/prompt",
json.dumps(data).encode("utf-8"),
)
req.add_header("Comfy-Usage-Source", "comfy-cli")
try:
resp = request.urlopen(req, timeout=self.timeout)
raw_body = resp.read()
Expand Down
19 changes: 15 additions & 4 deletions tests/comfy_cli/command/test_run.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -175,7 +175,7 @@ def test_queue_embeds_api_key_in_extra_data(self, workflow):
ex.queue()
req = mock_open.call_args[0][0]
body = json.loads(req.data)
assert body["extra_data"] == {"api_key_comfy_org": "sk-secret"}
assert body["extra_data"] == {"comfy_usage_source": "comfy-cli", "api_key_comfy_org": "sk-secret"}

def test_queue_does_not_send_x_api_key_header(self, workflow):
ex = self._make_exec(workflow, api_key="sk-secret")
Expand All @@ -185,15 +185,26 @@ def test_queue_does_not_send_x_api_key_header(self, workflow):
req = mock_open.call_args[0][0]
assert req.get_header("X-api-key") is None

def test_queue_omits_extra_data_when_no_api_key(self, workflow):
def test_queue_omits_api_key_when_not_set(self, workflow):
ex = self._make_exec(workflow)
with patch("comfy_cli.command.run.request.urlopen") as mock_open:
mock_open.return_value.read.return_value = json.dumps({"prompt_id": "abc"}).encode()
ex.queue()
req = mock_open.call_args[0][0]
body = json.loads(req.data)
assert "extra_data" not in body
assert body == {"prompt": workflow, "client_id": ex.client_id}
assert body == {
"prompt": workflow,
"client_id": ex.client_id,
"extra_data": {"comfy_usage_source": "comfy-cli"},
}

def test_queue_sends_usage_source_header(self, workflow):
ex = self._make_exec(workflow)
with patch("comfy_cli.command.run.request.urlopen") as mock_open:
mock_open.return_value.read.return_value = json.dumps({"prompt_id": "abc"}).encode()
ex.queue()
req = mock_open.call_args[0][0]
assert req.get_header("Comfy-usage-source") == "comfy-cli"
Comment on lines +201 to +207

@coderabbitai coderabbitai Bot Jun 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Verify that urllib.request.Request.get_header() is case-insensitive

python -c "
import urllib.request

req = urllib.request.Request('http://example.com')
req.add_header('Comfy-Usage-Source', 'test-value')

# Test various casings
print('Exact match:', req.get_header('Comfy-Usage-Source'))
print('Lowercase u:', req.get_header('Comfy-usage-source'))
print('All lowercase:', req.get_header('comfy-usage-source'))
print('All uppercase:', req.get_header('COMFY-USAGE-SOURCE'))

# Check if all variants return the same value
variants = [
    req.get_header('Comfy-Usage-Source'),
    req.get_header('Comfy-usage-source'),
    req.get_header('comfy-usage-source'),
    req.get_header('COMFY-USAGE-SOURCE')
]
if all(v == 'test-value' for v in variants):
    print('✓ get_header() is case-insensitive')
else:
    print('✗ get_header() is case-sensitive - test may fail')
"

Repository: Comfy-Org/comfy-cli

Length of output: 194

Align Comfy-* header casing in the test.

run.py sets "Comfy-Usage-Source", but the test asserts req.get_header("Comfy-usage-source"); urllib.request.Request.get_header() can require matching header-name casing (returns None otherwise), so this assertion is brittle. Update the test to use the exact same header name casing—case closed! 🎭

🧰 Tools
🪛 Pylint (4.0.5)

[convention] 201-201: Missing function or method docstring

(C0116)

[warning] 201-201: Redefining name 'workflow' from outer scope (line 21)

(W0621)

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/comfy_cli/command/test_run.py` around lines 201 - 207, The test
test_queue_sends_usage_source_header is asserting the wrong header casing;
update the assertion to match the exact header name used in run.py
("Comfy-Usage-Source") so it reads assert req.get_header("Comfy-Usage-Source")
== "comfy-cli" (or alternatively fetch headers case-insensitively from
req.header_items()), targeting the test function
test_queue_sends_usage_source_header and the code that calls
ex.queue()/urllib.request.Request to ensure casing matches.



class TestWatchExecution:
Expand Down
35 changes: 35 additions & 0 deletions uv.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading