Make it work

Author: dandavison

temporalio/client.py

@@ -5930,8 +5930,13 @@ async def _populate_start_workflow_execution_request(
         req.workflow_type.name = input.workflow
         req.task_queue.name = input.task_queue
         if input.args:
+            context = temporalio.converter.WorkflowSerializationContext(
+                namespace=self._client.namespace, workflow_id=input.id
+            )
             req.input.payloads.extend(
-                await self._client.data_converter.encode(input.args)
+                await self._client.data_converter._with_context(context).encode(
+                    input.args
+                )
             )
         if input.execution_timeout is not None:
             req.workflow_execution_timeout.FromTimedelta(input.execution_timeout)

temporalio/converter.py

@@ -1311,35 +1311,26 @@ async def decode_failure(
         return self.failure_converter.from_failure(failure, self.payload_converter)
 
     def _with_context(self, context: Optional[SerializationContext]) -> Self:
-        new_self = type(self).__new__(type(self))
-        setattr(
-            new_self,
-            "payload_converter",
-            (
-                self.payload_converter.with_context(context)
-                if isinstance(self.payload_converter, WithSerializationContext)
-                else self.payload_converter
-            ),
+        payload_converter = (
+            self.payload_converter.with_context(context)
+            if isinstance(self.payload_converter, WithSerializationContext)
+            else self.payload_converter
         )
-        setattr(
-            new_self,
-            "payload_codec",
-            (
-                self.payload_codec.with_context(context)
-                if isinstance(self.payload_codec, WithSerializationContext)
-                else self.payload_codec
-            ),
+        payload_codec = (
+            self.payload_codec.with_context(context)
+            if isinstance(self.payload_codec, WithSerializationContext)
+            else self.payload_codec
         )
-        setattr(
-            new_self,
-            "failure_converter",
-            (
-                self.failure_converter.with_context(context)
-                if isinstance(self.failure_converter, WithSerializationContext)
-                else self.failure_converter
-            ),
+        failure_converter = (
+            self.failure_converter.with_context(context)
+            if isinstance(self.failure_converter, WithSerializationContext)
+            else self.failure_converter
         )
-        return new_self
+        cloned = dataclasses.replace(self)
+        object.__setattr__(cloned, "payload_converter", payload_converter)
+        object.__setattr__(cloned, "payload_codec", payload_codec)
+        object.__setattr__(cloned, "failure_converter", failure_converter)
+        return cloned
 
 
 DefaultPayloadConverter.default_encoding_payload_converters = (

temporalio/worker/_workflow_instance.py

@@ -984,7 +984,14 @@ def _apply_initialize_workflow(
         async def run_workflow(input: ExecuteWorkflowInput) -> None:
             try:
                 result = await self._inbound.execute_workflow(input)
-                result_payloads = self._payload_converter.to_payloads([result])
+                converter = self._payload_converter
+                if isinstance(converter, temporalio.converter.WithSerializationContext):
+                    context = temporalio.converter.WorkflowSerializationContext(
+                        namespace=self._info.namespace,
+                        workflow_id=self._info.workflow_id,
+                    )
+                    converter = converter.with_context(context)
+                result_payloads = converter.to_payloads([result])
                 if len(result_payloads) != 1:
                     raise ValueError(
                         f"Expected 1 result payload, got {len(result_payloads)}"

tests/test_serialization_context.py

@@ -2,7 +2,7 @@
 
 import dataclasses
 import uuid
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any, Optional, Type
 
 from temporalio import workflow
@@ -21,14 +21,22 @@
 
 
 @dataclass
-class WorkflowData:
-    workflow_context: Optional[WorkflowSerializationContext] = None
+class PayloadConverterTraceData:
+    to_payload: Optional[WorkflowSerializationContext] = None
+    from_payload: Optional[WorkflowSerializationContext] = None
+
+
+@dataclass
+class TraceData:
+    workflow_context: PayloadConverterTraceData = field(
+        default_factory=PayloadConverterTraceData
+    )
 
 
 @workflow.defn
 class SerializationContextTestWorkflow:
     @workflow.run
-    async def run(self, input: WorkflowData) -> WorkflowData:
+    async def run(self, input: TraceData) -> TraceData:
         return input
 
 
@@ -51,9 +59,7 @@ def with_context(
         return SerializationContextTestEncodingPayloadConverter(context)
 
     def to_payload(self, value: Any) -> Optional[Payload]:
-        assert isinstance(value, WorkflowData)
-        assert isinstance(self.context, WorkflowSerializationContext)
-        value.workflow_context = self.context
+        value.workflow_context.to_payload = self.context
         return None
 
     def from_payload(self, payload: Payload, type_hint: Optional[Type] = None) -> Any:
@@ -62,11 +68,14 @@ def from_payload(self, payload: Payload, type_hint: Optional[Type] = None) -> An
 
 
 class SerializationContextTestPayloadConverter(CompositePayloadConverter):
-    def __init__(self):
-        super().__init__(
-            SerializationContextTestEncodingPayloadConverter(None),
-            *DefaultPayloadConverter.default_encoding_payload_converters,
-        )
+    def __init__(self, *converters):
+        # TODO: we cannot expect users to do this
+        if not converters:
+            converters = (
+                SerializationContextTestEncodingPayloadConverter(None),
+                *DefaultPayloadConverter.default_encoding_payload_converters,
+            )
+        super().__init__(*converters)
 
 
 data_converter = dataclasses.replace(
@@ -81,6 +90,10 @@ async def test_workflow_payload_conversion_can_be_given_access_to_serialization_
     workflow_id = str(uuid.uuid4())
     task_queue = str(uuid.uuid4())
 
+    config = client.config()
+    config["data_converter"] = data_converter
+    client = Client(**config)
+
     async with Worker(
         client,
         task_queue=task_queue,
@@ -89,12 +102,12 @@ async def test_workflow_payload_conversion_can_be_given_access_to_serialization_
     ):
         result = await client.execute_workflow(
             SerializationContextTestWorkflow.run,
-            WorkflowData(),
+            TraceData(),
             id=workflow_id,
             task_queue=task_queue,
         )
 
-        assert result.workflow_context == WorkflowSerializationContext(
+        assert result.workflow_context.to_payload == WorkflowSerializationContext(
             namespace="default",
             workflow_id=workflow_id,
         )