feat(extstore): Add support for Nexus task handling (#1676)

Author: jmaeagle99

CHANGELOG.md

@@ -52,6 +52,7 @@ to include examples, links to docs, or any other relevant information.
 - Added the experimental `Worker` `patch_activation_callback` option, allowing workers
   to decide whether a first non-replay `workflow.patched` call should activate a patch
   during rolling deployments.
+- Added external storage support to Nexus task handling.
 
 ### Changed
 

scripts/gen_payload_visitor.py

@@ -12,6 +12,7 @@
 sys.path.insert(0, str(base_dir))
 
 from temporalio.api.common.v1.message_pb2 import Payload, Payloads, SearchAttributes
+from temporalio.bridge.proto.nexus import NexusTaskCompletion
 from temporalio.bridge.proto.workflow_activation.workflow_activation_pb2 import (
     WorkflowActivation,
 )
@@ -425,9 +426,12 @@ def walk(self, desc: Descriptor) -> bool:
 def write_bridge_visitors() -> None:
     out_path = base_dir / "temporalio" / "bridge" / "_visitor.py"
 
+    # Build root descriptors: WorkflowActivation, WorkflowActivationCompletion,
+    # NexusTaskCompletion, and the system Nexus operation roots.
     roots: list[Descriptor] = [
         WorkflowActivation.DESCRIPTOR,
         WorkflowActivationCompletion.DESCRIPTOR,
+        NexusTaskCompletion.DESCRIPTOR,
     ] + discover_system_nexus_roots()
 
     code = VisitorGenerator().generate(roots)

temporalio/bridge/_visitor.py

@@ -548,6 +548,58 @@ async def _visit_coresdk_workflow_completion_WorkflowActivationCompletion(
         elif o.HasField("failed"):
             await self._visit_coresdk_workflow_completion_Failure(fs, o.failed)
 
+    async def _visit_temporal_api_nexus_v1_StartOperationResponse_Sync(
+        self, fs: VisitorFunctions, o: Any
+    ):
+        if o.HasField("payload"):
+            await self._visit_temporal_api_common_v1_Payload(fs, o.payload)
+
+    async def _visit_temporal_api_nexus_v1_Failure(self, fs: VisitorFunctions, o: Any):
+        if o.HasField("cause"):
+            await self._visit_temporal_api_nexus_v1_Failure(fs, o.cause)
+
+    async def _visit_temporal_api_nexus_v1_UnsuccessfulOperationError(
+        self, fs: VisitorFunctions, o: Any
+    ):
+        if o.HasField("failure"):
+            await self._visit_temporal_api_nexus_v1_Failure(fs, o.failure)
+
+    async def _visit_temporal_api_nexus_v1_StartOperationResponse(
+        self, fs: VisitorFunctions, o: Any
+    ):
+        if o.HasField("sync_success"):
+            await self._visit_temporal_api_nexus_v1_StartOperationResponse_Sync(
+                fs, o.sync_success
+            )
+        elif o.HasField("operation_error"):
+            await self._visit_temporal_api_nexus_v1_UnsuccessfulOperationError(
+                fs, o.operation_error
+            )
+        elif o.HasField("failure"):
+            await self._visit_temporal_api_failure_v1_Failure(fs, o.failure)
+
+    async def _visit_temporal_api_nexus_v1_Response(self, fs: VisitorFunctions, o: Any):
+        if o.HasField("start_operation"):
+            await self._visit_temporal_api_nexus_v1_StartOperationResponse(
+                fs, o.start_operation
+            )
+
+    async def _visit_temporal_api_nexus_v1_HandlerError(
+        self, fs: VisitorFunctions, o: Any
+    ):
+        if o.HasField("failure"):
+            await self._visit_temporal_api_nexus_v1_Failure(fs, o.failure)
+
+    async def _visit_coresdk_nexus_NexusTaskCompletion(
+        self, fs: VisitorFunctions, o: Any
+    ):
+        if o.HasField("completed"):
+            await self._visit_temporal_api_nexus_v1_Response(fs, o.completed)
+        elif o.HasField("error"):
+            await self._visit_temporal_api_nexus_v1_HandlerError(fs, o.error)
+        elif o.HasField("failure"):
+            await self._visit_temporal_api_failure_v1_Failure(fs, o.failure)
+
     async def _visit_temporal_api_common_v1_Header(self, fs: VisitorFunctions, o: Any):
         for v in o.fields.values():
             await self._visit_temporal_api_common_v1_Payload(fs, v)

temporalio/converter/_data_converter.py

@@ -13,9 +13,9 @@
 import temporalio.api.common.v1
 import temporalio.api.failure.v1
 import temporalio.common
-from temporalio.api.sdk.v1.external_storage_pb2 import ExternalStorageReference
 from temporalio.converter._extstore import (
     _REFERENCE_ENCODING,
+    _REFERENCE_MESSAGE_TYPE,
     ExternalStorage,
     StorageDriverStoreContext,
 )
@@ -35,8 +35,6 @@
     WithSerializationContext,
 )
 
-_REFERENCE_MESSAGE_TYPE = ExternalStorageReference.DESCRIPTOR.full_name.encode()
-
 
 def _is_reference_payload(p: temporalio.api.common.v1.Payload) -> bool:
     """Return True if *p* is an external-storage reference payload."""

temporalio/converter/_extstore.py

@@ -27,6 +27,7 @@
 _T = TypeVar("_T")
 
 _REFERENCE_ENCODING = b"json/external-storage-reference"
+_REFERENCE_MESSAGE_TYPE = ExternalStorageReference.DESCRIPTOR.full_name.encode()
 
 
 @dataclass
@@ -455,8 +456,6 @@ async def _store_payload_sequence(
 
     def _decode_reference(self, payload: Payload) -> ExternalStorageReference | None:
         """Decode an external storage reference from a payload."""
-        if len(payload.external_payloads) == 0:
-            return None
         encoding = payload.metadata.get("encoding", b"")
         if encoding == _REFERENCE_ENCODING:
             legacy = self._legacy_claim_converter.from_payload(
@@ -468,6 +467,11 @@ def _decode_reference(self, payload: Payload) -> ExternalStorageReference | None
                 driver_name=legacy.driver_name,
                 claim_data=legacy.driver_claim.claim_data,
             )
+        if not (
+            encoding == b"json/protobuf"
+            and payload.metadata.get("messageType") == _REFERENCE_MESSAGE_TYPE
+        ):
+            return None
         ref = self._claim_converter.from_payload(payload, ExternalStorageReference)
         return ref if isinstance(ref, ExternalStorageReference) else None
 

temporalio/worker/_nexus.py

@@ -6,7 +6,7 @@
 import concurrent.futures
 import contextvars
 import threading
-from collections.abc import Callable, Mapping, Sequence
+from collections.abc import Awaitable, Callable, Mapping, Sequence
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from functools import reduce
@@ -30,6 +30,8 @@
 import temporalio.common
 import temporalio.converter
 import temporalio.nexus
+from temporalio.bridge._visitor import PayloadVisitor
+from temporalio.bridge._visitor_functions import PayloadSequence, VisitorFunctions
 from temporalio.bridge.worker import PollShutdownError
 from temporalio.exceptions import (
     ApplicationError,
@@ -216,6 +218,19 @@ async def _complete_task(
     ):
         await asyncio.shield(self._bridge_worker().complete_nexus_task(completion))
 
+    async def _encode_completion(
+        self, completion: temporalio.bridge.proto.nexus.NexusTaskCompletion
+    ) -> None:
+        """Apply the payload codec then external storage to the completion's payloads."""
+        dc = self._data_converter
+        await PayloadVisitor(skip_search_attributes=True, skip_headers=True).visit(
+            _PayloadTransformVisitor(dc._encode_payload_sequence), completion
+        )
+        await PayloadVisitor(skip_search_attributes=True).visit(
+            _PayloadTransformVisitor(dc._external_store_payload_sequence),
+            completion,
+        )
+
     # TODO(nexus-preview): stack trace pruning. See sdk-typescript NexusHandler.execute
     # "Any call up to this function and including this one will be trimmed out of stack traces.""
 
@@ -260,6 +275,14 @@ async def _handle_cancel_operation_task(
         try:
             try:
                 await self._handler.cancel_operation(ctx, request.operation_token)
+                completion = temporalio.bridge.proto.nexus.NexusTaskCompletion(
+                    task_token=task_token,
+                    completed=temporalio.api.nexus.v1.Response(
+                        cancel_operation=temporalio.api.nexus.v1.CancelOperationResponse()
+                    ),
+                )
+                # No-op but keeps the cancel covered if it ever carries a payload.
+                await self._encode_completion(completion)
             except asyncio.CancelledError:
                 completion = temporalio.bridge.proto.nexus.NexusTaskCompletion(
                     task_token=task_token,
@@ -271,16 +294,12 @@ async def _handle_cancel_operation_task(
                 completion = temporalio.bridge.proto.nexus.NexusTaskCompletion(
                     task_token=task_token,
                 )
-                await self._data_converter.encode_failure(
-                    handler_error, completion.failure
-                )
-            else:
-                completion = temporalio.bridge.proto.nexus.NexusTaskCompletion(
-                    task_token=task_token,
-                    completed=temporalio.api.nexus.v1.Response(
-                        cancel_operation=temporalio.api.nexus.v1.CancelOperationResponse()
-                    ),
+                self._data_converter.failure_converter.to_failure(
+                    handler_error,
+                    self._data_converter.payload_converter,
+                    completion.failure,
                 )
+                await self._encode_completion(completion)
             await self._complete_task(completion)
         except Exception:
             logger.exception("Failed to send Nexus task completion")
@@ -315,6 +334,13 @@ async def _handle_start_operation_task(
                     request_deadline,
                     endpoint,
                 )
+                completion = temporalio.bridge.proto.nexus.NexusTaskCompletion(
+                    task_token=task_token,
+                    completed=temporalio.api.nexus.v1.Response(
+                        start_operation=start_response
+                    ),
+                )
+                await self._encode_completion(completion)
             except asyncio.CancelledError:
                 completion = temporalio.bridge.proto.nexus.NexusTaskCompletion(
                     task_token=task_token,
@@ -326,19 +352,15 @@ async def _handle_start_operation_task(
                     task_token=task_token,
                 )
                 handler_error = _exception_to_handler_error(err)
-                await self._data_converter.encode_failure(
-                    handler_error, completion.failure
+                self._data_converter.failure_converter.to_failure(
+                    handler_error,
+                    self._data_converter.payload_converter,
+                    completion.failure,
                 )
 
                 if isinstance(err, concurrent.futures.BrokenExecutor):
                     self._fail_worker_exception_queue.put_nowait(err)
-            else:
-                completion = temporalio.bridge.proto.nexus.NexusTaskCompletion(
-                    task_token=task_token,
-                    completed=temporalio.api.nexus.v1.Response(
-                        start_operation=start_response
-                    ),
-                )
+                await self._encode_completion(completion)
 
             await self._complete_task(completion)
         except Exception:
@@ -417,7 +439,9 @@ async def _start_operation(
                     )
                 )
             elif isinstance(result, nexusrpc.handler.StartOperationResultSync):
-                [payload] = await self._data_converter.encode([result.value])
+                [payload] = self._data_converter.payload_converter.to_payloads(
+                    [result.value]
+                )
                 return temporalio.api.nexus.v1.StartOperationResponse(
                     sync_success=temporalio.api.nexus.v1.StartOperationResponse.Sync(
                         payload=payload,
@@ -446,10 +470,41 @@ async def _start_operation(
                         ) from err.__cause__
             except FailureError as new_err:
                 response = temporalio.api.nexus.v1.StartOperationResponse()
-                await self._data_converter.encode_failure(new_err, response.failure)
+                self._data_converter.failure_converter.to_failure(
+                    new_err,
+                    self._data_converter.payload_converter,
+                    response.failure,
+                )
                 return response
 
 
+class _PayloadTransformVisitor(VisitorFunctions):
+    """Adapts a payload-sequence transform for use with :class:`PayloadVisitor`."""
+
+    def __init__(
+        self,
+        f: Callable[
+            [Sequence[temporalio.api.common.v1.Payload]],
+            Awaitable[list[temporalio.api.common.v1.Payload]],
+        ],
+    ) -> None:
+        self._f = f
+
+    async def visit_payload(self, payload: temporalio.api.common.v1.Payload) -> None:
+        new_payload = (await self._f([payload]))[0]
+        if new_payload is not payload:
+            payload.CopyFrom(new_payload)
+
+    async def visit_payloads(self, payloads: PayloadSequence) -> None:
+        if len(payloads) == 0:
+            return
+        new_payloads = await self._f(payloads)
+        if new_payloads is payloads:
+            return
+        del payloads[:]
+        payloads.extend(new_payloads)
+
+
 @dataclass
 class _DummyPayloadSerializer:
     data_converter: temporalio.converter.DataConverter
@@ -465,18 +520,35 @@ async def deserialize(
         content: nexusrpc.Content,  # type:ignore[reportUnusedParameter]
         as_type: type[Any] | None = None,
     ) -> Any:
-        payload = self.payload
-        if self.data_converter.payload_codec:
-            try:
-                [payload] = await self.data_converter.payload_codec.decode([payload])
-            except Exception as err:
-                raise nexusrpc.HandlerError(
-                    "Payload codec failed to decode Nexus operation input",
-                    type=nexusrpc.HandlerErrorType.INTERNAL,
-                ) from err
+        dc = self.data_converter
+        # The visitor mutates in place, so work on a copy to leave the request
+        # payload untouched.
+        payload = temporalio.api.common.v1.Payload()
+        payload.CopyFrom(self.payload)
+        try:
+            await PayloadVisitor(skip_search_attributes=True).visit(
+                _PayloadTransformVisitor(dc._external_retrieve_payload_sequence),
+                payload,
+            )
+        except Exception as err:
+            raise nexusrpc.HandlerError(
+                "Failed to retrieve Nexus operation input from external storage",
+                type=nexusrpc.HandlerErrorType.INTERNAL,
+                retryable_override=True,
+            ) from err
+
+        try:
+            await PayloadVisitor(skip_search_attributes=True, skip_headers=True).visit(
+                _PayloadTransformVisitor(dc._decode_payload_sequence), payload
+            )
+        except Exception as err:
+            raise nexusrpc.HandlerError(
+                "Payload codec failed to decode Nexus operation input",
+                type=nexusrpc.HandlerErrorType.INTERNAL,
+            ) from err
 
         try:
-            [input] = self.data_converter.payload_converter.from_payloads(
+            [input] = dc.payload_converter.from_payloads(
                 [payload],
                 type_hints=[as_type] if as_type else None,
             )