server.py

from mcp.server.fastmcp import FastMCP
from pathlib import Path
from typing import Dict
import src.loader as loader
import src.factory as factory
import sys

# Initialize FastMCP Server
mcp = FastMCP("ComfyUI")

# Global Registry
# Maps file_path -> WorkflowConfig
REGISTERED_CONFIGS: Dict[Path, loader.WorkflowConfig] = {}

def _generate_manifest() -> str:
    """Generates a text manifest of available workflows."""
    if not REGISTERED_CONFIGS:
        return "No workflows registered."
    
    lines = ["# Available Workflows\n"]
    for config in REGISTERED_CONFIGS.values():
        lines.append(f"## {config.name}")
        lines.append(f"_{config.description}_")
        lines.append("### Parameters:")
        for param in config.parameters:
            req = "(Required)" if param.required else "(Optional)"
            lines.append(f"- **{param.name}** ({param.type}): {param.description} {req}")
        lines.append("\n---")
    
    return "\n".join(lines)

@mcp.resource("comfy://list")
def list_resource() -> str:
    """Returns a list of available workflows and their parameters."""
    return _generate_manifest()

@mcp.tool()
def list_available_workflows() -> str:
    """
    Returns a list of all available workflows on this server.
    Use this to discover what you can do.
    """
    return _generate_manifest()

def register_workflows():
    """
    Scans the workflows directory and registers them as MCP tools.
    """
    global REGISTERED_CONFIGS
    workflows_dir = Path(__file__).parent / "workflows"
    
    # 1. Scan for valid configs
    print(f"Scanning for workflows in {workflows_dir}...")
    REGISTERED_CONFIGS = loader.scan_workflows(workflows_dir)
    
    if not REGISTERED_CONFIGS:
        print("No valid workflows found to register.")
        return

    # Register each config as a tool
    for file_path, config in REGISTERED_CONFIGS.items():
        try:
            # Create the dynamic function
            tool_func = factory.create_dynamic_tool(config, file_path)
            
            # Register with FastMCP
            # add_tool inspects the function signature and docstring
            mcp.add_tool(tool_func)
            
            print(f"Registered Tool: {config.name} (from {file_path.name})")
            
        except Exception as e:
            print(f"Failed to register tool for {file_path.name}: {e}")

# Register tools on module import (Essential for SSE/mcp run mode)
register_workflows()

if __name__ == "__main__":
    print("Starting ComfyUI MCP Server...")
    # Run the server
    mcp.run()