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

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).

aws_durable_execution_sdk_python/plugin.py

@@ -1,15 +1,16 @@
 from __future__ import annotations
 
 import contextlib
+import contextvars
 import copy
 import datetime
-import functools
+import inspect
 import logging
 from collections.abc import Mapping, Sequence
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Any, Callable, MutableMapping, cast
+from typing import Any, Callable, cast
 
 from aws_durable_execution_sdk_python.identifier import OperationIdentifier
 from aws_durable_execution_sdk_python.lambda_service import (
@@ -259,8 +260,8 @@ class InvocationInfo:
     )
     """EXPERIMENTAL: The deserialized execution input, when available.
 
-    Surfaced to instrumentation plugins that need to record it (e.g. Workflow
-    Insight). Mirrors the JS SDK's ``InvocationInfo.executionInput``.
+    Surfaced to instrumentation plugins that need to record it (e.g. Durable
+    Execution Telemetry). Mirrors the JS SDK's ``InvocationInfo.executionInput``.
 
     Excluded from ``repr`` on purpose: instrumentation logs hook infos wholesale
     (the bundled OTel plugins at debug level, the plugin example at info), so
@@ -381,9 +382,41 @@ class InvocationEndInfo(InvocationInfo):
         )
 
 
+@dataclass(frozen=True)
+class PropagationInput:
+    """SDK-owned identity for a new chained-invoke START propagation request.
+
+    This contract does not contain generated service-model or telemetry types.
+    """
+
+    execution_arn: str
+    operation_id: str
+    target_function_name: str
+    parent_operation_id: str | None = None
+
+
+@dataclass(frozen=True)
+class PropagationMetadata:
+    """Immutable plugin contribution; members are opaque to the core SDK."""
+
+    x_amzn_trace_id: str | None = None
+
+
 class DurableInstrumentationPlugin:
     """Base class for plugins. Override only the methods you need."""
 
+    def provide_propagation_metadata(
+        self,
+        info: PropagationInput,
+    ) -> PropagationMetadata | None:
+        """Synchronously provide metadata without performing a durable operation.
+
+        Called before checkpointing a new invoke START, never for its replay.
+        Existing plugins can inherit the no-op default. A failed, uncommitted
+        START may collect again on a later attempt.
+        """
+        return None
+
     def on_invocation_start(self, info: InvocationStartInfo) -> None:
         """Called when an invocation starts. This is called within the thread that runs user function handler.
 
@@ -466,6 +499,92 @@ class PluginExecutor:
         self._executor: ThreadPoolExecutor | None = None
         self._invocation_status: InvocationStartInfo | None = None
         self._operations_provider: Callable[[], Mapping[str, Operation]] | None = None
+        # Non-None only after a Start hook fails: later setup and invocation
+        # work use its pre-hook snapshot. End still uses each hook's token owner.
+        self._startup_context: contextvars.Context | None = None
+        self._invocation_contexts: list[contextvars.Context | None] = []
+
+    def provide_propagation_metadata(
+        self, info: PropagationInput
+    ) -> PropagationMetadata:
+        """Collect opaque metadata synchronously in configured plugin order.
+
+        The first non-blank value wins. Ordinary plugin failures are isolated,
+        matching lifecycle dispatch; BaseException cancellation/control signals
+        retain the existing propagation policy. The caller attaches the result
+        to its operation model before checkpoint serialization.
+        """
+        value: str | None = None
+        owner: str | None = None
+        conflict_count = 0
+        for plugin in self._plugins:
+            # Read type identity without invoking instance attribute access.
+            identity = "unknown plugin"
+            try:
+                plugin_type = type(plugin)
+                identity = type.__getattribute__(plugin_type, "__qualname__")
+                module = type.__getattribute__(plugin_type, "__module__")
+                if type(module) is str:
+                    identity = ".".join((module, identity))
+            except Exception:
+                pass
+            try:
+                metadata = plugin.provide_propagation_metadata(info)
+                if metadata is None:
+                    continue
+                if inspect.iscoroutine(metadata):
+                    metadata.close()
+                    raise TypeError("propagation metadata hook must be synchronous")
+                if not isinstance(metadata, PropagationMetadata):
+                    raise TypeError("expected PropagationMetadata or None")
+                candidate = metadata.x_amzn_trace_id
+                if candidate is not None:
+                    if not isinstance(candidate, str):
+                        raise TypeError("x_amzn_trace_id must be a string or None")
+                    # Treat str subclasses as data, not executable comparison
+                    # hooks during first-value/conflict aggregation.
+                    candidate = str.__str__(candidate)
+            except Exception:
+                self._log_propagation_diagnostic(
+                    logging.ERROR,
+                    "Plugin %s propagation metadata failure ignored",
+                    identity,
+                    exc_info=True,
+                )
+                continue
+            if candidate is None or not candidate.strip():
+                continue
+            if value is None:
+                value, owner = candidate, identity
+            elif candidate != value:
+                conflict_count += 1
+                self._log_propagation_diagnostic(
+                    logging.WARNING,
+                    "Propagation metadata conflict between %s and %s for "
+                    "x_amzn_trace_id; keeping first value (conflict_count=%d)",
+                    owner,
+                    identity,
+                    conflict_count,
+                )
+        return PropagationMetadata(x_amzn_trace_id=value)
+
+    @staticmethod
+    def _log_propagation_diagnostic(
+        level: int,
+        message: str,
+        *args: object,
+        exc_info: bool = False,
+    ) -> None:
+        """Keep ordinary diagnostic failures from changing propagation collection."""
+        try:
+            logger.log(level, message, *args, exc_info=exc_info)
+        except Exception:
+            pass
+
+    @property
+    def has_plugins(self) -> bool:
+        """Return True when at least one plugin is registered."""
+        return bool(self._plugins)
 
     @contextlib.contextmanager
     def run(self):
