From d65ac9629d974b9a003dd1f5d979ab831df5ee2c Mon Sep 17 00:00:00 2001 From: Alexander Piskun Date: Fri, 15 May 2026 10:08:32 +0300 Subject: [PATCH] feat: add X-API-Key header support Signed-off-by: Alexander Piskun --- comfy_cli/cmdline.py | 18 +++++++- comfy_cli/command/run.py | 25 ++++++++-- comfy_cli/tracking.py | 10 +++- tests/comfy_cli/command/test_command.py | 51 +++++++++++++++++++++ tests/comfy_cli/command/test_run.py | 61 +++++++++++++++++++++++++ tests/comfy_cli/test_tracking.py | 35 ++++++++++++++ 6 files changed, 193 insertions(+), 7 deletions(-) diff --git a/comfy_cli/cmdline.py b/comfy_cli/cmdline.py index 60cd6186..0cea8fbe 100644 --- a/comfy_cli/cmdline.py +++ b/comfy_cli/cmdline.py @@ -446,7 +446,23 @@ def run( int | None, typer.Option(help="The timeout in seconds for the workflow execution."), ] = 30, + api_key: Annotated[ + str | None, + typer.Option( + "--api-key", + envvar="COMFY_API_KEY", + help=( + "Comfy API key for API Nodes (Partner Nodes). " + "Embedded in the prompt body as extra_data.api_key_comfy_org on POST /prompt. " + "For scripting, prefer the COMFY_API_KEY environment variable so the secret " + "stays out of shell history." + ), + ), + ] = None, ): + if api_key: + api_key = api_key.strip() or None + config = ConfigManager() if host: @@ -470,7 +486,7 @@ def run( if not port: port = 8188 - run_inner.execute(workflow, host, port, wait, verbose, local_paths, timeout) + run_inner.execute(workflow, host, port, wait, verbose, local_paths, timeout, api_key=api_key) def validate_comfyui(_env_checker): diff --git a/comfy_cli/command/run.py b/comfy_cli/command/run.py index 4ae0f86a..f8b46857 100644 --- a/comfy_cli/command/run.py +++ b/comfy_cli/command/run.py @@ -65,7 +65,16 @@ def fetch_object_info(host: str, port: int, timeout: int) -> dict: raise typer.Exit(code=1) from e -def execute(workflow: str, host, port, wait=True, verbose=False, local_paths=False, timeout=30): +def execute( + workflow: str, + host, + port, + wait=True, + verbose=False, + local_paths=False, + timeout=30, + api_key: str | None = None, +): workflow_name = os.path.abspath(os.path.expanduser(workflow)) if not os.path.isfile(workflow): pprint( @@ -128,7 +137,7 @@ def execute(workflow: str, host, port, wait=True, verbose=False, local_paths=Fal else: print(f"Queuing workflow: {workflow_name}") - execution = WorkflowExecution(workflow, host, port, verbose, progress, local_paths, timeout) + execution = WorkflowExecution(workflow, host, port, verbose, progress, local_paths, timeout, api_key=api_key) try: if wait: @@ -182,7 +191,7 @@ def get_renderables(self): class WorkflowExecution: - def __init__(self, workflow, host, port, verbose, progress, local_paths, timeout=30): + def __init__(self, workflow, host, port, verbose, progress, local_paths, timeout=30, api_key: str | None = None): self.workflow = workflow self.host = host self.port = port @@ -201,14 +210,20 @@ def __init__(self, workflow, host, port, verbose, progress, local_paths, timeout self.prompt_id = None self.ws = None self.timeout = timeout + self.api_key = api_key def connect(self): self.ws = WebSocket() self.ws.connect(f"ws://{self.host}:{self.port}/ws?clientId={self.client_id}") def queue(self): - data = {"prompt": self.workflow, "client_id": self.client_id} - req = request.Request(f"http://{self.host}:{self.port}/prompt", json.dumps(data).encode("utf-8")) + data: dict = {"prompt": self.workflow, "client_id": self.client_id} + if 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"), + ) try: resp = request.urlopen(req) body = json.loads(resp.read()) diff --git a/comfy_cli/tracking.py b/comfy_cli/tracking.py index 713a3a90..570dee46 100644 --- a/comfy_cli/tracking.py +++ b/comfy_cli/tracking.py @@ -15,6 +15,10 @@ MIXPANEL_TOKEN = "93aeab8962b622d431ac19800ccc9f67" mp = Mixpanel(MIXPANEL_TOKEN) if MIXPANEL_TOKEN else None +# Kwargs whose values must never reach tracking system. +# The key is kept (with a redacted marker) so we can still see whether the option was supplied. +SENSITIVE_TRACKING_KEYS = frozenset({"api_key"}) + # Generate a unique tracing ID per command. config_manager = ConfigManager() cli_version = config_manager.get_cli_version() @@ -69,7 +73,11 @@ def wrapper(*args, **kwargs): # Copy kwargs to avoid mutating original dictionary # Remove context and ctx from the dictionary as they are not needed for tracking and not serializable. - filtered_kwargs = {k: v for k, v in kwargs.items() if k != "ctx" and k != "context"} + filtered_kwargs = { + k: ("" if v is not None else None) if k in SENSITIVE_TRACKING_KEYS else v + for k, v in kwargs.items() + if k != "ctx" and k != "context" + } logging.debug(f"Tracking command: {command_name} with arguments: {filtered_kwargs}") track_event(command_name, properties=filtered_kwargs) diff --git a/tests/comfy_cli/command/test_command.py b/tests/comfy_cli/command/test_command.py index 86486137..5c442941 100644 --- a/tests/comfy_cli/command/test_command.py +++ b/tests/comfy_cli/command/test_command.py @@ -59,3 +59,54 @@ def test_version(runner): result = runner.invoke(app, ["-v"]) assert result.exit_code == 0 assert "0.0.0" in result.stdout + + +@pytest.fixture +def mock_run_execute(): + with patch("comfy_cli.command.run.execute") as mock: + yield mock + + +def _write_workflow(tmp_path): + wf = tmp_path / "wf.json" + wf.write_text('{"1": {"class_type": "X", "inputs": {}}}') + return str(wf) + + +class TestRunApiKeyResolution: + """typer envvar resolution: --api-key + COMFY_API_KEY must reach run.execute().""" + + def test_envvar_is_picked_up(self, runner, mock_run_execute, tmp_path): + wf = _write_workflow(tmp_path) + result = runner.invoke(app, ["run", "--workflow", wf], env={"COMFY_API_KEY": "env-key-xyz"}) + assert result.exit_code == 0, result.output + assert mock_run_execute.call_args.kwargs["api_key"] == "env-key-xyz" + + def test_flag_overrides_envvar(self, runner, mock_run_execute, tmp_path): + wf = _write_workflow(tmp_path) + result = runner.invoke( + app, + ["run", "--workflow", wf, "--api-key", "flag-key-abc"], + env={"COMFY_API_KEY": "env-key-xyz"}, + ) + assert result.exit_code == 0, result.output + assert mock_run_execute.call_args.kwargs["api_key"] == "flag-key-abc" + + def test_absent_resolves_to_none(self, runner, mock_run_execute, tmp_path): + wf = _write_workflow(tmp_path) + # Explicit empty env to neutralize any host-level COMFY_API_KEY leak. + result = runner.invoke(app, ["run", "--workflow", wf], env={"COMFY_API_KEY": ""}) + assert result.exit_code == 0, result.output + assert mock_run_execute.call_args.kwargs["api_key"] is None + + def test_envvar_trailing_whitespace_is_stripped(self, runner, mock_run_execute, tmp_path): + wf = _write_workflow(tmp_path) + result = runner.invoke(app, ["run", "--workflow", wf], env={"COMFY_API_KEY": " sk-abc\n"}) + assert result.exit_code == 0, result.output + assert mock_run_execute.call_args.kwargs["api_key"] == "sk-abc" + + def test_whitespace_only_collapses_to_none(self, runner, mock_run_execute, tmp_path): + wf = _write_workflow(tmp_path) + result = runner.invoke(app, ["run", "--workflow", wf], env={"COMFY_API_KEY": " \n\t"}) + assert result.exit_code == 0, result.output + assert mock_run_execute.call_args.kwargs["api_key"] is None diff --git a/tests/comfy_cli/command/test_run.py b/tests/comfy_cli/command/test_run.py index 1a8b9a54..3b07db97 100644 --- a/tests/comfy_cli/command/test_run.py +++ b/tests/comfy_cli/command/test_run.py @@ -153,6 +153,51 @@ def test_invalid_json_exits_cleanly(self): assert exc_info.value.exit_code == 1 +class TestWorkflowExecutionAuth: + """X-API-Key is the credential the ComfyUI server forwards to Partner Nodes.""" + + def _make_exec(self, workflow, api_key=None): + progress = MagicMock() + progress.add_task.return_value = 0 + return WorkflowExecution( + workflow=workflow, + host="127.0.0.1", + port=8188, + verbose=False, + progress=progress, + local_paths=False, + timeout=30, + api_key=api_key, + ) + + def test_queue_embeds_api_key_in_extra_data(self, workflow): + ex = self._make_exec(workflow, api_key="sk-secret") + 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 body["extra_data"] == {"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") + 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("X-api-key") is None + + def test_queue_omits_extra_data_when_no_api_key(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} + + class TestWatchExecution: def test_successful_execution(self, mock_execution): prompt_id = "test-prompt" @@ -497,6 +542,22 @@ def test_ui_workflow_exits_cleanly_on_unexpected_converter_crash(self, ui_workfl assert exc_info.value.exit_code == 1 MockExec.assert_not_called() + def test_ui_workflow_plumbs_api_key_through_to_execution(self, ui_workflow_file): + with ( + patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), + patch("comfy_cli.command.run.fetch_object_info", return_value=self.OBJECT_INFO) as mock_fetch, + patch("comfy_cli.command.run.ExecutionProgress"), + patch("comfy_cli.command.run.WorkflowExecution") as MockExec, + ): + mock_exec = MagicMock() + MockExec.return_value = mock_exec + mock_exec.outputs = [] + + execute(ui_workflow_file, host="127.0.0.1", port=8188, wait=True, timeout=30, api_key="sk-test") + + mock_fetch.assert_called_once_with("127.0.0.1", 8188, 30) + assert MockExec.call_args.kwargs["api_key"] == "sk-test" + def test_ui_workflow_exits_when_conversion_yields_nothing(self): # All nodes are UI-only (Note/PrimitiveNode/Reroute/GetNode/SetNode) and # therefore stripped by the converter → execute() should bail before diff --git a/tests/comfy_cli/test_tracking.py b/tests/comfy_cli/test_tracking.py index 3b3502db..5cb7b020 100644 --- a/tests/comfy_cli/test_tracking.py +++ b/tests/comfy_cli/test_tracking.py @@ -63,6 +63,41 @@ def test_swallows_mixpanel_errors(self, tracking_module): tracking_module.mp.track.assert_called_once() +class TestTrackCommandRedaction: + """track_command must redact secret-bearing kwargs before they reach the tracking system.""" + + def test_api_key_value_is_redacted(self, tracking_module): + tracking_module.config_manager.set(constants.CONFIG_KEY_ENABLE_TRACKING, "True") + + @tracking_module.track_command() + def some_cmd(workflow, api_key=None): + return None + + some_cmd(workflow="wf.json", api_key="sk-supersecret") + + tracking_module.mp.track.assert_called_once() + _, kwargs = tracking_module.mp.track.call_args + props = kwargs["properties"] + assert props["api_key"] == "" + assert props["workflow"] == "wf.json" + assert "sk-supersecret" not in str(props) + + def test_api_key_none_stays_none(self, tracking_module): + # When the user didn't pass --api-key (or set $COMFY_API_KEY), we still + # want to be able to see in the analytics that it was absent — not a + # "" sentinel that would imply they did pass one. + tracking_module.config_manager.set(constants.CONFIG_KEY_ENABLE_TRACKING, "True") + + @tracking_module.track_command() + def some_cmd(workflow, api_key=None): + return None + + some_cmd(workflow="wf.json", api_key=None) + + _, kwargs = tracking_module.mp.track.call_args + assert kwargs["properties"]["api_key"] is None + + class TestInitTrackingRoundTrip: """End-to-end: init_tracking() writes the string "False"/"True", and track_event honors it.