From ae17a79bdf3e5fe994bce0883afbbce9f37c0cde Mon Sep 17 00:00:00 2001
From: chengyongru <chengyongru.ai@gmail.com>
Date: Thu, 4 Jun 2026 13:46:19 +0800
Subject: [PATCH] fix: harden custom image generation config

Maintainer edit: require providers.custom.apiBase before making custom image requests and allow unauthenticated local endpoints by omitting Authorization when no apiKey is configured.
---
 docs/image-generation.md                 |  4 ++--
 nanobot/providers/image_generation.py    | 15 ++++++++-------
 tests/providers/test_image_generation.py | 21 ++++++++++++++++++---
 3 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/docs/image-generation.md b/docs/image-generation.md
index f298042dd..55250c406 100644
--- a/docs/image-generation.md
+++ b/docs/image-generation.md
@@ -108,7 +108,7 @@ Configure:
 }
 ```
 
-The `apiBase` is required. The provider sends requests to `{apiBase}/images/generations` using the OpenAI Images API format with `response_format: "b64_json"`.
+The `apiBase` is required. The provider sends requests to `{apiBase}/images/generations` using the OpenAI Images API format with `response_format: "b64_json"`. The `apiKey` is optional for local or unauthenticated endpoints.
 
 ### AIHubMix
 
@@ -350,7 +350,7 @@ Use the reference image. Keep the same robot and composition, change the palette
 |---------|-------|
 | `generate_image` is not available | Set `tools.imageGeneration.enabled` to `true` and restart the gateway |
 | Missing API key error | Configure `providers.<provider>.apiKey`; if using `${VAR_NAME}`, confirm the environment variable is visible to the gateway process |
-| `unsupported image generation provider` | Use `openrouter`, `aihubmix`, `minimax`, `gemini`, `ollama`, `stepfun`, or `zhipu` |
+| `unsupported image generation provider` | Use `openrouter`, `custom`, `aihubmix`, `minimax`, `gemini`, `ollama`, `stepfun`, or `zhipu` |
 | AIHubMix says `Incorrect model ID` | Use `model: "gpt-image-2-free"`; nanobot expands it to the required `openai/gpt-image-2-free` model path internally |
 | Generation times out | Try a smaller/default image size, set AIHubMix `extraBody.quality` to `"low"`, or retry later |
 | Reference image rejected | Reference image paths must be inside the workspace or nanobot media directory and must be valid image files |
diff --git a/nanobot/providers/image_generation.py b/nanobot/providers/image_generation.py
index ccf2be2ba..18f62e26c 100644
--- a/nanobot/providers/image_generation.py
+++ b/nanobot/providers/image_generation.py
@@ -1037,8 +1037,8 @@ class CustomImageGenerationClient(ImageGenerationProvider):
     """OpenAI-compatible Images API for user-configured custom providers."""
 
     provider_name = "custom"
-    missing_key_message = (
-        "Custom image generation API key is not configured. Set providers.custom.apiKey."
+    missing_base_message = (
+        "Custom image generation API base is not configured. Set providers.custom.apiBase."
     )
 
     def _default_base_url(self) -> str:
@@ -1057,8 +1057,8 @@ class CustomImageGenerationClient(ImageGenerationProvider):
         aspect_ratio: str | None = None,
         image_size: str | None = None,
     ) -> GeneratedImageResponse:
-        if not self.api_key:
-            raise ImageGenerationError(self.missing_key_message)
+        if not self.api_base:
+            raise ImageGenerationError(self.missing_base_message)
 
         if reference_images:
             logger.warning(
@@ -1068,11 +1068,12 @@ class CustomImageGenerationClient(ImageGenerationProvider):
                 model,
             )
 
-        headers = {
-            "Authorization": f"Bearer {self.api_key}",
+        headers: dict[str, str] = {
             "Content-Type": "application/json",
-            **self.extra_headers,
         }
+        if self.api_key:
+            headers["Authorization"] = f"Bearer {self.api_key}"
+        headers.update(self.extra_headers)
 
         body: dict[str, Any] = {
             "model": model,
diff --git a/tests/providers/test_image_generation.py b/tests/providers/test_image_generation.py
index 8df0ac03b..b7bd29a4b 100644
--- a/tests/providers/test_image_generation.py
+++ b/tests/providers/test_image_generation.py
@@ -844,10 +844,25 @@ async def test_custom_generate_success() -> None:
 
 
 @pytest.mark.asyncio
-async def test_custom_generate_no_api_key() -> None:
-    client = CustomImageGenerationClient(api_key=None)
+async def test_custom_generate_without_api_key_omits_authorization() -> None:
+    fake = FakeClient(FakeResponse({"data": [{"b64_json": RAW_B64}]}))
+    client = CustomImageGenerationClient(
+        api_key=None,
+        api_base="http://localhost:7860/v1",
+        client=fake,  # type: ignore[arg-type]
+    )
 
-    with pytest.raises(ImageGenerationError, match="providers.custom.apiKey"):
+    response = await client.generate(prompt="draw", model="custom-image-model")
+
+    assert response.images == [PNG_DATA_URL]
+    assert "Authorization" not in fake.calls[0]["headers"]
+
+
+@pytest.mark.asyncio
+async def test_custom_generate_requires_api_base() -> None:
+    client = CustomImageGenerationClient(api_key="sk-custom-test")
+
+    with pytest.raises(ImageGenerationError, match="providers.custom.apiBase"):
         await client.generate(prompt="draw", model="custom-image-model")