Command-aware codec for payload visitor

Author: dandavison

temporalio/worker/_workflow.py

@@ -24,6 +24,7 @@
 
 import temporalio.activity
 import temporalio.api.common.v1
+import temporalio.bridge._visitor
 import temporalio.bridge.client
 import temporalio.bridge.proto.workflow_activation
 import temporalio.bridge.proto.workflow_completion
@@ -254,6 +255,7 @@ async def _handle_activation(
             temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion()
         )
         completion.successful.SetInParent()
+        workflow = None
         try:
             if LOG_PROTOS:
                 logger.debug("Received workflow activation:\n%s", act)
@@ -280,9 +282,13 @@ async def _handle_activation(
                 )
             )
             if data_converter.payload_codec:
+                workflow_instance = workflow.instance if workflow else None
+                payload_codec = _CommandAwarePayloadCodec(
+                    workflow_instance, data_converter.payload_codec
+                )
                 await temporalio.bridge.worker.decode_activation(
                     act,
-                    data_converter.payload_codec,
+                    payload_codec.decode,
                     decode_headers=self._encode_headers,
                 )
             if not workflow:
@@ -348,13 +354,17 @@ async def _handle_activation(
                 )
 
         completion.run_id = act.run_id
+        assert workflow
 
         # Encode completion
         if data_converter.payload_codec:
+            payload_codec = _CommandAwarePayloadCodec(
+                workflow.instance, data_converter.payload_codec
+            )
             try:
                 await temporalio.bridge.worker.encode_completion(
                     completion,
-                    data_converter.payload_codec,
+                    payload_codec.encode,
                     encode_headers=self._encode_headers,
                 )
             except Exception as err:
@@ -705,5 +715,47 @@ def attempt_deadlock_interruption(self) -> None:
                 )
 
 
+class _CommandAwarePayloadCodec(temporalio.converter.PayloadCodec):
+    """A payload codec that sets serialization context for the associated command.
+
+    This codec responds to the :py:data:`temporalio.bridge._visitor.current_command_seq` context
+    variable set by the payload visitor.
+    """
+
+    def __init__(
+        self,
+        instance: Optional[WorkflowInstance],
+        base_codec: temporalio.converter.PayloadCodec,
+    ):
+        self.instance = instance
+        self.base_codec = base_codec
+
+    async def encode(
+        self,
+        payloads: Sequence[temporalio.api.common.v1.Payload],
+    ) -> List[temporalio.api.common.v1.Payload]:
+        return await self._get_current_command_codec().encode(payloads)
+
+    async def decode(
+        self,
+        payloads: Sequence[temporalio.api.common.v1.Payload],
+    ) -> List[temporalio.api.common.v1.Payload]:
+        return await self._get_current_command_codec().decode(payloads)
+
+    def _get_current_command_codec(self) -> temporalio.converter.PayloadCodec:
+        codec = self.base_codec
+        if self.instance:
+            if isinstance(
+                self.base_codec, temporalio.converter.WithSerializationContext
+            ):
+                if seq := temporalio.bridge._visitor.current_command_seq.get():
+                    if (
+                        context
+                        := self.instance.get_pending_command_serialization_context(seq)
+                    ):
+                        codec = self.base_codec.with_context(context)
+        return codec
+
+
 class _InterruptDeadlockError(BaseException):
     pass

temporalio/worker/_workflow_instance.py

@@ -167,6 +167,20 @@ def activate(
         """
         raise NotImplementedError
 
+    @abstractmethod
+    def get_pending_command_serialization_context(
+        self, command_seq: int
+    ) -> Optional[temporalio.converter.SerializationContext]:
+        """Return the serialization context for a pending command.
+
+        Args:
+            command_seq: The sequence number of the command.
+
+        Returns:
+            The serialization context for the command, or None if not found.
+        """
+        raise NotImplementedError
+
     def get_thread_id(self) -> Optional[int]:
         """Return the thread identifier that this workflow is running on.
 
@@ -225,6 +239,7 @@ def __init__(self, det: WorkflowInstanceDetails) -> None:
                 self._context_free_failure_converter,
             )
         )
+        self._payload_codec = det.data_converter.payload_codec
 
         self._extern_functions = det.extern_functions
         self._disable_eager_activity_execution = det.disable_eager_activity_execution
@@ -2089,6 +2104,45 @@ def _converters_with_context(
                 failure_converter = failure_converter.with_context(context)
         return payload_converter, failure_converter
 
+    def get_pending_command_serialization_context(
+        self, command_seq: int
+    ) -> Optional[temporalio.converter.SerializationContext]:
+        if isinstance(
+            self._payload_codec, temporalio.converter.WithSerializationContext
+        ):
+            if command_seq in self._pending_activities:
+                handle = self._pending_activities[command_seq]
+                return temporalio.converter.ActivitySerializationContext(
+                    namespace=self._info.namespace,
+                    workflow_id=self._info.workflow_id,
+                    workflow_type=self._info.workflow_type,
+                    activity_type=handle._input.activity,
+                    activity_task_queue=(
+                        handle._input.task_queue or self._info.task_queue
+                        if isinstance(handle._input, StartActivityInput)
+                        else self._info.task_queue
+                    ),
+                    is_local=isinstance(handle._input, StartLocalActivityInput),
+                )
+
+            elif command_seq in self._pending_child_workflows:
+                handle = self._pending_child_workflows[command_seq]
+                return temporalio.converter.WorkflowSerializationContext(
+                    namespace=self._info.namespace,
+                    workflow_id=handle._input.id,
+                )
+
+            elif command_seq in self._pending_external_signals:
+                _, workflow_id = self._pending_external_signals[command_seq]
+                return temporalio.converter.WorkflowSerializationContext(
+                    namespace=self._info.namespace,
+                    workflow_id=workflow_id,
+                )
+
+            elif command_seq in self._pending_nexus_operations:
+                # We don't set any context for nexus operations
+                pass
+
     def _instantiate_workflow_object(self) -> Any:
         if not self._workflow_input:
             raise RuntimeError("Expected workflow input. This is a Python SDK bug.")