codex-app-server/cloud/imagegen_mcp.py

"""Stdio MCP server exposing image_generate to codex.

Protocol: JSON-RPC 2.0 over stdio, one JSON object per line.
Logs go to stderr only -- anything on stdout would corrupt the
JSON-RPC stream.

This file is deliberately thin: all provider logic lives in
``imagegen.py``. It can run two ways:
- inside the repo: ``python3 -m cloud.imagegen_mcp`` (imports
  ``cloud.imagegen``);
- copied flat into a tenant codex-home next to ``imagegen.py``
  (imports plain ``imagegen`` from its own directory).
The ``ensure_codex_home_imagegen`` helper in imagegen.py does the
copying; keys are read from the MCP server process environment
by imagegen.py, never from the tool arguments.
"""

import json
import os
import sys

try:
    from cloud.imagegen import ImageGenError, generate
except ImportError:  # flat copy in a tenant codex-home
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from imagegen import ImageGenError, generate  # noqa: E402

SERVER_NAME = "imagegen"
SERVER_VERSION = "0.1.0"

_TOOLS = [
    {
        "name": "image_generate",
        "description": (
            "Generate one image from a text prompt via the relay's "
            "image group and save it as a file in the current working "
            "directory (the task worktree). Returns the absolute path, "
            "format, dimensions when detectable, byte size and the "
            "relay usage block. The file is a real generated image, "
            "never a placeholder: on any failure the tool errors "
            "instead of writing a stand-in file. Verify the returned "
            "path with shell/PIL before reporting success."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "prompt": {"type": "string",
                           "description": "Image description in detail"},
                "output_filename": {"type": "string",
                                    "description": "Plain basename "
                                    "to save as, e.g. 'hero.png'. "
                                    "Default: generated-<epoch>.png"},
                "size": {"type": "string", "default": "1024x1024",
                         "description": "1024x1024, 1536x1024, "
                         "1024x1536, or auto"},
                "model": {"type": "string",
                          "description": "Image model override; "
                          "default is the IMAGE_GEN_MODEL env "
                          "or gpt-image-2.5-flare"},
            },
            "required": ["prompt"],
        },
    },
]


def _err(code, message):
    return {"code": code, "message": message}


def handle_request(request, generate_fn=None):
    """Handle one JSON-RPC request dict. Returns a response dict or
    None for notifications (no id). Pure function for testability:
    the stdio loop passes the real generate, tests pass fakes."""
    generate_fn = generate_fn or generate
    if not isinstance(request, dict):
        return {"jsonrpc": "2.0", "id": None,
                "error": _err(-32700, "request is not an object")}
    method = request.get("method")
    req_id = request.get("id")
    is_notification = "id" not in request

    if method == "initialize":
        if is_notification:
            return None
        return {"jsonrpc": "2.0", "id": req_id,
                "result": {
                    "protocolVersion": request.get("params", {}).get(
                        "protocolVersion", "2024-11-05"),
                    "capabilities": {"tools": {}},
                    "serverInfo": {"name": SERVER_NAME,
                                   "version": SERVER_VERSION}}}
    if method in ("notifications/initialized", "notifications/cancelled"):
        return None
    if method == "ping":
        if is_notification:
            return None
        return {"jsonrpc": "2.0", "id": req_id, "result": {}}
    if method == "tools/list":
        if is_notification:
            return None
        return {"jsonrpc": "2.0", "id": req_id,
                "result": {"tools": _TOOLS}}
    if method == "tools/call":
        if is_notification:
            return None
        params = request.get("params") or {}
        name = params.get("name")
        args = params.get("arguments") or {}
        try:
            if name == "image_generate":
                result = generate_fn(
                    args.get("prompt"),
                    output_filename=args.get("output_filename"),
                    size=args.get("size", "1024x1024"),
                    model=args.get("model"))
            else:
                return {"jsonrpc": "2.0", "id": req_id,
                        "error": _err(-32602, "unknown tool: %r" % name)}
        except Exception as exc:  # noqa: BLE001 - tool errors are data
            # ImageGenError messages are key-free by construction
            # in imagegen.py; other exceptions are rendered as
            # type + message, never with env or headers.
            return {"jsonrpc": "2.0", "id": req_id,
                    "result": {"content": [{"type": "text",
                                            "text": "error: %s" % exc}],
                               "isError": True}}
        return {"jsonrpc": "2.0", "id": req_id,
                "result": {"content": [{"type": "text",
                                        "text": json.dumps(
                                            result, ensure_ascii=False)}],
                           "isError": False}}
    if is_notification:
        return None
    return {"jsonrpc": "2.0", "id": req_id,
            "error": _err(-32601, "method not found: %r" % method)}


def main():
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            request = json.loads(line)
        except json.JSONDecodeError as exc:
            sys.stdout.write(json.dumps(
                {"jsonrpc": "2.0", "id": None,
                 "error": _err(-32700, "parse error: %s" % exc)}) + "\n")
            sys.stdout.flush()
            continue
        try:
            response = handle_request(request)
        except Exception as exc:  # noqa: BLE001 - server must not die
            print("imagegen_mcp: handler error: %s" % exc, file=sys.stderr)
            if "id" in request:
                response = {"jsonrpc": "2.0", "id": request["id"],
                            "error": _err(-32603, "internal error")}
            else:
                response = None
        if response is not None:
            sys.stdout.write(json.dumps(response) + "\n")
            sys.stdout.flush()


if __name__ == "__main__":
    main()