Allow passing input/output serdes

Changes:
 - We allow input and output serdes on the decorator. This is
   particularly useful for chained invoke to ensure callees
   automatically have the correct serialization

Author: FullyTyped

src/aws_durable_execution_sdk_python/__init__.py

@@ -18,14 +18,24 @@
 # Core decorator - used in every durable function
 from aws_durable_execution_sdk_python.execution import durable_execution
 
+# Serialization - for custom input/output serialization
+from aws_durable_execution_sdk_python.serdes import (
+    ExtendedTypeSerDes,
+    JsonSerDes,
+    SerDes,
+)
+
 # Essential context types - passed to user functions
 from aws_durable_execution_sdk_python.types import BatchResult, StepContext
 
 __all__ = [
     "BatchResult",
     "DurableContext",
     "DurableExecutionsError",
+    "ExtendedTypeSerDes",
     "InvocationError",
+    "JsonSerDes",
+    "SerDes",
     "StepContext",
     "ValidationError",
     "durable_execution",

src/aws_durable_execution_sdk_python/execution.py

@@ -27,6 +27,12 @@
     OperationType,
     OperationUpdate,
 )
+from aws_durable_execution_sdk_python.serdes import (
+    JsonSerDes,
+    SerDes,
+    deserialize,
+    serialize,
+)
 
 if TYPE_CHECKING:
     from collections.abc import Callable, MutableMapping
@@ -206,11 +212,18 @@ def durable_execution(
     func: Callable[[Any, DurableContext], Any] | None = None,
     *,
     boto3_client: boto3.client | None = None,
+    input_serdes: SerDes = JsonSerDes(),
+    output_serdes: SerDes = JsonSerDes(),
 ) -> Callable[[Any, LambdaContext], Any]:
     # Decorator called with parameters
     if func is None:
         logger.debug("Decorator called with parameters")
-        return functools.partial(durable_execution, boto3_client=boto3_client)
+        return functools.partial(
+            durable_execution,
+            boto3_client=boto3_client,
+            input_serdes=input_serdes,
+            output_serdes=output_serdes,
+        )
 
     logger.debug("Starting durable execution handler...")
 
@@ -250,18 +263,14 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
             invocation_input.initial_execution_state.get_input_payload()
         )
 
-        # Python RIC LambdaMarshaller just uses standard json deserialization for event
-        # https://github.com/aws/aws-lambda-python-runtime-interface-client/blob/main/awslambdaric/lambda_runtime_marshaller.py#L46
-        input_event: MutableMapping[str, Any] = {}
+        input_event: MutableMapping[str, Any] = None
         if raw_input_payload and raw_input_payload.strip():
-            try:
-                input_event = json.loads(raw_input_payload)
-            except json.JSONDecodeError:
-                logger.exception(
-                    "Failed to parse input payload as JSON: payload: %r",
-                    raw_input_payload,
-                )
-                raise
+            input_event = deserialize(
+                serdes=input_serdes,
+                data=raw_input_payload,
+                operation_id="EXECUTION",
+                durable_execution_arn=invocation_input.durable_execution_arn,
+            )
 
         execution_state: ExecutionState = ExecutionState(
             durable_execution_arn=invocation_input.durable_execution_arn,
@@ -310,7 +319,15 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
                     "%s exiting user-space...",
                     invocation_input.durable_execution_arn,
                 )
-                serialized_result = json.dumps(result)
+
+                # Serialize result using output_serdes if provided
+                serialized_result = serialize(
+                    serdes=output_serdes,
+                    value=result,
+                    operation_id="EXECUTION",
+                    durable_execution_arn=invocation_input.durable_execution_arn,
+                )
+
                 # large response handling here. Remember if checkpointing to complete, NOT to include
                 # payload in response
                 if (

tests/execution_test.py

@@ -3,9 +3,32 @@
 import datetime
 import json
 import time
+from decimal import Decimal
 from typing import Any
 from unittest.mock import Mock, patch
 
+
+from aws_durable_execution_sdk_python import (
+    DurableContext,
+    ExtendedTypeSerDes,
+    JsonSerDes,
+    durable_execution,
+)
+from aws_durable_execution_sdk_python.execution import (
+    DurableExecutionInvocationInputWithClient,
+    InitialExecutionState,
+)
+from aws_durable_execution_sdk_python.lambda_service import (
+    CheckpointOutput,
+    CheckpointUpdatedExecutionState,
+    DurableServiceClient,
+    ExecutionDetails,
+    Operation,
+    OperationStatus,
+    OperationType,
+)
+
+
 import pytest
 
 from aws_durable_execution_sdk_python.config import StepConfig, StepSemantics
@@ -946,7 +969,7 @@ def test_handler(event: Any, context: DurableContext) -> dict:
     lambda_context.invoked_function_arn = None
     lambda_context.tenant_id = None
 
-    with pytest.raises(json.JSONDecodeError):
+    with pytest.raises(ExecutionError):
         test_handler(invocation_input, lambda_context)
 
 
@@ -2047,7 +2070,6 @@ def test_handler(event: Any, context: DurableContext) -> dict:
 
 def test_durable_execution_with_non_dict_event_raises_error():
     """Test that invoking a durable function with a non-dict event raises a helpful error."""
-
     # GIVEN a durable function
     @durable_execution
     def test_handler(event: Any, context: DurableContext) -> dict:
@@ -2071,3 +2093,125 @@ def test_handler(event: Any, context: DurableContext) -> dict:
         match="The payload is not the correct Durable Function input",
     ):
         test_handler(non_dict_event, lambda_context)
+
+
+# region SERDES
+
+@pytest.mark.parametrize(
+    "input_serdes,output_serdes,input_data,expected_output",
+    [
+        # Both ExtendedTypeSerDes
+        (
+            ExtendedTypeSerDes(),
+            ExtendedTypeSerDes(),
+            {
+                "amount": Decimal("123.45"),
+                "timestamp": datetime.datetime(2025, 11, 20, 18, 0, 0),
+            },
+            {
+                "result_amount": Decimal("678.90"),
+                "processed_at": datetime.datetime(2025, 11, 20, 19, 0, 0),
+            },
+        ),
+        # Both JsonSerDes
+        (
+            JsonSerDes(),
+            JsonSerDes(),
+            {"name": "test", "value": 42},
+            {"result": "success", "count": 100},
+        ),
+        # Input ExtendedTypeSerDes, Output JsonSerDes
+        (
+            ExtendedTypeSerDes(),
+            JsonSerDes(),
+            {"amount": Decimal("123.45")},
+            {"result": "success"},
+        ),
+        # Input JsonSerDes, Output ExtendedTypeSerDes
+        (
+            JsonSerDes(),
+            ExtendedTypeSerDes(),
+            {"name": "test"},
+            {"timestamp": datetime.datetime(2025, 11, 20, 19, 0, 0)},
+        ),
+    ],
+)
+def test_durable_execution_with_serdes(input_serdes, output_serdes, input_data, expected_output):
+    """Test that input_serdes and output_serdes are invoked correctly."""
+    serialized_input = input_serdes.serialize(input_data, None)
+
+    mock_client = Mock(spec=DurableServiceClient)
+    mock_client.checkpoint.return_value = CheckpointOutput(
+        checkpoint_token="new-token",
+        new_execution_state=CheckpointUpdatedExecutionState(operations=[], next_marker=None),
+    )
+
+    execution_op = Operation(
+        operation_id="EXECUTION",
+        operation_type=OperationType.EXECUTION,
+        status=OperationStatus.STARTED,
+        execution_details=ExecutionDetails(input_payload=serialized_input),
+    )
+
+    invocation_input = DurableExecutionInvocationInputWithClient(
+        durable_execution_arn="arn:aws:lambda:us-east-1:123456789012:function:test",
+        checkpoint_token="initial-token",
+        initial_execution_state=InitialExecutionState(
+            operations=[execution_op],
+            next_marker="",
+        ),
+        is_local_runner=False,
+        service_client=mock_client,
+    )
+
+    with patch.object(input_serdes, "deserialize", wraps=input_serdes.deserialize) as mock_input_deser, patch.object(
+        output_serdes, "serialize", wraps=output_serdes.serialize
+    ) as mock_output_ser:
+
+        @durable_execution(input_serdes=input_serdes, output_serdes=output_serdes)
+        def handler(event, context: DurableContext):
+            return expected_output
+
+        result = handler(invocation_input, Mock())
+
+        mock_input_deser.assert_called_once()
+        mock_output_ser.assert_called_once_with(expected_output, mock_output_ser.call_args[0][1])
+
+    assert result["Status"] == "SUCCEEDED"
+    result_data = output_serdes.deserialize(result["Result"], None)
+    assert result_data == expected_output
+
+
+def test_durable_execution_with_none_input():
+    """Test that None input is handled correctly with serdes."""
+    mock_client = Mock(spec=DurableServiceClient)
+    mock_client.checkpoint.return_value = CheckpointOutput(
+        checkpoint_token="new-token",
+        new_execution_state=CheckpointUpdatedExecutionState(operations=[], next_marker=None),
+    )
+
+    execution_op = Operation(
+        operation_id="EXECUTION",
+        operation_type=OperationType.EXECUTION,
+        status=OperationStatus.STARTED,
+        execution_details=ExecutionDetails(input_payload=None),
+    )
+
+    invocation_input = DurableExecutionInvocationInputWithClient(
+        durable_execution_arn="arn:aws:lambda:us-east-1:123456789012:function:test",
+        checkpoint_token="initial-token",
+        initial_execution_state=InitialExecutionState(
+            operations=[execution_op],
+            next_marker="",
+        ),
+        is_local_runner=False,
+        service_client=mock_client,
+    )
+
+    @durable_execution(input_serdes=ExtendedTypeSerDes())
+    def handler(event, context: DurableContext):
+        assert event is None
+        return {"result": "success"}
+
+    result = handler(invocation_input, Mock())
+    assert result["Status"] == "SUCCEEDED"