aws-durable-execution-sdk-python-testing 1.2.0__py3-none-any.whl → 1.3.0__py3-none-any.whl

aws_durable_execution_sdk_python_testing/__about__.py

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

aws_durable_execution_sdk_python_testing/checkpoint/core.py

@@ -0,0 +1,171 @@
+"""The shared checkpoint write-transaction.
+
+Both checkpoint entry points apply a batch of updates and compute the
+delta of operations the handler has not yet seen. :class:`CheckpointCore`
+holds that common logic so the two callers differ only in their gate,
+locking, and response type.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, NamedTuple
+
+from aws_durable_execution_sdk_python_testing.checkpoint.validators.checkpoint import (
+    CheckpointValidator,
+)
+from aws_durable_execution_sdk_python_testing.execution import (
+    CheckpointIdempotencyRecord,
+    OperationPaginatorState,
+)
+from aws_durable_execution_sdk_python_testing.token import CheckpointToken
+
+if TYPE_CHECKING:
+    from datetime import datetime
+
+    from aws_durable_execution_sdk_python.lambda_service import (
+        Operation,
+        OperationUpdate,
+    )
+
+    from aws_durable_execution_sdk_python_testing.checkpoint.effects import (
+        CheckpointEffect,
+    )
+    from aws_durable_execution_sdk_python_testing.checkpoint.transformer import (
+        CheckpointRequestDispatcher,
+    )
+    from aws_durable_execution_sdk_python_testing.execution import Execution
+
+
+class CheckpointResult(NamedTuple):
+    """Outcome of applying a checkpoint: the new token, the operations to
+    return to the handler this round, and the lifecycle effects raised.
+
+    ``checkpoint_token`` is None when the execution is paused and the token
+    is withheld, telling the SDK this invocation may checkpoint no further.
+    """
+
+    checkpoint_token: str | None
+    operations: list[Operation]
+    effects: list[CheckpointEffect]
+
+
+class CheckpointCore:
+    """The checkpoint write-transaction shared by both entry points."""
+
+    @staticmethod
+    def match_cached(
+        execution: Execution,
+        checkpoint_token: str,
+        client_token: str | None,
+    ) -> CheckpointIdempotencyRecord | None:
+        """Return the cached idempotency record when the incoming
+        ``(client_token, checkpoint_token)`` matches the last checkpoint,
+        else ``None``.
+
+        A missing ``client_token`` cannot replay: idempotency requires an
+        explicit client identifier. The caller builds its own response
+        type from the returned record.
+        """
+        if not client_token:
+            return None
+        cached: CheckpointIdempotencyRecord | None = execution.last_checkpoint
+        if cached is None:
+            return None
+        if (
+            cached.client_token != client_token
+            or cached.inbound_checkpoint_token != checkpoint_token
+        ):
+            return None
+        return cached
+
+    @staticmethod
+    def apply(
+        execution: Execution,
+        checkpoint_token: str,
+        updates: list[OperationUpdate],
+        client_token: str | None,
+        dispatcher: CheckpointRequestDispatcher,
+        now: datetime,
+    ) -> CheckpointResult:
+        """Apply ``updates`` to ``execution`` and compute the response delta.
+
+        Advances ``token_sequence`` exactly once, returns the full set of
+        operations the handler has not yet seen unless paused, advances
+        ``handler_seen_seq`` only for returned operations, and records an
+        idempotency entry for a byte-identical replay of a retried call
+        only when returning a checkpoint token.
+
+        The caller is responsible for the invocation gate, locking,
+        persistence, and applying the returned effects.
+
+        ``now`` is the checkpoint's single source of "now", resolved from
+        the execution's clock, so all timestamps stamped by this apply
+        advance with modeled time under a skip clock.
+
+        Before computing the delta, completes every wait and step retry
+        whose scheduled moment has passed, so those transitions are
+        delivered in this response and a suspended branch resumes within
+        the running invocation instead of waiting for the next one.
+        """
+        effects: list[CheckpointEffect] = []
+        if updates:
+            CheckpointValidator.validate_input(updates, execution)
+            effects = dispatcher.apply_updates(
+                execution=execution,
+                updates=updates,
+                client_token=client_token,
+                touch=execution.touch_operation,
+                now=now,
+            )
+
+        new_token_sequence: int = execution.advance_token_sequence()
+
+        # Pick up scheduler-due completions before pinning the paginator:
+        # each transition touches its operation, so it lands in this
+        # response's unseen delta and the handler sees it immediately.
+        execution.complete_due_operations(now)
+
+        paginator: OperationPaginatorState = OperationPaginatorState.pin(execution)
+        # The checkpoint response returns the full unseen delta in a single
+        # response. Advance handler_seen_seq to cover every returned op so
+        # the next delta carries only operations touched after this response.
+        # Paused checkpoints persist updates without delivering any state.
+        response_ops: list[Operation] = (
+            [] if execution.is_paused else paginator.unseen_operations()
+        )
+        if response_ops:
+            highest_delivered_seq: int = max(
+                execution.operation_last_touched_seq[op.operation_id]
+                for op in response_ops
+            )
+            paginator.advance_handler_seen(highest_delivered_seq)
+
+        new_token: str = CheckpointToken(
+            execution_arn=execution.durable_execution_arn,
+            token_sequence=new_token_sequence,
+            invocation_id=execution.current_invocation_id,
+        ).to_str()
+
+        # A paused execution registers this checkpoint's updates but withholds
+        # the token: a response without one tells the SDK this invocation may
+        # checkpoint no further, so it reports PENDING at its next checkpoint
+        # rather than continuing, and owes a re-invoke once resumed.
+        #
+        # Checkpoints that were interrupted by a pause are not added to the
+        # idempotency record. A retry of this checkpoint will fail rather than be
+        # answered again.
+
+        outbound_token: str | None = new_token
+        if execution.is_paused:
+            outbound_token = None
+            execution.defer_invocation()
+        else:
+            execution.last_checkpoint = CheckpointIdempotencyRecord(
+                client_token=client_token or "",
+                inbound_checkpoint_token=checkpoint_token,
+                outbound_checkpoint_token=outbound_token,
+                operations=list(response_ops),
+                next_marker=None,
+            )
+
+        return CheckpointResult(outbound_token, response_ops, effects)

aws_durable_execution_sdk_python_testing/checkpoint/effects.py

@@ -0,0 +1,61 @@
+"""Lifecycle effects produced by applying checkpoint updates.
+
+Applying a batch of operation updates can imply actions on the execution
+beyond the state change itself: the execution completed, it failed, or a
+callback was created and now needs a timeout. These are returned as data
+so the caller can act on them after the write finishes, rather than
+triggering them in the middle of it.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from aws_durable_execution_sdk_python.lambda_service import (
+        CallbackOptions,
+        ErrorObject,
+    )
+
+    from aws_durable_execution_sdk_python_testing.token import CallbackToken
+
+
+@dataclass(frozen=True)
+class Completed:
+    """The execution completed successfully with ``result``."""
+
+    execution_arn: str
+    result: str | None = None
+
+
+@dataclass(frozen=True)
+class Failed:
+    """The execution failed with ``error``."""
+
+    execution_arn: str
+    error: ErrorObject | None
+
+
+@dataclass(frozen=True)
+class CallbackCreated:
+    """A callback was created and its timeout (if any) must be scheduled."""
+
+    execution_arn: str
+    operation_id: str
+    callback_options: CallbackOptions | None
+    callback_token: CallbackToken
+
+
+@dataclass(frozen=True)
+class ChainedInvokeStarted:
+    """A chained invoke was accepted and its target must be dispatched."""
+
+    execution_arn: str
+    operation_id: str
+    function_name: str
+    tenant_id: str | None
+    payload: str | None
+
+
+CheckpointEffect = Completed | Failed | CallbackCreated | ChainedInvokeStarted

