Wait for terminal progress callbacks before completing client execution

src/zmqruntime/execution/client.py

@@ -7,6 +7,7 @@
 import threading
 import uuid
 from abc import ABC, abstractmethod
+from dataclasses import replace
 from typing import Generic, TypeVar
 
 import zmq
@@ -18,6 +19,7 @@
 )
 from zmqruntime.execution.responses import (
     ExecutionSubmissionResponse,
+    ExecutionWaitResult,
     WireRequest,
     WireResponse,
 )
@@ -132,7 +134,20 @@ def wait_for_completion(
             known_server_process_is_alive=self.known_server_process_is_alive,
             owned_server_process_exit=self.owned_server_process_exit,
         )
-        return waiter.wait(execution_id, policy)
+        result = waiter.wait(execution_id, policy)
+        record = ExecutionWaitResult.from_wire(result).execution
+        if (
+            record is not None
+            and record.progress_sequence
+            and self._progress_registration is not None
+            and self._progress_stream is not None
+        ):
+            self._progress_stream.wait_for_delivery(
+                execution_id,
+                lambda: self._delivered_progress_through(execution_id, record.progress_sequence),
+                timeout=policy.status_timeout_ms / 1000,
+            )
+        return result
 
     def execute(
         self,
@@ -255,8 +270,8 @@ def disconnect(self):
             raise listener_error
 
     def _ensure_progress_subscription(self, *, timeout_ms: int = 5000) -> None:
-        self._start_progress_listener()
         if self._progress_registration is not None:
+            self._start_progress_listener()
             return
         response = self._send_control_request(
             {
@@ -269,9 +284,15 @@ def _ensure_progress_subscription(self, *, timeout_ms: int = 5000) -> None:
             raise RuntimeError(f"Progress registration response missing status: {response}")
         if response[MessageFields.STATUS] != ResponseType.OK.value:
             raise RuntimeError(f"Progress registration failed: {response}")
-        self._progress_registration = CallbackSubscription(
-            self._unregister_progress,
-        )
+        with self._progress_lock:
+            for execution_id, payload in response[MessageFields.PROGRESS_OBSERVATIONS].items():
+                baseline = ExecutionProgressObservation.from_wire(payload)
+                previous = self._progress_by_execution_id.get(execution_id)
+                if previous is not None:
+                    baseline = replace(baseline, delivery_error=previous.delivery_error)
+                self._progress_by_execution_id[execution_id] = baseline
+        self._progress_registration = CallbackSubscription(self._unregister_progress)
+        self._start_progress_listener()
 
     def _unregister_progress(self) -> bool:
         if not self.is_connected():
@@ -292,16 +313,44 @@ def enable_progress_stream(self) -> None:
 
     def _record_progress(self, data: dict) -> None:
         execution_id = data[MessageFields.EXECUTION_ID]
+        observation = ExecutionProgressObservation.from_wire(
+            {
+                "sequence": data[MessageFields.PROGRESS_SEQUENCE],
+                "event": data,
+            }
+        )
         with self._progress_lock:
             current = self._progress_by_execution_id.get(execution_id)
-            observation = (
-                ExecutionProgressObservation.first(data)
-                if current is None
-                else current.followed_by(data)
-            )
+            if current is not None and observation.sequence <= current.sequence:
+                return  # Events admitted before this subscription's baseline.
+            expected = 1 if current is None else current.sequence + 1
+            delivery_error = None if current is None else current.delivery_error
+            if observation.sequence != expected:
+                delivery_error = (
+                    f"Progress for {execution_id} skipped sequence {expected}; "
+                    f"received {observation.sequence}"
+                )
+            observation = replace(observation, delivery_error=delivery_error)
             self._progress_by_execution_id[execution_id] = observation
+        if delivery_error is not None:
+            raise RuntimeError(delivery_error)
         if self.progress_callback is not None:
-            self.progress_callback(data)
+            try:
+                self.progress_callback(data)
+            except Exception as error:
+                with self._progress_lock:
+                    self._progress_by_execution_id[execution_id] = replace(
+                        observation, delivery_error=f"Progress callback failed: {error}"
+                    )
+                raise
+
+    def _delivered_progress_through(self, execution_id: str, sequence: int) -> bool:
+        observation = self.progress_observation(execution_id)
+        if observation is None:
+            return False
+        if observation.delivery_error is not None:
+            raise RuntimeError(observation.delivery_error)
+        return observation.sequence >= sequence
 
     def progress_observation(
         self,

src/zmqruntime/execution/progress_stream.py

@@ -5,7 +5,6 @@
 import json
 import logging
 import threading
-import time
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass
 from types import MappingProxyType
@@ -16,7 +15,7 @@
     WireResponse,
     WireValue,  # noqa: F401 -- resolves WireResponse's recursive annotation in this module
 )
-from zmqruntime.messages import validate_progress_payload
+from zmqruntime.messages import MessageFields, validate_progress_payload
 
 logger = logging.getLogger(__name__)
 
@@ -43,6 +42,7 @@ class ExecutionProgressObservation:
 
     sequence: int
     event: WireResponse
+    delivery_error: str | None = None
 
     def __post_init__(self) -> None:
         if self.sequence < 1:
@@ -53,23 +53,13 @@ def __post_init__(self) -> None:
             _freeze_wire_value(self.event),
         )
 
-    @classmethod
-    def first(cls, event: WireResponse) -> ExecutionProgressObservation:
-        """Create the first retained observation for one execution."""
-
-        return cls(sequence=1, event=event)
-
-    def followed_by(self, event: WireResponse) -> ExecutionProgressObservation:
-        """Return the next immutable observation for the same execution."""
-
-        return type(self)(sequence=self.sequence + 1, event=event)
-
     def as_wire(self) -> dict:
         """Return a detached JSON-compatible projection."""
 
         return {
             type(self).sequence.__name__: self.sequence,
             type(self).event.__name__: _thaw_wire_value(self.event),
+            type(self).delivery_error.__name__: self.delivery_error,
         }
 
     @classmethod
@@ -82,7 +72,9 @@ def from_wire(cls, payload: WireResponse) -> ExecutionProgressObservation:
             raise TypeError("Execution progress sequence must be an integer")
         if not isinstance(event, Mapping):
             raise TypeError("Execution progress event must be a mapping")
-        return cls(sequence=sequence, event=event)
+        return cls(
+            sequence=sequence, event=event, delivery_error=payload.get(cls.delivery_error.__name__)
+        )
 
 
 class ProgressStreamSubscriber:
@@ -98,6 +90,8 @@ def __init__(
         self._thread: threading.Thread | None = None
         self._stop_event = threading.Event()
         self._lifecycle_lock = threading.RLock()
+        self._delivery_condition = threading.Condition()
+        self._dispatching_execution_id: str | None = None
 
     def start(self) -> None:
         with self._lifecycle_lock:
@@ -123,19 +117,37 @@ def stop(self, timeout: float = 2.0) -> None:
             if self._thread is thread:
                 self._thread = None
 
+    def wait_for_delivery(
+        self,
+        execution_id: str,
+        delivered: Callable[[], bool],
+        timeout: float,
+    ) -> None:
+        """Wait for this execution's callback dispatch, not merely its socket read."""
+        with self._delivery_condition:
+            complete = self._delivery_condition.wait_for(
+                lambda: self._dispatching_execution_id != execution_id and delivered(),
+                timeout=timeout,
+            )
+        if not complete:
+            raise TimeoutError(
+                f"Progress delivery for {execution_id} did not complete within {timeout:.3f}s"
+            )
+
     def _listen_loop(self) -> None:
         logger.info("Progress listener loop started")
         message_count = 0
         try:
             while not self._stop_event.is_set():
                 socket = self._socket_provider()
                 if socket is None:
-                    time.sleep(0.1)
+                    self._stop_event.wait(0.1)
+                    continue
+                if not socket.poll(timeout=50, flags=zmq.POLLIN):
                     continue
                 try:
                     message = socket.recv_string(zmq.NOBLOCK)
                 except zmq.Again:
-                    time.sleep(0.05)
                     continue
 
                 if self._dispatch_message(message):
@@ -155,7 +167,14 @@ def _dispatch_message(self, message: str) -> bool:
         try:
             data = json.loads(message)
             validate_progress_payload(data)
-            self._callback(data)
+            with self._delivery_condition:
+                self._dispatching_execution_id = data.get(MessageFields.EXECUTION_ID)
+            try:
+                self._callback(data)
+            finally:
+                with self._delivery_condition:
+                    self._dispatching_execution_id = None
+                    self._delivery_condition.notify_all()
         except Exception as error:
             logger.exception("Progress message dispatch failed: %s", error)
             return False

src/zmqruntime/execution/server.py

@@ -16,6 +16,7 @@
 
 from zmqruntime.config import ZMQConfig
 from zmqruntime.execution.lifecycle import InMemoryExecutionLifecycleEngine
+from zmqruntime.execution.progress_stream import ExecutionProgressObservation
 from zmqruntime.messages import (
     CancelRequest,
     ControlMessageType,
@@ -409,10 +410,20 @@ def _handle_register_progress(self, msg):
         request, error = self._validate_and_parse(msg, ProgressRegistrationRequest)
         if error:
             return error
-        self._progress_subscribers.add(request.client_id)
+        with self._progress_publish_lock:
+            self._progress_subscribers.add(request.client_id)
+            observations = {
+                execution_id: ExecutionProgressObservation(
+                    sequence=record.progress_sequence,
+                    event=record.progress_event,
+                ).as_wire()
+                for execution_id, record in self.active_executions.items()
+                if record.progress_event is not None
+            }
         return {
             MessageFields.STATUS: ResponseType.OK.value,
             MessageFields.MESSAGE: "Progress subscriber registered",
+            MessageFields.PROGRESS_OBSERVATIONS: observations,
             MessageFields.CLIENT_ID: request.client_id,
             MessageFields.PROGRESS_SUBSCRIBERS: len(self._progress_subscribers),
         }
@@ -433,7 +444,18 @@ def send_progress_update(self, progress_update: dict) -> None:
         from zmqruntime.messages import validate_progress_payload
 
         validate_progress_payload(progress_update)
-        self.progress_queue.put(progress_update)
+        execution_id = progress_update[MessageFields.EXECUTION_ID]
+        with self._progress_publish_lock:
+            record = self.active_executions[execution_id]
+            observation = ExecutionProgressObservation(
+                sequence=record.progress_sequence + 1,
+                event=progress_update,
+            )
+            event = observation.as_wire()["event"]
+            record.progress_sequence = observation.sequence
+            event[MessageFields.PROGRESS_SEQUENCE] = observation.sequence
+            record.progress_event = event
+            self.progress_queue.put(event)
 
     def _get_worker_info(self):
         try:

src/zmqruntime/execution/wait_policy.py

@@ -75,6 +75,7 @@ def wait(self, execution_id: str, policy: WaitPolicy) -> dict[str, WireValue]:
                         result = {
                             MessageFields.STATUS: ExecutionStatus.FAILED.value,
                             MessageFields.EXECUTION_ID: execution_id,
+                            MessageFields.EXECUTION: status_response[MessageFields.EXECUTION],
                         }
                         if record.error is not None:
                             result[MessageFields.MESSAGE] = record.error
@@ -83,6 +84,7 @@ def wait(self, execution_id: str, policy: WaitPolicy) -> dict[str, WireValue]:
                         return {
                             MessageFields.STATUS: ExecutionStatus.CANCELLED.value,
                             MessageFields.EXECUTION_ID: execution_id,
+                            MessageFields.EXECUTION: status_response[MessageFields.EXECUTION],
                             MessageFields.MESSAGE: "Execution was cancelled",
                         }
 

src/zmqruntime/messages.py

@@ -302,6 +302,8 @@ class MessageFields:
     METADATA = "metadata"
     CLIENT_ID = "client_id"
     PROGRESS_SUBSCRIBERS = "progress_subscribers"
+    PROGRESS_SEQUENCE = "progress_sequence"
+    PROGRESS_OBSERVATIONS = "progress_observations"
     PROCESS_IDENTITY = "process_identity"
     CONTROL_CAPABILITIES = "control_capabilities"
     APPLICATION = "application"
@@ -796,6 +798,14 @@ class ExecutionRecord:
     results_summary: Optional[Dict[str, Any]] = None
     compile_only: bool = False
     metadata: Dict[str, Any] = field(default_factory=dict)
+    progress_sequence: int = 0
+    progress_event: dict | None = field(default=None, repr=False)
+
+    def __post_init__(self) -> None:
+        if isinstance(self.progress_sequence, bool) or not isinstance(self.progress_sequence, int):
+            raise TypeError("Execution progress sequence must be an integer")
+        if self.progress_sequence < 0:
+            raise ValueError("Execution progress sequence cannot be negative")
 
     def set_extra(self, key: str, value: Any) -> None:
         self.metadata[key] = value
@@ -835,6 +845,7 @@ def _to_transport_value(value: Any) -> Any:
             MessageFields.END_TIME: _to_transport_value(self.end_time),
             MessageFields.ERROR: _to_transport_value(self.error),
             MessageFields.COMPILE_ONLY: self.compile_only,
+            MessageFields.PROGRESS_SEQUENCE: self.progress_sequence,
         }
         if self.traceback is not None:
             result[MessageFields.TRACEBACK] = _to_transport_value(self.traceback)
@@ -855,6 +866,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "ExecutionRecord":
             MessageFields.TRACEBACK,
             MessageFields.RESULTS_SUMMARY,
             MessageFields.COMPILE_ONLY,
+            MessageFields.PROGRESS_SEQUENCE,
         }
         metadata = {k: v for k, v in data.items() if k not in known}
         return cls(
@@ -869,6 +881,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "ExecutionRecord":
             results_summary=data.get(MessageFields.RESULTS_SUMMARY),
             compile_only=bool(data.get(MessageFields.COMPILE_ONLY, False)),
             metadata=metadata,
+            progress_sequence=data.get(MessageFields.PROGRESS_SEQUENCE, 0),
         )