fix: carry producer-owned ACK return routes through streaming batches

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "polystore"
-version = "0.2.19"
+version = "0.3.0"
 description = "Framework-agnostic multi-backend storage abstraction for ML and scientific computing"
 readme = "README.md"
 requires-python = ">=3.11"
@@ -45,7 +45,7 @@ dependencies = [
     "zarr>=3.0.10",  # Reads Zarr 2 and 3 stores; output format is backend-owned
     "numcodecs>=0.14.0",
     "ome-zarr>=0.12.2",   # Required for OME-ZARR HCS compliance
-    "zmqruntime>=0.2.21",
+    "zmqruntime>=0.3.0,<0.4",
     "platformdirs>=4.0.0",
 ]
 

src/polystore/streaming/_streaming_backend.py

@@ -22,6 +22,7 @@
 import zmq
 from zmqruntime.ack_listener import GlobalAckListener
 from zmqruntime.config import ZMQConfig
+from zmqruntime.messages import AckReturnRoute, ImageTransferIdentity
 from zmqruntime.viewer_protocol import (
     ViewerBatchItemWireField,
     ViewerBatchMessagePayload,
@@ -189,6 +190,7 @@ class StreamingBatchMessageRequest:
     file_paths: list[FilePath]
     stream_request: ViewerStreamRequest
     component_names_request: StreamingComponentNamesRequest
+    return_route: AckReturnRoute = field(kw_only=True)
     display_payload_extra: ViewerDisplayPayloadExtra = field(
         default_factory=ViewerDisplayPayloadExtra
     )
@@ -445,13 +447,15 @@ def prepare(
                     streaming_data_type=streaming_data_type,
                 )
             )
+            routed_payload = dict(item_payload.item_payload)
+            routed_payload.update(ImageTransferIdentity(image_id, request.return_route).to_dict())
 
             batch_images.append(
                 request.stream_request.producer.batch_item_payload(
                     ViewerStreamBatchItemSource.from_input(
                         ViewerStreamBatchItemInput(
                             stream_source=request.stream_request.source,
-                            item_payload=item_payload.item_payload,
+                            item_payload=routed_payload,
                             streaming_data_type=item_payload.streaming_data_type,
                             file_path=item_path.value,
                             index=index,
@@ -659,13 +663,6 @@ def _register_with_queue_tracker(
             port: Port number for tracker lookup
             image_ids: List of image IDs to register
         """
-        listener = GlobalAckListener()
-        listener.start(
-            port=transport_config.shared_ack_port,
-            transport_mode=transport_endpoint.transport_mode,
-            config=transport_config,
-        )
-
         from zmqruntime.queue_tracker import GlobalQueueTrackerRegistry
 
         registry = GlobalQueueTrackerRegistry()
@@ -736,20 +733,26 @@ def save_batch(
             return
 
         stream_request = ViewerStreamBackendKwargs.from_kwargs(kwargs).stream_request
+        transport_config = stream_request.transport_config.resolve(self._transport_config)
+        transport_endpoint = stream_request.viewer_transport
+        return_route = GlobalAckListener().start(
+            port=transport_config.shared_ack_port,
+            transport_mode=transport_endpoint.transport_mode,
+            config=transport_config,
+        )
         built_batch = StreamingBatchMessageBuilder.build(
             self,
             StreamingBatchMessageRequest(
                 data_list=data_list,
                 file_paths=file_paths,
                 stream_request=stream_request,
                 component_names_request=self.component_names_request(stream_request),
+                return_route=return_route,
                 display_payload_extra=self.display_payload_extra(stream_request),
             ),
         )
         self.after_batch_message_built(stream_request, built_batch)
 
-        transport_config = stream_request.transport_config.resolve(self._transport_config)
-        transport_endpoint = stream_request.viewer_transport
         self._register_with_queue_tracker(
             transport_endpoint,
             built_batch.image_ids,

src/polystore/streaming/handlers/fiji_rois.py

@@ -6,10 +6,11 @@
 
 import logging
 from collections.abc import Mapping, Sequence
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 from polystore.streaming.handlers import HandlerBase
 from polystore.streaming.base import HandlerContext
+from zmqruntime.messages import ImageTransferIdentity
 from zmqruntime.viewer_protocol import (
     ViewerBatchItemWireField,
     ViewerWireMapping,
@@ -44,13 +45,14 @@ def fiji_roi_component_value(value: ViewerWireValue) -> FijiROIComponentValue:
 
 @dataclass(frozen=True, slots=True)
 class FijiROIWireItem:
-    """Typed ROI wire item used by the legacy Fiji handler path."""
+    """Typed ROI wire item with its original producer ACK contract."""
 
     payload: ViewerWireMapping
+    transfer: ImageTransferIdentity | None = field(kw_only=True)
 
     @classmethod
     def from_payload(cls, payload: Mapping[str, ViewerWireValue]) -> "FijiROIWireItem":
-        return cls(payload)
+        return cls(payload, transfer=ImageTransferIdentity.from_item(payload))
 
     @property
     def rois(self) -> list[str]:
@@ -72,13 +74,6 @@ def metadata(self) -> ViewerWireMapping:
             raise TypeError("Fiji ROI item 'metadata' field must be a mapping.")
         return value
 
-    @property
-    def image_id(self) -> str | None:
-        field = ViewerBatchItemWireField.IMAGE_ID.value
-        if field not in self.payload or self.payload[field] is None:
-            return None
-        return str(self.payload[field])
-
     def component_value_tuple(
         self,
         component_names: Sequence[str],
@@ -180,8 +175,7 @@ def run(self):
         for roi_item in roi_items:
             rois_encoded = roi_item.rois
             if not rois_encoded:
-                if image_id := roi_item.image_id:
-                    context.server._send_ack(image_id, status="success")
+                context.server.send_ack(roi_item.transfer, status="success")
                 continue
 
             logger.info(f"🔬 FIJI ROI HANDLER: Processing {len(rois_encoded)} ROIs")
@@ -205,8 +199,7 @@ def run(self):
 
             total_rois_added += len(java_rois)
 
-            if image_id := roi_item.image_id:
-                context.server._send_ack(image_id, status="success")
+            context.server.send_ack(roi_item.transfer, status="success")
 
         logger.info(
             f"🔬 FIJI ROI HANDLER: Added {total_rois_added} ROIs to window '{context.window_key}'"

src/polystore/zmq_config.py

@@ -18,6 +18,6 @@ def _env_int(name: str, default: int) -> int:
     ipc_socket_dir=os.getenv("POLYSTORE_ZMQ_IPC_DIR", "ipc"),
     ipc_socket_prefix=os.getenv("POLYSTORE_ZMQ_IPC_PREFIX", "polystore-zmq"),
     ipc_socket_extension=os.getenv("POLYSTORE_ZMQ_IPC_EXT", ".sock"),
-    shared_ack_port=_env_int("POLYSTORE_ZMQ_ACK_PORT", 7555),
+    shared_ack_port=_env_int("POLYSTORE_ZMQ_ACK_PORT", ZMQConfig().shared_ack_port),
     app_name=os.getenv("POLYSTORE_ZMQ_APP_NAME", "polystore"),
 )

tests/test_streaming_metadata.py

@@ -4,6 +4,7 @@
 import pytest
 from zmqruntime import queue_tracker
 from zmqruntime.config import TransportMode
+from zmqruntime.messages import AckReturnRoute, ProcessIdentity
 from zmqruntime.viewer_protocol import ViewerAckPolicy, ViewerTransportEndpoint
 
 import polystore.streaming._streaming_backend as streaming_backend_module
@@ -92,11 +93,7 @@ def test_queue_tracker_uses_endpoint_transport_authority(monkeypatch) -> None:
         transport_config,
     )
 
-    assert listener_start == {
-        "port": 8111,
-        "transport_mode": endpoint.transport_mode,
-        "config": transport_config,
-    }
+    assert listener_start == {}  # Registration no longer invents a second destination.
     assert registered_image_ids == ["image-1", "image-2"]
 
 
@@ -167,6 +164,8 @@ def batch_message_request(data_list, file_paths, viewer_request):
     return StreamingBatchMessageRequest(
         data_list=data_list,
         file_paths=file_paths,
+        return_route=AckReturnRoute("tcp://127.0.0.1:8111", "00000000-0000-0000-0000-000000000001",
+                                    ProcessIdentity.current()),
         stream_request=viewer_request,
         component_names_request=(
             StreamingComponentNamesRequest.from_stream_request(viewer_request)

tests/test_streaming_return_route.py

@@ -0,0 +1,19 @@
+"""Guard the deleted ACK alias across every production streaming consumer.
+
+The historical handlers package cannot import with current metaclass-registry;
+its separate failure witness is retained rather than mocked away here. Actual
+producer/viewer delivery is exercised in OpenHCS's continuous source journey.
+"""
+
+import ast
+from pathlib import Path
+
+
+def test_no_consumer_retains_the_deleted_global_ack_alias():
+    root = Path(__file__).resolve().parents[1] / "src" / "polystore"
+    for source in root.rglob("*.py"):
+        tree = ast.parse(source.read_text(), filename=str(source))
+        assert not any(
+            isinstance(node, ast.Attribute) and node.attr in {"_send_ack", "_setup_ack_socket"}
+            for node in ast.walk(tree)
+        ), source