From 85cf047c930f2316e5e94bd9047bf0cb582e55d5 Mon Sep 17 00:00:00 2001 From: Alastair Beal Date: Mon, 29 Jun 2026 22:38:18 +0100 Subject: [PATCH] comfyui: route custom workflows through image/video selectors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The selector path previously hid the custom-workflow feature: video_selector filtered tools on per-operation readiness (bundled WAN models) and both selectors only chose ToolStatus.AVAILABLE providers, so comfyui_image/ comfyui_video — DEGRADED when bundled model metadata is missing — were dropped even when the ComfyUI server was up and the caller supplied a full workflow_json/workflow_path plus output_node. - Add a custom-workflow readiness path to both selectors: when a custom workflow is supplied, eligibility is based on server availability (status != UNAVAILABLE) for any provider advertising supports.custom_workflow, not on bundled-model readiness. A custom workflow also restricts routing to custom-workflow-capable providers, since the graph JSON is ComfyUI specific. - Expose workflow_json, workflow_path, output_node, workflow_name, workflow_model, and workflow_model_stack in both selector schemas so agents can discover the feature without bypassing the selectors. - image_selector only forwards the workflow inputs to providers that declare them. - Add contract tests for the new eligibility path and schema exposure. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/contracts/test_comfyui_tools.py | 107 ++++++++++++++++++++++++++ tools/graphics/image_selector.py | 74 +++++++++++++++++- tools/video/video_selector.py | 70 ++++++++++++++++- 3 files changed, 248 insertions(+), 3 deletions(-) diff --git a/tests/contracts/test_comfyui_tools.py b/tests/contracts/test_comfyui_tools.py index 05009beb..5d636bd0 100644 --- a/tests/contracts/test_comfyui_tools.py +++ b/tests/contracts/test_comfyui_tools.py @@ -18,6 +18,7 @@ from tools.base_tool import ( ToolTier, ) from tools.graphics.comfyui_image import ComfyUIImage +from tools.graphics.image_selector import ImageSelector from tools.tool_registry import ToolRegistry from tools.video.video_selector import VideoSelector from tools.video.comfyui_video import ComfyUIVideo @@ -543,3 +544,109 @@ class TestVideoOperationReadiness: assert rank_inputs["operation"] == "image_to_video" assert selector._filter_candidates(rank_inputs, candidates) == [] + +# ------------------------------------------------------------------ +# Custom-workflow selector eligibility +# ------------------------------------------------------------------ + +class _DegradedComfyVideo(BaseTool): + """Server reachable, but bundled WAN models missing -> DEGRADED, no + operation ready. Stands in for comfyui_video on a low-VRAM box.""" + + name = "comfyui_video" + capability = "video_generation" + provider = "comfyui" + supports = {"custom_workflow": True, "image_to_video": True} + input_schema = {"type": "object", "properties": {"workflow_json": {"type": "string"}}} + + def get_status(self): + return ToolStatus.DEGRADED + + def is_operation_available(self, operation): + return False + + def execute(self, inputs): + raise AssertionError("not used") + + +class _DegradedComfyImage(BaseTool): + name = "comfyui_image" + capability = "image_generation" + provider = "comfyui" + supports = {"custom_workflow": True} + input_schema = {"type": "object", "properties": {"workflow_json": {"type": "string"}}} + + def get_status(self): + return ToolStatus.DEGRADED + + def execute(self, inputs): + raise AssertionError("not used") + + +class TestCustomWorkflowSelectorEligibility: + + def test_video_selector_passes_degraded_tool_for_custom_workflow(self): + selector = VideoSelector() + candidates = [_DegradedComfyVideo()] + inputs = { + "prompt": "x", + "operation": "text_to_video", + "workflow_json": "{}", + "output_node": "14", + } + # Without the custom-workflow path this DEGRADED, operation-unready tool + # would be filtered out; with it, it is eligible and selectable. + filtered = selector._filter_candidates(inputs, candidates) + assert [t.name for t in filtered] == ["comfyui_video"] + assert selector._tool_selectable(candidates[0], inputs) is True + + def test_video_selector_custom_workflow_requires_output_node(self): + selector = VideoSelector() + candidates = [_DegradedComfyVideo()] + inputs = {"prompt": "x", "operation": "text_to_video", "workflow_json": "{}"} + # output_node missing -> not eligible -> filtered out. + assert selector._filter_candidates(inputs, candidates) == [] + assert selector._tool_selectable(candidates[0], inputs) is False + + def test_video_selector_custom_workflow_needs_server(self): + class _OfflineComfyVideo(_DegradedComfyVideo): + def get_status(self): + return ToolStatus.UNAVAILABLE + + selector = VideoSelector() + candidates = [_OfflineComfyVideo()] + inputs = { + "prompt": "x", + "operation": "text_to_video", + "workflow_json": "{}", + "output_node": "14", + } + assert selector._filter_candidates(inputs, candidates) == [] + + def test_image_selector_passes_degraded_tool_for_custom_workflow(self): + selector = ImageSelector() + candidates = [_DegradedComfyImage()] + inputs = {"prompt": "x", "workflow_json": "{}", "output_node": "13"} + filtered = selector._filter_candidates(inputs, candidates) + assert [t.name for t in filtered] == ["comfyui_image"] + assert selector._tool_selectable(candidates[0], inputs) is True + + def test_image_selector_custom_workflow_requires_output_node(self): + selector = ImageSelector() + candidates = [_DegradedComfyImage()] + inputs = {"prompt": "x", "workflow_json": "{}"} + assert selector._filter_candidates(inputs, candidates) == [] + assert selector._tool_selectable(candidates[0], inputs) is False + + def test_selector_schemas_expose_custom_workflow_inputs(self): + for selector in (VideoSelector(), ImageSelector()): + props = selector.input_schema["properties"] + for field in ( + "workflow_json", + "workflow_path", + "output_node", + "workflow_name", + "workflow_model", + "workflow_model_stack", + ): + assert field in props, f"{selector.name} missing {field}" diff --git a/tools/graphics/image_selector.py b/tools/graphics/image_selector.py index d92a60e5..0621f51e 100644 --- a/tools/graphics/image_selector.py +++ b/tools/graphics/image_selector.py @@ -94,6 +94,38 @@ class ImageSelector(BaseTool): "default": "generate", "description": "Operation mode. 'rank' returns scored provider rankings without generating.", }, + "workflow_json": { + "type": "string", + "description": ( + "Optional full ComfyUI workflow JSON. Routes to a custom-workflow-capable " + "provider (e.g. comfyui_image) based on server availability, not bundled " + "model readiness. Requires output_node." + ), + }, + "workflow_path": { + "type": "string", + "description": ( + "Optional path to a ComfyUI workflow JSON file. Routes to a custom-workflow-" + "capable provider based on server availability. Requires output_node." + ), + }, + "output_node": { + "type": "string", + "description": "ComfyUI output node ID for a custom workflow_json/workflow_path.", + }, + "workflow_name": { + "type": "string", + "description": "Optional human-readable provenance label for a custom workflow.", + }, + "workflow_model": { + "type": "string", + "description": "Optional model/provenance label for a custom workflow.", + }, + "workflow_model_stack": { + "type": "array", + "items": {"type": "object"}, + "description": "Optional provenance metadata for custom workflow dependencies.", + }, "output_path": {"type": "string"}, }, } @@ -184,6 +216,12 @@ class ImageSelector(BaseTool): "image_path", "image_urls", "image_paths", + "workflow_json", + "workflow_path", + "output_node", + "workflow_name", + "workflow_model", + "workflow_model_stack", ): if passthrough_key in adapted and passthrough_key not in props: stripped.append(f"{passthrough_key}={adapted.pop(passthrough_key)}") @@ -226,7 +264,7 @@ class ImageSelector(BaseTool): tool_by_provider: dict[str, BaseTool] = {} for tool in candidates: - if tool.provider not in tool_by_provider and tool.get_status() == ToolStatus.AVAILABLE: + if tool.provider not in tool_by_provider and self._tool_selectable(tool, inputs): tool_by_provider[tool.provider] = tool if preferred != "auto": @@ -277,6 +315,12 @@ class ImageSelector(BaseTool): return serialized def _filter_candidates(self, inputs: dict[str, Any], candidates: list[BaseTool]) -> list[BaseTool]: + # A caller-supplied custom workflow is provider-specific (ComfyUI graph + # JSON). Route it only to custom-workflow-capable providers whose server + # is reachable — bundled-model readiness is irrelevant in that case. + if self._has_custom_workflow(inputs): + return [t for t in candidates if self._custom_workflow_eligible(t, inputs)] + wants_edit = ( inputs.get("generation_mode") == "edit" or inputs.get("image_url") @@ -296,3 +340,31 @@ class ImageSelector(BaseTool): ): filtered.append(tool) return filtered or candidates + + @staticmethod + def _has_custom_workflow(inputs: dict[str, Any]) -> bool: + return bool(inputs.get("workflow_json") or inputs.get("workflow_path")) + + def _custom_workflow_eligible(self, tool: BaseTool, inputs: dict[str, Any]) -> bool: + """Whether a tool can run the caller-supplied custom workflow. + + Eligibility is based on server availability, not bundled-model readiness: + a provider qualifies when it advertises ``custom_workflow`` support, an + ``output_node`` is supplied, and its backend is reachable (status is not + UNAVAILABLE). + """ + if not self._has_custom_workflow(inputs): + return False + if not inputs.get("output_node"): + return False + supports = getattr(tool, "supports", {}) + if not supports.get("custom_workflow"): + return False + return tool.get_status() != ToolStatus.UNAVAILABLE + + def _tool_selectable(self, tool: BaseTool, inputs: dict[str, Any]) -> bool: + """A provider is selectable if it is AVAILABLE, or if it can serve a + caller-supplied custom workflow even while bundled models report DEGRADED.""" + if tool.get_status() == ToolStatus.AVAILABLE: + return True + return self._custom_workflow_eligible(tool, inputs) diff --git a/tools/video/video_selector.py b/tools/video/video_selector.py index 0ae77c68..8e9a9412 100644 --- a/tools/video/video_selector.py +++ b/tools/video/video_selector.py @@ -96,6 +96,38 @@ class VideoSelector(BaseTool): "type": "string", "description": "Resolution hint for providers that support named output resolutions.", }, + "workflow_json": { + "type": "string", + "description": ( + "Optional full ComfyUI workflow JSON. Routes to a custom-workflow-capable " + "provider (e.g. comfyui_video) based on server availability, not bundled " + "model readiness. Requires output_node." + ), + }, + "workflow_path": { + "type": "string", + "description": ( + "Optional path to a ComfyUI workflow JSON file. Routes to a custom-workflow-" + "capable provider based on server availability. Requires output_node." + ), + }, + "output_node": { + "type": "string", + "description": "ComfyUI output node ID for a custom workflow_json/workflow_path.", + }, + "workflow_name": { + "type": "string", + "description": "Optional human-readable provenance label for a custom workflow.", + }, + "workflow_model": { + "type": "string", + "description": "Optional model/provenance label for a custom workflow.", + }, + "workflow_model_stack": { + "type": "array", + "items": {"type": "object"}, + "description": "Optional provenance metadata for custom workflow dependencies.", + }, "output_path": {"type": "string"}, }, } @@ -231,10 +263,10 @@ class VideoSelector(BaseTool): rankings = rank_providers(candidates, task_context) - # Build tool lookup: provider → tool (first available per provider) + # Build tool lookup: provider → tool (first selectable per provider) tool_by_provider: dict[str, BaseTool] = {} for tool in candidates: - if tool.provider not in tool_by_provider and tool.get_status() == ToolStatus.AVAILABLE: + if tool.provider not in tool_by_provider and self._tool_selectable(tool, inputs): tool_by_provider[tool.provider] = tool # If a preferred provider is explicitly requested and available, @@ -298,6 +330,12 @@ class VideoSelector(BaseTool): inputs: dict[str, object], candidates: list[BaseTool], ) -> list[BaseTool]: + # A caller-supplied custom workflow is provider-specific (ComfyUI graph + # JSON). Route it only to custom-workflow-capable providers whose server + # is reachable — bundled-model readiness is irrelevant in that case. + if self._has_custom_workflow(inputs): + return [t for t in candidates if self._custom_workflow_eligible(t, inputs)] + operation = inputs.get("operation", "text_to_video") if operation == "rank": operation = inputs.get("target_operation", "text_to_video") @@ -333,3 +371,31 @@ class VideoSelector(BaseTool): if not callable(checker): return True return bool(checker(operation)) + + @staticmethod + def _has_custom_workflow(inputs: dict[str, object]) -> bool: + return bool(inputs.get("workflow_json") or inputs.get("workflow_path")) + + def _custom_workflow_eligible(self, tool: BaseTool, inputs: dict[str, object]) -> bool: + """Whether a tool can run the caller-supplied custom workflow. + + Eligibility is based on server availability, not bundled-model readiness: + a provider qualifies when it advertises ``custom_workflow`` support, an + ``output_node`` is supplied, and its backend is reachable (status is not + UNAVAILABLE). + """ + if not self._has_custom_workflow(inputs): + return False + if not inputs.get("output_node"): + return False + supports = getattr(tool, "supports", {}) + if not supports.get("custom_workflow"): + return False + return tool.get_status() != ToolStatus.UNAVAILABLE + + def _tool_selectable(self, tool: BaseTool, inputs: dict[str, object]) -> bool: + """A provider is selectable if it is AVAILABLE, or if it can serve a + caller-supplied custom workflow even while bundled models report DEGRADED.""" + if tool.get_status() == ToolStatus.AVAILABLE: + return True + return self._custom_workflow_eligible(tool, inputs)