feat: api_key on submit/run for partner (API) node auth

.github/CODEOWNERS

@@ -0,0 +1,4 @@
+# Default owners for everything in the repo. A pull request that touches any
+# file requires an approving review from at least one member of these teams
+# (enforced by the "Requirements to Merge to Main" ruleset).
+* @Comfy-Org/comfy-cloud-team @Comfy-Org/core-engine-team

README.md

@@ -66,6 +66,24 @@ requests aimed at the configured `base_url`'s own origin — a server-returned
 follow-up link (`job.urls.self`/`cancel`/`events`, or a redirected asset
 download) pointing anywhere else never receives it.
 
+## Partner (API) node auth
+
+Workflows that use partner/API nodes (Gemini, etc.) need a Comfy API key to
+authenticate them. Pass it per submit with `api_key=`. This is **not** the same
+as the `api_key` you construct `Comfy` with: the constructor key authenticates
+*you* to the server, while this one authenticates the partner nodes *inside* the
+workflow (it is often the same `comfyui-…` key):
+
+```python
+job = client.run(wf, api_key="comfyui-...")
+# or drive it yourself:
+job = client.submit(wf, api_key="comfyui-...")
+```
+
+The SDK sends it once as `extra_data.api_key_comfy_org` alongside the workflow —
+one key authenticates every partner node in the graph. It is never logged or
+persisted by the SDK. Omit `api_key` and no `extra_data` is sent at all.
+
 ## Assets and `core/ASSET`
 
 `client.assets.from_file(...)` / `from_bytes(...)` / `from_stream(...)` /

scripts/check_public_repo_hygiene.py

@@ -83,6 +83,15 @@
     "comfy-api-proxy",
     "comfy-cloud-mcp-server",
 }
+# CODEOWNERS team handles (`@Comfy-Org/<team>`) are inherently public on a
+# public repo -- GitHub renders the CODEOWNERS owners to anyone who can see the
+# repo, so listing them here is not a leak. These mirror the sibling repos'
+# CODEOWNERS (e.g. comfy-api-proxy). An `@Comfy-Org/<team>` handle NOT in this
+# set is still flagged, so a genuinely-internal team reference surfaces.
+PUBLIC_COMFY_ORG_TEAMS = {
+    "comfy-cloud-team",
+    "core-engine-team",
+}
 REPO_REF_RE = re.compile(r"Comfy-Org/([A-Za-z0-9_.-]+)")
 
 
@@ -119,10 +128,19 @@ def _check_file(rel: Path) -> list[str]:
                 )
 
         for match in REPO_REF_RE.finditer(line):
-            repo = match.group(1)
-            if repo not in PUBLIC_COMFY_ORG_REPOS:
+            name = match.group(1)
+            # A leading `@` makes this a CODEOWNERS team handle, not a repo ref.
+            if match.start() > 0 and line[match.start() - 1] == "@":
+                if name not in PUBLIC_COMFY_ORG_TEAMS:
+                    findings.append(
+                        f"{rel}:{lineno}: reference to @Comfy-Org/{name}, a team not in the "
+                        "known-public allowlist (scripts/check_public_repo_hygiene.py) -- "
+                        "confirm it's public and add it, or remove the reference"
+                    )
+                continue
+            if name not in PUBLIC_COMFY_ORG_REPOS:
                 findings.append(
-                    f"{rel}:{lineno}: reference to Comfy-Org/{repo}, which is not in the "
+                    f"{rel}:{lineno}: reference to Comfy-Org/{name}, which is not in the "
                     "known-public allowlist (scripts/check_public_repo_hygiene.py) -- "
                     "confirm it's public and add it, or remove the reference"
                 )

spec/openapi.yaml

@@ -362,6 +362,18 @@ paths:
                   type: object
                   description: API-format workflow graph, verbatim.
                   additionalProperties: true
+                extra_data:
+                  type: object
+                  description: >-
+                    Optional, sibling of `workflow` (never nested inside it).
+                    Omit entirely when there is no key to send — an empty
+                    `extra_data` object is not sent by the SDK and should not
+                    be sent by other clients either.
+                  additionalProperties: false
+                  properties:
+                    api_key_comfy_org:
+                      type: string
+                      description: API key for partner (API) nodes.
       responses:
         '201':
           description: Job created and queued.

