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,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()
|
||||
Reference in new issue
Block a user