Project native image ABI without discarding registry slice context

openhcs/processing/backends/lib_registry/unified_registry.py

@@ -1141,9 +1141,10 @@ def supports_measurement_image_batch(
     def main_flow_call_argument(
         self, callable_contract: "CallableContract", source_payload: Any,
     ) -> Any:
-        """Keep the declared plane domain until the PURE_2D slicer consumes it."""
-        del callable_contract
-        return source_payload
+        """Retain plane context only for the declared contract execution wrapper."""
+        if callable_contract.raw_processing_function is not None:
+            return source_payload
+        return super().main_flow_call_argument(callable_contract, source_payload)
 
     def execute(self, registry, func, image, *args, **kwargs):
         return registry.execute_pure_2d(func, image, *args, **kwargs)

tests/unit/test_unified_registry_runtime_payloads.py

@@ -19,12 +19,14 @@
 
 from openhcs.core.callable_contract import CallableContract, callable_request
 from openhcs.core.config import LazyDtypeConfig
+from openhcs.core.function_patterns import compile_function_pattern
 from openhcs.core.memory import MEMORY_TYPE_NUMPY
 from openhcs.core.memory.decorators import numpy
 from openhcs.core.runtime_array_values import RuntimeArrayPayload
 from openhcs.core.runtime_image_values import (
     ImageMetadataPayload,
     ImagePayloadMetadata,
+    ImagePayloadMetadataCarrier,
     image_payload_data,
     image_payload_mask,
     image_payload_metadata,
@@ -213,7 +215,7 @@ def add_one(image: ImageMetadataPayload) -> ImageMetadataPayload:
         payload,
     )
 
-    assert isinstance(result, ImageMetadataPayload)
+    assert isinstance(result, ImagePayloadMetadataCarrier)
     np.testing.assert_array_equal(image_payload_data(result), stack + 1)
     assert seen_paths == ["z0.tif", "z1.tif"]
     assert image_payload_metadata(result).source_image_provenance_planes.paths == (
@@ -222,6 +224,43 @@ def add_one(image: ImageMetadataPayload) -> ImageMetadataPayload:
     )
 
 
+@pytest.mark.parametrize("slice_by_slice", (True, False))
+def test_native_flexible_callable_receives_declared_array_abi_and_keeps_slice_controls(
+    slice_by_slice: bool,
+) -> None:
+    stack = np.arange(40, dtype=np.float32).reshape(2, 4, 5)
+    mask = np.ones_like(stack, dtype=bool)
+    payload = ImagePayloadMetadata(plane_axis=RuntimePlaneAxis.RUNTIME_SLICE).payload_with(stack, mask)
+    seen_shapes = []
+
+    @numpy
+    def native_image_operation(image: np.ndarray) -> np.ndarray:
+        assert isinstance(image, np.ndarray)
+        seen_shapes.append(image.shape)
+        return image + 1
+
+    invocation = compile_function_pattern(
+        (native_image_operation, {"slice_by_slice": slice_by_slice}), {}, {},
+    ).default_group.invocations[0]
+    result = invocation.runtime_callable(
+        invocation.main_flow_call_argument(payload), **dict(invocation.runtime_kwargs),
+    )
+    assert seen_shapes == ([(4, 5), (4, 5)] if slice_by_slice else [(2, 4, 5)])
+    np.testing.assert_array_equal(result, stack + 1)
+    assert image_payload_mask(payload) is mask
+
+
+def test_native_flexible_callable_preserves_explicitly_declared_carrier_abi() -> None:
+    payload = ImagePayloadMetadata().payload_with(np.ones((2, 4, 5)), np.ones((2, 4, 5), dtype=bool))
+
+    @numpy
+    def carrier_operation(image: RuntimeArrayPayload) -> RuntimeArrayPayload:
+        return image
+
+    invocation = compile_function_pattern(carrier_operation, {}, {}).default_group.invocations[0]
+    assert invocation.main_flow_call_argument(payload) is payload
+
+
 @pytest.mark.parametrize("slice_by_slice", (True, False))
 def test_flexible_runtime_argument_preserves_plane_context_until_mode_selection(
     slice_by_slice: bool,