ttulttul-ComfyUI-Modal/tests/test_modal_executor_node.py

"""Tests for dynamic Modal proxy nodes and local execution fallback."""

from __future__ import annotations

import asyncio
import copy
from concurrent.futures import Future, ThreadPoolExecutor
from dataclasses import replace
from datetime import datetime, timedelta, timezone
from decimal import Decimal
import hashlib
import importlib.util
import json
import pickle
import subprocess
import sys
import threading
import time
import types
from contextlib import nullcontext
from io import BytesIO, StringIO
from pathlib import Path
import logging
from typing import Any, Iterator

import pytest

DEFAULT_TEST_DEPLOYMENT_APP_NAME = "comfy-modal-sync-gpu-rtx-pro-6000"


class _FakeOriginalNode:
    """Simple fake legacy node for proxy signature mirroring."""

    RETURN_TYPES = ("IMAGE", "INT")
    RETURN_NAMES = ("image", "count")
    OUTPUT_IS_LIST = (False, False)
    FUNCTION = "run"

    def run(self, **kwargs: Any) -> tuple[Any, ...]:
        """Return a tuple that exposes the inputs for verification."""
        return (kwargs["value"], 1)


class _CloneableCacheValue:
    """Simple cloneable object used for loader cache tests."""

    def __init__(self, value: str) -> None:
        """Store an identifying value for later clone assertions."""
        self.value = value

    def clone(self) -> "_CloneableCacheValue":
        """Return a fresh object carrying the same value."""
        return _CloneableCacheValue(self.value)


class _FakeModelValue:
    """Simple cloneable stand-in for a non-transportable MODEL output."""

    def __init__(self, value: str) -> None:
        """Store an identifying value for later assertions."""
        self.value = value

    def clone(self) -> "_FakeModelValue":
        """Return a fresh model value carrying the same identifier."""
        return _FakeModelValue(self.value)


class _FakeModelLoaderNode:
    """Fake self-contained loader node that returns one MODEL-like object."""

    RETURN_TYPES = ("MODEL",)
    RETURN_NAMES = ("model",)
    OUTPUT_IS_LIST = (False,)
    FUNCTION = "load_checkpoint"

    def load_checkpoint(self, ckpt_name: str) -> tuple[_FakeModelValue]:
        """Return a deterministic model value derived from the checkpoint name."""
        return (_FakeModelValue(f"model::{ckpt_name}"),)


class _PromptMetadataSerializationNode:
    """Output node double that JSON-serializes ComfyUI's hidden PROMPT input."""

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("metadata",)
    OUTPUT_IS_LIST = (False,)
    OUTPUT_NODE = True
    FUNCTION = "run"

    @classmethod
    def INPUT_TYPES(cls) -> dict[str, dict[str, Any]]:
        """Declare one hydrated tensor input and the hidden prompt metadata input."""
        return {
            "required": {"image": ("IMAGE",)},
            "hidden": {"prompt": "PROMPT"},
        }

    def run(self, image: Any, prompt: dict[str, Any]) -> tuple[str]:
        """Serialize the prompt exactly as metadata-writing output nodes do."""
        del image
        return (json.dumps(prompt, sort_keys=True),)


class _FakeClipLoaderNode:
    """Fake self-contained loader node that returns one CLIP-like object."""

    RETURN_TYPES = ("CLIP",)
    RETURN_NAMES = ("clip",)
    OUTPUT_IS_LIST = (False,)
    FUNCTION = "load_clip"

    def load_clip(
        self,
        clip_name: str,
        type: str = "stable_diffusion",
        device: str = "default",
    ) -> tuple[str]:
        """Return a deterministic CLIP value derived from the loader inputs."""
        return (f"clip::{clip_name}:{type}:{device}",)


class _FakeSessionValueNode:
    """Fake node that produces one remote-only STRING value."""

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    OUTPUT_IS_LIST = (False,)
    FUNCTION = "execute"

    def execute(self) -> tuple[str]:
        """Return a deterministic value that can be stored in session state."""
        return ("shared-session-value",)


class _FakeSessionEchoNode:
    """Fake node that echoes a STRING input back to the caller."""

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    OUTPUT_IS_LIST = (False,)
    FUNCTION = "execute"

    def execute(self, text: str) -> tuple[str]:
        """Return the supplied input unchanged for session-ref resolution tests."""
        return (text,)


class _FakeRewriteRemoteModelNode:
    """Fake rewrite-time node that produces a non-transportable MODEL output."""

    RETURN_TYPES = ("MODEL",)
    RETURN_NAMES = ("model",)
    OUTPUT_IS_LIST = (False,)


class _FakeRewriteRemoteSamplerNode:
    """Fake rewrite-time node that produces a transportable LATENT output."""

    RETURN_TYPES = ("LATENT",)
    RETURN_NAMES = ("latent",)
    OUTPUT_IS_LIST = (False,)


class _FakeRewriteLatentSourceNode:
    """Fake local source used to feed LATENT values into remote proxies."""

    RETURN_TYPES = ("LATENT",)
    RETURN_NAMES = ("latent",)
    OUTPUT_IS_LIST = (False,)


class _FakeRewriteModalMapInputNode:
    """Fake rewrite-time Modal map marker node."""

    RETURN_TYPES = ("*",)
    RETURN_NAMES = ("value",)
    OUTPUT_IS_LIST = (False,)


class _FakeRewriteLocalSinkNode:
    """Fake local sink used to model downstream local work in rewrite tests."""

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    OUTPUT_IS_LIST = (False,)


class _FakeRewriteLocalFeedbackNode:
    """Fake local node that turns a remote latent into a transportable remote input."""

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    OUTPUT_IS_LIST = (False,)


class _FakeRewriteRemoteImageNode:
    """Fake remote node that produces one locally previewable image."""

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    OUTPUT_IS_LIST = (False,)


class _FakeRewriteRemoteTextNode:
    """Fake remote node that consumes an image and produces text."""

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    OUTPUT_IS_LIST = (False,)


class _FakeRewritePreviewImageNode:
    """Fake local PreviewImage output node."""

    RETURN_TYPES: tuple[str, ...] = ()
    RETURN_NAMES: tuple[str, ...] = ()
    OUTPUT_IS_LIST: tuple[bool, ...] = ()
    OUTPUT_NODE = True


def _current_remote_runtime_payload(remote_modal_app_module: Any) -> dict[str, Any]:
    """Return version metadata for a compatible remote-engine test double."""
    return {
        "protocol_version": remote_modal_app_module._REMOTE_APP_PROTOCOL_VERSION,
        "runtime_fingerprint": remote_modal_app_module._expected_remote_runtime_fingerprint(),
    }


def test_modal_proxy_waits_for_active_prompt_before_starting_next_prompt(
    modal_executor_module: Any,
) -> None:
    """A Modal proxy from a later prompt should wait until the active prompt's remote work drains."""
    fake_nodes_module = types.SimpleNamespace(
        NODE_CLASS_MAPPINGS={},
        NODE_DISPLAY_NAME_MAPPINGS={},
    )
    proxy_node_id = modal_executor_module.ensure_modal_component_proxy_node_registered(
        output_types=("STRING",),
        output_names=("value",),
        output_is_list=(False,),
        nodes_module=fake_nodes_module,
        is_output_node=False,
    )
    proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id]
    first_started = asyncio.Event()
    second_started = asyncio.Event()
    release_first = asyncio.Event()
    observed_events: list[str] = []

    class FakeClient:
        """Fake async remote client that records when each prompt reaches execution."""

        async def execute_payload_async(
            self,
            payload: dict[str, Any],
            kwargs: dict[str, Any],
        ) -> tuple[str]:
            """Block the first prompt so the second prompt has to wait at the gate."""
            component_id = str(payload["component_id"])
            observed_events.append(f"start:{component_id}")
            if component_id == "component-1":
                first_started.set()
                await release_first.wait()
            if component_id == "component-2":
                second_started.set()
            observed_events.append(f"finish:{component_id}")
            return (component_id,)

    async def run_scenario() -> tuple[Any, Any]:
        """Run two prompts concurrently and verify prompt-level serialization."""
        first_task = asyncio.create_task(
            proxy_class.execute(
                original_node_data={
                    "payload_kind": "subgraph",
                    "prompt_id": "prompt-1",
                    "component_id": "component-1",
                },
                unique_id="component-1",
            )
        )
        await first_started.wait()
        second_task = asyncio.create_task(
            proxy_class.execute(
                original_node_data={
                    "payload_kind": "subgraph",
                    "prompt_id": "prompt-2",
                    "component_id": "component-2",
                },
                unique_id="component-2",
            )
        )
        await asyncio.sleep(0.05)
        assert not second_started.is_set()
        release_first.set()
        return await asyncio.gather(first_task, second_task)

    modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient())
    try:
        first_output, second_output = asyncio.run(run_scenario())
    finally:
        modal_executor_module.set_remote_executor_client_factory(None)

    assert first_output.result == ("component-1",)
    assert second_output.result == ("component-2",)
    assert observed_events == [
        "start:component-1",
        "finish:component-1",
        "start:component-2",
        "finish:component-2",
    ]


def test_modal_proxy_allows_same_prompt_components_to_overlap(
    modal_executor_module: Any,
) -> None:
    """Modal proxies from the same prompt should be allowed to execute concurrently."""
    fake_nodes_module = types.SimpleNamespace(
        NODE_CLASS_MAPPINGS={},
        NODE_DISPLAY_NAME_MAPPINGS={},
    )
    proxy_node_id = modal_executor_module.ensure_modal_component_proxy_node_registered(
        output_types=("STRING",),
        output_names=("value",),
        output_is_list=(False,),
        nodes_module=fake_nodes_module,
        is_output_node=False,
    )
    proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id]
    started_components: set[str] = set()
    both_started = asyncio.Event()
    release_components = asyncio.Event()
    observed_events: list[str] = []

    class FakeClient:
        """Fake async remote client that waits until both same-prompt components start."""

        async def execute_payload_async(
            self,
            payload: dict[str, Any],
            kwargs: dict[str, Any],
        ) -> tuple[str]:
            """Record concurrent same-prompt starts before returning."""
            component_id = str(payload["component_id"])
            observed_events.append(f"start:{component_id}")
            started_components.add(component_id)
            if started_components == {"component-1", "component-2"}:
                both_started.set()
            await release_components.wait()
            observed_events.append(f"finish:{component_id}")
            return (component_id,)

    async def run_scenario() -> tuple[Any, Any, Any]:
        """Run two same-prompt proxy components concurrently."""
        first_task = asyncio.create_task(
            proxy_class.execute(
                original_node_data={
                    "payload_kind": "subgraph",
                    "prompt_id": "prompt-1",
                    "component_id": "component-1",
                },
                unique_id="component-1",
            )
        )
        second_task = asyncio.create_task(
            proxy_class.execute(
                original_node_data={
                    "payload_kind": "subgraph",
                    "prompt_id": "prompt-1",
                    "component_id": "component-2",
                },
                unique_id="component-2",
            )
        )
        await asyncio.wait_for(both_started.wait(), timeout=1.0)
        release_components.set()
        return await asyncio.gather(first_task, second_task)

    modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient())
    try:
        first_output, second_output = asyncio.run(run_scenario())
    finally:
        modal_executor_module.set_remote_executor_client_factory(None)

    assert first_output.result == ("component-1",)
    assert second_output.result == ("component-2",)
    assert observed_events[:2] == ["start:component-1", "start:component-2"]


def test_modal_proxy_cancellation_while_waiting_does_not_leak_workflow_gate(
    modal_executor_module: Any,
) -> None:
    """Cancelling a proxy before it gets the prompt gate should not block later prompts."""
    fake_nodes_module = types.SimpleNamespace(
        NODE_CLASS_MAPPINGS={},
        NODE_DISPLAY_NAME_MAPPINGS={},
    )
    proxy_node_id = modal_executor_module.ensure_modal_component_proxy_node_registered(
        output_types=("STRING",),
        output_names=("value",),
        output_is_list=(False,),
        nodes_module=fake_nodes_module,
        is_output_node=False,
    )
    proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id]
    first_started = asyncio.Event()
    release_first = asyncio.Event()
    observed_events: list[str] = []

    class FakeClient:
        """Fake async remote client that lets one prompt hold the workflow gate."""

        async def execute_payload_async(
            self,
            payload: dict[str, Any],
            kwargs: dict[str, Any],
        ) -> tuple[str]:
            """Record remote dispatch and optionally block the first prompt."""
            component_id = str(payload["component_id"])
            observed_events.append(f"start:{component_id}")
            if component_id == "component-1":
                first_started.set()
                await release_first.wait()
            observed_events.append(f"finish:{component_id}")
            return (component_id,)

    async def run_scenario() -> tuple[Any, Any]:
        """Cancel a waiting prompt, then prove the next prompt can still dispatch."""
        first_task = asyncio.create_task(
            proxy_class.execute(
                original_node_data={
                    "payload_kind": "subgraph",
                    "prompt_id": "prompt-1",
                    "component_id": "component-1",
                },
                unique_id="component-1",
            )
        )
        await first_started.wait()
        waiting_task = asyncio.create_task(
            proxy_class.execute(
                original_node_data={
                    "payload_kind": "subgraph",
                    "prompt_id": "prompt-2",
                    "component_id": "component-2",
                },
                unique_id="component-2",
            )
        )
        await asyncio.sleep(0.05)
        waiting_task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await waiting_task
        release_first.set()
        first_output = await first_task
        third_output = await asyncio.wait_for(
            proxy_class.execute(
                original_node_data={
                    "payload_kind": "subgraph",
                    "prompt_id": "prompt-3",
                    "component_id": "component-3",
                },
                unique_id="component-3",
            ),
            timeout=1.0,
        )
        return first_output, third_output

    modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient())
    try:
        first_output, third_output = asyncio.run(run_scenario())
    finally:
        modal_executor_module.set_remote_executor_client_factory(None)

    assert first_output.result == ("component-1",)
    assert third_output.result == ("component-3",)
    assert observed_events == [
        "start:component-1",
        "finish:component-1",
        "start:component-3",
        "finish:component-3",
    ]


def test_abandoned_modal_prompt_gate_unblocks_next_prompt(
    modal_executor_module: Any,
) -> None:
    """Cancelling an active Modal prompt should not leave later prompts stuck behind its gate."""
    fake_nodes_module = types.SimpleNamespace(
        NODE_CLASS_MAPPINGS={},
        NODE_DISPLAY_NAME_MAPPINGS={},
    )
    proxy_node_id = modal_executor_module.ensure_modal_component_proxy_node_registered(
        output_types=("STRING",),
        output_names=("value",),
        output_is_list=(False,),
        nodes_module=fake_nodes_module,
        is_output_node=False,
    )
    proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id]
    started_first_components: set[str] = set()
    first_components_started = asyncio.Event()
    second_started = asyncio.Event()
    release_first_components = asyncio.Event()
    observed_events: list[str] = []

    class FakeClient:
        """Fake async remote client that lets cancelled prompt components stay stuck."""

        async def execute_payload_async(
            self,
            payload: dict[str, Any],
            kwargs: dict[str, Any],
        ) -> tuple[str]:
            """Record remote dispatch while prompt one waits for cleanup."""
            del kwargs
            component_id = str(payload["component_id"])
            observed_events.append(f"start:{component_id}")
            if component_id in {"component-1a", "component-1b"}:
                started_first_components.add(component_id)
                if started_first_components == {"component-1a", "component-1b"}:
                    first_components_started.set()
                await release_first_components.wait()
            if component_id == "component-2":
                second_started.set()
            observed_events.append(f"finish:{component_id}")
            return (component_id,)

    async def run_scenario() -> tuple[Any, Any]:
        """Abandon prompt one, then prove prompt two can start before prompt one exits."""
        first_task_a = asyncio.create_task(
            proxy_class.execute(
                original_node_data={
                    "payload_kind": "subgraph",
                    "prompt_id": "prompt-1",
                    "component_id": "component-1a",
                },
                unique_id="component-1a",
            )
        )
        first_task_b = asyncio.create_task(
            proxy_class.execute(
                original_node_data={
                    "payload_kind": "subgraph",
                    "prompt_id": "prompt-1",
                    "component_id": "component-1b",
                },
                unique_id="component-1b",
            )
        )
        await first_components_started.wait()

        second_task = asyncio.create_task(
            proxy_class.execute(
                original_node_data={
                    "payload_kind": "subgraph",
                    "prompt_id": "prompt-2",
                    "component_id": "component-2",
                },
                unique_id="component-2",
            )
        )
        await asyncio.sleep(0.05)
        assert not second_started.is_set()

        modal_executor_module.abandon_modal_workflow_execution_prompt(
            "prompt-1",
            "test cancellation",
        )
        await asyncio.wait_for(second_started.wait(), timeout=1.0)
        second_output = await second_task

        release_first_components.set()
        first_output_a, first_output_b = await asyncio.gather(first_task_a, first_task_b)
        return first_output_a, first_output_b, second_output

    modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient())
    try:
        first_output_a, first_output_b, second_output = asyncio.run(run_scenario())
    finally:
        modal_executor_module.set_remote_executor_client_factory(None)

    assert first_output_a.result == ("component-1a",)
    assert first_output_b.result == ("component-1b",)
    assert second_output.result == ("component-2",)
    assert set(observed_events[:2]) == {"start:component-1a", "start:component-1b"}
    assert observed_events[2:4] == ["start:component-2", "finish:component-2"]
    assert set(observed_events[4:]) == {"finish:component-1a", "finish:component-1b"}


class _FakeImplicitBatchKSamplerNode:
    """Fake sampler node exposing ComfyUI-style LATENT and primitive socket types."""

    @classmethod
    def INPUT_TYPES(cls) -> dict[str, dict[str, tuple[str]]]:
        """Return the minimal socket schema needed for implicit batch target inspection."""
        return {
            "required": {
                "latent_image": ("LATENT",),
                "seed": ("INT",),
                "positive": ("CONDITIONING",),
            }
        }


def _assert_node_module_identity(node_class: type[Any], expected_module: str) -> None:
    """Assert a registered V3 node exposes its loader-assigned module identity."""
    assert node_class.RELATIVE_PYTHON_MODULE == expected_module
    assert node_class.GET_NODE_INFO_V1()["python_module"] == expected_module


def test_dynamic_proxy_registration_sets_module_identity_on_new_and_cached_class(
    modal_executor_module: Any,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Ordinary proxies should serialize a valid module after creation and cache reuse."""
    expected_module = "custom_nodes.ComfyUI-Modal"
    monkeypatch.setattr(
        modal_executor_module.ModalUniversalExecutor,
        "RELATIVE_PYTHON_MODULE",
        expected_module,
    )
    fake_nodes_module = types.SimpleNamespace(
        NODE_CLASS_MAPPINGS={"OriginalNodeForModuleIdentity": _FakeOriginalNode},
        NODE_DISPLAY_NAME_MAPPINGS={},
    )

    proxy_node_id = modal_executor_module.ensure_modal_proxy_node_registered(
        original_class_type="OriginalNodeForModuleIdentity",
        original_class=_FakeOriginalNode,
        nodes_module=fake_nodes_module,
    )
    proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id]
    _assert_node_module_identity(proxy_class, expected_module)

    proxy_class.RELATIVE_PYTHON_MODULE = None
    fake_nodes_module.NODE_CLASS_MAPPINGS.clear()
    fake_nodes_module.NODE_DISPLAY_NAME_MAPPINGS.clear()

    cached_proxy_node_id = modal_executor_module.ensure_modal_proxy_node_registered(
        original_class_type="OriginalNodeForModuleIdentity",
        original_class=_FakeOriginalNode,
        nodes_module=fake_nodes_module,
    )

    assert cached_proxy_node_id == proxy_node_id
    assert fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id] is proxy_class
    _assert_node_module_identity(proxy_class, expected_module)


def test_component_proxy_registration_sets_module_identity_on_new_and_cached_class(
    modal_executor_module: Any,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Component proxies should serialize a valid module after creation and cache reuse."""
    expected_module = "custom_nodes.ComfyUI-Modal"
    monkeypatch.setattr(
        modal_executor_module.ModalUniversalExecutor,
        "RELATIVE_PYTHON_MODULE",
        expected_module,
    )
    fake_nodes_module = types.SimpleNamespace(
        NODE_CLASS_MAPPINGS={},
        NODE_DISPLAY_NAME_MAPPINGS={},
    )
    registration_kwargs = {
        "output_types": ("MODULE_IDENTITY_COMPONENT",),
        "output_names": ("value",),
        "output_is_list": (False,),
        "nodes_module": fake_nodes_module,
        "is_output_node": False,
    }

    proxy_node_id = modal_executor_module.ensure_modal_component_proxy_node_registered(
        **registration_kwargs
    )
    proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id]
    _assert_node_module_identity(proxy_class, expected_module)

    proxy_class.RELATIVE_PYTHON_MODULE = None
    fake_nodes_module.NODE_CLASS_MAPPINGS.clear()
    fake_nodes_module.NODE_DISPLAY_NAME_MAPPINGS.clear()

    cached_proxy_node_id = (
        modal_executor_module.ensure_modal_component_proxy_node_registered(
            **registration_kwargs
        )
    )

    assert cached_proxy_node_id == proxy_node_id
    assert fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id] is proxy_class
    _assert_node_module_identity(proxy_class, expected_module)


@pytest.mark.parametrize(
    ("registration_function_name", "node_class_name", "node_id_constant_name"),
    [
        (
            "ensure_modal_parallel_local_passthrough_registered",
            "ModalParallelLocalPassthrough",
            "MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID",
        ),
        (
            "ensure_modal_local_bridge_materializer_registered",
            "ModalLocalBridgeMaterializer",
            "MODAL_LOCAL_BRIDGE_MATERIALIZER_NODE_ID",
        ),
        (
            "ensure_modal_artifact_finalizer_registered",
            "ModalArtifactFinalizer",
            "MODAL_ARTIFACT_FINALIZER_NODE_ID",
        ),
    ],
)
def test_internal_node_registration_sets_module_identity(
    modal_executor_module: Any,
    monkeypatch: pytest.MonkeyPatch,
    registration_function_name: str,
    node_class_name: str,
    node_id_constant_name: str,
) -> None:
    """Static internal nodes should serialize the startup node's module identity."""
    expected_module = "custom_nodes.ComfyUI-Modal"
    monkeypatch.setattr(
        modal_executor_module.ModalUniversalExecutor,
        "RELATIVE_PYTHON_MODULE",
        expected_module,
    )
    fake_nodes_module = types.SimpleNamespace(
        NODE_CLASS_MAPPINGS={},
        NODE_DISPLAY_NAME_MAPPINGS={},
    )

    registration_function = getattr(modal_executor_module, registration_function_name)
    node_class = getattr(modal_executor_module, node_class_name)
    monkeypatch.setattr(
        node_class,
        "RELATIVE_PYTHON_MODULE",
        node_class.RELATIVE_PYTHON_MODULE,
    )
    registration_function(fake_nodes_module)

    node_id = getattr(modal_executor_module, node_id_constant_name)
    assert fake_nodes_module.NODE_CLASS_MAPPINGS[node_id] is node_class
    _assert_node_module_identity(node_class, expected_module)


def test_dynamic_proxy_node_preserves_output_signature(
    modal_executor_module: Any,
) -> None:
    """Dynamic Modal proxies should mirror the original output count and names."""
    fake_nodes_module = type(
        "FakeNodesModule",
        (),
        {
            "NODE_CLASS_MAPPINGS": {"OriginalNode": _FakeOriginalNode},
            "NODE_DISPLAY_NAME_MAPPINGS": {},
        },
    )()

    proxy_id = modal_executor_module.ensure_modal_proxy_node_registered(
        original_class_type="OriginalNode",
        original_class=_FakeOriginalNode,
        nodes_module=fake_nodes_module,
    )

    proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_id]
    schema = proxy_class.GET_SCHEMA()

    assert schema.node_id == proxy_id
    assert [output.display_name for output in schema.outputs] == ["image", "count"]
    assert [output.io_type for output in schema.outputs] == ["IMAGE", "INT"]
    assert proxy_class.INPUT_IS_LIST is True


def test_component_proxy_cache_distinguishes_scheduler_list_outputs(
    modal_executor_module: Any,
) -> None:
    """Mapped list proxies must not reuse a cached scalar-output proxy class."""
    fake_nodes_module = types.SimpleNamespace(
        NODE_CLASS_MAPPINGS={},
        NODE_DISPLAY_NAME_MAPPINGS={},
    )

    scalar_proxy_id = modal_executor_module.ensure_modal_component_proxy_node_registered(
        output_types=("LATENT",),
        output_names=("samples",),
        output_is_list=(False,),
        nodes_module=fake_nodes_module,
        is_output_node=False,
    )
    list_proxy_id = modal_executor_module.ensure_modal_component_proxy_node_registered(
        output_types=("LATENT",),
        output_names=("samples",),
        output_is_list=(True,),
        nodes_module=fake_nodes_module,
        is_output_node=False,
    )

    assert scalar_proxy_id != list_proxy_id
    assert fake_nodes_module.NODE_CLASS_MAPPINGS[scalar_proxy_id].OUTPUT_IS_LIST == [False]
    assert fake_nodes_module.NODE_CLASS_MAPPINGS[list_proxy_id].OUTPUT_IS_LIST == [True]


def test_mapped_boundary_payload_requests_scheduler_list_output(
    api_intercept_module: Any,
) -> None:
    """Explicit mapped outputs should tell the local proxy to expose ComfyUI list semantics."""
    boundary_output = api_intercept_module.BoundaryOutputSpec(
        proxy_output_name="sampler_samples",
        source=api_intercept_module.LinkedOutputRef(node_id="12", output_index=0),
        io_type="LATENT",
        is_list=False,
    )

    payload = api_intercept_module._boundary_output_payload(
        boundary_output,
        mapped_output=True,
    )

    assert payload["mapped_output"] is True
    assert payload["scheduler_is_list"] is True
    assert api_intercept_module._proxy_boundary_output_is_list(payload) is True


def test_scheduler_list_proxy_wraps_singleton_image_tensor(
    modal_executor_module: Any,
) -> None:
    """A singleton mapped IMAGE must retain its batch dimension in ComfyUI."""
    torch = pytest.importorskip("torch")

    fake_nodes_module = types.SimpleNamespace(
        NODE_CLASS_MAPPINGS={},
        NODE_DISPLAY_NAME_MAPPINGS={},
    )
    proxy_node_id = modal_executor_module.ensure_modal_component_proxy_node_registered(
        output_types=("IMAGE",),
        output_names=("image",),
        output_is_list=(True,),
        nodes_module=fake_nodes_module,
        is_output_node=False,
    )
    proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id]
    image = torch.zeros((1, 8, 6, 3), dtype=torch.float32)

    class FakeClient:
        """Return one ordinary IMAGE tensor from a singleton remote execution."""

        async def execute_payload_async(
            self,
            payload: dict[str, Any],
            kwargs: dict[str, Any],
        ) -> tuple[Any, ...]:
            """Return the unwrapped tensor produced by the ordinary subgraph path."""
            assert payload["boundary_outputs"][0]["scheduler_is_list"] is True
            assert kwargs == {}
            return (image,)

    modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient())
    try:
        result = asyncio.run(
            proxy_class.execute(
                original_node_data={
                    "payload_kind": "subgraph",
                    "prompt_id": "singleton-image-prompt",
                    "component_id": "singleton-image-component",
                    "boundary_outputs": [
                        {
                            "proxy_output_name": "image",
                            "node_id": "11",
                            "output_index": 0,
                            "io_type": "IMAGE",
                            "is_list": False,
                            "scheduler_is_list": True,
                        }
                    ],
                },
                unique_id="singleton-image-component",
            )
        )
    finally:
        modal_executor_module.set_remote_executor_client_factory(None)

    assert isinstance(result.result[0], list)
    assert len(result.result[0]) == 1
    assert result.result[0][0] is image
    scheduler_items: list[Any] = []
    scheduler_items.extend(result.result[0])
    assert len(scheduler_items) == 1
    assert scheduler_items[0] is image
    assert tuple(scheduler_items[0].shape) == (1, 8, 6, 3)


def test_transportable_singleton_list_reaches_scalar_remote_consumer_as_item(
    modal_executor_module: Any,
) -> None:
    """A cached singleton list output should enter the next remote proxy as a scalar."""
    fake_nodes_module = types.SimpleNamespace(
        NODE_CLASS_MAPPINGS={},
        NODE_DISPLAY_NAME_MAPPINGS={},
    )
    producer_proxy_id = (
        modal_executor_module.ensure_modal_component_proxy_node_registered(
            output_types=("INT",),
            output_names=("seed",),
            output_is_list=(True,),
            nodes_module=fake_nodes_module,
            is_output_node=False,
        )
    )
    consumer_proxy_id = (
        modal_executor_module.ensure_modal_component_proxy_node_registered(
            output_types=("STRING",),
            output_names=("result",),
            output_is_list=(False,),
            nodes_module=fake_nodes_module,
            is_output_node=False,
        )
    )
    producer_proxy = fake_nodes_module.NODE_CLASS_MAPPINGS[producer_proxy_id]
    consumer_proxy = fake_nodes_module.NODE_CLASS_MAPPINGS[consumer_proxy_id]
    observed_consumer_inputs: list[Any] = []

    class FakeClient:
        """Return a singleton seed list and record the downstream normalized input."""

        async def execute_payload_async(
            self,
            payload: dict[str, Any],
            kwargs: dict[str, Any],
        ) -> tuple[Any, ...]:
            """Model the two remote components around ComfyUI's output cache."""
            if payload["component_id"] == "producer":
                assert kwargs == {}
                return ([123],)
            observed_consumer_inputs.append(kwargs["remote_input_0"])
            return (f"seed:{int(kwargs['remote_input_0'])}",)

    modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient())
    try:
        producer_result = asyncio.run(
            producer_proxy.execute(
                original_node_data={
                    "payload_kind": "subgraph",
                    "component_id": "producer",
                    "boundary_outputs": [
                        {
                            "proxy_output_name": "seed",
                            "node_id": "325",
                            "output_index": 0,
                            "io_type": "INT",
                            "is_list": True,
                        }
                    ],
                },
                unique_id="producer",
            )
        )
        cached_scheduler_items: list[Any] = []
        cached_scheduler_items.extend(producer_result.result[0])
        consumer_result = asyncio.run(
            consumer_proxy.execute(
                original_node_data={
                    "payload_kind": "subgraph",
                    "component_id": "consumer",
                    "boundary_outputs": [
                        {
                            "proxy_output_name": "result",
                            "node_id": "372",
                            "output_index": 0,
                            "io_type": "STRING",
                            "is_list": False,
                        }
                    ],
                },
                remote_input_0=cached_scheduler_items,
                unique_id="consumer",
            )
        )
    finally:
        modal_executor_module.set_remote_executor_client_factory(None)

    assert cached_scheduler_items == [123]
    assert observed_consumer_inputs == [123]
    assert consumer_result.result == ("seed:123",)


def test_component_proxy_emits_completion_token_after_remote_execution(
    modal_executor_module: Any,
) -> None:
    """Artifact-only component proxies should expose completion after remote work returns."""
    fake_nodes_module = types.SimpleNamespace(
        NODE_CLASS_MAPPINGS={},
        NODE_DISPLAY_NAME_MAPPINGS={},
    )
    proxy_node_id = modal_executor_module.ensure_modal_component_proxy_node_registered(
        output_types=(),
        output_names=(),
        output_is_list=(),
        nodes_module=fake_nodes_module,
        is_output_node=False,
        include_completion_output=True,
    )
    proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id]
    schema = proxy_class.GET_SCHEMA()

    assert [output.io_type for output in schema.outputs] == ["BOOLEAN"]
    assert [output.display_name for output in schema.outputs] == [
        modal_executor_module.MODAL_COMPONENT_COMPLETION_OUTPUT_NAME
    ]

    class FakeClient:
        """Return no remote boundary values for an artifact-only component."""

        async def execute_payload_async(
            self,
            payload: dict[str, Any],
            kwargs: dict[str, Any],
        ) -> tuple[()]:
            """Confirm the remote component ran before returning no boundary outputs."""
            assert payload["component_id"] == "artifact-component"
            assert kwargs == {}
            return ()

    modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient())
    try:
        result = asyncio.run(
            proxy_class.execute(
                original_node_data={
                    "payload_kind": "subgraph",
                    "prompt_id": "artifact-prompt",
                    "component_id": "artifact-component",
                },
                unique_id="artifact-component",
            )
        )
    finally:
        modal_executor_module.set_remote_executor_client_factory(None)

    assert result.result == (True,)


def test_modal_artifact_finalizer_requires_completed_components(
    modal_executor_module: Any,
) -> None:
    """The internal output sink should reject any missing completion signal."""
    schema = modal_executor_module.ModalArtifactFinalizer.GET_SCHEMA()

    assert schema.is_output_node is True
    assert schema.outputs == []
    assert modal_executor_module.ModalArtifactFinalizer.execute(
        {"component_0": True, "component_1": True}
    ).result is None
    with pytest.raises(RuntimeError, match="component_1"):
        modal_executor_module.ModalArtifactFinalizer.execute(
            {"component_0": True, "component_1": False}
        )


def test_modal_parallel_local_passthrough_releases_after_remote_dispatch(
    modal_executor_module: Any,
) -> None:
    """A local preview should unblock when its remote continuation starts."""
    schema = modal_executor_module.ModalParallelLocalPassthrough.GET_SCHEMA()
    dispatch_context = {
        "dispatch_group_id": "prompt-parallel",
        "component_ids": ["component-b"],
    }

    async def run_scenario() -> tuple[Any, ...]:
        """Prove the passthrough remains pending until remote dispatch is signalled."""
        passthrough_task = asyncio.create_task(
            modal_executor_module.ModalParallelLocalPassthrough.execute(
                "image-value",
                dispatch_context,
            )
        )
        await asyncio.sleep(0)
        assert passthrough_task.done() is False
        assert modal_executor_module._signal_parallel_local_dispatch(
            {
                "parallel_local_dispatch_group_id": "prompt-parallel",
                "component_id": "component-b",
                "signal_parallel_local_dispatch": True,
            }
        ) is True
        result = await asyncio.wait_for(passthrough_task, timeout=1.0)
        return result.result

    try:
        result = asyncio.run(run_scenario())
    finally:
        modal_executor_module._clear_parallel_dispatch_events("prompt-parallel")

    assert schema.is_output_node is False
    assert result == ("image-value",)


def test_modal_parallel_local_passthrough_runs_before_remote_result_returns(
    modal_executor_module: Any,
) -> None:
    """A local preview should run while its downstream remote call is in flight."""
    fake_nodes_module = types.SimpleNamespace(
        NODE_CLASS_MAPPINGS={},
        NODE_DISPLAY_NAME_MAPPINGS={},
    )
    proxy_node_id = modal_executor_module.ensure_modal_component_proxy_node_registered(
        output_types=("STRING",),
        output_names=("value",),
        output_is_list=(False,),
        nodes_module=fake_nodes_module,
        is_output_node=False,
    )
    proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_node_id]
    remote_started = asyncio.Event()
    release_remote = asyncio.Event()

    class FakeClient:
        """Fake async remote client that remains in flight until explicitly released."""

        async def execute_payload_async(
            self,
            payload: dict[str, Any],
            kwargs: dict[str, Any],
        ) -> tuple[str]:
            """Record dispatch, then keep the remote result pending."""
            remote_started.set()
            await release_remote.wait()
            return (str(payload["component_id"]),)

    async def run_scenario() -> tuple[tuple[Any, ...], tuple[Any, ...]]:
        """Prove local passthrough completion precedes the remote result."""
        dispatch_group_id = "prompt-parallel-in-flight"
        local_task = asyncio.create_task(
            modal_executor_module.ModalParallelLocalPassthrough.execute(
                "preview-image",
                {
                    "dispatch_group_id": dispatch_group_id,
                    "component_ids": ["component-b"],
                },
            )
        )
        remote_task = asyncio.create_task(
            proxy_class.execute(
                original_node_data={
                    "payload_kind": "subgraph",
                    "prompt_id": dispatch_group_id,
                    "component_id": "component-b",
                    "parallel_local_dispatch_group_id": dispatch_group_id,
                    "signal_parallel_local_dispatch": True,
                },
                unique_id="component-b",
            )
        )
        await asyncio.wait_for(remote_started.wait(), timeout=1.0)
        local_result = await asyncio.wait_for(local_task, timeout=1.0)
        assert remote_task.done() is False
        release_remote.set()
        remote_result = await asyncio.wait_for(remote_task, timeout=1.0)
        return local_result.result, remote_result.result

    modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient())
    try:
        local_result, remote_result = asyncio.run(run_scenario())
    finally:
        modal_executor_module.set_remote_executor_client_factory(None)
        modal_executor_module._clear_parallel_dispatch_events(
            "prompt-parallel-in-flight"
        )

    assert local_result == ("preview-image",)
    assert remote_result == ("component-b",)


def test_modal_local_bridge_materializer_downloads_without_blocking_event_loop(
    modal_executor_module: Any,
    remote_modal_app_module: Any,
    session_state_module: Any,
    monkeypatch: Any,
) -> None:
    """The local bridge node should perform its durable download in a worker thread."""
    materialization_started = threading.Event()
    release_materialization = threading.Event()

    def fake_materialize(ref_payload: dict[str, Any]) -> str:
        """Block the worker thread until the async test permits completion."""
        assert session_state_module.is_remote_session_bridge_ref_payload(ref_payload)
        materialization_started.set()
        assert release_materialization.wait(timeout=1.0)
        return "local-image"

    monkeypatch.setattr(
        remote_modal_app_module,
        "materialize_remote_session_bridge_ref_locally",
        fake_materialize,
    )
    bridge_ref = session_state_module.RemoteSessionBridgeRef(
        bridge_key="RSB_local_materializer",
        node_id="251",
        output_index=0,
        session_id="session-source",
    ).to_payload()

    async def run_scenario() -> tuple[Any, ...]:
        """Verify another coroutine runs while durable materialization is blocked."""
        materializer_task = asyncio.create_task(
            modal_executor_module.ModalLocalBridgeMaterializer.execute(bridge_ref)
        )
        await asyncio.to_thread(materialization_started.wait, 1.0)
        assert materializer_task.done() is False
        await asyncio.sleep(0)
        release_materialization.set()
        result = await asyncio.wait_for(materializer_task, timeout=1.0)
        return result.result

    assert asyncio.run(run_scenario()) == ("local-image",)


def test_local_bridge_materialization_restores_durable_inline_output(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """A local branch should restore the producer's durable serialized output."""
    import torch

    bridge_ref = remote_modal_app_module.RemoteSessionBridgeRef(
        bridge_key="RSB_local_inline_bridge",
        node_id="251",
        output_index=0,
        session_id="session-source",
    )
    image = torch.arange(12, dtype=torch.float32).reshape(1, 2, 2, 3)
    record = remote_modal_app_module.RemoteSessionBridgeRecord(
        bridge_key=bridge_ref.bridge_key,
        node_id=bridge_ref.node_id,
        output_index=bridge_ref.output_index,
        producer_payload={"component_id": "251"},
        producer_inputs={},
        serialized_output=remote_modal_app_module.serialize_value(image),
        serialized_output_io_type="IMAGE",
    )
    monkeypatch.setattr(
        remote_modal_app_module._REMOTE_SESSION_BRIDGE_STORE,
        "get_record",
        lambda bridge_key: record,
    )

    restored = remote_modal_app_module.materialize_remote_session_bridge_ref_locally(
        bridge_ref.to_payload()
    )

    assert torch.equal(restored, image)


def test_local_bridge_materialization_downloads_object_backed_output(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """A large bridge should download its content-addressed output object locally."""
    monkeypatch.delenv("MODAL_ENVIRONMENT", raising=False)
    bridge_ref = remote_modal_app_module.RemoteSessionBridgeRef(
        bridge_key="RSB_local_object_bridge",
        node_id="251",
        output_index=0,
        session_id="session-source",
    )
    stored_payload = remote_modal_app_module.serialize_node_outputs(("large-image",))
    object_ref = remote_modal_app_module.DurableObjectRef(
        object_path="bridge-outputs/sha256/value.bin",
        sha256=hashlib.sha256(stored_payload).hexdigest(),
        size_bytes=len(stored_payload),
    )
    record = remote_modal_app_module.RemoteSessionBridgeRecord(
        bridge_key=bridge_ref.bridge_key,
        node_id=bridge_ref.node_id,
        output_index=bridge_ref.output_index,
        producer_payload={"component_id": "251"},
        producer_inputs={},
        serialized_output_object=object_ref,
        serialized_output_io_type="IMAGE",
    )
    dict_calls: list[tuple[Any, ...]] = []
    volume_calls: list[tuple[Any, ...]] = []

    class FakeDict:
        """Expose the durable bridge record through Modal's shared Dict API."""

        @staticmethod
        def from_name(*args: Any, **kwargs: Any) -> Any:
            """Return a Dict handle that records bridge lookups."""
            dict_calls.append((*args, kwargs))
            return types.SimpleNamespace(get=lambda bridge_key: record.to_payload())

    class FakeVolume:
        """Expose the durable object through Modal's direct Volume API."""

        @staticmethod
        def from_name(*args: Any, **kwargs: Any) -> Any:
            """Return a Volume handle that records object downloads."""
            volume_calls.append((*args, kwargs))
            return types.SimpleNamespace(
                read_file=lambda volume_path: iter((stored_payload[:7], stored_payload[7:]))
            )

    monkeypatch.setattr(
        remote_modal_app_module,
        "get_settings",
        lambda: types.SimpleNamespace(
            execution_mode="remote",
            session_bridge_dict_name="bridge-dict",
            volume_name="bridge-volume",
        ),
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "modal",
        types.SimpleNamespace(Dict=FakeDict, Volume=FakeVolume),
    )
    remote_modal_app_module._MODAL_SESSION_BRIDGE_DICTS.clear()
    remote_modal_app_module._MODAL_DURABLE_VOLUMES.clear()

    restored = remote_modal_app_module.materialize_remote_session_bridge_ref_locally(
        bridge_ref.to_payload()
    )

    assert restored == "large-image"
    assert dict_calls == [
        (
            "bridge-dict",
            {"environment_name": None, "create_if_missing": True},
        )
    ]
    assert volume_calls == [
        (
            "bridge-volume",
            {"environment_name": None, "create_if_missing": True},
        )
    ]
    remote_modal_app_module._MODAL_SESSION_BRIDGE_DICTS.clear()
    remote_modal_app_module._MODAL_DURABLE_VOLUMES.clear()


def test_parallel_local_dispatch_frontier_stops_at_nearest_remote_component(
    api_intercept_module: Any,
) -> None:
    """A preview gate should not wait for remote descendants of a remote frontier."""
    fake_nodes_module = types.SimpleNamespace(
        NODE_CLASS_MAPPINGS={
            "RemoteProxy": _FakeRewriteRemoteSamplerNode,
            "LocalFeedback": _FakeRewriteLocalFeedbackNode,
        },
    )
    rewritten_prompt = {
        "source": {"class_type": "RemoteProxy", "inputs": {}},
        "feedback": {
            "class_type": "LocalFeedback",
            "inputs": {"value": ["source", 0]},
        },
        "nearest": {
            "class_type": "RemoteProxy",
            "inputs": {"value": ["feedback", 0]},
        },
        "later": {
            "class_type": "RemoteProxy",
            "inputs": {"value": ["nearest", 0]},
        },
    }
    consumers = api_intercept_module._build_consumer_map(rewritten_prompt)

    nearest_component_ids = (
        api_intercept_module._nearest_downstream_remote_component_ids(
            rewritten_prompt=rewritten_prompt,
            consumers=consumers,
            seed_targets=consumers[
                api_intercept_module.LinkedOutputRef("source", 0)
            ],
            remote_component_id_set={"source", "nearest", "later"},
            nodes_module=fake_nodes_module,
        )
    )

    assert nearest_component_ids == ["nearest"]


def test_proxy_execution_uses_injected_remote_client(
    modal_executor_module: Any,
) -> None:
    """Proxy execution should delegate to the configured remote client asynchronously."""
    fake_nodes_module = type(
        "FakeNodesModule",
        (),
        {
            "NODE_CLASS_MAPPINGS": {"OriginalNode": _FakeOriginalNode},
            "NODE_DISPLAY_NAME_MAPPINGS": {},
        },
    )()

    proxy_id = modal_executor_module.ensure_modal_proxy_node_registered(
        original_class_type="OriginalNode",
        original_class=_FakeOriginalNode,
        nodes_module=fake_nodes_module,
    )
    proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_id]

    class FakeClient:
        """Test client that returns deterministic outputs."""

        async def execute_payload_async(
            self,
            payload: dict[str, Any],
            kwargs: dict[str, Any],
        ) -> tuple[str, int]:
            """Return values derived from the proxied node payload."""
            return (f"{payload['class_type']}::{kwargs['value']}", 3)

    modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient())
    try:
        result = asyncio.run(
            proxy_class.execute(original_node_data={"class_type": "OriginalNode"}, value="payload")
        )
    finally:
        modal_executor_module.set_remote_executor_client_factory(None)

    assert result.result == ("OriginalNode::payload", 3)


def test_proxy_execution_normalizes_input_is_list_kwargs(
    modal_executor_module: Any,
) -> None:
    """Dynamic Modal proxies should unwrap singleton INPUT_IS_LIST wrappers before remote execution."""
    fake_nodes_module = type(
        "FakeNodesModule",
        (),
        {
            "NODE_CLASS_MAPPINGS": {"OriginalNode": _FakeOriginalNode},
            "NODE_DISPLAY_NAME_MAPPINGS": {},
        },
    )()

    proxy_id = modal_executor_module.ensure_modal_proxy_node_registered(
        original_class_type="OriginalNode",
        original_class=_FakeOriginalNode,
        nodes_module=fake_nodes_module,
    )
    proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_id]
    observed_kwargs: dict[str, Any] = {}

    class FakeClient:
        """Test client that records normalized proxy kwargs."""

        async def execute_payload_async(
            self,
            payload: dict[str, Any],
            kwargs: dict[str, Any],
        ) -> tuple[str, int]:
            """Capture the kwargs forwarded by the proxy."""
            observed_kwargs["payload_kind"] = payload["payload_kind"]
            observed_kwargs.update(kwargs)
            return ("ok", 1)

    modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient())
    try:
        result = asyncio.run(
            proxy_class.execute(
                original_node_data=[{"payload_kind": "subgraph", "class_type": "OriginalNode"}],
                scalar_value=[3],
                mapped_value=["a", "b", "c"],
            )
        )
    finally:
        modal_executor_module.set_remote_executor_client_factory(None)

    assert result.result == ("ok", 1)
    assert observed_kwargs == {
        "payload_kind": "subgraph",
        "scalar_value": 3,
        "mapped_value": ["a", "b", "c"],
    }


@pytest.mark.parametrize("hidden_delivery", ["v3_class_clone", "legacy_kwargs"])
def test_cache_friendly_proxy_payload_rehydrates_prompt_id_at_execution(
    modal_executor_module: Any,
    hidden_delivery: str,
) -> None:
    """V3 and legacy hidden inputs should select the correct overlapping prompt."""
    fake_nodes_module = type(
        "FakeNodesModule",
        (),
        {
            "NODE_CLASS_MAPPINGS": {"OriginalNode": _FakeOriginalNode},
            "NODE_DISPLAY_NAME_MAPPINGS": {},
        },
    )()

    proxy_id = modal_executor_module.ensure_modal_proxy_node_registered(
        original_class_type="OriginalNode",
        original_class=_FakeOriginalNode,
        nodes_module=fake_nodes_module,
    )
    proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_id]

    payload = modal_executor_module.register_cache_friendly_proxy_payload(
        "node-1",
        {
            "payload_kind": "subgraph",
            "component_id": "component-1",
            "prompt_id": "prompt-1",
            "boundary_outputs": [],
            "execute_node_ids": [],
        },
    )
    next_payload = modal_executor_module.register_cache_friendly_proxy_payload(
        "node-1",
        {
            "payload_kind": "subgraph",
            "component_id": "component-1",
            "prompt_id": "prompt-2",
            "boundary_outputs": [],
            "execute_node_ids": [],
        },
    )

    class FakeClient:
        """Test client that captures the rehydrated payload."""

        async def execute_payload_async(
            self,
            payload: dict[str, Any],
            kwargs: dict[str, Any],
        ) -> tuple[str, int]:
            """Return values derived from the restored prompt id."""
            return (str(payload.get("prompt_id")), len(kwargs))

    modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient())
    try:
        hidden_metadata = {
            modal_executor_module.MODAL_PROMPT_ID_EXTRA_PNGINFO_KEY: "prompt-1"
        }
        execution_proxy_class = proxy_class
        hidden_kwargs: dict[str, Any] = {}
        if hidden_delivery == "v3_class_clone":
            execution_proxy_class = proxy_class.PREPARE_CLASS_CLONE(
                {
                    "hidden_inputs": {
                        modal_executor_module.io.Hidden.unique_id: "node-1",
                        modal_executor_module.io.Hidden.extra_pnginfo: hidden_metadata,
                    }
                }
            )
        else:
            hidden_kwargs = {
                "unique_id": ["node-1"],
                "extra_pnginfo": [hidden_metadata],
            }
        result = asyncio.run(
            execution_proxy_class.execute(
                original_node_data=[payload],
                value=["payload"],
                **hidden_kwargs,
            )
        )
    finally:
        modal_executor_module.set_remote_executor_client_factory(None)

    assert "prompt_id" not in payload
    assert payload == next_payload
    assert result.result == ("prompt-1", 1)


def test_register_cache_friendly_proxy_payload_strips_session_fields_and_rehydrates_them(
    modal_executor_module: Any,
) -> None:
    """Session-backed proxy payloads should strip run-scoped fields from the local cache surface."""
    payload = modal_executor_module.register_cache_friendly_proxy_payload(
        "node-1",
        {
            "payload_kind": "subgraph",
            "component_id": "component-1",
            "prompt_id": "prompt-1",
            "remote_session": {
                "session_id": "session-1",
                "prompt_id": "prompt-1",
                "owner_component_id": "component-1",
            },
            "boundary_outputs": [],
            "execute_node_ids": [],
            "clear_remote_session": True,
            "extra_data": {
                "prompt_id": "prompt-1",
                "create_time": 1234567890,
                "modal": {"remote_node_ids": ["12"], "estimated_max_parallel_requests": 1},
            },
            "requires_volume_reload": True,
            "volume_reload_marker": "marker-1",
            "uploaded_volume_paths": ["/storage/assets/example.safetensors"],
        },
    )

    assert "prompt_id" not in payload
    assert "remote_session" not in payload
    assert "clear_remote_session" not in payload
    assert "extra_data" not in payload
    assert "requires_volume_reload" not in payload
    assert "volume_reload_marker" not in payload
    assert "uploaded_volume_paths" not in payload
    assert modal_executor_module._rehydrate_proxy_payload(payload, unique_id="node-1") == {
        "payload_kind": "subgraph",
        "component_id": "component-1",
        "prompt_id": "prompt-1",
        "remote_session": {
            "session_id": "session-1",
            "prompt_id": "prompt-1",
            "owner_component_id": "component-1",
        },
        "boundary_outputs": [],
        "execute_node_ids": [],
        "clear_remote_session": True,
        "extra_data": {
            "prompt_id": "prompt-1",
            "create_time": 1234567890,
            "modal": {"remote_node_ids": ["12"], "estimated_max_parallel_requests": 1},
        },
        "requires_volume_reload": True,
        "volume_reload_marker": "marker-1",
        "uploaded_volume_paths": ["/storage/assets/example.safetensors"],
    }


def test_cache_friendly_proxy_payload_rehydrates_without_hidden_unique_id(
    modal_executor_module: Any,
) -> None:
    """Cache-friendly proxy payloads should still rehydrate when ComfyUI omits hidden unique_id."""
    payload = modal_executor_module.register_cache_friendly_proxy_payload(
        "node-7",
        {
            "payload_kind": "subgraph",
            "component_id": "component-7",
            "prompt_id": "prompt-7",
            "remote_session": {
                "session_id": "session-7",
                "prompt_id": "prompt-7",
                "owner_component_id": "component-7",
            },
            "boundary_outputs": [],
            "execute_node_ids": [],
        },
    )

    assert modal_executor_module._rehydrate_proxy_payload(payload, unique_id=None) == {
        "payload_kind": "subgraph",
        "component_id": "component-7",
        "prompt_id": "prompt-7",
        "remote_session": {
            "session_id": "session-7",
            "prompt_id": "prompt-7",
            "owner_component_id": "component-7",
        },
        "boundary_outputs": [],
        "execute_node_ids": [],
    }


def test_cache_friendly_proxy_payload_ignores_volatile_queue_metadata(
    modal_executor_module: Any,
) -> None:
    """Identical proxy work should sanitize to one local cache surface across prompt runs."""
    first_payload = modal_executor_module.register_cache_friendly_proxy_payload(
        "node-9",
        {
            "payload_kind": "subgraph",
            "component_id": "component-9",
            "prompt_id": "prompt-1",
            "boundary_outputs": [],
            "execute_node_ids": ["12"],
            "extra_data": {
                "client_id": "client-1",
                "create_time": 1000,
                "modal": {"remote_component_ids": ["12"]},
            },
            "requires_volume_reload": True,
            "volume_reload_marker": "marker-1",
            "uploaded_volume_paths": ["/storage/assets/a.bin"],
        },
    )
    second_payload = modal_executor_module.register_cache_friendly_proxy_payload(
        "node-9",
        {
            "payload_kind": "subgraph",
            "component_id": "component-9",
            "prompt_id": "prompt-2",
            "boundary_outputs": [],
            "execute_node_ids": ["12"],
            "extra_data": {
                "client_id": "client-1",
                "create_time": 2000,
                "modal": {"remote_component_ids": ["12", "39"]},
            },
            "requires_volume_reload": False,
            "volume_reload_marker": None,
            "uploaded_volume_paths": [],
        },
    )

    assert first_payload == second_payload == {
        "payload_kind": "subgraph",
        "component_id": "component-9",
        "boundary_outputs": [],
        "execute_node_ids": ["12"],
        modal_executor_module._PROXY_CACHE_CONTEXT_ID_KEY: "node-9",
    }
    assert modal_executor_module._rehydrate_proxy_payload(
        first_payload,
        unique_id="node-9",
        prompt_id="prompt-1",
    ) == {
        "payload_kind": "subgraph",
        "component_id": "component-9",
        "prompt_id": "prompt-1",
        "boundary_outputs": [],
        "execute_node_ids": ["12"],
        "extra_data": {
            "client_id": "client-1",
            "create_time": 1000,
            "modal": {"remote_component_ids": ["12"]},
        },
        "requires_volume_reload": True,
        "volume_reload_marker": "marker-1",
        "uploaded_volume_paths": ["/storage/assets/a.bin"],
    }
    assert modal_executor_module._rehydrate_proxy_payload(
        second_payload,
        unique_id="node-9",
        prompt_id="prompt-2",
    ) == {
        "payload_kind": "subgraph",
        "component_id": "component-9",
        "prompt_id": "prompt-2",
        "boundary_outputs": [],
        "execute_node_ids": ["12"],
        "extra_data": {
            "client_id": "client-1",
            "create_time": 2000,
            "modal": {"remote_component_ids": ["12", "39"]},
        },
        "requires_volume_reload": False,
        "volume_reload_marker": None,
        "uploaded_volume_paths": [],
    }


def test_proxy_execution_wraps_sync_remote_clients(
    modal_executor_module: Any,
) -> None:
    """Async proxy execution should still support legacy sync remote clients."""
    fake_nodes_module = type(
        "FakeNodesModule",
        (),
        {
            "NODE_CLASS_MAPPINGS": {"OriginalNode": _FakeOriginalNode},
            "NODE_DISPLAY_NAME_MAPPINGS": {},
        },
    )()

    proxy_id = modal_executor_module.ensure_modal_proxy_node_registered(
        original_class_type="OriginalNode",
        original_class=_FakeOriginalNode,
        nodes_module=fake_nodes_module,
    )
    proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[proxy_id]

    class FakeSyncClient:
        """Legacy client that only exposes the blocking execution method."""

        def execute_payload(self, payload: dict[str, Any], kwargs: dict[str, Any]) -> tuple[str, int]:
            """Return values derived from the proxied node payload."""
            return (f"sync::{payload['class_type']}::{kwargs['value']}", 4)

    modal_executor_module.set_remote_executor_client_factory(lambda: FakeSyncClient())
    try:
        result = asyncio.run(
            proxy_class.execute(original_node_data={"class_type": "OriginalNode"}, value="payload")
        )
    finally:
        modal_executor_module.set_remote_executor_client_factory(None)

    assert result.result == ("sync::OriginalNode::payload", 4)


def test_modal_map_input_execute_boosts_exact_warmup_for_registered_context(
    modal_executor_module: Any,
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Local Modal Map Input execution should kick off exact warmup once the real list value is known."""
    observed_boost: dict[str, Any] = {}

    def fake_boost_mapped_component_warmup(
        payload: dict[str, Any],
        *,
        total_items: int,
        reason: str,
    ) -> tuple[int, int]:
        """Record the exact warmup boost request without touching Modal."""
        observed_boost["payload"] = dict(payload)
        observed_boost["total_items"] = total_items
        observed_boost["reason"] = reason
        return 2, 2

    monkeypatch.setattr(
        remote_modal_app_module,
        "boost_mapped_component_warmup",
        fake_boost_mapped_component_warmup,
    )
    with modal_executor_module._MODAL_MAP_WARMUP_CONTEXTS_LOCK:
        modal_executor_module._MODAL_MAP_WARMUP_CONTEXTS.clear()
    modal_executor_module.register_modal_map_input_warmup_context(
        "map-node-1",
        {
            "prompt_id": "prompt-1",
            "component_id": "mapped-component-1",
            "extra_data": {"modal": {"mapped_component_ids": ["mapped-component-1"]}},
        },
        "INT",
    )

    result = modal_executor_module.ModalMapInput.execute(
        value=[10, 11, 12],
        unique_id="map-node-1",
    )

    assert result.result == ([10, 11, 12],)
    assert observed_boost == {
        "payload": {
            "prompt_id": "prompt-1",
            "component_id": "mapped-component-1",
            "extra_data": {"modal": {"mapped_component_ids": ["mapped-component-1"]}},
        },
        "total_items": 3,
        "reason": "modal_map_input_execute",
    }


def test_modal_map_input_execute_accepts_scalar_for_single_item_warmup(
    modal_executor_module: Any,
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """A scalar Modal Map Input should pass through and warm one mapped item."""
    observed_boost: dict[str, Any] = {}

    def fake_boost_mapped_component_warmup(
        payload: dict[str, Any],
        *,
        total_items: int,
        reason: str,
    ) -> tuple[int, int]:
        """Record the scalar warmup boost request without touching Modal."""
        observed_boost["payload"] = dict(payload)
        observed_boost["total_items"] = total_items
        observed_boost["reason"] = reason
        return 1, 1

    monkeypatch.setattr(
        remote_modal_app_module,
        "boost_mapped_component_warmup",
        fake_boost_mapped_component_warmup,
    )
    with modal_executor_module._MODAL_MAP_WARMUP_CONTEXTS_LOCK:
        modal_executor_module._MODAL_MAP_WARMUP_CONTEXTS.clear()
    modal_executor_module.register_modal_map_input_warmup_context(
        "map-node-1",
        {
            "prompt_id": "prompt-1",
            "component_id": "mapped-component-1",
            "extra_data": {"modal": {"mapped_component_ids": ["mapped-component-1"]}},
        },
        "INT",
    )

    result = modal_executor_module.ModalMapInput.execute(
        value=7,
        unique_id="map-node-1",
    )

    assert result.result == (7,)
    assert observed_boost == {
        "payload": {
            "prompt_id": "prompt-1",
            "component_id": "mapped-component-1",
            "extra_data": {"modal": {"mapped_component_ids": ["mapped-component-1"]}},
        },
        "total_items": 1,
        "reason": "modal_map_input_execute",
    }


def test_local_remote_app_executes_original_node(
    remote_modal_app_module: Any,
    serialization_module: Any,
) -> None:
    """The local fallback remote app should execute a mapped legacy node."""
    payload = remote_modal_app_module.execute_node_locally(
        node_data={"class_type": "OriginalNode"},
        kwargs_payload='{"value": "hello"}',
        node_mapping={"OriginalNode": _FakeOriginalNode},
    )
    outputs = serialization_module.deserialize_node_outputs(payload)
    assert outputs == ("hello", 1)


def test_remote_modal_call_worker_count_is_bounded_independently_of_local_cpus(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Local CPU count should not determine expensive remote GPU concurrency."""
    monkeypatch.setattr(remote_modal_app_module.os, "cpu_count", lambda: 8)
    remote_modal_app_module.get_settings.cache_clear()

    assert remote_modal_app_module._remote_modal_call_worker_count() == 4


def test_remote_modal_call_worker_count_honors_explicit_inflight_limit(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Operators should be able to choose the local remote-call budget explicitly."""
    monkeypatch.setattr(remote_modal_app_module.os, "cpu_count", lambda: 8)
    monkeypatch.setenv("COMFY_MODAL_MAX_INFLIGHT_CALLS", "12")
    remote_modal_app_module.get_settings.cache_clear()
    try:
        assert remote_modal_app_module._remote_modal_call_worker_count() == 12
    finally:
        remote_modal_app_module.get_settings.cache_clear()


def test_ensure_remote_warm_capacity_deduplicates_prompt_slots(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Prompt-scoped proactive warmup should only schedule each target slot once."""
    submitted_tasks: list[tuple[Any, tuple[Any, ...]]] = []

    class FakeExecutor:
        """Minimal executor that records submitted warmup jobs."""

        def submit(self, fn: Any, *args: Any) -> Future[Any]:
            """Capture one scheduled warmup task without running it."""
            submitted_tasks.append((fn, args))
            return Future()

    monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "remote")
    monkeypatch.setattr(remote_modal_app_module, "modal", object())
    monkeypatch.setattr(remote_modal_app_module, "_REMOTE_MODAL_WARMUP_EXECUTOR", FakeExecutor())
    remote_modal_app_module.get_settings.cache_clear()
    with remote_modal_app_module._PROMPT_WARMUP_STATES_LOCK:
        remote_modal_app_module._PROMPT_WARMUP_STATES.clear()
        remote_modal_app_module._PROMPT_WARMUP_STATE_ORDER = None

    try:
        warmup_request = {"prompt_id": "prompt-1", "component_id": "component-1"}
        first_target = remote_modal_app_module.ensure_remote_warm_capacity(
            warmup_request,
            warmup_target=2,
            reason="queue_time",
        )
        second_target = remote_modal_app_module.ensure_remote_warm_capacity(
            warmup_request,
            warmup_target=2,
            reason="queue_time_repeat",
        )
        third_target = remote_modal_app_module.ensure_remote_warm_capacity(
            warmup_request,
            warmup_target=4,
            reason="runtime_top_up",
        )
    finally:
        remote_modal_app_module.get_settings.cache_clear()

    assert first_target == 2
    assert second_target == 2
    assert third_target == 4
    assert len(submitted_tasks) == 4
    assert [args[1] for _fn, args in submitted_tasks] == [0, 1, 2, 3]


def test_speculative_affinity_prewarm_is_distinct_and_deduplicated(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """A running component should schedule its next distinct affinity exactly once."""
    submitted_tasks: list[tuple[Any, tuple[Any, ...]]] = []

    class FakeExecutor:
        """Minimal executor that records speculative warmup jobs."""

        def submit(self, fn: Any, *args: Any) -> Future[Any]:
            """Capture one scheduled task without running it."""
            submitted_tasks.append((fn, args))
            return Future()

    monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "remote")
    monkeypatch.setenv("COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT", "false")
    monkeypatch.setattr(remote_modal_app_module, "modal", object())
    monkeypatch.setattr(
        remote_modal_app_module,
        "_REMOTE_MODAL_WARMUP_EXECUTOR",
        FakeExecutor(),
    )
    remote_modal_app_module.get_settings.cache_clear()
    with remote_modal_app_module._PROMPT_WARMUP_STATES_LOCK:
        remote_modal_app_module._PROMPT_WARMUP_STATES.clear()
        remote_modal_app_module._PROMPT_WARMUP_STATE_ORDER = None

    payload = {
        "prompt_id": "prompt-spec",
        "component_id": "llm-component",
        "execution_provider": "modal",
        "execution_environment_id": "modal:RTX-PRO-6000",
        "remote_worker_affinity_group": "llm",
        "speculative_remote_prewarm_target": {
            "prompt_id": "prompt-spec",
            "component_id": "comfy-component",
            "modal_gpu": "RTX-PRO-6000",
            "execution_provider": "modal",
            "execution_environment_id": "modal:RTX-PRO-6000",
            "remote_worker_affinity_group": "comfy",
            "remote_local_gap_pool": True,
            "subgraph_prompt": {
                "11": {
                    "class_type": "UNETLoader",
                    "inputs": {"unet_name": "video-model.safetensors"},
                }
            },
        },
    }

    try:
        assert remote_modal_app_module._schedule_speculative_affinity_prewarm(
            payload,
            reason="test_first_event",
        ) is True
        assert remote_modal_app_module._schedule_speculative_affinity_prewarm(
            payload,
            reason="test_duplicate_event",
        ) is False
    finally:
        remote_modal_app_module.get_settings.cache_clear()

    assert len(submitted_tasks) == 1
    scheduled_function, scheduled_args = submitted_tasks[0]
    assert scheduled_function is remote_modal_app_module._run_speculative_affinity_prewarm
    assert scheduled_args[0] == "prompt-spec"
    scheduled_identity = json.loads(scheduled_args[1])
    assert scheduled_identity["affinity"] == "worker-pool:comfy:slot:0"
    assert scheduled_identity["snapshot_variant"] == "direct"
    assert scheduled_args[2]["remote_worker_affinity_group"] == "comfy"
    assert scheduled_args[2]["remote_local_gap_pool"] is True
    assert [
        plan["class_type"] for plan in scheduled_args[2]["loader_prewarm_plans"]
    ] == ["UNETLoader"]


def test_speculative_affinity_prewarm_rejects_cross_provider_target(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """An SSH continuation must never stage its payload on a Modal worker."""
    submitted_tasks: list[tuple[Any, tuple[Any, ...]]] = []

    class FakeExecutor:
        """Executor double that records unexpected warmup jobs."""

        def submit(self, fn: Any, *args: Any) -> Future[Any]:
            """Capture one scheduled job."""
            submitted_tasks.append((fn, args))
            return Future()

    monkeypatch.setattr(
        remote_modal_app_module,
        "_REMOTE_MODAL_WARMUP_EXECUTOR",
        FakeExecutor(),
    )
    payload = {
        "prompt_id": "prompt-cross-provider",
        "component_id": "modal-component",
        "execution_provider": "modal",
        "execution_environment_id": "modal:H200",
        "speculative_remote_prewarm_target": {
            "prompt_id": "prompt-cross-provider",
            "component_id": "ssh-component",
            "execution_provider": "ssh_docker",
            "execution_environment_id": "lambda",
        },
    }

    assert (
        remote_modal_app_module._schedule_speculative_affinity_prewarm(
            payload,
            reason="provider_boundary",
        )
        is False
    )
    assert submitted_tasks == []


def test_local_gap_keepalive_is_bounded_and_stopped_by_next_remote_component(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """A completed producer should retain its slot only until downstream remote work starts."""
    submitted_tasks: list[tuple[Any, tuple[Any, ...], Future[Any]]] = []

    class FakeExecutor:
        """Executor double that captures but does not run the keepalive loop."""

        def submit(self, fn: Any, *args: Any) -> Future[Any]:
            """Record one keepalive job and return its pending future."""
            future: Future[Any] = Future()
            submitted_tasks.append((fn, args, future))
            return future

    monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "remote")
    monkeypatch.setattr(remote_modal_app_module, "modal", object())
    monkeypatch.setattr(
        remote_modal_app_module,
        "_REMOTE_MODAL_KEEPALIVE_EXECUTOR",
        FakeExecutor(),
    )
    remote_modal_app_module.get_settings.cache_clear()
    with remote_modal_app_module._LOCAL_GAP_KEEPALIVES_LOCK:
        remote_modal_app_module._LOCAL_GAP_KEEPALIVES.clear()
        producer_payload = {
            "prompt_id": "prompt-gap",
            "component_id": "component-a",
            "execution_provider": "modal",
            "execution_environment_id": "modal:RTX-PRO-6000",
            "keepalive_after_remote_component": True,
        }
        consumer_payload = {
            "prompt_id": "prompt-gap",
            "component_id": "component-b",
            "execution_provider": "modal",
            "execution_environment_id": "modal:RTX-PRO-6000",
            "stop_local_gap_keepalive_before_remote_component": True,
    }

    try:
        assert remote_modal_app_module._start_local_gap_keepalive(producer_payload) is True
        assert len(submitted_tasks) == 1
        _fn, args, _future = submitted_tasks[0]
        stop_event = args[1]
        assert args[2:] == (900.0, 15.0)
        assert stop_event.is_set() is False

        assert remote_modal_app_module._stop_local_gap_keepalive(
            consumer_payload,
            reason="test_next_component",
        ) is True
        assert stop_event.is_set() is True
    finally:
        remote_modal_app_module.get_settings.cache_clear()
        with remote_modal_app_module._LOCAL_GAP_KEEPALIVES_LOCK:
            remote_modal_app_module._LOCAL_GAP_KEEPALIVES.clear()


def test_await_prompt_warmup_slots_waits_for_inflight_futures(
    remote_modal_app_module: Any,
) -> None:
    """Mapped execution should be able to wait briefly for already scheduled warmup slots."""
    prompt_id = "prompt-head-start"
    with remote_modal_app_module._PROMPT_WARMUP_STATES_LOCK:
        remote_modal_app_module._PROMPT_WARMUP_STATES.clear()
        remote_modal_app_module._PROMPT_WARMUP_STATE_ORDER = None
        remote_modal_app_module._ensure_prompt_warmup_state(prompt_id)

    future: Future[Any] = Future()
    remote_modal_app_module._track_prompt_warmup_future(prompt_id, 0, future)

    def complete_future() -> None:
        """Complete the synthetic warmup slot after a short delay."""
        time.sleep(0.01)
        future.set_result({"ok": True})

    threading.Thread(target=complete_future, daemon=True).start()

    completed_count = asyncio.run(
        remote_modal_app_module._await_prompt_warmup_slots(
            prompt_id,
            [0],
            0.2,
        )
    )

    assert completed_count == 1


def test_build_prompt_warmup_request_includes_root_loader_prewarm_plans(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Warmup requests should synthesize one-node plans for root literal loader nodes only."""
    monkeypatch.setenv("COMFY_MODAL_ENABLE_LOADER_PREWARM", "true")
    monkeypatch.setenv("COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT", "false")
    remote_modal_app_module.get_settings.cache_clear()
    try:
        warmup_request = remote_modal_app_module._build_prompt_warmup_request(
            {
                "prompt_id": "prompt-1",
                "component_id": "component-1",
                "modal_gpu": "B300",
                "subgraph_prompt": {
                    "1": {
                        "class_type": "UNETLoader",
                        "inputs": {"unet_name": "model-a.safetensors", "weight_dtype": "default"},
                    },
                    "2": {
                        "class_type": "CLIPLoader",
                        "inputs": {"clip_name": "clip-a.safetensors", "type": "stable_diffusion"},
                    },
                    "3": {
                        "class_type": "DualCLIPLoader",
                        "inputs": {"clip_name1": "clip-a.safetensors", "clip_name2": "clip-b.safetensors", "type": "flux"},
                    },
                    "4": {
                        "class_type": "UNETLoader",
                        "inputs": {"unet_name": ["99", 0]},
                    },
                    "5": {
                        "class_type": "KSampler",
                        "inputs": {"model": ["1", 0]},
                    },
                },
            }
        )
    finally:
        remote_modal_app_module.get_settings.cache_clear()

    loader_plans = warmup_request["loader_prewarm_plans"]
    assert warmup_request["modal_gpu"] == "B300"
    assert [plan["node_id"] for plan in loader_plans] == ["1", "2", "3"]
    assert [plan["class_type"] for plan in loader_plans] == [
        "UNETLoader",
        "CLIPLoader",
        "DualCLIPLoader",
    ]
    assert all(plan["execute_node_ids"] == [plan["node_id"]] for plan in loader_plans)


def test_build_prompt_warmup_request_registers_snapshot_profile_when_gpu_snapshots_enabled(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Warmup requests should register one snapshot profile for GPU-snapshot loader plans."""

    class FakeSnapshotProfiles(dict[str, Any]):
        """Minimal modal.Dict shim for snapshot profile writes."""

    snapshot_profiles = FakeSnapshotProfiles()

    class FakeModal:
        """Minimal modal SDK double that exposes Dict.from_name."""

        class Dict:
            """Namespace for fake dict lookups."""

            @staticmethod
            def from_name(dict_name: str, create_if_missing: bool = False) -> Any:
                """Return the shared fake snapshot profile store."""
                assert dict_name == "comfy-modal-sync-snapshot-profiles"
                assert create_if_missing is True
                return snapshot_profiles

    monkeypatch.setattr(remote_modal_app_module, "modal", FakeModal)
    monkeypatch.setenv("COMFY_MODAL_ENABLE_LOADER_PREWARM", "true")
    monkeypatch.setenv("COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT", "true")
    remote_modal_app_module.get_settings.cache_clear()
    remote_modal_app_module._SNAPSHOT_PROFILE_RECORDS.clear()
    try:
        warmup_request = remote_modal_app_module._build_prompt_warmup_request(
            {
                "prompt_id": "prompt-1",
                "component_id": "component-1",
                "subgraph_prompt": {
                    "1": {
                        "class_type": "UNETLoader",
                        "inputs": {"unet_name": "model-a.safetensors", "weight_dtype": "default"},
                    },
                    "2": {
                        "class_type": "CLIPLoader",
                        "inputs": {"clip_name": "clip-a.safetensors", "type": "flux"},
                    },
                },
            }
        )
    finally:
        remote_modal_app_module.get_settings.cache_clear()
        remote_modal_app_module._SNAPSHOT_PROFILE_RECORDS.clear()

    snapshot_profile_key = warmup_request["snapshot_profile_key"]
    assert snapshot_profile_key.startswith("loader-profile:")
    assert snapshot_profile_key in snapshot_profiles
    assert snapshot_profiles[snapshot_profile_key]["loader_prewarm_plans"] == warmup_request["loader_prewarm_plans"]


def test_build_prompt_warmup_request_skips_generic_gpu_snapshot_warmup_without_profile(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Generic proactive warmup should be skipped when GPU snapshots are enabled without a loader profile."""
    monkeypatch.setenv("COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT", "true")
    monkeypatch.setenv("COMFY_MODAL_ENABLE_LOADER_PREWARM", "true")
    remote_modal_app_module.get_settings.cache_clear()
    try:
        warmup_request = remote_modal_app_module._build_prompt_warmup_request(
            {
                "prompt_id": "prompt-1",
                "component_id": "component-1",
                "subgraph_prompt": {
                    "5": {
                        "class_type": "KSampler",
                        "inputs": {"model": ["1", 0]},
                    }
                },
            }
        )
    finally:
        remote_modal_app_module.get_settings.cache_clear()

    assert warmup_request is None


def test_build_prompt_warmup_request_includes_llm_load_and_jit_plan(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """An LLM-only future component should warm even without a Comfy loader profile."""
    monkeypatch.setenv("COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT", "true")
    remote_modal_app_module.get_settings.cache_clear()
    try:
        warmup_request = remote_modal_app_module._build_prompt_warmup_request(
            {
                "prompt_id": "prompt-llm",
                "component_id": "component-llm",
                "remote_worker_affinity_group": "llm",
                "execute_node_ids": ["263"],
                "subgraph_prompt": {
                    "263": {
                        "class_type": "ModalLLM",
                        "inputs": {
                            "model_profile": "qwen-test",
                            "images": ["251", 0],
                        },
                    }
                },
            }
        )
    finally:
        remote_modal_app_module.get_settings.cache_clear()

    assert warmup_request is not None
    assert warmup_request["snapshot_profile_key"] == ""
    assert warmup_request["loader_prewarm_plans"] == []
    assert len(warmup_request["llm_prewarm_plans"]) == 1
    assert warmup_request["llm_prewarm_plans"][0]["model_profile"] == "qwen-test"
    assert warmup_request["llm_prewarm_plans"][0]["representative_request_count"] == 3


def test_build_prompt_warmup_request_ignores_llm_outside_executable_closure(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Disconnected nodes and serialized planning metadata must not warm an LLM."""
    monkeypatch.setenv("COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT", "true")
    remote_modal_app_module.get_settings.cache_clear()
    try:
        warmup_request = remote_modal_app_module._build_prompt_warmup_request(
            {
                "prompt_id": "prompt-comfy",
                "component_id": "251",
                "execute_node_ids": ["251"],
                "subgraph_prompt": {
                    "251": {
                        "class_type": "ImageFromBatch",
                        "inputs": {"image": ["250", 0]},
                    },
                    "250": {
                        "class_type": "VAEDecode",
                        "inputs": {"samples": ["14", 0]},
                    },
                    "263": {
                        "class_type": "ModalLLM",
                        "inputs": {"model_profile": "disconnected-profile"},
                    },
                },
                "extra_data": {
                    "modal": {
                        "future_component": {
                            "class_type": "ModalLLM",
                            "inputs": {"model_profile": "metadata-profile"},
                        }
                    }
                },
            }
        )
    finally:
        remote_modal_app_module.get_settings.cache_clear()

    assert warmup_request is None


def test_dispatch_joins_matching_speculative_warmup(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """A real dispatch should wait on the exact active speculative worker identity."""
    monkeypatch.setenv("COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT", "false")
    remote_modal_app_module.get_settings.cache_clear()
    payload = {
        "prompt_id": "prompt-join",
        "component_id": "component-llm",
        "remote_worker_affinity_group": "llm",
    }
    future: Future[Any] = Future()
    identity = remote_modal_app_module._speculative_warmup_identity(
        {**payload, "gpu_snapshot_variant": "direct"}
    )
    with remote_modal_app_module._PROMPT_WARMUP_STATES_LOCK:
        remote_modal_app_module._PROMPT_WARMUP_STATES.clear()
        state = remote_modal_app_module._ensure_prompt_warmup_state("prompt-join")
        state.speculative_affinity_futures[identity] = future

    timer = threading.Timer(0.03, lambda: future.set_result({"ready": True}))
    timer.start()
    started_at = time.perf_counter()
    try:
        remote_modal_app_module._await_matching_speculative_prewarm(payload, None)
    finally:
        timer.join()
        remote_modal_app_module.get_settings.cache_clear()
        with remote_modal_app_module._PROMPT_WARMUP_STATES_LOCK:
            remote_modal_app_module._PROMPT_WARMUP_STATES.clear()

    assert time.perf_counter() - started_at >= 0.02


def test_snapshot_policy_samples_both_variants_then_selects_faster(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Loader profiles should learn snapshot versus direct startup independently."""
    profile_key = "loader-profile:test"
    store: dict[str, Any] = {
        profile_key: {
            "snapshot_profile_key": profile_key,
            "loader_prewarm_plans": [],
            "snapshot_policy": {
                "selected_variant": None,
                "measurements": {"snapshot": [], "direct": []},
            },
        }
    }

    class FakeModal:
        """Expose the shared policy dict through Modal's named-object shape."""

        class Dict:
            """Return the in-memory profile store."""

            @staticmethod
            def from_name(name: str, create_if_missing: bool = False) -> Any:
                """Return the fake store."""
                del name, create_if_missing
                return store

    monkeypatch.setattr(remote_modal_app_module, "modal", FakeModal)
    monkeypatch.setenv("COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT", "true")
    remote_modal_app_module.get_settings.cache_clear()
    remote_modal_app_module._SNAPSHOT_PROFILE_RECORDS.clear()
    try:
        snapshot_payload = {"snapshot_profile_key": profile_key}
        assert remote_modal_app_module._select_gpu_snapshot_for_profile(
            snapshot_payload, profile_key
        ) is True
        remote_modal_app_module._record_snapshot_warmup_measurement(
            snapshot_payload, 80.0
        )
        direct_payload = {"snapshot_profile_key": profile_key}
        assert remote_modal_app_module._select_gpu_snapshot_for_profile(
            direct_payload, profile_key
        ) is False
        remote_modal_app_module._record_snapshot_warmup_measurement(
            direct_payload, 50.0
        )
        selected_payload = {"snapshot_profile_key": profile_key}
        assert remote_modal_app_module._select_gpu_snapshot_for_profile(
            selected_payload, profile_key
        ) is False
    finally:
        remote_modal_app_module.get_settings.cache_clear()
        remote_modal_app_module._SNAPSHOT_PROFILE_RECORDS.clear()

    assert store[profile_key]["snapshot_policy"]["selected_variant"] == "direct"


def test_post_deploy_seed_registers_a_joinable_profile_future(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Automatic deployment should seed the current important profile before dispatch."""
    submitted_tasks: list[tuple[Any, tuple[Any, ...]]] = []

    class FakeExecutor:
        """Capture the automatic seed without invoking Modal."""

        def submit(self, function: Any, *args: Any) -> Future[Any]:
            """Record one seed job and return its pending future."""
            submitted_tasks.append((function, args))
            return Future()

    monkeypatch.setattr(remote_modal_app_module, "modal", object())
    monkeypatch.setattr(
        remote_modal_app_module,
        "_REMOTE_MODAL_WARMUP_EXECUTOR",
        FakeExecutor(),
    )
    monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "remote")
    monkeypatch.setenv("COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT", "false")
    remote_modal_app_module.get_settings.cache_clear()
    with remote_modal_app_module._PROMPT_WARMUP_STATES_LOCK:
        remote_modal_app_module._PROMPT_WARMUP_STATES.clear()
    try:
        scheduled = remote_modal_app_module._schedule_post_deploy_runtime_seed(
            {
                "prompt_id": "prompt-deploy",
                "component_id": "llm-component",
                "remote_worker_affinity_group": "llm",
                "subgraph_prompt": {
                    "263": {
                        "class_type": "ModalLLM",
                        "inputs": {"model_profile": "qwen-test"},
                    }
                },
            }
        )
    finally:
        remote_modal_app_module.get_settings.cache_clear()

    assert scheduled is True
    assert len(submitted_tasks) == 1
    scheduled_function, scheduled_args = submitted_tasks[0]
    assert scheduled_function is remote_modal_app_module._run_speculative_affinity_prewarm
    assert scheduled_args[-1] == "post_deploy_runtime_seed"
    with remote_modal_app_module._PROMPT_WARMUP_STATES_LOCK:
        state = remote_modal_app_module._PROMPT_WARMUP_STATES["prompt-deploy"]
        assert len(state.speculative_affinity_futures) == 1
        remote_modal_app_module._PROMPT_WARMUP_STATES.clear()


def test_register_exact_component_parallelism_refines_prompt_target(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Mapped fan-out should raise the prompt-wide warmup target once exact item count is known."""
    monkeypatch.setenv("COMFY_MODAL_MAX_CONTAINERS", "6")
    remote_modal_app_module.get_settings.cache_clear()
    with remote_modal_app_module._PROMPT_WARMUP_STATES_LOCK:
        remote_modal_app_module._PROMPT_WARMUP_STATES.clear()
        remote_modal_app_module._PROMPT_WARMUP_STATE_ORDER = None

    try:
        payload = {
            "prompt_id": "prompt-2",
            "component_id": "component-a",
            "extra_data": {
                "modal": {
                    "component_execution_stages": [["component-a", "component-b"], ["component-c"]],
                    "mapped_component_ids": ["component-a"],
                    "estimated_max_parallel_requests": 2,
                }
            },
        }
        refined_target = remote_modal_app_module._register_exact_component_parallelism(payload, 5)
    finally:
        remote_modal_app_module.get_settings.cache_clear()

    assert refined_target == 6


def test_stable_modal_cloud_entry_imports_without_modal_sdk(
    modal_cloud_module: Any,
) -> None:
    """The stable Modal cloud module should stay importable when modal is unavailable."""
    assert modal_cloud_module.__name__ == "comfyui_modal_sync_cloud"
    assert hasattr(modal_cloud_module, "RemoteEngine")


def test_modal_cloud_exposes_affinity_parameter_and_local_gap_keepalive() -> None:
    """The deployed class must make affinity part of identity and expose a heartbeat."""
    source = (Path(__file__).resolve().parents[1] / "comfyui_modal_sync_cloud.py").read_text()

    assert "worker_affinity_key: str = modal.parameter" in source
    assert "def keepalive_for_local_gap" in source


def test_modal_cloud_image_environment_preserves_unique_app_name(
    modal_cloud_module: Any,
) -> None:
    """Remote workers should receive the same per-ComfyUI app name resolved locally."""
    settings = types.SimpleNamespace(
        app_name="comfy-modal-sync-AAECAwQFBgc",
        modal_gpu="B300",
        modal_secret_name="workflow-credentials",
        stream_event_queue_maxsize=256,
        bridge_inline_max_bytes=1024,
        invocation_result_inline_max_bytes=2048,
        execution_timeout_seconds=3600,
        startup_timeout_seconds=900,
        llm_vllm_execution_mode="throughput",
    )

    image_environment = modal_cloud_module._modal_image_environment(settings, "fingerprint-1")

    assert image_environment["COMFY_MODAL_APP_NAME"] == "comfy-modal-sync-AAECAwQFBgc"
    assert image_environment["COMFY_MODAL_GPU"] == "B300"
    assert image_environment["COMFY_MODAL_REMOTE_STORAGE_ROOT"] == "/storage"
    assert image_environment["COMFY_MODAL_REMOTE_WORKER"] == "1"
    assert image_environment["COMFY_MODAL_LLM_VLLM_EXECUTION_MODE"] == "throughput"
    assert image_environment["VLLM_CACHE_ROOT"].startswith(
        "/root/.cache/comfy-modal-llm/"
    )
    assert image_environment["TORCHINDUCTOR_CACHE_DIR"].startswith(
        "/root/.cache/comfy-modal-llm/"
    )
    assert image_environment["VLLM_USE_FLASHINFER_SAMPLER"] == "0"
    assert image_environment["COMFY_MODAL_LLM_MAX_RESIDENT_MODELS"] == "2"
    assert image_environment["COMFY_MODAL_LLM_MEMORY_RECOVERY_TIMEOUT_SECONDS"] == "15.0"
    assert image_environment["COMFY_MODAL_LLM_RESERVE_FREE_GB"] == "24.0"
    assert image_environment["COMFY_MODAL_SECRET_NAME"] == "workflow-credentials"
    assert image_environment["COMFY_MODAL_RUNTIME_FINGERPRINT"] == "fingerprint-1"


def test_model_stager_image_environment_preserves_deployment_identity(
    modal_cloud_module: Any,
) -> None:
    """The CPU protocol helper must report the exact deployed GPU runtime identity."""
    settings = types.SimpleNamespace(
        app_name="comfy-modal-sync-instance",
        modal_gpu="RTX-PRO-6000",
        remote_storage_root="/storage",
    )

    image_environment = modal_cloud_module._model_stager_image_environment(
        settings,
        "fingerprint-2",
    )

    assert image_environment == {
        "COMFY_MODAL_APP_NAME": "comfy-modal-sync-instance",
        "COMFY_MODAL_GPU": "RTX-PRO-6000",
        "COMFY_MODAL_REMOTE_STORAGE_ROOT": "/storage",
        "COMFY_MODAL_RUNTIME_FINGERPRINT": "fingerprint-2",
        "HF_HUB_DISABLE_TELEMETRY": "1",
    }


def test_modal_cloud_observes_distinct_workflows_for_llm_auto_mode(
    modal_cloud_module: Any,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """The container boundary should count prompts while excluding canary traffic."""
    observed_prompt_ids: list[str | None] = []
    forced_prompt_ids: list[str | None] = []
    runtime_module = types.ModuleType("modal_llm_runtime")
    runtime_module.observe_modal_workflow_execution = observed_prompt_ids.append
    runtime_module.force_modal_vllm_throughput_after_memory_recovery = (
        forced_prompt_ids.append
    )
    monkeypatch.setitem(sys.modules, "modal_llm_runtime", runtime_module)

    modal_cloud_module._observe_remote_workflow_for_llm_mode(
        {"prompt_id": "prompt-1", "payload_kind": "subgraph"}
    )
    modal_cloud_module._observe_remote_workflow_for_llm_mode(
        {"prompt_id": "prompt-1", "payload_kind": "mapped_subgraph"}
    )
    modal_cloud_module._observe_remote_workflow_for_llm_mode(
        {"prompt_id": "health-check", "payload_kind": "canary"}
    )
    modal_cloud_module._observe_remote_workflow_for_llm_mode(
        {"prompt_id": "prompt-2", "payload_kind": "subgraph"}
    )
    modal_cloud_module._observe_remote_workflow_for_llm_mode(
        {
            "prompt_id": "prompt-3",
            "payload_kind": "subgraph",
            "force_vllm_throughput_after_memory_recovery": True,
        }
    )

    assert observed_prompt_ids == ["prompt-1", "prompt-1", "prompt-2"]
    assert forced_prompt_ids == ["prompt-3"]


def test_modal_cloud_resolves_configured_named_secret(
    modal_cloud_module: Any,
) -> None:
    """Cloud app construction should reference the configured collection without reading .env."""
    observed_names: list[str] = []
    expected_secret = object()
    fake_modal = types.SimpleNamespace(
        Secret=types.SimpleNamespace(
            from_name=lambda name: observed_names.append(name) or expected_secret,
        )
    )

    resolved_secret = modal_cloud_module._modal_secret_from_settings(
        types.SimpleNamespace(modal_secret_name="workflow-credentials"),
        fake_modal,
    )

    assert resolved_secret is expected_secret
    assert observed_names == ["workflow-credentials"]


def test_modal_cloud_installs_timestamped_logger_handler(
    modal_cloud_module: Any,
) -> None:
    """The cloud runtime should install its own timestamped logger handler."""
    matching_handlers = [
        handler
        for handler in modal_cloud_module.logger.handlers
        if getattr(handler, "name", "") == modal_cloud_module._CLOUD_HANDLER_NAME
    ]

    assert len(matching_handlers) == 1
    assert modal_cloud_module.logger.propagate is False
    assert modal_cloud_module.logger.level == logging.INFO
    assert matching_handlers[0].stream is sys.stdout
    formatter = matching_handlers[0].formatter
    assert isinstance(formatter, logging.Formatter)
    assert "%(asctime)s" in formatter._fmt
    assert "%(relativeCreated)" in formatter._fmt


def test_modal_cloud_mirrors_phase_logs_to_stdout_in_modal_runtime(
    modal_cloud_module: Any,
    monkeypatch: Any,
    capsys: Any,
) -> None:
    """Timed cloud phases should write directly to stdout inside Modal containers."""
    monkeypatch.setenv("MODAL_IS_REMOTE", "1")

    with modal_cloud_module._timed_phase("phase_under_test", component="component-1"):
        pass

    captured = capsys.readouterr()
    assert "Starting phase_under_test component=component-1" in captured.out
    assert "Finished phase_under_test in " in captured.out
    assert "component=component-1" in captured.out


def test_modal_cloud_reuses_extracted_custom_nodes_bundle(
    modal_cloud_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """The remote runtime should avoid re-extracting an unchanged custom_nodes bundle."""
    storage_root = tmp_path / "storage"
    bundle_path = storage_root / "custom_nodes" / "bundle.zip"
    bundle_path.parent.mkdir(parents=True)

    import zipfile

    with zipfile.ZipFile(bundle_path, "w") as archive:
        archive.writestr("example/__init__.py", "NODE_CLASS_MAPPINGS = {}\n")

    monkeypatch.setenv("COMFY_MODAL_REMOTE_STORAGE_ROOT", str(storage_root))
    modal_cloud_module.get_settings.cache_clear()
    original_cache = dict(modal_cloud_module._EXTRACTED_CUSTOM_NODE_BUNDLES)
    modal_cloud_module._EXTRACTED_CUSTOM_NODE_BUNDLES.clear()
    try:
        first_root = modal_cloud_module._extract_custom_nodes_bundle("/custom_nodes/bundle.zip")
        monkeypatch.setattr(
            modal_cloud_module.zipfile,
            "ZipFile",
            lambda *args, **kwargs: (_ for _ in ()).throw(
                AssertionError("Expected cached extraction root to be reused.")
            ),
        )
        second_root = modal_cloud_module._extract_custom_nodes_bundle("/custom_nodes/bundle.zip")
    finally:
        modal_cloud_module.get_settings.cache_clear()
        modal_cloud_module._EXTRACTED_CUSTOM_NODE_BUNDLES.clear()
        modal_cloud_module._EXTRACTED_CUSTOM_NODE_BUNDLES.update(original_cache)

    assert first_root is not None
    assert second_root == first_root


def test_modal_cloud_extracts_custom_nodes_manifest_with_multiple_archives(
    modal_cloud_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """The remote runtime should resolve manifest-based custom_nodes bundles into several archives."""
    storage_root = tmp_path / "storage"
    entry_a_path = storage_root / "custom_nodes" / "entries" / "example_a" / "hash_a_bundle.zip"
    entry_b_path = storage_root / "custom_nodes" / "entries" / "example_b" / "hash_b_bundle.zip"
    manifest_path = storage_root / "custom_nodes" / "manifests" / "bundle_manifest.json"
    entry_a_path.parent.mkdir(parents=True, exist_ok=True)
    entry_b_path.parent.mkdir(parents=True, exist_ok=True)
    manifest_path.parent.mkdir(parents=True, exist_ok=True)

    import zipfile

    with zipfile.ZipFile(entry_a_path, "w") as archive:
        archive.writestr("example_a/__init__.py", "NODE_CLASS_MAPPINGS = {}\n")
    with zipfile.ZipFile(entry_b_path, "w") as archive:
        archive.writestr("example_b/__init__.py", "NODE_CLASS_MAPPINGS = {}\n")
    manifest_path.write_text(
        json.dumps(
            {
                "version": 1,
                "bundle_sha256": "bundle-hash",
                "entries": [
                    {
                        "entry_name": "example_a",
                        "display_name": "example_a",
                        "sha256": "hash-a",
                        "remote_path": "/custom_nodes/entries/example_a/hash_a_bundle.zip",
                    },
                    {
                        "entry_name": "example_b",
                        "display_name": "example_b",
                        "sha256": "hash-b",
                        "remote_path": "/custom_nodes/entries/example_b/hash_b_bundle.zip",
                    },
                ],
            },
            sort_keys=True,
        ),
        encoding="utf-8",
    )

    monkeypatch.setenv("COMFY_MODAL_REMOTE_STORAGE_ROOT", str(storage_root))
    modal_cloud_module.get_settings.cache_clear()
    original_cache = dict(modal_cloud_module._EXTRACTED_CUSTOM_NODE_BUNDLES)
    modal_cloud_module._EXTRACTED_CUSTOM_NODE_BUNDLES.clear()
    try:
        extraction_root = modal_cloud_module._extract_custom_nodes_bundle(
            "/custom_nodes/manifests/bundle_manifest.json"
        )
    finally:
        modal_cloud_module.get_settings.cache_clear()
        modal_cloud_module._EXTRACTED_CUSTOM_NODE_BUNDLES.clear()
        modal_cloud_module._EXTRACTED_CUSTOM_NODE_BUNDLES.update(original_cache)

    assert extraction_root is not None
    assert (extraction_root / "example_a" / "__init__.py").exists()
    assert (extraction_root / "example_b" / "__init__.py").exists()


def test_modal_cloud_links_version_two_custom_node_assets(
    modal_cloud_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """Version-two manifests should keep model assets on mounted storage."""
    storage_root = tmp_path / "storage"
    asset_sha256 = "a" * 64
    archive_path = storage_root / "custom_nodes/entries/example/code_bundle.zip"
    asset_path = storage_root / f"custom_nodes/assets/example/{asset_sha256}_model.pth"
    manifest_path = storage_root / "custom_nodes/manifests/bundle_manifest_v2.json"
    archive_path.parent.mkdir(parents=True)
    asset_path.parent.mkdir(parents=True)
    manifest_path.parent.mkdir(parents=True)
    asset_path.write_bytes(b"mounted-model")
    import zipfile

    with zipfile.ZipFile(archive_path, "w") as archive:
        archive.writestr("example/__init__.py", "NODE_CLASS_MAPPINGS = {}\n")
    manifest_path.write_text(
        json.dumps(
            {
                "version": 2,
                "bundle_sha256": "bundle-hash",
                "entries": [
                    {
                        "entry_name": "example",
                        "display_name": "example",
                        "sha256": "code-hash",
                        "remote_path": "/custom_nodes/entries/example/code_bundle.zip",
                        "assets": [
                            {
                                "relative_path": "example/checkpoints/model.pth",
                                "sha256": asset_sha256,
                                "size_bytes": len(b"mounted-model"),
                                "remote_path": (
                                    f"/custom_nodes/assets/example/{asset_sha256}_model.pth"
                                ),
                            }
                        ],
                    }
                ],
            },
            sort_keys=True,
        ),
        encoding="utf-8",
    )

    monkeypatch.setenv("COMFY_MODAL_REMOTE_STORAGE_ROOT", str(storage_root))
    modal_cloud_module.get_settings.cache_clear()
    original_cache = dict(modal_cloud_module._EXTRACTED_CUSTOM_NODE_BUNDLES)
    modal_cloud_module._EXTRACTED_CUSTOM_NODE_BUNDLES.clear()
    try:
        extraction_root = modal_cloud_module._extract_custom_nodes_bundle(
            "/custom_nodes/manifests/bundle_manifest_v2.json"
        )
    finally:
        modal_cloud_module.get_settings.cache_clear()
        modal_cloud_module._EXTRACTED_CUSTOM_NODE_BUNDLES.clear()
        modal_cloud_module._EXTRACTED_CUSTOM_NODE_BUNDLES.update(original_cache)

    assert extraction_root is not None
    linked_asset = extraction_root / "example/checkpoints/model.pth"
    assert linked_asset.is_symlink()
    assert linked_asset.read_bytes() == b"mounted-model"


def test_local_fallback_links_version_two_custom_node_assets(
    remote_modal_app_module: Any,
    tmp_path: Path,
) -> None:
    """The non-Modal extractor should preserve version-two package asset links."""
    storage_root = tmp_path / "storage"
    extraction_root = tmp_path / "extracted"
    asset_sha256 = "b" * 64
    asset_path = storage_root / f"custom_nodes/assets/example/{asset_sha256}_model.pth"
    manifest_path = storage_root / "custom_nodes/manifests/bundle_manifest_v2.json"
    asset_path.parent.mkdir(parents=True)
    manifest_path.parent.mkdir(parents=True)
    extraction_root.mkdir()
    asset_path.write_bytes(b"local-mounted-model")
    manifest_path.write_text(
        json.dumps(
            {
                "version": 2,
                "entries": [
                    {
                        "entry_name": "example",
                        "assets": [
                            {
                                "relative_path": "example/checkpoints/model.pth",
                                "sha256": asset_sha256,
                                "size_bytes": len(b"local-mounted-model"),
                                "remote_path": (
                                    f"/custom_nodes/assets/example/{asset_sha256}_model.pth"
                                ),
                            }
                        ],
                    }
                ],
            },
            sort_keys=True,
        ),
        encoding="utf-8",
    )

    remote_modal_app_module._materialize_local_custom_node_assets(
        manifest_path,
        storage_root,
        extraction_root,
    )

    linked_asset = extraction_root / "example/checkpoints/model.pth"
    assert linked_asset.is_symlink()
    assert linked_asset.read_bytes() == b"local-mounted-model"


def test_modal_cloud_traces_remote_node_execution_spans(
    modal_cloud_module: Any,
    monkeypatch: Any,
    capsys: Any,
) -> None:
    """The tracing prompt server should emit per-node timing lines."""
    prompt = {
        "7": {"class_type": "UNETLoader", "inputs": {}},
        "2": {"class_type": "KSampler", "inputs": {}},
    }
    monkeypatch.setenv("MODAL_IS_REMOTE", "1")
    server = modal_cloud_module._TracingPromptServer("component-1", prompt)

    server.send_sync("executing", {"node": "7"}, None)
    server.send_sync("executed", {"node": "7"}, None)
    server.send_sync("executing", {"node": "2"}, None)
    server.send_sync("execution_success", {"prompt_id": "component-1"}, None)

    captured = capsys.readouterr()
    assert "Remote node 7 class_type=UNETLoader role=model_load started" in captured.out
    assert "Remote node 7 class_type=UNETLoader role=model_load finished in " in captured.out
    assert "Remote node 2 class_type=KSampler role=sampling started" in captured.out
    assert "Remote node 2 class_type=KSampler role=sampling finished in " in captured.out


def test_modal_cloud_installs_headless_prompt_server_instance(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Remote custom-node init should get a minimal PromptServer.instance shim."""
    class FakeNodeReplaceManager:
        """Record replacement registrations made through the headless server."""

        def __init__(self) -> None:
            """Initialize an empty replacement registry."""
            self.registrations: list[Any] = []

        def register(self, replacement: Any) -> None:
            """Record one replacement registered by a ComfyUI extension."""
            self.registrations.append(replacement)

    fake_prompt_server_class = type("PromptServer", (), {})
    fake_server_module = types.SimpleNamespace(
        NodeReplaceManager=FakeNodeReplaceManager,
        PromptServer=fake_prompt_server_class,
    )

    monkeypatch.setitem(sys.modules, "server", fake_server_module)
    modal_cloud_module._ensure_headless_prompt_server_instance()

    instance = fake_prompt_server_class.instance
    assert instance is not None
    assert hasattr(instance, "routes")
    assert hasattr(instance, "app")
    assert isinstance(instance.node_replace_manager, FakeNodeReplaceManager)
    assert instance.supports == ["custom_nodes_from_web"]
    assert instance.client_id is None
    assert instance.last_node_id is None
    assert instance.number == 0
    assert instance.prompt_queue.currently_running == {}
    assert instance.prompt_queue.get_current_queue() == ([], [])

    instance.send_progress_text("width: 1024, height: 768", "104")

    instance.add_on_prompt_handler("handler")
    assert instance.on_prompt_handlers == ["handler"]

    instance.node_replace_manager.register("replacement")
    assert instance.node_replace_manager.registrations == ["replacement"]

    instance.prompt_queue.put((1, "prompt-id", {}, {}, []))
    assert instance.prompt_queue.get_current_queue() == (
        [],
        [(1, "prompt-id", {}, {}, [])],
    )
    assert instance.prompt_queue.get_tasks_remaining() == 1

    instance.prompt_queue.set_flag("free_memory", True)
    assert instance.prompt_queue.get_flags() == {"free_memory": True}
    assert instance.prompt_queue.get_flags() == {}


def test_modal_cloud_dynamic_prompt_preserves_thread_local_metadata_graph(
    modal_cloud_module: Any,
) -> None:
    """Hydrated execution graphs should not replace the JSON-safe hidden PROMPT graph."""

    class FakeDynamicPrompt:
        """Minimal ComfyUI DynamicPrompt stand-in."""

        def __init__(self, original_prompt: dict[str, Any]) -> None:
            """Retain the graph used for runtime dependency resolution."""
            self.original_prompt = original_prompt

        def get_original_prompt(self) -> dict[str, Any]:
            """Return the runtime graph when no metadata override is active."""
            return self.original_prompt

    fake_execution = types.SimpleNamespace(DynamicPrompt=FakeDynamicPrompt)
    hydrated_prompt = {"1": {"inputs": {"image": object()}}}
    metadata_prompt = {"1": {"inputs": {"image": ["upstream", 0]}}}

    modal_cloud_module._install_metadata_safe_dynamic_prompt_wrapper(fake_execution)
    with modal_cloud_module._temporary_prompt_metadata(metadata_prompt):
        dynamic_prompt = fake_execution.DynamicPrompt(hydrated_prompt)

    ordinary_prompt = fake_execution.DynamicPrompt(hydrated_prompt)

    assert dynamic_prompt.original_prompt is hydrated_prompt
    assert dynamic_prompt.get_original_prompt() is metadata_prompt
    assert ordinary_prompt.get_original_prompt() is hydrated_prompt


def test_modal_cloud_hidden_prompt_stays_json_safe_after_tensor_boundary_hydration(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Metadata-writing output nodes should receive links, not hydrated tensors, in PROMPT."""
    import torch

    monkeypatch.setitem(sys.modules, "torchsde", types.ModuleType("torchsde"))
    import comfy_execution.caching as comfy_caching

    modal_cloud_module._ensure_comfy_runtime_initialized(None)
    nodes_module = modal_cloud_module._load_nodes_module()
    monkeypatch.setitem(
        nodes_module.NODE_CLASS_MAPPINGS,
        "PromptMetadataSerializationNode",
        _PromptMetadataSerializationNode,
    )
    monkeypatch.setitem(
        comfy_caching.nodes.NODE_CLASS_MAPPINGS,
        "PromptMetadataSerializationNode",
        _PromptMetadataSerializationNode,
    )
    monkeypatch.setitem(
        nodes_module.NODE_DISPLAY_NAME_MAPPINGS,
        "PromptMetadataSerializationNode",
        "PromptMetadataSerializationNode",
    )
    monkeypatch.setattr(modal_cloud_module, "_node_output_cache_store", lambda: None)
    payload = {
        "payload_kind": "subgraph",
        "component_id": "metadata-output",
        "prompt_id": "prompt-metadata-output",
        "component_node_ids": ["output"],
        "subgraph_prompt": {
            "output": {
                "class_type": "PromptMetadataSerializationNode",
                "inputs": {"image": ["upstream", 0]},
            }
        },
        "boundary_inputs": [
            {
                "proxy_input_name": "remote_image",
                "io_type": "IMAGE",
                "source_signature": "image-signature",
                "targets": [{"node_id": "output", "input_name": "image"}],
            }
        ],
        "boundary_outputs": [
            {
                "proxy_output_name": "metadata",
                "node_id": "output",
                "output_index": 0,
                "io_type": "STRING",
                "is_list": False,
            }
        ],
        "execute_node_ids": ["output"],
        "extra_data": {},
    }

    outputs = modal_cloud_module._execute_subgraph_prompt(
        payload,
        {"remote_image": torch.zeros((1, 2, 2, 3), dtype=torch.float32)},
        None,
    )

    assert len(outputs) == 1
    assert json.loads(outputs[0]) == payload["subgraph_prompt"]


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(
                {
                    "event_type": "node_progress",
                    "node_id": "7",
                    "display_node_id": "7",
                    "value": 3,
                    "max": 10,
                }
            )
            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",
            "event_type": "node_progress",
            "node_id": "7",
            "display_node_id": "7",
            "value": 3,
            "max": 10,
        },
        {
            "kind": "progress",
            "phase": "execution_success",
        },
        {
            "kind": "result",
            "outputs": b"serialized-outputs",
        },
    ]


def test_modal_cloud_streams_remote_log_task_id_before_progress(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Remote streamed payloads should surface the current Modal task id for local log mirroring."""

    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"})
        return b"serialized-outputs"

    monkeypatch.setenv("MODAL_TASK_ID", "ta-remote-123")
    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 events[0] == {"kind": "remote_logs", "task_id": "ta-remote-123"}
    assert events[1:] == [
        {
            "kind": "progress",
            "phase": "executing",
            "active_node_id": "7",
        },
        {
            "kind": "result",
            "outputs": b"serialized-outputs",
        },
    ]


def test_modal_cloud_streams_tensor_safe_progress_and_result_events(
    modal_cloud_module: Any,
    monkeypatch: Any,
    serialization_module: Any,
) -> None:
    """Streamed Modal events should serialize stray tensor payloads before yielding them."""
    torch = pytest.importorskip("torch")
    tensor = torch.arange(4, dtype=torch.float32).reshape(2, 2)

    def fake_execute_subgraph_locally(
        payload: dict[str, Any],
        kwargs_payload: bytes,
        status_callback: Any = None,
    ) -> tuple[Any, ...]:
        if status_callback is not None:
            status_callback(
                {
                    "phase": "executing",
                    "active_node_id": "7",
                    "preview_tensor": tensor,
                }
            )
        return (tensor,)

    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 events[0]["kind"] == "progress"
    assert events[0]["phase"] == "executing"
    assert torch.equal(
        serialization_module.deserialize_value(events[0]["preview_tensor"]),
        tensor,
    )

    assert events[1]["kind"] == "result"
    decoded_outputs = serialization_module.deserialize_node_outputs(events[1]["outputs"])
    assert len(decoded_outputs) == 1
    assert torch.equal(decoded_outputs[0], tensor)


def test_modal_cloud_stream_buffer_coalesces_progress_and_preserves_terminal_events(
    modal_cloud_module: Any,
) -> None:
    """A slow stream consumer should have bounded progress memory without losing results."""
    event_buffer = modal_cloud_module._BoundedStreamEventBuffer(maxsize=4)
    for value in range(100):
        event_buffer.publish_progress({"value": value})

    terminal_thread = threading.Thread(
        target=lambda: (
            event_buffer.publish_terminal("result", b"outputs"),
            event_buffer.publish_terminal("done", None),
        )
    )
    terminal_thread.start()
    observed_events: list[tuple[str, Any]] = []
    while True:
        event = event_buffer.get()
        observed_events.append(event)
        if event[0] == "done":
            break
    terminal_thread.join(timeout=1.0)
    event_buffer.close()

    progress_values = [payload["value"] for kind, payload in observed_events if kind == "progress"]
    assert progress_values == [96, 97, 98, 99]
    assert observed_events[-2:] == [("result", b"outputs"), ("done", None)]
    assert event_buffer.dropped_progress_events == 96
    assert not terminal_thread.is_alive()


def test_modal_cloud_only_reloads_volume_for_requests_with_new_uploads(
    modal_cloud_module: Any,
) -> None:
    """Steady-state requests should skip Modal volume reload when queue-time sync uploaded nothing."""

    assert modal_cloud_module._should_reload_modal_volume({"requires_volume_reload": True}) is True
    assert modal_cloud_module._should_reload_modal_volume({"requires_volume_reload": False}) is False
    assert modal_cloud_module._should_reload_modal_volume({}) is True


def test_modal_cloud_skips_reload_when_uploaded_paths_are_already_visible(
    modal_cloud_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """Visible immutable uploaded paths should not force an extra Modal volume reload."""
    storage_root = tmp_path / "storage"
    storage_root.mkdir()
    uploaded_path = storage_root / "assets" / "hash_model.safetensors"
    uploaded_path.parent.mkdir(parents=True, exist_ok=True)
    uploaded_path.write_bytes(b"weights")
    recorded_markers: list[str] = []

    monkeypatch.setattr(
        modal_cloud_module,
        "get_settings",
        lambda: types.SimpleNamespace(remote_storage_root=str(storage_root)),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_record_modal_volume_reload_marker",
        lambda marker: recorded_markers.append(marker),
    )

    payload = {
        "requires_volume_reload": True,
        "volume_reload_marker": "marker-1",
        "uploaded_volume_paths": ["/assets/hash_model.safetensors"],
    }

    assert modal_cloud_module._should_reload_modal_volume(payload) is False
    assert recorded_markers == ["marker-1"]


def test_modal_cloud_reads_missing_committed_asset_without_reloading_volume(
    modal_cloud_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """A stale warm mount should read a committed asset into ephemeral storage."""

    class FakeVolume:
        """Modal Volume double that exposes committed file reads."""

        def __init__(self) -> None:
            """Record direct read paths."""
            self.read_paths: list[str] = []

        def read_file(self, path: str) -> Iterator[bytes]:
            """Yield one committed model file in chunks."""
            self.read_paths.append(path)
            yield b"model-"
            yield b"weights"

    storage_root = tmp_path / "storage"
    storage_root.mkdir()
    readthrough_root = tmp_path / "readthrough"
    monkeypatch.setattr(
        modal_cloud_module,
        "get_settings",
        lambda: types.SimpleNamespace(remote_storage_root=str(storage_root)),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_REMOTE_VOLUME_READTHROUGH_ROOT",
        readthrough_root,
    )
    payload = {
        "component_id": "component-1",
        "requires_volume_reload": False,
        "subgraph_prompt": {
            "14": {
                "class_type": "LoraLoaderModelOnly",
                "inputs": {"lora_name": "/assets/hash_style.safetensors"},
            }
        },
    }

    volume = FakeVolume()
    hydrated_paths = modal_cloud_module._hydrate_missing_payload_volume_paths(volume, payload)

    cached_path = readthrough_root / "assets" / "hash_style.safetensors"
    assert hydrated_paths == [cached_path]
    assert cached_path.read_bytes() == b"model-weights"
    assert volume.read_paths == ["assets/hash_style.safetensors"]
    assert modal_cloud_module._should_reload_modal_volume(payload) is False
    assert (
        modal_cloud_module._resolve_runtime_asset_path("/assets/hash_style.safetensors")
        == str(cached_path)
    )


def test_modal_cloud_reloads_when_reused_referenced_asset_is_still_missing(
    modal_cloud_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """A sync-index hit must not suppress recovery for a missing runtime asset."""
    storage_root = tmp_path / "storage"
    storage_root.mkdir()
    monkeypatch.setattr(
        modal_cloud_module,
        "get_settings",
        lambda: types.SimpleNamespace(remote_storage_root=str(storage_root)),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_REMOTE_VOLUME_READTHROUGH_ROOT",
        tmp_path / "readthrough",
    )
    payload = {
        "requires_volume_reload": False,
        "subgraph_prompt": {
            "14": {
                "class_type": "LoraLoaderModelOnly",
                "inputs": {"lora_name": "/assets/missing_style.safetensors"},
            }
        },
    }

    assert modal_cloud_module._should_reload_modal_volume(payload) is True


def test_modal_cloud_readthrough_hydrates_custom_node_manifest_dependencies(
    modal_cloud_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """Read-through hydration should include archives and assets named by a manifest."""
    manifest_payload = json.dumps(
        {
            "version": 2,
            "entries": [
                {
                    "remote_path": "/custom_nodes/packages/hash_package.zip",
                    "assets": [
                        {"remote_path": "/custom_nodes/assets/hash_model.safetensors"}
                    ],
                }
            ],
        }
    ).encode()
    committed_files = {
        "custom_nodes/manifests/hash_manifest.json": manifest_payload,
        "custom_nodes/packages/hash_package.zip": b"package",
        "custom_nodes/assets/hash_model.safetensors": b"model",
    }

    class FakeVolume:
        """Modal Volume double backed by committed in-memory files."""

        def read_file(self, path: str) -> Iterator[bytes]:
            """Yield the requested committed file."""
            yield committed_files[path]

    storage_root = tmp_path / "storage"
    storage_root.mkdir()
    readthrough_root = tmp_path / "readthrough"
    monkeypatch.setattr(
        modal_cloud_module,
        "get_settings",
        lambda: types.SimpleNamespace(remote_storage_root=str(storage_root)),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_REMOTE_VOLUME_READTHROUGH_ROOT",
        readthrough_root,
    )

    hydrated_paths = modal_cloud_module._hydrate_missing_payload_volume_paths(
        FakeVolume(),
        {"custom_nodes_bundle": "/custom_nodes/manifests/hash_manifest.json"},
    )

    assert hydrated_paths == [
        readthrough_root / "custom_nodes" / "manifests" / "hash_manifest.json",
        readthrough_root / "custom_nodes" / "assets" / "hash_model.safetensors",
        readthrough_root / "custom_nodes" / "packages" / "hash_package.zip",
    ]


def test_modal_cloud_preserves_worker_for_deterministic_remote_failures(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Ordinary node failures should not discard a healthy warm Modal worker."""
    scheduled_exits: list[tuple[float, int]] = []
    original_flag = modal_cloud_module._CONTAINER_TERMINATION_SCHEDULED
    monkeypatch.setattr(modal_cloud_module, "_CONTAINER_TERMINATION_SCHEDULED", False)
    monkeypatch.setattr(modal_cloud_module, "_is_modal_container_runtime", lambda: True)
    monkeypatch.setattr(
        modal_cloud_module,
        "_schedule_process_exit",
        lambda delay_seconds, exit_code: scheduled_exits.append((delay_seconds, exit_code)),
    )
    try:
        scheduled = modal_cloud_module._maybe_schedule_container_termination_on_error(
            {"component_id": "component-1", "terminate_container_on_error": True},
            RuntimeError("boom"),
        )
    finally:
        monkeypatch.setattr(modal_cloud_module, "_CONTAINER_TERMINATION_SCHEDULED", original_flag)

    assert scheduled is False
    assert scheduled_exits == []


def test_modal_cloud_retires_worker_for_poisoned_cuda_runtime(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """CUDA failures that can poison process state should retire the worker."""
    scheduled_exits: list[tuple[float, int]] = []
    original_flag = modal_cloud_module._CONTAINER_TERMINATION_SCHEDULED
    monkeypatch.setattr(modal_cloud_module, "_CONTAINER_TERMINATION_SCHEDULED", False)
    monkeypatch.setattr(modal_cloud_module, "_is_modal_container_runtime", lambda: True)
    monkeypatch.setattr(
        modal_cloud_module,
        "_schedule_process_exit",
        lambda delay_seconds, exit_code: scheduled_exits.append((delay_seconds, exit_code)),
    )
    try:
        scheduled = modal_cloud_module._maybe_schedule_container_termination_on_error(
            {"component_id": "component-1", "terminate_container_on_error": True},
            RuntimeError("CUDA error: an illegal memory access was encountered"),
        )
    finally:
        monkeypatch.setattr(modal_cloud_module, "_CONTAINER_TERMINATION_SCHEDULED", original_flag)

    assert scheduled is True
    assert scheduled_exits == [(modal_cloud_module._REMOTE_ERROR_CONTAINER_EXIT_DELAY_SECONDS, 1)]


def test_modal_cloud_retires_worker_after_llm_memory_recovery_timeout(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """An unrecovered eviction should retire dirty process state before retry."""
    scheduled_exits: list[tuple[float, int]] = []
    original_flag = modal_cloud_module._CONTAINER_TERMINATION_SCHEDULED
    monkeypatch.setattr(modal_cloud_module, "_CONTAINER_TERMINATION_SCHEDULED", False)
    monkeypatch.setattr(modal_cloud_module, "_is_modal_container_runtime", lambda: True)
    monkeypatch.setattr(
        modal_cloud_module,
        "_schedule_process_exit",
        lambda delay_seconds, exit_code: scheduled_exits.append((delay_seconds, exit_code)),
    )
    try:
        scheduled = modal_cloud_module._maybe_schedule_container_termination_on_error(
            {"component_id": "llm-1", "terminate_container_on_error": True},
            RuntimeError("[comfy-modal-llm-memory-recovery-exhausted] still low"),
        )
    finally:
        monkeypatch.setattr(modal_cloud_module, "_CONTAINER_TERMINATION_SCHEDULED", original_flag)

    assert scheduled is True
    assert scheduled_exits == [(modal_cloud_module._REMOTE_ERROR_CONTAINER_EXIT_DELAY_SECONDS, 1)]


def test_modal_cloud_does_not_schedule_container_exit_for_interruptions(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Expected interruption-style failures should not tear down the warm Modal container."""
    scheduled_exits: list[tuple[float, int]] = []
    original_flag = modal_cloud_module._CONTAINER_TERMINATION_SCHEDULED
    monkeypatch.setattr(modal_cloud_module, "_CONTAINER_TERMINATION_SCHEDULED", False)
    monkeypatch.setattr(modal_cloud_module, "_is_modal_container_runtime", lambda: True)
    monkeypatch.setattr(
        modal_cloud_module,
        "_schedule_process_exit",
        lambda delay_seconds, exit_code: scheduled_exits.append((delay_seconds, exit_code)),
    )
    try:
        scheduled = modal_cloud_module._maybe_schedule_container_termination_on_error(
            {"component_id": "component-1", "terminate_container_on_error": True},
            modal_cloud_module.RemoteSubgraphExecutionError(
                "Remote subgraph execution was interrupted."
            ),
        )
    finally:
        monkeypatch.setattr(modal_cloud_module, "_CONTAINER_TERMINATION_SCHEDULED", original_flag)

    assert scheduled is False
    assert scheduled_exits == []


def test_modal_cloud_does_not_schedule_container_exit_for_session_state_misses(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Prompt-scoped session misses are routing problems, not poisoned-container crashes."""
    scheduled_exits: list[tuple[float, int]] = []
    original_flag = modal_cloud_module._CONTAINER_TERMINATION_SCHEDULED
    monkeypatch.setattr(modal_cloud_module, "_CONTAINER_TERMINATION_SCHEDULED", False)
    monkeypatch.setattr(modal_cloud_module, "_is_modal_container_runtime", lambda: True)
    monkeypatch.setattr(
        modal_cloud_module,
        "_schedule_process_exit",
        lambda delay_seconds, exit_code: scheduled_exits.append((delay_seconds, exit_code)),
    )
    try:
        scheduled = modal_cloud_module._maybe_schedule_container_termination_on_error(
            {"component_id": "component-1", "terminate_container_on_error": True},
            modal_cloud_module.RemoteSessionStateError("Remote session 'abc' was not found."),
        )
    finally:
        monkeypatch.setattr(modal_cloud_module, "_CONTAINER_TERMINATION_SCHEDULED", original_flag)

    assert scheduled is False
    assert scheduled_exits == []


def test_modal_cloud_skips_duplicate_reload_markers_in_same_container(
    modal_cloud_module: Any,
) -> None:
    """One container should reload a given uploaded-asset marker only once."""

    class FakeVolume:
        """Simple Modal volume double that tracks reload calls."""

        def __init__(self) -> None:
            """Initialize the reload counter."""
            self.reload_calls = 0

        def reload(self) -> None:
            """Record one reload attempt."""
            self.reload_calls += 1

    original_marker_queue = modal_cloud_module._MODAL_VOLUME_RELOAD_MARKERS
    original_marker_set = set(modal_cloud_module._MODAL_VOLUME_RELOAD_MARKER_SET)
    modal_cloud_module._MODAL_VOLUME_RELOAD_MARKERS = None
    modal_cloud_module._MODAL_VOLUME_RELOAD_MARKER_SET.clear()
    try:
        payload = {"requires_volume_reload": True, "volume_reload_marker": "marker-1"}
        assert modal_cloud_module._should_reload_modal_volume(payload) is True

        volume = FakeVolume()
        modal_cloud_module._reload_modal_volume_for_request(
            volume,
            "component-1",
            reload_marker="marker-1",
        )

        assert volume.reload_calls == 1
        assert modal_cloud_module._should_reload_modal_volume(payload) is False
    finally:
        modal_cloud_module._MODAL_VOLUME_RELOAD_MARKERS = original_marker_queue
        modal_cloud_module._MODAL_VOLUME_RELOAD_MARKER_SET.clear()
        modal_cloud_module._MODAL_VOLUME_RELOAD_MARKER_SET.update(original_marker_set)


def test_modal_cloud_retries_volume_reload_after_clearing_warm_state(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Modal volume reload should retry after unloading warm caches when open files block it."""

    class FakeVolume:
        """Simple Modal volume double that fails twice before succeeding."""

        def __init__(self) -> None:
            """Initialize the reload attempt counter."""
            self.reload_calls = 0

        def reload(self) -> None:
            """Raise on the first two calls and succeed on the third."""
            self.reload_calls += 1
            if self.reload_calls < 3:
                raise RuntimeError("there are open files preventing the operation")

    prepare_calls: list[str] = []
    sleep_calls: list[float] = []
    monkeypatch.setattr(
        modal_cloud_module,
        "_prepare_for_modal_volume_reload",
        lambda: prepare_calls.append("prepared"),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_sleep_before_modal_volume_reload_retry",
        lambda delay_seconds: sleep_calls.append(delay_seconds),
    )

    volume = FakeVolume()
    modal_cloud_module._reload_modal_volume_for_request(volume, "component-1")

    assert volume.reload_calls == 3
    assert prepare_calls == ["prepared", "prepared"]
    assert sleep_calls == [0.25, 0.5]


def test_modal_cloud_raises_after_exhausting_open_file_reload_retries(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Modal volume reload should surface persistent open-file errors after bounded retries."""

    class FakeVolume:
        """Simple Modal volume double that always fails with open files."""

        def __init__(self) -> None:
            """Initialize the reload attempt counter."""
            self.reload_calls = 0

        def reload(self) -> None:
            """Always fail with the same open-file reload error."""
            self.reload_calls += 1
            raise RuntimeError("there are open files preventing the operation")

    prepare_calls: list[str] = []
    sleep_calls: list[float] = []
    monkeypatch.setattr(
        modal_cloud_module,
        "_prepare_for_modal_volume_reload",
        lambda: prepare_calls.append("prepared"),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_sleep_before_modal_volume_reload_retry",
        lambda delay_seconds: sleep_calls.append(delay_seconds),
    )

    volume = FakeVolume()
    with pytest.raises(RuntimeError, match="open files"):
        modal_cloud_module._reload_modal_volume_for_request(volume, "component-1")

    assert volume.reload_calls == len(
        modal_cloud_module._MODAL_VOLUME_RELOAD_OPEN_FILE_RETRY_DELAYS_SECONDS
    )
    assert prepare_calls == ["prepared"] * (
        len(modal_cloud_module._MODAL_VOLUME_RELOAD_OPEN_FILE_RETRY_DELAYS_SECONDS) - 1
    )
    assert sleep_calls == list(
        modal_cloud_module._MODAL_VOLUME_RELOAD_OPEN_FILE_RETRY_DELAYS_SECONDS[1:]
    )


def test_modal_cloud_proceeds_when_referenced_volume_paths_are_already_visible(
    modal_cloud_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """Persistent open-file reload errors may be ignored when the payload's mounted files are already visible."""

    class FakeVolume:
        """Simple Modal volume double that always fails with open files."""

        def __init__(self) -> None:
            """Initialize the reload attempt counter."""
            self.reload_calls = 0

        def reload(self) -> None:
            """Always fail with the same open-file reload error."""
            self.reload_calls += 1
            raise RuntimeError("there are open files preventing the operation")

    storage_root = tmp_path / "storage"
    storage_root.mkdir()
    asset_path = storage_root / "assets" / "hash_model.safetensors"
    asset_path.parent.mkdir(parents=True, exist_ok=True)
    asset_path.write_bytes(b"weights")
    bundle_path = storage_root / "custom_nodes" / "hash_bundle.zip"
    bundle_path.parent.mkdir(parents=True, exist_ok=True)
    bundle_path.write_bytes(b"bundle")

    recorded_markers: list[str] = []
    prepare_calls: list[str] = []
    sleep_calls: list[float] = []
    monkeypatch.setattr(
        modal_cloud_module,
        "get_settings",
        lambda: types.SimpleNamespace(remote_storage_root=str(storage_root)),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_prepare_for_modal_volume_reload",
        lambda: prepare_calls.append("prepared"),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_sleep_before_modal_volume_reload_retry",
        lambda delay_seconds: sleep_calls.append(delay_seconds),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_record_modal_volume_reload_marker",
        lambda marker: recorded_markers.append(marker),
    )

    payload = {
        "custom_nodes_bundle": "/custom_nodes/hash_bundle.zip",
        "subgraph_prompt": {
            "1": {
                "class_type": "CheckpointLoaderSimple",
                "inputs": {"ckpt_name": "/assets/hash_model.safetensors"},
            }
        },
    }

    volume = FakeVolume()
    modal_cloud_module._reload_modal_volume_for_request(
        volume,
        "component-1",
        reload_marker="marker-1",
        payload=payload,
    )

    assert volume.reload_calls == len(
        modal_cloud_module._MODAL_VOLUME_RELOAD_OPEN_FILE_RETRY_DELAYS_SECONDS
    )
    assert recorded_markers == ["marker-1"]
    assert prepare_calls == ["prepared"] * (
        len(modal_cloud_module._MODAL_VOLUME_RELOAD_OPEN_FILE_RETRY_DELAYS_SECONDS) - 1
    )
    assert sleep_calls == list(
        modal_cloud_module._MODAL_VOLUME_RELOAD_OPEN_FILE_RETRY_DELAYS_SECONDS[1:]
    )


def test_modal_cloud_logs_volume_reload_diagnostics_for_open_file_retries(
    modal_cloud_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """Open-file retries should log which uploaded and referenced volume paths matter."""

    class FakeVolume:
        """Simple Modal volume double that always fails with open files."""

        def reload(self) -> None:
            """Always fail with the same open-file reload error."""
            raise RuntimeError("there are open files preventing the operation")

    storage_root = tmp_path / "storage"
    storage_root.mkdir()
    uploaded_path = storage_root / "assets" / "missing_model.safetensors"
    logged_messages: list[tuple[str, tuple[Any, ...]]] = []

    monkeypatch.setattr(
        modal_cloud_module,
        "get_settings",
        lambda: types.SimpleNamespace(remote_storage_root=str(storage_root)),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_prepare_for_modal_volume_reload",
        lambda: None,
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_sleep_before_modal_volume_reload_retry",
        lambda delay_seconds: None,
    )
    monkeypatch.setattr(
        modal_cloud_module.logger,
        "info",
        lambda message, *args: logged_messages.append((message, args)),
    )

    payload = {
        "uploaded_volume_paths": ["/assets/missing_model.safetensors"],
        "subgraph_prompt": {
            "1": {
                "class_type": "CheckpointLoaderSimple",
                "inputs": {"ckpt_name": "/assets/missing_model.safetensors"},
            }
        },
    }

    with pytest.raises(RuntimeError, match="open files"):
        modal_cloud_module._reload_modal_volume_for_request(
            FakeVolume(),
            "component-1",
            payload=payload,
        )

    assert any(
        "Modal volume reload diagnostics for component=%s context=%s uploaded_paths=%s referenced_paths=%s visible_uploaded=%s visible_referenced=%s."
        in message
        and args[0] == "component-1"
        and args[1] == "open_files_retry"
        and args[2] == [str(uploaded_path)]
        and args[3] == [str(uploaded_path)]
        and args[4] is False
        and args[5] is False
        for message, args in logged_messages
    )


def test_modal_cloud_loader_cache_reuses_and_clones_outputs(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Cached loader wrappers should avoid repeated loads and clone cached outputs on hits."""

    class FakeLoader:
        """Simple loader with one expensive method."""

        def __init__(self) -> None:
            """Initialize call counter state."""
            self.calls = 0

        def load(self, model_name: str, device: str = "default") -> tuple[_CloneableCacheValue]:
            """Return a cloneable payload while counting underlying loads."""
            self.calls += 1
            return (_CloneableCacheValue(f"{model_name}:{device}:{self.calls}"),)

    original_cache = dict(modal_cloud_module._LOADER_OUTPUT_CACHE)
    original_wrapped = set(modal_cloud_module._LOADER_CACHE_WRAPPED_CLASSES)
    original_metrics = dict(modal_cloud_module._LOADER_CACHE_METRICS)
    modal_cloud_module._LOADER_OUTPUT_CACHE.clear()
    modal_cloud_module._LOADER_CACHE_WRAPPED_CLASSES.clear()
    modal_cloud_module._LOADER_CACHE_METRICS.clear()
    modal_cloud_module._LOADER_CACHE_METRICS.update({"hit": 0, "miss": 0})
    try:
        modal_cloud_module._wrap_loader_method_with_cache(
            "FakeLoader",
            FakeLoader,
            "load",
            lambda kwargs: modal_cloud_module._serialize_loader_cache_key(kwargs),
        )
        loader = FakeLoader()
        first = loader.load("model.safetensors", device="cpu")[0]
        second = loader.load("model.safetensors", device="cpu")[0]
        third = loader.load("other.safetensors", device="cpu")[0]
        metrics_snapshot = modal_cloud_module._loader_cache_metric_snapshot()
    finally:
        modal_cloud_module._LOADER_OUTPUT_CACHE.clear()
        modal_cloud_module._LOADER_OUTPUT_CACHE.update(original_cache)
        modal_cloud_module._LOADER_CACHE_WRAPPED_CLASSES.clear()
        modal_cloud_module._LOADER_CACHE_WRAPPED_CLASSES.update(original_wrapped)
        modal_cloud_module._LOADER_CACHE_METRICS.clear()
        modal_cloud_module._LOADER_CACHE_METRICS.update(original_metrics)

    assert loader.calls == 2
    assert first.value == "model.safetensors:cpu:1"
    assert second.value == "model.safetensors:cpu:1"
    assert third.value == "other.safetensors:cpu:2"
    assert first is not second
    assert metrics_snapshot == {"hit": 1, "miss": 2}


def test_modal_cloud_rehydrates_bridge_refs_from_warm_value_cache_without_replay(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Warm-worker bridge values should restore directly into a fresh session without replay."""
    bridge_key = "RSB_cached_bridge"
    target_handle = modal_cloud_module.RemoteSessionHandle(
        session_id="session-target",
        prompt_id="prompt-1",
        owner_component_id="component-1",
    )
    bridge_ref = modal_cloud_module.RemoteSessionBridgeRef(
        bridge_key=bridge_key,
        node_id="node-7",
        output_index=0,
        session_id="session-source",
    )
    original_cache = dict(modal_cloud_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE)
    original_order = list(modal_cloud_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER)
    try:
        seed_value = _CloneableCacheValue("warm-bridge-value")
        modal_cloud_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE.clear()
        modal_cloud_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.clear()
        modal_cloud_module._store_remote_session_bridge_value(bridge_key, seed_value)
        monkeypatch.setattr(
            modal_cloud_module,
            "_load_remote_session_bridge_record",
            lambda bridge_key: (_ for _ in ()).throw(
                AssertionError(f"warm bridge cache hit should skip record lookup for {bridge_key}")
            ),
        )
        resolution_stats = modal_cloud_module._RemoteSessionBridgeResolutionStats()

        restored_value = modal_cloud_module._rehydrate_remote_session_bridge_value(
            bridge_ref,
            target_session_handle=target_handle,
            custom_nodes_root=None,
            cancellation_event=None,
            interrupt_store=None,
            interrupt_flag_key=None,
            resolution_stats=resolution_stats,
        )

        stored_value = modal_cloud_module._REMOTE_SESSION_STORE.get_output(
            modal_cloud_module.RemoteSessionValueRef(
                session_id=target_handle.session_id,
                node_id="node-7",
                output_index=0,
            )
        )
    finally:
        modal_cloud_module._REMOTE_SESSION_STORE.clear_session(target_handle)
        modal_cloud_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE.clear()
        modal_cloud_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE.update(original_cache)
        modal_cloud_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.clear()
        modal_cloud_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.extend(original_order)

    assert isinstance(restored_value, _CloneableCacheValue)
    assert restored_value.value == "warm-bridge-value"
    assert restored_value is not seed_value
    assert stored_value is restored_value
    assert resolution_stats.bridge_cache_hits == 1
    assert resolution_stats.bridge_record_lookups == 0
    assert resolution_stats.replay_count == 0
    assert resolution_stats.session_restore_writes == 1


def test_modal_cloud_rehydrates_conditioning_bridge_refs_from_durable_record_without_replay(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Durably serialized CONDITIONING bridge values should restore without replay on fresh workers."""
    import torch

    target_handle = modal_cloud_module.RemoteSessionHandle(
        session_id="session-target",
        prompt_id="prompt-1",
        owner_component_id="component-1",
    )
    bridge_ref = modal_cloud_module.RemoteSessionBridgeRef(
        bridge_key="RSB_conditioning_bridge",
        node_id="node-9",
        output_index=0,
        session_id="session-source",
    )
    conditioning = [
        [
            torch.arange(6, dtype=torch.float32).reshape(1, 2, 3),
            {"pooled_output": torch.arange(4, dtype=torch.float32).reshape(1, 4)},
        ]
    ]
    monkeypatch.setattr(
        modal_cloud_module,
        "_load_remote_session_bridge_record",
        lambda bridge_key: modal_cloud_module.RemoteSessionBridgeRecord(
            bridge_key=bridge_key,
            node_id="node-9",
            output_index=0,
            producer_payload={"component_id": "should-not-replay"},
            producer_inputs={},
            serialized_output=modal_cloud_module.serialize_value(conditioning),
            serialized_output_io_type="CONDITIONING",
        ),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_execute_subgraph_prompt",
        lambda *args, **kwargs: (_ for _ in ()).throw(
            AssertionError("durable CONDITIONING restore should skip replay")
        ),
    )
    resolution_stats = modal_cloud_module._RemoteSessionBridgeResolutionStats()

    try:
        restored_value = modal_cloud_module._rehydrate_remote_session_bridge_value(
            bridge_ref,
            target_session_handle=target_handle,
            custom_nodes_root=None,
            cancellation_event=None,
            interrupt_store=None,
            interrupt_flag_key=None,
            resolution_stats=resolution_stats,
        )
        stored_value = modal_cloud_module._REMOTE_SESSION_STORE.get_output(
            modal_cloud_module.RemoteSessionValueRef(
                session_id=target_handle.session_id,
                node_id="node-9",
                output_index=0,
            )
        )
    finally:
        modal_cloud_module._REMOTE_SESSION_STORE.clear_session(target_handle)

    assert torch.equal(restored_value[0][0], conditioning[0][0])
    assert torch.equal(restored_value[0][1]["pooled_output"], conditioning[0][1]["pooled_output"])
    assert torch.equal(stored_value[0][0], conditioning[0][0])
    assert resolution_stats.bridge_cache_hits == 0
    assert resolution_stats.durable_bridge_hits == 1
    assert resolution_stats.bridge_record_lookups == 1
    assert resolution_stats.replay_count == 0
    assert resolution_stats.session_restore_writes == 1


@pytest.mark.parametrize(
    ("io_type", "node_id", "class_type", "node_inputs"),
    [
        ("NOISE", "15", "RandomNoise", {"noise_seed": 42}),
        ("SAMPLER", "17", "KSamplerSelect", {"sampler_name": "euler"}),
    ],
)
def test_modal_cloud_rehydrates_literal_sampling_strategy_bridges_without_replay(
    modal_cloud_module: Any,
    monkeypatch: Any,
    io_type: str,
    node_id: str,
    class_type: str,
    node_inputs: dict[str, Any],
) -> None:
    """Literal NOISE and SAMPLER strategy nodes should rebuild without producer replay."""
    target_handle = modal_cloud_module.RemoteSessionHandle(
        session_id="session-target",
        prompt_id="prompt-1",
        owner_component_id="component-1",
    )
    bridge_ref = modal_cloud_module.RemoteSessionBridgeRef(
        bridge_key=f"RSB_{io_type.lower()}_bridge",
        node_id=node_id,
        output_index=0,
        session_id="session-source",
    )
    record = modal_cloud_module._build_remote_session_bridge_record(
        payload={
            "component_id": "image-preview-component",
            "execute_node_ids": ["251"],
            "subgraph_prompt": {
                "14": {
                    "class_type": "SamplerCustomAdvanced",
                    "inputs": {"noise": ["15", 0], "sampler": ["17", 0]},
                },
                "15": {
                    "class_type": "RandomNoise",
                    "inputs": {"noise_seed": 42},
                },
                "17": {
                    "class_type": "KSamplerSelect",
                    "inputs": {"sampler_name": "euler"},
                },
                "250": {
                    "class_type": "VAEDecode",
                    "inputs": {"samples": ["14", 0]},
                },
                "251": {
                    "class_type": "ImageFromBatch",
                    "inputs": {"image": ["250", 0]},
                },
            },
        },
        hydrated_inputs={},
        node_id=node_id,
        output_index=0,
        io_type=io_type,
        output_value=object(),
    )
    assert record.recovery_kind is modal_cloud_module.RemoteSessionBridgeRecoveryKind.SINGLE_NODE_PLAN
    assert record.rehydration_plan == {
        "kind": "single_node_output",
        "node_data": {"class_type": class_type},
        "node_inputs": node_inputs,
    }
    execute_calls: list[tuple[dict[str, Any], dict[str, Any]]] = []
    monkeypatch.setattr(
        modal_cloud_module,
        "_load_remote_session_bridge_record",
        lambda bridge_key: record,
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_execute_subgraph_prompt",
        lambda *args, **kwargs: (_ for _ in ()).throw(
            AssertionError("literal strategy bridge should not replay its producer component")
        ),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_execute_node_locally_raw",
        lambda node_data, kwargs_payload, **kwargs: (
            execute_calls.append((dict(node_data), dict(kwargs_payload))),
            (f"restored-{io_type.lower()}",),
        )[1],
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_store_remote_session_bridge_value",
        lambda bridge_key, value: None,
    )
    resolution_stats = modal_cloud_module._RemoteSessionBridgeResolutionStats()

    try:
        restored_value = modal_cloud_module._rehydrate_remote_session_bridge_value(
            bridge_ref,
            target_session_handle=target_handle,
            custom_nodes_root=None,
            cancellation_event=None,
            interrupt_store=None,
            interrupt_flag_key=None,
            resolution_stats=resolution_stats,
        )
    finally:
        modal_cloud_module._REMOTE_SESSION_STORE.clear_session(target_handle)

    assert restored_value == f"restored-{io_type.lower()}"
    assert execute_calls == [({"class_type": class_type}, node_inputs)]
    assert resolution_stats.durable_bridge_hits == 1
    assert resolution_stats.replay_count == 0


def test_modal_cloud_rehydrates_sampler_latent_bridge_refs_from_durable_record_without_replay(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Durably serialized LATENT bridge values should restore without replaying sampler producers."""
    import torch

    target_handle = modal_cloud_module.RemoteSessionHandle(
        session_id="session-target",
        prompt_id="prompt-1",
        owner_component_id="component-1",
    )
    bridge_ref = modal_cloud_module.RemoteSessionBridgeRef(
        bridge_key="RSB_sampler_latent_bridge",
        node_id="sampler-1",
        output_index=0,
        session_id="session-source",
    )
    latent = {"samples": torch.arange(8, dtype=torch.float32).reshape(1, 2, 2, 2)}
    record = modal_cloud_module._build_remote_session_bridge_record(
        payload={
            "component_id": "sampler-component",
            "execute_node_ids": ["sampler-1"],
            "subgraph_prompt": {
                "sampler-1": {
                    "class_type": "KSampler",
                    "inputs": {},
                }
            },
        },
        hydrated_inputs={},
        node_id="sampler-1",
        output_index=0,
        io_type="LATENT",
        output_value=latent,
    )
    assert record.serialized_output is not None
    assert record.serialized_output_io_type == "LATENT"
    monkeypatch.setattr(
        modal_cloud_module,
        "_load_remote_session_bridge_record",
        lambda bridge_key: record,
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_execute_subgraph_prompt",
        lambda *args, **kwargs: (_ for _ in ()).throw(
            AssertionError("durable LATENT restore should skip sampler replay")
        ),
    )
    resolution_stats = modal_cloud_module._RemoteSessionBridgeResolutionStats()

    try:
        restored_value = modal_cloud_module._rehydrate_remote_session_bridge_value(
            bridge_ref,
            target_session_handle=target_handle,
            custom_nodes_root=None,
            cancellation_event=None,
            interrupt_store=None,
            interrupt_flag_key=None,
            resolution_stats=resolution_stats,
        )
        stored_value = modal_cloud_module._REMOTE_SESSION_STORE.get_output(
            modal_cloud_module.RemoteSessionValueRef(
                session_id=target_handle.session_id,
                node_id="sampler-1",
                output_index=0,
            )
        )
    finally:
        modal_cloud_module._REMOTE_SESSION_STORE.clear_session(target_handle)

    assert torch.equal(restored_value["samples"], latent["samples"])
    assert torch.equal(stored_value["samples"], latent["samples"])
    assert resolution_stats.bridge_cache_hits == 0
    assert resolution_stats.durable_bridge_hits == 1
    assert resolution_stats.bridge_record_lookups == 1
    assert resolution_stats.replay_count == 0
    assert resolution_stats.session_restore_writes == 1


def test_modal_cloud_rehydrates_model_bridge_refs_from_durable_plan_without_replay(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Durable MODEL bridge plans should rebuild one self-contained loader output without replay."""
    target_handle = modal_cloud_module.RemoteSessionHandle(
        session_id="session-target",
        prompt_id="prompt-1",
        owner_component_id="component-1",
    )
    bridge_ref = modal_cloud_module.RemoteSessionBridgeRef(
        bridge_key="RSB_model_bridge",
        node_id="node-5",
        output_index=0,
        session_id="session-source",
    )
    record = modal_cloud_module._build_remote_session_bridge_record(
        payload={
            "component_id": "component-seed",
            "subgraph_prompt": {
                "node-5": {
                    "class_type": "CheckpointLoaderSimple",
                    "inputs": {"ckpt_name": "model.safetensors"},
                }
            },
        },
        hydrated_inputs={},
        node_id="node-5",
        output_index=0,
        io_type="MODEL",
        output_value=_FakeModelValue("seed-model"),
    )
    execute_calls: list[tuple[dict[str, Any], dict[str, Any]]] = []
    monkeypatch.setattr(modal_cloud_module, "_load_remote_session_bridge_record", lambda bridge_key: record)
    monkeypatch.setattr(
        modal_cloud_module,
        "_execute_subgraph_prompt",
        lambda *args, **kwargs: (_ for _ in ()).throw(
            AssertionError("durable MODEL rehydration should skip replay")
        ),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_execute_node_locally_raw",
        lambda node_data, kwargs_payload, **kwargs: (
            execute_calls.append((dict(node_data), dict(kwargs_payload))),
            (_FakeModelValue("restored-model"),),
        )[1],
    )
    resolution_stats = modal_cloud_module._RemoteSessionBridgeResolutionStats()

    try:
        restored_value = modal_cloud_module._rehydrate_remote_session_bridge_value(
            bridge_ref,
            target_session_handle=target_handle,
            custom_nodes_root=None,
            cancellation_event=None,
            interrupt_store=None,
            interrupt_flag_key=None,
            resolution_stats=resolution_stats,
        )
        stored_value = modal_cloud_module._REMOTE_SESSION_STORE.get_output(
            modal_cloud_module.RemoteSessionValueRef(
                session_id=target_handle.session_id,
                node_id="node-5",
                output_index=0,
            )
        )
    finally:
        modal_cloud_module._REMOTE_SESSION_STORE.clear_session(target_handle)

    assert isinstance(restored_value, _FakeModelValue)
    assert restored_value.value == "restored-model"
    assert stored_value is restored_value
    assert execute_calls == [
        (
            {"class_type": "CheckpointLoaderSimple"},
            {"ckpt_name": "model.safetensors"},
        )
    ]
    assert resolution_stats.bridge_cache_hits == 0
    assert resolution_stats.durable_bridge_hits == 1
    assert resolution_stats.bridge_record_lookups == 1
    assert resolution_stats.replay_count == 0
    assert resolution_stats.session_restore_writes == 1


def test_modal_cloud_rehydrates_clip_bridge_refs_from_durable_plan_without_replay(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Durable CLIP bridge plans should rebuild one self-contained loader output without replay."""
    target_handle = modal_cloud_module.RemoteSessionHandle(
        session_id="session-target",
        prompt_id="prompt-1",
        owner_component_id="component-1",
    )
    bridge_ref = modal_cloud_module.RemoteSessionBridgeRef(
        bridge_key="RSB_clip_bridge",
        node_id="clip-1",
        output_index=0,
        session_id="session-source",
    )
    record = modal_cloud_module._build_remote_session_bridge_record(
        payload={
            "component_id": "sampler-component",
            "execute_node_ids": ["sampler-1"],
            "subgraph_prompt": {
                "clip-1": {
                    "class_type": "CLIPLoader",
                    "inputs": {
                        "clip_name": "text.safetensors",
                        "type": "flux",
                        "device": "default",
                    },
                },
                "sampler-1": {
                    "class_type": "KSampler",
                    "inputs": {"clip": ["clip-1", 0]},
                },
            },
        },
        hydrated_inputs={},
        node_id="clip-1",
        output_index=0,
        io_type="CLIP",
        output_value="seed-clip",
    )
    assert record.rehydration_plan is not None
    assert record.rehydration_plan_io_type == "CLIP"
    monkeypatch.setattr(modal_cloud_module, "_load_remote_session_bridge_record", lambda bridge_key: record)
    monkeypatch.setattr(
        modal_cloud_module,
        "_execute_subgraph_prompt",
        lambda *args, **kwargs: (_ for _ in ()).throw(
            AssertionError("durable CLIP rehydration should skip replay")
        ),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_execute_node_locally_raw",
        lambda node_data, kwargs_payload, **kwargs: (
            "clip::"
            f"{kwargs_payload['clip_name']}:"
            f"{kwargs_payload['type']}:"
            f"{kwargs_payload['device']}",
        ),
    )
    resolution_stats = modal_cloud_module._RemoteSessionBridgeResolutionStats()

    try:
        restored_value = modal_cloud_module._rehydrate_remote_session_bridge_value(
            bridge_ref,
            target_session_handle=target_handle,
            custom_nodes_root=None,
            cancellation_event=None,
            interrupt_store=None,
            interrupt_flag_key=None,
            resolution_stats=resolution_stats,
        )
    finally:
        modal_cloud_module._REMOTE_SESSION_STORE.clear_session(target_handle)

    assert restored_value == "clip::text.safetensors:flux:default"
    assert resolution_stats.durable_bridge_hits == 1
    assert resolution_stats.replay_count == 0


def test_modal_cloud_rehydrates_linked_model_bridge_with_non_sampler_subgraph_plan(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Linked MODEL bridge plans should rerun only the non-sampler model dependency closure."""
    target_handle = modal_cloud_module.RemoteSessionHandle(
        session_id="session-target",
        prompt_id="prompt-1",
        owner_component_id="component-1",
    )
    bridge_ref = modal_cloud_module.RemoteSessionBridgeRef(
        bridge_key="RSB_linked_model_bridge",
        node_id="lora-1",
        output_index=0,
        session_id="session-source",
    )
    record = modal_cloud_module._build_remote_session_bridge_record(
        payload={
            "component_id": "sampler-component",
            "execute_node_ids": ["sampler-1"],
            "subgraph_prompt": {
                "loader-1": {
                    "class_type": "UNETLoader",
                    "inputs": {"unet_name": "base.safetensors"},
                },
                "lora-1": {
                    "class_type": "LoraLoaderModelOnly",
                    "inputs": {"model": ["loader-1", 0], "lora_name": "style.safetensors"},
                },
                "sampler-1": {
                    "class_type": "KSampler",
                    "inputs": {"model": ["lora-1", 0]},
                },
            },
        },
        hydrated_inputs={},
        node_id="lora-1",
        output_index=0,
        io_type="MODEL",
        output_value=_FakeModelValue("lora-model"),
    )
    assert record.rehydration_plan is not None
    assert record.rehydration_plan["kind"] == "subgraph_output"
    monkeypatch.setattr(modal_cloud_module, "_load_remote_session_bridge_record", lambda bridge_key: record)
    observed_payloads: list[dict[str, Any]] = []

    def fake_execute_subgraph_prompt(
        payload: dict[str, Any],
        hydrated_inputs: dict[str, Any],
        *_args: Any,
    ) -> tuple[_FakeModelValue]:
        """Assert the durable plan avoids rerunning the sampler component."""
        del hydrated_inputs
        observed_payloads.append(copy.deepcopy(payload))
        assert payload["execute_node_ids"] == ["lora-1"]
        assert payload["remote_session"]["session_id"] == target_handle.session_id
        assert sorted(payload["subgraph_prompt"]) == ["loader-1", "lora-1"]
        assert "sampler-1" not in payload["subgraph_prompt"]
        return (_FakeModelValue("restored-lora-model"),)

    monkeypatch.setattr(modal_cloud_module, "_execute_subgraph_prompt", fake_execute_subgraph_prompt)
    resolution_stats = modal_cloud_module._RemoteSessionBridgeResolutionStats()

    try:
        restored_value = modal_cloud_module._rehydrate_remote_session_bridge_value(
            bridge_ref,
            target_session_handle=target_handle,
            custom_nodes_root=None,
            cancellation_event=None,
            interrupt_store=None,
            interrupt_flag_key=None,
            resolution_stats=resolution_stats,
        )
    finally:
        modal_cloud_module._REMOTE_SESSION_STORE.clear_session(target_handle)

    assert isinstance(restored_value, _FakeModelValue)
    assert restored_value.value == "restored-lora-model"
    assert len(observed_payloads) == 1
    assert resolution_stats.durable_bridge_hits == 1
    assert resolution_stats.replay_count == 0


def test_modal_cloud_refuses_sampler_ancestor_bridge_replay(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Producer replay should fail when a terminal node's dependency closure samples."""
    target_handle = modal_cloud_module.RemoteSessionHandle(
        session_id="session-target",
        prompt_id="prompt-1",
        owner_component_id="component-1",
    )
    bridge_ref = modal_cloud_module.RemoteSessionBridgeRef(
        bridge_key="RSB_sampler_bridge",
        node_id="unsupported-bridge",
        output_index=0,
        session_id="session-source",
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_load_remote_session_bridge_record",
        lambda bridge_key: modal_cloud_module.RemoteSessionBridgeRecord(
            bridge_key=bridge_key,
            node_id="unsupported-bridge",
            output_index=0,
            producer_payload={
                "component_id": "sampler-component",
                "execute_node_ids": ["251"],
                "subgraph_prompt": {
                    "14": {"class_type": "KSampler", "inputs": {}},
                    "250": {
                        "class_type": "VAEDecode",
                        "inputs": {"samples": ["14", 0]},
                    },
                    "251": {
                        "class_type": "ImageFromBatch",
                        "inputs": {"image": ["250", 0]},
                    },
                },
            },
            producer_inputs={},
        ),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_execute_subgraph_prompt",
        lambda *args, **kwargs: (_ for _ in ()).throw(
            AssertionError("sampler bridge replay should be blocked")
        ),
    )
    resolution_stats = modal_cloud_module._RemoteSessionBridgeResolutionStats()

    with pytest.raises(modal_cloud_module.RemoteSessionStateError, match="rerun a sampler"):
        modal_cloud_module._rehydrate_remote_session_bridge_value(
            bridge_ref,
            target_session_handle=target_handle,
            custom_nodes_root=None,
            cancellation_event=None,
            interrupt_store=None,
            interrupt_flag_key=None,
            resolution_stats=resolution_stats,
        )

    assert resolution_stats.bridge_record_lookups == 1
    assert resolution_stats.replay_count == 0


def test_modal_cloud_skips_seed_execution_when_session_outputs_are_already_restored(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Cloud seed payloads should no-op when bridge outputs are already in session memory."""
    session_handle = modal_cloud_module.RemoteSessionHandle(
        session_id="session-restored",
        prompt_id="prompt-1",
        owner_component_id="component-1",
    )
    restored_value = _FakeModelValue("restored-model")
    modal_cloud_module._REMOTE_SESSION_STORE.put_output(
        session_handle,
        node_id="5",
        output_index=0,
        value=restored_value,
    )
    payload = {
        "payload_kind": "subgraph",
        "component_id": "1__mapped::seed:0",
        "prompt_id": "prompt-1",
        "component_node_ids": ["5"],
        "subgraph_prompt": {
            "5": {
                "class_type": "CheckpointLoaderSimple",
                "inputs": {"ckpt_name": "model.safetensors"},
            }
        },
        "boundary_inputs": [
            {
                "proxy_input_name": "static_input_0",
                "io_type": "MODEL",
                "targets": [{"node_id": "5", "input_name": "model"}],
            }
        ],
        "boundary_outputs": [
            {
                "proxy_output_name": "static_input_0",
                "node_id": "5",
                "output_index": 0,
                "io_type": "MODEL",
                "is_list": False,
                "session_output": True,
            }
        ],
        "execute_node_ids": ["5"],
        "remote_session": session_handle.to_payload(),
    }
    hydrated_inputs = {
        "static_input_0": modal_cloud_module.RemoteSessionValueRef(
            session_id=session_handle.session_id,
            node_id="5",
            output_index=0,
        ).to_payload()
    }
    monkeypatch.setattr(
        modal_cloud_module,
        "_load_execution_module",
        lambda: (_ for _ in ()).throw(AssertionError("seed short-circuit should skip PromptExecutor")),
    )

    try:
        outputs = modal_cloud_module._execute_subgraph_prompt(
            payload,
            hydrated_inputs,
            None,
        )
    finally:
        modal_cloud_module._REMOTE_SESSION_STORE.clear_session(session_handle)

    assert len(outputs) == 1
    assert modal_cloud_module.is_remote_session_bridge_ref_payload(outputs[0])


def test_modal_cloud_logs_remote_session_resolution_summary(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Mapped-phase bridge resolution should emit one summary block with replay and loader deltas."""
    observed_logs: list[tuple[str, tuple[Any, ...]]] = []
    monkeypatch.setattr(
        modal_cloud_module,
        "_emit_cloud_info",
        lambda message, *args: observed_logs.append((message, args)),
    )

    modal_cloud_module._log_remote_session_resolution_summary(
        component_id="1::mapped",
        resolution_stats=modal_cloud_module._RemoteSessionBridgeResolutionStats(
            input_ref_count=2,
            live_session_hits=1,
            bridge_cache_hits=1,
            durable_bridge_hits=1,
            bridge_record_lookups=1,
            bridge_record_lookup_seconds=0.25,
            replay_count=1,
            replay_seconds=1.5,
            direct_restore_seconds=0.02,
            session_restore_writes=1,
        ),
        loader_cache_before={"hit": 3, "miss": 4},
        loader_cache_after={"hit": 5, "miss": 5},
    )

    assert observed_logs == [
        (
            "Remote session resolution summary component=%s refs=%d live_hits=%d warm_bridge_hits=%d durable_bridge_hits=%d bridge_record_lookups=%d bridge_record_lookup_seconds=%.3f replay_count=%d replay_seconds=%.3f direct_restore_seconds=%.3f session_restore_writes=%d loader_cache_hits=%d loader_cache_misses=%d",
            ("1::mapped", 2, 1, 1, 1, 1, 0.25, 1, 1.5, 0.02, 1, 2, 1),
        )
    ]


def test_modal_cloud_installs_loader_cache_wrappers_for_builtin_loaders(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """The runtime should patch the heavyweight built-in loaders once they are available."""
    fake_nodes_module = types.SimpleNamespace(
        NODE_CLASS_MAPPINGS={
            "CheckpointLoader": type("CheckpointLoader", (), {"load_checkpoint": lambda self, config_name, ckpt_name: (config_name, ckpt_name)}),
            "CheckpointLoaderSimple": type("CheckpointLoaderSimple", (), {"load_checkpoint": lambda self, ckpt_name: (ckpt_name,)}),
            "UNETLoader": type("UNETLoader", (), {"load_unet": lambda self, unet_name, weight_dtype="default": (unet_name,)}),
            "CLIPLoader": type("CLIPLoader", (), {"load_clip": lambda self, clip_name, type="stable_diffusion", device="default": (clip_name,)}),
            "DualCLIPLoader": type("DualCLIPLoader", (), {"load_clip": lambda self, clip_name1, clip_name2, type, device="default": (clip_name1, clip_name2)}),
            "VAELoader": type("VAELoader", (), {"load_vae": lambda self, vae_name: (vae_name,)}),
            "unCLIPCheckpointLoader": type("unCLIPCheckpointLoader", (), {"load_checkpoint": lambda self, ckpt_name, output_vae=True, output_clip=True: (ckpt_name,)}),
            "ImageOnlyCheckpointLoader": type("ImageOnlyCheckpointLoader", (), {"load_checkpoint": lambda self, ckpt_name, output_vae=True, output_clip=True: (ckpt_name,)}),
        }
    )

    original_wrapped = set(modal_cloud_module._LOADER_CACHE_WRAPPED_CLASSES)
    modal_cloud_module._LOADER_CACHE_WRAPPED_CLASSES.clear()
    monkeypatch.setattr(modal_cloud_module, "_load_nodes_module", lambda: fake_nodes_module)
    try:
        modal_cloud_module._install_loader_cache_wrappers()
        installed_wrappers = set(modal_cloud_module._LOADER_CACHE_WRAPPED_CLASSES)
    finally:
        modal_cloud_module._LOADER_CACHE_WRAPPED_CLASSES.clear()
        modal_cloud_module._LOADER_CACHE_WRAPPED_CLASSES.update(original_wrapped)

    assert {
        "CheckpointLoader",
        "CheckpointLoaderSimple",
        "UNETLoader",
        "CLIPLoader",
        "DualCLIPLoader",
        "VAELoader",
        "unCLIPCheckpointLoader",
        "ImageOnlyCheckpointLoader",
    } <= installed_wrappers


def test_modal_cloud_node_cache_key_hashes_boundary_tensors(
    modal_cloud_module: Any,
) -> None:
    """Boundary tensors inside ComfyUI cache signatures should produce stable cache keys."""
    torch = pytest.importorskip("torch")

    first_tensor = torch.arange(6, dtype=torch.float32).reshape(1, 2, 3)
    same_value_tensor = first_tensor.clone()
    different_tensor = first_tensor + 1
    signature = frozenset(
        {
            (
                12,
                frozenset(
                    {
                        (
                            4,
                            frozenset(
                                {
                                    (0, "latent_image"),
                                    (1, frozenset({("samples", first_tensor)})),
                                }
                            ),
                        )
                    }
                ),
            )
        }
    )
    same_signature = frozenset(
        {
            (
                12,
                frozenset(
                    {
                        (
                            4,
                            frozenset(
                                {
                                    (0, "latent_image"),
                                    (1, frozenset({("samples", same_value_tensor)})),
                                }
                            ),
                        )
                    }
                ),
            )
        }
    )
    different_signature = frozenset(
        {
            (
                12,
                frozenset(
                    {
                        (
                            4,
                            frozenset(
                                {
                                    (0, "latent_image"),
                                    (1, frozenset({("samples", different_tensor)})),
                                }
                            ),
                        )
                    }
                ),
            )
        }
    )

    first_key = modal_cloud_module._node_output_cache_key(signature)
    second_key = modal_cloud_module._node_output_cache_key(same_signature)
    different_key = modal_cloud_module._node_output_cache_key(different_signature)

    assert isinstance(first_key, str)
    assert first_key.startswith("NC_")
    assert second_key == first_key
    assert different_key != first_key


def test_modal_cloud_node_cache_key_rebuilds_input_signature_before_unhashable_conversion(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Distributed node caching should bypass ComfyUI's precomputed `Unhashable` data key."""
    torch = pytest.importorskip("torch")

    class FakeDynPrompt:
        """Minimal dynamic-prompt stub for input-signature reconstruction."""

        def __init__(self, node: dict[str, Any]) -> None:
            """Store one node payload under id `12`."""
            self._node = node

        def has_node(self, node_id: str) -> bool:
            """Return whether the requested node exists."""
            return str(node_id) == "12"

        def get_node(self, node_id: str) -> dict[str, Any]:
            """Return the stored node payload."""
            if not self.has_node(node_id):
                raise KeyError(node_id)
            return self._node

    class FakeUnhashable:
        """Stand-in for ComfyUI's `Unhashable` marker."""

    tensor = torch.arange(6, dtype=torch.float32).reshape(1, 2, 3)
    dynprompt = FakeDynPrompt(
        {
            "class_type": "FakeSampler",
            "inputs": {
                "latent_image": {"samples": tensor},
                "steps": 18,
            },
        }
    )
    cache_key_set = types.SimpleNamespace(
        dynprompt=dynprompt,
        is_changed_cache=types.SimpleNamespace(is_changed={"12": False}),
        get_ordered_ancestry=lambda current_dynprompt, node_id: ([], {}),
        include_node_id_in_input=lambda: False,
        get_data_key=lambda node_id: frozenset({("latent_image", FakeUnhashable())}),
    )

    monkeypatch.setattr(
        modal_cloud_module,
        "_load_nodes_module",
        lambda: types.SimpleNamespace(
            NODE_CLASS_MAPPINGS={"FakeSampler": type("FakeSampler", (), {})}
        ),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_include_unique_id_in_input_signature",
        lambda class_type: False,
    )

    rebuilt_key = modal_cloud_module._node_output_cache_key_from_key_set_sync(cache_key_set, "12")
    direct_bad_key = modal_cloud_module._node_output_cache_key(cache_key_set.get_data_key("12"))

    assert isinstance(rebuilt_key, str)
    assert rebuilt_key.startswith("NC_")
    assert direct_bad_key is None


def test_modal_cloud_node_cache_key_uses_boundary_source_signature_for_unhashable_inputs(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Boundary-fed runtime objects should hash by stable source provenance instead of object identity."""

    class FakeDynPrompt:
        """Minimal dynamic prompt wrapper backed by one mutable prompt dict."""

        def __init__(self, prompt: dict[str, Any]) -> None:
            """Store the prompt used by the cache-key rebuild."""
            self._prompt = prompt

        def has_node(self, node_id: str) -> bool:
            """Return whether the requested node exists in the prompt."""
            return str(node_id) in self._prompt

        def get_node(self, node_id: str) -> dict[str, Any]:
            """Return the stored node payload."""
            return self._prompt[str(node_id)]

    class FakeModelPatcher:
        """Stand-in for ComfyUI's unhashable `ModelPatcher` runtime object."""

    prompt_one = {
        "39": {
            "class_type": "FakeSampler",
            "inputs": {},
        }
    }
    prompt_two = {
        "39": {
            "class_type": "FakeSampler",
            "inputs": {},
        }
    }
    prompt_three = {
        "39": {
            "class_type": "FakeSampler",
            "inputs": {},
        }
    }
    boundary_spec_one = [
        {
            "proxy_input_name": "remote_input_0",
            "io_type": "MODEL",
            "source_signature": "SRC_same_model",
            "targets": [{"node_id": "39", "input_name": "model"}],
        }
    ]
    boundary_spec_three = [
        {
            "proxy_input_name": "remote_input_0",
            "io_type": "MODEL",
            "source_signature": "SRC_other_model",
            "targets": [{"node_id": "39", "input_name": "model"}],
        }
    ]
    modal_cloud_module._apply_boundary_inputs(
        prompt_one,
        boundary_spec_one,
        {"remote_input_0": FakeModelPatcher()},
    )
    modal_cloud_module._apply_boundary_inputs(
        prompt_two,
        boundary_spec_one,
        {"remote_input_0": FakeModelPatcher()},
    )
    modal_cloud_module._apply_boundary_inputs(
        prompt_three,
        boundary_spec_three,
        {"remote_input_0": FakeModelPatcher()},
    )

    monkeypatch.setattr(
        modal_cloud_module,
        "_load_nodes_module",
        lambda: types.SimpleNamespace(
            NODE_CLASS_MAPPINGS={"FakeSampler": type("FakeSampler", (), {})}
        ),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_include_unique_id_in_input_signature",
        lambda class_type: False,
    )

    def cache_key_for(prompt: dict[str, Any]) -> str | None:
        """Build one distributed cache key for the prepared prompt."""
        cache_key_set = types.SimpleNamespace(
            dynprompt=FakeDynPrompt(prompt),
            is_changed_cache=types.SimpleNamespace(is_changed={"39": False}),
            get_ordered_ancestry=lambda current_dynprompt, node_id: ([], {}),
            include_node_id_in_input=lambda: False,
            get_data_key=lambda node_id: None,
        )
        return modal_cloud_module._node_output_cache_key_from_key_set_sync(cache_key_set, "39")

    first_key = cache_key_for(prompt_one)
    second_key = cache_key_for(prompt_two)
    different_key = cache_key_for(prompt_three)

    assert isinstance(first_key, str)
    assert first_key.startswith("NC_")
    assert second_key == first_key
    assert different_key != first_key


def test_modal_cloud_node_cache_key_includes_remote_session_bridge_key(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Mapped remote-session bridge inputs should not collide in the distributed cache."""

    class FakeDynPrompt:
        """Minimal dynamic prompt wrapper backed by one mutable prompt dict."""

        def __init__(self, prompt: dict[str, Any]) -> None:
            """Store the prompt used by the cache-key rebuild."""
            self._prompt = prompt

        def has_node(self, node_id: str) -> bool:
            """Return whether the requested node exists in the prompt."""
            return str(node_id) in self._prompt

        def get_node(self, node_id: str) -> dict[str, Any]:
            """Return the stored node payload."""
            return self._prompt[str(node_id)]

    class FakeConditioning:
        """Stand-in for a resolved CONDITIONING object."""

    def build_prompt(bridge_key: str) -> dict[str, Any]:
        """Build one sampler prompt with a resolved conditioning and original bridge ref metadata."""
        prompt = {
            "507": {
                "class_type": "FakeSampler",
                "inputs": {},
            }
        }
        boundary_spec = [
            {
                "proxy_input_name": "phase_bridge_3",
                "io_type": "CONDITIONING",
                "source_signature": "SRC_node_508_output_0",
                "targets": [{"node_id": "507", "input_name": "positive"}],
            }
        ]
        bridge_ref = modal_cloud_module.RemoteSessionBridgeRef(
            bridge_key=bridge_key,
            node_id="508",
            output_index=0,
            session_id="session_1",
        ).to_payload()
        modal_cloud_module._apply_boundary_inputs(
            prompt,
            boundary_spec,
            {"phase_bridge_3": FakeConditioning()},
            cache_signature_inputs={"phase_bridge_3": bridge_ref},
        )
        return prompt

    monkeypatch.setattr(
        modal_cloud_module,
        "_load_nodes_module",
        lambda: types.SimpleNamespace(
            NODE_CLASS_MAPPINGS={"FakeSampler": type("FakeSampler", (), {})}
        ),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_include_unique_id_in_input_signature",
        lambda class_type: False,
    )

    def cache_key_for(prompt: dict[str, Any]) -> str | None:
        """Build one distributed cache key for the prepared prompt."""
        cache_key_set = types.SimpleNamespace(
            dynprompt=FakeDynPrompt(prompt),
            is_changed_cache=types.SimpleNamespace(is_changed={"507": False}),
            get_ordered_ancestry=lambda current_dynprompt, node_id: ([], {}),
            include_node_id_in_input=lambda: False,
            get_data_key=lambda node_id: None,
        )
        return modal_cloud_module._node_output_cache_key_from_key_set_sync(cache_key_set, "507")

    first_key = cache_key_for(build_prompt("RSB_first"))
    same_key = cache_key_for(build_prompt("RSB_first"))
    second_key = cache_key_for(build_prompt("RSB_second"))

    assert isinstance(first_key, str)
    assert first_key.startswith("NC_")
    assert same_key == first_key
    assert second_key != first_key


def test_modal_cloud_node_cache_key_treats_nested_two_item_lists_as_data(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Nested conditioning-style two-item lists should not be mistaken for prompt links."""

    class FakeDynPrompt:
        """Minimal dynamic prompt wrapper for one data-only node."""

        def has_node(self, node_id: str) -> bool:
            """Return whether the requested node exists."""
            return str(node_id) == "545"

        def get_node(self, node_id: str) -> dict[str, Any]:
            """Return a node with nested list data shaped like Comfy conditioning."""
            if not self.has_node(node_id):
                raise KeyError(node_id)
            return {
                "class_type": "AnythingToMarkdown",
                "inputs": {
                    "anything": [
                        [
                            {"pooled_output": "summary"},
                            [["not-a-node-id"], ["not-an-output-index"]],
                        ]
                    ]
                },
            }

    cache_key_set = types.SimpleNamespace(
        dynprompt=FakeDynPrompt(),
        is_changed_cache=types.SimpleNamespace(is_changed={"545": False}),
        get_ordered_ancestry=lambda current_dynprompt, node_id: ([], {}),
        include_node_id_in_input=lambda: False,
        get_data_key=lambda node_id: None,
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_load_nodes_module",
        lambda: types.SimpleNamespace(
            NODE_CLASS_MAPPINGS={"AnythingToMarkdown": type("AnythingToMarkdown", (), {})}
        ),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_include_unique_id_in_input_signature",
        lambda class_type: False,
    )

    cache_key = modal_cloud_module._node_output_cache_key_from_key_set_sync(cache_key_set, "545")

    assert isinstance(cache_key, str)
    assert cache_key.startswith("NC_")


def test_modal_cloud_ignores_heavy_comfyui_paths(
    modal_cloud_module: Any,
) -> None:
    """The Modal cloud module should skip heavyweight ComfyUI runtime artifacts."""
    from pathlib import Path

    assert modal_cloud_module._should_ignore_comfyui_path(Path("models/checkpoint.safetensors"))
    assert modal_cloud_module._should_ignore_comfyui_path(Path("custom_nodes/example/__init__.py"))
    assert modal_cloud_module._should_ignore_comfyui_path(Path("output/run/output.png"))
    assert modal_cloud_module._should_ignore_comfyui_path(
        Path(".cache/strings/acfca22bde9a1a1fee53fe6e1299f4fe54a78a6f1d306dbb6cac2e71cf35d2c2.txt")
    )
    assert modal_cloud_module._should_ignore_comfyui_path(Path("__pycache__/execution.pyc"))
    assert modal_cloud_module._should_ignore_comfyui_path(Path("False/checkpoints/model.pth"))
    assert modal_cloud_module._should_ignore_comfyui_path(Path("tests/test_execution.py"))
    assert modal_cloud_module._should_ignore_comfyui_path(Path("unexpected/code.py"))
    assert not modal_cloud_module._should_ignore_comfyui_path(Path("execution.py"))
    assert not modal_cloud_module._should_ignore_comfyui_path(Path("requirements.txt"))
    assert not modal_cloud_module._should_ignore_comfyui_path(Path("comfy/model_management.py"))
    assert not modal_cloud_module._should_ignore_comfyui_path(Path("comfy/ldm/models/diffusion/ddpm.py"))


def test_modal_cloud_installs_comfyui_runtime_packages(
    modal_cloud_module: Any,
) -> None:
    """The Modal cloud image should include the core packages ComfyUI imports at runtime."""
    assert modal_cloud_module._comfyui_apt_packages() == (
        "build-essential",
        "libgl1",
        "libglib2.0-0",
    )

    packages = modal_cloud_module._comfyui_runtime_packages()
    package_names = {package.split("==", maxsplit=1)[0] for package in packages}

    assert all("==" in package for package in packages)
    assert {
        "aiohttp",
        "alembic",
        "blake3",
        "comfy-angle",
        "comfy-aimdo",
        "comfy-kitchen",
        "hf-xet",
        "huggingface-hub",
        "kornia",
        "opencv-python-headless",
        "pydantic-settings",
        "psutil",
        "pyopengl",
        "sentencepiece",
        "simpleeval",
        "spandrel",
        "torchsde",
        "transformers",
    } <= package_names


def test_modal_cloud_collects_custom_node_runtime_packages(
    modal_cloud_module: Any,
    tmp_path: Path,
) -> None:
    """The Modal image should install requirements declared by bundled custom nodes."""
    custom_nodes_dir = tmp_path / "custom_nodes"
    first_package = custom_nodes_dir / "first"
    second_package = custom_nodes_dir / "second"
    first_package.mkdir(parents=True)
    second_package.mkdir(parents=True)
    (first_package / "requirements.txt").write_text(
        "\n".join(
            [
                "# comments are ignored",
                "omegaconf==2.3.0",
                "hydra-core==1.3.2 # inline comment",
                "--extra-index-url https://example.invalid/simple",
                "-r extra-requirements.txt",
            ]
        ),
        encoding="utf-8",
    )
    (first_package / "extra-requirements.txt").write_text(
        "diffusers>=0.37.0\n",
        encoding="utf-8",
    )
    (second_package / "requirements.txt").write_text(
        "omegaconf==2.3.0\nsoundfile\n",
        encoding="utf-8",
    )

    packages = modal_cloud_module._custom_node_runtime_packages(custom_nodes_dir)

    assert packages == (
        "omegaconf==2.3.0",
        "hydra-core==1.3.2",
        "diffusers>=0.37.0",
        "soundfile",
    )


def test_modal_cloud_returns_no_custom_node_packages_without_requirements(
    modal_cloud_module: Any,
    tmp_path: Path,
) -> None:
    """Custom-node package installation should be skipped when no requirements exist."""
    custom_nodes_dir = tmp_path / "custom_nodes"
    (custom_nodes_dir / "example").mkdir(parents=True)

    assert modal_cloud_module._custom_node_runtime_packages(custom_nodes_dir) == ()


def test_modal_cloud_selects_gpu_compatible_pytorch_stack(
    modal_cloud_module: Any,
) -> None:
    """Every Modal GPU image should expose the shared pinned PyTorch build."""
    default_build = modal_cloud_module._select_remote_torch_build("A100")
    b300_build = modal_cloud_module._select_remote_torch_build("B300")

    assert default_build.install_layers[0].packages == (
        "torch==2.13.0",
        "torchvision==0.28.0",
    )
    assert default_build.install_layers[0].index_url == "https://download.pytorch.org/whl/cu130"
    assert default_build == b300_build
    assert b300_build.install_layers[0].packages == (
        "torch==2.13.0",
        "torchvision==0.28.0",
    )
    assert b300_build.install_layers[0].index_url == "https://download.pytorch.org/whl/cu130"
    assert b300_build.install_layers[1].packages == ("torchaudio==2.11.0+cpu",)
    assert b300_build.install_layers[1].index_url == "https://download.pytorch.org/whl/cpu"
    assert b300_build.install_layers[1].extra_options == "--no-deps"


def test_modal_cloud_installs_and_validates_torch_layers_in_order(
    modal_cloud_module: Any,
) -> None:
    """The CPU TorchAudio layer must not replace the CUDA-enabled Torch package."""

    class RecordingImage:
        """Record Modal image package and command layers."""

        def __init__(self) -> None:
            """Initialize an empty call record."""
            self.calls: list[tuple[str, tuple[Any, ...], dict[str, Any]]] = []

        def pip_install(self, *packages: str, **options: Any) -> "RecordingImage":
            """Record one package installation layer."""
            self.calls.append(("pip_install", packages, options))
            return self

        def run_commands(self, *commands: str) -> "RecordingImage":
            """Record one image validation command."""
            self.calls.append(("run_commands", commands, {}))
            return self

    build = modal_cloud_module._select_remote_torch_build("B300")
    image = RecordingImage()

    result = modal_cloud_module._install_remote_torch_build(image, build)

    assert result is image
    assert image.calls[:2] == [
        (
            "pip_install",
            ("torch==2.13.0", "torchvision==0.28.0"),
            {
                "index_url": "https://download.pytorch.org/whl/cu130",
                "extra_options": "",
            },
        ),
        (
            "pip_install",
            ("torchaudio==2.11.0+cpu",),
            {
                "index_url": "https://download.pytorch.org/whl/cpu",
                "extra_options": "--no-deps",
            },
        ),
    ]
    assert image.calls[2][0] == "run_commands"
    assert "import torch, torchaudio, torchvision" in image.calls[2][1][0]


def test_modal_cloud_installs_and_validates_vllm_for_non_b300_gpu(
    modal_cloud_module: Any,
) -> None:
    """A lower-cost GPU image should install and import the pinned vLLM wheel."""

    class RecordingImage:
        """Record Modal image package and command layers."""

        def __init__(self) -> None:
            """Initialize an empty call record."""
            self.calls: list[tuple[str, tuple[Any, ...], dict[str, Any]]] = []

        def pip_install(self, *packages: str, **options: Any) -> "RecordingImage":
            """Record one package installation layer."""
            self.calls.append(("pip_install", packages, options))
            return self

        def run_commands(self, *commands: str) -> "RecordingImage":
            """Record one image validation command."""
            self.calls.append(("run_commands", commands, {}))
            return self

    image = RecordingImage()

    result = modal_cloud_module._install_remote_accelerator_packages(
        image,
        "RTX-PRO-6000",
    )

    assert result is image
    assert image.calls[0] == (
        "pip_install",
        (modal_cloud_module._remote_accelerator_packages("RTX-PRO-6000")[0],),
        {},
    )
    assert image.calls[1][0] == "run_commands"
    assert "import cv2, numpy, torch, vllm" in image.calls[1][1][0]


def test_modal_cloud_restores_runtime_pins_after_custom_node_packages(
    modal_cloud_module: Any,
) -> None:
    """Custom requirements must not remain authoritative for shared runtime packages."""

    class RecordingImage:
        """Record Modal image package layers."""

        def __init__(self) -> None:
            """Initialize an empty call record."""
            self.calls: list[tuple[str, ...]] = []

        def pip_install(self, *packages: str, **_options: Any) -> "RecordingImage":
            """Record one package installation layer."""
            self.calls.append(packages)
            return self

    image = RecordingImage()

    result = modal_cloud_module._install_custom_node_packages(
        image,
        ("numpy==1.26.4", "opencv-python-headless==4.11.0.86"),
    )

    assert result is image
    assert image.calls[0] == (
        "numpy==1.26.4",
        "opencv-python-headless==4.11.0.86",
    )
    assert "numpy==2.3.5" in image.calls[1]
    assert "opencv-python-headless==4.13.0.92" in image.calls[1]


def test_modal_cloud_missing_prompt_node_class_raises_clear_error(
    modal_cloud_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """Missing custom-node classes should fail before ComfyUI cache setup raises KeyError."""
    async def fake_init_external_custom_nodes() -> None:
        """Leave the missing custom node unregistered."""

    fake_nodes_module = types.SimpleNamespace(
        NODE_CLASS_MAPPINGS={"KnownNode": object},
        init_external_custom_nodes=fake_init_external_custom_nodes,
    )
    monkeypatch.setitem(sys.modules, "nodes", fake_nodes_module)
    custom_nodes_root = tmp_path / "custom_nodes"
    package_dir = custom_nodes_root / "Skoogeer-Noise" / "src"
    package_dir.mkdir(parents=True)
    (package_dir / "qwen_noise_nodes.py").write_text(
        "class KSamplerLoraSigmaInverse:\n    pass\n",
        encoding="utf-8",
    )

    with pytest.raises(modal_cloud_module.RemoteSubgraphExecutionError) as exc_info:
        modal_cloud_module._ensure_prompt_node_classes_registered(
            component_id="component-1",
            prompt={
                "1": {"class_type": "KnownNode", "inputs": {}},
                "2": {"class_type": "KSamplerLoraSigmaInverse", "inputs": {}},
            },
            custom_nodes_root=custom_nodes_root,
        )

    message = str(exc_info.value)
    assert "KSamplerLoraSigmaInverse" in message
    assert "custom-node sync is enabled" in message
    assert "Skoogeer-Noise/src/qwen_noise_nodes.py" in message
    assert "package=Skoogeer-Noise" in message


def test_modal_cloud_missing_prompt_node_class_reports_unavailable_bundle_path(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Missing-class diagnostics should distinguish absent payload metadata from unavailable storage."""
    fake_nodes_module = types.SimpleNamespace(NODE_CLASS_MAPPINGS={})
    monkeypatch.setitem(sys.modules, "nodes", fake_nodes_module)

    with pytest.raises(modal_cloud_module.RemoteSubgraphExecutionError) as exc_info:
        modal_cloud_module._ensure_prompt_node_classes_registered(
            component_id="component-1",
            prompt={"2": {"class_type": "KSamplerLoraSigmaInverse", "inputs": {}}},
            custom_nodes_root=None,
            custom_nodes_bundle_path="/custom_nodes/manifest.json",
        )

    message = str(exc_info.value)
    assert "custom_nodes_bundle='/custom_nodes/manifest.json'" in message
    assert "not available in Modal worker storage" in message


def test_modal_cloud_retries_custom_node_import_for_missing_prompt_class(
    modal_cloud_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """A missing prompt class should trigger one external custom-node import retry."""
    import_calls: list[str] = []
    custom_nodes_root = tmp_path / "custom_nodes_bundle"
    custom_nodes_root.mkdir()

    class RegisteredLaterNode:
        """Node type registered by the retry import."""

    async def fake_init_external_custom_nodes() -> None:
        """Register the custom class when the retry import runs."""
        import_calls.append("init")
        fake_nodes_module.NODE_CLASS_MAPPINGS["KSamplerLoraSigmaInverse"] = RegisteredLaterNode

    fake_nodes_module = types.SimpleNamespace(
        NODE_CLASS_MAPPINGS={},
        init_external_custom_nodes=fake_init_external_custom_nodes,
    )
    monkeypatch.setitem(sys.modules, "nodes", fake_nodes_module)
    monkeypatch.setattr(modal_cloud_module, "_register_custom_nodes_root", lambda path: import_calls.append(str(path)))
    monkeypatch.setattr(modal_cloud_module, "_install_loader_cache_wrappers", lambda: import_calls.append("wrappers"))

    resolved_mapping = modal_cloud_module._ensure_prompt_node_classes_registered(
        component_id="component-1",
        prompt={"2": {"class_type": "KSamplerLoraSigmaInverse", "inputs": {}}},
        custom_nodes_root=custom_nodes_root,
    )

    assert resolved_mapping["KSamplerLoraSigmaInverse"] is RegisteredLaterNode
    assert import_calls == [str(custom_nodes_root), "init", "wrappers"]


def test_modal_cloud_existing_app_guard_uses_non_creating_sdk_lookup(
    modal_cloud_module: Any,
) -> None:
    """Deploy-time app construction should fail if the configured Modal app already exists."""

    class FakeNotFoundError(Exception):
        """Stand-in for Modal app lookup misses."""

    lookup_calls: list[tuple[str, bool]] = []

    class FakeApp:
        """Minimal Modal App namespace with non-creating lookup support."""

        @staticmethod
        def lookup(app_name: str, create_if_missing: bool = True) -> object:
            """Record lookup arguments and return an existing app."""
            lookup_calls.append((app_name, create_if_missing))
            return object()

    fake_modal = types.SimpleNamespace(
        App=FakeApp,
        exception=types.SimpleNamespace(NotFoundError=FakeNotFoundError),
        is_local=lambda: True,
    )
    settings = types.SimpleNamespace(app_name="comfy-modal-sync")

    with pytest.raises(modal_cloud_module.ExistingModalAppError) as exc_info:
        modal_cloud_module._guard_against_existing_modal_app(settings, fake_modal)

    assert lookup_calls == [(DEFAULT_TEST_DEPLOYMENT_APP_NAME, False)]
    assert "Delete the existing app" in str(exc_info.value)
    assert "COMFY_MODAL_GPU" in str(exc_info.value)


def test_modal_cloud_existing_app_guard_allows_missing_sdk_lookup(
    modal_cloud_module: Any,
) -> None:
    """A missing Modal app should not block first-run app construction."""

    class FakeNotFoundError(Exception):
        """Stand-in for Modal app lookup misses."""

    class FakeApp:
        """Minimal Modal App namespace that reports a missing app."""

        @staticmethod
        def lookup(app_name: str, create_if_missing: bool = True) -> object:
            """Raise the SDK's not-found error without creating the app."""
            del app_name, create_if_missing
            raise FakeNotFoundError("app not found")

    fake_modal = types.SimpleNamespace(
        App=FakeApp,
        exception=types.SimpleNamespace(NotFoundError=FakeNotFoundError),
        is_local=lambda: True,
    )

    modal_cloud_module._guard_against_existing_modal_app(
        types.SimpleNamespace(app_name="comfy-modal-sync"),
        fake_modal,
    )


def test_modal_cloud_existing_app_guard_falls_back_to_cli_json(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """SDKs without a non-creating lookup should fall back to `modal app list --json`."""

    class FakeApp:
        """Modal App namespace whose lookup signature cannot safely check existence."""

        @staticmethod
        def lookup(app_name: str) -> object:
            """This must not be called because it may create the app."""
            raise AssertionError(f"unsafe lookup called for {app_name}")

    completed = subprocess.CompletedProcess(
        args=["modal", "app", "list", "--json"],
        returncode=0,
        stdout=json.dumps([{"name": "other"}, {"name": DEFAULT_TEST_DEPLOYMENT_APP_NAME}]),
        stderr="",
    )
    observed_commands: list[list[str]] = []

    def fake_run(command: list[str], **kwargs: Any) -> subprocess.CompletedProcess[str]:
        """Return a deterministic Modal CLI app list response."""
        del kwargs
        observed_commands.append(command)
        return completed

    monkeypatch.setattr(modal_cloud_module.shutil, "which", lambda name: f"/usr/bin/{name}")
    monkeypatch.setattr(modal_cloud_module.subprocess, "run", fake_run)
    fake_modal = types.SimpleNamespace(App=FakeApp, exception=types.SimpleNamespace(), is_local=lambda: True)

    with pytest.raises(modal_cloud_module.ExistingModalAppError):
        modal_cloud_module._guard_against_existing_modal_app(
            types.SimpleNamespace(app_name="comfy-modal-sync"),
            fake_modal,
        )

    assert observed_commands == [["/usr/bin/modal", "app", "list", "--json"]]


def test_modal_cloud_existing_app_guard_skips_remote_container_import(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """The guard should not run inside Modal worker containers that import the same module."""

    class FakeApp:
        """Modal App namespace that would fail if the guard tried a lookup."""

        @staticmethod
        def lookup(app_name: str, create_if_missing: bool = True) -> object:
            """Raise if remote-container imports accidentally perform local deploy checks."""
            del app_name, create_if_missing
            raise AssertionError("remote container import should not query Modal apps")

    monkeypatch.setenv("MODAL_TASK_ID", "ta-123")
    fake_modal = types.SimpleNamespace(App=FakeApp, exception=types.SimpleNamespace(), is_local=lambda: False)

    modal_cloud_module._guard_against_existing_modal_app(
        types.SimpleNamespace(app_name="comfy-modal-sync"),
        fake_modal,
    )


def test_modal_cloud_builds_snapshot_enabled_cls_options(
    modal_cloud_module: Any,
) -> None:
    """The remote engine should default to CPU memory snapshots and optional GPU snapshots."""
    base_settings = types.SimpleNamespace(
        remote_storage_root="/storage",
        enable_memory_snapshot=True,
        enable_gpu_memory_snapshot=False,
        modal_gpu="L40S",
        scaledown_window_seconds=600,
        min_containers=0,
    )

    modal_secret = object()
    options = modal_cloud_module._remote_engine_cls_options(
        base_settings,
        "volume",
        "image",
        modal_secret,
        "compile-cache-volume",
    )

    assert options["enable_memory_snapshot"] is True
    assert "experimental_options" not in options
    assert options["gpu"] == "L40S"
    assert options["volumes"] == {
        "/storage": "volume",
        "/root/.cache/comfy-modal-llm": "compile-cache-volume",
    }
    assert options["scaledown_window"] == 600
    assert options["min_containers"] == 0
    assert options["timeout"] == 3600
    assert options["startup_timeout"] == 900
    assert options["secrets"] == [modal_secret]

    gpu_snapshot_settings = types.SimpleNamespace(
        remote_storage_root="/storage",
        enable_memory_snapshot=True,
        enable_gpu_memory_snapshot=True,
        modal_gpu="A100",
        scaledown_window_seconds=900,
        min_containers=1,
    )
    gpu_snapshot_options = modal_cloud_module._remote_engine_cls_options(
        gpu_snapshot_settings,
        "volume",
        "image",
    )
    assert gpu_snapshot_options["experimental_options"] == {"enable_gpu_snapshot": True}
    assert gpu_snapshot_options["gpu"] == "A100"
    assert gpu_snapshot_options["scaledown_window"] == 900
    assert gpu_snapshot_options["min_containers"] == 1


def test_modal_cloud_prewarms_snapshot_state_without_gpu_runtime_by_default(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """CPU-only snapshot prewarm should avoid full ComfyUI runtime initialization."""
    calls: list[str] = []

    monkeypatch.setattr(
        modal_cloud_module,
        "_ensure_comfyui_support_packages",
        lambda: calls.append("support"),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_ensure_comfy_runtime_initialized",
        lambda custom_nodes_root: calls.append(f"runtime:{custom_nodes_root}"),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_load_execution_module",
        lambda: calls.append("execution"),
    )

    modal_cloud_module._prewarm_snapshot_state(
        gpu_snapshot_enabled=False,
    )

    assert calls == ["support"]


def test_modal_cloud_prewarms_snapshot_loader_profile_when_gpu_snapshots_enabled(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """GPU snapshot prewarm should execute registered loader plans for one snapshot profile."""
    calls: list[tuple[str, Any]] = []

    monkeypatch.setattr(
        modal_cloud_module,
        "_ensure_comfyui_support_packages",
        lambda: calls.append(("support", None)),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_ensure_comfy_runtime_initialized",
        lambda custom_nodes_root: calls.append(("runtime", custom_nodes_root)),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_load_execution_module",
        lambda: calls.append(("execution", None)),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_load_loader_snapshot_profile",
        lambda snapshot_profile_key: [
            {
                "signature": "loader-plan-1",
                "node_id": "7",
                "subgraph_prompt": {"7": {"class_type": "UNETLoader", "inputs": {"unet_name": "model.safetensors"}}},
                "execute_node_ids": ["7"],
            }
        ]
        if snapshot_profile_key == "loader-profile:abc"
        else [],
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_execute_loader_prewarm_plans",
        lambda *, component_id, loader_prewarm_plans, custom_nodes_root: calls.append(
            ("prewarm", component_id, tuple(plan["signature"] for plan in loader_prewarm_plans), custom_nodes_root)
        ),
    )

    modal_cloud_module._prewarm_snapshot_state(
        gpu_snapshot_enabled=True,
        snapshot_profile_key="loader-profile:abc",
    )

    assert calls == [
        ("support", None),
        ("runtime", None),
        ("execution", None),
        ("prewarm", "snapshot-profile:loader-profile:abc", ("loader-plan-1",), None),
    ]


def test_modal_cloud_skips_generic_gpu_snapshot_prewarm_without_profile(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """GPU snapshot prewarm should skip runtime init when no stable snapshot profile exists."""
    calls: list[str] = []

    monkeypatch.setattr(
        modal_cloud_module,
        "_ensure_comfyui_support_packages",
        lambda: calls.append("support"),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_ensure_comfy_runtime_initialized",
        lambda custom_nodes_root: calls.append(f"runtime:{custom_nodes_root}"),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_load_execution_module",
        lambda: calls.append("execution"),
    )

    modal_cloud_module._prewarm_snapshot_state(
        gpu_snapshot_enabled=True,
    )

    assert calls == ["support"]


def test_modal_cloud_prewarms_restored_runtime(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Post-restore prewarm should fully initialize the request-serving runtime."""
    calls: list[str] = []

    monkeypatch.setattr(
        modal_cloud_module,
        "_ensure_comfy_runtime_initialized",
        lambda custom_nodes_root: calls.append(f"runtime:{custom_nodes_root}"),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_load_execution_module",
        lambda: calls.append("execution"),
    )

    modal_cloud_module._prewarm_restored_runtime()

    assert calls == ["runtime:None", "execution"]


def test_modal_cloud_prepares_warm_container_for_request(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Warmup requests should prime volume visibility and extracted custom nodes without executing a payload."""
    calls: list[tuple[str, Any]] = []

    monkeypatch.setattr(modal_cloud_module, "_should_reload_modal_volume", lambda payload: True)
    monkeypatch.setattr(
        modal_cloud_module,
        "_reload_modal_volume_for_request",
        lambda volume, component_id, reload_marker=None, payload=None: calls.append(
            ("reload", component_id, reload_marker, payload.get("uploaded_volume_paths"))
        ),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_extract_custom_nodes_bundle",
        lambda bundle_path: Path("/tmp/extracted-bundle") if bundle_path else None,
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_register_custom_nodes_root",
        lambda custom_nodes_root: calls.append(("register", custom_nodes_root)),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_execute_loader_prewarm_plans",
        lambda *, component_id, loader_prewarm_plans, custom_nodes_root: calls.append(
            ("prewarm", component_id, tuple(plan["signature"] for plan in loader_prewarm_plans), custom_nodes_root)
        ),
    )
    monkeypatch.setenv("MODAL_TASK_ID", "task-123")

    result = modal_cloud_module._prepare_warm_container_for_request(
        object(),
        {
            "component_id": "component-1",
            "volume_reload_marker": "marker-1",
            "uploaded_volume_paths": ["/storage/example.bin"],
            "custom_nodes_bundle": "custom_nodes_bundle.zip",
            "warmup_slot_index": 2,
            "loader_prewarm_plans": [{"signature": "loader-plan-1"}],
        },
    )

    assert calls == [
        ("reload", "component-1", "marker-1", ["/storage/example.bin"]),
        ("register", Path("/tmp/extracted-bundle")),
        ("prewarm", "component-1", ("loader-plan-1",), Path("/tmp/extracted-bundle")),
    ]
    assert result == {
        "component_id": "component-1",
        "task_id": "task-123",
        "warmup_slot_index": 2,
        "reloaded_volume": True,
        "llm_prewarm_results": [],
    }


def test_modal_cloud_executes_loader_prewarm_plans_once_per_worker(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Worker-local loader prewarm plans should execute once and then be skipped on reuse."""
    original_plan_keys = set(modal_cloud_module._LOADER_PREWARM_PLAN_KEYS)
    modal_cloud_module._LOADER_PREWARM_PLAN_KEYS.clear()
    observed_calls: list[tuple[str, tuple[str, ...]]] = []

    monkeypatch.setattr(
        modal_cloud_module,
        "_ensure_comfy_runtime_initialized",
        lambda custom_nodes_root: observed_calls.append(("runtime", (str(custom_nodes_root),))),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_execute_subgraph_prompt",
        lambda payload, hydrated_inputs, custom_nodes_root, **kwargs: observed_calls.append(
            ("execute", (str(payload["component_id"]), str(tuple(payload["execute_node_ids"]))))
        ) or tuple(),
    )
    monkeypatch.setenv("COMFY_MODAL_ENABLE_LOADER_PREWARM", "true")
    modal_cloud_module.get_settings.cache_clear()
    try:
        plans = [
            {
                "signature": "loader-plan-1",
                "node_id": "7",
                "prompt_id": "prompt-1",
                "subgraph_prompt": {"7": {"class_type": "UNETLoader", "inputs": {"unet_name": "model.safetensors"}}},
                "execute_node_ids": ["7"],
            }
        ]
        modal_cloud_module._execute_loader_prewarm_plans(
            component_id="component-1",
            loader_prewarm_plans=plans,
            custom_nodes_root=Path("/tmp/extracted-bundle"),
        )
        modal_cloud_module._execute_loader_prewarm_plans(
            component_id="component-1",
            loader_prewarm_plans=plans,
            custom_nodes_root=Path("/tmp/extracted-bundle"),
        )
    finally:
        modal_cloud_module.get_settings.cache_clear()
        modal_cloud_module._LOADER_PREWARM_PLAN_KEYS.clear()
        modal_cloud_module._LOADER_PREWARM_PLAN_KEYS.update(original_plan_keys)

    assert observed_calls == [
        ("runtime", ("/tmp/extracted-bundle",)),
        ("execute", ("component-1::loader-prewarm:7", "('7',)")),
        ("runtime", ("/tmp/extracted-bundle",)),
    ]


def test_modal_cloud_parallelizes_independent_loader_prewarms(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Independent CLIP and UNET loads should overlap under bounded concurrency."""
    active_count = 0
    maximum_active_count = 0
    counter_lock = threading.Lock()
    original_plan_keys = set(modal_cloud_module._LOADER_PREWARM_PLAN_KEYS)
    monkeypatch.setattr(
        modal_cloud_module,
        "_ensure_comfy_runtime_initialized",
        lambda custom_nodes_root: None,
    )

    def execute_plan(*args: Any, **kwargs: Any) -> tuple[Any, ...]:
        """Record overlap while simulating one heavyweight loader."""
        del args, kwargs
        nonlocal active_count, maximum_active_count
        with counter_lock:
            active_count += 1
            maximum_active_count = max(maximum_active_count, active_count)
        time.sleep(0.04)
        with counter_lock:
            active_count -= 1
        return ()

    monkeypatch.setattr(modal_cloud_module, "_execute_subgraph_prompt", execute_plan)
    monkeypatch.setenv("COMFY_MODAL_ENABLE_LOADER_PREWARM", "true")
    monkeypatch.setenv("COMFY_MODAL_LOADER_PREWARM_WORKERS", "2")
    modal_cloud_module.get_settings.cache_clear()
    plans = [
        {
            "signature": f"parallel-loader-{index}",
            "node_id": str(index),
            "class_type": class_type,
            "subgraph_prompt": {
                str(index): {"class_type": class_type, "inputs": {}}
            },
            "execute_node_ids": [str(index)],
        }
        for index, class_type in enumerate(("CLIPLoader", "UNETLoader"), start=1)
    ]
    try:
        modal_cloud_module._execute_loader_prewarm_plans(
            component_id="parallel-component",
            loader_prewarm_plans=plans,
            custom_nodes_root=None,
        )
    finally:
        modal_cloud_module.get_settings.cache_clear()
        modal_cloud_module._LOADER_PREWARM_PLAN_KEYS.clear()
        modal_cloud_module._LOADER_PREWARM_PLAN_KEYS.update(original_plan_keys)

    assert maximum_active_count == 2


def test_modal_cloud_llm_prewarm_commits_content_addressed_manifest(
    modal_cloud_module: Any,
    modal_llm_runtime_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """A representative LLM cache miss should publish and commit its artifacts."""
    observed_requests: list[tuple[str, int, str | None]] = []
    compile_miss_signal = {"size": 0}

    def prewarm_profile(
        *,
        model_profile: str,
        representative_request_count: int,
        workflow_execution_id: str | None,
    ) -> dict[str, Any]:
        """Return deterministic warmup telemetry."""
        observed_requests.append(
            (model_profile, representative_request_count, workflow_execution_id)
        )
        compile_miss_signal["size"] += 100
        return {
            "profile_id": model_profile,
            "representative_request_count": representative_request_count,
        }

    class FakeVolume:
        """Count durable compile-cache commits."""

        def __init__(self) -> None:
            """Initialize counters."""
            self.commits = 0

        def commit(self) -> None:
            """Record one commit."""
            self.commits += 1

    cloud_llm_runtime_module = sys.modules.get("modal_llm_runtime")
    if cloud_llm_runtime_module is None:
        cloud_llm_runtime_module = modal_llm_runtime_module
        monkeypatch.setitem(
            sys.modules,
            "modal_llm_runtime",
            cloud_llm_runtime_module,
        )
    monkeypatch.setattr(
        cloud_llm_runtime_module,
        "prewarm_modal_llm_profile",
        prewarm_profile,
    )
    monkeypatch.setattr(
        cloud_llm_runtime_module,
        "triton_compile_miss_signal_size",
        lambda: compile_miss_signal["size"],
    )
    monkeypatch.setattr(
        cloud_llm_runtime_module,
        "triton_compile_listener_engine_pids",
        lambda: (321,),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_llm_compile_manifest_path",
        lambda signature: tmp_path / "manifests" / f"{signature}.json",
    )
    original_plan_keys = set(modal_cloud_module._LLM_PREWARM_PLAN_KEYS)
    modal_cloud_module._LLM_PREWARM_PLAN_KEYS.clear()
    volume = FakeVolume()
    try:
        results = modal_cloud_module._execute_llm_prewarm_plans(
            component_id="llm-component",
            prompt_id="prompt-1",
            llm_prewarm_plans=[
                {
                    "signature": "plan-signature",
                    "model_profile": "old-profile",
                    "representative_request_count": 3,
                    "prompt_node": {
                        "class_type": "ModalLLM",
                        "inputs": {"model_profile": "staged-profile"},
                    },
                }
            ],
            compile_cache_volume=volume,
        )
    finally:
        modal_cloud_module._LLM_PREWARM_PLAN_KEYS.clear()
        modal_cloud_module._LLM_PREWARM_PLAN_KEYS.update(original_plan_keys)

    assert observed_requests == [("staged-profile", 3, "prompt-1")]
    assert volume.commits == 1
    assert Path(results[0]["manifest_path"]).exists()
    assert results[0]["manifest_cache_hit"] is False
    assert results[0]["compile_cache_committed"] is True


def test_modal_cloud_llm_prewarm_does_not_commit_persistent_cache_hit(
    modal_cloud_module: Any,
    modal_llm_runtime_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """A successful representative request with only disk hits must not commit."""

    class FakeVolume:
        """Reject an unnecessary persistent cache commit."""

        def commit(self) -> None:
            """Fail if a disk-cache hit is committed."""
            raise AssertionError("persistent cache hits must not be committed")

    cloud_llm_runtime_module = sys.modules.get(
        "modal_llm_runtime",
        modal_llm_runtime_module,
    )
    monkeypatch.setitem(sys.modules, "modal_llm_runtime", cloud_llm_runtime_module)
    monkeypatch.setattr(
        cloud_llm_runtime_module,
        "prewarm_modal_llm_profile",
        lambda **_kwargs: {"profile_id": "cached-profile"},
    )
    monkeypatch.setattr(
        cloud_llm_runtime_module,
        "triton_compile_miss_signal_size",
        lambda: 50,
    )
    monkeypatch.setattr(
        cloud_llm_runtime_module,
        "triton_compile_listener_engine_pids",
        lambda: (654,),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_llm_compile_manifest_path",
        lambda signature: tmp_path / "manifests" / f"{signature}.json",
    )
    original_plan_keys = set(modal_cloud_module._LLM_PREWARM_PLAN_KEYS)
    modal_cloud_module._LLM_PREWARM_PLAN_KEYS.clear()
    try:
        results = modal_cloud_module._execute_llm_prewarm_plans(
            component_id="llm-component",
            prompt_id="prompt-hit",
            llm_prewarm_plans=[
                {
                    "signature": "cached-plan",
                    "model_profile": "cached-profile",
                    "representative_request_count": 1,
                }
            ],
            compile_cache_volume=FakeVolume(),
        )
    finally:
        modal_cloud_module._LLM_PREWARM_PLAN_KEYS.clear()
        modal_cloud_module._LLM_PREWARM_PLAN_KEYS.update(original_plan_keys)

    assert results[0]["compile_cache_committed"] is False


def test_modal_cloud_llm_compile_profiles_are_limited_to_executable_subgraph(
    modal_cloud_module: Any,
) -> None:
    """Nested metadata and disconnected LLM nodes must not trigger cache commits."""
    payload = {
        "payload_kind": "subgraph",
        "execute_node_ids": ["3"],
        "subgraph_prompt": {
            "1": {
                "class_type": "ModalLLM",
                "inputs": {"model_profile": "executed-profile"},
            },
            "2": {
                "class_type": "ModalLLM",
                "inputs": {"model_profile": "disconnected-profile"},
            },
            "3": {
                "class_type": "PreviewText",
                "inputs": {"text": ["1", 0]},
            },
        },
        "boundary_inputs": [
            {
                "source_signature": {
                    "class_type": "ModalLLM",
                    "inputs": {"model_profile": "metadata-profile"},
                }
            }
        ],
    }

    assert modal_cloud_module._llm_profiles_in_payload(payload) == (
        "executed-profile",
    )
    assert (
        modal_cloud_module._llm_profiles_in_payload(
            {**payload, "payload_kind": "canary"}
        )
        == ()
    )


def test_modal_cloud_commits_after_each_genuine_triton_cache_miss(
    modal_cloud_module: Any,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Repeated new specializations for one profile must each become durable."""
    signal = {"size": 10}
    runtime_module = types.ModuleType("modal_llm_runtime")
    runtime_module.triton_compile_miss_signal_size = lambda: signal["size"]
    runtime_module.triton_compile_listener_engine_pids = lambda: (1234,)
    monkeypatch.setitem(sys.modules, "modal_llm_runtime", runtime_module)
    payload = {
        "payload_kind": "subgraph",
        "execute_node_ids": ["llm"],
        "subgraph_prompt": {
            "llm": {
                "class_type": "ModalLLM",
                "inputs": {"model_profile": "same-profile"},
            }
        },
    }

    class FakeVolume:
        """Count explicit compile-cache commits."""

        def __init__(self) -> None:
            """Initialize an empty commit count."""
            self.commits = 0

        def commit(self) -> None:
            """Record one explicit commit."""
            self.commits += 1

    volume = FakeVolume()
    checkpoint = modal_cloud_module._llm_compile_miss_checkpoint(payload)
    assert checkpoint is not None
    assert (
        modal_cloud_module._commit_actual_llm_compile_cache(checkpoint, volume)
        is False
    )

    signal["size"] = 20
    assert modal_cloud_module._commit_actual_llm_compile_cache(checkpoint, volume)
    second_checkpoint = modal_cloud_module._llm_compile_miss_checkpoint(payload)
    signal["size"] = 35
    assert modal_cloud_module._commit_actual_llm_compile_cache(
        second_checkpoint,
        volume,
    )
    assert volume.commits == 2


def test_modal_cloud_does_not_claim_cache_hits_without_engine_listener(
    modal_cloud_module: Any,
    monkeypatch: pytest.MonkeyPatch,
    caplog: pytest.LogCaptureFixture,
) -> None:
    """Missing EngineCore telemetry must skip commits without reporting a hit."""
    runtime_module = types.ModuleType("modal_llm_runtime")
    runtime_module.triton_compile_miss_signal_size = lambda: 0
    runtime_module.triton_compile_listener_engine_pids = lambda: ()
    monkeypatch.setitem(sys.modules, "modal_llm_runtime", runtime_module)
    checkpoint = modal_cloud_module._LLMCompileMissCheckpoint(
        profiles=("profile",),
        signal_size=0,
        listener_engine_pids=(),
    )

    class FakeVolume:
        """Reject any commit without listener evidence."""

        def commit(self) -> None:
            """Fail if missing telemetry causes a commit."""
            raise AssertionError("listener-less execution cannot commit")

    with caplog.at_level(logging.WARNING):
        committed = modal_cloud_module._commit_actual_llm_compile_cache(
            checkpoint,
            FakeVolume(),
        )

    assert committed is False
    assert "no live vLLM EngineCore" in caplog.text
    assert "every Triton lookup hit" not in caplog.text


def test_modal_cloud_reloads_compile_cache_before_restored_runtime(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Restored workers should refresh the cache Volume before importing runtimes."""
    calls: list[str] = []

    class FakeVolume:
        """Record one reload."""

        def reload(self) -> None:
            """Record the refresh."""
            calls.append("reload")

    monkeypatch.setattr(
        modal_cloud_module,
        "_ensure_comfy_runtime_initialized",
        lambda custom_nodes_root: calls.append("runtime"),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_load_execution_module",
        lambda: calls.append("execution"),
    )

    modal_cloud_module._prewarm_restored_runtime(FakeVolume())

    assert calls == ["reload", "runtime", "execution"]


def test_modal_cloud_does_not_reload_compile_cache_during_request_warmup(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Do not reload a cache Volume after native libraries may be mapped."""

    class FakeCompileCacheVolume:
        """Fail if request-time warmup tries to reload the compile cache."""

        def reload(self) -> None:
            """Reject the unsafe request-time reload."""
            raise AssertionError("request-time compile-cache reload is unsafe")

    monkeypatch.setattr(
        modal_cloud_module,
        "_hydrate_missing_payload_volume_paths",
        lambda volume, payload: [],
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_should_reload_modal_volume",
        lambda payload: False,
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_emit_modal_volume_reload_skip",
        lambda component_id, payload: None,
    )

    result = modal_cloud_module._prepare_warm_container_for_request(
        object(),
        {"component_id": "llm-component"},
        FakeCompileCacheVolume(),
    )

    assert result["component_id"] == "llm-component"


def test_modal_cloud_finds_memory_mapped_compile_cache_files(
    modal_cloud_module: Any,
    tmp_path: Path,
) -> None:
    """Identify native cache files visible only in process memory maps."""
    volume_root = tmp_path / "compile-cache"
    mapped_library = volume_root / "triton" / "cuda_utils.so"
    mapped_library.parent.mkdir(parents=True)
    mapped_library.touch()
    proc_root = tmp_path / "proc"
    process_root = proc_root / "191"
    process_root.mkdir(parents=True)
    (process_root / "maps").write_text(
        f"7f000000-7f001000 r-xp 00000000 00:1e 144 {mapped_library}\n"
        f"7f001000-7f002000 r--p 00001000 00:1e 144 {mapped_library}\n",
        encoding="utf-8",
    )

    mapped_files = modal_cloud_module._mapped_process_files_under(
        volume_root,
        proc_root=proc_root,
    )

    assert mapped_files == ((191, str(mapped_library)),)


def test_remote_modal_auto_deploys_missing_app_by_default(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Remote mode should auto-deploy the stable Modal app on first lookup failure."""

    class FakeLookupError(Exception):
        """Stand-in for Modal deployed lookup failures."""

    deploy_calls: list[tuple[str | None, str | None]] = []
    stage_calls: list[list[str]] = []
    status_events: list[dict[str, Any]] = []

    class FakeExecuteMethod:
        """Minimal Modal method handle that records remote calls."""

        def remote(self, payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
            """Return deterministic remote bytes."""
            return b"remote-response"

    class FakeRemoteEngine:
        """Minimal deployed remote engine instance."""

        execute_payload = FakeExecuteMethod()
        runtime_version = types.SimpleNamespace(
            remote=lambda: _current_remote_runtime_payload(remote_modal_app_module)
        )

    class FakeModelStager:
        """Minimal CPU model stager used by the first deployed LLM request."""

        stage_profiles = types.SimpleNamespace(
            remote=lambda profile_ids: stage_calls.append(list(profile_ids))
            or [{"profile_id": profile_id} for profile_id in profile_ids]
        )

    class FakeApp:
        """Minimal deployable cloud app double."""

        def deploy(self, *, name: str | None = None, environment_name: str | None = None, **_: Any) -> "FakeApp":
            """Record the deploy request and mark the deployment available."""
            deploy_calls.append((name, environment_name))
            FakeModal.deployed = True
            return self

    class FakeModal:
        """Minimal modal SDK double with deployed lookup failure types."""

        deployed = False
        exception = types.SimpleNamespace(
            NotFoundError=FakeLookupError,
            ExecutionError=FakeLookupError,
            InvalidError=FakeLookupError,
        )

        class Cls:
            """Namespace for deployed class lookups."""

            @staticmethod
            def from_name(app_name: str, class_name: str) -> Any:
                """Return a deployed class after the first auto-deploy."""
                if not FakeModal.deployed:
                    raise FakeLookupError("not deployed")
                if class_name == "ModelStager":
                    return lambda: FakeModelStager()
                return lambda **kwargs: FakeRemoteEngine()

        @staticmethod
        def enable_output() -> Any:
            """Provide a no-op output context manager."""
            return nullcontext()

    monkeypatch.setattr(remote_modal_app_module, "modal", FakeModal)
    monkeypatch.setattr(
        remote_modal_app_module,
        "_emit_local_modal_status",
        lambda **event: status_events.append(event),
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_load_modal_cloud_module",
        lambda: types.SimpleNamespace(app=FakeApp()),
    )
    monkeypatch.setenv("COMFY_MODAL_AUTO_DEPLOY", "true")
    remote_modal_app_module.get_settings.cache_clear()
    remote_modal_app_module._MODAL_AUTO_DEPLOY_STATES.clear()
    remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()
    with remote_modal_app_module._STAGED_LLM_PROFILES_LOCK:
        remote_modal_app_module._STAGED_LLM_PROFILES.clear()
    try:
        response = remote_modal_app_module._invoke_modal_payload_blocking(
            {
                "component_id": "component-1",
                "component_node_ids": ["node-1"],
                "prompt_id": "prompt-1",
                "extra_data": {"client_id": "client-1"},
                "subgraph_prompt": {
                    "node-1": {
                        "class_type": "ModalLLM",
                        "inputs": {"model_profile": "smolvlm2-2.2b-instruct"},
                    }
                },
            },
            b"{}",
        )
    finally:
        remote_modal_app_module.get_settings.cache_clear()
        remote_modal_app_module._MODAL_AUTO_DEPLOY_STATES.clear()
        remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()
        with remote_modal_app_module._STAGED_LLM_PROFILES_LOCK:
            remote_modal_app_module._STAGED_LLM_PROFILES.clear()

    assert response == b"remote-response"
    assert deploy_calls == [(DEFAULT_TEST_DEPLOYMENT_APP_NAME, None)]
    assert stage_calls == [["smolvlm2-2.2b-instruct"]]
    assert [
        (event["phase"], event["status_message"])
        for event in status_events
        ] == [
            ("setup", "Rebuilding Modal app"),
            (
                "llm_staging",
                "Preparing LLM model snapshots on CPU; no GPU is allocated yet",
            ),
            (
                "llm_staged",
                "LLM staging complete (0.0 GiB downloaded); starting GPU worker",
            ),
            ("starting", "Starting remote component on Modal"),
        ]


def test_remote_modal_does_not_classify_remote_execution_error_as_missing_app(
    remote_modal_app_module: Any,
) -> None:
    """Remote execution tracebacks should surface directly instead of triggering redeploy."""
    remote_error = RuntimeError(
        "Could not deserialize remote exception due to local error:\n"
        "Here is the remote traceback:\n"
        "comfyui_modal_sync_cloud.RemoteSubgraphExecutionError: "
        "Object of type CLIP is not JSON serializable\n"
        "Lookup failed for Cls 'RemoteEngine' from the 'comfy-modal-sync' app: "
        "App 'comfy-modal-sync' not found in environment 'main'."
    )
    missing_lookup_error = RuntimeError(
        "Lookup failed for Cls 'RemoteEngine' from the 'comfy-modal-sync' app: "
        "App 'comfy-modal-sync' not found in environment 'main'."
    )

    assert not remote_modal_app_module._is_missing_modal_deployment_error(remote_error)
    assert remote_modal_app_module._is_missing_modal_deployment_error(missing_lookup_error)


def test_remote_modal_does_not_redeploy_after_remote_execution_error(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """The deployed invocation path should not auto-deploy after a remote runtime failure."""

    class FakeLookupError(Exception):
        """Stand-in for Modal errors that share the lookup error type."""

    class FakeExecuteMethod:
        """Minimal Modal method handle that raises a remote execution failure."""

        def remote(self, payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
            """Raise the wrapped remote traceback from the deployed call."""
            del payload, kwargs_payload
            raise FakeLookupError(
                "Could not deserialize remote exception due to local error:\n"
                "Here is the remote traceback:\n"
                "comfyui_modal_sync_cloud.RemoteSubgraphExecutionError: "
                "Object of type CLIP is not JSON serializable\n"
                "App 'comfy-modal-sync' not found in environment 'main'."
            )

    class FakeRemoteEngine:
        """Minimal deployed remote engine instance."""

        execute_payload = FakeExecuteMethod()
        runtime_version = types.SimpleNamespace(
            remote=lambda: _current_remote_runtime_payload(remote_modal_app_module)
        )

    class FakeModal:
        """Minimal Modal SDK double with shared lookup/runtime error types."""

        exception = types.SimpleNamespace(
            NotFoundError=FakeLookupError,
            ExecutionError=FakeLookupError,
            InvalidError=FakeLookupError,
        )

        class Cls:
            """Namespace for deployed class lookups."""

            @staticmethod
            def from_name(app_name: str, class_name: str) -> Any:
                """Return a deployed class successfully."""
                del app_name, class_name
                return lambda **kwargs: FakeRemoteEngine()

    def fail_load_cloud_module() -> Any:
        """Fail if runtime errors attempt auto-deploy."""
        raise AssertionError("remote execution errors must not trigger auto-deploy")

    monkeypatch.setattr(remote_modal_app_module, "modal", FakeModal)
    monkeypatch.setattr(remote_modal_app_module, "_load_modal_cloud_module", fail_load_cloud_module)
    monkeypatch.setenv("COMFY_MODAL_AUTO_DEPLOY", "true")
    remote_modal_app_module.get_settings.cache_clear()
    try:
        with pytest.raises(FakeLookupError, match="Object of type CLIP"):
            remote_modal_app_module._invoke_modal_payload_blocking(
                {"component_id": "component-1"},
                b"{}",
            )
    finally:
        remote_modal_app_module.get_settings.cache_clear()


def test_remote_modal_redeploys_when_cached_app_was_deleted(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """A stale in-process auto-deploy cache should not block redeploy after app deletion."""

    class FakeLookupError(Exception):
        """Stand-in for Modal deployed lookup failures."""

    deploy_calls: list[tuple[str | None, str | None]] = []

    class FakeExecuteMethod:
        """Minimal Modal method handle that records remote calls."""

        def remote(self, payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
            """Return deterministic remote bytes."""
            del payload, kwargs_payload
            return b"remote-response"

    class FakeRemoteEngine:
        """Minimal deployed remote engine instance."""

        execute_payload = FakeExecuteMethod()
        runtime_version = types.SimpleNamespace(
            remote=lambda: _current_remote_runtime_payload(remote_modal_app_module)
        )

    class FakeApp:
        """Minimal deployable cloud app double."""

        def deploy(self, *, name: str | None = None, environment_name: str | None = None, **_: Any) -> "FakeApp":
            """Record the redeploy request and mark the deployment available again."""
            deploy_calls.append((name, environment_name))
            FakeModal.deployed = True
            return self

    class FakeModal:
        """Minimal modal SDK double with deployed lookup failure types."""

        deployed = False
        exception = types.SimpleNamespace(
            NotFoundError=FakeLookupError,
            ExecutionError=FakeLookupError,
            InvalidError=FakeLookupError,
        )

        class Cls:
            """Namespace for deployed class lookups."""

            @staticmethod
            def from_name(app_name: str, class_name: str) -> Any:
                """Raise the same missing-app error until auto-deploy recreates the app."""
                del app_name, class_name
                if not FakeModal.deployed:
                    raise FakeLookupError(
                        "Lookup failed for Cls 'RemoteEngine' from the 'comfy-modal-sync' app: "
                        "App 'comfy-modal-sync' not found in environment 'main'."
                    )
                return lambda **kwargs: FakeRemoteEngine()

        @staticmethod
        def enable_output() -> Any:
            """Provide a no-op output context manager."""
            return nullcontext()

    monkeypatch.setattr(remote_modal_app_module, "modal", FakeModal)
    monkeypatch.setattr(
        remote_modal_app_module,
        "_load_modal_cloud_module",
        lambda: types.SimpleNamespace(app=FakeApp()),
    )
    monkeypatch.setenv("COMFY_MODAL_AUTO_DEPLOY", "true")
    monkeypatch.setenv("MODAL_ENVIRONMENT", "main")
    remote_modal_app_module.get_settings.cache_clear()
    remote_modal_app_module._MODAL_AUTO_DEPLOY_STATES.clear()
    remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()
    remote_modal_app_module._MODAL_AUTO_DEPLOY_STATES[
        (DEFAULT_TEST_DEPLOYMENT_APP_NAME, "main")
    ] = (
        remote_modal_app_module._ModalAutoDeployState(ready=True)
    )
    try:
        response = remote_modal_app_module._invoke_modal_payload_blocking(
            {"component_id": "component-1"},
            b"{}",
        )
    finally:
        remote_modal_app_module.get_settings.cache_clear()
        remote_modal_app_module._MODAL_AUTO_DEPLOY_STATES.clear()
        remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()

    assert response == b"remote-response"
    assert deploy_calls == [(DEFAULT_TEST_DEPLOYMENT_APP_NAME, "main")]


def test_remote_modal_redeploys_when_deployed_handle_disappears_during_payload_invoke(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """A stale deployed engine handle should auto-redeploy when invocation hydrate fails."""

    class FakeLookupError(Exception):
        """Stand-in for Modal deployed lookup failures."""

    deploy_calls: list[tuple[str | None, str | None]] = []

    class FakeRemoteEngine:
        """Minimal deployed remote engine instance with stale-handle behavior."""

        def __init__(self, *, stale: bool) -> None:
            """Record whether this engine should fail once as a stale handle."""
            self._stale = stale
            self.execute_payload = types.SimpleNamespace(remote=self._remote)
            self.runtime_version = types.SimpleNamespace(
                remote=lambda: _current_remote_runtime_payload(remote_modal_app_module)
            )

        def _remote(self, payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
            """Raise the missing-app error for stale handles and otherwise succeed."""
            del payload, kwargs_payload
            if self._stale:
                raise FakeLookupError(
                    "Lookup failed for Cls 'RemoteEngine' from the 'comfy-modal-sync' app: "
                    "App 'comfy-modal-sync' not found in environment 'main'."
                )
            return b"remote-response"

    class FakeApp:
        """Minimal deployable cloud app double."""

        def deploy(
            self,
            *,
            name: str | None = None,
            environment_name: str | None = None,
            **_: Any,
        ) -> "FakeApp":
            """Record the redeploy request and mark future handles as fresh."""
            deploy_calls.append((name, environment_name))
            FakeModal.redeployed = True
            return self

    class FakeModal:
        """Minimal modal SDK double with stale-handle invocation failures."""

        redeployed = False
        exception = types.SimpleNamespace(
            NotFoundError=FakeLookupError,
            ExecutionError=FakeLookupError,
            InvalidError=FakeLookupError,
        )

        class Cls:
            """Namespace for deployed class lookups."""

            @staticmethod
            def from_name(app_name: str, class_name: str) -> Any:
                """Return stale handles before redeploy and fresh handles after."""
                del app_name, class_name
                return lambda **kwargs: FakeRemoteEngine(stale=not FakeModal.redeployed)

        @staticmethod
        def enable_output() -> Any:
            """Provide a no-op output context manager."""
            return nullcontext()

    monkeypatch.setattr(remote_modal_app_module, "modal", FakeModal)
    monkeypatch.setattr(
        remote_modal_app_module,
        "_load_modal_cloud_module",
        lambda: types.SimpleNamespace(app=FakeApp()),
    )
    monkeypatch.setenv("COMFY_MODAL_AUTO_DEPLOY", "true")
    monkeypatch.setenv("MODAL_ENVIRONMENT", "main")
    remote_modal_app_module.get_settings.cache_clear()
    remote_modal_app_module._MODAL_AUTO_DEPLOY_STATES.clear()
    remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()
    try:
        response = remote_modal_app_module._invoke_modal_payload_blocking(
            {"component_id": "component-1"},
            b"{}",
        )
    finally:
        remote_modal_app_module.get_settings.cache_clear()
        remote_modal_app_module._MODAL_AUTO_DEPLOY_STATES.clear()
        remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()

    assert response == b"remote-response"
    assert deploy_calls == [(DEFAULT_TEST_DEPLOYMENT_APP_NAME, "main")]


def test_remote_modal_redeploys_when_deployed_handle_disappears_during_warmup(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """A stale deployed warmup handle should auto-redeploy instead of surfacing a warmup error."""

    class FakeLookupError(Exception):
        """Stand-in for Modal deployed lookup failures."""

    deploy_calls: list[tuple[str | None, str | None]] = []

    class FakeRemoteEngine:
        """Minimal deployed remote engine instance with stale-handle warmup behavior."""

        def __init__(self, *, stale: bool) -> None:
            """Expose a warmup method that fails for stale handles."""
            self._stale = stale
            self.warmup_for_request = types.SimpleNamespace(remote=self._remote)
            self.runtime_version = types.SimpleNamespace(
                remote=lambda: _current_remote_runtime_payload(remote_modal_app_module)
            )

        def _remote(self, payload: dict[str, Any]) -> dict[str, Any]:
            """Raise the missing-app error for stale handles and otherwise succeed."""
            if self._stale:
                raise FakeLookupError(
                    "Lookup failed for Cls 'RemoteEngine' from the 'comfy-modal-sync' app: "
                    "App 'comfy-modal-sync' not found in environment 'main'."
                )
            return {"component_id": str(payload.get("component_id"))}

    class FakeApp:
        """Minimal deployable cloud app double."""

        def deploy(
            self,
            *,
            name: str | None = None,
            environment_name: str | None = None,
            **_: Any,
        ) -> "FakeApp":
            """Record the redeploy request and mark future handles as fresh."""
            deploy_calls.append((name, environment_name))
            FakeModal.redeployed = True
            return self

    class FakeModal:
        """Minimal modal SDK double with stale-handle warmup failures."""

        redeployed = False
        exception = types.SimpleNamespace(
            NotFoundError=FakeLookupError,
            ExecutionError=FakeLookupError,
            InvalidError=FakeLookupError,
        )

        class Cls:
            """Namespace for deployed class lookups."""

            @staticmethod
            def from_name(app_name: str, class_name: str) -> Any:
                """Return stale handles before redeploy and fresh handles after."""
                del app_name, class_name
                return lambda **kwargs: FakeRemoteEngine(stale=not FakeModal.redeployed)

        @staticmethod
        def enable_output() -> Any:
            """Provide a no-op output context manager."""
            return nullcontext()

    monkeypatch.setattr(remote_modal_app_module, "modal", FakeModal)
    monkeypatch.setattr(
        remote_modal_app_module,
        "_load_modal_cloud_module",
        lambda: types.SimpleNamespace(app=FakeApp()),
    )
    monkeypatch.setenv("COMFY_MODAL_AUTO_DEPLOY", "true")
    monkeypatch.setenv("MODAL_ENVIRONMENT", "main")
    remote_modal_app_module.get_settings.cache_clear()
    remote_modal_app_module._MODAL_AUTO_DEPLOY_STATES.clear()
    remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()
    try:
        response = remote_modal_app_module._invoke_modal_warmup_blocking(
            {"component_id": "component-1::warmup:0"},
        )
    finally:
        remote_modal_app_module.get_settings.cache_clear()
        remote_modal_app_module._MODAL_AUTO_DEPLOY_STATES.clear()
        remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()

    assert response == {"component_id": "component-1::warmup:0"}
    assert deploy_calls == [(DEFAULT_TEST_DEPLOYMENT_APP_NAME, "main")]


def test_remote_modal_replaces_out_of_date_deployed_app(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """A deployed app with an old protocol should be stopped and auto-deployed again."""

    class FakeLookupError(Exception):
        """Stand-in for Modal deployed lookup failures."""

    deploy_calls: list[tuple[str | None, str | None]] = []
    stop_calls: list[str] = []

    class FakeExecuteMethod:
        """Minimal Modal method handle that records remote calls."""

        def remote(self, payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
            """Return deterministic remote bytes."""
            del payload, kwargs_payload
            return b"remote-response"

    class FakeRemoteEngine:
        """Minimal deployed remote engine with protocol metadata."""

        def __init__(self, *, current: bool) -> None:
            """Record whether this engine should report the current protocol."""
            version_payload = _current_remote_runtime_payload(remote_modal_app_module)
            if not current:
                version_payload["protocol_version"] = (
                    remote_modal_app_module._REMOTE_APP_PROTOCOL_VERSION - 1
                )
            self.execute_payload = FakeExecuteMethod()
            self.runtime_version = types.SimpleNamespace(
                remote=lambda: version_payload
            )

    class FakeApp:
        """Minimal deployable cloud app double."""

        def deploy(
            self,
            *,
            name: str | None = None,
            environment_name: str | None = None,
            **_: Any,
        ) -> "FakeApp":
            """Record the deploy request and make subsequent lookups current."""
            deploy_calls.append((name, environment_name))
            FakeModal.deployed = True
            FakeModal.current = True
            return self

    class FakeModal:
        """Minimal Modal SDK double with an initially stale deployed app."""

        deployed = True
        current = False
        exception = types.SimpleNamespace(
            NotFoundError=FakeLookupError,
            ExecutionError=FakeLookupError,
            InvalidError=FakeLookupError,
        )

        class Cls:
            """Namespace for deployed class lookups."""

            @staticmethod
            def from_name(app_name: str, class_name: str) -> Any:
                """Return the stale app first and the current replacement after deploy."""
                del app_name, class_name
                if not FakeModal.deployed:
                    raise FakeLookupError("Lookup failed for Cls 'RemoteEngine': not deployed")
                return lambda **kwargs: FakeRemoteEngine(current=FakeModal.current)

        @staticmethod
        def enable_output() -> Any:
            """Provide a no-op output context manager."""
            return nullcontext()

    def fake_stop_app(app_name: str) -> None:
        """Record that stale app replacement stopped the old deployment."""
        stop_calls.append(app_name)
        FakeModal.deployed = False

    monkeypatch.setattr(remote_modal_app_module, "modal", FakeModal)
    monkeypatch.setattr(remote_modal_app_module, "_stop_modal_app_for_replacement", fake_stop_app)
    monkeypatch.setattr(
        remote_modal_app_module,
        "_load_modal_cloud_module",
        lambda: types.SimpleNamespace(app=FakeApp()),
    )
    monkeypatch.setenv("COMFY_MODAL_AUTO_DEPLOY", "true")
    monkeypatch.setenv("MODAL_ENVIRONMENT", "main")
    remote_modal_app_module.get_settings.cache_clear()
    remote_modal_app_module._MODAL_AUTO_DEPLOY_STATES.clear()
    remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()
    try:
        response = remote_modal_app_module._invoke_modal_payload_blocking(
            {"component_id": "component-1"},
            b"{}",
        )
    finally:
        remote_modal_app_module.get_settings.cache_clear()
        remote_modal_app_module._MODAL_AUTO_DEPLOY_STATES.clear()
        remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()

    assert response == b"remote-response"
    assert stop_calls == [DEFAULT_TEST_DEPLOYMENT_APP_NAME]
    assert deploy_calls == [(DEFAULT_TEST_DEPLOYMENT_APP_NAME, "main")]


def test_remote_modal_stops_app_through_experimental_sdk(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Modal 1.4 should use its non-interactive experimental stop API."""
    stop_calls: list[tuple[str, str | None]] = []

    def fake_stop_app(app_name: str, *, environment_name: str | None = None) -> None:
        """Record one SDK stop request."""
        stop_calls.append((app_name, environment_name))

    fake_experimental_namespace = types.SimpleNamespace(stop_app=fake_stop_app)
    original_import_module = remote_modal_app_module.importlib.import_module

    def fake_import_module(module_name: str) -> Any:
        """Return the Modal experimental namespace and delegate other imports."""
        if module_name == "modal.experimental":
            return fake_experimental_namespace
        return original_import_module(module_name)

    monkeypatch.setattr(
        remote_modal_app_module,
        "modal",
        types.SimpleNamespace(exception=types.SimpleNamespace()),
    )
    monkeypatch.setattr(
        remote_modal_app_module.importlib,
        "import_module",
        fake_import_module,
    )
    monkeypatch.setenv("MODAL_ENVIRONMENT", "main")

    stopped = remote_modal_app_module._stop_modal_app_via_sdk("comfy-modal-sync-instance")

    assert stopped is True
    assert stop_calls == [("comfy-modal-sync-instance", "main")]


def test_remote_modal_cli_stop_is_noninteractive_and_environment_scoped(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """The CLI fallback should never block waiting for confirmation."""
    observed_calls: list[tuple[list[str], dict[str, Any]]] = []

    def fake_run(command: list[str], **kwargs: Any) -> Any:
        """Record the CLI command and return success."""
        observed_calls.append((command, kwargs))
        return types.SimpleNamespace(returncode=0, stderr="")

    monkeypatch.setattr(remote_modal_app_module.shutil, "which", lambda _: "/venv/bin/modal")
    monkeypatch.setattr(remote_modal_app_module.subprocess, "run", fake_run)
    monkeypatch.setenv("MODAL_ENVIRONMENT", "main")

    stopped = remote_modal_app_module._stop_modal_app_via_cli("comfy-modal-sync-instance")

    assert stopped is True
    assert observed_calls == [
        (
            [
                "/venv/bin/modal",
                "app",
                "stop",
                "comfy-modal-sync-instance",
                "--yes",
                "--env",
                "main",
            ],
            {
                "check": False,
                "capture_output": True,
                "text": True,
                "timeout": remote_modal_app_module._MODAL_APP_STOP_TIMEOUT_SECONDS,
            },
        )
    ]


def test_remote_modal_cli_stop_timeout_returns_controlled_failure(
    remote_modal_app_module: Any,
    monkeypatch: Any,
    caplog: Any,
) -> None:
    """A genuine CLI timeout should not leak TimeoutExpired into prompt execution."""
    command = ["/venv/bin/modal", "app", "stop", "comfy-modal-sync-instance", "--yes"]

    def fake_run(*_: Any, **__: Any) -> Any:
        """Simulate an unresponsive Modal CLI."""
        raise subprocess.TimeoutExpired(
            command,
            remote_modal_app_module._MODAL_APP_STOP_TIMEOUT_SECONDS,
        )

    monkeypatch.setattr(remote_modal_app_module.shutil, "which", lambda _: "/venv/bin/modal")
    monkeypatch.setattr(remote_modal_app_module.subprocess, "run", fake_run)
    monkeypatch.delenv("MODAL_ENVIRONMENT", raising=False)

    with caplog.at_level(logging.WARNING):
        stopped = remote_modal_app_module._stop_modal_app_via_cli(
            "comfy-modal-sync-instance"
        )

    assert stopped is False
    assert "Modal CLI app stop timed out" in caplog.text


def test_remote_modal_rejects_fingerprint_mismatch_when_auto_deploy_is_disabled(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """A same-protocol worker with different source or settings must not execute."""
    stale_version_payload = _current_remote_runtime_payload(remote_modal_app_module)
    stale_version_payload["runtime_fingerprint"] = "stale-runtime"
    remote_engine = types.SimpleNamespace(
        runtime_version=types.SimpleNamespace(remote=lambda: stale_version_payload)
    )
    monkeypatch.setenv("COMFY_MODAL_AUTO_DEPLOY", "false")
    remote_modal_app_module.get_settings.cache_clear()
    remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()
    try:
        with pytest.raises(
            remote_modal_app_module.ModalRemoteInvocationError,
            match="runtime fingerprint is out of date",
        ):
            remote_modal_app_module._ensure_remote_engine_protocol_current(
                remote_engine,
                {"component_id": "component-1"},
            )
    finally:
        remote_modal_app_module.get_settings.cache_clear()
        remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()


def test_remote_modal_rebinds_affinity_after_compatible_protocol_probe(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """A legacy-safe version probe should be replaced before real execution."""
    version_payload = _current_remote_runtime_payload(remote_modal_app_module)
    probe_engine = types.SimpleNamespace(
        runtime_version=types.SimpleNamespace(remote=lambda: version_payload)
    )
    affinity_engine = object()
    observed_payloads: list[dict[str, Any]] = []

    def fake_lookup(payload: dict[str, Any]) -> object:
        """Return the affinity-aware handle created after protocol validation."""
        observed_payloads.append(payload)
        return affinity_engine

    monkeypatch.setattr(
        remote_modal_app_module,
        "_lookup_deployed_remote_engine",
        fake_lookup,
    )
    remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()
    payload = {"component_id": "component-1"}
    try:
        result = remote_modal_app_module._ensure_remote_engine_protocol_current(
            probe_engine,
            payload,
        )
    finally:
        remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()

    assert result is affinity_engine
    assert observed_payloads == [payload]


def test_remote_modal_rebinds_affinity_after_cached_protocol_validation(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """A cached protocol result must not turn its parameterless probe into execution."""
    probe_engine = object()
    affinity_engine = object()
    observed_payloads: list[dict[str, Any]] = []

    def fake_lookup(payload: dict[str, Any]) -> object:
        """Return the affinity-aware handle created for actual execution."""
        observed_payloads.append(payload)
        return affinity_engine

    monkeypatch.setattr(
        remote_modal_app_module,
        "_lookup_deployed_remote_engine",
        fake_lookup,
    )
    payload = {
        "component_id": "component-llm",
        "remote_worker_affinity_group": "llm",
    }
    runtime_cache_key = remote_modal_app_module._modal_runtime_cache_key(payload)
    remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()
    remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.add(runtime_cache_key)
    try:
        result = remote_modal_app_module._ensure_remote_engine_protocol_current(
            probe_engine,
            payload,
        )
    finally:
        remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()

    assert result is affinity_engine
    assert observed_payloads == [payload]


def test_remote_modal_cached_protocol_skips_parameterless_probe_construction(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """A cached runtime must construct only the affinity-bound engine handle."""
    affinity_engine = object()
    observed_lookups: list[tuple[dict[str, Any], bool]] = []

    def fake_lookup(
        payload: dict[str, Any],
        *,
        affinity_key_override: str | None = None,
        protocol_probe: bool = False,
    ) -> object:
        """Record whether lookup attempted to allocate a parameterless probe."""
        del affinity_key_override
        observed_lookups.append((payload, protocol_probe))
        return affinity_engine

    monkeypatch.setattr(
        remote_modal_app_module,
        "_lookup_deployed_remote_engine",
        fake_lookup,
    )
    payload = {
        "component_id": "component-comfy::warmup:0",
        "remote_worker_affinity_group": "comfy",
    }
    runtime_cache_key = remote_modal_app_module._modal_runtime_cache_key(payload)
    remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()
    remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.add(runtime_cache_key)
    try:
        result = remote_modal_app_module._lookup_protocol_current_remote_engine(payload)
    finally:
        remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()

    assert result is affinity_engine
    assert observed_lookups == [(payload, False)]


def test_remote_modal_uncached_protocol_uses_cpu_stager_before_gpu_lookup(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """First validation should use CPU metadata and construct one bound GPU handle."""
    affinity_engine = object()
    observed_lookups: list[bool] = []
    version_payload = _current_remote_runtime_payload(remote_modal_app_module)

    def fake_lookup(
        payload: dict[str, Any],
        *,
        affinity_key_override: str | None = None,
        protocol_probe: bool = False,
    ) -> object:
        """Record whether lookup attempted to allocate a GPU protocol probe."""
        del payload, affinity_key_override
        observed_lookups.append(protocol_probe)
        return affinity_engine

    monkeypatch.setattr(
        remote_modal_app_module,
        "_remote_runtime_version_from_cpu_stager",
        lambda payload: version_payload,
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_lookup_deployed_remote_engine",
        fake_lookup,
    )
    payload = {
        "component_id": "component-llm",
        "remote_worker_affinity_group": "llm",
    }
    runtime_cache_key = remote_modal_app_module._modal_runtime_cache_key(payload)
    remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()
    try:
        result = remote_modal_app_module._lookup_protocol_current_remote_engine(payload)
        runtime_was_cached = (
            runtime_cache_key
            in remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK
        )
    finally:
        remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()

    assert result is affinity_engine
    assert runtime_was_cached is True
    assert observed_lookups == [False]


def test_workflow_gpu_changes_expected_remote_runtime_fingerprint(
    remote_modal_app_module: Any,
) -> None:
    """Each workflow GPU target should select a distinct deploy-time runtime identity."""
    a100_fingerprint = remote_modal_app_module._expected_remote_runtime_fingerprint(
        {"modal_gpu": "A100"}
    )
    b300_fingerprint = remote_modal_app_module._expected_remote_runtime_fingerprint(
        {"modal_gpu": "B300"}
    )

    assert a100_fingerprint != b300_fingerprint
    assert remote_modal_app_module._settings_for_payload(
        {"modal_gpu": "L40S"}
    ).modal_gpu == "L40S"
    configured_settings = remote_modal_app_module._settings_for_payload(
        {"modal_gpu": "H200", "modal_max_containers": 3}
    )
    assert configured_settings.modal_gpu == "H200"
    assert configured_settings.max_containers == 3
    assert remote_modal_app_module._modal_deploy_cache_key(
        {"modal_gpu": "A100"}
    )[0] == "comfy-modal-sync"
    assert remote_modal_app_module._modal_deploy_cache_key(
        {"modal_gpu": "B300"}
    )[0] == "comfy-modal-sync-gpu-b300"


def test_remote_modal_auto_deploy_is_shared_across_concurrent_first_run_callers(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Concurrent first-run callers should share one deploy and wait for lookup readiness."""

    class FakeLookupError(Exception):
        """Stand-in for Modal deployed lookup failures."""

    deploy_calls: list[tuple[str | None, str | None]] = []
    deployed_ready_event = threading.Event()
    lookup_lock = threading.Lock()
    ready_after_deploy_misses = 0

    class FakeExecuteMethod:
        """Minimal Modal method handle that returns deterministic bytes."""

        def remote(self, payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
            """Return deterministic remote bytes."""
            del payload, kwargs_payload
            return b"remote-response"

    class FakeWarmupMethod:
        """Minimal Modal warmup handle."""

        def remote(self, payload: dict[str, Any]) -> dict[str, Any]:
            """Return deterministic warmup metadata."""
            return {"component_id": str(payload.get("component_id"))}

    class FakeRemoteEngine:
        """Minimal deployed remote engine instance."""

        execute_payload = FakeExecuteMethod()
        warmup_for_request = FakeWarmupMethod()
        runtime_version = types.SimpleNamespace(
            remote=lambda: _current_remote_runtime_payload(remote_modal_app_module)
        )

    class FakeApp:
        """Minimal deployable cloud app double."""

        def deploy(self, *, name: str | None = None, environment_name: str | None = None, **_: Any) -> "FakeApp":
            """Record one deploy request and make the app discoverable shortly after."""
            nonlocal ready_after_deploy_misses
            deploy_calls.append((name, environment_name))
            with lookup_lock:
                FakeModal.deployed = True
                ready_after_deploy_misses = 2
            deployed_ready_event.set()
            return self

    class FakeModal:
        """Minimal modal SDK double with eventual-consistency lookup behavior."""

        deployed = False
        exception = types.SimpleNamespace(
            NotFoundError=FakeLookupError,
            ExecutionError=FakeLookupError,
            InvalidError=FakeLookupError,
        )

        class Cls:
            """Namespace for deployed class lookups."""

            @staticmethod
            def from_name(app_name: str, class_name: str) -> Any:
                """Return a deployed class after one deploy and a short readiness delay."""
                nonlocal ready_after_deploy_misses
                del app_name, class_name
                with lookup_lock:
                    if not FakeModal.deployed:
                        raise FakeLookupError("Lookup failed for Cls 'RemoteEngine': not deployed")
                    if ready_after_deploy_misses > 0:
                        ready_after_deploy_misses -= 1
                        raise FakeLookupError(
                            "Lookup failed for Cls 'RemoteEngine' from the 'comfy-modal-sync' app: "
                            "App 'comfy-modal-sync' not found in environment 'main'."
                        )
                return lambda **kwargs: FakeRemoteEngine()

        @staticmethod
        def enable_output() -> Any:
            """Provide a no-op output context manager."""
            return nullcontext()

    monkeypatch.setattr(remote_modal_app_module, "modal", FakeModal)
    monkeypatch.setattr(
        remote_modal_app_module,
        "_load_modal_cloud_module",
        lambda: types.SimpleNamespace(app=FakeApp()),
    )
    monkeypatch.setattr(remote_modal_app_module.time, "sleep", lambda _: None)
    monkeypatch.setenv("COMFY_MODAL_AUTO_DEPLOY", "true")
    monkeypatch.setenv("MODAL_ENVIRONMENT", "main")
    remote_modal_app_module.get_settings.cache_clear()
    remote_modal_app_module._MODAL_AUTO_DEPLOY_STATES.clear()
    remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()

    try:
        payload_response: list[bytes] = []
        warmup_response: list[dict[str, Any]] = []
        thread_errors: list[BaseException] = []

        def run_payload() -> None:
            """Invoke the payload path from one thread."""
            try:
                payload_response.append(
                    remote_modal_app_module._invoke_modal_payload_blocking(
                        {"component_id": "component-1"},
                        b"{}",
                    )
                )
            except Exception as exc:  # pragma: no cover - assertion surface
                thread_errors.append(exc)

        def run_warmup() -> None:
            """Invoke the warmup path from one thread."""
            try:
                warmup_response.append(
                    remote_modal_app_module._invoke_modal_warmup_blocking(
                        {"component_id": "component-1", "prompt_id": "prompt-1"},
                    )
                )
            except Exception as exc:  # pragma: no cover - assertion surface
                thread_errors.append(exc)

        payload_thread = threading.Thread(target=run_payload)
        warmup_thread = threading.Thread(target=run_warmup)
        payload_thread.start()
        assert deployed_ready_event.wait(1.0)
        warmup_thread.start()
        payload_thread.join(timeout=1.0)
        warmup_thread.join(timeout=1.0)
    finally:
        remote_modal_app_module.get_settings.cache_clear()
        remote_modal_app_module._MODAL_AUTO_DEPLOY_STATES.clear()
        remote_modal_app_module._MODAL_REMOTE_APP_VERSION_OK.clear()

    assert thread_errors == []
    assert payload_response == [b"remote-response"]
    assert warmup_response == [{"component_id": "component-1"}]
    assert deploy_calls == [(DEFAULT_TEST_DEPLOYMENT_APP_NAME, "main")]


def test_lookup_deployed_remote_engine_uses_affinity_as_modal_class_parameter(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Deployed Modal lookups should isolate reusable worker slots by class identity."""
    observed_kwargs: list[dict[str, Any]] = []

    class FakeModal:
        """Minimal modal SDK double that captures class construction kwargs."""

        class Cls:
            """Namespace for deployed class lookups."""

            @staticmethod
            def from_name(app_name: str, class_name: str) -> Any:
                """Return a factory that records the provided class parameters."""
                assert app_name == DEFAULT_TEST_DEPLOYMENT_APP_NAME
                assert class_name == "RemoteEngine"

                def build_remote_engine(**kwargs: Any) -> dict[str, Any]:
                    """Record one synthesized Modal class constructor call."""
                    observed_kwargs.append(kwargs)
                    return kwargs

                return build_remote_engine

    monkeypatch.setattr(remote_modal_app_module, "modal", FakeModal)

    result = remote_modal_app_module._lookup_deployed_remote_engine(
        {
            "component_id": "component-1",
            "remote_session": remote_modal_app_module.RemoteSessionHandle(
                session_id="session-123",
                prompt_id="prompt-1",
                owner_component_id="component-1",
            ).to_payload(),
        }
    )
    probe_result = remote_modal_app_module._lookup_deployed_remote_engine(
        {"component_id": "component-1"},
        protocol_probe=True,
    )

    assert result == {
        "gpu_snapshot_enabled": False,
        "worker_affinity_key": "worker-pool:slot:0",
    }
    assert probe_result == {"gpu_snapshot_enabled": False}
    assert observed_kwargs == [
        {
            "gpu_snapshot_enabled": False,
            "worker_affinity_key": "worker-pool:slot:0",
        },
        {"gpu_snapshot_enabled": False},
    ]


def test_lookup_deployed_remote_engine_uses_workflow_gpu_app_identity(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """A B300 workflow must miss or reuse its own app without contacting the A100 class."""
    observed_lookups: list[tuple[str, str]] = []

    class FakeModal:
        """Minimal Modal SDK double that captures deployed class identity."""

        class Cls:
            """Namespace for deployed class lookups."""

            @staticmethod
            def from_name(app_name: str, class_name: str) -> Any:
                """Return a no-op class factory for the requested deployment."""
                observed_lookups.append((app_name, class_name))
                return lambda **kwargs: kwargs

    monkeypatch.setattr(remote_modal_app_module, "modal", FakeModal)

    result = remote_modal_app_module._lookup_deployed_remote_engine(
        {"component_id": "component-1", "modal_gpu": "B300"}
    )

    assert result == {
        "gpu_snapshot_enabled": False,
        "worker_affinity_key": "worker-pool:slot:0",
    }
    assert observed_lookups == [("comfy-modal-sync-gpu-b300", "RemoteEngine")]


def test_lookup_deployed_remote_engine_extends_local_gap_pool_scaledown(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Sandwiched local work should use a longer independently autoscaled Modal pool."""
    observed_options: list[dict[str, Any]] = []
    observed_kwargs: list[dict[str, Any]] = []

    class FakeRemoteCls:
        """Minimal deployed class supporting Modal runtime option overrides."""

        def with_options(self, **kwargs: Any) -> "FakeRemoteCls":
            """Record the independent autoscaler configuration."""
            observed_options.append(kwargs)
            return self

        def __call__(self, **kwargs: Any) -> dict[str, Any]:
            """Record the parametrized class instance identity."""
            observed_kwargs.append(kwargs)
            return kwargs

    class FakeModal:
        """Minimal Modal SDK double for deployed class lookup."""

        class Cls:
            """Namespace for deployed class lookups."""

            @staticmethod
            def from_name(app_name: str, class_name: str) -> FakeRemoteCls:
                """Return the shared fake deployed class."""
                assert app_name == DEFAULT_TEST_DEPLOYMENT_APP_NAME
                assert class_name == "RemoteEngine"
                return FakeRemoteCls()

    monkeypatch.setattr(remote_modal_app_module, "modal", FakeModal)

    result = remote_modal_app_module._lookup_deployed_remote_engine(
        {
            "component_id": "component-1",
            "prompt_id": "prompt-1",
            "remote_local_gap_pool": True,
        }
    )

    assert observed_options == [{"scaledown_window": 900}]
    assert observed_kwargs == [result]
    assert result["worker_affinity_key"] == "worker-pool:slot:0"


def test_local_gap_components_share_one_affinity_slot(
    remote_modal_app_module: Any,
) -> None:
    """Sequential remote phases around local work must address the same worker pool."""
    parallelism_metadata = {
        "modal": {
            "component_execution_stages": [
                ["component-a", "independent-component"],
                ["component-b"],
            ],
            "estimated_max_parallel_requests": 2,
        }
    }
    first_payload = {
        "component_id": "component-a",
        "remote_local_gap_pool": True,
        "extra_data": parallelism_metadata,
    }
    second_payload = {
        "component_id": "component-b",
        "remote_local_gap_pool": True,
        "extra_data": parallelism_metadata,
    }

    assert remote_modal_app_module._component_pool_slot_index(first_payload) == 0
    assert remote_modal_app_module._component_pool_slot_index(second_payload) == 0
    assert remote_modal_app_module._remote_worker_affinity_key(first_payload) == (
        remote_modal_app_module._remote_worker_affinity_key(second_payload)
    )


def test_lookup_deployed_remote_engine_passes_snapshot_profile_parameter_for_gpu_snapshots(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Deployed Modal lookups should parameterize workers by snapshot profile when enabled."""
    observed_kwargs: list[dict[str, Any]] = []
    snapshot_profiles: dict[str, Any] = {}

    class FakeModal:
        """Minimal modal SDK double that captures class construction kwargs."""

        class Dict:
            """Namespace for fake dict lookups."""

            @staticmethod
            def from_name(dict_name: str, create_if_missing: bool = False) -> Any:
                """Return the shared fake snapshot profile store."""
                assert dict_name == "comfy-modal-sync-snapshot-profiles"
                assert create_if_missing is True
                return snapshot_profiles

        class Cls:
            """Namespace for deployed class lookups."""

            @staticmethod
            def from_name(app_name: str, class_name: str) -> Any:
                """Return a factory that records the provided class parameters."""
                assert app_name == DEFAULT_TEST_DEPLOYMENT_APP_NAME
                assert class_name == "RemoteEngine"

                def build_remote_engine(**kwargs: Any) -> dict[str, Any]:
                    """Record one synthesized Modal class constructor call."""
                    observed_kwargs.append(kwargs)
                    return kwargs

                return build_remote_engine

    monkeypatch.setattr(remote_modal_app_module, "modal", FakeModal)
    monkeypatch.setenv("COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT", "true")
    remote_modal_app_module.get_settings.cache_clear()
    remote_modal_app_module._SNAPSHOT_PROFILE_RECORDS.clear()
    payload = {
        "component_id": "component-1",
        "subgraph_prompt": {
            "1": {
                "class_type": "UNETLoader",
                "inputs": {"unet_name": "model-a.safetensors", "weight_dtype": "default"},
            }
        },
    }
    try:
        result = remote_modal_app_module._lookup_deployed_remote_engine(payload)
    finally:
        remote_modal_app_module.get_settings.cache_clear()
        remote_modal_app_module._SNAPSHOT_PROFILE_RECORDS.clear()

    assert result["gpu_snapshot_enabled"] is True
    assert result["snapshot_profile_key"].startswith("loader-profile:")
    assert observed_kwargs == [result]
    assert result["snapshot_profile_key"] in snapshot_profiles
    assert payload["snapshot_profile_key"] == result["snapshot_profile_key"]


def test_lookup_deployed_remote_engine_stores_existing_snapshot_profile_record_when_possible(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Deployed lookups should backfill snapshot records even when the payload already carries the key."""
    observed_kwargs: list[dict[str, Any]] = []
    snapshot_profiles: dict[str, Any] = {}

    class FakeModal:
        """Minimal modal SDK double that captures class construction kwargs."""

        class Dict:
            """Namespace for fake dict lookups."""

            @staticmethod
            def from_name(dict_name: str, create_if_missing: bool = False) -> Any:
                """Return the shared fake snapshot profile store."""
                assert dict_name == "comfy-modal-sync-snapshot-profiles"
                assert create_if_missing is True
                return snapshot_profiles

        class Cls:
            """Namespace for deployed class lookups."""

            @staticmethod
            def from_name(app_name: str, class_name: str) -> Any:
                """Return a factory that records the provided class parameters."""
                assert app_name == DEFAULT_TEST_DEPLOYMENT_APP_NAME
                assert class_name == "RemoteEngine"

                def build_remote_engine(**kwargs: Any) -> dict[str, Any]:
                    """Record one synthesized Modal class constructor call."""
                    observed_kwargs.append(kwargs)
                    return kwargs

                return build_remote_engine

    monkeypatch.setattr(remote_modal_app_module, "modal", FakeModal)
    monkeypatch.setenv("COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT", "true")
    remote_modal_app_module.get_settings.cache_clear()
    remote_modal_app_module._SNAPSHOT_PROFILE_RECORDS.clear()
    expected_snapshot_profile_key = remote_modal_app_module._loader_snapshot_profile_key(
        [
            {
                "signature": remote_modal_app_module._loader_prewarm_plan_signature(
                    "UNETLoader",
                    {"unet_name": "model-a.safetensors", "weight_dtype": "default"},
                )
            }
        ]
    )
    payload = {
        "component_id": "component-1",
        "snapshot_profile_key": expected_snapshot_profile_key,
        "subgraph_prompt": {
            "1": {
                "class_type": "UNETLoader",
                "inputs": {"unet_name": "model-a.safetensors", "weight_dtype": "default"},
            }
        },
    }
    try:
        result = remote_modal_app_module._lookup_deployed_remote_engine(payload)
    finally:
        remote_modal_app_module.get_settings.cache_clear()
        remote_modal_app_module._SNAPSHOT_PROFILE_RECORDS.clear()

    assert result["snapshot_profile_key"] == expected_snapshot_profile_key
    assert observed_kwargs == [result]
    assert expected_snapshot_profile_key in snapshot_profiles


def test_lookup_deployed_remote_engine_reuses_worker_pool_slots_across_prompt_sessions(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Different prompt sessions should reuse one shared deployed RemoteEngine identity."""
    observed_kwargs: list[dict[str, Any]] = []

    class FakeModal:
        """Minimal modal SDK double that captures class construction kwargs."""

        class Cls:
            """Namespace for deployed class lookups."""

            @staticmethod
            def from_name(app_name: str, class_name: str) -> Any:
                """Return a factory that records the provided class parameters."""
                assert app_name == DEFAULT_TEST_DEPLOYMENT_APP_NAME
                assert class_name == "RemoteEngine"

                def build_remote_engine(**kwargs: Any) -> dict[str, Any]:
                    """Record one synthesized Modal class constructor call."""
                    observed_kwargs.append(kwargs)
                    return kwargs

                return build_remote_engine

    monkeypatch.setattr(remote_modal_app_module, "modal", FakeModal)

    first_result = remote_modal_app_module._lookup_deployed_remote_engine(
        {
            "component_id": "component-1",
            "prompt_id": "prompt-1",
            "remote_session": remote_modal_app_module.RemoteSessionHandle(
                session_id="session-123",
                prompt_id="prompt-1",
                owner_component_id="component-1",
            ).to_payload(),
        }
    )
    second_result = remote_modal_app_module._lookup_deployed_remote_engine(
        {
            "component_id": "component-1",
            "prompt_id": "prompt-2",
            "remote_session": remote_modal_app_module.RemoteSessionHandle(
                session_id="session-456",
                prompt_id="prompt-2",
                owner_component_id="component-1",
            ).to_payload(),
        }
    )

    assert first_result == {
        "gpu_snapshot_enabled": False,
        "worker_affinity_key": "worker-pool:slot:0",
    }
    assert second_result == {
        "gpu_snapshot_enabled": False,
        "worker_affinity_key": "worker-pool:slot:0",
    }
    assert observed_kwargs == [
        {
            "gpu_snapshot_enabled": False,
            "worker_affinity_key": "worker-pool:slot:0",
        },
        {
            "gpu_snapshot_enabled": False,
            "worker_affinity_key": "worker-pool:slot:0",
        },
    ]


def test_lookup_deployed_remote_engine_isolates_profiled_lane_overrides(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Lane-specific affinity overrides should create stable separate Modal identities."""
    observed_kwargs: list[dict[str, Any]] = []
    snapshot_profiles: dict[str, Any] = {}

    class FakeModal:
        """Minimal modal SDK double that captures class construction kwargs."""

        class Dict:
            """Namespace for fake dict lookups."""

            @staticmethod
            def from_name(dict_name: str, create_if_missing: bool = False) -> Any:
                """Return the shared fake snapshot profile store."""
                assert dict_name == "comfy-modal-sync-snapshot-profiles"
                assert create_if_missing is True
                return snapshot_profiles

        class Cls:
            """Namespace for deployed class lookups."""

            @staticmethod
            def from_name(app_name: str, class_name: str) -> Any:
                """Return a factory that records the provided class parameters."""
                assert app_name == DEFAULT_TEST_DEPLOYMENT_APP_NAME
                assert class_name == "RemoteEngine"

                def build_remote_engine(**kwargs: Any) -> dict[str, Any]:
                    """Record one synthesized Modal class constructor call."""
                    observed_kwargs.append(kwargs)
                    return kwargs

                return build_remote_engine

    monkeypatch.setattr(remote_modal_app_module, "modal", FakeModal)
    monkeypatch.setenv("COMFY_MODAL_ENABLE_GPU_MEMORY_SNAPSHOT", "true")
    remote_modal_app_module.get_settings.cache_clear()
    remote_modal_app_module._SNAPSHOT_PROFILE_RECORDS.clear()
    payload = {
        "component_id": "component-1__mapped",
        "subgraph_prompt": {
            "1": {
                "class_type": "UNETLoader",
                "inputs": {"unet_name": "model-a.safetensors", "weight_dtype": "default"},
            }
        },
    }
    try:
        first_result = remote_modal_app_module._lookup_deployed_remote_engine(
            payload,
            affinity_key_override="worker-pool:slot:0",
        )
        second_result = remote_modal_app_module._lookup_deployed_remote_engine(
            payload,
            affinity_key_override="worker-pool:slot:3",
        )
    finally:
        remote_modal_app_module.get_settings.cache_clear()
        remote_modal_app_module._SNAPSHOT_PROFILE_RECORDS.clear()

    assert first_result != second_result
    assert first_result["worker_affinity_key"] == "worker-pool:slot:0"
    assert second_result["worker_affinity_key"] == "worker-pool:slot:3"
    assert first_result["snapshot_profile_key"].startswith("loader-profile:")
    assert observed_kwargs == [first_result, second_result]


def test_load_modal_cloud_module_reloads_stale_partial_module(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Stale partially imported cloud modules should be discarded and reloaded."""
    original_module = sys.modules.get(remote_modal_app_module._MODAL_CLOUD_MODULE_NAME)
    stale_module = types.SimpleNamespace(app=None)
    sys.modules[remote_modal_app_module._MODAL_CLOUD_MODULE_NAME] = stale_module

    loaded_module = types.SimpleNamespace(app="fresh-app")

    class FakeLoader:
        """Populate the fresh replacement module during exec."""

        def create_module(self, spec: Any) -> None:
            """Use the default module creation path."""
            del spec
            return None

        def exec_module(self, module: Any) -> None:
            """Install the expected deployable app onto the reloaded module."""
            module.app = loaded_module.app

    monkeypatch.setattr(
        remote_modal_app_module.importlib.util,
        "spec_from_file_location",
        lambda *args, **kwargs: importlib.util.spec_from_loader(
            remote_modal_app_module._MODAL_CLOUD_MODULE_NAME,
            FakeLoader(),
        ),
    )
    try:
        reloaded_module = remote_modal_app_module._load_modal_cloud_module()
    finally:
        sys.modules.pop(remote_modal_app_module._MODAL_CLOUD_MODULE_NAME, None)
        if original_module is not None:
            sys.modules[remote_modal_app_module._MODAL_CLOUD_MODULE_NAME] = original_module

    assert reloaded_module is not stale_module
    assert getattr(reloaded_module, "app", None) == "fresh-app"


def test_load_modal_cloud_module_reloads_for_workflow_gpu_change(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Cloud app construction should receive the workflow-selected GPU settings."""
    module_name = remote_modal_app_module._MODAL_CLOUD_MODULE_NAME
    original_module = sys.modules.get(module_name)
    sys.modules[module_name] = types.SimpleNamespace(
        app="a100-app",
        __comfy_modal_gpu__="A100",
    )
    observed_gpu_values: list[str] = []

    class FakeLoader:
        """Populate a replacement cloud module from its injected settings."""

        def create_module(self, spec: Any) -> None:
            """Use the default module creation path."""
            del spec
            return None

        def exec_module(self, module: Any) -> None:
            """Capture the settings override and expose a deployable app."""
            settings_override = module.__comfy_modal_settings_override__
            observed_gpu_values.append(settings_override.modal_gpu)
            module.app = "b300-app"

    monkeypatch.setattr(
        remote_modal_app_module.importlib.util,
        "spec_from_file_location",
        lambda *args, **kwargs: importlib.util.spec_from_loader(module_name, FakeLoader()),
    )
    settings = remote_modal_app_module.settings_for_modal_gpu(
        remote_modal_app_module.get_settings(),
        "B300",
    )
    try:
        with remote_modal_app_module._modal_cloud_settings_override(settings):
            reloaded_module = remote_modal_app_module._load_modal_cloud_module()
    finally:
        sys.modules.pop(module_name, None)
        if original_module is not None:
            sys.modules[module_name] = original_module

    assert observed_gpu_values == ["B300"]
    assert reloaded_module.app == "b300-app"
    assert reloaded_module.__comfy_modal_gpu__ == "B300"
    assert reloaded_module.__comfy_modal_app_name__ == "comfy-modal-sync-gpu-b300"


def test_load_modal_cloud_module_reloads_for_secret_name_change(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Changing only the secret name should rebuild the cloud module with new settings."""
    module_name = remote_modal_app_module._MODAL_CLOUD_MODULE_NAME
    original_module = sys.modules.get(module_name)
    base_settings = remote_modal_app_module.get_settings()
    sys.modules[module_name] = types.SimpleNamespace(
        app="old-secret-app",
        __comfy_modal_gpu__=base_settings.modal_gpu,
        __comfy_modal_app_name__=remote_modal_app_module.modal_deployment_app_name(
            base_settings
        ),
        __comfy_modal_secret_name__="comfy",
    )
    observed_secret_names: list[str] = []

    class FakeLoader:
        """Populate a replacement cloud module from its injected settings."""

        def create_module(self, spec: Any) -> None:
            """Use the default module creation path."""
            del spec
            return None

        def exec_module(self, module: Any) -> None:
            """Capture the secret setting and expose a deployable app."""
            settings_override = module.__comfy_modal_settings_override__
            observed_secret_names.append(settings_override.modal_secret_name)
            module.app = "new-secret-app"

    monkeypatch.setattr(
        remote_modal_app_module.importlib.util,
        "spec_from_file_location",
        lambda *args, **kwargs: importlib.util.spec_from_loader(module_name, FakeLoader()),
    )
    new_settings = replace(base_settings, modal_secret_name="workflow-credentials")
    try:
        with remote_modal_app_module._modal_cloud_settings_override(new_settings):
            reloaded_module = remote_modal_app_module._load_modal_cloud_module()
    finally:
        sys.modules.pop(module_name, None)
        if original_module is not None:
            sys.modules[module_name] = original_module

    assert observed_secret_names == ["workflow-credentials"]
    assert reloaded_module.app == "new-secret-app"
    assert reloaded_module.__comfy_modal_secret_name__ == "workflow-credentials"


def test_remote_modal_installs_cloud_exception_compatibility_module(
    remote_modal_app_module: Any,
) -> None:
    """Deployed remote exceptions should have importable local definitions."""
    module_name = remote_modal_app_module._MODAL_CLOUD_MODULE_NAME
    original_module = sys.modules.pop(module_name, None)
    serialization_module = types.ModuleType(module_name)
    serialization_module.RemoteSubgraphExecutionError = type(
        "RemoteSubgraphExecutionError",
        (RuntimeError,),
        {"__module__": module_name},
    )
    sys.modules[module_name] = serialization_module
    serialized_error = pickle.dumps(
        serialization_module.RemoteSubgraphExecutionError("remote failure")
    )
    sys.modules.pop(module_name)
    try:
        remote_modal_app_module._install_modal_cloud_exception_compatibility_module()

        compatibility_module = sys.modules[module_name]
        assert (
            compatibility_module.RemoteSubgraphExecutionError
            is remote_modal_app_module.RemoteSubgraphExecutionError
        )
        assert issubclass(compatibility_module.RemoteInvocationAbandonedError, RuntimeError)
        restored_error = pickle.loads(serialized_error)
        assert isinstance(restored_error, remote_modal_app_module.RemoteSubgraphExecutionError)
        assert str(restored_error) == "remote failure"
    finally:
        sys.modules.pop(module_name, None)
        if original_module is not None:
            sys.modules[module_name] = original_module


def test_load_modal_cloud_module_clears_failed_import_from_sys_modules(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Failed cloud-module imports should not leave a poisoned cache entry behind."""
    original_module = sys.modules.pop(
        remote_modal_app_module._MODAL_CLOUD_MODULE_NAME,
        None,
    )

    class FakeLoader:
        """Raise during module execution to simulate a partial import failure."""

        def create_module(self, spec: Any) -> None:
            """Use the default module creation path."""
            del spec
            return None

        def exec_module(self, module: Any) -> None:
            """Fail while the module is being initialized."""
            module.app = None
            raise RuntimeError("boom")

    monkeypatch.setattr(
        remote_modal_app_module.importlib.util,
        "spec_from_file_location",
        lambda *args, **kwargs: importlib.util.spec_from_loader(
            remote_modal_app_module._MODAL_CLOUD_MODULE_NAME,
            FakeLoader(),
        ),
    )
    try:
        with pytest.raises(RuntimeError, match="boom"):
            remote_modal_app_module._load_modal_cloud_module()
        assert remote_modal_app_module._MODAL_CLOUD_MODULE_NAME not in sys.modules
    finally:
        sys.modules.pop(remote_modal_app_module._MODAL_CLOUD_MODULE_NAME, None)
        if original_module is not None:
            sys.modules[remote_modal_app_module._MODAL_CLOUD_MODULE_NAME] = original_module


def test_remote_modal_consumes_streamed_progress_and_result(
    remote_modal_app_module: Any,
    monkeypatch: Any,
    caplog: pytest.LogCaptureFixture,
) -> 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",
        "invocation_id": "RIV_stream_timing",
        "component_node_ids": ["7", "8"],
        "modal_gpu": "B300",
        "extra_data": {"client_id": "client-1"},
    }
    caplog.set_level(logging.INFO, logger=remote_modal_app_module.__name__)
    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",
                    "event_type": "node_progress",
                    "node_id": "7",
                    "display_node_id": "7",
                    "value": 4,
                    "max": 20,
                },
                {
                    "kind": "progress",
                    "phase": "execution_success",
                },
                {
                    "kind": "result",
                    "outputs": b"serialized-outputs",
                },
            ]
        ),
    )

    assert result == b"serialized-outputs"
    assert "Starting local Modal stream consumption component=component-1" in caplog.text
    assert "Received streamed Modal result component=component-1" in caplog.text
    assert "result_bytes=18" in caplog.text
    assert "event_count=4 progress_event_count=3" in caplog.text
    assert "Finished local Modal result stream close" in caplog.text
    assert "Finished local Modal stream consumption" in caplog.text
    assert prompt_server.messages == [
        (
            "modal_status",
            {
                "phase": "executing",
                "prompt_id": "prompt-1",
                "node_ids": ["7", "8"],
                "modal_gpu": "B300",
                "active_node_id": "7",
                "active_node_class_type": "UNETLoader",
                "active_node_role": "model_load",
            },
            "client-1",
        ),
        (
            "modal_progress",
            {
                "prompt_id": "prompt-1",
                "node_id": "7",
                "display_node_id": "7",
                "value": 4.0,
                "max": 20.0,
            },
            "client-1",
        ),
        (
            "modal_status",
            {
                "phase": "finalizing",
                "prompt_id": "prompt-1",
                "node_ids": ["7", "8"],
                "modal_gpu": "B300",
                    "status_message": "Receiving remote outputs from Modal",
            },
            "client-1",
        ),
    ]


def test_remote_stream_first_event_triggers_speculative_prewarm_once(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """The current worker must emit an event before its future affinity is prepared."""
    scheduled_payloads: list[tuple[dict[str, Any], str]] = []
    monkeypatch.setattr(
        remote_modal_app_module,
        "_schedule_speculative_affinity_prewarm",
        lambda payload, *, reason: scheduled_payloads.append((dict(payload), reason))
        or True,
    )

    payload = {
        "prompt_id": "prompt-spec",
        "component_id": "component-current",
        "speculative_remote_prewarm_target": {
            "component_id": "component-next",
        },
    }
    result = remote_modal_app_module._consume_remote_payload_stream(
        payload,
        iter(
            [
                {"kind": "remote_logs", "task_id": "ta-spec"},
                {"kind": "progress", "phase": "executing"},
                {"kind": "result", "outputs": b"serialized-outputs"},
            ]
        ),
    )

    assert result == b"serialized-outputs"
    assert scheduled_payloads == [(payload, "current_remote_stream_started")]


def test_emit_local_remote_dispatch_status_marks_component_starting(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Dispatching a remote component should immediately tell the local UI it is starting."""

    class FakePromptServer:
        """Capture websocket events emitted by local dispatch status."""

        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)

    remote_modal_app_module._emit_local_remote_dispatch_status(
        {
            "prompt_id": "prompt-1",
            "component_id": "component-1",
            "component_node_ids": ["7", "8"],
            "modal_gpu": "B300",
            "extra_data": {"client_id": "client-1"},
        }
    )

    assert prompt_server.messages == [
        (
            "modal_status",
            {
                "phase": "starting",
                "prompt_id": "prompt-1",
                "node_ids": ["7", "8"],
                "modal_gpu": "B300",
                    "status_message": "Starting remote component on Modal",
            },
            "client-1",
        )
    ]


def test_emit_local_mapped_lane_progress_start_marks_lane_as_setup_only(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Provisioning a mapped lane should emit a dedicated setup-only lane progress event."""
    emitted_progress: list[dict[str, Any]] = []

    monkeypatch.setattr(
        remote_modal_app_module,
        "_emit_local_modal_progress",
        lambda **kwargs: emitted_progress.append(kwargs),
    )

    remote_modal_app_module._emit_local_mapped_lane_progress_start(
        {
            "prompt_id": "prompt-1",
            "component_id": "component-1",
            "extra_data": {"client_id": "client-1"},
        },
        lane_index=3,
    )

    assert emitted_progress == [
        {
            "prompt_id": "prompt-1",
            "client_id": "client-1",
            "node_id": "component-1",
            "value": 0.0,
            "max_value": 1.0,
            "display_node_id": "component-1",
            "lane_id": "3",
            "item_index": None,
            "setup_only": True,
        }
    ]


def test_remote_modal_consumes_remote_log_stream_events_with_retain_release(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Stream consumers should retain and release one container-log watcher per remote payload."""
    log_stream_calls: list[tuple[str, str]] = []

    monkeypatch.setattr(
        remote_modal_app_module,
        "get_settings",
        lambda: types.SimpleNamespace(stream_remote_container_logs=True),
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_retain_remote_container_log_stream",
        lambda task_id: log_stream_calls.append(("retain", task_id)) or task_id,
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_release_remote_container_log_stream",
        lambda task_id: log_stream_calls.append(("release", task_id)),
    )

    result = remote_modal_app_module._consume_remote_payload_stream(
        {
            "prompt_id": "prompt-1",
            "component_id": "component-1",
            "component_node_ids": ["7"],
            "extra_data": {"client_id": "client-1"},
        },
        iter(
            [
                {"kind": "remote_logs", "task_id": "ta-123"},
                {"kind": "result", "outputs": b"serialized-outputs"},
            ]
        ),
    )

    assert result == b"serialized-outputs"
    assert log_stream_calls == [("retain", "ta-123"), ("release", "ta-123")]


def test_list_active_modal_containers_filters_and_classifies_managed_apps(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Container status should include every active GPU app for this ComfyUI instance."""
    settings = remote_modal_app_module.get_settings()
    b300_app_name = remote_modal_app_module.modal_deployment_app_name(
        remote_modal_app_module.settings_for_modal_gpu(settings, "B300")
    )
    request_environment_names: list[str] = []
    request_thread_ids: list[int] = []
    request_event_loops: list[Any] = []
    stopped_container_ids: list[str] = []
    caller_thread_id = threading.get_ident()
    sdk_event_loop = asyncio.new_event_loop()

    class FakeStub:
        """Return a mixed active-container list."""

        async def TaskList(self, request: Any) -> Any:
            """Capture the environment and return managed plus unrelated tasks."""
            request_environment_names.append(request.environment_name)
            request_thread_ids.append(threading.get_ident())
            request_event_loops.append(asyncio.get_running_loop())
            return types.SimpleNamespace(
                tasks=[
                    types.SimpleNamespace(
                        task_id="ta-starting",
                        app_id="ap-managed",
                        app_description=b300_app_name,
                        started_at=0.0,
                        enqueued_at=100.0,
                    ),
                    types.SimpleNamespace(
                        task_id="ta-running",
                        app_id="ap-managed",
                        app_description=b300_app_name,
                        started_at=101.0,
                        enqueued_at=99.0,
                    ),
                    types.SimpleNamespace(
                        task_id="ta-unrelated",
                        app_id="ap-other",
                        app_description="another-modal-app",
                        started_at=102.0,
                        enqueued_at=98.0,
                    ),
                ]
            )

        async def ContainerStop(self, request: Any) -> Any:
            """Capture one exact Modal task termination request."""
            stopped_container_ids.append(request.task_id)
            request_thread_ids.append(threading.get_ident())
            request_event_loops.append(asyncio.get_running_loop())
            return types.SimpleNamespace()

    class FakeClient:
        """Expose the synchronized Modal TaskList stub."""

        def __init__(self) -> None:
            """Create the fake TaskList stub."""
            self.stub = FakeStub()

    class FakeClientFactory:
        """Create the fake synchronized Modal client."""

        @staticmethod
        async def from_env() -> FakeClient:
            """Return one fake authenticated client."""
            return FakeClient()

    def synchronize_api(async_callable: Any) -> Any:
        """Run the fake SDK operation on one persistent event loop."""

        def blocking_callable(*args: Any) -> Any:
            """Synchronously marshal one call onto the fake SDK loop."""
            return sdk_event_loop.run_until_complete(async_callable(*args))

        return blocking_callable

    original_import_module = remote_modal_app_module.importlib.import_module

    def fake_import_module(name: str) -> Any:
        """Supply the Modal SDK modules used by the container status query."""
        if name == "modal._object":
            return types.SimpleNamespace(_get_environment_name=lambda _environment: "test-env")
        if name == "modal.environments":
            return types.SimpleNamespace(ensure_env=lambda environment: environment or "test-env")
        if name == "modal.client":
            return types.SimpleNamespace(_Client=FakeClientFactory)
        if name == "modal._utils.async_utils":
            return types.SimpleNamespace(synchronize_api=synchronize_api)
        if name == "modal.exception":
            return types.SimpleNamespace(Error=RuntimeError)
        if name == "modal_proto.api_pb2":
            return types.SimpleNamespace(
                TaskListRequest=lambda **kwargs: types.SimpleNamespace(**kwargs),
                ContainerStopRequest=lambda **kwargs: types.SimpleNamespace(**kwargs),
            )
        return original_import_module(name)

    monkeypatch.setattr(remote_modal_app_module, "modal", object())
    monkeypatch.setattr(remote_modal_app_module.importlib, "import_module", fake_import_module)

    first_containers = asyncio.run(
        remote_modal_app_module.list_active_modal_containers(settings)
    )
    containers = asyncio.run(remote_modal_app_module.list_active_modal_containers(settings))
    stopped = asyncio.run(
        remote_modal_app_module.stop_managed_modal_container("ta-running", settings)
    )

    assert first_containers == containers
    assert [container.container_id for container in containers] == ["ta-running", "ta-starting"]
    assert [container.state for container in containers] == ["running", "starting"]
    assert all(container.modal_gpu == "B300" for container in containers)
    assert containers[0].as_dict()["started_at"] == 101.0
    assert containers[0].as_dict()["estimated_gpu_cost_per_second"] == 0.001972
    assert stopped
    assert stopped_container_ids == ["ta-running"]
    assert request_environment_names == ["test-env", "test-env", "test-env"]
    assert request_thread_ids
    assert all(thread_id != caller_thread_id for thread_id in request_thread_ids)
    assert request_event_loops == [sdk_event_loop] * 4
    sdk_event_loop.close()


def test_stop_managed_modal_container_verifies_ownership_before_exact_stop(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Container termination must target only a currently listed managed task."""
    stopped_container_ids: list[str] = []

    async def fake_list_active(_settings: Any) -> list[Any]:
        """Return one container already filtered to this ComfyUI installation."""
        return [
            remote_modal_app_module.ModalContainerStatus(
                container_id="ta-managed",
                app_id="ap-managed",
                app_name="comfy-modal-sync-B300",
                modal_gpu="B300",
                estimated_gpu_cost_per_second=0.001972,
                state="running",
                enqueued_at=100.0,
                started_at=101.0,
            )
        ]

    def fake_stop(
        _client_module: Any,
        _api_pb2: Any,
        container_id: str,
    ) -> None:
        """Record the exact task passed to the synchronized SDK bridge."""
        stopped_container_ids.append(container_id)

    original_import_module = remote_modal_app_module.importlib.import_module

    def fake_import_module(name: str) -> Any:
        """Supply the minimal modules needed after ownership verification."""
        if name == "modal.client":
            return types.SimpleNamespace()
        if name == "modal.exception":
            return types.SimpleNamespace(Error=RuntimeError)
        if name == "modal_proto.api_pb2":
            return types.SimpleNamespace()
        return original_import_module(name)

    monkeypatch.setattr(
        remote_modal_app_module,
        "list_active_modal_containers",
        fake_list_active,
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_stop_modal_task_synchronously",
        fake_stop,
    )
    monkeypatch.setattr(remote_modal_app_module.importlib, "import_module", fake_import_module)

    assert asyncio.run(
        remote_modal_app_module.stop_managed_modal_container(
            "ta-managed",
            remote_modal_app_module.get_settings(),
        )
    )
    assert not asyncio.run(
        remote_modal_app_module.stop_managed_modal_container(
            "ta-unrelated",
            remote_modal_app_module.get_settings(),
        )
    )
    assert stopped_container_ids == ["ta-managed"]


def test_modal_gpu_estimated_rates_cover_supported_aliases(
    remote_modal_app_module: Any,
) -> None:
    """Published GPU estimates should cover every selectable billed GPU identity."""
    rates = remote_modal_app_module.MODAL_GPU_ESTIMATED_USD_PER_SECOND

    assert set(rates) == set(remote_modal_app_module.MODAL_GPU_TYPES)
    assert rates["A100"] == rates["A100-40GB"]
    assert rates["H100!"] == rates["H100"]
    assert rates["B200+"] == rates["B200"]
    assert remote_modal_app_module.MODAL_GPU_PRICING_EFFECTIVE_DATE == "2026-08-13"


def test_completed_modal_billing_interval_uses_hourly_resolution_and_buffer(
    remote_modal_app_module: Any,
) -> None:
    """Billing should use the latest full UTC hour after a collection buffer."""
    interval_start, interval_end, next_refresh_at = (
        remote_modal_app_module._completed_modal_billing_interval(
            datetime(2026, 8, 19, 8, 5, tzinfo=timezone.utc)
        )
    )

    assert interval_start == datetime(2026, 8, 19, 6, 0, tzinfo=timezone.utc)
    assert interval_end == datetime(2026, 8, 19, 7, 0, tzinfo=timezone.utc)
    assert next_refresh_at == datetime(2026, 8, 19, 8, 10, tzinfo=timezone.utc)


def test_get_hourly_modal_app_billing_filters_and_caches_gpu_app(
    remote_modal_app_module: Any,
    monkeypatch: Any,
    caplog: pytest.LogCaptureFixture,
) -> None:
    """Hourly billing should quietly select and cache one GPU app/environment."""
    settings = remote_modal_app_module.get_settings()
    selected_settings = remote_modal_app_module.settings_for_modal_gpu(
        settings,
        "B300",
    )
    app_name = remote_modal_app_module.modal_deployment_app_name(selected_settings)
    interval_start = datetime(2026, 8, 19, 7, 0, tzinfo=timezone.utc)
    report_calls: list[dict[str, Any]] = []

    def workspace_billing_report(**kwargs: Any) -> list[dict[str, Any]]:
        """Return matching and unrelated hourly billing rows."""
        report_calls.append(kwargs)
        return [
            {
                "object_id": "ap-selected",
                "description": app_name,
                "environment_name": "main",
                "interval_start": interval_start,
                "cost": Decimal("0.11799566"),
            },
            {
                "object_id": "ap-other-env",
                "description": app_name,
                "environment_name": "dev",
                "interval_start": interval_start,
                "cost": Decimal("9.99"),
            },
            {
                "object_id": "ap-other-app",
                "description": "unrelated",
                "environment_name": "main",
                "interval_start": interval_start,
                "cost": Decimal("8.88"),
            },
        ]

    original_import_module = remote_modal_app_module.importlib.import_module

    def fake_import_module(name: str) -> Any:
        """Supply the public Modal billing and environment SDK surfaces."""
        if name == "modal._object":
            return types.SimpleNamespace(_get_environment_name=lambda _environment: "main")
        if name == "modal.environments":
            return types.SimpleNamespace(ensure_env=lambda environment: environment or "main")
        if name == "modal.billing":
            return types.SimpleNamespace(
                workspace_billing_report=workspace_billing_report
            )
        if name == "modal.exception":
            return types.SimpleNamespace(Error=RuntimeError)
        return original_import_module(name)

    monkeypatch.setattr(remote_modal_app_module, "modal", object())
    monkeypatch.setattr(
        remote_modal_app_module.importlib,
        "import_module",
        fake_import_module,
    )
    remote_modal_app_module._MODAL_HOURLY_BILLING_CACHE.clear()
    remote_modal_app_module._MODAL_HOURLY_BILLING_ERROR_CACHE.clear()
    caplog.set_level(logging.INFO, logger=remote_modal_app_module.__name__)
    now = datetime(2026, 8, 19, 8, 15, tzinfo=timezone.utc)

    first_status = asyncio.run(
        remote_modal_app_module.get_hourly_modal_app_billing(
            "B300",
            settings,
            now=now,
        )
    )
    cached_status = asyncio.run(
        remote_modal_app_module.get_hourly_modal_app_billing(
            "B300",
            settings,
            now=now + timedelta(minutes=30),
        )
    )

    assert first_status is cached_status
    assert first_status.app_id == "ap-selected"
    assert first_status.app_name == app_name
    assert first_status.modal_gpu == "B300"
    assert first_status.environment_name == "main"
    assert first_status.app_cost_usd_before_credits == Decimal("0.11799566")
    assert first_status.has_usage is True
    assert first_status.interval_start == interval_start
    assert first_status.interval_end == datetime(
        2026,
        8,
        19,
        8,
        tzinfo=timezone.utc,
    )
    assert first_status.next_refresh_at == datetime(
        2026,
        8,
        19,
        9,
        10,
        tzinfo=timezone.utc,
    )
    assert report_calls == [
        {
            "start": interval_start,
            "end": datetime(2026, 8, 19, 8, 0, tzinfo=timezone.utc),
            "resolution": "h",
        }
    ]
    assert "Fetched Modal hourly billing" not in caplog.text


def test_get_hourly_modal_app_billing_sums_historical_ids_in_default_environment(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Redeployments within one implicit environment should sum every app ID."""
    settings = remote_modal_app_module.get_settings()
    selected_settings = remote_modal_app_module.settings_for_modal_gpu(settings, "L4")
    app_name = remote_modal_app_module.modal_deployment_app_name(selected_settings)
    interval_start = datetime(2026, 8, 19, 7, 0, tzinfo=timezone.utc)
    original_import_module = remote_modal_app_module.importlib.import_module

    def fake_import_module(name: str) -> Any:
        """Supply an implicit environment with two historical app identities."""
        if name == "modal._object":
            return types.SimpleNamespace(_get_environment_name=lambda _environment: None)
        if name == "modal.environments":
            return types.SimpleNamespace(ensure_env=lambda _environment: None)
        if name == "modal.billing":
            return types.SimpleNamespace(
                workspace_billing_report=lambda **_kwargs: [
                    {
                        "object_id": "ap-workspace-default-old",
                        "description": app_name,
                        "environment_name": "main",
                        "interval_start": interval_start,
                        "cost": Decimal("0.17"),
                    },
                    {
                        "object_id": "ap-workspace-default-new",
                        "description": app_name,
                        "environment_name": "main",
                        "interval_start": interval_start,
                        "cost": Decimal("0.25"),
                    },
                ]
            )
        if name == "modal.exception":
            return types.SimpleNamespace(Error=RuntimeError)
        return original_import_module(name)

    monkeypatch.setattr(remote_modal_app_module, "modal", object())
    monkeypatch.setattr(
        remote_modal_app_module.importlib,
        "import_module",
        fake_import_module,
    )
    remote_modal_app_module._MODAL_HOURLY_BILLING_CACHE.clear()
    remote_modal_app_module._MODAL_HOURLY_BILLING_ERROR_CACHE.clear()

    status = asyncio.run(
        remote_modal_app_module.get_hourly_modal_app_billing(
            "L4",
            settings,
            now=datetime(2026, 8, 19, 8, 15, tzinfo=timezone.utc),
        )
    )

    assert status.app_id is None
    assert status.environment_name == "main"
    assert status.app_cost_usd_before_credits == Decimal("0.42")


def test_modal_hourly_billing_rejects_multiple_implicit_environments(
    remote_modal_app_module: Any,
) -> None:
    """An implicit environment should reject only genuinely distinct environments."""
    interval_start = datetime(2026, 8, 19, 7, 0, tzinfo=timezone.utc)
    rows = [
        {
            "object_id": "ap-main",
            "description": "shared-app-name",
            "environment_name": "main",
            "interval_start": interval_start,
            "cost": Decimal("0.17"),
        },
        {
            "object_id": "ap-dev",
            "description": "shared-app-name",
            "environment_name": "dev",
            "interval_start": interval_start,
            "cost": Decimal("0.25"),
        },
    ]

    with pytest.raises(
        remote_modal_app_module.ModalBillingStatusError,
        match="multiple environments",
    ):
        remote_modal_app_module._matching_modal_hourly_billing_rows(
            rows,
            app_name="shared-app-name",
            environment_name=None,
            interval_start=interval_start,
        )


def test_get_hourly_modal_app_billing_reports_zero_for_no_usage(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """An absent app row should render as zero cost for that completed hour."""
    original_import_module = remote_modal_app_module.importlib.import_module

    def fake_import_module(name: str) -> Any:
        """Supply an empty public Modal billing report."""
        if name == "modal._object":
            return types.SimpleNamespace(_get_environment_name=lambda _environment: "main")
        if name == "modal.environments":
            return types.SimpleNamespace(ensure_env=lambda environment: environment or "main")
        if name == "modal.billing":
            return types.SimpleNamespace(workspace_billing_report=lambda **_kwargs: [])
        if name == "modal.exception":
            return types.SimpleNamespace(Error=RuntimeError)
        return original_import_module(name)

    monkeypatch.setattr(remote_modal_app_module, "modal", object())
    monkeypatch.setattr(
        remote_modal_app_module.importlib,
        "import_module",
        fake_import_module,
    )
    remote_modal_app_module._MODAL_HOURLY_BILLING_CACHE.clear()
    remote_modal_app_module._MODAL_HOURLY_BILLING_ERROR_CACHE.clear()

    status = asyncio.run(
        remote_modal_app_module.get_hourly_modal_app_billing(
            "L4",
            now=datetime(2026, 8, 19, 8, 15, tzinfo=timezone.utc),
        )
    )

    assert status.app_cost_usd_before_credits == Decimal("0")
    assert status.has_usage is False
    assert status.as_dict()["resolution"] == "hour"
    assert status.as_dict()["collection_buffer_seconds"] == 600


def test_get_hourly_modal_app_billing_caches_report_failures(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Fast status polling should attempt a failed billing interval only once."""
    report_calls = 0

    def workspace_billing_report(**_kwargs: Any) -> list[dict[str, Any]]:
        """Simulate a workspace that cannot expose billing reports."""
        nonlocal report_calls
        report_calls += 1
        raise RuntimeError("billing access denied")

    original_import_module = remote_modal_app_module.importlib.import_module

    def fake_import_module(name: str) -> Any:
        """Supply a failing public Modal billing report."""
        if name == "modal._object":
            return types.SimpleNamespace(_get_environment_name=lambda _environment: "main")
        if name == "modal.environments":
            return types.SimpleNamespace(ensure_env=lambda environment: environment or "main")
        if name == "modal.billing":
            return types.SimpleNamespace(
                workspace_billing_report=workspace_billing_report
            )
        if name == "modal.exception":
            return types.SimpleNamespace(Error=RuntimeError)
        return original_import_module(name)

    monkeypatch.setattr(remote_modal_app_module, "modal", object())
    monkeypatch.setattr(
        remote_modal_app_module.importlib,
        "import_module",
        fake_import_module,
    )
    remote_modal_app_module._MODAL_HOURLY_BILLING_CACHE.clear()
    remote_modal_app_module._MODAL_HOURLY_BILLING_ERROR_CACHE.clear()
    now = datetime(2026, 8, 19, 8, 15, tzinfo=timezone.utc)

    for _attempt in range(2):
        with pytest.raises(
            remote_modal_app_module.ModalBillingStatusError,
            match="billing access denied",
        ):
            asyncio.run(
                remote_modal_app_module.get_hourly_modal_app_billing(
                    "L4",
                    now=now,
                )
            )

    assert report_calls == 1


def test_remote_modal_stops_consuming_stream_after_terminal_result(
    remote_modal_app_module: Any,
) -> None:
    """The local stream consumer should not wait for extra events after the final result arrives."""

    class FakeStreamEvents:
        """Iterator that fails if the consumer requests post-result events."""

        def __init__(self) -> None:
            """Initialize the deterministic event sequence."""
            self._events = [
                {"kind": "remote_logs", "task_id": "ta-123"},
                {"kind": "result", "outputs": b"serialized-outputs"},
            ]
            self._index = 0
            self.close_calls = 0

        def __iter__(self) -> "FakeStreamEvents":
            """Return the iterator itself."""
            return self

        def __next__(self) -> dict[str, Any]:
            """Return the next event and reject any post-result polling."""
            if self._index >= len(self._events):
                raise AssertionError("The stream consumer requested events after the terminal result.")
            next_event = self._events[self._index]
            self._index += 1
            return next_event

        def close(self) -> None:
            """Record one best-effort close from the local consumer."""
            self.close_calls += 1

    stream_events = FakeStreamEvents()

    result = remote_modal_app_module._consume_remote_payload_stream(
        {
            "prompt_id": "prompt-1",
            "component_id": "component-1",
            "component_node_ids": ["7"],
            "extra_data": {"client_id": "client-1"},
        },
        stream_events,
    )

    assert result == b"serialized-outputs"
    assert stream_events.close_calls == 1


def test_remote_modal_cli_log_stream_mirrors_lines_to_stderr(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """CLI-backed container log streaming should mirror complete prefixed lines into local stderr."""

    class FakeStdout:
        """Expose a deterministic binary read interface for the fake Modal CLI process."""

        def __init__(self, chunks: list[bytes]) -> None:
            """Store the binary chunks that should be returned to the caller."""
            self._chunks = chunks
            self._index = 0

        def read(self, size: int = -1) -> bytes:
            """Return one queued binary chunk or EOF when exhausted."""
            del size
            if self._index >= len(self._chunks):
                return b""
            chunk = self._chunks[self._index]
            self._index += 1
            return chunk

    class FakeProcess:
        """Minimal subprocess stub used to emulate `modal container logs -f`."""

        def __init__(self) -> None:
            """Expose a stdout pipe and process lifecycle methods."""
            self.stdout = FakeStdout([b"session reuse\n", b"session miss\n"])
            self.returncode: int | None = None

        def poll(self) -> int | None:
            """Report process completion after all fake lines have been read."""
            if self.stdout._index >= len(self.stdout._chunks):
                self.returncode = 0
            return self.returncode

        def terminate(self) -> None:
            """Mark the fake process as stopped."""
            self.returncode = 0

        def wait(self, timeout: float | None = None) -> int:
            """Return the final process exit code."""
            self.returncode = 0
            return 0

        def kill(self) -> None:
            """Mark the fake process as killed."""
            self.returncode = 0

    stderr_buffer = StringIO()
    monkeypatch.setattr(remote_modal_app_module.shutil, "which", lambda name: "/usr/bin/modal")
    monkeypatch.setattr(remote_modal_app_module.subprocess, "Popen", lambda *args, **kwargs: FakeProcess())
    monkeypatch.setattr(
        remote_modal_app_module.select,
        "select",
        lambda streams, _write, _error, _timeout: (
            streams if streams[0]._index < len(streams[0]._chunks) else [],
            [],
            [],
        ),
    )
    monkeypatch.setattr(remote_modal_app_module.sys, "stderr", stderr_buffer)

    result = remote_modal_app_module._stream_remote_container_logs_via_modal_cli(
        "ta-123",
        threading.Event(),
    )

    assert result is True
    assert stderr_buffer.getvalue() == (
        "[modal:ta-123] session reuse\n"
        "[modal:ta-123] session miss\n"
    )


def test_remote_modal_log_stream_prefers_cli_before_sdk(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """The local log watcher should avoid the SDK path when the Modal CLI is available."""
    backend_calls: list[str] = []

    monkeypatch.setattr(
        remote_modal_app_module,
        "_stream_remote_container_logs_via_modal_cli",
        lambda task_id, stop_event: backend_calls.append(f"cli:{task_id}") or True,
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_stream_remote_container_logs_via_modal_sdk",
        lambda task_id, stop_event: backend_calls.append(f"sdk:{task_id}") or True,
    )

    remote_modal_app_module._run_remote_container_log_stream("ta-123", threading.Event())

    assert backend_calls == ["cli:ta-123"]


def test_remote_modal_log_stream_survives_short_container_reuse_gap(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """A reused task id should keep one watcher and avoid replaying container history."""
    created_threads: list[Any] = []
    created_timers: list[Any] = []

    class FakeThread:
        """Record watcher lifecycle without starting a real thread."""

        def __init__(self, **kwargs: Any) -> None:
            """Store thread construction arguments."""
            self.kwargs = kwargs
            self.started = False
            self.join_calls: list[float | None] = []
            created_threads.append(self)

        def start(self) -> None:
            """Mark the watcher as alive."""
            self.started = True

        def is_alive(self) -> bool:
            """Return whether the fake watcher has started."""
            return self.started

        def join(self, timeout: float | None = None) -> None:
            """Record one bounded watcher join."""
            self.join_calls.append(timeout)

    class FakeTimer:
        """Expose a manually fired daemon idle timer."""

        def __init__(
            self,
            interval: float,
            function: Any,
            args: tuple[Any, ...],
        ) -> None:
            """Store the delayed callback without scheduling it."""
            self.interval = interval
            self.function = function
            self.args = args
            self.daemon = False
            self.started = False
            self.cancelled = False
            created_timers.append(self)

        def start(self) -> None:
            """Mark the timer as scheduled."""
            self.started = True

        def cancel(self) -> None:
            """Mark the scheduled callback as cancelled."""
            self.cancelled = True

        def fire(self) -> None:
            """Run the callback when it has not been cancelled."""
            if not self.cancelled:
                self.function(*self.args)

    monkeypatch.setattr(remote_modal_app_module, "_REMOTE_CONTAINER_LOG_STREAMS", {})
    monkeypatch.setattr(remote_modal_app_module.threading, "Thread", FakeThread)
    monkeypatch.setattr(remote_modal_app_module.threading, "Timer", FakeTimer)

    assert remote_modal_app_module._retain_remote_container_log_stream("ta-reused") == "ta-reused"
    stream_state = remote_modal_app_module._REMOTE_CONTAINER_LOG_STREAMS["ta-reused"]
    remote_modal_app_module._release_remote_container_log_stream("ta-reused")

    assert len(created_threads) == 1
    assert len(created_timers) == 1
    assert created_timers[0].daemon is True
    assert stream_state.stop_event.is_set() is False

    remote_modal_app_module._retain_remote_container_log_stream("ta-reused")

    assert len(created_threads) == 1
    assert created_timers[0].cancelled is True
    assert stream_state.refcount == 1

    remote_modal_app_module._release_remote_container_log_stream("ta-reused")
    assert len(created_timers) == 2
    created_timers[1].fire()

    assert stream_state.stop_event.is_set() is True
    assert stream_state.thread.join_calls == [0.2]
    assert "ta-reused" not in remote_modal_app_module._REMOTE_CONTAINER_LOG_STREAMS


def test_remote_modal_consumes_streamed_executed_outputs_and_previews(
    remote_modal_app_module: Any,
    monkeypatch: Any,
    serialization_module: Any,
) -> None:
    """The local Modal client should forward streamed executed outputs and previews."""
    from PIL import Image
    from protocol import BinaryEventTypes

    class FakePromptServer:
        """Capture websocket events emitted by streamed remote UI output updates."""

        def __init__(self) -> None:
            """Initialize the event sink."""
            self.messages: list[tuple[Any, Any, str | None]] = []

        def send_sync(self, event: Any, data: Any, sid: str | None) -> None:
            """Record one emitted websocket message."""
            self.messages.append((event, data, sid))

    preview_buffer = BytesIO()
    Image.new("RGB", (2, 2), color="red").save(preview_buffer, format="PNG")

    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"],
        "extra_data": {"client_id": "client-1"},
    }
    result = remote_modal_app_module._consume_remote_payload_stream(
        payload,
        iter(
            [
                {
                    "kind": "progress",
                    "event_type": "executed",
                    "node_id": "7",
                    "display_node_id": "7",
                    "output": {
                        "images": [
                            {
                                "filename": "preview.png",
                                "subfolder": "",
                                "type": "temp",
                            }
                        ]
                    },
                },
                {
                    "kind": "progress",
                    "event_type": "preview",
                    "node_id": "7",
                    "display_node_id": "7",
                    "parent_node_id": None,
                    "real_node_id": "7",
                    "image_type": "PNG",
                    "image_bytes": serialization_module.serialize_value(preview_buffer.getvalue()),
                    "max_size": 256,
                },
                {
                    "kind": "result",
                    "outputs": b"serialized-outputs",
                },
            ]
        ),
    )

    assert result == b"serialized-outputs"
    assert prompt_server.messages[0] == (
        "executed",
        {
            "prompt_id": "prompt-1",
            "node": "7",
            "display_node": "7",
            "output": {
                "images": [
                    {
                        "filename": "preview.png",
                        "subfolder": "",
                        "type": "temp",
                    }
                ]
            },
        },
        "client-1",
    )

    preview_event, preview_payload, preview_sid = prompt_server.messages[1]
    preview_image, preview_metadata = preview_payload
    assert preview_event == BinaryEventTypes.PREVIEW_IMAGE_WITH_METADATA
    assert preview_sid == "client-1"
    assert preview_image[0] == "PNG"
    assert preview_image[2] == 256
    assert preview_image[1].size == (2, 2)
    assert preview_metadata == {
        "node_id": "7",
        "prompt_id": "prompt-1",
        "display_node_id": "7",
        "real_node_id": "7",
    }


def test_remote_modal_consumes_streamed_boundary_output_preview_targets(
    remote_modal_app_module: Any,
    monkeypatch: Any,
    serialization_module: Any,
) -> None:
    """A streamed remote boundary IMAGE should synthesize local PreviewImage executed events."""
    torch = pytest.importorskip("torch")
    image_tensor = torch.zeros((1, 8, 8, 3), dtype=torch.float32)

    class FakePromptServer:
        """Capture websocket events emitted by boundary preview synthesis."""

        def __init__(self) -> None:
            """Initialize the event sink."""
            self.messages: list[tuple[Any, Any, str | None]] = []

        def send_sync(self, event: Any, data: Any, sid: str | None) -> None:
            """Record one emitted websocket message."""
            self.messages.append((event, data, sid))

    class FakePreviewImage:
        """Minimal PreviewImage double that returns deterministic UI payloads."""

        def save_images(self, images: Any) -> dict[str, Any]:
            """Return a fake UI payload for the supplied image tensor."""
            assert torch.equal(images, image_tensor)
            return {
                "ui": {
                    "images": [
                        {
                            "filename": "temp_preview.png",
                            "subfolder": "",
                            "type": "temp",
                        }
                    ]
                }
            }

    prompt_server = FakePromptServer()
    monkeypatch.setattr(remote_modal_app_module, "_lookup_local_prompt_server", lambda: prompt_server)
    monkeypatch.setitem(sys.modules, "nodes", types.SimpleNamespace(PreviewImage=FakePreviewImage))

    payload = {
        "prompt_id": "prompt-1",
        "component_id": "component-1",
        "component_node_ids": ["7"],
        "extra_data": {"client_id": "client-1"},
    }
    result = remote_modal_app_module._consume_remote_payload_stream(
        payload,
        iter(
            [
                {
                    "kind": "progress",
                    "event_type": "boundary_output",
                    "node_id": "7",
                    "output_index": 0,
                    "io_type": "IMAGE",
                    "is_list": False,
                    "preview_target_node_ids": ["9"],
                    "value": serialization_module.serialize_value(image_tensor),
                },
                {
                    "kind": "result",
                    "outputs": b"serialized-outputs",
                },
            ]
        ),
    )

    assert result == b"serialized-outputs"
    assert prompt_server.messages == [
        (
            "executed",
            {
                "prompt_id": "prompt-1",
                "node": "9",
                "display_node": "9",
                "output": {
                    "images": [
                        {
                            "filename": "temp_preview.png",
                            "subfolder": "",
                            "type": "temp",
                        }
                    ]
                },
            },
            "client-1",
        )
    ]


def test_modal_cloud_tracing_prompt_server_emits_numeric_node_progress(
    modal_cloud_module: Any,
) -> None:
    """The cloud tracing prompt server should forward active-node numeric progress updates."""
    observed_updates: list[dict[str, Any]] = []
    server = modal_cloud_module._TracingPromptServer(
        "component-1",
        {"7": {"class_type": "KSampler", "inputs": {}}},
        status_callback=observed_updates.append,
    )

    server.send_sync("executing", {"node": "7"}, None)
    server.send_sync(
        "progress_state",
        {
            "prompt_id": "component-1",
            "nodes": {
                "7": {
                    "node_id": "7",
                    "display_node_id": "7",
                    "real_node_id": "7",
                    "state": "running",
                    "value": 5,
                    "max": 20,
                }
            },
        },
        None,
    )

    assert observed_updates[0]["phase"] == "executing"
    assert observed_updates[1] == {
        "event_type": "node_progress",
        "node_id": "7",
        "display_node_id": "7",
        "real_node_id": "7",
        "value": 5.0,
        "max": 20.0,
    }


def test_modal_cloud_tracing_prompt_server_ignores_trivial_node_progress(
    modal_cloud_module: Any,
) -> None:
    """The cloud tracing prompt server should ignore 0/1 progress updates from non-progress nodes."""
    observed_updates: list[dict[str, Any]] = []
    server = modal_cloud_module._TracingPromptServer(
        "component-1",
        {"18": {"class_type": "CLIPTextEncode", "inputs": {}}},
        status_callback=observed_updates.append,
    )

    server.send_sync("executing", {"node": "18"}, None)
    server.send_sync(
        "progress_state",
        {
            "prompt_id": "component-1",
            "nodes": {
                "18": {
                    "node_id": "18",
                    "display_node_id": "18",
                    "real_node_id": "18",
                    "state": "running",
                    "value": 0,
                    "max": 1,
                }
            },
        },
        None,
    )

    assert observed_updates == [
        {
            "phase": "executing",
            "active_node_id": "18",
            "active_node_class_type": "CLIPTextEncode",
            "active_node_role": "conditioning",
        }
    ]


def test_modal_cloud_tracing_prompt_server_preserves_llm_progress_metadata(
    modal_cloud_module: Any,
) -> None:
    """Structured LLM stages and token telemetry should survive remote tracing."""
    observed_updates: list[dict[str, Any]] = []
    server = modal_cloud_module._TracingPromptServer(
        "component-1",
        {"9": {"class_type": "ModalLLM", "inputs": {}}},
        status_callback=observed_updates.append,
    )

    server.send_sync(
        "modal_llm_progress",
        {
            "node_id": "9",
            "stage": "generating",
            "message": "Generating",
            "value": 23,
            "max": 128,
            "unit": "tokens",
            "time_to_first_token_seconds": 1.25,
            "tokens_per_second": 17.5,
        },
        None,
    )

    assert observed_updates == [
        {
            "event_type": "node_progress",
            "node_id": "9",
            "display_node_id": "9",
            "real_node_id": "9",
            "value": 23.0,
            "max": 128.0,
            "stage": "generating",
            "message": "Generating",
            "indeterminate": False,
            "unit": "tokens",
            "time_to_first_token_seconds": 1.25,
            "tokens_per_second": 17.5,
        }
    ]


def test_modal_cloud_tracing_prompt_server_emits_executed_outputs(
    modal_cloud_module: Any,
) -> None:
    """The cloud tracing prompt server should stream node UI outputs as executed events."""
    observed_updates: list[dict[str, Any]] = []
    server = modal_cloud_module._TracingPromptServer(
        "component-1",
        {"7": {"class_type": "PreviewImage", "inputs": {}}},
        status_callback=observed_updates.append,
    )

    server.send_sync(
        "executed",
        {
            "node": "7",
            "display_node": "7",
            "output": {"images": [{"filename": "preview.png"}]},
        },
        None,
    )

    assert observed_updates == [
        {
            "event_type": "executed",
            "node_id": "7",
            "display_node_id": "7",
            "output": {"images": [{"filename": "preview.png"}]},
        }
    ]


def test_modal_cloud_emits_restored_node_cache_events(
    modal_cloud_module: Any,
) -> None:
    """Persisted node-cache restores should surface one cached-node marker per restored node."""
    observed_updates: list[dict[str, Any]] = []

    modal_cloud_module._emit_restored_node_cache_events(observed_updates.append, ["7", "9"])

    assert observed_updates == [
        {
            "event_type": "node_cached",
            "node_id": "7",
            "display_node_id": "7",
            "real_node_id": "7",
        },
        {
            "event_type": "node_cached",
            "node_id": "9",
            "display_node_id": "9",
            "real_node_id": "9",
        },
    ]


def test_modal_cloud_tracing_prompt_server_emits_boundary_image_outputs(
    modal_cloud_module: Any,
) -> None:
    """The cloud tracing prompt server should stream configured boundary IMAGE outputs once cached."""
    torch = pytest.importorskip("torch")
    image_tensor = torch.zeros((1, 4, 4, 3), dtype=torch.float32)
    observed_updates: list[dict[str, Any]] = []
    server = modal_cloud_module._TracingPromptServer(
        "component-1",
        {
            "7": {"class_type": "VAEDecode", "inputs": {}},
            "8": {"class_type": "OtherNode", "inputs": {}},
        },
        status_callback=observed_updates.append,
    )
    cache_entries = {
        "7": types.SimpleNamespace(outputs=[image_tensor]),
    }
    server.configure_boundary_output_stream(
        boundary_outputs=[
            {
                "node_id": "7",
                "output_index": 0,
                "io_type": "IMAGE",
                "is_list": False,
                "preview_target_node_ids": ["9"],
            }
        ],
        lookup_cache_entry=lambda node_id: cache_entries.get(node_id),
    )

    server.send_sync("executing", {"node": "7"}, None)
    server.send_sync("executing", {"node": "8"}, None)

    assert observed_updates[0]["phase"] == "executing"
    assert observed_updates[1] == {
        "event_type": "boundary_output",
        "node_id": "7",
        "output_index": 0,
        "io_type": "IMAGE",
        "is_list": False,
        "preview_target_node_ids": ["9"],
        "value": image_tensor,
    }
    assert observed_updates[2]["phase"] == "executing"


def test_remote_modal_consumes_streamed_tensor_result_payload(
    remote_modal_app_module: Any,
    serialization_module: Any,
) -> None:
    """The local stream consumer should accept JSON-safe serialized tensor outputs."""
    torch = pytest.importorskip("torch")
    tensor = torch.arange(5, dtype=torch.float32)

    result = remote_modal_app_module._consume_remote_payload_stream(
        {
            "prompt_id": "prompt-1",
            "component_id": "component-1",
            "component_node_ids": ["7"],
            "extra_data": {"client_id": "client-1"},
        },
        iter(
            [
                {
                    "kind": "result",
                    "outputs": [serialization_module.serialize_value(tensor)],
                },
            ]
        ),
    )

    decoded_outputs = serialization_module.deserialize_node_outputs(result)
    assert len(decoded_outputs) == 1
    assert torch.equal(decoded_outputs[0], tensor)


def test_remote_modal_requires_manual_deploy_when_auto_deploy_disabled(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Remote mode should fail clearly when auto-deploy and ephemeral fallback are both disabled."""

    class FakeLookupError(Exception):
        """Stand-in for Modal deployed lookup failures."""

    class FakeModal:
        """Minimal modal SDK double with deployed lookup failure types."""

        exception = types.SimpleNamespace(
            NotFoundError=FakeLookupError,
            ExecutionError=FakeLookupError,
            InvalidError=FakeLookupError,
        )

        class Cls:
            """Namespace for deployed class lookups."""

            @staticmethod
            def from_name(app_name: str, class_name: str) -> Any:
                """Simulate a missing deployed app."""
                raise FakeLookupError("not deployed")

    monkeypatch.setattr(remote_modal_app_module, "modal", FakeModal)
    monkeypatch.setenv("COMFY_MODAL_AUTO_DEPLOY", "false")
    monkeypatch.setenv("COMFY_MODAL_ALLOW_EPHEMERAL_FALLBACK", "false")
    remote_modal_app_module.get_settings.cache_clear()
    try:
        try:
            remote_modal_app_module._invoke_modal_payload_blocking(
                {"component_id": "component-1"},
                b"{}",
            )
        except remote_modal_app_module.ModalRemoteInvocationError as exc:
            message = str(exc)
        else:
            raise AssertionError("Expected ModalRemoteInvocationError to be raised.")
    finally:
        remote_modal_app_module.get_settings.cache_clear()

    assert "requires a deployed Modal app or a successful first-run auto-deploy" in message
    assert "COMFY_MODAL_ALLOW_EPHEMERAL_FALLBACK=true" in message


def test_invoke_remote_engine_propagates_local_interrupt_to_modal(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """The local proxy should propagate ComfyUI interrupts to the remote Modal call."""

    class FakeInterrupt(Exception):
        """Stand-in for ComfyUI's InterruptProcessingException."""

    observed_cancellation_events: list[threading.Event] = []
    interrupt_checks = iter([False, True, True])

    def fake_blocking_invoke(
        payload: dict[str, Any],
        kwargs_payload: bytes,
        cancellation_event: threading.Event | None = None,
    ) -> bytes:
        assert cancellation_event is not None
        observed_cancellation_events.append(cancellation_event)
        while not cancellation_event.is_set():
            time.sleep(0.01)
        raise RuntimeError("remote interrupted")

    def fake_local_processing_interrupted() -> bool:
        return next(interrupt_checks, True)

    monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "remote")
    monkeypatch.setattr(remote_modal_app_module, "modal", object())
    monkeypatch.setattr(remote_modal_app_module, "_invoke_modal_payload_blocking", fake_blocking_invoke)
    monkeypatch.setattr(
        remote_modal_app_module,
        "_local_processing_interrupted",
        fake_local_processing_interrupted,
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_raise_local_interrupt",
        lambda: (_ for _ in ()).throw(FakeInterrupt()),
    )

    with pytest.raises(FakeInterrupt):
        remote_modal_app_module.invoke_remote_engine(
            {"component_id": "component-1", "payload_kind": "subgraph"},
            b"{}",
        )

    assert len(observed_cancellation_events) == 1
    assert observed_cancellation_events[0].is_set()


def test_invoke_remote_engine_async_propagates_local_interrupt_to_modal(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """The async local proxy should also propagate ComfyUI interrupts to the remote Modal call."""

    class FakeInterrupt(Exception):
        """Stand-in for ComfyUI's InterruptProcessingException."""

    observed_cancellation_events: list[threading.Event] = []
    interrupt_checks = iter([False, True, True])

    def fake_blocking_invoke(
        payload: dict[str, Any],
        kwargs_payload: bytes,
        cancellation_event: threading.Event | None = None,
    ) -> bytes:
        assert cancellation_event is not None
        observed_cancellation_events.append(cancellation_event)
        while not cancellation_event.is_set():
            time.sleep(0.01)
        raise RuntimeError("remote interrupted")

    def fake_local_processing_interrupted() -> bool:
        return next(interrupt_checks, True)

    monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "remote")
    monkeypatch.setattr(remote_modal_app_module, "modal", object())
    monkeypatch.setattr(remote_modal_app_module, "_invoke_modal_payload_blocking", fake_blocking_invoke)
    monkeypatch.setattr(
        remote_modal_app_module,
        "_local_processing_interrupted",
        fake_local_processing_interrupted,
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_raise_local_interrupt",
        lambda: (_ for _ in ()).throw(FakeInterrupt()),
    )

    with pytest.raises(FakeInterrupt):
        asyncio.run(
            remote_modal_app_module.invoke_remote_engine_async(
                {"component_id": "component-1", "payload_kind": "subgraph"},
                b"{}",
            )
        )

    assert len(observed_cancellation_events) == 1
    assert observed_cancellation_events[0].is_set()


def test_invoke_remote_engine_releases_prompt_when_cancelled_during_deploy(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Cancellation during Modal deploy/provisioning should release the local prompt promptly."""

    class FakeInterrupt(Exception):
        """Stand-in for ComfyUI's InterruptProcessingException."""

    release_blocking_call = threading.Event()
    observed_cancellation_events: list[threading.Event] = []
    interrupt_checks = iter([False, True, True, True])
    remote_interrupt_payloads: list[dict[str, Any]] = []

    def fake_blocking_invoke(
        payload: dict[str, Any],
        kwargs_payload: bytes,
        cancellation_event: threading.Event | None = None,
    ) -> bytes:
        """Simulate a worker stuck in Modal deployment until the test releases it."""
        del payload, kwargs_payload
        assert cancellation_event is not None
        observed_cancellation_events.append(cancellation_event)
        release_blocking_call.wait(timeout=5.0)
        return b"late-response"

    def fake_local_processing_interrupted() -> bool:
        """Report a local interrupt after the first outer wait poll."""
        return next(interrupt_checks, True)

    monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "remote")
    monkeypatch.setenv("COMFY_MODAL_REMOTE_CANCEL_GRACE_SECONDS", "0.01")
    remote_modal_app_module.get_settings.cache_clear()
    monkeypatch.setattr(remote_modal_app_module, "modal", object())
    monkeypatch.setattr(remote_modal_app_module, "_invoke_modal_payload_blocking", fake_blocking_invoke)
    monkeypatch.setattr(
        remote_modal_app_module,
        "_local_processing_interrupted",
        fake_local_processing_interrupted,
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_request_remote_interrupt",
        lambda payload: remote_interrupt_payloads.append(dict(payload)) or True,
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_raise_local_interrupt",
        lambda: (_ for _ in ()).throw(FakeInterrupt()),
    )

    try:
        with pytest.raises(FakeInterrupt):
            remote_modal_app_module.invoke_remote_engine(
                {"prompt_id": "prompt-1", "component_id": "component-1", "payload_kind": "subgraph"},
                b"{}",
            )
    finally:
        release_blocking_call.set()
        remote_modal_app_module.get_settings.cache_clear()

    assert len(observed_cancellation_events) == 1
    assert observed_cancellation_events[0].is_set()
    assert [payload["component_id"] for payload in remote_interrupt_payloads] == ["component-1"]


def test_invoke_remote_engine_payload_refuses_dispatch_after_prestart_cancel(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """A cancellation observed during deploy should prevent the later remote payload call."""
    cancellation_event = threading.Event()
    cancellation_event.set()
    remote_interrupt_payloads: list[dict[str, Any]] = []

    class FakeExecuteMethod:
        """Remote method double that must not be called after cancellation."""

        def remote(self, payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
            """Fail if the cancelled payload reaches Modal execution."""
            del payload, kwargs_payload
            raise AssertionError("cancelled payload should not dispatch")

    monkeypatch.setattr(
        remote_modal_app_module,
        "_request_remote_interrupt",
        lambda payload: remote_interrupt_payloads.append(dict(payload)) or True,
    )

    with pytest.raises(remote_modal_app_module.ModalRemoteInvocationError, match="cancelled before"):
        remote_modal_app_module._invoke_remote_engine_payload(
            types.SimpleNamespace(execute_payload=FakeExecuteMethod()),
            {"prompt_id": "prompt-1", "component_id": "component-1"},
            b"{}",
            cancellation_event,
        )

    assert [payload["component_id"] for payload in remote_interrupt_payloads] == ["component-1"]


def test_remote_modal_interrupt_callback_writes_shared_control_flag(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """The local interrupt callback should write to the shared Modal Dict control store."""

    class FakeInterruptStore:
        """Simple Modal Dict double that records written interruption flags."""

        def __init__(self) -> None:
            """Initialize captured writes."""
            self.put_calls: list[tuple[str, Any]] = []

        def put(self, key: str, value: Any, *, skip_if_exists: bool = False) -> bool:
            """Record one interrupt flag write."""
            self.put_calls.append((key, value))
            return True

    interrupt_store = FakeInterruptStore()

    class FakeModalDict:
        """Minimal modal.Dict shim that returns the fake interrupt store."""

        @staticmethod
        def from_name(
            name: str,
            *,
            environment_name: str | None = None,
            create_if_missing: bool = False,
            client: Any | None = None,
        ) -> FakeInterruptStore:
            return interrupt_store

    monkeypatch.setattr(
        remote_modal_app_module,
        "modal",
        types.SimpleNamespace(Dict=FakeModalDict),
    )
    monkeypatch.setenv("COMFY_MODAL_INTERRUPT_DICT_NAME", "shared-interrupts")
    remote_modal_app_module.get_settings.cache_clear()
    remote_modal_app_module._MODAL_INTERRUPT_DICTS.clear()
    try:
        callback = remote_modal_app_module._build_remote_interrupt_callback(
            object(),
            {"prompt_id": "prompt-1", "component_id": "component-2"},
        )
        assert callback is not None
        callback()
    finally:
        remote_modal_app_module.get_settings.cache_clear()
        remote_modal_app_module._MODAL_INTERRUPT_DICTS.clear()

    assert len(interrupt_store.put_calls) == 1
    interrupt_key, interrupt_value = interrupt_store.put_calls[0]
    assert interrupt_key == "prompt-1:component-2"
    assert isinstance(interrupt_value["requested_at"], float)


def test_request_remote_interrupt_async_uses_shared_control_async_put(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Async local cancellation should use Modal Dict.put.aio instead of blocking put."""

    class FakePut:
        """Modal Dict put double that records sync and async writes separately."""

        def __init__(self) -> None:
            """Initialize captured writes."""
            self.sync_calls: list[tuple[str, Any]] = []
            self.async_calls: list[tuple[str, Any]] = []

        def __call__(self, key: str, value: Any, *, skip_if_exists: bool = False) -> bool:
            """Record one blocking write."""
            del skip_if_exists
            self.sync_calls.append((key, value))
            return True

        async def aio(self, key: str, value: Any, *, skip_if_exists: bool = False) -> bool:
            """Record one async write."""
            del skip_if_exists
            self.async_calls.append((key, value))
            return True

    class FakeInterruptStore:
        """Simple Modal Dict double exposing put.aio."""

        def __init__(self) -> None:
            """Initialize the fake async put handle."""
            self.put = FakePut()

    interrupt_store = FakeInterruptStore()

    class FakeModalDict:
        """Minimal modal.Dict shim that returns the fake interrupt store."""

        @staticmethod
        def from_name(
            name: str,
            *,
            environment_name: str | None = None,
            create_if_missing: bool = False,
            client: Any | None = None,
        ) -> FakeInterruptStore:
            return interrupt_store

    monkeypatch.setattr(
        remote_modal_app_module,
        "modal",
        types.SimpleNamespace(Dict=FakeModalDict),
    )
    monkeypatch.setenv("COMFY_MODAL_INTERRUPT_DICT_NAME", "shared-interrupts")
    remote_modal_app_module.get_settings.cache_clear()
    remote_modal_app_module._MODAL_INTERRUPT_DICTS.clear()
    try:
        wrote_interrupt = asyncio.run(
            remote_modal_app_module._request_remote_interrupt_async(
                {"prompt_id": "prompt-1", "component_id": "component-2"}
            )
        )
    finally:
        remote_modal_app_module.get_settings.cache_clear()
        remote_modal_app_module._MODAL_INTERRUPT_DICTS.clear()

    assert wrote_interrupt is True
    assert interrupt_store.put.sync_calls == []
    assert len(interrupt_store.put.async_calls) == 1
    interrupt_key, interrupt_value = interrupt_store.put.async_calls[0]
    assert interrupt_key == "prompt-1:component-2"
    assert isinstance(interrupt_value["requested_at"], float)


def test_request_remote_modal_prompt_interrupt_cancels_active_components(
    remote_modal_app_module: Any,
) -> None:
    """Prompt-level cancellation should interrupt every active remote component."""
    cancellation_event = threading.Event()
    interrupt_calls: list[str] = []

    with remote_modal_app_module._registered_active_remote_invocation(
        {"prompt_id": "prompt-1", "component_id": "component-1"},
        cancellation_event,
        lambda: interrupt_calls.append("component-1"),
    ):
        assert remote_modal_app_module.active_remote_modal_prompt_ids() == {"prompt-1"}
        assert remote_modal_app_module.request_remote_modal_prompt_interrupt("prompt-1") is True

    assert cancellation_event.is_set()
    assert interrupt_calls == ["component-1"]
    assert remote_modal_app_module.active_remote_modal_prompt_ids() == set()


def test_invoke_remote_engine_payload_stream_detects_local_interrupt_without_outer_sync(
    remote_modal_app_module: Any,
    serialization_module: Any,
    monkeypatch: Any,
) -> None:
    """The blocking streamed Modal bridge should propagate interrupts without relying on the outer wrapper loop."""
    cancellation_event = threading.Event()
    remote_release_event = threading.Event()
    interrupt_calls: list[str] = []
    interrupt_checks = iter([False, True, True])

    def fake_local_processing_interrupted() -> bool:
        """Report a local interrupt after the first poll interval."""
        return next(interrupt_checks, True)

    def fake_interrupt_remote_call() -> None:
        """Record the propagated remote interrupt and let the fake stream finish."""
        interrupt_calls.append("interrupt")
        remote_release_event.set()

    def fake_stream_events() -> Iterator[dict[str, Any]]:
        """Block until the local bridge requests cancellation, then yield one final result."""
        while not remote_release_event.is_set():
            time.sleep(0.01)
        yield {
            "kind": "result",
            "outputs": serialization_module.serialize_node_outputs(("done",)),
        }

    class FakeStreamMethod:
        """Minimal Modal stream method shim."""

        def remote_gen(self, payload: dict[str, Any], kwargs_payload: bytes) -> Iterator[dict[str, Any]]:
            """Return the fake delayed stream for this request."""
            del payload, kwargs_payload
            return fake_stream_events()

    monkeypatch.setattr(
        remote_modal_app_module,
        "_local_processing_interrupted",
        fake_local_processing_interrupted,
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_build_remote_interrupt_callback",
        lambda remote_engine, payload: fake_interrupt_remote_call,
    )

    response = remote_modal_app_module._invoke_remote_engine_payload(
        types.SimpleNamespace(execute_payload_stream=FakeStreamMethod()),
        {
            "component_id": "component-1",
            "payload_kind": "subgraph",
            "prompt_id": "prompt-1",
            "component_node_ids": ["1"],
            "extra_data": {"client_id": "client-1"},
        },
        b"{}",
        cancellation_event,
    )

    assert serialization_module.deserialize_node_outputs(response) == ("done",)
    assert cancellation_event.is_set()
    assert interrupt_calls == ["interrupt"]


def test_invoke_remote_engine_payload_releases_local_prompt_after_cancel_grace(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """The local proxy should not wait forever for a cancelled Modal call to unwind."""
    cancellation_event = threading.Event()
    interrupt_calls: list[str] = []
    interrupt_checks = iter([False, True, True, True])

    def fake_local_processing_interrupted() -> bool:
        """Report a local interrupt after the first poll interval."""
        return next(interrupt_checks, True)

    def fake_interrupt_remote_call() -> None:
        """Record the propagated remote interrupt."""
        interrupt_calls.append("interrupt")

    def never_returns() -> bytes:
        """Simulate a Modal call stuck in remote work after cancellation."""
        while True:
            time.sleep(0.01)

    monkeypatch.setenv("COMFY_MODAL_REMOTE_CANCEL_GRACE_SECONDS", "0.05")
    remote_modal_app_module.get_settings.cache_clear()
    monkeypatch.setattr(
        remote_modal_app_module,
        "_local_processing_interrupted",
        fake_local_processing_interrupted,
    )

    try:
        with pytest.raises(remote_modal_app_module.ModalRemoteInvocationError):
            remote_modal_app_module._invoke_remote_call_with_interrupts(
                payload={"prompt_id": "prompt-1", "component_id": "component-1"},
                invoke_remote_call=never_returns,
                interrupt_remote_call=fake_interrupt_remote_call,
                cancellation_event=cancellation_event,
            )
    finally:
        remote_modal_app_module.get_settings.cache_clear()

    assert cancellation_event.is_set()
    assert interrupt_calls == ["interrupt"]


def test_invoke_mapped_remote_engine_async_runs_explicit_mapped_phase_items(
    remote_modal_app_module: Any,
    serialization_module: Any,
    monkeypatch: Any,
) -> None:
    """Mapped remote execution should run the explicit mapped phase once per item in order."""
    observed_calls: list[tuple[str, dict[str, Any]]] = []
    progress_updates: list[dict[str, Any]] = []

    def fake_execute_subgraph_prompt(
        payload: dict[str, Any],
        hydrated_inputs: dict[str, Any],
        node_mapping: Any = None,
    ) -> tuple[str]:
        assert payload["payload_kind"] == "subgraph"
        assert payload["suppress_status_stream"] is True
        observed_calls.append((str(payload["component_id"]), dict(hydrated_inputs)))
        return (f"done:{hydrated_inputs['remote_input_0']}",)

    monkeypatch.setattr(
        remote_modal_app_module,
        "_execute_subgraph_prompt",
        fake_execute_subgraph_prompt,
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_emit_local_modal_progress",
        lambda **kwargs: progress_updates.append(kwargs),
    )
    payload = {
        "payload_kind": "mapped_subgraph",
        "component_id": "6",
        "prompt_id": "prompt-1",
        "mapped_input": {"proxy_input_name": "remote_input_0", "io_type": "STRING"},
        "boundary_outputs": [
            {
                "proxy_output_name": "7_text",
                "node_id": "7",
                "output_index": 0,
                "io_type": "STRING",
                "is_list": False,
                "mapped_output": True,
            }
        ],
        "static_to_mapped_boundaries": [],
        "static_phase": {
            "component_node_ids": [],
            "subgraph_prompt": {},
            "boundary_inputs": [],
            "boundary_outputs": [],
            "execute_node_ids": [],
        },
        "mapped_phase": {
            "component_node_ids": ["7"],
            "subgraph_prompt": {
                "7": {
                    "class_type": "RemoteStringEcho",
                    "inputs": {"text": ["remote_input_0", 0]},
                }
            },
            "boundary_inputs": [
                {
                    "proxy_input_name": "remote_input_0",
                    "io_type": "STRING",
                    "targets": [{"node_id": "7", "input_name": "text"}],
                }
            ],
            "boundary_outputs": [
                {
                    "proxy_output_name": "7_text",
                    "node_id": "7",
                    "output_index": 0,
                    "io_type": "STRING",
                    "is_list": False,
                    "mapped_output": True,
                }
            ],
            "execute_node_ids": ["7"],
        },
        "extra_data": {"client_id": "client-1"},
    }
    response = asyncio.run(
        remote_modal_app_module._invoke_mapped_remote_engine_async(
            payload,
            serialization_module.serialize_node_inputs(
                {"remote_input_0": ["a", "b", "c", "d"]}
            ),
        )
    )

    assert serialization_module.deserialize_node_outputs(response) == (
        ["done:a", "done:b", "done:c", "done:d"],
    )
    assert observed_calls == [
        ("6::item:0", {"remote_input_0": "a"}),
        ("6::item:1", {"remote_input_0": "b"}),
        ("6::item:2", {"remote_input_0": "c"}),
        ("6::item:3", {"remote_input_0": "d"}),
    ]
    assert progress_updates[0]["value"] == 0.0
    assert progress_updates[0].get("lane_id") is None
    assert progress_updates[-1]["value"] == 4.0


@pytest.mark.parametrize(
    ("module_fixture_name", "aggregate_function_name"),
    [
        ("remote_modal_app_module", "_aggregate_mapped_outputs"),
        ("modal_cloud_module", "_aggregate_mapped_phase_outputs"),
    ],
)
def test_mapped_latent_aggregation_preserves_scheduler_items_when_shapes_differ(
    request: Any,
    module_fixture_name: str,
    aggregate_function_name: str,
) -> None:
    """Both execution paths should return heterogeneous LATENT values as scheduler items."""
    torch = pytest.importorskip("torch")
    target_module = request.getfixturevalue(module_fixture_name)
    aggregate_outputs = getattr(target_module, aggregate_function_name)
    first_latent = {"samples": torch.zeros((1, 4, 32, 32), dtype=torch.float32)}
    second_latent = {"samples": torch.zeros((1, 4, 35, 35), dtype=torch.float32)}
    payload = {
        "boundary_outputs": [
            {
                "io_type": "LATENT",
                "is_list": False,
                "mapped_output": True,
                "scheduler_is_list": True,
            }
        ]
    }

    result = aggregate_outputs([(first_latent,), (second_latent,)], payload)

    assert len(result) == 1
    assert len(result[0]) == 2
    assert result[0][0] is first_latent
    assert result[0][1] is second_latent


def test_implicit_mapping_propagates_scheduler_list_through_downstream_component(
    api_intercept_module: Any,
    modal_executor_module: Any,
    remote_modal_app_module: Any,
    serialization_module: Any,
    settings_module: Any,
    sync_engine_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """Heterogeneous implicit LATENT outputs should map into an ordinary local consumer."""
    torch = pytest.importorskip("torch")

    class PromptSeedList:
        """Expose integer seeds as one scheduler list output."""

        RETURN_TYPES = ("INT",)
        RETURN_NAMES = ("seeds",)
        OUTPUT_IS_LIST = (True,)

    class ModalMapInput:
        """Represent the explicit queue-time mapped input marker."""

        RETURN_TYPES = ("*",)
        RETURN_NAMES = ("value",)
        OUTPUT_IS_LIST = (False,)

    class RemoteSeed:
        """Publish an inexpensive mapped integer across a remote component boundary."""

        RETURN_TYPES = ("INT",)
        RETURN_NAMES = ("seed",)
        OUTPUT_IS_LIST = (False,)

    class RemoteLatent:
        """Publish one latent whose shape depends on its mapped seed."""

        RETURN_TYPES = ("LATENT",)
        RETURN_NAMES = ("samples",)
        OUTPUT_IS_LIST = (False,)

    class LocalVAEDecode:
        """Model the ordinary mapping contract used by ComfyUI's VAE decoder."""

        RETURN_TYPES = ("IMAGE",)
        RETURN_NAMES = ("image",)
        OUTPUT_IS_LIST = (False,)

        @staticmethod
        def decode(samples: dict[str, Any]) -> int:
            """Require a LATENT mapping and return its width."""
            return int(samples["samples"].shape[-1])

    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=tmp_path / "storage",
        remote_storage_root="/storage",
        custom_nodes_archive_name="custom_nodes_bundle.zip",
        comfyui_root=None,
        custom_nodes_dir=tmp_path / "custom_nodes",
    )
    sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings)
    fake_nodes_module = type(
        "FakeNodesModule",
        (),
        {
            "NODE_CLASS_MAPPINGS": {
                "PromptSeedList": PromptSeedList,
                "ModalMapInput": ModalMapInput,
                "RemoteSeed": RemoteSeed,
                "RemoteLatent": RemoteLatent,
                "VAEDecode": LocalVAEDecode,
            },
            "NODE_DISPLAY_NAME_MAPPINGS": {},
        },
    )()
    workflow = {
        "nodes": [
            {"id": 49, "properties": {"is_modal_remote": False}},
            {"id": 50, "properties": {"is_modal_remote": False}},
            {"id": 51, "properties": {"is_modal_remote": True}},
            {"id": 1, "properties": {"is_modal_remote": True}},
            {"id": 11, "properties": {"is_modal_remote": False}},
        ]
    }
    prompt = {
        "49": {"class_type": "PromptSeedList", "inputs": {}},
        "50": {"class_type": "ModalMapInput", "inputs": {"value": ["49", 0]}},
        "51": {"class_type": "RemoteSeed", "inputs": {"seed": ["50", 0]}},
        "1": {"class_type": "RemoteLatent", "inputs": {"seed": ["51", 0]}},
        "11": {"class_type": "VAEDecode", "inputs": {"samples": ["1", 0]}},
    }

    rewritten_prompt, summary = api_intercept_module.rewrite_prompt_for_modal(
        prompt=prompt,
        workflow=workflow,
        sync_engine=sync_engine,
        settings=settings,
        nodes_module=fake_nodes_module,
        extra_data={"prompt_id": "implicit-list-regression"},
    )

    downstream_proxy_id = rewritten_prompt["1"]["class_type"]
    downstream_proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[downstream_proxy_id]
    downstream_payload = rewritten_prompt["1"]["inputs"]["original_node_data"]
    downstream_execution_payload = modal_executor_module._rehydrate_proxy_payload(
        downstream_payload,
        unique_id="1",
    )

    assert summary.remote_component_ids == ["51", "1"]
    assert summary.mapped_component_ids == ["1", "51"]
    assert (
        downstream_execution_payload["boundary_outputs"][0]["scheduler_is_list"]
        is True
    )
    assert downstream_proxy_class.OUTPUT_IS_LIST == [True, False]
    assert rewritten_prompt["11"]["inputs"]["samples"] == ["1", 0]

    async def fake_invoke_remote_engine_async(
        payload: dict[str, Any],
        kwargs_payload: bytes,
    ) -> bytes:
        """Return distinct latent shapes for the downstream mapped items."""
        if not payload.get("execute_node_ids"):
            return serialization_module.serialize_node_outputs(())
        hydrated_inputs = serialization_module.deserialize_node_inputs(
            kwargs_payload
        )
        seed = int(hydrated_inputs["remote_input_0"])
        latent_size = 32 if seed == 10 else 35
        latent = {
            "samples": torch.zeros(
                (1, 4, latent_size, latent_size),
                dtype=torch.float32,
            )
        }
        return serialization_module.serialize_node_outputs((latent,))

    monkeypatch.setattr(
        remote_modal_app_module,
        "invoke_remote_engine_async",
        fake_invoke_remote_engine_async,
    )

    class ImplicitMappingClient:
        """Run the downstream payload through implicit mapped aggregation."""

        async def execute_payload_async(
            self,
            payload: dict[str, Any],
            kwargs: dict[str, Any],
        ) -> tuple[Any, ...]:
            """Return deserialized outputs from the implicit mapped subgraph runner."""
            invoke_implicitly_mapped = (
                remote_modal_app_module._invoke_implicitly_mapped_subgraph_async
            )
            response = await invoke_implicitly_mapped(
                payload,
                serialization_module.serialize_node_inputs(kwargs),
            )
            return serialization_module.deserialize_node_outputs(response)

    modal_executor_module.set_remote_executor_client_factory(
        lambda: ImplicitMappingClient()
    )
    try:
        proxy_result = asyncio.run(
            downstream_proxy_class.execute(
                original_node_data=downstream_payload,
                remote_input_0=[10, 11],
                unique_id="1",
            )
        )
    finally:
        modal_executor_module.set_remote_executor_client_factory(None)

    latent_items = proxy_result.result[0]
    assert isinstance(latent_items, list)
    assert [LocalVAEDecode.decode(latent) for latent in latent_items] == [32, 35]


def test_invoke_mapped_remote_engine_async_splits_int_inputs_for_direct_targets(
    remote_modal_app_module: Any,
    serialization_module: Any,
    monkeypatch: Any,
) -> None:
    """Mapped remote execution should itemize list(INT) inputs even when ModalMapInput stays local."""
    observed_calls: list[tuple[str, dict[str, Any]]] = []

    def fake_execute_subgraph_prompt(
        payload: dict[str, Any],
        hydrated_inputs: dict[str, Any],
        node_mapping: Any = None,
    ) -> tuple[str]:
        observed_calls.append((str(payload["component_id"]), dict(hydrated_inputs)))
        return (f"seed:{hydrated_inputs['remote_input_0']}",)

    monkeypatch.setattr(
        remote_modal_app_module,
        "_execute_subgraph_prompt",
        fake_execute_subgraph_prompt,
    )

    payload = {
        "payload_kind": "mapped_subgraph",
        "component_id": "12",
        "prompt_id": "prompt-1",
        "mapped_input": {"proxy_input_name": "remote_input_0", "io_type": "INT"},
        "boundary_outputs": [
            {
                "proxy_output_name": "12_latent",
                "node_id": "12",
                "output_index": 0,
                "io_type": "STRING",
                "is_list": False,
                "mapped_output": True,
            }
        ],
        "static_to_mapped_boundaries": [],
        "static_phase": {
            "component_node_ids": [],
            "subgraph_prompt": {},
            "boundary_inputs": [],
            "boundary_outputs": [],
            "execute_node_ids": [],
        },
        "mapped_phase": {
            "component_node_ids": ["12"],
            "subgraph_prompt": {
                "12": {
                    "class_type": "RemoteSampler",
                    "inputs": {"seed": ["remote_input_0", 0]},
                }
            },
            "boundary_inputs": [
                {
                    "proxy_input_name": "remote_input_0",
                    "io_type": "INT",
                    "targets": [{"node_id": "12", "input_name": "seed"}],
                }
            ],
            "boundary_outputs": [
                {
                    "proxy_output_name": "12_latent",
                    "node_id": "12",
                    "output_index": 0,
                    "io_type": "STRING",
                    "is_list": False,
                    "mapped_output": True,
                }
            ],
            "execute_node_ids": ["12"],
        },
        "extra_data": {"client_id": "client-1"},
    }

    response = asyncio.run(
        remote_modal_app_module._invoke_mapped_remote_engine_async(
            payload,
            serialization_module.serialize_node_inputs(
                {"remote_input_0": [10, 11, 12]}
            ),
        )
    )

    assert serialization_module.deserialize_node_outputs(response) == (
        [
            "seed:10",
            "seed:11",
            "seed:12",
        ],
    )
    assert observed_calls == [
        ("12::item:0", {"remote_input_0": 10}),
        ("12::item:1", {"remote_input_0": 11}),
        ("12::item:2", {"remote_input_0": 12}),
    ]


def test_invoke_mapped_remote_engine_async_executes_static_branch_once(
    remote_modal_app_module: Any,
    serialization_module: Any,
    monkeypatch: Any,
) -> None:
    """Mapped remote execution should run the explicit static phase once and inject its bridge outputs."""
    observed_execute_node_ids: list[tuple[str, tuple[str, ...]]] = []

    def fake_execute_subgraph_prompt(
        payload: dict[str, Any],
        hydrated_inputs: dict[str, Any],
        node_mapping: Any = None,
    ) -> tuple[str, ...]:
        observed_execute_node_ids.append(
            (
                str(payload["component_id"]),
                tuple(str(node_id) for node_id in payload.get("execute_node_ids", [])),
            )
        )
        if str(payload["component_id"]).endswith("::static"):
            assert tuple(payload.get("execute_node_ids", [])) == ("1", "3")
            assert [output["proxy_output_name"] for output in payload.get("boundary_outputs", [])] == [
                "3_text",
                "static_input_0",
            ]
            return ("static-output", "shared-model")

        assert tuple(payload.get("execute_node_ids", [])) == ("7",)
        assert hydrated_inputs["static_input_0"] == "shared-model"
        assert [output["proxy_output_name"] for output in payload.get("boundary_outputs", [])] == ["7_text"]
        return (f"mapped:{hydrated_inputs['remote_input_1']}",)

    monkeypatch.setattr(
        remote_modal_app_module,
        "_execute_subgraph_prompt",
        fake_execute_subgraph_prompt,
    )

    payload = {
        "payload_kind": "mapped_subgraph",
        "component_id": "1",
        "prompt_id": "prompt-1",
        "mapped_input": {"proxy_input_name": "remote_input_1", "io_type": "STRING"},
        "static_to_mapped_boundaries": [
            {
                "proxy_name": "static_input_0",
                "node_id": "1",
                "output_index": 0,
                "io_type": "MODEL",
                "is_list": False,
                "targets": [{"node_id": "7", "input_name": "model"}],
            }
        ],
        "static_phase": {
            "component_node_ids": ["1", "3"],
            "subgraph_prompt": {
                "1": {"class_type": "RemoteModel", "inputs": {}},
                "3": {"class_type": "RemoteSampler", "inputs": {"model": ["1", 0]}},
            },
            "boundary_inputs": [],
            "boundary_outputs": [
                {
                    "proxy_output_name": "3_text",
                    "node_id": "3",
                    "output_index": 0,
                    "io_type": "STRING",
                    "is_list": False,
                },
                {
                    "proxy_output_name": "static_input_0",
                    "node_id": "1",
                    "output_index": 0,
                    "io_type": "MODEL",
                    "is_list": False,
                },
            ],
            "execute_node_ids": ["1", "3"],
        },
        "mapped_phase": {
            "component_node_ids": ["6", "7"],
            "subgraph_prompt": {
                "6": {"class_type": "ModalMapInput", "inputs": {"value": ["remote_input_1", 0]}},
                "7": {
                    "class_type": "RemoteSampler",
                    "inputs": {"model": ["static_input_0", 0], "latent": ["6", 0]},
                },
            },
            "boundary_inputs": [
                {
                    "proxy_input_name": "remote_input_1",
                    "io_type": "STRING",
                    "targets": [{"node_id": "6", "input_name": "value"}],
                },
                {
                    "proxy_input_name": "static_input_0",
                    "io_type": "MODEL",
                    "targets": [{"node_id": "7", "input_name": "model"}],
                },
            ],
            "boundary_outputs": [
                {
                    "proxy_output_name": "7_text",
                    "node_id": "7",
                    "output_index": 0,
                    "io_type": "STRING",
                    "is_list": False,
                    "mapped_output": True,
                }
            ],
            "execute_node_ids": ["7"],
        },
        "boundary_outputs": [
            {
                "proxy_output_name": "3_text",
                "node_id": "3",
                "output_index": 0,
                "io_type": "STRING",
                "is_list": False,
                "mapped_output": False,
            },
            {
                "proxy_output_name": "7_text",
                "node_id": "7",
                "output_index": 0,
                "io_type": "STRING",
                "is_list": False,
                "mapped_output": True,
            },
        ],
        "execute_node_ids": ["3", "7"],
        "static_execute_node_ids": ["1", "3"],
        "mapped_execute_node_ids": ["7"],
        "extra_data": {"client_id": "client-1"},
    }

    response = asyncio.run(
        remote_modal_app_module._invoke_mapped_remote_engine_async(
            payload,
            serialization_module.serialize_node_inputs(
                {"remote_input_1": ["a", "b"]}
            ),
        )
    )

    assert serialization_module.deserialize_node_outputs(response) == (
        "static-output",
        ["mapped:a", "mapped:b"],
    )
    assert observed_execute_node_ids[0] == ("1::static", ("1", "3"))
    assert observed_execute_node_ids[1:] == [
        ("1::item:0", ("7",)),
        ("1::item:1", ("7",)),
    ]


def test_modal_cloud_execute_mapped_subgraph_payload_injects_static_bridges(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """The cloud runtime should execute the static phase once and feed its outputs into each mapped item."""
    observed_calls: list[tuple[str, tuple[str, ...], dict[str, Any]]] = []

    def fake_execute_subgraph_prompt(
        payload: dict[str, Any],
        hydrated_inputs: dict[str, Any],
        custom_nodes_root: Any,
        status_callback: Any = None,
        cancellation_event: Any = None,
        interrupt_store: Any = None,
        interrupt_flag_key: Any = None,
    ) -> tuple[str, ...]:
        observed_calls.append(
            (
                str(payload["component_id"]),
                tuple(str(node_id) for node_id in payload.get("execute_node_ids", [])),
                dict(hydrated_inputs),
            )
        )
        if str(payload["component_id"]).endswith("::static"):
            return ("static-output", "shared-model")
        return (f"mapped:{hydrated_inputs['remote_input_1']}:{hydrated_inputs['static_input_0']}",)

    monkeypatch.setattr(
        modal_cloud_module,
        "_execute_subgraph_prompt",
        fake_execute_subgraph_prompt,
    )

    payload = {
        "payload_kind": "mapped_subgraph",
        "component_id": "cloud-1",
        "prompt_id": "prompt-1",
        "mapped_input": {"proxy_input_name": "remote_input_1", "io_type": "STRING"},
        "static_to_mapped_boundaries": [
            {
                "proxy_name": "static_input_0",
                "node_id": "1",
                "output_index": 0,
                "io_type": "MODEL",
                "is_list": False,
                "targets": [{"node_id": "7", "input_name": "model"}],
            }
        ],
        "static_phase": {
            "component_node_ids": ["1", "3"],
            "subgraph_prompt": {
                "1": {"class_type": "RemoteModel", "inputs": {}},
                "3": {"class_type": "RemoteSampler", "inputs": {"model": ["1", 0]}},
            },
            "boundary_inputs": [],
            "boundary_outputs": [
                {
                    "proxy_output_name": "3_text",
                    "node_id": "3",
                    "output_index": 0,
                    "io_type": "STRING",
                    "is_list": False,
                },
                {
                    "proxy_output_name": "static_input_0",
                    "node_id": "1",
                    "output_index": 0,
                    "io_type": "MODEL",
                    "is_list": False,
                },
            ],
            "execute_node_ids": ["1", "3"],
        },
        "mapped_phase": {
            "component_node_ids": ["6", "7"],
            "subgraph_prompt": {
                "6": {"class_type": "ModalMapInput", "inputs": {"value": ["remote_input_1", 0]}},
                "7": {
                    "class_type": "RemoteSampler",
                    "inputs": {"model": ["static_input_0", 0], "latent": ["6", 0]},
                },
            },
            "boundary_inputs": [
                {
                    "proxy_input_name": "remote_input_1",
                    "io_type": "STRING",
                    "targets": [{"node_id": "6", "input_name": "value"}],
                },
                {
                    "proxy_input_name": "static_input_0",
                    "io_type": "MODEL",
                    "targets": [{"node_id": "7", "input_name": "model"}],
                },
            ],
            "boundary_outputs": [
                {
                    "proxy_output_name": "7_text",
                    "node_id": "7",
                    "output_index": 0,
                    "io_type": "STRING",
                    "is_list": False,
                    "mapped_output": True,
                }
            ],
            "execute_node_ids": ["7"],
        },
        "boundary_outputs": [
            {
                "proxy_output_name": "3_text",
                "node_id": "3",
                "output_index": 0,
                "io_type": "STRING",
                "is_list": False,
                "mapped_output": False,
            },
            {
                "proxy_output_name": "7_text",
                "node_id": "7",
                "output_index": 0,
                "io_type": "STRING",
                "is_list": False,
                "mapped_output": True,
            },
        ],
        "extra_data": {"client_id": "client-1"},
    }

    outputs = modal_cloud_module._execute_mapped_subgraph_payload(
        payload,
        {"remote_input_1": ["a", "b"]},
        None,
    )

    assert outputs == (
        "static-output",
        ["mapped:a:shared-model", "mapped:b:shared-model"],
    )
    assert observed_calls == [
        ("cloud-1::static", ("1", "3"), {}),
        ("cloud-1::item:0", ("7",), {"static_input_0": "shared-model", "remote_input_1": "a"}),
        ("cloud-1::item:1", ("7",), {"static_input_0": "shared-model", "remote_input_1": "b"}),
    ]


def test_modal_cloud_execute_mapped_subgraph_payload_preserves_assigned_lane_id(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Mapped cloud progress should keep the caller-assigned lane id instead of collapsing to `0`."""
    observed_progress_events: list[dict[str, Any]] = []

    def fake_execute_subgraph_prompt(
        payload: dict[str, Any],
        hydrated_inputs: dict[str, Any],
        custom_nodes_root: Any,
        status_callback: Any = None,
        cancellation_event: Any = None,
        interrupt_store: Any = None,
        interrupt_flag_key: Any = None,
    ) -> tuple[str, ...]:
        del custom_nodes_root, cancellation_event, interrupt_store, interrupt_flag_key
        if status_callback is not None:
            status_callback(
                {
                    "event_type": "node_progress",
                    "node_id": "12",
                    "display_node_id": "12",
                    "real_node_id": "12",
                    "value": 3.0,
                    "max": 9.0,
                }
            )
        return (f"mapped:{hydrated_inputs['remote_input_1']}",)

    monkeypatch.setattr(
        modal_cloud_module,
        "_execute_subgraph_prompt",
        fake_execute_subgraph_prompt,
    )

    payload = {
        "payload_kind": "mapped_subgraph",
        "component_id": "cloud-2",
        "prompt_id": "prompt-1",
        "mapped_progress_lane_id": "3",
        "mapped_input": {"proxy_input_name": "remote_input_1", "io_type": "STRING"},
        "static_to_mapped_boundaries": [],
        "mapped_phase": {
            "component_node_ids": ["12", "39"],
            "subgraph_prompt": {
                "39": {
                    "class_type": "RemoteSampler",
                    "inputs": {"latent": ["remote_input_1", 0]},
                },
            },
            "boundary_inputs": [
                {
                    "proxy_input_name": "remote_input_1",
                    "io_type": "STRING",
                    "targets": [{"node_id": "39", "input_name": "latent"}],
                }
            ],
            "boundary_outputs": [
                {
                    "proxy_output_name": "39_text",
                    "node_id": "39",
                    "output_index": 0,
                    "io_type": "STRING",
                    "is_list": False,
                    "mapped_output": True,
                }
            ],
            "execute_node_ids": ["39"],
        },
        "boundary_outputs": [
            {
                "proxy_output_name": "39_text",
                "node_id": "39",
                "output_index": 0,
                "io_type": "STRING",
                "is_list": False,
                "mapped_output": True,
            },
        ],
    }

    outputs = modal_cloud_module._execute_mapped_subgraph_payload(
        payload,
        {"remote_input_1": ["a"]},
        None,
        status_callback=lambda event: observed_progress_events.append(dict(event)),
    )

    assert outputs == (["mapped:a"],)
    assert any(
        event.get("event_type") == "node_progress"
        and event.get("real_node_id") == "12"
        and event.get("lane_id") == "3"
        for event in observed_progress_events
    )
    assert any(
        event.get("event_type") == "node_progress"
        and event.get("clear") is True
        and event.get("lane_id") == "3"
        for event in observed_progress_events
    )


def test_execute_subgraph_locally_round_trips_remote_session_bridge_refs(
    remote_modal_app_module: Any,
    serialization_module: Any,
    session_state_module: Any,
) -> None:
    """Split proxy subgraphs should return durable bridge refs that survive lost sessions."""
    static_payload = {
        "payload_kind": "subgraph",
        "component_id": "1",
        "prompt_id": "prompt-1",
        "component_node_ids": ["1"],
        "subgraph_prompt": {
            "1": {
                "class_type": "SessionValueNode",
                "inputs": {},
            }
        },
        "boundary_inputs": [],
        "boundary_outputs": [
            {
                "proxy_output_name": "static_input_0",
                "node_id": "1",
                "output_index": 0,
                "io_type": "MODEL",
                "is_list": False,
                "session_output": True,
            }
        ],
        "execute_node_ids": ["1"],
        "remote_session": session_state_module.RemoteSessionHandle(
            session_id="session-runtime-1",
            prompt_id="prompt-1",
            owner_component_id="1",
        ).to_payload(),
    }

    static_outputs = serialization_module.deserialize_node_outputs(
        remote_modal_app_module.execute_subgraph_locally(
            static_payload,
            serialization_module.serialize_node_inputs({}),
            node_mapping={"SessionValueNode": _FakeSessionValueNode},
        )
    )
    assert session_state_module.is_remote_session_bridge_ref_payload(static_outputs[0])
    remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(
        session_state_module.RemoteSessionHandle.from_payload(static_payload["remote_session"])
    )

    mapped_payload = {
        "payload_kind": "subgraph",
        "component_id": "1__mapped",
        "prompt_id": "prompt-1",
        "component_node_ids": ["7"],
        "subgraph_prompt": {
            "7": {
                "class_type": "SessionEchoNode",
                "inputs": {"text": ["static_input_0", 0]},
            }
        },
        "boundary_inputs": [
            {
                "proxy_input_name": "static_input_0",
                "io_type": "MODEL",
                "targets": [{"node_id": "7", "input_name": "text"}],
            }
        ],
        "boundary_outputs": [
            {
                "proxy_output_name": "7_text",
                "node_id": "7",
                "output_index": 0,
                "io_type": "STRING",
                "is_list": False,
            }
        ],
        "execute_node_ids": ["7"],
        "remote_session": session_state_module.RemoteSessionHandle(
            session_id="session-runtime-2",
            prompt_id="prompt-1",
            owner_component_id="1",
        ).to_payload(),
        "clear_remote_session": True,
    }

    mapped_outputs = serialization_module.deserialize_node_outputs(
        remote_modal_app_module.execute_subgraph_locally(
            mapped_payload,
            serialization_module.serialize_node_inputs({"static_input_0": static_outputs[0]}),
            node_mapping={
                "SessionEchoNode": _FakeSessionEchoNode,
                "SessionValueNode": _FakeSessionValueNode,
            },
        )
    )

    assert mapped_outputs == ("shared-session-value",)
    replay_outputs = remote_modal_app_module._execute_subgraph_prompt(
        {
            **mapped_payload,
            "remote_session": session_state_module.RemoteSessionHandle(
                session_id="session-runtime-3",
                prompt_id="prompt-1",
                owner_component_id="1",
            ).to_payload(),
            "clear_remote_session": False,
        },
        {"static_input_0": static_outputs[0]},
        {
            "SessionEchoNode": _FakeSessionEchoNode,
            "SessionValueNode": _FakeSessionValueNode,
        },
    )
    assert replay_outputs == ("shared-session-value",)


def test_local_fallback_skips_seed_execution_when_session_outputs_are_already_restored(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """Local fallback seed payloads should no-op when bridge outputs are already in session memory."""
    session_handle = remote_modal_app_module.RemoteSessionHandle(
        session_id="session-restored",
        prompt_id="prompt-1",
        owner_component_id="component-1",
    )
    restored_value = _FakeModelValue("restored-model")
    remote_modal_app_module._REMOTE_SESSION_STORE.put_output(
        session_handle,
        node_id="5",
        output_index=0,
        value=restored_value,
    )
    payload = {
        "payload_kind": "subgraph",
        "component_id": "1__mapped::seed:0",
        "prompt_id": "prompt-1",
        "component_node_ids": ["5"],
        "subgraph_prompt": {
            "5": {
                "class_type": "ShouldNotRun",
                "inputs": {"model": ["static_input_0", 0]},
            }
        },
        "boundary_inputs": [
            {
                "proxy_input_name": "static_input_0",
                "io_type": "MODEL",
                "targets": [{"node_id": "5", "input_name": "model"}],
            }
        ],
        "boundary_outputs": [
            {
                "proxy_output_name": "static_input_0",
                "node_id": "5",
                "output_index": 0,
                "io_type": "MODEL",
                "is_list": False,
                "session_output": True,
            }
        ],
        "execute_node_ids": ["5"],
        "remote_session": session_handle.to_payload(),
    }
    hydrated_inputs = {
        "static_input_0": remote_modal_app_module.RemoteSessionValueRef(
            session_id=session_handle.session_id,
            node_id="5",
            output_index=0,
        ).to_payload()
    }
    monkeypatch.setattr(
        remote_modal_app_module,
        "_invoke_original_node",
        lambda *args, **kwargs: (_ for _ in ()).throw(
            AssertionError("seed short-circuit should skip local node execution")
        ),
    )

    try:
        outputs = remote_modal_app_module._execute_subgraph_prompt(
            payload,
            hydrated_inputs,
            node_mapping={"ShouldNotRun": type("ShouldNotRun", (), {})},
        )
    finally:
        remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(session_handle)

    assert len(outputs) == 1
    assert remote_modal_app_module.is_remote_session_bridge_ref_payload(outputs[0])


def test_local_remote_app_rehydrates_bridge_refs_from_warm_value_cache_without_replay(
    remote_modal_app_module: Any,
) -> None:
    """The local fallback bridge fast path should restore retained values without replaying producers."""
    bridge_key = "RSB_local_cached_bridge"
    target_handle = remote_modal_app_module.RemoteSessionHandle(
        session_id="session-target",
        prompt_id="prompt-1",
        owner_component_id="component-1",
    )
    bridge_ref = remote_modal_app_module.RemoteSessionBridgeRef(
        bridge_key=bridge_key,
        node_id="node-7",
        output_index=0,
        session_id="session-source",
    )
    original_cache = dict(remote_modal_app_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE)
    original_order = list(remote_modal_app_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER)
    try:
        seed_value = _CloneableCacheValue("warm-local-bridge-value")
        remote_modal_app_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE.clear()
        remote_modal_app_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.clear()
        remote_modal_app_module._store_remote_session_bridge_value(bridge_key, seed_value)
        resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats()

        restored_value = remote_modal_app_module._rehydrate_remote_session_bridge_value(
            bridge_ref,
            target_session_handle=target_handle,
            node_mapping=None,
            resolution_stats=resolution_stats,
        )

        stored_value = remote_modal_app_module._REMOTE_SESSION_STORE.get_output(
            remote_modal_app_module.RemoteSessionValueRef(
                session_id=target_handle.session_id,
                node_id="node-7",
                output_index=0,
            )
        )
    finally:
        remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(target_handle)
        remote_modal_app_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE.clear()
        remote_modal_app_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE.update(original_cache)
        remote_modal_app_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.clear()
        remote_modal_app_module._REMOTE_SESSION_BRIDGE_VALUE_CACHE_ORDER.extend(original_order)

    assert isinstance(restored_value, _CloneableCacheValue)
    assert restored_value.value == "warm-local-bridge-value"
    assert restored_value is not seed_value
    assert stored_value is restored_value
    assert resolution_stats.bridge_cache_hits == 1
    assert resolution_stats.bridge_record_lookups == 0
    assert resolution_stats.replay_count == 0
    assert resolution_stats.session_restore_writes == 1


def test_local_remote_app_rehydrates_conditioning_bridge_refs_from_durable_record_without_replay(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """The local fallback should restore durably serialized CONDITIONING bridge values without replay."""
    import torch

    target_handle = remote_modal_app_module.RemoteSessionHandle(
        session_id="session-target",
        prompt_id="prompt-1",
        owner_component_id="component-1",
    )
    bridge_ref = remote_modal_app_module.RemoteSessionBridgeRef(
        bridge_key="RSB_local_conditioning_bridge",
        node_id="node-9",
        output_index=0,
        session_id="session-source",
    )
    conditioning = [
        [
            torch.arange(6, dtype=torch.float32).reshape(1, 2, 3),
            {"pooled_output": torch.arange(4, dtype=torch.float32).reshape(1, 4)},
        ]
    ]
    monkeypatch.setattr(
        remote_modal_app_module._REMOTE_SESSION_BRIDGE_STORE,
        "get_record",
        lambda bridge_key: remote_modal_app_module.RemoteSessionBridgeRecord(
            bridge_key=bridge_key,
            node_id="node-9",
            output_index=0,
            producer_payload={"component_id": "should-not-replay"},
            producer_inputs={},
            serialized_output=remote_modal_app_module.serialize_value(conditioning),
            serialized_output_io_type="CONDITIONING",
        ),
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_execute_subgraph_prompt",
        lambda *args, **kwargs: (_ for _ in ()).throw(
            AssertionError("durable CONDITIONING restore should skip local replay")
        ),
    )
    resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats()

    try:
        restored_value = remote_modal_app_module._rehydrate_remote_session_bridge_value(
            bridge_ref,
            target_session_handle=target_handle,
            node_mapping=None,
            resolution_stats=resolution_stats,
        )
        stored_value = remote_modal_app_module._REMOTE_SESSION_STORE.get_output(
            remote_modal_app_module.RemoteSessionValueRef(
                session_id=target_handle.session_id,
                node_id="node-9",
                output_index=0,
            )
        )
    finally:
        remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(target_handle)

    assert torch.equal(restored_value[0][0], conditioning[0][0])
    assert torch.equal(restored_value[0][1]["pooled_output"], conditioning[0][1]["pooled_output"])
    assert torch.equal(stored_value[0][0], conditioning[0][0])
    assert resolution_stats.bridge_cache_hits == 0
    assert resolution_stats.durable_bridge_hits == 1
    assert resolution_stats.bridge_record_lookups == 1
    assert resolution_stats.replay_count == 0
    assert resolution_stats.session_restore_writes == 1


def test_remote_session_input_resolution_handles_nested_conditioning_bridge_refs(
    remote_modal_app_module: Any,
) -> None:
    """Boundary input resolution should resolve bridge refs nested inside list outputs."""
    source_handle = remote_modal_app_module.RemoteSessionHandle(
        session_id="session-source",
        prompt_id="prompt-1",
        owner_component_id="component-1",
    )
    first_ref = remote_modal_app_module._REMOTE_SESSION_STORE.put_bridge_output(
        source_handle,
        bridge_key="RSB_a",
        node_id="508:item:0",
        output_index=0,
        value=[["conditioning-a", {"pooled_output": "pool-a"}]],
    )
    second_ref = remote_modal_app_module._REMOTE_SESSION_STORE.put_bridge_output(
        source_handle,
        bridge_key="RSB_b",
        node_id="508:item:1",
        output_index=0,
        value=[["conditioning-b", {"pooled_output": "pool-b"}]],
    )
    resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats()
    try:
        resolved_inputs = remote_modal_app_module._resolve_remote_session_inputs(
            {
                "positive": [
                    remote_modal_app_module.RemoteSessionBridgeRef(
                        bridge_key="RSB_a",
                        node_id=first_ref.node_id,
                        output_index=first_ref.output_index,
                        session_id=first_ref.session_id,
                    ).to_payload(),
                    remote_modal_app_module.RemoteSessionBridgeRef(
                        bridge_key="RSB_b",
                        node_id=second_ref.node_id,
                        output_index=second_ref.output_index,
                        session_id=second_ref.session_id,
                    ).to_payload(),
                ]
            },
            component_id="sampler-component",
            resolution_stats=resolution_stats,
        )
    finally:
        remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(source_handle)

    assert resolved_inputs["positive"] == [
        [["conditioning-a", {"pooled_output": "pool-a"}]],
        [["conditioning-b", {"pooled_output": "pool-b"}]],
    ]
    assert resolution_stats.input_ref_count == 2


def test_local_remote_app_rehydrates_sampler_latent_bridge_refs_from_durable_record_without_replay(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """The local fallback should restore serialized LATENT bridge values without sampler replay."""
    import torch

    target_handle = remote_modal_app_module.RemoteSessionHandle(
        session_id="session-target",
        prompt_id="prompt-1",
        owner_component_id="component-1",
    )
    bridge_ref = remote_modal_app_module.RemoteSessionBridgeRef(
        bridge_key="RSB_local_sampler_latent_bridge",
        node_id="sampler-1",
        output_index=0,
        session_id="session-source",
    )
    latent = {"samples": torch.arange(8, dtype=torch.float32).reshape(1, 2, 2, 2)}
    record = remote_modal_app_module._build_remote_session_bridge_record(
        payload={
            "component_id": "sampler-component",
            "execute_node_ids": ["sampler-1"],
            "subgraph_prompt": {
                "sampler-1": {
                    "class_type": "KSampler",
                    "inputs": {},
                }
            },
        },
        hydrated_inputs={},
        node_id="sampler-1",
        output_index=0,
        io_type="LATENT",
        output_value=latent,
    )
    assert record.serialized_output is not None
    assert record.serialized_output_io_type == "LATENT"
    monkeypatch.setattr(
        remote_modal_app_module._REMOTE_SESSION_BRIDGE_STORE,
        "get_record",
        lambda bridge_key: record,
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_execute_subgraph_prompt",
        lambda *args, **kwargs: (_ for _ in ()).throw(
            AssertionError("durable LATENT restore should skip local sampler replay")
        ),
    )
    resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats()

    try:
        restored_value = remote_modal_app_module._rehydrate_remote_session_bridge_value(
            bridge_ref,
            target_session_handle=target_handle,
            node_mapping=None,
            resolution_stats=resolution_stats,
        )
        stored_value = remote_modal_app_module._REMOTE_SESSION_STORE.get_output(
            remote_modal_app_module.RemoteSessionValueRef(
                session_id=target_handle.session_id,
                node_id="sampler-1",
                output_index=0,
            )
        )
    finally:
        remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(target_handle)

    assert torch.equal(restored_value["samples"], latent["samples"])
    assert torch.equal(stored_value["samples"], latent["samples"])
    assert resolution_stats.bridge_cache_hits == 0
    assert resolution_stats.durable_bridge_hits == 1
    assert resolution_stats.bridge_record_lookups == 1
    assert resolution_stats.replay_count == 0
    assert resolution_stats.session_restore_writes == 1


@pytest.mark.parametrize(
    ("io_type", "node_id", "class_type", "node_inputs"),
    [
        ("NOISE", "15", "RandomNoise", {"noise_seed": 42}),
        ("SAMPLER", "17", "KSamplerSelect", {"sampler_name": "euler"}),
    ],
)
def test_local_remote_app_rehydrates_literal_sampling_strategy_bridges_without_replay(
    remote_modal_app_module: Any,
    monkeypatch: Any,
    io_type: str,
    node_id: str,
    class_type: str,
    node_inputs: dict[str, Any],
) -> None:
    """The local fallback should rebuild literal sampling strategies without replay."""
    target_handle = remote_modal_app_module.RemoteSessionHandle(
        session_id="session-target",
        prompt_id="prompt-1",
        owner_component_id="component-1",
    )
    bridge_ref = remote_modal_app_module.RemoteSessionBridgeRef(
        bridge_key=f"RSB_local_{io_type.lower()}_bridge",
        node_id=node_id,
        output_index=0,
        session_id="session-source",
    )
    record = remote_modal_app_module._build_remote_session_bridge_record(
        payload={
            "component_id": "image-preview-component",
            "execute_node_ids": ["251"],
            "subgraph_prompt": {
                "14": {
                    "class_type": "SamplerCustomAdvanced",
                    "inputs": {"noise": ["15", 0], "sampler": ["17", 0]},
                },
                "15": {
                    "class_type": "RandomNoise",
                    "inputs": {"noise_seed": 42},
                },
                "17": {
                    "class_type": "KSamplerSelect",
                    "inputs": {"sampler_name": "euler"},
                },
                "250": {
                    "class_type": "VAEDecode",
                    "inputs": {"samples": ["14", 0]},
                },
                "251": {
                    "class_type": "ImageFromBatch",
                    "inputs": {"image": ["250", 0]},
                },
            },
        },
        hydrated_inputs={},
        node_id=node_id,
        output_index=0,
        io_type=io_type,
        output_value=object(),
    )
    assert (
        record.recovery_kind
        is remote_modal_app_module.RemoteSessionBridgeRecoveryKind.SINGLE_NODE_PLAN
    )
    assert record.rehydration_plan == {
        "kind": "single_node_output",
        "node_data": {"class_type": class_type},
        "node_inputs": node_inputs,
    }
    execute_calls: list[tuple[dict[str, Any], dict[str, Any]]] = []
    monkeypatch.setattr(
        remote_modal_app_module._REMOTE_SESSION_BRIDGE_STORE,
        "get_record",
        lambda bridge_key: record,
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_execute_subgraph_prompt",
        lambda *args, **kwargs: (_ for _ in ()).throw(
            AssertionError("literal strategy bridge should not replay its producer component")
        ),
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_execute_node_locally_raw",
        lambda node_data, kwargs_payload, **kwargs: (
            execute_calls.append((dict(node_data), dict(kwargs_payload))),
            (f"restored-{io_type.lower()}",),
        )[1],
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_store_remote_session_bridge_value",
        lambda bridge_key, value: None,
    )
    resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats()

    try:
        restored_value = remote_modal_app_module._rehydrate_remote_session_bridge_value(
            bridge_ref,
            target_session_handle=target_handle,
            node_mapping=None,
            resolution_stats=resolution_stats,
        )
    finally:
        remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(target_handle)

    assert restored_value == f"restored-{io_type.lower()}"
    assert execute_calls == [({"class_type": class_type}, node_inputs)]
    assert resolution_stats.durable_bridge_hits == 1
    assert resolution_stats.replay_count == 0


def test_local_remote_app_rehydrates_model_bridge_refs_from_durable_plan_without_replay(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """The local fallback should rebuild one MODEL bridge output from a durable plan without replay."""
    target_handle = remote_modal_app_module.RemoteSessionHandle(
        session_id="session-target",
        prompt_id="prompt-1",
        owner_component_id="component-1",
    )
    bridge_ref = remote_modal_app_module.RemoteSessionBridgeRef(
        bridge_key="RSB_local_model_bridge",
        node_id="node-5",
        output_index=0,
        session_id="session-source",
    )
    record = remote_modal_app_module._build_remote_session_bridge_record(
        payload={
            "component_id": "component-seed",
            "subgraph_prompt": {
                "node-5": {
                    "class_type": "CheckpointLoaderSimple",
                    "inputs": {"ckpt_name": "model.safetensors"},
                }
            },
        },
        hydrated_inputs={},
        node_id="node-5",
        output_index=0,
        io_type="MODEL",
        output_value=_FakeModelValue("seed-model"),
    )
    monkeypatch.setattr(remote_modal_app_module._REMOTE_SESSION_BRIDGE_STORE, "get_record", lambda bridge_key: record)
    monkeypatch.setattr(
        remote_modal_app_module,
        "_execute_subgraph_prompt",
        lambda *args, **kwargs: (_ for _ in ()).throw(
            AssertionError("durable MODEL rehydration should skip local replay")
        ),
    )
    resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats()

    try:
        restored_value = remote_modal_app_module._rehydrate_remote_session_bridge_value(
            bridge_ref,
            target_session_handle=target_handle,
            node_mapping={"CheckpointLoaderSimple": _FakeModelLoaderNode},
            resolution_stats=resolution_stats,
        )
        stored_value = remote_modal_app_module._REMOTE_SESSION_STORE.get_output(
            remote_modal_app_module.RemoteSessionValueRef(
                session_id=target_handle.session_id,
                node_id="node-5",
                output_index=0,
            )
        )
    finally:
        remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(target_handle)

    assert isinstance(restored_value, _FakeModelValue)
    assert restored_value.value == "model::model.safetensors"
    assert stored_value is restored_value
    assert resolution_stats.bridge_cache_hits == 0
    assert resolution_stats.durable_bridge_hits == 1
    assert resolution_stats.bridge_record_lookups == 1
    assert resolution_stats.replay_count == 0
    assert resolution_stats.session_restore_writes == 1


def test_local_remote_app_rehydrates_clip_bridge_refs_from_durable_plan_without_replay(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """The local fallback should rebuild one CLIP bridge output from a durable plan."""
    target_handle = remote_modal_app_module.RemoteSessionHandle(
        session_id="session-target",
        prompt_id="prompt-1",
        owner_component_id="component-1",
    )
    bridge_ref = remote_modal_app_module.RemoteSessionBridgeRef(
        bridge_key="RSB_local_clip_bridge",
        node_id="clip-1",
        output_index=0,
        session_id="session-source",
    )
    record = remote_modal_app_module._build_remote_session_bridge_record(
        payload={
            "component_id": "sampler-component",
            "execute_node_ids": ["sampler-1"],
            "subgraph_prompt": {
                "clip-1": {
                    "class_type": "CLIPLoader",
                    "inputs": {
                        "clip_name": "text.safetensors",
                        "type": "flux",
                        "device": "default",
                    },
                },
                "sampler-1": {
                    "class_type": "KSampler",
                    "inputs": {"clip": ["clip-1", 0]},
                },
            },
        },
        hydrated_inputs={},
        node_id="clip-1",
        output_index=0,
        io_type="CLIP",
        output_value="seed-clip",
    )
    assert record.rehydration_plan is not None
    assert record.rehydration_plan_io_type == "CLIP"
    monkeypatch.setattr(remote_modal_app_module._REMOTE_SESSION_BRIDGE_STORE, "get_record", lambda bridge_key: record)
    monkeypatch.setattr(
        remote_modal_app_module,
        "_execute_subgraph_prompt",
        lambda *args, **kwargs: (_ for _ in ()).throw(
            AssertionError("durable CLIP rehydration should skip local replay")
        ),
    )
    resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats()

    try:
        restored_value = remote_modal_app_module._rehydrate_remote_session_bridge_value(
            bridge_ref,
            target_session_handle=target_handle,
            node_mapping={"CLIPLoader": _FakeClipLoaderNode},
            resolution_stats=resolution_stats,
        )
    finally:
        remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(target_handle)

    assert restored_value == "clip::text.safetensors:flux:default"
    assert resolution_stats.durable_bridge_hits == 1
    assert resolution_stats.replay_count == 0


def test_local_remote_app_rehydrates_linked_model_bridge_with_non_sampler_subgraph_plan(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """The local fallback should rehydrate linked MODEL bridges without sampler replay."""
    target_handle = remote_modal_app_module.RemoteSessionHandle(
        session_id="session-target",
        prompt_id="prompt-1",
        owner_component_id="component-1",
    )
    bridge_ref = remote_modal_app_module.RemoteSessionBridgeRef(
        bridge_key="RSB_local_linked_model_bridge",
        node_id="lora-1",
        output_index=0,
        session_id="session-source",
    )
    record = remote_modal_app_module._build_remote_session_bridge_record(
        payload={
            "component_id": "sampler-component",
            "execute_node_ids": ["sampler-1"],
            "subgraph_prompt": {
                "loader-1": {
                    "class_type": "UNETLoader",
                    "inputs": {"unet_name": "base.safetensors"},
                },
                "lora-1": {
                    "class_type": "LoraLoaderModelOnly",
                    "inputs": {"model": ["loader-1", 0], "lora_name": "style.safetensors"},
                },
                "sampler-1": {
                    "class_type": "KSampler",
                    "inputs": {"model": ["lora-1", 0]},
                },
            },
        },
        hydrated_inputs={},
        node_id="lora-1",
        output_index=0,
        io_type="MODEL",
        output_value=_FakeModelValue("lora-model"),
    )
    assert record.rehydration_plan is not None
    assert record.rehydration_plan["kind"] == "subgraph_output"
    monkeypatch.setattr(remote_modal_app_module._REMOTE_SESSION_BRIDGE_STORE, "get_record", lambda bridge_key: record)
    observed_payloads: list[dict[str, Any]] = []

    def fake_execute_subgraph_prompt(
        payload: dict[str, Any],
        hydrated_inputs: dict[str, Any],
        _node_mapping: dict[str, type[Any]] | None = None,
    ) -> tuple[_FakeModelValue]:
        """Assert the durable plan avoids rerunning the sampler component."""
        del hydrated_inputs, _node_mapping
        observed_payloads.append(copy.deepcopy(payload))
        assert payload["execute_node_ids"] == ["lora-1"]
        assert payload["remote_session"]["session_id"] == target_handle.session_id
        assert sorted(payload["subgraph_prompt"]) == ["loader-1", "lora-1"]
        assert "sampler-1" not in payload["subgraph_prompt"]
        return (_FakeModelValue("restored-lora-model"),)

    monkeypatch.setattr(remote_modal_app_module, "_execute_subgraph_prompt", fake_execute_subgraph_prompt)
    resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats()

    try:
        restored_value = remote_modal_app_module._rehydrate_remote_session_bridge_value(
            bridge_ref,
            target_session_handle=target_handle,
            node_mapping=None,
            resolution_stats=resolution_stats,
        )
    finally:
        remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(target_handle)

    assert isinstance(restored_value, _FakeModelValue)
    assert restored_value.value == "restored-lora-model"
    assert len(observed_payloads) == 1
    assert resolution_stats.durable_bridge_hits == 1
    assert resolution_stats.replay_count == 0


def test_local_remote_app_refuses_sampler_ancestor_bridge_replay(
    remote_modal_app_module: Any,
    monkeypatch: Any,
) -> None:
    """The local fallback should reject replay when a terminal depends on sampling."""
    target_handle = remote_modal_app_module.RemoteSessionHandle(
        session_id="session-target",
        prompt_id="prompt-1",
        owner_component_id="component-1",
    )
    bridge_ref = remote_modal_app_module.RemoteSessionBridgeRef(
        bridge_key="RSB_local_sampler_bridge",
        node_id="unsupported-bridge",
        output_index=0,
        session_id="session-source",
    )
    monkeypatch.setattr(
        remote_modal_app_module._REMOTE_SESSION_BRIDGE_STORE,
        "get_record",
        lambda bridge_key: remote_modal_app_module.RemoteSessionBridgeRecord(
            bridge_key=bridge_key,
            node_id="unsupported-bridge",
            output_index=0,
            producer_payload={
                "component_id": "sampler-component",
                "execute_node_ids": ["251"],
                "subgraph_prompt": {
                    "14": {"class_type": "KSampler", "inputs": {}},
                    "250": {
                        "class_type": "VAEDecode",
                        "inputs": {"samples": ["14", 0]},
                    },
                    "251": {
                        "class_type": "ImageFromBatch",
                        "inputs": {"image": ["250", 0]},
                    },
                },
            },
            producer_inputs={},
        ),
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_execute_subgraph_prompt",
        lambda *args, **kwargs: (_ for _ in ()).throw(
            AssertionError("sampler bridge replay should be blocked")
        ),
    )
    resolution_stats = remote_modal_app_module._RemoteSessionBridgeResolutionStats()

    try:
        with pytest.raises(remote_modal_app_module.RemoteSessionStateError, match="rerun a sampler"):
            remote_modal_app_module._rehydrate_remote_session_bridge_value(
                bridge_ref,
                target_session_handle=target_handle,
                node_mapping=None,
                resolution_stats=resolution_stats,
            )
    finally:
        remote_modal_app_module._REMOTE_SESSION_STORE.clear_session(target_handle)

    assert resolution_stats.bridge_record_lookups == 1
    assert resolution_stats.replay_count == 0


def test_local_phase_payload_builder_preserves_remote_session_and_snapshot_profile(
    remote_modal_app_module: Any,
) -> None:
    """Explicit local mapped phase payloads should keep remote_session context and snapshot profile."""
    payload = {
        "prompt_id": "prompt-1",
        "extra_data": {"client_id": "c-1"},
        "snapshot_profile_key": "loader-profile:abc",
        "remote_session": {
            "__comfy_modal_remote_session_handle__": True,
            "session_id": "session-1",
            "prompt_id": "prompt-1",
            "owner_component_id": "component-1",
        },
        "clear_remote_session": True,
        "static_phase": {
            "component_node_ids": ["1"],
            "subgraph_prompt": {},
            "boundary_inputs": [],
            "boundary_outputs": [],
            "execute_node_ids": ["1"],
        },
    }

    phase_payload = remote_modal_app_module._build_phase_subgraph_payload(
        payload,
        "static_phase",
        "component-1::static",
        suppress_status_stream=True,
    )

    assert phase_payload["remote_session"]["session_id"] == "session-1"
    assert phase_payload["clear_remote_session"] is True
    assert phase_payload["snapshot_profile_key"] == "loader-profile:abc"


def test_cloud_phase_payload_builder_preserves_remote_session_and_snapshot_profile(
    modal_cloud_module: Any,
) -> None:
    """Explicit cloud mapped phase payloads should keep remote_session context and snapshot profile."""
    payload = {
        "prompt_id": "prompt-1",
        "extra_data": {"client_id": "c-1"},
        "snapshot_profile_key": "loader-profile:abc",
        "remote_session": {
            "__comfy_modal_remote_session_handle__": True,
            "session_id": "session-1",
            "prompt_id": "prompt-1",
            "owner_component_id": "component-1",
        },
        "clear_remote_session": True,
        "mapped_phase": {
            "component_node_ids": ["7"],
            "subgraph_prompt": {},
            "boundary_inputs": [],
            "boundary_outputs": [],
            "execute_node_ids": ["7"],
        },
    }

    phase_payload = modal_cloud_module._build_phase_subgraph_payload(
        payload,
        "mapped_phase",
        "component-1::mapped",
    )

    assert phase_payload["remote_session"]["session_id"] == "session-1"
    assert phase_payload["clear_remote_session"] is True
    assert phase_payload["snapshot_profile_key"] == "loader-profile:abc"


def test_split_hybrid_proxies_allow_local_downstream_work_before_mapped_completion(
    api_intercept_module: Any,
    modal_executor_module: Any,
    session_state_module: Any,
    settings_module: Any,
    sync_engine_module: Any,
    tmp_path: Path,
) -> None:
    """A local consumer of the static proxy should be able to run while the mapped proxy is still in flight."""
    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=tmp_path / "storage",
        remote_storage_root="/storage",
        custom_nodes_archive_name="custom_nodes_bundle.zip",
        comfyui_root=None,
        custom_nodes_dir=tmp_path / "custom_nodes",
    )
    settings.custom_nodes_dir.mkdir()
    sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings)
    fake_nodes_module = type(
        "FakeNodesModule",
        (),
        {
            "NODE_CLASS_MAPPINGS": {
                "RemoteModel": _FakeRewriteRemoteModelNode,
                "RemoteSampler": _FakeRewriteRemoteSamplerNode,
                "LatentSource": _FakeRewriteLatentSourceNode,
                "ModalMapInput": _FakeRewriteModalMapInputNode,
                "LocalSink": _FakeRewriteLocalSinkNode,
            },
            "NODE_DISPLAY_NAME_MAPPINGS": {},
        },
    )()
    workflow = {
        "nodes": [
            {"id": 1, "properties": {"is_modal_remote": True}},
            {"id": 2, "properties": {"is_modal_remote": False}},
            {"id": 3, "properties": {"is_modal_remote": True}},
            {"id": 4, "properties": {"is_modal_remote": False}},
            {"id": 5, "properties": {"is_modal_remote": False}},
            {"id": 6, "properties": {"is_modal_remote": True}},
            {"id": 7, "properties": {"is_modal_remote": True}},
            {"id": 8, "properties": {"is_modal_remote": False}},
        ]
    }
    prompt = {
        "1": {
            "class_type": "RemoteModel",
            "inputs": {},
            "_meta": {"title": "Shared Model"},
        },
        "2": {
            "class_type": "LatentSource",
            "inputs": {},
            "_meta": {"title": "Single Latent"},
        },
        "3": {
            "class_type": "RemoteSampler",
            "inputs": {"model": ["1", 0], "latent": ["2", 0]},
            "_meta": {"title": "Unmapped Sampler"},
        },
        "4": {
            "class_type": "LocalSink",
            "inputs": {"image": ["3", 0]},
            "_meta": {"title": "Local Sink 1"},
        },
        "5": {
            "class_type": "LatentSource",
            "inputs": {},
            "_meta": {"title": "Batch Latent Source"},
        },
        "6": {
            "class_type": "ModalMapInput",
            "inputs": {"value": ["5", 0]},
            "_meta": {"title": "Map Input"},
        },
        "7": {
            "class_type": "RemoteSampler",
            "inputs": {"model": ["1", 0], "latent": ["6", 0]},
            "_meta": {"title": "Mapped Sampler"},
        },
        "8": {
            "class_type": "LocalSink",
            "inputs": {"image": ["7", 0]},
            "_meta": {"title": "Local Sink 2"},
        },
    }

    rewritten_prompt, _summary = api_intercept_module.rewrite_prompt_for_modal(
        prompt=prompt,
        workflow=workflow,
        sync_engine=sync_engine,
        settings=settings,
        nodes_module=fake_nodes_module,
    )
    static_proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[rewritten_prompt["1"]["class_type"]]
    mapped_proxy_class = fake_nodes_module.NODE_CLASS_MAPPINGS[rewritten_prompt["1__mapped"]["class_type"]]
    static_payload = rewritten_prompt["1"]["inputs"]["original_node_data"]
    mapped_payload = rewritten_prompt["1__mapped"]["inputs"]["original_node_data"]

    assert rewritten_prompt["4"]["inputs"]["image"] == ["1", 0]
    assert rewritten_prompt["8"]["inputs"]["image"] == ["1__mapped", 0]
    assert mapped_payload["static_phase"]["execute_node_ids"] == ["1"]
    assert mapped_payload["static_to_mapped_boundaries"] == [
        {
            "proxy_name": "static_input_0",
            "node_id": "1",
            "output_index": 0,
            "io_type": "MODEL",
            "is_list": False,
            "targets": [{"node_id": "7", "input_name": "model"}],
        }
    ]

    observed_order: list[str] = []
    mapped_started = asyncio.Event()
    release_mapped = asyncio.Event()

    class FakeClient:
        """Fake async remote client that blocks the mapped proxy until released."""

        async def execute_payload_async(
            self,
            payload: dict[str, Any],
            kwargs: dict[str, Any],
        ) -> tuple[Any, ...]:
            """Return deterministic outputs for the split static and mapped proxies."""
            if str(payload.get("component_id")) == "1":
                observed_order.append("static_proxy_finish")
                return (
                    "static-latent",
                    session_state_module.RemoteSessionBridgeRef(
                        bridge_key="RSB_static_model",
                        node_id="1",
                        output_index=0,
                        session_id=str(payload["remote_session"]["session_id"]),
                    ).to_payload(),
                )

            if str(payload.get("component_id")) == "1__mapped":
                observed_order.append("mapped_proxy_start")
                mapped_started.set()
                await release_mapped.wait()
                observed_order.append("mapped_proxy_finish")
                return ("mapped-latent",)

            raise AssertionError(f"Unexpected proxy payload: {payload!r}")

    async def run_scenario() -> tuple[Any, ...]:
        """Run the split static and mapped proxies with a local consumer in between."""
        static_result = await static_proxy_class.execute(
            original_node_data=static_payload,
            unique_id="1",
            remote_input_0="single-latent",
        )
        mapped_task = asyncio.create_task(
            mapped_proxy_class.execute(
                original_node_data=mapped_payload,
                unique_id="1__mapped",
                remote_input_1="batched-latent",
                static_input_0=static_result.result[1],
            )
        )
        await mapped_started.wait()
        observed_order.append(f"local_sink:{static_result.result[0]}")
        assert not mapped_task.done()
        release_mapped.set()
        mapped_result = await mapped_task
        return static_result.result, mapped_result.result

    modal_executor_module.set_remote_executor_client_factory(lambda: FakeClient())
    try:
        static_outputs, mapped_outputs = asyncio.run(run_scenario())
    finally:
        modal_executor_module.set_remote_executor_client_factory(None)

    assert static_outputs[0] == "static-latent"
    assert session_state_module.is_remote_session_bridge_ref_payload(static_outputs[1])
    assert static_outputs[2] is True
    assert mapped_outputs == ("mapped-latent", True)
    assert observed_order == [
        "static_proxy_finish",
        "mapped_proxy_start",
        "local_sink:static-latent",
        "mapped_proxy_finish",
    ]


def test_mixed_remote_and_preview_fanout_uses_bridge_and_local_materializer(
    api_intercept_module: Any,
    modal_executor_module: Any,
    settings_module: Any,
    sync_engine_module: Any,
    tmp_path: Path,
) -> None:
    """A remote continuation should use a bridge while its preview materializes locally."""
    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=tmp_path / "storage",
        remote_storage_root="/storage",
        custom_nodes_archive_name="custom_nodes_bundle.zip",
        comfyui_root=None,
        custom_nodes_dir=tmp_path / "custom_nodes",
    )
    settings.custom_nodes_dir.mkdir()
    sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings)
    fake_nodes_module = types.SimpleNamespace(
        NODE_CLASS_MAPPINGS={
            "RemoteImage": _FakeRewriteRemoteImageNode,
            "RemoteText": _FakeRewriteRemoteTextNode,
            "PreviewImage": _FakeRewritePreviewImageNode,
            "LocalSink": _FakeRewriteLocalSinkNode,
        },
        NODE_DISPLAY_NAME_MAPPINGS={},
    )
    workflow = {
        "nodes": [
            {"id": 1, "properties": {"is_modal_remote": True}},
            {"id": 2, "properties": {"is_modal_remote": True}},
            {"id": 3, "properties": {"is_modal_remote": False}},
            {"id": 4, "properties": {"is_modal_remote": False}},
        ]
    }
    prompt = {
        "1": {
            "class_type": "RemoteImage",
            "inputs": {},
            "_meta": {"title": "Generated Image"},
        },
        "2": {
            "class_type": "RemoteText",
            "inputs": {"image": ["1", 0]},
            "_meta": {"title": "Remote LLM"},
        },
        "3": {
            "class_type": "PreviewImage",
            "inputs": {"images": ["1", 0]},
            "_meta": {"title": "Local Preview"},
        },
        "4": {
            "class_type": "LocalSink",
            "inputs": {"text": ["2", 0]},
            "_meta": {"title": "Response"},
        },
    }

    rewritten_prompt, summary = api_intercept_module.rewrite_prompt_for_modal(
        prompt=prompt,
        workflow=workflow,
        sync_engine=sync_engine,
        settings=settings,
        nodes_module=fake_nodes_module,
        extra_data={"prompt_id": "prompt-mixed-fanout"},
    )

    assert summary.remote_component_ids == ["1", "2"]
    assert summary.component_execution_stages == [["1"], ["2"]]
    producer_payload = modal_executor_module._rehydrate_proxy_payload(
        rewritten_prompt["1"]["inputs"]["original_node_data"],
        unique_id="1",
    )
    consumer_payload = modal_executor_module._rehydrate_proxy_payload(
        rewritten_prompt["2"]["inputs"]["original_node_data"],
        unique_id="2",
    )
    assert producer_payload["component_node_ids"] == ["1"]
    assert producer_payload["boundary_outputs"][0]["session_output"] is True
    assert consumer_payload["component_node_ids"] == ["2"]
    assert rewritten_prompt["2"]["inputs"]["remote_input_0"] == ["1", 0]

    materializer_node_ids = [
        node_id
        for node_id, prompt_node in rewritten_prompt.items()
        if prompt_node["class_type"]
        == api_intercept_module.MODAL_LOCAL_BRIDGE_MATERIALIZER_NODE_ID
    ]
    assert len(materializer_node_ids) == 1
    materializer_node_id = materializer_node_ids[0]
    assert len(summary.parallel_local_branch_node_ids) == 1
    dispatch_gate_node_id = summary.parallel_local_branch_node_ids[0]
    assert rewritten_prompt[materializer_node_id]["inputs"]["bridge_ref"] == [
        dispatch_gate_node_id,
        0,
    ]
    assert rewritten_prompt[dispatch_gate_node_id]["inputs"]["value"] == ["1", 0]
    assert rewritten_prompt["3"]["inputs"]["images"] == [materializer_node_id, 0]


def test_mapped_component_with_local_reentry_rewrites_to_ordered_acyclic_proxies(
    api_intercept_module: Any,
    modal_executor_module: Any,
    settings_module: Any,
    sync_engine_module: Any,
    tmp_path: Path,
) -> None:
    """A mapped input should not collapse local feedback back into the same remote proxy."""
    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=tmp_path / "storage",
        remote_storage_root="/storage",
        custom_nodes_archive_name="custom_nodes_bundle.zip",
        comfyui_root=None,
        custom_nodes_dir=tmp_path / "custom_nodes",
    )
    settings.custom_nodes_dir.mkdir()
    sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings)
    fake_nodes_module = type(
        "FakeNodesModule",
        (),
        {
            "NODE_CLASS_MAPPINGS": {
                "RemoteModel": _FakeRewriteRemoteModelNode,
                "RemoteSampler": _FakeRewriteRemoteSamplerNode,
                "LatentSource": _FakeRewriteLatentSourceNode,
                "ModalMapInput": _FakeRewriteModalMapInputNode,
                "LocalFeedback": _FakeRewriteLocalFeedbackNode,
                "LocalSink": _FakeRewriteLocalSinkNode,
            },
            "NODE_DISPLAY_NAME_MAPPINGS": {},
        },
    )()
    workflow = {
        "nodes": [
            {"id": 1, "properties": {"is_modal_remote": True}},
            {"id": 2, "properties": {"is_modal_remote": False}},
            {"id": 3, "properties": {"is_modal_remote": True}},
            {"id": 4, "properties": {"is_modal_remote": False}},
            {"id": 5, "properties": {"is_modal_remote": False}},
            {"id": 6, "properties": {"is_modal_remote": False}},
            {"id": 7, "properties": {"is_modal_remote": True}},
            {"id": 8, "properties": {"is_modal_remote": False}},
            {"id": 9, "properties": {"is_modal_remote": False}},
        ]
    }
    prompt = {
        "1": {
            "class_type": "RemoteModel",
            "inputs": {},
            "_meta": {"title": "Shared Model"},
        },
        "2": {
            "class_type": "LatentSource",
            "inputs": {},
            "_meta": {"title": "Single Latent"},
        },
        "3": {
            "class_type": "RemoteSampler",
            "inputs": {"model": ["1", 0], "latent": ["2", 0]},
            "_meta": {"title": "First Sampler"},
        },
        "4": {
            "class_type": "LocalFeedback",
            "inputs": {"image": ["3", 0]},
            "_meta": {"title": "Local Feedback"},
        },
        "5": {
            "class_type": "LatentSource",
            "inputs": {},
            "_meta": {"title": "Batch Values"},
        },
        "6": {
            "class_type": "ModalMapInput",
            "inputs": {"value": ["5", 0]},
            "_meta": {"title": "Map Input"},
        },
        "7": {
            "class_type": "RemoteSampler",
            "inputs": {"model": ["1", 0], "latent": ["6", 0], "prompt": ["4", 0]},
            "_meta": {"title": "Mapped Sampler"},
        },
        "8": {
            "class_type": "LocalSink",
            "inputs": {"image": ["7", 0]},
            "_meta": {"title": "Local Sink"},
        },
        "9": {
            "class_type": "LocalSink",
            "inputs": {"image": ["3", 0]},
            "_meta": {"title": "Independent Local Preview"},
        },
    }

    rewritten_prompt, summary = api_intercept_module.rewrite_prompt_for_modal(
        prompt=prompt,
        workflow=workflow,
        sync_engine=sync_engine,
        settings=settings,
        nodes_module=fake_nodes_module,
    )

    assert api_intercept_module._find_prompt_dependency_cycles(rewritten_prompt) == []
    assert summary.mapped_component_ids == ["7"]
    assert summary.component_execution_stages == [["3"], ["7"]]
    assert rewritten_prompt["4"]["inputs"]["image"] == ["3", 0]
    assert rewritten_prompt["7"]["inputs"]["remote_input_2"] == ["4", 0]
    assert len(summary.parallel_local_branch_node_ids) == 1
    passthrough_node_id = summary.parallel_local_branch_node_ids[0]
    passthrough_node = rewritten_prompt[passthrough_node_id]
    assert passthrough_node["class_type"] == (
        api_intercept_module.MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID
    )
    assert passthrough_node["inputs"]["value"] == ["3", 0]
    dispatch_context = passthrough_node["inputs"]["dispatch_context"]
    assert dispatch_context["component_ids"] == ["7"]
    assert rewritten_prompt["9"]["inputs"]["image"] == [passthrough_node_id, 0]
    first_payload = rewritten_prompt["3"]["inputs"]["original_node_data"]
    mapped_payload = rewritten_prompt["7"]["inputs"]["original_node_data"]
    assert first_payload["component_node_ids"] == ["1", "3"]
    assert mapped_payload["component_node_ids"] == ["7"]
    assert first_payload["remote_local_gap_pool"] is True
    assert first_payload["keepalive_after_remote_component"] is True
    assert "stop_local_gap_keepalive_before_remote_component" not in first_payload
    assert mapped_payload["remote_local_gap_pool"] is True
    assert "keepalive_after_remote_component" not in mapped_payload
    assert mapped_payload["stop_local_gap_keepalive_before_remote_component"] is True
    assert mapped_payload["parallel_local_dispatch_group_id"] == (
        dispatch_context["dispatch_group_id"]
    )
    assert mapped_payload["signal_parallel_local_dispatch"] is True
    rehydrated_mapped_payload = modal_executor_module._rehydrate_proxy_payload(
        mapped_payload,
        unique_id="7",
    )
    assert rehydrated_mapped_payload["signal_parallel_local_dispatch"] is True
    assert {
        boundary_input["proxy_input_name"]
        for boundary_input in mapped_payload["boundary_inputs"]
    } == {"phase_bridge_0", "remote_input_1", "remote_input_2"}


def test_invoke_implicitly_mapped_subgraph_async_zips_batched_boundary_inputs(
    remote_modal_app_module: Any,
    serialization_module: Any,
    monkeypatch: Any,
) -> None:
    """Ordinary remote subgraphs should fan out when multiple boundary inputs arrive batched."""
    observed_inputs: list[dict[str, Any]] = []

    async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
        assert payload["payload_kind"] == "subgraph"
        assert payload["suppress_status_stream"] is True
        hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
        observed_inputs.append(hydrated_inputs)
        return serialization_module.serialize_node_outputs(
            (f"{hydrated_inputs['remote_input_0']}:{hydrated_inputs['remote_input_1']}",)
        )

    monkeypatch.setattr(
        remote_modal_app_module,
        "invoke_remote_engine_async",
        fake_invoke_remote_engine_async,
    )

    payload = {
        "payload_kind": "subgraph",
        "component_id": "12",
        "prompt_id": "prompt-1",
        "execute_node_ids": ["12"],
        "subgraph_prompt": {
            "12": {
                "class_type": "KSampler",
                "inputs": {
                    "latent_image": ["remote_input_0", 0],
                    "seed": ["remote_input_1", 0],
                },
            }
        },
        "boundary_inputs": [
            {
                "proxy_input_name": "remote_input_0",
                "io_type": "LATENT",
                "targets": [{"node_id": "12", "input_name": "latent_image"}],
            },
            {
                "proxy_input_name": "remote_input_1",
                "io_type": "INT",
                "targets": [{"node_id": "12", "input_name": "seed"}],
            },
        ],
        "boundary_outputs": [{"node_id": "12", "io_type": "STRING", "is_list": False}],
        "extra_data": {"client_id": "client-1"},
    }

    response = asyncio.run(
        remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
            payload,
            serialization_module.serialize_node_inputs(
                {
                    "remote_input_0": ["latent-a", "latent-b"],
                    "remote_input_1": [10, 11],
                }
            ),
        )
    )

    assert observed_inputs == [
        {"remote_input_0": "latent-a", "remote_input_1": 10},
        {"remote_input_0": "latent-b", "remote_input_1": 11},
    ]
    assert serialization_module.deserialize_node_outputs(response) == (
        ["latent-a:10", "latent-b:11"],
    )


def test_implicit_mapping_preserves_create_video_frame_sequence(
    remote_modal_app_module: Any,
    caplog: pytest.LogCaptureFixture,
) -> None:
    """CreateVideo should receive one complete IMAGE frame sequence."""
    torch = pytest.importorskip("torch")
    frame_sequence = torch.zeros((124, 8, 8, 3), dtype=torch.float32)
    payload = {
        "payload_kind": "subgraph",
        "component_id": "105:91",
        "subgraph_prompt": {
            "105:91": {
                "class_type": "CreateVideo",
                "inputs": {"images": ["105:10", 0]},
            }
        },
        "boundary_inputs": [
            {
                "proxy_input_name": "remote_input_0",
                "io_type": "IMAGE",
                "targets": [{"node_id": "105:91", "input_name": "images"}],
            }
        ],
    }
    caplog.set_level(logging.INFO)

    split_inputs = remote_modal_app_module._split_batch_boundary_inputs(
        payload,
        {"remote_input_0": frame_sequence},
    )

    assert split_inputs is None
    assert "target sockets consume the complete tensor batch" in caplog.text
    assert "105:91.images" in caplog.text


def test_implicitly_mapped_subgraph_shared_model_keeps_unbatched_sampler_single_run(
    remote_modal_app_module: Any,
    serialization_module: Any,
    monkeypatch: Any,
) -> None:
    """A shared MODEL with mixed batch-size INT seeds should run sampler 4 once and sampler 12 four times."""
    observed_calls: list[tuple[str, tuple[str, ...], dict[str, Any]]] = []

    async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
        hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
        execute_node_ids = tuple(str(node_id) for node_id in payload.get("execute_node_ids", []))
        observed_calls.append((str(payload["component_id"]), execute_node_ids, hydrated_inputs))

        if execute_node_ids == ("4",):
            return serialization_module.serialize_node_outputs(("sampler-4",))
        if execute_node_ids == ("12",):
            return serialization_module.serialize_node_outputs((f"sampler-12:{hydrated_inputs['remote_input_0']}",))
        raise AssertionError(f"Unexpected execute nodes for implicit mapped regression: {execute_node_ids!r}")

    monkeypatch.setattr(
        remote_modal_app_module,
        "invoke_remote_engine_async",
        fake_invoke_remote_engine_async,
    )

    payload = {
        "payload_kind": "subgraph",
        "component_id": "17",
        "prompt_id": "prompt-1",
        "component_node_ids": ["4", "12", "17"],
        "execute_node_ids": ["4", "12"],
        "subgraph_prompt": {
            "17": {"class_type": "LoraLoaderModelOnly", "inputs": {}},
            "4": {"class_type": "KSampler", "inputs": {"model": ["17", 0], "seed": 0}},
            "12": {
                "class_type": "KSampler",
                "inputs": {"model": ["17", 0], "seed": 0},
            },
        },
        "boundary_inputs": [
            {
                "proxy_input_name": "remote_input_0",
                "io_type": "INT",
                "targets": [{"node_id": "12", "input_name": "seed"}],
            },
            {
                "proxy_input_name": "remote_input_1",
                "io_type": "INT",
                "targets": [{"node_id": "4", "input_name": "seed"}],
            },
        ],
        "boundary_outputs": [
            {"node_id": "4", "io_type": "STRING", "is_list": False},
            {"node_id": "12", "io_type": "STRING", "is_list": False},
        ],
        "extra_data": {"client_id": "client-1"},
    }

    response = asyncio.run(
        remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
            payload,
            serialization_module.serialize_node_inputs(
                {
                    "remote_input_0": [10, 11, 12, 13],
                    "remote_input_1": [28],
                }
            ),
        )
    )

    assert serialization_module.deserialize_node_outputs(response) == (
        "sampler-4",
        ["sampler-12:10", "sampler-12:11", "sampler-12:12", "sampler-12:13"],
    )
    assert observed_calls == [
        ("17::static", ("4",), {"remote_input_1": [28]}),
        ("17::item:0", ("12",), {"remote_input_0": 10, "remote_input_1": [28]}),
        ("17::item:1", ("12",), {"remote_input_0": 11, "remote_input_1": [28]}),
        ("17::item:2", ("12",), {"remote_input_0": 12, "remote_input_1": [28]}),
        ("17::item:3", ("12",), {"remote_input_0": 13, "remote_input_1": [28]}),
    ]


def test_implicitly_mapped_subgraph_seeds_remote_lanes_before_item_dispatch(
    remote_modal_app_module: Any,
    serialization_module: Any,
    session_state_module: Any,
    monkeypatch: Any,
) -> None:
    """Mapped remote scheduling should seed one bound worker lane before sending per-item calls there."""
    observed_calls: list[tuple[str, str, bool, tuple[str, ...], dict[str, Any]]] = []
    observed_lane_setup_starts: list[tuple[str, int]] = []

    monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "remote")
    monkeypatch.setattr(remote_modal_app_module, "modal", object())
    monkeypatch.setattr(remote_modal_app_module, "_mapped_execution_parallelism", lambda total_items: 2)
    monkeypatch.setattr(remote_modal_app_module, "ensure_remote_warm_capacity", lambda *args, **kwargs: 0)
    monkeypatch.setattr(
        remote_modal_app_module,
        "_lookup_deployed_remote_engine",
        lambda payload, affinity_key_override=None: (
            f"engine:{affinity_key_override or remote_modal_app_module._remote_worker_affinity_key(payload)}"
        ),
    )

    async def fake_invoke_bound_remote_engine_async(
        remote_engine: Any,
        payload: dict[str, Any],
        kwargs_payload: bytes,
    ) -> bytes:
        hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
        execute_node_ids = tuple(str(node_id) for node_id in payload.get("execute_node_ids", []))
        observed_calls.append(
            (
                str(remote_engine),
                str(payload["component_id"]),
                bool(payload.get("clear_remote_session")),
                execute_node_ids,
                hydrated_inputs,
            )
        )
        if str(payload["component_id"]).startswith("17::seed:"):
            return serialization_module.serialize_node_outputs(())
        if str(payload["component_id"]).startswith("17::cleanup:"):
            return serialization_module.serialize_node_outputs(())
        if execute_node_ids == ("12",):
            await asyncio.sleep(0)
            return serialization_module.serialize_node_outputs((f"{remote_engine}:{hydrated_inputs['remote_input_0']}",))
        raise AssertionError(f"Unexpected seeded-lane payload: {payload!r}")

    monkeypatch.setattr(
        remote_modal_app_module,
        "_invoke_bound_remote_engine_async",
        fake_invoke_bound_remote_engine_async,
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_emit_local_mapped_lane_progress_start",
        lambda payload, lane_index, item_index=None: observed_lane_setup_starts.append(
            (str(payload["component_id"]), int(lane_index))
        ),
    )

    payload = {
        "payload_kind": "subgraph",
        "component_id": "17",
        "prompt_id": "prompt-1",
        "component_node_ids": ["12"],
        "execute_node_ids": ["12"],
        "subgraph_prompt": {
            "12": {
                "class_type": "KSampler",
                "inputs": {"model": 0, "seed": 0},
            },
        },
        "boundary_inputs": [
            {
                "proxy_input_name": "remote_input_0",
                "io_type": "INT",
                "targets": [{"node_id": "12", "input_name": "seed"}],
            },
            {
                "proxy_input_name": "static_input_0",
                "io_type": "MODEL",
                "targets": [{"node_id": "12", "input_name": "model"}],
            },
        ],
        "boundary_outputs": [
            {"node_id": "12", "io_type": "STRING", "is_list": False},
        ],
        "extra_data": {"client_id": "client-1"},
        "remote_session": session_state_module.RemoteSessionHandle(
            session_id="session-1",
            prompt_id="prompt-1",
            owner_component_id="17",
        ).to_payload(),
        "clear_remote_session": True,
        "static_to_mapped_boundaries": [
            {
                "proxy_name": "static_input_0",
                "node_id": "4",
                "output_index": 0,
                "io_type": "MODEL",
                "is_list": False,
                "targets": [{"node_id": "12", "input_name": "model"}],
            }
        ],
        "static_phase": {
            "component_node_ids": ["4"],
            "subgraph_prompt": {
                "4": {"class_type": "LoraLoaderModelOnly", "inputs": {}},
            },
            "boundary_inputs": [],
            "boundary_outputs": [
                {
                    "proxy_output_name": "static_input_0",
                    "node_id": "4",
                    "output_index": 0,
                    "io_type": "MODEL",
                    "is_list": False,
                    "session_output": True,
                    "preview_target_node_ids": [],
                }
            ],
            "execute_node_ids": ["4"],
        },
    }

    response = asyncio.run(
        remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
            payload,
            serialization_module.serialize_node_inputs(
                {
                    "remote_input_0": [10, 11, 12, 13],
                    "static_input_0": {"__comfy_modal_remote_session_bridge_ref__": True},
                }
            ),
        )
    )

    seed_calls = [call for call in observed_calls if call[1].startswith("17::seed:")]
    item_calls = [call for call in observed_calls if call[1].startswith("17::item:")]
    cleanup_calls = [call for call in observed_calls if call[1].startswith("17::cleanup:")]

    assert sorted(seed_calls) == [
        (
            "engine:worker-pool:slot:0",
            "17::seed:0",
            False,
            ("4",),
            {"static_input_0": {"__comfy_modal_remote_session_bridge_ref__": True}},
        ),
        (
            "engine:worker-pool:slot:1",
            "17::seed:1",
            False,
            ("4",),
            {"static_input_0": {"__comfy_modal_remote_session_bridge_ref__": True}},
        ),
    ]
    assert {call[0] for call in item_calls} <= {"engine:worker-pool:slot:0", "engine:worker-pool:slot:1"}
    assert {call[0] for call in item_calls} <= {call[0] for call in seed_calls}
    assert sorted(cleanup_calls) == [
        ("engine:worker-pool:slot:0", "17::cleanup:0", True, (), {}),
        ("engine:worker-pool:slot:1", "17::cleanup:1", True, (), {}),
    ]
    assert sorted(observed_lane_setup_starts) == [("17", 0), ("17", 1)]
    response_outputs = serialization_module.deserialize_node_outputs(response)[0]
    assert len(response_outputs) == 4
    assert [output.rsplit(":", 1)[-1] for output in response_outputs] == ["10", "11", "12", "13"]


def test_implicitly_mapped_subgraph_allows_ready_lanes_to_start_before_slowest_seed_finishes(
    remote_modal_app_module: Any,
    serialization_module: Any,
    session_state_module: Any,
    monkeypatch: Any,
) -> None:
    """A fast seeded lane should start item execution before another lane finishes seeding."""
    observed_events: list[tuple[str, str, str]] = []

    monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "remote")
    monkeypatch.setattr(remote_modal_app_module, "modal", object())
    monkeypatch.setattr(remote_modal_app_module, "_mapped_execution_parallelism", lambda total_items: 2)
    monkeypatch.setattr(remote_modal_app_module, "ensure_remote_warm_capacity", lambda *args, **kwargs: 0)
    monkeypatch.setattr(
        remote_modal_app_module,
        "_lookup_deployed_remote_engine",
        lambda payload, affinity_key_override=None: (
            f"engine:{affinity_key_override or remote_modal_app_module._remote_worker_affinity_key(payload)}"
        ),
    )

    async def fake_invoke_bound_remote_engine_async(
        remote_engine: Any,
        payload: dict[str, Any],
        kwargs_payload: bytes,
    ) -> bytes:
        hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
        component_id = str(payload["component_id"])
        if component_id == "17::seed:0":
            observed_events.append(("seed_start", component_id, str(remote_engine)))
            await asyncio.sleep(0.05)
            observed_events.append(("seed_done", component_id, str(remote_engine)))
            return serialization_module.serialize_node_outputs(())
        if component_id == "17::seed:1":
            observed_events.append(("seed_start", component_id, str(remote_engine)))
            await asyncio.sleep(0)
            observed_events.append(("seed_done", component_id, str(remote_engine)))
            return serialization_module.serialize_node_outputs(())
        if component_id.startswith("17::cleanup:"):
            observed_events.append(("cleanup", component_id, str(remote_engine)))
            return serialization_module.serialize_node_outputs(())
        if component_id.startswith("17::item:"):
            observed_events.append(("item", component_id, str(remote_engine)))
            return serialization_module.serialize_node_outputs(
                (f"{remote_engine}:{hydrated_inputs['remote_input_0']}",)
            )
        raise AssertionError(f"Unexpected seeded-lane payload: {payload!r}")

    monkeypatch.setattr(
        remote_modal_app_module,
        "_invoke_bound_remote_engine_async",
        fake_invoke_bound_remote_engine_async,
    )

    payload = {
        "payload_kind": "subgraph",
        "component_id": "17",
        "prompt_id": "prompt-1",
        "component_node_ids": ["12"],
        "execute_node_ids": ["12"],
        "subgraph_prompt": {
            "12": {
                "class_type": "KSampler",
                "inputs": {"model": 0, "seed": 0},
            },
        },
        "boundary_inputs": [
            {
                "proxy_input_name": "remote_input_0",
                "io_type": "INT",
                "targets": [{"node_id": "12", "input_name": "seed"}],
            },
            {
                "proxy_input_name": "static_input_0",
                "io_type": "MODEL",
                "targets": [{"node_id": "12", "input_name": "model"}],
            },
        ],
        "boundary_outputs": [
            {"node_id": "12", "io_type": "STRING", "is_list": False},
        ],
        "extra_data": {"client_id": "client-1"},
        "remote_session": session_state_module.RemoteSessionHandle(
            session_id="session-1",
            prompt_id="prompt-1",
            owner_component_id="17",
        ).to_payload(),
        "clear_remote_session": True,
        "static_to_mapped_boundaries": [
            {
                "proxy_name": "static_input_0",
                "node_id": "4",
                "output_index": 0,
                "io_type": "MODEL",
                "is_list": False,
                "targets": [{"node_id": "12", "input_name": "model"}],
            }
        ],
        "static_phase": {
            "component_node_ids": ["4"],
            "subgraph_prompt": {
                "4": {"class_type": "LoraLoaderModelOnly", "inputs": {}},
            },
            "boundary_inputs": [],
            "boundary_outputs": [
                {
                    "proxy_output_name": "static_input_0",
                    "node_id": "4",
                    "output_index": 0,
                    "io_type": "MODEL",
                    "is_list": False,
                    "session_output": True,
                    "preview_target_node_ids": [],
                }
            ],
            "execute_node_ids": ["4"],
        },
    }

    response = asyncio.run(
        remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
            payload,
            serialization_module.serialize_node_inputs(
                {
                    "remote_input_0": [10, 11, 12, 13],
                    "static_input_0": {"__comfy_modal_remote_session_bridge_ref__": True},
                }
            ),
        )
    )

    assert len(serialization_module.deserialize_node_outputs(response)[0]) == 4
    assert ("seed_done", "17::seed:1", "engine:worker-pool:slot:1") in observed_events
    first_item_index = next(
        index for index, event in enumerate(observed_events) if event[0] == "item"
    )
    fast_seed_done_index = observed_events.index(
        ("seed_done", "17::seed:1", "engine:worker-pool:slot:1")
    )
    assert first_item_index > fast_seed_done_index


def test_implicitly_mapped_subgraph_returns_once_all_items_finish_without_waiting_for_unused_lane(
    remote_modal_app_module: Any,
    serialization_module: Any,
    session_state_module: Any,
    monkeypatch: Any,
) -> None:
    """Prompt completion should not wait for a slow seeded lane after other lanes finish all mapped items."""
    observed_events: list[tuple[str, str, str]] = []

    monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "remote")
    monkeypatch.setattr(remote_modal_app_module, "modal", object())
    monkeypatch.setattr(remote_modal_app_module, "_mapped_execution_parallelism", lambda total_items: 2)
    monkeypatch.setattr(remote_modal_app_module, "ensure_remote_warm_capacity", lambda *args, **kwargs: 0)
    monkeypatch.setattr(
        remote_modal_app_module,
        "_lookup_deployed_remote_engine",
        lambda payload, affinity_key_override=None: (
            f"engine:{affinity_key_override or remote_modal_app_module._remote_worker_affinity_key(payload)}"
        ),
    )

    async def fake_invoke_bound_remote_engine_async(
        remote_engine: Any,
        payload: dict[str, Any],
        kwargs_payload: bytes,
    ) -> bytes:
        hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
        component_id = str(payload["component_id"])
        if component_id == "17::seed:0":
            observed_events.append(("seed_start", component_id, str(remote_engine)))
            await asyncio.sleep(1.0)
            observed_events.append(("seed_done", component_id, str(remote_engine)))
            return serialization_module.serialize_node_outputs(())
        if component_id == "17::seed:1":
            observed_events.append(("seed_start", component_id, str(remote_engine)))
            await asyncio.sleep(0)
            observed_events.append(("seed_done", component_id, str(remote_engine)))
            return serialization_module.serialize_node_outputs(())
        if component_id.startswith("17::cleanup:"):
            observed_events.append(("cleanup", component_id, str(remote_engine)))
            return serialization_module.serialize_node_outputs(())
        if component_id.startswith("17::item:"):
            observed_events.append(("item", component_id, str(remote_engine)))
            return serialization_module.serialize_node_outputs(
                (f"{remote_engine}:{hydrated_inputs['remote_input_0']}",)
            )
        raise AssertionError(f"Unexpected seeded-lane payload: {payload!r}")

    monkeypatch.setattr(
        remote_modal_app_module,
        "_invoke_bound_remote_engine_async",
        fake_invoke_bound_remote_engine_async,
    )

    payload = {
        "payload_kind": "subgraph",
        "component_id": "17",
        "prompt_id": "prompt-1",
        "component_node_ids": ["12"],
        "execute_node_ids": ["12"],
        "subgraph_prompt": {
            "12": {
                "class_type": "KSampler",
                "inputs": {"model": 0, "seed": 0},
            },
        },
        "boundary_inputs": [
            {
                "proxy_input_name": "remote_input_0",
                "io_type": "INT",
                "targets": [{"node_id": "12", "input_name": "seed"}],
            },
            {
                "proxy_input_name": "static_input_0",
                "io_type": "MODEL",
                "targets": [{"node_id": "12", "input_name": "model"}],
            },
        ],
        "boundary_outputs": [
            {"node_id": "12", "io_type": "STRING", "is_list": False},
        ],
        "extra_data": {"client_id": "client-1"},
        "remote_session": session_state_module.RemoteSessionHandle(
            session_id="session-1",
            prompt_id="prompt-1",
            owner_component_id="17",
        ).to_payload(),
        "clear_remote_session": True,
        "static_to_mapped_boundaries": [
            {
                "proxy_name": "static_input_0",
                "node_id": "4",
                "output_index": 0,
                "io_type": "MODEL",
                "is_list": False,
                "targets": [{"node_id": "12", "input_name": "model"}],
            }
        ],
        "static_phase": {
            "component_node_ids": ["4"],
            "subgraph_prompt": {
                "4": {"class_type": "LoraLoaderModelOnly", "inputs": {}},
            },
            "boundary_inputs": [],
            "boundary_outputs": [
                {
                    "proxy_output_name": "static_input_0",
                    "node_id": "4",
                    "output_index": 0,
                    "io_type": "MODEL",
                    "is_list": False,
                    "session_output": True,
                    "preview_target_node_ids": [],
                }
            ],
            "execute_node_ids": ["4"],
        },
    }

    started_at = time.perf_counter()
    response = asyncio.run(
        remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
            payload,
            serialization_module.serialize_node_inputs(
                {
                    "remote_input_0": [10, 11],
                    "static_input_0": {"__comfy_modal_remote_session_bridge_ref__": True},
                }
            ),
        )
    )
    elapsed = time.perf_counter() - started_at

    assert serialization_module.deserialize_node_outputs(response)[0] == [
        "engine:worker-pool:slot:1:10",
        "engine:worker-pool:slot:1:11",
    ]
    assert elapsed < 0.5
    assert ("item", "17::item:0", "engine:worker-pool:slot:1") in observed_events
    assert ("cleanup", "17::cleanup:1", "engine:worker-pool:slot:1") in observed_events
    assert ("cleanup", "17::cleanup:0", "engine:worker-pool:slot:0") not in observed_events
    assert ("seed_done", "17::seed:0", "engine:worker-pool:slot:0") not in observed_events


def test_implicitly_mapped_subgraph_ignores_late_lane_failure_after_all_items_complete(
    remote_modal_app_module: Any,
    serialization_module: Any,
    session_state_module: Any,
    monkeypatch: Any,
) -> None:
    """A detached late lane must not fail the prompt after all mapped outputs are already complete."""
    observed_events: list[tuple[str, str, str]] = []

    monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "remote")
    monkeypatch.setattr(remote_modal_app_module, "modal", object())
    monkeypatch.setattr(remote_modal_app_module, "_mapped_execution_parallelism", lambda total_items: 2)
    monkeypatch.setattr(remote_modal_app_module, "ensure_remote_warm_capacity", lambda *args, **kwargs: 0)
    monkeypatch.setattr(
        remote_modal_app_module,
        "_lookup_deployed_remote_engine",
        lambda payload, affinity_key_override=None: (
            f"engine:{affinity_key_override or remote_modal_app_module._remote_worker_affinity_key(payload)}"
        ),
    )

    async def fake_invoke_bound_remote_engine_async(
        remote_engine: Any,
        payload: dict[str, Any],
        kwargs_payload: bytes,
    ) -> bytes:
        hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
        component_id = str(payload["component_id"])
        if component_id == "17::seed:0":
            observed_events.append(("seed_start", component_id, str(remote_engine)))
            await asyncio.sleep(0.05)
            raise RuntimeError("late lane failed after prompt completion")
        if component_id == "17::seed:1":
            observed_events.append(("seed_start", component_id, str(remote_engine)))
            await asyncio.sleep(0)
            observed_events.append(("seed_done", component_id, str(remote_engine)))
            return serialization_module.serialize_node_outputs(())
        if component_id.startswith("17::cleanup:"):
            observed_events.append(("cleanup", component_id, str(remote_engine)))
            return serialization_module.serialize_node_outputs(())
        if component_id.startswith("17::item:"):
            observed_events.append(("item", component_id, str(remote_engine)))
            return serialization_module.serialize_node_outputs(
                (f"{remote_engine}:{hydrated_inputs['remote_input_0']}",)
            )
        raise AssertionError(f"Unexpected seeded-lane payload: {payload!r}")

    monkeypatch.setattr(
        remote_modal_app_module,
        "_invoke_bound_remote_engine_async",
        fake_invoke_bound_remote_engine_async,
    )

    payload = {
        "payload_kind": "subgraph",
        "component_id": "17",
        "prompt_id": "prompt-1",
        "component_node_ids": ["12"],
        "execute_node_ids": ["12"],
        "subgraph_prompt": {
            "12": {
                "class_type": "KSampler",
                "inputs": {"model": 0, "seed": 0},
            },
        },
        "boundary_inputs": [
            {
                "proxy_input_name": "remote_input_0",
                "io_type": "INT",
                "targets": [{"node_id": "12", "input_name": "seed"}],
            },
            {
                "proxy_input_name": "static_input_0",
                "io_type": "MODEL",
                "targets": [{"node_id": "12", "input_name": "model"}],
            },
        ],
        "boundary_outputs": [
            {"node_id": "12", "io_type": "STRING", "is_list": False},
        ],
        "extra_data": {"client_id": "client-1"},
        "remote_session": session_state_module.RemoteSessionHandle(
            session_id="session-1",
            prompt_id="prompt-1",
            owner_component_id="17",
        ).to_payload(),
        "clear_remote_session": True,
        "static_to_mapped_boundaries": [
            {
                "proxy_name": "static_input_0",
                "node_id": "4",
                "output_index": 0,
                "io_type": "MODEL",
                "is_list": False,
                "targets": [{"node_id": "12", "input_name": "model"}],
            }
        ],
        "static_phase": {
            "component_node_ids": ["4"],
            "subgraph_prompt": {
                "4": {"class_type": "LoraLoaderModelOnly", "inputs": {}},
            },
            "boundary_inputs": [],
            "boundary_outputs": [
                {
                    "proxy_output_name": "static_input_0",
                    "node_id": "4",
                    "output_index": 0,
                    "io_type": "MODEL",
                    "is_list": False,
                    "session_output": True,
                    "preview_target_node_ids": [],
                }
            ],
            "execute_node_ids": ["4"],
        },
    }

    response = asyncio.run(
        remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
            payload,
            serialization_module.serialize_node_inputs(
                {
                    "remote_input_0": [10, 11],
                    "static_input_0": {"__comfy_modal_remote_session_bridge_ref__": True},
                }
            ),
        )
    )

    assert serialization_module.deserialize_node_outputs(response)[0] == [
        "engine:worker-pool:slot:1:10",
        "engine:worker-pool:slot:1:11",
    ]
    assert ("cleanup", "17::cleanup:1", "engine:worker-pool:slot:1") in observed_events
    assert ("cleanup", "17::cleanup:0", "engine:worker-pool:slot:0") not in observed_events


def test_implicitly_mapped_subgraph_skips_outer_fanout_for_input_is_list_targets(
    remote_modal_app_module: Any,
    serialization_module: Any,
    monkeypatch: Any,
) -> None:
    """Implicit fan-out should not split components whose list boundary lands on INPUT_IS_LIST nodes."""
    observed_calls: list[tuple[str, tuple[str, ...], dict[str, Any]]] = []

    fake_nodes_module = type(
        "FakeNodesModule",
        (),
        {
            "NODE_CLASS_MAPPINGS": {
                "ImplicitBatchListSource": _ImplicitBatchListSourceNode,
                "ImplicitBatchScalarConsumer": _ImplicitBatchScalarConsumerNode,
                "ImplicitBatchListConsumer": _ImplicitBatchListConsumerNode,
            }
        },
    )()
    monkeypatch.setattr(remote_modal_app_module, "_load_nodes_module", lambda: fake_nodes_module)

    async def fake_invoke_remote_engine_async(
        payload: dict[str, Any],
        kwargs_payload: bytes,
        *,
        allow_implicit_mapping: bool = True,
    ) -> bytes:
        del allow_implicit_mapping
        hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
        execute_node_ids = tuple(str(node_id) for node_id in payload.get("execute_node_ids", []))
        observed_calls.append((str(payload["component_id"]), execute_node_ids, hydrated_inputs))
        return serialization_module.serialize_node_outputs(
            ("scalar-output", ["list-output:0", "list-output:1", "list-output:2"])
        )

    monkeypatch.setattr(
        remote_modal_app_module,
        "invoke_remote_engine_async",
        fake_invoke_remote_engine_async,
    )

    payload = {
        "payload_kind": "subgraph",
        "component_id": "17",
        "prompt_id": "prompt-1",
        "component_node_ids": ["4", "12", "17"],
        "execute_node_ids": ["4", "12"],
        "subgraph_prompt": {
            "17": {"class_type": "ImplicitBatchListSource", "inputs": {"values": 0}},
            "4": {
                "class_type": "ImplicitBatchScalarConsumer",
                "inputs": {"value": ["17", 0]},
            },
            "12": {
                "class_type": "ImplicitBatchListConsumer",
                "inputs": {"values": ["17", 1]},
            },
        },
        "boundary_inputs": [
            {
                "proxy_input_name": "remote_input_0",
                "io_type": "INT",
                "targets": [{"node_id": "17", "input_name": "values"}],
            }
        ],
        "boundary_outputs": [
            {"node_id": "4", "io_type": "STRING", "is_list": False},
            {"node_id": "12", "io_type": "STRING", "is_list": True},
        ],
        "extra_data": {"client_id": "client-1"},
    }

    response = asyncio.run(
        remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
            payload,
            serialization_module.serialize_node_inputs({"remote_input_0": [10, 11, 12]}),
        )
    )

    assert serialization_module.deserialize_node_outputs(response) == (
        "scalar-output",
        ["list-output:0", "list-output:1", "list-output:2"],
    )
    assert observed_calls == [
        ("17", ("4", "12"), {"remote_input_0": [10, 11, 12]}),
    ]


def test_invoke_remote_engine_async_bypasses_implicit_mapping_once_for_input_is_list_targets(
    remote_modal_app_module: Any,
    serialization_module: Any,
    monkeypatch: Any,
) -> None:
    """The async dispatcher should not recurse when implicit fan-out is suppressed."""
    observed_calls: list[tuple[bool, dict[str, Any]]] = []

    fake_nodes_module = type(
        "FakeNodesModule",
        (),
        {
            "NODE_CLASS_MAPPINGS": {
                "ImplicitBatchListSource": _ImplicitBatchListSourceNode,
                "ImplicitBatchScalarConsumer": _ImplicitBatchScalarConsumerNode,
                "ImplicitBatchListConsumer": _ImplicitBatchListConsumerNode,
            }
        },
    )()
    monkeypatch.setattr(remote_modal_app_module, "_load_nodes_module", lambda: fake_nodes_module)

    def fake_invoke_remote_engine(
        payload: dict[str, Any],
        kwargs_payload: bytes,
        *,
        allow_implicit_mapping: bool = True,
    ) -> bytes:
        hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
        observed_calls.append((allow_implicit_mapping, hydrated_inputs))
        return serialization_module.serialize_node_outputs(("ordinary-subgraph",))

    monkeypatch.setattr(
        remote_modal_app_module,
        "invoke_remote_engine",
        fake_invoke_remote_engine,
    )

    payload = {
        "payload_kind": "subgraph",
        "component_id": "17",
        "prompt_id": "prompt-1",
        "component_node_ids": ["4", "12", "17"],
        "execute_node_ids": ["4", "12"],
        "subgraph_prompt": {
            "17": {"class_type": "ImplicitBatchListSource", "inputs": {"values": 0}},
            "4": {
                "class_type": "ImplicitBatchScalarConsumer",
                "inputs": {"value": ["17", 0]},
            },
            "12": {
                "class_type": "ImplicitBatchListConsumer",
                "inputs": {"values": ["17", 1]},
            },
        },
        "boundary_inputs": [
            {
                "proxy_input_name": "remote_input_0",
                "io_type": "INT",
                "targets": [{"node_id": "17", "input_name": "values"}],
            }
        ],
        "boundary_outputs": [
            {"node_id": "4", "io_type": "STRING", "is_list": False},
        ],
        "extra_data": {"client_id": "client-1"},
    }

    response = asyncio.run(
        remote_modal_app_module.invoke_remote_engine_async(
            payload,
            serialization_module.serialize_node_inputs({"remote_input_0": [10, 11, 12]}),
        )
    )

    assert serialization_module.deserialize_node_outputs(response) == ("ordinary-subgraph",)
    assert observed_calls == [
        (False, {"remote_input_0": [10, 11, 12]}),
    ]


def test_implicitly_mapped_subgraph_clears_remote_session_once_after_all_items_finish(
    remote_modal_app_module: Any,
    serialization_module: Any,
    session_state_module: Any,
    monkeypatch: Any,
) -> None:
    """Implicit mapped execution should reserve remote-session cleanup for one final cleanup payload."""
    observed_calls: list[tuple[str, bool, tuple[str, ...], dict[str, Any]]] = []
    session_cleared = False

    monkeypatch.setattr(remote_modal_app_module, "_mapped_execution_parallelism", lambda total_items: 2)

    async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
        nonlocal session_cleared
        hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
        execute_node_ids = tuple(str(node_id) for node_id in payload.get("execute_node_ids", []))
        clear_remote_session = bool(payload.get("clear_remote_session"))
        observed_calls.append(
            (str(payload["component_id"]), clear_remote_session, execute_node_ids, hydrated_inputs)
        )

        if clear_remote_session:
            session_cleared = True
        elif session_cleared:
            raise session_state_module.RemoteSessionStateError(
                "Remote session 'session-1' was not found."
            )

        if execute_node_ids == ("4",):
            return serialization_module.serialize_node_outputs(("sampler-4",))
        if execute_node_ids == ("12",):
            return serialization_module.serialize_node_outputs((f"sampler-12:{hydrated_inputs['remote_input_0']}",))
        if execute_node_ids == ():
            return serialization_module.serialize_node_outputs(())
        raise AssertionError(f"Unexpected execute nodes for implicit cleanup regression: {execute_node_ids!r}")

    monkeypatch.setattr(
        remote_modal_app_module,
        "invoke_remote_engine_async",
        fake_invoke_remote_engine_async,
    )

    payload = {
        "payload_kind": "subgraph",
        "component_id": "17",
        "prompt_id": "prompt-1",
        "component_node_ids": ["4", "12", "17"],
        "execute_node_ids": ["4", "12"],
        "subgraph_prompt": {
            "17": {"class_type": "LoraLoaderModelOnly", "inputs": {}},
            "4": {"class_type": "KSampler", "inputs": {"model": ["17", 0], "seed": 0}},
            "12": {
                "class_type": "KSampler",
                "inputs": {"model": ["17", 0], "seed": 0},
            },
        },
        "boundary_inputs": [
            {
                "proxy_input_name": "remote_input_0",
                "io_type": "INT",
                "targets": [{"node_id": "12", "input_name": "seed"}],
            },
            {
                "proxy_input_name": "remote_input_1",
                "io_type": "INT",
                "targets": [{"node_id": "4", "input_name": "seed"}],
            },
        ],
        "boundary_outputs": [
            {"node_id": "4", "io_type": "STRING", "is_list": False},
            {"node_id": "12", "io_type": "STRING", "is_list": False},
        ],
        "extra_data": {"client_id": "client-1"},
        "remote_session": session_state_module.RemoteSessionHandle(
            session_id="session-1",
            prompt_id="prompt-1",
            owner_component_id="17",
        ).to_payload(),
        "clear_remote_session": True,
    }

    response = asyncio.run(
        remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
            payload,
            serialization_module.serialize_node_inputs(
                {
                    "remote_input_0": [10, 11, 12],
                    "remote_input_1": [28],
                }
            ),
        )
    )

    assert serialization_module.deserialize_node_outputs(response) == (
        "sampler-4",
        ["sampler-12:10", "sampler-12:11", "sampler-12:12"],
    )
    assert [
        (component_id, execute_node_ids)
        for component_id, _, execute_node_ids, _ in observed_calls
    ] == [
        ("17::static", ("4",)),
        ("17::item:0", ("12",)),
        ("17::item:1", ("12",)),
        ("17::item:2", ("12",)),
        ("17::cleanup", ()),
    ]
    assert all(
        not clear_remote_session
        for component_id, clear_remote_session, _, _ in observed_calls
        if component_id != "17::cleanup"
    )
    assert observed_calls[-1] == ("17::cleanup", True, (), {})


def test_implicitly_mapped_subgraph_stops_queued_items_after_local_interrupt(
    remote_modal_app_module: Any,
    serialization_module: Any,
    session_state_module: Any,
    monkeypatch: Any,
) -> None:
    """Implicit mapped cancellation should not dispatch queued items or cleanup work."""

    class FakeInterrupt(Exception):
        """Stand-in for ComfyUI's InterruptProcessingException."""

    observed_calls: list[str] = []
    local_interrupted = False

    async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
        nonlocal local_interrupted
        del kwargs_payload
        component_id = str(payload["component_id"])
        observed_calls.append(component_id)
        if component_id.endswith("::cleanup"):
            raise AssertionError("interrupted mapped execution should not run cleanup")
        local_interrupted = True
        return serialization_module.serialize_node_outputs((component_id,))

    def fake_local_processing_interrupted() -> bool:
        return local_interrupted

    monkeypatch.setattr(
        remote_modal_app_module,
        "boost_mapped_component_warmup",
        lambda payload, *, total_items, reason: (1, 1),
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "invoke_remote_engine_async",
        fake_invoke_remote_engine_async,
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_local_processing_interrupted",
        fake_local_processing_interrupted,
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_raise_local_interrupt",
        lambda: (_ for _ in ()).throw(FakeInterrupt()),
    )

    payload = {
        "payload_kind": "subgraph",
        "component_id": "17",
        "prompt_id": "prompt-1",
        "component_node_ids": ["12"],
        "execute_node_ids": ["12"],
        "subgraph_prompt": {
            "12": {
                "class_type": "KSampler",
                "inputs": {"seed": 0},
            },
        },
        "boundary_inputs": [
            {
                "proxy_input_name": "remote_input_0",
                "io_type": "INT",
                "targets": [{"node_id": "12", "input_name": "seed"}],
            },
        ],
        "boundary_outputs": [{"node_id": "12", "io_type": "STRING", "is_list": False}],
        "extra_data": {"client_id": "client-1"},
        "remote_session": session_state_module.RemoteSessionHandle(
            session_id="session-1",
            prompt_id="prompt-1",
            owner_component_id="17",
        ).to_payload(),
        "clear_remote_session": True,
    }

    with pytest.raises(FakeInterrupt):
        asyncio.run(
            remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
                payload,
                serialization_module.serialize_node_inputs({"remote_input_0": [10, 11, 12]}),
            )
        )

    assert observed_calls == ["17::item:0"]


def test_implicitly_mapped_subgraph_keeps_conditioning_lists_broadcast(
    remote_modal_app_module: Any,
    serialization_module: Any,
    monkeypatch: Any,
) -> None:
    """Implicit mapped execution must not split list-backed CONDITIONING inputs per item."""
    conditioning = [
        ["cond-a", {"pooled_output": "pool-a"}],
        ["cond-b", {"pooled_output": "pool-b"}],
    ]
    observed_calls: list[tuple[str, tuple[str, ...], dict[str, Any]]] = []

    async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
        hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
        execute_node_ids = tuple(str(node_id) for node_id in payload.get("execute_node_ids", []))
        observed_calls.append((str(payload["component_id"]), execute_node_ids, hydrated_inputs))

        if execute_node_ids != ("12",):
            raise AssertionError(
                f"Unexpected execute nodes for conditioning implicit mapped regression: {execute_node_ids!r}"
            )
        return serialization_module.serialize_node_outputs(
            (f"{hydrated_inputs['remote_input_1']}:{len(hydrated_inputs['remote_input_0'])}",)
        )

    monkeypatch.setattr(
        remote_modal_app_module,
        "invoke_remote_engine_async",
        fake_invoke_remote_engine_async,
    )

    payload = {
        "payload_kind": "subgraph",
        "component_id": "17",
        "prompt_id": "prompt-1",
        "component_node_ids": ["12"],
        "execute_node_ids": ["12"],
        "subgraph_prompt": {
            "12": {
                "class_type": "KSampler",
                "inputs": {
                    "positive": 0,
                    "seed": 0,
                },
            },
        },
        "boundary_inputs": [
            {
                "proxy_input_name": "remote_input_0",
                "io_type": "CONDITIONING",
                "targets": [{"node_id": "12", "input_name": "positive"}],
            },
            {
                "proxy_input_name": "remote_input_1",
                "io_type": "INT",
                "targets": [{"node_id": "12", "input_name": "seed"}],
            },
        ],
        "boundary_outputs": [
            {"node_id": "12", "io_type": "STRING", "is_list": False},
        ],
        "extra_data": {"client_id": "client-1"},
    }

    response = asyncio.run(
        remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
            payload,
            serialization_module.serialize_node_inputs(
                {
                    "remote_input_0": conditioning,
                    "remote_input_1": [10, 11, 12, 13],
                }
            ),
        )
    )

    assert serialization_module.deserialize_node_outputs(response) == (
        ["10:2", "11:2", "12:2", "13:2"],
    )
    assert observed_calls == [
        ("17::item:0", ("12",), {"remote_input_0": conditioning, "remote_input_1": 10}),
        ("17::item:1", ("12",), {"remote_input_0": conditioning, "remote_input_1": 11}),
        ("17::item:2", ("12",), {"remote_input_0": conditioning, "remote_input_1": 12}),
        ("17::item:3", ("12",), {"remote_input_0": conditioning, "remote_input_1": 13}),
    ]


def test_implicitly_mapped_subgraph_splits_mapped_conditioning_outputs(
    remote_modal_app_module: Any,
    serialization_module: Any,
    monkeypatch: Any,
) -> None:
    """Mapped CONDITIONING aggregates should feed downstream mapped items in lockstep."""
    conditioning_items = [
        [["cond-a", {"pooled_output": "pool-a"}]],
        [["cond-b", {"pooled_output": "pool-b"}]],
        [["cond-c", {"pooled_output": "pool-c"}]],
    ]
    mapped_conditioning = serialization_module.join_mapped_values(
        conditioning_items,
        "CONDITIONING",
        is_list=False,
    )
    latent_items = ["latent-a", "latent-b", "latent-c"]
    observed_calls: list[tuple[str, dict[str, Any]]] = []

    async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
        hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
        observed_calls.append((str(payload["component_id"]), hydrated_inputs))
        return serialization_module.serialize_node_outputs(
            (f"{hydrated_inputs['remote_input_0'][0][0]}:{hydrated_inputs['remote_input_1']}",)
        )

    monkeypatch.setattr(
        remote_modal_app_module,
        "invoke_remote_engine_async",
        fake_invoke_remote_engine_async,
    )

    payload = {
        "payload_kind": "subgraph",
        "component_id": "17",
        "prompt_id": "prompt-1",
        "component_node_ids": ["12"],
        "execute_node_ids": ["12"],
        "subgraph_prompt": {
            "12": {
                "class_type": "KSampler",
                "inputs": {
                    "positive": 0,
                    "latent_image": 0,
                },
            },
        },
        "boundary_inputs": [
            {
                "proxy_input_name": "remote_input_0",
                "io_type": "CONDITIONING",
                "targets": [{"node_id": "12", "input_name": "positive"}],
            },
            {
                "proxy_input_name": "remote_input_1",
                "io_type": "LATENT",
                "targets": [{"node_id": "12", "input_name": "latent_image"}],
            },
        ],
        "boundary_outputs": [
            {"node_id": "12", "io_type": "STRING", "is_list": False},
        ],
        "extra_data": {"client_id": "client-1"},
    }

    response = asyncio.run(
        remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
            payload,
            serialization_module.serialize_node_inputs(
                {
                    "remote_input_0": mapped_conditioning,
                    "remote_input_1": latent_items,
                }
            ),
        )
    )

    assert serialization_module.deserialize_node_outputs(response) == (
        ["cond-a:latent-a", "cond-b:latent-b", "cond-c:latent-c"],
    )
    assert observed_calls == [
        (
            "17::item:0",
            {"remote_input_0": conditioning_items[0], "remote_input_1": latent_items[0]},
        ),
        (
            "17::item:1",
            {"remote_input_0": conditioning_items[1], "remote_input_1": latent_items[1]},
        ),
        (
            "17::item:2",
            {"remote_input_0": conditioning_items[2], "remote_input_1": latent_items[2]},
        ),
    ]


def test_implicitly_mapped_subgraph_splits_wildcard_latent_lists_for_latent_targets(
    remote_modal_app_module: Any,
    serialization_module: Any,
    monkeypatch: Any,
) -> None:
    """Wildcard boundary lists of LATENT-like mappings should still itemize for LATENT sockets."""
    latent_items = [
        {"samples": "latent-a", "batch_index": [0]},
        {"samples": "latent-b", "batch_index": [1]},
        {"samples": "latent-c", "batch_index": [2]},
    ]
    observed_calls: list[tuple[str, tuple[str, ...], dict[str, Any]]] = []

    fake_nodes_module = type(
        "FakeNodesModule",
        (),
        {"NODE_CLASS_MAPPINGS": {"KSampler": _FakeImplicitBatchKSamplerNode}},
    )()
    monkeypatch.setattr(remote_modal_app_module, "_load_nodes_module", lambda: fake_nodes_module)

    async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
        hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
        execute_node_ids = tuple(str(node_id) for node_id in payload.get("execute_node_ids", []))
        observed_calls.append((str(payload["component_id"]), execute_node_ids, hydrated_inputs))

        if execute_node_ids != ("12",):
            raise AssertionError(
                f"Unexpected execute nodes for wildcard LATENT implicit mapped regression: {execute_node_ids!r}"
            )
        return serialization_module.serialize_node_outputs(
            (f"{hydrated_inputs['remote_input_1']}:{hydrated_inputs['remote_input_0']['samples']}",)
        )

    monkeypatch.setattr(
        remote_modal_app_module,
        "invoke_remote_engine_async",
        fake_invoke_remote_engine_async,
    )

    payload = {
        "payload_kind": "subgraph",
        "component_id": "17",
        "prompt_id": "prompt-1",
        "component_node_ids": ["12"],
        "execute_node_ids": ["12"],
        "subgraph_prompt": {
            "12": {
                "class_type": "KSampler",
                "inputs": {
                    "latent_image": 0,
                    "seed": 0,
                },
            },
        },
        "boundary_inputs": [
            {
                "proxy_input_name": "remote_input_0",
                "io_type": "*",
                "targets": [{"node_id": "12", "input_name": "latent_image"}],
            },
            {
                "proxy_input_name": "remote_input_1",
                "io_type": "INT",
                "targets": [{"node_id": "12", "input_name": "seed"}],
            },
        ],
        "boundary_outputs": [
            {"node_id": "12", "io_type": "STRING", "is_list": False},
        ],
        "extra_data": {"client_id": "client-1"},
    }

    response = asyncio.run(
        remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
            payload,
            serialization_module.serialize_node_inputs(
                {
                    "remote_input_0": latent_items,
                    "remote_input_1": [10, 11, 12],
                }
            ),
        )
    )

    assert serialization_module.deserialize_node_outputs(response) == (
        ["10:latent-a", "11:latent-b", "12:latent-c"],
    )
    assert observed_calls == [
        ("17::item:0", ("12",), {"remote_input_0": latent_items[0], "remote_input_1": 10}),
        ("17::item:1", ("12",), {"remote_input_0": latent_items[1], "remote_input_1": 11}),
        ("17::item:2", ("12",), {"remote_input_0": latent_items[2], "remote_input_1": 12}),
    ]


def test_implicitly_mapped_subgraph_splits_session_ref_lists_for_nontransportable_inputs(
    remote_modal_app_module: Any,
    serialization_module: Any,
    monkeypatch: Any,
) -> None:
    """Implicit mapped execution should itemize lists of remote session refs even for MODEL/CONDITIONING."""
    model_ref = {
        "__comfy_modal_remote_session_value_ref__": True,
        "session_id": "session-1",
        "node_id": "31",
        "output_index": 0,
    }
    conditioning_ref = {
        "__comfy_modal_remote_session_value_ref__": True,
        "session_id": "session-1",
        "node_id": "2",
        "output_index": 0,
    }
    observed_calls: list[tuple[str, tuple[str, ...], dict[str, Any]]] = []

    async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
        hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
        execute_node_ids = tuple(str(node_id) for node_id in payload.get("execute_node_ids", []))
        observed_calls.append((str(payload["component_id"]), execute_node_ids, hydrated_inputs))

        if execute_node_ids != ("12",):
            raise AssertionError(
                f"Unexpected execute nodes for implicit session-ref regression: {execute_node_ids!r}"
            )
        return serialization_module.serialize_node_outputs((f"seed:{hydrated_inputs['remote_input_2']}",))

    monkeypatch.setattr(
        remote_modal_app_module,
        "invoke_remote_engine_async",
        fake_invoke_remote_engine_async,
    )

    payload = {
        "payload_kind": "subgraph",
        "component_id": "17",
        "prompt_id": "prompt-1",
        "component_node_ids": ["12"],
        "execute_node_ids": ["12"],
        "subgraph_prompt": {
            "12": {
                "class_type": "KSampler",
                "inputs": {
                    "model": 0,
                    "positive": 0,
                    "seed": 0,
                },
            },
        },
        "boundary_inputs": [
            {
                "proxy_input_name": "remote_input_0",
                "io_type": "MODEL",
                "targets": [{"node_id": "12", "input_name": "model"}],
            },
            {
                "proxy_input_name": "remote_input_1",
                "io_type": "CONDITIONING",
                "targets": [{"node_id": "12", "input_name": "positive"}],
            },
            {
                "proxy_input_name": "remote_input_2",
                "io_type": "INT",
                "targets": [{"node_id": "12", "input_name": "seed"}],
            },
        ],
        "boundary_outputs": [
            {"node_id": "12", "io_type": "STRING", "is_list": False},
        ],
        "extra_data": {"client_id": "client-1"},
    }

    response = asyncio.run(
        remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
            payload,
            serialization_module.serialize_node_inputs(
                {
                    "remote_input_0": [model_ref, model_ref, model_ref],
                    "remote_input_1": [conditioning_ref, conditioning_ref, conditioning_ref],
                    "remote_input_2": [10, 11, 12],
                }
            ),
        )
    )

    assert serialization_module.deserialize_node_outputs(response) == (
        ["seed:10", "seed:11", "seed:12"],
    )
    assert observed_calls == [
        (
            "17::item:0",
            ("12",),
            {"remote_input_0": model_ref, "remote_input_1": conditioning_ref, "remote_input_2": 10},
        ),
        (
            "17::item:1",
            ("12",),
            {"remote_input_0": model_ref, "remote_input_1": conditioning_ref, "remote_input_2": 11},
        ),
        (
            "17::item:2",
            ("12",),
            {"remote_input_0": model_ref, "remote_input_1": conditioning_ref, "remote_input_2": 12},
        ),
    ]


def test_implicitly_mapped_subgraph_splits_singleton_wrapped_bridge_ref_lists(
    remote_modal_app_module: Any,
    serialization_module: Any,
    monkeypatch: Any,
) -> None:
    """Comfy output wrappers around bridge-ref lists should still itemize mapped conditioning."""
    first_conditioning_ref = {
        "__comfy_modal_remote_session_bridge_ref__": True,
        "bridge_key": "RSB_first",
        "session_id": "session-1",
        "node_id": "508",
        "output_index": 0,
    }
    second_conditioning_ref = {
        "__comfy_modal_remote_session_bridge_ref__": True,
        "bridge_key": "RSB_second",
        "session_id": "session-1",
        "node_id": "508",
        "output_index": 0,
    }
    observed_calls: list[tuple[str, dict[str, Any]]] = []

    async def fake_invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: bytes) -> bytes:
        hydrated_inputs = serialization_module.deserialize_node_inputs(kwargs_payload)
        observed_calls.append((str(payload["component_id"]), hydrated_inputs))
        return serialization_module.serialize_node_outputs((f"seed:{hydrated_inputs['remote_input_1']}",))

    monkeypatch.setattr(
        remote_modal_app_module,
        "invoke_remote_engine_async",
        fake_invoke_remote_engine_async,
    )

    payload = {
        "payload_kind": "subgraph",
        "component_id": "507",
        "prompt_id": "prompt-1",
        "component_node_ids": ["507"],
        "execute_node_ids": ["507"],
        "subgraph_prompt": {
            "507": {
                "class_type": "KSampler",
                "inputs": {
                    "positive": 0,
                    "seed": 0,
                },
            },
        },
        "boundary_inputs": [
            {
                "proxy_input_name": "phase_bridge_3",
                "io_type": "CONDITIONING",
                "targets": [{"node_id": "507", "input_name": "positive"}],
            },
            {
                "proxy_input_name": "remote_input_1",
                "io_type": "INT",
                "targets": [{"node_id": "507", "input_name": "seed"}],
            },
        ],
        "boundary_outputs": [
            {"node_id": "507", "io_type": "STRING", "is_list": False},
        ],
        "extra_data": {"client_id": "client-1"},
    }

    response = asyncio.run(
        remote_modal_app_module._invoke_implicitly_mapped_subgraph_async(
            payload,
            serialization_module.serialize_node_inputs(
                {
                    "phase_bridge_3": [[first_conditioning_ref, second_conditioning_ref]],
                    "remote_input_1": [10, 11],
                }
            ),
        )
    )

    assert serialization_module.deserialize_node_outputs(response) == (
        ["seed:10", "seed:11"],
    )
    assert observed_calls == [
        (
            "507::item:0",
            {"phase_bridge_3": first_conditioning_ref, "remote_input_1": 10},
        ),
        (
            "507::item:1",
            {"phase_bridge_3": second_conditioning_ref, "remote_input_1": 11},
        ),
    ]


@pytest.mark.parametrize(
    ("module_fixture_name",),
    [
        ("remote_modal_app_module",),
        ("modal_cloud_module",),
    ],
)
def test_trim_subgraph_payload_to_required_nodes_drops_unrelated_mapped_branch(
    request: Any,
    module_fixture_name: str,
) -> None:
    """Static or per-item sub-runs should exclude unrelated nodes from the mapped sibling branch."""
    target_module = request.getfixturevalue(module_fixture_name)
    payload = {
        "component_id": "1::static",
        "component_node_ids": ["1", "2", "3", "7"],
        "subgraph_prompt": {
            "1": {"class_type": "LoadDiffusionModel", "inputs": {}},
            "2": {"class_type": "ModalMapInput", "inputs": {"value": ["remote_input_1", 0]}},
            "3": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "steps": 20}},
            "7": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "latent_image": ["2", 0]}},
        },
        "boundary_inputs": [
            {
                "proxy_input_name": "remote_input_1",
                "targets": [{"node_id": "2", "input_name": "value"}],
            }
        ],
        "boundary_outputs": [
            {"node_id": "3", "output_index": 0, "io_type": "LATENT", "is_list": False},
        ],
        "execute_node_ids": ["3"],
        "mapped_execute_node_ids": ["7"],
        "static_execute_node_ids": ["3"],
    }

    trimmed_payload = target_module._trim_subgraph_payload_to_required_nodes(payload)

    assert trimmed_payload["component_node_ids"] == ["1", "3"]
    assert list(trimmed_payload["subgraph_prompt"].keys()) == ["1", "3"]
    assert trimmed_payload["boundary_inputs"] == []
    assert trimmed_payload["boundary_outputs"] == payload["boundary_outputs"]
    assert trimmed_payload["execute_node_ids"] == ["3"]
    assert trimmed_payload["mapped_execute_node_ids"] == []
    assert trimmed_payload["static_execute_node_ids"] == ["3"]


@pytest.mark.parametrize(
    ("module_fixture_name",),
    [
        ("remote_modal_app_module",),
        ("modal_cloud_module",),
    ],
)
def test_trim_subgraph_payload_to_required_nodes_drops_stale_execute_targets(
    request: Any,
    module_fixture_name: str,
) -> None:
    """Trimmed subgraph payloads should ignore execute targets that are absent from the current prompt."""
    target_module = request.getfixturevalue(module_fixture_name)
    payload = {
        "component_id": "1::static",
        "component_node_ids": ["1", "3"],
        "subgraph_prompt": {
            "1": {"class_type": "LoadDiffusionModel", "inputs": {}},
            "3": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "steps": 20}},
        },
        "boundary_inputs": [],
        "boundary_outputs": [
            {"node_id": "3", "output_index": 0, "io_type": "LATENT", "is_list": False},
        ],
        "execute_node_ids": ["3", "5"],
        "mapped_execute_node_ids": [],
        "static_execute_node_ids": ["3", "5"],
    }

    trimmed_payload = target_module._trim_subgraph_payload_to_required_nodes(payload)

    assert trimmed_payload["component_node_ids"] == ["1", "3"]
    assert list(trimmed_payload["subgraph_prompt"].keys()) == ["1", "3"]
    assert trimmed_payload["execute_node_ids"] == ["3"]
    assert trimmed_payload["mapped_execute_node_ids"] == []
    assert trimmed_payload["static_execute_node_ids"] == ["3"]


def test_consume_remote_payload_stream_suppresses_status_but_keeps_boundary_previews(
    remote_modal_app_module: Any,
    serialization_module: Any,
    monkeypatch: Any,
) -> None:
    """Mapped per-item remote calls should suppress status chatter but still forward previews and lane progress."""
    progress_calls: list[dict[str, Any]] = []
    status_calls: list[dict[str, Any]] = []
    preview_calls: list[dict[str, Any]] = []

    monkeypatch.setattr(
        remote_modal_app_module,
        "_emit_local_modal_progress",
        lambda **kwargs: progress_calls.append(kwargs),
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_emit_local_modal_status",
        lambda **kwargs: status_calls.append(kwargs),
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_emit_local_preview_boundary_output",
        lambda **kwargs: preview_calls.append(kwargs),
    )

    payload = {
        "component_id": "6::item:0",
        "prompt_id": "prompt-1",
        "component_node_ids": ["6", "7"],
        "extra_data": {"client_id": "client-1"},
        "suppress_status_stream": True,
        "mapped_progress_lane_id": "1",
        "mapped_progress_display_node_id": "6",
        "map_item_index": 0,
    }
    stream_events = iter(
        [
            {
                "kind": "progress",
                "event_type": "node_progress",
                "node_id": "7",
                "value": 1.0,
                "max": 4.0,
            },
            {
                "kind": "progress",
                "event_type": "boundary_output",
                "node_id": "7",
                "preview_target_node_ids": ["9"],
                "value": serialization_module.serialize_value(["preview"]),
            },
            {
                "kind": "progress",
                "phase": "executing",
                "active_node_id": "7",
            },
            {
                "kind": "result",
                "outputs": serialization_module.serialize_node_outputs(("done",)),
            },
        ]
    )

    response = remote_modal_app_module._consume_remote_payload_stream(payload, stream_events)

    assert serialization_module.deserialize_node_outputs(response) == ("done",)
    assert progress_calls == [
        {
            "prompt_id": "prompt-1",
            "client_id": "client-1",
            "node_id": "7",
            "value": 1.0,
            "max_value": 4.0,
            "display_node_id": "7",
            "real_node_id": None,
            "lane_id": "1",
            "clear": False,
            "item_index": 0,
            "aggregate_only": False,
        }
    ]
    assert status_calls == []
    assert preview_calls == [
        {
            "prompt_id": "prompt-1",
            "client_id": "client-1",
            "preview_target_node_ids": ["9"],
            "image_value": ["preview"],
        }
    ]


def test_consume_remote_payload_stream_keeps_static_execute_node_progress_when_status_is_suppressed(
    remote_modal_app_module: Any,
    serialization_module: Any,
    monkeypatch: Any,
) -> None:
    """Static sub-runs should still forward real execute-node progress under suppressed status streams."""
    progress_calls: list[dict[str, Any]] = []

    monkeypatch.setattr(
        remote_modal_app_module,
        "_emit_local_modal_progress",
        lambda **kwargs: progress_calls.append(kwargs),
    )

    payload = {
        "component_id": "1::static",
        "prompt_id": "prompt-1",
        "component_node_ids": ["1", "2", "12", "4"],
        "execute_node_ids": ["12", "4"],
        "boundary_outputs": [
            {"node_id": "4", "output_index": 0, "io_type": "LATENT", "is_list": False}
        ],
        "extra_data": {"client_id": "client-1"},
        "suppress_status_stream": True,
    }
    stream_events = iter(
        [
            {
                "kind": "progress",
                "event_type": "node_progress",
                "node_id": "1",
                "display_node_id": "1",
                "real_node_id": "12",
                "value": 5.0,
                "max": 20.0,
            },
            {
                "kind": "progress",
                "event_type": "node_progress",
                "node_id": "2",
                "display_node_id": "2",
                "real_node_id": "2",
                "value": 1.0,
                "max": 10.0,
            },
            {
                "kind": "result",
                "outputs": serialization_module.serialize_node_outputs(("done",)),
            },
        ]
    )

    response = remote_modal_app_module._consume_remote_payload_stream(payload, stream_events)

    assert serialization_module.deserialize_node_outputs(response) == ("done",)
    assert progress_calls == [
        {
            "prompt_id": "prompt-1",
            "client_id": "client-1",
            "node_id": "1",
            "value": 5.0,
            "max_value": 20.0,
            "display_node_id": "1",
            "real_node_id": "12",
            "lane_id": None,
            "clear": False,
            "item_index": None,
            "aggregate_only": False,
        }
    ]


def test_consume_remote_payload_stream_marks_progress_node_ancestors_complete(
    remote_modal_app_module: Any,
    serialization_module: Any,
    monkeypatch: Any,
) -> None:
    """Streamed progress should tell the UI which upstream remote nodes can be completed."""
    progress_calls: list[dict[str, Any]] = []

    monkeypatch.setattr(
        remote_modal_app_module,
        "_emit_local_modal_progress",
        lambda **kwargs: progress_calls.append(kwargs),
    )

    payload = {
        "component_id": "3",
        "prompt_id": "prompt-1",
        "component_node_ids": ["1", "2", "3"],
        "extra_data": {"client_id": "client-1"},
        "subgraph_prompt": {
            "1": {"class_type": "Loader", "inputs": {}},
            "2": {"class_type": "Condition", "inputs": {"model": [[["1", 0]]]}},
            "3": {"class_type": "Sampler", "inputs": {"conditioning": ["2", 0]}},
        },
    }
    stream_events = iter(
        [
            {
                "kind": "progress",
                "event_type": "node_progress",
                "node_id": "3",
                "value": 1.0,
                "max": 4.0,
            },
            {
                "kind": "result",
                "outputs": serialization_module.serialize_node_outputs(("done",)),
            },
        ]
    )

    response = remote_modal_app_module._consume_remote_payload_stream(payload, stream_events)

    assert serialization_module.deserialize_node_outputs(response) == ("done",)
    assert progress_calls[0]["node_id"] == "3"
    assert progress_calls[0]["completed_ancestor_node_ids"] == ["1", "2"]


def test_consume_remote_payload_stream_clears_static_execute_node_progress_on_suppressed_completion(
    remote_modal_app_module: Any,
    serialization_module: Any,
    monkeypatch: Any,
) -> None:
    """Suppressed static sub-runs should emit an explicit clear for lane-less node progress on completion."""
    progress_calls: list[dict[str, Any]] = []

    monkeypatch.setattr(
        remote_modal_app_module,
        "_emit_local_modal_progress",
        lambda **kwargs: progress_calls.append(kwargs),
    )

    payload = {
        "component_id": "1::static",
        "prompt_id": "prompt-1",
        "component_node_ids": ["1", "2", "12", "4"],
        "execute_node_ids": ["12", "4"],
        "boundary_outputs": [
            {"node_id": "4", "output_index": 0, "io_type": "LATENT", "is_list": False}
        ],
        "extra_data": {"client_id": "client-1"},
        "suppress_status_stream": True,
    }
    stream_events = iter(
        [
            {
                "kind": "progress",
                "event_type": "node_progress",
                "node_id": "1",
                "display_node_id": "1",
                "real_node_id": "12",
                "value": 5.0,
                "max": 20.0,
            },
            {
                "kind": "progress",
                "phase": "execution_success",
            },
            {
                "kind": "result",
                "outputs": serialization_module.serialize_node_outputs(("done",)),
            },
        ]
    )

    response = remote_modal_app_module._consume_remote_payload_stream(payload, stream_events)

    assert serialization_module.deserialize_node_outputs(response) == ("done",)
    assert progress_calls == [
        {
            "prompt_id": "prompt-1",
            "client_id": "client-1",
            "node_id": "1",
            "value": 5.0,
            "max_value": 20.0,
            "display_node_id": "1",
            "real_node_id": "12",
            "lane_id": None,
            "clear": False,
            "item_index": None,
            "aggregate_only": False,
        },
        {
            "prompt_id": "prompt-1",
            "client_id": "client-1",
            "node_id": "1",
            "value": 0.0,
            "max_value": 1.0,
            "display_node_id": "1",
            "real_node_id": "12",
            "clear": True,
        },
    ]


def test_consume_remote_payload_stream_filters_static_sibling_ui_events_from_mapped_items(
    remote_modal_app_module: Any,
    serialization_module: Any,
    monkeypatch: Any,
) -> None:
    """Mapped item streams should not forward executed or preview events for static sibling nodes."""
    executed_calls: list[dict[str, Any]] = []
    preview_calls: list[dict[str, Any]] = []

    monkeypatch.setattr(
        remote_modal_app_module,
        "_emit_local_executed_output",
        lambda **kwargs: executed_calls.append(kwargs),
    )
    monkeypatch.setattr(
        remote_modal_app_module,
        "_emit_local_preview_image",
        lambda **kwargs: preview_calls.append(kwargs),
    )

    payload = {
        "component_id": "6::item:0",
        "prompt_id": "prompt-1",
        "component_node_ids": ["3", "6", "7"],
        "execute_node_ids": ["7"],
        "boundary_outputs": [
            {"node_id": "7", "output_index": 0, "io_type": "IMAGE", "is_list": False}
        ],
        "extra_data": {"client_id": "client-1"},
        "suppress_status_stream": True,
        "mapped_progress_lane_id": "0",
        "mapped_progress_display_node_id": "6",
        "map_item_index": 0,
    }
    preview_bytes = serialization_module.serialize_value(b"preview-bytes")
    stream_events = iter(
        [
            {
                "kind": "progress",
                "event_type": "executed",
                "node_id": "3",
                "display_node_id": "3",
                "output": serialization_module.serialize_value({"images": ["static"]}),
            },
            {
                "kind": "progress",
                "event_type": "preview",
                "node_id": "3",
                "display_node_id": "3",
                "image_type": "PNG",
                "image_bytes": preview_bytes,
            },
            {
                "kind": "progress",
                "event_type": "executed",
                "node_id": "7",
                "display_node_id": "7",
                "output": serialization_module.serialize_value({"images": ["mapped"]}),
            },
            {
                "kind": "progress",
                "event_type": "preview",
                "node_id": "7",
                "display_node_id": "7",
                "image_type": "PNG",
                "image_bytes": preview_bytes,
            },
            {
                "kind": "result",
                "outputs": serialization_module.serialize_node_outputs(("done",)),
            },
        ]
    )

    response = remote_modal_app_module._consume_remote_payload_stream(payload, stream_events)

    assert serialization_module.deserialize_node_outputs(response) == ("done",)
    assert executed_calls == [
        {
            "prompt_id": "prompt-1",
            "client_id": "client-1",
            "node_id": "7",
            "display_node_id": "7",
            "output_payload": {"images": ["mapped"]},
        }
    ]
    assert preview_calls == [
        {
            "prompt_id": "prompt-1",
            "client_id": "client-1",
            "node_id": "7",
            "display_node_id": "7",
            "parent_node_id": None,
            "real_node_id": None,
            "image_type": "PNG",
            "image_bytes": b"preview-bytes",
            "max_size": None,
        }
    ]


def test_consume_remote_payload_stream_forwards_cached_node_markers(
    remote_modal_app_module: Any,
    serialization_module: Any,
    monkeypatch: Any,
) -> None:
    """Streamed node-cache hits should mark the matching UI node as cached without fake progress."""
    progress_calls: list[dict[str, Any]] = []

    monkeypatch.setattr(
        remote_modal_app_module,
        "_emit_local_modal_progress",
        lambda **kwargs: progress_calls.append(kwargs),
    )

    payload = {
        "component_id": "1::static",
        "prompt_id": "prompt-1",
        "component_node_ids": ["1", "2", "12", "4"],
        "execute_node_ids": ["12", "4"],
        "extra_data": {"client_id": "client-1"},
        "suppress_status_stream": True,
    }
    stream_events = iter(
        [
            {
                "kind": "progress",
                "event_type": "node_cached",
                "node_id": "1",
                "display_node_id": "1",
                "real_node_id": "12",
            },
            {
                "kind": "result",
                "outputs": serialization_module.serialize_node_outputs(("done",)),
            },
        ]
    )

    response = remote_modal_app_module._consume_remote_payload_stream(payload, stream_events)

    assert serialization_module.deserialize_node_outputs(response) == ("done",)
    assert progress_calls == [
        {
            "prompt_id": "prompt-1",
            "client_id": "client-1",
            "node_id": "1",
            "value": 0.0,
            "max_value": 1.0,
            "display_node_id": "1",
            "real_node_id": "12",
            "cached_hit": True,
        }
    ]


def test_modal_cloud_initializes_remote_comfy_runtime_once_per_custom_node_root(
    modal_cloud_module: Any,
    monkeypatch: Any,
    tmp_path: Any,
) -> None:
    """The remote runtime should load built-in extras once and custom bundles per extracted root."""
    init_calls: list[tuple[Any, ...]] = []
    folder_path_calls: list[tuple[str, str, bool]] = []

    fake_nodes_module = types.SimpleNamespace(
        NODE_CLASS_MAPPINGS={},
        NODE_DISPLAY_NAME_MAPPINGS={},
    )

    async def fake_init_extra_nodes(init_custom_nodes: bool = True, init_api_nodes: bool = True) -> None:
        init_calls.append(("extra", init_custom_nodes, init_api_nodes))

    async def fake_init_external_custom_nodes() -> None:
        init_calls.append(("external",))

    fake_nodes_module.init_extra_nodes = fake_init_extra_nodes
    fake_nodes_module.init_external_custom_nodes = fake_init_external_custom_nodes

    fake_folder_paths_module = types.SimpleNamespace(
        add_model_folder_path=lambda folder_name, full_folder_path, is_default=False: folder_path_calls.append(
            (folder_name, full_folder_path, is_default)
        )
    )

    monkeypatch.setitem(sys.modules, "nodes", fake_nodes_module)
    monkeypatch.setitem(sys.modules, "folder_paths", fake_folder_paths_module)

    original_base_initialized = modal_cloud_module._COMFY_RUNTIME_BASE_INITIALIZED
    original_custom_node_roots = set(modal_cloud_module._COMFY_RUNTIME_CUSTOM_NODE_ROOTS)
    modal_cloud_module._COMFY_RUNTIME_BASE_INITIALIZED = False
    modal_cloud_module._COMFY_RUNTIME_CUSTOM_NODE_ROOTS.clear()
    try:
        custom_nodes_root = tmp_path / "custom_nodes"
        custom_nodes_root.mkdir()

        modal_cloud_module._ensure_comfy_runtime_initialized(None)
        modal_cloud_module._ensure_comfy_runtime_initialized(custom_nodes_root)
        modal_cloud_module._ensure_comfy_runtime_initialized(custom_nodes_root)
    finally:
        modal_cloud_module._COMFY_RUNTIME_BASE_INITIALIZED = original_base_initialized
        modal_cloud_module._COMFY_RUNTIME_CUSTOM_NODE_ROOTS.clear()
        modal_cloud_module._COMFY_RUNTIME_CUSTOM_NODE_ROOTS.update(original_custom_node_roots)

    assert init_calls == [("extra", False, True), ("external",)]
    assert folder_path_calls == [("custom_nodes", str(custom_nodes_root), True)]


def test_modal_cloud_uses_comfy_prompt_executor_cache_defaults(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """The remote worker should mirror ComfyUI's prompt executor cache configuration."""
    fake_args = types.SimpleNamespace(
        cache_lru=0,
        cache_ram=4.0,
        cache_none=False,
    )
    fake_cli_args_module = types.SimpleNamespace(args=fake_args)
    fake_execution_module = types.SimpleNamespace(
        CacheType=types.SimpleNamespace(
            CLASSIC="classic",
            LRU="lru",
            RAM_PRESSURE="ram-pressure",
            NONE="none",
        )
    )
    monkeypatch.setitem(sys.modules, "comfy.cli_args", fake_cli_args_module)

    cache_type, cache_args = modal_cloud_module._prompt_executor_cache_config(fake_execution_module)

    assert cache_type == "ram-pressure"
    assert cache_args == {"lru": 0, "ram": 4.0}


def test_modal_cloud_uses_current_comfy_ram_pressure_cache_defaults(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Current list-valued cache arguments should use ComfyUI's calculated RAM thresholds."""
    fake_args = types.SimpleNamespace(
        cache_classic=False,
        cache_lru=0,
        cache_ram=[],
        cache_none=False,
    )
    fake_cli_args_module = types.SimpleNamespace(args=fake_args)
    fake_model_management_module = types.SimpleNamespace(total_ram=64 * 1024)
    fake_execution_module = types.SimpleNamespace(
        CacheType=types.SimpleNamespace(
            CLASSIC="classic",
            LRU="lru",
            RAM_PRESSURE="ram-pressure",
            NONE="none",
        )
    )
    monkeypatch.setitem(sys.modules, "comfy.cli_args", fake_cli_args_module)
    monkeypatch.setitem(
        sys.modules,
        "comfy.model_management",
        fake_model_management_module,
    )

    cache_type, cache_args = modal_cloud_module._prompt_executor_cache_config(fake_execution_module)

    assert cache_type == "ram-pressure"
    assert cache_args == {"lru": 0, "ram": 6.4, "ram_inactive": 64.0}


def test_modal_cloud_uses_current_comfy_explicit_ram_pressure_thresholds(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Two current cache-ram values should flow into active and inactive thresholds."""
    fake_args = types.SimpleNamespace(
        cache_classic=False,
        cache_lru=0,
        cache_ram=[3.0, 20.0],
        cache_none=False,
    )
    fake_cli_args_module = types.SimpleNamespace(args=fake_args)
    fake_model_management_module = types.SimpleNamespace(total_ram=64 * 1024)
    fake_execution_module = types.SimpleNamespace(
        CacheType=types.SimpleNamespace(
            CLASSIC="classic",
            LRU="lru",
            RAM_PRESSURE="ram-pressure",
            NONE="none",
        )
    )
    monkeypatch.setitem(sys.modules, "comfy.cli_args", fake_cli_args_module)
    monkeypatch.setitem(
        sys.modules,
        "comfy.model_management",
        fake_model_management_module,
    )

    cache_type, cache_args = modal_cloud_module._prompt_executor_cache_config(fake_execution_module)

    assert cache_type == "ram-pressure"
    assert cache_args == {"lru": 0, "ram": 3.0, "ram_inactive": 20.0}


def test_modal_cloud_awaits_async_prompt_executor_api(
    modal_cloud_module: Any,
) -> None:
    """An asynchronous PromptExecutor compatibility API should finish before cache reads."""
    observed_calls: list[tuple[dict[str, Any], str, dict[str, Any], list[str]]] = []

    class FakeExecutor:
        """Minimal executor exposing a coroutine-based execute method."""

        async def execute(
            self,
            *,
            prompt: dict[str, Any],
            prompt_id: str,
            extra_data: dict[str, Any],
            execute_outputs: list[str],
        ) -> None:
            """Record execution after yielding once."""
            await asyncio.sleep(0)
            observed_calls.append((prompt, prompt_id, extra_data, execute_outputs))

    modal_cloud_module._execute_prompt_executor_compat(
        FakeExecutor(),
        prompt={"1": {"class_type": "Example", "inputs": {}}},
        prompt_id="prompt-1",
        extra_data={"client_id": "client-1"},
        execute_outputs=["1"],
    )

    assert observed_calls == [
        (
            {"1": {"class_type": "Example", "inputs": {}}},
            "prompt-1",
            {"client_id": "client-1"},
            ["1"],
        )
    ]


def test_modal_cloud_sync_cache_reads_use_current_get_local_api(
    modal_cloud_module: Any,
) -> None:
    """Synchronous progress paths should avoid creating current cache get coroutines."""
    cache_entry = types.SimpleNamespace(outputs=[[512, 512]])

    class FakeOutputsCache:
        """Current cache double with async get and synchronous get_local."""

        async def get(self, node_id: str) -> Any:
            """Fail if the synchronous compatibility path calls async get."""
            del node_id
            raise AssertionError("async get should not be created")

        def get_local(self, node_id: str) -> Any:
            """Return the already materialized local cache entry."""
            assert node_id == "115"
            return cache_entry

    resolved_entry = modal_cloud_module._prompt_executor_cache_get_sync(
        FakeOutputsCache(),
        "115",
    )

    assert resolved_entry is cache_entry


def test_modal_cloud_class_options_do_not_use_deprecated_concurrency_flag(
    modal_cloud_module: Any,
) -> None:
    """The deployed Modal class options should avoid deprecated concurrency flags."""
    fake_settings = types.SimpleNamespace(
        modal_gpu="A100",
        remote_storage_root="/vol/data",
        scaledown_window_seconds=60,
        min_containers=0,
        max_containers=4,
        buffer_containers=1,
        execution_timeout_seconds=1800,
        startup_timeout_seconds=600,
        enable_memory_snapshot=True,
        enable_gpu_memory_snapshot=False,
    )

    options = modal_cloud_module._remote_engine_cls_options(
        fake_settings,
        vol=object(),
        image=object(),
    )

    assert "allow_concurrent_inputs" not in options
    assert options["max_containers"] == 4
    assert options["buffer_containers"] == 1
    assert options["timeout"] == 1800
    assert options["startup_timeout"] == 600
    module_source = Path(modal_cloud_module.__file__).read_text(encoding="utf-8")
    assert "@modal.concurrent(max_inputs=1)" in module_source


def test_modal_cloud_registered_execution_preserves_pre_start_interrupt_flags(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Remote execution registration should not delete a valid cancel flag on entry."""

    class FakeInterruptFlags:
        """Simple Modal Dict double that records cleared keys."""

        def __init__(self) -> None:
            """Initialize captured pop calls."""
            self.pop_calls: list[tuple[str, Any]] = []

        def pop(self, key: str, default: Any = None) -> Any:
            """Record one cleared interrupt flag."""
            self.pop_calls.append((key, default))
            return None

    interrupt_flags = FakeInterruptFlags()
    monkeypatch.setattr(modal_cloud_module, "modal", object())
    monkeypatch.setattr(modal_cloud_module, "interrupt_flags", interrupt_flags, raising=False)

    with modal_cloud_module._registered_remote_execution(
        {"prompt_id": "prompt-1", "component_id": "component-2"}
    ) as execution_control:
        assert execution_control.interrupt_flag_key == "prompt-1:component-2"
        assert interrupt_flags.pop_calls == []

    assert interrupt_flags.pop_calls == [
        ("prompt-1:component-2", None),
    ]


def test_modal_cloud_interrupt_monitor_consumes_shared_cancel_flag(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """The remote interrupt monitor should trip when the shared Modal Dict flag appears."""

    class FakeInterruptFlags:
        """Simple Modal Dict double that exposes one shared cancel flag."""

        def __init__(self) -> None:
            """Initialize the backing key set."""
            self.keys = {"prompt-1:component-2"}
            self.contains_calls = 0
            self.pop_calls: list[tuple[str, Any]] = []

        def contains(self, key: str) -> bool:
            """Report whether the shared interrupt flag exists."""
            self.contains_calls += 1
            return key in self.keys

        def pop(self, key: str, default: Any = None) -> Any:
            """Remove the shared interrupt flag once consumed."""
            self.pop_calls.append((key, default))
            self.keys.discard(key)
            return None

    interrupt_calls: list[str] = []
    cancellation_event = threading.Event()
    monkeypatch.setitem(
        sys.modules,
        "nodes",
        types.SimpleNamespace(interrupt_processing=lambda: interrupt_calls.append("interrupt")),
    )

    with modal_cloud_module._temporary_remote_interrupt_monitor(
        "component-2",
        cancellation_event,
        interrupt_store=FakeInterruptFlags(),
        interrupt_flag_key="prompt-1:component-2",
    ):
        deadline = time.time() + 1.0
        while not interrupt_calls and time.time() < deadline:
            time.sleep(0.01)

    assert interrupt_calls == ["interrupt"]
    assert cancellation_event.is_set()


def test_modal_cloud_interrupt_monitor_schedules_container_restart_after_cancel(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Remote cancellation should restart the worker if the prompt ignores ComfyUI interrupt."""

    class FakeInterruptFlags:
        """Simple Modal Dict double that exposes one shared cancel flag."""

        def __init__(self) -> None:
            """Initialize the backing key set."""
            self.keys = {"prompt-1:component-2"}

        def contains(self, key: str) -> bool:
            """Report whether the shared interrupt flag exists."""
            return key in self.keys

        def pop(self, key: str, default: Any = None) -> Any:
            """Remove the shared interrupt flag once consumed."""
            del default
            self.keys.discard(key)
            return None

    scheduled_restarts: list[dict[str, Any]] = []
    interrupt_calls: list[str] = []
    cancellation_event = threading.Event()

    def fake_schedule_process_exit_unless_cancelled(**kwargs: Any) -> None:
        """Record the delayed restart request without exiting pytest."""
        scheduled_restarts.append(kwargs)

    monkeypatch.setenv("COMFY_MODAL_REMOTE_CANCEL_RESTART_SECONDS", "0.25")
    modal_cloud_module.get_settings.cache_clear()
    monkeypatch.setattr(modal_cloud_module, "_is_modal_container_runtime", lambda: True)
    monkeypatch.setattr(
        modal_cloud_module,
        "_schedule_process_exit_unless_cancelled",
        fake_schedule_process_exit_unless_cancelled,
    )
    monkeypatch.setitem(
        sys.modules,
        "nodes",
        types.SimpleNamespace(interrupt_processing=lambda: interrupt_calls.append("interrupt")),
    )

    try:
        with modal_cloud_module._temporary_remote_interrupt_monitor(
            "component-2",
            cancellation_event,
            interrupt_store=FakeInterruptFlags(),
            interrupt_flag_key="prompt-1:component-2",
        ):
            deadline = time.time() + 1.0
            while not scheduled_restarts and time.time() < deadline:
                time.sleep(0.01)
    finally:
        modal_cloud_module.get_settings.cache_clear()

    assert interrupt_calls == ["interrupt"]
    assert len(scheduled_restarts) == 1
    assert scheduled_restarts[0]["delay_seconds"] == 0.25
    assert scheduled_restarts[0]["exit_code"] == 0
    assert scheduled_restarts[0]["cancel_event"].is_set()
    assert "component=component-2" in scheduled_restarts[0]["reason"]


def test_modal_cloud_interrupt_monitor_ignores_modal_client_shutdown(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """The remote interrupt monitor should exit quietly if Modal tears down the client first."""

    class FakeClientClosedError(RuntimeError):
        """Stand-in for `modal.exception.ClientClosed`."""

    class FakeInterruptFlags:
        """Simple Modal Dict double that fails once the client is already shutting down."""

        def __init__(self) -> None:
            """Initialize the poll counter."""
            self.contains_calls = 0
            self.pop_calls: list[tuple[str, Any]] = []

        def contains(self, key: str) -> bool:
            """Raise the same client-shutdown error Modal emits during teardown."""
            del key
            self.contains_calls += 1
            raise FakeClientClosedError("client closed")

        def pop(self, key: str, default: Any = None) -> Any:
            """Record unexpected cleanup attempts."""
            self.pop_calls.append((key, default))
            return None

    interrupt_calls: list[str] = []
    thread_exceptions: list[BaseException] = []
    cancellation_event = threading.Event()
    fake_modal_module = types.ModuleType("modal")
    fake_modal_exception_module = types.ModuleType("modal.exception")
    fake_modal_exception_module.ClientClosed = FakeClientClosedError
    fake_modal_module.exception = fake_modal_exception_module
    monkeypatch.setitem(sys.modules, "modal", fake_modal_module)
    monkeypatch.setitem(sys.modules, "modal.exception", fake_modal_exception_module)
    monkeypatch.setitem(
        sys.modules,
        "nodes",
        types.SimpleNamespace(interrupt_processing=lambda: interrupt_calls.append("interrupt")),
    )
    monkeypatch.setattr(
        threading,
        "excepthook",
        lambda args: thread_exceptions.append(args.exc_value),
    )
    interrupt_store = FakeInterruptFlags()

    with modal_cloud_module._temporary_remote_interrupt_monitor(
        "component-2",
        cancellation_event,
        interrupt_store=interrupt_store,
        interrupt_flag_key="prompt-1:component-2",
    ):
        deadline = time.time() + 1.0
        while interrupt_store.contains_calls == 0 and time.time() < deadline:
            time.sleep(0.01)

    assert interrupt_store.contains_calls >= 1
    assert interrupt_store.pop_calls == []
    assert interrupt_calls == []
    assert not cancellation_event.is_set()
    assert thread_exceptions == []


def test_modal_cloud_reuses_prompt_executor_for_same_cache_scope(
    modal_cloud_module: Any,
    tmp_path: Path,
) -> None:
    """Warm-container subgraph runs should reuse one PromptExecutor per cache scope."""

    class FakePromptExecutor:
        """Simple PromptExecutor double that records how many instances were created."""

        instances_created = 0

        def __init__(self, server: Any, cache_type: Any = False, cache_args: Any = None) -> None:
            """Capture initialization state for later assertions."""
            type(self).instances_created += 1
            self.server = server
            self.cache_type = cache_type
            self.cache_args = cache_args
            self.status_messages = [("stale", {})]
            self.success = False
            self.history_result = {"stale": True}

    fake_execution_module = types.SimpleNamespace(PromptExecutor=FakePromptExecutor)
    first_server = types.SimpleNamespace(client_id="first", last_node_id="node-1")
    second_server = types.SimpleNamespace(client_id="second", last_node_id="node-2")

    original_states = dict(modal_cloud_module._PROMPT_EXECUTOR_STATES)
    modal_cloud_module._PROMPT_EXECUTOR_STATES.clear()
    try:
        first_state = modal_cloud_module._get_or_create_prompt_executor_state(
            execution=fake_execution_module,
            prompt_server=first_server,
            cache_type="classic",
            cache_args={"lru": 0, "ram": 0.0},
            custom_nodes_root=tmp_path / "bundle-a",
        )
        modal_cloud_module._reset_prompt_executor_request_state(first_state.executor, first_server)
        second_state = modal_cloud_module._get_or_create_prompt_executor_state(
            execution=fake_execution_module,
            prompt_server=second_server,
            cache_type="classic",
            cache_args={"lru": 0, "ram": 0.0},
            custom_nodes_root=tmp_path / "bundle-a",
        )
        modal_cloud_module._reset_prompt_executor_request_state(second_state.executor, second_server)
    finally:
        modal_cloud_module._PROMPT_EXECUTOR_STATES.clear()
        modal_cloud_module._PROMPT_EXECUTOR_STATES.update(original_states)

    assert FakePromptExecutor.instances_created == 1
    assert first_state is second_state
    assert second_state.executor.server is second_server
    assert second_state.executor.status_messages == []
    assert second_state.executor.success is True
    assert second_state.executor.history_result == {}
    assert second_server.client_id is None
    assert second_server.last_node_id is None


def test_modal_cloud_separates_prompt_executor_cache_scopes_by_custom_nodes_root(
    modal_cloud_module: Any,
    tmp_path: Path,
) -> None:
    """Different custom-node bundle roots should not share a PromptExecutor cache scope."""

    class FakePromptExecutor:
        """Simple PromptExecutor double used to count cache-scope creations."""

        instances_created = 0

        def __init__(self, server: Any, cache_type: Any = False, cache_args: Any = None) -> None:
            """Capture initialization state for later assertions."""
            type(self).instances_created += 1
            self.server = server
            self.cache_type = cache_type
            self.cache_args = cache_args
            self.status_messages = []
            self.success = True
            self.history_result = {}

    fake_execution_module = types.SimpleNamespace(PromptExecutor=FakePromptExecutor)

    original_states = dict(modal_cloud_module._PROMPT_EXECUTOR_STATES)
    modal_cloud_module._PROMPT_EXECUTOR_STATES.clear()
    try:
        first_state = modal_cloud_module._get_or_create_prompt_executor_state(
            execution=fake_execution_module,
            prompt_server=types.SimpleNamespace(client_id=None, last_node_id=None),
            cache_type="classic",
            cache_args={"lru": 0, "ram": 0.0},
            custom_nodes_root=tmp_path / "bundle-a",
        )
        second_state = modal_cloud_module._get_or_create_prompt_executor_state(
            execution=fake_execution_module,
            prompt_server=types.SimpleNamespace(client_id=None, last_node_id=None),
            cache_type="classic",
            cache_args={"lru": 0, "ram": 0.0},
            custom_nodes_root=tmp_path / "bundle-b",
        )
    finally:
        modal_cloud_module._PROMPT_EXECUTOR_STATES.clear()
        modal_cloud_module._PROMPT_EXECUTOR_STATES.update(original_states)

    assert FakePromptExecutor.instances_created == 2
    assert first_state is not second_state


class _PersistentCacheNode:
    """Simple node used to verify persisted node-cache reuse across prompt runs."""

    RETURN_TYPES = ("INT",)
    RETURN_NAMES = ("value",)
    OUTPUT_IS_LIST = (False,)
    FUNCTION = "run"
    invocation_count = 0

    @classmethod
    def INPUT_TYPES(cls) -> dict[str, dict[str, tuple[str]]]:
        """Return the minimal V1 schema needed for cache-key generation."""
        return {"required": {"value": ("INT",)}}

    def run(self, value: int) -> tuple[int]:
        """Count real executions so persisted cache hits are visible to the test."""
        type(self).invocation_count += 1
        return (value + 1,)


def test_modal_cloud_serializes_only_small_transport_safe_node_outputs(
    modal_cloud_module: Any,
) -> None:
    """Persisted node-cache records should keep small tensor outputs and skip oversized ones."""
    import torch

    execution = modal_cloud_module._load_execution_module()
    small_entry = execution.CacheEntry(ui=None, outputs=[[torch.zeros((8,), dtype=torch.float32)]])
    large_entry = execution.CacheEntry(ui=None, outputs=[[torch.zeros((512,), dtype=torch.float32)]])

    small_record = modal_cloud_module._serialize_node_output_cache_entry(
        small_entry,
        max_bytes=1024,
    )
    large_record = modal_cloud_module._serialize_node_output_cache_entry(
        large_entry,
        max_bytes=1024,
    )

    assert small_record is not None
    assert large_record is None


def test_modal_cloud_awaits_async_node_output_cache_writes(
    modal_cloud_module: Any,
) -> None:
    """Persisted cache writes should use Modal's non-blocking Dict interface."""
    observed_writes: list[tuple[str, dict[str, Any]]] = []

    class AsyncPut:
        """Expose the callable shape used by Modal's synchronized methods."""

        async def aio(self, cache_key: str, record: dict[str, Any]) -> None:
            """Record one asynchronous cache write."""
            observed_writes.append((cache_key, record))

    class AsyncCacheStore:
        """Reject synchronous assignment while exposing an asynchronous put method."""

        put = AsyncPut()

        def __setitem__(self, cache_key: str, record: dict[str, Any]) -> None:
            """Fail if persistence regresses to Modal's blocking operator interface."""
            del cache_key, record
            raise AssertionError("cache write must use put.aio")

    record = {"version": 1, "outputs_zlib": b"payload"}

    asyncio.run(
        modal_cloud_module._node_output_cache_store_put(
            AsyncCacheStore(),
            "NC_example",
            record,
        )
    )

    assert observed_writes == [("NC_example", record)]


def test_modal_cloud_restores_persisted_node_cache_across_prompt_executor_instances(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """A fresh PromptExecutor cache should round-trip persisted node outputs through Modal Dict state."""
    monkeypatch.setitem(sys.modules, "torchsde", types.ModuleType("torchsde"))
    modal_cloud_module._ensure_comfy_runtime_initialized(None)
    import comfy_execution.caching as comfy_caching

    execution = modal_cloud_module._load_execution_module()
    nodes_module = modal_cloud_module._load_nodes_module()
    cache_store: dict[str, Any] = {}
    prompt = {
        "node_1": {
            "class_type": "PersistentCacheNode",
            "inputs": {"value": 4},
            "_meta": {},
        }
    }

    _PersistentCacheNode.invocation_count = 0
    monkeypatch.setitem(nodes_module.NODE_CLASS_MAPPINGS, "PersistentCacheNode", _PersistentCacheNode)
    monkeypatch.setitem(
        comfy_caching.nodes.NODE_CLASS_MAPPINGS,
        "PersistentCacheNode",
        _PersistentCacheNode,
    )
    monkeypatch.setitem(
        nodes_module.NODE_DISPLAY_NAME_MAPPINGS,
        "PersistentCacheNode",
        "PersistentCacheNode",
    )
    cache_entry = execution.CacheEntry(ui={"output": {"value": [5]}}, outputs=[[5]])
    first_executor = execution.PromptExecutor(
        modal_cloud_module._NullPromptServer(),
        cache_type=execution.CacheType.CLASSIC,
        cache_args={"lru": 0, "ram": 0.0},
    )
    restored_first = asyncio.run(
        modal_cloud_module._restore_persisted_node_output_cache_entries(
            execution,
            first_executor,
            prompt_id="prompt-a",
            prompt=copy.deepcopy(prompt),
            cache_store=cache_store,
        )
    )
    asyncio.run(
        modal_cloud_module._await_maybe(
            first_executor.caches.outputs.set("node_1", cache_entry)
        )
    )
    persisted_nodes = asyncio.run(
        modal_cloud_module._persist_node_output_cache_entries(
            first_executor,
            prompt=copy.deepcopy(prompt),
            cache_store=cache_store,
        )
    )

    second_executor = execution.PromptExecutor(
        modal_cloud_module._NullPromptServer(),
        cache_type=execution.CacheType.CLASSIC,
        cache_args={"lru": 0, "ram": 0.0},
    )
    restored_cache_keys_by_node_id: dict[str, str] = {}
    restored_second = asyncio.run(
        modal_cloud_module._restore_persisted_node_output_cache_entries(
            execution,
            second_executor,
            prompt_id="prompt-b",
            prompt=copy.deepcopy(prompt),
            cache_store=cache_store,
            restored_cache_keys_by_node_id=restored_cache_keys_by_node_id,
        )
    )
    restored_entry = asyncio.run(
        modal_cloud_module._await_maybe(second_executor.caches.outputs.get("node_1"))
    )

    assert restored_first == []
    assert persisted_nodes == ["node_1"]
    assert restored_second == ["node_1"]
    assert restored_cache_keys_by_node_id == {"node_1": next(iter(cache_store))}
    assert list(cache_store) and all(key.startswith("NC_") for key in cache_store)
    assert restored_entry == cache_entry


def test_modal_cloud_installs_persisted_cache_restore_after_live_set_prompt(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Persisted-cache restore should run after PromptExecutor prepares the active outputs cache."""

    class FakeOutputsCache:
        """Minimal outputs-cache stub with a mutable cache-key-set marker."""

        def __init__(self) -> None:
            """Initialize the fake cache-key-set marker."""
            self.cache_key_set = None

        async def set_prompt(self, dynprompt: Any, node_ids: Any, is_changed_cache: Any) -> None:
            """Simulate ComfyUI assigning the live cache-key set during prompt setup."""
            del dynprompt, node_ids, is_changed_cache
            self.cache_key_set = "live-cache-key-set"

    outputs_cache = FakeOutputsCache()
    executor = types.SimpleNamespace(caches=types.SimpleNamespace(outputs=outputs_cache))
    observed_events: list[tuple[str, Any]] = []

    async def fake_restore(
        execution: Any,
        prepared_outputs_cache: Any,
        *,
        prompt: dict[str, Any],
        cache_store: Any,
        required_materialized_node_ids: Any = None,
        restored_cache_keys_by_node_id: dict[str, str] | None = None,
    ) -> list[str]:
        """Record the cache-key-set marker visible at restore time."""
        del execution, required_materialized_node_ids
        if restored_cache_keys_by_node_id is not None:
            restored_cache_keys_by_node_id["12"] = "NC_example"
        observed_events.append(
            (
                "restore",
                prepared_outputs_cache.cache_key_set,
                tuple(sorted(prompt)),
                cache_store,
            )
        )
        return ["12"]

    monkeypatch.setattr(
        modal_cloud_module,
        "_restore_persisted_node_output_cache_entries_into_prepared_cache",
        fake_restore,
    )

    restore_state = (
        modal_cloud_module._install_prompt_executor_persisted_cache_restore(
            object(),
            executor,
            component_id="component-1",
            prompt={"12": {"class_type": "PersistentCacheNode", "inputs": {}}},
            cache_store={"NC_example": {"version": 1}},
        )
    )

    try:
        asyncio.run(outputs_cache.set_prompt(object(), ["12"], object()))
    finally:
        restore_state.restore_original_method()

    assert restore_state.restored_node_ids == ["12"]
    assert restore_state.restored_cache_keys_by_node_id == {"12": "NC_example"}
    assert observed_events == [
        ("restore", "live-cache-key-set", ("12",), {"NC_example": {"version": 1}})
    ]
    assert outputs_cache.set_prompt.__func__ is FakeOutputsCache.set_prompt


def test_modal_cloud_skips_rewriting_restored_distributed_cache_entries(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Persist should skip distributed cache entries that were restored unchanged this run."""
    monkeypatch.setitem(sys.modules, "torchsde", types.ModuleType("torchsde"))
    modal_cloud_module._ensure_comfy_runtime_initialized(None)

    execution = modal_cloud_module._load_execution_module()
    cache_entry = execution.CacheEntry(ui={"output": {"value": [5]}}, outputs=[[5]])
    cache_key = "NC_existing"
    cache_store: dict[str, Any] = {cache_key: {"version": 1, "outputs_zlib": b"old"}}
    observed_logs: list[tuple[Any, ...]] = []

    class FakeOutputsCache:
        """Minimal outputs cache stub for persist-phase tests."""

        def __init__(self) -> None:
            """Populate one persistent cache entry."""
            self.cache_key_set = object()

        async def get(self, node_id: str) -> Any:
            """Return the prepared cache entry for the target node only."""
            if node_id == "node_1":
                return cache_entry
            return None

    executor = types.SimpleNamespace(caches=types.SimpleNamespace(outputs=FakeOutputsCache()))

    monkeypatch.setattr(
        modal_cloud_module,
        "_node_output_cache_key_from_key_set_sync",
        lambda cache_key_set, node_id: cache_key if node_id == "node_1" else None,
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_emit_cloud_info",
        lambda message, *args: observed_logs.append((message, *args)),
    )

    persisted_nodes = asyncio.run(
        modal_cloud_module._persist_node_output_cache_entries(
            executor,
            prompt={"node_1": {"class_type": "PersistentCacheNode", "inputs": {"value": 4}}},
            cache_store=cache_store,
            restored_cache_keys_by_node_id={"node_1": cache_key},
        )
    )

    assert persisted_nodes == []
    assert cache_store == {cache_key: {"version": 1, "outputs_zlib": b"old"}}
    assert observed_logs[-1] == (
        "Node output cache write node=%s key_prefix=%s result=skip reason=restored-hit",
        "node_1",
        "NC_existing",
    )


def test_modal_cloud_skips_restored_hit_before_cache_entry_serialization(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Persist should avoid cache-entry serialization entirely for restored-hit keys."""
    cache_entry = object()

    class FakeOutputsCache:
        """Minimal outputs cache stub for persist fast-path tests."""

        def __init__(self) -> None:
            """Expose a cache-key-set marker for sync key generation."""
            self.cache_key_set = object()

        def get(self, node_id: str) -> Any:
            """Return a placeholder cache entry for the target node."""
            if node_id == "node_1":
                return cache_entry
            return None

    executor = types.SimpleNamespace(caches=types.SimpleNamespace(outputs=FakeOutputsCache()))
    monkeypatch.setattr(
        modal_cloud_module,
        "_node_output_cache_key_from_key_set_sync",
        lambda cache_key_set, node_id: "NC_existing" if node_id == "node_1" else None,
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_serialize_node_output_cache_entry",
        lambda cache_entry, max_bytes: (_ for _ in ()).throw(
            AssertionError("restored-hit should skip before serialization")
        ),
    )

    persisted_nodes = asyncio.run(
        modal_cloud_module._persist_node_output_cache_entries(
            executor,
            prompt={"node_1": {"class_type": "PersistentCacheNode", "inputs": {"value": 4}}},
            cache_store={},
            restored_cache_keys_by_node_id={"node_1": "NC_existing"},
        )
    )

    assert persisted_nodes == []


def test_modal_cloud_restores_persisted_node_cache_entries_in_parallel(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Prepared-cache restore should overlap independent distributed lookups."""
    in_flight_gets = 0
    max_in_flight_gets = 0
    restored_values: dict[str, Any] = {}

    class FakeOutputsCache:
        """Minimal outputs cache stub for restore concurrency tests."""

        def __init__(self) -> None:
            """Expose the cache-key-set marker read by restore."""
            self.cache_key_set = object()

        async def get(self, node_id: str) -> Any:
            """Return any previously restored entry."""
            return restored_values.get(node_id)

        async def set(self, node_id: str, cache_entry: Any) -> None:
            """Record one restored cache entry."""
            restored_values[node_id] = cache_entry

    async def fake_key_from_key_set(cache_key_set: Any, node_id: str) -> str:
        """Yield briefly so multiple node lookups can queue together."""
        del cache_key_set
        await asyncio.sleep(0)
        return f"NC_{node_id}"

    async def fake_store_get(cache_store: Any, cache_key: str) -> Any:
        """Track how many distributed cache reads overlap."""
        nonlocal in_flight_gets, max_in_flight_gets
        del cache_store
        in_flight_gets += 1
        max_in_flight_gets = max(max_in_flight_gets, in_flight_gets)
        await asyncio.sleep(0.01)
        in_flight_gets -= 1
        return {"cache_key": cache_key}

    monkeypatch.setattr(
        modal_cloud_module,
        "_node_output_cache_key_from_key_set_async",
        fake_key_from_key_set,
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_node_output_cache_store_get",
        fake_store_get,
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_deserialize_node_output_cache_entry",
        lambda execution, record: record,
    )
    monkeypatch.setattr(modal_cloud_module, "_emit_cloud_info", lambda *args: None)

    outputs_cache = FakeOutputsCache()
    restored_node_ids = asyncio.run(
        modal_cloud_module._restore_persisted_node_output_cache_entries_into_prepared_cache(
            object(),
            outputs_cache,
            prompt={
                "node_1": {"class_type": "PersistentCacheNode", "inputs": {"value": 1}},
                "node_2": {"class_type": "PersistentCacheNode", "inputs": {"value": 2}},
                "node_3": {"class_type": "PersistentCacheNode", "inputs": {"value": 3}},
            },
            cache_store={},
        )
    )

    assert restored_node_ids == ["node_1", "node_2", "node_3"]
    assert max_in_flight_gets >= 2
    assert restored_values == {
        "node_1": {"cache_key": "NC_node_1"},
        "node_2": {"cache_key": "NC_node_2"},
        "node_3": {"cache_key": "NC_node_3"},
    }


def test_modal_cloud_skips_downstream_cache_hit_when_boundary_ancestor_is_missing(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """A downstream distributed hit must not bypass a missing required boundary producer."""
    restored_values: dict[str, Any] = {}
    observed_logs: list[tuple[Any, ...]] = []

    class FakeOutputsCache:
        """Minimal outputs cache stub for boundary-aware restore tests."""

        def __init__(self) -> None:
            """Expose the cache-key-set marker read by restore."""
            self.cache_key_set = object()

        def get(self, node_id: str) -> Any:
            """Return any previously restored entry."""
            return restored_values.get(node_id)

        def set(self, node_id: str, cache_entry: Any) -> None:
            """Record one restored cache entry."""
            restored_values[node_id] = cache_entry

    async def fake_key_from_key_set(cache_key_set: Any, node_id: str) -> str:
        """Return stable fake distributed cache keys."""
        del cache_key_set
        return f"NC_{node_id}"

    async def fake_store_get(cache_store: Any, cache_key: str) -> Any:
        """Only the downstream execute target has a distributed cache hit."""
        del cache_store
        if cache_key == "NC_11":
            return {"cache_key": cache_key}
        return None

    monkeypatch.setattr(
        modal_cloud_module,
        "_node_output_cache_key_from_key_set_async",
        fake_key_from_key_set,
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_node_output_cache_store_get",
        fake_store_get,
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_deserialize_node_output_cache_entry",
        lambda execution, record: record,
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_node_output_cache_ancestor_ids",
        lambda cache_key_set, node_id: {"14"} if node_id == "11" else set(),
    )
    monkeypatch.setattr(
        modal_cloud_module,
        "_emit_cloud_info",
        lambda message, *args: observed_logs.append((message, *args)),
    )

    restored_node_ids = asyncio.run(
        modal_cloud_module._restore_persisted_node_output_cache_entries_into_prepared_cache(
            object(),
            FakeOutputsCache(),
            prompt={
                "11": {"class_type": "VAEDecode", "inputs": {"samples": ["14", 0]}},
                "14": {"class_type": "LoraLoaderModelOnly", "inputs": {}},
            },
            cache_store={},
            required_materialized_node_ids={"14"},
        )
    )

    assert restored_node_ids == []
    assert restored_values == {}
    assert (
        "Node output cache lookup node=%s key_prefix=%s result=skip reason=missing-required-boundary-ancestors ancestors=%s",
        "11",
        "NC_11",
        ["14"],
    ) in observed_logs


def test_modal_cloud_materializes_synced_asset_paths(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """Remote asset references should resolve to absolute files under the storage root."""
    monkeypatch.setenv("COMFY_MODAL_REMOTE_STORAGE_ROOT", "/storage")
    modal_cloud_module.get_settings.cache_clear()
    try:
        assert modal_cloud_module._materialize_remote_asset_path("/assets/model.safetensors") == (
            "/storage/assets/model.safetensors"
        )
        assert modal_cloud_module._rewrite_modal_asset_references(
            {"clip_name": "/assets/model.safetensors", "nested": ["/assets/other.pt", 3]}
        ) == {
            "clip_name": "/storage/assets/model.safetensors",
            "nested": ["/storage/assets/other.pt", 3],
        }
    finally:
        modal_cloud_module.get_settings.cache_clear()


def test_modal_cloud_summarizes_suspicious_wrapped_prompt_inputs(
    modal_cloud_module: Any,
) -> None:
    """Remote failure diagnostics should flag remaining singleton-list prompt wrappers."""
    prompt = {
        "12": {
            "class_type": "ExampleNode",
            "inputs": {
                "steps": [20],
                "latent": ["7", [0]],
                "ok_link": ["8", 0],
            },
        }
    }

    findings = modal_cloud_module._summarize_suspicious_prompt_inputs(prompt)

    assert findings == [
        "12.steps=[20]",
        "12.latent=['7', [0]]",
    ]


def test_modal_cloud_accepts_absolute_asset_paths_in_folder_lookup(
    modal_cloud_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """Patched folder lookups should return already-materialized absolute asset paths."""
    remote_storage_root = tmp_path / "storage"
    asset_path = remote_storage_root / "assets" / "clip.safetensors"
    asset_path.parent.mkdir(parents=True)
    asset_path.write_bytes(b"clip")

    fake_folder_paths_module = types.SimpleNamespace(
        get_full_path=lambda folder_name, filename: None,
        get_full_path_or_raise=lambda folder_name, filename: (_ for _ in ()).throw(
            FileNotFoundError(filename)
        ),
    )
    monkeypatch.setitem(sys.modules, "folder_paths", fake_folder_paths_module)
    monkeypatch.setenv("COMFY_MODAL_REMOTE_STORAGE_ROOT", str(remote_storage_root))
    modal_cloud_module.get_settings.cache_clear()
    try:
        with modal_cloud_module._patched_folder_paths_absolute_lookup():
            resolved = fake_folder_paths_module.get_full_path(
                "text_encoders",
                "/assets/clip.safetensors",
            )
            assert resolved == str(asset_path)
            assert (
                fake_folder_paths_module.get_full_path_or_raise(
                    "text_encoders",
                    "/assets/clip.safetensors",
                )
                == str(asset_path)
            )
    finally:
        modal_cloud_module.get_settings.cache_clear()


def test_modal_cloud_preserves_absolute_lookup_until_overlapping_contexts_exit(
    modal_cloud_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """One loader finishing must not remove the lookup patch from another loader."""
    asset_path = tmp_path / "unet.safetensors"
    asset_path.write_bytes(b"unet")

    def original_get_full_path(folder_name: str, filename: str) -> None:
        """Represent ComfyUI's original lookup missing the absolute asset."""
        del folder_name, filename
        return None

    def original_get_full_path_or_raise(folder_name: str, filename: str) -> str:
        """Represent ComfyUI's original raising lookup."""
        del folder_name
        raise FileNotFoundError(filename)

    fake_folder_paths_module = types.SimpleNamespace(
        get_full_path=original_get_full_path,
        get_full_path_or_raise=original_get_full_path_or_raise,
    )
    monkeypatch.setitem(sys.modules, "folder_paths", fake_folder_paths_module)
    first_entered = threading.Event()
    second_entered = threading.Event()
    release_first = threading.Event()
    first_exited = threading.Event()
    release_second = threading.Event()
    results: list[str | None] = []

    def first_loader() -> None:
        """Enter first and exit while the second lookup context remains active."""
        with modal_cloud_module._patched_folder_paths_absolute_lookup():
            first_entered.set()
            assert release_first.wait(timeout=2.0)
        first_exited.set()

    def second_loader() -> None:
        """Resolve an absolute model after the first overlapping loader exits."""
        assert first_entered.wait(timeout=2.0)
        with modal_cloud_module._patched_folder_paths_absolute_lookup():
            second_entered.set()
            assert first_exited.wait(timeout=2.0)
            results.append(
                fake_folder_paths_module.get_full_path("diffusion_models", str(asset_path))
            )
            assert release_second.wait(timeout=2.0)

    with ThreadPoolExecutor(max_workers=2) as executor:
        first_future = executor.submit(first_loader)
        second_future = executor.submit(second_loader)
        assert second_entered.wait(timeout=2.0)
        release_first.set()
        assert first_exited.wait(timeout=2.0)
        first_future.result(timeout=2.0)
        release_second.set()
        second_future.result(timeout=2.0)

    assert results == [str(asset_path)]
    assert fake_folder_paths_module.get_full_path is original_get_full_path
    assert fake_folder_paths_module.get_full_path_or_raise is original_get_full_path_or_raise


def test_modal_cloud_aliases_flux_rms_norm_weight_keys_for_model_detection(
    modal_cloud_module: Any,
) -> None:
    """Saved Flux models using RMSNorm `.weight` keys should gain `.scale` aliases."""
    state_dict = {
        "model.diffusion_model.double_blocks.0.img_attn.norm.key_norm.weight": object(),
        "model.diffusion_model.double_blocks.0.img_attn.norm.query_norm.weight": object(),
        "model.diffusion_model.double_blocks.0.txt_attn.norm.key_norm.weight": object(),
        "model.diffusion_model.double_blocks.0.txt_attn.norm.query_norm.weight": object(),
        "model.diffusion_model.img_in.weight": object(),
    }

    alias_count = modal_cloud_module._alias_flux_rms_norm_weight_keys(state_dict)

    assert alias_count == 4
    assert (
        state_dict["model.diffusion_model.double_blocks.0.img_attn.norm.key_norm.scale"]
        is state_dict["model.diffusion_model.double_blocks.0.img_attn.norm.key_norm.weight"]
    )
    assert (
        state_dict["model.diffusion_model.double_blocks.0.img_attn.norm.query_norm.scale"]
        is state_dict["model.diffusion_model.double_blocks.0.img_attn.norm.query_norm.weight"]
    )
    assert (
        state_dict["model.diffusion_model.double_blocks.0.txt_attn.norm.key_norm.scale"]
        is state_dict["model.diffusion_model.double_blocks.0.txt_attn.norm.key_norm.weight"]
    )
    assert (
        state_dict["model.diffusion_model.double_blocks.0.txt_attn.norm.query_norm.scale"]
        is state_dict["model.diffusion_model.double_blocks.0.txt_attn.norm.query_norm.weight"]
    )
    assert modal_cloud_module._alias_flux_rms_norm_weight_keys(state_dict) == 0


def test_modal_cloud_model_state_dict_wrapper_forwards_new_comfyui_arguments(
    modal_cloud_module: Any,
    monkeypatch: Any,
) -> None:
    """The compatibility wrapper should preserve evolving ComfyUI loader arguments."""
    calls: list[tuple[dict[str, Any], tuple[Any, ...], dict[str, Any]]] = []

    def original_loader(
        state_dict: dict[str, Any],
        *args: Any,
        **kwargs: Any,
    ) -> str:
        """Record the arguments forwarded by the compatibility wrapper."""
        calls.append((state_dict, args, kwargs))
        return "loaded"

    fake_comfy_module = types.ModuleType("comfy")
    fake_comfy_module.__path__ = []
    fake_sd_module = types.ModuleType("comfy.sd")
    fake_sd_module.load_diffusion_model_state_dict = original_loader
    fake_comfy_module.sd = fake_sd_module
    monkeypatch.setitem(sys.modules, "comfy", fake_comfy_module)
    monkeypatch.setitem(sys.modules, "comfy.sd", fake_sd_module)
    monkeypatch.setattr(modal_cloud_module, "_MODEL_STATE_DICT_COMPAT_WRAPPED", False)

    modal_cloud_module._install_model_state_dict_compatibility_wrappers()
    state_dict = {
        "model.diffusion_model.double_blocks.0.img_attn.norm.key_norm.weight": object()
    }

    result = fake_sd_module.load_diffusion_model_state_dict(
        state_dict,
        {"dtype": "default"},
        metadata={"source": "test"},
        disable_dynamic=True,
        future_option="preserved",
    )

    assert result == "loaded"
    assert calls == [
        (
            state_dict,
            ({"dtype": "default"},),
            {
                "metadata": {"source": "test"},
                "disable_dynamic": True,
                "future_option": "preserved",
            },
        )
    ]
    assert (
        state_dict["model.diffusion_model.double_blocks.0.img_attn.norm.key_norm.scale"]
        is state_dict["model.diffusion_model.double_blocks.0.img_attn.norm.key_norm.weight"]
    )


def test_modal_cloud_force_imports_comfyui_utils_package(
    modal_cloud_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """ComfyUI's utils package should override a shadowing non-package module."""
    package_root = tmp_path / "comfyui"
    utils_dir = package_root / "utils"
    utils_dir.mkdir(parents=True)
    (utils_dir / "__init__.py").write_text("SENTINEL = 'comfy-utils'\n", encoding="utf-8")

    shadow_module = types.ModuleType("utils")
    shadow_module.__file__ = str(tmp_path / "utils.py")
    monkeypatch.setitem(sys.modules, "utils", shadow_module)

    modal_cloud_module._force_import_package_from_root("utils", package_root)

    imported_module = sys.modules["utils"]
    assert getattr(imported_module, "SENTINEL", None) == "comfy-utils"
    assert list(getattr(imported_module, "__path__", [])) == [str(utils_dir)]


def test_modal_cloud_creates_default_custom_nodes_dir_when_missing(
    modal_cloud_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """The remote runtime should create an empty default custom_nodes directory for ComfyUI."""
    comfyui_root = tmp_path / "comfyui"
    comfyui_root.mkdir()
    monkeypatch.setattr(modal_cloud_module, "_REMOTE_COMFYUI_ROOT", comfyui_root)
    monkeypatch.setattr(modal_cloud_module, "_LOCAL_COMFYUI_ROOT", tmp_path / "missing-local")

    custom_nodes_dir = modal_cloud_module._ensure_default_custom_nodes_dir()

    assert custom_nodes_dir == comfyui_root / "custom_nodes"
    assert custom_nodes_dir is not None
    assert custom_nodes_dir.exists()
    assert custom_nodes_dir.is_dir()


class _BoundarySourceNode:
    """Simple source node used for subgraph execution tests."""

    RETURN_TYPES = ("INT",)
    RETURN_NAMES = ("value",)
    OUTPUT_IS_LIST = (False,)
    FUNCTION = "run"

    @classmethod
    def INPUT_TYPES(cls) -> dict[str, dict[str, tuple[str]]]:
        """Return the minimal V1 input schema."""
        return {"required": {"value": ("INT",)}}

    def run(self, value: int) -> tuple[int]:
        """Increment the boundary input."""
        return (value + 1,)


class _BoundarySinkNode:
    """Simple downstream node used for subgraph execution tests."""

    RETURN_TYPES = ("INT",)
    RETURN_NAMES = ("value",)
    OUTPUT_IS_LIST = (False,)
    FUNCTION = "run"

    @classmethod
    def INPUT_TYPES(cls) -> dict[str, dict[str, tuple[str]]]:
        """Return the minimal V1 input schema."""
        return {"required": {"value": ("INT",)}}

    def run(self, value: int) -> tuple[int]:
        """Double the upstream value."""
        return (value * 2,)


class _PrimitiveEchoNode:
    """Simple node used to verify primitive widget coercion."""

    RETURN_TYPES = ("INT", "FLOAT", "BOOLEAN", "STRING")
    RETURN_NAMES = ("steps", "cfg", "enabled", "label")
    OUTPUT_IS_LIST = (False, False, False, False)
    FUNCTION = "run"

    @classmethod
    def INPUT_TYPES(cls) -> dict[str, dict[str, tuple[str]]]:
        """Return one primitive input of each V1 widget type."""
        return {
            "required": {
                "steps": ("INT",),
                "cfg": ("FLOAT",),
                "enabled": ("BOOLEAN",),
                "label": ("STRING",),
            }
        }

    def run(
        self,
        steps: int,
        cfg: float,
        enabled: bool,
        label: str,
    ) -> tuple[int, float, bool, str]:
        """Echo primitive inputs after asserting their coerced Python types."""
        assert isinstance(steps, int)
        assert isinstance(cfg, float)
        assert isinstance(enabled, bool)
        assert isinstance(label, str)
        return (steps, cfg, enabled, label)


def _v3_batch_images_node_class() -> type[Any]:
    """Return a dependency-light V3 Batch Images node with Autogrow inputs."""
    class _V3BatchImagesNode:
        """Minimal V3 node matching ComfyUI's Batch Images input schema."""

        @classmethod
        def INPUT_TYPES(cls) -> dict[str, dict[str, tuple[Any, ...]]]:
            """Return the raw V3 Autogrow schema emitted by ComfyUI."""
            return {
                "required": {
                    "images": (
                        "COMFY_AUTOGROW_V3",
                        {
                            "template": {
                                "input": {"required": {"image": ("IMAGE", {})}},
                                "prefix": "image",
                                "min": 1,
                                "max": 50,
                            }
                        },
                    )
                }
            }

    return _V3BatchImagesNode


def _install_fake_v3_input_finalizer(monkeypatch: pytest.MonkeyPatch) -> None:
    """Install a narrow ComfyUI V3 finalizer double for Autogrow unit tests."""
    comfy_api_module = types.ModuleType("comfy_api")
    comfy_api_module.__path__ = []
    comfy_api_latest_module = types.ModuleType("comfy_api.latest")

    def get_finalized_class_inputs(
        raw_input_types: dict[str, Any],
        live_inputs: dict[str, Any],
    ) -> tuple[dict[str, Any], None, dict[str, Any]]:
        """Expand the test Autogrow schema using ComfyUI's public prompt shape."""
        _, config = raw_input_types["required"]["images"]
        template = config["template"]
        template_input = template["input"]["required"]["image"]
        names = [f"{template['prefix']}{index}" for index in range(template["max"])]
        finalized: dict[str, dict[str, Any]] = {"required": {}, "optional": {}}
        dynamic_paths: dict[str, str] = {}
        for index, name in enumerate(names):
            expanded_name = f"images.{name}"
            section_name = "required" if index < template["min"] else "optional"
            finalized[section_name][expanded_name] = template_input
            if expanded_name in live_inputs:
                dynamic_paths[expanded_name] = expanded_name
        return finalized, None, {"dynamic_paths": dynamic_paths}

    comfy_api_latest_module._io = types.SimpleNamespace(
        get_finalized_class_inputs=get_finalized_class_inputs
    )
    comfy_api_module.latest = comfy_api_latest_module
    monkeypatch.setitem(sys.modules, "comfy_api", comfy_api_module)
    monkeypatch.setitem(sys.modules, "comfy_api.latest", comfy_api_latest_module)


class _ImplicitBatchListSourceNode:
    """Fake node that consumes a whole list once and emits scalar and list outputs."""

    RETURN_TYPES = ("INT", "INT")
    RETURN_NAMES = ("first_value", "all_values")
    OUTPUT_IS_LIST = (False, True)
    INPUT_IS_LIST = True
    FUNCTION = "run"


class _ImplicitBatchScalarConsumerNode:
    """Fake scalar consumer used for outer implicit batch regressions."""

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    OUTPUT_IS_LIST = (False,)
    FUNCTION = "run"


class _ImplicitBatchListConsumerNode:
    """Fake list-aware consumer used for outer implicit batch regressions."""

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    OUTPUT_IS_LIST = (True,)
    INPUT_IS_LIST = True
    FUNCTION = "run"


def test_local_remote_app_executes_subgraph_payload(
    remote_modal_app_module: Any,
    serialization_module: Any,
) -> None:
    """The local fallback remote app should execute rewritten subgraph payloads."""
    payload = remote_modal_app_module.execute_subgraph_locally(
        payload={
            "payload_kind": "subgraph",
            "component_id": "component-1",
            "subgraph_prompt": {
                "remote_1": {
                    "class_type": "BoundarySource",
                    "inputs": {"value": 0},
                    "_meta": {},
                },
                "remote_2": {
                    "class_type": "BoundarySink",
                    "inputs": {"value": ["remote_1", 0]},
                    "_meta": {},
                },
            },
            "boundary_inputs": [
                {
                    "proxy_input_name": "remote_input_0",
                    "targets": [{"node_id": "remote_1", "input_name": "value"}],
                }
            ],
            "boundary_outputs": [
                {
                    "proxy_output_name": "remote_2_value",
                    "node_id": "remote_2",
                    "output_index": 0,
                    "io_type": "INT",
                    "is_list": False,
                }
            ],
            "execute_node_ids": ["remote_2"],
            "extra_data": {},
            "custom_nodes_bundle": None,
        },
        kwargs_payload='{"remote_input_0": 4}',
        node_mapping={
            "BoundarySource": _BoundarySourceNode,
            "BoundarySink": _BoundarySinkNode,
        },
    )
    outputs = serialization_module.deserialize_node_outputs(payload)
    assert outputs == (10,)


def test_local_remote_app_normalizes_wrapped_subgraph_link_indexes(
    remote_modal_app_module: Any,
    serialization_module: Any,
) -> None:
    """The local fallback runner should canonicalize singleton-list prompt link indexes."""
    payload = remote_modal_app_module.execute_subgraph_locally(
        payload={
            "payload_kind": "subgraph",
            "component_id": "component-1",
            "subgraph_prompt": {
                "remote_1": {
                    "class_type": "BoundarySource",
                    "inputs": {"value": 0},
                    "_meta": {},
                },
                "remote_2": {
                    "class_type": "BoundarySink",
                    "inputs": {"value": [[["remote_1", [0]]]]},
                    "_meta": {},
                },
            },
            "boundary_inputs": [
                {
                    "proxy_input_name": "remote_input_0",
                    "targets": [{"node_id": "remote_1", "input_name": "value"}],
                }
            ],
            "boundary_outputs": [
                {
                    "proxy_output_name": "remote_2_value",
                    "node_id": "remote_2",
                    "output_index": [[0]],
                    "io_type": "INT",
                    "is_list": False,
                }
            ],
            "execute_node_ids": ["remote_2"],
            "extra_data": {},
            "custom_nodes_bundle": None,
        },
        kwargs_payload='{"remote_input_0": 4}',
        node_mapping={
            "BoundarySource": _BoundarySourceNode,
            "BoundarySink": _BoundarySinkNode,
        },
    )
    outputs = serialization_module.deserialize_node_outputs(payload)
    assert outputs == (10,)


def test_local_remote_app_normalizes_wrapped_scalar_prompt_inputs(
    remote_modal_app_module: Any,
    serialization_module: Any,
) -> None:
    """The local fallback runner should unwrap singleton-list scalar prompt inputs."""
    payload = remote_modal_app_module.execute_subgraph_locally(
        payload={
            "payload_kind": "subgraph",
            "component_id": "component-1",
            "subgraph_prompt": {
                "remote_1": {
                    "class_type": "BoundarySource",
                    "inputs": {"value": [4]},
                    "_meta": {},
                }
            },
            "boundary_inputs": [],
            "boundary_outputs": [
                {
                    "proxy_output_name": "remote_1_value",
                    "node_id": "remote_1",
                    "output_index": 0,
                    "io_type": "INT",
                    "is_list": False,
                }
            ],
            "execute_node_ids": [["remote_1"]],
            "extra_data": {},
            "custom_nodes_bundle": None,
        },
        kwargs_payload="{}",
        node_mapping={
            "BoundarySource": _BoundarySourceNode,
        },
    )
    outputs = serialization_module.deserialize_node_outputs(payload)
    assert outputs == (5,)


@pytest.mark.parametrize(
    ("module_fixture_name",),
    [
        ("remote_modal_app_module",),
        ("modal_cloud_module",),
    ],
)
def test_apply_boundary_inputs_normalizes_wrapped_scalar_values(
    request: Any,
    module_fixture_name: str,
) -> None:
    """Boundary input hydration should unwrap singleton-list scalar wrappers before PromptExecutor sees them."""
    target_module = request.getfixturevalue(module_fixture_name)
    prompt = {
        "remote_1": {
            "class_type": "BoundarySource",
            "inputs": {"value": 0},
            "_meta": {},
        }
    }

    target_module._apply_boundary_inputs(
        prompt=prompt,
        boundary_input_specs=[
            {
                "proxy_input_name": "remote_input_0",
                "targets": [{"node_id": "remote_1", "input_name": "value"}],
            }
        ],
        hydrated_inputs={"remote_input_0": [4]},
    )

    assert prompt["remote_1"]["inputs"]["value"] == 4


@pytest.mark.parametrize(
    ("module_fixture_name",),
    [
        ("remote_modal_app_module",),
        ("modal_cloud_module",),
    ],
)
def test_apply_boundary_inputs_preserves_singleton_conditioning_lists(
    request: Any,
    module_fixture_name: str,
) -> None:
    """Boundary input hydration must not flatten singleton CONDITIONING payloads."""
    target_module = request.getfixturevalue(module_fixture_name)
    conditioning = [["cond", {"pooled_output": None}]]
    prompt = {
        "remote_1": {
            "class_type": "BoundarySource",
            "inputs": {"value": 0},
            "_meta": {},
        }
    }

    target_module._apply_boundary_inputs(
        prompt=prompt,
        boundary_input_specs=[
            {
                "proxy_input_name": "remote_input_0",
                "io_type": "CONDITIONING",
                "targets": [{"node_id": "remote_1", "input_name": "value"}],
            }
        ],
        hydrated_inputs={"remote_input_0": conditioning},
    )

    assert prompt["remote_1"]["inputs"]["value"] == conditioning


@pytest.mark.parametrize(
    ("module_fixture_name",),
    [
        ("remote_modal_app_module",),
        ("modal_cloud_module",),
    ],
)
def test_validate_required_prompt_inputs_reports_missing_latent_image(
    request: Any,
    module_fixture_name: str,
) -> None:
    """Remote execution should fail before PromptExecutor calls samplers with missing required inputs."""
    target_module = request.getfixturevalue(module_fixture_name)
    prompt = {
        "3": {
            "class_type": "KSamplerLoraSigmaInverse",
            "inputs": {
                "seed": 1,
                "steps": 24,
                "positive": [["cond", {"reference_latents": ["reference"]}]],
            },
        }
    }

    with pytest.raises(target_module.RemoteSubgraphExecutionError) as exc_info:
        target_module._validate_required_prompt_inputs(
            prompt,
            {"KSamplerLoraSigmaInverse": _FakeImplicitBatchKSamplerNode},
        )

    message = str(exc_info.value)
    assert "missing required node inputs" in message
    assert "latent_image" in message
    assert "available_inputs=['positive', 'seed', 'steps']" in message


@pytest.mark.parametrize(
    ("module_fixture_name",),
    [
        ("remote_modal_app_module",),
        ("modal_cloud_module",),
    ],
)
def test_validate_required_prompt_inputs_expands_v3_autogrow_inputs(
    request: Any,
    module_fixture_name: str,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """V3 Autogrow prompts should validate their expanded `group.name` sockets."""
    target_module = request.getfixturevalue(module_fixture_name)
    _install_fake_v3_input_finalizer(monkeypatch)
    batch_images_node = _v3_batch_images_node_class()
    live_inputs = {
        "images.image0": ["101", 0],
        "images.image1": ["102", 0],
    }
    prompt = {
        "153": {
            "class_type": "ModalTestV3BatchImagesNode",
            "inputs": live_inputs,
        }
    }

    target_module._validate_required_prompt_inputs(
        prompt,
        {"ModalTestV3BatchImagesNode": batch_images_node},
    )

    assert target_module._node_required_input_names(
        batch_images_node,
        live_inputs,
    ) == {"images.image0"}
    input_type_map = target_module._node_input_type_map(batch_images_node, live_inputs)
    assert input_type_map["images.image0"] == "IMAGE"
    assert input_type_map["images.image1"] == "IMAGE"
    assert "images" not in input_type_map


@pytest.mark.parametrize(
    ("module_fixture_name",),
    [
        ("remote_modal_app_module",),
        ("modal_cloud_module",),
    ],
)
def test_validate_prompt_input_shapes_rejects_list_on_primitive_socket(
    request: Any,
    module_fixture_name: str,
) -> None:
    """Prepared remote prompts should fail early when primitive widget inputs still carry raw lists."""
    target_module = request.getfixturevalue(module_fixture_name)
    prompt = {
        "remote_1": {
            "class_type": "BoundarySource",
            "inputs": {"value": [4, 5]},
            "_meta": {},
        }
    }

    with pytest.raises(target_module.RemoteSubgraphExecutionError, match="input_name='value'"):
        target_module._validate_prompt_input_shapes(
            prompt,
            {"BoundarySource": _BoundarySourceNode},
        )


@pytest.mark.parametrize(
    ("module_fixture_name",),
    [
        ("remote_modal_app_module",),
        ("modal_cloud_module",),
    ],
)
def test_validate_prompt_input_shapes_allows_boundary_supplied_list_on_primitive_socket(
    request: Any,
    module_fixture_name: str,
) -> None:
    """Boundary-fed primitive lists should pass through so ComfyUI can map over them normally."""
    target_module = request.getfixturevalue(module_fixture_name)
    prompt = {
        "remote_1": {
            "class_type": "BoundarySource",
            "inputs": {"value": [4, 5]},
            "_meta": {},
        }
    }

    target_module._validate_prompt_input_shapes(
        prompt,
        {"BoundarySource": _BoundarySourceNode},
        [
            {
                "proxy_input_name": "remote_input_0",
                "targets": [{"node_id": "remote_1", "input_name": "value"}],
            }
        ],
    )


@pytest.mark.parametrize(
    ("module_fixture_name",),
    [
        ("remote_modal_app_module",),
        ("modal_cloud_module",),
    ],
)
def test_coerce_prompt_primitive_input_values_matches_comfyui_semantics(
    request: Any,
    module_fixture_name: str,
) -> None:
    """Remote runtimes should coerce primitive prompt literals the same way ComfyUI does."""
    target_module = request.getfixturevalue(module_fixture_name)
    prompt = {
        "remote_1": {
            "class_type": "PrimitiveEcho",
            "inputs": {
                "steps": 18.0,
                "cfg": 5,
                "enabled": 1,
                "label": 7,
            },
            "_meta": {},
        }
    }

    target_module._coerce_prompt_primitive_input_values(
        prompt,
        {"PrimitiveEcho": _PrimitiveEchoNode},
    )

    assert prompt["remote_1"]["inputs"] == {
        "steps": 18,
        "cfg": 5.0,
        "enabled": True,
        "label": "7",
    }
    assert isinstance(prompt["remote_1"]["inputs"]["steps"], int)
    assert isinstance(prompt["remote_1"]["inputs"]["cfg"], float)
    assert isinstance(prompt["remote_1"]["inputs"]["enabled"], bool)
    assert isinstance(prompt["remote_1"]["inputs"]["label"], str)


def test_local_remote_app_coerces_primitive_widget_literals_before_execution(
    remote_modal_app_module: Any,
    serialization_module: Any,
) -> None:
    """The local fallback runner should coerce primitive widget literals before executing nodes."""
    payload = remote_modal_app_module.execute_subgraph_locally(
        payload={
            "payload_kind": "subgraph",
            "component_id": "component-primitive-coercion",
            "subgraph_prompt": {
                "remote_1": {
                    "class_type": "PrimitiveEcho",
                    "inputs": {
                        "steps": 18.0,
                        "cfg": 5,
                        "enabled": 1,
                        "label": 7,
                    },
                    "_meta": {},
                }
            },
            "boundary_inputs": [],
            "boundary_outputs": [
                {
                    "proxy_output_name": "steps",
                    "node_id": "remote_1",
                    "output_index": 0,
                    "io_type": "INT",
                    "is_list": False,
                },
                {
                    "proxy_output_name": "cfg",
                    "node_id": "remote_1",
                    "output_index": 1,
                    "io_type": "FLOAT",
                    "is_list": False,
                },
                {
                    "proxy_output_name": "enabled",
                    "node_id": "remote_1",
                    "output_index": 2,
                    "io_type": "BOOLEAN",
                    "is_list": False,
                },
                {
                    "proxy_output_name": "label",
                    "node_id": "remote_1",
                    "output_index": 3,
                    "io_type": "STRING",
                    "is_list": False,
                },
            ],
            "execute_node_ids": ["remote_1"],
            "extra_data": {},
            "custom_nodes_bundle": None,
        },
        kwargs_payload="{}",
        node_mapping={"PrimitiveEcho": _PrimitiveEchoNode},
    )

    outputs = serialization_module.deserialize_node_outputs(payload)
    assert outputs == (18, 5.0, True, "7")


@pytest.mark.parametrize(
    ("module_fixture_name",),
    [
        ("remote_modal_app_module",),
        ("modal_cloud_module",),
    ],
)
def test_format_prompt_executor_error_payload_includes_node_context(
    request: Any,
    module_fixture_name: str,
) -> None:
    """PromptExecutor failure formatting should surface the failing node and current inputs."""
    target_module = request.getfixturevalue(module_fixture_name)

    message = target_module._format_prompt_executor_error_payload(
        {
            "exception_message": "int() argument must be a string, a bytes-like object or a real number, not 'list'",
            "node_id": "12",
            "node_type": "KSampler",
            "current_inputs": [{"input_name": "steps", "value": [4, 5]}],
        }
    )

    assert "node_id='12'" in message
    assert "node_type='KSampler'" in message
    assert "current_inputs=" in message