aws_durable_execution_sdk_python_testing/checkpoint/processor.py

@@ -1,4 +1,10 @@
-"""Main checkpoint processor that orchestrates operation transformations."""
+"""In-process checkpoint flow.
+
+Orchestrates the full checkpoint request for in-process callers
+(``InMemoryServiceClient``). Mirrors the flow inside
+``Executor.checkpoint_execution`` used by the HTTP path, so both
+entry points share identical delta + watermark semantics.
+"""
 
 from __future__ import annotations
 
@@ -7,81 +13,128 @@ from typing import TYPE_CHECKING
 from aws_durable_execution_sdk_python.lambda_service import (
     CheckpointOutput,
     CheckpointUpdatedExecutionState,
-    OperationUpdate,
     StateOutput,
 )
 
+from aws_durable_execution_sdk_python_testing.checkpoint.core import CheckpointCore
 from aws_durable_execution_sdk_python_testing.checkpoint.transformer import (
-    OperationTransformer,
-)
-from aws_durable_execution_sdk_python_testing.checkpoint.validators.checkpoint import (
-    CheckpointValidator,
+    CheckpointRequestDispatcher,
 )
+from aws_durable_execution_sdk_python_testing.clock import RealClock
 from aws_durable_execution_sdk_python_testing.exceptions import (
     InvalidParameterValueException,
 )
