Nexus serialization context workflows activities

changelog/added/wandering-wombat.md

@@ -0,0 +1,2 @@
+Nexus operation handlers now propagate their `NexusSerializationContext` to the workflows and standalone activities they start, which encode their results and failures (and a standalone activity's heartbeat details) with it because these are returned to the Nexus caller. Clients decode them with the context the server returns, including handles obtained by ID and standalone Nexus operation handles. Requires a server that propagates the context.
+Added `activity.Info.nexus_serialization_context`. To complete a Nexus-started standalone activity asynchronously, pass it to `AsyncActivityHandle.with_context`.

temporalio/activity.py

@@ -141,6 +141,20 @@ class Info:
     activity_run_id: str | None = None
     """Run ID of this activity. None for workflow activities."""
 
+    nexus_serialization_context: (
+        temporalio.converter.NexusSerializationContext | None
+    ) = None
+    """Serialization context of the Nexus operation handler that started this activity, if any.
+
+    Only set for standalone activities. The activity's result, failure and heartbeat details are
+    encoded with it because they are returned to the Nexus caller. To complete the activity
+    asynchronously, store this context with the task token and pass it to
+    :py:meth:`temporalio.client.AsyncActivityHandle.with_context`.
+
+    .. warning::
+        This API is experimental and unstable.
+    """
+
     @property
     def in_workflow(self) -> bool:
         """Was this activity started by a workflow?"""

temporalio/client/_activity.py

@@ -28,6 +28,7 @@
 import temporalio.common
 import temporalio.converter
 import temporalio.converter._search_attributes
+import temporalio.converter._serialization_context
 from temporalio.converter import (
     ActivitySerializationContext,
     DataConverter,
@@ -437,6 +438,18 @@ def _from_resp(
             **kwargs,
         )
 
+    @functools.cached_property
+    def _caller_facing_data_converter(self) -> DataConverter:
+        # An activity started by a Nexus operation handler encodes its result, failures and
+        # heartbeat details with the handler's Nexus context, which the server returns with the
+        # activity info. Its input keeps the activity context.
+        context = temporalio.converter._serialization_context._propagated_nexus_serialization_context(
+            self.raw_info
+        )
+        if context is None:
+            return self.data_converter
+        return self.data_converter.with_context(context)
+
     def has_heartbeat_details(self) -> bool:
         """True if heartbeat details are available. Use `heartbeat_details` to retrieve them.
 
@@ -453,7 +466,7 @@ async def heartbeat_details(
         Type hints can be provided to aid data conversion.
         """
         return (
-            await self.data_converter.decode_wrapper(
+            await self._caller_facing_data_converter.decode_wrapper(
                 self.raw_info.heartbeat_details, type_hints
             )
             if self.has_heartbeat_details()
@@ -473,7 +486,9 @@ async def last_failure(self) -> BaseException | None:
         Always `None` if ``include_last_failure`` was false in the `ActivityHandle.describe` call.
         """
         return (
-            await self.data_converter.decode_failure(self.raw_info.last_failure)
+            await self._caller_facing_data_converter.decode_failure(
+                self.raw_info.last_failure
+            )
             if self.has_last_failure()
             else None
         )
@@ -515,7 +530,7 @@ async def result(self, type_hint: type | None = None) -> Any | None:
         if self.raw_outcome is None or not self.raw_outcome.HasField("result"):
             return None
         type_hints = [type_hint] if type_hint is not None else None
-        results = await self.data_converter.decode_wrapper(
+        results = await self._caller_facing_data_converter.decode_wrapper(
             self.raw_outcome.result, type_hints
         )
         if not results:
@@ -543,7 +558,9 @@ async def outcome_failure(self) -> BaseException | None:
         Always false if ``include_outcome`` was false in the `ActivityHandle.describe` call.
         """
         return (
-            await self.data_converter.decode_failure(self.raw_outcome.failure)
+            await self._caller_facing_data_converter.decode_failure(
+                self.raw_outcome.failure
+            )
             if self.raw_outcome is not None and self.raw_outcome.HasField("failure")
             else None
         )
@@ -942,6 +959,11 @@ def with_context(self, context: SerializationContext) -> Self:
         with :py:class:`ActivitySerializationContext` set as context. If you are using a custom data
         converter that makes use of this context then you can use this method to supply matching
         context data to the data converter used to serialize and encode the outbound payloads.
+
+        A standalone activity started by a Nexus operation handler encodes its result, failure and
+        heartbeat details with the handler's
+        :py:class:`temporalio.converter.NexusSerializationContext`. To complete or heartbeat such an
+        activity, pass :py:attr:`temporalio.activity.Info.nexus_serialization_context`.
         """
         data_converter = self._client.data_converter.with_context(context)
         if data_converter is self._client.data_converter:
@@ -978,6 +1000,9 @@ def __init__(
         self._known_outcome: (
             temporalio.api.activity.v1.ActivityExecutionOutcome | None
         ) = None
+        self._known_outcome_nexus_context: (
+            temporalio.converter.NexusSerializationContext | None
+        ) = None
 
     @functools.cached_property
     def _data_converter(self) -> temporalio.converter.DataConverter:
@@ -1034,18 +1059,22 @@ async def result(
             rpc_metadata=rpc_metadata, rpc_timeout=rpc_timeout
         )
 
-        # Convert outcome to failure or value
+        # Convert outcome to failure or value. An activity started by a Nexus operation handler
+        # encodes them with the handler's Nexus context, which the server returns with the outcome.
         assert self._known_outcome
+        data_converter = (
+            self._client.data_converter.with_context(self._known_outcome_nexus_context)
+            if self._known_outcome_nexus_context
+            else self._data_converter
+        )
         if self._known_outcome.HasField("failure"):
             raise ActivityFailureError(
-                cause=await self._data_converter.decode_failure(
-                    self._known_outcome.failure
-                ),
+                cause=await data_converter.decode_failure(self._known_outcome.failure),
             )
         if not self._known_outcome.result.payloads:
             return None  # type: ignore
         type_hints = [self._result_type] if self._result_type else None
-        results = await self._data_converter.decode(
+        results = await data_converter.decode(
             self._known_outcome.result.payloads, type_hints
         )
         if not results:
@@ -1080,6 +1109,9 @@ async def _poll_until_outcome(
                 )
                 if res.HasField("outcome"):
                     self._known_outcome = res.outcome
+                    self._known_outcome_nexus_context = temporalio.converter._serialization_context._propagated_nexus_serialization_context(
+                        res
+                    )
                     return
             except RPCError as err:
                 if err.status == RPCStatusCode.DEADLINE_EXCEEDED:

temporalio/client/_impl.py

@@ -29,6 +29,7 @@
 import temporalio.api.workflowservice.v1
 import temporalio.common
 import temporalio.converter
+import temporalio.converter._serialization_context
 import temporalio.exceptions
 import temporalio.nexus
 import temporalio.nexus._operation_context
@@ -264,6 +265,9 @@ async def _build_update_with_start_start_workflow_execution_request(
     ) -> temporalio.api.workflowservice.v1.StartWorkflowExecutionRequest:
         req = temporalio.api.workflowservice.v1.StartWorkflowExecutionRequest()
         await self._populate_start_workflow_execution_request(req, input)
+        temporalio.nexus._operation_context._apply_nexus_serialization_context_to_start_workflow_request(
+            req
+        )
         return req
 
     async def _populate_start_workflow_execution_request(
@@ -1676,10 +1680,8 @@ async def get_nexus_operation_result(
     ) -> Any:
         """Poll for nexus operation result until it's available."""
         data_converter = self._client.data_converter
-        # These three are set together or not at all: a handle that started the operation has all
-        # of them, and a handle obtained by operation ID alone has none and defaults them to "".
-        # An empty endpoint therefore means "no operation to build a context from", not "an
-        # operation named the empty string".
+        # Fallback for older servers that don't return the context when polling. Handles obtained
+        # by operation ID alone have no start request context.
         if input.endpoint and input.service and input.operation:
             data_converter = data_converter.with_context(
                 temporalio.converter.NexusSerializationContext(
@@ -1706,6 +1708,15 @@ async def get_nexus_operation_result(
                         timeout=input.rpc_timeout,
                     )
                 )
+                # The context the server returns also identifies an operation whose handle was
+                # obtained by operation ID
+                server_context = temporalio.converter._serialization_context._propagated_nexus_serialization_context(
+                    res
+                )
+                if server_context:
+                    data_converter = self._client.data_converter.with_context(
+                        server_context
+                    )
                 match res.WhichOneof("outcome"):
                     case "result":
                         type_hints = [input.result_type] if input.result_type else None

temporalio/client/_workflow.py

@@ -37,6 +37,7 @@
 import temporalio.common
 import temporalio.converter
 import temporalio.converter._search_attributes
+import temporalio.converter._serialization_context
 import temporalio.exceptions
 import temporalio.workflow
 from temporalio.converter import (
@@ -136,6 +137,17 @@ def _data_converter(self) -> temporalio.converter.DataConverter:
             )
         )
 
+    def _close_event_data_converter(
+        self, events: WorkflowHistoryEventAsyncIterator
+    ) -> temporalio.converter.DataConverter:
+        # A workflow started by a Nexus operation handler encodes its result and closing failure
+        # with the handler's Nexus context, which the server returns with the history. An
+        # interceptor may return its own iterator, which then falls back to the workflow context.
+        context = getattr(events, "_propagated_nexus_serialization_context", None)
+        if context is None:
+            return self._data_converter
+        return self._client.data_converter.with_context(context)
+
     @property
     def id(self) -> str:
         """ID of the workflow."""
@@ -225,14 +237,16 @@ async def result(
         # executions
         hist_run_id = self._result_run_id
         while True:
-            async for event in self._fetch_history_events_for_run(
+            events = self._fetch_history_events_for_run(
                 hist_run_id,
                 wait_new_event=True,
                 event_filter_type=WorkflowHistoryEventFilterType.CLOSE_EVENT,
                 skip_archival=True,
                 rpc_metadata=rpc_metadata,
                 rpc_timeout=rpc_timeout,
-            ):
+            )
+            async for event in events:
+                data_converter = self._close_event_data_converter(events)
                 if event.HasField("workflow_execution_completed_event_attributes"):
                     complete_attr = event.workflow_execution_completed_event_attributes
                     # Follow execution
@@ -241,7 +255,7 @@ async def result(
                         break
                     # Ignoring anything after the first response like TypeScript
                     type_hints = [self._result_type] if self._result_type else None
-                    results = await self._data_converter.decode_wrapper(
+                    results = await data_converter.decode_wrapper(
                         complete_attr.result,
                         type_hints,
                     )
@@ -257,32 +271,22 @@ async def result(
                         hist_run_id = fail_attr.new_execution_run_id
                         break
                     raise WorkflowFailureError(
-                        cause=await self._data_converter.decode_failure(
-                            fail_attr.failure
-                        ),
+                        cause=await data_converter.decode_failure(fail_attr.failure),
                     )
                 elif event.HasField("workflow_execution_canceled_event_attributes"):
                     cancel_attr = event.workflow_execution_canceled_event_attributes
                     raise WorkflowFailureError(
                         cause=temporalio.exceptions.CancelledError(
                             "Workflow cancelled",
-                            *(
-                                await self._data_converter.decode_wrapper(
-                                    cancel_attr.details
-                                )
-                            ),
+                            *(await data_converter.decode_wrapper(cancel_attr.details)),
                         )
                     )
                 elif event.HasField("workflow_execution_terminated_event_attributes"):
                     term_attr = event.workflow_execution_terminated_event_attributes
                     raise WorkflowFailureError(
                         cause=temporalio.exceptions.TerminatedError(
                             term_attr.reason or "Workflow terminated",
-                            *(
-                                await self._data_converter.decode_wrapper(
-                                    term_attr.details
-                                )
-                            ),
+                            *(await data_converter.decode_wrapper(term_attr.details)),
                         ),
                     )
                 elif event.HasField("workflow_execution_timed_out_event_attributes"):
@@ -1760,6 +1764,9 @@ def __init__(
             None | (Sequence[temporalio.api.history.v1.HistoryEvent])
         ) = None
         self._current_page_index = 0
+        self._propagated_nexus_serialization_context: (
+            temporalio.converter.NexusSerializationContext | None
+        ) = None
 
     @property
     def current_page_index(self) -> int:
@@ -1811,6 +1818,13 @@ async def fetch_next_page(self, *, page_size: int | None = None) -> None:  # typ
         self._current_page = list(resp.history.events)
         self._current_page_index = 0
         self._next_page_token = resp.next_page_token or None
+        # Pages after the first may omit the context, such as the page holding the close event
+        # after a long poll times out
+        context = temporalio.converter._serialization_context._propagated_nexus_serialization_context(
+            resp
+        )
+        if context is not None:
+            self._propagated_nexus_serialization_context = context
 
     def __aiter__(self) -> WorkflowHistoryEventAsyncIterator:
         """Return self as the iterator."""

temporalio/converter/_serialization_context.py

@@ -5,6 +5,7 @@
 from abc import ABC
 from dataclasses import dataclass
 
+import google.protobuf.message
 from typing_extensions import Self
 
 
@@ -102,15 +103,17 @@ class ActivitySerializationContext(SerializationContext):
 class NexusSerializationContext(SerializationContext):
     """Serialization context for Nexus operation payloads.
 
-    Callers receive this context when encoding inputs and decoding results or failures. The context
-    is not propagated to a handler that completes an asynchronous operation. Handlers receive it
-    when decoding inputs, encoding synchronous results, and encoding failures produced while
-    handling a Nexus task.
+    Callers receive this context when encoding inputs and decoding results or failures. Handlers
+    receive it when decoding inputs, encoding synchronous results, and encoding failures produced
+    while handling a Nexus task.
 
-    A standalone operation handle retains the context used to start the operation and uses it to
-    decode the result, including when the start request returns an existing operation. A handle
-    created with :py:meth:`temporalio.client.Client.get_nexus_operation_handle` has no endpoint,
-    service, or operation information and therefore decodes without Nexus context.
+    Workflows and standalone activities started by a handler receive this context through the
+    server and use it for payloads returned to the caller: results, failures, and standalone
+    activity heartbeat details. Clients decode these with the context the server returns. See
+    :py:attr:`temporalio.activity.Info.nexus_serialization_context` for asynchronous completion.
+
+    A standalone operation handle decodes with the context the server returns, falling back to the
+    context used to start the operation on older servers.
 
     A failure encoded by a handler is later decoded by a caller. Because some operation paths may
     lack this context, contextual encodings must be self-describing and decoders must continue to
@@ -130,6 +133,24 @@ class NexusSerializationContext(SerializationContext):
     """Nexus operation name."""
 
 
+def _propagated_nexus_serialization_context(  # pyright: ignore[reportUnusedFunction]
+    message: google.protobuf.message.Message,
+) -> NexusSerializationContext | None:
+    """Return the Nexus serialization context the server propagated on a message, if any.
+
+    The server sets ``propagated_nexus_serialization_context`` on executions started by a Nexus
+    operation handler.
+    """
+    if not message.HasField("propagated_nexus_serialization_context"):
+        return None
+    context = getattr(message, "propagated_nexus_serialization_context")
+    return NexusSerializationContext(
+        endpoint=context.endpoint,
+        service=context.service,
+        operation=context.operation,
+    )
+
+
 class WithSerializationContext(ABC):
     """Interface for classes that can use serialization context.