Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -131,7 +131,7 @@ docker run --rm ghcr.io/hybridindie/comfyui-mcp:latest --help
| `upscale_image` | Upscale an image using a model-based upscaler. Params: image (filename), upscale_model (default: RealESRGAN_x4plus.pth). |
| `run_workflow` | Submit arbitrary ComfyUI workflow JSON. Inspected for dangerous nodes before execution. Set `wait=True` to block until complete and return outputs. |
| `summarize_workflow` | Summarize a workflow's structure, data flow, models, and parameters. Supports `format="text"` (default) or `format="mermaid"` for diagram markup. |
| `create_workflow` | Create a workflow from a template (txt2img, img2img, upscale, inpaint, txt2vid_animatediff, txt2vid_wan) with parameter overrides. |
| `create_workflow` | Create a workflow from templates including txt2img/img2img/upscale/inpaint, txt2vid_animatediff/txt2vid_wan, controlnet_canny/controlnet_depth/controlnet_openpose, ip_adapter, lora_stack, face_restore, flux_txt2img, and sdxl_txt2img. |
| `modify_workflow` | Apply batch operations (add_node, remove_node, set_input, connect, disconnect) to a workflow. |
| `validate_workflow` | Validate workflow structure, server compatibility, and security. |

Expand Down
2 changes: 1 addition & 1 deletion src/comfyui_mcp/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,7 @@ def _register_all_tools(
model_checker=model_checker,
sanitizer=sanitizer,
)
register_workflow_tools(server, client, audit, rate_limiters["read"], inspector)
register_workflow_tools(server, client, audit, rate_limiters["read"], inspector, sanitizer)
register_model_tools(
mcp=server,
client=client,
Expand Down
41 changes: 38 additions & 3 deletions src/comfyui_mcp/tools/workflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,17 +11,44 @@
from comfyui_mcp.client import ComfyUIClient
from comfyui_mcp.security.inspector import WorkflowInspector
from comfyui_mcp.security.rate_limit import RateLimiter
from comfyui_mcp.security.sanitizer import PathSanitizer, PathValidationError
from comfyui_mcp.workflow.operations import apply_operations
from comfyui_mcp.workflow.templates import create_from_template
from comfyui_mcp.workflow.validation import validate_workflow as _validate_workflow

_PATH_LIKE_TEMPLATE_PARAMS = {
"model",
"model_name",
"motion_module",
"controlnet_model",
"ipadapter_model",
"clip_vision_model",
"lora_name",
"face_restore_model",
"image",
"mask",
}


def _sanitize_template_params(
param_dict: dict[str, Any], sanitizer: PathSanitizer
) -> dict[str, Any]:
"""Sanitize filename-like template params to block traversal/null-byte inputs."""
sanitized = dict(param_dict)
for key in _PATH_LIKE_TEMPLATE_PARAMS:
value = sanitized.get(key)
if isinstance(value, str):
sanitized[key] = sanitizer.validate_path_segment(value, label=key)
return sanitized


def register_workflow_tools(
mcp: FastMCP,
client: ComfyUIClient,
audit: AuditLogger,
limiter: RateLimiter,
inspector: WorkflowInspector,
sanitizer: PathSanitizer,
) -> dict[str, Any]:
"""Register workflow composition tools."""
tool_fns: dict[str, Any] = {}
Expand All @@ -30,13 +57,16 @@ def register_workflow_tools(
async def create_workflow(template: str, params: str = "{}") -> str:
"""Create a ComfyUI workflow from a template with optional parameter overrides.
Available templates: txt2img, img2img, upscale, inpaint, txt2vid_animatediff, txt2vid_wan.
Available templates: txt2img, img2img, upscale, inpaint, txt2vid_animatediff,
txt2vid_wan, controlnet_canny, controlnet_depth, controlnet_openpose,
ip_adapter, lora_stack, face_restore, flux_txt2img, sdxl_txt2img.
Args:
template: Template name (e.g. 'txt2img', 'img2img')
params: Optional JSON string of parameter overrides.
Common params: prompt, negative_prompt, width, height,
steps, cfg, model, denoise.
steps, cfg, model, denoise, controlnet_model,
control_strength, lora_name, lora_strength.
Comment on lines +60 to +69
Copy link

Copilot AI Mar 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

create_workflow now documents additional params that are effectively filenames/model identifiers (e.g. controlnet_model, lora_name, clip_vision_model, ipadapter_model). The tool currently forwards user-provided strings directly into the workflow JSON without any path/filename sanitization. Per project security rules, tools that accept filenames/subfolders should validate/sanitize these inputs (e.g. via PathSanitizer) before returning a workflow, so callers can’t accidentally (or intentionally) generate workflows that reference absolute paths / traversal sequences / null bytes.

Copilot uses AI. Check for mistakes.
"""
limiter.check("create_workflow")
try:
Expand All @@ -47,7 +77,12 @@ async def create_workflow(template: str, params: str = "{}") -> str:
if not isinstance(param_dict, dict):
raise ValueError('params must be a JSON object (e.g. {"key": "value"})')

wf = create_from_template(template, param_dict)
try:
clean_params = _sanitize_template_params(param_dict, sanitizer)
except PathValidationError as e:
raise ValueError(str(e)) from e

wf = create_from_template(template, clean_params)
audit.log(
tool="create_workflow",
action="created",
Expand Down
Loading
Loading