fix(dynamicprompts): prevent hang on unknown wildcards (invoke-ai#9307)

* fix(dynamicprompts): prevent hang on unknown wildcards

Referencing an unknown wildcard previously hung the combinatorial generator
forever: its not-found fallback yields the wrapped wildcard infinitely, and
the variant dedup logic discards those duplicates without ever advancing.
This froze the UI prompt preview. Unknown wildcards are now detected up front
and reported as a clear error instead of attempting generation.

* fix(dynamicprompts): narrow unknown-wildcard guard to combinatorial variant values

Only an unknown wildcard used as a variant value hangs the combinatorial
generator; a bare `a __nope__ b` and the random generator both handle unknown
wildcards fine. Restrict detection to variant-nested wildcards and gate the
check on the combinatorial path so previously-working prompts no longer error.

Author: Pfannkuchensack

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