@@ -479,12 +598,14 @@ class PluginExecutor:
         finally:
             self._invocation_status = None
             self._operations_provider = None
+            self._startup_context = None
+            self._invocation_contexts.clear()
             # Shut down the thread pool, waiting for pending tasks to complete.
             if self._executor:
                 self._executor.shutdown(wait=True)
 
     @staticmethod
-    def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None:
+    def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> bool:
         """Invoke the appropriate plugin callback. Runs inside the thread pool."""
         try:
             match info:
@@ -507,18 +628,54 @@ class PluginExecutor:
         except Exception:
             # log and ignore the exception
             logger.exception("Plugin %s exception ignored", plugin.__class__.__name__)
+            return False
+        return True
 
     def execute_plugins(self, info, sync):
         if not self._executor:
             return
-        for plugin in self._plugins:
-            if sync:
-                # this is called synchronously, so plugins will be able to manipulate thread local objects
-                self._dispatch_plugin(plugin, info)
+        if sync and isinstance(info, InvocationStartInfo):
+            self._startup_context = None
+            self._invocation_contexts.clear()
+        for index, plugin in enumerate(self._plugins):
+            if sync and isinstance(info, InvocationStartInfo):
+                owner = self._startup_context
+                before = (
+                    owner.copy() if owner is not None else contextvars.copy_context()
+                )
+                self._invocation_contexts.append(owner)
+                succeeded = (
+                    owner.run(self._dispatch_plugin, plugin, info)
+                    if owner is not None
+                    else self._dispatch_plugin(plugin, info)
+                )
+                if not succeeded:
+                    # A failing hook may have left new bindings with no reset token.
+                    # Continue setup and the handler in the pre-hook snapshot.
+                    self._startup_context = before
+            elif sync:
+                # End hooks must reset tokens in the Context that created them,
+                # even when a failed start hook moved later setup to a snapshot.
+                owner = (
+                    self._invocation_contexts[index]
+                    if isinstance(info, InvocationEndInfo)
+                    and index < len(self._invocation_contexts)
+                    else None
+                )
+                if owner is not None:
+                    owner.run(self._dispatch_plugin, plugin, info)
+                else:
+                    self._dispatch_plugin(plugin, info)
             else:
                 # this is called asynchronously, so plugins cannot manipulate thread local objects
                 self._executor.submit(self._dispatch_plugin, plugin, info)
 
