Fix command context awareness to apply to all of DataConverter (#1387)

Author: jmaeagle99

temporalio/worker/_workflow.py

@@ -289,18 +289,13 @@ async def _handle_activation(
                 workflow_id=workflow_id,
             )
             data_converter = self._data_converter.with_context(workflow_context)
-            if self._data_converter.payload_codec:
-                assert data_converter.payload_codec
-                if workflow:
-                    data_converter = dataclasses.replace(
-                        data_converter,
-                        payload_codec=_CommandAwarePayloadCodec(
-                            workflow.instance,
-                            context_free_payload_codec=self._data_converter.payload_codec,
-                            workflow_context_payload_codec=data_converter.payload_codec,
-                            workflow_context=workflow_context,
-                        ),
-                    )
+            if workflow:
+                data_converter = _CommandAwareDataConverter.create(
+                    instance=workflow.instance,
+                    context_free_dc=self._data_converter,
+                    workflow_context_dc=data_converter,
+                    workflow_context=workflow_context,
+                )
             download_metrics = await temporalio.bridge.worker.decode_activation(
                 act,
                 data_converter,
@@ -395,19 +390,16 @@ async def _handle_activation(
         completion.run_id = act.run_id
 
         # Encode completion
-        if self._data_converter.payload_codec and workflow:
-            assert data_converter.payload_codec
-            data_converter = dataclasses.replace(
-                data_converter,
-                payload_codec=_CommandAwarePayloadCodec(
-                    workflow.instance,
-                    context_free_payload_codec=self._data_converter.payload_codec,
-                    workflow_context_payload_codec=data_converter.payload_codec,
-                    workflow_context=temporalio.converter.WorkflowSerializationContext(
-                        namespace=self._namespace,
-                        workflow_id=workflow.workflow_id,
-                    ),
-                ),
+        if workflow:
+            workflow_context = temporalio.converter.WorkflowSerializationContext(
+                namespace=self._namespace,
+                workflow_id=workflow.workflow_id,
+            )
+            data_converter = _CommandAwareDataConverter.create(
+                instance=workflow.instance,
+                context_free_dc=self._data_converter,
+                workflow_context_dc=self._data_converter.with_context(workflow_context),
+                workflow_context=workflow_context,
             )
 
         upload_metrics = temporalio.converter._extstore.StorageOperationMetrics()
@@ -837,45 +829,74 @@ def attempt_deadlock_interruption(self) -> None:
 
 
 @dataclass(frozen=True)
-class _CommandAwarePayloadCodec(temporalio.converter.PayloadCodec):
-    """A payload codec that sets serialization context for the command associated with each payload.
+class _CommandAwareDataConverter(temporalio.converter.DataConverter):
+    """Data converter that resolves serialization context per-command.
 
-    This codec responds to the context variable set by
+    Responds to the context variable set by
     :py:class:`_command_aware_visitor.CommandAwarePayloadVisitor`.
     """
 
-    instance: WorkflowInstance
-    context_free_payload_codec: temporalio.converter.PayloadCodec
-    workflow_context_payload_codec: temporalio.converter.PayloadCodec
-    workflow_context: temporalio.converter.WorkflowSerializationContext
-
-    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)
+    _ca_instance: WorkflowInstance = dataclasses.field(
+        default=None,
+        repr=False,
+        compare=False,  # type: ignore[assignment]
+    )
+    _ca_context_free_dc: temporalio.converter.DataConverter = dataclasses.field(
+        default=None,
+        repr=False,
+        compare=False,  # type: ignore[assignment]
+    )
+    _ca_workflow_context_dc: temporalio.converter.DataConverter = dataclasses.field(
+        default=None,
+        repr=False,
+        compare=False,  # type: ignore[assignment]
+    )
+    _ca_workflow_context: temporalio.converter.WorkflowSerializationContext = (
+        dataclasses.field(
+            default=None,
+            repr=False,
+            compare=False,  # type: ignore[assignment]
+        )
+    )
 
-    def _get_current_command_codec(self) -> temporalio.converter.PayloadCodec:
-        if not isinstance(
-            self.context_free_payload_codec,
-            temporalio.converter.WithSerializationContext,
-        ):
-            return self.context_free_payload_codec
+    @staticmethod
+    def create(
+        instance: WorkflowInstance,
+        context_free_dc: temporalio.converter.DataConverter,
+        workflow_context_dc: temporalio.converter.DataConverter,
+        workflow_context: temporalio.converter.WorkflowSerializationContext,
+    ) -> _CommandAwareDataConverter:
+        return _CommandAwareDataConverter(
+            payload_converter_class=workflow_context_dc.payload_converter_class,
+            payload_codec=workflow_context_dc.payload_codec,
+            failure_converter_class=workflow_context_dc.failure_converter_class,
+            payload_limits=workflow_context_dc.payload_limits,
+            external_storage=workflow_context_dc.external_storage,
+            _ca_instance=instance,
+            _ca_context_free_dc=context_free_dc,
+            _ca_workflow_context_dc=workflow_context_dc,
+            _ca_workflow_context=workflow_context,
+        )
 
-        if context := self.instance.get_serialization_context(
+    def _get_current_dc(self) -> temporalio.converter.DataConverter:
+        context = self._ca_instance.get_serialization_context(
             _command_aware_visitor.current_command_info.get(),
-        ):
-            if context == self.workflow_context:
-                return self.workflow_context_payload_codec
-            return self.context_free_payload_codec.with_context(context)
+        )
+        if context is None:
+            return self._ca_context_free_dc
+        if context == self._ca_workflow_context:
+            return self._ca_workflow_context_dc
+        return self._ca_context_free_dc.with_context(context)
+
+    async def _encode_payload_sequence(
+        self, payloads: Sequence[temporalio.api.common.v1.Payload]
+    ) -> list[temporalio.api.common.v1.Payload]:
+        return await self._get_current_dc()._encode_payload_sequence(payloads)
 
-        return self.context_free_payload_codec
+    async def _decode_payload_sequence(
+        self, payloads: Sequence[temporalio.api.common.v1.Payload]
+    ) -> list[temporalio.api.common.v1.Payload]:
+        return await self._get_current_dc()._decode_payload_sequence(payloads)
 
 
 class _InterruptDeadlockError(BaseException):

tests/test_serialization_context.py

@@ -40,10 +40,15 @@
     DefaultFailureConverter,
     DefaultPayloadConverter,
     EncodingPayloadConverter,
+    ExternalStorage,
     JSONPlainPayloadConverter,
     PayloadCodec,
     PayloadConverter,
     SerializationContext,
+    StorageDriver,
+    StorageDriverClaim,
+    StorageDriverRetrieveContext,
+    StorageDriverStoreContext,
     WithSerializationContext,
     WorkflowSerializationContext,
 )
@@ -1914,3 +1919,84 @@ async def test_user_customization_of_default_payload_converter(
                 id=wf_id,
                 task_queue=task_queue,
             )
+
+
+# Child workflow external storage context test
+
+
+class ContextTrackingStorageDriver(StorageDriver):
+    """In-memory driver that records the serialization context on each store/retrieve."""
+
+    def __init__(self) -> None:
+        self._storage: dict[str, bytes] = {}
+        self.store_contexts: list[SerializationContext | None] = []
+
+    def name(self) -> str:
+        return "context-tracking"
+
+    async def store(
+        self,
+        context: StorageDriverStoreContext,
+        payloads: Sequence[temporalio.api.common.v1.Payload],
+    ) -> list[StorageDriverClaim]:
+        self.store_contexts.append(context.serialization_context)
+        claims: list[StorageDriverClaim] = []
+        for payload in payloads:
+            key = f"payload-{len(self._storage)}"
+            self._storage[key] = payload.SerializeToString()
+            claims.append(StorageDriverClaim(claim_data={"key": key}))
+        return claims
+
+    async def retrieve(
+        self,
+        context: StorageDriverRetrieveContext,
+        claims: Sequence[StorageDriverClaim],
+    ) -> list[temporalio.api.common.v1.Payload]:
+        results: list[temporalio.api.common.v1.Payload] = []
+        for claim in claims:
+            payload = temporalio.api.common.v1.Payload()
+            payload.ParseFromString(self._storage[claim.claim_data["key"]])
+            results.append(payload)
+        return results
+
+
+async def test_child_workflow_external_storage_with_context(client: Client):
+    """External storage should receive the child workflow's context, not the parent's."""
+    workflow_id = str(uuid.uuid4())
+    child_workflow_id = f"{workflow_id}-child"
+    task_queue = str(uuid.uuid4())
+
+    driver = ContextTrackingStorageDriver()
+    config = client.config()
+    config["data_converter"] = dataclasses.replace(
+        DataConverter.default,
+        external_storage=ExternalStorage(
+            drivers=[driver],
+            payload_size_threshold=None,
+        ),
+    )
+    client = Client(**config)
+
+    async with Worker(
+        client,
+        task_queue=task_queue,
+        workflows=[ChildWorkflowCodecTestWorkflow, EchoWorkflow],
+        workflow_runner=UnsandboxedWorkflowRunner(),
+    ):
+        await client.execute_workflow(
+            ChildWorkflowCodecTestWorkflow.run,
+            TraceData(),
+            id=workflow_id,
+            task_queue=task_queue,
+        )
+
+    child_context = WorkflowSerializationContext(
+        namespace=client.namespace,
+        workflow_id=child_workflow_id,
+    )
+    # store_contexts[0]: parent input encode        → parent context
+    # store_contexts[1]: child workflow input encode → child context
+    # store_contexts[2]: child workflow result encode → child context
+    # store_contexts[3]: parent result encode        → parent context
+    child_context_count = sum(1 for c in driver.store_contexts if c == child_context)
+    assert child_context_count == 2