-from aws_durable_execution_sdk_python_testing.observer import ExecutionNotifier
+from aws_durable_execution_sdk_python_testing.observer import apply_effects
 from aws_durable_execution_sdk_python_testing.token import CheckpointToken
 
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from aws_durable_execution_sdk_python.lambda_service import (
+        ErrorObject,
+        OperationUpdate,
+    )
+
+    from aws_durable_execution_sdk_python_testing.clock import Clock
     from aws_durable_execution_sdk_python_testing.execution import Execution
+    from aws_durable_execution_sdk_python_testing.observer import ExecutionObserver
     from aws_durable_execution_sdk_python_testing.scheduler import Scheduler
     from aws_durable_execution_sdk_python_testing.stores.base import ExecutionStore
 
 
+# Canonical invocation-page byte budget, shared with the Executor.
+DEFAULT_MAX_INVOCATION_PAGE_BYTES = 5 * 1024 * 1024
+
+
 class CheckpointProcessor:
-    """Handle OperationUpdate transformations and execution state updates."""
+    """In-process checkpoint flow used by ``InMemoryServiceClient``.
 
-    def __init__(self, store: ExecutionStore, scheduler: Scheduler):
-        self._store = store
-        self._scheduler = scheduler
-        self._notifier = ExecutionNotifier()
-        self._transformer = OperationTransformer()
+    Uses the same :class:`CheckpointRequestDispatcher` + paginator
+    primitives as ``Executor.checkpoint_execution`` so observable
+    behaviour is identical across the two entry points.
+    """
 
-    def add_execution_observer(self, observer) -> None:
+    def __init__(
+        self,
+        store: ExecutionStore,
+        scheduler: Scheduler,  # noqa: ARG002 — kept for backward-compatible signature
+        clock: Clock | None = None,
+    ):
+        self._store = store
+        self._observers: list[ExecutionObserver] = []
+        self._dispatcher = CheckpointRequestDispatcher()
+        # The runner's clock, shared with the Executor so stamping uses
+        # the same "now" as timer arming. Defaults to a real clock for
+        # standalone construction (e.g. direct unit tests).
+        self._clock: Clock = clock if clock is not None else RealClock()
+
+    def add_execution_observer(self, observer: ExecutionObserver) -> None:
         """Add observer for execution events."""
-        self._notifier.add_observer(observer)
+        self._observers.append(observer)
+
+    def set_chained_invoke_preflight(
+        self, preflight: Callable[[str], ErrorObject | None]
+    ) -> None:
+        """Resolve chained-invoke targets at checkpoint time with ``preflight``.
+
+        A target it fails comes back FAILED in the checkpoint response
+        instead of being dispatched (see ``ChainedInvokeProcessor``).
+        """
+        self._dispatcher = CheckpointRequestDispatcher(
+            chained_invoke_preflight=preflight
+        )
 
     def process_checkpoint(
         self,
         checkpoint_token: str,
         updates: list[OperationUpdate],
-        client_token: str | None,  # noqa: ARG002
+        client_token: str | None,
     ) -> CheckpointOutput:
