TD-P001: Python workers do not prove unsupported payload codecs fail closed (#27)

Author: durable-workflow-ops

src/durable_workflow/worker.py

@@ -70,6 +70,16 @@ def _command_payload_codec(codec: object) -> str:
     return codec if isinstance(codec, str) and codec in serializer.SUPPORTED_CODECS else serializer.AVRO_CODEC
 
 
+def _validate_payload_codec(codec: object) -> str | None:
+    if codec is None:
+        return None
+    if isinstance(codec, str) and codec in serializer.SUPPORTED_CODECS:
+        return codec
+    raise ValueError(
+        f"Unsupported payload codec {codec!r}; this SDK supports {serializer.SUPPORTED_CODECS!r}."
+    )
+
+
 def _activity_name(fn: Callable[..., Any]) -> str:
     return getattr(fn, "__activity_name__", fn.__name__)
 
@@ -414,9 +424,9 @@ async def _run_workflow_task_core(self, task: dict[str, Any]) -> list[dict[str,
 
         start_input: list[Any] = []
         codec = task.get("payload_codec")
-        command_codec = _command_payload_codec(codec)
         raw_args = task.get("arguments")
         try:
+            codec = _validate_payload_codec(codec)
             decoded = serializer.decode_envelope(raw_args, codec=codec)
             if decoded is not None:
                 start_input = decoded if isinstance(decoded, list) else [decoded]
@@ -456,6 +466,7 @@ async def _run_workflow_task_core(self, task: dict[str, Any]) -> list[dict[str,
             return None
 
         run_id: str = task.get("run_id", "")
+        command_codec = _command_payload_codec(codec)
 
         cls = self.workflows.get(wf_type)
         if cls is None:
@@ -583,9 +594,9 @@ async def _run_activity_task(self, task: dict[str, Any]) -> str:
         activity_type: str = task.get("activity_type", "")
         attempt_number: int = task.get("attempt_number", 1)
         raw_args = task.get("arguments")
-        inbound_codec = task.get("payload_codec") or serializer.JSON_CODEC
-        result_codec = inbound_codec if inbound_codec in serializer.SUPPORTED_CODECS else serializer.AVRO_CODEC
+        inbound_codec = task.get("payload_codec")
         try:
+            inbound_codec = _validate_payload_codec(inbound_codec) or serializer.JSON_CODEC
             args = serializer.decode_envelope(raw_args, codec=inbound_codec) or []
         except AvroNotInstalledError as e:
             log.exception("activity %s arguments Avro decode failed (avro dependency unavailable)", task_id)
@@ -625,6 +636,7 @@ async def _run_activity_task(self, task: dict[str, Any]) -> str:
             return "decode_error"
         if not isinstance(args, list):
             args = [args]
+        result_codec = inbound_codec
 
         fn = self.activities.get(activity_type)
         if fn is None:
@@ -783,6 +795,19 @@ async def _run_query_task_core(self, task: dict[str, Any]) -> str:
         wf_type: str = task.get("workflow_type", "")
         query_name: str = task.get("query_name", "")
         codec = task.get("payload_codec")
+        try:
+            codec = _validate_payload_codec(codec)
+        except ValueError as e:
+            await self._fail_query_task(
+                query_task_id,
+                attempt,
+                f"cannot decode query payload with codec {codec!r}: {e}.",
+                reason="query_payload_decode_failed",
+                failure_type=type(e).__name__,
+                stack_trace=traceback.format_exc(),
+            )
+            return "failed"
+
         result_codec = _command_payload_codec(codec)
         history = task.get("history_events", [])
 

tests/test_worker.py

@@ -878,9 +878,7 @@ async def test_workflow_with_envelope_arguments(self, mock_client: AsyncMock) ->
 
 
 class TestCodecDecodeFailures:
-    """TD-P012 / #370 regression: codec decode failures at the task boundary
-    must turn into a deterministic fail_{workflow,activity}_task call so the
-    lease does not sit until timeout."""
+    """Codec decode failures at the task boundary must fail tasks deterministically."""
 
     @pytest.mark.asyncio
     async def test_activity_json_decode_failure_fails_task(self, mock_client: AsyncMock) -> None:
@@ -919,6 +917,29 @@ async def test_activity_avro_decode_failure_fails_task(self, mock_client: AsyncM
         assert call_kwargs["non_retryable"] is True
         mock_client.complete_activity_task.assert_not_called()
 
+    @pytest.mark.asyncio
+    async def test_activity_unsupported_payload_codec_fails_before_handler(
+        self, mock_client: AsyncMock
+    ) -> None:
+        worker = Worker(mock_client, task_queue="q1", workflows=[], activities=[echo_activity])
+        task = {
+            "task_id": "at-unsupported-codec",
+            "activity_attempt_id": "aa-unsupported-codec",
+            "activity_type": "test-act",
+            "arguments": {"codec": "json", "blob": '["hello"]'},
+            "payload_codec": "zstd",
+        }
+
+        outcome = await worker._run_activity_task(task)
+
+        assert outcome == "decode_error"
+        mock_client.fail_activity_task.assert_called_once()
+        call_kwargs = mock_client.fail_activity_task.call_args.kwargs
+        assert "Unsupported payload codec 'zstd'" in call_kwargs["message"]
+        assert call_kwargs["failure_type"] == "ValueError"
+        assert call_kwargs["non_retryable"] is True
+        mock_client.complete_activity_task.assert_not_called()
+
     @pytest.mark.asyncio
     async def test_activity_avro_missing_dependency_fails_task(
         self, mock_client: AsyncMock, monkeypatch: pytest.MonkeyPatch
@@ -968,6 +989,55 @@ async def test_workflow_json_decode_failure_fails_task(self, mock_client: AsyncM
         assert "json" in call_kwargs["message"]
         mock_client.complete_workflow_task.assert_not_called()
 
+    @pytest.mark.asyncio
+    async def test_workflow_unsupported_payload_codec_fails_before_replay(
+        self, mock_client: AsyncMock
+    ) -> None:
+        worker = Worker(mock_client, task_queue="q1", workflows=[TestWorkflow], activities=[])
+        task = {
+            "task_id": "t-unsupported-codec",
+            "workflow_type": "test-wf",
+            "workflow_task_attempt": 1,
+            "history_events": [],
+            "arguments": {"codec": "json", "blob": '["hello"]'},
+            "payload_codec": "zstd",
+        }
+
+        commands = await worker._run_workflow_task(task)
+
+        assert commands is None
+        mock_client.fail_workflow_task.assert_called_once()
+        call_kwargs = mock_client.fail_workflow_task.call_args.kwargs
+        assert "Unsupported payload codec 'zstd'" in call_kwargs["message"]
+        assert call_kwargs["failure_type"] == "ValueError"
+        mock_client.complete_workflow_task.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_query_unsupported_payload_codec_fails_before_query_handler(
+        self, mock_client: AsyncMock
+    ) -> None:
+        worker = Worker(mock_client, task_queue="q1", workflows=[QueryWorkflow], activities=[])
+        task = {
+            "query_task_id": "qt-unsupported-codec",
+            "query_task_attempt": 1,
+            "workflow_type": "query-wf",
+            "query_name": "status",
+            "history_events": [],
+            "workflow_arguments": serializer.envelope([], codec="json"),
+            "query_arguments": serializer.envelope([], codec="json"),
+            "payload_codec": "zstd",
+        }
+
+        outcome = await worker._run_query_task(task)
+
+        assert outcome == "failed"
+        mock_client.fail_query_task.assert_called_once()
+        call_kwargs = mock_client.fail_query_task.call_args.kwargs
+        assert call_kwargs["reason"] == "query_payload_decode_failed"
+        assert "Unsupported payload codec 'zstd'" in call_kwargs["message"]
+        assert call_kwargs["failure_type"] == "ValueError"
+        mock_client.complete_query_task.assert_not_called()
+
     @pytest.mark.asyncio
     async def test_workflow_avro_missing_dependency_fails_task(
         self, mock_client: AsyncMock, monkeypatch: pytest.MonkeyPatch