feat: Implement json serdes for Operation

- Implement json compatiable serdes for Operation and InvocationInput class

Author: Alex Wang

src/aws_durable_execution_sdk_python/execution.py

@@ -59,6 +59,16 @@ def from_dict(input_dict: MutableMapping[str, Any]) -> InitialExecutionState:
             next_marker=input_dict.get("NextMarker", ""),
         )
 
+    @staticmethod
+    def from_json_dict(input_dict: MutableMapping[str, Any]) -> InitialExecutionState:
+        operations = []
+        if input_operations := input_dict.get("Operations"):
+            operations = [Operation.from_json_dict(op) for op in input_operations]
+        return InitialExecutionState(
+            operations=operations,
+            next_marker=input_dict.get("NextMarker", ""),
+        )
+
     def get_execution_operation(self) -> Operation | None:
         if not self.operations:
             # Due to payload size limitations we may have an empty operations list.
@@ -91,6 +101,12 @@ def to_dict(self) -> MutableMapping[str, Any]:
             "NextMarker": self.next_marker,
         }
 
+    def to_json_dict(self) -> MutableMapping[str, Any]:
+        return {
+            "Operations": [op.to_json_dict() for op in self.operations],
+            "NextMarker": self.next_marker,
+        }
+
 
 @dataclass(frozen=True)
 class DurableExecutionInvocationInput:
@@ -110,13 +126,32 @@ def from_dict(
             ),
         )
 
+    @staticmethod
+    def from_json_dict(
+        input_dict: MutableMapping[str, Any],
+    ) -> DurableExecutionInvocationInput:
+        return DurableExecutionInvocationInput(
+            durable_execution_arn=input_dict["DurableExecutionArn"],
+            checkpoint_token=input_dict["CheckpointToken"],
+            initial_execution_state=InitialExecutionState.from_json_dict(
+                input_dict.get("InitialExecutionState", {})
+            ),
+        )
+
     def to_dict(self) -> MutableMapping[str, Any]:
         return {
             "DurableExecutionArn": self.durable_execution_arn,
             "CheckpointToken": self.checkpoint_token,
             "InitialExecutionState": self.initial_execution_state.to_dict(),
         }
 
+    def to_json_dict(self) -> MutableMapping[str, Any]:
+        return {
+            "DurableExecutionArn": self.durable_execution_arn,
+            "CheckpointToken": self.checkpoint_token,
+            "InitialExecutionState": self.initial_execution_state.to_json_dict(),
+        }
+
 
 @dataclass(frozen=True)
 class DurableExecutionInvocationInputWithClient(DurableExecutionInvocationInput):

src/aws_durable_execution_sdk_python/lambda_service.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
+import copy
 import datetime
 import logging
 from dataclasses import dataclass, field
+from datetime import UTC
 from enum import Enum
 from typing import TYPE_CHECKING, Any, Protocol, TypeAlias
 
@@ -805,9 +807,11 @@ def to_dict(self) -> MutableMapping[str, Any]:
                 step_dict["Error"] = self.step_details.error.to_dict()
             result["StepDetails"] = step_dict
         if self.wait_details:
-            result["WaitDetails"] = {
-                "ScheduledEndTimestamp": self.wait_details.scheduled_end_timestamp
-            }
+            result["WaitDetails"] = (
+                {"ScheduledEndTimestamp": self.wait_details.scheduled_end_timestamp}
+                if self.wait_details.scheduled_end_timestamp
+                else {}
+            )
         if self.callback_details:
             callback_dict: MutableMapping[str, Any] = {
                 "CallbackId": self.callback_details.callback_id
@@ -826,6 +830,93 @@ def to_dict(self) -> MutableMapping[str, Any]:
             result["ChainedInvokeDetails"] = invoke_dict
         return result
 
+    @classmethod
+    def _dt_to_ms(cls, dt: datetime.datetime | None) -> int | None:
+        """Convert datetime to millisecond timestamp."""
+        if not dt:
+            return None
+        return int(dt.timestamp() * 1000)
+
+    @classmethod
+    def _ms_to_dt(cls, ms: int | None) -> datetime.datetime | None:
+        """Convert millisecond timestamp to datetime."""
+        if not ms:
+            return None
+        return datetime.datetime.fromtimestamp(ms / 1000, tz=UTC)
+
+    def to_json_dict(self) -> MutableMapping[str, Any]:
+        """Convert the Operation to a JSON-serializable dictionary.
+
+        Converts datetime objects to millisecond timestamps for JSON compatibility.
+
+        Returns:
+            A dictionary with JSON-serializable values
+        """
+        # Start with the regular to_dict output
+        result = self.to_dict()
+
+        # Convert datetime objects to millisecond timestamps
+        if "StartTimestamp" in result and isinstance(
+            result["StartTimestamp"], datetime.datetime
+        ):
+            result["StartTimestamp"] = self._dt_to_ms(result["StartTimestamp"])
+
+        if "EndTimestamp" in result and isinstance(
+            result["EndTimestamp"], datetime.datetime
+        ):
+            result["EndTimestamp"] = self._dt_to_ms(result["EndTimestamp"])
+
+        if self.step_details and self.step_details.next_attempt_timestamp:
+            result["StepDetails"]["NextAttemptTimestamp"] = self._dt_to_ms(
+                result["StepDetails"].get("NextAttemptTimestamp")
+            )
+
+        if self.wait_details and self.wait_details.scheduled_end_timestamp:
+            result["WaitDetails"]["ScheduledEndTimestamp"] = self._dt_to_ms(
+                result["WaitDetails"]["ScheduledEndTimestamp"]
+            )
+
+        return result
+
+    @classmethod
+    def from_json_dict(cls, data: MutableMapping[str, Any]) -> Operation:
+        """Create an Operation from a JSON-serializable dictionary.
+
+        Converts millisecond timestamps back to datetime objects.
+
+        Args:
+            data: Dictionary with JSON-serializable values (millisecond timestamps)
+
+        Returns:
+            An Operation instance with datetime objects
+        """
+        # Make a copy to avoid modifying the original data
+        data_copy = copy.deepcopy(data)
+
+        # Convert millisecond timestamps back to datetime objects
+        if "StartTimestamp" in data_copy and isinstance(
+            data_copy["StartTimestamp"], int
+        ):
+            data_copy["StartTimestamp"] = cls._ms_to_dt(data_copy["StartTimestamp"])
+
+        if "EndTimestamp" in data_copy and isinstance(data_copy["EndTimestamp"], int):
+            data_copy["EndTimestamp"] = cls._ms_to_dt(data_copy["EndTimestamp"])
+
+        step_details = data_copy.get("StepDetails")
+        if step_details and isinstance(step_details.get("NextAttemptTimestamp"), int):
+            step_details["NextAttemptTimestamp"] = cls._ms_to_dt(
+                step_details["NextAttemptTimestamp"]
+            )
+
+        wait_details = data_copy.get("WaitDetails")
+        if wait_details and isinstance(wait_details.get("ScheduledEndTimestamp"), int):
+            wait_details["ScheduledEndTimestamp"] = cls._ms_to_dt(
+                wait_details["ScheduledEndTimestamp"]
+            )
+
+        # Use the existing from_dict method with the converted data
+        return cls.from_dict(data_copy)
+
 
 @dataclass(frozen=True)
 class CheckpointUpdatedExecutionState: