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

aws_durable_execution_sdk_python_testing/checkpoint/processors/invoke.py

@@ -0,0 +1,139 @@
+"""Chained invoke operation processor for handling CHAINED_INVOKE operation updates."""
+
+from __future__ import annotations
+
+import datetime
+from dataclasses import replace
+from typing import TYPE_CHECKING
+
+from aws_durable_execution_sdk_python.lambda_service import (
+    ChainedInvokeDetails,
+    ErrorObject,
+    Operation,
+    OperationAction,
+    OperationStatus,
+    OperationUpdate,
+)
+
+from aws_durable_execution_sdk_python_testing.checkpoint.processors.base import (
+    OperationProcessor,
+)
+from aws_durable_execution_sdk_python_testing.exceptions import (
+    InvalidParameterValueException,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from aws_durable_execution_sdk_python_testing.observer import ExecutionNotifier
+
+
+class ChainedInvokeProcessor(OperationProcessor):
+    """Processes CHAINED_INVOKE operation updates.
+
+    The service resolves the target before it schedules anything. Its
+    checkpoint response then carries the operation STARTED, because the
+    scheduling event is written as part of completing the checkpoint,
+    or FAILED when the target could not be resolved, so the handler
+    sees such a failure in the response to its own checkpoint and
+    raises without suspending. This processor does the same with
+    ``preflight``: a START update whose target fails the preflight
+    comes back FAILED with that error; otherwise it comes back STARTED
+    and raises a :class:`ChainedInvokeStarted` effect for the caller to
+    dispatch the target once the checkpoint write completes. The
+    operation holds STARTED until its terminal transition, so a handler
+    observing it sees exactly one status change; the target's terminal
+    state completes it. Completion is never a handler checkpoint, so
+    START is the only accepted action.
+    """
+
+    def __init__(self, preflight: Callable[[str], ErrorObject | None] | None = None):
+        """``preflight`` resolves a target at checkpoint time; without it
+        every target is dispatched and any failure arrives later."""
+        self._preflight = preflight
+
+    def process(
+        self,
+        update: OperationUpdate,
+        current_op: Operation | None,
+        notifier: ExecutionNotifier,
+        execution_arn: str,
+        now: datetime.datetime,
+    ) -> Operation:
+        """Process CHAINED_INVOKE operation update."""
+        match update.action:
+            case OperationAction.START:
+                options = update.chained_invoke_options
+                if options is None:
+                    msg_options_required: str = (
+                        "Update for CHAINED_INVOKE operation requires "
+                        "ChainedInvokeOptions."
+                    )
+                    raise InvalidParameterValueException(msg_options_required)
+
+                start_timestamp: datetime.datetime | None = self._get_start_time(
+                    current_op, now
+                )
+                error: ErrorObject | None = (
+                    self._preflight(options.function_name)
+                    if self._preflight is not None
+                    else None
+                )
+                if error is not None:
+                    return Operation(
+                        operation_id=update.operation_id,
+                        parent_id=update.parent_id,
+                        name=update.name,
+                        start_timestamp=start_timestamp,
+                        end_timestamp=now,
+                        operation_type=update.operation_type,
+                        status=OperationStatus.FAILED,
+                        sub_type=update.sub_type,
+                        chained_invoke_details=ChainedInvokeDetails(
+                            result=None, error=error
+                        ),
+                    )
+
+                operation: Operation = Operation(
+                    operation_id=update.operation_id,
+                    parent_id=update.parent_id,
+                    name=update.name,
+                    start_timestamp=start_timestamp,
+                    end_timestamp=None,
+                    operation_type=update.operation_type,
+                    status=OperationStatus.STARTED,
+                    sub_type=update.sub_type,
+                    chained_invoke_details=ChainedInvokeDetails(
+                        result=None, error=None
+                    ),
+                )
+
+                notifier.notify_chained_invoke_started(
+                    execution_arn=execution_arn,
+                    operation_id=update.operation_id,
+                    function_name=options.function_name,
+                    tenant_id=options.tenant_id,
+                    payload=update.payload,
+                )
+                return operation
+            case _:
+                msg_invalid_action: str = "Invalid action for CHAINED_INVOKE operation."
+                raise InvalidParameterValueException(msg_invalid_action)
+
+    def stored_update(
+        self,
+        update: OperationUpdate,
+        current_op: Operation | None,
+        updated_op: Operation,
+    ) -> OperationUpdate:
+        """Keep no input for a chained invoke that failed before starting.
+
+        The service stores no input payload when the target cannot be
+        resolved: it records a payload size of zero and an empty payload
+        location, and its history shows the function name and the error
+        without the input. So the runner stores the update without its
+        payload, which also keeps it out of the operation's size.
+        """
+        if current_op is None and updated_op.status is OperationStatus.FAILED:
+            return replace(update, payload=None)
+        return update

aws_durable_execution_sdk_python_testing/checkpoint/processors/step.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from datetime import UTC, datetime, timedelta
+from datetime import datetime, timedelta
 from typing import TYPE_CHECKING
 
 from aws_durable_execution_sdk_python.lambda_service import (
@@ -34,6 +34,7 @@ class StepProcessor(OperationProcessor):
         current_op: Operation | None,
         notifier: ExecutionNotifier,
         execution_arn: str,
+        now: datetime,
     ) -> Operation:
         """Process STEP operation update with scheduler integration for retries."""
         match update.action:
@@ -42,6 +43,7 @@ class StepProcessor(OperationProcessor):
                     update=update,
                     current_operation=current_op,
                     status=OperationStatus.STARTED,
+                    now=now,
                 )
             case OperationAction.RETRY:
                 # set Status=PENDING, next attempt time, attempt count + 1
