feat(image-gen): support Codex reference-image edits on raw SSE path

plugins/image_gen/openai-codex/__init__.py

@@ -19,9 +19,12 @@
 
 from __future__ import annotations
 
+import base64
 import json
 import logging
-from typing import Any, Dict, List, Optional, Tuple
+import mimetypes
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional, Tuple
 
 from agent.image_gen_provider import (
     DEFAULT_ASPECT_RATIO,
@@ -80,6 +83,16 @@
     "using the image_generation tool when provided."
 )
 
+_MAX_REFERENCE_IMAGES = 8
+_MAX_REFERENCE_IMAGE_BYTES = 20 * 1024 * 1024
+_MAX_DATA_IMAGE_URL_LENGTH = 30 * 1024 * 1024
+_SUPPORTED_IMAGE_MIMES = {
+    "image/png",
+    "image/jpeg",
+    "image/webp",
+    "image/gif",
+}
+
 
 # ---------------------------------------------------------------------------
 # Config + auth helpers
@@ -143,26 +156,107 @@ def _read_codex_access_token() -> Optional[str]:
         return None
 
 
-def _build_responses_payload(*, prompt: str, size: str, quality: str) -> Dict[str, Any]:
+def _detect_image_mime(raw: bytes) -> Optional[str]:
+    """Detect common web image formats from magic bytes."""
+    if raw.startswith(b"\x89PNG\r\n\x1a\n"):
+        return "image/png"
+    if raw.startswith(b"\xff\xd8\xff"):
+        return "image/jpeg"
+    if raw.startswith((b"GIF87a", b"GIF89a")):
+        return "image/gif"
+    if len(raw) >= 12 and raw[:4] == b"RIFF" and raw[8:12] == b"WEBP":
+        return "image/webp"
+    return None
+
+
+def _reference_image_to_url(ref: str) -> str:
+    """Return a Responses-compatible image URL for a local path, URL, or data URL."""
+    value = str(ref).strip()
+    if not value:
+        raise ValueError("reference image path/URL must be non-empty")
+    if value.startswith(("http://", "https://")):
+        return value
+    if value.startswith("data:image/"):
+        if len(value) > _MAX_DATA_IMAGE_URL_LENGTH:
+            raise ValueError("reference image data URL is too large")
+        header, sep, data = value.partition(",")
+        if not sep or ";base64" not in header:
+            raise ValueError("reference image data URL must be base64-encoded")
+        mime = header.removeprefix("data:").split(";", 1)[0]
+        if mime not in _SUPPORTED_IMAGE_MIMES:
+            raise ValueError(f"unsupported reference image MIME type: {mime}")
+        try:
+            base64.b64decode(data, validate=True)
+        except Exception as exc:
+            raise ValueError("reference image data URL contains invalid base64") from exc
+        return value
+
+    path = Path(value).expanduser()
+    if not path.is_file():
+        raise ValueError(f"reference image not found: {path}")
+    try:
+        if path.stat().st_size > _MAX_REFERENCE_IMAGE_BYTES:
+            raise ValueError("local reference image is too large")
+        data = path.read_bytes()
+    except OSError as exc:
+        raise ValueError(f"could not read reference image: {path}") from exc
+    mime = _detect_image_mime(data) or mimetypes.guess_type(str(path))[0]
+    if mime not in _SUPPORTED_IMAGE_MIMES:
+        raise ValueError(f"unsupported or invalid reference image: {path}")
+    return f"data:{mime};base64,{base64.b64encode(data).decode('ascii')}"
+
+
+def _normalize_reference_images(reference_images: Optional[Iterable[str]]) -> List[str]:
+    """Validate reference count and convert refs into Responses image URLs."""
+    if not reference_images:
+        return []
+    if isinstance(reference_images, str):
+        reference_images = [reference_images]
+    refs = [str(ref) for ref in reference_images if str(ref).strip()]
+    if len(refs) > _MAX_REFERENCE_IMAGES:
+        raise ValueError(f"at most {_MAX_REFERENCE_IMAGES} reference images are supported")
+    return [_reference_image_to_url(ref) for ref in refs]
+
+
+def _build_input_content(prompt: str, reference_images: Optional[Iterable[str]]) -> List[Dict[str, str]]:
+    """Build Responses message content with optional image references."""
+    content: List[Dict[str, str]] = [{"type": "input_text", "text": prompt}]
+    for image_url in _normalize_reference_images(reference_images):
+        content.append({"type": "input_image", "image_url": image_url})
+    return content
+
+
+def _build_responses_payload(
+    *,
+    prompt: str,
+    size: str,
+    quality: str,
+    reference_images: Optional[Iterable[str]] = None,
+) -> Dict[str, Any]:
     """Build the Codex Responses request body for an image_generation call."""
