diff --git a/README.md b/README.md index 286975c..dc4a4aa 100644 --- a/README.md +++ b/README.md @@ -66,6 +66,15 @@ 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. +The SDK identifies itself via a `User-Agent` header (for support and usage +analytics) — this is request metadata only; no other data is collected. Pass +`client_info="my-app"` to append an `app/my-app` token so an integration can +attribute its own traffic: + +```python +client = Comfy("https://api.comfy.org", api_key="ck_...", client_info="my-app") +``` + ## Partner (API) node auth Workflows that use partner/API nodes (Gemini, etc.) need a Comfy API key to diff --git a/src/comfy_low/transport.py b/src/comfy_low/transport.py index 55fd563..1a894ff 100644 --- a/src/comfy_low/transport.py +++ b/src/comfy_low/transport.py @@ -17,11 +17,16 @@ from __future__ import annotations import asyncio +import platform import secrets +import sys from collections.abc import AsyncIterator, Iterator from contextlib import asynccontextmanager, contextmanager +from datetime import datetime, timedelta, timezone +from importlib.metadata import PackageNotFoundError +from importlib.metadata import version as _pkg_version from typing import Any, BinaryIO -from urllib.parse import urlsplit +from urllib.parse import parse_qs, urlsplit import httpx @@ -44,6 +49,30 @@ _DEFAULT_PORTS = {"http": 80, "https": 443} +def _build_user_agent(client_info: str | None) -> str: + """SDK identity sent on every request. This is request metadata (not + telemetry — no phone-home), so adoption is measurable server-side from + request logs. An optional caller-set ``app/{client_info}`` token lets an + integration attribute its own traffic. + """ + try: + ver = _pkg_version("comfy-sdk") + except PackageNotFoundError: + ver = "0" + ua = ( + f"comfy-sdk-python/{ver} " + f"({platform.python_implementation()} {sys.version_info.major}.{sys.version_info.minor})" + ) + if client_info: + # A caller-set token goes verbatim into a header value; reject CR/LF so + # it can never split/inject headers (httpx would reject it anyway, but + # fail fast with a clear message at construction). + if "\r" in client_info or "\n" in client_info: + raise ValueError("client_info must not contain CR or LF characters") + ua += f" app/{client_info}" + return ua + + def _retry_after(resp: httpx.Response) -> int | None: raw = resp.headers.get("Retry-After") if raw is None: @@ -67,10 +96,11 @@ def _origin(url: str) -> tuple[str, str, int | None]: class _Prepared: """Sans-IO request building shared by both transports.""" - def __init__(self, base_url: str, api_key: str | None) -> None: + def __init__(self, base_url: str, api_key: str | None, client_info: str | None = None) -> None: self.base_url = base_url.rstrip("/") self.api_key = api_key self._base_origin = _origin(self.base_url) + self._user_agent = _build_user_agent(client_info) def url(self, path: str) -> str: # A path may be an absolute follow-up link (job.urls.*) or an API path. @@ -81,7 +111,9 @@ def url(self, path: str) -> str: return self.base_url + _API + path def headers(self, url: str, extra: dict[str, str] | None = None) -> dict[str, str]: - h: dict[str, str] = {} + # User-Agent goes on every request (all origins — it is not a + # credential); a caller-supplied one in ``extra`` still wins. + 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: @@ -107,6 +139,29 @@ def parse_or_raise(self, resp: httpx.Response, ok: tuple[int, ...]) -> dict[str, raise error_from_envelope(resp.status_code, body, retry_after=_retry_after(resp)) +def parse_expiry(url: str) -> datetime | None: + """Expiry of a GCS V4 signed URL from its ``X-Goog-Date``/``X-Goog-Expires`` + query params — mirrors the server's own signed-URL-expiry computation. + + Returns ``None`` when either param is missing or malformed, which is the + normal case for a same-origin content URL (self-hosted backend) rather than + a signed one. + """ + query = parse_qs(urlsplit(url).query) + date_vals = query.get("X-Goog-Date") + expires_vals = query.get("X-Goog-Expires") + if not date_vals or not expires_vals: + return None + try: + issued = datetime.strptime(date_vals[0], "%Y%m%dT%H%M%SZ").replace(tzinfo=timezone.utc) + seconds = int(expires_vals[0]) + return issued + timedelta(seconds=seconds) + except (ValueError, OverflowError): + # A malformed date/expires, or an ``expires`` so large the resulting + # datetime overflows, is treated as "no expiry" rather than raising. + return None + + def _new_boundary() -> str: return "----comfy" + secrets.token_hex(16) @@ -138,8 +193,9 @@ def __init__( *, client: httpx.Client | None = None, timeout: float | None = 30.0, + client_info: str | None = None, ) -> None: - self._p = _Prepared(base_url, api_key) + self._p = _Prepared(base_url, api_key, client_info) self._own_client = client is None self._client = client or httpx.Client(timeout=timeout, follow_redirects=True) @@ -304,6 +360,40 @@ def get_asset_content( self._p.parse_or_raise(resp, (200, 206)) yield resp + def get_asset_content_url( + self, asset_id: str, *, timeout: Any = _UNSET + ) -> tuple[str, datetime | None]: + """GET /api/v2/assets/{id}/content without following a redirect. + + A Cloud/serverless backend answers with a redirect to a short-lived + signed URL (object storage); a self-hosted backend serves the bytes + inline (200/206), with no redirect at all. Either way this never reads + the body — only headers are needed to resolve the URL, so a large + asset's bytes are never pulled through this call. + + ``follow_redirects=False`` is passed per-request (the shared client + otherwise follows redirects for every other call) because the + ``Location`` *is* the answer here — following it would fetch bytes we + don't want, and would also run the client's auth-header logic against + the redirect target, which must never see our bearer token. + """ + path = f"/assets/{asset_id}/content" + url = self._p.url(path) + kw: dict[str, Any] = {} + if timeout is not _UNSET: + kw["timeout"] = timeout + with self._client.stream( + "GET", url, headers=self._p.headers(url), follow_redirects=False, **kw + ) as resp: + if resp.has_redirect_location: + location = resp.headers["Location"] + return location, parse_expiry(location) + if resp.status_code in (200, 206): + return url, None + resp.read() + self._p.parse_or_raise(resp, (200, 206)) + raise AssertionError("unreachable: parse_or_raise always raises for a non-ok status") + # -- jobs ------------------------------------------------------------- def post_jobs( self, @@ -379,8 +469,9 @@ def __init__( *, client: httpx.AsyncClient | None = None, timeout: float | None = 30.0, + client_info: str | None = None, ) -> None: - self._p = _Prepared(base_url, api_key) + self._p = _Prepared(base_url, api_key, client_info) self._own_client = client is None self._client = client or httpx.AsyncClient(timeout=timeout, follow_redirects=True) @@ -541,6 +632,27 @@ async def get_asset_content( self._p.parse_or_raise(resp, (200, 206)) yield resp + async def get_asset_content_url( + self, asset_id: str, *, timeout: Any = _UNSET + ) -> tuple[str, datetime | None]: + """See the sync ``get_asset_content_url`` for the redirect/inline split.""" + path = f"/assets/{asset_id}/content" + url = self._p.url(path) + kw: dict[str, Any] = {} + if timeout is not _UNSET: + kw["timeout"] = timeout + async with self._client.stream( + "GET", url, headers=self._p.headers(url), follow_redirects=False, **kw + ) as resp: + if resp.has_redirect_location: + location = resp.headers["Location"] + return location, parse_expiry(location) + if resp.status_code in (200, 206): + return url, None + await resp.aread() + self._p.parse_or_raise(resp, (200, 206)) + raise AssertionError("unreachable: parse_or_raise always raises for a non-ok status") + async def post_jobs( self, workflow: dict[str, Any], diff --git a/src/comfy_sdk/__init__.py b/src/comfy_sdk/__init__.py index c1a2985..d308429 100644 --- a/src/comfy_sdk/__init__.py +++ b/src/comfy_sdk/__init__.py @@ -51,7 +51,7 @@ WorkflowFormatUi, ) from .jobs import AsyncJob, Job -from .outputs import AsyncOutput, Output +from .outputs import AsyncOutput, DownloadUrl, Output from .workflows import Workflow, WorkflowFactory __version__ = "0.1.0" @@ -71,6 +71,7 @@ "AsyncJob", "Output", "AsyncOutput", + "DownloadUrl", # events "Event", "Progress", diff --git a/src/comfy_sdk/client.py b/src/comfy_sdk/client.py index 48abb9c..d6778bd 100644 --- a/src/comfy_sdk/client.py +++ b/src/comfy_sdk/client.py @@ -52,8 +52,9 @@ def __init__( api_key: str | None = None, *, timeout: float | None = 30.0, + client_info: str | None = None, ) -> None: - self._low = ComfyLow(base_url, api_key, timeout=timeout) + self._low = ComfyLow(base_url, api_key, timeout=timeout, client_info=client_info) self.assets = AssetFactory(self._low) self.workflows = WorkflowFactory() self.jobs = JobFactory(self._low) @@ -146,8 +147,9 @@ def __init__( api_key: str | None = None, *, timeout: float | None = 30.0, + client_info: str | None = None, ) -> None: - self._low = AsyncComfyLow(base_url, api_key, timeout=timeout) + self._low = AsyncComfyLow(base_url, api_key, timeout=timeout, client_info=client_info) self.assets = AsyncAssetFactory(self._low) self.workflows = WorkflowFactory() self.jobs = AsyncJobFactory(self._low) diff --git a/src/comfy_sdk/outputs.py b/src/comfy_sdk/outputs.py index 097b1c1..921e621 100644 --- a/src/comfy_sdk/outputs.py +++ b/src/comfy_sdk/outputs.py @@ -2,11 +2,14 @@ An output is an asset: the bytes are retrievable via ``getAssetContent`` (which serves directly or ``302``-redirects to a signed URL) for as long as the job is -retained. ``to_file`` streams to disk in chunks; ``to_bytes`` buffers. +retained. ``to_file`` streams to disk in chunks; ``to_bytes`` buffers; +``get_download_url`` resolves a fetchable URL without transferring any bytes. """ from __future__ import annotations +from dataclasses import dataclass +from datetime import datetime from os import PathLike from pathlib import Path from typing import BinaryIO @@ -17,6 +20,20 @@ _CHUNK = 64 * 1024 +@dataclass(frozen=True) +class DownloadUrl: + """A directly-fetchable URL for one output, e.g. to hand to a downstream + consumer without streaming the bytes through the caller. + + ``url`` is a short-lived, self-authorizing bearer credential in its own + right on backends that hand out signed URLs — whoever holds it can read + the asset until ``expires_at``, no separate auth required. + """ + + url: str + expires_at: datetime | None + + class Output: """A single job output, downloadable via the sync transport.""" @@ -71,6 +88,19 @@ def to_bytes(self, *, range: tuple[int, int] | None = None) -> bytes: buf.extend(chunk) return bytes(buf) + def get_download_url(self) -> DownloadUrl: + """A directly-fetchable URL for this output — never throws. + + On a Cloud/serverless backend this is a short-lived, self-authorizing + signed URL for object storage: anyone holding it can read the bytes + until ``expires_at``. On a self-hosted backend it is the content + endpoint itself (the same auth as every other call still applies), and + ``expires_at`` is ``None``. (A genuine failure — e.g. an unknown output + id — still raises the same typed error as any other call.) + """ + url, expires_at = self._low.get_asset_content_url(self._model.id) + return DownloadUrl(url=url, expires_at=expires_at) + def __repr__(self) -> str: return f"Output(node_id={self.node_id!r}, name={self.name!r}, id={self.id!r})" @@ -123,5 +153,10 @@ async def to_bytes(self, *, range: tuple[int, int] | None = None) -> bytes: buf.extend(chunk) return bytes(buf) + async def get_download_url(self) -> DownloadUrl: + """See the sync ``Output.get_download_url`` for the redirect/inline split.""" + url, expires_at = await self._low.get_asset_content_url(self._model.id) + return DownloadUrl(url=url, expires_at=expires_at) + def __repr__(self) -> str: return f"AsyncOutput(node_id={self.node_id!r}, name={self.name!r}, id={self.id!r})" diff --git a/tests/conftest.py b/tests/conftest.py index 0e17a40..4b8b472 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -47,6 +47,10 @@ class ServerState: # If set, GET /assets/{id}/content responds 302 to this URL instead of # serving bytes directly (simulates a signed-URL redirect to another host). redirect_content_to: str | None = None + # Outputs a succeeded job reports; None means the single default `_OUTPUT`. + # Set to a list of output dicts (see `_output_json`) to test a job with + # multiple outputs, each backed by a distinct asset id. + job_outputs: list[dict] | None = None # --- counters the tests assert on --- upload_count: int = 0 @@ -68,6 +72,7 @@ class ServerState: # Authorization header seen on the most recent request (any method). # `None` before any request; `""` if a request arrived without one. last_auth_header: str | None = None + last_user_agent: str | None = None def _asset_json(asset_id: str, hash_: str, created_new: bool, size: int) -> dict: @@ -105,17 +110,21 @@ def _job_json(job_id: str, status: str, outputs: list[dict] | None = None) -> di } -_OUTPUT = { - "node_id": "13", - "name": "out.png", - "type": "image", - "content_type": "image/png", - "size_bytes": 33, - "id": "asset_out_01", - "hash": None, - "url": "http://example.invalid/out", - "url_expires_at": "2026-07-10T19:20:00Z", -} +def _output_json(node_id: str, asset_id: str) -> dict: + return { + "node_id": node_id, + "name": f"{asset_id}.png", + "type": "image", + "content_type": "image/png", + "size_bytes": 33, + "id": asset_id, + "hash": None, + "url": "http://example.invalid/out", + "url_expires_at": "2026-07-10T19:20:00Z", + } + + +_OUTPUT = _output_json("13", "asset_out_01") def _make_handler(state: ServerState): @@ -142,6 +151,7 @@ def _auth_ok(self) -> bool: # required) so a test can assert the bearer token was — or was # not — attached to a given request. state.last_auth_header = self.headers.get("Authorization", "") + state.last_user_agent = self.headers.get("User-Agent", "") if not state.require_auth: return True return bool(state.last_auth_header) @@ -220,7 +230,10 @@ def _serve_job(self, job_id: str) -> None: state.job_poll_count += 1 if state.job_poll_count >= state.polls_to_succeed: status = state.terminal_status - outputs = [_OUTPUT] if status == "succeeded" else [] + if status == "succeeded": + outputs = state.job_outputs if state.job_outputs is not None else [_OUTPUT] + else: + outputs = [] else: status = "running" outputs = [] diff --git a/tests/test_get_download_url.py b/tests/test_get_download_url.py new file mode 100644 index 0000000..58b8f06 --- /dev/null +++ b/tests/test_get_download_url.py @@ -0,0 +1,170 @@ +"""``get_download_url()`` — a directly-fetchable URL without streaming bytes. + +Covers both shapes a backend can answer with (see ``comfy_low.transport. +get_asset_content_url``): a redirect to a signed URL (Cloud/serverless, backed +by object storage) and an inline 200/206 body (self-hosted, no redirect at +all) — plus that each output on a multi-output job resolves its own distinct +URL, and that the bearer token is never sent to the redirect target because +the redirect is never followed. +""" + +from __future__ import annotations + +from datetime import datetime, timezone + +from conftest import _output_json + +from comfy_low.transport import AsyncComfyLow, ComfyLow, parse_expiry +from comfy_sdk import AsyncComfy, Comfy, DownloadUrl + +_SIGNED_URL = ( + "https://storage.googleapis.com/bucket/object" + "?X-Goog-Algorithm=GOOG4-RSA-SHA256" + "&X-Goog-Credential=example%2F20260710%2Fauto%2Fstorage%2Fgoog4_request" + "&X-Goog-Date=20260710T180000Z" + "&X-Goog-Expires=3600" + "&X-Goog-SignedHeaders=host" + "&X-Goog-Signature=deadbeef" +) + + +def _wf(client: Comfy): + return client.workflows.from_json({"3": {"class_type": "KSampler", "inputs": {}}}) + + +# -- parse_expiry ----------------------------------------------------------- + + +def test_parse_expiry_reads_goog_date_and_expires() -> None: + assert parse_expiry(_SIGNED_URL) == datetime(2026, 7, 10, 19, 0, 0, tzinfo=timezone.utc) + + +def test_parse_expiry_returns_none_when_params_absent() -> None: + assert parse_expiry("https://example.invalid/assets/x/content") is None + + +def test_parse_expiry_returns_none_on_malformed_values() -> None: + bad_date = "https://example.invalid?X-Goog-Date=not-a-date&X-Goog-Expires=60" + bad_expires = "https://example.invalid?X-Goog-Date=20260710T180000Z&X-Goog-Expires=x" + assert parse_expiry(bad_date) is None + assert parse_expiry(bad_expires) is None + + +def test_parse_expiry_returns_none_on_overflowing_expires() -> None: + # A huge expires parses as an int but overflows datetime addition — treat + # it as "no expiry" rather than raising OverflowError. + huge = "https://example.invalid?X-Goog-Date=20260710T180000Z&X-Goog-Expires=999999999999999999" + assert parse_expiry(huge) is None + + +# -- cloud path: redirect to a signed URL ------------------------------------ + + +def test_cloud_redirect_returns_signed_url_and_computed_expiry(server) -> None: + server.state.redirect_content_to = _SIGNED_URL + with ComfyLow(server.base_url) as low: + url, expires_at = low.get_asset_content_url("asset_out_01") + assert url == _SIGNED_URL + assert expires_at == datetime(2026, 7, 10, 19, 0, 0, tzinfo=timezone.utc) + + +async def test_async_cloud_redirect_returns_signed_url_and_computed_expiry(server) -> None: + server.state.redirect_content_to = _SIGNED_URL + async with AsyncComfyLow(server.base_url) as low: + url, expires_at = await low.get_asset_content_url("asset_out_01") + assert url == _SIGNED_URL + assert expires_at == datetime(2026, 7, 10, 19, 0, 0, tzinfo=timezone.utc) + + +def test_output_get_download_url_on_cloud_redirect(server) -> None: + server.state.redirect_content_to = _SIGNED_URL + with Comfy(server.base_url) as client: + job = client.run(_wf(client)) + out = job.get_outputs("13")[0] + download = out.get_download_url() + assert isinstance(download, DownloadUrl) + assert download.url == _SIGNED_URL + assert download.expires_at == datetime(2026, 7, 10, 19, 0, 0, tzinfo=timezone.utc) + + +# -- self-hosted path: inline 200, no redirect ------------------------------- + + +def test_self_hosted_returns_content_url_and_no_expiry(server) -> None: + # No `redirect_content_to` configured -> the stub serves bytes inline + # (200), the never-throws / works-everywhere path. + with ComfyLow(server.base_url) as low: + url, expires_at = low.get_asset_content_url("asset_out_01") + assert url == f"{server.base_url}/api/v2/assets/asset_out_01/content" + assert expires_at is None + + +async def test_async_self_hosted_returns_content_url_and_no_expiry(server) -> None: + async with AsyncComfyLow(server.base_url) as low: + url, expires_at = await low.get_asset_content_url("asset_out_01") + assert url == f"{server.base_url}/api/v2/assets/asset_out_01/content" + assert expires_at is None + + +def test_output_get_download_url_on_self_hosted(server) -> None: + with Comfy(server.base_url) as client: + job = client.run(_wf(client)) + out = job.get_outputs("13")[0] + download = out.get_download_url() # must not throw + assert download.url == f"{server.base_url}/api/v2/assets/asset_out_01/content" + assert download.expires_at is None + + +# -- multiple outputs: each resolves its own distinct URL -------------------- + + +def test_multi_output_job_each_get_download_url_is_distinct(server) -> None: + server.state.job_outputs = [ + _output_json("13", "asset_out_01"), + _output_json("14", "asset_out_02"), + ] + with Comfy(server.base_url) as client: + job = client.run(_wf(client)) + urls = {out.id: out.get_download_url().url for out in job.outputs} + + assert len(job.outputs) == 2 + assert urls["asset_out_01"] == f"{server.base_url}/api/v2/assets/asset_out_01/content" + assert urls["asset_out_02"] == f"{server.base_url}/api/v2/assets/asset_out_02/content" + assert urls["asset_out_01"] != urls["asset_out_02"] + + +async def test_async_multi_output_job_each_get_download_url_is_distinct(server) -> None: + server.state.job_outputs = [ + _output_json("13", "asset_out_01"), + _output_json("14", "asset_out_02"), + ] + async with AsyncComfy(server.base_url) as client: + job = await client.run(_wf(client)) + urls = {} + for out in job.outputs: + download = await out.get_download_url() + urls[out.id] = download.url + + assert len(job.outputs) == 2 + assert urls["asset_out_01"] != urls["asset_out_02"] + + +# -- the redirect is never followed: no bearer token reaches the other host -- + + +def test_redirect_target_never_contacted_bearer_never_leaked(server, second_server) -> None: + server.state.require_auth = True + redirect_url = f"{second_server.base_url}/api/v2/assets/asset_out_01/content" + server.state.redirect_content_to = redirect_url + + with Comfy(server.base_url, api_key="ck_super_secret") as client: + job = client.run(_wf(client)) + out = job.get_outputs("13")[0] + download = out.get_download_url() + + assert download.url == redirect_url + # The original request reached `server` with the bearer token as always... + assert server.state.last_auth_header == "Bearer ck_super_secret" + # ...but the redirect target was never even contacted (it isn't followed), + # so it never saw a request, let alone the credential. + assert second_server.state.last_auth_header is None diff --git a/tests/test_user_agent.py b/tests/test_user_agent.py new file mode 100644 index 0000000..3ef8269 --- /dev/null +++ b/tests/test_user_agent.py @@ -0,0 +1,41 @@ +"""The SDK identifies itself via ``User-Agent`` on every request (request +metadata, not telemetry — no phone-home), so adoption is measurable +server-side. An optional ``client_info`` lets an integration attribute its +own traffic via an ``app/{name}`` token, without clobbering the base identity. +""" + +from __future__ import annotations + +import pytest + +from comfy_sdk import Comfy + + +def _wf(client: Comfy): + return client.workflows.from_json({"3": {"class_type": "KSampler", "inputs": {}}}) + + +def test_user_agent_identifies_the_sdk(server) -> None: + with Comfy(server.base_url) as client: + job = client.submit(_wf(client)) + job.refresh() + ua = server.state.last_user_agent or "" + assert ua.startswith("comfy-sdk-python/") + assert "(" in ua and ")" in ua # runtime segment + assert "app/" not in ua # no client_info set + + +def test_client_info_appends_app_token(server) -> None: + with Comfy(server.base_url, client_info="glary-bot") as client: + job = client.submit(_wf(client)) + job.refresh() + ua = server.state.last_user_agent or "" + assert ua.startswith("comfy-sdk-python/") + assert ua.endswith(" app/glary-bot") + + +def test_client_info_rejects_crlf() -> None: + # A CR/LF in the caller token must never reach the header (no injection). + for bad in ("evil\r\nX-Injected: 1", "line\nbreak", "carriage\rreturn"): + with pytest.raises(ValueError): + Comfy("https://api.comfy.org", client_info=bad)