diff --git a/CHANGELOG.md b/CHANGELOG.md index 4662bda..086a580 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -114,6 +114,43 @@ notes for each version. ### Changed +- **Breaking (wire): `client.models.run` now posts to Comfy Router.** It sends + `POST {COMFY_ROUTER_BASE_URL}/v1/models/{provider}/{model}` — the route + `spec/router-openapi.yaml` declares as `runRouterModel` — with the partner + model's **own native JSON input** as the body, forwarded to the provider + unchanged. It previously posted `{COMFY_BASE_URL}/api/v2/models/run` with a + `{"model": ..., "arguments": {...}}` envelope, which nothing serves: the + `/api/v2` surface is jobs and assets, and the model-ID-addressed invocation + routes are Router's. The Python method signature is unchanged + (`run(model, arguments, *, idempotency_key=None, timeout=...)`), the result is + still the provider's payload returned as-is, and the `Idempotency-Key` and + retry behaviour are unchanged — what moved is the URL and the body shape. + **Anyone who pointed `COMFY_BASE_URL` at a Router host to make model runs work + must now point `COMFY_ROUTER_BASE_URL` there instead**, and set + `COMFY_BASE_URL` back at their v2 deployment (or unset it for Comfy Cloud). +- The `model` argument to `client.models.run` is now the canonical + `{provider}/{model}` id, because it *is* the two path segments the route is + addressed by. Exactly two non-empty segments are accepted; a one-segment id, a + three-segment `{provider}/{model}/{variant}` id (that form is not addressable + on this route yet), and any `.`/`..` segment now raise `ValueError` locally + before a request is made, and a non-string raises `TypeError`. Each segment is + percent-encoded into exactly one path segment. Previously any string was + accepted and travelled in the body. Mirrors the TypeScript SDK's + `parseModelId`, so the two SDKs accept and reject the same ids. +- `COMFY_ROUTER_BASE_URL` (default `https://api.comfy.org`) selects the Router + deployment, deliberately separate from `COMFY_BASE_URL` — one variable + pointed at both surfaces would send jobs to Router or model runs to the v2 + API. Same validation and read-per-construction rules as `COMFY_BASE_URL`, and + the same name as the TypeScript SDK's. `COMFY_ROUTER_BASE_URL` and + `ROUTER_BASE_URL_ENV_VAR` are exported from `comfy_sdk`. +- `client.models.base_url` (and the namespace's `repr`) now reports the Router + base URL rather than the client's `COMFY_BASE_URL`, since that is the host + model runs actually reach. The client's own `base_url` is unchanged. +- The client's API key is attached to *both* of its configured origins — the + `COMFY_BASE_URL` deployment and the `COMFY_ROUTER_BASE_URL` Router — and to no + third origin. Both are set by the caller (a constant or an environment + variable), never by a server response, so a server-returned absolute + follow-up link pointing anywhere else still receives no credential. - A client targeting a deployment named by `COMFY_BASE_URL` is unchanged: with no key resolved it is still built without one and still sends no credentials, which is what a self-hosted ComfyUI behind the API proxy needs. Only the Comfy diff --git a/README.md b/README.md index 6b59c01..ad05441 100644 --- a/README.md +++ b/README.md @@ -185,6 +185,21 @@ export COMFY_BASE_URL="http://127.0.0.1:8189" # self-hosted proxy It is read each time a client is constructed, must be an `http(s)` URL, and an unset or blank value (including whitespace-only) means Comfy Cloud. +`COMFY_BASE_URL` selects the **jobs and assets** surface. `client.models` talks +to a different one — Comfy Router, `https://api.comfy.org` — and follows its own +variable, `COMFY_ROUTER_BASE_URL`, with the same rules (read per construction, +`http(s)` only, blank means the default): + +```bash +export COMFY_ROUTER_BASE_URL="https://api.comfy.org" # the default; set it to redirect model runs +``` + +Two variables rather than one because they are genuinely two hosts: the v2 +surface serves `/api/v2/jobs` and `/api/v2/assets`, Router serves +`/v1/models/{provider}/{model}`, and neither serves the other's routes. Your +`COMFY_API_KEY` is the credential for both — the SDK attaches it to those two +configured origins and to no third one. + Upgrading from an earlier version: `Comfy("", "")` becomes `Comfy(api_key="")` with `COMFY_BASE_URL` set. `api_key` is keyword-only, so the old positional call raises `TypeError` rather than reading a URL as a @@ -362,27 +377,51 @@ Model operations live in a namespace on the client you already constructed — ```python client = Comfy(api_key="comfyui-...") -client.models.base_url # the client's own base URL, where model requests go +client.models.base_url # Comfy Router's base URL, where model runs go client.models.timeout # the client's own HTTP timeout ``` The namespace is bound to that client's transport, so it uses the client's -credentials, base URL, connection pool and timeout, and a configuration change -made on the client afterwards applies through `models` as well — there is no -second set of settings to keep in sync. `AsyncComfy` carries the same `models` -namespace, and nothing extra is imported or constructed for it: +credentials, connection pool and timeout, and a configuration change made on +the client afterwards applies through `models` as well — there is no second set +of settings to keep in sync. `AsyncComfy` carries the same `models` namespace, +and nothing extra is imported or constructed for it: `from comfy_sdk import Comfy` stays the only entry point. -`base_url` and `timeout` are a read-only view of that shared configuration; -model operations are added to this namespace as they land. +The one setting it does **not** share is the target host: `client.models.base_url` +reports `COMFY_ROUTER_BASE_URL` (`https://api.comfy.org` by default), not the +client's `COMFY_BASE_URL`. See +[Targeting another deployment](#targeting-another-deployment) for why they are +two variables. + +`base_url` and `timeout` are a read-only view of that configuration; model +operations are added to this namespace as they land. ### `models.run` — one call, one result ```python -result = client.models.run("acme/flux/dev", {"prompt": "a cat", "steps": 4}) +result = client.models.run("fal-ai/flux-pro", {"prompt": "a cat", "steps": 4}) result["images"][0]["url"] ``` +That call is `POST https://api.comfy.org/v1/models/fal-ai/flux-pro` with +`{"prompt": "a cat", "steps": 4}` as the body. + +Three things follow from that, and they are the whole contract of this method: + +- **It targets Comfy Router** — `https://api.comfy.org`, redirected by + `COMFY_ROUTER_BASE_URL`, not by `COMFY_BASE_URL`. Your API key goes with it. +- **The first argument is the model's canonical id**, `{provider}/{model}` — + exactly the two segments that address the route, and exactly what Router's + model catalog lists. It must be two non-empty segments: `"fal-ai"` alone, + `"fal-ai/flux-pro/fp8"` (the three-segment variant form, which this route does + not take yet), and anything containing a `.` or `..` segment all raise + `ValueError` locally, before any request. A non-string raises `TypeError`. +- **The second argument is the model's own native input**, forwarded to the + provider unchanged. There is no Comfy-shaped envelope around it: whatever the + partner's own API documents as the request body is what you pass here, so you + can move between the partner's API and Router by changing the host. + `run` returns when the generation is **complete**. There is no submit step and nothing to poll: where the platform has to submit-and-poll an upstream provider, that happens server side inside this one call. The value you get back @@ -393,7 +432,7 @@ The awaitable form is the **async client**, not a differently-named method: ```python async with AsyncComfy(api_key="comfyui-...") as client: - result = await client.models.run("acme/flux/dev", {"prompt": "a cat"}) + result = await client.models.run("fal-ai/flux-pro", {"prompt": "a cat"}) ``` There is no `run_async()`, and there will not be one — one operation, one name, @@ -467,10 +506,11 @@ arrives at Comfy's own ten-minute deadline, so a single-window budget would already be spent when it lands and the collect attempt it exists for would never start. Nothing else pays for that room. -A note on what the default trades: `POST /models/run` is in neither vendored -spec, so a deployment may apply the v2 rule instead and keep the key claimed -across the `504`. There the collect resend comes back `422 -idempotency_key_reuse` in place of the real `504`. Set +A note on what the default trades: the collect rule is +`spec/router-openapi.yaml`'s, so it binds Comfy Router — but +`COMFY_ROUTER_BASE_URL` can name a deployment that applies the v2 rule instead +and keeps the key claimed across the `504`. There the collect resend comes back +`422 idempotency_key_reuse` in place of the real `504`. Set `retry_collectable=False` on such a deployment. Other 5xx responses and client-side timeouts are the cases left out by default, diff --git a/scripts/check_drift.py b/scripts/check_drift.py index f96c951..611491d 100755 --- a/scripts/check_drift.py +++ b/scripts/check_drift.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 """Fail if the committed tree drifts from the vendored specs. -Two checks, one job: +Three checks, one job: 1. **Models vs ``spec/openapi.yaml``.** Regenerates the models into a temp file and diffs against the committed one, so a spec edit without a regen (or a @@ -14,8 +14,16 @@ for, which would reach callers as an untyped ``RouterError``. This compares the spec's ``x-comfy-error-types`` list against ``ROUTER_ERROR_TYPES``, which is what makes the next vendored Router sync a real diff review. +3. **The bound model-run route vs ``spec/router-openapi.yaml``.** Same reason, + different artifact: ``comfy_low.transport`` posts a model run to a + hand-written path constant and a hand-written host constant. The spec + declares both -- the path whose ``post.operationId`` is ``runRouterModel``, + and ``servers[0].url``. A sync that *moves* the route (the ``/v1`` -> ``/v2`` + move already on the roadmap) while those constants stay put would leave the + SDK posting to a route the contract no longer declares, with nothing else in + CI noticing. -``tests/test_router_spec_contract.py`` asserts the same thing from the test +``tests/test_router_spec_contract.py`` asserts the same things from the test suite. Both exist on purpose: the suite is where a contributor sees it, and this script is the job that fails a spec-only PR that never ran pytest. """ @@ -117,6 +125,99 @@ def _declared_router_error_types() -> list[str]: return values +def _declared_run_route() -> tuple[str, str]: + """The spec's ``(runRouterModel path, servers[0].url)``. + + Same failure policy as :func:`_declared_router_error_types`: every way the + file can be unusable becomes a ``ValueError`` with a sentence someone can + act on, rather than a traceback that reads like a bug in the checker. + """ + try: + import yaml + except ImportError as exc: # pragma: no cover - depends on the install extra + raise ValueError( + f"PyYAML is not installed, so {ROUTER_SPEC.name} cannot be read " + "(pip install -e '.[dev]')" + ) from exc + + try: + doc = yaml.safe_load(ROUTER_SPEC.read_text(encoding="utf-8")) + except OSError as exc: + raise ValueError(f"{ROUTER_SPEC.name} could not be read: {exc}") from exc + except yaml.YAMLError as exc: + raise ValueError(f"{ROUTER_SPEC.name} is not valid YAML: {exc}") from exc + + if not isinstance(doc, dict): + raise ValueError(f"{ROUTER_SPEC.name} is not a mapping at the top level") + paths = doc.get("paths") + if not isinstance(paths, dict): + raise ValueError(f"{ROUTER_SPEC.name} has no paths object") + # Searched by operationId rather than looked up by the path we expect: a + # lookup would silently find nothing the day the path moves, which is the + # one day this check exists for. + declared = [ + path + for path, item in paths.items() + if isinstance(item, dict) + and isinstance(item.get("post"), dict) + and item["post"].get("operationId") == "runRouterModel" + ] + if len(declared) != 1: + raise ValueError( + f"{ROUTER_SPEC.name} declares {len(declared)} paths with " + f"post.operationId 'runRouterModel' (expected exactly 1): {declared}" + ) + servers = doc.get("servers") + if not isinstance(servers, list) or not servers or not isinstance(servers[0], dict): + raise ValueError(f"{ROUTER_SPEC.name} has no servers[0]") + host = servers[0].get("url") + if not isinstance(host, str) or not host: + raise ValueError(f"{ROUTER_SPEC.name}'s servers[0].url is not a non-empty string") + return declared[0], host + + +def _check_router_run_route() -> int: + if not ROUTER_SPEC.exists(): + print(f"ERROR: {ROUTER_SPEC.name} is missing from spec/", file=sys.stderr) + return 1 + sys.path.insert(0, str(ROOT / "src")) + try: + from comfy_low.transport import _MODEL_RUN_PATH_TEMPLATE, ROUTER_BASE_URL + except Exception as exc: + print(f"ERROR: comfy_low.transport does not import: {exc!r}", file=sys.stderr) + return 1 + + try: + declared_path, declared_host = _declared_run_route() + except ValueError as exc: + print(f"ERROR: {exc}", file=sys.stderr) + return 1 + + failed = False + if declared_path != _MODEL_RUN_PATH_TEMPLATE: + print( + f"ERROR: the bound model-run route has drifted from {ROUTER_SPEC.name}.\n" + f" spec (runRouterModel): {declared_path}\n" + f" sdk (_MODEL_RUN_PATH_TEMPLATE): {_MODEL_RUN_PATH_TEMPLATE}\n" + " Update comfy_low.transport._MODEL_RUN_PATH_TEMPLATE to the spec's path.", + file=sys.stderr, + ) + failed = True + if declared_host != ROUTER_BASE_URL: + print( + f"ERROR: the default router host has drifted from {ROUTER_SPEC.name}.\n" + f" spec (servers[0].url): {declared_host}\n" + f" sdk (ROUTER_BASE_URL): {ROUTER_BASE_URL}\n" + " Update comfy_low.transport.ROUTER_BASE_URL to the spec's server URL.", + file=sys.stderr, + ) + failed = True + if failed: + return 1 + print(f"OK: the SDK posts a model run to {declared_host}{declared_path}, as the spec declares") + return 0 + + def _check_models() -> int: if not COMMITTED.exists(): print("ERROR: committed models missing; run scripts/gen_models.sh", file=sys.stderr) @@ -208,12 +309,13 @@ def _run(name: str, check: Callable[[], int]) -> int: def main() -> int: - # Both run every time: reporting only the first failure would hide the - # second one behind a fix for the first. `max` over both results rather - # than a short-circuit for the same reason. + # All three run every time: reporting only the first failure would hide the + # others behind a fix for it. `max` over the results rather than a + # short-circuit for the same reason. return max( _run("models", _check_models), _run("router error types", _check_router_error_types), + _run("router run route", _check_router_run_route), ) diff --git a/src/comfy_low/errors.py b/src/comfy_low/errors.py index 3e4f194..3fc587c 100644 --- a/src/comfy_low/errors.py +++ b/src/comfy_low/errors.py @@ -158,7 +158,7 @@ def error_from_envelope( well-formed envelope (so a bare ``401`` with no JSON still maps to ``Unauthorized``). - Not every route answers in the envelope shape. ``POST /api/v2/models/run`` + Not every route answers in the envelope shape. ``POST {router}/v1/models/{provider}/{model}`` is fronted by Router, whose error body is ``{detail, error_type}`` and which repeats the same coarse bucket on the ``X-Comfy-Error-Type`` header (``spec/router-openapi.yaml``). Reading only ``error["code"]`` would collapse diff --git a/src/comfy_low/transport.py b/src/comfy_low/transport.py index b84ef4b..843191b 100644 --- a/src/comfy_low/transport.py +++ b/src/comfy_low/transport.py @@ -10,11 +10,18 @@ * **per-request timeout / abort** — every method takes ``timeout`` and the raw httpx cancellation applies. -One binding is *not* backed by an ``operationId``: ``post_model_run``. The -vendored contract declares no model routes yet, so it is hand-written against -the agreed wire shape, kept out of ``comfy_low.OPERATION_IDS``, and confined to -``model_run_request`` / ``_MODEL_RUN_PATH`` so vendoring the real route later is -a one-place change. +One binding is *not* backed by an ``operationId`` *in this module's sense*: +``post_model_run``. It targets a different surface — Comfy Router, on its own +host (:data:`ROUTER_BASE_URL`) — rather than the ``/api/v2`` deployment the rest +of these methods speak to, and it is declared by a *second* vendored contract, +``spec/router-openapi.yaml`` (``operationId: runRouterModel``, path +``/v1/models/{provider}/{model}``). Nothing is generated from that second file +yet, so the binding is still hand-written; what changed is that it is no longer +hand-*invented*. It stays out of ``comfy_low.OPERATION_IDS`` (which is the +``spec/openapi.yaml`` set, exactly) and stays confined to ``model_run_request`` +/ :data:`_MODEL_RUN_PATH_TEMPLATE`, so a Router spec sync is a one-place change +— and ``tests/test_router_spec_contract.py`` plus ``scripts/check_drift.py`` +fail if that constant and the vendored path disagree. This layer contains no orchestration, retries, hashing, or reconnection — those live in ``comfy_sdk``. @@ -32,7 +39,7 @@ from importlib.metadata import PackageNotFoundError from importlib.metadata import version as _pkg_version from typing import Any, BinaryIO -from urllib.parse import parse_qs, urlsplit, urlunsplit +from urllib.parse import parse_qs, quote, urlsplit, urlunsplit import httpx @@ -60,18 +67,78 @@ #: ``connect`` stays short: an unreachable host is not a slow generation. MODEL_RUN_TIMEOUT = httpx.Timeout(600.0, connect=10.0) -#: Route for a model run. NOT an ``operationId`` from ``spec/openapi.yaml`` — -#: the vendored v2 contract declares no model routes, so this binding is -#: hand-written and is deliberately absent from ``comfy_low.OPERATION_IDS`` -#: (the spec-coverage test asserts that set equals the spec's, exactly). +#: Base URL of Comfy Router — the surface ``post_model_run`` targets, and the +#: ``servers[0].url`` of ``spec/router-openapi.yaml``. It is a *different host* +#: from the ``/api/v2`` deployment ``base_url`` names: the v2 surface serves +#: jobs and assets, Router serves the model-ID-addressed invocation routes. A +#: caller redirects it with ``COMFY_ROUTER_BASE_URL`` (see +#: :data:`comfy_sdk.client.ROUTER_BASE_URL_ENV_VAR`), which is deliberately a +#: second variable rather than a reuse of ``COMFY_BASE_URL`` — pointing one +#: variable at both would send jobs to Router or model runs to the v2 API. +ROUTER_BASE_URL = "https://api.comfy.org" + +#: Route for a model run, verbatim from ``spec/router-openapi.yaml`` — the path +#: whose ``post.operationId`` is ``runRouterModel``. It is NOT an ``operationId`` +#: from ``spec/openapi.yaml`` (a different contract), so this binding is +#: deliberately absent from ``comfy_low.OPERATION_IDS`` (the spec-coverage test +#: asserts that set equals *that* spec's, exactly) and is still hand-bound +#: rather than generated — nothing generates from the Router spec yet. #: Everything about the wire shape is confined to this constant and -#: :func:`model_run_request` so it is one place to reconcile when the route is -#: vendored into the spec and the models are regenerated from it. -_MODEL_RUN_PATH = "/models/run" +#: :func:`model_run_request`, and ``tests/test_router_spec_contract.py`` / +#: ``scripts/check_drift.py`` fail when the vendored spec's path moves and this +#: constant does not follow it. +_MODEL_RUN_PATH_TEMPLATE = "/v1/models/{provider}/{model}" _DEFAULT_PORTS = {"http": 80, "https": 443} +def parse_model_id(model: str) -> tuple[str, str]: + """Split a canonical ``{provider}/{model}`` id into its two path segments. + + Mirrors the TypeScript SDK's ``parseModelId`` so the same id is accepted, + and rejected, identically on both. The id is not an opaque string here: it + *is* the tail of the request path (``/v1/models/{provider}/{model}``), so a + malformed one has to fail locally rather than be pasted into a URL and + answered by whatever route it happens to land on. + + Raises ``TypeError`` when ``model`` is not a ``str`` (a wrong *type*) and + ``ValueError`` when it is a string of the wrong shape (a wrong *value*) — + Python's own split of the two, so ``except ValueError`` around a call that + formats user input does not also swallow a plain programming error. + """ + if not isinstance(model, str): + raise TypeError(f"model id must be a str, got {type(model).__name__}") + segments = model.split("/") + shape = ( + f"model id must be '{{provider}}/{{model}}' — exactly two non-empty " + f"segments separated by '/'; got {model!r}" + ) + # Emptiness first: 'a//b' splits into three segments, and reporting it as + # the variant case below would send the caller looking for a variant they + # never wrote. + if not all(segments): + raise ValueError(shape) + if len(segments) == 3: + raise ValueError( + f"model id {model!r} carries a variant segment, which is not addressable on " + f"POST {_MODEL_RUN_PATH_TEMPLATE} yet — that route takes the two-segment " + f"'{{provider}}/{{model}}' id. Pass the id the catalog lists." + ) + if len(segments) != 2: + raise ValueError(shape) + provider, name = segments + # Refused rather than encoded: `quote` leaves `.` alone (it is unreserved), + # so a `.`/`..` segment would survive into the path and let an id walk the + # route — `acme/..` resolving to `/v1/models/acme` on any intermediary that + # normalizes dot segments, which most do. + if provider in (".", "..") or name in (".", ".."): + raise ValueError( + f"model id segments must not be '.' or '..' — they would traverse the " + f"request path rather than name a model; got {model!r}" + ) + return provider, name + + def model_run_request( model: str, arguments: Mapping[str, Any], @@ -79,18 +146,28 @@ def model_run_request( ) -> tuple[str, dict[str, Any], dict[str, str]]: """Sans-IO ``(path, json_body, headers)`` for one model run. - The model id travels in the *body*, not the path: ids are commonly - provider-namespaced and contain ``/`` (``vendor/family/variant``), which in - a path segment needs percent-encoding that intermediaries normalize - inconsistently. A named body field also leaves room for sibling fields - later without moving the route. + The model id *addresses* the request: it is the two path segments of + ``/v1/models/{provider}/{model}``, not a body field. That is Router's + contract — the request body is the partner model's OWN native JSON input, + forwarded to the provider unchanged, so there is no room in it for a + Comfy-shaped ``{model, arguments}`` envelope. A caller can therefore move + between the partner's API and Router by changing the host, which is the + whole point of the surface. + + Each segment is percent-encoded with ``safe=""`` so nothing in it can add a + path segment, a query, or a fragment. The path is still the one the + vendored spec declares — see :data:`_MODEL_RUN_PATH_TEMPLATE`. ``arguments`` is copied into a plain dict so any ``Mapping`` is accepted and the caller's object is never handed to the JSON encoder directly. """ - body: dict[str, Any] = {"model": model, "arguments": dict(arguments)} + provider, name = parse_model_id(model) + path = _MODEL_RUN_PATH_TEMPLATE.format( + provider=quote(provider, safe=""), model=quote(name, safe="") + ) + body: dict[str, Any] = dict(arguments) headers = {"Idempotency-Key": idempotency_key} if idempotency_key else {} - return _MODEL_RUN_PATH, body, headers + return path, body, headers def _build_user_agent(client_info: str | None) -> str: @@ -183,7 +260,13 @@ def redact_userinfo(url: str) -> str: class _Prepared: """Sans-IO request building shared by both transports.""" - def __init__(self, base_url: str, api_key: str | None, client_info: str | None = None) -> None: + def __init__( + self, + base_url: str, + api_key: str | None, + client_info: str | None = None, + router_base_url: str = ROUTER_BASE_URL, + ) -> None: self.base_url = base_url.rstrip("/") self.api_key = api_key #: ``base_url`` with any userinfo redacted — what every ``repr`` shows. @@ -191,6 +274,20 @@ def __init__(self, base_url: str, api_key: str | None, client_info: str | None = self._base_origin = origin(self.base_url) parts = urlsplit(self.base_url) self._origin_url = f"{parts.scheme}://{parts.netloc}" + #: Comfy Router's base URL — a *second* target, used only by + #: ``post_model_run``, which builds an absolute URL against it rather + #: than going through :meth:`url` (that method prepends ``/api/v2``, + #: which is the other surface's mount prefix, not Router's). + self.router_base_url = router_base_url.rstrip("/") + # Checked here rather than left to httpx: `url()` passes a string + # through only when it starts with `http`, so a non-http router base + # would silently fall through to the `base_url + /api/v2 + ...` branch + # and send the run to the wrong surface under a mangled URL instead of + # failing with a sentence that names the setting. + if not self.router_base_url.startswith(("http://", "https://")): + raise ValueError(f"router_base_url must be an http(s) URL; got {router_base_url!r}") + self.safe_router_base_url = redact_userinfo(self.router_base_url) + self._router_origin = origin(self.router_base_url) self._user_agent = _build_user_agent(client_info) def __repr__(self) -> str: @@ -219,11 +316,21 @@ def headers(self, url: str, extra: dict[str, str] | None = None) -> dict[str, st h: dict[str, str] = {"User-Agent": self._user_agent} # Only authenticate when a key is set: a local proxy fronts a ComfyUI # with no auth, so we never leak credentials it does not want. And only - # attach it when the resolved request URL is same-origin as base_url: - # server-returned absolute follow-up links (job.urls.self/cancel/events) - # must not carry the key to a different scheme/host/port. Relative paths - # are always resolved under base_url, so they are unaffected. - if self.api_key and origin(url) == self._base_origin: + # attach it when the resolved request URL is same-origin as one of this + # client's own *configured* targets: server-returned absolute follow-up + # links (job.urls.self/cancel/events) must not carry the key to a + # different scheme/host/port. Relative paths are always resolved under + # base_url, so they are unaffected. + # + # There are two configured targets rather than one because the SDK + # speaks to two surfaces: the ``/api/v2`` deployment (`base_url`) and + # Comfy Router (`router_base_url`), which is a different host by + # default. Both come from the client's own construction — a constant, or + # an environment variable the operator set — never from a server + # response, so trusting the router origin does not widen what a + # malicious server can point the credential at. A third origin still + # gets nothing; ``tests/test_transport_security.py`` asserts it. + if self.api_key and origin(url) in (self._base_origin, self._router_origin): h["Authorization"] = f"Bearer {self.api_key}" if extra: h.update(extra) @@ -322,8 +429,9 @@ def __init__( client: httpx.Client | None = None, timeout: float | None = 30.0, client_info: str | None = None, + router_base_url: str = ROUTER_BASE_URL, ) -> None: - self._p = _Prepared(base_url, api_key, client_info) + self._p = _Prepared(base_url, api_key, client_info, router_base_url) self._own_client = client is None self._client = client or httpx.Client(timeout=timeout, follow_redirects=True) @@ -342,6 +450,21 @@ def safe_base_url(self) -> str: """:attr:`base_url` with any userinfo redacted — the form safe to log.""" return self._p.safe_base_url + @property + def router_base_url(self) -> str: + """Comfy Router's base URL — where :meth:`post_model_run` sends its request. + + A second target, not a view of :attr:`base_url`: model runs are + model-ID-addressed routes on Router's own host, while jobs and assets + are ``/api/v2`` routes on the deployment :attr:`base_url` names. + """ + return self._p.router_base_url + + @property + def safe_router_base_url(self) -> str: + """:attr:`router_base_url` with any userinfo redacted — safe to log.""" + return self._p.safe_router_base_url + @property def timeout(self) -> httpx.Timeout: """The httpx client's default timeout. A per-request ``timeout=`` still wins.""" @@ -649,7 +772,12 @@ def post_model_run( idempotency_key: str | None = None, timeout: Any = MODEL_RUN_TIMEOUT, ) -> dict[str, Any]: - """POST /api/v2/models/run — run a model, awaited server-side. + """POST ``{router_base_url}/v1/models/{provider}/{model}`` — awaited server-side. + + Addressed to Comfy Router, not to the ``/api/v2`` deployment + :attr:`base_url` names: the URL is built absolute against + :attr:`router_base_url` precisely so it does not pick up the ``/api/v2`` + prefix ``_Prepared.url`` adds to a relative path. One request, one response: the server does not answer until the generation is complete, so the decoded body *is* the finished result. @@ -657,12 +785,19 @@ def post_model_run( side, inside this call, which is why the default ``timeout`` is :data:`MODEL_RUN_TIMEOUT` rather than the client's own. - The body is returned verbatim — the provider's native payload, with no - model class layered over it. This is not a spec operation; see - :data:`_MODEL_RUN_PATH`. + ``arguments`` is sent as the body verbatim (the partner model's native + JSON input) and the response body is returned verbatim (its native + output), with no model class layered over either. This is not an + ``operationId`` of ``spec/openapi.yaml``; it is ``runRouterModel`` of + ``spec/router-openapi.yaml``, hand-bound — see + :data:`_MODEL_RUN_PATH_TEMPLATE`. + + Raises ``TypeError``/``ValueError`` from :func:`parse_model_id` before + any request when ``model`` is not a ``{provider}/{model}`` id. """ path, body, headers = model_run_request(model, arguments, idempotency_key) - resp = self.raw_request("POST", path, headers=headers, json=body, timeout=timeout) + url = self._p.router_base_url + path + resp = self.raw_request("POST", url, headers=headers, json=body, timeout=timeout) return self._p.parse_or_raise(resp, (200, 201)) @@ -677,8 +812,9 @@ def __init__( client: httpx.AsyncClient | None = None, timeout: float | None = 30.0, client_info: str | None = None, + router_base_url: str = ROUTER_BASE_URL, ) -> None: - self._p = _Prepared(base_url, api_key, client_info) + self._p = _Prepared(base_url, api_key, client_info, router_base_url) self._own_client = client is None self._client = client or httpx.AsyncClient(timeout=timeout, follow_redirects=True) @@ -693,6 +829,16 @@ def safe_base_url(self) -> str: """:attr:`base_url` with any userinfo redacted — the form safe to log.""" return self._p.safe_base_url + @property + def router_base_url(self) -> str: + """Comfy Router's base URL — mirrors :attr:`ComfyLow.router_base_url`.""" + return self._p.router_base_url + + @property + def safe_router_base_url(self) -> str: + """:attr:`router_base_url` with any userinfo redacted — safe to log.""" + return self._p.safe_router_base_url + @property def timeout(self) -> httpx.Timeout: """The httpx client's default timeout. A per-request ``timeout=`` still wins.""" @@ -971,7 +1117,8 @@ async def post_model_run( ) -> dict[str, Any]: """Async :meth:`ComfyLow.post_model_run`.""" path, body, headers = model_run_request(model, arguments, idempotency_key) - resp = await self.raw_request("POST", path, headers=headers, json=body, timeout=timeout) + url = self._p.router_base_url + path + resp = await self.raw_request("POST", url, headers=headers, json=body, timeout=timeout) return self._p.parse_or_raise(resp, (200, 201)) diff --git a/src/comfy_sdk/__init__.py b/src/comfy_sdk/__init__.py index efb9ec5..b2ea25d 100644 --- a/src/comfy_sdk/__init__.py +++ b/src/comfy_sdk/__init__.py @@ -32,7 +32,15 @@ from importlib.metadata import version as _pkg_version from .assets import Asset, AssetFactory, AsyncAsset, AsyncAssetFactory -from .client import API_KEY_ENV_VAR, BASE_URL_ENV_VAR, COMFY_CLOUD_BASE_URL, AsyncComfy, Comfy +from .client import ( + API_KEY_ENV_VAR, + BASE_URL_ENV_VAR, + COMFY_CLOUD_BASE_URL, + COMFY_ROUTER_BASE_URL, + ROUTER_BASE_URL_ENV_VAR, + AsyncComfy, + Comfy, +) from .events import ( Event, Log, @@ -75,6 +83,8 @@ "Comfy", "COMFY_CLOUD_BASE_URL", "BASE_URL_ENV_VAR", + "COMFY_ROUTER_BASE_URL", + "ROUTER_BASE_URL_ENV_VAR", "API_KEY_ENV_VAR", "AsyncComfy", # assets / workflows / jobs / outputs diff --git a/src/comfy_sdk/client.py b/src/comfy_sdk/client.py index a3b714e..3cbe9a0 100644 --- a/src/comfy_sdk/client.py +++ b/src/comfy_sdk/client.py @@ -11,6 +11,14 @@ serverless one — is selected through the ``COMFY_BASE_URL`` environment variable; there is no base-URL constructor parameter. +``client.models`` is the exception, and deliberately so: model runs are +model-ID-addressed routes on **Comfy Router**, a different host, so they follow +``COMFY_ROUTER_BASE_URL`` (default ``https://api.comfy.org``) instead. One +variable pointed at both surfaces would send jobs to Router or model runs to +the v2 API; neither serves the other's routes. The credential is the same one +either way — this client's own bearer token, attached to both of its configured +origins and to no third one. + Credentials resolve in a fixed order at construction: the explicit ``api_key`` argument, then the ``COMFY_API_KEY`` environment variable, then — targeting Comfy Cloud, which always requires a key — a local :class:`MissingApiKey` @@ -36,7 +44,7 @@ from urllib.parse import urlsplit from comfy_low.errors import ApiError -from comfy_low.transport import AsyncComfyLow, ComfyLow, origin +from comfy_low.transport import ROUTER_BASE_URL, AsyncComfyLow, ComfyLow, origin from . import _core from .assets import AssetFactory, AsyncAssetFactory @@ -52,6 +60,20 @@ COMFY_CLOUD_BASE_URL = "https://cloud.comfy.org" #: Environment variable that redirects a client at another deployment. BASE_URL_ENV_VAR = "COMFY_BASE_URL" +#: Base URL of Comfy Router — where ``client.models`` sends its requests. A +#: *separate* target from :data:`COMFY_CLOUD_BASE_URL`: the ``/api/v2`` surface +#: serves jobs and assets, Router serves the model-ID-addressed invocation +#: routes, and they are different hosts. The literal lives in +#: :mod:`comfy_low.transport` (the layer that owns the wire) and is re-exported +#: here so it sits beside the constant it is most often confused with. +COMFY_ROUTER_BASE_URL = ROUTER_BASE_URL +#: Environment variable that redirects *model runs* at another Router +#: deployment. Deliberately a second variable rather than a reuse of +#: :data:`BASE_URL_ENV_VAR`: one variable pointed at both surfaces would send +#: jobs to Router, or model runs to the v2 API, and neither serves the other's +#: routes. Spelled the same as the TypeScript SDK's, so a test environment sets +#: one variable for both. +ROUTER_BASE_URL_ENV_VAR = "COMFY_ROUTER_BASE_URL" #: Environment variable read when no ``api_key`` is passed to the constructor. API_KEY_ENV_VAR = "COMFY_API_KEY" @@ -77,17 +99,21 @@ def _retry_delay(exc: ApiError, deadline: float) -> float | None: return max(0.0, min(raw_delay, remaining)) -def _resolve_base_url() -> str: - """Comfy Cloud, unless ``COMFY_BASE_URL`` names another deployment. +def _resolve_env_url(var: str, default: str) -> str: + """``var``'s value if it names a usable target, else ``default``. Read per construction rather than at import so a process can point successive clients at different deployments. An unset-or-blank variable - means Comfy Cloud, so ``COMFY_BASE_URL=`` in a shell profile or ``.env`` + means the default, so ``COMFY_BASE_URL=`` in a shell profile or ``.env`` is not an error. + + Shared by both targets so the two cannot validate differently — a rule that + held for the v2 base URL and not for the Router one would be a rule nobody + could state. """ - raw = os.environ.get(BASE_URL_ENV_VAR, "").strip() + raw = os.environ.get(var, "").strip() if not raw: - return COMFY_CLOUD_BASE_URL + return default parsed = urlsplit(raw) try: # urlsplit defers the port check, so a non-numeric or out-of-range one @@ -106,12 +132,29 @@ def _resolve_base_url() -> str: valid = False if not valid: raise ValueError( - f"{BASE_URL_ENV_VAR} must be an http(s) URL with no query or fragment " + f"{var} must be an http(s) URL with no query or fragment " f"(e.g. 'http://127.0.0.1:8189'); got {raw!r}" ) return raw +def _resolve_base_url() -> str: + """Comfy Cloud, unless ``COMFY_BASE_URL`` names another deployment.""" + return _resolve_env_url(BASE_URL_ENV_VAR, COMFY_CLOUD_BASE_URL) + + +def _resolve_router_base_url() -> str: + """Comfy Router, unless ``COMFY_ROUTER_BASE_URL`` names another one. + + Same validation as :func:`_resolve_base_url`, and the same + read-per-construction rule. The trailing slash is stripped here as well as + in the transport, because this value is *concatenated* with a path that + already starts with ``/`` — ``https://api.comfy.org//v1/models/...`` is a + different path to an origin server than the one the spec declares. + """ + return _resolve_env_url(ROUTER_BASE_URL_ENV_VAR, COMFY_ROUTER_BASE_URL).rstrip("/") + + def _same_deployment(url: str, other: str) -> bool: """Whether two base URLs name the same deployment. @@ -198,7 +241,13 @@ def __init__( ) -> None: base_url = _resolve_base_url() key = _resolve_api_key(api_key, base_url) - self._low = ComfyLow(base_url, key, timeout=timeout, client_info=client_info) + self._low = ComfyLow( + base_url, + key, + timeout=timeout, + client_info=client_info, + router_base_url=_resolve_router_base_url(), + ) self.assets = AssetFactory(self._low) self.workflows = WorkflowFactory() self.jobs = JobFactory(self._low) @@ -324,7 +373,13 @@ def __init__( ) -> None: base_url = _resolve_base_url() key = _resolve_api_key(api_key, base_url) - self._low = AsyncComfyLow(base_url, key, timeout=timeout, client_info=client_info) + self._low = AsyncComfyLow( + base_url, + key, + timeout=timeout, + client_info=client_info, + router_base_url=_resolve_router_base_url(), + ) self.assets = AsyncAssetFactory(self._low) self.workflows = WorkflowFactory() self.jobs = AsyncJobFactory(self._low) diff --git a/src/comfy_sdk/models.py b/src/comfy_sdk/models.py index 9e5def3..93df39a 100644 --- a/src/comfy_sdk/models.py +++ b/src/comfy_sdk/models.py @@ -2,9 +2,17 @@ Reached from a client you already constructed (``Comfy().models`` / ``AsyncComfy().models``) rather than built on its own, so it uses that client's -credentials, base URL, transport and timeout: one connection pool, one -credential, one place to configure both. A separate client object for model -operations would fork all of that, which is what namespacing avoids. +credentials, transport and timeout: one connection pool, one credential, one +place to configure both. A separate client object for model operations would +fork all of that, which is what namespacing avoids. + +The one setting it does *not* share is the target host. Model runs go to Comfy +Router — ``POST {router_base_url}/v1/models/{provider}/{model}``, the partner +model's native JSON straight through — while the client's ``base_url`` names +the ``/api/v2`` deployment serving jobs and assets. So ``models.base_url`` +reports ``COMFY_ROUTER_BASE_URL`` (default ``https://api.comfy.org``), and a +``model`` argument is the canonical two-segment ``{provider}/{model}`` id that +addresses the route. The namespace holds the host client's transport itself — not a copy of its settings — so a change made on the client after construction (a rotated key, a @@ -75,8 +83,15 @@ class _ModelsBase: @property def base_url(self) -> str: - """The host client's base URL — where model requests are sent.""" - return self._low.base_url + """Comfy Router's base URL — where model requests are sent. + + Not the host client's ``base_url``: a model run is a model-ID-addressed + route on Comfy Router (``https://api.comfy.org`` by default, redirected + by ``COMFY_ROUTER_BASE_URL``), while the client's own ``base_url`` names + the ``/api/v2`` deployment that serves jobs and assets. Read live off + the shared transport, exactly as :attr:`timeout` is. + """ + return self._low.router_base_url @property def timeout(self) -> httpx.Timeout: @@ -91,8 +106,11 @@ def retry(self) -> RetryPolicy: def __repr__(self) -> str: # Redacted like the host client's repr: a base URL may carry proxy # credentials in its userinfo, and this namespace lands in the same - # tracebacks and CI logs the client does. - return f"{type(self).__name__}(base_url={self._low.safe_base_url!r})" + # tracebacks and CI logs the client does. It renders the *router* base + # URL because that is the target this namespace actually talks to — + # showing the client's `/api/v2` base URL here would name a host no + # `models` call ever reaches. + return f"{type(self).__name__}(base_url={self._low.safe_router_base_url!r})" class Models(_ModelsBase): @@ -116,6 +134,20 @@ def run( ) -> dict[str, Any]: """Run ``model`` with ``arguments`` and return the completed result. + ``model`` is the canonical ``{provider}/{model}`` id — exactly the two + path segments that address the run on Comfy Router + (``POST {router_base_url}/v1/models/{provider}/{model}``), and exactly + what the model catalog lists. It must be two non-empty segments: a + one-segment id, the three-segment ``{provider}/{model}/{variant}`` form + (not addressable on this route yet), and any ``.``/``..`` segment each + raise ``ValueError`` locally, before any request is made; a non-string + raises ``TypeError``. + + ``arguments`` is the partner model's **own native JSON input**, + forwarded to the provider unchanged — there is no Comfy-shaped envelope + around it, so whatever the partner documents as its request body is + what goes here. + One call, one result. It blocks until the generation is finished — including for a provider the platform has to submit-and-poll, where the polling happens server side inside this call, invisible to the caller. @@ -233,9 +265,10 @@ async def run( """Awaitable :meth:`Models.run` — same arguments, same result shape. This *is* the async form of ``run``: awaiting it on ``AsyncComfy`` is - the whole difference from the sync client — including the retry policy, - the one-key-per-call rule, and the ``.idempotency_key`` every exception - it raises carries for the replay. See :meth:`Models.run`. + the whole difference from the sync client — including the model-id + rule, the retry policy, the one-key-per-call rule, and the + ``.idempotency_key`` every exception it raises carries for the replay. + See :meth:`Models.run`. """ low = cast(AsyncComfyLow, self._low) key = idempotency_key or new_idempotency_key() diff --git a/src/comfy_sdk/retry.py b/src/comfy_sdk/retry.py index 5829a1b..ac07329 100644 --- a/src/comfy_sdk/retry.py +++ b/src/comfy_sdk/retry.py @@ -27,10 +27,11 @@ the server cannot characterise ("an upstream timeout or 5xx where the job may or may not have been created") keeps it claimed. -``POST /models/run`` is not itself in that spec, so its key semantics are not -that spec's to state — and for one failure the *router* contract states them -directly. ``spec/router-openapi.yaml``'s ``deadline_exceeded`` bucket says to -"retry it with the SAME ``Idempotency-Key``: when the provider had already +The model-run route (``POST /v1/models/{provider}/{model}``) is not in that +spec at all — it is a different surface, declared by ``spec/router-openapi.yaml`` +— so its key semantics are not the v2 spec's to state, and for one failure the +*router* contract states them directly. That spec's ``deadline_exceeded`` +bucket says to "retry it with the SAME ``Idempotency-Key``: when the provider had already accepted the generation, the retry collects that generation rather than dispatching another", and pins the ``Retry-After`` it carries to "seconds to wait before retrying the SAME request with the SAME ``Idempotency-Key``". So @@ -302,9 +303,10 @@ def error_bucket_of(exc: BaseException) -> str | None: Read by attribute rather than by type, for the same reason :func:`retry_after_of` is: one failure reaches this module modelled by two different layers. A typed router error carries the wire ``error_type`` - (``RouterError.error_type``), while ``POST /models/run`` today raises the - protocol :class:`~comfy_low.errors.ApiError`, whose envelope names the same - thing ``code``. Reading only ``error_type`` would make every bucket-keyed + (``RouterError.error_type``), while ``POST /v1/models/{provider}/{model}`` + today raises the protocol :class:`~comfy_low.errors.ApiError`, whose + envelope names the same thing ``code``. Reading only ``error_type`` would + make every bucket-keyed rule below unreachable on the route those rules were written for — a retry that is a silent no-op with no test failing, which is the failure mode this module has already been bitten by twice. diff --git a/tests/conftest.py b/tests/conftest.py index 6710100..53244e9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,7 +3,14 @@ Keeps the SDK's own test suite independent of a real v2 server or proxy. Each test configures ``server.state`` to drive a specific scenario (dedup hit, hash mismatch, queue-full-then-ok, SSE reconnect, ...); the ``server`` fixture points -the SDK at the stub by setting ``COMFY_BASE_URL``. +the SDK at the stub by setting ``COMFY_BASE_URL`` *and* +``COMFY_ROUTER_BASE_URL``. + +Both, because the SDK speaks to two surfaces: the ``/api/v2`` deployment (jobs, +assets) and Comfy Router (``/v1/models/{provider}/{model}``), which is a +different host in production. This one stub answers both route families, so a +test that exercises either gets a single server — while a test that is *about* +the two being separate points ``COMFY_ROUTER_BASE_URL`` at ``second_server``. """ from __future__ import annotations @@ -15,10 +22,11 @@ from dataclasses import dataclass, field from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Any +from urllib.parse import unquote import pytest -from comfy_sdk import API_KEY_ENV_VAR, BASE_URL_ENV_VAR +from comfy_sdk import API_KEY_ENV_VAR, BASE_URL_ENV_VAR, ROUTER_BASE_URL_ENV_VAR @dataclass @@ -83,7 +91,7 @@ class ServerState: job_workflow_format: str = "api" job_workflow_not_found: bool = False - # --- POST /models/run (the awaited model run) --- + # --- POST /v1/models/{provider}/{model} (the awaited model run) --- # The provider's native payload the run resolves to. Deliberately not a # Comfy-shaped envelope: the SDK must hand it back untouched. model_run_result: dict[str, Any] = field( @@ -99,7 +107,7 @@ class ServerState: model_run_delay: float = 0.0 # (status, code) answered instead of the result. model_run_error: tuple[int, str] | None = None - # POST /models/run fails this many times before serving the result — the + # A model run fails this many times before serving the result — the # transient-failure-then-success path a retry policy exists for. Decremented # per request, and checked *before* `model_run_error`, which is the # permanent-failure knob. @@ -146,7 +154,7 @@ class ServerState: model_run_request_id: str | None = None # Answer model-run failures in Router's own error shape -- the coarse bucket # on `X-Comfy-Error-Type` plus a `{detail, error_type}` body -- instead of - # the v2 `{error: {code, message}}` envelope. `POST /api/v2/models/run` is + # the v2 `{error: {code, message}}` envelope. The model-run route is # fronted by Router, so this is the shape a real deployment's 504 arrives # in, and the bucket-keyed collect rule has to read it. model_run_router_error_shape: bool = False @@ -175,12 +183,22 @@ class ServerState: last_auth_header: str | None = None last_user_agent: str | None = None model_run_count: int = 0 + # The raw JSON body of the last model run — the partner model's *native* + # input, with no `{model, arguments}` envelope around it, exactly as Router + # forwards it upstream. last_model_run_body: dict[str, Any] | None = None - # Every Idempotency-Key seen on POST /models/run, in arrival order (`None` + # The two path segments of the last model run, percent-DECODED, so a test + # asserts the id the caller passed rather than a particular encoding of it. + last_model_run_provider: str | None = None + last_model_run_model: str | None = None + # ...and the raw, still-encoded request path, for the tests that are about + # the encoding itself. + last_model_run_path: str | None = None + # Every Idempotency-Key seen on a model run, in arrival order (`None` # records a run that arrived without the header at all). model_run_idempotency_keys: list[str | None] = field(default_factory=list) - # Keys POST /models/run has *claimed*, so a reuse can be rejected exactly - # as POST /jobs rejects one. Kept apart from `idempotency` only so a model + # Keys a model run has *claimed*, so a reuse can be rejected exactly as + # POST /jobs rejects one. Kept apart from `idempotency` only so a model # test cannot perturb a workflow test's bookkeeping. model_run_idempotency: dict[str, str] = field(default_factory=dict) # Idempotency-Key -> the result recorded for it under @@ -470,8 +488,14 @@ def do_POST(self) -> None: if self.path == "/api/v2/jobs": self._post_jobs() return - if self.path == "/api/v2/models/run": - self._post_model_run() + # Comfy Router's invocation route — a different surface from the + # `/api/v2` paths above (a different host in production; the same + # stub here, with `COMFY_ROUTER_BASE_URL` pointed at it). The two + # segments are the model id, so they are matched rather than + # compared to a fixed string. + m = re.match(r"/v1/models/([^/]+)/([^/]+)$", self.path) + if m: + self._post_model_run(m.group(1), m.group(2)) return m = re.match(r"/api/v2/jobs/([^/]+)/cancel$", self.path) if m: @@ -497,8 +521,16 @@ def _post_from_hash(self) -> None: else: self._err(404, "blob_not_found", "no such blob") - def _post_model_run(self) -> None: + def _post_model_run(self, provider: str, model: str) -> None: state.model_run_count += 1 + # Decoded, because the SDK percent-encodes each segment and a real + # origin server decodes it before routing — asserting the encoded + # form everywhere would pin tests to an encoding rather than to the + # id the caller passed. `last_model_run_path` keeps the raw form + # for the tests that are about the encoding. + state.last_model_run_provider = unquote(provider) + state.last_model_run_model = unquote(model) + state.last_model_run_path = self.path state.last_model_run_body = json.loads(self._read_body() or b"{}") key = self.headers.get("Idempotency-Key") state.model_run_idempotency_keys.append(key) @@ -697,6 +729,10 @@ def _no_ambient_base_url(request, monkeypatch): if "integration" in request.path.parts: return monkeypatch.delenv(BASE_URL_ENV_VAR, raising=False) + # The Router target too: a developer with `COMFY_ROUTER_BASE_URL` exported + # would otherwise send every `models.run` test at their own host — and the + # default-value assertions would pass or fail on their shell, not the code. + monkeypatch.delenv(ROUTER_BASE_URL_ENV_VAR, raising=False) @pytest.fixture(autouse=True) @@ -719,6 +755,12 @@ def server(monkeypatch): # Clients read their target from the environment, so pointing them at the # stub is part of standing it up: tests just construct ``Comfy()``. monkeypatch.setenv(BASE_URL_ENV_VAR, srv.base_url) + # Both targets, because the SDK has two: jobs and assets resolve under + # `COMFY_BASE_URL`, model runs under `COMFY_ROUTER_BASE_URL`. The one stub + # serves both route families, so pointing both here keeps a `models.run` + # test a single-server test — the *separate*-origin cases point this second + # variable at `second_server` themselves. + monkeypatch.setenv(ROUTER_BASE_URL_ENV_VAR, srv.base_url) try: yield srv finally: diff --git a/tests/test_api_key.py b/tests/test_api_key.py index 0e98ca6..b36fb55 100644 --- a/tests/test_api_key.py +++ b/tests/test_api_key.py @@ -27,6 +27,7 @@ API_KEY_ENV_VAR, BASE_URL_ENV_VAR, COMFY_CLOUD_BASE_URL, + ROUTER_BASE_URL_ENV_VAR, AsyncComfy, Comfy, ComfyError, @@ -299,8 +300,13 @@ def test_base_url_userinfo_is_redacted_in_every_repr(monkeypatch, client_cls) -> out of. Only what is *rendered* is redacted — the transport still resolves requests against the URL it was given, so the proxy credential keeps working. + + Both targets carry the URL, because both can carry userinfo: a fronting + proxy is as legitimate in front of Comfy Router as in front of the v2 + deployment, and ``repr(client.models)`` renders the *router* URL. """ monkeypatch.setenv(BASE_URL_ENV_VAR, PROXY_BASE_URL) + monkeypatch.setenv(ROUTER_BASE_URL_ENV_VAR, PROXY_BASE_URL) with constructed(client_cls) as client: for rendered in ( repr(client), @@ -313,15 +319,18 @@ def test_base_url_userinfo_is_redacted_in_every_repr(monkeypatch, client_cls) -> assert "proxy-user" not in rendered assert "***@proxy.example" in rendered assert client._low.base_url == PROXY_BASE_URL + assert client._low.router_base_url == PROXY_BASE_URL @CLIENTS def test_a_base_url_without_userinfo_is_rendered_unchanged(monkeypatch, client_cls) -> None: """Redaction stays invisible in the ordinary case — the target reads plainly.""" monkeypatch.setenv(BASE_URL_ENV_VAR, LOCAL) + monkeypatch.setenv(ROUTER_BASE_URL_ENV_VAR, LOCAL) with constructed(client_cls) as client: assert f"base_url={LOCAL!r}" in repr(client) assert f"base_url={LOCAL!r}" in repr(client._low) + assert f"base_url={LOCAL!r}" in repr(client.models) def test_environment_key_reaches_the_server(server, monkeypatch) -> None: diff --git a/tests/test_error_mapping.py b/tests/test_error_mapping.py index 17a732d..c6d29e1 100644 --- a/tests/test_error_mapping.py +++ b/tests/test_error_mapping.py @@ -23,7 +23,7 @@ def test_404_codes_map_to_notfound(code: str) -> None: # --- the two error shapes one client has to read --- # -# `POST /api/v2/models/run` is fronted by Router, whose error body is +# `POST /v1/models/{provider}/{model}` is Router's own route, whose error body is # `{detail, error_type}` with the coarse bucket repeated on `X-Comfy-Error-Type` # -- not the v2 `{error: {code, message}}` envelope every other route answers in. # `error_from_envelope` has to read both, because `comfy_sdk.retry` keys its diff --git a/tests/test_models_namespace.py b/tests/test_models_namespace.py index cecd45d..b8197da 100644 --- a/tests/test_models_namespace.py +++ b/tests/test_models_namespace.py @@ -39,14 +39,19 @@ async def test_async_models_holds_the_host_clients_transport(server) -> None: assert client.models._low is client._low -def test_models_reports_the_host_clients_base_url(server) -> None: +def test_models_reports_the_routers_base_url(server) -> None: + # The namespace's own target, read live off the shared transport — Router, + # not the client's `/api/v2` deployment. The `server` fixture points both + # at the one stub, so `client._low.router_base_url` is what is asserted + # rather than the coincidence of the two being equal here; the separate-host + # case lives in tests/test_models_run.py. with Comfy() as client: - assert client.models.base_url == client._low.base_url == server.base_url + assert client.models.base_url == client._low.router_base_url == server.base_url -async def test_async_models_reports_the_host_clients_base_url(server) -> None: +async def test_async_models_reports_the_routers_base_url(server) -> None: async with AsyncComfy() as client: - assert client.models.base_url == client._low.base_url == server.base_url + assert client.models.base_url == client._low.router_base_url == server.base_url def test_a_timeout_change_on_the_client_is_visible_through_models(server) -> None: @@ -92,7 +97,7 @@ def test_two_clients_get_independent_namespaces(server) -> None: assert one.models._low is not two.models._low -def test_models_repr_names_the_shared_base_url(server) -> None: +def test_models_repr_names_the_routers_base_url(server) -> None: with Comfy() as client: assert repr(client.models) == f"Models(base_url={server.base_url!r})" diff --git a/tests/test_models_run.py b/tests/test_models_run.py index 973488b..f9aeace 100644 --- a/tests/test_models_run.py +++ b/tests/test_models_run.py @@ -3,7 +3,9 @@ Covers the whole contract of the headline model API on both clients: the sync call blocks and returns the finished result, the async call awaits to the same shape, the awaitable form is the *async client* rather than a suffixed method -(asserted, not merely absent), the wait is sized for a server that polls +(asserted, not merely absent), the run is addressed to Comfy Router by the +model's two id segments with the model's own native JSON as the body, the +result is that model's native output, the wait is sized for a server that polls upstream inside the call, an ``Idempotency-Key`` is plumbed onto the wire and rides out on whatever the call raises, and the result handed back is the provider's own payload rather than a wrapper. @@ -22,8 +24,22 @@ import httpx import pytest -from comfy_low.transport import MODEL_RUN_TIMEOUT, AsyncComfyLow, ComfyLow, model_run_request -from comfy_sdk import NO_RETRY, AsyncComfy, Comfy, RetryPolicy +from comfy_low.transport import ( + MODEL_RUN_TIMEOUT, + ROUTER_BASE_URL, + AsyncComfyLow, + ComfyLow, + model_run_request, + parse_model_id, +) +from comfy_sdk import ( + COMFY_ROUTER_BASE_URL, + NO_RETRY, + ROUTER_BASE_URL_ENV_VAR, + AsyncComfy, + Comfy, + RetryPolicy, +) from comfy_sdk.exceptions import ComfyError, NotFound, Unauthorized from comfy_sdk.models import AsyncModels, Models from comfy_sdk.router_exceptions import ( @@ -33,7 +49,10 @@ error_from_response, ) -MODEL = "acme/flux/dev" +#: The canonical two-segment ``{provider}/{model}`` id the route is addressed +#: by — the same shape ``spec/router-openapi.yaml``'s ``RouterModelId`` pattern +#: declares, and the id the catalog lists. +MODEL = "acme/flux-dev" ARGS = {"prompt": "a cat", "steps": 4} @@ -78,10 +97,111 @@ def test_a_created_shaped_success_is_also_a_result(server) -> None: assert client.models.run(MODEL, ARGS) == server.state.model_run_result -def test_run_sends_the_model_and_arguments(server) -> None: +# --- the model id addresses the route; the body is the model's own input ---- + + +def test_run_addresses_the_model_by_path_and_sends_the_native_body(server) -> None: + # The wire shape Router declares: the id is the two path segments of + # `/v1/models/{provider}/{model}`, and the body is the partner model's OWN + # native JSON input, forwarded unchanged — no `{model, arguments}` envelope. with Comfy() as client: client.models.run(MODEL, ARGS) - assert server.state.last_model_run_body == {"model": MODEL, "arguments": ARGS} + assert server.state.last_model_run_path == "/v1/models/acme/flux-dev" + assert server.state.last_model_run_provider == "acme" + assert server.state.last_model_run_model == "flux-dev" + assert server.state.last_model_run_body == ARGS + assert "model" not in server.state.last_model_run_body + assert "arguments" not in server.state.last_model_run_body + + +async def test_the_async_client_addresses_the_route_the_same_way(server) -> None: + async with AsyncComfy() as client: + await client.models.run(MODEL, ARGS) + assert server.state.last_model_run_path == "/v1/models/acme/flux-dev" + assert server.state.last_model_run_body == ARGS + + +def test_the_sans_io_request_builder_agrees_with_the_wire() -> None: + # The one place the wire shape is decided, asserted directly so a change to + # it cannot hide behind the stub's own routing. + path, body, headers = model_run_request(MODEL, ARGS, "k-1") + assert path == "/v1/models/acme/flux-dev" + assert body == ARGS + assert headers == {"Idempotency-Key": "k-1"} + + +@pytest.mark.parametrize( + "model, path", + [ + # `.`, `_` and `-` are all legal *inside* a segment per the spec's + # `RouterModelSegment` pattern, and none of them is percent-encoded: + # they are unreserved (or sub-delims) in a path segment, so the URL the + # caller reads in a log is the id they passed. + ("fal-ai/flux-pro", "/v1/models/fal-ai/flux-pro"), + ("acme/sd_xl.turbo", "/v1/models/acme/sd_xl.turbo"), + ("acme_labs/v1.5", "/v1/models/acme_labs/v1.5"), + # ...while anything that would change the *structure* of the path is + # encoded rather than passed through. + ("acme/a b", "/v1/models/acme/a%20b"), + ("acme/a?b", "/v1/models/acme/a%3Fb"), + ("acme/a#b", "/v1/models/acme/a%23b"), + ], +) +def test_each_segment_is_percent_encoded_into_exactly_one_path_segment( + model: str, path: str +) -> None: + assert model_run_request(model, {}, None)[0] == path + + +@pytest.mark.parametrize( + "bad", + [ + "flux-dev", # one segment — no provider + "acme/flux/dev", # three — the variant form, not addressable here + "acme/flux/dev/fp8", # four + "acme/..", # traversal + "../flux-dev", + "./flux-dev", + "acme/.", + "a//b", # an empty middle segment + "/flux-dev", # empty provider + "acme/", # empty model + "", + "/", + ], +) +def test_a_malformed_model_id_is_refused_locally(bad: str) -> None: + # Refused before any request: the id *is* the path, so a malformed one + # would otherwise be pasted into a URL and answered by whatever route it + # landed on — a 404 that looks like "no such model" rather than "you passed + # a bad id". + with pytest.raises(ValueError): + model_run_request(bad, {}, None) + with pytest.raises(ValueError): + parse_model_id(bad) + + +def test_a_three_segment_id_says_the_variant_is_not_addressable_yet() -> None: + # The message matters: a `{provider}/{model}/{variant}` id is a real id + # shape, just not one this route takes, and "invalid model id" would send + # the caller looking for a typo. + with pytest.raises(ValueError, match="variant"): + parse_model_id("acme/flux/dev") + + +def test_a_non_string_model_id_is_a_type_error_not_a_value_error() -> None: + # Python's own split: a wrong *type* is a programming error, and folding it + # into ValueError would let `except ValueError` around user input swallow it. + for bad in (None, 3, ["acme", "flux-dev"]): + with pytest.raises(TypeError): + parse_model_id(bad) # type: ignore[arg-type] + + +def test_a_malformed_id_never_reaches_the_server(server) -> None: + with Comfy() as client: + with pytest.raises(ValueError): + client.models.run("acme/flux/dev", ARGS) + assert server.state.model_run_count == 0 def test_run_accepts_any_mapping_and_does_not_alias_the_callers_object(server) -> None: @@ -90,12 +210,76 @@ def test_run_accepts_any_mapping_and_does_not_alias_the_callers_object(server) - caller_args = {"prompt": "a dog"} with Comfy() as client: client.models.run(MODEL, MappingProxyType(caller_args)) - assert server.state.last_model_run_body == {"model": MODEL, "arguments": {"prompt": "a dog"}} + assert server.state.last_model_run_body == {"prompt": "a dog"} _path, body, _headers = model_run_request(MODEL, caller_args, None) - body["arguments"]["prompt"] = "mutated" + body["prompt"] = "mutated" assert caller_args == {"prompt": "a dog"} +# --- which host the run is addressed to --------------------------------- + + +def test_the_default_router_base_url_is_the_public_one() -> None: + assert COMFY_ROUTER_BASE_URL == ROUTER_BASE_URL == "https://api.comfy.org" + assert ROUTER_BASE_URL_ENV_VAR == "COMFY_ROUTER_BASE_URL" + + +def test_a_client_defaults_to_the_public_router(monkeypatch) -> None: + # No `server` fixture here on purpose: that fixture is what points the + # router at the stub, so this asserts the *unconfigured* default. + monkeypatch.delenv(ROUTER_BASE_URL_ENV_VAR, raising=False) + with Comfy(api_key="comfyui-test") as client: + assert client.models.base_url == "https://api.comfy.org" + assert client._low.router_base_url == "https://api.comfy.org" + + +def test_the_router_env_var_redirects_model_runs(monkeypatch, server) -> None: + monkeypatch.setenv(ROUTER_BASE_URL_ENV_VAR, "http://127.0.0.1:9/router") + with Comfy() as client: + assert client.models.base_url == "http://127.0.0.1:9/router" + + +@pytest.mark.parametrize("blank", ["", " "]) +def test_a_blank_router_env_var_means_the_default(monkeypatch, blank: str) -> None: + monkeypatch.setenv(ROUTER_BASE_URL_ENV_VAR, blank) + with Comfy(api_key="comfyui-test") as client: + assert client.models.base_url == COMFY_ROUTER_BASE_URL + + +def test_a_trailing_slash_on_the_router_env_var_is_stripped(monkeypatch) -> None: + # It is concatenated with a path that already starts with `/`, so a kept + # slash would request `//v1/models/...` — a different path to an origin + # server than the one the vendored spec declares. + monkeypatch.setenv(ROUTER_BASE_URL_ENV_VAR, "https://router.example/") + with Comfy(api_key="comfyui-test") as client: + assert client.models.base_url == "https://router.example" + + +@pytest.mark.parametrize("bad", ["not-a-url", "ftp://h", "https://h?x=1", "https://h#f"]) +def test_a_malformed_router_env_var_is_rejected(monkeypatch, bad: str) -> None: + monkeypatch.setenv(ROUTER_BASE_URL_ENV_VAR, bad) + with pytest.raises(ValueError, match=ROUTER_BASE_URL_ENV_VAR): + Comfy(api_key="comfyui-test") + + +def test_the_models_namespace_reports_the_router_not_the_v2_deployment(monkeypatch, server) -> None: + # The distinction this whole binding rests on: jobs and assets resolve + # under COMFY_BASE_URL, model runs under COMFY_ROUTER_BASE_URL, and they + # are different hosts by default. + monkeypatch.setenv(ROUTER_BASE_URL_ENV_VAR, "https://router.example") + with Comfy() as client: + assert client.models.base_url == "https://router.example" + assert client._low.base_url == server.base_url + assert repr(client.models) == "Models(base_url='https://router.example')" + + +async def test_the_async_namespace_reports_the_router_too(monkeypatch, server) -> None: + monkeypatch.setenv(ROUTER_BASE_URL_ENV_VAR, "https://router.example") + async with AsyncComfy() as client: + assert client.models.base_url == "https://router.example" + assert repr(client.models) == "AsyncModels(base_url='https://router.example')" + + # --- the awaitable form is AsyncClient, not a run_async() suffix --------- _SUFFIXED = re.compile(r"(^async_|_async$|_sync$)") @@ -225,6 +409,32 @@ def test_run_carries_the_host_clients_credentials(server) -> None: assert server.state.last_auth_header == "Bearer k-run" +def test_the_key_reaches_the_router_when_it_is_a_separate_origin( + monkeypatch, server, second_server +) -> None: + # The production shape: `COMFY_BASE_URL` and `COMFY_ROUTER_BASE_URL` are + # different hosts. The client's own credential goes to *both* of its + # configured targets — otherwise every real model run would be a 401. + monkeypatch.setenv(ROUTER_BASE_URL_ENV_VAR, second_server.base_url) + with Comfy(api_key="k-router") as client: + client.models.run(MODEL, ARGS) + assert second_server.state.last_auth_header == "Bearer k-router" + assert second_server.state.model_run_count == 1 + # ...and the run went to the router, not to the v2 deployment. + assert server.state.model_run_count == 0 + + +def test_the_key_is_not_sent_to_a_third_origin(monkeypatch, server, second_server) -> None: + # Neither configured target: a server-returned absolute follow-up link + # pointing at `second_server` must still get nothing, and widening the + # credential rule to cover the router origin must not have widened it to + # "any absolute URL". + monkeypatch.setenv(ROUTER_BASE_URL_ENV_VAR, "https://router.example") + with Comfy(api_key="k-not-yours") as client: + client._low.get_job(f"{second_server.base_url}/api/v2/jobs/whatever") + assert second_server.state.last_auth_header == "" + + # --- errors stay on the SDK's own surface -------------------------------- diff --git a/tests/test_models_run_retry.py b/tests/test_models_run_retry.py index 851d199..9cba636 100644 --- a/tests/test_models_run_retry.py +++ b/tests/test_models_run_retry.py @@ -124,7 +124,8 @@ async def post_model_run( # type: ignore[override] return self._attempt(arguments, idempotency_key) -MODEL = "acme/flux/dev" +#: The canonical two-segment `{provider}/{model}` id the route is addressed by. +MODEL = "acme/flux-dev" ARGS = {"prompt": "a cat", "steps": 4} #: Retries the suite can afford to sit through: the schedule is asserted @@ -1078,7 +1079,7 @@ def test_a_paced_hash_mismatch_409_is_still_a_refusal(server) -> None: def test_a_deadline_504_in_routers_own_error_shape_is_still_collected(server) -> None: - # `POST /api/v2/models/run` is fronted by Router, whose error body is + # The model-run route is served by Router, whose error body is # `{detail, error_type}` with the bucket repeated on `X-Comfy-Error-Type` -- # not the v2 `{error: {code}}` envelope. Reading only the envelope would # collapse every real 504 to the status-derived default, and the whole @@ -1126,10 +1127,11 @@ def test_a_router_error_body_keeps_its_detail_as_the_message(server) -> None: def test_the_default_collect_loop_against_a_non_collecting_deployment(server) -> None: - # The cost of defaulting this on, asserted rather than assumed. `POST - # /models/run` is in neither vendored spec, so a deployment may well apply - # the v2 rule instead -- key stays claimed across an unknown-outcome 5xx, - # and the resend is rejected. The caller then sees `422 + # The cost of defaulting this on, asserted rather than assumed. The collect + # rule is `spec/router-openapi.yaml`'s, so it binds Comfy Router -- but + # `COMFY_ROUTER_BASE_URL` can name a deployment that applies the v2 rule + # instead: key stays claimed across an unknown-outcome 5xx, and the resend + # is rejected. The caller then sees `422 # idempotency_key_reuse` in place of the real 504. That is the trade the # default makes; `retry_collectable=False` is the way out of it, and this # pins both halves so neither can change silently. @@ -1189,4 +1191,6 @@ def post_model_run(self, model: str, args: Mapping[str, Any], **kw: Any) -> Any: models = Models(cast(Any, _MutatingLow(client._low)), FAST) assert models.run(MODEL, arguments) == server.state.model_run_result assert server.state.model_run_count == 2 - assert server.state.last_model_run_body["arguments"]["config"] == {"steps": 4} + # The body is the model's native input, unwrapped — the nested value the + # caller mutated between attempts is *not* what the resend carried. + assert server.state.last_model_run_body["config"] == {"steps": 4} diff --git a/tests/test_router_spec_contract.py b/tests/test_router_spec_contract.py index 3d6f4d5..2c408b7 100644 --- a/tests/test_router_spec_contract.py +++ b/tests/test_router_spec_contract.py @@ -1,11 +1,22 @@ -"""The router exception table is the vendored contract's, not this repo's. - -``spec/router-openapi.yaml`` declares the closed error set as -``x-comfy-error-types`` -- one entry per wire value, each with the ``meaning`` -prose the exception docstrings reproduce. Everything here reads that list and -compares it against :mod:`comfy_sdk.router_exceptions`, so the two cannot drift: -the failure this guards against is a Router spec sync landing a new bucket that -then reaches callers as an untyped ``RouterError`` with nothing going red. +"""The router binding is the vendored contract's, not this repo's. + +Two things are read straight out of ``spec/router-openapi.yaml`` rather than +restated here: + +* the closed error set it declares as ``x-comfy-error-types`` -- one entry per + wire value, each with the ``meaning`` prose the exception docstrings + reproduce -- compared against :mod:`comfy_sdk.router_exceptions`; +* the **route** ``post_model_run`` is bound to -- the path whose + ``post.operationId`` is ``runRouterModel``, and the ``servers[0].url`` it is + addressed against -- compared against + :data:`comfy_low.transport._MODEL_RUN_PATH_TEMPLATE` and + :data:`comfy_sdk.COMFY_ROUTER_BASE_URL`. + +Neither is generated, so a Router spec sync is the moment they can drift. The +failures guarded against are a sync landing a new bucket that then reaches +callers as an untyped ``RouterError``, and a sync **moving the path** (the +``/v1`` -> ``/v2`` move already on the roadmap) while the constant the SDK +posts to stays where it was -- with nothing going red either time. That is also why the assertions are written against the file rather than against a list copied out of it. A test that restated the set would pass a sync @@ -23,6 +34,8 @@ import pytest import yaml +from comfy_low.transport import _MODEL_RUN_PATH_TEMPLATE +from comfy_sdk import COMFY_ROUTER_BASE_URL from comfy_sdk.router_exceptions import ( ROUTER_ERROR_TYPES, ROUTER_EXCEPTIONS, @@ -142,3 +155,42 @@ def test_the_spec_states_a_meaning_and_a_tier_for_every_bucket() -> None: assert isinstance(entry, dict), f"x-comfy-error-types entry is not a mapping: {entry!r}" assert entry.get("tier") in {"request", "transport"}, entry assert isinstance(entry.get("meaning"), str) and entry["meaning"].strip(), entry + + +# --- the route the SDK posts a model run to ----------------------------- + + +def test_run_path_matches_vendored_spec() -> None: + """The bound path and host are the spec's, read out of it rather than restated. + + Written as a search for the ``operationId`` rather than a lookup of the + path we expect: looking the path up by name would pass vacuously the day a + sync moves it, which is the one day this test exists for. + """ + doc = yaml.safe_load(ROUTER_SPEC.read_text(encoding="utf-8")) + declared = [ + path + for path, item in (doc.get("paths") or {}).items() + if isinstance(item, dict) + and isinstance(item.get("post"), dict) + and item["post"].get("operationId") == "runRouterModel" + ] + assert declared == [_MODEL_RUN_PATH_TEMPLATE], ( + f"the vendored spec declares runRouterModel at {declared} and the SDK posts to " + f"{_MODEL_RUN_PATH_TEMPLATE!r} -- update comfy_low.transport._MODEL_RUN_PATH_TEMPLATE" + ) + servers = doc.get("servers") or [] + declared_host = servers[0].get("url") if servers else None + assert declared_host == COMFY_ROUTER_BASE_URL, ( + f"the vendored spec's servers[0].url is {declared_host!r} and the SDK defaults to " + f"{COMFY_ROUTER_BASE_URL!r} -- update comfy_low.transport.ROUTER_BASE_URL" + ) + + +def test_the_bound_path_has_exactly_the_two_segments_the_binding_fills() -> None: + # `model_run_request` fills `{provider}` and `{model}` by name; a sync that + # renamed or added a template variable would silently KeyError at call time + # rather than here. + assert _MODEL_RUN_PATH_TEMPLATE.count("{") == 2 + assert "{provider}" in _MODEL_RUN_PATH_TEMPLATE + assert "{model}" in _MODEL_RUN_PATH_TEMPLATE diff --git a/tests/test_sync_async_parity.py b/tests/test_sync_async_parity.py index 380ff8d..ed21e24 100644 --- a/tests/test_sync_async_parity.py +++ b/tests/test_sync_async_parity.py @@ -54,7 +54,7 @@ import comfy_low import comfy_sdk from comfy_sdk import AsyncComfy, Comfy -from comfy_sdk.client import BASE_URL_ENV_VAR +from comfy_sdk.client import BASE_URL_ENV_VAR, ROUTER_BASE_URL_ENV_VAR #: The packages whose public classes make up the surface under test. _PACKAGES = (comfy_sdk, comfy_low) @@ -207,13 +207,19 @@ def _default_deployment() -> Iterator[None]: fixture could adjust the environment. Inheriting an ambient value would let a stray export in a developer's shell turn this whole module into a collection error that has nothing to do with parity. + + Both target variables, for exactly the same reason: + ``_resolve_router_base_url()`` raises the same way on a malformed + ``COMFY_ROUTER_BASE_URL``, and a client construction reads both. """ - saved = os.environ.pop(BASE_URL_ENV_VAR, None) + names = (BASE_URL_ENV_VAR, ROUTER_BASE_URL_ENV_VAR) + saved = {name: os.environ.pop(name, None) for name in names} try: yield finally: - if saved is not None: - os.environ[BASE_URL_ENV_VAR] = saved + for name, value in saved.items(): + if value is not None: + os.environ[name] = value def _run(coro: Any) -> None: diff --git a/tests/test_transport_security.py b/tests/test_transport_security.py index 67b5d81..4f89ed6 100644 --- a/tests/test_transport_security.py +++ b/tests/test_transport_security.py @@ -7,10 +7,21 @@ attached ``Authorization: Bearer `` to *any* absolute URL with no origin check — a malicious or misconfigured server could point a job's follow-up link at an attacker-controlled host and have the client hand it the credential. + +The rule is "one of this client's **configured** targets", not "base_url", and +there are two of them: the ``/api/v2`` deployment and Comfy Router +(``router_base_url``), which is a different host by default and would answer +every model run ``401`` without the key. Both come from the client's own +construction — a constant or an environment variable the operator set — never +from a server response, which is what keeps the second target from widening +what a malicious server can aim the credential at. A third origin still gets +nothing, and that is asserted here in both directions. """ from __future__ import annotations +import pytest + from comfy_low.transport import ComfyLow @@ -43,3 +54,53 @@ def test_relative_path_resolved_against_base_url_still_gets_token(server) -> Non with ComfyLow(server.base_url, api_key="ck_test") as low: low.get_job("whatever") assert server.state.last_auth_header == "Bearer ck_test" + + +# --- the second configured target: Comfy Router -------------------------- + +MODEL = "acme/flux-dev" + + +def test_the_router_origin_gets_the_token_even_when_it_is_not_base_url( + server, second_server +) -> None: + # The production shape: two different hosts, one credential. Without this + # the headline model API would be unauthenticated against a correctly + # configured client. + second_server.state.require_auth = True + with ComfyLow( + server.base_url, api_key="ck_test", router_base_url=second_server.base_url + ) as low: + low.post_model_run(MODEL, {"prompt": "a cat"}) + assert second_server.state.last_auth_header == "Bearer ck_test" + assert second_server.state.last_model_run_path == "/v1/models/acme/flux-dev" + + +def test_a_third_origin_still_gets_no_token_once_the_router_is_trusted( + server, second_server +) -> None: + # The regression this pairs with: broadening the origin check from one + # configured target to two must not broaden it to "any absolute URL". + # `second_server` is neither `base_url` nor `router_base_url` here. + with ComfyLow( + server.base_url, api_key="ck_super_secret", router_base_url="https://router.example" + ) as low: + low.get_job(f"{second_server.base_url}/api/v2/jobs/whatever") + assert second_server.state.last_auth_header == "" + + +def test_the_run_url_does_not_pick_up_the_v2_prefix(server) -> None: + # `_Prepared.url()` prepends `/api/v2` to a relative path; the model-run URL + # is built absolute against `router_base_url` precisely so it does not. + with ComfyLow(server.base_url, api_key="ck_test", router_base_url=server.base_url) as low: + low.post_model_run(MODEL, {"prompt": "a cat"}) + assert server.state.last_model_run_path == "/v1/models/acme/flux-dev" + + +@pytest.mark.parametrize("bad", ["ftp://h", "router.example", "//router.example", ""]) +def test_a_non_http_router_base_url_is_refused_at_construction(bad: str) -> None: + # Otherwise `url()` — which passes a string through only when it starts with + # `http` — would fall through to the `base_url + /api/v2 + ...` branch and + # send an authenticated run at a mangled URL on the wrong surface. + with pytest.raises(ValueError, match="router_base_url"): + ComfyLow("https://cloud.comfy.org", api_key="ck_test", router_base_url=bad)