src/comfy_low/transport.py

@@ -302,17 +302,28 @@ def post_jobs(
         workflow: dict[str, Any],
         *,
         idempotency_key: str | None = None,
+        extra_data: dict[str, Any] | None = None,
         timeout: Any = _UNSET,
     ) -> Job:
-        """POST /api/v2/jobs."""
+        """POST /api/v2/jobs.
+
+        ``extra_data`` (e.g. ``{"api_key_comfy_org": "..."}``) is a sibling of
+        ``workflow`` in the body, never nested inside it. Omitted from the
+        request entirely when ``None`` — the server rejects an empty
+        ``extra_data`` object, and a caller with no partner key should never
+        send one.
+        """
         headers: dict[str, str] = {}
         if idempotency_key:
             headers["Idempotency-Key"] = idempotency_key
+        body: dict[str, Any] = {"workflow": workflow}
+        if extra_data:
+            body["extra_data"] = extra_data
         resp = self.raw_request(
             "POST",
             "/jobs",
             headers=headers,
-            json={"workflow": workflow},
+            json=body,
             timeout=timeout,
         )
         data = self._p.parse_or_raise(resp, (201,))
@@ -525,16 +536,21 @@ async def post_jobs(
         workflow: dict[str, Any],
         *,
         idempotency_key: str | None = None,
+        extra_data: dict[str, Any] | None = None,
         timeout: Any = _UNSET,
     ) -> Job:
+        """POST /api/v2/jobs — see the sync ``post_jobs`` for ``extra_data``."""
         headers: dict[str, str] = {}
         if idempotency_key:
             headers["Idempotency-Key"] = idempotency_key
+        body: dict[str, Any] = {"workflow": workflow}
+        if extra_data:
+            body["extra_data"] = extra_data
         resp = await self.raw_request(
             "POST",
             "/jobs",
             headers=headers,
-            json={"workflow": workflow},
+            json=body,
             timeout=timeout,
         )
         data = self._p.parse_or_raise(resp, (201,))

src/comfy_sdk/_core.py

@@ -23,6 +23,20 @@ def new_idempotency_key() -> str:
     return str(uuid.uuid4())
 
 
+def extra_data_for(api_key: str | None) -> dict[str, Any] | None:
+    """Build the wire ``extra_data`` object carrying the partner API key.
+
+    ``api_key`` authenticates partner (API) nodes (e.g. Gemini) embedded in a
+    workflow — it is unrelated to the client's own ``Authorization`` bearer
+    token. Returns ``None`` when no key is supplied, so callers omit
+    ``extra_data`` from the request entirely rather than sending an empty
+    object.
+    """
+    if not api_key:
+        return None
+    return {"api_key_comfy_org": api_key}
+
+
 def is_terminal(status: str) -> bool:
     return status in TERMINAL
 

src/comfy_sdk/client.py

@@ -5,9 +5,13 @@
 awaiting methods are duplicated; the rules (idempotency, 429 backoff, asset
 materialization, UI-format detection) live in ``_core`` and are called from both.
 
-Per-surface key behavior is inherited from ``comfy_low``: pass ``api_key`` for
-Comfy Cloud / serverless; leave it unset for a self-hosted proxy that has no auth
-(no credentials are then sent).
+Per-surface key behavior is inherited from ``comfy_low``: pass ``api_key`` to
+the constructor for Comfy Cloud / serverless; leave it unset for a self-hosted
+proxy that has no auth (no credentials are then sent). That constructor key is
+this client's own credential (sent as the ``Authorization`` bearer token) and
+is unrelated to the ``api_key`` accepted by ``submit``/``run``, which
+authenticates partner (API) nodes embedded in a workflow (e.g. Gemini) and is
+sent in the request body instead.
 """
 
 from __future__ import annotations
@@ -72,7 +76,13 @@ def _materialize(self, workflow: Workflow) -> dict[str, Any]:
             refs[id(h)] = h.as_reference()
         return _core.substitute_asset_handles(workflow.json, refs)
 
-    def submit(self, workflow: Workflow, *, idempotency_key: str | None = None) -> Job:
+    def submit(
+        self,
+        workflow: Workflow,
+        *,
+        api_key: str | None = None,
+        idempotency_key: str | None = None,
+    ) -> Job:
         """Submit a workflow. Retries ``queue_full`` with ``Retry-After``.
 
         Sends an auto-generated ``Idempotency-Key`` so the server rejects an
