fix: comfy generate refresh fetches /openapi with /openapi.yml fallback (BE-3391)

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]")
 
 

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)

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

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()