Respect run payload codec for workflow commands

Author: durable-workflow-ops

src/durable_workflow/worker.py

@@ -23,6 +23,10 @@
 log = logging.getLogger("durable_workflow.worker")
 
 
+def _command_payload_codec(codec: object) -> str:
+    return codec if isinstance(codec, str) and codec in serializer.SUPPORTED_CODECS else serializer.AVRO_CODEC
+
+
 def _activity_name(fn: Callable[..., Any]) -> str:
     return getattr(fn, "__activity_name__", fn.__name__)
 
@@ -172,6 +176,7 @@ async def _run_workflow_task(self, task: dict[str, Any]) -> list[dict[str, Any]]
 
         start_input: list[Any] = []
         codec = task.get("payload_codec")
+        command_codec = _command_payload_codec(codec)
         raw_args = task.get("arguments")
         try:
             decoded = serializer.decode_envelope(raw_args, codec=codec)
@@ -262,7 +267,10 @@ async def _run_workflow_task(self, task: dict[str, Any]) -> list[dict[str, Any]]
                 log.warning("failed to report replay failure: %s", fe)
             return None
 
-        commands = [c.to_server_command(self.task_queue) for c in outcome.commands]
+        commands = [
+            c.to_server_command(self.task_queue, payload_codec=command_codec)
+            for c in outcome.commands
+        ]
         log.info(
             "completing workflow task %s with %d command(s): %s",
             task_id,

src/durable_workflow/workflow.py

@@ -36,11 +36,13 @@ class ScheduleActivity:
     arguments: list[Any]
     queue: str | None = None
 
-    def to_server_command(self, task_queue: str) -> dict[str, Any]:
+    def to_server_command(
+        self, task_queue: str, *, payload_codec: str = serializer.AVRO_CODEC
+    ) -> dict[str, Any]:
         return {
             "type": "schedule_activity",
             "activity_type": self.activity_type,
-            "arguments": serializer.envelope(self.arguments),
+            "arguments": serializer.envelope(self.arguments, codec=payload_codec),
             "queue": self.queue or task_queue,
         }
 
@@ -49,7 +51,9 @@ def to_server_command(self, task_queue: str) -> dict[str, Any]:
 class StartTimer:
     delay_seconds: int
 
-    def to_server_command(self, task_queue: str) -> dict[str, Any]:
+    def to_server_command(
+        self, task_queue: str, *, payload_codec: str = serializer.AVRO_CODEC
+    ) -> dict[str, Any]:
         return {
             "type": "start_timer",
             "delay_seconds": self.delay_seconds,
@@ -60,10 +64,12 @@ def to_server_command(self, task_queue: str) -> dict[str, Any]:
 class CompleteWorkflow:
     result: Any
 
-    def to_server_command(self, task_queue: str) -> dict[str, Any]:
+    def to_server_command(
+        self, task_queue: str, *, payload_codec: str = serializer.AVRO_CODEC
+    ) -> dict[str, Any]:
         return {
             "type": "complete_workflow",
-            "result": serializer.envelope(self.result),
+            "result": serializer.envelope(self.result, codec=payload_codec),
         }
 
 
@@ -73,7 +79,9 @@ class FailWorkflow:
     exception_type: str | None = None
     non_retryable: bool = False
 
-    def to_server_command(self, task_queue: str) -> dict[str, Any]:
+    def to_server_command(
+        self, task_queue: str, *, payload_codec: str = serializer.AVRO_CODEC
+    ) -> dict[str, Any]:
         cmd: dict[str, Any] = {
             "type": "fail_workflow",
             "message": self.message,
@@ -91,11 +99,13 @@ class ContinueAsNew:
     arguments: list[Any] = field(default_factory=list)
     task_queue: str | None = None
 
-    def to_server_command(self, task_queue: str) -> dict[str, Any]:
+    def to_server_command(
+        self, task_queue: str, *, payload_codec: str = serializer.AVRO_CODEC
+    ) -> dict[str, Any]:
         cmd: dict[str, Any] = {"type": "continue_as_new"}
         if self.workflow_type is not None:
             cmd["workflow_type"] = self.workflow_type
-        cmd["arguments"] = serializer.envelope(self.arguments)
+        cmd["arguments"] = serializer.envelope(self.arguments, codec=payload_codec)
         cmd["queue"] = self.task_queue or task_queue
         return cmd
 
@@ -104,10 +114,12 @@ def to_server_command(self, task_queue: str) -> dict[str, Any]:
 class RecordSideEffect:
     result: Any
 
-    def to_server_command(self, task_queue: str) -> dict[str, Any]:
+    def to_server_command(
+        self, task_queue: str, *, payload_codec: str = serializer.AVRO_CODEC
+    ) -> dict[str, Any]:
         return {
             "type": "record_side_effect",
-            "result": serializer.envelope(self.result),
+            "result": serializer.encode(self.result, codec=payload_codec),
         }
 
 
@@ -118,11 +130,13 @@ class StartChildWorkflow:
     task_queue: str | None = None
     parent_close_policy: str | None = None
 
-    def to_server_command(self, task_queue: str) -> dict[str, Any]:
+    def to_server_command(
+        self, task_queue: str, *, payload_codec: str = serializer.AVRO_CODEC
+    ) -> dict[str, Any]:
         cmd: dict[str, Any] = {
             "type": "start_child_workflow",
             "workflow_type": self.workflow_type,
-            "arguments": serializer.envelope(self.arguments),
+            "arguments": serializer.envelope(self.arguments, codec=payload_codec),
         }
         if self.task_queue is not None:
             cmd["queue"] = self.task_queue
@@ -140,7 +154,9 @@ class RecordVersionMarker:
     min_supported: int
     max_supported: int
 
-    def to_server_command(self, task_queue: str) -> dict[str, Any]:
+    def to_server_command(
+        self, task_queue: str, *, payload_codec: str = serializer.AVRO_CODEC
+    ) -> dict[str, Any]:
         return {
             "type": "record_version_marker",
             "change_id": self.change_id,
@@ -154,7 +170,9 @@ def to_server_command(self, task_queue: str) -> dict[str, Any]:
 class UpsertSearchAttributes:
     attributes: dict[str, Any]
 
-    def to_server_command(self, task_queue: str) -> dict[str, Any]:
+    def to_server_command(
+        self, task_queue: str, *, payload_codec: str = serializer.AVRO_CODEC
+    ) -> dict[str, Any]:
         return {
             "type": "upsert_search_attributes",
             "attributes": self.attributes,

tests/integration/test_polyglot.py

@@ -89,7 +89,7 @@ async def test_python_workflow_calls_php_activity(server_url: str, server_token:
         outcome = replay(PolyglotPythonWorkflow, history, start_input, run_id=wf_task.get("run_id", ""))
         assert len(outcome.commands) == 1
         cmd = outcome.commands[0]
-        server_cmd = cmd.to_server_command(task_queue)
+        server_cmd = cmd.to_server_command(task_queue, payload_codec=codec)
         assert server_cmd["type"] == "schedule_activity"
         assert server_cmd["activity_type"] == "tests.polyglot.php-activity"
 
@@ -166,7 +166,8 @@ async def test_python_workflow_calls_php_activity(server_url: str, server_token:
         outcome2 = replay(PolyglotPythonWorkflow, history2, start_input2, run_id=wf_task2.get("run_id", ""))
         assert len(outcome2.commands) == 1
         cmd2 = outcome2.commands[0]
-        server_cmd2 = cmd2.to_server_command(task_queue)
+        codec2 = wf_task2.get("payload_codec")
+        server_cmd2 = cmd2.to_server_command(task_queue, payload_codec=codec2)
 
         # Debug: inspect command before asserting
         print("\n=== Replay outcome ===")

tests/integration/test_smoke.py

@@ -87,7 +87,7 @@ async def test_greeter_workflow_end_to_end(server_url: str, server_token: str) -
         outcome = replay(SmokeGreeterWorkflow, history, start_input, run_id=wf_task.get("run_id", ""))
         assert len(outcome.commands) == 1
         cmd = outcome.commands[0]
-        server_cmd = cmd.to_server_command(task_queue)
+        server_cmd = cmd.to_server_command(task_queue, payload_codec=codec)
         assert server_cmd["type"] == "schedule_activity"
 
         # 5. Complete workflow task with ScheduleActivity command
@@ -137,7 +137,8 @@ async def test_greeter_workflow_end_to_end(server_url: str, server_token: str) -
         outcome2 = replay(SmokeGreeterWorkflow, history2, start_input2, run_id=wf_task2.get("run_id", ""))
         assert len(outcome2.commands) == 1
         cmd2 = outcome2.commands[0]
-        server_cmd2 = cmd2.to_server_command(task_queue)
+        codec2 = wf_task2.get("payload_codec")
+        server_cmd2 = cmd2.to_server_command(task_queue, payload_codec=codec2)
         assert server_cmd2["type"] == "complete_workflow"
 
         # 11. Complete the final workflow task

tests/test_replay.py

@@ -102,6 +102,14 @@ def test_server_command_shape(self) -> None:
         assert server_cmd["type"] == "schedule_activity"
         assert server_cmd["activity_type"] == "greet"
         assert server_cmd["queue"] == "default-queue"
+        assert server_cmd["arguments"]["codec"] == "avro"
+
+    def test_server_command_uses_payload_codec(self) -> None:
+        outcome = replay(OneActivity, [], ["world"])
+        cmd = outcome.commands[0]
+        server_cmd = cmd.to_server_command("default-queue", payload_codec="json")
+        assert server_cmd["arguments"]["codec"] == "json"
+        assert serializer.decode(server_cmd["arguments"]["blob"], codec="json") == ["world"]
 
 
 class TestTwoActivities:
@@ -188,6 +196,12 @@ def test_server_command(self) -> None:
         assert server_cmd["result"]["codec"] == "avro"
         assert serializer.decode(server_cmd["result"]["blob"], codec="avro") == {"key": "val"}
 
+    def test_server_command_uses_payload_codec(self) -> None:
+        cmd = CompleteWorkflow(result={"key": "val"})
+        server_cmd = cmd.to_server_command("q", payload_codec="json")
+        assert server_cmd["result"]["codec"] == "json"
+        assert serializer.decode(server_cmd["result"]["blob"], codec="json") == {"key": "val"}
+
 
 @workflow.defn(name="continue-as-new-wf")
 class ContinueAsNewWorkflow:
@@ -289,8 +303,12 @@ def test_server_command_shape(self) -> None:
         cmd = RecordSideEffect(result={"key": "val"})
         sc = cmd.to_server_command("q")
         assert sc["type"] == "record_side_effect"
-        assert sc["result"]["codec"] == "avro"
-        assert serializer.decode(sc["result"]["blob"], codec="avro") == {"key": "val"}
+        assert serializer.decode(sc["result"], codec="avro") == {"key": "val"}
+
+    def test_server_command_uses_payload_codec(self) -> None:
+        cmd = RecordSideEffect(result={"key": "val"})
+        sc = cmd.to_server_command("q", payload_codec="json")
+        assert serializer.decode(sc["result"], codec="json") == {"key": "val"}
 
 
 class TestWorkflowContext:

tests/test_worker.py

@@ -6,7 +6,7 @@
 
 import pytest
 
-from durable_workflow import activity, workflow
+from durable_workflow import activity, serializer, workflow
 from durable_workflow.client import (
     CONTROL_PLANE_REQUEST_CONTRACT_SCHEMA,
     CONTROL_PLANE_REQUEST_CONTRACT_VERSION,
@@ -165,6 +165,8 @@ async def test_schedule_activity_on_first_replay(self, mock_client: AsyncMock) -
         assert len(commands) == 1
         assert commands[0]["type"] == "schedule_activity"
         assert commands[0]["activity_type"] == "test-act"
+        assert commands[0]["arguments"]["codec"] == "json"
+        assert serializer.decode(commands[0]["arguments"]["blob"], codec="json") == ["hello"]
 
     @pytest.mark.asyncio
     async def test_complete_on_resolved_activity(self, mock_client: AsyncMock) -> None:
@@ -183,6 +185,8 @@ async def test_complete_on_resolved_activity(self, mock_client: AsyncMock) -> No
         mock_client.complete_workflow_task.assert_called_once()
         commands = mock_client.complete_workflow_task.call_args.kwargs["commands"]
         assert commands[0]["type"] == "complete_workflow"
+        assert commands[0]["result"]["codec"] == "json"
+        assert serializer.decode(commands[0]["result"]["blob"], codec="json") == "done"
 
     @pytest.mark.asyncio
     async def test_unknown_workflow_type_fails_task(self, mock_client: AsyncMock) -> None: