From 85ddd9d042e93a751ee463a1ae500314d07ec2b2 Mon Sep 17 00:00:00 2001 From: Ken Simpson Date: Sat, 18 Apr 2026 11:54:23 -0700 Subject: [PATCH] feat: add modal remote expansion context menu --- README.md | 2 + api_intercept.py | 302 +++++++++++++++++++++++++++++++-- docs/LEARNINGS.md | 1 + tests/test_api_intercept.py | 153 +++++++++++++++++ tests/test_web_modal_toggle.py | 26 +++ web/modal_toggle.js | 252 +++++++++++++++++++++++++++ 6 files changed, 725 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 76552ed..f169d45 100644 --- a/README.md +++ b/README.md @@ -138,6 +138,8 @@ For each connected region you want to offload: 2. Enable it on every node that should be part of the same remote island. 3. Confirm the node shows the blue remote-execution border. +If queue-time validation tells you a remote node still depends on local-only runtime objects, you no longer have to hunt those upstream nodes manually. Right-click the node and use `Modal: Include Required Upstream Nodes` to ask the backend which extra nodes must join that remote island. If multiple nodes are selected in the current graph, the same context menu expands the whole selection at once. + The toggle stores `properties.is_modal_remote = true` in workflow metadata. The visible graph is not rewritten in the editor. Rewrite happens only at queue time. ### 4. Queue the workflow diff --git a/api_intercept.py b/api_intercept.py index 1f6e38a..e44070e 100644 --- a/api_intercept.py +++ b/api_intercept.py @@ -94,6 +94,33 @@ class RewriteSummary: custom_nodes_bundle: SyncedAsset | None = None +@dataclass(frozen=True) +class RemoteExpansionReason: + """Describe why one upstream node had to join a remote component.""" + + node_id: str + class_type: str + required_by_node_id: str + required_by_class_type: str + output_index: int + io_type: str + + +@dataclass +class RemoteNodeAnalysis: + """Structured dry-run result for context-menu remote expansion.""" + + requested_node_ids: list[str] = field(default_factory=list) + requested_workflow_node_paths: list[str] = field(default_factory=list) + current_remote_node_ids: list[str] = field(default_factory=list) + current_remote_workflow_node_paths: list[str] = field(default_factory=list) + resolved_remote_node_ids: list[str] = field(default_factory=list) + resolved_workflow_node_paths: list[str] = field(default_factory=list) + added_node_ids: list[str] = field(default_factory=list) + added_workflow_node_paths: list[str] = field(default_factory=list) + reasons: list[RemoteExpansionReason] = field(default_factory=list) + + class ModalPromptValidationError(ValueError): """Raised when a prompt cannot be executed with the current Modal transport.""" @@ -289,6 +316,118 @@ def _resolve_prompt_node_ids_for_workflow_node( return resolved_prompt_node_ids +def _workflow_node_path(workflow_node_id: str, ancestor_node_ids: tuple[str, ...]) -> str: + """Return one workflow node's composed path, including subgraph ancestors.""" + if not ancestor_node_ids: + return workflow_node_id + return ":".join((*ancestor_node_ids, workflow_node_id)) + + +def _extract_marked_workflow_node_paths( + workflow: dict[str, Any] | None, + settings: ModalSyncSettings | None = None, +) -> set[str]: + """Return composed workflow paths for nodes explicitly marked remote in metadata.""" + if workflow is None: + return set() + + marker = (settings or get_settings()).marker_property + marked_workflow_node_paths: set[str] = set() + for node, ancestor_node_ids in _iter_workflow_nodes_with_ancestors(workflow): + properties = node.get("properties") or {} + if not properties.get(marker): + continue + marked_workflow_node_paths.add( + _workflow_node_path(str(node.get("id")), ancestor_node_ids) + ) + return marked_workflow_node_paths + + +def _build_workflow_prompt_resolution_maps( + workflow: dict[str, Any] | None, + prompt_node_ids: set[str], +) -> tuple[dict[str, set[str]], dict[str, set[str]]]: + """Return bidirectional mappings between workflow paths and queued prompt ids.""" + workflow_path_to_prompt_node_ids: dict[str, set[str]] = {} + prompt_node_id_to_workflow_paths: dict[str, set[str]] = defaultdict(set) + + if workflow is None: + for prompt_node_id in prompt_node_ids: + workflow_path_to_prompt_node_ids[prompt_node_id] = {prompt_node_id} + prompt_node_id_to_workflow_paths[prompt_node_id].add(prompt_node_id) + return workflow_path_to_prompt_node_ids, prompt_node_id_to_workflow_paths + + for node, ancestor_node_ids in _iter_workflow_nodes_with_ancestors(workflow): + workflow_node_id = str(node.get("id")) + workflow_node_path = _workflow_node_path(workflow_node_id, ancestor_node_ids) + resolved_prompt_node_ids = _resolve_prompt_node_ids_for_workflow_node( + workflow_node_id, + ancestor_node_ids, + prompt_node_ids, + ) + if workflow_node_id in prompt_node_ids: + resolved_prompt_node_ids.add(workflow_node_id) + if workflow_node_path in prompt_node_ids: + resolved_prompt_node_ids.add(workflow_node_path) + if not resolved_prompt_node_ids: + continue + + workflow_path_to_prompt_node_ids[workflow_node_path] = resolved_prompt_node_ids + for prompt_node_id in resolved_prompt_node_ids: + prompt_node_id_to_workflow_paths[prompt_node_id].add(workflow_node_path) + + for prompt_node_id in prompt_node_ids: + workflow_path_to_prompt_node_ids.setdefault(prompt_node_id, {prompt_node_id}) + prompt_node_id_to_workflow_paths[prompt_node_id].add(prompt_node_id) + + return workflow_path_to_prompt_node_ids, prompt_node_id_to_workflow_paths + + +def _resolve_requested_prompt_node_ids( + requested_workflow_node_paths: set[str], + prompt_node_ids: set[str], + workflow_path_to_prompt_node_ids: dict[str, set[str]], +) -> set[str]: + """Resolve requested workflow-node paths to the prompt ids that queue-time rewrite sees.""" + requested_prompt_node_ids: set[str] = set() + for requested_workflow_node_path in requested_workflow_node_paths: + if requested_workflow_node_path in workflow_path_to_prompt_node_ids: + requested_prompt_node_ids.update( + workflow_path_to_prompt_node_ids[requested_workflow_node_path] + ) + continue + if requested_workflow_node_path in prompt_node_ids: + requested_prompt_node_ids.add(requested_workflow_node_path) + return requested_prompt_node_ids + + +def _best_workflow_path_for_prompt_node( + prompt_node_id: str, + prompt_node_id_to_workflow_paths: dict[str, set[str]], +) -> str: + """Choose the most specific workflow path for one queued prompt node id.""" + workflow_node_paths = prompt_node_id_to_workflow_paths.get(prompt_node_id, {prompt_node_id}) + return max( + workflow_node_paths, + key=lambda workflow_node_path: ( + workflow_node_path.count(":"), + len(workflow_node_path), + workflow_node_path, + ), + ) + + +def _resolve_workflow_node_paths_for_prompt_nodes( + prompt_node_ids: set[str], + prompt_node_id_to_workflow_paths: dict[str, set[str]], +) -> set[str]: + """Map queued prompt ids back to the workflow node paths the UI can mark remote.""" + return { + _best_workflow_path_for_prompt_node(prompt_node_id, prompt_node_id_to_workflow_paths) + for prompt_node_id in prompt_node_ids + } + + def extract_remote_node_ids( workflow: dict[str, Any] | None, settings: ModalSyncSettings | None = None, @@ -404,10 +543,11 @@ def _expand_remote_node_ids_for_non_transportable_inputs( prompt: dict[str, Any], remote_node_ids: set[str], nodes_module: Any, -) -> set[str]: +) -> tuple[set[str], list[RemoteExpansionReason]]: """Grow the remote set upstream until non-transportable inputs stay inside the remote island.""" expanded_remote_node_ids = set(remote_node_ids) - added_node_ids: list[str] = [] + added_node_ids: set[str] = set() + reasons: list[RemoteExpansionReason] = [] changed = True while changed: @@ -442,14 +582,23 @@ def _expand_remote_node_ids_for_non_transportable_inputs( continue expanded_remote_node_ids.add(upstream_node_id) - added_node_ids.append(upstream_node_id) + reason = RemoteExpansionReason( + node_id=upstream_node_id, + class_type=upstream_class_type, + required_by_node_id=node_id, + required_by_class_type=str(prompt_node["class_type"]), + output_index=output_index, + io_type=io_type, + ) + reasons.append(reason) + added_node_ids.add(upstream_node_id) changed = True logger.info( "Auto-expanded remote execution upstream: added node %s (%s) because node %s depends on non-transportable type '%s'.", - upstream_node_id, - upstream_class_type, - node_id, - io_type, + reason.node_id, + reason.class_type, + reason.required_by_node_id, + reason.io_type, ) break if changed: @@ -460,9 +609,71 @@ def _expand_remote_node_ids_for_non_transportable_inputs( "Expanded remote node set from %d to %d nodes by absorbing upstream non-transportable dependencies: %s", len(remote_node_ids), len(expanded_remote_node_ids), - sorted(set(added_node_ids)), + sorted(added_node_ids), ) - return expanded_remote_node_ids + return expanded_remote_node_ids, reasons + + +def analyze_remote_node_selection( + prompt: dict[str, Any], + workflow: dict[str, Any] | None, + seed_workflow_node_paths: list[str], + settings: ModalSyncSettings | None = None, + nodes_module: Any | None = None, +) -> RemoteNodeAnalysis: + """Return the nodes the UI should mark remote for one context-menu expansion request.""" + resolved_settings = settings or get_settings() + resolved_nodes_module = nodes_module or _get_nodes_module() + prompt_node_ids = {str(node_id) for node_id in prompt.keys()} + requested_workflow_node_paths = { + str(seed_workflow_node_path) + for seed_workflow_node_path in seed_workflow_node_paths + if str(seed_workflow_node_path) + } + current_remote_node_ids = extract_remote_node_ids( + workflow, + resolved_settings, + prompt_node_ids, + ) + current_remote_workflow_node_paths = _extract_marked_workflow_node_paths( + workflow, + resolved_settings, + ) + workflow_path_to_prompt_node_ids, prompt_node_id_to_workflow_paths = ( + _build_workflow_prompt_resolution_maps(workflow, prompt_node_ids) + ) + requested_node_ids = _resolve_requested_prompt_node_ids( + requested_workflow_node_paths, + prompt_node_ids, + workflow_path_to_prompt_node_ids, + ) + initial_remote_node_ids = current_remote_node_ids | requested_node_ids + resolved_remote_node_ids, reasons = _expand_remote_node_ids_for_non_transportable_inputs( + prompt=prompt, + remote_node_ids=initial_remote_node_ids, + nodes_module=resolved_nodes_module, + ) + resolved_workflow_node_paths = ( + _resolve_workflow_node_paths_for_prompt_nodes( + resolved_remote_node_ids, + prompt_node_id_to_workflow_paths, + ) + | current_remote_workflow_node_paths + ) + added_node_ids = resolved_remote_node_ids - current_remote_node_ids + added_workflow_node_paths = resolved_workflow_node_paths - current_remote_workflow_node_paths + + return RemoteNodeAnalysis( + requested_node_ids=sorted(requested_node_ids), + requested_workflow_node_paths=sorted(requested_workflow_node_paths), + current_remote_node_ids=sorted(current_remote_node_ids), + current_remote_workflow_node_paths=sorted(current_remote_workflow_node_paths), + resolved_remote_node_ids=sorted(resolved_remote_node_ids), + resolved_workflow_node_paths=sorted(resolved_workflow_node_paths), + added_node_ids=sorted(added_node_ids), + added_workflow_node_paths=sorted(added_workflow_node_paths), + reasons=reasons, + ) def _build_component_plan( @@ -845,7 +1056,7 @@ def rewrite_prompt_for_modal( resolved_nodes_module = nodes_module or _get_nodes_module() resolved_sync_engine = sync_engine or ModalAssetSyncEngine.from_environment(resolved_settings) rewritten_prompt = copy.deepcopy(prompt) - expanded_remote_node_ids = _expand_remote_node_ids_for_non_transportable_inputs( + expanded_remote_node_ids, _ = _expand_remote_node_ids_for_non_transportable_inputs( prompt=rewritten_prompt, remote_node_ids=remote_node_ids, nodes_module=resolved_nodes_module, @@ -960,6 +1171,13 @@ async def _queue_prompt_json(prompt_server: Any, json_data: dict[str, Any]) -> w ) +def _analysis_route_path(route_path: str) -> str: + """Return the sibling HTTP route used for dry-run remote-node expansion.""" + if route_path.endswith("/queue_prompt"): + return f"{route_path.removesuffix('/queue_prompt')}/analyze_remote_nodes" + return f"{route_path.rstrip('/')}/analyze_remote_nodes" + + def setup_modal_queue_route( prompt_server: Any | None = None, sync_engine: ModalAssetSyncEngine | None = None, @@ -983,6 +1201,64 @@ def setup_modal_queue_route( return resolved_sync_engine = sync_engine or ModalAssetSyncEngine.from_environment(resolved_settings) + analysis_route_path = _analysis_route_path(resolved_settings.route_path) + + @prompt_server.routes.post(analysis_route_path) + async def modal_analyze_remote_nodes(request: web.Request) -> web.Response: + """Analyze which workflow nodes should be marked remote for the current graph.""" + logger.info("Received Modal remote-node analysis request.") + try: + request_started_at = time.perf_counter() + json_data = await request.json() + prompt = json_data.get("prompt") + if not isinstance(prompt, dict): + raise ValueError("Modal remote-node analysis requires a 'prompt' object.") + + workflow = json_data.get("workflow") + seed_node_ids = json_data.get("seed_node_ids") or [] + if not isinstance(seed_node_ids, list): + raise ValueError("Modal remote-node analysis requires 'seed_node_ids' to be a list.") + + analysis = analyze_remote_node_selection( + prompt=prompt, + workflow=workflow if isinstance(workflow, dict) else None, + seed_workflow_node_paths=[str(seed_node_id) for seed_node_id in seed_node_ids], + settings=resolved_settings, + ) + logger.info( + "Modal remote-node analysis finished in %.3fs with %d requested nodes and %d additions.", + time.perf_counter() - request_started_at, + len(analysis.requested_workflow_node_paths), + len(analysis.added_workflow_node_paths), + ) + return web.json_response( + { + "requested_node_ids": analysis.requested_node_ids, + "requested_workflow_node_paths": analysis.requested_workflow_node_paths, + "current_remote_node_ids": analysis.current_remote_node_ids, + "current_remote_workflow_node_paths": ( + analysis.current_remote_workflow_node_paths + ), + "resolved_remote_node_ids": analysis.resolved_remote_node_ids, + "resolved_workflow_node_paths": analysis.resolved_workflow_node_paths, + "added_node_ids": analysis.added_node_ids, + "added_workflow_node_paths": analysis.added_workflow_node_paths, + "reasons": [ + { + "node_id": reason.node_id, + "class_type": reason.class_type, + "required_by_node_id": reason.required_by_node_id, + "required_by_class_type": reason.required_by_class_type, + "output_index": reason.output_index, + "io_type": reason.io_type, + } + for reason in analysis.reasons + ], + } + ) + except (TypeError, ValueError) as exc: + logger.warning("Modal remote-node analysis request was invalid: %s", exc) + return web.json_response({"error": str(exc), "node_errors": []}, status=400) @prompt_server.routes.post(resolved_settings.route_path) async def modal_queue_prompt(request: web.Request) -> web.Response: @@ -1095,4 +1371,8 @@ def setup_modal_queue_route( return web.json_response({"error": str(exc), "node_errors": []}, status=500) _ROUTE_REGISTERED = True - logger.info("Registered Modal queue route at %s", resolved_settings.route_path) + logger.info( + "Registered Modal queue route at %s and analysis route at %s", + resolved_settings.route_path, + analysis_route_path, + ) diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index 01aa874..92fb49b 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -6,6 +6,7 @@ - ComfyUI registry setup is mostly metadata, not code. The repo needs a real `[tool.comfy]` block in `pyproject.toml`, a stable immutable node-pack id, repository URLs, and an optional GitHub publish workflow; choosing a software license should stay explicit instead of being guessed during packaging work. - `comfy node pack` is not git-aware by default; it uses `.gitignore` as its practical packaging filter. For a clean registry artifact, this repo needs an allowlist-style `.gitignore` so untracked local files like logs, caches, virtualenvs, and scratch docs cannot leak into `node.zip`. - Remote per-node phase events are not enough to mirror sampler progress in the local UI. The Modal worker also needs a headless equivalent of ComfyUI's global progress hook so `ProgressBar` updates feed `progress_state`, and the local relay should forward those numeric updates as a separate high-frequency UI event rather than overloading coarse execution-phase messages. +- Queue-time remote-boundary fixes should come from the backend's actual rewrite logic, not from parsing error text in the browser. A small dry-run analysis route can reuse the same prompt-id resolution and upstream auto-expansion rules, then return workflow-node paths for the frontend to mark directly from a context-menu action, including nested subgraph nodes. ## 2026-04-17 diff --git a/tests/test_api_intercept.py b/tests/test_api_intercept.py index c6d3f1d..262b2c9 100644 --- a/tests/test_api_intercept.py +++ b/tests/test_api_intercept.py @@ -582,6 +582,159 @@ def test_rewrite_auto_expands_upstream_non_transportable_dependencies( assert rewritten_prompt["4"]["inputs"]["latent"] == ["1", 0] +def test_analyze_remote_node_selection_returns_nodes_to_mark_and_reasons( + api_intercept_module: Any, + settings_module: Any, +) -> None: + """Dry-run analysis should surface the clicked node plus required upstream nodes.""" + settings = settings_module.ModalSyncSettings( + app_name="app", + auto_deploy=True, + allow_ephemeral_fallback=False, + enable_memory_snapshot=True, + enable_gpu_memory_snapshot=False, + execution_mode="local", + sync_custom_nodes=False, + volume_name="volume", + route_path="/modal/queue_prompt", + marker_property="is_modal_remote", + local_storage_root=Path("/tmp/storage"), + remote_storage_root="/storage", + custom_nodes_archive_name="custom_nodes_bundle.zip", + comfyui_root=None, + custom_nodes_dir=Path("/tmp/custom_nodes"), + ) + fake_nodes_module = type( + "FakeNodesModule", + (), + { + "NODE_CLASS_MAPPINGS": { + "ModelSource": _FakeRemoteModelNode, + "ConditioningSource": _FakeRemoteConditioningNode, + "RemoteConsumer": _FakeRemoteSamplerNode, + }, + "NODE_DISPLAY_NAME_MAPPINGS": {}, + }, + )() + workflow = { + "nodes": [ + {"id": 1, "properties": {"is_modal_remote": False}}, + {"id": 2, "properties": {"is_modal_remote": False}}, + {"id": 3, "properties": {"is_modal_remote": False}}, + ] + } + prompt = { + "1": {"class_type": "ModelSource", "inputs": {}, "_meta": {"title": "Model"}}, + "2": { + "class_type": "ConditioningSource", + "inputs": {}, + "_meta": {"title": "Conditioning"}, + }, + "3": { + "class_type": "RemoteConsumer", + "inputs": {"model": ["1", 0], "conditioning": ["2", 0]}, + "_meta": {"title": "Remote Consumer"}, + }, + } + + analysis = api_intercept_module.analyze_remote_node_selection( + prompt=prompt, + workflow=workflow, + seed_workflow_node_paths=["3"], + settings=settings, + nodes_module=fake_nodes_module, + ) + + assert analysis.requested_node_ids == ["3"] + assert analysis.requested_workflow_node_paths == ["3"] + assert analysis.current_remote_node_ids == [] + assert analysis.current_remote_workflow_node_paths == [] + assert analysis.resolved_remote_node_ids == ["1", "2", "3"] + assert analysis.resolved_workflow_node_paths == ["1", "2", "3"] + assert analysis.added_node_ids == ["1", "2", "3"] + assert analysis.added_workflow_node_paths == ["1", "2", "3"] + assert [(reason.node_id, reason.required_by_node_id) for reason in analysis.reasons] == [ + ("1", "3"), + ("2", "3"), + ] + + +def test_analyze_remote_node_selection_prefers_nested_workflow_paths( + api_intercept_module: Any, + settings_module: Any, +) -> None: + """Nested prompt ids should map back to the specific inner workflow node path.""" + settings = settings_module.ModalSyncSettings( + app_name="app", + auto_deploy=True, + allow_ephemeral_fallback=False, + enable_memory_snapshot=True, + enable_gpu_memory_snapshot=False, + execution_mode="local", + sync_custom_nodes=False, + volume_name="volume", + route_path="/modal/queue_prompt", + marker_property="is_modal_remote", + local_storage_root=Path("/tmp/storage"), + remote_storage_root="/storage", + custom_nodes_archive_name="custom_nodes_bundle.zip", + comfyui_root=None, + custom_nodes_dir=Path("/tmp/custom_nodes"), + ) + fake_nodes_module = type( + "FakeNodesModule", + (), + { + "NODE_CLASS_MAPPINGS": { + "ModelSource": _FakeRemoteModelNode, + "RemoteConsumer": _FakeRemoteSamplerNode, + }, + "NODE_DISPLAY_NAME_MAPPINGS": {}, + }, + )() + workflow = { + "nodes": [ + { + "id": 24, + "properties": {"is_modal_remote": False}, + "subgraph": { + "nodes": [ + {"id": 23, "properties": {"is_modal_remote": True}}, + ] + }, + }, + {"id": 30, "properties": {"is_modal_remote": False}}, + ] + } + prompt = { + "30": {"class_type": "ModelSource", "inputs": {}, "_meta": {"title": "Model"}}, + "24:23": { + "class_type": "RemoteConsumer", + "inputs": {"model": ["30", 0]}, + "_meta": {"title": "Nested Consumer"}, + }, + } + + analysis = api_intercept_module.analyze_remote_node_selection( + prompt=prompt, + workflow=workflow, + seed_workflow_node_paths=["24:23"], + settings=settings, + nodes_module=fake_nodes_module, + ) + + assert analysis.requested_node_ids == ["24:23"] + assert analysis.current_remote_node_ids == ["24:23"] + assert analysis.current_remote_workflow_node_paths == ["24:23"] + assert analysis.resolved_remote_node_ids == ["24:23", "30"] + assert analysis.resolved_workflow_node_paths == ["24:23", "30"] + assert analysis.added_node_ids == ["30"] + assert analysis.added_workflow_node_paths == ["30"] + assert [(reason.node_id, reason.required_by_node_id) for reason in analysis.reasons] == [ + ("30", "24:23"), + ] + + def test_rewrite_rejects_non_transportable_remote_outputs( api_intercept_module: Any, settings_module: Any, diff --git a/tests/test_web_modal_toggle.py b/tests/test_web_modal_toggle.py index 50c9172..8b72c3b 100644 --- a/tests/test_web_modal_toggle.py +++ b/tests/test_web_modal_toggle.py @@ -87,6 +87,32 @@ def test_streamed_modal_node_progress_updates_active_overlay() -> None: assert 'ctx.fillText(`${Math.round(progressRatio * 100)}%`' in source +def test_modal_context_menu_can_expand_required_upstream_nodes() -> None: + """Right-clicking a node should offer a dry-run expansion action backed by the backend.""" + source = _modal_toggle_source() + + assert 'const MODAL_ANALYZE_ROUTE = MODAL_ROUTE.replace(/\\/queue_prompt$/, "/analyze_remote_nodes");' in source + assert "function workflowNodePath(node)" in source + assert "function findNodeByWorkflowPath(workflowPath)" in source + assert "function selectedWorkflowNodePaths(node)" in source + assert "function analyzeAndMarkRequiredRemoteNodes(node)" in source + assert 'api.fetchApi(MODAL_ANALYZE_ROUTE, {' in source + assert '"Modal: Include Required Upstream Nodes"' in source + assert '"Modal: Include Required Upstream Nodes for Selection"' in source + + +def test_modal_context_menu_marks_nodes_across_subgraphs() -> None: + """The UI action should be able to resolve and mark nested workflow-node paths.""" + source = _modal_toggle_source() + + assert "function rootGraph()" in source + assert "function findSomethingInAllSubgraphs(matcher)" in source + assert "function findContainingSubgraphNode(subgraphId)" in source + assert "candidate.subgraph?.id === subgraphId" in source + assert "function markWorkflowNodePathsRemote(workflowNodePaths)" in source + assert "setRemoteFlag(node, true);" in source + + def test_prompt_cleanup_prunes_orphaned_global_status_entries() -> None: """Workflow cleanup should remove stale global badge states once prompt state is gone.""" source = _modal_toggle_source() diff --git a/web/modal_toggle.js b/web/modal_toggle.js index c3b9dde..91c1fcd 100644 --- a/web/modal_toggle.js +++ b/web/modal_toggle.js @@ -3,6 +3,7 @@ import { PromptExecutionError, api } from "../../scripts/api.js"; const REMOTE_PROPERTY = "is_modal_remote"; const MODAL_ROUTE = "/modal/queue_prompt"; +const MODAL_ANALYZE_ROUTE = MODAL_ROUTE.replace(/\/queue_prompt$/, "/analyze_remote_nodes"); const INTERNAL_NODE_PREFIX = "ModalUniversalExecutor"; const IDLE_BORDER_COLOR = "#1d9bf0"; @@ -643,6 +644,229 @@ function extractRemoteNodeIds(workflow) { return remoteNodeIds; } +/** + * Return the root workflow graph, including subgraphs when available. + * @returns {LGraph | null} + */ +function rootGraph() { + return app.rootGraph ?? app.graph?.rootGraph ?? app.graph ?? null; +} + +/** + * Look up a graph-local node id without assuming whether ids are numeric or strings. + * @param {LGraph | null | undefined} graph + * @param {string} id + * @returns {LGraphNode | null} + */ +function getGraphNodeById(graph, id) { + if (!graph || id == null) { + return null; + } + const directMatch = graph.getNodeById?.(id) ?? null; + if (directMatch) { + return directMatch; + } + const numericId = Number(id); + if (Number.isFinite(numericId)) { + return graph.getNodeById?.(numericId) ?? null; + } + return null; +} + +/** + * Search every live subgraph for one matching value. + * @param {(graph: LGraph) => any} matcher + * @returns {any} + */ +function findSomethingInAllSubgraphs(matcher) { + const graph = rootGraph(); + if (!graph) { + return null; + } + + const subgraphs = [graph]; + if (graph.subgraphs?.values) { + subgraphs.push(...graph.subgraphs.values()); + } + for (const subgraph of subgraphs) { + const match = matcher(subgraph); + if (match) { + return match; + } + } + return null; +} + +/** + * Return the workflow node that owns one nested subgraph graph id. + * @param {number | string | undefined} subgraphId + * @returns {LGraphNode | null} + */ +function findContainingSubgraphNode(subgraphId) { + if (subgraphId == null) { + return null; + } + return ( + findSomethingInAllSubgraphs((graph) => + (graph?.nodes ?? []).find( + (candidate) => + typeof candidate?.isSubgraphNode === "function" && + candidate.isSubgraphNode() && + candidate.subgraph?.id === subgraphId, + ), + ) ?? null + ); +} + +/** + * Return one node's composed workflow path, including any subgraph ancestors. + * @param {LGraphNode} node + * @returns {string} + */ +function workflowNodePath(node) { + const pathSegments = [String(node?.id ?? "")]; + let currentGraph = node?.graph ?? null; + const currentRootGraph = rootGraph(); + + while (currentGraph && currentRootGraph && currentGraph !== currentRootGraph) { + const parentNode = findContainingSubgraphNode(currentGraph.id); + if (!parentNode) { + break; + } + pathSegments.unshift(String(parentNode.id)); + currentGraph = parentNode.graph ?? null; + } + + return pathSegments.filter(Boolean).join(":"); +} + +/** + * Resolve a composed workflow path like `24:23` to a live node instance. + * @param {string} workflowPath + * @returns {LGraphNode | null} + */ +function findNodeByWorkflowPath(workflowPath) { + const pathSegments = String(workflowPath) + .split(":") + .filter(Boolean); + if (pathSegments.length === 0) { + return null; + } + + let currentGraph = rootGraph(); + let currentNode = null; + for (const pathSegment of pathSegments) { + currentNode = getGraphNodeById(currentGraph, pathSegment); + if (!currentNode) { + return null; + } + currentGraph = currentNode.subgraph ?? null; + } + return currentNode; +} + +/** + * Return the workflow-node paths that a context-menu action should expand from. + * @param {LGraphNode} node + * @returns {string[]} + */ +function selectedWorkflowNodePaths(node) { + const selectedNodes = Object.values(app.canvas?.selected_nodes ?? {}).filter( + (candidate) => candidate?.graph === node?.graph && isEligibleNode(candidate), + ); + if (selectedNodes.some((candidate) => candidate === node) && selectedNodes.length > 1) { + return selectedNodes.map((candidate) => workflowNodePath(candidate)); + } + return [workflowNodePath(node)]; +} + +/** + * Return the graph snapshot shape ComfyUI already uses for queue submission. + * @returns {Promise<{ output: object, workflow: object }>} + */ +async function serializeCurrentGraphForModal() { + if (typeof app.graphToPrompt !== "function") { + throw new Error("ComfyUI graph serialization is unavailable."); + } + const prompt = await app.graphToPrompt(); + if (!prompt?.output || !prompt?.workflow) { + throw new Error("ComfyUI did not return prompt and workflow data."); + } + return { + output: prompt.output, + workflow: prompt.workflow, + }; +} + +/** + * Show a short frontend notification without taking over the whole UI. + * @param {string} value + */ +function notifyModal(value) { + dispatchSyntheticApiEvent("notification", { + id: `modal-analysis-${Date.now()}`, + value, + }); +} + +/** + * Apply the remote marker to the workflow nodes named by composed workflow paths. + * @param {string[]} workflowNodePaths + * @returns {number} + */ +function markWorkflowNodePathsRemote(workflowNodePaths) { + let appliedCount = 0; + for (const workflowPath of workflowNodePaths) { + const node = findNodeByWorkflowPath(workflowPath); + if (!node) { + console.warn("Unable to find Modal workflow node path in the live graph.", workflowPath); + continue; + } + setRemoteFlag(node, true); + appliedCount += 1; + } + return appliedCount; +} + +/** + * Request required upstream nodes from the backend and mark them remote in the UI. + * @param {LGraphNode} node + */ +async function analyzeAndMarkRequiredRemoteNodes(node) { + const seedNodeIds = selectedWorkflowNodePaths(node); + const graphSnapshot = await serializeCurrentGraphForModal(); + const response = await api.fetchApi(MODAL_ANALYZE_ROUTE, { + method: "POST", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify({ + prompt: graphSnapshot.output, + workflow: graphSnapshot.workflow, + seed_node_ids: seedNodeIds, + }), + }); + if (response.status !== 200) { + throw new PromptExecutionError(await response.json()); + } + + const result = await response.json(); + const resolvedWorkflowNodePaths = result.resolved_workflow_node_paths ?? []; + const addedWorkflowNodePaths = result.added_workflow_node_paths ?? []; + const appliedCount = markWorkflowNodePathsRemote(resolvedWorkflowNodePaths); + app.graph?.setDirtyCanvas(true, true); + + if (addedWorkflowNodePaths.length > 0) { + notifyModal(`Marked ${addedWorkflowNodePaths.length} node${addedWorkflowNodePaths.length === 1 ? "" : "s"} for Modal.`); + return; + } + notifyModal( + appliedCount > 0 + ? "No extra upstream Modal nodes were required." + : "Modal analysis finished, but no matching live nodes were found to update.", + ); +} + /** * Update the node state and redraw the canvas. * @param {LGraphNode} node @@ -799,6 +1023,34 @@ function decorateNode(node) { originalDrawForeground?.apply(this, arguments); drawRemoteNodeDecoration(this, ctx); }; + + const originalGetExtraMenuOptions = node.getExtraMenuOptions; + node.getExtraMenuOptions = function getExtraMenuOptionsWithModalAnalysis(_, options) { + const menuOptions = originalGetExtraMenuOptions?.apply(this, arguments) ?? options ?? []; + const targetOptions = options ?? menuOptions; + if (!Array.isArray(targetOptions)) { + return menuOptions; + } + + const selectedNodePaths = selectedWorkflowNodePaths(this); + const menuItemLabel = + selectedNodePaths.length > 1 + ? "Modal: Include Required Upstream Nodes for Selection" + : "Modal: Include Required Upstream Nodes"; + if (!targetOptions.some((option) => option?.content === menuItemLabel)) { + targetOptions.push(null, { + content: menuItemLabel, + callback: () => { + void analyzeAndMarkRequiredRemoteNodes(this).catch((error) => { + console.error("Modal remote-node analysis failed.", error); + notifyModal(`Modal remote-node analysis failed: ${String(error?.message ?? error)}`); + }); + }, + }); + } + + return menuOptions; + }; } /**