Fix payload codec usage

- Revert "Store data_converter on WorkflowInstanceDetails instead of converter classes"
- Don't require payload codec on workflow instance
- Don't stack contexts

Revert "Store data_converter on WorkflowInstanceDetails instead of converter classes"

This reverts commit f869af87ba5bffb891f8abb3189e655921247ca3.

Author: dandavison

temporalio/worker/_workflow.py

@@ -249,13 +249,12 @@ async def _handle_activation(
             await self._handle_cache_eviction(act, cache_remove_job)
             return
 
-        data_converter = self._data_converter
         # Build default success completion (e.g. remove-job-only activations)
         completion = (
             temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion()
         )
         completion.successful.SetInParent()
-        workflow = None
+        workflow = workflow_id = None
         try:
             if LOG_PROTOS:
                 logger.debug("Received workflow activation:\n%s", act)
@@ -275,17 +274,14 @@ async def _handle_activation(
                         "Cache already exists for activation with initialize job"
                     )
 
-            data_converter = self._data_converter._with_context(
-                temporalio.converter.WorkflowSerializationContext(
-                    namespace=self._namespace,
-                    workflow_id=workflow_id,
-                )
-            )
-            if data_converter.payload_codec:
+            if self._data_converter.payload_codec:
                 if not workflow:
-                    payload_codec = data_converter.payload_codec
+                    payload_codec = self._data_converter.payload_codec
                 else:
-                    payload_codec = _CommandAwarePayloadCodec(workflow.instance)
+                    payload_codec = _CommandAwarePayloadCodec(
+                        workflow.instance,
+                        self._data_converter.payload_codec,
+                    )
                 await temporalio.bridge.worker.decode_activation(
                     act,
                     payload_codec.decode,
@@ -339,6 +335,14 @@ async def _handle_activation(
 
             completion.failed.failure.SetInParent()
             try:
+                data_converter = self._data_converter
+                if workflow_id:
+                    data_converter = data_converter._with_context(
+                        temporalio.converter.WorkflowSerializationContext(
+                            namespace=self._namespace,
+                            workflow_id=workflow_id,
+                        )
+                    )
                 data_converter.failure_converter.to_failure(
                     err,
                     data_converter.payload_converter,
@@ -356,8 +360,11 @@ async def _handle_activation(
         completion.run_id = act.run_id
 
         # Encode completion
-        if data_converter.payload_codec and workflow:
-            payload_codec = _CommandAwarePayloadCodec(workflow.instance)
+        if self._data_converter.payload_codec and workflow:
+            payload_codec = _CommandAwarePayloadCodec(
+                workflow.instance,
+                self._data_converter.payload_codec,
+            )
             try:
                 await temporalio.bridge.worker.encode_completion(
                     completion,
@@ -572,7 +579,8 @@ def _create_workflow_instance(
 
         # Create instance from details
         det = WorkflowInstanceDetails(
-            data_converter=self._data_converter,
+            payload_converter_class=self._data_converter.payload_converter_class,
+            failure_converter_class=self._data_converter.failure_converter_class,
             interceptor_classes=self._interceptor_classes,
             defn=defn,
             info=info,
@@ -722,8 +730,10 @@ class _CommandAwarePayloadCodec(temporalio.converter.PayloadCodec):
     def __init__(
         self,
         instance: WorkflowInstance,
+        context_free_payload_codec: temporalio.converter.PayloadCodec,
     ):
         self.instance = instance
+        self.context_free_payload_codec = context_free_payload_codec
 
     async def encode(
         self,
@@ -738,10 +748,10 @@ async def decode(
         return await self._get_current_command_codec().decode(payloads)
 
     def _get_current_command_codec(self) -> temporalio.converter.PayloadCodec:
-        seq = temporalio.bridge._visitor.current_command_seq.get()
-        codec = self.instance.get_payload_codec(seq)
-        assert codec, "Payload codec must be set on the data converter"
-        return codec
+        return self.instance.get_payload_codec_with_context(
+            self.context_free_payload_codec,
+            temporalio.bridge._visitor.current_command_seq.get(),
+        )
 
 
 class _InterruptDeadlockError(BaseException):

temporalio/worker/_workflow_instance.py

@@ -135,7 +135,8 @@ def set_worker_level_failure_exception_types(
 class WorkflowInstanceDetails:
     """Immutable details for creating a workflow instance."""
 
-    data_converter: temporalio.converter.DataConverter
+    payload_converter_class: Type[temporalio.converter.PayloadConverter]
+    failure_converter_class: Type[temporalio.converter.FailureConverter]
     interceptor_classes: Sequence[Type[WorkflowInboundInterceptor]]
     defn: temporalio.workflow._Definition
     info: temporalio.workflow.Info
@@ -168,9 +169,11 @@ def activate(
         raise NotImplementedError
 
     @abstractmethod
-    def get_payload_codec(
-        self, command_seq: Optional[int]
-    ) -> Optional[temporalio.converter.PayloadCodec]:
+    def get_payload_codec_with_context(
+        self,
+        payload_codec: temporalio.converter.PayloadCodec,
+        command_seq: Optional[int],
+    ) -> temporalio.converter.PayloadCodec:
         """Return a payload codec with appropriate serialization context.
 
         Args:
@@ -224,13 +227,8 @@ def __init__(self, det: WorkflowInstanceDetails) -> None:
         self._defn = det.defn
         self._workflow_input: Optional[ExecuteWorkflowInput] = None
         self._info = det.info
-        self._context_free_payload_codec = det.data_converter.payload_codec
-        self._context_free_payload_converter = (
-            det.data_converter.payload_converter_class()
-        )
-        self._context_free_failure_converter = (
-            det.data_converter.failure_converter_class()
-        )
+        self._context_free_payload_converter = det.payload_converter_class()
+        self._context_free_failure_converter = det.failure_converter_class()
         self._payload_converter, self._failure_converter = (
             self._converters_with_context(
                 temporalio.converter.WorkflowSerializationContext(
@@ -2099,10 +2097,11 @@ def _converters_with_context(
         return payload_converter, failure_converter
 
     # _WorkflowInstanceImpl.get_pending_command_serialization_context
-    def get_payload_codec(
-        self, command_seq: Optional[int]
-    ) -> Optional[temporalio.converter.PayloadCodec]:
-        payload_codec = self._context_free_payload_codec
+    def get_payload_codec_with_context(
+        self,
+        payload_codec: temporalio.converter.PayloadCodec,
+        command_seq: Optional[int],
+    ) -> temporalio.converter.PayloadCodec:
         if not isinstance(
             payload_codec,
             temporalio.converter.WithSerializationContext,

temporalio/worker/workflow_sandbox/_in_sandbox.py

@@ -81,8 +81,10 @@ def activate(
         """Send activation to this instance."""
         return self.instance.activate(act)
 
-    def get_payload_codec(
-        self, command_seq: Optional[int]
-    ) -> Optional[temporalio.converter.PayloadCodec]:
-        """Get payload codec."""
-        return self.instance.get_payload_codec(command_seq)
+    def get_payload_codec_with_context(
+        self,
+        payload_codec: temporalio.converter.PayloadCodec,
+        command_seq: Optional[int],
+    ) -> temporalio.converter.PayloadCodec:
+        """Get payload codec with context."""
+        return self.instance.get_payload_codec_with_context(payload_codec, command_seq)

temporalio/worker/workflow_sandbox/_runner.py

@@ -77,7 +77,8 @@ def prepare_workflow(self, defn: temporalio.workflow._Definition) -> None:
         # Just create with fake info which validates
         self.create_instance(
             WorkflowInstanceDetails(
-                data_converter=temporalio.converter.DataConverter.default,
+                payload_converter_class=temporalio.converter.DataConverter.default.payload_converter_class,
+                failure_converter_class=temporalio.converter.DataConverter.default.failure_converter_class,
                 interceptor_classes=[],
                 defn=defn,
                 # Just use fake info during validation
@@ -185,16 +186,19 @@ def _run_code(self, code: str, **extra_globals: Any) -> None:
     def get_thread_id(self) -> Optional[int]:
         return self._current_thread_id
 
-    def get_payload_codec(
-        self, command_seq: Optional[int]
-    ) -> Optional[temporalio.converter.PayloadCodec]:
+    def get_payload_codec_with_context(
+        self,
+        payload_codec: temporalio.converter.PayloadCodec,
+        command_seq: Optional[int],
+    ) -> temporalio.converter.PayloadCodec:
         # Forward call to the sandboxed instance
         self.importer.restriction_context.is_runtime = True
         try:
             self._run_code(
                 "with __temporal_importer.applied():\n"
-                "  __temporal_codec = __temporal_in_sandbox.get_payload_codec(__temporal_command_seq)\n",
+                "  __temporal_codec = __temporal_in_sandbox.get_payload_codec_with_context(__temporal_payload_codec, __temporal_command_seq)\n",
                 __temporal_importer=self.importer,
+                __temporal_payload_codec=payload_codec,
                 __temporal_command_seq=command_seq,
             )
             return self.globals_and_locals.pop("__temporal_codec", None)  # type: ignore

tests/worker/test_workflow.py

@@ -1610,10 +1610,14 @@ def activate(self, act: WorkflowActivation) -> WorkflowActivationCompletion:
         self._runner._pairs.append((act, comp))
         return comp
 
-    def get_payload_codec(
-        self, command_seq: Optional[int]
-    ) -> Optional[temporalio.converter.PayloadCodec]:
-        return self._unsandboxed.get_payload_codec(command_seq)
+    def get_payload_codec_with_context(
+        self,
+        payload_codec: temporalio.converter.PayloadCodec,
+        command_seq: Optional[int],
+    ) -> temporalio.converter.PayloadCodec:
+        return self._unsandboxed.get_payload_codec_with_context(
+            payload_codec, command_seq
+        )
 
 
 async def test_workflow_with_custom_runner(client: Client):