+    input_content = _build_input_content(prompt, reference_images)
+    tool_options = {
+        "type": "image_generation",
+        "model": API_MODEL,
+        "size": size,
+        "quality": quality,
+        "output_format": "png",
+        "background": "opaque",
+        "partial_images": 1,
+    }
+    if len(input_content) > 1:
+        tool_options["action"] = "edit"
+
     return {
         "model": _CODEX_CHAT_MODEL,
         "store": False,
         "instructions": _CODEX_INSTRUCTIONS,
         "input": [{
             "type": "message",
             "role": "user",
-            "content": [{"type": "input_text", "text": prompt}],
-        }],
-        "tools": [{
-            "type": "image_generation",
-            "model": API_MODEL,
-            "size": size,
-            "quality": quality,
-            "output_format": "png",
-            "background": "opaque",
-            "partial_images": 1,
+            "content": input_content,
         }],
+        "tools": [tool_options],
         "tool_choice": {
             "type": "allowed_tools",
             "mode": "required",
@@ -242,7 +336,14 @@ def flush():
         yield payload
 
 
-def _collect_image_b64(token: str, *, prompt: str, size: str, quality: str) -> Optional[str]:
+def _collect_image_b64(
+    token: str,
+    *,
+    prompt: str,
+    size: str,
+    quality: str,
+    reference_images: Optional[Iterable[str]] = None,
+) -> Optional[str]:
     """Stream a Codex Responses image_generation call and return the b64 image."""
     import httpx
     from agent.auxiliary_client import _codex_cloudflare_headers
@@ -253,7 +354,12 @@ def _collect_image_b64(token: str, *, prompt: str, size: str, quality: str) -> O
         "Authorization": f"Bearer {token}",
         "Content-Type": "application/json",
     })
-    payload = _build_responses_payload(prompt=prompt, size=size, quality=quality)
+    payload = _build_responses_payload(
+        prompt=prompt,
+        size=size,
+        quality=quality,
+        reference_images=reference_images,
+    )
     timeout = httpx.Timeout(300.0, connect=30.0, read=300.0, write=30.0, pool=30.0)
 
     image_b64: Optional[str] = None
@@ -382,12 +488,28 @@ def generate(
                 aspect_ratio=aspect,
             )
 
+        reference_images = kwargs.get("reference_images")
+        if isinstance(reference_images, str):
+            reference_images = [reference_images]
+        try:
+            normalized_references = _normalize_reference_images(reference_images)
+        except (TypeError, ValueError) as exc:
+            return error_response(
+                error=f"Invalid reference_images: {exc}",
+                error_type="invalid_argument",
+                provider="openai-codex",
+                model=tier_id,
+                prompt=prompt,
+                aspect_ratio=aspect,
+            )
+
         try:
             b64 = _collect_image_b64(
                 token,
                 prompt=prompt,
                 size=size,
                 quality=meta["quality"],
+                reference_images=normalized_references,
             )
         except Exception as exc:
             logger.debug("Codex image generation failed", exc_info=True)

tests/plugins/image_gen/test_openai_codex_provider.py

@@ -8,6 +8,7 @@
 
 from __future__ import annotations
 
+import base64
 import importlib
 from pathlib import Path
 
@@ -129,11 +130,12 @@ def test_codex_stream_request_shape(self, provider, monkeypatch):
 
         captured = {}
 
-        def _collect(token, *, prompt, size, quality):
+        def _collect(token, *, prompt, size, quality, reference_images=None):
             captured.update(codex_plugin._build_responses_payload(
                 prompt=prompt,
                 size=size,
                 quality=quality,
+                reference_images=reference_images,
             ))
             return _b64_png()
 
@@ -159,6 +161,87 @@ def _collect(token, *, prompt, size, quality):
         assert tool["output_format"] == "png"
         assert tool["background"] == "opaque"
         assert tool["partial_images"] == 1
