feat: allow --comfy-api-base target ephemeral testenvs

Signed-off-by: bigcat88 <bigcat88@icloud.com>

comfy/comfy_api_env.py

@@ -0,0 +1,46 @@
+"""Runtime config the frontend reads from /features to follow --comfy-api-base.
+
+For a non-prod comfy.org backend (staging or an ephemeral preview env), "/features" exposes the api and
+platform base so the frontend talks to it without a rebuild; the frontend picks the Firebase project from the api base.
+Prod bases are left alone and keep their build-time defaults.
+"""
+
+from typing import Any
+from urllib.parse import urlparse
+
+from comfy.cli_args import args
+
+# Staging and the ephemeral preview envs ("testenvs") are one tier: same dev Firebase project and platform.
+_STAGING_API_HOST = "stagingapi.comfy.org"
+_TESTENV_HOST_SUFFIX = ".testenvs.comfy.org"
+_STAGING_PLATFORM_BASE_URL = "https://stagingplatform.comfy.org"
+
+
+def _is_staging_tier(host: str) -> bool:
+    return host == _STAGING_API_HOST or host.endswith(_TESTENV_HOST_SUFFIX)
+
+
+def normalize_comfy_api_base(url: str) -> str:
+    """Rewrite a testenv's friendly main host to its comfy-api '-registry' sibling."""
+    parsed = urlparse(url)
+    host = parsed.hostname or ""
+    if not host.endswith(_TESTENV_HOST_SUFFIX):
+        return url
+    label = host[: -len(_TESTENV_HOST_SUFFIX)]
+    if label.endswith("-registry"):
+        return url
+    return f"{parsed.scheme or 'https'}://{label}-registry{_TESTENV_HOST_SUFFIX}"
+
+
+def frontend_config_for_base(base_url: str) -> dict[str, Any] | None:
+    """The /features overrides for a staging-tier base, or None for prod."""
+    if not _is_staging_tier(urlparse(base_url).hostname or ""):
+        return None
+    return {
+        "comfy_api_base_url": normalize_comfy_api_base(base_url).rstrip("/"),
+        "comfy_platform_base_url": _STAGING_PLATFORM_BASE_URL,
+    }
+
+
+def get_frontend_config() -> dict[str, Any] | None:
+    return frontend_config_for_base(getattr(args, "comfy_api_base", "") or "")

comfy_api/feature_flags.py

@@ -9,6 +9,7 @@ import logging
 from typing import Any, TypedDict
 
 from comfy.cli_args import args
+from comfy.comfy_api_env import get_frontend_config
 
 
 class FeatureFlagInfo(TypedDict):
