diff --git a/comfy_cli/cql/__init__.py b/comfy_cli/cql/__init__.py index e035a987..02fccd05 100644 --- a/comfy_cli/cql/__init__.py +++ b/comfy_cli/cql/__init__.py @@ -1,10 +1,10 @@ -"""object_info loader — normalize ComfyUI's ``/object_info`` into typed rows. +"""CQL public surface. Public entry points: - load_graph(input_path=..., host=..., port=...) -> dict + resilient_load_object_info(mode=..., host=..., port=..., input_path=...) -> dict """ from comfy_cli.cql.errors import CQLRuntimeError -from comfy_cli.cql.loader import load_graph +from comfy_cli.cql.loader import resilient_load_object_info -__all__ = ["CQLRuntimeError", "load_graph"] +__all__ = ["CQLRuntimeError", "resilient_load_object_info"] diff --git a/comfy_cli/cql/loader.py b/comfy_cli/cql/loader.py index f432882e..58d51fce 100644 --- a/comfy_cli/cql/loader.py +++ b/comfy_cli/cql/loader.py @@ -1,19 +1,21 @@ -"""Build a CQL-shaped graph dict from sources. - -Sources, in priority order: - -1. A local file (``--input path``). May be: - - A raw ``object_info`` JSON dump (the response from ``/object_info``). - - An API-format workflow JSON. - - An already-shaped CQL graph (``{"nodes": [...], "inputs": [...]}``). -2. A local ComfyUI server's ``/object_info`` endpoint (``--host`` / ``--port``). - -The loader is intentionally permissive: anything dict-shaped that looks like -one of those formats is normalized into ``{"nodes": [...], "inputs": [...], -"categories": [...]}`` so the engine can run uniformly. - -This module performs only local I/O. Network calls hit ``http://host:port`` -and are short-circuited when no host is provided. +"""Shape and load CQL ``object_info`` graphs. + +This module contains two things: + +- ``normalize`` — turn any supported input (a raw ``object_info`` dump, an + API-format workflow, or an already-shaped CQL graph) into the uniform + ``{"nodes": [...], "inputs": [...], "categories": [...]}`` dict the engine + runs on. It is intentionally permissive: anything dict-shaped that looks + like one of those formats is accepted. +- ``resilient_load_object_info`` — a cache + refresh-retry + stale-fallback + wrapper over the engine's loaders (``comfy_cli.cql.engine._load_from_file`` + / ``_load_from_target``). It auto-caches every successful fetch per host, + retries once after a token refresh on failure, and falls back to the cached + dump (with a stderr warning) when the retry still fails. + +The live network fetch and its security guards (loopback check, no-redirect +opener, byte cap, cloud HTTPS+auth) live in ``comfy_cli.cql.engine`` — this +module never opens a socket itself. """ from __future__ import annotations @@ -22,95 +24,10 @@ import json import os import sys -import urllib.error -import urllib.parse -import urllib.request from pathlib import Path from typing import Any -from comfy_cli.cql._net import is_loopback_host from comfy_cli.cql.errors import CQLRuntimeError -from comfy_cli.http import NoRedirectHandler - -# Cap raw bytes read from disk or the network. Real `object_info` dumps are a -# few MB; anything past 256 MiB is almost certainly a wrong path or a hostile -# server and would just OOM the CLI before json.loads even fails. -MAX_INPUT_BYTES = 256 * 1024 * 1024 - - -_LOADER_OPENER = urllib.request.build_opener(NoRedirectHandler()) - - -def load_graph( - *, - input_path: str | None = None, - host: str | None = None, - port: int | None = None, - timeout: float = 5.0, -) -> dict[str, Any]: - if input_path: - return _load_from_file(input_path) - if host and port: - return _load_from_server(host, int(port), timeout=timeout) - raise CQLRuntimeError( - "no graph source available", - details={"hint": "pass --input or --host/--port pointing at a ComfyUI server"}, - ) - - -def _load_from_file(path: str) -> dict[str, Any]: - p = Path(path).expanduser() - try: - size = p.stat().st_size - except OSError as e: - raise CQLRuntimeError(f"cannot stat {p}: {e}") from e - if size > MAX_INPUT_BYTES: - raise CQLRuntimeError( - f"{p} is {size} bytes, exceeds MAX_INPUT_BYTES={MAX_INPUT_BYTES}", - details={"hint": "shrink the input or raise MAX_INPUT_BYTES in cql.loader"}, - ) - try: - raw = p.read_text(encoding="utf-8") - except OSError as e: - raise CQLRuntimeError(f"cannot read {p}: {e}") from e - try: - data = json.loads(raw) - except json.JSONDecodeError as e: - raise CQLRuntimeError(f"{p} is not valid JSON: {e}") from e - return normalize(data) - - -def _load_from_server(host: str, port: int, *, timeout: float) -> dict[str, Any]: - url = f"http://{host}:{port}/object_info" - # Refuse anything that isn't a localhost-ish target — we don't want CQL - # silently sending traffic to a remote box. (Cloud CQL goes through its - # own path; this loader is local-only by design.) - parsed = urllib.parse.urlsplit(url) - hostname = (parsed.hostname or "").strip().lower() - if not is_loopback_host(hostname): - raise CQLRuntimeError( - f"refusing non-loopback CQL server target: {host}", - details={"hint": "pass --input for remote object_info dumps"}, - ) - try: - with _LOADER_OPENER.open(url, timeout=timeout) as resp: - # Bounded read so a misbehaving server can't OOM us. - raw = resp.read(MAX_INPUT_BYTES + 1) - if len(raw) > MAX_INPUT_BYTES: - raise CQLRuntimeError( - f"server response exceeds MAX_INPUT_BYTES={MAX_INPUT_BYTES}", - details={"host": host, "port": port}, - ) - data = json.loads(raw) - except urllib.error.URLError as e: - raise CQLRuntimeError( - f"failed to reach {url}: {e.reason if hasattr(e, 'reason') else e}", - details={"host": host, "port": port}, - ) from e - except (json.JSONDecodeError, OSError) as e: - raise CQLRuntimeError(f"server returned invalid object_info: {e}") from e - return normalize(data) - # ---- normalization -------------------------------------------------------- diff --git a/tests/comfy_cli/cql/test_loader.py b/tests/comfy_cli/cql/test_loader.py index d5b9aa23..543a8ef7 100644 --- a/tests/comfy_cli/cql/test_loader.py +++ b/tests/comfy_cli/cql/test_loader.py @@ -2,13 +2,10 @@ from __future__ import annotations -import json - import pytest -from comfy_cli.cql import loader from comfy_cli.cql.errors import CQLRuntimeError -from comfy_cli.cql.loader import _load_from_server, load_graph, normalize +from comfy_cli.cql.loader import normalize OBJECT_INFO = { "KSampler": { @@ -97,53 +94,6 @@ def test_normalize_preshaped_graph_pass_through(): assert g["nodes"][0]["name"] == "Foo" -def test_load_graph_from_file(tmp_path): - p = tmp_path / "object_info.json" - p.write_text(json.dumps(OBJECT_INFO)) - g = load_graph(input_path=str(p)) - assert {n["name"] for n in g["nodes"]} == {"KSampler", "CheckpointLoaderSimple"} - - -def test_load_graph_missing_source_raises(): - with pytest.raises(CQLRuntimeError): - load_graph() - - -def test_load_graph_bad_json(tmp_path): - p = tmp_path / "broken.json" - p.write_text("{ not json") - with pytest.raises(CQLRuntimeError): - load_graph(input_path=str(p)) - - def test_normalize_rejects_garbage(): with pytest.raises(CQLRuntimeError): normalize({"foo": 1, "bar": "baz"}) - - -class _FakeResp: - def __init__(self, payload: bytes): - self._payload = payload - - def __enter__(self): - return self - - def __exit__(self, *exc): - return False - - def read(self, _n=None): - return self._payload - - -def test_load_from_server_refuses_non_loopback_host(): - # SSRF guard: a public host must never be fetched by the local loader. - with pytest.raises(CQLRuntimeError, match="non-loopback"): - _load_from_server("example.com", 8188, timeout=0.1) - - -def test_load_from_server_accepts_loopback(monkeypatch): - # 127.0.0.1 passes the guard and proceeds to the fetch (mocked here). - payload = json.dumps(OBJECT_INFO).encode("utf-8") - monkeypatch.setattr(loader._LOADER_OPENER, "open", lambda *a, **k: _FakeResp(payload)) - g = _load_from_server("127.0.0.1", 8188, timeout=0.1) - assert {n["name"] for n in g["nodes"]} == {"KSampler", "CheckpointLoaderSimple"}