+        assert "action" not in tool
+
+    def test_reference_images_are_sent_as_edit_inputs(self, provider, monkeypatch, tmp_path):
+        monkeypatch.setattr(codex_plugin, "_read_codex_access_token", lambda: "codex-token")
+        ref_path = tmp_path / "ref.png"
+        ref_path.write_bytes(bytes.fromhex(_PNG_HEX))
+
+        captured = {}
+
+        def _collect(token, *, prompt, size, quality, reference_images=None):
+            captured.update(codex_plugin._build_responses_payload(
+                prompt=prompt,
+                size=size,
+                quality=quality,
+                reference_images=reference_images,
+            ))
+            return _b64_png()
+
+        monkeypatch.setattr(codex_plugin, "_collect_image_b64", _collect)
+
+        result = provider.generate("same person", reference_images=[str(ref_path)])
+        assert result["success"] is True
+
+        content = captured["input"][0]["content"]
+        assert content[0] == {"type": "input_text", "text": "same person"}
+        assert content[1]["type"] == "input_image"
+        assert content[1]["image_url"].startswith("data:image/png;base64,")
+        assert base64.b64decode(content[1]["image_url"].split(",", 1)[1]) == bytes.fromhex(_PNG_HEX)
+        assert captured["tools"][0]["action"] == "edit"
+
+    def test_reference_image_url_and_data_url_passthrough(self):
+        refs = [
+            "https://example.com/ref.png",
+            "data:image/png;base64,YWJj",
+        ]
+        payload = codex_plugin._build_responses_payload(
+            prompt="same object",
+            size="1024x1024",
+            quality="medium",
+            reference_images=refs,
+        )
+
+        content = payload["input"][0]["content"]
+        assert [item["image_url"] for item in content[1:]] == refs
+        assert payload["tools"][0]["action"] == "edit"
+
+    def test_reference_images_are_limited(self):
+        refs = ["https://example.com/ref.png"] * (codex_plugin._MAX_REFERENCE_IMAGES + 1)
+        with pytest.raises(ValueError, match="at most"):
+            codex_plugin._build_responses_payload(
+                prompt="same object",
+                size="1024x1024",
+                quality="medium",
+                reference_images=refs,
+            )
+
+    def test_invalid_data_url_rejected(self):
+        with pytest.raises(ValueError, match="invalid base64"):
+            codex_plugin._build_responses_payload(
+                prompt="same object",
+                size="1024x1024",
+                quality="medium",
+                reference_images=["data:image/png;base64,not valid"],
+            )
+
+    def test_invalid_local_reference_returns_invalid_argument(
+        self, provider, monkeypatch, tmp_path
+    ):
+        monkeypatch.setattr(codex_plugin, "_read_codex_access_token", lambda: "codex-token")
+
+        def _collect(*args, **kwargs):  # pragma: no cover - should not be called
+            raise AssertionError("invalid refs should fail before API request")
+
+        monkeypatch.setattr(codex_plugin, "_collect_image_b64", _collect)
+        bad_ref = tmp_path / "not-image.txt"
+        bad_ref.write_text("not an image")
+
+        result = provider.generate("same person", reference_images=[str(bad_ref)])
+        assert result["success"] is False
+        assert result["error_type"] == "invalid_argument"
+        assert "Invalid reference_images" in result["error"]
 
     def test_partial_image_event_used_when_done_missing(self):
         """If output_item.done is missing, partial_image_b64 is accepted."""

tests/tools/test_image_generation.py

@@ -363,11 +363,11 @@ def test_empty_aspect_defaults_to_landscape(self, image_tool):
 
 class TestRegistryIntegration:
 
-    def test_schema_exposes_only_prompt_and_aspect_ratio_to_agent(self, image_tool):
+    def test_schema_exposes_prompt_aspect_ratio_and_reference_images_to_agent(self, image_tool):
         """The agent-facing schema must stay tight — model selection is a
-        user-level config choice, not an agent-level arg."""
+        user-level config choice, while references are request inputs."""
         props = image_tool.IMAGE_GENERATE_SCHEMA["parameters"]["properties"]
-        assert set(props.keys()) == {"prompt", "aspect_ratio"}
+        assert set(props.keys()) == {"prompt", "aspect_ratio", "reference_images"}
 
     def test_aspect_ratio_enum_is_three_values(self, image_tool):
         enum = image_tool.IMAGE_GENERATE_SCHEMA["parameters"]["properties"]["aspect_ratio"]["enum"]

tests/tools/test_image_generation_plugin_dispatch.py

@@ -15,11 +15,14 @@ def _reset_registry():
 
 
 class _FakeCodexProvider(ImageGenProvider):
+    last_kwargs = None
+
     @property
     def name(self) -> str:
         return "codex"
 
     def generate(self, prompt, aspect_ratio="landscape", **kwargs):
+        type(self).last_kwargs = kwargs
         return {
             "success": True,
             "image": "/tmp/codex-test.png",
@@ -52,6 +55,29 @@ def test_dispatch_routes_to_codex_provider(self, monkeypatch, tmp_path):
         assert payload["image"] == "/tmp/codex-test.png"
         assert payload["aspect_ratio"] == "square"
 
+    def test_dispatch_passes_reference_images_to_plugin_provider(self, monkeypatch, tmp_path):
+        from tools import image_generation_tool
+        from agent import image_gen_registry as registry_module
+        from hermes_cli import plugins as plugins_module
+
+        monkeypatch.setenv("HERMES_HOME", str(tmp_path))
+        (tmp_path / "config.yaml").write_text("image_gen:\n  provider: codex\n")
+        _FakeCodexProvider.last_kwargs = None
+
+        monkeypatch.setattr(image_generation_tool, "_read_configured_image_provider", lambda: "codex")
+        monkeypatch.setattr(plugins_module, "_ensure_plugins_discovered", lambda: None)
+        monkeypatch.setattr(registry_module, "get_provider", lambda name: _FakeCodexProvider() if name == "codex" else None)
+
+        dispatched = image_generation_tool._dispatch_to_plugin_provider(
+            "draw cat",
+            "square",
+            reference_images=["/tmp/ref.png"],
+        )
+        payload = json.loads(dispatched)
+
+        assert payload["success"] is True
+        assert _FakeCodexProvider.last_kwargs == {"reference_images": ["/tmp/ref.png"]}
+
     def test_dispatch_reports_missing_registered_provider(self, monkeypatch, tmp_path):
         from tools import image_generation_tool
         from hermes_cli import plugins as plugins_module