"""Node introspection tools for ComfyUI MCP Server.
Thin MCP wrapper around NodeOrchestrator for node discovery operations.
Enables AI agents to understand available nodes and their parameters.
Primary Use Cases:
- Troubleshoot user-provided workflows by validating node configurations
- Extend existing workflows by discovering compatible nodes
- Build new workflows from scratch by understanding available capabilities
- Validate workflow JSON before execution to catch errors early
"""
from fastmcp import Context, FastMCP
from src.orchestrators.node import NodeOrchestrator
from src.utils import get_global_logger
logger = get_global_logger("MCP_Server.tools.node")
def register_node_tools(mcp: FastMCP, node_orchestrator: NodeOrchestrator):
"""Register node introspection tools with the MCP server.
Args:
mcp: FastMCP server instance
node_orchestrator: Orchestrator for node operations
"""
# ========================================================================
# RESOURCES - Allow agents to inspect node catalog
# ========================================================================
@mcp.resource("nodes://catalog")
async def get_node_catalog_resource(ctx: Context = None) -> str:
"""List all available ComfyUI nodes with categories.
Returns a formatted catalog organized by category.
"""
nodes = await node_orchestrator.get_all_nodes(ctx=ctx)
categories = await node_orchestrator.get_categories(ctx=ctx)
lines = ["# Available ComfyUI Nodes\n"]
lines.append(f"Total nodes: {len(nodes)}\n")
lines.append(f"Categories: {len(categories)}\n")
# Organize by category
for category in categories:
cat_nodes = [n for n in nodes.values() if n.category == category]
if cat_nodes:
lines.append(f"\n## {category} ({len(cat_nodes)} nodes)\n")
for node in sorted(cat_nodes, key=lambda n: n.class_type):
lines.append(f"- **{node.class_type}**: {node.display_name}")
# Uncategorized nodes
uncategorized = [n for n in nodes.values() if not n.category]
if uncategorized:
lines.append(f"\n## Uncategorized ({len(uncategorized)} nodes)\n")
for node in sorted(uncategorized, key=lambda n: n.class_type):
lines.append(f"- **{node.class_type}**")
return "\n".join(lines)
@mcp.resource("nodes://{class_type}/definition")
async def get_node_definition_resource(class_type: str, ctx: Context = None) -> dict:
"""Get complete definition for a specific node.
Args:
class_type: Node class identifier (e.g., "KSampler")
Returns:
Complete node definition with inputs, outputs, and metadata
"""
node = await node_orchestrator.get_node_by_class(class_type, ctx=ctx)
if not node:
return {"error": f"Node '{class_type}' not found"}
return node.to_dict()
# ========================================================================
# TOOLS - Node discovery and search
# ========================================================================
@mcp.tool()
async def list_available_nodes(
category: str | None = None,
force_refresh: bool = False,
ctx: Context = None,
) -> dict:
"""List all available ComfyUI nodes.
Returns a catalog of nodes organized by category, with basic metadata
for each node. Use this to discover what nodes are available in the
ComfyUI environment.
Args:
category: Optional filter by category (e.g., "sampling", "conditioning")
force_refresh: Force refresh node cache (default: False)
Returns:
Dictionary with:
- nodes: List of node summaries
- categories: List of all categories
- total_count: Total number of nodes
Example:
>>> # Discover available nodes when building a workflow
>>> list_available_nodes()
{
"nodes": [
{
"class_type": "KSampler",
"display_name": "KSampler",
"category": "sampling",
"input_count": 8,
"output_count": 1
},
...
],
"categories": ["sampling", "conditioning", ...],
"total_count": 247
}
>>> # Filter to specific category when extending workflow
>>> list_available_nodes(category="upscale")
# Returns only upscaling nodes
"""
logger.info(f"Listing nodes (category={category}, force_refresh={force_refresh})")
if ctx:
await ctx.info(f"Listing nodes (category={category})")
nodes = await node_orchestrator.get_all_nodes(force_refresh=force_refresh, ctx=ctx)
# Filter by category if specified
if category:
nodes = {k: v for k, v in nodes.items() if v.category == category}
logger.debug(f"Filtered to {len(nodes)} nodes in category '{category}'")
# Build summary list
node_list = [
{
"class_type": node.class_type,
"display_name": node.display_name,
"category": node.category,
"description": node.description[:100] + "..."
if len(node.description) > 100
else node.description,
"input_count": len(node.get_all_inputs()),
"output_count": len(node.output_types),
"python_module": node.python_module,
}
for node in sorted(nodes.values(), key=lambda n: n.class_type)
]
categories = await node_orchestrator.get_categories(ctx=ctx)
return {
"nodes": node_list,
"categories": categories,
"total_count": len(nodes),
"filtered": category is not None,
}
@mcp.tool()
async def get_node_definition(class_type: str, ctx: Context = None) -> dict:
"""Get detailed schema for a specific ComfyUI node.
Returns complete information about a node including all input parameters,
their types, default values, output types, and metadata. Use this to
understand how to configure a specific node in a workflow.
Args:
class_type: Node class identifier (e.g., "KSampler", "CLIPTextEncode")
Returns:
Dictionary with complete node definition or error if not found
Example:
>>> # Troubleshoot: User's workflow uses unknown node
>>> get_node_definition("KSampler")
{
"class_type": "KSampler",
"display_name": "KSampler",
"category": "sampling",
"input_types": {
"required": [
{
"name": "model",
"type": "MODEL",
"required": true,
"description": "..."
},
...
]
},
"output_types": [{"index": 0, "type": "LATENT"}],
...
}
>>> # Use output_types to find compatible next nodes
>>> # LATENT output → search for nodes accepting LATENT input
"""
logger.info(f"Getting definition for node '{class_type}'")
if ctx:
await ctx.info(f"Getting definition for node '{class_type}'")
node = await node_orchestrator.get_node_by_class(class_type, ctx=ctx)
if not node:
logger.warning(f"Node '{class_type}' not found")
return {"error": f"Node '{class_type}' not found"}
definition = node.to_dict()
logger.debug(
f"Returned definition with {len(node.get_all_inputs())} inputs, "
f"{len(node.output_types)} outputs"
)
return definition
@mcp.tool()
async def search_nodes(
query: str | None = None,
category: str | None = None,
output_type: str | None = None,
ctx: Context = None,
) -> dict:
"""Search for ComfyUI nodes by query, category, or output type.
Performs flexible search across node names, categories, and capabilities.
Use this to find nodes that match specific criteria for workflow construction.
Args:
query: Search string (matches class_type, display_name, description)
category: Filter by category (e.g., "sampling", "conditioning")
output_type: Filter by output type (e.g., "IMAGE", "LATENT", "MODEL")
Returns:
Dictionary with search results
Example:
>>> # Find sampling nodes for workflow construction
>>> search_nodes(query="sampler")
{
"results": [
{
"class_type": "KSampler",
"display_name": "KSampler",
"category": "sampling",
"match_reason": "name"
},
...
],
"count": 5
}
>>> # Find nodes that output images (for workflow extension)
>>> search_nodes(output_type="IMAGE")
# Returns: VAEDecode, LoadImage, ImageScale, etc.
>>> # Find nodes in specific category
>>> search_nodes(category="conditioning")
# Returns: CLIPTextEncode, CLIPSetLastLayer, etc.
"""
logger.info(
f"Searching nodes (query={query}, category={category}, " f"output_type={output_type})"
)
if ctx:
await ctx.info(
f"Searching nodes (query={query}, category={category}, output_type={output_type})"
)
# Use orchestrator search methods
if output_type:
results = await node_orchestrator.get_nodes_by_output_type(output_type, ctx=ctx)
match_reason = f"output_type={output_type}"
else:
results = await node_orchestrator.search_nodes(query=query, category=category, ctx=ctx)
match_reason = "query/category match"
# Format results
result_list = [
{
"class_type": node.class_type,
"display_name": node.display_name,
"category": node.category,
"description": node.description[:100] + "..."
if len(node.description) > 100
else node.description,
"input_count": len(node.get_all_inputs()),
"output_count": len(node.output_types),
"output_types": node.return_types,
"match_reason": match_reason,
}
for node in results
]
logger.info(f"Search returned {len(result_list)} results")
return {
"results": result_list,
"count": len(result_list),
"query": query,
"category": category,
"output_type": output_type,
}
@mcp.tool()
async def get_node_inputs(class_type: str, ctx: Context = None) -> dict:
"""Get input parameter specifications for a specific node.
Returns detailed information about all input parameters (required and optional)
including types, default values, and valid options. Use this to understand
what parameters are needed to configure a node.
Args:
class_type: Node class identifier (e.g., "KSampler")
Returns:
Dictionary with required and optional inputs
Example:
>>> # Troubleshoot: User's workflow missing required parameters
>>> get_node_inputs("KSampler")
{
"class_type": "KSampler",
"required": [
{
"name": "seed",
"type": "INT",
"default": 0,
"options": null,
"description": "..."
},
{
"name": "sampler_name",
"type": "select",
"options": ["euler", "euler_ancestral", "dpm_2", ...],
"description": "..."
},
...
],
"optional": [...],
"required_count": 8,
"optional_count": 0
}
>>> # Agent identifies: workflow missing 'sampler_name' parameter
>>> # Agent provides valid options from 'options' field
"""
logger.info(f"Getting inputs for node '{class_type}'")
if ctx:
await ctx.info(f"Getting inputs for node '{class_type}'")
node = await node_orchestrator.get_node_by_class(class_type, ctx=ctx)
if not node:
logger.warning(f"Node '{class_type}' not found")
return {"error": f"Node '{class_type}' not found"}
required = [
{
"name": inp.name,
"type": inp.type if isinstance(inp.type, str) else inp.type.__name__,
"required": inp.required,
"default": inp.default,
"options": inp.options,
"description": inp.description,
}
for inp in node.get_required_inputs()
]
optional = [
{
"name": inp.name,
"type": inp.type if isinstance(inp.type, str) else inp.type.__name__,
"required": inp.required,
"default": inp.default,
"options": inp.options,
"description": inp.description,
}
for inp in node.get_optional_inputs()
]
logger.debug(f"Returned {len(required)} required, {len(optional)} optional inputs")
return {
"class_type": class_type,
"required": required,
"optional": optional,
"required_count": len(required),
"optional_count": len(optional),
}
