Keep converter serialization context when external storage is enabled

changelog/fixed/contextual-courier-crab.md

@@ -0,0 +1 @@
+Fix the payload converter and failure converter losing their serialization context when external storage is configured. A converter that implements WithSerializationContext now gets the workflow, activity or other context whether or not external storage is enabled.

temporalio/converter/_data_converter.py

@@ -219,10 +219,18 @@ def _with_store_context(
         """Return an instance with ``store_ctx`` bound into :attr:`external_storage`."""
         if self.external_storage is None:
             return self
-        return dataclasses.replace(
-            self,
-            external_storage=self.external_storage._with_store_context(store_ctx),
+        # dataclasses.replace reruns __post_init__, which builds new payload and
+        # failure converters and drops the context set by with_context. Keep the
+        # current ones, like with_context does.
+        cloned = dataclasses.replace(self)
+        object.__setattr__(cloned, "payload_converter", self.payload_converter)
+        object.__setattr__(cloned, "failure_converter", self.failure_converter)
+        object.__setattr__(
+            cloned,
+            "external_storage",
+            self.external_storage._with_store_context(store_ctx),
         )
+        return cloned
 
     def _with_contexts(
         self,

tests/test_serialization_context.py

@@ -44,6 +44,7 @@
     DefaultFailureConverter,
     DefaultPayloadConverter,
     EncodingPayloadConverter,
+    ExternalStorage,
     JSONPlainPayloadConverter,
     NexusSerializationContext,
     PayloadCodec,
@@ -56,6 +57,7 @@
 from temporalio.testing import WorkflowEnvironment
 from temporalio.worker import Replayer, Worker
 from temporalio.worker._workflow_instance import UnsandboxedWorkflowRunner
+from tests.test_extstore import InMemoryTestDriver
 
 
 @dataclass
@@ -192,8 +194,10 @@ async def run(self, data: TraceData) -> TraceData:
         return data
 
 
+@pytest.mark.parametrize("use_external_storage", [False, True])
 async def test_payload_conversion_calls_follow_expected_sequence_and_contexts(
     client: Client,
+    use_external_storage: bool,
 ):
     workflow_id = str(uuid.uuid4())
     task_queue = str(uuid.uuid4())
@@ -202,6 +206,12 @@ async def test_payload_conversion_calls_follow_expected_sequence_and_contexts(
     config["data_converter"] = dataclasses.replace(
         DataConverter.default,
         payload_converter_class=SerializationContextCompositePayloadConverter,
+        # External storage must not drop the context
+        external_storage=(
+            ExternalStorage(drivers=[InMemoryTestDriver()])
+            if use_external_storage
+            else None
+        ),
     )
     client = Client(**config)