@@ -82,14 +92,21 @@ def submit(self, workflow: Workflow, *, idempotency_key: str | None = None) -> J
         idempotent, pass an explicit ``idempotency_key`` and reuse it. Note a
         reused key is *rejected*, not replayed: on reuse, catch the error and
         poll/list for the job the first attempt already created.
+
+        ``api_key`` authenticates partner (API) nodes embedded in the workflow
+        (e.g. Gemini) — unrelated to idempotency and unrelated to the bearer
+        token this client was constructed with. It is never persisted or
+        logged by the SDK, and is sent as ``extra_data.api_key_comfy_org``
+        only when supplied; omitted from the request entirely otherwise.
         """
         _guard_ui_format(workflow)
         graph = self._materialize(workflow)
         key = idempotency_key or _core.new_idempotency_key()
+        extra_data = _core.extra_data_for(api_key)
         deadline = time.monotonic() + _QUEUE_RETRY_BUDGET
         while True:
             try:
-                model = self._low.post_jobs(graph, idempotency_key=key)
+                model = self._low.post_jobs(graph, idempotency_key=key, extra_data=extra_data)
                 return Job(self._low, model)
             except ApiError as exc:
                 err = to_sdk_error(exc)
@@ -98,9 +115,15 @@ def submit(self, workflow: Workflow, *, idempotency_key: str | None = None) -> J
                     continue
                 raise err from exc
 
-    def run(self, workflow: Workflow, *, timeout: float | None = None) -> Job:
+    def run(
+        self,
+        workflow: Workflow,
+        *,
+        api_key: str | None = None,
+        timeout: float | None = None,
+    ) -> Job:
         """Submit, then poll to terminal (authoritative). Raises on failure."""
-        job = self.submit(workflow)
+        job = self.submit(workflow, api_key=api_key)
         return job.result() if timeout is None else _run_with_timeout(job, timeout)
 
 
@@ -146,16 +169,24 @@ async def _materialize(self, workflow: Workflow) -> dict[str, Any]:
             refs[id(h)] = await h.as_reference()
         return _core.substitute_asset_handles(workflow.json, refs)
 
-    async def submit(self, workflow: Workflow, *, idempotency_key: str | None = None) -> AsyncJob:
+    async def submit(
+        self,
+        workflow: Workflow,
+        *,
+        api_key: str | None = None,
+        idempotency_key: str | None = None,
+    ) -> AsyncJob:
+        """Mirrors :meth:`Comfy.submit` — see there for ``api_key`` details."""
         import asyncio
 
         _guard_ui_format(workflow)
         graph = await self._materialize(workflow)
         key = idempotency_key or _core.new_idempotency_key()
+        extra_data = _core.extra_data_for(api_key)
         deadline = time.monotonic() + _QUEUE_RETRY_BUDGET
         while True:
             try:
-                model = await self._low.post_jobs(graph, idempotency_key=key)
+                model = await self._low.post_jobs(graph, idempotency_key=key, extra_data=extra_data)
                 return AsyncJob(self._low, model)
             except ApiError as exc:
                 err = to_sdk_error(exc)
@@ -164,11 +195,17 @@ async def submit(self, workflow: Workflow, *, idempotency_key: str | None = None
                     continue
                 raise err from exc
 
-    async def run(self, workflow: Workflow, *, timeout: float | None = None) -> AsyncJob:
+    async def run(
+        self,
+        workflow: Workflow,
+        *,
+        api_key: str | None = None,
+        timeout: float | None = None,
+    ) -> AsyncJob:
         from ._core import SUCCESS
         from .exceptions import JobFailed
 
-        job = await self.submit(workflow)
+        job = await self.submit(workflow, api_key=api_key)
         if timeout is None:
             return await job.result()
         await job.wait(timeout=timeout)

tests/conftest.py

@@ -52,6 +52,9 @@ class ServerState:
     events_connect_count: int = 0
     submit_count: int = 0
     last_workflow: dict[str, Any] | None = None
+    # The full POST /jobs body (so a test can assert `extra_data` is present
+    # with the right value, or absent entirely, without touching last_workflow).
+    last_jobs_body: dict[str, Any] | None = None
     # Idempotency-Key -> job id of the first (accepted) request, so a reuse of
     # the same key can be detected and rejected (single-use, no replay).
     idempotency: dict[str, str] = field(default_factory=dict)
@@ -283,6 +286,7 @@ def _post_jobs(self) -> None:
             state.submit_count += 1
             body = json.loads(self._read_body() or b"{}")
             state.last_workflow = body.get("workflow")
+            state.last_jobs_body = body
             key = self.headers.get("Idempotency-Key")
 
             if key and key in state.idempotency:

tests/test_async.py

@@ -125,6 +125,23 @@ async def test_async_core_asset_substitution(server, tmp_path) -> None:
     assert ref["info"]["id"] == "asset_uploaded_01"
 
 
+async def test_async_submit_with_api_key_sends_extra_data(server) -> None:
+    # Mirrors the sync `test_submit_with_api_key_sends_extra_data_sibling_of_workflow`.
+    async with AsyncComfy(server.base_url) as client:
+        await client.submit(_wf(client), api_key="comfyui-secret-key")
+    body = server.state.last_jobs_body
+    assert body is not None
+    assert body["extra_data"] == {"api_key_comfy_org": "comfyui-secret-key"}
+
+
+async def test_async_submit_without_api_key_omits_extra_data(server) -> None:
+    async with AsyncComfy(server.base_url) as client:
+        await client.submit(_wf(client))
+    body = server.state.last_jobs_body
+    assert body is not None
+    assert "extra_data" not in body
+
+
 async def test_async_queue_full_retries_with_retry_after(server) -> None:
     server.state.queue_full_times = 2  # 429 twice, then 201
     async with AsyncComfy(server.base_url) as client:

tests/test_jobs.py

@@ -128,6 +128,46 @@ def test_authorized_when_key_present(server) -> None:
     assert job.id.startswith("job_")
 
 
+def test_submit_with_api_key_sends_extra_data_sibling_of_workflow(server) -> None:
+    # The partner-node API key must ride alongside `workflow` as `extra_data`,
+    # not nested inside it, and use the exact wire key `api_key_comfy_org`.
+    with Comfy(server.base_url) as client:
+        client.submit(_wf(client), api_key="comfyui-secret-key")
+    body = server.state.last_jobs_body
+    assert body is not None
+    assert body["extra_data"] == {"api_key_comfy_org": "comfyui-secret-key"}
+    assert "workflow" in body
+    assert "api_key_comfy_org" not in body["workflow"]  # sibling, never nested
+
+
+def test_submit_without_api_key_omits_extra_data_entirely(server) -> None:
+    # No key supplied -> no `extra_data` key at all (never an empty object).
+    with Comfy(server.base_url) as client:
+        client.submit(_wf(client))
+    body = server.state.last_jobs_body
+    assert body is not None
+    assert "extra_data" not in body
+
+
+def test_submit_with_empty_string_api_key_omits_extra_data(server) -> None:
+    # An empty string is "no key": no `extra_data` on the wire. Pinned so the
+    # TypeScript SDK stays in lockstep with this behavior.
+    with Comfy(server.base_url) as client:
+        client.submit(_wf(client), api_key="")
+    body = server.state.last_jobs_body
+    assert body is not None
+    assert "extra_data" not in body
+
+
+def test_run_forwards_api_key_to_submit(server) -> None:
+    # `run()` is submit-then-wait; the api_key must still reach the wire.
+    with Comfy(server.base_url) as client:
+        client.run(_wf(client), api_key="comfyui-secret-key")
+    body = server.state.last_jobs_body
+    assert body is not None
+    assert body["extra_data"] == {"api_key_comfy_org": "comfyui-secret-key"}
+
+
 def test_cancel_reaches_server_and_marks_canceling(server) -> None:
     # cancel() hits the server and reflects its `canceling` response, which is
     # deliberately NOT a terminal state.