Add image_edit tool for Codex image editing

plugins/image_gen/openai-codex/__init__.py

@@ -19,7 +19,10 @@
 
 from __future__ import annotations
 
+import base64
 import logging
+import mimetypes
+from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple
 
 from agent.image_gen_provider import (
@@ -161,9 +164,54 @@ def _build_codex_client():
         return None
 
 
-def _collect_image_b64(client: Any, *, prompt: str, size: str, quality: str) -> Optional[str]:
+def _local_image_to_data_url(value: str) -> str:
+    """Return a Responses-compatible image URL for a URL, data URL, or local path."""
+    value = (value or "").strip()
+    if value.startswith(("http://", "https://", "data:")):
+        return value
+
+    path = Path(value).expanduser()
+    raw = path.read_bytes()
+    mime = mimetypes.guess_type(path.name)[0] or "image/png"
+    encoded = base64.b64encode(raw).decode("ascii")
+    return f"data:{mime};base64,{encoded}"
+
+
+def _image_content_part(value: str) -> Dict[str, Any]:
+    return {"type": "input_image", "image_url": _local_image_to_data_url(value)}
+
+
+def _collect_image_b64(
+    client: Any,
+    *,
+    prompt: str,
+    size: str,
+    quality: str,
+    image_url: Optional[str] = None,
+    mask_url: Optional[str] = None,
+    action: str = "auto",
+) -> Optional[str]:
     """Stream a Codex Responses image_generation call and return the b64 image."""
     image_b64: Optional[str] = None
+    content: List[Dict[str, Any]] = [{"type": "input_text", "text": prompt}]
+    if image_url:
+        content.append(_image_content_part(image_url))
+
+    tool: Dict[str, Any] = {
+        "type": "image_generation",
+        "model": API_MODEL,
+        "size": size,
+        "quality": quality,
+        "output_format": "png",
+        "background": "opaque",
+        "partial_images": 1,
+    }
+    if action and action != "auto":
+        tool["action"] = action
+    elif image_url:
+        tool["action"] = "edit"
+    if mask_url:
+        tool["input_image_mask"] = _local_image_to_data_url(mask_url)
 
     with client.responses.stream(
         model=_CODEX_CHAT_MODEL,
@@ -172,17 +220,9 @@ def _collect_image_b64(client: Any, *, prompt: str, size: str, quality: str) ->
         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": content,
         }],