-        """Process checkpoint updates and return result with updated execution state."""
-        # 1. Get current execution state
-        token: CheckpointToken = CheckpointToken.from_str(checkpoint_token)
+        """Apply ``updates`` and return the delta since the handler's
+        last observation.
+
+        Advances ``token_sequence`` exactly once; bumps
+        ``seq_counter`` once per accepted update via
+        ``touch_operation``; advances ``handler_seen_seq`` only for
+        operations actually returned in the response.
+        """
+        token = CheckpointToken.from_str(checkpoint_token)
         execution: Execution = self._store.load(token.execution_arn)
 
-        # 2. Validate checkpoint token
-        if execution.is_complete or token.token_sequence != execution.token_sequence:
-            msg: str = "Invalid checkpoint token"
-
+        # Idempotency first, before the token-sequence check.
+        cached = _maybe_replay_cached(execution, checkpoint_token, client_token)
+        if cached is not None:
+            return cached
+
+        if (
+            execution.is_complete
+            or token.token_sequence != execution.token_sequence
+            or token.invocation_id != execution.current_invocation_id
+        ):
+            msg = "Invalid checkpoint token"
             raise InvalidParameterValueException(msg)
 
-        # 3. Validate all updates, state transitions are valid, sizes etc.
-        CheckpointValidator.validate_input(updates, execution)
-
-        # 4. Transform OperationUpdate -> Operation and schedule future replays
-        updated_operations, all_updates = self._transformer.process_updates(
-            updates=updates,
-            current_operations=execution.operations,
-            notifier=self._notifier,
-            execution_arn=token.execution_arn,
+        now = self._clock.now()
+        result = CheckpointCore.apply(
+            execution,
+            checkpoint_token,
+            updates,
+            client_token,
+            self._dispatcher,
+            now,
         )
 
-        # 5. Generate a new checkpoint token and save updated operations
-        new_checkpoint_token = execution.get_new_checkpoint_token()
-        execution.operations = updated_operations
-        execution.updates.extend(all_updates)
         self._store.update(execution)
 
-        # 6. Return checkpoint result
+        for observer in self._observers:
+            apply_effects(result.effects, observer)
+
         return CheckpointOutput(
-            checkpoint_token=new_checkpoint_token,
+            checkpoint_token=result.checkpoint_token,
             new_execution_state=CheckpointUpdatedExecutionState(
-                operations=execution.get_navigable_operations(), next_marker=None
+                operations=result.operations,
+                next_marker=None,
             ),
         )
 
@@ -91,11 +144,37 @@ class CheckpointProcessor:
         next_marker: str,  # noqa: ARG002
         max_items: int = 1000,  # noqa: ARG002
     ) -> StateOutput:
-        """Get current execution state."""
-        token: CheckpointToken = CheckpointToken.from_str(checkpoint_token)
-        execution: Execution = self._store.load(token.execution_arn)
-
-        # TODO: paging when size or max
+        """Get current execution state.
+
+        Returns the full navigable operation list with a null marker.
+        Marker round-tripping and the invocation-state gate are
+        enforced on the HTTP path in :class:`Executor`.
+        """
+        token = CheckpointToken.from_str(checkpoint_token)
+        execution = self._store.load(token.execution_arn)
         return StateOutput(
-            operations=execution.get_navigable_operations(), next_marker=None
+            operations=execution.get_navigable_operations(),
+            next_marker=None,
         )
+
+
+def _maybe_replay_cached(
+    execution: Execution,
+    checkpoint_token: str,
+    client_token: str | None,
+) -> CheckpointOutput | None:
+    """Replay the cached response when the incoming
+    ``(client_token, checkpoint_token)`` matches the last checkpoint.
+
+    Returns ``None`` when there is no matching cached record.
+    """
+    cached = CheckpointCore.match_cached(execution, checkpoint_token, client_token)
+    if cached is None:
+        return None
+    return CheckpointOutput(
+        checkpoint_token=cached.outbound_checkpoint_token,
+        new_execution_state=CheckpointUpdatedExecutionState(
+            operations=list(cached.operations),
+            next_marker=cached.next_marker,
+        ),
+    )

