aws-durable-execution-sdk-python 2.0.0__py3-none-any.whl → 2.1.0__py3-none-any.whl

aws_durable_execution_sdk_python/execution.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
 import contextlib
+import contextvars
 import functools
 import json
 import logging
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import Future, ThreadPoolExecutor, wait
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any
 
@@ -188,9 +189,9 @@ def durable_execution(
 
     logger.debug("Starting durable execution handler...")
 
-    plugin_executor = PluginExecutor(load_configured_plugins(plugins))
+    configured_plugins = load_configured_plugins(plugins)
+    plugin_executor = PluginExecutor(configured_plugins)
 
-    @plugin_executor.handle_durable_output
     def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
         invocation_input: DurableExecutionInvocationInput
         service_client: DurableServiceClient
@@ -279,88 +280,75 @@ def durable_execution(
             ),
         )
 
-        # Use ThreadPoolExecutor for concurrent execution of user code and background checkpoint processing
-        with (
-            ThreadPoolExecutor(
-                max_workers=2, thread_name_prefix="dex-handler"
-            ) as executor,
-            contextlib.closing(execution_state) as execution_state,
-        ):
-            execution_operation = execution_state.get_execution_operation()
-
-            # execute the plugins
-            plugin_executor.on_invocation_start(
-                execution_arn=invocation_input.durable_execution_arn,
-                lambda_context=context,
-                execution_start_time=(
-                    execution_operation.start_timestamp
-                    if execution_operation is not None
-                    else None
-                ),
-                is_first_invocation=not has_prior_operations,
-                execution_input=input_event,
-                # Read the map through a callable rather than snapshotting it
-                # here: the invocation-end hook needs the state as of the end of
-                # the invocation, and neither hook pays for the conversion until
-                # a plugin actually reads it.
-                operations_provider=lambda: execution_state.operations,
-                updated_operation_ids=invocation_input.updated_operation_ids,
-            )
-            # Thread 1: Run background checkpoint processing
-            executor.submit(execution_state.checkpoint_batches_forever)
+        def invoke_handler() -> MutableMapping[str, Any]:
+            output: MutableMapping[str, Any]
 
-            # Thread 2: Execute user function
             logger.debug(
                 "%s entering user-space...", invocation_input.durable_execution_arn
             )
-            user_future = executor.submit(func, input_event, durable_context)
-
-            logger.debug(
-                "%s waiting for user code completion...",
-                invocation_input.durable_execution_arn,
-            )
-
             try:
                 # Background checkpointing errors will propagate through CompletionEvent.wait() as BackgroundThreadError
-                result = user_future.result()
+                result = func(input_event, durable_context)
 
                 # done with userland
                 logger.debug(
                     "%s exiting user-space...",
                     invocation_input.durable_execution_arn,
                 )
-                serialized_result = json.dumps(result)
-                # large response handling here. Remember if checkpointing to complete, NOT to include
-                # payload in response
-                if (
-                    serialized_result
-                    and len(serialized_result) > LAMBDA_RESPONSE_SIZE_LIMIT
-                ):
+
+                # The handler may finish while background checkpoint work is
+                # still running for async operations it did not await, such as
+                # early-completing map or parallel branches. If one of those
+                # checkpoint responses omitted the next checkpoint token, this
+                # invocation can no longer record anything else with the
+                # service, including the handler's final result.
+                #
+                # In that case, return PENDING instead of SUCCEEDED. The next
+                # invocation will replay from the last durable checkpoint and
+                # produce the result again. This check catches revocation
+                # already observed before we build the terminal response; the
+                # post-with check below catches revocation discovered during
+                # cleanup of in-flight checkpoint work.
+                if execution_state.is_checkpoint_token_revoked:
                     logger.debug(
-                        "Response size (%s bytes) exceeds Lambda limit (%s) bytes). Checkpointing result.",
-                        len(serialized_result),
-                        LAMBDA_RESPONSE_SIZE_LIMIT,
-                    )
-                    success_operation = OperationUpdate.create_execution_succeed(
-                        payload=serialized_result
+                        "Checkpoint token revoked; ending invocation with "
+                        "PENDING instead of SUCCEEDED."
                     )
-                    # Checkpoint large result with blocking (is_sync=True, default).
-                    # Must ensure the result is persisted before returning to Lambda.
-                    # Large results exceed Lambda response limits and must be stored durably
-                    # before the execution completes.
-                    try:
-                        execution_state.create_checkpoint(
-                            success_operation, is_sync=True
+                    output = DurableExecutionInvocationOutput.create_pending().to_dict()
+                else:
+                    serialized_result = json.dumps(result)
+                    # large response handling here. Remember if checkpointing to complete, NOT to include
+                    # payload in response
+                    if (
+                        serialized_result
+                        and len(serialized_result) > LAMBDA_RESPONSE_SIZE_LIMIT
+                    ):
+                        logger.debug(
+                            "Response size (%s bytes) exceeds Lambda limit (%s) bytes). Checkpointing result.",
+                            len(serialized_result),
+                            LAMBDA_RESPONSE_SIZE_LIMIT,
                         )
-                    except CheckpointError as e:
-                        return handle_checkpoint_error(e).to_dict()
-                    return DurableExecutionInvocationOutput.create_succeeded(
-                        result=""
-                    ).to_dict()
-
-                return DurableExecutionInvocationOutput.create_succeeded(
-                    result=serialized_result
-                ).to_dict()
+                        success_operation = OperationUpdate.create_execution_succeed(
+                            payload=serialized_result
+                        )
+                        # Checkpoint large result with blocking (is_sync=True, default).
+                        # Must ensure the result is persisted before returning to Lambda.
+                        # Large results exceed Lambda response limits and must be stored durably
+                        # before the execution completes.
+                        try:
+                            execution_state.create_checkpoint(
+                                success_operation, is_sync=True
+                            )
+                        except CheckpointError as e:
+                            output = handle_checkpoint_error(e).to_dict()
+                        else:
+                            output = DurableExecutionInvocationOutput.create_succeeded(
+                                result=""
+                            ).to_dict()
+                    else:
+                        output = DurableExecutionInvocationOutput.create_succeeded(
+                            result=serialized_result
+                        ).to_dict()
 
             except BackgroundThreadError as bg_error:
                 # Background checkpoint system failed - propagated through CompletionEvent
@@ -378,85 +366,202 @@ def durable_execution(
                             "without retry.",
                             extra=bg_error.source_exception.build_logger_extras(),
                         )
-                        return DurableExecutionInvocationOutput(
+                        output = DurableExecutionInvocationOutput(
                             status=InvocationStatus.FAILED,
                             error=ErrorObject.from_exception(bg_error.source_exception),
                         ).to_dict()
+                    else:
+                        raise bg_error.source_exception from bg_error
                 else:
                     logger.exception("Checkpoint processing failed")
-                raise bg_error.source_exception from bg_error
+                    raise bg_error.source_exception from bg_error
 
             except SuspendExecution:
                 # User code suspended - stop background checkpointing thread
                 logger.debug("Suspending execution...")
-                return DurableExecutionInvocationOutput(
-                    status=InvocationStatus.PENDING
-                ).to_dict()
+                output = DurableExecutionInvocationOutput.create_pending().to_dict()
 
             except CheckpointError as e:
-                # Checkpoint system is broken - stop background thread and exit immediately
-                logger.exception(
-                    "Checkpoint system failed",
-                    extra=e.build_logger_extras(),
-                )
-                return handle_checkpoint_error(e).to_dict()
+                # Checkpoint system is broken - stop background thread and exit immediately.
+                if execution_state.is_checkpoint_token_revoked:
+                    output = DurableExecutionInvocationOutput.create_pending().to_dict()
+                else:
+                    logger.exception(
+                        "Checkpoint system failed",
+                        extra=e.build_logger_extras(),
+                    )
+                    output = handle_checkpoint_error(e).to_dict()
             except InvocationError as e:
-                # Non-retryable Durable API errors (e.g., customer configuration issues,
-                # 4xx client errors) will never succeed on retry — fail the execution immediately.
-                if not e.is_retryable():
+                if execution_state.is_checkpoint_token_revoked:
+                    output = DurableExecutionInvocationOutput.create_pending().to_dict()
+                else:
+                    # Non-retryable Durable API errors (e.g., customer configuration issues,
+                    # 4xx client errors) will never succeed on retry — fail the execution immediately.
+                    if not e.is_retryable():
+                        logger.exception(
+                            "Non-retryable Durable API error. Must fail execution without retry.",
+                            extra=e.build_logger_extras(),  # type: ignore[attr-defined]
+                        )
+                        output = DurableExecutionInvocationOutput(
+                            status=InvocationStatus.FAILED,
+                            error=ErrorObject.from_exception(e),
+                        ).to_dict()
+                    else:
+                        logger.exception("Invocation error. Must terminate.")
+                        # Throw the error to trigger Lambda retry
+                        raise
+            except ExecutionError as e:
+                if execution_state.is_checkpoint_token_revoked:
+                    output = DurableExecutionInvocationOutput.create_pending().to_dict()
+                else:
                     logger.exception(
-                        "Non-retryable Durable API error. Must fail execution without retry.",
-                        extra=e.build_logger_extras(),  # type: ignore[attr-defined]
+                        "Execution error. Must fail execution without retry."
                     )
-                    return DurableExecutionInvocationOutput(
+                    output = DurableExecutionInvocationOutput(
                         status=InvocationStatus.FAILED,
                         error=ErrorObject.from_exception(e),
                     ).to_dict()
-                logger.exception("Invocation error. Must terminate.")
-                # Throw the error to trigger Lambda retry
-                raise
-            except ExecutionError as e:
-                logger.exception("Execution error. Must fail execution without retry.")
-                return DurableExecutionInvocationOutput(
-                    status=InvocationStatus.FAILED,
-                    error=ErrorObject.from_exception(e),
-                ).to_dict()
             except Exception as e:
                 # all user-space errors go here
-                logger.exception("Execution failed")
+                if execution_state.is_checkpoint_token_revoked:
+                    output = DurableExecutionInvocationOutput.create_pending().to_dict()
+                else:
+                    logger.exception("Execution failed")
 
-                result = DurableExecutionInvocationOutput(
-                    status=InvocationStatus.FAILED, error=ErrorObject.from_exception(e)
-                ).to_dict()
+                    result = DurableExecutionInvocationOutput(
+                        status=InvocationStatus.FAILED,
+                        error=ErrorObject.from_exception(e),
+                    ).to_dict()
 
-                serialized_result = json.dumps(result)
+                    serialized_result = json.dumps(result)
 
-                if (
-                    serialized_result
-                    and len(serialized_result) > LAMBDA_RESPONSE_SIZE_LIMIT
-                ):
-                    logger.debug(
-                        "Response size (%s bytes) exceeds Lambda limit (%s) bytes). Checkpointing result.",
-                        len(serialized_result),
-                        LAMBDA_RESPONSE_SIZE_LIMIT,
+                    if (
+                        serialized_result
+                        and len(serialized_result) > LAMBDA_RESPONSE_SIZE_LIMIT
+                    ):
+                        logger.debug(
+                            "Response size (%s bytes) exceeds Lambda limit (%s) bytes). Checkpointing result.",
+                            len(serialized_result),
+                            LAMBDA_RESPONSE_SIZE_LIMIT,
+                        )
+                        failed_operation = OperationUpdate.create_execution_fail(
+                            error=ErrorObject.from_exception(e)
+                        )
+
+                        # Checkpoint large result with blocking (is_sync=True, default).
+                        # Must ensure the result is persisted before returning to Lambda.
+                        # Large results exceed Lambda response limits and must be stored durably
+                        # before the execution completes.
+                        try:
+                            execution_state.create_checkpoint_sync(failed_operation)
+                        except SuspendExecution:
+                            # Revocation can occur after the check at the top of this handler and
+                            # before this checkpoint. In that case create_checkpoint_sync raises
+                            # SuspendExecution. So return PENDING instead of FAILED.
+                            output = DurableExecutionInvocationOutput.create_pending().to_dict()
+                        except CheckpointError as checkpoint_error:
+                            output = handle_checkpoint_error(checkpoint_error).to_dict()
+                        else:
+                            output = DurableExecutionInvocationOutput(
+                                status=InvocationStatus.FAILED
+                            ).to_dict()
+                    else:
+                        output = result
+
+            return output
+
+        def invoke_worker() -> MutableMapping[str, Any]:
+            # This task owns Start, user cleanup, output preparation, resource
+            # cleanup and End. The caller owns the executor that runs this task.
+            with plugin_executor.run():
+                checkpoint_future: Future[None] | None = None
+
+                def wait_for_checkpoint() -> None:
+                    if checkpoint_future is not None:
+                        # Match the former executor join: observed checkpoint
+                        # errors retain their existing classification paths.
+                        wait((checkpoint_future,))
+
+                try:
+                    with contextlib.ExitStack() as resources:
+                        # LIFO cleanup: branches join with checkpointing alive;
+                        # then close stops checkpointing and the wait completes.
+                        resources.callback(wait_for_checkpoint)
+                        resources.callback(
+                            plugin_executor._run_in_invocation_context,
+                            execution_state.close,
+                        )
+                        execution_operation = execution_state.get_execution_operation()
+
+                        # execute the plugins
+                        plugin_executor.on_invocation_start(
+                            execution_arn=invocation_input.durable_execution_arn,
+                            lambda_context=context,
+                            execution_start_time=(
+                                execution_operation.start_timestamp
+                                if execution_operation is not None
+                                else None
+                            ),
+                            is_first_invocation=not has_prior_operations,
+                            execution_input=input_event,
+                            # Read the map through a callable rather than snapshotting it
+                            # here: the invocation-end hook needs the state as of the end of
+                            # the invocation, and neither hook pays for the conversion until
+                            # a plugin actually reads it.
+                            operations_provider=lambda: execution_state.operations,
+                            updated_operation_ids=invocation_input.updated_operation_ids,
+                        )
+                        # No checkpoint work starts until all Start hooks finish.
+                        checkpoint_future = executor.submit(
+                            execution_state.checkpoint_batches_forever
+                        )
+                        output = plugin_executor._run_in_invocation_context(
+                            invoke_handler
+                        )
+                    # The ExitStack has now closed execution_state and joined the
+                    # checkpoint thread. So any revocation observed while draining
+                    # in-flight checkpoint work is visible here. A terminal status
+                    # built before that point can no longer be recorded with the
+                    # service. So report PENDING instead; the next invocation
+                    # replays from the last durable checkpoint.
+                    if (
+                        output.get("Status")
+                        in {
+                            InvocationStatus.SUCCEEDED.value,
+                            InvocationStatus.FAILED.value,
+                        }
+                        and execution_state.is_checkpoint_token_revoked
+                    ):
+                        logger.debug(
+                            "Checkpoint token revoked during invocation cleanup; ending "
+                            "invocation with PENDING instead of terminal status."
+                        )
+                        output = (
+                            DurableExecutionInvocationOutput.create_pending().to_dict()
+                        )
+                    plugin_executor.on_invocation_end(
+                        DurableExecutionInvocationOutput.from_dict(output)
                     )
-                    failed_operation = OperationUpdate.create_execution_fail(
-                        error=ErrorObject.from_exception(e)
+                    return output
+                except Exception as error:
+                    plugin_executor.on_invocation_end(
+                        DurableExecutionInvocationOutput.create_retry(
+                            ErrorObject.from_exception(error)
+                        )
                     )
-
-                    # Checkpoint large result with blocking (is_sync=True, default).
-                    # Must ensure the result is persisted before returning to Lambda.
-                    # Large results exceed Lambda response limits and must be stored durably
-                    # before the execution completes.
-                    try:
-                        execution_state.create_checkpoint_sync(failed_operation)
-                    except CheckpointError as e:
-                        return handle_checkpoint_error(e).to_dict()
-                    return DurableExecutionInvocationOutput(
-                        status=InvocationStatus.FAILED
-                    ).to_dict()
-
-                return result
+                    raise
+
+        with ThreadPoolExecutor(
+            max_workers=2, thread_name_prefix="dex-handler"
+        ) as executor:
+            if configured_plugins:
+                invocation_future = executor.submit(
+                    contextvars.copy_context().run, invoke_worker
+                )
+            else:
+                # Preserve the original fresh-worker context without plugins.
+                invocation_future = executor.submit(invoke_worker)
+            return invocation_future.result()
 
     return wrapper
 

aws_durable_execution_sdk_python/identifier.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import hashlib
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 from aws_durable_execution_sdk_python.exceptions import (
     NonDeterministicExecutionError,
@@ -34,17 +34,17 @@ class OperationIdNamespace:
 
 @dataclass(frozen=True)
 class OperationIdentifier:
-    """Container for operation id, type, subtype, hierarchy, and name."""
+    """Operation identity with an explicit type independent of its subtype."""
 
     operation_id: str
     sub_type: OperationSubType
     parent_id: str | None = None
     name: str | None = None
-    operation_type: OperationType | None = None
+    operation_type: OperationType = field(kw_only=True)
 
     @property
     def type(self) -> OperationType:
-        return self.operation_type or OperationType.from_sub_type(self.sub_type)
+        return self.operation_type
 
     def validate_checkpoint(self, checkpoint: Operation | None) -> None:
         """Ensure replay history belongs to this operation before it is consumed."""

aws_durable_execution_sdk_python/lambda_service.py

@@ -75,29 +75,6 @@ class OperationType(Enum):
     CALLBACK = "CALLBACK"
     CHAINED_INVOKE = "CHAINED_INVOKE"
 
-    @classmethod
-    def from_sub_type(cls, sub_type: OperationSubType) -> OperationType:
-        match sub_type:
-            case OperationSubType.STEP | OperationSubType.WAIT_FOR_CONDITION:
-                return OperationType.STEP
-            case OperationSubType.WAIT:
-                return OperationType.WAIT
-            case OperationSubType.CHAINED_INVOKE:
-                return OperationType.CHAINED_INVOKE
-            case OperationSubType.CALLBACK:
-                return OperationType.CALLBACK
-            case (
-                OperationSubType.WAIT_FOR_CALLBACK
-                | OperationSubType.RUN_IN_CHILD_CONTEXT
-                | OperationSubType.MAP
-                | OperationSubType.MAP_ITERATION
-                | OperationSubType.PARALLEL
-                | OperationSubType.PARALLEL_BRANCH
-            ):
-                return OperationType.CONTEXT
-            case _:
-                raise ValueError(f"Unknown operation sub-type {sub_type}")
-
 
 class CallbackTimeoutType(Enum):
     TIMEOUT = "Callback.Timeout"
@@ -181,6 +158,11 @@ class DurableExecutionInvocationOutput:
         """Create a failed invocation output."""
         return cls(status=InvocationStatus.RETRY, error=error)
 
+    @classmethod
+    def create_pending(cls) -> DurableExecutionInvocationOutput:
+        """Create a pending invocation output without a result or error."""
+        return cls(status=InvocationStatus.PENDING)
+
 
 @dataclass(frozen=True)
 class ExecutionDetails:
@@ -453,19 +435,18 @@ class CallbackOptions:
 
 @dataclass(frozen=True)
 class ChainedInvokeOptions:
-    """
-    As of 2025/10/27:
-     - Chained invoke options only contains a function name
-    """
+    """Target and optional per-operation trace context for a chained invocation."""
 
     function_name: str
     tenant_id: str | None = None
+    x_amzn_trace_id: str | None = None
 
     @classmethod
     def from_dict(cls, data: MutableMapping[str, Any]) -> ChainedInvokeOptions:
         return cls(
             function_name=data["FunctionName"],
             tenant_id=data.get("TenantId"),
+            x_amzn_trace_id=data.get("XAmznTraceId"),
         )
 
     def to_dict(self) -> MutableMapping[str, Any]:
@@ -474,6 +455,8 @@ class ChainedInvokeOptions:
         }
         if self.tenant_id is not None:
             result["TenantId"] = self.tenant_id
+        if self.x_amzn_trace_id is not None:
+            result["XAmznTraceId"] = self.x_amzn_trace_id
 
         return result
 

aws_durable_execution_sdk_python/operation/child.py

@@ -10,11 +10,13 @@ from aws_durable_execution_sdk_python.exceptions import (
     ChildContextError,
     ExecutionError,
     InvocationError,
+    NonDeterministicExecutionError,
     SuspendExecution,
 )
 from aws_durable_execution_sdk_python.lambda_service import (
     ContextOptions,
     ErrorObject,
+    Operation,
     OperationStatus,
     OperationSubType,
     OperationUpdate,
@@ -71,6 +73,28 @@ class ChildOperationExecutor(OperationExecutor[T]):
         self.is_virtual: bool = config.is_virtual
         self.sub_type = config.sub_type or OperationSubType.RUN_IN_CHILD_CONTEXT
 
+    @staticmethod
+    def validate_virtual_checkpoint(
+        operation_identifier: OperationIdentifier,
+        checkpoint: Operation | None,
+    ) -> None:
+        """Reject a checkpoint at the operation id of a virtual child context.
+
+        A virtual child never checkpoints itself, so a checkpoint at its id was
+        written by a non-virtual child in an earlier invocation. Consuming it
+        would skip the virtual body and leave its inner operations unvalidated.
+        """
+        if not isinstance(checkpoint, Operation):
+            return
+        msg = (
+            "Non-deterministic child context at "
+            f"id={operation_identifier.operation_id!r}: checkpoint contains a "
+            "non-virtual child context but current child context is virtual"
+        )
+        raise NonDeterministicExecutionError(
+            msg, step_id=operation_identifier.operation_id
+        )
+
     def check_result_status(self) -> CheckResult[T]:
         """Check operation status and create START checkpoint if needed.
 
@@ -82,8 +106,14 @@ class ChildOperationExecutor(OperationExecutor[T]):
 
         Raises:
             ChildContextError: For FAILED operations
+            NonDeterministicExecutionError: For a virtual child whose operation
+                id already has a checkpoint
         """
         checkpointed_result = self._get_checkpoint_result()
+        if self.is_virtual:
+            self.validate_virtual_checkpoint(
+                self.operation_identifier, checkpointed_result.operation
+            )
 
         # Terminal success without replay_children - deserialize and return
         if (

aws_durable_execution_sdk_python/operation/invoke.py

@@ -123,12 +123,18 @@ class InvokeOperationExecutor(OperationExecutor[R]):
                 operation_id=self.operation_identifier.operation_id,
                 durable_execution_arn=self.state.durable_execution_arn,
             )
+            # Only a new START contributes outbound context. Replayed pending or
+            # terminal invokes keep their checkpointed identity and outcome.
+            propagation = self.state.provide_propagation_metadata(
+                self.operation_identifier, self.function_name
+            )
             start_operation: OperationUpdate = OperationUpdate.create_invoke_start(
                 identifier=self.operation_identifier,
                 payload=serialized_payload,
                 chained_invoke_options=ChainedInvokeOptions(
                     function_name=self.function_name,
                     tenant_id=self.config.tenant_id,
+                    x_amzn_trace_id=propagation.x_amzn_trace_id,
                 ),
             )
             # Checkpoint invoke START with blocking (is_sync=True).