feat: stream remote node progress to ui

This commit is contained in:
Ken Simpson
2026-04-17 19:23:40 -07:00
parent 6239e9be5a
commit ced705a25f
9 changed files with 467 additions and 18 deletions
+2
View File
@@ -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.
+12
View File
@@ -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 {})
+109 -4
View File
@@ -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", "<unknown>"))
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
+1
View File
@@ -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.
+152 -11
View File
@@ -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"),
+10 -2
View File
@@ -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",
+129
View File
@@ -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,
+10
View File
@@ -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
+42 -1
View File
@@ -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);
}