Compute context, not payload codec

Author: dandavison

temporalio/worker/_workflow.py

@@ -290,7 +290,6 @@ async def _handle_activation(
                     payload_codec = _CommandAwarePayloadCodec(
                         workflow.instance,
                         context_free_payload_codec=self._data_converter.payload_codec,
-                        workflow_context_payload_codec=data_converter.payload_codec,
                     )
                 await temporalio.bridge.worker.decode_activation(
                     act,
@@ -367,7 +366,6 @@ async def _handle_activation(
             payload_codec = _CommandAwarePayloadCodec(
                 workflow.instance,
                 context_free_payload_codec=self._data_converter.payload_codec,
-                workflow_context_payload_codec=data_converter.payload_codec,
             )
             try:
                 await temporalio.bridge.worker.encode_completion(
@@ -734,7 +732,6 @@ class _CommandAwarePayloadCodec(temporalio.converter.PayloadCodec):
 
     instance: WorkflowInstance
     context_free_payload_codec: temporalio.converter.PayloadCodec
-    workflow_context_payload_codec: temporalio.converter.PayloadCodec
 
     async def encode(
         self,
@@ -749,11 +746,18 @@ async def decode(
         return await self._get_current_command_codec().decode(payloads)
 
     def _get_current_command_codec(self) -> temporalio.converter.PayloadCodec:
-        return self.instance.get_payload_codec_with_context(
+        if not isinstance(
             self.context_free_payload_codec,
-            self.workflow_context_payload_codec,
+            temporalio.converter.WithSerializationContext,
+        ):
+            return self.context_free_payload_codec
+
+        if context := self.instance.get_serialization_context(
             temporalio.bridge._visitor.current_command_info.get(),
-        )
+        ):
+            return self.context_free_payload_codec.with_context(context)
+
+        return self.context_free_payload_codec
 
 
 class _InterruptDeadlockError(BaseException):

temporalio/worker/_workflow_instance.py

@@ -170,22 +170,18 @@ def activate(
         raise NotImplementedError
 
     @abstractmethod
-    def get_payload_codec_with_context(
+    def get_serialization_context(
         self,
-        base_payload_codec: temporalio.converter.PayloadCodec,
-        workflow_context_payload_codec: temporalio.converter.PayloadCodec,
         command_info: Optional[temporalio.bridge._visitor.CommandInfo],
-    ) -> temporalio.converter.PayloadCodec:
-        """Return a payload codec with appropriate serialization context.
+    ) -> Optional[temporalio.converter.SerializationContext]:
+        """Return appropriate serialization context.
 
         Args:
-            base_payload_codec: The base payload codec to apply context to.
-            workflow_context_payload_codec: A payload codec that already has workflow context set.
             command_info: Optional information identifying the associated command. If set, the payload
                          codec will have serialization context set appropriately for that command.
 
         Returns:
-            The payload codec.
+            The serialization context, or None if no context should be set.
         """
         raise NotImplementedError
 
@@ -2116,22 +2112,18 @@ def _converters_with_context(
             failure_converter = failure_converter.with_context(context)
         return payload_converter, failure_converter
 
-    def get_payload_codec_with_context(
+    def get_serialization_context(
         self,
-        base_payload_codec: temporalio.converter.PayloadCodec,
-        workflow_context_payload_codec: temporalio.converter.PayloadCodec,
         command_info: Optional[temporalio.bridge._visitor.CommandInfo],
-    ) -> temporalio.converter.PayloadCodec:
-        if not isinstance(
-            base_payload_codec,
-            temporalio.converter.WithSerializationContext,
-        ):
-            return base_payload_codec
-
+    ) -> Optional[temporalio.converter.SerializationContext]:
+        workflow_context = temporalio.converter.WorkflowSerializationContext(
+            namespace=self._info.namespace,
+            workflow_id=self._info.workflow_id,
+        )
         if command_info is None:
             # Use payload codec with workflow context by default (i.e. for payloads not associated
             # with a pending command)
-            return workflow_context_payload_codec
+            return workflow_context
 
         if (
             command_info.command_type
@@ -2140,22 +2132,18 @@ def get_payload_codec_with_context(
         ):
             # Use the activity's context
             activity_handle = self._pending_activities[command_info.command_seq]
-            return base_payload_codec.with_context(
-                temporalio.converter.ActivitySerializationContext(
-                    namespace=self._info.namespace,
-                    workflow_id=self._info.workflow_id,
-                    workflow_type=self._info.workflow_type,
-                    activity_type=activity_handle._input.activity,
-                    activity_task_queue=(
-                        activity_handle._input.task_queue
-                        if isinstance(activity_handle._input, StartActivityInput)
-                        and activity_handle._input.task_queue
-                        else self._info.task_queue
-                    ),
-                    is_local=isinstance(
-                        activity_handle._input, StartLocalActivityInput
-                    ),
-                )
+            return temporalio.converter.ActivitySerializationContext(
+                namespace=self._info.namespace,
+                workflow_id=self._info.workflow_id,
+                workflow_type=self._info.workflow_type,
+                activity_type=activity_handle._input.activity,
+                activity_task_queue=(
+                    activity_handle._input.task_queue
+                    if isinstance(activity_handle._input, StartActivityInput)
+                    and activity_handle._input.task_queue
+                    else self._info.task_queue
+                ),
+                is_local=isinstance(activity_handle._input, StartLocalActivityInput),
             )
 
         elif (
@@ -2165,11 +2153,9 @@ def get_payload_codec_with_context(
         ):
             # Use the child workflow's context
             child_wf_handle = self._pending_child_workflows[command_info.command_seq]
-            return base_payload_codec.with_context(
-                temporalio.converter.WorkflowSerializationContext(
-                    namespace=self._info.namespace,
-                    workflow_id=child_wf_handle._input.id,
-                )
+            return temporalio.converter.WorkflowSerializationContext(
+                namespace=self._info.namespace,
+                workflow_id=child_wf_handle._input.id,
             )
 
         elif (
@@ -2181,11 +2167,9 @@ def get_payload_codec_with_context(
             _, target_workflow_id = self._pending_external_signals[
                 command_info.command_seq
             ]
-            return base_payload_codec.with_context(
-                temporalio.converter.WorkflowSerializationContext(
-                    namespace=self._info.namespace,
-                    workflow_id=target_workflow_id,
-                )
+            return temporalio.converter.WorkflowSerializationContext(
+                namespace=self._info.namespace,
+                workflow_id=target_workflow_id,
             )
 
         elif (
@@ -2196,11 +2180,11 @@ def get_payload_codec_with_context(
             # Use empty context for nexus operations: users will never want to encrypt using a
             # key derived from caller workflow context because the caller workflow context is
             # not available on the handler side for decryption.
-            return base_payload_codec
+            return None
 
         else:
             # Use payload codec with workflow context for all other payloads
-            return workflow_context_payload_codec
+            return workflow_context
 
     def _instantiate_workflow_object(self) -> Any:
         if not self._workflow_input:

temporalio/worker/workflow_sandbox/_in_sandbox.py

@@ -82,15 +82,9 @@ def activate(
         """Send activation to this instance."""
         return self.instance.activate(act)
 
-    def get_payload_codec_with_context(
+    def get_serialization_context(
         self,
-        base_payload_codec: temporalio.converter.PayloadCodec,
-        workflow_context_payload_codec: temporalio.converter.PayloadCodec,
         command_info: Optional[temporalio.bridge._visitor.CommandInfo],
-    ) -> temporalio.converter.PayloadCodec:
-        """Get payload codec with context."""
-        return self.instance.get_payload_codec_with_context(
-            base_payload_codec,
-            workflow_context_payload_codec,
-            command_info,
-        )
+    ) -> Optional[temporalio.converter.SerializationContext]:
+        """Get serialization context."""
+        return self.instance.get_serialization_context(command_info)

temporalio/worker/workflow_sandbox/_runner.py

@@ -187,23 +187,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_with_context(
+    def get_serialization_context(
         self,
-        base_payload_codec: temporalio.converter.PayloadCodec,
-        workflow_context_payload_codec: temporalio.converter.PayloadCodec,
         command_info: Optional[temporalio.bridge._visitor.CommandInfo],
-    ) -> temporalio.converter.PayloadCodec:
+    ) -> Optional[temporalio.converter.SerializationContext]:
         # 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_with_context(__temporal_base_payload_codec, __temporal_workflow_context_payload_codec, __temporal_command_info)\n",
+                "  __temporal_context = __temporal_in_sandbox.get_serialization_context(__temporal_command_info)\n",
                 __temporal_importer=self.importer,
-                __temporal_base_payload_codec=base_payload_codec,
-                __temporal_workflow_context_payload_codec=workflow_context_payload_codec,
                 __temporal_command_info=command_info,
             )
-            return self.globals_and_locals.pop("__temporal_codec", None)  # type: ignore
+            return self.globals_and_locals.pop("__temporal_context", None)  # type: ignore
         finally:
             self.importer.restriction_context.is_runtime = False

tests/worker/test_workflow.py

@@ -1611,17 +1611,11 @@ def activate(self, act: WorkflowActivation) -> WorkflowActivationCompletion:
         self._runner._pairs.append((act, comp))
         return comp
 
-    def get_payload_codec_with_context(
+    def get_serialization_context(
         self,
-        base_payload_codec: temporalio.converter.PayloadCodec,
-        workflow_context_payload_codec: temporalio.converter.PayloadCodec,
         command_info: Optional[temporalio.bridge._visitor.CommandInfo],
-    ) -> temporalio.converter.PayloadCodec:
-        return self._unsandboxed.get_payload_codec_with_context(
-            base_payload_codec,
-            workflow_context_payload_codec,
-            command_info,
-        )
+    ) -> Optional[temporalio.converter.SerializationContext]:
+        return self._unsandboxed.get_serialization_context(command_info)
 
 
 async def test_workflow_with_custom_runner(client: Client):