+        tools=[tool],
         tool_choice={
             "type": "allowed_tools",
             "mode": "required",
@@ -318,12 +358,29 @@ def generate(
                 aspect_ratio=aspect,
             )
 
+        image_url = kwargs.get("image_url") or kwargs.get("input_image")
+        mask_url = kwargs.get("mask_url") or kwargs.get("input_image_mask")
+        action = kwargs.get("action") or ("edit" if image_url else "auto")
+
+        if action == "edit" and not image_url:
+            return error_response(
+                error="image_url is required for image edit requests",
+                error_type="invalid_argument",
+                provider="openai-codex",
+                model=tier_id,
+                prompt=prompt,
+                aspect_ratio=aspect,
+            )
+
         try:
             b64 = _collect_image_b64(
                 client,
                 prompt=prompt,
                 size=size,
                 quality=meta["quality"],
+                image_url=image_url,
+                mask_url=mask_url,
+                action=action,
             )
         except Exception as exc:
             logger.debug("Codex image generation failed", exc_info=True)
@@ -364,7 +421,7 @@ def generate(
             prompt=prompt,
             aspect_ratio=aspect,
             provider="openai-codex",
-            extra={"size": size, "quality": meta["quality"]},
+            extra={"size": size, "quality": meta["quality"], "action": action},
         )
 
 

tests/plugins/image_gen/test_openai_codex_provider.py

@@ -199,6 +199,85 @@ def _stream(**kwargs):
         assert tool["background"] == "opaque"
         assert tool["partial_images"] == 1
 
+    def test_codex_stream_edit_request_shape_with_local_image(self, provider, monkeypatch, tmp_path):
+        monkeypatch.setattr(codex_plugin, "_read_codex_access_token", lambda: "codex-token")
+
+        image = tmp_path / "reference.png"
+        image.write_bytes(bytes.fromhex(_PNG_HEX))
+        captured = {}
+
+        def _stream(**kwargs):
+            captured.update(kwargs)
+            output_item = SimpleNamespace(
+                type="image_generation_call",
+                status="generating",
+                id="ig_test",
+                result=_b64_png(),
+            )
+            done_event = SimpleNamespace(type="response.output_item.done", item=output_item)
+            final_response = SimpleNamespace(output=[], status="completed", output_text="")
+            return _FakeStream([done_event], final_response)
+
+        fake_client = SimpleNamespace(responses=SimpleNamespace(stream=_stream))
+        monkeypatch.setattr(codex_plugin, "_build_codex_client", lambda: fake_client)
+
+        result = provider.generate(
+            "enhance the background but preserve UI text",
+            aspect_ratio="portrait",
+            image_url=str(image),
+            action="edit",
+        )
+        assert result["success"] is True
+        assert result["action"] == "edit"
+
+        content = captured["input"][0]["content"]
+        assert content[0] == {
+            "type": "input_text",
+            "text": "enhance the background but preserve UI text",
+        }
+        assert content[1]["type"] == "input_image"
+        assert content[1]["image_url"].startswith("data:image/png;base64,")
+
+        tool = captured["tools"][0]
+        assert tool["action"] == "edit"
+        assert tool["model"] == "gpt-image-2"
+        assert tool["size"] == "1024x1536"
+
+    def test_codex_stream_edit_accepts_mask_url(self, provider, monkeypatch):
+        monkeypatch.setattr(codex_plugin, "_read_codex_access_token", lambda: "codex-token")
+        captured = {}
+
+        def _stream(**kwargs):
+            captured.update(kwargs)
+            output_item = SimpleNamespace(
+                type="image_generation_call",
+                status="generating",
+                id="ig_test",
+                result=_b64_png(),
+            )
+            done_event = SimpleNamespace(type="response.output_item.done", item=output_item)
+            final_response = SimpleNamespace(output=[], status="completed", output_text="")
+            return _FakeStream([done_event], final_response)
+
+        fake_client = SimpleNamespace(responses=SimpleNamespace(stream=_stream))
+        monkeypatch.setattr(codex_plugin, "_build_codex_client", lambda: fake_client)
+
+        result = provider.generate(
+            "replace the masked logo",
+            image_url="https://example.com/input.png",
+            mask_url="https://example.com/mask.png",
+            action="edit",
+        )
+        assert result["success"] is True
+        assert captured["tools"][0]["input_image_mask"] == "https://example.com/mask.png"
+
+    def test_edit_action_requires_input_image(self, provider, monkeypatch):
+        monkeypatch.setattr(codex_plugin, "_read_codex_access_token", lambda: "codex-token")
+        result = provider.generate("enhance this", action="edit")
+        assert result["success"] is False
+        assert result["error_type"] == "invalid_argument"
+        assert "image_url" in result["error"]
+
     def test_partial_image_event_used_when_done_missing(self, provider, monkeypatch):
         """If the stream never emits output_item.done, fall back to the
         partial_image event so users at least get the latest preview frame."""

tests/tools/test_image_generation.py

@@ -363,12 +363,16 @@ 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):
-        """The agent-facing schema must stay tight — model selection is a
-        user-level config choice, not an agent-level arg."""
+    def test_generate_schema_exposes_only_prompt_and_aspect_ratio_to_agent(self, image_tool):
+        """Generation stays tight — edits use the separate image_edit tool."""
         props = image_tool.IMAGE_GENERATE_SCHEMA["parameters"]["properties"]
         assert set(props.keys()) == {"prompt", "aspect_ratio"}
 
+    def test_edit_schema_exposes_input_image_and_optional_mask(self, image_tool):
+        props = image_tool.IMAGE_EDIT_SCHEMA["parameters"]["properties"]
+        assert set(props.keys()) == {"prompt", "image_url", "aspect_ratio", "mask_url"}
+        assert image_tool.IMAGE_EDIT_SCHEMA["parameters"]["required"] == ["prompt", "image_url"]
+
     def test_aspect_ratio_enum_is_three_values(self, image_tool):
         enum = image_tool.IMAGE_GENERATE_SCHEMA["parameters"]["properties"]["aspect_ratio"]["enum"]
         assert set(enum) == {"landscape", "square", "portrait"}

tools/image_generation_tool.py

@@ -953,6 +953,40 @@ def check_image_generation_requirements() -> bool:
     },
 }
 
+IMAGE_EDIT_SCHEMA = {
+    "name": "image_edit",
+    "description": (
+        "Edit or enhance an existing raster image using the configured image "
+        "generation backend. Requires an input image URL, data URL, or local "
+        "absolute path. Best for reference-image edits, inpainting, background "
+        "enhancement, and ASO screenshot polish while preserving core content."
+    ),
+    "parameters": {
+        "type": "object",
+        "properties": {
+            "prompt": {
+                "type": "string",
+                "description": "Edit instructions. Be explicit about what must be preserved and what may change.",
+            },
+            "image_url": {
+                "type": "string",
+                "description": "Input image as an HTTP(S) URL, data URL, or local absolute file path.",
+            },
+            "aspect_ratio": {
+                "type": "string",
+                "enum": list(VALID_ASPECT_RATIOS),
+                "description": "Target aspect ratio for the edited output.",
+                "default": DEFAULT_ASPECT_RATIO,
+            },
+            "mask_url": {
+                "type": "string",
+                "description": "Optional mask image as URL, data URL, or local path. Backend-specific; for GPT Image masks should match input dimensions and use alpha.",
+            },
+        },
+        "required": ["prompt", "image_url"],
+    },
+}
+
 
 def _read_configured_image_model():
     """Return the value of ``image_gen.model`` from config.yaml, or None."""
