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

aws_durable_execution_sdk_python/__about__.py

@@ -1,4 +1,4 @@
 # SPDX-FileCopyrightText: 2025-present Amazon.com, Inc. or its affiliates.
 #
 # SPDX-License-Identifier: Apache-2.0
-__version__ = "2.0.1"
+__version__ = "2.1.0"

aws_durable_execution_sdk_python/concurrency/executor.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import contextvars
 import heapq
 import logging
 import queue
@@ -44,7 +45,7 @@ from aws_durable_execution_sdk_python.identifier import (
     OperationIdentifier,
     OperationIdNamespace,
 )
-from aws_durable_execution_sdk_python.lambda_service import ErrorObject
+from aws_durable_execution_sdk_python.lambda_service import ErrorObject, OperationType
 from aws_durable_execution_sdk_python.operation.child import child_handler
 
 
@@ -183,6 +184,7 @@ class ConcurrentExecutor(Generic[CallableType, ResultType]):
                 executable.index
             ),
             sub_type=self.sub_type_iteration,
+            operation_type=OperationType.CONTEXT,
             parent_id=executor_context._parent_id,  # noqa: SLF001
             name=self.get_iteration_name(executable.index),
         )
@@ -291,13 +293,25 @@ class ConcurrentExecutor(Generic[CallableType, ResultType]):
 
         def submit(branch: Branch[CallableType, ResultType]) -> None:
             branch.start()
-            pool.submit(
-                self._branch_worker,
-                execution_state,
-                executor_context,
-                events,
-                branch.executable,
-            )
+            if execution_state._plugin_executor._plugins:  # noqa: SLF001
+                # Every admission/resume gets its own Context. Branch bindings
+                # cannot leak into siblings, the coordinator, or reused workers.
+                pool.submit(
+                    contextvars.copy_context().run,
+                    self._branch_worker,
+                    execution_state,
+                    executor_context,
+                    events,
+                    branch.executable,
+                )
+            else:
+                pool.submit(
+                    self._branch_worker,
+                    execution_state,
+                    executor_context,
+                    events,
+                    branch.executable,
+                )
 
         try:
             # Only rebuild the items snapshot after a terminal event changes

aws_durable_execution_sdk_python/context.py

@@ -51,7 +51,10 @@ from aws_durable_execution_sdk_python.operation.callback import (
     CallbackOperationExecutor,
     wait_for_callback_handler,
 )
-from aws_durable_execution_sdk_python.operation.child import child_handler
+from aws_durable_execution_sdk_python.operation.child import (
+    ChildOperationExecutor,
+    child_handler,
+)
 from aws_durable_execution_sdk_python.operation.invoke import InvokeOperationExecutor
 from aws_durable_execution_sdk_python.operation.map import map_handler
 from aws_durable_execution_sdk_python.operation.parallel import parallel_handler
