Add intermediate model payload conversion hooks

temporalio/activity.py

@@ -238,9 +238,17 @@ def payload_converter(self) -> temporalio.converter.PayloadConverter:
                 self.payload_converter_class_or_instance,
                 temporalio.converter.PayloadConverter,
             ):
-                self._payload_converter = self.payload_converter_class_or_instance
+                self._payload_converter = (
+                    temporalio.converter.TemporalIntermediatePayloadConverter.wrap(
+                        self.payload_converter_class_or_instance
+                    )
+                )
             else:
-                self._payload_converter = self.payload_converter_class_or_instance()
+                self._payload_converter = (
+                    temporalio.converter.TemporalIntermediatePayloadConverter.wrap(
+                        self.payload_converter_class_or_instance()
+                    )
+                )
         return self._payload_converter
 
     @property

temporalio/converter/__init__.py

@@ -33,6 +33,7 @@
     JSONTypeConverter,
     JSONTypeConverterUnhandled,
     PayloadConverter,
+    TemporalIntermediatePayloadConverter,
     value_to_type,
 )
 from temporalio.converter._payload_limits import (
@@ -83,6 +84,7 @@
     "PayloadLimitsConfig",
     "PayloadSizeWarning",
     "SerializationContext",
+    "TemporalIntermediatePayloadConverter",
     "WithSerializationContext",
     "WorkflowSerializationContext",
     "decode_search_attributes",

temporalio/converter/_data_converter.py

@@ -29,6 +29,7 @@
 )
 from temporalio.converter._payload_converter import (
     PayloadConverter,
+    TemporalIntermediatePayloadConverter,
 )
 from temporalio.converter._payload_limits import (
     PayloadLimitsConfig,
@@ -103,7 +104,11 @@ class DataConverter(WithSerializationContext):
     """Server-reported limits for payloads."""
 
     def __post_init__(self) -> None:  # noqa: D105
-        object.__setattr__(self, "payload_converter", self.payload_converter_class())
+        object.__setattr__(
+            self,
+            "payload_converter",
+            TemporalIntermediatePayloadConverter.wrap(self.payload_converter_class()),
+        )
         object.__setattr__(self, "failure_converter", self.failure_converter_class())
 
     async def encode(

temporalio/converter/_payload_converter.py

@@ -514,6 +514,83 @@ def from_payload(
             raise RuntimeError("Failed parsing") from err
 
 
+class TemporalIntermediatePayloadConverter(PayloadConverter, WithSerializationContext):
+    """Payload converter wrapper for generated Temporal intermediate hooks.
+
+    Values with a ``_temporal_to_intermediate`` method are first converted to
+    their intermediate value, then encoded by the wrapped payload converter. When
+    decoding to a type with ``_temporal_from_intermediate``, the wrapped
+    converter first decodes the payload to the intermediate value and this
+    wrapper constructs the requested user-facing type from it.
+    """
+
+    _inner_payload_converter: PayloadConverter
+
+    def __init__(self, inner_payload_converter: PayloadConverter) -> None:
+        """Create a Temporal intermediate payload converter."""
+        self._inner_payload_converter = inner_payload_converter
+
+    @staticmethod
+    def wrap(payload_converter: PayloadConverter) -> PayloadConverter:
+        """Wrap a payload converter unless it is already wrapped."""
+        if isinstance(payload_converter, TemporalIntermediatePayloadConverter):
+            return payload_converter
+        return TemporalIntermediatePayloadConverter(payload_converter)
+
+    def to_payloads(
+        self, values: Sequence[Any]
+    ) -> list[temporalio.api.common.v1.Payload]:
+        """See base class."""
+        intermediate_values: list[Any] = []
+        for value in values:
+            to_intermediate = getattr(value, "_temporal_to_intermediate", None)
+            if to_intermediate is not None:
+                value = to_intermediate()
+            intermediate_values.append(value)
+        return self._inner_payload_converter.to_payloads(intermediate_values)
+
+    def from_payloads(
+        self,
+        payloads: Sequence[temporalio.api.common.v1.Payload],
+        type_hints: list[type] | None = None,
+    ) -> list[Any]:
+        """See base class."""
+        if type_hints is None:
+            return self._inner_payload_converter.from_payloads(payloads, None)
+        normalized_type_hints: list[type | None] = list(type_hints)
+        if len(normalized_type_hints) < len(payloads):
+            normalized_type_hints.extend([None] * (len(payloads) - len(type_hints)))
+        inner_type_hints = [
+            None
+            if getattr(type_hint, "_temporal_from_intermediate", None) is not None
+            else type_hint
+            for type_hint in normalized_type_hints
+        ]
+        values = self._inner_payload_converter.from_payloads(
+            payloads, typing.cast("list[type]", inner_type_hints)
+        )
+        return [
+            from_intermediate(value)
+            if (
+                from_intermediate := getattr(
+                    type_hint, "_temporal_from_intermediate", None
+                )
+            )
+            is not None
+            else value
+            for value, type_hint in zip(values, normalized_type_hints)
+        ]
+
+    def with_context(self, context: SerializationContext) -> Self:
+        """Return a new instance with context set on the inner converter."""
+        if not isinstance(self._inner_payload_converter, WithSerializationContext):
+            return self
+        inner_payload_converter = self._inner_payload_converter.with_context(context)
+        if inner_payload_converter is self._inner_payload_converter:
+            return self
+        return type(self)(inner_payload_converter)
+
+
 class AdvancedJSONEncoder(json.JSONEncoder):
     """Advanced JSON encoder.
 

temporalio/nexus/system/__init__.py

@@ -2,22 +2,82 @@
 
 from __future__ import annotations
 
+import contextlib
+import contextvars
+from collections.abc import Iterator, Sequence
+from typing import Any
+
 import temporalio.api.common.v1
 import temporalio.converter
 from temporalio.bridge._visitor_functions import VisitorFunctions
 from temporalio.converter import BinaryProtoPayloadConverter, CompositePayloadConverter
 
 TEMPORAL_SYSTEM_ENDPOINT = "__temporal_system"
+_user_payload_converter: contextvars.ContextVar[
+    temporalio.converter.PayloadConverter | None
+] = contextvars.ContextVar("temporal-system-nexus-user-payload-converter", default=None)
 
 
-class SystemNexusPayloadConverter(CompositePayloadConverter):
-    """Payload converter for system Nexus outer envelopes."""
+@contextlib.contextmanager
+def user_payload_converter_context(
+    payload_converter: temporalio.converter.PayloadConverter,
+) -> Iterator[None]:
+    """Set the user payload converter for system Nexus model conversion."""
+    token = _user_payload_converter.set(payload_converter)
+    try:
+        yield
+    finally:
+        _user_payload_converter.reset(token)
+
+
+def current_user_payload_converter() -> temporalio.converter.PayloadConverter:
+    """Return the active user payload converter for system Nexus model conversion."""
+    payload_converter = _user_payload_converter.get()
+    if payload_converter is None:
+        raise RuntimeError("System Nexus user payload converter context is not active")
+    return payload_converter
+
+
+class _SystemNexusOuterPayloadConverter(CompositePayloadConverter):
+    """Payload converter for system Nexus outer proto envelopes."""
 
     def __init__(self) -> None:
         """Create a payload converter for system Nexus outer envelopes."""
         super().__init__(BinaryProtoPayloadConverter())
 
 
+class SystemNexusPayloadConverter(temporalio.converter.PayloadConverter):
+    """Payload converter for system Nexus outer envelopes."""
+
+    _user_payload_converter: temporalio.converter.PayloadConverter
+    _outer_payload_converter: temporalio.converter.PayloadConverter
+
+    def __init__(self, user_payload_converter: temporalio.converter.PayloadConverter) -> None:
+        """Create a payload converter for system Nexus outer envelopes."""
+        self._user_payload_converter = user_payload_converter
+        self._outer_payload_converter = (
+            temporalio.converter.TemporalIntermediatePayloadConverter.wrap(
+                _SystemNexusOuterPayloadConverter()
+            )
+        )
+
+    def to_payloads(
+        self, values: Sequence[Any]
+    ) -> list[temporalio.api.common.v1.Payload]:
+        """See base class."""
+        with user_payload_converter_context(self._user_payload_converter):
+            return self._outer_payload_converter.to_payloads(values)
+
+    def from_payloads(
+        self,
+        payloads: Sequence[temporalio.api.common.v1.Payload],
+        type_hints: list[type] | None = None,
+    ) -> list[Any]:
+        """See base class."""
+        with user_payload_converter_context(self._user_payload_converter):
+            return self._outer_payload_converter.from_payloads(payloads, type_hints)
+
+
 def is_system_endpoint(endpoint: str) -> bool:
     """Return whether a Nexus endpoint is the Temporal system endpoint."""
     return endpoint == TEMPORAL_SYSTEM_ENDPOINT
@@ -33,7 +93,7 @@ async def maybe_visit_payload(
     if not is_system_endpoint(endpoint):
         return None
 
-    payload_converter = get_payload_converter()
+    payload_converter = _SystemNexusOuterPayloadConverter()
     value = payload_converter.from_payload(payload)
     from ._payload_visitor import PayloadVisitor
 
@@ -43,15 +103,19 @@ async def maybe_visit_payload(
     return payload_converter.to_payload(value)
 
 
-def get_payload_converter() -> temporalio.converter.PayloadConverter:
+def get_payload_converter(
+    user_payload_converter: temporalio.converter.PayloadConverter,
+) -> temporalio.converter.PayloadConverter:
     """Return the fixed payload converter for system Nexus outer envelopes."""
-    return SystemNexusPayloadConverter()
+    return SystemNexusPayloadConverter(user_payload_converter)
 
 
 __all__ = [
     "TEMPORAL_SYSTEM_ENDPOINT",
+    "current_user_payload_converter",
     "get_payload_converter",
     "is_system_endpoint",
     "maybe_visit_payload",
     "SystemNexusPayloadConverter",
+    "user_payload_converter_context",
 ]

temporalio/worker/_workflow_instance.py

@@ -246,7 +246,11 @@ def __init__(self, det: WorkflowInstanceDetails) -> None:
         self._defn = det.defn
         self._workflow_input: ExecuteWorkflowInput | None = None
         self._info = det.info
-        self._context_free_payload_converter = det.payload_converter_class()
+        self._context_free_payload_converter = (
+            temporalio.converter.TemporalIntermediatePayloadConverter.wrap(
+                det.payload_converter_class()
+            )
+        )
         self._context_free_failure_converter = det.failure_converter_class()
         workflow_context = temporalio.converter.WorkflowSerializationContext(
             namespace=det.info.namespace,
@@ -2089,7 +2093,9 @@ async def operation_handle_fn() -> OutputT:
                         t.uncancel()  # type: ignore[union-attr]
 
         payload_converter = (
-            temporalio.nexus.system.get_payload_converter()
+            temporalio.nexus.system.get_payload_converter(
+                self._workflow_context_payload_converter
+            )
             if temporalio.nexus.system.is_system_endpoint(input.endpoint)
             else self._context_free_payload_converter
         )

tests/nexus/test_temporal_system_nexus.py

@@ -186,7 +186,9 @@ def _new_system_nexus_request_payload() -> temporalio.api.common.v1.Payload:
     assert nested_payload is not None
     request = workflowservice_pb2.SignalWithStartWorkflowExecutionRequest()
     request.input.payloads.add().CopyFrom(nested_payload)
-    payload = nexus_system.get_payload_converter().to_payload(request)
+    payload = nexus_system.get_payload_converter(
+        temporalio.converter.PayloadConverter.default
+    ).to_payload(request)
     assert payload is not None
     return payload
 
@@ -201,7 +203,9 @@ async def test_schedule_system_nexus_endpoint_ignores_operation_registry() -> No
     await PayloadVisitor().visit(visitor, completion)
 
     schedule = completion.successful.commands[0].schedule_nexus_operation
-    decoded = nexus_system.get_payload_converter().from_payload(schedule.input)
+    decoded = nexus_system.get_payload_converter(
+        temporalio.converter.PayloadConverter.default
+    ).from_payload(schedule.input)
     assert isinstance(
         decoded, workflowservice_pb2.SignalWithStartWorkflowExecutionRequest
     )
@@ -339,7 +343,9 @@ def _field_is_repeated(field: FieldDescriptor) -> bool:
     ],
 )
 def test_system_nexus_proto_roundtrip(message_type: type[Message]) -> None:
-    payload_converter = nexus_system.get_payload_converter()
+    payload_converter = nexus_system.get_payload_converter(
+        temporalio.converter.PayloadConverter.default
+    )
     proto_value = _build_proto_sample(message_type)
     payload = payload_converter.to_payload(proto_value)
     assert payload is not None

tests/test_converter.py

@@ -44,6 +44,8 @@
     JSONTypeConverter,
     JSONTypeConverterUnhandled,
     PayloadCodec,
+    PayloadConverter,
+    TemporalIntermediatePayloadConverter,
     decode_search_attributes,
     encode_search_attribute_values,
     value_to_type,
@@ -254,6 +256,52 @@ def test_binary_proto():
     assert decoded == proto
 
 
+@dataclass
+class TemporalIntermediateModel:
+    value: str
+
+    @classmethod
+    def _temporal_from_intermediate(
+        cls,
+        intermediate: temporalio.api.common.v1.WorkflowExecution,
+    ) -> TemporalIntermediateModel:
+        return cls(value=intermediate.workflow_id)
+
+    def _temporal_to_intermediate(self) -> temporalio.api.common.v1.WorkflowExecution:
+        return temporalio.api.common.v1.WorkflowExecution(
+            workflow_id=self.value,
+            run_id="run-id",
+        )
+
+
+class CustomDefaultPayloadConverter(DefaultPayloadConverter):
+    pass
+
+
+def test_temporal_intermediate_payload_converter_wraps_user_converter():
+    data_converter = DataConverter(
+        payload_converter_class=CustomDefaultPayloadConverter
+    )
+    converter = data_converter.payload_converter
+    assert isinstance(converter, TemporalIntermediatePayloadConverter)
+    value = TemporalIntermediateModel("workflow-id")
+
+    payload = converter.to_payload(value)
+
+    assert payload.metadata["encoding"] == b"json/protobuf"
+    assert (
+        payload.metadata["messageType"] == b"temporal.api.common.v1.WorkflowExecution"
+    )
+    assert all("temporal-wire" not in key for key in payload.metadata)
+    assert all(b"temporal-wire" not in value for value in payload.metadata.values())
+    assert converter.from_payload(payload, TemporalIntermediateModel) == value
+
+    plain_proto_payload = converter.to_payload(
+        temporalio.api.common.v1.WorkflowExecution(workflow_id="id1", run_id="id2")
+    )
+    assert plain_proto_payload.metadata["encoding"] == b"json/protobuf"
+
+
 def test_encode_search_attribute_values():
     with pytest.raises(TypeError, match="of type tuple not one of"):
         encode_search_attribute_values([("bad type",)])  # type: ignore[arg-type]

tests/worker/test_visitor.py

@@ -357,7 +357,9 @@ async def _visit(self) -> None:
             finally:
                 active_visits -= 1
 
-    payload_converter = nexus_system.get_payload_converter()
+    payload_converter = nexus_system.get_payload_converter(
+        temporalio.converter.PayloadConverter.default
+    )
     system_request = workflowservice_pb2.SignalWithStartWorkflowExecutionRequest(
         input=Payloads(payloads=[Payload(data=b"workflow-input")]),
         signal_input=Payloads(payloads=[Payload(data=b"signal-input")]),