diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py index b541a6e7..c0ada9a6 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py +++ b/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) diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py index a8076a58..96bfb373 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py +++ b/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,6 +377,7 @@ 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( @@ -356,9 +385,39 @@ def on_invocation_start( 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", diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_payload_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_payload_int_test.py new file mode 100644 index 00000000..4055b8f1 --- /dev/null +++ b/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] diff --git a/packages/aws-durable-execution-sdk-python/tests/execution_test.py b/packages/aws-durable-execution-sdk-python/tests/execution_test.py index 9d4c1c32..4badbc47 100644 --- a/packages/aws-durable-execution-sdk-python/tests/execution_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/execution_test.py @@ -4,6 +4,7 @@ import json import time import warnings +from copy import deepcopy from typing import Any from unittest.mock import Mock, patch @@ -2675,13 +2676,13 @@ def test_from_dict_leaves_timestamps_as_integers(): # ============================================================================= -def _make_invocation_input(mock_client, next_marker=""): +def _make_invocation_input(mock_client, next_marker="", input_payload="{}"): """Helper to create a standard test invocation input.""" operation = Operation( operation_id="exec1", operation_type=OperationType.EXECUTION, status=OperationStatus.STARTED, - execution_details=ExecutionDetails(input_payload="{}"), + execution_details=ExecutionDetails(input_payload=input_payload), ) return DurableExecutionInvocationInputWithClient( durable_execution_arn="arn:test:execution/exec1", @@ -2855,6 +2856,12 @@ class _RecordingPlugin(DurableInstrumentationPlugin): def __init__(self) -> None: self.calls: list[str] = [] + # Payload surfaces observed on the invocation hooks, so tests can assert + # durable_execution() actually forwards them (not just that the + # dataclasses can hold them). + self.start_execution_inputs: list[Any] = [] + self.end_execution_inputs: list[Any] = [] + self.end_execution_results: list[str | None] = [] def on_execution_start(self, info): self.calls.append("execution_start") @@ -2864,9 +2871,12 @@ def on_execution_end(self, info): def on_invocation_start(self, info): self.calls.append("invocation_start") + self.start_execution_inputs.append(info.execution_input) def on_invocation_end(self, info): self.calls.append(f"invocation_end:{info.status.value}") + self.end_execution_inputs.append(info.execution_input) + self.end_execution_results.append(info.execution_result) def on_operation_start(self, info): self.calls.append(f"operation_start:{info.operation_id}") @@ -2957,6 +2967,134 @@ def test_handler(event: Any, context: DurableContext) -> dict: assert "invocation_end:SUCCEEDED" in plugin.calls +def test_durable_execution_forwards_execution_input_to_plugins(): + """durable_execution() must hand the deserialized input to both hooks. + + Guards the one production line that surfaces real input + (`execution_input=input_event`). A non-empty payload is used deliberately: + the field's own default is None and an empty payload deserializes to {}, so + only a populated payload distinguishes real forwarding from either. + """ + mock_client = Mock(spec=DurableServiceClient) + mock_client.checkpoint.return_value = CheckpointOutput( + checkpoint_token="new_token", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState(), + ) + + plugin = _RecordingPlugin() + + @durable_execution(plugins=[plugin]) + def test_handler(event: Any, context: DurableContext) -> dict: + return {"echoed": event["name"]} + + result = test_handler( + _make_invocation_input(mock_client, input_payload='{"name": "World"}'), + _make_lambda_context(), + ) + + assert result["Status"] == InvocationStatus.SUCCEEDED.value + # The invocation-start hook sees the deserialized input, not the raw payload. + assert plugin.start_execution_inputs == [{"name": "World"}] + # The end info inherits the same input from the start info. + assert plugin.end_execution_inputs == [{"name": "World"}] + # And the end hook carries the serialized result. + assert plugin.end_execution_results == ['{"echoed": "World"}'] + + +def test_durable_execution_surfaces_empty_input_as_empty_mapping(): + """An empty payload reaches plugins as {}, never None.""" + mock_client = Mock(spec=DurableServiceClient) + mock_client.checkpoint.return_value = CheckpointOutput( + checkpoint_token="new_token", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState(), + ) + + plugin = _RecordingPlugin() + + @durable_execution(plugins=[plugin]) + def test_handler(event: Any, context: DurableContext) -> str: + return "ok" + + test_handler( + _make_invocation_input(mock_client, input_payload=""), + _make_lambda_context(), + ) + + assert plugin.start_execution_inputs == [{}] + + +def test_durable_execution_isolates_execution_input_from_handler(): + """Plugin and handler must not observe each other's input mutations. + + 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. + """ + mock_client = Mock(spec=DurableServiceClient) + mock_client.checkpoint.return_value = CheckpointOutput( + checkpoint_token="new_token", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState(), + ) + + observed: dict[str, Any] = {} + + class _MutatingPlugin(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + # Direction A: plugin mutates its view before the handler runs. + info.execution_input["injected_by_plugin"] = True + + def on_invocation_end(self, info): + observed["end_input"] = dict(info.execution_input) + + @durable_execution(plugins=[_MutatingPlugin()]) + def test_handler(event: Any, context: DurableContext) -> dict: + observed["handler_saw"] = dict(event) + # Direction B: handler mutates its event after the start hook fired. + event["injected_by_handler"] = True + return {"ok": True} + + test_handler( + _make_invocation_input(mock_client, input_payload='{"name": "World"}'), + _make_lambda_context(), + ) + + # Direction A: the plugin's mutation must not reach the handler. + assert observed["handler_saw"] == {"name": "World"} + # Direction B: the handler's mutation must not reach the end hook, which + # still reports the plugin-side snapshot taken at invocation-start. + assert "injected_by_handler" not in observed["end_input"] + assert observed["end_input"] == {"name": "World", "injected_by_plugin": True} + + +def test_durable_execution_isolates_nested_execution_input(): + """Isolation must be deep, not just a top-level copy.""" + mock_client = Mock(spec=DurableServiceClient) + mock_client.checkpoint.return_value = CheckpointOutput( + checkpoint_token="new_token", # noqa: S106 + new_execution_state=CheckpointUpdatedExecutionState(), + ) + + observed: dict[str, Any] = {} + + class _NestedMutatingPlugin(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + info.execution_input["outer"]["inner"].append("from_plugin") + + @durable_execution(plugins=[_NestedMutatingPlugin()]) + def test_handler(event: Any, context: DurableContext) -> dict: + observed["handler_saw"] = deepcopy(event) + return {"ok": True} + + test_handler( + _make_invocation_input( + mock_client, input_payload='{"outer": {"inner": ["original"]}}' + ), + _make_lambda_context(), + ) + + assert observed["handler_saw"] == {"outer": {"inner": ["original"]}} + + def test_durable_execution_with_plugins_failure(): """Test that plugins receive invocation end and execution end on user error.""" mock_client = Mock(spec=DurableServiceClient) diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py index 322e1e3f..556dac0a 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -73,6 +73,7 @@ execution_arn="arn:aws:lambda:us-east-1:123:durable:abc", execution_start_time=START_TS, is_first_invocation=True, + execution_input={"name": "World"}, ) INVOCATION_END_INFO = InvocationEndInfo( request_id="req-1", @@ -81,6 +82,8 @@ status=InvocationStatus.FAILED, error=ERROR, is_first_invocation=False, + execution_input={"name": "World"}, + execution_result='"Hello, World!"', ) USER_FUNCTION_START_INFO = UserFunctionStartInfo( @@ -164,6 +167,16 @@ def test_invocation_start_info(self): ) self.assertEqual(INVOCATION_START_INFO.execution_start_time, START_TS) self.assertTrue(INVOCATION_START_INFO.is_first_invocation) + self.assertEqual(INVOCATION_START_INFO.execution_input, {"name": "World"}) + + def test_invocation_info_execution_input_defaults_to_none(self): + info = InvocationStartInfo( + request_id="req-1", + execution_arn="arn:test", + execution_start_time=START_TS, + is_first_invocation=True, + ) + self.assertIsNone(info.execution_input) def test_invocation_end_info(self): self.assertEqual(INVOCATION_END_INFO.request_id, "req-1") @@ -172,6 +185,23 @@ def test_invocation_end_info(self): self.assertFalse(INVOCATION_END_INFO.is_first_invocation) self.assertEqual(INVOCATION_END_INFO.status, InvocationStatus.FAILED) self.assertEqual(INVOCATION_END_INFO.error.message, "boom") + self.assertEqual(INVOCATION_END_INFO.execution_input, {"name": "World"}) + self.assertEqual(INVOCATION_END_INFO.execution_result, '"Hello, World!"') + + def test_invocation_end_info_from_invocation_output_carries_input_and_result(self): + output = DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, + result='"Hello, World!"', + ) + end_info = InvocationEndInfo.from_durable_execution_invocation_output( + INVOCATION_START_INFO, output + ) + self.assertEqual(end_info.request_id, INVOCATION_START_INFO.request_id) + self.assertEqual(end_info.execution_arn, INVOCATION_START_INFO.execution_arn) + self.assertEqual(end_info.execution_input, {"name": "World"}) + self.assertEqual(end_info.execution_result, '"Hello, World!"') + self.assertEqual(end_info.status, InvocationStatus.SUCCEEDED) + self.assertIsNone(end_info.error) def test_user_function_start_info(self): self.assertEqual(USER_FUNCTION_START_INFO.operation_id, "op-1")