aws_durable_execution_sdk_python_testing/checkpoint/processors/base.py

@@ -34,22 +34,42 @@ class OperationProcessor:
         current_op: Operation | None,
         notifier: ExecutionNotifier,
         execution_arn: str,
+        now: datetime.datetime,
     ) -> Operation | None:
-        """Process an operation update and return the transformed operation."""
+        """Process an operation update and return the transformed operation.
+
+        ``now`` is the checkpoint's single source of "now", resolved from
+        the execution's clock so timestamps advance with modeled time
+        under a skip clock and match wall-clock under a real clock.
+        """
         raise NotImplementedError
 
+    def stored_update(
+        self,
+        update: OperationUpdate,
+        current_op: Operation | None,  # noqa: ARG002
+        updated_op: Operation,  # noqa: ARG002
+    ) -> OperationUpdate:
+        """The update to keep in the execution's record for history.
+
+        Most processors keep the update as sent. A processor overrides
+        this when the service keeps less than it was sent.
+        """
+        return update
+
     def _get_start_time(
-        self, current_operation: Operation | None
+        self, current_operation: Operation | None, now: datetime.datetime
     ) -> datetime.datetime | None:
         start_time: datetime.datetime | None = (
-            current_operation.start_timestamp
-            if current_operation
-            else datetime.datetime.now(tz=datetime.UTC)
+            current_operation.start_timestamp if current_operation else now
         )
         return start_time
 
     def _get_end_time(
-        self, current_operation: Operation | None, status: OperationStatus
+        self,
+        current_operation: Operation | None,
+        status: OperationStatus,
+        now: datetime.datetime,
     ) -> datetime.datetime | None:
         """Get end timestamp for operation based on current state and status."""
         if current_operation and current_operation.end_timestamp:
@@ -61,7 +81,7 @@ class OperationProcessor:
             OperationStatus.TIMED_OUT,
             OperationStatus.STOPPED,
         }:
-            return datetime.datetime.now(tz=datetime.UTC)
+            return now
         return None
 
     def _create_execution_details(
@@ -151,11 +171,14 @@ class OperationProcessor:
         update: OperationUpdate,
         current_operation: Operation | None,
         status: OperationStatus,
+        now: datetime.datetime,
     ) -> Operation:
         """Transform OperationUpdate to Operation, always creating new Operation."""
-        start_time: datetime.datetime | None = self._get_start_time(current_operation)
+        start_time: datetime.datetime | None = self._get_start_time(
+            current_operation, now
+        )
         end_time: datetime.datetime | None = self._get_end_time(
-            current_operation, status
+            current_operation, status, now
         )
 
         execution_details = self._create_execution_details(update)
@@ -163,7 +186,7 @@ class OperationProcessor:
         step_details = self._create_step_details(update, current_operation)
         callback_details = self._create_callback_details(update)
         invoke_details = self._create_invoke_details(update)
-        wait_details = self._create_wait_details(update, current_operation)
+        wait_details = self._create_wait_details(update, current_operation, now)
 
         return Operation(
             operation_id=update.operation_id,
@@ -183,7 +206,10 @@ class OperationProcessor:
         )
 
     def _create_wait_details(
-        self, update: OperationUpdate, current_operation: Operation | None
+        self,
+        update: OperationUpdate,
+        current_operation: Operation | None,
+        now: datetime.datetime,
     ) -> WaitDetails | None:
         """Create WaitDetails from OperationUpdate."""
         if update.operation_type == OperationType.WAIT and update.wait_options:
@@ -192,8 +218,8 @@ class OperationProcessor:
                     current_operation.wait_details.scheduled_end_timestamp
                 )
             else:
-                scheduled_end_timestamp = datetime.datetime.now(
-                    tz=datetime.UTC
-                ) + timedelta(seconds=update.wait_options.wait_seconds)
+                scheduled_end_timestamp = now + timedelta(
+                    seconds=update.wait_options.wait_seconds
+                )
             return WaitDetails(scheduled_end_timestamp=scheduled_end_timestamp)
         return None