@@ -50,7 +52,7 @@ class StepProcessor(OperationProcessor):
                     if update.step_options
                     else 0
                 )
-                next_attempt_time = datetime.now(UTC) + timedelta(seconds=delay)
+                next_attempt_time = now + timedelta(seconds=delay)
 
                 # Build new step_details with incremented attempt
                 current_attempt = (
@@ -61,16 +63,8 @@ class StepProcessor(OperationProcessor):
                 new_step_details = StepDetails(
                     attempt=current_attempt + 1,
                     next_attempt_timestamp=next_attempt_time,
-                    result=(
-                        current_op.step_details.result
-                        if current_op and current_op.step_details
-                        else None
-                    ),
-                    error=(
-                        current_op.step_details.error
-                        if current_op and current_op.step_details
-                        else None
-                    ),
+                    result=update.payload,
+                    error=update.error,
                 )
 
                 # Create new operation with updated step_details
@@ -80,9 +74,7 @@ class StepProcessor(OperationProcessor):
                     status=OperationStatus.PENDING,
                     parent_id=update.parent_id,
                     name=update.name,
-                    start_timestamp=(
-                        current_op.start_timestamp if current_op else datetime.now(UTC)
-                    ),
+                    start_timestamp=(current_op.start_timestamp if current_op else now),
                     end_timestamp=None,
                     sub_type=update.sub_type,
                     execution_details=current_op.execution_details
@@ -99,24 +91,26 @@ class StepProcessor(OperationProcessor):
                     else None,
                 )
 
-                # Schedule step retry timer to fire after delay
-                notifier.notify_step_retry_scheduled(
-                    execution_arn=execution_arn,
-                    operation_id=update.operation_id,
-                    delay=delay,
-                )
+                # Timer scheduling is handled centrally by
+                # Executor._schedule_earliest_pending. The
+                # processor only needs to set
+                # step_details.next_attempt_timestamp; the executor
+                # finds the earliest pending wake across all ops and
+                # arms a single timer.
                 return retry_operation
             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 STEP operation."

aws_durable_execution_sdk_python_testing/checkpoint/processors/wait.py

@@ -2,9 +2,7 @@
 
 from __future__ import annotations
 
-import logging
-import os
-from datetime import UTC, datetime, timedelta
+from datetime import datetime, timedelta
 from typing import TYPE_CHECKING
 
 from aws_durable_execution_sdk_python.lambda_service import (
@@ -36,6 +34,7 @@ class WaitProcessor(OperationProcessor):
         current_op: Operation | None,
         notifier: ExecutionNotifier,
         execution_arn: str,
+        now: datetime,
     ) -> Operation:
         """Process WAIT operation update with scheduler integration for timers."""
         match update.action:
@@ -43,13 +42,8 @@ class WaitProcessor(OperationProcessor):
                 wait_seconds = (
                     update.wait_options.wait_seconds if update.wait_options else 0
                 )
