[FROZEN][Bridge] Let Workflow code wait on a Nexus operation's progress.

temporalio/worker/_workflow_instance.py

@@ -4,6 +4,7 @@
 
 import asyncio
 import collections
+import contextlib
 import contextvars
 import inspect
 import json
@@ -13,6 +14,7 @@
 import threading
 import time
 import traceback
+import types
 import warnings
 from abc import ABC, abstractmethod
 from collections import deque
@@ -47,6 +49,7 @@
 import temporalio.activity
 import temporalio.api.common.v1
 import temporalio.api.enums.v1
+import temporalio.api.nexus.v1
 import temporalio.api.sdk.v1
 import temporalio.api.stream.v1
 import temporalio.bridge.proto.activity_result
@@ -1184,6 +1187,34 @@ def _apply_resolve_nexus_operation(
         else:
             raise RuntimeError("Nexus operation did not have a result")
 
+    def _apply_nexus_operation_progress(
+        self, progress: temporalio.api.nexus.v1.NexusOperationProgress
+    ) -> None:
+        if not progress.HasField("operation_key"):
+            return
+        key = progress.operation_key
+        # Only a started operation that has not resolved has a pending handle
+        # with a token. A resolution in the same task comes before the
+        # attachments, so its progress finds no handle and is dropped. A token
+        # is unique only within one handler, so two operations with one key are
+        # one handler operation, and both get the progress.
+        for handle in self._pending_nexus_operations.values():
+            if (
+                handle._input.endpoint == key.endpoint
+                and handle._input.service == key.service
+                and handle._input.operation_name == key.operation
+                and handle.operation_token == key.operation_token
+            ):
+                handle._resolve_progress(
+                    temporalio.workflow.NexusOperationProgress(
+                        position=progress.position,
+                        counter=progress.counter,
+                        # Read-only, so Workflow code can't change what the next
+                        # reader of latest_progress sees.
+                        metadata=types.MappingProxyType(dict(progress.metadata)),
+                    )
+                )
+
     def _apply_resolve_request_cancel_external_workflow(
         self,
         job: temporalio.bridge.proto.workflow_activation.ResolveRequestCancelExternalWorkflow,
@@ -1286,11 +1317,18 @@ def _apply_workflow_task_attachments(
         # Other attachments belong to features this Worker does not run, so
         # they are left alone.
         woke = False
+        progress: list[temporalio.api.nexus.v1.NexusOperationProgress] = []
         for attachment in job.attachments:
             wake = temporalio.api.stream.v1.StreamWake()
             if attachment.Unpack(wake):
                 self._stream_wakes.append(wake)
                 woke = True
+                continue
+            one = temporalio.api.nexus.v1.NexusOperationProgress()
+            if attachment.Unpack(one):
+                progress.append(one)
+        for one in fold_nexus_operation_progress(progress):
+            self._apply_nexus_operation_progress(one)
         if woke:
             # A wake resolves no future, and wait conditions are checked only
             # after a ready callback ran, so give the loop one.
@@ -3894,6 +3932,13 @@ def __init__(
         self._payload_converter = payload_converter
         self._failure_converter = failure_converter
         self._event_group_markers = _capture_event_group_markers(input.event_groups)
+        self._latest_progress: temporalio.workflow.NexusOperationProgress | None = None
+        self._progress_waiters: list[
+            tuple[
+                int, asyncio.Future[temporalio.workflow.NexusOperationProgress | None]
+            ]
+        ] = []
+        self._result_fut.add_done_callback(lambda _: self._end_progress_waits())
 
     @property
     def operation_token(self) -> str | None:
@@ -3902,6 +3947,56 @@ def operation_token(self) -> str | None:
         except BaseException:
             return None
 
+    @property
+    def latest_progress(self) -> temporalio.workflow.NexusOperationProgress | None:
+        return self._latest_progress
+
+    async def progress(
+        self, *, after_counter: int = 0
+    ) -> temporalio.workflow.NexusOperationProgress | None:
+        latest = self._latest_progress
+        if latest is not None and latest.counter > after_counter:
+            return latest
+        if self._result_fut.done():
+            return None
+        waiter: asyncio.Future[temporalio.workflow.NexusOperationProgress | None] = (
+            self._instance.create_future()
+        )
+        entry = (after_counter, waiter)
+        self._progress_waiters.append(entry)
+        try:
+            return await waiter
+        finally:
+            # A wait given up on, such as one under a timeout in a loop, would
+            # otherwise stay in the list until the next progress arrives.
+            with contextlib.suppress(ValueError):
+                self._progress_waiters.remove(entry)
+
+    def _resolve_progress(
+        self, progress: temporalio.workflow.NexusOperationProgress
+    ) -> None:
+        if (
+            self._latest_progress is not None
+            and progress.counter <= self._latest_progress.counter
+        ):
+            return
+        self._latest_progress = progress
+        still_waiting = []
+        for after_counter, waiter in self._progress_waiters:
+            if waiter.done():
+                continue
+            if progress.counter > after_counter:
+                waiter.set_result(progress)
+            else:
+                still_waiting.append((after_counter, waiter))
+        self._progress_waiters = still_waiting
+
+    def _end_progress_waits(self) -> None:
+        for _, waiter in self._progress_waiters:
+            if not waiter.done():
+                waiter.set_result(None)
+        self._progress_waiters = []
+
     def __await__(self) -> Generator[Any, Any, OutputT]:
         return self._task.__await__()
 
@@ -4032,6 +4127,31 @@ def _apply_command(self) -> None:
         command.event_group_markers.extend(self._event_group_markers)
 
 
+def fold_nexus_operation_progress(
+    progress: list[temporalio.api.nexus.v1.NexusOperationProgress],
+) -> list[temporalio.api.nexus.v1.NexusOperationProgress]:
+    """Fold the progress of one Workflow Task to one per operation key.
+
+    The key is the endpoint, service, operation and token. The progress with
+    the highest counter for a key stays, and the first of equal counters
+    stays. Keys keep the order of their first progress. Progress with no key
+    is dropped. This is the fold rule that every SDK shares, and the
+    sdk-streams vectors ``testdata/nexus_progress_folds.json`` check it.
+    """
+    folded: dict[
+        tuple[str, str, str, str], temporalio.api.nexus.v1.NexusOperationProgress
+    ] = {}
+    for one in progress:
+        if not one.HasField("operation_key"):
+            continue
+        key = one.operation_key
+        id = (key.endpoint, key.service, key.operation, key.operation_token)
+        held = folded.get(id)
+        if held is None or one.counter > held.counter:
+            folded[id] = one
+    return list(folded.values())
+
+
 def _encode_search_attributes(
     attributes: (
         temporalio.common.SearchAttributes | temporalio.common.TypedSearchAttributes

temporalio/workflow/__init__.py

@@ -137,6 +137,7 @@
     NexusClient,
     NexusOperationCancellationType,
     NexusOperationHandle,
+    NexusOperationProgress,
     _NexusClient,
     create_nexus_client,
 )
@@ -253,6 +254,7 @@
     "NexusClient",
     "NexusOperationCancellationType",
     "NexusOperationHandle",
+    "NexusOperationProgress",
     "create_nexus_client",
     "LoggerAdapter",
     "SandboxImportNotificationPolicy",

temporalio/workflow/_nexus.py

@@ -2,6 +2,7 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import Awaitable, Callable, Generator, Mapping, Sequence
+from dataclasses import dataclass
 from datetime import timedelta
 from enum import IntEnum
 from typing import Any, Generic, overload
@@ -21,10 +22,36 @@
     "NexusClient",
     "NexusOperationCancellationType",
     "NexusOperationHandle",
+    "NexusOperationProgress",
     "create_nexus_client",
 ]
 
 
+@dataclass(frozen=True)
+class NexusOperationProgress:
+    """Progress a started Nexus operation reported to its caller.
+
+    Progress says that the operation's output moved, not what the output is.
+    The Workflow reads the output itself, starting from :py:attr:`position`.
+
+    .. warning::
+       This API is experimental.
+    """
+
+    position: str
+    """Where the operation's output stands, in the handler's terms, such as a
+    stream cursor."""
+
+    counter: int
+    """Orders the progress of one operation. It only grows."""
+
+    metadata: Mapping[str, str]
+    """Small details from the handler."""
+
+    # The metadata mapping isn't hashable, so a hash would always raise.
+    __hash__ = None  # type: ignore[assignment]
+
+
 class NexusOperationHandle(Generic[OutputT]):
     """Handle for interacting with a Nexus operation."""
 
@@ -43,6 +70,37 @@ def operation_token(self) -> str | None:
         """The operation token for this handle."""
         raise NotImplementedError
 
+    @property
+    def latest_progress(self) -> NexusOperationProgress | None:
+        """The latest progress the operation reported, or ``None`` if none yet.
+
+        .. warning::
+           This API is experimental.
+        """
+        raise NotImplementedError
+
+    async def progress(
+        self,
+        *,
+        after_counter: int = 0,  # type: ignore[reportUnusedParameter]
+    ) -> NexusOperationProgress | None:
+        """Wait for progress newer than ``after_counter``.
+
+        Returns the latest progress as soon as its counter is higher than
+        ``after_counter``, which may be at once. Returns ``None`` if the
+        operation resolves first, since no newer progress can follow. Progress
+        comes from History, so a replay returns the same values at the same
+        points.
+
+        Progress is a hint that the operation's output moved to
+        :py:attr:`NexusOperationProgress.position`. A burst of progress folds
+        to its highest counter, so intermediate counters may be skipped.
+
+        .. warning::
+           This API is experimental.
+        """
+        raise NotImplementedError
+
 
 class NexusOperationCancellationType(IntEnum):
     """Defines behavior of a Nexus operation when the caller workflow initiates cancellation.

tests/contrib/streams/_support.py

@@ -8,7 +8,9 @@
 import subprocess
 import uuid
 from pathlib import Path
-from typing import Any
+from typing import Any, TypeVar
+
+from google.protobuf.message import Message
 
 from temporalio.api.workflowservice.v1 import (
     DescribeStreamNotifierRequest,
@@ -19,6 +21,9 @@
 from temporalio.contrib.streams._plugin import call, store_for_client
 from temporalio.contrib.streams.memory import MemoryStreams
 from temporalio.converter import DataConverter
+from temporalio.service import RPCError, RPCStatusCode
+
+ResponseT = TypeVar("ResponseT", bound=Message)
 
 
 async def connect_with(
@@ -57,6 +62,40 @@ async def describe_stream_notifier(
     return await call(service.describe_stream_notifier(request))
 
 
+def _serialize(message: Message) -> bytes:
+    return message.SerializeToString()
+
+
+async def call_workflow_service(
+    client: Client, rpc: str, request: Message, response_type: type[ResponseT]
+) -> ResponseT:
+    """Call a workflow service RPC that the client has no stub for, by its method path.
+
+    The stream runtime calls the notifier RPCs this way. A test that plays a
+    producer or a handler calls them over a plain gRPC channel to the client's
+    server, and gets a failure as the ``RPCError`` that a stub raises.
+    """
+    # Imported here, since the Workflow sandbox imports this module again.
+    import grpc
+
+    async with grpc.aio.insecure_channel(
+        client.service_client.config.target_host
+    ) as channel:
+        method = channel.unary_unary(
+            f"/temporal.api.workflowservice.v1.WorkflowService/{rpc}",
+            request_serializer=_serialize,
+            response_deserializer=response_type.FromString,
+        )
+        try:
+            return await method(
+                request, metadata=(("temporal-namespace", client.namespace),)
+            )
+        except grpc.aio.AioRpcError as err:
+            raise RPCError(
+                err.details() or "", RPCStatusCode(err.code().value[0]), b""
+            ) from None
+
+
 @functools.cache
 def streams_crate_dir() -> Path:
     """The checkout of the temporalio-streams crate the bridge builds against.