From cd88f9cd101ea3f7993ba54056bbdcf09929bd2c Mon Sep 17 00:00:00 2001
From: ethanfel <ethan.fel@ts-pc.fr>
Date: Sun, 28 Jun 2026 13:47:09 +0200
Subject: [PATCH] feat: FastMCP server build + background lifecycle

---
 src/bridge/mcp_server.py              | 51 +++++++++++++++++++++++++++
 tests/bridge/test_mcp_server_smoke.py | 15 ++++++++
 2 files changed, 66 insertions(+)
 create mode 100644 src/bridge/mcp_server.py
 create mode 100644 tests/bridge/test_mcp_server_smoke.py

diff --git a/src/bridge/mcp_server.py b/src/bridge/mcp_server.py
new file mode 100644
index 0000000..46d65cc
--- /dev/null
+++ b/src/bridge/mcp_server.py
@@ -0,0 +1,51 @@
+"""FastMCP server wiring + background lifecycle for the ComfyUI agent bridge."""
+import os
+import threading
+
+from mcp.server.fastmcp import FastMCP
+
+from src.bridge import mcp_tools
+
+_started = False
+_lock = threading.Lock()
+
+
+def build_server(host: str = "127.0.0.1", port: int = 9188) -> FastMCP:
+    mcp = FastMCP("comfy-bridge", host=host, port=port)
+    mcp.tool()(mcp_tools.comfy_pull)
+    mcp.tool()(mcp_tools.comfy_push)
+    mcp.tool()(mcp_tools.comfy_list_channels)
+    mcp.tool()(mcp_tools.comfy_run_workflow)
+    mcp.tool()(mcp_tools.comfy_get_result)
+    return mcp
+
+
+def registered_tool_names(mcp: FastMCP) -> list:
+    # FastMCP exposes a tool manager; fall back across SDK versions.
+    mgr = getattr(mcp, "_tool_manager", None)
+    if mgr is not None and hasattr(mgr, "_tools"):
+        return list(mgr._tools.keys())
+    return [t.name for t in getattr(mcp, "list_tools", lambda: [])()]
+
+
+def start_in_background(host: str = "127.0.0.1", port: int = None,
+                        _test_no_serve: bool = False) -> None:
+    global _started
+    with _lock:
+        if _started:
+            return
+        _started = True
+    port = port if port is not None else int(
+        os.environ.get("COMFY_BRIDGE_MCP_PORT", "9188"))
+    mcp = build_server(host=host, port=port)
+    if _test_no_serve:
+        return
+
+    def _serve():
+        mcp.run(transport="streamable-http")
+
+    threading.Thread(target=_serve, name="comfy-bridge-mcp",
+                     daemon=True).start()
+    print(f"[comfyui-nodes-agents] MCP bridge on "
+          f"http://{host}:{port}/mcp  (claude mcp add --transport http "
+          f"comfy http://{host}:{port}/mcp)")
diff --git a/tests/bridge/test_mcp_server_smoke.py b/tests/bridge/test_mcp_server_smoke.py
new file mode 100644
index 0000000..c087a41
--- /dev/null
+++ b/tests/bridge/test_mcp_server_smoke.py
@@ -0,0 +1,15 @@
+from src.bridge import mcp_server
+
+
+def test_build_server_registers_tools():
+    mcp = mcp_server.build_server()
+    # FastMCP stores tools; assert our names are present
+    tool_names = mcp_server.registered_tool_names(mcp)
+    for name in ["comfy_pull", "comfy_push", "comfy_list_channels",
+                 "comfy_run_workflow", "comfy_get_result"]:
+        assert name in tool_names
+
+
+def test_start_is_idempotent():
+    mcp_server.start_in_background(port=0, _test_no_serve=True)
+    mcp_server.start_in_background(port=0, _test_no_serve=True)  # no crash