fix(dynamicprompts): prevent hang on unknown wildcards

invokeai/app/api/routers/utilities.py

@@ -17,6 +17,7 @@
 from invokeai.app.api.routers._access import assert_image_read_access
 from invokeai.app.services.image_files.image_files_common import ImageFileNotFoundException
 from invokeai.app.services.model_records.model_records_base import UnknownModelException
+from invokeai.app.util.dynamicprompts import find_missing_wildcards
 from invokeai.backend.llava_onevision_pipeline import LlavaOnevisionPipeline
 from invokeai.backend.model_manager.taxonomy import ModelType
 from invokeai.backend.text_llm_pipeline import DEFAULT_SYSTEM_PROMPT, TextLLMPipeline
@@ -52,8 +53,19 @@ async def parse_dynamicprompts(
     """Creates a batch process"""
     max_prompts = min(max_prompts, 10000)
     generator: Union[RandomPromptGenerator, CombinatorialPromptGenerator]
+    error: Optional[str] = None
+
+    # An unknown wildcard used as a variant value sends the combinatorial generator into an infinite
+    # loop, so bail out early with a clear message instead of hanging the request (and with it the UI
+    # preview). The random generator handles unknown wildcards gracefully, so only the combinatorial
+    # path is guarded.
+    if combinatorial:
+        missing_wildcards = find_missing_wildcards(prompt)
+        if missing_wildcards:
+            wildcards = ", ".join(missing_wildcards)
+            return DynamicPromptsResponse(prompts=[prompt], error=f"No values found for wildcard(s): {wildcards}")
+
     try:
-        error: Optional[str] = None
         if combinatorial:
             generator = CombinatorialPromptGenerator()
             prompts = generator.generate(prompt, max_prompts=max_prompts)

invokeai/app/invocations/prompt.py

@@ -9,6 +9,7 @@
 from invokeai.app.invocations.fields import InputField, UIComponent
 from invokeai.app.invocations.primitives import StringCollectionOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.app.util.dynamicprompts import find_missing_wildcards
 
 
 @invocation(
@@ -31,6 +32,13 @@ class DynamicPromptInvocation(BaseInvocation):
 
     def invoke(self, context: InvocationContext) -> StringCollectionOutput:
         if self.combinatorial:
+            # An unknown wildcard used as a variant value sends the combinatorial generator into an
+            # infinite loop, so fail fast with a clear message instead of hanging the invocation. The
+            # random generator handles unknown wildcards gracefully and needs no guard.
+            missing_wildcards = find_missing_wildcards(self.prompt)
+            if missing_wildcards:
+                raise ValueError(f"No values found for wildcard(s): {', '.join(missing_wildcards)}")
+
             generator = CombinatorialPromptGenerator()
             prompts = generator.generate(self.prompt, max_prompts=self.max_prompts)
         else:

invokeai/app/util/dynamicprompts.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+from collections.abc import Iterator
+
+from dynamicprompts.commands import (
+    Command,
+    SequenceCommand,
+    VariantCommand,
+    WildcardCommand,
+    WrapCommand,
+)
+from dynamicprompts.parser.parse import parse
+from dynamicprompts.wildcards import WildcardManager
+from pyparsing import ParseException
+
+
+def _iter_wildcard_names(command: Command, in_variant: bool = False) -> Iterator[str]:
+    """Recursively yield the statically-known wildcard names that appear as (part of) a variant value.
+
+    Only wildcards reachable from a `VariantCommand` value are yielded: those are the references that
+    hang the combinatorial generator (see `find_missing_wildcards`). The same wildcard used as plain
+    literal text (e.g. `a __nope__ b`) generates fine, so it is intentionally not reported.
+    """
+    if isinstance(command, WildcardCommand):
+        # The wildcard name may itself be a dynamic Command (e.g. `__${var}__`). Only plain string
+        # names can be validated ahead of time, so the dynamic case is intentionally skipped.
+        if in_variant and isinstance(command.wildcard, str):
+            yield command.wildcard
+    elif isinstance(command, SequenceCommand):
+        for token in command.tokens:
+            yield from _iter_wildcard_names(token, in_variant)
+    elif isinstance(command, VariantCommand):
+        # Everything below a variant value is a variant-nested reference, even across sequences.
+        for value in command.values:
+            yield from _iter_wildcard_names(value, in_variant=True)
+    elif isinstance(command, WrapCommand):
+        yield from _iter_wildcard_names(command.wrapper, in_variant)
+        yield from _iter_wildcard_names(command.inner, in_variant)
+    # LiteralCommand and variable commands reference no wildcards we can resolve statically.
+
+
+def find_missing_wildcards(prompt: str, wildcard_manager: WildcardManager | None = None) -> list[str]:
+    """Return the unique unknown wildcard names in `prompt` that hang the combinatorial generator.
+
+    Referencing an unknown wildcard *as a variant value* (e.g. `{__nope__|x}`) makes dynamicprompts'
+    combinatorial generator loop forever: its not-found fallback (`get_wildcard_not_found_fallback`)
+    yields the wrapped wildcard infinitely, and the combinatorial variant logic dedupes those
+    duplicates away without ever advancing. Detecting these names up front lets callers report a clear
+    error instead of hanging. Only the combinatorial generator is affected, and only for wildcards
+    nested in a variant — a bare `a __nope__ b` generates fine and is not reported.
+
+    Without a configured `wildcard_manager`, an empty one is used so that every referenced wildcard is
+    treated as missing (wildcards are not resolved against any files here).
+    """
+    if wildcard_manager is None:
+        wildcard_manager = WildcardManager()
+
+    try:
+        tree = parse(prompt)
+    except ParseException:
+        # Malformed prompts are surfaced separately by the generators; nothing to validate here.
+        return []
+
+    missing: list[str] = []
+    for name in _iter_wildcard_names(tree):
+        if name not in missing and not wildcard_manager.get_values(name):
+            missing.append(name)
+    return missing

tests/app/routers/test_utilities.py

@@ -67,6 +67,50 @@ def test_dynamicprompts_works_for_user(client: TestClient, user1_token: str):
     assert "prompts" in body
 
 
+def test_dynamicprompts_unknown_wildcard_returns_error_without_hanging(client: TestClient, user1_token: str):
+    """An unknown wildcard used as a variant value would otherwise loop forever in the combinatorial generator.
+
+    The endpoint must instead return promptly with a clear error and the original prompt echoed back.
+    """
+    r = client.post(
+        "/api/v1/utilities/dynamicprompts",
+        json={"prompt": "{__random__8chan|fenster|stuff}"},
+        headers={"Authorization": f"Bearer {user1_token}"},
+    )
+    assert r.status_code == status.HTTP_200_OK
+    body = r.json()
+    assert body["error"] is not None
+    assert "random" in body["error"]
+    assert body["prompts"] == ["{__random__8chan|fenster|stuff}"]
+
+
+def test_dynamicprompts_bare_unknown_wildcard_still_generates(client: TestClient, user1_token: str):
+    """A wildcard used as plain literal text (not a variant value) does not hang and must not error."""
+    r = client.post(
+        "/api/v1/utilities/dynamicprompts",
+        json={"prompt": "a photo, __my_style__"},
+        headers={"Authorization": f"Bearer {user1_token}"},
+    )
+    assert r.status_code == status.HTTP_200_OK
+    body = r.json()
+    assert body["error"] is None
+    assert body["prompts"]  # non-empty
+    assert all(p == "a photo, __my_style__" for p in body["prompts"])
+
+
+def test_dynamicprompts_random_generator_ignores_unknown_wildcard(client: TestClient, user1_token: str):
+    """The random generator handles unknown wildcards gracefully, so the guard must not fire for it."""
+    r = client.post(
+        "/api/v1/utilities/dynamicprompts",
+        json={"prompt": "{__random__8chan|fenster|stuff}", "combinatorial": False, "seed": 0},
+        headers={"Authorization": f"Bearer {user1_token}"},
+    )
+    assert r.status_code == status.HTTP_200_OK
+    body = r.json()
+    assert body["error"] is None
+    assert body["prompts"]  # non-empty; the random generator expanded the variant
+
+
 # ----------------------------- image_to_prompt: ownership / read-access -----------------------------
 
 

tests/app/util/test_dynamicprompts.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+import pytest
+
+from invokeai.app.util.dynamicprompts import find_missing_wildcards
+
+
+def test_find_missing_wildcards_detects_unknown_wildcard_in_variant() -> None:
+    # Regression: `__random__` inside a variant is parsed as a wildcard reference. Left unchecked it
+    # sends the combinatorial generator into an infinite loop, so it must be reported up front.
+    assert find_missing_wildcards("{__random__8chan|fenster|stuff}") == ["random"]
+
+
+def test_find_missing_wildcards_detects_unknown_wildcard_nested_in_sequence_in_variant() -> None:
+    # The wildcard hangs the generator even when wrapped in other text inside the variant value.
+    assert find_missing_wildcards("{a __nope__|b}") == ["nope"]
+
+
+@pytest.mark.parametrize("prompt", ["a __nope__ b", "__nope__", "a photo, __my_style__"])
+def test_find_missing_wildcards_ignores_wildcards_outside_variants(prompt: str) -> None:
+    # A wildcard used as plain literal text generates fine (no hang), so it must not be reported.
+    assert find_missing_wildcards(prompt) == []
+
+
+@pytest.mark.parametrize("prompt", ["plain text", "{a|b|c}", "a {2$$x|y|z}"])
+def test_find_missing_wildcards_ignores_prompts_without_wildcards(prompt: str) -> None:
+    assert find_missing_wildcards(prompt) == []
+
+
+def test_find_missing_wildcards_dedupes_repeated_unknown_wildcards() -> None:
+    assert find_missing_wildcards("{__nope__|a} {__nope__|b} {__other__|c}") == ["nope", "other"]