feat: add plugin interface

Author: zhongkechen

src/aws_durable_execution_sdk_python/execution.py

@@ -26,7 +26,9 @@
     Operation,
     OperationType,
     OperationUpdate,
+    InvocationStatus,
 )
+from aws_durable_execution_sdk_python.plugin import DurableExecutionPlugin, PluginExecutor
 from aws_durable_execution_sdk_python.state import ExecutionState, ReplayStatus
 
 if TYPE_CHECKING:
@@ -149,12 +151,6 @@ 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.
@@ -204,22 +200,63 @@ 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)
 
 # endregion Invocation models
 
 
+def handle_plugins(plugin_executor: PluginExecutor):
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(event: Any, context: LambdaContext):
+            with plugin_executor.start():
+                durable_execution_arn = event.get("DurableExecutionArn")
+                try:
+                    output = func(event, context)
+                    plugin_executor.on_invocation_end(
+                        durable_execution_arn=durable_execution_arn,
+                        context=context,
+                        output=output,
+                    )
+                except Exception as e:
+                    plugin_executor.on_invocation_end(
+                        durable_execution_arn=durable_execution_arn,
+                        context=context,
+                        output=DurableExecutionInvocationOutput.create_retry(ErrorObject.from_exception(e)),
+                    )
+
+        return wrapper
+    return decorator
+
+
 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 +292,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 +344,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,15 @@ 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 ExecutionDetails:
     input_payload: str | None = None