diff --git a/tests/tools/test_image_generation_image_to_image.py b/tests/tools/test_image_generation_image_to_image.py index 4e9d457a49f..60f8d3ca680 100644 --- a/tests/tools/test_image_generation_image_to_image.py +++ b/tests/tools/test_image_generation_image_to_image.py @@ -79,6 +79,40 @@ class TestFalEditPayload: assert FAL_MODELS["fal-ai/nano-banana-pro"].get("edit_endpoint") +class TestMandatoryKeysSurviveWhitelist: + """A model whose whitelist forgets the mandatory keys must not produce a + request with the prompt / source images silently stripped.""" + + _SIZES = {"square": "1024x1024", "landscape": "1536x1024", "portrait": "1024x1536"} + + def test_edit_keeps_prompt_and_image_urls(self, monkeypatch): + from tools import image_generation_tool as t + + fake = { + "size_style": "image_size_preset", + "sizes": self._SIZES, + "edit_supports": {"seed"}, # intentionally omits prompt + image_urls + } + monkeypatch.setitem(t.FAL_MODELS, "test/edit-model", fake) + payload = t._build_fal_edit_payload( + "test/edit-model", "make it blue", ["https://x/y.png"], "square", + ) + assert payload["prompt"] == "make it blue" + assert payload["image_urls"] == ["https://x/y.png"] + + def test_text_keeps_prompt(self, monkeypatch): + from tools import image_generation_tool as t + + fake = { + "size_style": "image_size_preset", + "sizes": self._SIZES, + "supports": {"seed"}, # intentionally omits prompt + } + monkeypatch.setitem(t.FAL_MODELS, "test/text-model", fake) + payload = t._build_fal_payload("test/text-model", "a cat", aspect_ratio="square") + assert payload["prompt"] == "a cat" + + class TestFalRouting: def _patch_submit(self, monkeypatch, image_tool, capture: dict): class _Handler: diff --git a/tools/image_generation_tool.py b/tools/image_generation_tool.py index 3213068ddd9..101b000db2a 100644 --- a/tools/image_generation_tool.py +++ b/tools/image_generation_tool.py @@ -607,7 +607,13 @@ def _build_fal_payload( payload[k] = v supports = meta["supports"] - return {k: v for k, v in payload.items() if k in supports} + # ``prompt`` is required by every FAL text-to-image endpoint; keep it even + # if a model's ``supports`` whitelist omits it, so a missing whitelist entry + # can't silently strip the prompt and send an empty request. + return { + k: v for k, v in payload.items() + if k in supports or k == "prompt" + } def _build_fal_edit_payload( @@ -656,7 +662,15 @@ def _build_fal_edit_payload( if v is not None: payload[k] = v - return {k: v for k, v in payload.items() if k in edit_supports} + # ``prompt`` and ``image_urls`` are required by every FAL edit endpoint; + # keep them even if a model's ``edit_supports`` whitelist omits them, so a + # missing whitelist entry can't silently drop the prompt or the source + # images and send a broken edit request. + _required = {"prompt", "image_urls"} + return { + k: v for k, v in payload.items() + if k in edit_supports or k in _required + } # ---------------------------------------------------------------------------