feat(plugin): surface execution input and result

packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py

@@ -298,6 +298,7 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]:
                     else None
                 ),
                 is_first_invocation=not has_prior_operations,
+                execution_input=input_event,
             )
             # Thread 1: Run background checkpoint processing
             executor.submit(execution_state.checkpoint_batches_forever)

packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import contextlib
+import copy
 import datetime
 import functools
 import logging
@@ -175,6 +176,20 @@ class InvocationInfo:
     execution_arn: str | None
     is_first_invocation: bool
     execution_start_time: datetime.datetime | None = None
+    execution_input: Any = field(
+        default=None,
+        kw_only=True,
+        metadata={"experimental": True},
+    )
+    """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``.
+
+    Defaults to ``None`` only when the field is not populated (a hook info built
+    without it); ``durable_execution()`` always populates it with the
+    deserialized input payload, which is ``{}`` when the payload is empty.
+    """
 
 
 @dataclass(frozen=True)
@@ -190,6 +205,17 @@ class InvocationEndInfo(InvocationInfo):
         metadata={"experimental": True},
     )
     """EXPERIMENTAL: The invocation error, when available."""
+    execution_result: str | None = field(
+        default=None,
+        kw_only=True,
+        metadata={"experimental": True},
+    )
+    """EXPERIMENTAL: The serialized execution result, when available.
+
+    A JSON string, or ``""`` when the result was checkpointed out-of-band for a
+    large payload. Mirrors the JS SDK's ``InvocationEndInfo.executionResult``.
+    ``None`` on failure or suspend.
+    """
 
     @classmethod
     def from_durable_execution_invocation_output(
@@ -202,8 +228,10 @@ def from_durable_execution_invocation_output(
             execution_arn=invocation_start_info.execution_arn,
             is_first_invocation=invocation_start_info.is_first_invocation,
             execution_start_time=invocation_start_info.execution_start_time,
+            execution_input=invocation_start_info.execution_input,
             status=output.status,
             error=output.error,
+            execution_result=output.result,
         )
 
 
@@ -349,16 +377,47 @@ def on_invocation_start(
         is_first_invocation: bool,
         execution_start_time: datetime.datetime | None,
         lambda_context: LambdaContext | None,
+        execution_input: Any = None,
     ) -> None:
         aws_request_id = lambda_context.aws_request_id if lambda_context else None
         self._invocation_status = InvocationStartInfo(
             execution_arn=execution_arn,
             request_id=aws_request_id,
             is_first_invocation=is_first_invocation,
             execution_start_time=execution_start_time,
+            execution_input=self._snapshot_execution_input(execution_input),
         )
         self.execute_plugins(self._invocation_status, sync=True)
 
+    def _snapshot_execution_input(self, execution_input: Any) -> Any:
+        """Deep-copy the execution input so the plugin view is isolated.
+
+        ``durable_execution()`` hands the same mutable object to the user handler
+        and to this hook. Without a copy the aliasing runs both ways: a plugin
+        mutating ``info.execution_input`` would change the handler's event and so
+        alter execution behaviour, and a handler mutating its event would change
+        what this frozen info -- and the invocation-end info derived from it --
+        reports afterwards.
+
+        The copy is eager rather than deferred: the handler starts running
+        immediately after this hook, so a lazily-taken snapshot could already
+        have observed the handler's mutations. It is skipped when no plugins are
+        registered, so non-plugin executions pay nothing.
+
+        The snapshot is shared by all plugins for this invocation; plugins should
+        still treat it as read-only with respect to each other.
+        """
+        if not self._plugins or execution_input is None:
+            return execution_input
+        try:
+            return copy.deepcopy(execution_input)
+        except Exception:
+            # Preserve handler isolation if a snapshot cannot be created.
+            logger.exception(
+                "Failed to copy execution input for plugins; omitting plugin input"
+            )
+            return None
+
     def on_invocation_end(
         self,
         output: "DurableExecutionInvocationOutput",

packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_payload_int_test.py

@@ -0,0 +1,256 @@
+"""Integration tests for the plugin invocation payload surfaces.
+
+Exercises `InvocationInfo.execution_input` / `InvocationEndInfo.execution_result`
+through complete `durable_execution()` invocations -- across the decorator, the
+plugin executor, and the invocation hooks -- including a suspend/replay pair
+where the suspending invocation has no result and the replay carries the
+terminal one, and the isolation guarantee between the plugin view and the user
+handler's event.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+from unittest.mock import Mock, patch
+
+from aws_durable_execution_sdk_python.config import Duration
+from aws_durable_execution_sdk_python.context import DurableContext
+from aws_durable_execution_sdk_python.execution import (
+    InvocationStatus,
+    durable_execution,
+)
+from aws_durable_execution_sdk_python.lambda_service import (
+    CheckpointOutput,
+    CheckpointUpdatedExecutionState,
+    Operation,
+    OperationStatus,
+    OperationType,
+)
+from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin
+from tests.test_helpers import operation_id_sequence
+
+
+class _PayloadRecordingPlugin(DurableInstrumentationPlugin):
+    """Records the payload surfaces seen on each invocation hook."""
+
+    def __init__(self) -> None:
+        self.starts: list[Any] = []
+        self.ends: list[tuple[str, Any, str | None]] = []
+
+    def on_invocation_start(self, info) -> None:
+        self.starts.append(info.execution_input)
+
+    def on_invocation_end(self, info) -> None:
+        self.ends.append(
+            (info.status.value, info.execution_input, info.execution_result)
+        )
+
+
+def _lambda_context() -> Mock:
+    ctx = Mock()
+    ctx.aws_request_id = "test-request-id"
+    ctx.client_context = None
+    ctx.identity = None
+    ctx._epoch_deadline_time_in_ms = 0  # noqa: SLF001
+    ctx.invoked_function_arn = "test-arn"
+    ctx.tenant_id = None
+    return ctx
+
+
+def _event(input_payload: str, extra_operations: list[dict] | None = None) -> dict:
+    """Build an invocation event carrying the given execution input payload."""
+    execution_operation = {
+        "Id": "execution-1",
+        "Type": "EXECUTION",
+        "Status": "STARTED",
+        "ExecutionDetails": {"InputPayload": input_payload},
+    }
+    return {
+        "DurableExecutionArn": "test-arn/execution-1",
+        "CheckpointToken": "test-token",
+        "InitialExecutionState": {
+            "Operations": [execution_operation, *(extra_operations or [])],
+            "NextMarker": "",
+        },
+        "LocalRunner": True,
+    }
+
+
+def _tracking_checkpoint(initial_operations: list[Operation] | None = None):
+    """Checkpoint mock that accumulates operations, as the service would.
+
+    A stub returning an empty execution state is not enough for suspending
+    paths: after the WAIT START is checkpointed the SDK re-reads the operation
+    from the returned state, so it must be present.
+    """
+    operations: list[Operation] = list(initial_operations or [])
+
+    def mock_checkpoint(
+        durable_execution_arn,  # noqa: ARG001
+        checkpoint_token,  # noqa: ARG001
+        updates,
+        client_token="token",  # noqa: S107, ARG001
+    ) -> CheckpointOutput:
+        for update in updates:
+            operations.append(
+                Operation(
+                    operation_id=update.operation_id,
+                    operation_type=update.operation_type,
+                    status=OperationStatus.STARTED,
+                    parent_id=update.parent_id,
+                )
+            )
+        return CheckpointOutput(
+            checkpoint_token="new_token",  # noqa: S106
+            new_execution_state=CheckpointUpdatedExecutionState(
+                operations=operations.copy()
+            ),
+        )
+
+    return mock_checkpoint
+
+
+def test_plugin_sees_execution_input_and_result_end_to_end():
+    """A completing invocation surfaces the input on both hooks and the result."""
+    plugin = _PayloadRecordingPlugin()
+
+    @durable_execution(plugins=[plugin])
+    def my_handler(event: Any, context: DurableContext) -> dict:  # noqa: ARG001
+        return {"greeting": f"Hello, {event['name']}!"}
+
+    with patch(
+        "aws_durable_execution_sdk_python.execution.LambdaClient"
+    ) as mock_client_class:
+        mock_client = Mock()
+        mock_client.checkpoint = _tracking_checkpoint()
+        mock_client_class.initialize_client.return_value = mock_client
+
+        result = my_handler(_event('{"name": "World"}'), _lambda_context())
+
+    assert result["Status"] == InvocationStatus.SUCCEEDED.value
+
+    # Start hook: the deserialized input, not the raw payload string.
+    assert plugin.starts == [{"name": "World"}]
+
+    # End hook: the same input, plus the serialized result.
+    assert len(plugin.ends) == 1
+    status, end_input, end_result = plugin.ends[0]
+    assert status == InvocationStatus.SUCCEEDED.value
+    assert end_input == {"name": "World"}
+    assert end_result == '{"greeting": "Hello, World!"}'
+
+
+def test_plugin_payload_surfaces_on_suspending_invocation():
+    """A suspending invocation carries the input but no execution result."""
+    plugin = _PayloadRecordingPlugin()
+
+    @durable_execution(plugins=[plugin])
+    def my_handler(event: Any, context: DurableContext) -> str:
+        context.wait(Duration.from_seconds(60))
+        return f"done-{event['name']}"
+
+    with patch(
+        "aws_durable_execution_sdk_python.execution.LambdaClient"
+    ) as mock_client_class:
+        mock_client = Mock()
+        mock_client.checkpoint = _tracking_checkpoint()
+        mock_client_class.initialize_client.return_value = mock_client
+
+        result = my_handler(_event('{"name": "World"}'), _lambda_context())
+
+    assert result["Status"] == InvocationStatus.PENDING.value
+    assert plugin.starts == [{"name": "World"}]
+
+    status, end_input, end_result = plugin.ends[0]
+    assert status == InvocationStatus.PENDING.value
+    # The input is still reported on a non-terminal invocation-end.
+    assert end_input == {"name": "World"}
+    # But a suspending invocation produced no execution result.
+    assert end_result is None
+
+
+def test_plugin_payload_surfaces_on_replay_invocation():
+    """A replay past a completed wait carries the input and the terminal result."""
+    plugin = _PayloadRecordingPlugin()
+
+    @durable_execution(plugins=[plugin])
+    def my_handler(event: Any, context: DurableContext) -> str:
+        context.wait(Duration.from_seconds(60))
+        return f"done-{event['name']}"
+
+    # The wait completed externally while the execution was suspended.
+    completed_wait = {
+        "Id": next(operation_id_sequence()),
+        "Type": OperationType.WAIT.value,
+        "Status": OperationStatus.SUCCEEDED.value,
+    }
+
+    with patch(
+        "aws_durable_execution_sdk_python.execution.LambdaClient"
+    ) as mock_client_class:
+        mock_client = Mock()
+        mock_client.checkpoint = _tracking_checkpoint()
+        mock_client_class.initialize_client.return_value = mock_client
+
+        result = my_handler(
+            _event('{"name": "World"}', extra_operations=[completed_wait]),
+            _lambda_context(),
+        )
+
+    assert result["Status"] == InvocationStatus.SUCCEEDED.value
+    # The input is carried identically across invocations of one execution.
+    assert plugin.starts == [{"name": "World"}]
+
+    status, end_input, end_result = plugin.ends[0]
+    assert status == InvocationStatus.SUCCEEDED.value
+    assert end_input == {"name": "World"}
+    assert end_result == '"done-World"'
+
+
+def test_plugin_execution_input_is_isolated_from_handler_end_to_end():
+    """The plugin's input view and the handler's event must not alias.
+
+    durable_execution() hands one mutable object to both, so the plugin view is
+    deep-copied. Without that a plugin could alter execution behaviour, and a
+    handler could retroactively change what the frozen hook info reports.
+    """
+
+    class _MutatingPlugin(DurableInstrumentationPlugin):
+        def __init__(self) -> None:
+            self.end_inputs: list[Any] = []
+
+        def on_invocation_start(self, info) -> None:
+            info.execution_input["injected_by_plugin"] = True
+            info.execution_input["nested"]["items"].append("from_plugin")
+
+        def on_invocation_end(self, info) -> None:
+            self.end_inputs.append(info.execution_input)
+
+    plugin = _MutatingPlugin()
+    handler_saw: dict[str, Any] = {}
+
+    @durable_execution(plugins=[plugin])
+    def my_handler(event: Any, context: DurableContext) -> str:  # noqa: ARG001
+        handler_saw.update(
+            {"top": dict(event), "nested_items": list(event["nested"]["items"])}
+        )
+        event["injected_by_handler"] = True
+        return "ok"
+
+    with patch(
+        "aws_durable_execution_sdk_python.execution.LambdaClient"
+    ) as mock_client_class:
+        mock_client = Mock()
+        mock_client.checkpoint = _tracking_checkpoint()
+        mock_client_class.initialize_client.return_value = mock_client
+
+        my_handler(
+            _event('{"name": "World", "nested": {"items": ["original"]}}'),
+            _lambda_context(),
+        )
+
+    # The plugin's mutations never reached the handler, at any depth.
+    assert "injected_by_plugin" not in handler_saw["top"]
+    assert handler_saw["nested_items"] == ["original"]
+    # The handler's mutation never reached the end hook.
+    assert "injected_by_handler" not in plugin.end_inputs[0]