host
e13a95067c
- imagegen.edit：POST multipart /v1/images/edits，输入图 必须是工作区内 plain basename 且存在（上限 20MB），产出 必须是新文件、绝不覆盖输入；响应解析与 generate 共用 _save_image_response（url 或 b64_json、下载落盘、格式探测）。 - imagegen_mcp 新增 image_edit 工具；image-workflow 技能 改图分支从“无专用工具”改为用 image_edit。 - 测试：multipart 体含提示词与原图字节、输入缺失/同名/ 路径逃逸拒绝、MCP 列表含 image_edit；全量相关 16 项过。 实测依据（2026-10-30）：中转站图片分组 key 调 gpt-image-2.5-flare edits 返回 200，URL 图为真实 PNG 1024x1024、909K，usage 含输入图 token（1659 image tokens）。
203 lines
8.7 KiB
Python
203 lines
8.7 KiB
Python
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_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",
|
|
"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_build_mcp_env_resolves_local_gateway_base(self):
|
|
# Regression: first production verification failed 403
|
|
# Filtered because MCP env had no base URL and imagegen.py
|
|
# fell back to the public relay domain, which is not on
|
|
# the Tinyproxy whitelist. The local gateway origin must
|
|
# be resolved here and land in NO_PROXY.
|
|
env = imagegen.build_mcp_env(
|
|
{"SUB2API_BASE_URL": "http://172.17.0.1:28084",
|
|
"IMAGE_GEN_API_KEY": "k"},
|
|
{"HTTP_PROXY": "http://10.20.20.1:8888",
|
|
"NO_PROXY": "localhost,127.0.0.1"})
|
|
self.assertEqual(env["IMAGE_GEN_BASE_URL"],
|
|
"http://172.17.0.1:28084")
|
|
self.assertIn("172.17.0.1", env["NO_PROXY"])
|
|
self.assertEqual(env["IMAGE_GEN_API_KEY"], "k")
|
|
self.assertNotIn("SUB2API_API_KEY", env)
|
|
|
|
def test_build_mcp_env_explicit_base_wins(self):
|
|
env = imagegen.build_mcp_env(
|
|
{"IMAGE_GEN_BASE_URL": "https://relay.example",
|
|
"SUB2API_BASE_URL": "http://172.17.0.1:28084"})
|
|
self.assertEqual(env["IMAGE_GEN_BASE_URL"],
|
|
"https://relay.example")
|
|
self.assertIn("relay.example", env["NO_PROXY"])
|
|
|
|
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()
|