Merge pull request #167 from openqj/feat-imagegen-mcp

feat: 图片生成 MCP 工具 image_generate（中转站图片分组）

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()

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()