@@ -162,4 +163,11 @@ def get_server_features() -> dict[str, Any]:
     Returns:
         Dictionary of server feature flags
     """
-    return SERVER_FEATURE_FLAGS.copy()
+    features = SERVER_FEATURE_FLAGS.copy()
+    # When --comfy-api-base targets a staging-tier comfy.org backend (the staging api host or an ephemeral testenv),
+    # surface the api + platform base so the frontend can reach it without a rebuild
+    # (it derives the Firebase project from the api base). Prod / self-hosted bases keep build-time defaults.
+    overrides = get_frontend_config()
+    if overrides:
+        features.update(overrides)
+    return features

comfy_api_nodes/nodes_gemini.py

@@ -5,7 +5,6 @@ See: https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/infer
 
 import base64
 import os
-from enum import Enum
 from fnmatch import fnmatch
 from io import BytesIO
 from typing import Any, Literal
@@ -78,15 +77,6 @@ GEMINI_IMAGE_2_PRICE_BADGE = IO.PriceBadge(
 )
 
 
-class GeminiImageModel(str, Enum):
-    """
-    Gemini Image Model Names allowed by comfy-api
-    """
-
-    gemini_2_5_flash_image_preview = "gemini-2.5-flash-image-preview"
-    gemini_2_5_flash_image = "gemini-2.5-flash-image"
-
-
 async def create_image_parts(
     cls: type[IO.ComfyNode],
     images: Input.Image | list[Input.Image],
@@ -243,21 +233,15 @@ def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | N
     if not response.modelVersion:
         return None
     # Define prices (Cost per 1,000,000 tokens), see https://cloud.google.com/vertex-ai/generative-ai/pricing
-    if response.modelVersion in ("gemini-2.5-pro-preview-05-06", "gemini-2.5-pro"):
+    if response.modelVersion == "gemini-2.5-pro":
         input_tokens_price = 1.25
         output_text_tokens_price = 10.0
         output_image_tokens_price = 0.0
-    elif response.modelVersion in (
-        "gemini-2.5-flash-preview-04-17",
-        "gemini-2.5-flash",
-    ):
+    elif response.modelVersion == "gemini-2.5-flash":
         input_tokens_price = 0.30
         output_text_tokens_price = 2.50
         output_image_tokens_price = 0.0
-    elif response.modelVersion in (
-        "gemini-2.5-flash-image-preview",
-        "gemini-2.5-flash-image",
-    ):
+    elif response.modelVersion == "gemini-2.5-flash-image":
         input_tokens_price = 0.30
         output_text_tokens_price = 2.50
         output_image_tokens_price = 30.0
@@ -455,8 +439,6 @@ class GeminiNode(IO.ComfyNode):
                 IO.Combo.Input(
                     "model",
                     options=[
-                        "gemini-2.5-pro-preview-05-06",
-                        "gemini-2.5-flash-preview-04-17",
                         "gemini-2.5-pro",
                         "gemini-2.5-flash",
                         "gemini-3-pro-preview",
@@ -904,8 +886,7 @@ class GeminiImage(IO.ComfyNode):
                 ),
                 IO.Combo.Input(
                     "model",
-                    options=GeminiImageModel,
-                    default=GeminiImageModel.gemini_2_5_flash_image,
+                    options=["gemini-2.5-flash-image"],
                     tooltip="The Gemini model to use for generating responses.",
                 ),
                 IO.Int.Input(

comfy_api_nodes/util/_helpers.py

@@ -11,6 +11,7 @@ from io import BytesIO
 from yarl import URL
 
 from comfy.cli_args import args
+from comfy.comfy_api_env import normalize_comfy_api_base
 from comfy.deploy_environment import get_deploy_environment
 from comfy.model_management import processing_interrupted
 from comfy_api.latest import IO
@@ -63,7 +64,7 @@ def get_comfy_api_headers(node_cls: type[IO.ComfyNode]) -> dict[str, str]:
 
 
 def default_base_url() -> str:
-    return getattr(args, "comfy_api_base", "https://api.comfy.org")
+    return normalize_comfy_api_base(getattr(args, "comfy_api_base", "https://api.comfy.org"))
 
 
 async def sleep_with_interrupt(

tests-unit/feature_flags_test.py

@@ -11,6 +11,10 @@ from comfy_api.feature_flags import (
     _coerce_flag_value,
     _parse_cli_feature_flags,
 )
+from comfy.comfy_api_env import (
+    frontend_config_for_base,
+    normalize_comfy_api_base,
+)
 
 
 class TestFeatureFlags:
@@ -181,3 +185,50 @@ class TestCliFeatureFlagRegistry:
             assert "type" in info, f"{key} missing 'type'"
             assert "default" in info, f"{key} missing 'default'"
             assert "description" in info, f"{key} missing 'description'"
+
+
+class TestComfyApiEnv:
+    """--comfy-api-base staging-tier detection + testenv main-host -> -registry rewrite."""
+
+    @pytest.mark.parametrize(
+        "url, expected",
+        [
+            # testenv friendly main host -> comfy-api -registry sibling (slash trimmed)
+            ("https://pr-4398.testenvs.comfy.org", "https://pr-4398-registry.testenvs.comfy.org"),
+            ("https://pr-4398.testenvs.comfy.org/", "https://pr-4398-registry.testenvs.comfy.org"),
+            ("https://pr-4398-registry.testenvs.comfy.org", "https://pr-4398-registry.testenvs.comfy.org"),
+            # staging + everything else -> unchanged (no -registry split)
+            ("https://stagingapi.comfy.org", "https://stagingapi.comfy.org"),
+            ("https://api.comfy.org", "https://api.comfy.org"),
+            ("https://pr-1.testenvs.comfy.org.evil.com", "https://pr-1.testenvs.comfy.org.evil.com"),
+            ("", ""),
+        ],
+    )
+    def test_normalize_comfy_api_base(self, url, expected):
+        assert normalize_comfy_api_base(url) == expected
+
+    def test_config_for_staging_tier_else_none(self):
+        # ephemeral testenv: friendly main host -> -registry, staging platform
+        eph = frontend_config_for_base("https://pr-1234.testenvs.comfy.org/")
+        assert eph["comfy_api_base_url"] == "https://pr-1234-registry.testenvs.comfy.org"
+        assert eph["comfy_platform_base_url"] == "https://stagingplatform.comfy.org"
+        # staging api host: emitted as-is
+        stg = frontend_config_for_base("https://stagingapi.comfy.org")
+        assert stg["comfy_api_base_url"] == "https://stagingapi.comfy.org"
+        assert stg["comfy_platform_base_url"] == "https://stagingplatform.comfy.org"
+        # prod / unknown: nothing
+        assert frontend_config_for_base("https://api.comfy.org") is None
+
+    def test_server_features_merge_only_for_staging_tier(self, monkeypatch):
+        def set_base(url):
+            monkeypatch.setattr(
+                "comfy.comfy_api_env.args",
+                type("Args", (), {"comfy_api_base": url})(),
+            )
+
+        set_base("https://stagingapi.comfy.org")
+        assert "comfy_api_base_url" in get_server_features()
+        set_base("https://pr-7.testenvs.comfy.org")
+        assert "comfy_api_base_url" in get_server_features()
+        set_base("https://api.comfy.org")
+        assert "comfy_api_base_url" not in get_server_features()