diff --git a/cloud/imagegen.py b/cloud/imagegen.py index 652d320..933bf8e 100644 --- a/cloud/imagegen.py +++ b/cloud/imagegen.py @@ -112,17 +112,7 @@ def _download(url, download=None): raise ImageGenError("image download failed: %s" % exc) -def generate(prompt, output_filename=None, size="1024x1024", - model=None, request=None, download=None, cwd=None): - """Generate one image and save it to the current working - directory (the task worktree when run as a Codex MCP server). - - ``output_filename`` must be a plain basename (no separators, - no ``..``); default is ``generated-.png``. Returns a - dict with path, format, dimensions (PNG only), byte size and - the relay's usage block. Raises ImageGenError on any failure; - the message never contains the API key. - """ +def _validate_common(prompt, size, output_filename): if not prompt or not str(prompt).strip(): raise ImageGenError("prompt is required") if size not in ALLOWED_SIZES: @@ -138,20 +128,20 @@ def generate(prompt, output_filename=None, size="1024x1024", or "/" in name or "\\" in name): raise ImageGenError( "output_filename must be a plain basename, got %r" % name) - use_model = model or _default_model() - status, data = _post_json( - _base_url() + "/v1/images/generations", key, - {"model": use_model, "prompt": str(prompt), "size": size}, - request=request) + return key, name + + +def _save_image_response(status, data, name, use_model, size, + download, cwd, verb): if status != 200: err = (data or {}).get("error") or {} msg = err.get("message") if isinstance(err, dict) else str(err) raise ImageGenError( - "image generation failed: HTTP %s: %s" - % (status, (msg or data or "unknown error"))) + "image %s failed: HTTP %s: %s" + % (verb, status, (msg or data or "unknown error"))) items = (data or {}).get("data") or [] if not items or not isinstance(items[0], dict): - raise ImageGenError("image generation returned no image data") + raise ImageGenError("image %s returned no image data" % verb) first = items[0] b64 = first.get("b64_json") if b64: @@ -163,8 +153,8 @@ def generate(prompt, output_filename=None, size="1024x1024", raw = _download(first["url"], download=download) else: raise ImageGenError( - "image generation returned neither url nor b64_json " - "(keys: %s)" % ", ".join(sorted(first.keys()))) + "image %s returned neither url nor b64_json " + "(keys: %s)" % (verb, ", ".join(sorted(first.keys())))) if not raw: raise ImageGenError("downloaded image is empty") fmt, width, height = _detect_format(raw) @@ -182,6 +172,130 @@ def generate(prompt, output_filename=None, size="1024x1024", return result +def generate(prompt, output_filename=None, size="1024x1024", + model=None, request=None, download=None, cwd=None): + """Generate one image and save it to the current working + directory (the task worktree when run as a Codex MCP server). + + ``output_filename`` must be a plain basename (no separators, + no ``..``); default is ``generated-.png``. Returns a + dict with path, format, dimensions (PNG only), byte size and + the relay's usage block. Raises ImageGenError on any failure; + the message never contains the API key. + """ + key, name = _validate_common(prompt, size, output_filename) + use_model = model or _default_model() + status, data = _post_json( + _base_url() + "/v1/images/generations", key, + {"model": use_model, "prompt": str(prompt), "size": size}, + request=request) + return _save_image_response(status, data, name, use_model, size, + download, cwd, "generation") + + +def _content_type_for(filename): + lower = filename.lower() + if lower.endswith(".jpg") or lower.endswith(".jpeg"): + return "image/jpeg" + if lower.endswith(".webp"): + return "image/webp" + return "image/png" + + +def _build_multipart(fields, file_field, filename, file_bytes): + """Build a multipart/form-data body. Returns (body, boundary). + Pure function for testability: the real post uses it too.""" + import uuid + + boundary = "----codeximagegen" + uuid.uuid4().hex + body = b"" + for name, value in fields: + body += ("--" + boundary + "\r\n").encode() + body += ('Content-Disposition: form-data; ' + 'name="%s"\r\n\r\n' % name).encode() + body += value.encode() + b"\r\n" + body += ("--" + boundary + "\r\n").encode() + body += ('Content-Disposition: form-data; name="%s"; ' + 'filename="%s"\r\n' % (file_field, filename)).encode() + body += ("Content-Type: %s\r\n\r\n" + % _content_type_for(filename)).encode() + body += file_bytes + b"\r\n" + body += ("--" + boundary + "--\r\n").encode() + return body, boundary + + +def _post_multipart(url, key, body, content_type, + request_multipart=None): + if request_multipart is not None: + return request_multipart(url, key, body, content_type) + req = urllib.request.Request(url, data=body, method="POST", headers={ + "Authorization": "Bearer " + key, + "Content-Type": content_type}) + try: + with urllib.request.urlopen(req, timeout=API_TIMEOUT) as resp: + return resp.status, json.loads(resp.read().decode() or "{}") + except urllib.error.HTTPError as exc: + raw = exc.read().decode() + try: + payload = json.loads(raw or "{}") + except ValueError: + payload = {"raw": raw[:400]} + return exc.code, payload + except Exception as exc: # noqa: BLE001 - surfaced as tool error + return 0, {"error": str(exc)[:300]} + + +MAX_INPUT_BYTES = 20 * 1024 * 1024 + + +def edit(prompt, input_filename, output_filename=None, + size="1024x1024", model=None, request_multipart=None, + download=None, cwd=None): + """Edit an existing image in the working directory and save + the result as a NEW file (the input is never overwritten). + + ``input_filename`` and ``output_filename`` must both be plain + basenames in the working directory. The relay's + ``/v1/images/edits`` endpoint is multipart/form-data (verified + 2026-10-30: 200, URL PNG 1024x1024, input image tokens billed + in usage). Returns the same dict shape as ``generate``. + """ + key, name = _validate_common(prompt, size, output_filename) + if (os.path.basename(input_filename or "") != input_filename + or input_filename in (".", "..") + or "/" in (input_filename or "") + or "\\" in (input_filename or "")): + raise ImageGenError( + "input_filename must be a plain basename, got %r" + % input_filename) + workdir = os.path.abspath(cwd or os.getcwd()) + input_path = os.path.join(workdir, input_filename) + if not os.path.isfile(input_path): + raise ImageGenError("input image not found: %r" % input_filename) + with open(input_path, "rb") as fh: + file_bytes = fh.read() + if len(file_bytes) > MAX_INPUT_BYTES: + raise ImageGenError( + "input image too large: %d bytes (max %d)" + % (len(file_bytes), MAX_INPUT_BYTES)) + if os.path.abspath(os.path.join(workdir, name)) == input_path: + raise ImageGenError( + "output_filename must differ from input_filename") + use_model = model or _default_model() + body, boundary = _build_multipart( + [("model", use_model), ("prompt", str(prompt)), + ("size", size)], + "image", input_filename, file_bytes) + status, data = _post_multipart( + _base_url() + "/v1/images/edits", key, body, + "multipart/form-data; boundary=" + boundary, + request_multipart=request_multipart) + result = _save_image_response(status, data, name, use_model, size, + download, cwd, "edit") + result["input_filename"] = input_filename + return result + + # ------------------------------------------------------- codex-home glue def ensure_codex_home_imagegen(codex_home_dir, command_args=None, diff --git a/cloud/imagegen_mcp.py b/cloud/imagegen_mcp.py index 098873e..e67b3ea 100644 --- a/cloud/imagegen_mcp.py +++ b/cloud/imagegen_mcp.py @@ -20,10 +20,10 @@ import os import sys try: - from cloud.imagegen import ImageGenError, generate + 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, generate # noqa: E402 + from imagegen import ImageGenError, edit, generate # noqa: E402 SERVER_NAME = "imagegen" SERVER_VERSION = "0.1.0" @@ -61,6 +61,42 @@ _TOOLS = [ "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"], + }, + }, ] @@ -68,11 +104,12 @@ def _err(code, message): return {"code": code, "message": message} -def handle_request(request, generate_fn=None): +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, tests pass fakes.""" + 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")} @@ -114,6 +151,13 @@ def handle_request(request, generate_fn=None): 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)} diff --git a/cloud/skills/image-workflow/SKILL.md b/cloud/skills/image-workflow/SKILL.md index 633fe18..abdea5d 100644 --- a/cloud/skills/image-workflow/SKILL.md +++ b/cloud/skills/image-workflow/SKILL.md @@ -37,9 +37,11 @@ description: 查看、裁剪、缩放、格式转换已有图片，或在图片 缩放），绝不能： - 把网上找来的图、占位图、已有图片说成是刚生成的； - 编造一个图片路径或图片内容描述。 - 编辑已有图片（按描述改图）目前没有专用工具：不要 - 把 `image_generate` 说成改图工具，改图需求先按“生成 - 一张新图”处理并向用户说明这是新图不是原图修改。 + 编辑已有图片：用 `image_edit` 工具，参数 `prompt`（改什么）、 + `input_filename`（工作区内已有图的文件名）、 + `output_filename`（新文件名，必须与输入不同，原图绝不 + 被覆盖）、`size`。改图产出是新文件，不是原图修改； + 同样要用 shell / PIL 重新打开验证后再交给用户。 ## 步骤（已有图片处理） @@ -49,10 +51,12 @@ description: 查看、裁剪、缩放、格式转换已有图片，或在图片 3. 验证产出：重新打开产出图，确认格式、尺寸符合要求， 文件大小不是 0。 -## 步骤（生成，仅当 image_generate 工具在位时） +## 步骤（生成 / 改图，仅当对应工具在位时） 1. 把用户描述原样记下来，不要自己加没说的元素。 -2. 调 `image_generate`，记录提示词、尺寸、产出路径。 +2. 生成调 `image_generate`，改图调 `image_edit`（输入图 + 必须是工作区内已存在的文件，产出文件名与输入不同）， + 记录提示词、尺寸、产出路径。 3. 验证产出：用 shell / PIL 重新打开返回路径的图，确认 格式、尺寸符合要求，文件大小不是 0，再交给用户； 失败时贴出工具报错，不要换一张别的图顶替。 diff --git a/cloud/tests/test_imagegen.py b/cloud/tests/test_imagegen.py index 34dec90..6ccaa9e 100644 --- a/cloud/tests/test_imagegen.py +++ b/cloud/tests/test_imagegen.py @@ -87,11 +87,55 @@ class ImageGenTests(unittest.TestCase): self.assertIn("404", str(ctx.exception)) self.assertNotIn("k", str(ctx.exception)) + def test_edit_multipart_writes_new_file(self): + with open(os.path.join(self._tmp.name, "in.png"), "wb") as fh: + fh.write(PNG_1x1) + seen = {} + + def fake_multipart(url, key, body, content_type): + seen["url"] = url + seen["body"] = body + seen["content_type"] = content_type + self.assertTrue(url.endswith("/v1/images/edits")) + self.assertIn(b"Change the background", body) + self.assertIn(PNG_1x1, body) + return 200, {"data": [{"url": "https://cdn.example/e.png", + "revised_prompt": "blue bg"}], + "usage": {"total_tokens": 2727}} + + out = imagegen.edit( + "Change the background to blue", "in.png", + output_filename="out.png", + request_multipart=fake_multipart, + download=lambda url: PNG_1x1, cwd=self._tmp.name) + self.assertEqual(out["filename"], "out.png") + self.assertEqual(out["input_filename"], "in.png") + self.assertEqual(out["format"], "png") + self.assertTrue(seen["content_type"].startswith( + "multipart/form-data; boundary=")) + # Input untouched, output is a separate file. + with open(os.path.join(self._tmp.name, "in.png"), "rb") as fh: + self.assertEqual(fh.read(), PNG_1x1) + self.assertTrue(os.path.exists( + os.path.join(self._tmp.name, "out.png"))) + + def test_edit_missing_input_and_same_output_rejected(self): + with self.assertRaises(ImageGenError): + imagegen.edit("x", "nope.png", cwd=self._tmp.name) + with open(os.path.join(self._tmp.name, "in.png"), "wb") as fh: + fh.write(PNG_1x1) + with self.assertRaises(ImageGenError): + imagegen.edit("x", "in.png", output_filename="in.png", + cwd=self._tmp.name) + with self.assertRaises(ImageGenError): + imagegen.edit("x", "../in.png", cwd=self._tmp.name) + def test_mcp_tools_list_and_call(self): listed = handle_request({"jsonrpc": "2.0", "id": 1, "method": "tools/list"}) names = [t["name"] for t in listed["result"]["tools"]] self.assertIn("image_generate", names) + self.assertIn("image_edit", names) called = handle_request( {"jsonrpc": "2.0", "id": 2, "method": "tools/call", "params": {"name": "image_generate",