-                time_scale = float(os.getenv("DURABLE_EXECUTION_TIME_SCALE", "1.0"))
-                logging.info("Using DURABLE_EXECUTION_TIME_SCALE: %f", time_scale)
-                scaled_wait_seconds = wait_seconds * time_scale
 
-                scheduled_end_timestamp = datetime.now(UTC) + timedelta(
-                    seconds=scaled_wait_seconds
-                )
+                scheduled_end_timestamp = now + timedelta(seconds=wait_seconds)
 
                 # Create WaitDetails with scheduled timestamp
                 wait_details = WaitDetails(
@@ -63,7 +57,7 @@ class WaitProcessor(OperationProcessor):
                     status=OperationStatus.STARTED,
                     parent_id=update.parent_id,
                     name=update.name,
-                    start_timestamp=datetime.now(UTC),
+                    start_timestamp=now,
                     end_timestamp=None,
                     sub_type=update.sub_type,
                     execution_details=None,
@@ -74,12 +68,11 @@ class WaitProcessor(OperationProcessor):
                     chained_invoke_details=None,
                 )
 
-                # Schedule wait timer to complete after delay
-                notifier.notify_wait_timer_scheduled(
-                    execution_arn=execution_arn,
-                    operation_id=update.operation_id,
-                    delay=scaled_wait_seconds,
-                )
+                # Timer scheduling is handled centrally by
+                # Executor._schedule_earliest_pending. The processor
+                # only needs to set wait_details.scheduled_end_timestamp;
+                # the executor finds the earliest pending wake across
+                # all ops and arms a single timer.
                 return wait_operation
             case OperationAction.CANCEL:
                 # TODO: need to cancel the WAIT in the executor
@@ -88,6 +81,7 @@ class WaitProcessor(OperationProcessor):
                     update=update,
                     current_operation=current_op,
                     status=OperationStatus.CANCELLED,
+                    now=now,
                 )
             case _:
                 msg: str = "Invalid action for WAIT operation."

aws_durable_execution_sdk_python_testing/checkpoint/transformer.py

