feat: add plugin interface

Author: zhongkechen

src/aws_durable_execution_sdk_python/execution.py

@@ -26,6 +26,13 @@
     Operation,
     OperationType,
     OperationUpdate,
+    InvocationStatus,
+    DurableExecutionInvocationOutput,
+)
+from aws_durable_execution_sdk_python.plugin import (
+    DurableExecutionPlugin,
+    PluginExecutor,
+    handle_plugins,
 )
 from aws_durable_execution_sdk_python.state import ExecutionState, ReplayStatus
 
@@ -149,77 +156,36 @@ def from_durable_execution_invocation_input(
         )
 
 
-class InvocationStatus(Enum):
-    SUCCEEDED = "SUCCEEDED"
-    FAILED = "FAILED"
-    PENDING = "PENDING"
-
-
-@dataclass(frozen=True)
-class DurableExecutionInvocationOutput:
-    """Representation the DurableExecutionInvocationOutput. This is what the Durable lambda handler returns.
-
-    If the execution has been already completed via an update to the EXECUTION operation via CheckpointDurableExecution,
-    payload must be empty for SUCCEEDED/FAILED status.
-    """
-
-    status: InvocationStatus
-    result: str | None = None
-    error: ErrorObject | None = None
-
-    @classmethod
-    def from_dict(
-        cls, data: MutableMapping[str, Any]
-    ) -> DurableExecutionInvocationOutput:
-        """Create an instance from a dictionary.
-
-        Args:
-            data: Dictionary with camelCase keys matching the original structure
-
-        Returns:
-            A DurableExecutionInvocationOutput instance
-        """
-        status = InvocationStatus(data.get("Status"))
-        error = ErrorObject.from_dict(data["Error"]) if data.get("Error") else None
-        return cls(status=status, result=data.get("Result"), error=error)
-
-    def to_dict(self) -> MutableMapping[str, Any]:
-        """Convert to a dictionary with the original field names.
-
-        Returns:
-            Dictionary with the original camelCase keys
-        """
-        result: MutableMapping[str, Any] = {"Status": self.status.value}
-
-        if self.result is not None:
-            # large payloads return "", because checkpointed already
-            result["Result"] = self.result
-        if self.error:
-            result["Error"] = self.error.to_dict()
-
-        return result
-
-    @classmethod
-    def create_succeeded(cls, result: str) -> DurableExecutionInvocationOutput:
-        """Create a succeeded invocation output."""
-        return cls(status=InvocationStatus.SUCCEEDED, result=result)
-
-
 # endregion Invocation models
 
 
 def durable_execution(
     func: Callable[[Any, DurableContext], Any] | None = None,
     *,
     boto3_client: Boto3LambdaClient | None = None,
+    plugins: list[DurableExecutionPlugin] | None = None,
 ) -> Callable[[Any, LambdaContext], Any]:
+    """
+    Decorator to create a durable execution handler.
+
+    Args:
+        func: The user function to decorate
+        boto3_client: Optional boto3 Lambda client to use
+        plugins: Optional list of plugins to use (EXPERIMENTAL: This
+            parameter may change or be removed.)
+    """
     # 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, plugins=plugins
+        )
 
     logger.debug("Starting durable execution handler...")
 
+    plugin_executor = PluginExecutor(plugins)
+
+    @handle_plugins(plugin_executor)
     def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
         invocation_input: DurableExecutionInvocationInput
         service_client: DurableServiceClient
@@ -255,6 +221,7 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
             operations={},
             service_client=service_client,
             replay_status=ReplayStatus.NEW,
+            plugin_executor=plugin_executor,
         )
 
         try:
@@ -306,6 +273,13 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
             ) as executor,
             contextlib.closing(execution_state) as execution_state,
         ):
+            # execute the plugins
+            plugin_executor.on_invocation_start(
+                durable_execution_arn=invocation_input.durable_execution_arn,
+                context=context,
+                execution_operation=execution_state.get_execution_operation(),
+                is_replaying=execution_state.is_replaying(),
+            )
             # Thread 1: Run background checkpoint processing
             executor.submit(execution_state.checkpoint_batches_forever)
 

src/aws_durable_execution_sdk_python/lambda_service.py

@@ -105,6 +105,70 @@ class OperationSubType(Enum):
     CHAINED_INVOKE = "ChainedInvoke"
 
 
+class InvocationStatus(Enum):
+    SUCCEEDED = "SUCCEEDED"
+    FAILED = "FAILED"
+    PENDING = "PENDING"
+
+    # Used internally only: the invocation failed and the backend will retry
+    RETRY = "RETRY"
+
+
+@dataclass(frozen=True)
+class DurableExecutionInvocationOutput:
+    """Representation the DurableExecutionInvocationOutput. This is what the Durable lambda handler returns.
+
+    If the execution has been already completed via an update to the EXECUTION operation via CheckpointDurableExecution,
+    payload must be empty for SUCCEEDED/FAILED status.
+    """
+
+    status: InvocationStatus
+    result: str | None = None
+    error: ErrorObject | None = None
+
+    @classmethod
+    def from_dict(
+        cls, data: MutableMapping[str, Any]
+    ) -> DurableExecutionInvocationOutput:
+        """Create an instance from a dictionary.
+
+        Args:
+            data: Dictionary with camelCase keys matching the original structure
+
+        Returns:
+            A DurableExecutionInvocationOutput instance
+        """
+        status = InvocationStatus(data.get("Status"))
+        error = ErrorObject.from_dict(data["Error"]) if data.get("Error") else None
+        return cls(status=status, result=data.get("Result"), error=error)
+
+    def to_dict(self) -> MutableMapping[str, Any]:
+        """Convert to a dictionary with the original field names.
+
+        Returns:
+            Dictionary with the original camelCase keys
+        """
+        result: MutableMapping[str, Any] = {"Status": self.status.value}
+
+        if self.result is not None:
+            # large payloads return "", because checkpointed already
+            result["Result"] = self.result
+        if self.error:
+            result["Error"] = self.error.to_dict()
+
+        return result
+
+    @classmethod
+    def create_succeeded(cls, result: str) -> DurableExecutionInvocationOutput:
+        """Create a succeeded invocation output."""
+        return cls(status=InvocationStatus.SUCCEEDED, result=result)
+
+    @classmethod
+    def create_retry(cls, error: ErrorObject) -> DurableExecutionInvocationOutput:
+        """Create a failed invocation output."""
+        return cls(status=InvocationStatus.RETRY, error=error)
+
+
 @dataclass(frozen=True)
 class ExecutionDetails:
     input_payload: str | None = None