+    def _run_in_invocation_context(self, invoke: Callable[[], Any]) -> Any:
+        """Continue in the pre-hook Context only after a failed Start hook."""
+        if self._startup_context is not None:
+            return self._startup_context.run(invoke)
+        return invoke()
+
     def _snapshot_operation_infos(
         self,
         operations_provider: Callable[[], Mapping[str, Operation]] | None,
@@ -763,7 +920,9 @@ class PluginExecutor:
             operations: full operation map after the update, when available.
             previous_operations: operation map before the update, when available.
         """
-        if operation_or_operations is None:
+        # Only plugins read the info built below, and the change info converts
+        # every operation. So without a plugin, return before building any of it.
+        if operation_or_operations is None or not self.has_plugins:
             return
 
         updated_operations: list[Operation] = (
@@ -836,28 +995,3 @@ class PluginExecutor:
             OperationStatus.CANCELLED,
             OperationStatus.STOPPED,
         ]
-
-    @property
-    def handle_durable_output(self):
-        def decorator(func: Callable[[Any, LambdaContext], MutableMapping[str, Any]]):
-            @functools.wraps(func)
-            def wrapper(event: Any, context: LambdaContext):
-                with self.run():
-                    try:
-                        output = func(event, context)
-
-                        self.on_invocation_end(
-                            output=DurableExecutionInvocationOutput.from_dict(output),
-                        )
-                        return output
-                    except Exception as e:
-                        self.on_invocation_end(
-                            output=DurableExecutionInvocationOutput.create_retry(
-                                ErrorObject.from_exception(e)
-                            ),
-                        )
-                        raise
-
-            return wrapper
-
-        return decorator

aws_durable_execution_sdk_python/plugin_discovery.py

@@ -121,6 +121,9 @@ def _create_plugin(
         ) from error
 
     if type(plugin) is not provider.plugin_type:
+        # A wrong-type result may already own constructor resources, but it is
+        # never appended to the accepted list cleaned up by the outer boundary.
+        _notify_registration_result([plugin], registered=False)
         raise PluginLoadError(
             f"Durable instrumentation plugin provider '{plugin_name}' returned "
             f"{_qualified_type_name(plugin)}; expected "
@@ -130,6 +133,101 @@ def _create_plugin(
     return plugin
 
 
+def _registration_api_owner(plugin_type: type[object]) -> type[object] | None:
+    """Find the first explicit marker in the MRO without metaclass descriptors.
+
+    A subclass inherits its declaring class's registration contract. Its legacy
+    same-named fields/helpers are not part of that contract unless it repeats
+    the marker. An explicit unsupported marker shadows any inherited opt-in.
+    """
+    mro: tuple[type[object], ...] = type.__dict__["__mro__"].__get__(
+        plugin_type, type(plugin_type)
+    )
+    for declaring_type in mro:
+        namespace = type.__dict__["__dict__"].__get__(
+            declaring_type, type(declaring_type)
+        )
+        if "__durable_registration_api__" in namespace:
+            version = namespace["__durable_registration_api__"]
+            return declaring_type if type(version) is int and version == 1 else None
+    return None
+
+
+def _registration_result_callback(
+    plugin: DurableInstrumentationPlugin, owner: type[object]
+) -> object | None:
+    """Bind the opted-in contract without calling a legacy subclass helper."""
+    if owner is type(plugin):
+        return getattr(plugin, "on_registration_result", None)
+    mro: tuple[type[object], ...] = type.__dict__["__mro__"].__get__(owner, type(owner))
+    for declaring_type in mro:
+        namespace = type.__dict__["__dict__"].__get__(
+            declaring_type, type(declaring_type)
+        )
+        if "on_registration_result" in namespace:
+            callback = namespace["on_registration_result"]
+            bind = getattr(callback, "__get__", None)
+            return bind(plugin, type(plugin)) if bind is not None else callback
+    return None
+
+
+def _exclusive_groups(plugin_type: type[DurableInstrumentationPlugin]) -> list[str]:
+    """Accumulate opted-in MRO contracts without interpreting legacy fields."""
+    mro: tuple[type[object], ...] = type.__dict__["__mro__"].__get__(
+        plugin_type, type(plugin_type)
+    )
+    groups: list[str] = []
+    for declaring_type in mro:
+        namespace = type.__dict__["__dict__"].__get__(
+            declaring_type, type(declaring_type)
+        )
+        version = namespace.get("__durable_registration_api__")
+        if type(version) is not int or version != 1:
+            continue
+        try:
+            group = getattr(declaring_type, "exclusive_group", None)
+        except Exception as error:
+            raise PluginLoadError(
+                f"Cannot read exclusive_group for {_qualified_class_name(plugin_type)}."
+            ) from error
+        if group is None:
+            continue
+        if type(group) is not str or not group.strip():
+            raise PluginLoadError(
+                f"Durable instrumentation plugin {_qualified_class_name(plugin_type)} "
+                "must declare exclusive_group as None or a non-empty string."
+            )
+        # Repeating a marker or a group does not register this one plugin twice.
+        if group not in groups:
+            groups.append(group)
+    return groups
+
+
+def _validate_exclusive_groups(
+    plugin_types: Sequence[type[DurableInstrumentationPlugin]],
+) -> None:
+    """Reject competing instrumentation before any lifecycle hooks run."""
+    groups: dict[str, type[DurableInstrumentationPlugin]] = {}
+    for plugin_type in plugin_types:
+        if _registration_api_owner(plugin_type) is None:
+            continue
+        for group in _exclusive_groups(plugin_type):
+            if (previous := groups.get(group)) is not None:
+                if previous is plugin_type:
+                    raise PluginLoadError(
+                        f"Durable instrumentation plugin {_qualified_class_name(plugin_type)} "
+                        f"is registered more than once in exclusive group '{group}'. "
+                        "Register this plugin only once in plugins."
+                    )
+                raise PluginLoadError(
+                    f"Durable instrumentation plugins {_qualified_class_name(previous)} "
+                    f"and {_qualified_class_name(plugin_type)} are mutually exclusive "
+                    f"(group '{group}'). Keep only one plugin from this group in "
+                    "plugins and DURABLE_EXECUTION_PLUGINS."
+                )
+            groups[group] = plugin_type
+
+
 def load_configured_plugins(
     explicit_plugins: Sequence[DurableInstrumentationPlugin] | None,
     *,
@@ -144,9 +242,45 @@ def load_configured_plugins(
     """
 
     resolved_plugins = list(explicit_plugins or [])
+    try:
+        result = _resolve_configured_plugins(resolved_plugins, environment=environment)
+    except PluginLoadError:
+        _notify_registration_result(resolved_plugins, registered=False)
+        raise
+    _notify_registration_result(result, registered=True)
+    return result
+
+
+def _notify_registration_result(
+    plugins: Sequence[DurableInstrumentationPlugin],
+    *,
+    registered: bool,
+) -> None:
+    """Let a plugin release rejected constructor resources without failing loading."""
+    for plugin in plugins:
+        owner = _registration_api_owner(type(plugin))
+        if owner is None:
+            continue
+        try:
+            callback = _registration_result_callback(plugin, owner)
+            if callable(callback):
+                callback(registered)
+        except Exception:
+            try:
+                logger.exception("Plugin registration-result callback failed")
+            except Exception:
+                pass
+
+
+def _resolve_configured_plugins(
+    resolved_plugins: list[DurableInstrumentationPlugin],
+    *,
+    environment: Mapping[str, str] | None,
+) -> list[DurableInstrumentationPlugin]:
     resolved_environment = os.environ if environment is None else environment
     plugin_names = _parse_configured_plugin_names(resolved_environment)
     if not plugin_names:
+        _validate_exclusive_groups([type(plugin) for plugin in resolved_plugins])
         return resolved_plugins
 
     try:
@@ -167,6 +301,9 @@ def load_configured_plugins(
         type(plugin): "the decorator's plugins argument" for plugin in resolved_plugins
     }
 
+    selected_providers: list[
+        tuple[str, metadata.EntryPoint, DurableInstrumentationPluginProvider]
+    ] = []
     for plugin_name in plugin_names:
         matching_entry_points = entry_points_by_name.get(plugin_name, [])
         if not matching_entry_points:
@@ -202,8 +339,15 @@ def load_configured_plugins(
             )
             continue
 
-        plugin = _create_plugin(plugin_name, entry_point, provider)
-        resolved_plugins.append(plugin)
+        selected_providers.append((plugin_name, entry_point, provider))
         registered_types[provider.plugin_type] = f"dynamic provider '{plugin_name}'"
 
+    # Factories can install instrumentation globally. Validate every selected
+    # type first so a rejected configuration leaves no discarded plugin behind.
+    _validate_exclusive_groups(
+        [type(plugin) for plugin in resolved_plugins]
+        + [provider.plugin_type for _, _, provider in selected_providers]
+    )
+    for plugin_name, entry_point, provider in selected_providers:
+        resolved_plugins.append(_create_plugin(plugin_name, entry_point, provider))
     return resolved_plugins