From ee898fbdf53696faacef6d76725f3e9bd4a33e8c Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Mon, 10 Aug 2026 22:36:04 +0000 Subject: [PATCH 1/5] feat(plugin): expose operation maps on invocation hooks Invocation-level plugin infos carried no view of the execution's operation state, so Python plugins could not see what the JS SDK exposes on every invocation hook. Add to the invocation infos: - InvocationInfo.operations, the checkpointed operation map converted to OperationInfo, on both invocation-start and invocation-end - InvocationStartInfo.updated_operations, the subset named by the invocation input's UpdatedOperationIds, i.e. operations completed externally while the execution was suspended Both are kw-only with empty-map defaults, so existing constructor calls and plugins are unaffected. The end hook re-reads the map so it reports end-of-invocation state rather than the start snapshot. The map is snapshotted when the hook fires but converted to OperationInfo only on first access, so an operation-heavy execution does not pay per invocation for a view no plugin reads, while a plugin that stashes the info still sees the state as of its hook. Refs #617 --- .../execution.py | 6 + .../plugin.py | 168 +++++++++++- .../tests/plugin_test.py | 255 +++++++++++++++++- 3 files changed, 421 insertions(+), 8 deletions(-) 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 c0ada9a6..8ca5ef90 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 @@ -299,6 +299,12 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: ), is_first_invocation=not has_prior_operations, execution_input=input_event, + # Read the map through a callable rather than snapshotting it + # here: the invocation-end hook needs the state as of the end of + # the invocation, and neither hook pays for the conversion until + # a plugin actually reads it. + operations_provider=lambda: execution_state.operations, + updated_operation_ids=invocation_input.updated_operation_ids, ) # Thread 1: Run background checkpoint processing executor.submit(execution_state.checkpoint_batches_forever) 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 7ccb8d50..72eae8eb 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 @@ -9,7 +9,8 @@ from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from enum import Enum -from typing import Any, Callable, MutableMapping, cast +from threading import Lock +from typing import Any, Callable, Iterator, MutableMapping, cast from aws_durable_execution_sdk_python.identifier import OperationIdentifier from aws_durable_execution_sdk_python.lambda_service import ( @@ -107,6 +108,82 @@ def from_operation( ) +def _to_operation_info_map( + operations: Mapping[str, Operation], +) -> dict[str, OperationInfo]: + """Convert a map of checkpointed operations to the plugin ``OperationInfo`` view. + + ``is_replayed`` is left at its default ``False``: these entries describe the + stored state of an operation, not a replay event for it. Replay is signalled + through the dedicated operation hooks. + """ + return { + operation_id: OperationInfo.from_operation(operation) + for operation_id, operation in operations.items() + } + + +class _LazyOperationInfoMap(Mapping[str, OperationInfo]): + """Read-only ``operation id -> OperationInfo`` map built on first access. + + Invocation hook infos carry the execution's whole operation map. Converting + it eagerly would charge every invocation of an operation-heavy execution for + a view most plugins never read, so the conversion is deferred to the first + mapping operation and then cached. The underlying operations are snapshotted + before this map is handed out, so deferring the conversion does not move the + point in time the map describes. Behaves like a plain read-only ``dict`` + (iteration, ``len``, ``in``, ``get``, ``==`` against any mapping). + """ + + __slots__ = ("_provider", "_lock", "_resolved") + + def __init__(self, provider: Callable[[], dict[str, OperationInfo]] | None) -> None: + self._provider = provider + self._lock = Lock() + self._resolved: dict[str, OperationInfo] | None = None + + def _resolve(self) -> dict[str, OperationInfo]: + with self._lock: + if self._resolved is None: + if self._provider is None: + self._resolved = {} + else: + try: + self._resolved = self._provider() + except Exception: + # A plugin-facing view must never break the execution. + logger.exception( + "Failed to build plugin operations map; using empty map" + ) + self._resolved = {} + return self._resolved + + def __getitem__(self, key: str) -> OperationInfo: + return self._resolve()[key] + + def __iter__(self) -> Iterator[str]: + return iter(self._resolve()) + + def __len__(self) -> int: + return len(self._resolve()) + + def __eq__(self, other: object) -> bool: + if isinstance(other, Mapping): + return self._resolve() == dict(other) + return NotImplemented + + def __ne__(self, other: object) -> bool: + result = self.__eq__(other) + if result is NotImplemented: + return result + return not result + + __hash__ = None # type: ignore[assignment] # mutable-by-materialization view + + def __repr__(self) -> str: + return f"{type(self).__name__}({self._resolve()!r})" + + @dataclass(frozen=True) class OperationStartInfo(OperationInfo): pass @@ -204,11 +281,29 @@ class InvocationInfo: without it); ``durable_execution()`` always populates it with the deserialized input payload, which is ``{}`` when the payload is empty. """ + operations: Mapping[str, OperationInfo] = field(default_factory=dict, kw_only=True) + """Checkpointed operations for this execution, keyed by operation id. + + A point-in-time view of the execution's operation map: as observed at the + start of the invocation on ``on_invocation_start``, and as observed at the + end of the invocation on ``on_invocation_end``. Empty on the very first + invocation-start, before any operation has been checkpointed. + """ @dataclass(frozen=True) class InvocationStartInfo(InvocationInfo): - pass + updated_operations: Mapping[str, OperationInfo] = field( + default_factory=dict, kw_only=True + ) + """Operations updated externally while this execution was suspended. + + A wait timer that expired, a callback that was delivered, or a chained + invoke that completed between the previous invocation and this one. This is + the subset of :attr:`InvocationInfo.operations` named by the durable + invocation input's ``UpdatedOperationIds``, so it is empty on the first + invocation. + """ @dataclass(frozen=True) @@ -244,6 +339,7 @@ def from_durable_execution_invocation_output( cls, invocation_start_info: InvocationStartInfo, output: "DurableExecutionInvocationOutput", + operations: Mapping[str, OperationInfo] | None = None, ): return InvocationEndInfo( request_id=invocation_start_info.request_id, @@ -251,6 +347,13 @@ def from_durable_execution_invocation_output( 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, + # Default to the start-of-invocation view when the caller has no + # fresher snapshot to offer. + operations=( + operations + if operations is not None + else invocation_start_info.operations + ), status=output.status, error=output.error, execution_result=output.result, @@ -341,6 +444,7 @@ def __init__(self, plugins: list[DurableInstrumentationPlugin] | None): self._plugins = plugins or [] self._executor: ThreadPoolExecutor | None = None self._invocation_status: InvocationStartInfo | None = None + self._operations_provider: Callable[[], Mapping[str, Operation]] | None = None @contextlib.contextmanager def run(self): @@ -353,6 +457,7 @@ def run(self): yield finally: self._invocation_status = None + self._operations_provider = None # Shut down the thread pool, waiting for pending tasks to complete. if self._executor: self._executor.shutdown(wait=True) @@ -393,6 +498,28 @@ def execute_plugins(self, info, sync): # this is called asynchronously, so plugins cannot manipulate thread local objects self._executor.submit(self._dispatch_plugin, plugin, info) + @staticmethod + def _snapshot_operation_infos( + operations_provider: Callable[[], Mapping[str, Operation]] | None, + ) -> Mapping[str, OperationInfo]: + """Capture the operation map now; defer the ``OperationInfo`` conversion. + + Copying the map is cheap, building an ``OperationInfo`` per operation is + not. Snapshotting eagerly pins the point in time the hook reports, so a + plugin that stashes the info and reads it later still sees the state as + of its hook; deferring the conversion keeps operation-heavy executions + from paying for a view no plugin reads. + """ + if operations_provider is None: + return _LazyOperationInfoMap(None) + try: + snapshot = dict(operations_provider()) + except Exception: + # A plugin-facing view must never break the execution. + logger.exception("Failed to snapshot operations for plugin hook") + return _LazyOperationInfoMap(None) + return _LazyOperationInfoMap(lambda: _to_operation_info_map(snapshot)) + def on_invocation_start( self, execution_arn: str, @@ -400,14 +527,40 @@ def on_invocation_start( execution_start_time: datetime.datetime | None, lambda_context: LambdaContext | None, execution_input: Any = None, + operations_provider: Callable[[], Mapping[str, Operation]] | None = None, + updated_operation_ids: Sequence[str] | None = None, ) -> None: + """Fire the invocation-start hook. + + Args: + execution_arn: ARN of the durable execution. + is_first_invocation: False when prior operations exist (a replay). + execution_start_time: Start timestamp of the execution operation. + lambda_context: Lambda context, for the request id. + execution_input: The deserialized execution input event. + operations_provider: Returns the current checkpointed operation map. + Called once here to snapshot it; the conversion to the plugin's + ``OperationInfo`` view is deferred to first access. + updated_operation_ids: Operation ids from the invocation input's + ``UpdatedOperationIds`` -- those updated while suspended. + """ aws_request_id = lambda_context.aws_request_id if lambda_context else None + self._operations_provider = operations_provider + operations = self._snapshot_operation_infos(operations_provider) 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), + operations=operations, + updated_operations=_LazyOperationInfoMap( + lambda: { + operation_id: operations[operation_id] + for operation_id in (updated_operation_ids or []) + if operation_id in operations + } + ), ) self.execute_plugins(self._invocation_status, sync=True) @@ -448,9 +601,13 @@ def on_invocation_end( # on_invocation_start not called, skip return + # Re-read the operation map so the end hook sees the state as of the end + # of this invocation, not the snapshot taken at its start. invocation_end_info = ( InvocationEndInfo.from_durable_execution_invocation_output( - self._invocation_status, output + self._invocation_status, + output, + operations=self._snapshot_operation_infos(self._operations_provider), ) ) self.execute_plugins(invocation_end_info, sync=True) @@ -629,10 +786,7 @@ def on_operation_update( operation.operation_id: OperationInfo.from_operation(operation) for operation in changed_operations }, - operations={ - operation_id: OperationInfo.from_operation(operation) - for operation_id, operation in operations.items() - }, + operations=_to_operation_info_map(operations), ), sync=True, ) 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 89de9c68..19fe9665 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -2,7 +2,7 @@ import logging import unittest from dataclasses import fields -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch from aws_durable_execution_sdk_python.identifier import OperationIdentifier from aws_durable_execution_sdk_python.lambda_service import ( @@ -330,6 +330,69 @@ def test_invocation_end_info_from_invocation_output_carries_input_and_result(sel self.assertEqual(end_info.status, InvocationStatus.SUCCEEDED) self.assertIsNone(end_info.error) + def test_invocation_info_operation_maps_default_to_empty(self): + """Both maps default to empty so existing constructor calls keep working.""" + self.assertEqual({}, INVOCATION_START_INFO.operations) + self.assertEqual({}, INVOCATION_START_INFO.updated_operations) + self.assertEqual({}, INVOCATION_END_INFO.operations) + + def test_invocation_start_info_carries_operation_maps(self): + info = InvocationStartInfo( + request_id="req-1", + execution_arn="arn:test", + is_first_invocation=False, + operations={"op-1": OPERATION_END_INFO, "op-2": OPERATION_START_INFO}, + updated_operations={"op-1": OPERATION_END_INFO}, + ) + + self.assertEqual( + {"op-1": OPERATION_END_INFO, "op-2": OPERATION_START_INFO}, + info.operations, + ) + self.assertEqual({"op-1": OPERATION_END_INFO}, info.updated_operations) + + def test_invocation_end_info_inherits_start_operations(self): + """The end factory reuses the start snapshot when given no override.""" + start = InvocationStartInfo( + request_id="req-1", + execution_arn="arn:test", + is_first_invocation=True, + operations={"op-1": OPERATION_END_INFO}, + updated_operations={"op-1": OPERATION_END_INFO}, + ) + output = DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, result=None, error=None + ) + + end = InvocationEndInfo.from_durable_execution_invocation_output(start, output) + + self.assertEqual({"op-1": OPERATION_END_INFO}, end.operations) + # updated_operations is a start-hook-only surface. + self.assertFalse(hasattr(end, "updated_operations")) + + def test_invocation_end_info_accepts_fresher_operations(self): + """An explicit map wins, letting the end hook see end-of-invocation state.""" + start = InvocationStartInfo( + request_id="req-1", + execution_arn="arn:test", + is_first_invocation=True, + operations={"op-1": OPERATION_START_INFO}, + ) + output = DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, result=None, error=None + ) + + end = InvocationEndInfo.from_durable_execution_invocation_output( + start, + output, + operations={"op-1": OPERATION_END_INFO, "op-2": OPERATION_START_INFO}, + ) + + self.assertEqual( + {"op-1": OPERATION_END_INFO, "op-2": OPERATION_START_INFO}, + end.operations, + ) + def test_user_function_start_info(self): self.assertEqual(USER_FUNCTION_START_INFO.operation_id, "op-1") self.assertEqual(USER_FUNCTION_START_INFO.operation_type, OperationType.STEP) @@ -706,6 +769,196 @@ def test_pending_fires_invocation_end(self): self.assertIn("invocation_end:req-1", self.plugin.calls) +class TestInvocationHookOperationMaps(unittest.TestCase): + """Tests for the operations / updated_operations maps on invocation hooks.""" + + def setUp(self): + self.captured: list[object] = [] + + class _CapturingPlugin(DurableInstrumentationPlugin): + def on_invocation_start(_self, info): # noqa: N805 + self.captured.append(info) + + def on_invocation_end(_self, info): # noqa: N805 + self.captured.append(info) + + self.executor = PluginExecutor(plugins=[_CapturingPlugin()]) + + @staticmethod + def _operation(operation_id, status=OperationStatus.SUCCEEDED): + return Operation( + operation_id=operation_id, + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name=f"name-{operation_id}", + parent_id="root", + status=status, + start_timestamp=START_TS, + end_timestamp=END_TS, + step_details=StepDetails(attempt=1), + ) + + def _start(self, *, operations=None, updated_operation_ids=None): + self.executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=False, + operations_provider=(lambda: operations) + if operations is not None + else None, + updated_operation_ids=updated_operation_ids, + ) + + def test_start_info_converts_operations_to_operation_info(self): + operations = {"op-1": self._operation("op-1")} + + with self.executor.run(): + self._start(operations=operations) + + (start_info,) = self.captured + self.assertEqual(["op-1"], list(start_info.operations)) + converted = start_info.operations["op-1"] + self.assertIsInstance(converted, OperationInfo) + self.assertEqual("op-1", converted.operation_id) + self.assertEqual(OperationStatus.SUCCEEDED, converted.status) + self.assertEqual(START_TS, converted.start_time) + self.assertEqual(END_TS, converted.end_time) + self.assertEqual(1, converted.attempt) + # Stored state, not a replay event for the operation. + self.assertFalse(converted.is_replayed) + + def test_updated_operations_is_the_subset_named_by_updated_ids(self): + operations = { + "op-1": self._operation("op-1"), + "op-2": self._operation("op-2"), + } + + with self.executor.run(): + self._start(operations=operations, updated_operation_ids=["op-2"]) + + (start_info,) = self.captured + self.assertEqual(2, len(start_info.operations)) + self.assertEqual(["op-2"], list(start_info.updated_operations)) + # Same converted instance as in the full map, not a second conversion. + self.assertIs( + start_info.operations["op-2"], start_info.updated_operations["op-2"] + ) + + def test_updated_ids_absent_from_the_map_are_skipped(self): + operations = {"op-1": self._operation("op-1")} + + with self.executor.run(): + self._start( + operations=operations, updated_operation_ids=["op-1", "op-missing"] + ) + + (start_info,) = self.captured + self.assertEqual(["op-1"], list(start_info.updated_operations)) + + def test_first_invocation_has_empty_maps(self): + with self.executor.run(): + self._start(operations={}, updated_operation_ids=[]) + + (start_info,) = self.captured + self.assertEqual({}, start_info.operations) + self.assertEqual({}, start_info.updated_operations) + + def test_missing_provider_yields_empty_maps(self): + """Callers that supply no provider still get a usable, empty map.""" + with self.executor.run(): + self._start() + self.executor.on_invocation_end( + output=DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, result=None, error=None + ) + ) + + start_info, end_info = self.captured + self.assertEqual({}, start_info.operations) + self.assertEqual({}, start_info.updated_operations) + self.assertEqual({}, end_info.operations) + + def test_operation_info_conversion_is_deferred_and_cached(self): + conversions: list[str] = [] + real_from_operation = OperationInfo.from_operation + + def counting(operation, **kwargs): + conversions.append(operation.operation_id) + return real_from_operation(operation, **kwargs) + + with ( + patch.object(OperationInfo, "from_operation", staticmethod(counting)), + self.executor.run(), + ): + self._start(operations={"op-1": self._operation("op-1")}) + # Nothing read the map inside the hook, so nothing was converted. + self.assertEqual([], conversions) + + (start_info,) = self.captured + self.assertEqual(1, len(start_info.operations)) + self.assertEqual(["op-1"], conversions) + # Repeated reads reuse the cached conversion. + self.assertEqual(["op-1"], list(start_info.operations)) + self.assertIn("op-1", start_info.operations) + self.assertEqual(["op-1"], conversions) + + def test_end_info_reflects_operations_added_during_the_invocation(self): + operations = {"op-1": self._operation("op-1")} + + with self.executor.run(): + self.executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + operations_provider=lambda: operations, + ) + operations["op-2"] = self._operation("op-2") + self.executor.on_invocation_end( + output=DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, result=None, error=None + ) + ) + + start_info, end_info = self.captured + self.assertEqual(["op-1", "op-2"], sorted(end_info.operations)) + # The start map is snapshotted at its hook, so a later checkpoint does + # not retroactively change what the start hook reported. + self.assertEqual(["op-1"], sorted(start_info.operations)) + + def test_failing_provider_degrades_to_an_empty_map(self): + def provider(): + raise RuntimeError("boom") + + with ( + self.executor.run(), + self.assertLogs( + "aws_durable_execution_sdk_python.plugin", level="ERROR" + ) as logs, + ): + self.executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + operations_provider=provider, + updated_operation_ids=["op-1"], + ) + + self.assertTrue(any("Failed to snapshot operations" in m for m in logs.output)) + (start_info,) = self.captured + self.assertEqual(0, len(start_info.operations)) + self.assertEqual(0, len(start_info.updated_operations)) + + def test_provider_is_released_when_the_run_scope_exits(self): + with self.executor.run(): + self._start(operations={"op-1": self._operation("op-1")}) + self.assertIsNotNone(self.executor._operations_provider) # noqa: SLF001 + + self.assertIsNone(self.executor._operations_provider) # noqa: SLF001 + + class TestPluginExecutorOnOperationAction(unittest.TestCase): """Tests for PluginExecutor.on_operation_action.""" From 877c16481ebf63626c67619433201ce0cd70b782 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Tue, 11 Aug 2026 18:51:08 +0000 Subject: [PATCH 2/5] fix(plugin): keep invocation info deepcopy-safe Addresses the Codex review comment on #623. The thread lock in the lazy operation map made the enclosing frozen info uncopyable: dataclasses.asdict() deep-copies any field that is not a dict, list, tuple or dataclass, and a lock cannot be copied. Both asdict() and deepcopy() raised TypeError, and since plugin exceptions are swallowed a plugin that serialized an info would have its hook silently dropped. Drop the lock: building the map is pure and idempotent, so a race can only duplicate work, never corrupt the result, and the assignment is atomic. Add __deepcopy__ so the view materializes into an ordinary dict when copied, which also keeps the provider closure out of the copy. Adds regression tests using asdict() and deepcopy() over both hook infos and both maps, including the not-yet-read case. Refs #617 --- .../plugin.py | 53 +++++++--- .../tests/plugin_test.py | 100 +++++++++++++++++- 2 files changed, 135 insertions(+), 18 deletions(-) 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 72eae8eb..3ca19764 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 @@ -9,7 +9,6 @@ from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from enum import Enum -from threading import Lock from typing import Any, Callable, Iterator, MutableMapping, cast from aws_durable_execution_sdk_python.identifier import OperationIdentifier @@ -132,31 +131,37 @@ class _LazyOperationInfoMap(Mapping[str, OperationInfo]): mapping operation and then cached. The underlying operations are snapshotted before this map is handed out, so deferring the conversion does not move the point in time the map describes. Behaves like a plain read-only ``dict`` - (iteration, ``len``, ``in``, ``get``, ``==`` against any mapping). + (iteration, ``len``, ``in``, ``get``, ``==`` against any mapping), and + materializes into an ordinary ``dict`` when deep-copied so the enclosing + info stays copyable and ``dataclasses.asdict()``-able. """ - __slots__ = ("_provider", "_lock", "_resolved") + __slots__ = ("_provider", "_resolved") def __init__(self, provider: Callable[[], dict[str, OperationInfo]] | None) -> None: self._provider = provider - self._lock = Lock() self._resolved: dict[str, OperationInfo] | None = None def _resolve(self) -> dict[str, OperationInfo]: - with self._lock: - if self._resolved is None: - if self._provider is None: + # Deliberately lock-free. Building the map is pure and idempotent, so a + # race can only duplicate the work, never corrupt the result, and the + # assignment below is atomic. A lock here would make the enclosing + # frozen info undeepcopyable, which breaks the plugins this view exists + # to serve: dataclasses.asdict() deep-copies non-dict fields, and a + # thread lock cannot be copied. + if self._resolved is None: + if self._provider is None: + self._resolved = {} + else: + try: + self._resolved = self._provider() + except Exception: + # A plugin-facing view must never break the execution. + logger.exception( + "Failed to build plugin operations map; using empty map" + ) self._resolved = {} - else: - try: - self._resolved = self._provider() - except Exception: - # A plugin-facing view must never break the execution. - logger.exception( - "Failed to build plugin operations map; using empty map" - ) - self._resolved = {} - return self._resolved + return self._resolved def __getitem__(self, key: str) -> OperationInfo: return self._resolve()[key] @@ -180,6 +185,20 @@ def __ne__(self, other: object) -> bool: __hash__ = None # type: ignore[assignment] # mutable-by-materialization view + def __deepcopy__(self, memo: dict[int, Any]) -> dict[str, OperationInfo]: + """Materialize into a plain ``dict`` when copied. + + Copying this view has no reason to preserve its laziness, and a plain + dict is what callers actually want: it makes + ``dataclasses.asdict(info)`` yield an ordinary mapping instead of an + opaque object, and keeps the provider closure out of the copy. + """ + resolved = { + key: copy.deepcopy(value, memo) for key, value in self._resolve().items() + } + memo[id(self)] = resolved + return resolved + def __repr__(self) -> str: return f"{type(self).__name__}({self._resolve()!r})" 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 19fe9665..e632bc05 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -1,7 +1,8 @@ import datetime import logging import unittest -from dataclasses import fields +from copy import deepcopy +from dataclasses import asdict, fields from unittest.mock import MagicMock, patch from aws_durable_execution_sdk_python.identifier import OperationIdentifier @@ -959,6 +960,103 @@ def test_provider_is_released_when_the_run_scope_exits(self): self.assertIsNone(self.executor._operations_provider) # noqa: SLF001 +class TestInvocationInfoCopySafety(unittest.TestCase): + """Invocation infos carrying the lazy operation map must stay copyable. + + ``dataclasses.asdict()`` deep-copies any field that is not a dict/list/tuple + or dataclass, so anything uncopyable embedded in the lazy map would make + every plugin that serializes an info raise -- and because plugin exceptions + are swallowed, the hook would be silently dropped. + """ + + def setUp(self): + self.operation = Operation( + operation_id="op-1", + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name="name-op-1", + parent_id="root", + status=OperationStatus.SUCCEEDED, + start_timestamp=START_TS, + end_timestamp=END_TS, + step_details=StepDetails(attempt=1), + ) + self.captured: list[object] = [] + + class _CapturingPlugin(DurableInstrumentationPlugin): + def on_invocation_start(_self, info): # noqa: N805 + self.captured.append(info) + + def on_invocation_end(_self, info): # noqa: N805 + self.captured.append(info) + + self.executor = PluginExecutor(plugins=[_CapturingPlugin()]) + + def _fire_hooks(self): + with self.executor.run(): + self.executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=False, + operations_provider=lambda: {"op-1": self.operation}, + updated_operation_ids=["op-1"], + ) + self.executor.on_invocation_end( + output=DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, result=None, error=None + ) + ) + return self.captured + + def test_asdict_on_invocation_start_info(self): + start_info, _ = self._fire_hooks() + + as_dict = asdict(start_info) + + # Both lazy maps materialize as ordinary dicts. + self.assertIsInstance(as_dict["operations"], dict) + self.assertIsInstance(as_dict["updated_operations"], dict) + self.assertEqual(["op-1"], list(as_dict["operations"])) + self.assertEqual(["op-1"], list(as_dict["updated_operations"])) + + def test_asdict_on_invocation_end_info(self): + _, end_info = self._fire_hooks() + + as_dict = asdict(end_info) + + self.assertIsInstance(as_dict["operations"], dict) + self.assertEqual(["op-1"], list(as_dict["operations"])) + + def test_deepcopy_on_invocation_infos(self): + start_info, end_info = self._fire_hooks() + + for info in (start_info, end_info): + copied = deepcopy(info) + self.assertIsInstance(copied.operations, dict) + self.assertEqual(["op-1"], list(copied.operations)) + + def test_deepcopy_is_independent_of_the_original(self): + start_info, _ = self._fire_hooks() + + copied = deepcopy(start_info) + copied.operations["op-2"] = OPERATION_END_INFO + + self.assertEqual(["op-1"], list(start_info.operations)) + + def test_unresolved_map_is_still_copyable(self): + """Copying must not require the map to have been read first.""" + start_info, _ = self._fire_hooks() + fresh = InvocationStartInfo( + request_id=start_info.request_id, + execution_arn=start_info.execution_arn, + is_first_invocation=start_info.is_first_invocation, + operations=start_info.operations, + ) + + self.assertIsInstance(asdict(fresh)["operations"], dict) + + class TestPluginExecutorOnOperationAction(unittest.TestCase): """Tests for PluginExecutor.on_operation_action.""" From ca7c6adf7e4b8c1daba99d25d84595c4b5bb320f Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Tue, 11 Aug 2026 20:58:01 +0000 Subject: [PATCH 3/5] fix(plugin): make operation maps additive too The operation maps had the same defect #627 fixed for the payload fields, and #627's test caught it as soon as these commits were rebased onto main: a mapping-valued field in the generated __hash__ makes a previously hashable InvocationStartInfo raise TypeError, and the map entries carry operation results and errors that instrumentation logs wholesale via repr. Set repr=False, compare=False, hash=False on operations and updated_operations, and extend the field-declaration tests to cover them alongside hashability, equality and repr assertions. Refs #617 --- .../plugin.py | 22 ++++++++++-- .../tests/plugin_test.py | 36 +++++++++++++++++++ 2 files changed, 56 insertions(+), 2 deletions(-) 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 3ca19764..beea9c7c 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 @@ -300,20 +300,35 @@ class InvocationInfo: without it); ``durable_execution()`` always populates it with the deserialized input payload, which is ``{}`` when the payload is empty. """ - operations: Mapping[str, OperationInfo] = field(default_factory=dict, kw_only=True) + operations: Mapping[str, OperationInfo] = field( + default_factory=dict, + kw_only=True, + repr=False, + compare=False, + hash=False, + ) """Checkpointed operations for this execution, keyed by operation id. A point-in-time view of the execution's operation map: as observed at the start of the invocation on ``on_invocation_start``, and as observed at the end of the invocation on ``on_invocation_end``. Empty on the very first invocation-start, before any operation has been checkpointed. + + Excluded from ``repr``, ``__eq__`` and ``__hash__`` for the same reasons as + :attr:`execution_input`: the entries carry operation results and errors that + instrumentation would otherwise log wholesale, and a mapping-valued field + would make a previously hashable info unhashable. """ @dataclass(frozen=True) class InvocationStartInfo(InvocationInfo): updated_operations: Mapping[str, OperationInfo] = field( - default_factory=dict, kw_only=True + default_factory=dict, + kw_only=True, + repr=False, + compare=False, + hash=False, ) """Operations updated externally while this execution was suspended. @@ -322,6 +337,9 @@ class InvocationStartInfo(InvocationInfo): the subset of :attr:`InvocationInfo.operations` named by the durable invocation input's ``UpdatedOperationIds``, so it is empty on the first invocation. + + Excluded from ``repr``, ``__eq__`` and ``__hash__`` like + :attr:`InvocationInfo.operations`. """ 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 e632bc05..edcb769f 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -230,6 +230,42 @@ def test_payload_fields_are_excluded_from_equality(self): ), ) + def test_operation_map_fields_are_additive(self): + """The operation maps must be as additive as the payload fields. + + They are mapping-valued, so including them in the generated methods + would make a previously hashable info unhashable, and their entries + carry operation results and errors that instrumentation logs wholesale. + """ + start_fields = {f.name: f for f in fields(InvocationStartInfo)} + end_fields = {f.name: f for f in fields(InvocationEndInfo)} + + for holder, name in ( + (start_fields, "operations"), + (start_fields, "updated_operations"), + (end_fields, "operations"), + ): + self.assertFalse(holder[name].repr, name) + self.assertFalse(holder[name].compare, name) + self.assertIs(holder[name].hash, False, name) + + base = { + "request_id": "req-1", + "execution_arn": "arn:test", + "is_first_invocation": True, + } + with_maps = InvocationStartInfo( + **base, + operations={"op-1": OPERATION_END_INFO}, + updated_operations={"op-1": OPERATION_END_INFO}, + ) + + # Hashable despite the mapping fields, and equal to the map-free info. + self.assertEqual(hash(InvocationStartInfo(**base)), hash(with_maps)) + self.assertEqual(InvocationStartInfo(**base), with_maps) + # Operation results do not leak through the info's repr. + self.assertNotIn("op-1", repr(with_maps)) + def test_payload_fields_are_declared_non_compare(self): """Pin the declarations, not just the observed behaviour.""" start_fields = {f.name: f for f in fields(InvocationStartInfo)} From d506cbc223fc9acf4c093f8340eb867f6cdf9c7b Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Tue, 11 Aug 2026 22:29:31 +0000 Subject: [PATCH 4/5] fix(plugin): isolate errors and gate the operations provider Addresses the review comments on #629. Error isolation: OperationInfo.from_operation reused the checkpoint's ErrorObject, whose stack_trace is a mutable list handed to user code on replay. A plugin reading info.operations could append to or clear that list and change the error the execution later raises. Reproduced, then fixed by cloning the error and its stack_trace at the conversion, so every plugin-facing OperationInfo is isolated, not just the ones in the new maps. Only the list needs cloning; the other fields are strings. Provider gating: durable_execution() passes an operations provider unconditionally, so a plugin-free execution invoked it at both hooks. Snapshotting and provider retention are now gated on self._plugins. The reviewer also asked to drop the snapshot copy since ExecutionState.operations already returns one. I kept it: the copy is what pins the point in time, and an existing test proved that removing it lets a provider returning a live mapping leak later mutations into an already-reported hook. It is a shallow pointer copy and now only runs when plugins are registered. Marks operations and updated_operations experimental, matching the payload fields. Adds unit tests for error isolation, the no-plugin gate and snapshot pinning, plus e2e coverage under tests/e2e/ for the start and end maps and UpdatedOperationIds across a real suspend/replay pair. Refs #617 --- .../plugin.py | 56 +++- .../plugin_invocation_operations_int_test.py | 258 ++++++++++++++++++ .../tests/plugin_test.py | 92 +++++++ 3 files changed, 395 insertions(+), 11 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py 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 beea9c7c..b8987650 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 @@ -98,7 +98,7 @@ def from_operation( start_time=operation.start_timestamp, end_time=operation.end_timestamp, result=_extract_result(operation), - error=_extract_error(operation), + error=_copy_error(_extract_error(operation)), attempt=( operation.step_details.attempt if operation.step_details else None ), @@ -107,6 +107,27 @@ def from_operation( ) +def _copy_error(error: ErrorObject | None) -> ErrorObject | None: + """Return a plugin-owned copy of an operation error. + + The checkpointed ``ErrorObject`` is handed straight to user code on replay, + and its ``stack_trace`` is a mutable list. Without a copy a plugin reading + ``info.operations`` could append to (or clear) that list and change the error + the execution later raises. Only the list needs cloning -- the other fields + are immutable strings -- so this is cheaper than a full deep copy. + """ + if error is None: + return None + return ErrorObject( + message=error.message, + type=error.type, + data=error.data, + stack_trace=( + list(error.stack_trace) if error.stack_trace is not None else None + ), + ) + + def _to_operation_info_map( operations: Mapping[str, Operation], ) -> dict[str, OperationInfo]: @@ -306,8 +327,9 @@ class InvocationInfo: repr=False, compare=False, hash=False, + metadata={"experimental": True}, ) - """Checkpointed operations for this execution, keyed by operation id. + """EXPERIMENTAL: Checkpointed operations for this execution, keyed by id. A point-in-time view of the execution's operation map: as observed at the start of the invocation on ``on_invocation_start``, and as observed at the @@ -329,8 +351,9 @@ class InvocationStartInfo(InvocationInfo): repr=False, compare=False, hash=False, + metadata={"experimental": True}, ) - """Operations updated externally while this execution was suspended. + """EXPERIMENTAL: Operations updated externally while this execution was suspended. A wait timer that expired, a callback that was delivered, or a chained invoke that completed between the previous invocation and this one. This is @@ -535,19 +558,30 @@ def execute_plugins(self, info, sync): # this is called asynchronously, so plugins cannot manipulate thread local objects self._executor.submit(self._dispatch_plugin, plugin, info) - @staticmethod def _snapshot_operation_infos( + self, operations_provider: Callable[[], Mapping[str, Operation]] | None, ) -> Mapping[str, OperationInfo]: """Capture the operation map now; defer the ``OperationInfo`` conversion. - Copying the map is cheap, building an ``OperationInfo`` per operation is - not. Snapshotting eagerly pins the point in time the hook reports, so a - plugin that stashes the info and reads it later still sees the state as - of its hook; deferring the conversion keeps operation-heavy executions - from paying for a view no plugin reads. + Snapshotting eagerly pins the point in time the hook reports, so a plugin + that stashes the info and reads it later still sees the state as of its + hook; deferring the conversion keeps operation-heavy executions from + paying for a view no plugin reads. + + Skipped entirely when no plugins are registered -- ``durable_execution()`` + passes a provider unconditionally, so without this gate a plugin-free + execution would still invoke it at both hooks. + + The returned mapping is copied even though the SDK's provider + (``ExecutionState.operations``) already returns a copy: the copy is what + pins the point in time, and relying on every provider to hand over a + mapping it will never touch again would make that invariant contingent on + an external contract. It is a shallow copy of pointers, and now only + happens when plugins are registered, so it is far cheaper than the + ``OperationInfo`` conversion it guards. """ - if operations_provider is None: + if not self._plugins or operations_provider is None: return _LazyOperationInfoMap(None) try: snapshot = dict(operations_provider()) @@ -582,7 +616,7 @@ def on_invocation_start( ``UpdatedOperationIds`` -- those updated while suspended. """ aws_request_id = lambda_context.aws_request_id if lambda_context else None - self._operations_provider = operations_provider + self._operations_provider = operations_provider if self._plugins else None operations = self._snapshot_operation_infos(operations_provider) self._invocation_status = InvocationStartInfo( execution_arn=execution_arn, diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py new file mode 100644 index 00000000..26f1b77a --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_operations_int_test.py @@ -0,0 +1,258 @@ +"""Integration tests for the plugin invocation operation maps. + +Exercises `InvocationInfo.operations` and +`InvocationStartInfo.updated_operations` through complete +`durable_execution()` invocations -- across the decorator, invocation-input +parsing, the checkpoint path and the invocation hooks -- including a +suspend/replay pair where the replay reports the externally-completed wait via +`UpdatedOperationIds`. +""" + +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 _MapRecordingPlugin(DurableInstrumentationPlugin): + """Records the operation maps seen on each invocation hook.""" + + def __init__(self) -> None: + self.starts: list[tuple[list[str], list[str]]] = [] + self.ends: list[tuple[str, list[str]]] = [] + + def on_invocation_start(self, info) -> None: + self.starts.append((sorted(info.operations), sorted(info.updated_operations))) + + def on_invocation_end(self, info) -> None: + self.ends.append((info.status.value, sorted(info.operations))) + + +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( + extra_operations: list[dict] | None = None, + updated_operation_ids: list[str] | None = None, +) -> dict: + execution_operation = { + "Id": "execution-1", + "Type": "EXECUTION", + "Status": "STARTED", + "ExecutionDetails": {"InputPayload": '{"name": "World"}'}, + } + event: dict[str, Any] = { + "DurableExecutionArn": "test-arn/execution-1", + "CheckpointToken": "test-token", + "InitialExecutionState": { + "Operations": [execution_operation, *(extra_operations or [])], + "NextMarker": "", + }, + "LocalRunner": True, + } + if updated_operation_ids is not None: + event["UpdatedOperationIds"] = updated_operation_ids + return event + + +def _tracking_checkpoint(initial_operations: list[Operation] | None = None): + """Checkpoint mock that accumulates operations, as the service would.""" + 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_operation_maps_on_a_completing_invocation(): + """The start map holds the prior state; the end map sees the step added.""" + plugin = _MapRecordingPlugin() + + @durable_execution(plugins=[plugin]) + def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + return context.step(lambda _ctx: "stepped", name="greet") + + 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(), _lambda_context()) + + assert result["Status"] == InvocationStatus.SUCCEEDED.value + + # First invocation: only the EXECUTION operation exists at start, and + # nothing was updated externally. + (start_operations, start_updated) = plugin.starts[0] + assert start_operations == ["execution-1"] + assert start_updated == [] + + # The end hook re-reads the map, so it sees the step checkpointed during + # this invocation -- that is the point of re-reading rather than reusing + # the start snapshot. + status, end_operations = plugin.ends[0] + assert status == InvocationStatus.SUCCEEDED.value + assert len(end_operations) > len(start_operations) + assert "execution-1" in end_operations + + +def test_operation_maps_across_suspend_and_replay(): + """The replay start hook reports the externally-completed wait. + + Invocation 1 suspends on a wait. Invocation 2 replays with the wait already + SUCCEEDED and its id in ``UpdatedOperationIds``, which is exactly what + ``updated_operations`` is derived from. + """ + wait_id = next(operation_id_sequence()) + + # --- Invocation 1: the wait starts and the execution suspends. + first = _MapRecordingPlugin() + + @durable_execution(plugins=[first]) + def suspending_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + context.wait(Duration.from_seconds(60)) + return "done" + + 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 + + first_result = suspending_handler(_event(), _lambda_context()) + + assert first_result["Status"] == InvocationStatus.PENDING.value + start_operations, start_updated = first.starts[0] + assert start_operations == ["execution-1"] + assert start_updated == [] + # The suspending end hook already sees the wait that was just checkpointed. + _, end_operations = first.ends[0] + assert wait_id in end_operations + + # --- Invocation 2: replay with the wait completed externally. + replay = _MapRecordingPlugin() + + @durable_execution(plugins=[replay]) + def replayed_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + context.wait(Duration.from_seconds(60)) + return "done" + + completed_wait = { + "Id": wait_id, + "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 + + replay_result = replayed_handler( + _event(extra_operations=[completed_wait], updated_operation_ids=[wait_id]), + _lambda_context(), + ) + + assert replay_result["Status"] == InvocationStatus.SUCCEEDED.value + + start_operations, start_updated = replay.starts[0] + # The replay start map carries the prior state, including the wait. + assert sorted(["execution-1", wait_id]) == start_operations + # And updated_operations is the UpdatedOperationIds subset of it. + assert start_updated == [wait_id] + + status, end_operations = replay.ends[0] + assert status == InvocationStatus.SUCCEEDED.value + assert wait_id in end_operations + + +def test_updated_operations_ignores_ids_absent_from_the_map(): + """An id the execution state does not carry must not appear in the subset.""" + plugin = _MapRecordingPlugin() + + @durable_execution(plugins=[plugin]) + def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + 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(updated_operation_ids=["execution-1", "never-checkpointed"]), + _lambda_context(), + ) + + _, start_updated = plugin.starts[0] + assert start_updated == ["execution-1"] + + +def test_plugin_free_execution_still_completes(): + """The provider gate must not disturb an execution without plugins.""" + + @durable_execution + def my_handler(event: Any, context: DurableContext) -> str: # noqa: ARG001 + 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 + + result = my_handler(_event(), _lambda_context()) + + assert result["Status"] == InvocationStatus.SUCCEEDED.value 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 edcb769f..f1753626 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -916,6 +916,98 @@ def test_missing_provider_yields_empty_maps(self): self.assertEqual({}, start_info.updated_operations) self.assertEqual({}, end_info.operations) + def test_plugin_cannot_mutate_checkpointed_error(self): + """A plugin must not be able to alter the error replayed to user code. + + ``ErrorObject.stack_trace`` is a mutable list and the checkpointed error + is handed to user code on replay, so the plugin-facing view gets its own + copy. + """ + checkpoint_error = ErrorObject( + message="boom", type="Error", data=None, stack_trace=["frame-A"] + ) + operation = Operation( + operation_id="op-1", + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name="failing", + parent_id="root", + status=OperationStatus.FAILED, + start_timestamp=START_TS, + end_timestamp=END_TS, + step_details=StepDetails(attempt=1, error=checkpoint_error), + ) + + info = OperationInfo.from_operation(operation) + + # A distinct ErrorObject, and a distinct stack_trace list. + self.assertIsNot(checkpoint_error, info.error) + self.assertIsNot(checkpoint_error.stack_trace, info.error.stack_trace) + self.assertEqual(["frame-A"], info.error.stack_trace) + + # Mutating the plugin's view leaves the checkpointed error untouched. + info.error.stack_trace.append("injected-by-plugin") + info.error.stack_trace.clear() + self.assertEqual(["frame-A"], checkpoint_error.stack_trace) + + def test_no_plugins_never_invokes_the_operations_provider(self): + """durable_execution() always passes a provider; an empty executor must + not call it, at either hook.""" + calls = [] + + def provider(): + calls.append(1) + return {} + + executor = PluginExecutor(plugins=[]) + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + operations_provider=provider, + updated_operation_ids=["op-1"], + ) + executor.on_invocation_end( + output=DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, result=None, error=None + ) + ) + # The retained provider is dropped too, not just the snapshot. + self.assertIsNone(executor._operations_provider) # noqa: SLF001 + + self.assertEqual([], calls) + + def test_snapshot_is_pinned_against_a_live_provider_mapping(self): + """The snapshot must pin the point in time, whatever the provider returns. + + The SDK's own provider returns a fresh copy, but the eager snapshot is + what actually pins the hook's view, so it must not depend on that. + """ + live = {"op-1": self._operation("op-1")} + seen: list[object] = [] + + class _CapturingPlugin(DurableInstrumentationPlugin): + def on_invocation_start(_self, info): # noqa: N805 + seen.append(info) + + executor = PluginExecutor(plugins=[_CapturingPlugin()]) + with executor.run(): + executor.on_invocation_start( + execution_arn="arn:exec", + lambda_context=LAMBDA_CTX, + execution_start_time=START_TS, + is_first_invocation=True, + operations_provider=lambda: live, + ) + # Mutating the provider's own mapping after the hook must not change + # what that hook reported. + live["op-2"] = self._operation("op-2") + + (start_info,) = seen + self.assertEqual(["op-1"], list(start_info.operations)) + def test_operation_info_conversion_is_deferred_and_cached(self): conversions: list[str] = [] real_from_operation = OperationInfo.from_operation From 6d51dc7af491b9d03bc599068a282c61aba44c21 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Wed, 12 Aug 2026 21:36:27 +0000 Subject: [PATCH 5/5] fix(plugin): use plain dicts for the operation maps Addresses the remaining review comments on #629, and reverses my own earlier design decision. The lazy Mapping view produced three separate findings in a row -- an embedded lock made the info undeepcopyable, a captured closure made it unpicklable, and asdict() does not recurse through arbitrary Mappings so values stayed OperationInfo instances. Each fix added machinery to make a custom Mapping behave like a dict in a public dataclass field. Replace it with a plain dict built eagerly, matching OperationChangeInfo in the same module, which never had any of these problems. asdict() now recurses to nested dicts, pickle round-trips, deepcopy works, and the point-in-time guarantee is exact rather than argued. The trade-off: a plugin-registered execution now converts the map at both hooks even if nothing reads it. Executions with no plugins still pay nothing thanks to the existing gate, and serializability is the whole point of these fields, so paying the conversion is the right side of the trade. Also corrects the operations docstring: the initial execution state already carries the EXECUTION operation, so the map is not empty on a first invocation -- callers should use is_first_invocation. Tests: asdict yields nested dicts, pickle round-trips for populated and empty maps, and conversion happens at hook time. Refs #617 --- .../plugin.py | 149 ++++-------------- .../tests/plugin_test.py | 74 +++++++-- 2 files changed, 97 insertions(+), 126 deletions(-) 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 b8987650..ffc5a65f 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 @@ -9,7 +9,7 @@ from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from enum import Enum -from typing import Any, Callable, Iterator, MutableMapping, cast +from typing import Any, Callable, MutableMapping, cast from aws_durable_execution_sdk_python.identifier import OperationIdentifier from aws_durable_execution_sdk_python.lambda_service import ( @@ -143,87 +143,6 @@ def _to_operation_info_map( } -class _LazyOperationInfoMap(Mapping[str, OperationInfo]): - """Read-only ``operation id -> OperationInfo`` map built on first access. - - Invocation hook infos carry the execution's whole operation map. Converting - it eagerly would charge every invocation of an operation-heavy execution for - a view most plugins never read, so the conversion is deferred to the first - mapping operation and then cached. The underlying operations are snapshotted - before this map is handed out, so deferring the conversion does not move the - point in time the map describes. Behaves like a plain read-only ``dict`` - (iteration, ``len``, ``in``, ``get``, ``==`` against any mapping), and - materializes into an ordinary ``dict`` when deep-copied so the enclosing - info stays copyable and ``dataclasses.asdict()``-able. - """ - - __slots__ = ("_provider", "_resolved") - - def __init__(self, provider: Callable[[], dict[str, OperationInfo]] | None) -> None: - self._provider = provider - self._resolved: dict[str, OperationInfo] | None = None - - def _resolve(self) -> dict[str, OperationInfo]: - # Deliberately lock-free. Building the map is pure and idempotent, so a - # race can only duplicate the work, never corrupt the result, and the - # assignment below is atomic. A lock here would make the enclosing - # frozen info undeepcopyable, which breaks the plugins this view exists - # to serve: dataclasses.asdict() deep-copies non-dict fields, and a - # thread lock cannot be copied. - if self._resolved is None: - if self._provider is None: - self._resolved = {} - else: - try: - self._resolved = self._provider() - except Exception: - # A plugin-facing view must never break the execution. - logger.exception( - "Failed to build plugin operations map; using empty map" - ) - self._resolved = {} - return self._resolved - - def __getitem__(self, key: str) -> OperationInfo: - return self._resolve()[key] - - def __iter__(self) -> Iterator[str]: - return iter(self._resolve()) - - def __len__(self) -> int: - return len(self._resolve()) - - def __eq__(self, other: object) -> bool: - if isinstance(other, Mapping): - return self._resolve() == dict(other) - return NotImplemented - - def __ne__(self, other: object) -> bool: - result = self.__eq__(other) - if result is NotImplemented: - return result - return not result - - __hash__ = None # type: ignore[assignment] # mutable-by-materialization view - - def __deepcopy__(self, memo: dict[int, Any]) -> dict[str, OperationInfo]: - """Materialize into a plain ``dict`` when copied. - - Copying this view has no reason to preserve its laziness, and a plain - dict is what callers actually want: it makes - ``dataclasses.asdict(info)`` yield an ordinary mapping instead of an - opaque object, and keeps the provider closure out of the copy. - """ - resolved = { - key: copy.deepcopy(value, memo) for key, value in self._resolve().items() - } - memo[id(self)] = resolved - return resolved - - def __repr__(self) -> str: - return f"{type(self).__name__}({self._resolve()!r})" - - @dataclass(frozen=True) class OperationStartInfo(OperationInfo): pass @@ -321,7 +240,7 @@ class InvocationInfo: without it); ``durable_execution()`` always populates it with the deserialized input payload, which is ``{}`` when the payload is empty. """ - operations: Mapping[str, OperationInfo] = field( + operations: dict[str, OperationInfo] = field( default_factory=dict, kw_only=True, repr=False, @@ -333,8 +252,13 @@ class InvocationInfo: A point-in-time view of the execution's operation map: as observed at the start of the invocation on ``on_invocation_start``, and as observed at the - end of the invocation on ``on_invocation_end``. Empty on the very first - invocation-start, before any operation has been checkpointed. + end of the invocation on ``on_invocation_end``. + + Not a reliable signal of whether this is the first invocation: the initial + execution state already carries the ``EXECUTION`` operation, so even a first + invocation-start sees a non-empty map. Use + :attr:`is_first_invocation` for that. What a first invocation lacks is prior + non-execution operations. Excluded from ``repr``, ``__eq__`` and ``__hash__`` for the same reasons as :attr:`execution_input`: the entries carry operation results and errors that @@ -345,7 +269,7 @@ class InvocationInfo: @dataclass(frozen=True) class InvocationStartInfo(InvocationInfo): - updated_operations: Mapping[str, OperationInfo] = field( + updated_operations: dict[str, OperationInfo] = field( default_factory=dict, kw_only=True, repr=False, @@ -399,7 +323,7 @@ def from_durable_execution_invocation_output( cls, invocation_start_info: InvocationStartInfo, output: "DurableExecutionInvocationOutput", - operations: Mapping[str, OperationInfo] | None = None, + operations: dict[str, OperationInfo] | None = None, ): return InvocationEndInfo( request_id=invocation_start_info.request_id, @@ -561,35 +485,31 @@ def execute_plugins(self, info, sync): def _snapshot_operation_infos( self, operations_provider: Callable[[], Mapping[str, Operation]] | None, - ) -> Mapping[str, OperationInfo]: - """Capture the operation map now; defer the ``OperationInfo`` conversion. + ) -> dict[str, OperationInfo]: + """Build the plugin ``OperationInfo`` view of the current operation map. - Snapshotting eagerly pins the point in time the hook reports, so a plugin - that stashes the info and reads it later still sees the state as of its - hook; deferring the conversion keeps operation-heavy executions from - paying for a view no plugin reads. + Returns a plain ``dict``, matching :class:`OperationChangeInfo`. That + matters beyond consistency: ``dataclasses.asdict()`` and ``pickle`` only + traverse real dicts, so a custom ``Mapping`` here would leave the + enclosing hook info unserializable for the very plugins these fields + exist to serve. + + Built eagerly, which also pins the point in time the hook reports: a + plugin that stashes the info and reads it later still sees the state as + of its own hook. Skipped entirely when no plugins are registered -- ``durable_execution()`` passes a provider unconditionally, so without this gate a plugin-free - execution would still invoke it at both hooks. - - The returned mapping is copied even though the SDK's provider - (``ExecutionState.operations``) already returns a copy: the copy is what - pins the point in time, and relying on every provider to hand over a - mapping it will never touch again would make that invariant contingent on - an external contract. It is a shallow copy of pointers, and now only - happens when plugins are registered, so it is far cheaper than the - ``OperationInfo`` conversion it guards. + execution would pay for a view nothing can read. """ if not self._plugins or operations_provider is None: - return _LazyOperationInfoMap(None) + return {} try: - snapshot = dict(operations_provider()) + return _to_operation_info_map(operations_provider()) except Exception: # A plugin-facing view must never break the execution. logger.exception("Failed to snapshot operations for plugin hook") - return _LazyOperationInfoMap(None) - return _LazyOperationInfoMap(lambda: _to_operation_info_map(snapshot)) + return {} def on_invocation_start( self, @@ -609,9 +529,8 @@ def on_invocation_start( execution_start_time: Start timestamp of the execution operation. lambda_context: Lambda context, for the request id. execution_input: The deserialized execution input event. - operations_provider: Returns the current checkpointed operation map. - Called once here to snapshot it; the conversion to the plugin's - ``OperationInfo`` view is deferred to first access. + operations_provider: Returns the current checkpointed operation map, + converted here into the plugin's ``OperationInfo`` view. updated_operation_ids: Operation ids from the invocation input's ``UpdatedOperationIds`` -- those updated while suspended. """ @@ -625,13 +544,11 @@ def on_invocation_start( execution_start_time=execution_start_time, execution_input=self._snapshot_execution_input(execution_input), operations=operations, - updated_operations=_LazyOperationInfoMap( - lambda: { - operation_id: operations[operation_id] - for operation_id in (updated_operation_ids or []) - if operation_id in operations - } - ), + updated_operations={ + operation_id: operations[operation_id] + for operation_id in (updated_operation_ids or []) + if operation_id in operations + }, ) self.execute_plugins(self._invocation_status, sync=True) 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 f1753626..c33c370a 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -1,5 +1,6 @@ import datetime import logging +import pickle import unittest from copy import deepcopy from dataclasses import asdict, fields @@ -1008,7 +1009,63 @@ def on_invocation_start(_self, info): # noqa: N805 (start_info,) = seen self.assertEqual(["op-1"], list(start_info.operations)) - def test_operation_info_conversion_is_deferred_and_cached(self): + def test_maps_are_plain_dicts_for_serialization(self): + """The maps must be real dicts, not a custom Mapping. + + ``dataclasses.asdict()`` and ``pickle`` only traverse real dicts, so a + custom Mapping here would leave the hook info unserializable for exactly + the plugins these fields exist to serve. + """ + with self.executor.run(): + self._start( + operations={"op-1": self._operation("op-1")}, + updated_operation_ids=["op-1"], + ) + self.executor.on_invocation_end( + output=DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, result=None, error=None + ) + ) + + start_info, end_info = self.captured + for info in (start_info, end_info): + self.assertIs(dict, type(info.operations)) + self.assertIs(dict, type(start_info.updated_operations)) + + # asdict recurses all the way down, so the values become nested dicts. + as_dict = asdict(start_info) + self.assertIs(dict, type(as_dict["operations"])) + self.assertIs(dict, type(as_dict["operations"]["op-1"])) + self.assertEqual("op-1", as_dict["operations"]["op-1"]["operation_id"]) + self.assertIs(dict, type(as_dict["updated_operations"]["op-1"])) + + def test_hook_infos_survive_a_pickle_round_trip(self): + """Plugins hand infos to multiprocessing queues and caches.""" + with self.executor.run(): + self._start( + operations={"op-1": self._operation("op-1")}, + updated_operation_ids=["op-1"], + ) + self.executor.on_invocation_end( + output=DurableExecutionInvocationOutput( + status=InvocationStatus.SUCCEEDED, result=None, error=None + ) + ) + + start_info, end_info = self.captured + for info in (start_info, end_info): + restored = pickle.loads(pickle.dumps(info)) # noqa: S301 + self.assertEqual(["op-1"], list(restored.operations)) + self.assertEqual("op-1", restored.operations["op-1"].operation_id) + + # And with the empty maps a plugin-free-style info carries. + empty = InvocationStartInfo( + request_id="req-1", execution_arn="arn:test", is_first_invocation=True + ) + self.assertEqual({}, pickle.loads(pickle.dumps(empty)).operations) # noqa: S301 + + def test_conversion_happens_at_hook_time(self): + """The map is built during the hook, pinning the point in time.""" conversions: list[str] = [] real_from_operation = OperationInfo.from_operation @@ -1021,17 +1078,14 @@ def counting(operation, **kwargs): self.executor.run(), ): self._start(operations={"op-1": self._operation("op-1")}) - # Nothing read the map inside the hook, so nothing was converted. - self.assertEqual([], conversions) - - (start_info,) = self.captured - self.assertEqual(1, len(start_info.operations)) - self.assertEqual(["op-1"], conversions) - # Repeated reads reuse the cached conversion. - self.assertEqual(["op-1"], list(start_info.operations)) - self.assertIn("op-1", start_info.operations) + # Converted while the hook ran, not on first read afterwards. self.assertEqual(["op-1"], conversions) + (start_info,) = self.captured + self.assertEqual(["op-1"], list(start_info.operations)) + # Reading again does not reconvert. + self.assertEqual(["op-1"], conversions) + def test_end_info_reflects_operations_added_during_the_invocation(self): operations = {"op-1": self._operation("op-1")}