diff --git a/README.md b/README.md index 9283100..dfaece4 100644 --- a/README.md +++ b/README.md @@ -183,6 +183,8 @@ While that is happening, the frontend now shows remote-node state directly on th - green border with light green shading: the remote component is actively executing - red border: queue-time or execution failure detected +During remote execution, the Modal runtime now streams progress events back over the same invocation. Those streamed events identify the currently active remote node, so the local canvas can brighten the specific original node that is running inside the remote component instead of only showing the component as generically active. + Because the Modal queue route can spend noticeable time hashing, syncing, and creating a remote runtime before the prompt is formally queued, the frontend also emits a temporary synthetic running state into ComfyUI's normal queue/execution UI. Those events now mirror ComfyUI's native websocket payload shapes, including the temporary "Waiting for a machine" initialization notification, so the built-in queue indicators stay alive during the preparatory phase instead of looking idle until the backend finally returns the queued prompt id. The extension also now renders its own global activity badge in the ComfyUI UI during Modal setup and remote execution. That badge stays visible even when ComfyUI's stock queue state does not fully reflect the pre-queue setup phase. diff --git a/api_intercept.py b/api_intercept.py index 0cc54da..51059c5 100644 --- a/api_intercept.py +++ b/api_intercept.py @@ -106,6 +106,9 @@ def _emit_modal_status( prompt_id: str | None, node_ids: list[str], component_node_ids_by_representative: dict[str, list[str]] | None = None, + active_node_id: str | None = None, + active_node_class_type: str | None = None, + active_node_role: str | None = None, error_message: str | None = None, ) -> None: """Send a Modal execution status event to the active websocket client.""" @@ -128,6 +131,12 @@ def _emit_modal_status( component_node_ids_by_representative.items() ) ] + if active_node_id is not None: + payload["active_node_id"] = active_node_id + if active_node_class_type is not None: + payload["active_node_class_type"] = active_node_class_type + if active_node_role is not None: + payload["active_node_role"] = active_node_role if error_message is not None: payload["error_message"] = error_message @@ -576,6 +585,8 @@ def _build_component_payload( payload = { "payload_kind": "subgraph", "component_id": component.representative_node_id, + "prompt_id": (extra_data or {}).get("prompt_id"), + "component_node_ids": list(component.node_ids), "subgraph_prompt": component_prompt, "boundary_inputs": [ { @@ -852,6 +863,7 @@ def setup_modal_queue_route( json_data = await request.json() json_data.setdefault("prompt_id", str(uuid.uuid4())) json_data.setdefault("extra_data", {}) + json_data["extra_data"]["prompt_id"] = json_data["prompt_id"] if json_data.get("client_id") is not None: json_data["extra_data"]["client_id"] = json_data["client_id"] extra_pnginfo = ((json_data.get("extra_data") or {}).get("extra_pnginfo") or {}) diff --git a/comfyui_modal_sync_cloud.py b/comfyui_modal_sync_cloud.py index ac4d9ea..2b53316 100644 --- a/comfyui_modal_sync_cloud.py +++ b/comfyui_modal_sync_cloud.py @@ -10,6 +10,7 @@ import json import inspect import logging import os +import queue import sys import tempfile import threading @@ -19,7 +20,7 @@ from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path -from typing import Any, Iterator +from typing import Any, Callable, Iterator _REPO_ROOT = Path(__file__).resolve().parent _REMOTE_REPO_ROOT = Path("/root/comfyui_modal_sync_repo") @@ -89,11 +90,17 @@ class _NullPromptServer: class _TracingPromptServer(_NullPromptServer): """PromptExecutor server stub that records coarse per-node execution timings.""" - def __init__(self, prompt_id: str, prompt: dict[str, Any]) -> None: + def __init__( + self, + prompt_id: str, + prompt: dict[str, Any], + status_callback: Callable[[dict[str, Any]], None] | None = None, + ) -> None: """Initialize timing state for a specific prompt execution.""" super().__init__() self.prompt_id = prompt_id self.prompt = prompt + self._status_callback = status_callback self._active_node_id: str | None = None self._active_node_started_at: float | None = None @@ -139,6 +146,15 @@ class _TracingPromptServer(_NullPromptServer): node_info = self.prompt.get(str(next_node_id), {}) class_type = str(node_info.get("class_type", "")) role = self._classify_node_role(class_type) + if self._status_callback is not None: + self._status_callback( + { + "phase": "executing", + "active_node_id": str(next_node_id), + "active_node_class_type": class_type, + "active_node_role": role, + } + ) self._active_node_id = str(next_node_id) self._active_node_started_at = time.perf_counter() _emit_cloud_info( @@ -157,6 +173,8 @@ class _TracingPromptServer(_NullPromptServer): if event in {"execution_error", "execution_interrupted", "execution_success"}: self._log_node_finish(reason=event) + if self._status_callback is not None: + self._status_callback({"phase": event}) return super().send_sync(event, data, client_id) @@ -769,6 +787,7 @@ def _execute_subgraph_prompt( payload: dict[str, Any], hydrated_inputs: dict[str, Any], custom_nodes_root: Path | None, + status_callback: Callable[[dict[str, Any]], None] | None = None, ) -> tuple[Any, ...]: """Execute a remote component prompt and return its exported outputs.""" component_id = str(payload.get("component_id", "modal-subgraph")) @@ -784,7 +803,11 @@ def _execute_subgraph_prompt( cache_type, cache_args = _prompt_executor_cache_config(execution) with _temporary_node_mapping(None), _patched_folder_paths_absolute_lookup(): - prompt_server = _TracingPromptServer(component_id, prompt) + prompt_server = _TracingPromptServer( + component_id, + prompt, + status_callback=status_callback, + ) with _timed_phase("create_prompt_executor", component=component_id): executor_state = _get_or_create_prompt_executor_state( execution=execution, @@ -840,6 +863,7 @@ def _execute_subgraph_prompt( def execute_subgraph_locally( payload: dict[str, Any], kwargs_payload: bytes | bytearray | str | dict[str, Any], + status_callback: Callable[[dict[str, Any]], None] | None = None, ) -> bytes: """Execute a rewritten remote component in-process and return serialized outputs.""" component_id = str(payload.get("component_id", "modal-subgraph")) @@ -850,12 +874,70 @@ def execute_subgraph_locally( hydrated_inputs = deserialize_node_inputs(kwargs_payload) with _timed_phase("subgraph_worker_roundtrip", component=component_id): with ThreadPoolExecutor(max_workers=1) as executor: - future = executor.submit(_execute_subgraph_prompt, payload, hydrated_inputs, custom_nodes_root) + future = executor.submit( + _execute_subgraph_prompt, + payload, + hydrated_inputs, + custom_nodes_root, + status_callback, + ) outputs = future.result() with _timed_phase("serialize_boundary_outputs", component=component_id): return serialize_node_outputs(outputs) +def _stream_remote_payload_events( + payload: dict[str, Any], + kwargs_payload: bytes | bytearray | str | dict[str, Any], +) -> Iterator[dict[str, Any]]: + """Yield progress and result events for one remote payload execution.""" + event_queue: queue.Queue[tuple[str, Any]] = queue.Queue() + + def publish_status(progress_state: dict[str, Any]) -> None: + """Queue a progress envelope for the remote caller.""" + event_queue.put(("progress", dict(progress_state))) + + def execute_payload() -> None: + """Run the payload in a worker thread and enqueue the terminal outcome.""" + try: + if payload.get("payload_kind") == "subgraph": + outputs = execute_subgraph_locally( + payload, + kwargs_payload, + status_callback=publish_status, + ) + else: + outputs = execute_node_locally(payload, kwargs_payload) + except Exception as exc: # pragma: no cover - exercised through generator consumer tests. + event_queue.put(("error", exc)) + else: + event_queue.put(("result", outputs)) + finally: + event_queue.put(("done", None)) + + worker_thread = threading.Thread( + target=execute_payload, + name=f"modal-stream-{payload.get('component_id', 'payload')}", + daemon=True, + ) + worker_thread.start() + try: + while True: + event_kind, event_payload = event_queue.get() + if event_kind == "progress": + yield {"kind": "progress", **event_payload} + continue + if event_kind == "result": + yield {"kind": "result", "outputs": event_payload} + continue + if event_kind == "error": + raise event_payload + if event_kind == "done": + return + finally: + worker_thread.join(timeout=1.0) + + def _should_ignore_repo_path(path: Path) -> bool: """Return whether a local repo path should be omitted from the Modal image mount.""" parts = set(path.parts) @@ -1028,6 +1110,29 @@ if modal is not None: # pragma: no branch - remote entrypoint configuration. return execute_subgraph_locally(payload, kwargs_payload) return execute_node_locally(payload, kwargs_payload) + @modal.method() + def execute_payload_stream( + self, + payload: dict[str, Any], + kwargs_payload: bytes, + ) -> Iterator[dict[str, Any]]: + """Stream progress envelopes and a final serialized result for one payload.""" + component_id = payload.get("component_id", "single-node") + with _timed_phase( + "remote_engine_execute_payload", + component=component_id, + payload_kind=payload.get("payload_kind"), + ): + if _should_reload_modal_volume(payload): + with _timed_phase("modal_volume_reload", component=component_id): + vol.reload() + else: + _emit_cloud_info( + "Skipping modal_volume_reload for component=%s because no new assets were uploaded for this request.", + component_id, + ) + yield from _stream_remote_payload_events(payload, kwargs_payload) + else: app = None RemoteEngine = None diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index f1474af..1fe0c26 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -38,6 +38,7 @@ - Importing `nodes.py` alone is not enough inside the Modal container. The remote worker also has to run `nodes.init_extra_nodes(...)` once so built-in `comfy_extras` nodes such as `EmptyFlux2LatentImage` exist in `NODE_CLASS_MAPPINGS`, and then run `init_external_custom_nodes()` for each extracted custom-node bundle root it exposes. - A bare `PromptExecutor(_NullPromptServer())` is not equivalent to real ComfyUI startup. The headless worker must pass the same cache configuration that `main.py` uses, or execution eventually crashes when the executor polls caches with `self.cache_args["ram"]`. - If `/modal/queue_prompt` does expensive preparatory work before it returns, the frontend can keep ComfyUI's own queue/execution indicators alive by dispatching temporary synthetic `status`, `execution_start`, and `executing` events, then letting the real websocket events take over once the prompt is actually queued. +- For remote node-level UI feedback, a streamed Modal generator method is cleaner than a side-channel file or volume poller. Yielding `progress` envelopes directly from the remote execution method lets the local client forward active-node updates into ComfyUI's websocket stream in real time without extra storage commits, reloads, or cleanup logic. - The synthetic frontend events have to match ComfyUI's websocket adapter shapes exactly, not just semantically. In practice that means `status` detail must be the inner status payload, `executing` detail should be the display node id string, and queue-list initialization state benefits from a matching `"Waiting for a machine"` notification payload keyed by prompt id. - Even with correctly-shaped synthetic events, ComfyUI's built-in queue UI may still not present an obvious "workflow running" state during long pre-queue Modal setup. The extension should own a separate global activity badge so users have explicit feedback while the remote runtime is still being prepared. - Synced asset paths cannot stay as abstract `/assets/...` markers inside the Modal worker. Remote execution needs those markers materialized to real absolute files under the mounted storage root, and ComfyUI's `folder_paths` helpers need an absolute-path fast path so loader nodes can still resolve them. diff --git a/remote/modal_app.py b/remote/modal_app.py index 7f16812..a38a02e 100644 --- a/remote/modal_app.py +++ b/remote/modal_app.py @@ -477,8 +477,125 @@ def _load_modal_cloud_module() -> Any: return cloud_module -def _lookup_deployed_modal_method(payload: dict[str, Any]) -> Any: - """Look up the deployed Modal method used to execute the remote runtime.""" +def _lookup_local_prompt_server() -> Any | None: + """Return the live local ComfyUI PromptServer instance when available.""" + try: + import server + except ModuleNotFoundError: + return None + + return getattr(server.PromptServer, "instance", None) + + +def _emit_local_modal_status( + *, + prompt_id: str | None, + client_id: str | None, + phase: str, + node_ids: list[str], + active_node_id: str | None = None, + active_node_class_type: str | None = None, + active_node_role: str | None = None, + error_message: str | None = None, +) -> None: + """Forward remote execution progress into the local ComfyUI websocket stream.""" + if client_id is None: + return + + prompt_server = _lookup_local_prompt_server() + if prompt_server is None: + return + + payload: dict[str, Any] = { + "phase": phase, + "prompt_id": prompt_id, + "node_ids": list(node_ids), + } + if active_node_id is not None: + payload["active_node_id"] = active_node_id + if active_node_class_type is not None: + payload["active_node_class_type"] = active_node_class_type + if active_node_role is not None: + payload["active_node_role"] = active_node_role + if error_message is not None: + payload["error_message"] = error_message + prompt_server.send_sync("modal_status", payload, client_id) + + +def _should_stream_remote_progress(payload: dict[str, Any]) -> bool: + """Return whether the local client has enough context to mirror remote node progress.""" + extra_data = payload.get("extra_data") or {} + return ( + payload.get("payload_kind") == "subgraph" + and isinstance(payload.get("prompt_id"), str) + and bool(payload.get("prompt_id")) + and isinstance(extra_data.get("client_id"), str) + and bool(extra_data.get("client_id")) + and isinstance(payload.get("component_node_ids"), list) + and len(payload.get("component_node_ids")) > 0 + ) + + +def _consume_remote_payload_stream( + payload: dict[str, Any], + stream_events: Iterator[dict[str, Any]], +) -> bytes: + """Forward remote progress events into the local UI and return the final payload bytes.""" + prompt_id = str(payload.get("prompt_id")) if payload.get("prompt_id") is not None else None + extra_data = payload.get("extra_data") or {} + client_id = str(extra_data.get("client_id")) if extra_data.get("client_id") is not None else None + node_ids = [str(node_id) for node_id in payload.get("component_node_ids", [])] + result_payload: bytes | bytearray | None = None + + for stream_event in stream_events: + event_kind = str(stream_event.get("kind", "")) + if event_kind == "progress": + _emit_local_modal_status( + prompt_id=prompt_id, + client_id=client_id, + phase=str(stream_event.get("phase", "executing")), + node_ids=node_ids, + active_node_id=( + str(stream_event["active_node_id"]) + if stream_event.get("active_node_id") is not None + else None + ), + active_node_class_type=( + str(stream_event["active_node_class_type"]) + if stream_event.get("active_node_class_type") is not None + else None + ), + active_node_role=( + str(stream_event["active_node_role"]) + if stream_event.get("active_node_role") is not None + else None + ), + ) + continue + if event_kind == "result": + candidate_outputs = stream_event.get("outputs") + if not isinstance(candidate_outputs, (bytes, bytearray)): + raise ModalRemoteInvocationError( + "Modal streamed payload result did not include serialized output bytes." + ) + result_payload = candidate_outputs + continue + + logger.debug( + "Ignoring unexpected streamed Modal event kind=%s for component=%s.", + event_kind, + payload.get("component_id"), + ) + + if result_payload is None: + raise ModalRemoteInvocationError( + f"Modal streamed payload for component={payload.get('component_id')!r} did not yield a final result." + ) + return bytes(result_payload) + + +def _lookup_deployed_remote_engine(payload: dict[str, Any]) -> Any: + """Look up the deployed Modal runtime class instance.""" if modal is None: raise ModalRemoteInvocationError("Modal SDK is unavailable.") @@ -490,8 +607,7 @@ def _lookup_deployed_modal_method(payload: dict[str, Any]) -> Any: payload.get("component_id"), ) remote_cls = modal.Cls.from_name(settings.app_name, "RemoteEngine") - remote_engine = remote_cls() - return remote_engine.execute_payload + return remote_cls() def _modal_environment_name() -> str | None: @@ -564,24 +680,36 @@ def _invoke_modal_payload_blocking(payload: dict[str, Any], kwargs_payload: byte settings = get_settings() if lookup_error_types: try: - remote_method = _lookup_deployed_modal_method(payload) + remote_engine = _lookup_deployed_remote_engine(payload) logger.info( "Using deployed Modal app %s for component %s.", settings.app_name, payload.get("component_id"), ) - return remote_method.remote(payload, kwargs_payload) + stream_method = getattr(remote_engine, "execute_payload_stream", None) + if _should_stream_remote_progress(payload) and hasattr(stream_method, "remote_gen"): + return _consume_remote_payload_stream( + payload, + stream_method.remote_gen(payload, kwargs_payload), + ) + return remote_engine.execute_payload.remote(payload, kwargs_payload) except lookup_error_types as exc: if settings.auto_deploy: _auto_deploy_modal_app(payload, exc) try: - remote_method = _lookup_deployed_modal_method(payload) + remote_engine = _lookup_deployed_remote_engine(payload) logger.info( "Using auto-deployed Modal app %s for component %s.", settings.app_name, payload.get("component_id"), ) - return remote_method.remote(payload, kwargs_payload) + stream_method = getattr(remote_engine, "execute_payload_stream", None) + if _should_stream_remote_progress(payload) and hasattr(stream_method, "remote_gen"): + return _consume_remote_payload_stream( + payload, + stream_method.remote_gen(payload, kwargs_payload), + ) + return remote_engine.execute_payload.remote(payload, kwargs_payload) except lookup_error_types as retry_exc: exc = retry_exc if not settings.allow_ephemeral_fallback: @@ -598,13 +726,19 @@ def _invoke_modal_payload_blocking(payload: dict[str, Any], kwargs_payload: byte exc, ) else: - remote_method = _lookup_deployed_modal_method(payload) + remote_engine = _lookup_deployed_remote_engine(payload) logger.info( "Using deployed Modal app %s for component %s.", settings.app_name, payload.get("component_id"), ) - return remote_method.remote(payload, kwargs_payload) + stream_method = getattr(remote_engine, "execute_payload_stream", None) + if _should_stream_remote_progress(payload) and hasattr(stream_method, "remote_gen"): + return _consume_remote_payload_stream( + payload, + stream_method.remote_gen(payload, kwargs_payload), + ) + return remote_engine.execute_payload.remote(payload, kwargs_payload) if "app" not in globals() or "RemoteEngine" not in globals(): logger.debug("Local module Modal runtime objects are unavailable; loading stable cloud entry module.") @@ -623,7 +757,14 @@ def _invoke_modal_payload_blocking(payload: dict[str, Any], kwargs_payload: byte run_context = cloud_app.run() if hasattr(cloud_app, "run") else nullcontext() with run_context: remote_engine = cloud_remote_engine() - result = remote_engine.execute_payload.remote(payload, kwargs_payload) + stream_method = getattr(remote_engine, "execute_payload_stream", None) + if _should_stream_remote_progress(payload) and hasattr(stream_method, "remote_gen"): + result = _consume_remote_payload_stream( + payload, + stream_method.remote_gen(payload, kwargs_payload), + ) + else: + result = remote_engine.execute_payload.remote(payload, kwargs_payload) logger.info( "Ephemeral Modal app.run() invocation completed for component %s.", payload.get("component_id"), diff --git a/tests/test_api_intercept.py b/tests/test_api_intercept.py index c82a871..9a2fb17 100644 --- a/tests/test_api_intercept.py +++ b/tests/test_api_intercept.py @@ -129,6 +129,8 @@ def test_rewrite_groups_connected_remote_nodes_into_single_proxy( payload = rewritten_node["inputs"]["original_node_data"] assert rewritten_node["class_type"].startswith("ModalUniversalExecutor_") assert payload["payload_kind"] == "subgraph" + assert payload["prompt_id"] is None + assert payload["component_node_ids"] == ["1", "2"] assert payload["subgraph_prompt"]["1"]["inputs"]["model_name"].startswith("/assets/") assert payload["execute_node_ids"] == ["2"] assert payload["requires_volume_reload"] is True @@ -416,20 +418,26 @@ def test_emit_modal_status_targets_prompt_client( prompt_server = FakePromptServer() api_intercept_module._emit_modal_status( prompt_server=prompt_server, - phase="setup", + phase="executing", client_id="client-1", prompt_id="prompt-1", node_ids=["4", "5"], component_node_ids_by_representative={"4": ["4", "5"]}, + active_node_id="5", + active_node_class_type="KSampler", + active_node_role="sampling", ) assert prompt_server.messages == [ ( "modal_status", { - "phase": "setup", + "phase": "executing", "prompt_id": "prompt-1", "node_ids": ["4", "5"], + "active_node_id": "5", + "active_node_class_type": "KSampler", + "active_node_role": "sampling", "components": [ { "representative_node_id": "4", diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index 6a503a6..5d55c9c 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -220,6 +220,60 @@ def test_modal_cloud_traces_remote_node_execution_spans( assert "Remote node 2 class_type=KSampler role=sampling finished in " in captured.out +def test_modal_cloud_streams_progress_and_result_events( + modal_cloud_module: Any, + monkeypatch: Any, +) -> None: + """The cloud runtime should stream progress envelopes before the final result.""" + progress_callbacks: list[dict[str, Any]] = [] + + def fake_execute_subgraph_locally( + payload: dict[str, Any], + kwargs_payload: bytes, + status_callback: Any = None, + ) -> bytes: + if status_callback is not None: + status_callback( + { + "phase": "executing", + "active_node_id": "7", + "active_node_class_type": "UNETLoader", + "active_node_role": "model_load", + } + ) + status_callback({"phase": "execution_success"}) + progress_callbacks.append({"component_id": payload["component_id"], "kwargs": kwargs_payload}) + return b"serialized-outputs" + + monkeypatch.setattr(modal_cloud_module, "execute_subgraph_locally", fake_execute_subgraph_locally) + + events = list( + modal_cloud_module._stream_remote_payload_events( + {"payload_kind": "subgraph", "component_id": "component-1"}, + b"{}", + ) + ) + + assert progress_callbacks == [{"component_id": "component-1", "kwargs": b"{}"}] + assert events == [ + { + "kind": "progress", + "phase": "executing", + "active_node_id": "7", + "active_node_class_type": "UNETLoader", + "active_node_role": "model_load", + }, + { + "kind": "progress", + "phase": "execution_success", + }, + { + "kind": "result", + "outputs": b"serialized-outputs", + }, + ] + + def test_modal_cloud_only_reloads_volume_for_requests_with_new_uploads( modal_cloud_module: Any, ) -> None: @@ -553,6 +607,81 @@ def test_remote_modal_auto_deploys_missing_app_by_default( assert deploy_calls == [("comfy-modal-sync", None)] +def test_remote_modal_consumes_streamed_progress_and_result( + remote_modal_app_module: Any, + monkeypatch: Any, +) -> None: + """The local Modal client should forward streamed progress events into the UI websocket.""" + + class FakePromptServer: + """Capture websocket events emitted by streamed remote progress.""" + + def __init__(self) -> None: + """Initialize the event sink.""" + self.messages: list[tuple[str, dict[str, Any], str | None]] = [] + + def send_sync(self, event: str, data: dict[str, Any], sid: str | None) -> None: + """Record one emitted websocket message.""" + self.messages.append((event, data, sid)) + + prompt_server = FakePromptServer() + monkeypatch.setattr(remote_modal_app_module, "_lookup_local_prompt_server", lambda: prompt_server) + + payload = { + "prompt_id": "prompt-1", + "component_id": "component-1", + "component_node_ids": ["7", "8"], + "extra_data": {"client_id": "client-1"}, + } + result = remote_modal_app_module._consume_remote_payload_stream( + payload, + iter( + [ + { + "kind": "progress", + "phase": "executing", + "active_node_id": "7", + "active_node_class_type": "UNETLoader", + "active_node_role": "model_load", + }, + { + "kind": "progress", + "phase": "execution_success", + }, + { + "kind": "result", + "outputs": b"serialized-outputs", + }, + ] + ), + ) + + assert result == b"serialized-outputs" + assert prompt_server.messages == [ + ( + "modal_status", + { + "phase": "executing", + "prompt_id": "prompt-1", + "node_ids": ["7", "8"], + "active_node_id": "7", + "active_node_class_type": "UNETLoader", + "active_node_role": "model_load", + }, + "client-1", + ), + ( + "modal_status", + { + "phase": "execution_success", + "prompt_id": "prompt-1", + "node_ids": ["7", "8"], + }, + "client-1", + ), + ] + + def test_remote_modal_requires_manual_deploy_when_auto_deploy_disabled( remote_modal_app_module: Any, monkeypatch: Any, diff --git a/tests/test_web_modal_toggle.py b/tests/test_web_modal_toggle.py index 7c8df84..e45c937 100644 --- a/tests/test_web_modal_toggle.py +++ b/tests/test_web_modal_toggle.py @@ -37,3 +37,13 @@ def test_global_modal_status_badge_is_installed() -> None: assert "Modal setup running for" in source assert "Modal workflow running on" in source assert "installGlobalStatusStyles()" in source + + +def test_remote_modal_status_tracks_active_node_ids() -> None: + """The frontend should track and highlight the currently active remote node.""" + source = _modal_toggle_source() + + assert "activeNodeId: null" in source + assert "function setPromptActiveNode(promptId, activeNodeId)" in source + assert "detail.active_node_id" in source + assert "isActiveRemoteNode" in source diff --git a/web/modal_toggle.js b/web/modal_toggle.js index fc305b6..d4f3a40 100644 --- a/web/modal_toggle.js +++ b/web/modal_toggle.js @@ -245,6 +245,7 @@ function ensurePromptState(promptId) { startedAt: nowMs(), remoteNodeIds: [], componentsByRepresentative: new Map(), + activeNodeId: null, }); } return modalPromptStates.get(promptId); @@ -372,6 +373,7 @@ function registerPromptComponents(promptId, remoteNodeIds, components) { const promptState = ensurePromptState(promptId); promptState.remoteNodeIds = [...remoteNodeIds]; promptState.componentsByRepresentative.clear(); + promptState.activeNodeId = null; for (const component of components) { promptState.componentsByRepresentative.set( String(component.representative_node_id), @@ -380,6 +382,16 @@ function registerPromptComponents(promptId, remoteNodeIds, components) { } } +/** + * Record the currently active remote node inside one prompt. + * @param {string} promptId + * @param {string | null} activeNodeId + */ +function setPromptActiveNode(promptId, activeNodeId) { + const promptState = ensurePromptState(promptId); + promptState.activeNodeId = activeNodeId ? String(activeNodeId) : null; +} + /** * Return remote component node ids for an executing proxy node event. * @param {string} promptId @@ -426,7 +438,15 @@ function setRemoteFlag(node, value) { * @returns {{ phase: string, promptId: string } | null} */ function getRemoteVisualState(node) { - return modalNodeStates.get(nodeId(node)) ?? null; + const state = modalNodeStates.get(nodeId(node)) ?? null; + if (!state?.promptId) { + return state; + } + const promptState = modalPromptStates.get(state.promptId); + return { + ...state, + isActiveRemoteNode: promptState?.activeNodeId === nodeId(node), + }; } /** @@ -460,6 +480,11 @@ function drawRemoteNodeDecoration(node, ctx) { borderColor = ACTIVE_BORDER_COLOR; shadowColor = "rgba(34, 197, 94, 0.35)"; fillColor = `rgba(134, 239, 172, ${0.12 + pulse * 0.08})`; + if (state.isActiveRemoteNode) { + borderColor = "#4ade80"; + shadowColor = `rgba(74, 222, 128, ${0.38 + pulse * 0.28})`; + fillColor = `rgba(187, 247, 208, ${0.24 + pulse * 0.12})`; + } } else if (state?.phase === STATE_COMPLETE) { borderColor = ACTIVE_BORDER_COLOR; shadowColor = "rgba(34, 197, 94, 0.28)"; @@ -543,10 +568,18 @@ function handleModalStatus(event) { if (detail.phase === STATE_SETUP) { beginSyntheticExecutionUi(promptId, (detail.node_ids ?? []).map((value) => String(value))); setGlobalStatusPhase(promptId, STATE_SETUP, (detail.node_ids ?? []).length); + setPromptActiveNode(promptId, null); } else if (detail.phase === STATE_ERROR) { endSyntheticExecutionUi(promptId, true); setGlobalStatusPhase(promptId, STATE_ERROR, (detail.node_ids ?? []).length); setTimeout(() => clearGlobalStatusPhase(promptId), ERROR_CLEAR_DELAY_MS); + setPromptActiveNode(promptId, null); + } else if (detail.phase === STATE_EXECUTING) { + setGlobalStatusPhase(promptId, STATE_EXECUTING, (detail.node_ids ?? []).length); + setPromptActiveNode( + promptId, + detail.active_node_id != null ? String(detail.active_node_id) : null, + ); } const nodeIds = (detail.node_ids ?? []).map((value) => String(value)); @@ -560,6 +593,11 @@ function handleModalStatus(event) { return; } + if (detail.phase === STATE_EXECUTING) { + setNodesPhase(nodeIds, STATE_EXECUTING, promptId); + return; + } + if (detail.phase === STATE_ERROR) { setNodesPhase(nodeIds, STATE_ERROR, promptId, detail.error_message); } @@ -587,6 +625,9 @@ function handleExecutionPhase(event, phase) { } else if (phase === STATE_ERROR) { setGlobalStatusPhase(promptId, STATE_ERROR, componentNodeIds.length); setTimeout(() => clearGlobalStatusPhase(promptId), ERROR_CLEAR_DELAY_MS); + setPromptActiveNode(promptId, null); + } else { + setPromptActiveNode(promptId, null); } setNodesPhase(componentNodeIds, phase, promptId, detail.exception_message); }