@@ -1,13 +1,17 @@
-"""Operation transformer for converting OperationUpdates to Operations."""
+"""Checkpoint request dispatcher.
+
+Routes each ``OperationUpdate`` in a checkpoint to its type-specific
+processor (step, wait, callback, context, execution) and applies the
+result to the execution.
+"""
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+import json
+from typing import TYPE_CHECKING, ClassVar
 
 from aws_durable_execution_sdk_python.lambda_service import (
-    Operation,
     OperationType,
-    OperationUpdate,
 )
 
 from aws_durable_execution_sdk_python_testing.checkpoint.processors.callback import (
@@ -19,6 +23,9 @@ from aws_durable_execution_sdk_python_testing.checkpoint.processors.context impo
 from aws_durable_execution_sdk_python_testing.checkpoint.processors.execution import (
     ExecutionProcessor,
 )
+from aws_durable_execution_sdk_python_testing.checkpoint.processors.invoke import (
+    ChainedInvokeProcessor,
+)
 from aws_durable_execution_sdk_python_testing.checkpoint.processors.step import (
     StepProcessor,
 )
@@ -28,20 +35,37 @@ from aws_durable_execution_sdk_python_testing.checkpoint.processors.wait import
 from aws_durable_execution_sdk_python_testing.exceptions import (
     InvalidParameterValueException,
 )
+from aws_durable_execution_sdk_python_testing.observer import ExecutionNotifier
 
 
 if TYPE_CHECKING:
-    from collections.abc import MutableMapping
+    from collections.abc import Callable, MutableMapping
+    from datetime import datetime
 
+    from aws_durable_execution_sdk_python.lambda_service import (
+        ErrorObject,
+        OperationUpdate,
+    )
+
+    from aws_durable_execution_sdk_python_testing.checkpoint.effects import (
+        CheckpointEffect,
+    )
     from aws_durable_execution_sdk_python_testing.checkpoint.processors.base import (
         OperationProcessor,
     )
+    from aws_durable_execution_sdk_python_testing.execution import Execution
 
-from typing import ClassVar
 
+class CheckpointRequestDispatcher:
+    """Apply a batch of ``OperationUpdate``\\ s to an :class:`Execution`.
 
-class OperationTransformer:
-    """Transforms OperationUpdates to Operations while maintaining order and triggering scheduler actions."""
+    Dispatches each update to the per-type processor (step, wait,
+    callback, context, execution), upserts the resulting operation
+    into ``execution.operations``, records per-op payload size in the
+    ``execution.operation_size_bytes`` sidecar dict (``Operation`` is
+    immutable, so size is tracked out-of-band), and calls the supplied
+    ``touch`` callback once per accepted update.
+    """
 
     _DEFAULT_PROCESSORS: ClassVar[dict[OperationType, OperationProcessor]] = {
         OperationType.STEP: StepProcessor(),
@@ -49,56 +73,129 @@ class OperationTransformer:
         OperationType.CONTEXT: ContextProcessor(),
         OperationType.CALLBACK: CallbackProcessor(),
         OperationType.EXECUTION: ExecutionProcessor(),
+        OperationType.CHAINED_INVOKE: ChainedInvokeProcessor(),
     }
 
     def __init__(
         self,
         processors: MutableMapping[OperationType, OperationProcessor] | None = None,
+        *,
+        chained_invoke_preflight: Callable[[str], ErrorObject | None] | None = None,
     ):
-        self.processors = processors if processors else self._DEFAULT_PROCESSORS
-
-    def process_updates(
+        """``chained_invoke_preflight`` resolves a chained-invoke target at
+        checkpoint time (see :class:`ChainedInvokeProcessor`); it replaces
+        the default CHAINED_INVOKE processor and is ignored when explicit
+        ``processors`` are given."""
+        if processors:
+            self.processors = processors
+        elif chained_invoke_preflight is not None:
+            self.processors = dict(self._DEFAULT_PROCESSORS)
+            self.processors[OperationType.CHAINED_INVOKE] = ChainedInvokeProcessor(
+                chained_invoke_preflight
+            )
+        else:
+            self.processors = self._DEFAULT_PROCESSORS
+
+    def apply_updates(
         self,
+        execution: Execution,
         updates: list[OperationUpdate],
-        current_operations: list[Operation],
-        notifier,
-        execution_arn: str,
-    ) -> tuple[list[Operation], list[OperationUpdate]]:
-        """Transform updates maintaining operation order and return (operations, updates)."""
-        op_map = {op.operation_id: op for op in current_operations}
-
-        # Start with copy of current operations list
-        result_operations = current_operations.copy()
+        client_token: str | None,  # noqa: ARG002 — reserved for future idempotency diagnostics
+        touch: Callable[[str], None],
+        now: datetime,
+    ) -> list[CheckpointEffect]:
+        """Apply ``updates`` to ``execution`` in place.
+
+        Callers are responsible for running
+        :class:`CheckpointValidator.validate_input` before calling this
+        method. The dispatcher does not re-validate — it assumes each
+        update is well-formed so that per-type processors can focus on
+        state transitions.
+
+        Each accepted update:
+
+        * is dispatched to the per-type processor via
+          ``execution.operations`` upsert semantics (existing op with
+          matching ``operation_id`` is replaced in place; new op is
+          appended).
+        * records a payload size estimate in
+          ``execution.operation_size_bytes[op_id]`` for later paging.
+        * triggers ``touch(op_id)`` exactly once, which (in the
+          production caller) bumps :attr:`Execution.seq_counter` and
+          sets ``operation_last_touched_seq[op_id]``.
+
+        Returns the lifecycle effects raised by the updates (completion,
+        failure, callback creation) for the caller to apply once the
+        write is done. Response construction is the responsibility of the
+        checkpoint orchestrator.
+        """
+        collector = ExecutionNotifier()
+        op_map = {op.operation_id: op for op in execution.operations}
+        stored_updates: list[OperationUpdate] = []
 
         for update in updates:
             processor = self.processors.get(update.operation_type)
-            if processor:
-                current_op = op_map.get(update.operation_id)
-                updated_op = processor.process(
-                    update=update,
-                    current_op=current_op,
-                    notifier=notifier,
-                    execution_arn=execution_arn,
-                )
-
-                if updated_op is not None:
-                    if update.operation_id in op_map:
-                        # Update existing operation in-place
-                        for i, op in enumerate(result_operations):  # pragma: no branch
-                            # no branch coverage because result_operation empty not reachable here
-                            if op.operation_id == update.operation_id:
-                                result_operations[i] = updated_op
-                                break
-                    else:
-                        # Append new operation to end
-                        result_operations.append(updated_op)
-
-                    # Update map for future lookups
-                    op_map[update.operation_id] = updated_op
-            else:
-                msg: str = (
-                    f"Checkpoint for {update.operation_type} is not implemented yet."
-                )
+            if processor is None:
+                msg = f"Checkpoint for {update.operation_type} is not implemented yet."
                 raise InvalidParameterValueException(msg)
 
-        return result_operations, updates
+            current_op = op_map.get(update.operation_id)
+            updated_op = processor.process(
+                update=update,
+                current_op=current_op,
+                notifier=collector,
+                execution_arn=execution.durable_execution_arn,
+                now=now,
+            )
+            if updated_op is None:
+                stored_updates.append(update)
+                continue
+
+            stored = processor.stored_update(update, current_op, updated_op)
+            stored_updates.append(stored)
+            if update.operation_id in op_map:
+                for i, op in enumerate(execution.operations):  # pragma: no branch
+                    if op.operation_id == update.operation_id:
+                        execution.operations[i] = updated_op
+                        break
+            else:
+                execution.operations.append(updated_op)
+
+            op_map[update.operation_id] = updated_op
+            execution.operation_size_bytes[update.operation_id] = (
+                _estimate_payload_size(stored)
+            )
+            touch(update.operation_id)
+
+        execution.updates.extend(stored_updates)
+        execution.update_timestamps.extend(now for _ in stored_updates)
+
+        return collector.effects
+
+
+def _estimate_payload_size(update: OperationUpdate) -> int:
+    """Estimate the on-wire size of an ``OperationUpdate``'s payload.
+
+    Approximate — paging decisions need this to be consistent and
+    reproducible, not exact. Sums JSON-ish lengths of ``payload`` and
+    ``error`` fields when present.
+    """
+    size = 0
+    payload = update.payload
+    if payload is not None:
+        size += _byte_length(payload)
+    error = update.error
+    if error is not None:
+        try:
+            size += len(json.dumps(error.to_dict()))
+        except (AttributeError, TypeError):  # pragma: no cover — defensive
+            size += len(str(error))
+    return size
+
+
+def _byte_length(payload: object) -> int:
+    if isinstance(payload, bytes):
+        return len(payload)
+    if isinstance(payload, str):
+        return len(payload.encode())
+    return len(str(payload))

aws_durable_execution_sdk_python_testing/checkpoint/validators/checkpoint.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import json
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, ClassVar
 
 from aws_durable_execution_sdk_python.lambda_service import (
     OperationAction,
@@ -13,24 +13,33 @@ from aws_durable_execution_sdk_python.lambda_service import (
 
 from aws_durable_execution_sdk_python_testing.checkpoint.validators.operations.callback import (
     CallbackOperationValidator,
+    VALID_ACTIONS_FOR_CALLBACK,
 )
 from aws_durable_execution_sdk_python_testing.checkpoint.validators.operations.context import (
     ContextOperationValidator,
+    VALID_ACTIONS_FOR_CONTEXT,
 )
 from aws_durable_execution_sdk_python_testing.checkpoint.validators.operations.execution import (
     ExecutionOperationValidator,
+    VALID_ACTIONS_FOR_EXECUTION,
 )
 from aws_durable_execution_sdk_python_testing.checkpoint.validators.operations.invoke import (
     ChainedInvokeOperationValidator,
+    VALID_ACTIONS_FOR_INVOKE,
 )
 from aws_durable_execution_sdk_python_testing.checkpoint.validators.operations.step import (
     StepOperationValidator,
+    VALID_ACTIONS_FOR_STEP,
 )
 from aws_durable_execution_sdk_python_testing.checkpoint.validators.operations.wait import (
     WaitOperationValidator,
+    VALID_ACTIONS_FOR_WAIT,
 )
-from aws_durable_execution_sdk_python_testing.checkpoint.validators.transitions import (
-    ValidActionsByOperationTypeValidator,
+from aws_durable_execution_sdk_python_testing.child_dispatcher import (
+    CHAINED_INVOKE_INPUT_TOO_LARGE_MESSAGE,
+    CHILD_EXECUTION_OUTPUT_TOO_LARGE_MESSAGE,
+    MAX_CHAINED_INVOKE_PAYLOAD_BYTES,
+    payload_size_bytes,
 )
 from aws_durable_execution_sdk_python_testing.exceptions import (
     InvalidParameterValueException,
@@ -48,6 +57,17 @@ MAX_ERROR_PAYLOAD_SIZE_BYTES = 32768
 class CheckpointValidator:
     """Validates checkpoint input based on current state."""
 
+    _VALID_ACTIONS_BY_OPERATION_TYPE: ClassVar[
+        dict[OperationType, frozenset[OperationAction]]
+    ] = {
+        OperationType.STEP: VALID_ACTIONS_FOR_STEP,
+        OperationType.CONTEXT: VALID_ACTIONS_FOR_CONTEXT,
+        OperationType.WAIT: VALID_ACTIONS_FOR_WAIT,
+        OperationType.CALLBACK: VALID_ACTIONS_FOR_CALLBACK,
+        OperationType.CHAINED_INVOKE: VALID_ACTIONS_FOR_INVOKE,
+        OperationType.EXECUTION: VALID_ACTIONS_FOR_EXECUTION,
+    }
+
     @staticmethod
     def validate_input(updates: list[OperationUpdate], execution: Execution) -> None:
         """Perform validation on the given input based on the current state."""
@@ -85,20 +105,55 @@ class CheckpointValidator:
     ) -> None:
         """Validate a single operation update."""
         CheckpointValidator._validate_inconsistent_operation_metadata(update, execution)
-        CheckpointValidator._validate_payload_sizes(update)
-        ValidActionsByOperationTypeValidator.validate(
+        CheckpointValidator._validate_payload_sizes(update, execution)
+        CheckpointValidator._validate_valid_action_for_type(
             update.operation_type, update.action
         )
         CheckpointValidator._validate_operation_status_transition(update, execution)
 
     @staticmethod
-    def _validate_payload_sizes(update: OperationUpdate) -> None:
-        """Validate that operation payload sizes are not too large."""
+    def _validate_payload_sizes(update: OperationUpdate, execution: Execution) -> None:
+        """Validate that operation payload sizes are not too large.
+
+        Besides the error-object bound, two chained-invoke bounds apply:
+        the input a parent sends to a target, and the result a child
+        execution returns to its parent, are each capped at 1 MiB.
+        """
         if update.error is not None:
             payload = json.dumps(update.error.to_dict())
             if len(payload) > MAX_ERROR_PAYLOAD_SIZE_BYTES:
                 msg: str = f"Error object size must be less than {MAX_ERROR_PAYLOAD_SIZE_BYTES} bytes."
                 raise InvalidParameterValueException(msg)
+        if update.payload is None:
+            return
+        if (
+            update.operation_type == OperationType.CHAINED_INVOKE
+            and payload_size_bytes(update.payload) > MAX_CHAINED_INVOKE_PAYLOAD_BYTES
+        ):
+            raise InvalidParameterValueException(CHAINED_INVOKE_INPUT_TOO_LARGE_MESSAGE)
+        if (
+            update.operation_type == OperationType.EXECUTION
+            and execution.parent_execution_arn is not None
+            and payload_size_bytes(update.payload) > MAX_CHAINED_INVOKE_PAYLOAD_BYTES
+        ):
+            raise InvalidParameterValueException(
+                CHILD_EXECUTION_OUTPUT_TOO_LARGE_MESSAGE
+            )
+
+    @staticmethod
+    def _validate_valid_action_for_type(
+        operation_type: OperationType, action: OperationAction
+    ) -> None:
+        """Reject the checkpoint if the action is not valid for the type."""
+        valid_actions: frozenset[OperationAction] | None = (
+            CheckpointValidator._VALID_ACTIONS_BY_OPERATION_TYPE.get(operation_type)
+        )
+        if valid_actions is None:
+            msg_unknown_op: str = "Unknown operation type."
+            raise InvalidParameterValueException(msg_unknown_op)
+        if action not in valid_actions:
+            msg_invalid_action: str = "Invalid action for the given operation type."
+            raise InvalidParameterValueException(msg_invalid_action)
 
     @staticmethod
     def _validate_operation_status_transition(
@@ -121,7 +176,9 @@ class CheckpointValidator:
             case OperationType.CALLBACK:
                 CallbackOperationValidator.validate(current_state, update)
             case OperationType.CHAINED_INVOKE:
-                ChainedInvokeOperationValidator.validate(current_state, update)
+                ChainedInvokeOperationValidator.validate(
+                    current_state, update, execution
+                )
             case OperationType.EXECUTION:
                 ExecutionOperationValidator.validate(update)
             case _:  # pragma: no cover