"""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, edit, 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, edit, 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-.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"], }, }, { "name": "image_edit", "description": ( "Edit an existing image in the current working directory " "from a text prompt and save the result as a NEW file; " "the input file is never overwritten. Returns the " "absolute output path, format, dimensions when " "detectable, byte size and the relay usage block. " "Verify the returned path with shell/PIL before " "reporting success." ), "inputSchema": { "type": "object", "properties": { "prompt": {"type": "string", "description": "What to change, in detail"}, "input_filename": {"type": "string", "description": "Plain basename " "of the existing image in the " "working directory"}, "output_filename": {"type": "string", "description": "Plain basename " "to save the edited image as. " "Must differ from input_filename. " "Default: generated-.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", "input_filename"], }, }, ] def _err(code, message): return {"code": code, "message": message} def handle_request(request, generate_fn=None, edit_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/edit, tests pass fakes.""" generate_fn = generate_fn or generate edit_fn = edit_fn or edit 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")) elif name == "image_edit": result = edit_fn( args.get("prompt"), args.get("input_filename"), 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()