diff --git a/comfy_cli/command/generate/app.py b/comfy_cli/command/generate/app.py index ab6ab7fb..9b867bfc 100644 --- a/comfy_cli/command/generate/app.py +++ b/comfy_cli/command/generate/app.py @@ -501,16 +501,43 @@ def _schema(extra_args: list[str]) -> None: _show_schema_help(ep) +def _fetch_spec(url: str) -> httpx.Response: + with httpx.Client(timeout=30.0, follow_redirects=True) as cli: + r = cli.get(url, headers={"Comfy-Env": "comfy-cli", "User-Agent": "comfy-cli/api"}) + r.raise_for_status() + return r + + def _refresh() -> None: - url = spec.base_url() + "/openapi.yml" + base = spec.base_url() + # The live spec is served at ``/openapi`` (no extension, JSON body). Older / + # custom ``COMFY_API_BASE_URL`` deployments may still serve ``/openapi.yml``, + # so fall back to it on a 404 to keep those working. + primary, fallback = base + "/openapi", base + "/openapi.yml" + fetched_from = primary try: - with httpx.Client(timeout=30.0, follow_redirects=True) as cli: - r = cli.get(url, headers={"Comfy-Env": "comfy-cli", "User-Agent": "comfy-cli/api"}) - r.raise_for_status() + try: + r = _fetch_spec(primary) + except httpx.HTTPStatusError as e: + if e.response.status_code != 404: + raise + fetched_from = fallback + r = _fetch_spec(fallback) except httpx.HTTPError as e: - rprint(f"[bold red]Failed to fetch {url}: {e}[/bold red]") + rprint(f"[bold red]Failed to fetch {fetched_from}: {e}[/bold red]") raise typer.Exit(code=1) - path = spec.write_cache(r.text) + + # Validate before caching so a 200-with-garbage response never poisons the + # ~/.comfy/openapi-cache.yml cache (used for CACHE_TTL_SECONDS by every + # subsequent `comfy generate`). + body = r.text + try: + spec.validate_spec_text(body) + except spec.SpecError as e: + rprint(f"[bold red]Refusing to cache spec from {fetched_from}: {e}[/bold red]") + raise typer.Exit(code=1) + + path = spec.write_cache(body) rprint(f"[bold green]Refreshed model catalog at {path}[/bold green]") diff --git a/comfy_cli/command/generate/spec.py b/comfy_cli/command/generate/spec.py index f0052e37..ce8e0d72 100644 --- a/comfy_cli/command/generate/spec.py +++ b/comfy_cli/command/generate/spec.py @@ -412,6 +412,26 @@ def _unknown_endpoint_message(endpoint_id: str) -> str: return msg +def validate_spec_text(text: str) -> dict[str, Any]: + """Parse a raw openapi spec body with the same loader ``load_raw_spec`` uses + and require a top-level ``paths`` mapping. + + The body may be YAML or JSON — JSON is a subset of YAML 1.2, and + ``_YamlLoader`` only restricts bool resolution, so a JSON spec (as served at + ``api.comfy.org/openapi``) parses. Raises :class:`SpecError` if the body does + not parse or lacks ``paths``; callers use this to avoid caching a + 200-with-garbage response, which would poison the on-disk cache for + ``CACHE_TTL_SECONDS``. + """ + try: + parsed = yaml.load(text, Loader=_YamlLoader) + except yaml.YAMLError as e: + raise SpecError(f"spec did not parse: {e}") from e + if not isinstance(parsed, dict) or not isinstance(parsed.get("paths"), dict): + raise SpecError("spec has no top-level 'paths' mapping") + return parsed + + def write_cache(yaml_text: str) -> Path: """Write `yaml_text` to the user cache, ensuring the parent dir exists.""" _USER_CACHE.parent.mkdir(parents=True, exist_ok=True) diff --git a/tests/comfy_cli/command/generate/test_app.py b/tests/comfy_cli/command/generate/test_app.py index 9fd7e3be..751eab94 100644 --- a/tests/comfy_cli/command/generate/test_app.py +++ b/tests/comfy_cli/command/generate/test_app.py @@ -23,6 +23,22 @@ def disable_tracking_prompt(monkeypatch): monkeypatch.setattr("comfy_cli.tracking.track_event", lambda *a, **kw: None) +@pytest.fixture(autouse=True) +def _isolate_spec_caches(): + """spec.load_raw_spec / spec._registry are process-global lru_caches keyed off + the module-level _USER_CACHE path. The refresh tests below monkeypatch that + path to a temp cache; base_url() loads it into those caches, so without a + teardown clear a temp spec leaks into later tests (empty registry → spurious + 'No models match' failures). Clear around every test to keep them isolated.""" + from comfy_cli.command.generate import spec as _spec + + _spec.load_raw_spec.cache_clear() + _spec._registry.cache_clear() + yield + _spec.load_raw_spec.cache_clear() + _spec._registry.cache_clear() + + @pytest.fixture def runner(): return CliRunner() @@ -388,9 +404,21 @@ def test_resume_with_download(runner, api_key, tmp_path, monkeypatch): # ─── refresh ───────────────────────────────────────────────────────────── +# A minimal but structurally-valid openapi spec served as a JSON body — mirrors +# how the live api.comfy.org/openapi endpoint responds (JSON, not YAML). It +# carries a real curated endpoint so the cache round-trips through the registry. +_VALID_JSON_SPEC = ( + '{"openapi":"3.1.0","servers":[{"url":"https://api.comfy.org"}],' + '"paths":{"/proxy/openai/images/generations":{"post":{"summary":"Create image",' + '"requestBody":{"content":{"application/json":{"schema":{"type":"object",' + '"properties":{"prompt":{"type":"string"}}}}}},' + '"responses":{"200":{"content":{"application/json":{"schema":{"type":"object"}}}}}}}}}' +) -def test_refresh_writes_cache(runner, monkeypatch, tmp_path): - captured = {} + +def _fake_client_factory(responder, calls): + """Build a stand-in httpx.Client whose ``get`` delegates to ``responder(url)`` + and records each requested URL in ``calls``.""" class FakeClient: def __init__(self, *a, **kw): @@ -403,39 +431,74 @@ def __exit__(self, *a): pass def get(self, url, headers=None): - captured["url"] = url - captured["headers"] = headers or {} - return httpx.Response( - 200, - text="openapi: 3.0.0\n", - request=httpx.Request("GET", url), - ) - - monkeypatch.setattr(gen_app.httpx, "Client", FakeClient) + calls.append((url, headers or {})) + return responder(url) + + return FakeClient + + +def test_refresh_hits_openapi_and_writes_cache(runner, monkeypatch, tmp_path): + """(a) refresh fetches ``/openapi`` and caches a valid JSON spec body.""" + calls = [] + + def responder(url): + return httpx.Response(200, text=_VALID_JSON_SPEC, request=httpx.Request("GET", url)) + + monkeypatch.setattr(gen_app.httpx, "Client", _fake_client_factory(responder, calls)) monkeypatch.setattr("comfy_cli.command.generate.spec._USER_CACHE", tmp_path / "openapi-cache.yml") r = runner.invoke(cli_app, ["generate", "refresh"]) assert r.exit_code == 0, r.stdout assert "Refreshed" in r.stdout assert (tmp_path / "openapi-cache.yml").exists() - assert captured["headers"].get("Comfy-Env") == "comfy-cli" + # First (and only) fetch targets the extension-less /openapi path. + assert len(calls) == 1 + assert calls[0][0].endswith("/openapi") + assert calls[0][1].get("Comfy-Env") == "comfy-cli" -def test_refresh_network_failure(runner, monkeypatch): - class FakeClient: - def __init__(self, *a, **kw): - pass +def test_refresh_falls_back_to_openapi_yml_on_404(runner, monkeypatch, tmp_path): + """(b) a 404 on ``/openapi`` falls back to ``/openapi.yml``.""" + calls = [] - def __enter__(self): - return self + def responder(url): + if url.endswith("/openapi"): + return httpx.Response(404, text='{"message":"Not Found"}', request=httpx.Request("GET", url)) + return httpx.Response(200, text=_VALID_JSON_SPEC, request=httpx.Request("GET", url)) + + monkeypatch.setattr(gen_app.httpx, "Client", _fake_client_factory(responder, calls)) + monkeypatch.setattr("comfy_cli.command.generate.spec._USER_CACHE", tmp_path / "openapi-cache.yml") + + r = runner.invoke(cli_app, ["generate", "refresh"]) + assert r.exit_code == 0, r.stdout + assert (tmp_path / "openapi-cache.yml").exists() + assert [c[0].rsplit("/", 1)[-1] for c in calls] == ["openapi", "openapi.yml"] - def __exit__(self, *a): - pass - def get(self, *a, **kw): - raise httpx.ConnectError("no net") +@pytest.mark.parametrize("body", ["not: [valid: yaml", '{"openapi":"3.1.0"}']) +def test_refresh_invalid_body_exits_and_leaves_cache_untouched(runner, monkeypatch, tmp_path, body): + """(c) a non-parsing body OR a 200 with no ``paths`` exits 1 without writing.""" + cache = tmp_path / "openapi-cache.yml" + cache.write_text("openapi: 3.0.0\npaths: {}\n", encoding="utf-8") # pre-existing good cache + before = cache.read_text(encoding="utf-8") + + def responder(url): + return httpx.Response(200, text=body, request=httpx.Request("GET", url)) + + monkeypatch.setattr(gen_app.httpx, "Client", _fake_client_factory(responder, [])) + monkeypatch.setattr("comfy_cli.command.generate.spec._USER_CACHE", cache) + + r = runner.invoke(cli_app, ["generate", "refresh"]) + assert r.exit_code == 1 + assert "Refusing to cache" in r.stdout + assert cache.read_text(encoding="utf-8") == before # untouched + + +def test_refresh_network_failure(runner, monkeypatch): + def responder(url): + raise httpx.ConnectError("no net") - monkeypatch.setattr(gen_app.httpx, "Client", FakeClient) + monkeypatch.setattr(gen_app.httpx, "Client", _fake_client_factory(responder, [])) r = runner.invoke(cli_app, ["generate", "refresh"]) assert r.exit_code == 1 assert "Failed to fetch" in r.stdout diff --git a/tests/comfy_cli/command/generate/test_spec.py b/tests/comfy_cli/command/generate/test_spec.py index 2ba407b7..6e6c5e79 100644 --- a/tests/comfy_cli/command/generate/test_spec.py +++ b/tests/comfy_cli/command/generate/test_spec.py @@ -3,6 +3,49 @@ from comfy_cli.command.generate import spec +# A minimal JSON spec body — the shape api.comfy.org/openapi actually serves +# (JSON, not YAML). JSON is a subset of YAML 1.2 so it loads via _YamlLoader. +_VALID_JSON_SPEC = ( + '{"openapi":"3.1.0","servers":[{"url":"https://api.comfy.org"}],' + '"paths":{"/proxy/openai/images/generations":{"post":{"summary":"Create image",' + '"requestBody":{"content":{"application/json":{"schema":{"type":"object",' + '"properties":{"prompt":{"type":"string"}}}}}},' + '"responses":{"200":{"content":{"application/json":{"schema":{"type":"object"}}}}}}}}}' +) + + +def test_validate_spec_text_accepts_json_and_rejects_bad_bodies(): + parsed = spec.validate_spec_text(_VALID_JSON_SPEC) + assert isinstance(parsed["paths"], dict) + for bad in ("not: [valid: yaml", '{"openapi":"3.1.0"}', "null", "[]"): + try: + spec.validate_spec_text(bad) + except spec.SpecError: + pass + else: + raise AssertionError(f"expected SpecError for {bad!r}") + + +def test_cached_json_spec_round_trips(monkeypatch, tmp_path): + """(d) a cached JSON body round-trips through load_raw_spec() and + get_endpoint() resolves an endpoint from it.""" + cache = tmp_path / "openapi-cache.yml" + cache.write_text(_VALID_JSON_SPEC, encoding="utf-8") + monkeypatch.setattr(spec, "_USER_CACHE", cache) + spec.load_raw_spec.cache_clear() + spec._registry.cache_clear() + try: + assert spec.active_spec_path() == cache + raw = spec.load_raw_spec() + assert isinstance(raw["paths"], dict) + ep = spec.get_endpoint("openai/images/generations") + assert ep.path == "/proxy/openai/images/generations" + assert ep.method == "post" + finally: + # Don't leak the temp spec into other tests via the module-level caches. + spec.load_raw_spec.cache_clear() + spec._registry.cache_clear() + def test_registry_loads_and_has_entries(): eps = spec.list_endpoints()