Moving decoding logic to the nexus operation handle

Author: JoshuaFrenchwood

temporalio/client/_impl.py

@@ -114,7 +114,6 @@
     NexusOperationExecutionAsyncIterator,
     NexusOperationExecutionCount,
     NexusOperationExecutionDescription,
-    NexusOperationFailureError,
     NexusOperationHandle,
 )
 from ._schedule import (
@@ -1659,13 +1658,8 @@ async def describe_nexus_operation(
 
     async def get_nexus_operation_result(
         self, input: GetNexusOperationResultInput
-    ) -> Any:
+    ) -> temporalio.api.workflowservice.v1.PollNexusOperationExecutionResponse:
         """Poll for nexus operation result until it's available."""
-        data_converter = self._client.data_converter
-        if input._nexus_serialization_context is not None:
-            data_converter = data_converter.with_context(
-                input._nexus_serialization_context
-            )
         req = temporalio.api.workflowservice.v1.PollNexusOperationExecutionRequest(
             namespace=self._client.namespace,
             operation_id=input.operation_id,
@@ -1684,20 +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 data_converter.decode([res.result], type_hints)
-                        return result
-
-                    case "failure":
-                        raise NexusOperationFailureError(
-                            cause=await 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,10 +18,7 @@
 import temporalio.api.common.v1
 import temporalio.api.workflowservice.v1
 import temporalio.common
-from temporalio.converter import (
-    DataConverter,
-    NexusSerializationContext,
-)
+from temporalio.converter import DataConverter
 
 if TYPE_CHECKING:
     from ._activity import (
@@ -660,8 +657,6 @@ class GetNexusOperationResultInput:
     run_id: str | None
     rpc_metadata: Mapping[str, str | bytes]
     rpc_timeout: timedelta | None
-    result_type: type[Any] | None
-    _nexus_serialization_context: NexusSerializationContext | None = None
 
 
 @dataclass
@@ -1005,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

@@ -1138,21 +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,
-                        _nexus_serialization_context=self._nexus_serialization_context,
-                    )
+            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:

tests/nexus/test_standalone_operations.py

@@ -22,6 +22,7 @@
 )
 
 import temporalio.api.enums.v1
+import temporalio.api.workflowservice.v1
 from temporalio import nexus, workflow
 from temporalio.client import (
     CancelNexusOperationInput,
@@ -867,7 +868,7 @@ async def describe_nexus_operation(
 
     async def get_nexus_operation_result(
         self, input: GetNexusOperationResultInput
-    ) -> Any:
+    ) -> temporalio.api.workflowservice.v1.PollNexusOperationExecutionResponse:
         self._parent.result_calls.append(input)
         return await super().get_nexus_operation_result(input)
 
@@ -980,7 +981,6 @@ async def test_interceptor_receives_inputs(client: Client, env: WorkflowEnvironm
         result_input = interceptor.result_calls[0]
         assert isinstance(result_input, GetNexusOperationResultInput)
         assert result_input.operation_id == op_id
-        assert result_input.result_type == EchoOutput
 
         # Start another so we can terminate it
         previous_start_count = len(interceptor.start_calls)