fix: track expansion output consumers and gate expected_outputs behind LAZY_OUTPUTS

Signed-off-by: bigcat88 <bigcat88@icloud.com>

tests/execution/test_execution.py

@@ -672,6 +672,46 @@ class TestExecution:
         assert numpy.array(images0[0]).min() == 255, "Output 0 should be white"
         assert numpy.array(images2[0]).min() == 255, "Output 2 should be white"
 
+    def test_expected_outputs_expansion_output_mapping(self, client: ComfyClient, builder: GraphBuilder):
+        """A socket consumed only via an expansion's parent-output mapping must still
+        be in the inner LAZY_OUTPUTS node's expected_outputs (white, not black)."""
+        g = builder
+        expander = g.node("TestExpectedOutputsExpansion", height=80, width=80)
+        output = g.node("PreviewImage", images=expander.out(0))
+
+        result = client.run(g)
+
+        images = result.get_images(output)
+        assert len(images) == 1, "Should have 1 image"
+        assert numpy.array(images[0]).min() == 255, (
+            "Inner node skipped an output that is consumed via the expansion's "
+            "parent-output mapping (expected white, got black)"
+        )
+
+    def test_expected_outputs_requires_opt_in(self, client: ComfyClient, builder: GraphBuilder, server):
+        """Nodes without LAZY_OUTPUTS must see expected_outputs=None: their cache key
+        ignores topology, so a skipped output would be served stale after rewiring."""
+        g = builder
+        node = g.node("TestExpectedOutputsNotOptedIn", height=96, width=96)
+        output0 = g.node("PreviewImage", images=node.out(0))
+
+        # Only output 0 connected: correct gating -> node sees None, computes all
+        result1 = client.run(g)
+        assert numpy.array(result1.get_images(output0)[0]).min() == 255
+
+        # Connect output 1: key unchanged -> cache hit must still serve correct data
+        output1 = g.node("PreviewImage", images=node.out(1))
+        result2 = client.run(g)
+
+        if server["should_cache_results"]:
+            assert not result2.did_run(node), "Node should be a cache hit (key ignores topology)"
+        images1 = result2.get_images(output1)
+        assert len(images1) == 1, "Should have 1 image for output1"
+        assert numpy.array(images1[0]).min() == 255, (
+            "Non-opted-in node observed expected_outputs and skipped output 1; "
+            "the stale skipped value was then served from cache"
+        )
+
     def test_parallel_sleep_nodes(self, client: ComfyClient, builder: GraphBuilder, skip_timing_checks):
         # Warmup execution to ensure server is fully initialized
         run_warmup(client)

tests/execution/testing_nodes/testing-pack/specific_tests.py

@@ -6,7 +6,7 @@ from .tools import VariantSupport
 from comfy_execution.graph_utils import GraphBuilder
 from comfy.comfy_types.node_typing import ComfyNodeABC
 from comfy.comfy_types import IO
-from comfy_execution.utils import get_executing_context
+from comfy_execution.utils import get_executing_context, is_output_needed
 
 class TestLazyMixImages:
     @classmethod
@@ -510,27 +510,76 @@ class TestExpectedOutputs:
     CATEGORY = "_for_testing"
 
     def execute(self, height, width):
-        ctx = get_executing_context()
-
-        # Default: assume all outputs are expected (backwards compatibility)
-        output0_expected = True
-        output1_expected = True
-        output2_expected = True
-
-        if ctx is not None and ctx.expected_outputs is not None:
-            output0_expected = 0 in ctx.expected_outputs
-            output1_expected = 1 in ctx.expected_outputs
-            output2_expected = 2 in ctx.expected_outputs
-
         # Return white image if expected, black if not
         # This allows tests to verify which outputs were expected via pixel values
         white = torch.ones(1, height, width, 3)
         black = torch.zeros(1, height, width, 3)
 
         return (
-            white if output0_expected else black,
-            white if output1_expected else black,
-            white if output2_expected else black,
+            white if is_output_needed(0) else black,
+            white if is_output_needed(1) else black,
+            white if is_output_needed(2) else black,
+        )
+
+
+class TestExpectedOutputsExpansion:
+    """Expands into an inner LAZY_OUTPUTS node whose output 1 is consumed ONLY via
+    the parent-output mapping (no input link anywhere). If that mapping is not part
+    of the expected-outputs map, the inner node wrongly skips it -> black not white.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "height": ("INT", {"default": 64, "min": 1, "max": 1024}),
+                "width": ("INT", {"default": 64, "min": 1, "max": 1024}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "execute"
+    CATEGORY = "_for_testing"
+
+    def execute(self, height, width):
+        g = GraphBuilder()
+        inner = g.node("TestExpectedOutputs", height=height, width=width)
+        return {"result": (inner.out(1),), "expand": g.finalize()}
+
+
+class TestExpectedOutputsNotOptedIn:
+    """Reads expected_outputs WITHOUT declaring LAZY_OUTPUTS; the executor must pass
+    None (such nodes have no cache-key protection against output rewiring). Outputs
+    are white when the node correctly sees None, otherwise they encode membership.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "height": ("INT", {"default": 64, "min": 1, "max": 1024}),
+                "width": ("INT", {"default": 64, "min": 1, "max": 1024}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE", "IMAGE")
+    RETURN_NAMES = ("output0", "output1")
+    FUNCTION = "execute"
+    CATEGORY = "_for_testing"
+
+    def execute(self, height, width):
+        # Raw context access (not is_output_needed): must distinguish None from a set
+        ctx = get_executing_context()
+        expected = ctx.expected_outputs if ctx is not None else None
+
+        white = torch.ones(1, height, width, 3)
+        black = torch.zeros(1, height, width, 3)
+
+        if expected is None:
+            return (white, white.clone())
+        return (
+            white if 0 in expected else black,
+            white if 1 in expected else black,
         )
 
 
@@ -551,6 +600,8 @@ TEST_NODE_CLASS_MAPPINGS = {
     "TestParallelSleep": TestParallelSleep,
     "TestOutputNodeWithSocketOutput": TestOutputNodeWithSocketOutput,
     "TestExpectedOutputs": TestExpectedOutputs,
+    "TestExpectedOutputsExpansion": TestExpectedOutputsExpansion,
+    "TestExpectedOutputsNotOptedIn": TestExpectedOutputsNotOptedIn,
 }
 
 TEST_NODE_DISPLAY_NAME_MAPPINGS = {
@@ -570,4 +621,6 @@ TEST_NODE_DISPLAY_NAME_MAPPINGS = {
     "TestParallelSleep": "Test Parallel Sleep",
     "TestOutputNodeWithSocketOutput": "Test Output Node With Socket Output",
     "TestExpectedOutputs": "Test Expected Outputs",
+    "TestExpectedOutputsExpansion": "Test Expected Outputs Expansion",
+    "TestExpectedOutputsNotOptedIn": "Test Expected Outputs Not Opted In",
 }