feat(generate): derive partner model enums from the active openapi spec (BE-3392)

comfy_cli/command/generate/adapters.py

@@ -24,12 +24,14 @@
 import base64
 import mimetypes
 from collections.abc import Callable
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from pathlib import Path
 from typing import Any
+from urllib.parse import quote
 
 import httpx
 
+from comfy_cli.command.generate import spec
 from comfy_cli.command.generate.client import ApiError
 from comfy_cli.command.generate.schema import FlagDef
 
@@ -44,6 +46,12 @@ class Adapter:
 
 # ── Gemini / nano-banana ──────────────────────────────────────────────────
 
+_GEMINI_ENDPOINT_ID = "vertexai/gemini/{model}"
+
+# Fallback only — the active openapi spec's enum wins when it carries one
+# (see _spec_model_flags). Gemini's model lives in the URL path, not the
+# request body, so the spec has no body enum for it today and this tuple is
+# the effective list.
 GEMINI_IMAGE_MODELS = (
     "gemini-2.5-flash-image",
     "gemini-2.5-flash-image-preview",
@@ -128,7 +136,7 @@ def _gemini_decode_sync(body: dict, download: str, request_id: str) -> list[Path
             name="model",
             kind="enum",
             required=False,
-            default="gemini-2.5-flash-image",
+            default=GEMINI_IMAGE_MODELS[0],
             description="Gemini image-model variant.",
             enum=list(GEMINI_IMAGE_MODELS),
         ),
@@ -141,6 +149,11 @@ def _gemini_decode_sync(body: dict, download: str, request_id: str) -> list[Path
 
 # ── Seedance ──────────────────────────────────────────────────────────────
 
+_SEEDANCE_ENDPOINT_ID = "byteplus/api/v3/contents/generations/tasks"
+
+# Fallback only — the active openapi spec's request-body enum wins when it
+# carries one (see _spec_model_flags), so new Seedance releases need a spec
+# refresh, not a CLI release.
 SEEDANCE_MODELS = (
     "seedance-1-0-pro-250528",
     "seedance-1-0-pro-fast-251015",
@@ -176,13 +189,20 @@ def _seedance_image_url(value: str, api_key: str) -> str:
     return upload.upload_path(Path(value).expanduser(), api_key).url
 
 
+def _seedance_default_model() -> str:
+    """Default model when the user didn't pass ``--model`` — the pinned default
+    if the active spec's enum still lists it, else the enum's first entry."""
+    models = spec.model_enum(_SEEDANCE_ENDPOINT_ID) or list(SEEDANCE_MODELS)
+    return SEEDANCE_MODELS[0] if SEEDANCE_MODELS[0] in models else models[0]
+
+
 def _seedance_build_body(values: dict, api_key: str) -> dict[str, Any]:
     content: list[dict[str, Any]] = [{"type": "text", "text": _seedance_text(values)}]
     image = values.get("image")
     if image:
         content.append({"type": "image_url", "image_url": {"url": _seedance_image_url(str(image), api_key)}})
     body: dict[str, Any] = {
-        "model": values.get("model") or SEEDANCE_MODELS[0],
+        "model": values.get("model") or _seedance_default_model(),
         "content": content,
     }
     if "generate_audio" in values:
@@ -206,7 +226,7 @@ def _seedance_build_body(values: dict, api_key: str) -> dict[str, Any]:
             name="model",
             kind="enum",
             required=False,
-            default="seedance-1-0-pro-250528",
+            default=SEEDANCE_MODELS[0],
             description="Seedance model variant.",
             enum=list(SEEDANCE_MODELS),
         ),
@@ -242,13 +262,37 @@ def _seedance_build_body(values: dict, api_key: str) -> dict[str, Any]:
 
 
 _ADAPTERS: dict[str, Adapter] = {
-    "vertexai/gemini/{model}": _gemini_adapter,
-    "byteplus/api/v3/contents/generations/tasks": _seedance_adapter,
+    _GEMINI_ENDPOINT_ID: _gemini_adapter,
+    _SEEDANCE_ENDPOINT_ID: _seedance_adapter,
 }
 
 
+def _spec_model_flags(endpoint_id: str, flags: list[FlagDef]) -> list[FlagDef]:
+    """Return ``flags`` with the ``model`` enum refreshed from the active spec.
+
+    Resolved lazily at lookup time (not import) so a refreshed openapi cache
+    surfaces new partner models with zero code changes; the hardcoded tuples
+    above stay as the fallback when the spec carries no enum. The pinned
+    default is kept while the derived enum still lists it; otherwise the first
+    enum entry takes over (a spec that drops a deprecated default must not
+    break the CLI)."""
+    derived = spec.model_enum(endpoint_id)
+    if not derived:
+        return flags
+    out: list[FlagDef] = []
+    for f in flags:
+        if f.name == "model" and f.kind == "enum":
+            default = f.default if f.default in derived else derived[0]
+            f = replace(f, enum=list(derived), default=default)
+        out.append(f)
+    return out
+
+
 def get(endpoint_id: str) -> Adapter | None:
-    return _ADAPTERS.get(endpoint_id)
+    adapter = _ADAPTERS.get(endpoint_id)
+    if adapter is None:
+        return None
+    return replace(adapter, flags=_spec_model_flags(endpoint_id, adapter.flags))
 
 
 def resolve_path(template: str, values: dict, adapter: Adapter) -> str:
@@ -264,4 +308,11 @@ def resolve_path(template: str, values: dict, adapter: Adapter) -> str:
                 break
     if not val:
         raise ApiError(0, "", f"Missing --{adapter.path_param}: required to fill in the URL path.")
-    return template.replace("{" + adapter.path_param + "}", str(val))
+    # The value may come from a spec-derived enum (refreshable cache), so pin it
+    # to a single path segment: percent-encode reserved characters ("/", "?",
+    # "#", …) and reject dot segments outright — a tampered spec must not be
+    # able to redirect the proxied request via path traversal.
+    val = str(val)
+    if val in (".", ".."):
+        raise ApiError(0, "", f"Invalid --{adapter.path_param} value: {val!r}.")
+    return template.replace("{" + adapter.path_param + "}", quote(val, safe=""))

comfy_cli/command/generate/spec.py

@@ -399,6 +399,88 @@ def get_endpoint(endpoint_id: str) -> Endpoint:
     raise SpecError(_unknown_endpoint_message(endpoint_id))
 
 
+def _extract_enum(prop: dict[str, Any]) -> list[str] | None:
+    """Pull a string enum out of a resolved property schema — directly, from
+    ``items`` (array-typed fields), or from ``anyOf``/``oneOf``/``allOf``
+    variants. ``anyOf``/``oneOf`` branches are unioned (a spec may split the
+    model set across variants) while ``allOf`` branches are intersected (every
+    constraint must hold). Numeric members are coerced to their string form so
+    an unquoted YAML value like ``3.5`` isn't silently dropped. Returns None
+    when no non-empty string enum is found."""
+    enum = prop.get("enum")
+    if isinstance(enum, list):
+        values = [str(v) if isinstance(v, int | float) and not isinstance(v, bool) else v for v in enum]
+        values = [v for v in values if isinstance(v, str)]
+        if values:
+            return values
+    items = prop.get("items")
+    if isinstance(items, dict):
+        found = _extract_enum(items)
+        if found:
+            return found
+    for key in ("anyOf", "oneOf"):
+        variants = prop.get(key)
+        if not isinstance(variants, list):
+            continue
+        merged: list[str] = []
+        for variant in variants:
+            if isinstance(variant, dict):
+                found = _extract_enum(variant)
+                if found:
+                    merged.extend(v for v in found if v not in merged)
+        if merged:
+            return merged
+    all_of = prop.get("allOf")
+    if isinstance(all_of, list):
+        branch_enums = [e for v in all_of if isinstance(v, dict) if (e := _extract_enum(v))]
+        if branch_enums:
+            intersected = [v for v in branch_enums[0] if all(v in b for b in branch_enums[1:])]
+            if intersected:
+                return intersected
+    return None
+
+
+def model_enum(endpoint_id: str, field: str = "model") -> list[str] | None:
+    """Return the model-variant enum the active spec carries for
+    ``endpoint_id``'s ``field`` request property, or None when the spec has no
+    enum there (callers fall back to their hardcoded lists).
+
+    The request schema is already ``$ref``-resolved by ``_resolve``, so a plain
+    property walk suffices. Reading from the active spec (user cache when
+    fresh, else the vendored copy) means a spec refresh surfaces new partner
+    models with zero code changes."""
+    try:
+        endpoint = get_endpoint(endpoint_id)
+    except SpecError:
+        return None
+    prop = _find_property(endpoint.request_schema or {}, field)
+    if prop is None:
+        return None
+    return _extract_enum(prop)
+
+
+def _find_property(schema: dict[str, Any], field: str) -> dict[str, Any] | None:
+    """Locate ``field`` in ``schema['properties']``, descending into top-level
+    ``allOf``/``anyOf``/``oneOf`` composition when the schema carries no direct
+    match — a composed request body must not silently defeat the spec-derived
+    enum and fall back to the hardcoded list."""
+    props = schema.get("properties")
+    if isinstance(props, dict):
+        prop = props.get(field)
+        if isinstance(prop, dict):
+            return prop
+    for key in ("allOf", "anyOf", "oneOf"):
+        variants = schema.get(key)
+        if not isinstance(variants, list):
+            continue
+        for variant in variants:
+            if isinstance(variant, dict):
+                found = _find_property(variant, field)
+                if found is not None:
+                    return found
+    return None
+
+
 def _unknown_endpoint_message(endpoint_id: str) -> str:
     """Build a helpful error suggesting close matches."""
     import difflib

tests/comfy_cli/command/generate/test_adapters.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import base64
+import json
 
 import httpx
 import pytest
@@ -127,6 +128,18 @@ def test_gemini_resolve_path_substitutes_model():
     assert url == "/proxy/vertexai/gemini/gemini-2.5-flash-image"
 
 
+def test_resolve_path_pins_value_to_one_segment():
+    """A path-param value can come from a refreshable spec cache — reserved
+    characters are percent-encoded and dot segments rejected, so a tampered
+    enum can't redirect the proxied request."""
+    ep = spec.get_endpoint("nano-banana")
+    adapter = adapters.get(ep.id)
+    url = adapters.resolve_path(ep.path, {"model": "../../admin?x=1#f"}, adapter)
+    assert url == "/proxy/vertexai/gemini/..%2F..%2Fadmin%3Fx%3D1%23f"
+    with pytest.raises(adapters.ApiError):
+        adapters.resolve_path(ep.path, {"model": ".."}, adapter)
+
+
 def test_gemini_send_request_hits_substituted_path(monkeypatch):
     captured = {}
 
@@ -256,6 +269,108 @@ def fake_post(url, *, json=None, headers=None, timeout=None, **_kw):
     assert captured["json"]["content"][0]["type"] == "text"
 
 
+# ── Spec-derived model enums ──────────────────────────────────────────────
+
+
+def _seedance_fixture_spec(model_prop: dict) -> str:
+    """A minimal JSON spec carrying the byteplus tasks endpoint with the given
+    ``model`` property schema — what a refreshed api.comfy.org spec looks like."""
+    return json.dumps(
+        {
+            "openapi": "3.1.0",
+            "servers": [{"url": "https://api.comfy.org"}],
+            "paths": {
+                "/proxy/byteplus/api/v3/contents/generations/tasks": {
+                    "post": {
+                        "summary": "Seedance video generation tasks",
+                        "requestBody": {
+                            "content": {
+                                "application/json": {
+                                    "schema": {
+                                        "type": "object",
+                                        "properties": {
+                                            "model": model_prop,
+                                            "content": {"type": "array", "items": {"type": "object"}},
+                                        },
+                                        "required": ["model", "content"],
+                                    }
+                                }
+                            }
+                        },
+                        "responses": {"200": {"content": {"application/json": {"schema": {"type": "object"}}}}},
+                    }
+                }
+            },
+        }
+    )
+
+
+@pytest.fixture
+def _spec_cache(monkeypatch, tmp_path):
+    """Redirect the user spec cache to tmp and clean module caches afterwards."""
+    monkeypatch.setattr(spec, "_USER_CACHE", tmp_path / "openapi-cache.yml")
+    try:
+        yield
+    finally:
+        spec.load_raw_spec.cache_clear()
+        spec._registry.cache_clear()
+
+
+def test_seedance_models_follow_refreshed_spec(_spec_cache):
+    """After a spec refresh (simulated via write_cache), the seedance schema
+    lists the spec's models — a new partner release needs zero code changes."""
+    new_models = ["seedance-2-0-260128", "seedance-2-0-fast-260128", "seedance-1-0-pro-250528"]
+    spec.write_cache(_seedance_fixture_spec({"type": "string", "enum": new_models}))
+    ep = spec.get_endpoint("seedance")
+    model = next(f for f in schema.flags_for(ep) if f.name == "model")
+    assert model.enum == new_models
+    # The pinned default survives while the spec still lists it.
+    assert model.default == "seedance-1-0-pro-250528"
+
+
+def test_seedance_default_falls_back_when_pinned_model_dropped(_spec_cache):
+    """A spec that drops the deprecated pinned default must not break: the
+    first enum entry takes over, in the flag default and in build_body."""
+    new_models = ["seedance-2-0-260128", "seedance-2-0-fast-260128"]
+    spec.write_cache(_seedance_fixture_spec({"type": "string", "enum": new_models}))
+    ep = spec.get_endpoint("seedance")
+    model = next(f for f in schema.flags_for(ep) if f.name == "model")
+    assert model.default == "seedance-2-0-260128"
+    body = adapters._seedance_build_body({"prompt": "a wave"}, api_key="k")
+    assert body["model"] == "seedance-2-0-260128"
+
+
+def test_seedance_models_fall_back_to_hardcoded_without_spec_enum(_spec_cache):
+    """A spec whose model property carries no enum keeps the hardcoded tuple."""
+    spec.write_cache(_seedance_fixture_spec({"type": "string"}))
+    ep = spec.get_endpoint("seedance")
+    model = next(f for f in schema.flags_for(ep) if f.name == "model")
+    assert model.enum == list(adapters.SEEDANCE_MODELS)
+    assert model.default == adapters.SEEDANCE_MODELS[0]
+    body = adapters._seedance_build_body({"prompt": "a wave"}, api_key="k")
+    assert body["model"] == adapters.SEEDANCE_MODELS[0]
+
+
+def test_validation_error_lists_spec_derived_models(_spec_cache):
+    """Flag validation rejects an unknown model naming the derived enum, so
+    agents introspecting the error see the fresh list."""
+    new_models = ["seedance-2-0-260128", "seedance-2-0-fast-260128"]
+    spec.write_cache(_seedance_fixture_spec({"type": "string", "enum": new_models}))
+    ep = spec.get_endpoint("seedance")
+    flags = schema.flags_for(ep)
+    with pytest.raises(schema.SchemaError, match="seedance-2-0-260128"):
+        schema.parse_args(flags, ["--prompt", "x", "--model", "bogus"])
+
+
+def test_gemini_models_keep_hardcoded_fallback():
+    """Gemini's model is a URL path param — the spec has no body enum, so the
+    hardcoded tuple (and its pinned default) stays in effect."""
+    ep = spec.get_endpoint("nano-banana")
+    model = next(f for f in schema.flags_for(ep) if f.name == "model")
+    assert model.enum == list(adapters.GEMINI_IMAGE_MODELS)
+    assert model.default == adapters.GEMINI_IMAGE_MODELS[0]
+
+
 # ── Seedance polling ──────────────────────────────────────────────────────