@@ -515,12 +518,18 @@ class DurableContext(DurableContextProtocol):
 
     @contextmanager
     def _replay_aware(
-        self, operation_identifier: OperationIdentifier | None = None
+        self,
+        operation_identifier: OperationIdentifier | None = None,
+        *,
+        is_virtual: bool = False,
     ) -> Iterator[None]:
         """Wrap a single operation with replay-boundary detection.
 
         The operation kind is inferred from its own checkpoint (when one
-        exists), so no per-call-site flag is needed.
+        exists). The one per-call-site input is `is_virtual`: a virtual child
+        context never checkpoints itself, so with `is_virtual=True` any
+        checkpoint at its id is rejected right after the identity check,
+        before replay and update hooks run.
 
         The boundary has these parts:
 
@@ -555,6 +564,10 @@ class DurableContext(DurableContextProtocol):
         if was_replaying and next_operation is not None:
             if operation_identifier is not None:
                 operation_identifier.validate_checkpoint(next_operation)
+                if is_virtual:
+                    ChildOperationExecutor.validate_virtual_checkpoint(
+                        operation_identifier, next_operation
+                    )
 
             # While replaying, an operation that already has a checkpoint was
             # observed in a prior invocation. If the backend says the operation
@@ -599,7 +612,9 @@ class DurableContext(DurableContextProtocol):
         self,
         sub_type: OperationSubType,
         name: str | None = None,
-        operation_type: OperationType | None = None,
+        *,
+        operation_type: OperationType,
+        is_virtual: bool = False,
     ) -> Iterator[OperationIdentifier]:
         """Allocate an operation and validate its replay identity before hooks."""
         operation_identifier = OperationIdentifier(
@@ -609,7 +624,7 @@ class DurableContext(DurableContextProtocol):
             name=name,
             operation_type=operation_type,
         )
-        with self._replay_aware(operation_identifier):
+        with self._replay_aware(operation_identifier, is_virtual=is_virtual):
             yield operation_identifier
 
     # endregion replay status
@@ -636,7 +651,7 @@ class DurableContext(DurableContextProtocol):
         if not config:
             config = CallbackConfig()
         with self._operation_replay_aware(
-            OperationSubType.CALLBACK, name
+            OperationSubType.CALLBACK, name, operation_type=OperationType.CALLBACK
         ) as operation_identifier:
             executor: CallbackOperationExecutor = CallbackOperationExecutor(
                 state=self.state,
@@ -672,7 +687,9 @@ class DurableContext(DurableContextProtocol):
         if not config:
             config = InvokeConfig[P, R]()
         with self._operation_replay_aware(
-            OperationSubType.CHAINED_INVOKE, name
+            OperationSubType.CHAINED_INVOKE,
+            name,
+            operation_type=OperationType.CHAINED_INVOKE,
         ) as operation_identifier:
             executor: InvokeOperationExecutor[R] = InvokeOperationExecutor(
                 function_name=function_name,
@@ -700,7 +717,7 @@ class DurableContext(DurableContextProtocol):
             config.completion_config._validate_for_total(len(inputs))
 
         with self._operation_replay_aware(
-            OperationSubType.MAP, map_name
+            OperationSubType.MAP, map_name, operation_type=OperationType.CONTEXT
         ) as operation_identifier:
             operation_id = operation_identifier.operation_id
             map_context = self.create_child_context(operation_id=operation_id)
@@ -751,7 +768,7 @@ class DurableContext(DurableContextProtocol):
             config.completion_config._validate_for_total(len(functions))
 
         with self._operation_replay_aware(
-            OperationSubType.PARALLEL, name
+            OperationSubType.PARALLEL, name, operation_type=OperationType.CONTEXT
         ) as operation_identifier:
             operation_id = operation_identifier.operation_id
             parallel_context = self.create_child_context(operation_id=operation_id)
@@ -811,14 +828,15 @@ class DurableContext(DurableContextProtocol):
             if config and config.sub_type
             else OperationSubType.RUN_IN_CHILD_CONTEXT
         )
+        is_virtual: bool = config.is_virtual if config else False
 
         with self._operation_replay_aware(
             sub_type,
             step_name,
             operation_type=OperationType.CONTEXT,
+            is_virtual=is_virtual,
         ) as operation_identifier:
             operation_id = operation_identifier.operation_id
-            is_virtual: bool = config.is_virtual if config else False
 
             def callable_with_child_context():
                 return func(
@@ -845,7 +863,7 @@ class DurableContext(DurableContextProtocol):
         if not config:
             config = StepConfig()
         with self._operation_replay_aware(
-            OperationSubType.STEP, step_name
+            OperationSubType.STEP, step_name, operation_type=OperationType.STEP
         ) as operation_identifier:
             executor: StepOperationExecutor[T] = StepOperationExecutor(
                 func=func,
@@ -868,7 +886,7 @@ class DurableContext(DurableContextProtocol):
             msg = "duration must be at least 1 second"
             raise ValidationError(msg)
         with self._operation_replay_aware(
-            OperationSubType.WAIT, name
+            OperationSubType.WAIT, name, operation_type=OperationType.WAIT
         ) as operation_identifier:
             wait_seconds = duration.seconds
             executor: WaitOperationExecutor = WaitOperationExecutor(
@@ -938,7 +956,7 @@ class DurableContext(DurableContextProtocol):
             raise ValidationError(msg)
 
         with self._operation_replay_aware(
-            OperationSubType.WAIT_FOR_CONDITION, name
+            OperationSubType.WAIT_FOR_CONDITION, name, operation_type=OperationType.STEP
         ) as operation_identifier:
             executor: WaitForConditionOperationExecutor[T] = (
                 WaitForConditionOperationExecutor(

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