feat: 图片生成 MCP 工具 image_generate（中转站图片分组）
- 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。
This commit is contained in:
host committed
1 parent
4bbe6a99e3
commit
33bd14638f
6 files changed
+788
-10
No files matched your search
@@ -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-<epoch>.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
|
||||
Reference in new issue
Block a user