aws_durable_execution_sdk_python_testing/checkpoint/processors/callback.py

@@ -35,6 +35,7 @@ class CallbackProcessor(OperationProcessor):
         current_op: Operation | None,
         notifier: ExecutionNotifier,  # noqa: ARG002
         execution_arn: str,  # noqa: ARG002
+        now: datetime.datetime,
     ) -> Operation:
         """Process CALLBACK operation update with scheduler integration for activities."""
         match update.action:
@@ -58,10 +59,12 @@ class CallbackProcessor(OperationProcessor):
 
                 status: OperationStatus = OperationStatus.STARTED
 
-                start_time: datetime.datetime | None = self._get_start_time(current_op)
+                start_time: datetime.datetime | None = self._get_start_time(
+                    current_op, now
+                )
 
                 end_time: datetime.datetime | None = self._get_end_time(
-                    current_op, status
+                    current_op, status, now
                 )
 
                 operation: Operation = Operation(

aws_durable_execution_sdk_python_testing/checkpoint/processors/context.py

@@ -20,6 +20,8 @@ from aws_durable_execution_sdk_python_testing.exceptions import (
 
 
 if TYPE_CHECKING:
+    from datetime import datetime
+
     from aws_durable_execution_sdk_python_testing.observer import ExecutionNotifier
 
 
@@ -32,6 +34,7 @@ class ContextProcessor(OperationProcessor):
         current_op: Operation | None,
         notifier: ExecutionNotifier,  # noqa: ARG002
         execution_arn: str,  # noqa: ARG002
+        now: datetime,
     ) -> Operation:
         """Process CONTEXT operation update for context state transitions."""
         match update.action:
@@ -41,18 +44,21 @@ class ContextProcessor(OperationProcessor):
                     update=update,
                     current_operation=current_op,
                     status=OperationStatus.STARTED,
+                    now=now,
                 )
             case OperationAction.SUCCEED:
                 return self._translate_update_to_operation(
                     update=update,
                     current_operation=current_op,
                     status=OperationStatus.SUCCEEDED,
+                    now=now,
                 )
             case OperationAction.FAIL:
                 return self._translate_update_to_operation(
                     update=update,
                     current_operation=current_op,
                     status=OperationStatus.FAILED,
+                    now=now,
                 )
             case _:
                 msg: str = "Invalid action for CONTEXT operation."

aws_durable_execution_sdk_python_testing/checkpoint/processors/execution.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 from typing import TYPE_CHECKING
 
 from aws_durable_execution_sdk_python.lambda_service import (
-    ErrorObject,
     Operation,
     OperationAction,
     OperationUpdate,
@@ -17,6 +16,8 @@ from aws_durable_execution_sdk_python_testing.checkpoint.processors.base import
 
 
 if TYPE_CHECKING:
+    from datetime import datetime
+
     from aws_durable_execution_sdk_python_testing.observer import ExecutionNotifier
 
 
@@ -29,6 +30,7 @@ class ExecutionProcessor(OperationProcessor):
         current_op: Operation | None,  # noqa: ARG002
         notifier: ExecutionNotifier,
         execution_arn: str,
+        now: datetime,  # noqa: ARG002
     ) -> Operation | None:
         """Process EXECUTION operation update for workflow completion/failure."""
         match update.action:
@@ -38,15 +40,8 @@ class ExecutionProcessor(OperationProcessor):
                 )
             case _:
                 # intentional. actual service will fail any EXECUTION update that is not SUCCEED.
-                error = (
-                    update.error
-                    if update.error
-                    else ErrorObject.from_message(
-                        "There is no error details but EXECUTION checkpoint action is not SUCCEED."
-                    )
-                )
                 # All EXECUTION failures go through normal fail path
                 # Timeout/Stop status is set by executor based on the operation that caused it
-                notifier.notify_failed(execution_arn=execution_arn, error=error)
+                notifier.notify_failed(execution_arn=execution_arn, error=update.error)
         # TODO: Svc doesn't actually create checkpoint for EXECUTION. might have to for localrunner though.
         return None