@@ -990,7 +1024,7 @@ def _read_configured_image_provider():
     return None
 
 
-def _dispatch_to_plugin_provider(prompt: str, aspect_ratio: str):
+def _dispatch_to_plugin_provider(prompt: str, aspect_ratio: str, **provider_kwargs):
     """Route the call to a plugin-registered provider when one is selected.
 
     Returns a JSON string on dispatch, or ``None`` to fall through to the
@@ -1044,6 +1078,7 @@ def _dispatch_to_plugin_provider(prompt: str, aspect_ratio: str):
 
     try:
         kwargs = {"prompt": prompt, "aspect_ratio": aspect_ratio}
+        kwargs.update({k: v for k, v in provider_kwargs.items() if v not in (None, "")})
         if configured_model:
             kwargs["model"] = configured_model
         result = provider.generate(**kwargs)
@@ -1068,14 +1103,12 @@ def _dispatch_to_plugin_provider(prompt: str, aspect_ratio: str):
     return json.dumps(result)
 
 
-def _handle_image_generate(args, **kw):
+def _handle_image_generate(args: Dict[str, Any], task_id=None):
     prompt = args.get("prompt", "")
     if not prompt:
         return tool_error("prompt is required for image generation")
     aspect_ratio = args.get("aspect_ratio", DEFAULT_ASPECT_RATIO)
 
-    # Route to a plugin-registered provider if one is active (and it's
-    # not the in-tree FAL path).
     dispatched = _dispatch_to_plugin_provider(prompt, aspect_ratio)
     if dispatched is not None:
         return dispatched
@@ -1086,6 +1119,39 @@ def _handle_image_generate(args, **kw):
     )
 
 
+def _handle_image_edit(args: Dict[str, Any], task_id=None):
+    prompt = args.get("prompt", "")
+    if not prompt:
+        return tool_error("prompt is required for image_edit", error_type="invalid_argument")
+    image_url = args.get("image_url", "")
+    aspect_ratio = args.get("aspect_ratio", DEFAULT_ASPECT_RATIO)
+    mask_url = args.get("mask_url")
+
+    if not image_url:
+        return tool_error("image_url is required for image_edit", error_type="invalid_argument")
+
+    dispatched = _dispatch_to_plugin_provider(
+        prompt,
+        aspect_ratio,
+        image_url=image_url,
+        mask_url=mask_url,
+        action="edit",
+    )
+    if dispatched is not None:
+        return dispatched
+
+    return json.dumps({
+        "success": False,
+        "image": None,
+        "error": (
+            "image_edit requires a plugin image provider that supports input "
+            "images. Configure one with `hermes tools` → Image Generation "
+            "(for example OpenAI (Codex auth))."
+        ),
+        "error_type": "unsupported_provider",
+    })
+
+
 registry.register(
     name="image_generate",
     toolset="image_gen",
@@ -1096,3 +1162,14 @@ def _handle_image_generate(args, **kw):
     is_async=False,   # sync fal_client API to avoid "Event loop is closed" in gateway
     emoji="🎨",
 )
+
+registry.register(
+    name="image_edit",
+    toolset="image_gen",
+    schema=IMAGE_EDIT_SCHEMA,
+    handler=_handle_image_edit,
+    check_fn=check_image_generation_requirements,
+    requires_env=[],
+    is_async=False,
+    emoji="🎨",
+)

toolsets.py

@@ -36,7 +36,7 @@
     # File manipulation
     "read_file", "write_file", "patch", "search_files",
     # Vision + image generation
-    "vision_analyze", "image_generate",
+    "vision_analyze", "image_generate", "image_edit",
     # Skills
     "skills_list", "skill_view", "skill_manage",
     # Browser automation
@@ -114,7 +114,7 @@
 
     "image_gen": {
         "description": "Creative generation tools (images)",
-        "tools": ["image_generate"],
+        "tools": ["image_generate", "image_edit"],
         "includes": []
     },
 
@@ -357,7 +357,7 @@
             # File manipulation
             "read_file", "write_file", "patch", "search_files",
             # Vision + image generation
-            "vision_analyze", "image_generate",
+            "vision_analyze", "image_generate", "image_edit",
             # Skills
             "skills_list", "skill_view", "skill_manage",
             # Browser automation