runtime.py•5.19 kB
"""FastMCP runtime integration for the ComfyUI MCP server."""
from __future__ import annotations
import argparse
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any, Dict, Mapping, MutableMapping, Optional, Sequence
from mcp.server.fastmcp import Context, FastMCP
from .config import ComfyUISettings, load_settings
from .server import ComfyUIMCPServer
def create_mcp_app(
*,
settings: Optional[ComfyUISettings] = None,
config_path: Optional[Path] = None,
host: str = "127.0.0.1",
port: int = 8000,
instructions: Optional[str] = None,
) -> FastMCP:
"""Create a FastMCP application exposing ComfyUI workflow tooling."""
server_state: MutableMapping[str, Any] = {}
resolved_config: Optional[Path] = config_path.expanduser().resolve() if config_path else None
@asynccontextmanager
async def lifespan(app: FastMCP): # noqa: D401 - FastMCP expects this signature
"""Manage the lifecycle of the ComfyUI backend server."""
cfg = settings or load_settings(resolved_config)
async with ComfyUIMCPServer(cfg).lifecycle() as backend:
server_state["settings"] = cfg
server_state["backend"] = backend
try:
yield {"settings": cfg}
finally:
server_state.clear()
app = FastMCP(
name="comfyui-mcp",
instructions=instructions
or "Manage ComfyUI workflows: list templates, describe nodes, customise parameters, and execute prompts.",
dependencies=["comfyui-mcp"],
host=host,
port=port,
lifespan=lifespan,
)
def _require_backend() -> ComfyUIMCPServer:
backend = server_state.get("backend")
if backend is None:
raise RuntimeError("ComfyUI MCP backend is not initialised")
return backend
def _coerce_changes(changes: Mapping[str, Any] | None) -> Mapping[str, Any]:
if changes is None:
return {}
if isinstance(changes, Mapping):
return dict(changes)
raise TypeError("changes must be a mapping of parameter names to values")
@app.tool(name="list_workflows", description="List available workflow templates.", structured_output=True)
async def list_workflows(context: Context) -> list[dict[str, Any]]: # noqa: D401 - FastMCP tool signature
backend = _require_backend()
return await backend.list_workflows()
@app.tool(
name="describe_workflow",
description="Describe a workflow including semantic roles.",
structured_output=True,
)
async def describe_workflow(name: str, context: Context) -> dict[str, Any]:
backend = _require_backend()
return await backend.describe_workflow(name)
@app.tool(
name="list_assets",
description="Enumerate discovered checkpoints, LoRAs, VAEs, and related assets.",
structured_output=True,
)
async def list_assets(context: Context) -> dict[str, list[str]]:
backend = _require_backend()
return await backend.list_assets()
@app.tool(
name="customize_workflow",
description="Apply high-level changes to a workflow template without executing it.",
structured_output=True,
)
async def customize_workflow(name: str, changes: Optional[Mapping[str, Any]] = None, context: Context | None = None) -> dict[str, Any]:
backend = _require_backend()
payload = await backend.customize_workflow(name, _coerce_changes(changes))
return payload
@app.tool(
name="execute_workflow",
description="Optionally mutate and execute a workflow via ComfyUI.",
structured_output=True,
)
async def execute_workflow(
name: str,
changes: Optional[Mapping[str, Any]] = None,
stream_updates: bool = False,
context: Context | None = None,
) -> dict[str, Any]:
backend = _require_backend()
result = await backend.execute_workflow(name, _coerce_changes(changes), stream_updates=stream_updates)
return result
return app
def main(argv: Optional[Sequence[str]] = None) -> None:
"""Run the FastMCP server with configurable transport."""
parser = argparse.ArgumentParser(description="Run the ComfyUI MCP FastMCP server")
parser.add_argument("--config", type=Path, help="Path to a configuration TOML file")
parser.add_argument("--host", default="127.0.0.1", help="Host binding for HTTP-based transports")
parser.add_argument("--port", type=int, default=8000, help="Port for HTTP-based transports")
parser.add_argument(
"--transport",
choices=["stdio", "sse", "streamable-http"],
default="stdio",
help="Transport protocol exposed to MCP clients",
)
parser.add_argument("--instructions", help="Override the instructions string advertised to clients")
args = parser.parse_args(argv)
app = create_mcp_app(
config_path=args.config,
host=args.host,
port=args.port,
instructions=args.instructions,
)
app.run(transport=args.transport)
__all__ = ["create_mcp_app", "main"]