feat(run): add --api-key / COMFY_API_KEY for API Nodes

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):

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())

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: ("<redacted>" 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)

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

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

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"] == "<redacted>"
+        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
+        # "<redacted>" 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.