NexusSerializationContext for data/failure converters

CHANGELOG.md

@@ -27,6 +27,10 @@ to include examples, links to docs, or any other relevant information.
 - Added GCP Cloud Run serverless-worker OpenTelemetry plugin in `temporalio.contrib.opentelemetry`.
 - Added new options to ActivityHandle.describe() to retrieve associated payloads, such as activity input and outcome.
 - New properties and methods in ActivityExecution and ActivityExecutionDescription.
+- Added experimental `temporalio.converter.NexusSerializationContext` support for Nexus callers
+  and handlers. Callers use it for inputs, results, and failures; handlers use it for inputs,
+  synchronous results, and failures. Asynchronous handler results and detached standalone handles
+  are not yet supported. Standalone `USE_EXISTING` handles use their start request's context.
 
 ### Changed
 

temporalio/client/_impl.py

@@ -114,7 +114,6 @@
     NexusOperationExecutionAsyncIterator,
     NexusOperationExecutionCount,
     NexusOperationExecutionDescription,
-    NexusOperationFailureError,
     NexusOperationHandle,
 )
 from ._schedule import (
@@ -1549,6 +1548,12 @@ async def start_nexus_operation(
         self, input: StartNexusOperationInput
     ) -> NexusOperationHandle[Any]:
         """Start a nexus operation and return a handle to it."""
+        nexus_context = temporalio.converter.NexusSerializationContext(
+            endpoint=input.endpoint,
+            service=input.service,
+            operation=input.operation,
+        )
+        data_converter = self._client.data_converter.with_context(nexus_context)
         req = temporalio.api.workflowservice.v1.StartNexusOperationExecutionRequest(
             namespace=self._client.namespace,
             identity=self._client.identity,
@@ -1575,7 +1580,7 @@ async def start_nexus_operation(
             req.start_to_close_timeout.FromTimedelta(input.start_to_close_timeout)
 
         # Set input payload
-        encoded = await self._client.data_converter.encode([input.arg])
+        encoded = await data_converter.encode([input.arg])
         if encoded:
             req.input.CopyFrom(encoded[0])
 
@@ -1620,6 +1625,7 @@ async def start_nexus_operation(
             result_type=input.result_type,
             endpoint=input.endpoint,
             service=input.service,
+            _nexus_serialization_context=nexus_context,
         )
 
     async def describe_nexus_operation(
@@ -1637,14 +1643,22 @@ async def describe_nexus_operation(
             metadata=input.rpc_metadata,
             timeout=input.rpc_timeout,
         )
+        nexus_context = temporalio.converter.NexusSerializationContext(
+            endpoint=resp.info.endpoint,
+            service=resp.info.service,
+            operation=resp.info.operation,
+        )
         return await NexusOperationExecutionDescription._from_execution_info(
             info=resp.info,
             data_converter=self._client.data_converter,
+            failure_data_converter=self._client.data_converter.with_context(
+                nexus_context
+            ),
         )
 
     async def get_nexus_operation_result(
         self, input: GetNexusOperationResultInput
-    ) -> Any:
+    ) -> temporalio.api.workflowservice.v1.PollNexusOperationExecutionResponse:
         """Poll for nexus operation result until it's available."""
         req = temporalio.api.workflowservice.v1.PollNexusOperationExecutionRequest(
             namespace=self._client.namespace,
@@ -1664,24 +1678,8 @@ async def get_nexus_operation_result(
                         timeout=input.rpc_timeout,
                     )
                 )
-                match res.WhichOneof("outcome"):
-                    case "result":
-                        type_hints = [input.result_type] if input.result_type else None
-                        [result] = await self._client.data_converter.decode(
-                            [res.result], type_hints
-                        )
-                        return result
-
-                    case "failure":
-                        raise NexusOperationFailureError(
-                            cause=await self._client.data_converter.decode_failure(
-                                res.failure
-                            )
-                        )
-
-                    case None:
-                        # poll again
-                        pass
+                if res.WhichOneof("outcome") is not None:
+                    return res
             except RPCError as err:
                 match err.status:
                     case RPCStatusCode.DEADLINE_EXCEEDED:

temporalio/client/_interceptor.py

@@ -18,9 +18,7 @@
 import temporalio.api.common.v1
 import temporalio.api.workflowservice.v1
 import temporalio.common
-from temporalio.converter import (
-    DataConverter,
-)
+from temporalio.converter import DataConverter
 
 if TYPE_CHECKING:
     from ._activity import (
@@ -659,7 +657,6 @@ class GetNexusOperationResultInput:
     run_id: str | None
     rpc_metadata: Mapping[str, str | bytes]
     rpc_timeout: timedelta | None
-    result_type: type[Any] | None
 
 
 @dataclass
@@ -1003,9 +1000,11 @@ async def describe_nexus_operation(
 
     async def get_nexus_operation_result(
         self, input: GetNexusOperationResultInput
-    ) -> Any:
+    ) -> temporalio.api.workflowservice.v1.PollNexusOperationExecutionResponse:
         """Called for every :py:meth:`NexusOperationHandle.result` call.
 
+        The raw response is decoded by the handle after interception.
+
         .. warning::
            This API is experimental and unstable.
         """

temporalio/client/_nexus.py

@@ -291,8 +291,10 @@ async def _from_execution_info(
         cls,
         info: temporalio.api.nexus.v1.NexusOperationExecutionInfo,
         data_converter: temporalio.converter.DataConverter,
+        failure_data_converter: temporalio.converter.DataConverter | None = None,
     ) -> Self:
         """Create from raw proto nexus operation execution info."""
+        failure_data_converter = failure_data_converter or data_converter
         return cls(
             _data_converter=data_converter,
             operation_id=info.operation_id,
@@ -360,7 +362,9 @@ async def _from_execution_info(
             last_attempt_failure=(
                 cast(
                     BaseException | None,
-                    await data_converter.decode_failure(info.last_attempt_failure),
+                    await failure_data_converter.decode_failure(
+                        info.last_attempt_failure
+                    ),
                 )
                 if info.HasField("last_attempt_failure")
                 else None
@@ -376,7 +380,7 @@ async def _from_execution_info(
             identity=info.identity,
             cancellation_info=(
                 await NexusOperationExecutionCancellationInfo._from_cancellation_info(
-                    info.cancellation_info, data_converter
+                    info.cancellation_info, failure_data_converter
                 )
                 if info.HasField("cancellation_info")
                 else None
@@ -1066,6 +1070,9 @@ def __init__(
         result_type: type | None = None,
         endpoint: str = "",
         service: str = "",
+        _nexus_serialization_context: (
+            temporalio.converter.NexusSerializationContext | None
+        ) = None,
     ) -> None:
         """Create nexus operation handle."""
         self._client = client
@@ -1074,6 +1081,7 @@ def __init__(
         self._result_type = result_type
         self._endpoint = endpoint
         self._service = service
+        self._nexus_serialization_context = _nexus_serialization_context
         # the default value is `_arg_unset` because ReturnType could be None
         self._known_outcome: ReturnType | NexusOperationFailureError | object = (
             temporalio.common._arg_unset
@@ -1130,22 +1138,38 @@ async def result(
             RPCError: Operation result could not be fetched for some reason.
         """
         if self._known_outcome is temporalio.common._arg_unset:
-            try:
-                self._known_outcome = (
-                    await self._client._impl.get_nexus_operation_result(
-                        GetNexusOperationResultInput(
-                            operation_id=self._operation_id,
-                            run_id=self._run_id,
-                            result_type=self._result_type,
-                            rpc_metadata=rpc_metadata,
-                            rpc_timeout=rpc_timeout,
-                        )
-                    )
+            response = await self._client._impl.get_nexus_operation_result(
+                GetNexusOperationResultInput(
+                    operation_id=self._operation_id,
+                    run_id=self._run_id,
+                    rpc_metadata=rpc_metadata,
+                    rpc_timeout=rpc_timeout,
+                )
+            )
+
+            data_converter = self._client.data_converter
+            if self._nexus_serialization_context is not None:
+                data_converter = data_converter.with_context(
+                    self._nexus_serialization_context
                 )
-                return cast(ReturnType, self._known_outcome)
-            except NexusOperationFailureError as failure:
-                self._known_outcome = failure
-                raise
+            match response.WhichOneof("outcome"):
+                case "result":
+                    type_hints = [self._result_type] if self._result_type else None
+                    [result] = await data_converter.decode(
+                        [response.result], type_hints
+                    )
+                    self._known_outcome = result
+                    return cast(ReturnType, result)
+                case "failure":
+                    operation_failure = NexusOperationFailureError(
+                        cause=await data_converter.decode_failure(response.failure)
+                    )
+                    self._known_outcome = operation_failure
+                    raise operation_failure
+                case None:
+                    raise RuntimeError(
+                        "Nexus operation result response did not contain an outcome"
+                    )
         elif isinstance(self._known_outcome, NexusOperationFailureError):
             raise self._known_outcome
         else:

temporalio/converter/__init__.py

@@ -50,6 +50,7 @@
 )
 from temporalio.converter._serialization_context import (
     ActivitySerializationContext,
+    NexusSerializationContext,
     SerializationContext,
     WithSerializationContext,
     WorkflowSerializationContext,
@@ -82,6 +83,7 @@
     "JSONProtoPayloadConverter",
     "JSONTypeConverter",
     "JSONTypeConverterUnhandled",
+    "NexusSerializationContext",
     "PayloadCodec",
     "PayloadConverter",
     "SerializationContext",

temporalio/converter/_serialization_context.py

@@ -28,6 +28,10 @@ class SerializationContext(ABC):
         context type is :py:class:`ActivitySerializationContext` and the workflow ID is that of the
         currently-executing workflow. ActivitySerializationContext is also set on data converter
         operations in the activity context.
+
+        When operating on a Nexus operation payload, the context type is
+        :py:class:`NexusSerializationContext` and identifies the Nexus endpoint, service, and
+        resolved operation name.
     """
 
     pass
@@ -94,6 +98,35 @@ class ActivitySerializationContext(SerializationContext):
     """Whether the activity is a local activity started from a workflow."""
 
 
+@dataclass(frozen=True)
+class NexusSerializationContext(SerializationContext):
+    """Serialization context for Nexus operation payloads.
+
+    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.
+
+    The context is not propagated to the eventual result of an asynchronous operation. Standalone
+    operation handles use the context of their start request, including when an existing operation
+    is returned, while handles created without starting an operation do not receive it.
+
+    Callers and handlers receive this context on opposite sides of failure conversion. Contextual
+    encodings should therefore be self-describing and support legacy payloads without context.
+
+    .. warning::
+        This API is experimental and unstable.
+    """
+
+    endpoint: str
+    """Nexus endpoint name."""
+
+    service: str
+    """Nexus service name."""
+
+    operation: str
+    """Nexus operation name."""
+
+
 class WithSerializationContext(ABC):
     """Interface for classes that can use serialization context.
 

temporalio/worker/_command_aware_visitor.py

@@ -24,6 +24,7 @@
     ScheduleNexusOperation,
     SignalExternalWorkflowExecution,
     StartChildWorkflowExecution,
+    WorkflowCommand,
 )
 
 
@@ -115,6 +116,18 @@ async def _visit_coresdk_workflow_commands_ScheduleNexusOperation(
         with current_command(CommandType.COMMAND_TYPE_SCHEDULE_NEXUS_OPERATION, o.seq):
             await super()._visit_coresdk_workflow_commands_ScheduleNexusOperation(fs, o)
 
+    async def _visit_coresdk_workflow_commands_WorkflowCommand(
+        self, fs: VisitorFunctions, o: WorkflowCommand
+    ) -> None:
+        if o.HasField("schedule_nexus_operation"):
+            with current_command(
+                CommandType.COMMAND_TYPE_SCHEDULE_NEXUS_OPERATION,
+                o.schedule_nexus_operation.seq,
+            ):
+                await super()._visit_coresdk_workflow_commands_WorkflowCommand(fs, o)
+        else:
+            await super()._visit_coresdk_workflow_commands_WorkflowCommand(fs, o)
+
     # Workflow activation jobs with payloads
     async def _visit_coresdk_workflow_activation_ResolveActivity(
         self, fs: VisitorFunctions, o: ResolveActivity