From 33bd14638f7fc7db96bee6e50923952e503cb860 Mon Sep 17 00:00:00 2001 From: host Date: Fri, 30 Oct 2026 02:53:30 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=9B=BE=E7=89=87=E7=94=9F=E6=88=90=20?= =?UTF-8?q?MCP=20=E5=B7=A5=E5=85=B7=20image=5Fgenerate=EF=BC=88=E4=B8=AD?= =?UTF-8?q?=E8=BD=AC=E7=AB=99=E5=9B=BE=E7=89=87=E5=88=86=E7=BB=84=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - cloud/imagegen.py：POST /v1/images/generations，解析 url 或 b64_json，下载落盘到工作区；basename 校验防路径逃逸； 缺 IMAGE_GEN_API_KEY 时明确报未配置，绝不回退任务自己 的分组 2 key（分组 2 allow_image=False，回退只会 404）。 - cloud/imagegen_mcp.py：stdio MCP，工具 image_generate。 - worker --imagegen-mcp：租户 codex-home 播种脚本与 [mcp_servers.imagegen]，容器启动前写 env（系统级 IMAGE_GEN_* + 代理 + NO_PROXY 网关 host）。 - image-workflow 技能：有 image_generate 时用它生成， 无工具或报未配置时明说无法生成；改图暂无专用工具。 - 测试：imagegen 9 项 + worker 播种/env 3 项。 实测依据（2026-10-30）：分组 4 key 调 gpt-image-2.5-flare 返回 200，URL 图为真实 PNG 1024x1024；此前分组 4 账号 schedulable=False 时全模型 404，调度开启后恢复。 部署注意（主机侧，不在代码里）：env 需配 IMAGE_GEN_API_KEY； Tinyproxy 白名单需加 qinglanapi.top（下载返回图 URL 用）； autoscaler 需加 --imagegen-mcp。 --- cloud/imagegen.py | 308 +++++++++++++++++++++++++++ cloud/imagegen_mcp.py | 167 +++++++++++++++ cloud/skills/image-workflow/SKILL.md | 30 ++- cloud/tests/test_imagegen.py | 133 ++++++++++++ cloud/tests/test_imagegen_worker.py | 96 +++++++++ cloud/worker.py | 64 +++++- 6 files changed, 788 insertions(+), 10 deletions(-) create mode 100644 cloud/imagegen.py create mode 100644 cloud/imagegen_mcp.py create mode 100644 cloud/tests/test_imagegen.py create mode 100644 cloud/tests/test_imagegen_worker.py diff --git a/cloud/imagegen.py b/cloud/imagegen.py new file mode 100644 index 0000000..97c22ec --- /dev/null +++ b/cloud/imagegen.py @@ -0,0 +1,308 @@ +"""Image generation through the sub2api relay (image group). + +Measured 2026-10-30: the relay's image group (group 4) serves +``gpt-image-2.5-flare`` at ``/v1/images/generations`` and returns a +download URL (not ``b64_json``) for the finished PNG. Production +task keys live in group 2 (``codex-system``), whose +``allow_image_generation`` is False, so a task's own gateway key +cannot generate images -- this tool uses a separate system-level +image key, ``IMAGE_GEN_API_KEY``, like the search keys. Never fall +back to the task user's gateway key: it 404s with +``model_not_found`` and the failure looks like a broken tool. + +Env (read by the MCP process, never printed in errors): +- ``IMAGE_GEN_BASE_URL``: relay origin, no trailing ``/v1``. + Defaults to ``SUB2API_BASE_URL`` (local gateway origin, + reachable from task containers via NO_PROXY), else + ``https://sub2.diante.me``. +- ``IMAGE_GEN_API_KEY``: system image-group key. Missing key is a + configuration error, reported as such, never as a model error. +- ``IMAGE_GEN_MODEL``: default model, default + ``gpt-image-2.5-flare`` (verified 2026-10-30: 200, real PNG + 1024x1024, 829K, via the relay). + +Network notes for task containers: the API call to the local +gateway origin must bypass the egress proxy (gateway host in +NO_PROXY); downloading the returned image URL goes through the +egress proxy, so the image CDN host (observed: +``*.qinglanapi.top``) must be on the Tinyproxy whitelist, or the +download step fails after a successful generation -- the tool +reports that distinctly. +""" + +import base64 +import json +import os +import struct +import time +import urllib.error +import urllib.request + +DEFAULT_MODEL = "gpt-image-2.5-flare" +DEFAULT_BASE_URL = "https://sub2.diante.me" +ALLOWED_SIZES = ("1024x1024", "1536x1024", "1024x1536", "auto") +API_TIMEOUT = 180 +DOWNLOAD_TIMEOUT = 120 + + +class ImageGenError(RuntimeError): + """Image generation failure. Messages never contain API keys.""" + + +def _base_url(): + base = (os.environ.get("IMAGE_GEN_BASE_URL", "").strip() + or os.environ.get("SUB2API_BASE_URL", "").strip() + or DEFAULT_BASE_URL) + return base.rstrip("/") + + +def _api_key(): + return os.environ.get("IMAGE_GEN_API_KEY", "").strip() + + +def _default_model(): + return os.environ.get("IMAGE_GEN_MODEL", "").strip() or DEFAULT_MODEL + + +def _detect_format(raw): + """Return (format, width, height) from magic bytes; width/height + only for PNG (IHDR), otherwise None -- the caller verifies with + PIL in shell when it needs exact dimensions.""" + if raw[:8] == b"\x89PNG\r\n\x1a\n" and len(raw) >= 24: + width, height = struct.unpack(">II", raw[16:24]) + return "png", width, height + if raw[:3] == b"\xff\xd8\xff": + return "jpeg", None, None + if raw[:4] == b"RIFF" and raw[8:12] == b"WEBP": + return "webp", None, None + if raw[:6] in (b"GIF87a", b"GIF89a"): + return "gif", None, None + return "unknown", None, None + + +def _post_json(url, key, body, request=None): + if request is not None: + return request("POST", url, body, key) + data = json.dumps(body).encode() + req = urllib.request.Request(url, data=data, method="POST", headers={ + "Authorization": "Bearer " + key, "Content-Type": "application/json"}) + 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]} + + +def _download(url, download=None): + if download is not None: + return download(url) + req = urllib.request.Request(url, headers={ + "User-Agent": "codex-imagegen/1.0"}) + try: + with urllib.request.urlopen(req, timeout=DOWNLOAD_TIMEOUT) as resp: + return resp.read() + except Exception as exc: # noqa: BLE001 - surfaced as tool error + 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. + """ + if not prompt or not str(prompt).strip(): + raise ImageGenError("prompt is required") + if size not in ALLOWED_SIZES: + raise ImageGenError("unsupported size: %r (allowed: %s)" + % (size, ", ".join(ALLOWED_SIZES))) + key = _api_key() + if not key: + raise ImageGenError( + "image generation not configured: IMAGE_GEN_API_KEY " + "is not set in the MCP environment") + name = output_filename or ("generated-%d.png" % int(time.time())) + if (os.path.basename(name) != name or name in (".", "..") + 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) + 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"))) + items = (data or {}).get("data") or [] + if not items or not isinstance(items[0], dict): + raise ImageGenError("image generation returned no image data") + first = items[0] + b64 = first.get("b64_json") + if b64: + try: + raw = base64.b64decode(b64) + except Exception as exc: # noqa: BLE001 + raise ImageGenError("invalid b64_json in response: %s" % exc) + elif first.get("url"): + raw = _download(first["url"], download=download) + else: + raise ImageGenError( + "image generation returned neither url nor b64_json " + "(keys: %s)" % ", ".join(sorted(first.keys()))) + if not raw: + raise ImageGenError("downloaded image is empty") + fmt, width, height = _detect_format(raw) + workdir = os.path.abspath(cwd or os.getcwd()) + path = os.path.join(workdir, name) + with open(path, "wb") as fh: + fh.write(raw) + result = {"path": path, "filename": name, "format": fmt, + "bytes": len(raw), "model": use_model, "size": size, + "revised_prompt": first.get("revised_prompt"), + "usage": data.get("usage")} + if width is not None: + result["width"] = width + result["height"] = height + return result + + +# ------------------------------------------------------- codex-home glue + +def ensure_codex_home_imagegen(codex_home_dir, command_args=None, + script_dir=None): + """Seed imagegen MCP files + config into a tenant codex-home. + + Mirrors ``ensure_codex_home_websearch``: copies ``imagegen.py`` + and ``imagegen_mcp.py`` only when missing or identical (a + tenant-modified script is never silently overwritten, it is + reported in ``script_differs``), and appends the + ``[mcp_servers.imagegen]`` section only when absent. Keys are + never written here -- they arrive via + ``ensure_imagegen_mcp_env`` from worker env. + """ + import shutil + + home = os.path.abspath(codex_home_dir) + os.makedirs(home, exist_ok=True) + here = os.path.abspath(script_dir or os.path.dirname(__file__)) + script_changed = False + script_differs = [] + for name in ("imagegen.py", "imagegen_mcp.py"): + src = os.path.join(here, name) + dst = os.path.join(home, name) + if not os.path.exists(src): + continue + if not os.path.exists(dst): + shutil.copyfile(src, dst) + script_changed = True + else: + with open(src, "rb") as fh: + src_bytes = fh.read() + with open(dst, "rb") as fh: + dst_bytes = fh.read() + if src_bytes != dst_bytes: + script_differs.append(name) + continue + config_path = os.path.join(home, "config.toml") + if command_args is None: + command_args = ["python3", + os.path.join(home, "imagegen_mcp.py")] + section = ("\n[mcp_servers.imagegen]\n" + "command = %s\n" + "args = [%s]\n" + % (json.dumps(command_args[0]), + ", ".join(json.dumps(a) for a in command_args[1:]))) + config_changed = False + existing = "" + if os.path.exists(config_path): + with open(config_path, encoding="utf-8") as fh: + existing = fh.read() + if "[mcp_servers.imagegen]" not in existing: + with open(config_path, "a", encoding="utf-8") as fh: + if existing and not existing.endswith("\n"): + fh.write("\n") + fh.write(section) + config_changed = True + return {"config_path": config_path, + "script_path": os.path.join(home, "imagegen_mcp.py"), + "config_changed": config_changed, + "script_changed": script_changed, + "script_differs": script_differs} + + +def ensure_imagegen_mcp_env(codex_home_dir, env): + """Write [mcp_servers.imagegen.env] into a tenant config.toml. + + Same rationale as ``ensure_websearch_mcp_env``: Codex's MCP + subprocess does not inherit the Codex process env, so the env + must live in config.toml. Idempotent in-place replacement of + the env subsection; other sections untouched. Returns True + when the file was written. + """ + if not env: + return False + home = os.path.abspath(codex_home_dir) + config_path = os.path.join(home, "config.toml") + try: + existing = "" + if os.path.exists(config_path): + with open(config_path, encoding="utf-8") as fh: + existing = fh.read() + if "[mcp_servers.imagegen]" not in existing: + existing += ("\n[mcp_servers.imagegen]\n" + "command = \"python3\"\n" + "args = [\"/var/lib/codex-app-server/" + "codex-home/imagegen_mcp.py\"]\n") + env_lines = ["[mcp_servers.imagegen.env]"] + for key in sorted(env): + env_lines.append("%s = %s" % (key, json.dumps(env[key]))) + env_block = "\n".join(env_lines) + "\n" + lines = existing.splitlines(keepends=True) + start = end = None + for i, line in enumerate(lines): + if line.strip() == "[mcp_servers.imagegen.env]": + start = i + elif start is not None and line.startswith("["): + end = i + break + if start is not None: + if end is None: + end = len(lines) + lines[start:end] = [env_block] + new_text = "".join(lines) + else: + insert_at = None + in_section = False + for i, line in enumerate(lines): + if line.strip() == "[mcp_servers.imagegen]": + in_section = True + elif in_section and line.startswith("["): + insert_at = i + break + if insert_at is None: + insert_at = len(lines) + lines.insert(insert_at, env_block) + new_text = "".join(lines) + if new_text == existing: + return False + with open(config_path, "w", encoding="utf-8") as fh: + fh.write(new_text) + return True + except Exception: + return False diff --git a/cloud/imagegen_mcp.py b/cloud/imagegen_mcp.py new file mode 100644 index 0000000..098873e --- /dev/null +++ b/cloud/imagegen_mcp.py @@ -0,0 +1,167 @@ +"""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-.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() diff --git a/cloud/skills/image-workflow/SKILL.md b/cloud/skills/image-workflow/SKILL.md index dbd4147..633fe18 100644 --- a/cloud/skills/image-workflow/SKILL.md +++ b/cloud/skills/image-workflow/SKILL.md @@ -22,13 +22,24 @@ description: 查看、裁剪、缩放、格式转换已有图片，或在图片 1. 查看和处理已有图片：查 `python3 -c "import PIL"`， 以及 `convert` / `magick`（ImageMagick）是否存在。 有其一即可做裁剪、缩放、转格式；都没有就明说。 -2. 生成 / 编辑图片：先查任务环境里是否已配置图片生成 - 工具或 MCP / API（包括需要什么 key、走哪个 endpoint）。 - 没有配置时，回答必须是“本环境未配置图片生成工具， - 无法按描述生成”，并给出替代方案（例如用 shell 画示意图、 - 或用户提供已有图片后做裁剪/缩放），绝不能： +2. 生成图片：本环境已配置图片生成 MCP 工具 + `image_generate`（中转站图片分组，走系统级 key，不用 + 任务自己的网关 key）。调用它生成，不要自己去调 + `/v1/images/generations`，也不要编造 key 或 endpoint。 + 参数：`prompt`（必填）、`output_filename`（工作区内的 + 文件名，如 `hero.png`）、`size`（1024x1024 / 1536x1024 / + 1024x1536，默认 1024x1024）。工具返回绝对路径、格式、 + 尺寸、字节数；失败时工具直接报错，不要换图顶替。 + 如果工具列表里没有 `image_generate`，或调用它报 + “image generation not configured”，回答必须是“本环境 + 未配置图片生成工具，无法按描述生成”，并给出替代方案 + （例如用 shell 画示意图、或用户提供已有图片后做裁剪/ + 缩放），绝不能： - 把网上找来的图、占位图、已有图片说成是刚生成的； - 编造一个图片路径或图片内容描述。 + 编辑已有图片（按描述改图）目前没有专用工具：不要 + 把 `image_generate` 说成改图工具，改图需求先按“生成 + 一张新图”处理并向用户说明这是新图不是原图修改。 ## 步骤（已有图片处理） @@ -38,12 +49,13 @@ description: 查看、裁剪、缩放、格式转换已有图片，或在图片 3. 验证产出：重新打开产出图，确认格式、尺寸符合要求， 文件大小不是 0。 -## 步骤（生成 / 编辑，仅当工具已配置时） +## 步骤（生成，仅当 image_generate 工具在位时） 1. 把用户描述原样记下来，不要自己加没说的元素。 -2. 调已配置的生成工具，记录工具名、提示词、产出路径。 -3. 验证产出图存在且能打开，再交给用户；失败时贴出 - 工具报错，不要换一张别的图顶替。 +2. 调 `image_generate`，记录提示词、尺寸、产出路径。 +3. 验证产出：用 shell / PIL 重新打开返回路径的图，确认 + 格式、尺寸符合要求，文件大小不是 0，再交给用户； + 失败时贴出工具报错，不要换一张别的图顶替。 ## 输出要求 diff --git a/cloud/tests/test_imagegen.py b/cloud/tests/test_imagegen.py new file mode 100644 index 0000000..78a9145 --- /dev/null +++ b/cloud/tests/test_imagegen.py @@ -0,0 +1,133 @@ +import json +import os +import sys +import tempfile +import unittest + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..")) + +from cloud import imagegen +from cloud.imagegen import ImageGenError, generate +from cloud.imagegen_mcp import handle_request + +PNG_1x1 = (b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01" + b"\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde") + + +def _ok_request(body=None, key="k"): + def req(method, url, payload, token): + assert token == key + assert url.endswith("/v1/images/generations") + assert payload["model"] == "gpt-image-2.5-flare" + return 200, {"data": [body or {"url": "https://cdn.example/x.png", + "revised_prompt": "a red apple"}], + "usage": {"total_tokens": 1067}} + return req + + +class ImageGenTests(unittest.TestCase): + def setUp(self): + self._old = dict(os.environ) + os.environ["IMAGE_GEN_API_KEY"] = "k" + os.environ["IMAGE_GEN_BASE_URL"] = "https://relay.example" + os.environ.pop("IMAGE_GEN_MODEL", None) + self._tmp = tempfile.TemporaryDirectory() + self.addCleanup(self._tmp.cleanup) + + def tearDown(self): + os.environ.clear() + os.environ.update(self._old) + + def test_url_download_writes_png_with_dimensions(self): + out = generate("a red apple", output_filename="apple.png", + request=_ok_request(), + download=lambda url: PNG_1x1, + cwd=self._tmp.name) + self.assertEqual(out["path"], + os.path.join(self._tmp.name, "apple.png")) + self.assertEqual(out["format"], "png") + self.assertEqual((out["width"], out["height"]), (1, 1)) + self.assertEqual(out["bytes"], len(PNG_1x1)) + with open(out["path"], "rb") as fh: + self.assertEqual(fh.read(), PNG_1x1) + + def test_b64_response_decoded(self): + import base64 + body = {"b64_json": base64.b64encode(PNG_1x1).decode()} + out = generate("a red apple", request=_ok_request(body), + cwd=self._tmp.name) + self.assertTrue(out["filename"].startswith("generated-")) + self.assertEqual(out["format"], "png") + + def test_missing_key_is_configuration_error(self): + os.environ.pop("IMAGE_GEN_API_KEY") + with self.assertRaises(ImageGenError) as ctx: + generate("x", cwd=self._tmp.name) + self.assertIn("not configured", str(ctx.exception)) + + def test_path_traversal_rejected(self): + for bad in ("../evil.png", "sub/apple.png", "/abs/apple.png"): + with self.assertRaises(ImageGenError): + generate("x", output_filename=bad, + request=_ok_request(), cwd=self._tmp.name) + + def test_bad_size_rejected_before_network(self): + def boom(*args, **kwargs): + raise AssertionError("must not call network") + with self.assertRaises(ImageGenError): + generate("x", size="999x999", request=boom, + cwd=self._tmp.name) + + def test_relay_error_surfaced_without_key(self): + def req(method, url, payload, token): + return 404, {"error": {"message": "Model not supported", + "type": "model_not_found"}} + with self.assertRaises(ImageGenError) as ctx: + generate("x", request=req, cwd=self._tmp.name) + self.assertIn("404", str(ctx.exception)) + self.assertNotIn("k", str(ctx.exception)) + + 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) + called = handle_request( + {"jsonrpc": "2.0", "id": 2, "method": "tools/call", + "params": {"name": "image_generate", + "arguments": {"prompt": "a red apple"}}}, + generate_fn=lambda *a, **k: {"path": "/w/apple.png"}) + text = called["result"]["content"][0]["text"] + self.assertEqual(json.loads(text)["path"], "/w/apple.png") + self.assertFalse(called["result"]["isError"]) + + def test_mcp_unknown_tool(self): + resp = handle_request({"jsonrpc": "2.0", "id": 3, + "method": "tools/call", + "params": {"name": "nope", + "arguments": {}}}) + self.assertEqual(resp["error"]["code"], -32602) + + def test_ensure_home_seeds_and_env(self): + home = os.path.join(self._tmp.name, "home") + os.makedirs(home) + with open(os.path.join(home, "config.toml"), "w") as fh: + fh.write("") + res = imagegen.ensure_codex_home_imagegen( + home, command_args=["python3", "/c/imagegen_mcp.py"]) + self.assertTrue(res["config_changed"]) + self.assertTrue(imagegen.ensure_imagegen_mcp_env( + home, {"IMAGE_GEN_API_KEY": "k", + "IMAGE_GEN_MODEL": "gpt-image-2.5-flare"})) + with open(os.path.join(home, "config.toml")) as fh: + text = fh.read() + self.assertIn("[mcp_servers.imagegen]", text) + self.assertIn("IMAGE_GEN_API_KEY", text) + # Idempotent: second env write with same values is a no-op. + self.assertFalse(imagegen.ensure_imagegen_mcp_env( + home, {"IMAGE_GEN_API_KEY": "k", + "IMAGE_GEN_MODEL": "gpt-image-2.5-flare"})) + + +if __name__ == "__main__": + unittest.main() diff --git a/cloud/tests/test_imagegen_worker.py b/cloud/tests/test_imagegen_worker.py new file mode 100644 index 0000000..163b9ba --- /dev/null +++ b/cloud/tests/test_imagegen_worker.py @@ -0,0 +1,96 @@ +"""Worker wiring for image generation MCP: tenant seeding + env. + +No Docker, no network: _tenant_dirs runs against temp dirs, and +the MCP env helper only edits the tenant config.toml. +""" + +import os +import sys +import tempfile +import unittest +from pathlib import Path + +try: + from cloud.server import init_db + from cloud.worker import Worker + from cloud.imagegen import ensure_imagegen_mcp_env +except ImportError: + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + from cloud.server import init_db + from cloud.worker import Worker + from cloud.imagegen import ensure_imagegen_mcp_env + + +class WorkerImagegenTestCase(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.db_path = os.path.join(self.tmp.name, "control.db") + init_db(self.db_path) + self.template_home = os.path.join(self.tmp.name, "tpl-home") + os.makedirs(self.template_home) + open(os.path.join(self.template_home, "config.toml"), + "w").write('[model_provider.xcmapi]\n' + 'base_url="https://xcmapi.org/v1"\n') + self.template_repo = os.path.join(self.tmp.name, "tpl-repo") + os.makedirs(os.path.join(self.template_repo, ".git")) + open(os.path.join(self.template_repo, "sample.py"), + "w").write("def add(a, b):\n return a + b\n") + + def tearDown(self): + self.tmp.cleanup() + + def _worker(self, **kwargs): + kwargs.setdefault("env", {"XCM_API_KEY": "test-placeholder"}) + kwargs.setdefault("tenants_root", + os.path.join(self.tmp.name, "tenants")) + kwargs.setdefault("tenant_template_home", self.template_home) + kwargs.setdefault("tenant_template_repo", self.template_repo) + kwargs.setdefault("tenant_isolation", True) + return Worker(self.db_path, "/fake/host.py", **kwargs) + + def test_flag_off_is_old_behaviour(self): + worker = self._worker() + self.assertFalse(worker.imagegen_mcp) + dirs = worker._tenant_dirs("projA") + config = open(os.path.join(dirs["codex_home"], + "config.toml")).read() + self.assertNotIn("[mcp_servers.imagegen]", config) + self.assertFalse(os.path.exists( + os.path.join(dirs["codex_home"], "imagegen_mcp.py"))) + + def test_flag_on_seeds_container_path_and_scripts(self): + worker = self._worker(imagegen_mcp=True) + dirs = worker._tenant_dirs("projA") + home = dirs["codex_home"] + config = open(os.path.join(home, "config.toml")).read() + self.assertIn("[mcp_servers.imagegen]", config) + self.assertIn("/var/lib/codex-app-server/codex-home/" + "imagegen_mcp.py", config) + self.assertNotIn(self.tmp.name, config) + self.assertTrue(os.path.exists( + os.path.join(home, "imagegen.py"))) + self.assertTrue(os.path.exists( + os.path.join(home, "imagegen_mcp.py"))) + script = open(os.path.join(home, "imagegen_mcp.py")).read() + self.assertIn("image_generate", script) + + def test_env_helper_writes_imagegen_section_only(self): + worker = self._worker(imagegen_mcp=True) + dirs = worker._tenant_dirs("projA") + home = dirs["codex_home"] + self.assertTrue(ensure_imagegen_mcp_env( + home, {"IMAGE_GEN_API_KEY": "k", + "IMAGE_GEN_MODEL": "gpt-image-2.5-flare", + "NO_PROXY": "localhost,127.0.0.1"})) + config = open(os.path.join(home, "config.toml")).read() + self.assertIn("[mcp_servers.imagegen.env]", config) + self.assertIn("gpt-image-2.5-flare", config) + # Idempotent: same values again is a no-op. + self.assertFalse(ensure_imagegen_mcp_env( + home, {"IMAGE_GEN_API_KEY": "k", + "IMAGE_GEN_MODEL": "gpt-image-2.5-flare", + "NO_PROXY": "localhost,127.0.0.1"})) + + +if __name__ == "__main__": + unittest.main() diff --git a/cloud/worker.py b/cloud/worker.py index 0012990..ce6be6d 100644 --- a/cloud/worker.py +++ b/cloud/worker.py @@ -167,6 +167,8 @@ from cloud.peer_messages import row_get as _row_get from cloud.sub2api import ensure_user_gateway_key as _ensure_sub2api_key from cloud.websearch import ensure_codex_home_websearch as _ensure_websearch_mcp from cloud.websearch import ensure_websearch_mcp_env as _ensure_websearch_mcp_env +from cloud.imagegen import ensure_codex_home_imagegen as _ensure_imagegen_mcp +from cloud.imagegen import ensure_imagegen_mcp_env as _ensure_imagegen_mcp_env from cloud.skill_seed import ensure_codex_home_skills as _ensure_skills from cloud.sub2api import gateway_name as _gateway_name from cloud.sub2api import gateway_origin as _gateway_origin @@ -415,6 +417,7 @@ class Worker: tenant_template_home=None, tenant_template_repo=None, egress_proxy_port=0, claim_executor_only=False, websearch_mcp=False, seed_skills=False, + imagegen_mcp=False, host_name=None, pr_jobs_only=False, upload_dir=None, digest_interval_seconds=3600, @@ -567,6 +570,16 @@ class Worker: # [mcp_servers.websearch]（容器内路径，见 _tenant_dirs）。 # 老租户下一次任务自动合并，已有 section 不重写。 self.websearch_mcp = bool(websearch_mcp) + # 图片生成 MCP：默认关（老行为：租户 codex-home 只有 + # 模板 config.toml + 可选的 websearch MCP）。开启后 + # _tenant_dirs 在 websearch 播种之后调用 imagegen 的 + # 播种助手，把 imagegen.py / imagegen_mcp.py 拷进租户 + # codex-home 并追加 [mcp_servers.imagegen]；_prepare_tenant + # 在容器启动前写 [mcp_servers.imagegen.env]（系统级 + # IMAGE_GEN_API_KEY + 代理 + NO_PROXY 网关 host）。 + # 未配 IMAGE_GEN_API_KEY 时工具会明确报错“未配置”， + # 不会静默回退到任务自己的分组 2 key（那会 404）。 + self.imagegen_mcp = bool(imagegen_mcp) # 技能播种：默认关（老行为：租户 codex-home 只有模板 # config.toml + 可选的 websearch MCP）。开启后 # _tenant_dirs 在 websearch 播种之后调用 skill_seed @@ -2419,6 +2432,17 @@ class Worker: "python3", "/var/lib/codex-app-server/codex-home/" "websearch_mcp.py"]) + if self.imagegen_mcp: + # 与 websearch 同款：播种失败直接抛错、任务不派发 + # ——静默无生图比 loud 失败更糟，尤其对已宣传 + # “能生图”的任务。key 不在这里写，只写 + # command/args；key 由 _prepare_tenant 的 env 块给。 + _ensure_imagegen_mcp( + dirs["codex_home"], + command_args=[ + "python3", + "/var/lib/codex-app-server/codex-home/" + "imagegen_mcp.py"]) if self.seed_skills: # 技能播种失败绝不阻塞任务：技能是可选的流程 # 说明，不是 websearch 那种“已宣传就必须有”的 @@ -2846,6 +2870,35 @@ class Worker: except Exception as exc: self.logger.warning( "websearch MCP env write failed: %s" % exc) + if self.tenant_isolation and self.imagegen_mcp: + # 图片生成 MCP env：系统级 IMAGE_GEN_*（worker 环境， + # 即主机 env 文件）+ 代理（供下载返回的图片 URL）+ + # NO_PROXY（供直连本地网关 origin，不走 egress 代理）。 + # 未配 IMAGE_GEN_API_KEY 时仍写 base/model，工具会 + # 明确报“未配置”；绝不把任务自己的分组 2 key 塞进来。 + try: + imagegen_env = dict(proxy_env) if proxy_env else {} + for key in ("IMAGE_GEN_BASE_URL", "IMAGE_GEN_API_KEY", + "IMAGE_GEN_MODEL"): + if os.environ.get(key): + imagegen_env[key] = os.environ[key] + base = (os.environ.get("IMAGE_GEN_BASE_URL", "") + or os.environ.get("SUB2API_BASE_URL", "")) + host = base.split("://", 1)[-1].split("/", 1)[0] + host = host.split(":", 1)[0] + no_proxy = imagegen_env.get("NO_PROXY", "") + parts = [p for p in no_proxy.split(",") if p] + for item in ("localhost", "127.0.0.1", host): + if item and item not in parts: + parts.append(item) + imagegen_env["NO_PROXY"] = ",".join(parts) + imagegen_env["no_proxy"] = ",".join(parts) + home = self._tenant_dirs( + task.get("project_id"))["codex_home"] + _ensure_imagegen_mcp_env(home, imagegen_env) + except Exception as exc: + self.logger.warning( + "imagegen MCP env write failed: %s" % exc) return (True, tenant_mounts, tenant_network, proxy_env, None) def _run_fork(self, task, quota_check=None): @@ -5897,6 +5950,14 @@ def main(argv=None): " 拷进租户 codex-home/skills/；" "默认关，老租户不受影响；租户改过的" "技能文件绝不覆盖") + parser.add_argument("--imagegen-mcp", action="store_true", + help="租户 codex-home 自动播种图片生成 MCP：" + "_tenant_dirs 追加 [mcp_servers.imagegen] 并拷入" + " imagegen.py / imagegen_mcp.py，_prepare_tenant" + " 写 [mcp_servers.imagegen.env]（系统级" + " IMAGE_GEN_API_KEY + 代理 + NO_PROXY）；" + "默认关，老租户不受影响；未配 key 时工具" + "明确报未配置，不回退任务自己的网关 key") parser.add_argument("--task-timeout", type=float, default=1800, help="每任务截止时间（秒，默认 1800）；0 表示不限时" "（与 host --turn-timeout 0 同约定，worker 级" @@ -5959,7 +6020,8 @@ def main(argv=None): egress_proxy_port=args.egress_proxy_port, claim_executor_only=args.claim_executor_only, websearch_mcp=args.websearch_mcp, - seed_skills=args.seed_skills) + seed_skills=args.seed_skills, + imagegen_mcp=args.imagegen_mcp) # CLI 启动清理兜底（常驻与单次都做）：清掉上一轮残留容器。 worker.cleanup_stale_containers() # CLI 起 worker 前先建表/迁移：全新库直接起循环不应因缺表而炸。