From 7d89a4e513041e51b7e242129fed22d5a6a0c08f Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 2 Oct 2026 23:22:03 +0000 Subject: [PATCH 01/24] fix: preserve handler worker context --- .../e2e/test_invocation_wait_resume_int.py | 131 ++++++++++++++++++ .../execution.py | 7 +- .../e2e/handler_worker_context_int_test.py | 125 +++++++++++++++++ 3 files changed, 262 insertions(+), 1 deletion(-) create mode 100644 packages/aws-durable-execution-sdk-python/tests/e2e/handler_worker_context_int_test.py diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py index 635634938..37dc53394 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py @@ -31,6 +31,9 @@ from aws_durable_execution_sdk_python_otel.execution_plugin import ExecutionOtelPlugin from aws_durable_execution_sdk_python_otel.invocation_plugin import InvocationOtelPlugin from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig +from opentelemetry import context as otel_context +from opentelemetry import trace +from opentelemetry.propagators.aws.aws_xray_propagator import AwsXRayPropagator from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter @@ -253,3 +256,131 @@ def handler_impl(_event: Any, context: DurableContext) -> str: assert completed_wait_span.end_time is not None assert after_resume.start_time is not None assert completed_wait_span.end_time <= after_resume.start_time + + +@pytest.mark.parametrize("plugin_type", [InvocationOtelPlugin, ExecutionOtelPlugin]) +@pytest.mark.parametrize("fail_after_resume", [False, True]) +def test_handler_user_spans_inherit_context_across_resume_and_failure( + monkeypatch: pytest.MonkeyPatch, + plugin_type: type[InvocationOtelPlugin] | type[ExecutionOtelPlugin], + fail_after_resume: bool, +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + monkeypatch.setenv("_X_AMZN_TRACE_ID", XRAY_TRACE_HEADER) + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = plugin_type( + OtelPluginConfig(tracer_provider=provider, enrich_logger=False) + ) + tracer = provider.get_tracer("customer") + before_context = otel_context.get_current() + calls: list[str] = [] + + def user_span(name: str) -> None: + # Ordinary instrumentation: the SDK/caller supplies the active parent. + span = tracer.start_span(name) + span.end() + + def step_body(_step_context: Any) -> str: + calls.append("step") + user_span("step-user") + return "saved" + + def handler_body(_event: Any, context: DurableContext) -> str: + user_span("handler-entry") + saved = context.step(step_body, name="before-wait") + user_span("handler-after-step") + context.wait(Duration.from_seconds(1), name="context-wait") + user_span("handler-after-resume") + if fail_after_resume: + raise ValueError("handler failed after resume") + return saved + + handler = durable_execution(handler_body, plugins=[plugin]) + remote = AwsXRayPropagator().extract({"X-Amzn-Trace-Id": XRAY_TRACE_HEADER}) + assert trace.get_current_span(remote).get_span_context().trace_id == XRAY_TRACE_ID + initial_operations = [_execution_operation()] + checkpoint, operations = _checkpoint_store(initial_operations) + ambient_ids: list[int] = [] + try: + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as client_class: + client = Mock() + client.checkpoint = checkpoint + client_class.initialize_client.return_value = client + # Standard host instrumentation supplies a same-trace Lambda span. + with tracer.start_as_current_span("lambda-first", context=remote) as host: + ambient_ids.append(host.get_span_context().span_id) + first = handler(_event(initial_operations), _lambda_context()) + assert trace.get_current_span() is host + assert first["Status"] == InvocationStatus.PENDING.value + assert otel_context.get_current() == before_context + resumed_operations = [ + replace( + operation, + status=OperationStatus.SUCCEEDED, + end_timestamp=datetime.now(UTC), + ) + if operation.name == "context-wait" + else operation + for operation in operations.values() + ] + wait_id = next( + operation.operation_id + for operation in resumed_operations + if operation.name == "context-wait" + ) + checkpoint, _ = _checkpoint_store(resumed_operations) + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as client_class: + client = Mock() + client.checkpoint = checkpoint + client_class.initialize_client.return_value = client + with tracer.start_as_current_span("lambda-resume", context=remote) as host: + ambient_ids.append(host.get_span_context().span_id) + resumed = handler( + _event(resumed_operations, updated_operation_ids=[wait_id]), + _lambda_context(), + ) + assert trace.get_current_span() is host + assert resumed["Status"] == ( + InvocationStatus.FAILED.value + if fail_after_resume + else InvocationStatus.SUCCEEDED.value + ) + assert calls == ["step"] + assert otel_context.get_current() == before_context + spans = exporter.get_finished_spans() + expected_parents = ( + [derive_workflow_span_id(EXECUTION_ARN)] * 2 + if plugin_type is ExecutionOtelPlugin + else ambient_ids + ) + for name in ("handler-entry", "handler-after-step"): + users = [span for span in spans if span.name == name] + assert len(users) == 2 + assert [span.parent.span_id if span.parent else None for span in users] == ( + expected_parents + ) + assert all( + span.context is not None and span.context.trace_id == XRAY_TRACE_ID + for span in users + ) + after_resume = next( + span for span in spans if span.name == "handler-after-resume" + ) + assert after_resume.parent is not None + assert after_resume.parent.span_id == expected_parents[1] + step_user = next(span for span in spans if span.name == "step-user") + assert step_user.parent is not None + assert any( + span.name == "before-wait attempt 1" + and span.context is not None + and span.context.span_id == step_user.parent.span_id + for span in spans + ) + finally: + provider.shutdown() 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 8ca5ef903..a52e54319 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 @@ -1,6 +1,7 @@ from __future__ import annotations import contextlib +import contextvars import functools import json import logging @@ -313,7 +314,11 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: logger.debug( "%s entering user-space...", invocation_input.durable_execution_arn ) - user_future = executor.submit(func, input_event, durable_context) + # Invocation-start hooks can establish tracing and other contextvars. + # Context.run restores worker bindings on both return and failure. + user_future = executor.submit( + contextvars.copy_context().run, func, input_event, durable_context + ) logger.debug( "%s waiting for user code completion...", diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/handler_worker_context_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/handler_worker_context_int_test.py new file mode 100644 index 000000000..d9b487c63 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/handler_worker_context_int_test.py @@ -0,0 +1,125 @@ +"""Public handler dispatch carries context without leaking worker mutations.""" + +from __future__ import annotations + +import contextvars +from collections.abc import Callable +from concurrent.futures import Future, ThreadPoolExecutor +from typing import Any +from unittest.mock import Mock + +import pytest + +from aws_durable_execution_sdk_python.context import DurableContext +from aws_durable_execution_sdk_python.exceptions import InvocationError +from aws_durable_execution_sdk_python.execution import durable_execution +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationEndInfo, + InvocationStartInfo, + InvocationStatus, +) + + +@pytest.mark.parametrize("outcome", ["success", "failure", "retry"]) +def test_handler_worker_preserves_context_and_restores_its_caller( + monkeypatch: pytest.MonkeyPatch, outcome: str +) -> None: + marker = contextvars.ContextVar("handler-worker-context", default="worker-empty") + seen: list[str] = [] + worker_boundaries: list[tuple[str, str]] = [] + statuses: list[InvocationStatus] = [] + + class ClaimPlugin(DurableInstrumentationPlugin): + token: contextvars.Token[str] | None = None + + def on_invocation_start(self, info: InvocationStartInfo) -> None: + self.token = marker.set("invocation-start") + + def on_invocation_end(self, info: InvocationEndInfo) -> None: + statuses.append(info.status) + assert self.token is not None + marker.reset(self.token) + self.token = None + + def body(_event: Any, _context: DurableContext) -> str: + seen.append(marker.get()) + marker.set("worker-mutation") + if outcome == "failure": + raise ValueError("handler failure") + if outcome == "retry": + raise InvocationError("handler retry") + return "ok" + + class ObservingExecutor(ThreadPoolExecutor): + def submit( + self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any + ) -> Future[Any]: + if fn is not body and not (args and args[0] is body): + return super().submit(fn, *args, **kwargs) + + def observe() -> Any: + before = marker.get() + try: + return fn(*args, **kwargs) + finally: + # This runs outside Context.run, on the actual SDK worker. + worker_boundaries.append((before, marker.get())) + + return super().submit(observe) + + monkeypatch.setattr( + "aws_durable_execution_sdk_python.execution.ThreadPoolExecutor", + ObservingExecutor, + ) + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + client = Mock() + handler = durable_execution(body, boto3_client=client, plugins=[ClaimPlugin()]) + event = { + "DurableExecutionArn": "test-arn/handler-context", + "CheckpointToken": "test-token", + "InitialExecutionState": { + "Operations": [ + { + "Id": "handler-context", + "Type": "EXECUTION", + "Status": "STARTED", + "ExecutionDetails": {"InputPayload": "{}"}, + } + ], + "NextMarker": "", + }, + } + lambda_context = Mock() + lambda_context.aws_request_id = "context-request" + lambda_context.client_context = None + lambda_context.identity = None + lambda_context._epoch_deadline_time_in_ms = 0 + lambda_context.invoked_function_arn = "test-arn" + lambda_context.tenant_id = None + token = marker.set("caller") + try: + if outcome == "retry": + with pytest.raises(InvocationError, match="handler retry"): + handler(event, lambda_context) + else: + result = handler(event, lambda_context) + assert result["Status"] == ( + "SUCCEEDED" if outcome == "success" else "FAILED" + ) + assert marker.get() == "caller" + finally: + marker.reset(token) + assert seen == ["invocation-start"] + assert len(worker_boundaries) == 1 + worker_before, worker_after = worker_boundaries[0] + assert worker_after == worker_before + assert worker_after != "worker-mutation" + assert statuses == [ + { + "success": InvocationStatus.SUCCEEDED, + "failure": InvocationStatus.FAILED, + "retry": InvocationStatus.RETRY, + }[outcome] + ] + client.checkpoint_durable_execution.assert_not_called() From 0ea7ea2598e399bc90587b60910c5a63002b8111 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 2 Oct 2026 23:40:11 +0000 Subject: [PATCH 02/24] test: classify handler dispatch unit coverage --- ...orker_context_int_test.py => handler_worker_context_test.py} | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) rename packages/aws-durable-execution-sdk-python/tests/{e2e/handler_worker_context_int_test.py => handler_worker_context_test.py} (98%) diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/handler_worker_context_int_test.py b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py similarity index 98% rename from packages/aws-durable-execution-sdk-python/tests/e2e/handler_worker_context_int_test.py rename to packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py index d9b487c63..1cf3fa482 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/handler_worker_context_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py @@ -1,4 +1,4 @@ -"""Public handler dispatch carries context without leaking worker mutations.""" +"""Focused handler-dispatch tests with a mock service and real worker thread.""" from __future__ import annotations From 95322aceb1f764247a945b7fe1af6bda639c29ea Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Sat, 3 Oct 2026 01:56:35 +0000 Subject: [PATCH 03/24] fix: bind missing handler trace context --- .../README.md | 9 +++++ .../invocation_plugin.py | 32 +++++++++++---- .../e2e/test_invocation_wait_resume_int.py | 36 +++++++++++++++-- .../tests/test_invocation_plugin.py | 40 +++++++++++++++++-- 4 files changed, 103 insertions(+), 14 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 5e78f2a9c..d6ab68103 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -156,6 +156,15 @@ lambda_.Function( ) ``` +### Handler context propagation + +A core SDK with handler-worker context propagation carries the context established +by invocation-start hooks into the handler. Invocation view preserves an active +ambient span on the canonical execution trace; when that context is absent or +belongs to a different trace, it makes the Invocation span current until invocation +finalization. The caller context is restored before Workflow export and flushing. +Existing plugin registration, factory lifetime, and checkpoint formats are unchanged. + ### 3. In your Lambda handler (index.py) ```python diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py index 012f17f2f..55cd5d4d0 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py @@ -159,6 +159,7 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: # owns is released through _detach_context so the plugin never leaves a # scope on the context stack. self._context_tokens: dict[str, tuple[int, object]] = {} + self._handler_context_token: tuple[int, object] | None = None self._operation_spans_lock = threading.RLock() self._tracing_enabled = False @@ -277,6 +278,12 @@ def _detach_remaining_contexts(self) -> None: for key in keys: self._detach_context(key) + def _detach_handler_context(self) -> None: + entry = self._handler_context_token + self._handler_context_token = None + if entry is not None and entry[0] == threading.get_ident(): + context.detach(entry[1]) # type: ignore[arg-type] + def get_current_span_context(self) -> SpanContext | None: """Return the span context to use for log correlation. @@ -286,12 +293,9 @@ def get_current_span_context(self) -> SpanContext | None: context this is the active context span (attached in on_user_function_start). Unrelated ambient spans are ignored so logs stay correlated to the durable execution trace. - 2. The invocation span from the plugin registry. This is the path used - for top-level handler code: the invocation span is never attached to - the worker thread's context, so the registry is the only way to - resolve it. It also covers code between top-level operations, where - detaching the operation scope restores a context with no durable - span. + 2. The invocation span from the plugin registry, including lifecycle + phases before the canonical handler context is attached. A core with + context propagation carries that context into the handler worker. Returns: A valid SpanContext, or None if no span is active. @@ -573,11 +577,19 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: self._start_workflow_span(info) - self._start_span( + invocation_span = self._start_span( operation_id=None, name="Invocation", attributes=self._extract_attributes(info), ) + ambient = trace.get_current_span().get_span_context() + if not ambient.is_valid or ambient.trace_id != self._execution_trace_id: + # Establish the fallback before the core copies context into its + # handler worker. Keep a valid same-trace caller span unchanged. + self._handler_context_token = ( + threading.get_ident(), + context.attach(trace.set_span_in_context(invocation_span)), + ) def _start_workflow_span(self, info: InvocationStartInfo) -> None: """Install a non-recording placeholder for the execution-scoped Workflow span. @@ -655,6 +667,11 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: self._reset_state() return + # User execution has finished. Restore the caller before finalization, + # so Workflow parent selection cannot adopt our own Invocation span. + self._detach_remaining_contexts() + self._detach_handler_context() + # Spans are registered parent-first, so close pending spans in reverse # order to keep every child contained within its parent. with self._operation_spans_lock: @@ -709,6 +726,7 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: def _reset_state(self) -> None: """Clear per-invocation state for warm Lambda environment reuse.""" self._detach_remaining_contexts() + self._detach_handler_context() self._execution_arn = "" self._execution_trace_id = None self._extracted_context = None diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py index 37dc53394..a75a4c18c 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py @@ -2,6 +2,7 @@ from __future__ import annotations +from contextlib import nullcontext from dataclasses import replace from datetime import UTC, datetime from typing import Any @@ -260,10 +261,12 @@ def handler_impl(_event: Any, context: DurableContext) -> str: @pytest.mark.parametrize("plugin_type", [InvocationOtelPlugin, ExecutionOtelPlugin]) @pytest.mark.parametrize("fail_after_resume", [False, True]) +@pytest.mark.parametrize("ambient_kind", ["same", "unrelated", "absent"]) def test_handler_user_spans_inherit_context_across_resume_and_failure( monkeypatch: pytest.MonkeyPatch, plugin_type: type[InvocationOtelPlugin] | type[ExecutionOtelPlugin], fail_after_resume: bool, + ambient_kind: str, ) -> None: monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) monkeypatch.setenv("_X_AMZN_TRACE_ID", XRAY_TRACE_HEADER) @@ -303,6 +306,7 @@ def handler_body(_event: Any, context: DurableContext) -> str: initial_operations = [_execution_operation()] checkpoint, operations = _checkpoint_store(initial_operations) ambient_ids: list[int] = [] + host_context = remote if ambient_kind == "same" else otel_context.Context() try: with patch( "aws_durable_execution_sdk_python.execution.LambdaClient" @@ -311,10 +315,19 @@ def handler_body(_event: Any, context: DurableContext) -> str: client.checkpoint = checkpoint client_class.initialize_client.return_value = client # Standard host instrumentation supplies a same-trace Lambda span. - with tracer.start_as_current_span("lambda-first", context=remote) as host: + host_scope = ( + tracer.start_as_current_span("lambda-first", context=host_context) + if ambient_kind != "absent" + else nullcontext() + ) + with host_scope: + host = trace.get_current_span() ambient_ids.append(host.get_span_context().span_id) first = handler(_event(initial_operations), _lambda_context()) - assert trace.get_current_span() is host + assert ( + trace.get_current_span().get_span_context() + == host.get_span_context() + ) assert first["Status"] == InvocationStatus.PENDING.value assert otel_context.get_current() == before_context resumed_operations = [ @@ -339,13 +352,22 @@ def handler_body(_event: Any, context: DurableContext) -> str: client = Mock() client.checkpoint = checkpoint client_class.initialize_client.return_value = client - with tracer.start_as_current_span("lambda-resume", context=remote) as host: + host_scope = ( + tracer.start_as_current_span("lambda-resume", context=host_context) + if ambient_kind != "absent" + else nullcontext() + ) + with host_scope: + host = trace.get_current_span() ambient_ids.append(host.get_span_context().span_id) resumed = handler( _event(resumed_operations, updated_operation_ids=[wait_id]), _lambda_context(), ) - assert trace.get_current_span() is host + assert ( + trace.get_current_span().get_span_context() + == host.get_span_context() + ) assert resumed["Status"] == ( InvocationStatus.FAILED.value if fail_after_resume @@ -358,6 +380,12 @@ def handler_body(_event: Any, context: DurableContext) -> str: [derive_workflow_span_id(EXECUTION_ARN)] * 2 if plugin_type is ExecutionOtelPlugin else ambient_ids + if ambient_kind == "same" + else [ + span.context.span_id + for span in spans + if span.name == "Invocation" and span.context is not None + ] ) for name in ("handler-entry", "handler-after-step"): users = [span for span in spans if span.name == name] diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py index f108ef728..bbd9ee372 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py @@ -721,6 +721,8 @@ def test_operation_end_without_start_links_previous_logical_operation(): span.attributes["durable.operation.status"] == OperationStatus.SUCCEEDED.value ) + plugin.on_invocation_end(_invocation_end_info()) + def test_continuation_span_uses_current_start_and_end_times(): """Continuation spans use current times within the invocation.""" @@ -750,6 +752,8 @@ def test_continuation_span_uses_current_start_and_end_times(): assert invocation_span.start_time <= span.start_time assert before_callback <= span.start_time <= span.end_time <= after_callback + plugin.on_invocation_end(_invocation_end_info()) + def test_resume_operation_timestamps_do_not_precede_current_invocation(): plugin, exporter = _create_plugin() @@ -809,6 +813,8 @@ def test_resume_operation_timestamps_do_not_precede_current_invocation(): assert after_resume_span.parent is not None assert after_resume_span.parent.span_id == invocation_span.context.span_id + plugin.on_invocation_end(_invocation_end_info()) + def test_ordered_timestamps_are_thread_safe(): plugin, _ = _create_plugin() @@ -875,6 +881,8 @@ def test_retried_operation_uses_fresh_id_and_links_previous_logical_operation(): operation_id_to_span_id(EXECUTION_ARN, operation_id), } + plugin.on_invocation_end(_invocation_end_info()) + def test_step_operation_span_parents_attempt_span(): """STEP operations have a logical span with attempt spans beneath it.""" @@ -1025,6 +1033,8 @@ def test_user_function_callbacks_emit_attempt_span_attributes(): ) assert "durable.operation.status" not in span.attributes + plugin.on_invocation_end(_invocation_end_info()) + def test_step_attempt_span_name_includes_attempt_number(): """Step attempt spans include the attempt number in the display name.""" @@ -1067,6 +1077,8 @@ def test_step_attempt_span_name_includes_attempt_number(): span = exporter.get_finished_spans()[0] assert span.name == "fetch-user attempt 2" + plugin.on_invocation_end(_invocation_end_info()) + def test_step_attempt_span_name_defaults_to_first_attempt(): """Step attempt spans default to attempt 1 when no attempt is provided.""" @@ -1109,6 +1121,8 @@ def test_step_attempt_span_name_defaults_to_first_attempt(): span = exporter.get_finished_spans()[0] assert span.name == "fetch-user attempt 1" + plugin.on_invocation_end(_invocation_end_info()) + @pytest.mark.parametrize( ("outcome", "terminal_status", "error", "expected_span_status"), @@ -1210,6 +1224,8 @@ def test_context_span_waits_for_terminal_status_and_omits_attempt_attributes( assert "durable.attempt.outcome" not in span.attributes assert span.status.status_code is expected_span_status + plugin.on_invocation_end(_invocation_end_info()) + def test_span_registry_helpers_can_be_called_from_multiple_threads(): """Verify active span registry helpers are safe under concurrent access.""" @@ -1256,6 +1272,8 @@ def test_user_function_end_restores_enclosing_context(): assert plugin._context_tokens == {} assert plugin.get_current_span_context().span_id == invocation_span_id + plugin.on_invocation_end(_invocation_end_info()) + def test_user_function_start_preserves_baggage_in_current_context(): plugin, _ = _create_plugin() @@ -1290,6 +1308,8 @@ def test_user_function_end_restores_enclosing_context_on_failure(): assert otel_context.get_current() == enclosing_context assert plugin._context_tokens == {} + plugin.on_invocation_end(_invocation_end_info()) + def test_user_function_end_restores_enclosing_context_across_multiple_steps(): """Verify sequential steps do not accumulate context scopes.""" @@ -1308,6 +1328,8 @@ def test_user_function_end_restores_enclosing_context_across_multiple_steps(): assert plugin._context_tokens == {} assert plugin.get_current_span_context().span_id == invocation_span_id + plugin.on_invocation_end(_invocation_end_info()) + # ---------------------------------------------------------------------- # get_current_span_context resolves the right span context @@ -1329,6 +1351,8 @@ def test_get_current_span_context_returns_invocation_span_at_top_level(): assert span_context is not None assert span_context.span_id == invocation_span.get_span_context().span_id + plugin.on_invocation_end(_invocation_end_info()) + def test_get_current_span_context_returns_operation_span_inside_step(): """Verify code inside a step resolves to the attempt span context.""" @@ -1360,6 +1384,8 @@ def test_get_current_span_context_returns_invocation_span_between_steps(): assert span_context is not None assert span_context.span_id == invocation_span.get_span_context().span_id + plugin.on_invocation_end(_invocation_end_info()) + # ---------------------------------------------------------------------- # on_user_function_end restores the ENCLOSING operation span (nested case) @@ -1437,12 +1463,14 @@ def test_top_level_step_end_falls_back_to_invocation_for_correlation(): plugin.on_user_function_start(_user_function_start_info(operation_id)) plugin.on_user_function_end(_user_function_end_info(operation_id)) - # No durable span is attached at the top level, so the registry fallback - # supplies the invocation span for log correlation. + # With no compatible ambient span, the handler inherits the Invocation + # context; completing a step restores that enclosing context. assert otel_context.get_current() == enclosing_context - assert not trace.get_current_span().get_span_context().is_valid + assert trace.get_current_span().get_span_context().span_id == invocation_span_id assert plugin.get_current_span_context().span_id == invocation_span_id + plugin.on_invocation_end(_invocation_end_info()) + def test_get_current_span_context_returns_context_span_between_nested_steps(): """Verify between-step code inside a child context resolves to that context. @@ -1706,6 +1734,8 @@ def test_replayed_context_span_links_previous_logical_operation(): operation_id_to_span_id(EXECUTION_ARN, operation_id), } + plugin.on_invocation_end(_invocation_end_info()) + def test_checkpointed_context_first_span_uses_deterministic_id(): plugin, exporter = _create_plugin() @@ -1865,6 +1895,8 @@ def test_child_context_end_restores_context_active_before_it(): assert otel_context.get_current() == enclosing_context assert plugin._context_tokens == {} + plugin.on_invocation_end(_invocation_end_info()) + def test_nested_scopes_are_released_without_accumulating(): """Verify a child context and its inner step unwind to their entry contexts.""" @@ -1895,6 +1927,8 @@ def test_nested_scopes_are_released_without_accumulating(): assert otel_context.get_current() == before_context assert plugin._context_tokens == {} + plugin.on_invocation_end(_invocation_end_info()) + def test_invocation_end_releases_scope_of_suspended_user_function(): """Verify a user function that never ends does not leak its scope. From 812f112d4b403bf37a9c4bdaab8185019545bd57 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Sat, 3 Oct 2026 02:21:01 +0000 Subject: [PATCH 04/24] fix: confine handler context to worker scopes --- .../README.md | 7 +- .../invocation_plugin.py | 44 ++++++------ .../e2e/test_invocation_wait_resume_int.py | 32 ++++++++- .../tests/test_invocation_plugin.py | 40 +---------- .../execution.py | 6 +- .../plugin.py | 54 +++++++++++++- .../tests/handler_worker_context_test.py | 71 ++++++++++++++++++- 7 files changed, 188 insertions(+), 66 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index d6ab68103..fd6f38eae 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -161,8 +161,11 @@ lambda_.Function( A core SDK with handler-worker context propagation carries the context established by invocation-start hooks into the handler. Invocation view preserves an active ambient span on the canonical execution trace; when that context is absent or -belongs to a different trace, it makes the Invocation span current until invocation -finalization. The caller context is restored before Workflow export and flushing. +belongs to a different trace, its optional `handler_context` scope makes the Invocation +span current only while the handler runs. The scope closes on the same worker in +reverse plugin order, including on failure and suspension, without changing the +invocation-hook caller. Older cores ignore this optional scope and retain their +existing behavior; install the updated core as well to get handler context propagation. Existing plugin registration, factory lifetime, and checkpoint formats are unchanged. ### 3. In your Lambda handler (index.py) diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py index 55cd5d4d0..beabcfeb7 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py @@ -2,9 +2,11 @@ from __future__ import annotations +import contextlib import datetime import logging import threading +from collections.abc import Iterator from typing import Any from aws_durable_execution_sdk_python.plugin import ( @@ -159,7 +161,6 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: # owns is released through _detach_context so the plugin never leaves a # scope on the context stack. self._context_tokens: dict[str, tuple[int, object]] = {} - self._handler_context_token: tuple[int, object] | None = None self._operation_spans_lock = threading.RLock() self._tracing_enabled = False @@ -278,12 +279,6 @@ def _detach_remaining_contexts(self) -> None: for key in keys: self._detach_context(key) - def _detach_handler_context(self) -> None: - entry = self._handler_context_token - self._handler_context_token = None - if entry is not None and entry[0] == threading.get_ident(): - context.detach(entry[1]) # type: ignore[arg-type] - def get_current_span_context(self) -> SpanContext | None: """Return the span context to use for log correlation. @@ -294,8 +289,7 @@ def get_current_span_context(self) -> SpanContext | None: on_user_function_start). Unrelated ambient spans are ignored so logs stay correlated to the durable execution trace. 2. The invocation span from the plugin registry, including lifecycle - phases before the canonical handler context is attached. A core with - context propagation carries that context into the handler worker. + phases outside the optional handler-worker context scope. Returns: A valid SpanContext, or None if no span is active. @@ -577,19 +571,29 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: self._start_workflow_span(info) - invocation_span = self._start_span( + self._start_span( operation_id=None, name="Invocation", attributes=self._extract_attributes(info), ) + + @contextlib.contextmanager + def handler_context(self, info: InvocationStartInfo) -> Iterator[None]: + """Bind the fallback only inside the SDK-owned handler worker scope.""" ambient = trace.get_current_span().get_span_context() - if not ambient.is_valid or ambient.trace_id != self._execution_trace_id: - # Establish the fallback before the core copies context into its - # handler worker. Keep a valid same-trace caller span unchanged. - self._handler_context_token = ( - threading.get_ident(), - context.attach(trace.set_span_in_context(invocation_span)), - ) + invocation_span = self._get_span(None) + token = None + if ( + self._tracing_enabled + and invocation_span is not None + and (not ambient.is_valid or ambient.trace_id != self._execution_trace_id) + ): + token = context.attach(trace.set_span_in_context(invocation_span)) + try: + yield + finally: + if token is not None: + context.detach(token) def _start_workflow_span(self, info: InvocationStartInfo) -> None: """Install a non-recording placeholder for the execution-scoped Workflow span. @@ -667,10 +671,9 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: self._reset_state() return - # User execution has finished. Restore the caller before finalization, - # so Workflow parent selection cannot adopt our own Invocation span. + # User execution has finished; the worker has already closed its handler + # context scope without modifying the invocation-hook caller. self._detach_remaining_contexts() - self._detach_handler_context() # Spans are registered parent-first, so close pending spans in reverse # order to keep every child contained within its parent. @@ -726,7 +729,6 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: def _reset_state(self) -> None: """Clear per-invocation state for warm Lambda environment reuse.""" self._detach_remaining_contexts() - self._detach_handler_context() self._execution_arn = "" self._execution_trace_id = None self._extracted_context = None diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py index a75a4c18c..374305bfb 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py @@ -259,7 +259,14 @@ def handler_impl(_event: Any, context: DurableContext) -> str: assert completed_wait_span.end_time <= after_resume.start_time -@pytest.mark.parametrize("plugin_type", [InvocationOtelPlugin, ExecutionOtelPlugin]) +@pytest.mark.parametrize( + ("plugin_type", "extra_context_plugin"), + [ + (InvocationOtelPlugin, False), + (InvocationOtelPlugin, True), + (ExecutionOtelPlugin, False), + ], +) @pytest.mark.parametrize("fail_after_resume", [False, True]) @pytest.mark.parametrize("ambient_kind", ["same", "unrelated", "absent"]) def test_handler_user_spans_inherit_context_across_resume_and_failure( @@ -267,6 +274,7 @@ def test_handler_user_spans_inherit_context_across_resume_and_failure( plugin_type: type[InvocationOtelPlugin] | type[ExecutionOtelPlugin], fail_after_resume: bool, ambient_kind: str, + extra_context_plugin: bool, ) -> None: monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) monkeypatch.setenv("_X_AMZN_TRACE_ID", XRAY_TRACE_HEADER) @@ -291,6 +299,8 @@ def step_body(_step_context: Any) -> str: return "saved" def handler_body(_event: Any, context: DurableContext) -> str: + if extra_context_plugin: + assert baggage.get_baggage("customer") == "present" user_span("handler-entry") saved = context.step(step_body, name="before-wait") user_span("handler-after-step") @@ -300,7 +310,25 @@ def handler_body(_event: Any, context: DurableContext) -> str: raise ValueError("handler failed after resume") return saved - handler = durable_execution(handler_body, plugins=[plugin]) + # An unrelated plugin may own a caller-thread OTel baggage scope. The + # invocation-view fallback must never become part of its saved token. + from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin + from opentelemetry import baggage + + class BaggagePlugin(DurableInstrumentationPlugin): + token: Any = None + + def on_invocation_start(self, _info: Any) -> None: + self.token = otel_context.attach(baggage.set_baggage("customer", "present")) + + def on_invocation_end(self, _info: Any) -> None: + otel_context.detach(self.token) + self.token = None + + plugins: list[DurableInstrumentationPlugin] = [plugin] + if extra_context_plugin: + plugins.append(BaggagePlugin()) + handler = durable_execution(handler_body, plugins=plugins) remote = AwsXRayPropagator().extract({"X-Amzn-Trace-Id": XRAY_TRACE_HEADER}) assert trace.get_current_span(remote).get_span_context().trace_id == XRAY_TRACE_ID initial_operations = [_execution_operation()] diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py index bbd9ee372..f108ef728 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py @@ -721,8 +721,6 @@ def test_operation_end_without_start_links_previous_logical_operation(): span.attributes["durable.operation.status"] == OperationStatus.SUCCEEDED.value ) - plugin.on_invocation_end(_invocation_end_info()) - def test_continuation_span_uses_current_start_and_end_times(): """Continuation spans use current times within the invocation.""" @@ -752,8 +750,6 @@ def test_continuation_span_uses_current_start_and_end_times(): assert invocation_span.start_time <= span.start_time assert before_callback <= span.start_time <= span.end_time <= after_callback - plugin.on_invocation_end(_invocation_end_info()) - def test_resume_operation_timestamps_do_not_precede_current_invocation(): plugin, exporter = _create_plugin() @@ -813,8 +809,6 @@ def test_resume_operation_timestamps_do_not_precede_current_invocation(): assert after_resume_span.parent is not None assert after_resume_span.parent.span_id == invocation_span.context.span_id - plugin.on_invocation_end(_invocation_end_info()) - def test_ordered_timestamps_are_thread_safe(): plugin, _ = _create_plugin() @@ -881,8 +875,6 @@ def test_retried_operation_uses_fresh_id_and_links_previous_logical_operation(): operation_id_to_span_id(EXECUTION_ARN, operation_id), } - plugin.on_invocation_end(_invocation_end_info()) - def test_step_operation_span_parents_attempt_span(): """STEP operations have a logical span with attempt spans beneath it.""" @@ -1033,8 +1025,6 @@ def test_user_function_callbacks_emit_attempt_span_attributes(): ) assert "durable.operation.status" not in span.attributes - plugin.on_invocation_end(_invocation_end_info()) - def test_step_attempt_span_name_includes_attempt_number(): """Step attempt spans include the attempt number in the display name.""" @@ -1077,8 +1067,6 @@ def test_step_attempt_span_name_includes_attempt_number(): span = exporter.get_finished_spans()[0] assert span.name == "fetch-user attempt 2" - plugin.on_invocation_end(_invocation_end_info()) - def test_step_attempt_span_name_defaults_to_first_attempt(): """Step attempt spans default to attempt 1 when no attempt is provided.""" @@ -1121,8 +1109,6 @@ def test_step_attempt_span_name_defaults_to_first_attempt(): span = exporter.get_finished_spans()[0] assert span.name == "fetch-user attempt 1" - plugin.on_invocation_end(_invocation_end_info()) - @pytest.mark.parametrize( ("outcome", "terminal_status", "error", "expected_span_status"), @@ -1224,8 +1210,6 @@ def test_context_span_waits_for_terminal_status_and_omits_attempt_attributes( assert "durable.attempt.outcome" not in span.attributes assert span.status.status_code is expected_span_status - plugin.on_invocation_end(_invocation_end_info()) - def test_span_registry_helpers_can_be_called_from_multiple_threads(): """Verify active span registry helpers are safe under concurrent access.""" @@ -1272,8 +1256,6 @@ def test_user_function_end_restores_enclosing_context(): assert plugin._context_tokens == {} assert plugin.get_current_span_context().span_id == invocation_span_id - plugin.on_invocation_end(_invocation_end_info()) - def test_user_function_start_preserves_baggage_in_current_context(): plugin, _ = _create_plugin() @@ -1308,8 +1290,6 @@ def test_user_function_end_restores_enclosing_context_on_failure(): assert otel_context.get_current() == enclosing_context assert plugin._context_tokens == {} - plugin.on_invocation_end(_invocation_end_info()) - def test_user_function_end_restores_enclosing_context_across_multiple_steps(): """Verify sequential steps do not accumulate context scopes.""" @@ -1328,8 +1308,6 @@ def test_user_function_end_restores_enclosing_context_across_multiple_steps(): assert plugin._context_tokens == {} assert plugin.get_current_span_context().span_id == invocation_span_id - plugin.on_invocation_end(_invocation_end_info()) - # ---------------------------------------------------------------------- # get_current_span_context resolves the right span context @@ -1351,8 +1329,6 @@ def test_get_current_span_context_returns_invocation_span_at_top_level(): assert span_context is not None assert span_context.span_id == invocation_span.get_span_context().span_id - plugin.on_invocation_end(_invocation_end_info()) - def test_get_current_span_context_returns_operation_span_inside_step(): """Verify code inside a step resolves to the attempt span context.""" @@ -1384,8 +1360,6 @@ def test_get_current_span_context_returns_invocation_span_between_steps(): assert span_context is not None assert span_context.span_id == invocation_span.get_span_context().span_id - plugin.on_invocation_end(_invocation_end_info()) - # ---------------------------------------------------------------------- # on_user_function_end restores the ENCLOSING operation span (nested case) @@ -1463,14 +1437,12 @@ def test_top_level_step_end_falls_back_to_invocation_for_correlation(): plugin.on_user_function_start(_user_function_start_info(operation_id)) plugin.on_user_function_end(_user_function_end_info(operation_id)) - # With no compatible ambient span, the handler inherits the Invocation - # context; completing a step restores that enclosing context. + # No durable span is attached at the top level, so the registry fallback + # supplies the invocation span for log correlation. assert otel_context.get_current() == enclosing_context - assert trace.get_current_span().get_span_context().span_id == invocation_span_id + assert not trace.get_current_span().get_span_context().is_valid assert plugin.get_current_span_context().span_id == invocation_span_id - plugin.on_invocation_end(_invocation_end_info()) - def test_get_current_span_context_returns_context_span_between_nested_steps(): """Verify between-step code inside a child context resolves to that context. @@ -1734,8 +1706,6 @@ def test_replayed_context_span_links_previous_logical_operation(): operation_id_to_span_id(EXECUTION_ARN, operation_id), } - plugin.on_invocation_end(_invocation_end_info()) - def test_checkpointed_context_first_span_uses_deterministic_id(): plugin, exporter = _create_plugin() @@ -1895,8 +1865,6 @@ def test_child_context_end_restores_context_active_before_it(): assert otel_context.get_current() == enclosing_context assert plugin._context_tokens == {} - plugin.on_invocation_end(_invocation_end_info()) - def test_nested_scopes_are_released_without_accumulating(): """Verify a child context and its inner step unwind to their entry contexts.""" @@ -1927,8 +1895,6 @@ def test_nested_scopes_are_released_without_accumulating(): assert otel_context.get_current() == before_context assert plugin._context_tokens == {} - plugin.on_invocation_end(_invocation_end_info()) - def test_invocation_end_releases_scope_of_suspended_user_function(): """Verify a user function that never ends does not leak its scope. 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 a52e54319..1d3d49e0b 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 @@ -317,7 +317,11 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: # Invocation-start hooks can establish tracing and other contextvars. # Context.run restores worker bindings on both return and failure. user_future = executor.submit( - contextvars.copy_context().run, func, input_event, durable_context + contextvars.copy_context().run, + plugin_executor.run_handler, + func, + input_event, + durable_context, ) logger.debug( 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 e9549d308..78188d641 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 @@ -5,11 +5,11 @@ import datetime import functools import logging -from collections.abc import Mapping, Sequence +from collections.abc import Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from enum import Enum -from typing import Any, Callable, MutableMapping, cast +from typing import Any, Callable, ContextManager, MutableMapping, cast from aws_durable_execution_sdk_python.identifier import OperationIdentifier from aws_durable_execution_sdk_python.lambda_service import ( @@ -400,6 +400,18 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: """ pass + def handler_context(self, info: InvocationStartInfo) -> ContextManager[None]: + """Optional scope around the top-level handler on its worker thread. + + The worker already carries a copy of the caller's context after the + invocation-start hooks. Scopes enter in plugin order and close in reverse + order on success, failure, or suspension. They are for context binding, + not exception handling: cleanup receives no handler exception and cannot + suppress or replace the handler's outcome. Existing invocation hooks keep + their original thread and ordering. Older cores ignore this optional hook. + """ + return contextlib.nullcontext() + def on_operation_start(self, info: OperationStartInfo) -> None: """ Called before an operation's START checkpoint is queued, or when a @@ -519,6 +531,44 @@ 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) + @contextlib.contextmanager + def _safe_handler_context( + self, plugin: DurableInstrumentationPlugin, info: InvocationStartInfo + ) -> Iterator[None]: + # Old plugin objects may not inherit this core's new optional method. + scope = None + try: + factory = getattr(plugin, "handler_context", None) + if factory is not None: + scope = factory(info) + scope.__enter__() + except Exception: + scope = None + logger.exception( + "Plugin %s handler context failed", plugin.__class__.__name__ + ) + try: + yield + finally: + if scope is not None: + try: + scope.__exit__(None, None, None) + except Exception: + logger.exception( + "Plugin %s handler context cleanup failed", + plugin.__class__.__name__, + ) + + def run_handler(self, handler: Callable[..., Any], *args: Any) -> Any: + """Run the user handler inside optional, balanced plugin context scopes.""" + with contextlib.ExitStack() as scopes: + if self._invocation_status is not None: + for plugin in self._plugins: + scopes.enter_context( + self._safe_handler_context(plugin, self._invocation_status) + ) + return handler(*args) + def _snapshot_operation_infos( self, operations_provider: Callable[[], Mapping[str, Operation]] | None, diff --git a/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py index 1cf3fa482..91e0e5474 100644 --- a/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py @@ -55,7 +55,7 @@ class ObservingExecutor(ThreadPoolExecutor): def submit( self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any ) -> Future[Any]: - if fn is not body and not (args and args[0] is body): + if fn is not body and body not in args: return super().submit(fn, *args, **kwargs) def observe() -> Any: @@ -123,3 +123,72 @@ def observe() -> Any: }[outcome] ] client.checkpoint_durable_execution.assert_not_called() + + +@pytest.mark.parametrize("failure", [None, "body", "enter", "exit"]) +def test_optional_handler_scopes_are_balanced_and_cannot_change_outcome( + failure: str | None, +) -> None: + from contextlib import contextmanager + from datetime import UTC, datetime + from collections.abc import Iterator + from aws_durable_execution_sdk_python.plugin import PluginExecutor + + marker = contextvars.ContextVar("handler-scope", default="caller") + events: list[str] = [] + + class ScopePlugin(DurableInstrumentationPlugin): + def __init__(self, name: str): + self.name = name + + @contextmanager + def handler_context(self, info: InvocationStartInfo) -> Iterator[None]: + assert info.execution_arn == "handler-scope" + events.append("enter-" + self.name) + if self.name == "inner" and failure == "enter": + raise ValueError("plugin entry failure") + token = marker.set(self.name) + try: + yield + finally: + marker.reset(token) + events.append("exit-" + self.name) + if self.name == "inner" and failure == "exit": + raise ValueError("plugin cleanup failure") + + executor = PluginExecutor([ScopePlugin("outer"), ScopePlugin("inner")]) + + def body() -> str: + assert marker.get() == ("outer" if failure == "enter" else "inner") + events.append("body") + if failure in ("body", "exit"): + raise RuntimeError("original handler failure") + return "ok" + + with executor.run(): + executor.on_invocation_start("handler-scope", True, datetime.now(UTC), None) + if failure in ("body", "exit"): + with pytest.raises(RuntimeError, match="original handler failure"): + executor.run_handler(body) + else: + assert executor.run_handler(body) == "ok" + assert marker.get() == "caller" + assert events == ["enter-outer", "enter-inner", "body"] + ( + ["exit-outer"] if failure == "enter" else ["exit-inner", "exit-outer"] + ) + + +def test_handler_accepts_plugin_without_optional_scope() -> None: + from datetime import UTC, datetime + from types import SimpleNamespace + from typing import cast + from aws_durable_execution_sdk_python.plugin import PluginExecutor + + legacy = cast( + DurableInstrumentationPlugin, + SimpleNamespace(on_invocation_start=lambda info: None), + ) + executor = PluginExecutor([legacy]) + with executor.run(): + executor.on_invocation_start("legacy", True, datetime.now(UTC), None) + assert executor.run_handler(lambda: "unchanged") == "unchanged" From 370295fbec57b748928781b68744726e367f714c Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Sat, 3 Oct 2026 02:55:43 +0000 Subject: [PATCH 05/24] test: retain published-core context expectations --- .../e2e/test_invocation_wait_resume_int.py | 24 ++++++++++++++----- 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py index 374305bfb..35e3f7756 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py @@ -26,6 +26,7 @@ OperationType, StepDetails, ) +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( derive_workflow_span_id, ) @@ -278,6 +279,10 @@ def test_handler_user_spans_inherit_context_across_resume_and_failure( ) -> None: monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) monkeypatch.setenv("_X_AMZN_TRACE_ID", XRAY_TRACE_HEADER) + # The documented PyPI compatibility environment deliberately uses an older + # core. Keep exercising its supported operation tracing and lifecycle while + # asserting the new handler contract only when that core exposes the scope. + supports_handler_context = hasattr(DurableInstrumentationPlugin, "handler_context") exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) @@ -300,7 +305,9 @@ def step_body(_step_context: Any) -> str: def handler_body(_event: Any, context: DurableContext) -> str: if extra_context_plugin: - assert baggage.get_baggage("customer") == "present" + assert baggage.get_baggage("customer") == ( + "present" if supports_handler_context else None + ) user_span("handler-entry") saved = context.step(step_body, name="before-wait") user_span("handler-after-step") @@ -312,7 +319,6 @@ def handler_body(_event: Any, context: DurableContext) -> str: # An unrelated plugin may own a caller-thread OTel baggage scope. The # invocation-view fallback must never become part of its saved token. - from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin from opentelemetry import baggage class BaggagePlugin(DurableInstrumentationPlugin): @@ -405,7 +411,9 @@ def on_invocation_end(self, _info: Any) -> None: assert otel_context.get_current() == before_context spans = exporter.get_finished_spans() expected_parents = ( - [derive_workflow_span_id(EXECUTION_ARN)] * 2 + [None, None] + if not supports_handler_context + else [derive_workflow_span_id(EXECUTION_ARN)] * 2 if plugin_type is ExecutionOtelPlugin else ambient_ids if ambient_kind == "same" @@ -422,14 +430,18 @@ def on_invocation_end(self, _info: Any) -> None: expected_parents ) assert all( - span.context is not None and span.context.trace_id == XRAY_TRACE_ID + span.context is not None + and ( + (span.context.trace_id == XRAY_TRACE_ID) == supports_handler_context + ) for span in users ) after_resume = next( span for span in spans if span.name == "handler-after-resume" ) - assert after_resume.parent is not None - assert after_resume.parent.span_id == expected_parents[1] + assert ( + after_resume.parent.span_id if after_resume.parent else None + ) == expected_parents[1] step_user = next(span for span in spans if span.name == "step-user") assert step_user.parent is not None assert any( From e1f2cc9a479cb7a97aec53e5995359e351fd34ec Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Sat, 3 Oct 2026 02:58:08 +0000 Subject: [PATCH 06/24] test: isolate the published-core OTel environment --- pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 428d680e3..bab6bff5f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -128,6 +128,8 @@ dependencies = [ test = "pytest packages/aws-durable-execution-sdk-python-examples/test {args}" [tool.hatch.envs.test-pypi-otel] +# Do not inherit the default workspace core when validating a published release. +workspace.members = ["packages/aws-durable-execution-sdk-python-otel"] dependencies = [ "aws-durable-execution-sdk-python>=2.0.0", "opentelemetry-sdk>=1.20.0", From bdb6aca293f5e77614400c2b620752cd2da9a082 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Sat, 3 Oct 2026 03:09:52 +0000 Subject: [PATCH 07/24] test: pin legacy OTel compatibility coverage --- CONTRIBUTING.md | 6 +++++- .../tests/test_package_metadata.py | 4 ++-- pyproject.toml | 2 +- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 3c750e8df..58a152125 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -75,10 +75,14 @@ hatch run dev-examples:test # run examples tests only To verify packages work against the published PyPI version of the core SDK (rather than the local workspace): ```bash -hatch run test-pypi-otel:test # test otel against PyPI core SDK +hatch run test-pypi-otel:test # test otel against the minimum supported core (2.0.0) hatch run test-pypi-examples:test # test examples against PyPI core SDK ``` +The OTel PyPI environment excludes the local core and pins the minimum supported +release, so newer PyPI releases cannot remove legacy compatibility coverage. Use +`hatch run dev-otel:test` for the current workspace core. + ### Package-level commands Some commands still run from within a package directory: diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py index 22b416384..cc4c01ee9 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py @@ -74,9 +74,9 @@ def test_test_environments_install_layer_provided_dependencies() -> None: assert TEST_OTEL_DEPENDENCIES <= set(environments["types"]["extra-dependencies"]) -def test_pypi_compatibility_environment_uses_compatible_core_sdk() -> None: +def test_pypi_compatibility_environment_pins_minimum_supported_core() -> None: dependencies = _load_pyproject(REPOSITORY_ROOT / "pyproject.toml")["tool"]["hatch"][ "envs" ]["test-pypi-otel"]["dependencies"] - assert CORE_DEPENDENCY in dependencies + assert "aws-durable-execution-sdk-python==2.0.0" in dependencies diff --git a/pyproject.toml b/pyproject.toml index bab6bff5f..3628f4ae4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -131,7 +131,7 @@ test = "pytest packages/aws-durable-execution-sdk-python-examples/test {args}" # Do not inherit the default workspace core when validating a published release. workspace.members = ["packages/aws-durable-execution-sdk-python-otel"] dependencies = [ - "aws-durable-execution-sdk-python>=2.0.0", + "aws-durable-execution-sdk-python==2.0.0", "opentelemetry-sdk>=1.20.0", "opentelemetry-propagator-aws-xray", "pytest", From 687f8f091f69a252d0bc2a8866a71ea3b98eb8de Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Fri, 2 Oct 2026 20:55:24 -0700 Subject: [PATCH 08/24] ci: verify minimum-core OTel compatibility --- .github/workflows/ci.yml | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 63c4b44c6..0bd1f6bfc 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -60,6 +60,18 @@ jobs: run: hatch run types:check - name: Run tests + coverage run: hatch run test:cov + - name: Test OTel with the minimum supported core + run: | + hatch run test-pypi-otel:python - <<'PYTHON' + from importlib.metadata import version + from pathlib import Path + import aws_durable_execution_sdk_python.plugin as plugin + + assert version("aws-durable-execution-sdk-python") == "2.0.0" + assert "site-packages" in Path(plugin.__file__).parts + assert not hasattr(plugin.DurableInstrumentationPlugin, "handler_context") + PYTHON + hatch run test-pypi-otel:test - name: Build distribution run: | for pkg in packages/*/; do From d459e8e107ec33f2dbe6b2ca9c506aa8b08805b8 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Sat, 3 Oct 2026 04:51:04 +0000 Subject: [PATCH 09/24] test: cover OTel handler context scopes directly --- .../tests/test_invocation_plugin.py | 68 +++++++++++++++++++ 1 file changed, 68 insertions(+) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py index f108ef728..7e8fb2065 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py @@ -2198,3 +2198,71 @@ def test_nested_suspension_unwinds_scopes_in_reverse_order(): assert plugin._context_tokens == {} plugin.on_invocation_end(_invocation_end_info()) + + +@pytest.mark.parametrize("ambient_kind", ["same", "unrelated", "absent"]) +@pytest.mark.parametrize("raises", [False, True]) +@pytest.mark.parametrize("sampler", [ALWAYS_ON, ALWAYS_OFF]) +def test_handler_context_preserves_or_replaces_parent_and_restores_baggage( + ambient_kind: str, raises: bool, sampler: Sampler +) -> None: + plugin, _ = _create_plugin_with_sampler(sampler) + info = _invocation_start_info() + plugin.on_invocation_start(info) + invocation = plugin.get_current_span_context() + assert invocation is not None and invocation.is_valid + caller = baggage.set_baggage("tenant", "scope-test", Context()) + expected = invocation + if ambient_kind != "absent": + ambient = SpanContext( + trace_id=( + invocation.trace_id + if ambient_kind == "same" + else (1 if invocation.trace_id != 1 else 2) + ), + span_id=0x42, + is_remote=False, + trace_flags=invocation.trace_flags, + ) + caller = trace.set_span_in_context(NonRecordingSpan(ambient), caller) + if ambient_kind == "same": + expected = ambient + token = otel_context.attach(caller) + error = ValueError("handler error") + try: + + def body() -> None: + with plugin.handler_context(info): + active = trace.get_current_span().get_span_context() + assert active == expected + assert active.is_valid + assert baggage.get_baggage("tenant") == "scope-test" + if raises: + raise error + + if raises: + with pytest.raises(ValueError) as caught: + body() + assert caught.value is error + else: + body() + assert otel_context.get_current() is caller + assert baggage.get_baggage("tenant") == "scope-test" + finally: + otel_context.detach(token) + plugin.on_invocation_end(_invocation_end_info()) + + +@pytest.mark.parametrize("completed", [False, True]) +def test_handler_context_without_live_invocation_leaves_context_unchanged( + completed: bool, +) -> None: + plugin, _ = _create_plugin() + info = _invocation_start_info() + if completed: + plugin.on_invocation_start(info) + plugin.on_invocation_end(_invocation_end_info()) + caller = otel_context.get_current() + with plugin.handler_context(info): + assert otel_context.get_current() is caller + assert otel_context.get_current() is caller From 7d7a1a3540129cac5146ed349397c6d7bbe8486d Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Sat, 3 Oct 2026 18:30:29 -0700 Subject: [PATCH 10/24] fix: isolate invocation plugin context bindings --- .../README.md | 11 +++ .../e2e/test_invocation_wait_resume_int.py | 14 ++- .../plugin.py | 14 ++- .../handler_invocation_context_int_test.py | 79 ++++++++++++++++ .../tests/handler_worker_context_test.py | 89 +++++++++++++++++++ 5 files changed, 204 insertions(+), 3 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python/tests/e2e/handler_invocation_context_int_test.py diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index fd6f38eae..c697dcc13 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -309,6 +309,17 @@ The resolved decision is applied to Workflow, Invocation, operation, and attempt spans. This avoids independently querying stateful or ratio-based samplers for each durable span in the same invocation. +Invocation hooks retain their caller thread and registration order. With the +updated core, invocation-local context-variable bindings are isolated from the +host: hooks see the incoming context and the handler receives their resulting +context, while invocation exit restores the host's original bindings even if a +plugin fails during setup or cleanup. Plugins must not use invocation context +bindings to mutate the host context after the invocation has returned. Older +supported cores retain their existing lifecycle behavior, including the +execution-view limitation when later plugins open invocation context scopes. +The new isolation applies only when plugins are registered. + + ### Log Correlation When `enrich_logger=True` (the default), the plugin installs a logging filter on diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py index 35e3f7756..0a45dba2e 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py @@ -266,8 +266,17 @@ def handler_impl(_event: Any, context: DurableContext) -> str: (InvocationOtelPlugin, False), (InvocationOtelPlugin, True), (ExecutionOtelPlugin, False), - ], + ] + # Execution-view caller isolation requires the coordinated newer core. + # Released 2.0.x retains the pre-existing same-order teardown limitation; + # the legacy lane continues checking its supported combinations above. + + ( + [(ExecutionOtelPlugin, True)] + if hasattr(DurableInstrumentationPlugin, "handler_context") + else [] + ), ) +@pytest.mark.parametrize("reverse_plugins", [False, True]) @pytest.mark.parametrize("fail_after_resume", [False, True]) @pytest.mark.parametrize("ambient_kind", ["same", "unrelated", "absent"]) def test_handler_user_spans_inherit_context_across_resume_and_failure( @@ -276,6 +285,7 @@ def test_handler_user_spans_inherit_context_across_resume_and_failure( fail_after_resume: bool, ambient_kind: str, extra_context_plugin: bool, + reverse_plugins: bool, ) -> None: monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) monkeypatch.setenv("_X_AMZN_TRACE_ID", XRAY_TRACE_HEADER) @@ -334,6 +344,8 @@ def on_invocation_end(self, _info: Any) -> None: plugins: list[DurableInstrumentationPlugin] = [plugin] if extra_context_plugin: plugins.append(BaggagePlugin()) + if reverse_plugins: + plugins.reverse() handler = durable_execution(handler_body, plugins=plugins) remote = AwsXRayPropagator().extract({"X-Amzn-Trace-Id": XRAY_TRACE_HEADER}) assert trace.get_current_span(remote).get_span_context().trace_id == XRAY_TRACE_ID 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 78188d641..5ae5beb51 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 contextvars import copy import datetime import functools @@ -890,8 +891,7 @@ def _is_terminal_status(status): @property def handle_durable_output(self): def decorator(func: Callable[[Any, LambdaContext], MutableMapping[str, Any]]): - @functools.wraps(func) - def wrapper(event: Any, context: LambdaContext): + def invoke(event: Any, context: LambdaContext): with self.run(): try: output = func(event, context) @@ -908,6 +908,16 @@ def wrapper(event: Any, context: LambdaContext): ) raise + @functools.wraps(func) + def wrapper(event: Any, context: LambdaContext): + if not self._plugins: + return invoke(event, context) + # Keep hooks on their existing caller thread and in registration + # order, but isolate their context bindings from the host. Two + # plugins can otherwise restore a stale predecessor when their + # invocation-end hooks close scopes in the original order. + return contextvars.copy_context().run(invoke, event, context) + return wrapper return decorator diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/handler_invocation_context_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/handler_invocation_context_int_test.py new file mode 100644 index 000000000..59f5c9b52 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/handler_invocation_context_int_test.py @@ -0,0 +1,79 @@ +"""Invocation context isolation across real suspension and replay.""" + +from __future__ import annotations + +import contextvars +import json +from typing import Any + +import pytest +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner + +from aws_durable_execution_sdk_python import DurableContext, durable_execution +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationEndInfo, + InvocationStartInfo, +) + + +@pytest.mark.parametrize("fail_after_resume", [False, True]) +@pytest.mark.parametrize("reverse", [False, True]) +def test_invocation_plugins_restore_host_context_across_resume( + monkeypatch: pytest.MonkeyPatch, fail_after_resume: bool, reverse: bool +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + marker = contextvars.ContextVar("invocation-host", default="host") + boundaries: list[tuple[str, str]] = [] + body_calls: list[str] = [] + + class ScopePlugin(DurableInstrumentationPlugin): + def __init__(self, name: str) -> None: + self.name = name + self.token: contextvars.Token[str] | None = None + + def on_invocation_start(self, _info: InvocationStartInfo) -> None: + self.token = marker.set(self.name) + + def on_invocation_end(self, _info: InvocationEndInfo) -> None: + assert self.token is not None + marker.reset(self.token) + self.token = None + + names = ["first", "second"] + if reverse: + names.reverse() + + def body(_event: Any, context: DurableContext) -> str: + assert marker.get() == names[-1] + + def step(_step_context: Any) -> str: + body_calls.append("step") + return "saved" + + saved = context.step(step, name="before-wait") + context.wait(Duration.from_seconds(1), name="resume") + assert marker.get() == names[-1] + if fail_after_resume: + raise ValueError("failure after resume") + return saved + + durable_handler = durable_execution( + body, plugins=[ScopePlugin(name) for name in names] + ) + + def host(event: Any, context: Any) -> Any: + before = marker.get() + try: + return durable_handler(event, context) + finally: + boundaries.append((before, marker.get())) + + with DurableFunctionTestRunner(handler=host) as runner: + result = runner.run(input="{}", timeout=15) + assert result.status.value == ("FAILED" if fail_after_resume else "SUCCEEDED") + if not fail_after_resume: + assert json.loads(result.result) == "saved" + assert body_calls == ["step"] + assert boundaries == [("host", "host"), ("host", "host")] diff --git a/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py index 91e0e5474..5f5eb8a86 100644 --- a/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py @@ -192,3 +192,92 @@ def test_handler_accepts_plugin_without_optional_scope() -> None: with executor.run(): executor.on_invocation_start("legacy", True, datetime.now(UTC), None) assert executor.run_handler(lambda: "unchanged") == "unchanged" + + +@pytest.mark.parametrize("outcome", ["SUCCEEDED", "PENDING", "FAILED", "retry"]) +@pytest.mark.parametrize("reverse", [False, True]) +@pytest.mark.parametrize("hook_failure", [None, "start", "end"]) +def test_invocation_context_scopes_do_not_escape_to_host( + outcome: str, reverse: bool, hook_failure: str | None +) -> None: + """Legacy hook order must not leave an already-ended plugin scope current.""" + from datetime import UTC, datetime + import threading + from aws_durable_execution_sdk_python.plugin import PluginExecutor + + marker = contextvars.ContextVar("invocation-scope", default="host") + events: list[tuple[str, str, str, int]] = [] + caller_thread = threading.get_ident() + + class ScopePlugin(DurableInstrumentationPlugin): + def __init__(self, name: str) -> None: + self.name = name + self.token: contextvars.Token[str] | None = None + + def on_invocation_start(self, _info: InvocationStartInfo) -> None: + events.append(("start", self.name, marker.get(), threading.get_ident())) + self.token = marker.set(self.name) + if self.name == names[0] and hook_failure == "start": + raise ValueError("plugin initialization failed") + + def on_invocation_end(self, _info: InvocationEndInfo) -> None: + events.append(("end", self.name, marker.get(), threading.get_ident())) + if self.name == names[0] and hook_failure == "end": + raise ValueError("plugin finalization failed") + assert self.token is not None + marker.reset(self.token) + self.token = None + + names = ["first", "second"] + if reverse: + names.reverse() + executor = PluginExecutor([ScopePlugin(name) for name in names]) + handler_failure = InvocationError("retry") + expected_output = {"Status": outcome} + + @executor.handle_durable_output + def invoke(_event: Any, _context: Any) -> dict[str, str]: + executor.on_invocation_start("test", True, datetime.now(UTC), None) + assert marker.get() == names[-1] + if outcome == "retry": + raise handler_failure + return expected_output + + token = marker.set("incoming") + try: + for _ in range(2): + if outcome == "retry": + with pytest.raises(InvocationError, match="retry") as caught: + invoke({}, None) + assert caught.value is handler_failure + else: + assert invoke({}, None) is expected_output + assert marker.get() == "incoming" + finally: + marker.reset(token) + assert [(kind, name) for kind, name, _, _ in events] == [ + (kind, name) for _ in range(2) for kind in ("start", "end") for name in names + ] + assert all(thread == caller_thread for _, _, _, thread in events) + assert [ + value for kind, name, value, _ in events if kind == "start" and name == names[0] + ] == ["incoming", "incoming"] + + +def test_no_plugin_invocation_keeps_original_caller_context_semantics() -> None: + from aws_durable_execution_sdk_python.plugin import PluginExecutor + + marker = contextvars.ContextVar("no-plugin-caller", default="host") + executor = PluginExecutor([]) + + @executor.handle_durable_output + def invoke(_event: Any, _context: Any) -> dict[str, str]: + marker.set("caller-side-change") + return {"Status": "SUCCEEDED"} + + token = marker.set("incoming") + try: + invoke({}, None) + assert marker.get() == "caller-side-change" + finally: + marker.reset(token) From d62289f65c927a9d44d5ea8753be7671819e56fd Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Sat, 3 Oct 2026 18:56:33 -0700 Subject: [PATCH 11/24] fix: preserve uninstrumented worker context --- .../execution.py | 26 ++++++++------ .../tests/handler_worker_context_test.py | 35 ++++++++++++------- 2 files changed, 39 insertions(+), 22 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 1d3d49e0b..732f6d6e8 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 @@ -189,7 +189,8 @@ def durable_execution( logger.debug("Starting durable execution handler...") - plugin_executor = PluginExecutor(load_configured_plugins(plugins)) + configured_plugins = load_configured_plugins(plugins) + plugin_executor = PluginExecutor(configured_plugins) @plugin_executor.handle_durable_output def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: @@ -314,15 +315,20 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: logger.debug( "%s entering user-space...", invocation_input.durable_execution_arn ) - # Invocation-start hooks can establish tracing and other contextvars. - # Context.run restores worker bindings on both return and failure. - user_future = executor.submit( - contextvars.copy_context().run, - plugin_executor.run_handler, - func, - input_event, - durable_context, - ) + if configured_plugins: + # Invocation-start hooks can establish tracing and other contextvars. + # Context.run restores worker bindings on both return and failure. + user_future = executor.submit( + contextvars.copy_context().run, + plugin_executor.run_handler, + func, + input_event, + durable_context, + ) + else: + # Preserve the original fresh-worker context for uninstrumented + # handlers, including the absence of caller ContextVar bindings. + user_future = executor.submit(func, input_event, durable_context) logger.debug( "%s waiting for user code completion...", diff --git a/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py index 5f5eb8a86..850341351 100644 --- a/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py @@ -22,8 +22,9 @@ @pytest.mark.parametrize("outcome", ["success", "failure", "retry"]) +@pytest.mark.parametrize("with_plugin", [False, True]) def test_handler_worker_preserves_context_and_restores_its_caller( - monkeypatch: pytest.MonkeyPatch, outcome: str + monkeypatch: pytest.MonkeyPatch, outcome: str, with_plugin: bool ) -> None: marker = contextvars.ContextVar("handler-worker-context", default="worker-empty") seen: list[str] = [] @@ -74,7 +75,9 @@ def observe() -> Any: ) monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) client = Mock() - handler = durable_execution(body, boto3_client=client, plugins=[ClaimPlugin()]) + handler = durable_execution( + body, boto3_client=client, plugins=[ClaimPlugin()] if with_plugin else [] + ) event = { "DurableExecutionArn": "test-arn/handler-context", "CheckpointToken": "test-token", @@ -110,18 +113,26 @@ def observe() -> Any: assert marker.get() == "caller" finally: marker.reset(token) - assert seen == ["invocation-start"] assert len(worker_boundaries) == 1 worker_before, worker_after = worker_boundaries[0] - assert worker_after == worker_before - assert worker_after != "worker-mutation" - assert statuses == [ - { - "success": InvocationStatus.SUCCEEDED, - "failure": InvocationStatus.FAILED, - "retry": InvocationStatus.RETRY, - }[outcome] - ] + if with_plugin: + assert seen == ["invocation-start"] + assert worker_after == worker_before + assert worker_after != "worker-mutation" + assert statuses == [ + { + "success": InvocationStatus.SUCCEEDED, + "failure": InvocationStatus.FAILED, + "retry": InvocationStatus.RETRY, + }[outcome] + ] + else: + # No plugin means the original direct worker call: caller bindings are + # absent, and mutations belong to the worker's own context. + assert seen == [worker_before] + assert seen != ["caller"] + assert worker_after == "worker-mutation" + assert statuses == [] client.checkpoint_durable_execution.assert_not_called() From 7f558a5d3f8fd22ae34893441d3f0456da2048db Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Mon, 5 Oct 2026 20:15:58 -0700 Subject: [PATCH 12/24] fix: isolate failed plugin context setup --- .../README.md | 6 +- .../plugin.py | 77 ++++++++-- .../tests/handler_worker_context_test.py | 141 +++++++++++++++++- 3 files changed, 203 insertions(+), 21 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index c697dcc13..8aa433a1b 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -317,7 +317,11 @@ plugin fails during setup or cleanup. Plugins must not use invocation context bindings to mutate the host context after the invocation has returned. Older supported cores retain their existing lifecycle behavior, including the execution-view limitation when later plugins open invocation context scopes. -The new isolation applies only when plugins are registered. +The new isolation applies only when plugins are registered. If an invocation-start +hook or handler-scope entry raises, subsequent setup and the handler retain the +bindings from before that hook. Successful scopes still clean up in their original +context, preserving token ownership. This isolates context-variable bindings; +it does not undo a plugin's mutations to shared objects or external side effects. ### Log Correlation 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 5ae5beb51..fd19fab47 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 @@ -409,7 +409,9 @@ def handler_context(self, info: InvocationStartInfo) -> ContextManager[None]: order on success, failure, or suspension. They are for context binding, not exception handling: cleanup receives no handler exception and cannot suppress or replace the handler's outcome. Existing invocation hooks keep - their original thread and ordering. Older cores ignore this optional hook. + their original thread and ordering. Failed setup bindings are discarded; + successful scopes are closed in the Context where they entered so their + ContextVar tokens remain valid. Older cores ignore this optional hook. """ return contextlib.nullcontext() @@ -479,6 +481,8 @@ def __init__(self, plugins: list[DurableInstrumentationPlugin] | None): self._executor: ThreadPoolExecutor | None = None self._invocation_status: InvocationStartInfo | None = None self._operations_provider: Callable[[], Mapping[str, Operation]] | None = None + self._startup_context: contextvars.Context | None = None + self._invocation_contexts: list[contextvars.Context | None] = [] @contextlib.contextmanager def run(self): @@ -492,12 +496,14 @@ def run(self): finally: self._invocation_status = None self._operations_provider = None + self._startup_context = None + self._invocation_contexts.clear() # Shut down the thread pool, waiting for pending tasks to complete. if self._executor: self._executor.shutdown(wait=True) @staticmethod - def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None: + def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> bool: """Invoke the appropriate plugin callback. Runs inside the thread pool.""" try: match info: @@ -520,14 +526,44 @@ def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> None: except Exception: # log and ignore the exception logger.exception("Plugin %s exception ignored", plugin.__class__.__name__) + return False + return True def execute_plugins(self, info, sync): if not self._executor: return - for plugin in self._plugins: - if sync: - # this is called synchronously, so plugins will be able to manipulate thread local objects - self._dispatch_plugin(plugin, info) + if sync and isinstance(info, InvocationStartInfo): + self._startup_context = None + self._invocation_contexts.clear() + for index, plugin in enumerate(self._plugins): + if sync and isinstance(info, InvocationStartInfo): + owner = self._startup_context + before = ( + owner.copy() if owner is not None else contextvars.copy_context() + ) + self._invocation_contexts.append(owner) + succeeded = ( + owner.run(self._dispatch_plugin, plugin, info) + if owner is not None + else self._dispatch_plugin(plugin, info) + ) + if not succeeded: + # A failing hook may have left new bindings with no reset token. + # Continue setup and the handler in the pre-hook snapshot. + self._startup_context = before + elif sync: + # End hooks must reset tokens in the Context that created them, + # even when a failed start hook moved later setup to a snapshot. + owner = ( + self._invocation_contexts[index] + if isinstance(info, InvocationEndInfo) + and index < len(self._invocation_contexts) + else None + ) + if owner is not None: + owner.run(self._dispatch_plugin, plugin, info) + else: + self._dispatch_plugin(plugin, info) else: # this is called asynchronously, so plugins cannot manipulate thread local objects self._executor.submit(self._dispatch_plugin, plugin, info) @@ -535,9 +571,10 @@ def execute_plugins(self, info, sync): @contextlib.contextmanager def _safe_handler_context( self, plugin: DurableInstrumentationPlugin, info: InvocationStartInfo - ) -> Iterator[None]: + ) -> Iterator[bool]: # Old plugin objects may not inherit this core's new optional method. scope = None + succeeded = True try: factory = getattr(plugin, "handler_context", None) if factory is not None: @@ -545,11 +582,12 @@ def _safe_handler_context( scope.__enter__() except Exception: scope = None + succeeded = False logger.exception( "Plugin %s handler context failed", plugin.__class__.__name__ ) try: - yield + yield succeeded finally: if scope is not None: try: @@ -562,13 +600,24 @@ def _safe_handler_context( def run_handler(self, handler: Callable[..., Any], *args: Any) -> Any: """Run the user handler inside optional, balanced plugin context scopes.""" - with contextlib.ExitStack() as scopes: - if self._invocation_status is not None: - for plugin in self._plugins: - scopes.enter_context( - self._safe_handler_context(plugin, self._invocation_status) - ) + if not self._plugins or self._invocation_status is None: return handler(*args) + owner = ( + self._startup_context.copy() + if self._startup_context is not None + else contextvars.copy_context() + ) + with contextlib.ExitStack() as scopes: + for plugin in self._plugins: + before = owner.copy() + scope = self._safe_handler_context(plugin, self._invocation_status) + succeeded = owner.run(scope.__enter__) + # Keep finalizers with their entry Context: ContextVar tokens + # cannot be reset in a copy, even when its bindings are identical. + scopes.callback(owner.run, scope.__exit__, None, None, None) + if not succeeded: + owner = before + return owner.run(handler, *args) def _snapshot_operation_infos( self, diff --git a/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py index 850341351..ed8e83429 100644 --- a/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py @@ -22,9 +22,9 @@ @pytest.mark.parametrize("outcome", ["success", "failure", "retry"]) -@pytest.mark.parametrize("with_plugin", [False, True]) +@pytest.mark.parametrize("plugin_mode", ["none", "healthy", "partial-failure"]) def test_handler_worker_preserves_context_and_restores_its_caller( - monkeypatch: pytest.MonkeyPatch, outcome: str, with_plugin: bool + monkeypatch: pytest.MonkeyPatch, outcome: str, plugin_mode: str ) -> None: marker = contextvars.ContextVar("handler-worker-context", default="worker-empty") seen: list[str] = [] @@ -36,6 +36,8 @@ class ClaimPlugin(DurableInstrumentationPlugin): def on_invocation_start(self, info: InvocationStartInfo) -> None: self.token = marker.set("invocation-start") + if plugin_mode == "partial-failure": + raise ValueError("partial plugin setup") def on_invocation_end(self, info: InvocationEndInfo) -> None: statuses.append(info.status) @@ -76,7 +78,9 @@ def observe() -> Any: monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) client = Mock() handler = durable_execution( - body, boto3_client=client, plugins=[ClaimPlugin()] if with_plugin else [] + body, + boto3_client=client, + plugins=[ClaimPlugin()] if plugin_mode != "none" else [], ) event = { "DurableExecutionArn": "test-arn/handler-context", @@ -115,8 +119,10 @@ def observe() -> Any: marker.reset(token) assert len(worker_boundaries) == 1 worker_before, worker_after = worker_boundaries[0] - if with_plugin: - assert seen == ["invocation-start"] + if plugin_mode != "none": + assert seen == [ + "caller" if plugin_mode == "partial-failure" else "invocation-start" + ] assert worker_after == worker_before assert worker_after != "worker-mutation" assert statuses == [ @@ -249,7 +255,7 @@ def on_invocation_end(self, _info: InvocationEndInfo) -> None: @executor.handle_durable_output def invoke(_event: Any, _context: Any) -> dict[str, str]: executor.on_invocation_start("test", True, datetime.now(UTC), None) - assert marker.get() == names[-1] + assert executor.run_handler(marker.get) == names[-1] if outcome == "retry": raise handler_failure return expected_output @@ -292,3 +298,126 @@ def invoke(_event: Any, _context: Any) -> dict[str, str]: assert marker.get() == "caller-side-change" finally: marker.reset(token) + + +@pytest.mark.parametrize("stage", ["start", "factory", "enter"]) +@pytest.mark.parametrize("bad_first", [False, True]) +@pytest.mark.parametrize("outcome", ["SUCCEEDED", "PENDING", "FAILED", "retry"]) +def test_failed_plugin_setup_discards_partial_context_bindings( + stage: str, bad_first: bool, outcome: str, caplog: pytest.LogCaptureFixture +) -> None: + """Keep healthy bindings, unset bindings, hook order, and token ownership.""" + from contextlib import nullcontext + from datetime import UTC, datetime + from typing import ContextManager + from aws_durable_execution_sdk_python.plugin import PluginExecutor + + marker = contextvars.ContextVar("partial-setup", default="default") + new_binding = contextvars.ContextVar[str]("partial-setup-no-default") + events: list[tuple[str, str]] = [] + cleanup: list[str] = [] + healthy_inputs: list[tuple[str, str | None]] = [] + + class Scope: + def __init__(self, plugin: SetupPlugin) -> None: + self.plugin = plugin + + def __enter__(self) -> None: + self.plugin.bind("enter") + + def __exit__(self, *_args: Any) -> None: + self.plugin.reset("exit") + + class SetupPlugin(DurableInstrumentationPlugin): + def __init__(self, name: str) -> None: + self.name = name + self.token: contextvars.Token[str] | None = None + + def bind(self, where: str) -> None: + events.append((where, self.name)) + if self.name == "healthy": + healthy_inputs.append((marker.get(), new_binding.get(None))) + self.token = marker.set(self.name) + if self.name == "bad": + new_binding.set("partial") + raise ValueError("partial plugin setup") + + def reset(self, where: str) -> None: + events.append((where, self.name)) + assert self.token is not None + marker.reset(self.token) + self.token = None + cleanup.append(self.name) + + def on_invocation_start(self, _info: InvocationStartInfo) -> None: + if stage == "start": + self.bind("start") + + def on_invocation_end(self, _info: InvocationEndInfo) -> None: + if stage == "start": + self.reset("end") + + def handler_context(self, _info: InvocationStartInfo) -> ContextManager[None]: + if stage == "start": + return nullcontext() + if self.name == "bad" and stage == "factory": + self.bind("factory") + return Scope(self) + + names = ["bad", "healthy"] if bad_first else ["healthy", "bad"] + executor = PluginExecutor([SetupPlugin(name) for name in names]) + output = {"Status": outcome} + failure = InvocationError("original retry") + handler_failure = RuntimeError("original handler error") + + def body() -> dict[str, str]: + assert marker.get() == "healthy" + with pytest.raises(LookupError): + new_binding.get() + if outcome == "retry": + raise failure + if outcome == "FAILED": + raise handler_failure + return output + + @executor.handle_durable_output + def invoke(_event: Any, _context: Any) -> dict[str, str]: + executor.on_invocation_start("partial-setup", True, datetime.now(UTC), None) + try: + return executor.run_handler(body) + except RuntimeError as error: + assert error is handler_failure + return output + + token = marker.set("incoming") + try: + for _ in range(2): + if outcome == "retry": + with pytest.raises(InvocationError) as caught: + invoke({}, None) + assert caught.value is failure + else: + assert invoke({}, None) is output + assert marker.get() == "incoming" + with pytest.raises(LookupError): + new_binding.get() + finally: + marker.reset(token) + assert healthy_inputs == [("incoming", None)] * 2 + assert cleanup == (names if stage == "start" else ["healthy"]) * 2 + setup = "start" if stage == "start" else "enter" + expected_setup = [ + ("factory" if name == "bad" and stage == "factory" else setup, name) + for name in names + ] + expected_cleanup = ( + [("end", name) for name in names] if stage == "start" else [("exit", "healthy")] + ) + assert events == (expected_setup + expected_cleanup) * 2 + # A token reset in the wrong Context is caught by the SDK, so explicitly + # check diagnostics and successful cleanup rather than relying on raises. + errors = [ + record.exc_info for record in caplog.records if record.exc_info is not None + ] + assert len(errors) == 2 + assert all(str(error[1]) == "partial plugin setup" for error in errors) From b5cdeac9cef0f149ae37e201167d86aeeddfba77 Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Mon, 5 Oct 2026 20:38:54 -0700 Subject: [PATCH 13/24] fix: require explicit handler scope opt-in --- .github/workflows/ci.yml | 1 + .../README.md | 5 + .../execution_plugin.py | 2 + .../invocation_plugin.py | 2 + .../e2e/test_invocation_wait_resume_int.py | 13 +- .../README.md | 19 +++ .../plugin.py | 37 +++-- .../tests/handler_worker_context_test.py | 139 ++++++++++++++++++ 8 files changed, 200 insertions(+), 18 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0bd1f6bfc..174a998bd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -70,6 +70,7 @@ jobs: assert version("aws-durable-execution-sdk-python") == "2.0.0" assert "site-packages" in Path(plugin.__file__).parts assert not hasattr(plugin.DurableInstrumentationPlugin, "handler_context") + assert not hasattr(plugin, "DURABLE_INSTRUMENTATION_HANDLER_CONTEXT_API_VERSION") PYTHON hatch run test-pypi-otel:test - name: Build distribution diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 8aa433a1b..9199dd726 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -166,6 +166,11 @@ span current only while the handler runs. The scope closes on the same worker in reverse plugin order, including on failure and suspension, without changing the invocation-hook caller. Older cores ignore this optional scope and retain their existing behavior; install the updated core as well to get handler context propagation. +The bundled OTel classes explicitly opt in with `__durable_handler_context_api__ = 1`. +Custom subclasses must repeat that literal marker on their own concrete class to +use the scope; inherited or instance markers are ignored. Unopted legacy helpers +and properties with the same name are never inspected. Older cores ignore the +marker without importing any new core API. Existing plugin registration, factory lifetime, and checkpoint formats are unchanged. ### 3. In your Lambda handler (index.py) diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py index 1a46cafcb..66f9ad0be 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py @@ -123,6 +123,8 @@ class ExecutionOtelPlugin(DurableInstrumentationPlugin): span). """ + __durable_handler_context_api__ = 1 + def __init__(self, config: OtelPluginConfig | None = None) -> None: self._config = config or OtelPluginConfig() self._context_extractor: ContextExtractor = ( diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py index beabcfeb7..980764f3e 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py @@ -106,6 +106,8 @@ class InvocationOtelPlugin(DurableInstrumentationPlugin): provider installed by the ADOT Lambda layer). """ + __durable_handler_context_api__ = 1 + DEFAULT_INSTRUMENT_NAME = "aws-durable-execution-sdk-python" def __init__(self, config: OtelPluginConfig | None = None) -> None: diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py index 0a45dba2e..89b88d599 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py @@ -26,6 +26,7 @@ OperationType, StepDetails, ) +from aws_durable_execution_sdk_python import plugin as core_plugin_api from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( derive_workflow_span_id, @@ -272,7 +273,10 @@ def handler_impl(_event: Any, context: DurableContext) -> str: # the legacy lane continues checking its supported combinations above. + ( [(ExecutionOtelPlugin, True)] - if hasattr(DurableInstrumentationPlugin, "handler_context") + if getattr( + core_plugin_api, "DURABLE_INSTRUMENTATION_HANDLER_CONTEXT_API_VERSION", None + ) + == 1 else [] ), ) @@ -292,7 +296,12 @@ def test_handler_user_spans_inherit_context_across_resume_and_failure( # The documented PyPI compatibility environment deliberately uses an older # core. Keep exercising its supported operation tracing and lifecycle while # asserting the new handler contract only when that core exposes the scope. - supports_handler_context = hasattr(DurableInstrumentationPlugin, "handler_context") + supports_handler_context = ( + getattr( + core_plugin_api, "DURABLE_INSTRUMENTATION_HANDLER_CONTEXT_API_VERSION", None + ) + == 1 + ) exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) diff --git a/packages/aws-durable-execution-sdk-python/README.md b/packages/aws-durable-execution-sdk-python/README.md index bf7776f45..bfaf92157 100644 --- a/packages/aws-durable-execution-sdk-python/README.md +++ b/packages/aws-durable-execution-sdk-python/README.md @@ -77,6 +77,25 @@ Provider names must be unique across installed distributions. Missing, ambiguous, incompatible, or invalid providers raise `PluginLoadError` during handler initialization with the provider and distribution details. +### Optional handler context scopes + +A plugin can declare `__durable_handler_context_api__ = 1` directly on its +concrete class and implement `handler_context(info)` returning a context manager. +The updated core enters these scopes around the top-level handler on its worker +thread, in registration order, and closes them in reverse order. Cleanup receives +no handler exception and cannot suppress or replace its outcome. Invocation hooks +retain their original thread and order; failed setup bindings are discarded and +successful scope cleanup stays in the context that owns its tokens. + +The marker must be the literal integer `1`; instance and inherited markers do not +opt in. A subclass must redeclare the marker to adopt this new hook. Unopted legacy +helpers, properties and dynamic attributes named `handler_context` are untouched. +The generic plugin base supplies neither a marker nor a default method. Core +support is advertised by module constant +`aws_durable_execution_sdk_python.plugin.DURABLE_INSTRUMENTATION_HANDLER_CONTEXT_API_VERSION`. +Older cores ignore this optional API. The provider API version and dependency +requirements are unchanged. + ## 🚀 Quick Start Install the execution SDK: 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 fd19fab47..f9e202017 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 @@ -10,7 +10,7 @@ from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from enum import Enum -from typing import Any, Callable, ContextManager, 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 ( @@ -30,6 +30,7 @@ logger = logging.getLogger(__name__) DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION = 1 +DURABLE_INSTRUMENTATION_HANDLER_CONTEXT_API_VERSION = 1 class InvocationStatus(Enum): @@ -401,20 +402,6 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: """ pass - def handler_context(self, info: InvocationStartInfo) -> ContextManager[None]: - """Optional scope around the top-level handler on its worker thread. - - The worker already carries a copy of the caller's context after the - invocation-start hooks. Scopes enter in plugin order and close in reverse - order on success, failure, or suspension. They are for context binding, - not exception handling: cleanup receives no handler exception and cannot - suppress or replace the handler's outcome. Existing invocation hooks keep - their original thread and ordering. Failed setup bindings are discarded; - successful scopes are closed in the Context where they entered so their - ContextVar tokens remain valid. Older cores ignore this optional hook. - """ - return contextlib.nullcontext() - def on_operation_start(self, info: OperationStartInfo) -> None: """ Called before an operation's START checkpoint is queued, or when a @@ -475,6 +462,18 @@ class DurableInstrumentationPluginProvider: plugin_api_version: int +def _handler_context_api_enabled( + plugin_type: type[DurableInstrumentationPlugin], +) -> bool: + """Read only the concrete class namespace, bypassing metaclass descriptors.""" + namespace = type.__dict__["__dict__"].__get__(plugin_type, type(plugin_type)) + version = namespace.get("__durable_handler_context_api__") + return ( + type(version) is int + and version == DURABLE_INSTRUMENTATION_HANDLER_CONTEXT_API_VERSION + ) + + class PluginExecutor: def __init__(self, plugins: list[DurableInstrumentationPlugin] | None): self._plugins = plugins or [] @@ -576,7 +575,13 @@ def _safe_handler_context( scope = None succeeded = True try: - factory = getattr(plugin, "handler_context", None) + # Old plugins may have an unrelated helper/property with this name. + # Never even inspect it unless this concrete class explicitly opts in. + factory = ( + getattr(plugin, "handler_context", None) + if _handler_context_api_enabled(type(plugin)) + else None + ) if factory is not None: scope = factory(info) scope.__enter__() diff --git a/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py index ed8e83429..215624ea3 100644 --- a/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py @@ -155,6 +155,8 @@ def test_optional_handler_scopes_are_balanced_and_cannot_change_outcome( events: list[str] = [] class ScopePlugin(DurableInstrumentationPlugin): + __durable_handler_context_api__ = 1 + def __init__(self, name: str): self.name = name @@ -329,6 +331,8 @@ def __exit__(self, *_args: Any) -> None: self.plugin.reset("exit") class SetupPlugin(DurableInstrumentationPlugin): + __durable_handler_context_api__ = 1 + def __init__(self, name: str) -> None: self.name = name self.token: contextvars.Token[str] | None = None @@ -421,3 +425,138 @@ def invoke(_event: Any, _context: Any) -> dict[str, str]: ] assert len(errors) == 2 assert all(str(error[1]) == "partial plugin setup" for error in errors) + + +@pytest.mark.parametrize("shape", ["helper", "property", "dynamic"]) +def test_unopted_legacy_handler_context_is_never_inspected(shape: str) -> None: + from contextlib import contextmanager + from collections.abc import Iterator + from datetime import UTC, datetime + from aws_durable_execution_sdk_python.plugin import PluginExecutor + + calls: list[str] = [] + + @contextmanager + def helper(_info: InvocationStartInfo) -> Iterator[None]: + calls.append("helper") + yield + + def property_getter(_self: Any) -> Any: + calls.append("property") + return helper + + def dynamic_getter(_self: Any, name: str) -> Any: + if name == "handler_context": + calls.append("dynamic") + return "legacy-business-value" + raise AttributeError(name) + + members_by_shape: dict[str, dict[str, Any]] = { + "helper": {"handler_context": staticmethod(helper)}, + "property": {"handler_context": property(property_getter)}, + "dynamic": {"__getattr__": dynamic_getter}, + } + plugin_type = type( + "Legacy", (DurableInstrumentationPlugin,), members_by_shape[shape] + ) + plugin = plugin_type() + if shape == "dynamic": + assert plugin.handler_context == "legacy-business-value" + calls.clear() + executor = PluginExecutor([plugin]) + with executor.run(): + executor.on_invocation_start("legacy-helper", True, datetime.now(UTC), None) + assert executor.run_handler(lambda: list(calls)) == [] + assert calls == [] + + +@pytest.mark.parametrize("marker", [None, 0, 2, True, "1", property(lambda _: 1)]) +def test_handler_scope_requires_literal_class_local_version(marker: Any) -> None: + from contextlib import nullcontext + from datetime import UTC, datetime + from aws_durable_execution_sdk_python.plugin import PluginExecutor + + calls: list[str] = [] + + def helper(_self: Any, _info: InvocationStartInfo) -> Any: + calls.append("scope") + return nullcontext() + + plugin_type = type( + "Legacy", + (DurableInstrumentationPlugin,), + {"__durable_handler_context_api__": marker, "handler_context": helper}, + ) + executor = PluginExecutor([plugin_type()]) + with executor.run(): + executor.on_invocation_start("invalid-marker", True, datetime.now(UTC), None) + assert executor.run_handler(lambda: "ok") == "ok" + assert calls == [] + + +def test_handler_scope_opt_in_is_not_inherited_or_taken_from_instance() -> None: + from contextlib import contextmanager + from collections.abc import Iterator + from datetime import UTC, datetime + from aws_durable_execution_sdk_python.plugin import PluginExecutor + + marker = contextvars.ContextVar("explicit-scope", default="outside") + events: list[str] = [] + + class OptedPlugin(DurableInstrumentationPlugin): + __durable_handler_context_api__ = 1 + + @contextmanager + def handler_context(self, _info: InvocationStartInfo) -> Iterator[None]: + events.append("enter") + token = marker.set("inside") + try: + yield + finally: + marker.reset(token) + events.append("exit") + + class LegacySubclass(OptedPlugin): + pass + + class ExplicitSubclass(OptedPlugin): + __durable_handler_context_api__ = 1 + + for plugin, expected in [ + (OptedPlugin(), "inside"), + (LegacySubclass(), "outside"), + (ExplicitSubclass(), "inside"), + ]: + # Assigning a marker to an instance cannot accidentally enable the hook. + plugin.__durable_handler_context_api__ = 1 + executor = PluginExecutor([plugin]) + with executor.run(): + executor.on_invocation_start("subclass", True, datetime.now(UTC), None) + assert executor.run_handler(marker.get) == expected + assert marker.get() == "outside" + assert events == ["enter", "exit", "enter", "exit"] + + +def test_handler_opt_in_does_not_trigger_legacy_metaclass_descriptors() -> None: + from datetime import UTC, datetime + from aws_durable_execution_sdk_python.plugin import PluginExecutor + + reads: list[str] = [] + + def namespace(_cls: Any) -> Any: + reads.append("metaclass-dict") + raise RuntimeError("legacy namespace") + + def helper(_self: Any, _info: Any) -> Any: + reads.append("legacy-helper") + raise RuntimeError("legacy helper") + + meta = type("LegacyMeta", (type,), {"__dict__": property(namespace)}) + plugin_type = meta( + "Legacy", (DurableInstrumentationPlugin,), {"handler_context": helper} + ) + executor = PluginExecutor([plugin_type()]) + with executor.run(): + executor.on_invocation_start("metaclass", True, datetime.now(UTC), None) + assert executor.run_handler(lambda: "ok") == "ok" + assert reads == [] From 40d9edf552c82fbfd79763e53a04b7324cac61a1 Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Mon, 5 Oct 2026 21:01:06 -0700 Subject: [PATCH 14/24] fix: bind execution view in the handler worker --- .github/workflows/cloud-tests.yml | 4 +- .../plugin/plugin_wait_replay_flag.py | 9 ++- .../tests/test_plugin_wait_replay_flag.py | 72 +++++++++++++++++++ .../README.md | 4 ++ .../execution_plugin.py | 20 ++++++ .../e2e/test_invocation_wait_resume_int.py | 28 +++++++- .../tests/test_execution_plugin.py | 71 ++++++++++++++++++ 7 files changed, 203 insertions(+), 5 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/tests/test_plugin_wait_replay_flag.py diff --git a/.github/workflows/cloud-tests.yml b/.github/workflows/cloud-tests.yml index 3d3b69950..fc4b0f57e 100644 --- a/.github/workflows/cloud-tests.yml +++ b/.github/workflows/cloud-tests.yml @@ -114,7 +114,9 @@ jobs: echo "Could not resolve the latest ADOT Python layer for $AWS_REGION" exit 1 fi - aws lambda get-layer-version-by-arn \ + # Parallel jobs can throttle this read; keep retries local to the lookup. + AWS_RETRY_MODE=standard AWS_MAX_ATTEMPTS=8 \ + aws lambda get-layer-version-by-arn \ --arn "$ADOT_LAYER_ARN" \ --region "$AWS_REGION" \ --query LayerVersionArn \ diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_replay_flag.py b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_replay_flag.py index eef6d3eed..2998740d8 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_replay_flag.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/handlers/plugin/plugin_wait_replay_flag.py @@ -18,6 +18,7 @@ """ import json +from threading import Lock from typing import Any from aws_durable_execution_sdk_python.config import Duration, ParallelConfig @@ -31,13 +32,19 @@ ) +_log_lock = Lock() + + def _emit(record: dict[str, Any], execution_arn: str | None) -> None: # Prefix every plugin record with the execution ARN as a top-level field so # the conformance runner's CloudWatch JSON filter can scope logs to a single # execution. Omit the field when the ARN is unset (never invent a value). if execution_arn: record = {"durableExecutionArn": execution_arn, **record} - print(json.dumps(record), flush=True) + # Start and end hooks can run on different threads. Keep print's separate + # body/newline writes together so the runner receives one JSON per line. + with _log_lock: + print(json.dumps(record), flush=True) class WaitReplayFlagPlugin(DurableInstrumentationPlugin): diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests/tests/test_plugin_wait_replay_flag.py b/packages/aws-durable-execution-sdk-python-conformance-tests/tests/test_plugin_wait_replay_flag.py new file mode 100644 index 000000000..0c6006043 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests/tests/test_plugin_wait_replay_flag.py @@ -0,0 +1,72 @@ +"""Concurrent plugin callbacks must emit separate parseable JSON records.""" + +from __future__ import annotations + +import importlib.util +import json +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from threading import Event +from typing import Any + +import pytest + +import aws_durable_execution_sdk_python.execution as execution + + +def test_concurrent_wait_hooks_keep_complete_stdout_records( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Import the real fixture without constructing a Lambda client: this test + # exercises its stdout producer, not a deployed durable invocation. + monkeypatch.setattr(execution, "durable_execution", lambda **_: lambda fn: fn) + path = ( + Path(__file__).resolve().parents[1] + / "handlers/plugin/plugin_wait_replay_flag.py" + ) + spec = importlib.util.spec_from_file_location("wait_replay_log_fixture", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + first_body = Event() + second_ready = Event() + second_done = Event() + chunks: list[str] = [] + + class FragmentingStdout: + def write(self, text: str) -> int: + chunks.append(text) + if '"operation-start"' in text: + first_body.set() + # print writes its body and newline separately. Permit the + # other real hook to run between them unless _emit serializes it. + second_done.wait(0.2) + return len(text) + + def flush(self) -> None: + pass + + records: list[dict[str, Any]] = [ + {"plugin": "CONFPLUGIN", "hook": "operation-start", "name": "long"}, + {"plugin": "CONFPLUGIN", "hook": "operation-end", "name": "short"}, + ] + + def emit_end() -> None: + second_ready.set() + assert first_body.wait(2) + module._emit(records[1], "execution-arn") + second_done.set() + + with monkeypatch.context() as capture: + capture.setattr("sys.stdout", FragmentingStdout()) + with ThreadPoolExecutor(max_workers=2) as executor: + end = executor.submit(emit_end) + assert second_ready.wait(2) + start = executor.submit(module._emit, records[0], "execution-arn") + start.result(timeout=2) + end.result(timeout=2) + actual = [json.loads(line) for line in "".join(chunks).splitlines() if line] + assert actual == [ + {"durableExecutionArn": "execution-arn", **record} for record in records + ] diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 9199dd726..b99ec4dd7 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -171,6 +171,10 @@ Custom subclasses must repeat that literal marker on their own concrete class to use the scope; inherited or instance markers are ignored. Unopted legacy helpers and properties with the same name are never inspected. Older cores ignore the marker without importing any new core API. +Execution view similarly restores the Workflow span inside the handler scope if +another invocation-start hook clears the active span or switches to an unrelated +trace. Both views retain valid same-trace parents and baggage, and restore the +worker's previous context when the scope ends. Existing plugin registration, factory lifetime, and checkpoint formats are unchanged. ### 3. In your Lambda handler (index.py) diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py index 66f9ad0be..c6bc2584a 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py @@ -31,9 +31,11 @@ from __future__ import annotations +import contextlib import datetime import logging import threading +from collections.abc import Iterator from typing import Any from aws_durable_execution_sdk_python.plugin import ( @@ -489,6 +491,24 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: ), ) + @contextlib.contextmanager + def handler_context(self, info: InvocationStartInfo) -> Iterator[None]: + """Keep handler instrumentation on this execution's trace in its worker.""" + ambient = trace.get_current_span().get_span_context() + workflow = self._workflow_span + token = None + if ( + self._tracing_enabled + and workflow is not None + and (not ambient.is_valid or ambient.trace_id != self._execution_trace_id) + ): + token = otel_context.attach(trace.set_span_in_context(workflow)) + try: + yield + finally: + if token is not None: + otel_context.detach(token) + def _start_workflow_span(self, info: InvocationStartInfo) -> None: """Install a non-recording placeholder for the execution-scoped Workflow span. diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py index 89b88d599..427854750 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py @@ -273,6 +273,7 @@ def handler_impl(_event: Any, context: DurableContext) -> str: # the legacy lane continues checking its supported combinations above. + ( [(ExecutionOtelPlugin, True)] + + [(ExecutionOtelPlugin, kind) for kind in ("same", "unrelated", "absent")] if getattr( core_plugin_api, "DURABLE_INSTRUMENTATION_HANDLER_CONTEXT_API_VERSION", None ) @@ -288,7 +289,7 @@ def test_handler_user_spans_inherit_context_across_resume_and_failure( plugin_type: type[InvocationOtelPlugin] | type[ExecutionOtelPlugin], fail_after_resume: bool, ambient_kind: str, - extra_context_plugin: bool, + extra_context_plugin: bool | str, reverse_plugins: bool, ) -> None: monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) @@ -344,7 +345,23 @@ class BaggagePlugin(DurableInstrumentationPlugin): token: Any = None def on_invocation_start(self, _info: Any) -> None: - self.token = otel_context.attach(baggage.set_baggage("customer", "present")) + current = baggage.set_baggage("customer", "present") + if isinstance(extra_context_plugin, str): + # A real third-party invocation hook can bind or clear a span. + # User functions below still use ordinary implicit parenting. + parent = trace.SpanContext( + trace_id=(XRAY_TRACE_ID if extra_context_plugin == "same" else 1), + span_id=0xCAFE, + is_remote=False, + trace_flags=trace.TraceFlags(1), + ) + current = trace.set_span_in_context( + trace.INVALID_SPAN + if extra_context_plugin == "absent" + else trace.NonRecordingSpan(parent), + current, + ) + self.token = otel_context.attach(current) def on_invocation_end(self, _info: Any) -> None: otel_context.detach(self.token) @@ -434,7 +451,12 @@ def on_invocation_end(self, _info: Any) -> None: expected_parents = ( [None, None] if not supports_handler_context - else [derive_workflow_span_id(EXECUTION_ARN)] * 2 + else [ + 0xCAFE + if extra_context_plugin == "same" and not reverse_plugins + else derive_workflow_span_id(EXECUTION_ARN) + ] + * 2 if plugin_type is ExecutionOtelPlugin else ambient_ids if ambient_kind == "same" diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py index 5c3cc3eca..f9043c816 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py @@ -28,6 +28,7 @@ from opentelemetry import baggage, trace from opentelemetry.context import Context from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.sampling import ALWAYS_ON, ALWAYS_OFF, Sampler from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import ( @@ -1620,3 +1621,73 @@ def test_nested_suspension_unwinds_scopes_in_reverse_order(): plugin.on_invocation_end(_invocation_end_info()) assert plugin._context_tokens == {} + + +@pytest.mark.parametrize("ambient_kind", ["same", "unrelated", "absent"]) +@pytest.mark.parametrize("raises", [False, True]) +@pytest.mark.parametrize("sampler", [ALWAYS_ON, ALWAYS_OFF]) +def test_handler_scope_keeps_execution_trace_and_restores_context( + ambient_kind: str, + raises: bool, + sampler: Sampler, +) -> None: + provider = TracerProvider(sampler=sampler) + plugin = ExecutionOtelPlugin( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + info = _invocation_start_info() + plugin.on_invocation_start(info) + workflow = trace.get_current_span().get_span_context() + assert workflow.is_valid + assert workflow.span_id == derive_workflow_span_id(EXECUTION_ARN) + caller = baggage.set_baggage("customer", "preserved", Context()) + expected = workflow + if ambient_kind != "absent": + ambient = SpanContext( + trace_id=workflow.trace_id if ambient_kind == "same" else 1, + span_id=0xCAFE, + is_remote=False, + trace_flags=workflow.trace_flags, + ) + caller = trace.set_span_in_context(NonRecordingSpan(ambient), caller) + if ambient_kind == "same": + expected = ambient + token = otel_context.attach(caller) + error = ValueError("user failure") + try: + + def body() -> None: + with plugin.handler_context(info): + assert trace.get_current_span().get_span_context() == expected + assert baggage.get_baggage("customer") == "preserved" + if raises: + raise error + + if raises: + with pytest.raises(ValueError) as caught: + body() + assert caught.value is error + else: + body() + assert otel_context.get_current() is caller + finally: + otel_context.detach(token) + plugin.on_invocation_end(_invocation_end_info()) + provider.shutdown() + + +@pytest.mark.parametrize("completed", [False, True]) +def test_handler_scope_without_live_workflow_is_noop(completed: bool) -> None: + plugin, _ = _create_plugin() + info = _invocation_start_info() + if completed: + plugin.on_invocation_start(info) + plugin.on_invocation_end(_invocation_end_info()) + caller = otel_context.get_current() + with plugin.handler_context(info): + assert otel_context.get_current() is caller + assert otel_context.get_current() is caller From 66f86d3222546f5dac5aafba6dcfde3b592c7d06 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 7 Oct 2026 00:04:24 +0000 Subject: [PATCH 15/24] ci: route OTel conformance to CodeBuild --- .../scripts/tests/test_opentelemetry_conformance_workflow.py | 1 + .github/workflows/opentelemetry-conformance-tests.yml | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py index 891cfd1be..8afaba36d 100644 --- a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py +++ b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py @@ -24,6 +24,7 @@ def test_opentelemetry_conformance_caller_uses_current_workflow_contract() -> No for configuration in ( "language: python", + "runs_on: codebuild-github-actions-runner-${{ github.run_id }}-${{ github.run_attempt }}", "resource_prefix: p", "sdk_repository: aws/aws-durable-execution-sdk-python", "sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }}", diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index b31fff281..8f01e56a6 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -59,9 +59,10 @@ jobs: actions: write contents: read id-token: write - uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@a628f5589bbf067a441696c792ab3023c0d0899b + uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@3542852f918f515d7c6a53af1eb0d5aa83d1eb1e with: language: python + runs_on: codebuild-github-actions-runner-${{ github.run_id }}-${{ github.run_attempt }} resource_prefix: p sdk_repository: aws/aws-durable-execution-sdk-python sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }} From a02876989f3c15f78ce6d07901b4802362130697 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 7 Oct 2026 00:36:36 +0000 Subject: [PATCH 16/24] ci: preserve queued conformance runs --- ...test_opentelemetry_conformance_workflow.py | 41 ++++++++++++++++--- .github/workflows/cloud-tests.yml | 1 + .github/workflows/conformance-tests.yml | 1 + .../opentelemetry-conformance-tests.yml | 7 ++++ 4 files changed, 45 insertions(+), 5 deletions(-) diff --git a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py index 8afaba36d..034be671e 100644 --- a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py +++ b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py @@ -1,12 +1,12 @@ from pathlib import Path +import yaml + WORKFLOW_PATH = ( Path(__file__).parents[2] / "workflows" / "opentelemetry-conformance-tests.yml" ) -EXAMPLES_DIR = ( - ".build/durable-sdk/packages/aws-durable-execution-sdk-python-conformance-tests-otel" -) +EXAMPLES_DIR = ".build/durable-sdk/packages/aws-durable-execution-sdk-python-conformance-tests-otel" def test_opentelemetry_conformance_caller_uses_current_workflow_contract() -> None: @@ -67,8 +67,39 @@ def test_opentelemetry_conformance_runs_when_the_handlers_change() -> None: workflow = WORKFLOW_PATH.read_text() trigger_path = ( - " - " - '"packages/aws-durable-execution-sdk-python-conformance-tests-otel/**"' + ' - "packages/aws-durable-execution-sdk-python-conformance-tests-otel/**"' ) # Once for pull_request, once for push. assert workflow.count(trigger_path) == 2 + + +def test_opentelemetry_conformance_queues_complete_runs() -> None: + workflow = yaml.safe_load(WORKFLOW_PATH.read_text()) + + assert workflow["concurrency"] == { + "group": "otel-conformance-tests", + "cancel-in-progress": False, + "queue": "max", + } + + +def test_cloud_tests_queue_each_shared_runtime_stack() -> None: + workflow = yaml.safe_load(WORKFLOW_PATH.with_name("cloud-tests.yml").read_text()) + + assert workflow["jobs"]["example-tests"]["concurrency"] == { + "group": "cloud-tests-${{ matrix.python-prefix }}", + "cancel-in-progress": False, + "queue": "max", + } + + +def test_general_conformance_keeps_all_pending_runs() -> None: + workflow = yaml.safe_load( + WORKFLOW_PATH.with_name("conformance-tests.yml").read_text() + ) + + assert workflow["concurrency"] == { + "group": "conformance-tests-global", + "cancel-in-progress": False, + "queue": "max", + } diff --git a/.github/workflows/cloud-tests.yml b/.github/workflows/cloud-tests.yml index cd3a5cb96..da80a43e2 100644 --- a/.github/workflows/cloud-tests.yml +++ b/.github/workflows/cloud-tests.yml @@ -46,6 +46,7 @@ jobs: concurrency: group: cloud-tests-${{ matrix.python-prefix }} cancel-in-progress: false + queue: max strategy: fail-fast: false matrix: diff --git a/.github/workflows/conformance-tests.yml b/.github/workflows/conformance-tests.yml index 2e9b73d7b..2d660987e 100644 --- a/.github/workflows/conformance-tests.yml +++ b/.github/workflows/conformance-tests.yml @@ -32,6 +32,7 @@ concurrency: # run against the shared stack. group: conformance-tests-global cancel-in-progress: false + queue: max permissions: contents: read diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index 8f01e56a6..188aa882e 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -48,6 +48,13 @@ on: default: main type: string +# Backend stacks are shared across PRs. Queue whole runs so reusable +# backend jobs do not displace another PR from their pending slots. +concurrency: + group: otel-conformance-tests + cancel-in-progress: false + queue: max + permissions: {} jobs: From daae0921df879a04ad7ffea506fb61bf7aa8eb95 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 7 Oct 2026 01:26:25 +0000 Subject: [PATCH 17/24] ci: adopt shared backend queue preservation --- .github/workflows/opentelemetry-conformance-tests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index 188aa882e..216490be3 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -66,7 +66,7 @@ jobs: actions: write contents: read id-token: write - uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@3542852f918f515d7c6a53af1eb0d5aa83d1eb1e + uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@a66037abbbfa55fde97f714e30f0bc262edefd63 with: language: python runs_on: codebuild-github-actions-runner-${{ github.run_id }}-${{ github.run_attempt }} From 74c5486bd8779fdf4c15aafc1c7507867f0bc9f8 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 7 Oct 2026 02:23:05 +0000 Subject: [PATCH 18/24] fix: build new js example workspace dependencies --- .github/scripts/js_examples/run.py | 16 ++++++- .github/scripts/tests/test_js_examples.py | 52 +++++++++++++++++++++++ 2 files changed, 67 insertions(+), 1 deletion(-) diff --git a/.github/scripts/js_examples/run.py b/.github/scripts/js_examples/run.py index e2fa359a8..7dac7e829 100755 --- a/.github/scripts/js_examples/run.py +++ b/.github/scripts/js_examples/run.py @@ -61,15 +61,24 @@ TESTING_SOURCE = REPO_ROOT / "packages/aws-durable-execution-sdk-python-testing/src" JS_SDK_URL = "https://github.com/aws/aws-durable-execution-sdk-js.git" EXAMPLES_REL = Path("packages/aws-durable-execution-sdk-js-examples") -# The examples depend on these workspaces by "*". The root "npm run build" +# The examples depend on these local workspaces. The root "npm run build" # also builds the insight tools and the VS Code extension, which the harness # does not need, so the script builds only these, in dependency order. JS_BUILD_WORKSPACES = ( "packages/aws-durable-execution-sdk-js", "packages/aws-durable-execution-sdk-js-testing", "packages/aws-durable-execution-sdk-js-otel", + "packages/aws-durable-execution-sdk-js-extras", + "packages/aws-durable-execution-sdk-js-microvm-worker", "packages/aws-durable-execution-sdk-js-examples", ) +# Older JS refs predate these packages; their examples do not need them. +OPTIONAL_JS_BUILD_WORKSPACES = frozenset( + { + "packages/aws-durable-execution-sdk-js-extras", + "packages/aws-durable-execution-sdk-js-microvm-worker", + } +) # otel examples export spans to an OpenTelemetry collector. The harness does # not run one, so these examples cannot pass here. They are not selected. EXCLUDED_DIRS = ("/otel/",) @@ -202,6 +211,11 @@ def build_js_sdk(js_dir: Path, *, force: bool) -> None: [npm, "ci", "--no-audit", "--no-fund"], cwd=js_dir, env=env, check=True ) for workspace in JS_BUILD_WORKSPACES: + if ( + workspace in OPTIONAL_JS_BUILD_WORKSPACES + and not (js_dir / workspace / "package.json").is_file() + ): + continue log(f"npm run build -w {workspace}") subprocess.run( [npm, "run", "build", "-w", workspace], cwd=js_dir, env=env, check=True diff --git a/.github/scripts/tests/test_js_examples.py b/.github/scripts/tests/test_js_examples.py index 15522209b..dd56d143e 100644 --- a/.github/scripts/tests/test_js_examples.py +++ b/.github/scripts/tests/test_js_examples.py @@ -623,6 +623,58 @@ def test_source_fingerprint_changes_with_local_edits(tmp_path: Path) -> None: assert run.source_fingerprint(tmp_path) != untracked +@pytest.mark.parametrize("has_microvm_packages", [False, True]) +def test_build_produces_example_workspace_dependencies( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + has_microvm_packages: bool, +) -> None: + core = "packages/aws-durable-execution-sdk-js" + testing = "packages/aws-durable-execution-sdk-js-testing" + otel = "packages/aws-durable-execution-sdk-js-otel" + extras = "packages/aws-durable-execution-sdk-js-extras" + worker = "packages/aws-durable-execution-sdk-js-microvm-worker" + examples_package = "packages/aws-durable-execution-sdk-js-examples" + packages = [core, testing, otel, examples_package] + if has_microvm_packages: + packages.extend([extras, worker]) + for package in packages: + directory = tmp_path / package + directory.mkdir(parents=True) + (directory / "package.json").write_text("{}") + + built: list[str] = [] + + def execute(args: list[str], **_kwargs: Any) -> subprocess.CompletedProcess[str]: + if args[1] == "ci": + return subprocess.CompletedProcess(args, 0) + workspace = args[-1] + assert (tmp_path / workspace / "package.json").is_file() + if workspace == worker: + assert testing in built + if workspace == examples_package and has_microvm_packages: + # The generator imports extras/microvm and typechecks the worker. + # A successful examples build requires both artifacts to exist. + assert extras in built + assert worker in built + built.append(workspace) + return subprocess.CompletedProcess(args, 0) + + stamp = tmp_path / "built-stamp" + monkeypatch.setattr(run, "_head", lambda _: "test-head") + monkeypatch.setattr(run, "state_file", lambda *_: stamp) + monkeypatch.setattr(run, "source_fingerprint", lambda _: "fingerprint") + monkeypatch.setattr(run.shutil, "which", lambda _: "/fake/npm") + monkeypatch.setattr(run.subprocess, "run", execute) + + run.build_js_sdk(tmp_path, force=True) + + assert built[-1] == examples_package + assert stamp.read_text() == "fingerprint\n" + if not has_microvm_packages: + assert built == [core, testing, otel, examples_package] + + def test_proxy_creates_the_dump_directory(tmp_path: Path) -> None: dump_dir = tmp_path / "new" / "dump" proxy = invoke_proxy.ProxyServer(0, "http://127.0.0.1:1", dump_dir) From 176026b221e2c5aec495719efa481a424c07bbf5 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 7 Oct 2026 22:54:57 +0000 Subject: [PATCH 19/24] ci: pin published conformance workflow --- .github/workflows/opentelemetry-conformance-tests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index 216490be3..4b4f2215a 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -66,7 +66,7 @@ jobs: actions: write contents: read id-token: write - uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@a66037abbbfa55fde97f714e30f0bc262edefd63 + uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@f18bd0b5f28c5c90e288d0fb8bca08a849b51863 with: language: python runs_on: codebuild-github-actions-runner-${{ github.run_id }}-${{ github.run_attempt }} From b989dad343da6aa7f96af6993af2fe7cbb2baf0b Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 18:37:17 +0000 Subject: [PATCH 20/24] docs: clarify failed plugin setup context --- .../src/aws_durable_execution_sdk_python/plugin.py | 3 +++ 1 file changed, 3 insertions(+) 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 f9e202017..d2ff95703 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 @@ -480,6 +480,9 @@ def __init__(self, plugins: list[DurableInstrumentationPlugin] | None): self._executor: ThreadPoolExecutor | None = None self._invocation_status: InvocationStartInfo | None = None self._operations_provider: Callable[[], Mapping[str, Operation]] | None = None + # Non-None only after a start hook fails: the pre-hook snapshot excludes + # its partial bindings from later setup and the handler. Otherwise this + # stays None, and run_handler uses a fresh copy_context(). self._startup_context: contextvars.Context | None = None self._invocation_contexts: list[contextvars.Context | None] = [] From d8de5f37b185e1d1fdbdb4634fceca82d6556fb6 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 20:41:01 +0000 Subject: [PATCH 21/24] refactor: simplify optional handler context scopes --- .github/workflows/ci.yml | 3 +- .../README.md | 10 +- .../execution_plugin.py | 1 - .../invocation_plugin.py | 2 - .../e2e/test_invocation_wait_resume_int.py | 23 +-- .../README.md | 30 ++- .../plugin.py | 25 +-- .../tests/handler_worker_context_test.py | 173 +++++++----------- 8 files changed, 97 insertions(+), 170 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4205c04c9..a3839308a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -95,8 +95,7 @@ jobs: assert version("aws-durable-execution-sdk-python") == "2.0.0" assert "site-packages" in Path(plugin.__file__).parts - assert not hasattr(plugin.DurableInstrumentationPlugin, "handler_context") - assert not hasattr(plugin, "DURABLE_INSTRUMENTATION_HANDLER_CONTEXT_API_VERSION") + assert not hasattr(plugin.PluginExecutor, "run_handler") PYTHON hatch run test-pypi-otel-minimum:test - name: Build distribution diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 2e69cd01d..fcc0f77e9 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -190,11 +190,11 @@ span current only while the handler runs. The scope closes on the same worker in reverse plugin order, including on failure and suspension, without changing the invocation-hook caller. Older cores ignore this optional scope and retain their existing behavior; install the updated core as well to get handler context propagation. -The bundled OTel classes explicitly opt in with `__durable_handler_context_api__ = 1`. -Custom subclasses must repeat that literal marker on their own concrete class to -use the scope; inherited or instance markers are ignored. Unopted legacy helpers -and properties with the same name are never inspected. Older cores ignore the -marker without importing any new core API. +The bundled OTel classes provide this optional method directly; custom subclasses +inherit it normally and can override it using standard Python method resolution. +The updated core invokes a callable `handler_context` attribute when present; +there is no separate plugin opt-in marker. Older cores ignore the method without +importing any new core API. Execution view similarly restores the Workflow span inside the handler scope if another invocation-start hook clears the active span or switches to an unrelated trace. Both views retain valid same-trace parents and baggage, and restore the diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py index b05a72a5e..9c384ed62 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py @@ -129,7 +129,6 @@ class ExecutionOtelPlugin(DurableInstrumentationPlugin): span). """ - __durable_handler_context_api__ = 1 __durable_registration_api__: ClassVar[int] = 1 exclusive_group: ClassVar[str | None] = "aws-durable-execution-otel-view" diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py index 73e4ac7c5..68c2893a4 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py @@ -110,8 +110,6 @@ class InvocationOtelPlugin(DurableInstrumentationPlugin): provider installed by the ADOT Lambda layer). """ - __durable_handler_context_api__ = 1 - DEFAULT_INSTRUMENT_NAME = "aws-durable-execution-sdk-python" __durable_registration_api__: ClassVar[int] = 1 diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py index 427854750..2e95429a7 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py @@ -261,6 +261,14 @@ def handler_impl(_event: Any, context: DurableContext) -> str: assert completed_wait_span.end_time <= after_resume.start_time +class InheritedInvocationPlugin(InvocationOtelPlugin): + pass + + +class InheritedExecutionPlugin(ExecutionOtelPlugin): + pass + + @pytest.mark.parametrize( ("plugin_type", "extra_context_plugin"), [ @@ -274,10 +282,8 @@ def handler_impl(_event: Any, context: DurableContext) -> str: + ( [(ExecutionOtelPlugin, True)] + [(ExecutionOtelPlugin, kind) for kind in ("same", "unrelated", "absent")] - if getattr( - core_plugin_api, "DURABLE_INSTRUMENTATION_HANDLER_CONTEXT_API_VERSION", None - ) - == 1 + + [(InheritedInvocationPlugin, False), (InheritedExecutionPlugin, "absent")] + if hasattr(core_plugin_api.PluginExecutor, "run_handler") else [] ), ) @@ -297,12 +303,7 @@ def test_handler_user_spans_inherit_context_across_resume_and_failure( # The documented PyPI compatibility environment deliberately uses an older # core. Keep exercising its supported operation tracing and lifecycle while # asserting the new handler contract only when that core exposes the scope. - supports_handler_context = ( - getattr( - core_plugin_api, "DURABLE_INSTRUMENTATION_HANDLER_CONTEXT_API_VERSION", None - ) - == 1 - ) + supports_handler_context = hasattr(core_plugin_api.PluginExecutor, "run_handler") exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) @@ -457,7 +458,7 @@ def on_invocation_end(self, _info: Any) -> None: else derive_workflow_span_id(EXECUTION_ARN) ] * 2 - if plugin_type is ExecutionOtelPlugin + if issubclass(plugin_type, ExecutionOtelPlugin) else ambient_ids if ambient_kind == "same" else [ diff --git a/packages/aws-durable-execution-sdk-python/README.md b/packages/aws-durable-execution-sdk-python/README.md index 8429e196e..bfe54384f 100644 --- a/packages/aws-durable-execution-sdk-python/README.md +++ b/packages/aws-durable-execution-sdk-python/README.md @@ -114,22 +114,20 @@ lifecycle order are unchanged. ### Optional handler context scopes -A plugin can declare `__durable_handler_context_api__ = 1` directly on its -concrete class and implement `handler_context(info)` returning a context manager. -The updated core enters these scopes around the top-level handler on its worker -thread, in registration order, and closes them in reverse order. Cleanup receives -no handler exception and cannot suppress or replace its outcome. Invocation hooks -retain their original thread and order; failed setup bindings are discarded and -successful scope cleanup stays in the context that owns its tokens. - -The marker must be the literal integer `1`; instance and inherited markers do not -opt in. A subclass must redeclare the marker to adopt this new hook. Unopted legacy -helpers, properties and dynamic attributes named `handler_context` are untouched. -The generic plugin base supplies neither a marker nor a default method. Core -support is advertised by module constant -`aws_durable_execution_sdk_python.plugin.DURABLE_INSTRUMENTATION_HANDLER_CONTEXT_API_VERSION`. -Older cores ignore this optional API. The provider API version and dependency -requirements are unchanged. +A plugin can provide `handler_context(info)` returning a context manager. The +updated core looks up this optional method normally, so inherited methods work +without extra declarations. A missing or non-callable attribute is ignored. +The core enters these scopes around the top-level handler on its worker thread, +in registration order, and closes them in reverse order. Cleanup receives no +handler exception and cannot suppress or replace its outcome. Invocation hooks +retain their original thread and order; failed lookup, setup or entry bindings +are discarded, and successful scope cleanup stays in the context that owns its +tokens. + +`handler_context` is an optional plugin API name: callable methods or attributes +with that name are invoked. The generic plugin base does not require a default +method. Older cores ignore this optional API and retain their existing behavior. +The provider API version and dependency requirements are unchanged. ## 🚀 Quick Start 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 d2ff95703..7b3ccf0bc 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 @@ -30,7 +30,6 @@ logger = logging.getLogger(__name__) DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION = 1 -DURABLE_INSTRUMENTATION_HANDLER_CONTEXT_API_VERSION = 1 class InvocationStatus(Enum): @@ -462,18 +461,6 @@ class DurableInstrumentationPluginProvider: plugin_api_version: int -def _handler_context_api_enabled( - plugin_type: type[DurableInstrumentationPlugin], -) -> bool: - """Read only the concrete class namespace, bypassing metaclass descriptors.""" - namespace = type.__dict__["__dict__"].__get__(plugin_type, type(plugin_type)) - version = namespace.get("__durable_handler_context_api__") - return ( - type(version) is int - and version == DURABLE_INSTRUMENTATION_HANDLER_CONTEXT_API_VERSION - ) - - class PluginExecutor: def __init__(self, plugins: list[DurableInstrumentationPlugin] | None): self._plugins = plugins or [] @@ -574,18 +561,12 @@ def execute_plugins(self, info, sync): def _safe_handler_context( self, plugin: DurableInstrumentationPlugin, info: InvocationStartInfo ) -> Iterator[bool]: - # Old plugin objects may not inherit this core's new optional method. + # Plugins may omit this optional scope method. scope = None succeeded = True try: - # Old plugins may have an unrelated helper/property with this name. - # Never even inspect it unless this concrete class explicitly opts in. - factory = ( - getattr(plugin, "handler_context", None) - if _handler_context_api_enabled(type(plugin)) - else None - ) - if factory is not None: + factory = getattr(plugin, "handler_context", None) + if callable(factory): scope = factory(info) scope.__enter__() except Exception: diff --git a/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py index 215624ea3..c62377261 100644 --- a/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py @@ -155,8 +155,6 @@ def test_optional_handler_scopes_are_balanced_and_cannot_change_outcome( events: list[str] = [] class ScopePlugin(DurableInstrumentationPlugin): - __durable_handler_context_api__ = 1 - def __init__(self, name: str): self.name = name @@ -331,8 +329,6 @@ def __exit__(self, *_args: Any) -> None: self.plugin.reset("exit") class SetupPlugin(DurableInstrumentationPlugin): - __durable_handler_context_api__ = 1 - def __init__(self, name: str) -> None: self.name = name self.token: contextvars.Token[str] | None = None @@ -427,85 +423,21 @@ def invoke(_event: Any, _context: Any) -> dict[str, str]: assert all(str(error[1]) == "partial plugin setup" for error in errors) -@pytest.mark.parametrize("shape", ["helper", "property", "dynamic"]) -def test_unopted_legacy_handler_context_is_never_inspected(shape: str) -> None: - from contextlib import contextmanager - from collections.abc import Iterator - from datetime import UTC, datetime - from aws_durable_execution_sdk_python.plugin import PluginExecutor - - calls: list[str] = [] - - @contextmanager - def helper(_info: InvocationStartInfo) -> Iterator[None]: - calls.append("helper") - yield - - def property_getter(_self: Any) -> Any: - calls.append("property") - return helper - - def dynamic_getter(_self: Any, name: str) -> Any: - if name == "handler_context": - calls.append("dynamic") - return "legacy-business-value" - raise AttributeError(name) - - members_by_shape: dict[str, dict[str, Any]] = { - "helper": {"handler_context": staticmethod(helper)}, - "property": {"handler_context": property(property_getter)}, - "dynamic": {"__getattr__": dynamic_getter}, - } - plugin_type = type( - "Legacy", (DurableInstrumentationPlugin,), members_by_shape[shape] - ) - plugin = plugin_type() - if shape == "dynamic": - assert plugin.handler_context == "legacy-business-value" - calls.clear() - executor = PluginExecutor([plugin]) - with executor.run(): - executor.on_invocation_start("legacy-helper", True, datetime.now(UTC), None) - assert executor.run_handler(lambda: list(calls)) == [] - assert calls == [] - - -@pytest.mark.parametrize("marker", [None, 0, 2, True, "1", property(lambda _: 1)]) -def test_handler_scope_requires_literal_class_local_version(marker: Any) -> None: - from contextlib import nullcontext - from datetime import UTC, datetime - from aws_durable_execution_sdk_python.plugin import PluginExecutor - - calls: list[str] = [] - - def helper(_self: Any, _info: InvocationStartInfo) -> Any: - calls.append("scope") - return nullcontext() - - plugin_type = type( - "Legacy", - (DurableInstrumentationPlugin,), - {"__durable_handler_context_api__": marker, "handler_context": helper}, - ) - executor = PluginExecutor([plugin_type()]) - with executor.run(): - executor.on_invocation_start("invalid-marker", True, datetime.now(UTC), None) - assert executor.run_handler(lambda: "ok") == "ok" - assert calls == [] - - -def test_handler_scope_opt_in_is_not_inherited_or_taken_from_instance() -> None: +@pytest.mark.parametrize("inherited", [False, True]) +@pytest.mark.parametrize("fail_handler", [False, True]) +def test_handler_scope_uses_direct_or_inherited_optional_method( + inherited: bool, fail_handler: bool +) -> None: from contextlib import contextmanager from collections.abc import Iterator from datetime import UTC, datetime from aws_durable_execution_sdk_python.plugin import PluginExecutor - marker = contextvars.ContextVar("explicit-scope", default="outside") + marker = contextvars.ContextVar("inherited-scope", default="outside") events: list[str] = [] + error = RuntimeError("original handler failure") - class OptedPlugin(DurableInstrumentationPlugin): - __durable_handler_context_api__ = 1 - + class ScopePlugin(DurableInstrumentationPlugin): @contextmanager def handler_context(self, _info: InvocationStartInfo) -> Iterator[None]: events.append("enter") @@ -516,47 +448,66 @@ def handler_context(self, _info: InvocationStartInfo) -> Iterator[None]: marker.reset(token) events.append("exit") - class LegacySubclass(OptedPlugin): + class InheritedPlugin(ScopePlugin): pass - class ExplicitSubclass(OptedPlugin): - __durable_handler_context_api__ = 1 - - for plugin, expected in [ - (OptedPlugin(), "inside"), - (LegacySubclass(), "outside"), - (ExplicitSubclass(), "inside"), - ]: - # Assigning a marker to an instance cannot accidentally enable the hook. - plugin.__durable_handler_context_api__ = 1 - executor = PluginExecutor([plugin]) - with executor.run(): - executor.on_invocation_start("subclass", True, datetime.now(UTC), None) - assert executor.run_handler(marker.get) == expected - assert marker.get() == "outside" - assert events == ["enter", "exit", "enter", "exit"] - - -def test_handler_opt_in_does_not_trigger_legacy_metaclass_descriptors() -> None: + plugin = InheritedPlugin() if inherited else ScopePlugin() + executor = PluginExecutor([plugin]) + + def handler() -> str: + assert marker.get() == "inside" + events.append("body") + if fail_handler: + raise error + return "ok" + + with executor.run(): + for _ in range(2): + executor.on_invocation_start("inherited", True, datetime.now(UTC), None) + if fail_handler: + with pytest.raises(RuntimeError) as caught: + executor.run_handler(handler) + assert caught.value is error + else: + assert executor.run_handler(handler) == "ok" + assert marker.get() == "outside" + assert events == ["enter", "body", "exit"] * 2 + + +@pytest.mark.parametrize("value", [None, "not a method", 42]) +def test_noncallable_optional_handler_scope_is_ignored(value: Any) -> None: from datetime import UTC, datetime + from types import SimpleNamespace + from typing import cast from aws_durable_execution_sdk_python.plugin import PluginExecutor - reads: list[str] = [] + plugin = cast( + DurableInstrumentationPlugin, + SimpleNamespace(on_invocation_start=lambda info: None, handler_context=value), + ) + executor = PluginExecutor([plugin]) + with executor.run(): + executor.on_invocation_start("noncallable", True, datetime.now(UTC), None) + assert executor.run_handler(lambda: "unchanged") == "unchanged" - def namespace(_cls: Any) -> Any: - reads.append("metaclass-dict") - raise RuntimeError("legacy namespace") - def helper(_self: Any, _info: Any) -> Any: - reads.append("legacy-helper") - raise RuntimeError("legacy helper") +def test_handler_scope_lookup_failure_discards_partial_bindings( + caplog: pytest.LogCaptureFixture, +) -> None: + from datetime import UTC, datetime + from aws_durable_execution_sdk_python.plugin import PluginExecutor - meta = type("LegacyMeta", (type,), {"__dict__": property(namespace)}) - plugin_type = meta( - "Legacy", (DurableInstrumentationPlugin,), {"handler_context": helper} - ) - executor = PluginExecutor([plugin_type()]) + marker = contextvars.ContextVar("lookup-failure", default="outside") + + class BrokenPlugin(DurableInstrumentationPlugin): + @property + def handler_context(self) -> Any: + marker.set("partial") + raise ValueError("scope lookup failed") + + executor = PluginExecutor([BrokenPlugin()]) with executor.run(): - executor.on_invocation_start("metaclass", True, datetime.now(UTC), None) - assert executor.run_handler(lambda: "ok") == "ok" - assert reads == [] + executor.on_invocation_start("lookup", True, datetime.now(UTC), None) + assert executor.run_handler(marker.get) == "outside" + assert marker.get() == "outside" + assert "scope lookup failed" in caplog.text From f225d88065922df85fa919fff8e7b64dd862ac68 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 23:29:03 +0000 Subject: [PATCH 22/24] fix(otel): unify invocation lifecycle on handler worker --- .github/scripts/install_otel_test_wheels.py | 101 ++ .../otel_lifecycle_compatibility_test.py | 162 +++ .github/workflows/ci.yml | 18 +- CONTRIBUTING.md | 30 +- .../README.md | 74 +- .../pyproject.toml | 2 +- .../execution_plugin.py | 20 - .../invocation_plugin.py | 38 +- .../e2e/test_invocation_wait_resume_int.py | 93 +- .../tests/test_execution_plugin.py | 77 +- .../tests/test_invocation_plugin.py | 130 +-- .../tests/test_package_metadata.py | 69 +- .../README.md | 49 +- .../execution.py | 129 ++- .../plugin.py | 99 +- .../tests/handler_worker_context_test.py | 944 ++++++++++-------- pyproject.toml | 66 +- 17 files changed, 1156 insertions(+), 945 deletions(-) create mode 100644 .github/scripts/install_otel_test_wheels.py create mode 100644 .github/tests/otel_lifecycle_compatibility_test.py diff --git a/.github/scripts/install_otel_test_wheels.py b/.github/scripts/install_otel_test_wheels.py new file mode 100644 index 000000000..938996aca --- /dev/null +++ b/.github/scripts/install_otel_test_wheels.py @@ -0,0 +1,101 @@ +"""Install and verify built artifacts without silently using editable sources.""" + +from __future__ import annotations + +import argparse +import ast +import hashlib +import importlib.metadata +import json +import subprocess +import sys +import zipfile +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[2] +CORE = "aws-durable-execution-sdk-python" +OTEL = CORE + "-otel" + + +def built_wheel(package: str) -> Path: + directory = ROOT / "packages" / package + module = package.replace("-", "_") + about = ast.parse((directory / "src" / module / "__about__.py").read_text()) + version = next( + ast.literal_eval(node.value) + for node in about.body + if isinstance(node, ast.Assign) + and any( + isinstance(target, ast.Name) and target.id == "__version__" + for target in node.targets + ) + ) + wheels = list((directory / "dist").glob(f"{module}-{version}-*.whl")) + if len(wheels) != 1: + raise ValueError(f"Build exactly one {package} {version} wheel first: {wheels}") + return wheels[0] + + +def verify(wheel: Path, package: str) -> None: + installed = importlib.metadata.distribution(package) + direct = json.loads(installed.read_text("direct_url.json") or "{}") + assert not direct.get("dir_info", {}).get("editable"), direct + digest = hashlib.sha256(wheel.read_bytes()).hexdigest() + assert direct["archive_info"]["hashes"]["sha256"] == digest, direct + module = package.replace("-", "_") + with zipfile.ZipFile(wheel) as archive: + sources = [ + name + for name in archive.namelist() + if name.startswith(module + "/") and name.endswith(".py") + ] + assert sources + for name in sources: + path = Path(installed.locate_file(name)).resolve() + assert "site-packages" in path.parts, path + assert path.read_bytes() == archive.read(name), path + print( + json.dumps( + { + "package": package, + "version": installed.version, + "wheel": str(wheel), + "sha256": digest, + "verified_sources": len(sources), + } + ) + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--legacy-plugin", action="store_true") + args = parser.parse_args() + packages = [CORE] if args.legacy_plugin else [CORE, OTEL] + wheels = [built_wheel(package) for package in packages] + subprocess.run( + [ + sys.executable, + "-m", + "pip", + "install", + "--no-index", + "--no-deps", + "--force-reinstall", + *map(str, wheels), + ], + check=True, + ) + for wheel, package in zip(wheels, packages, strict=True): + verify(wheel, package) + if args.legacy_plugin: + assert importlib.metadata.version(OTEL) == "1.0.0" + import aws_durable_execution_sdk_python_otel as otel + + assert "site-packages" in Path(otel.__file__).resolve().parts + subprocess.run([sys.executable, "-m", "pip", "check"], check=True) + + +if __name__ == "__main__": + main() diff --git a/.github/tests/otel_lifecycle_compatibility_test.py b/.github/tests/otel_lifecycle_compatibility_test.py new file mode 100644 index 000000000..8e9ce857b --- /dev/null +++ b/.github/tests/otel_lifecycle_compatibility_test.py @@ -0,0 +1,162 @@ +"""Exercise real installed version pairs through the public durable handler.""" + +from __future__ import annotations + +import contextvars +from datetime import UTC, datetime +from importlib.metadata import version +import os +from pathlib import Path +import threading +from types import SimpleNamespace + +import pytest +from aws_durable_execution_sdk_python import durable_execution +from aws_durable_execution_sdk_python import execution as core_execution +from aws_durable_execution_sdk_python.lambda_service import ( + ExecutionDetails, + Operation, + OperationStatus, + OperationType, +) +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from aws_durable_execution_sdk_python_otel.execution_plugin import ExecutionOtelPlugin +from aws_durable_execution_sdk_python_otel.invocation_plugin import InvocationOtelPlugin +from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig +from opentelemetry import baggage, context, trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from packaging.version import Version + + +class NoNetworkClient: + def __getattr__(self, name: str): + raise AssertionError(f"Unexpected service access: {name}") + + +@pytest.mark.parametrize("view", [InvocationOtelPlugin, ExecutionOtelPlugin]) +@pytest.mark.parametrize("order", ["alone", "baggage-first", "baggage-last"]) +@pytest.mark.parametrize("ambient_present", [False, True]) +def test_installed_pair_preserves_context_and_documents_legacy_fallback( + view, order, ambient_present +): + legacy = os.environ.get("OTEL_COMPAT_LEGACY") == "1" + assert Version(version("aws-durable-execution-sdk-python")) >= Version("2.1.0") + assert version("aws-durable-execution-sdk-python-otel") == ( + "1.0.0" if legacy else "1.1.0" + ) + assert "site-packages" in Path(core_execution.__file__).resolve().parts + import aws_durable_execution_sdk_python_otel as installed_otel + + assert "site-packages" in Path(installed_otel.__file__).resolve().parts + + def run() -> None: + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + tracer = provider.get_tracer("installed-lifecycle") + phases = [] + worker_threads = [] + caller_thread = threading.get_ident() + + class BaggagePlugin(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + self.token = context.attach(baggage.set_baggage("customer", "present")) + phases.append("baggage-start") + worker_threads.append(threading.get_ident()) + + def on_invocation_end(self, info): + context.detach(self.token) + phases.append("baggage-end") + worker_threads.append(threading.get_ident()) + + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + enrich_logger=False, + context_extractor=lambda _: None, + ) + ) + plugins = { + "alone": [plugin], + "baggage-first": [BaggagePlugin(), plugin], + "baggage-last": [plugin, BaggagePlugin()], + }[order] + body_parents = [] + + @durable_execution(plugins=plugins, boto3_client=NoNetworkClient()) + def handler(event, durable_context): + phases.append("body") + worker_threads.append(threading.get_ident()) + assert baggage.get_baggage("incoming") == "keep" + assert baggage.get_baggage("customer") == ( + None if order == "alone" else "present" + ) + parent = trace.get_current_span().get_span_context() + body_parents.append(parent) + with tracer.start_as_current_span("customer-span"): + pass + return "ok" + + operation = Operation( + operation_id="installed", + operation_type=OperationType.EXECUTION, + status=OperationStatus.STARTED, + start_timestamp=datetime(2026, 10, 8, tzinfo=UTC), + execution_details=ExecutionDetails(input_payload="{}"), + ) + event = { + "DurableExecutionArn": "test-arn/installed", + "CheckpointToken": "token", + "InitialExecutionState": { + "Operations": [operation.to_json_dict()], + "NextMarker": "", + }, + } + lambda_context = SimpleNamespace( + aws_request_id="installed", + client_context=None, + identity=None, + _epoch_deadline_time_in_ms=0, + invoked_function_arn="test-arn", + tenant_id=None, + ) + ambient = tracer.start_span("host") if ambient_present else trace.INVALID_SPAN + host = baggage.set_baggage( + "incoming", "keep", trace.set_span_in_context(ambient) + ) + token = context.attach(host) + try: + for _ in range(2): + assert handler(event, lambda_context)["Status"] == "SUCCEEDED" + assert context.get_current() is host + assert caller_thread not in worker_threads + assert phases == ( + ["body"] * 2 + if order == "alone" + else ["baggage-start", "body", "baggage-end"] * 2 + ) + spans = exporter.get_finished_spans() + if legacy and view is InvocationOtelPlugin: + # The released plugin does not attach an Invocation fallback. + assert body_parents == [ambient.get_span_context()] * 2 + else: + name = "Invocation" if view is InvocationOtelPlugin else "Workflow" + contexts = [span.context for span in spans if span.name == name] + assert contexts + assert all( + parent.is_valid and parent in contexts for parent in body_parents + ) + users = [span for span in spans if span.name == "customer-span"] + assert len(users) == 2 + assert [span.parent for span in users] == [ + parent if parent.is_valid else None for parent in body_parents + ] + assert plugin._context_tokens == {} + finally: + context.detach(token) + ambient.end() + provider.shutdown() + + contextvars.Context().run(run) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a3839308a..6cbca4dd8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -84,20 +84,6 @@ jobs: run: hatch run types:check - name: Run tests + coverage run: hatch run test:cov - - name: Verify supported legacy core compatibility - run: hatch run test-pypi-otel-legacy:test - - name: Test OTel with the minimum supported core - run: | - hatch run test-pypi-otel-minimum:python - <<'PYTHON' - from importlib.metadata import version - from pathlib import Path - import aws_durable_execution_sdk_python.plugin as plugin - - assert version("aws-durable-execution-sdk-python") == "2.0.0" - assert "site-packages" in Path(plugin.__file__).parts - assert not hasattr(plugin.PluginExecutor, "run_handler") - PYTHON - hatch run test-pypi-otel-minimum:test - name: Build distribution run: | for pkg in packages/*/; do @@ -108,6 +94,10 @@ jobs: cd "$GITHUB_WORKSPACE" fi done + - name: Test installed core and OTel wheels + run: hatch run test-wheel-otel:test + - name: Test released OTel with the new core wheel + run: hatch run test-wheel-otel-legacy:test - name: Verify OTel wheel dependency contract run: | OTEL_WHEEL=$(find packages/aws-durable-execution-sdk-python-otel/dist \ diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 15b865f78..a3b6bc2f7 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -70,23 +70,25 @@ hatch run dev-otel:typecheck # type check otel only hatch run dev-examples:test # run examples tests only ``` -### PyPI release testing +### Installed package compatibility testing -To verify packages work against the published PyPI version of the core SDK (rather than the local workspace): +Build the core and OTel distributions with `hatch build` in each package, then run +these commands from the repository root: ```bash -hatch run test-pypi-otel:test # test new OTel capabilities against capable installed core -hatch run test-pypi-otel-legacy:test # valid registrations/lifecycles on supported core 2.0.x -hatch run test-pypi-otel-minimum:test # all prior OTel tests and valid registrations on core 2.0.0 -hatch run test-pypi-examples:test # test examples against PyPI core SDK -``` - -The OTel minimum-core environment excludes the local core and pins 2.0.0, so -newer PyPI releases cannot remove legacy compatibility coverage. It retains all -pre-existing OTel tests plus valid registration/wait-resume cases. Exclusivity -validation requires the newer core and is exercised by the complete workspace -suite and the capable-core environment. Use `hatch run dev-otel:test` for the -current workspace core. +hatch run test-wheel-otel:test # full OTel suite on the two built wheels +hatch run test-wheel-otel-legacy:test # released OTel 1.0.0 with the built core +hatch run test-pypi-examples:test # examples against the published core +``` + +The wheel lanes have no editable workspace members. They verify installed source +bytes and artifact hashes before exercising public handlers. OTel 1.1 requires +the redesigned core 2.1.0 lifecycle; it no longer claims compatibility with core +2.0.x. Publish core first. Building both wheels lets CI verify the intended pair +before that minimum is available on PyPI. The legacy-plugin lane documents the +actual core-only upgrade: host isolation is provided by the new core, while old +Invocation OTel does not gain the new fallback. Workspace tests continue to cover +the complete current implementation with `hatch run dev-otel:test`. ### Package-level commands diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index fcc0f77e9..f42b5986c 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -55,9 +55,11 @@ DURABLE_EXECUTION_PLUGINS=otel-execution cold start, so the handler does not need to import or explicitly register the plugin. -Automatic mutual-exclusion validation requires core SDK 2.1.0 or later together -with OTel 1.1.0 or later. OTel 1.1 remains compatible with core 2.0.x for existing -valid registrations; those older cores do not enforce the new group metadata. +OTel 1.1.0 requires core SDK 2.1.0 or later for the invocation-worker lifecycle +and automatic mutual-exclusion validation. Publish the redesigned core first, +then the plugin. Core 2.0.x runs invocation hooks on the caller and cannot provide +the new plugin's handler propagation and host-context isolation; that version +pair is not supported by OTel 1.1.0. Configure only one OTel view on every core version. `InvocationOtelPlugin` shows work within each Lambda invocation; `ExecutionOtelPlugin` shows logical operations across the whole execution. They emit overlapping telemetry and @@ -182,24 +184,31 @@ lambda_.Function( ### Handler context propagation -A core SDK with handler-worker context propagation carries the context established -by invocation-start hooks into the handler. Invocation view preserves an active -ambient span on the canonical execution trace; when that context is absent or -belongs to a different trace, its optional `handler_context` scope makes the Invocation -span current only while the handler runs. The scope closes on the same worker in -reverse plugin order, including on failure and suspension, without changing the -invocation-hook caller. Older cores ignore this optional scope and retain their -existing behavior; install the updated core as well to get handler context propagation. -The bundled OTel classes provide this optional method directly; custom subclasses -inherit it normally and can override it using standard Python method resolution. -The updated core invokes a callable `handler_context` attribute when present; -there is no separate plugin opt-in marker. Older cores ignore the method without -importing any new core API. -Execution view similarly restores the Workflow span inside the handler scope if -another invocation-start hook clears the active span or switches to an unrelated -trace. Both views retain valid same-trace parents and baggage, and restore the -worker's previous context when the scope ends. -Existing plugin registration, factory lifetime, and checkpoint formats are unchanged. +The core runs the existing Start hooks, handler, output preparation, resource +cleanup and End hooks on one invocation worker. Start finishes before checkpoint +processing begins. End follows registered-branch joins and checkpoint shutdown, +and output serialization or checkpoint errors retain their normal classification. +There is no separate handler-context plugin API. + +During Start, Invocation view preserves a valid active span on the canonical +execution trace; when it is absent or unrelated, it attaches the Invocation span. +Execution view attaches the Workflow span, including when an incoming span is on +the same trace. The Invocation span's own ambient parenting is separate from the +active context supplied to handler instrumentation. Both views preserve baggage. + +Start hooks run in registration order. A later successful plugin that deliberately +sets or clears the active span wins; OTel does not apply a second correction pass. +Place OTel after a span-replacing plugin when OTel's view-specific context is +desired. Baggage-only plugins that extend the current context can appear on either +side. End hooks also retain registration order and reset tokens in their owning +Context; they do not promise a reverse-stack observation of other plugins' spans. +Existing registration, factory lifetime and checkpoint formats are unchanged. + +Upgrade both core to 2.1+ and OTel to 1.1+ for these guarantees. The released +OTel 1.0 plugin can run on the new core with worker/host isolation, but its +Invocation view does not attach the new fallback; a core-only upgrade does not +supply that plugin behavior. CI tests the new pair as installed wheels and tests +the actual released plugin separately, including before core 2.1 is on PyPI. ### 3. In your Lambda handler (index.py) @@ -353,19 +362,14 @@ remains in `durable.operation.status`. ### Invocation context isolation -Invocation hooks retain their caller thread and registration order. With the -updated core, invocation-local context-variable bindings are isolated from the -host: hooks see the incoming context and the handler receives their resulting -context, while invocation exit restores the host's original bindings even if a -plugin fails during setup or cleanup. Plugins must not use invocation context -bindings to mutate the host context after the invocation has returned. Older -supported cores retain their existing lifecycle behavior, including the -execution-view limitation when later plugins open invocation context scopes. -The new isolation applies only when plugins are registered. If an invocation-start -hook or handler-scope entry raises, subsequent setup and the handler retain the -bindings from before that hook. Successful scopes still clean up in their original -context, preserving token ownership. This isolates context-variable bindings; -it does not undo a plugin's mutations to shared objects or external side effects. +With core 2.1+, invocation hooks and the handler run on one worker in an +invocation-local Context initialized from the host's bindings. Successful Start +bindings are visible to later hooks, the handler and resource cleanup. The host's +bindings remain unchanged after return, even if a plugin fails during Start or +End. A failed Start's bindings are discarded for subsequent work, while End runs +in that Start's original Context to preserve token ownership. This isolates +context-variable bindings; it does not undo mutations to shared objects or +external side effects. The isolation applies when plugins are registered. ### Log Correlation @@ -467,7 +471,7 @@ setups. ## Requirements - Python >= 3.11 -- `aws-durable-execution-sdk-python` >= 2.0.0 (core >= 2.1.0 with OTel >= 1.1.0 for automatic view-exclusivity validation) +- `aws-durable-execution-sdk-python` >= 2.1.0 (release the redesigned core before OTel 1.1.0) - An ADOT/community OpenTelemetry Lambda layer, or the `standalone` extra ## License diff --git a/packages/aws-durable-execution-sdk-python-otel/pyproject.toml b/packages/aws-durable-execution-sdk-python-otel/pyproject.toml index cbc2f97a2..461718714 100644 --- a/packages/aws-durable-execution-sdk-python-otel/pyproject.toml +++ b/packages/aws-durable-execution-sdk-python-otel/pyproject.toml @@ -22,7 +22,7 @@ classifiers = [ "Programming Language :: Python :: Implementation :: PyPy", ] dependencies = [ - "aws-durable-execution-sdk-python>=2.0.0", + "aws-durable-execution-sdk-python>=2.1.0", ] [project.entry-points."aws_durable_execution.plugins"] diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py index 9c384ed62..40a62ca84 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py @@ -31,11 +31,9 @@ from __future__ import annotations -import contextlib import datetime import logging import threading -from collections.abc import Iterator from typing import Any, ClassVar from aws_durable_execution_sdk_python.plugin import ( @@ -511,24 +509,6 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: if self._config.enrich_logger: install_log_filter(self) - @contextlib.contextmanager - def handler_context(self, info: InvocationStartInfo) -> Iterator[None]: - """Keep handler instrumentation on this execution's trace in its worker.""" - ambient = trace.get_current_span().get_span_context() - workflow = self._workflow_span - token = None - if ( - self._tracing_enabled - and workflow is not None - and (not ambient.is_valid or ambient.trace_id != self._execution_trace_id) - ): - token = otel_context.attach(trace.set_span_in_context(workflow)) - try: - yield - finally: - if token is not None: - otel_context.detach(token) - def _start_workflow_span(self, info: InvocationStartInfo) -> None: """Install a non-recording placeholder for the execution-scoped Workflow span. diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py index 68c2893a4..0fb974fe1 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py @@ -2,11 +2,9 @@ from __future__ import annotations -import contextlib import datetime import logging import threading -from collections.abc import Iterator from typing import Any, ClassVar from aws_durable_execution_sdk_python.plugin import ( @@ -75,6 +73,7 @@ {InvocationStatus.SUCCEEDED, InvocationStatus.FAILED} ) _TIMESTAMP_STEP_NANOS = 1_000 +_INVOCATION_CONTEXT_KEY = "__invocation_context__" _SpanAttributes = dict[str, str | bool | int] @@ -302,7 +301,7 @@ def get_current_span_context(self) -> SpanContext | None: on_user_function_start). Unrelated ambient spans are ignored so logs stay correlated to the durable execution trace. 2. The invocation span from the plugin registry, including lifecycle - phases outside the optional handler-worker context scope. + phases where another plugin changed the active context. Returns: A valid SpanContext, or None if no span is active. @@ -591,27 +590,20 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: attributes=self._extract_attributes(info), ) - # Cover handlers installed after construction as well. - if self._enrich_logger: - install_log_filter(self) - - @contextlib.contextmanager - def handler_context(self, info: InvocationStartInfo) -> Iterator[None]: - """Bind the fallback only inside the SDK-owned handler worker scope.""" + # Start and End share the invocation worker and token-owning Context. + # Retain a valid same-trace ambient parent; otherwise bind this invocation. ambient = trace.get_current_span().get_span_context() invocation_span = self._get_span(None) - token = None - if ( - self._tracing_enabled - and invocation_span is not None - and (not ambient.is_valid or ambient.trace_id != self._execution_trace_id) + if invocation_span is not None and ( + not ambient.is_valid or ambient.trace_id != self._execution_trace_id ): - token = context.attach(trace.set_span_in_context(invocation_span)) - try: - yield - finally: - if token is not None: - context.detach(token) + self._attach_context( + _INVOCATION_CONTEXT_KEY, trace.set_span_in_context(invocation_span) + ) + + # Cover handlers installed after construction as well. + if self._enrich_logger: + install_log_filter(self) def _start_workflow_span(self, info: InvocationStartInfo) -> None: """Install a non-recording placeholder for the execution-scoped Workflow span. @@ -689,8 +681,8 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: self._reset_state() return - # User execution has finished; the worker has already closed its handler - # context scope without modifying the invocation-hook caller. + # User work and checkpoint cleanup have finished. Release Start bindings + # in the Context that created their tokens before closing the spans. self._detach_remaining_contexts() # Spans are registered parent-first, so close pending spans in reverse diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py index 2e95429a7..fa7036920 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py @@ -26,7 +26,6 @@ OperationType, StepDetails, ) -from aws_durable_execution_sdk_python import plugin as core_plugin_api from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( derive_workflow_span_id, @@ -272,20 +271,11 @@ class InheritedExecutionPlugin(ExecutionOtelPlugin): @pytest.mark.parametrize( ("plugin_type", "extra_context_plugin"), [ - (InvocationOtelPlugin, False), - (InvocationOtelPlugin, True), - (ExecutionOtelPlugin, False), + (view, extra) + for view in (InvocationOtelPlugin, ExecutionOtelPlugin) + for extra in (False, True, "same", "unrelated", "absent") ] - # Execution-view caller isolation requires the coordinated newer core. - # Released 2.0.x retains the pre-existing same-order teardown limitation; - # the legacy lane continues checking its supported combinations above. - + ( - [(ExecutionOtelPlugin, True)] - + [(ExecutionOtelPlugin, kind) for kind in ("same", "unrelated", "absent")] - + [(InheritedInvocationPlugin, False), (InheritedExecutionPlugin, "absent")] - if hasattr(core_plugin_api.PluginExecutor, "run_handler") - else [] - ), + + [(InheritedInvocationPlugin, False), (InheritedExecutionPlugin, "absent")], ) @pytest.mark.parametrize("reverse_plugins", [False, True]) @pytest.mark.parametrize("fail_after_resume", [False, True]) @@ -300,10 +290,6 @@ def test_handler_user_spans_inherit_context_across_resume_and_failure( ) -> None: monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) monkeypatch.setenv("_X_AMZN_TRACE_ID", XRAY_TRACE_HEADER) - # The documented PyPI compatibility environment deliberately uses an older - # core. Keep exercising its supported operation tracing and lifecycle while - # asserting the new handler contract only when that core exposes the scope. - supports_handler_context = hasattr(core_plugin_api.PluginExecutor, "run_handler") exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) @@ -326,9 +312,7 @@ def step_body(_step_context: Any) -> str: def handler_body(_event: Any, context: DurableContext) -> str: if extra_context_plugin: - assert baggage.get_baggage("customer") == ( - "present" if supports_handler_context else None - ) + assert baggage.get_baggage("customer") == "present" user_span("handler-entry") saved = context.step(step_body, name="before-wait") user_span("handler-after-step") @@ -338,8 +322,8 @@ def handler_body(_event: Any, context: DurableContext) -> str: raise ValueError("handler failed after resume") return saved - # An unrelated plugin may own a caller-thread OTel baggage scope. The - # invocation-view fallback must never become part of its saved token. + # Start hooks share the invocation worker, in registration order. A later + # successful plugin may deliberately replace or clear the active span. from opentelemetry import baggage class BaggagePlugin(DurableInstrumentationPlugin): @@ -449,37 +433,54 @@ def on_invocation_end(self, _info: Any) -> None: assert calls == ["step"] assert otel_context.get_current() == before_context spans = exporter.get_finished_spans() - expected_parents = ( - [None, None] - if not supports_handler_context - else [ - 0xCAFE - if extra_context_plugin == "same" and not reverse_plugins - else derive_workflow_span_id(EXECUTION_ARN) - ] - * 2 - if issubclass(plugin_type, ExecutionOtelPlugin) - else ambient_ids - if ambient_kind == "same" - else [ - span.context.span_id - for span in spans - if span.name == "Invocation" and span.context is not None + expected_parents: list[int | None] + if isinstance(extra_context_plugin, str) and not reverse_plugins: + # No second OTel correction pass: the later successful Start wins. + expected_parents = [ + None if extra_context_plugin == "absent" else 0xCAFE + ] * 2 + expected_trace_id = {"same": XRAY_TRACE_ID, "unrelated": 1, "absent": None}[ + extra_context_plugin ] - ) + elif issubclass(plugin_type, ExecutionOtelPlugin): + expected_parents = [derive_workflow_span_id(EXECUTION_ARN)] * 2 + expected_trace_id = XRAY_TRACE_ID + elif reverse_plugins and isinstance(extra_context_plugin, str): + expected_parents = ( + [0xCAFE] * 2 + if extra_context_plugin == "same" + else [ + span.context.span_id + for span in spans + if span.name == "Invocation" and span.context is not None + ] + ) + expected_trace_id = XRAY_TRACE_ID + else: + expected_parents = ( + [*ambient_ids] + if ambient_kind == "same" + else [ + span.context.span_id + for span in spans + if span.name == "Invocation" and span.context is not None + ] + ) + expected_trace_id = XRAY_TRACE_ID for name in ("handler-entry", "handler-after-step"): users = [span for span in spans if span.name == name] assert len(users) == 2 assert [span.parent.span_id if span.parent else None for span in users] == ( expected_parents ) - assert all( - span.context is not None - and ( - (span.context.trace_id == XRAY_TRACE_ID) == supports_handler_context + assert all(span.context is not None for span in users) + if expected_trace_id is None: + assert all( + span.parent is None and span.context.trace_id != XRAY_TRACE_ID + for span in users ) - for span in users - ) + else: + assert all(span.context.trace_id == expected_trace_id for span in users) after_resume = next( span for span in spans if span.name == "handler-after-resume" ) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py index 23c7394ba..214a3a89d 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py @@ -1628,10 +1628,8 @@ def test_nested_suspension_unwinds_scopes_in_reverse_order(): @pytest.mark.parametrize("ambient_kind", ["same", "unrelated", "absent"]) @pytest.mark.parametrize("raises", [False, True]) @pytest.mark.parametrize("sampler", [ALWAYS_ON, ALWAYS_OFF]) -def test_handler_scope_keeps_execution_trace_and_restores_context( - ambient_kind: str, - raises: bool, - sampler: Sampler, +def test_invocation_hooks_bind_parent_and_restore_baggage( + ambient_kind: str, raises: bool, sampler: Sampler ) -> None: provider = TracerProvider(sampler=sampler) plugin = ExecutionOtelPlugin( @@ -1641,55 +1639,50 @@ def test_handler_scope_keeps_execution_trace_and_restores_context( enrich_logger=False, ) ) - info = _invocation_start_info() - plugin.on_invocation_start(info) - workflow = trace.get_current_span().get_span_context() - assert workflow.is_valid - assert workflow.span_id == derive_workflow_span_id(EXECUTION_ARN) - caller = baggage.set_baggage("customer", "preserved", Context()) - expected = workflow + caller = baggage.set_baggage("tenant", "hook-test", Context()) + ambient = None if ambient_kind != "absent": ambient = SpanContext( - trace_id=workflow.trace_id if ambient_kind == "same" else 1, - span_id=0xCAFE, + trace_id=_to_otel_trace_id(EXECUTION_ARN, START_TIME) + if ambient_kind == "same" + else 1, + span_id=0x42, is_remote=False, - trace_flags=workflow.trace_flags, + trace_flags=TraceFlags(1), ) caller = trace.set_span_in_context(NonRecordingSpan(ambient), caller) - if ambient_kind == "same": - expected = ambient token = otel_context.attach(caller) - error = ValueError("user failure") + error = ValueError("handler error") try: + plugin.on_invocation_start(_invocation_start_info()) + expected = trace.get_current_span().get_span_context() + assert expected.trace_id == _to_otel_trace_id(EXECUTION_ARN, START_TIME) + assert expected.span_id == derive_workflow_span_id(EXECUTION_ARN) def body() -> None: - with plugin.handler_context(info): - assert trace.get_current_span().get_span_context() == expected - assert baggage.get_baggage("customer") == "preserved" - if raises: - raise error - - if raises: - with pytest.raises(ValueError) as caught: + active = trace.get_current_span().get_span_context() + assert active == expected + assert active.is_valid + assert baggage.get_baggage("tenant") == "hook-test" + if raises: + raise error + + try: + if raises: + with pytest.raises(ValueError) as caught: + body() + assert caught.value is error + else: body() - assert caught.value is error - else: - body() + finally: + plugin.on_invocation_end( + _invocation_end_info( + InvocationStatus.FAILED if raises else InvocationStatus.SUCCEEDED + ) + ) + assert plugin._context_tokens == {} assert otel_context.get_current() is caller + assert baggage.get_baggage("tenant") == "hook-test" finally: otel_context.detach(token) - plugin.on_invocation_end(_invocation_end_info()) provider.shutdown() - - -@pytest.mark.parametrize("completed", [False, True]) -def test_handler_scope_without_live_workflow_is_noop(completed: bool) -> None: - plugin, _ = _create_plugin() - info = _invocation_start_info() - if completed: - plugin.on_invocation_start(info) - plugin.on_invocation_end(_invocation_end_info()) - caller = otel_context.get_current() - with plugin.handler_context(info): - assert otel_context.get_current() is caller - assert otel_context.get_current() is caller diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py index 723ad41ef..ba373ad8d 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py @@ -724,6 +724,7 @@ def test_operation_end_without_start_links_previous_logical_operation(): assert ( span.attributes["durable.operation.status"] == OperationStatus.SUCCEEDED.value ) + plugin.on_invocation_end(_invocation_end_info()) def test_continuation_span_uses_current_start_and_end_times(): @@ -753,6 +754,7 @@ def test_continuation_span_uses_current_start_and_end_times(): span = exporter.get_finished_spans()[0] assert invocation_span.start_time <= span.start_time assert before_callback <= span.start_time <= span.end_time <= after_callback + plugin.on_invocation_end(_invocation_end_info()) def test_resume_operation_timestamps_do_not_precede_current_invocation(): @@ -812,6 +814,7 @@ def test_resume_operation_timestamps_do_not_precede_current_invocation(): assert invocation_span.start_time <= after_resume_span.start_time assert after_resume_span.parent is not None assert after_resume_span.parent.span_id == invocation_span.context.span_id + plugin.on_invocation_end(_invocation_end_info()) def test_ordered_timestamps_are_thread_safe(): @@ -878,6 +881,7 @@ def test_retried_operation_uses_fresh_id_and_links_previous_logical_operation(): derive_workflow_span_id(EXECUTION_ARN), operation_id_to_span_id(EXECUTION_ARN, operation_id), } + plugin.on_invocation_end(_invocation_end_info()) def test_step_operation_span_parents_attempt_span(): @@ -1028,6 +1032,7 @@ def test_user_function_callbacks_emit_attempt_span_attributes(): == UserFunctionOutcome.SUCCEEDED.value ) assert "durable.operation.status" not in span.attributes + plugin.on_invocation_end(_invocation_end_info()) def test_step_attempt_span_name_includes_attempt_number(): @@ -1070,6 +1075,7 @@ def test_step_attempt_span_name_includes_attempt_number(): span = exporter.get_finished_spans()[0] assert span.name == "fetch-user attempt 2" + plugin.on_invocation_end(_invocation_end_info()) def test_step_attempt_span_name_defaults_to_first_attempt(): @@ -1112,6 +1118,7 @@ def test_step_attempt_span_name_defaults_to_first_attempt(): span = exporter.get_finished_spans()[0] assert span.name == "fetch-user attempt 1" + plugin.on_invocation_end(_invocation_end_info()) @pytest.mark.parametrize( @@ -1213,6 +1220,7 @@ def test_context_span_waits_for_terminal_status_and_omits_attempt_attributes( assert "durable.attempt.number" not in span.attributes assert "durable.attempt.outcome" not in span.attributes assert span.status.status_code is expected_span_status + plugin.on_invocation_end(_invocation_end_info()) def test_span_registry_helpers_can_be_called_from_multiple_threads(): @@ -1257,8 +1265,9 @@ def test_user_function_end_restores_enclosing_context(): # After the step, the enclosing context is restored and no scope is left # behind. Log correlation resolves the invocation span from the registry. assert otel_context.get_current() == enclosing_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} assert plugin.get_current_span_context().span_id == invocation_span_id + plugin.on_invocation_end(_invocation_end_info()) def test_user_function_start_preserves_baggage_in_current_context(): @@ -1292,7 +1301,8 @@ def test_user_function_end_restores_enclosing_context_on_failure(): ) assert otel_context.get_current() == enclosing_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} + plugin.on_invocation_end(_invocation_end_info()) def test_user_function_end_restores_enclosing_context_across_multiple_steps(): @@ -1309,8 +1319,9 @@ def test_user_function_end_restores_enclosing_context_across_multiple_steps(): # Between each step the context is back to where it started, and log # correlation still resolves the invocation span. assert otel_context.get_current() == enclosing_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} assert plugin.get_current_span_context().span_id == invocation_span_id + plugin.on_invocation_end(_invocation_end_info()) # ---------------------------------------------------------------------- @@ -1332,6 +1343,7 @@ def test_get_current_span_context_returns_invocation_span_at_top_level(): invocation_span = plugin._get_span(None) assert span_context is not None assert span_context.span_id == invocation_span.get_span_context().span_id + plugin.on_invocation_end(_invocation_end_info()) def test_get_current_span_context_returns_operation_span_inside_step(): @@ -1363,6 +1375,7 @@ def test_get_current_span_context_returns_invocation_span_between_steps(): invocation_span = plugin._get_span(None) assert span_context is not None assert span_context.span_id == invocation_span.get_span_context().span_id + plugin.on_invocation_end(_invocation_end_info()) # ---------------------------------------------------------------------- @@ -1441,11 +1454,12 @@ def test_top_level_step_end_falls_back_to_invocation_for_correlation(): plugin.on_user_function_start(_user_function_start_info(operation_id)) plugin.on_user_function_end(_user_function_end_info(operation_id)) - # No durable span is attached at the top level, so the registry fallback - # supplies the invocation span for log correlation. + # The invocation Start binding is restored after the step, including + # between-step user instrumentation and log correlation. assert otel_context.get_current() == enclosing_context - assert not trace.get_current_span().get_span_context().is_valid + assert trace.get_current_span().get_span_context().span_id == invocation_span_id assert plugin.get_current_span_context().span_id == invocation_span_id + plugin.on_invocation_end(_invocation_end_info()) def test_get_current_span_context_returns_context_span_between_nested_steps(): @@ -1709,6 +1723,7 @@ def test_replayed_context_span_links_previous_logical_operation(): derive_workflow_span_id(EXECUTION_ARN), operation_id_to_span_id(EXECUTION_ARN, operation_id), } + plugin.on_invocation_end(_invocation_end_info()) def test_checkpointed_context_first_span_uses_deterministic_id(): @@ -1867,7 +1882,8 @@ def test_child_context_end_restores_context_active_before_it(): ) assert otel_context.get_current() == enclosing_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} + plugin.on_invocation_end(_invocation_end_info()) def test_nested_scopes_are_released_without_accumulating(): @@ -1891,13 +1907,14 @@ def test_nested_scopes_are_released_without_accumulating(): ) # The inner step restored the child-context scope, not a copy of it. assert otel_context.get_current() == inside_context - assert set(plugin._context_tokens) == {context_id} + assert set(plugin._context_tokens) == {context_id, "__invocation_context__"} plugin.on_user_function_end( _user_function_end_info(context_id, operation_type=OperationType.CONTEXT) ) assert otel_context.get_current() == before_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} + plugin.on_invocation_end(_invocation_end_info()) def test_invocation_end_releases_scope_of_suspended_user_function(): @@ -1981,7 +1998,7 @@ def test_detach_ignores_token_attached_on_another_thread(): plugin._detach_context("step-1:attempt:1") - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} assert otel_context.get_current() == before_context plugin.on_invocation_end(_invocation_end_info()) @@ -2072,7 +2089,7 @@ def test_reentered_step_attempt_releases_the_previous_scope(): plugin.on_user_function_end(_user_function_end_info(operation_id)) assert otel_context.get_current() == before_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} plugin.on_invocation_end(_invocation_end_info()) @@ -2166,7 +2183,7 @@ def test_nested_suspension_unwinds_scopes_in_reverse_order(): ) ) assert otel_context.get_current() == before_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} # Neither span is ended: both operations are still in flight. assert not exporter.get_finished_spans() @@ -2199,7 +2216,7 @@ def test_nested_suspension_unwinds_scopes_in_reverse_order(): _user_function_end_info("ctx-outer", operation_type=OperationType.CONTEXT) ) assert otel_context.get_current() == before_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} plugin.on_invocation_end(_invocation_end_info()) @@ -2207,66 +2224,61 @@ def test_nested_suspension_unwinds_scopes_in_reverse_order(): @pytest.mark.parametrize("ambient_kind", ["same", "unrelated", "absent"]) @pytest.mark.parametrize("raises", [False, True]) @pytest.mark.parametrize("sampler", [ALWAYS_ON, ALWAYS_OFF]) -def test_handler_context_preserves_or_replaces_parent_and_restores_baggage( +def test_invocation_hooks_bind_parent_and_restore_baggage( ambient_kind: str, raises: bool, sampler: Sampler ) -> None: - plugin, _ = _create_plugin_with_sampler(sampler) - info = _invocation_start_info() - plugin.on_invocation_start(info) - invocation = plugin.get_current_span_context() - assert invocation is not None and invocation.is_valid - caller = baggage.set_baggage("tenant", "scope-test", Context()) - expected = invocation + provider = TracerProvider(sampler=sampler) + plugin = InvocationOtelPlugin( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + caller = baggage.set_baggage("tenant", "hook-test", Context()) + ambient = None if ambient_kind != "absent": ambient = SpanContext( - trace_id=( - invocation.trace_id - if ambient_kind == "same" - else (1 if invocation.trace_id != 1 else 2) - ), + trace_id=_to_otel_trace_id(EXECUTION_ARN, START_TIME) + if ambient_kind == "same" + else 1, span_id=0x42, is_remote=False, - trace_flags=invocation.trace_flags, + trace_flags=TraceFlags(1), ) caller = trace.set_span_in_context(NonRecordingSpan(ambient), caller) - if ambient_kind == "same": - expected = ambient token = otel_context.attach(caller) error = ValueError("handler error") try: + plugin.on_invocation_start(_invocation_start_info()) + expected = ( + ambient if ambient_kind == "same" else plugin.get_current_span_context() + ) def body() -> None: - with plugin.handler_context(info): - active = trace.get_current_span().get_span_context() - assert active == expected - assert active.is_valid - assert baggage.get_baggage("tenant") == "scope-test" - if raises: - raise error - - if raises: - with pytest.raises(ValueError) as caught: + active = trace.get_current_span().get_span_context() + assert active == expected + assert active.is_valid + assert baggage.get_baggage("tenant") == "hook-test" + if raises: + raise error + + try: + if raises: + with pytest.raises(ValueError) as caught: + body() + assert caught.value is error + else: body() - assert caught.value is error - else: - body() + finally: + plugin.on_invocation_end( + _invocation_end_info( + InvocationStatus.FAILED if raises else InvocationStatus.SUCCEEDED + ) + ) + assert plugin._context_tokens == {} assert otel_context.get_current() is caller - assert baggage.get_baggage("tenant") == "scope-test" + assert baggage.get_baggage("tenant") == "hook-test" finally: otel_context.detach(token) - plugin.on_invocation_end(_invocation_end_info()) - - -@pytest.mark.parametrize("completed", [False, True]) -def test_handler_context_without_live_invocation_leaves_context_unchanged( - completed: bool, -) -> None: - plugin, _ = _create_plugin() - info = _invocation_start_info() - if completed: - plugin.on_invocation_start(info) - plugin.on_invocation_end(_invocation_end_info()) - caller = otel_context.get_current() - with plugin.handler_context(info): - assert otel_context.get_current() is caller - assert otel_context.get_current() is caller + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py index 28477a460..b7a34ebda 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py @@ -6,8 +6,7 @@ PACKAGE_ROOT = Path(__file__).resolve().parents[1] REPOSITORY_ROOT = PACKAGE_ROOT.parents[1] -CORE_DEPENDENCY = "aws-durable-execution-sdk-python>=2.0.0" -EXCLUSIVITY_TEST_CORE = "aws-durable-execution-sdk-python>=2.1.0" +CORE_DEPENDENCY = "aws-durable-execution-sdk-python>=2.1.0" TEST_OTEL_DEPENDENCIES = { "opentelemetry-sdk>=1.20.0", "opentelemetry-propagator-aws-xray", @@ -68,7 +67,7 @@ def test_test_environments_install_layer_provided_dependencies() -> None: "test", "dev-otel", "dev-examples", - "test-pypi-otel", + "test-wheel-otel", "test-pypi-examples", ): assert TEST_OTEL_DEPENDENCIES <= set( @@ -77,22 +76,7 @@ def test_test_environments_install_layer_provided_dependencies() -> None: assert TEST_OTEL_DEPENDENCIES <= set(environments["types"]["extra-dependencies"]) -def test_pypi_compatibility_environment_requires_exclusivity_core() -> None: - dependencies = _load_pyproject(REPOSITORY_ROOT / "pyproject.toml")["tool"]["hatch"][ - "envs" - ]["test-pypi-otel"]["dependencies"] - - assert EXCLUSIVITY_TEST_CORE in dependencies - - -def test_pypi_otel_environment_installs_lifecycle_test_runner() -> None: - dependencies = _load_pyproject(REPOSITORY_ROOT / "pyproject.toml")["tool"]["hatch"][ - "envs" - ]["test-pypi-otel"]["dependencies"] - assert "aws-durable-execution-sdk-python-testing>=1.2.1" in dependencies - - -def test_core_dependency_preserves_previously_supported_releases() -> None: +def test_core_dependency_requires_worker_lifecycle_release() -> None: dependencies = _load_pyproject(PACKAGE_ROOT / "pyproject.toml")["project"][ "dependencies" ] @@ -101,42 +85,23 @@ def test_core_dependency_preserves_previously_supported_releases() -> None: for value in dependencies if Requirement(value).name == "aws-durable-execution-sdk-python" ) - assert requirement.specifier.contains("2.0.0") - assert requirement.specifier.contains("2.0.1") + assert not requirement.specifier.contains("2.0.0") + assert not requirement.specifier.contains("2.0.1") assert requirement.specifier.contains("2.1.0") -def test_pypi_otel_environment_does_not_shadow_installed_core() -> None: - environment = _load_pyproject(REPOSITORY_ROOT / "pyproject.toml")["tool"]["hatch"][ +def test_wheel_lanes_do_not_shadow_installed_artifacts() -> None: + environments = _load_pyproject(REPOSITORY_ROOT / "pyproject.toml")["tool"]["hatch"][ "envs" - ]["test-pypi-otel"] - assert environment["workspace"]["members"] == [ - "packages/aws-durable-execution-sdk-python-otel" ] - - -def test_legacy_lane_retains_supported_core_20_range() -> None: - environment = _load_pyproject(REPOSITORY_ROOT / "pyproject.toml")["tool"]["hatch"][ - "envs" - ]["test-pypi-otel-legacy"] - requirement = next( - Requirement(value) - for value in environment["dependencies"] - if Requirement(value).name == "aws-durable-execution-sdk-python" + for name in ("test-wheel-otel", "test-wheel-otel-legacy"): + assert environments[name]["workspace"]["members"] == [] + assert environments[name]["detached"] is True + assert ( + "aws-durable-execution-sdk-python-otel==1.0.0" + in environments["test-wheel-otel-legacy"]["dependencies"] + ) + assert ( + "aws-durable-execution-sdk-python-testing>=1.2.1" + in environments["test-wheel-otel"]["dependencies"] ) - assert requirement.specifier.contains("2.0.0") - assert requirement.specifier.contains("2.0.1") - assert not requirement.specifier.contains("2.1.0") - assert environment["workspace"]["members"] == [ - "packages/aws-durable-execution-sdk-python-otel" - ] - - -def test_minimum_core_lane_retains_exact_supported_floor() -> None: - environment = _load_pyproject(REPOSITORY_ROOT / "pyproject.toml")["tool"]["hatch"][ - "envs" - ]["test-pypi-otel-minimum"] - assert "aws-durable-execution-sdk-python==2.0.0" in environment["dependencies"] - assert environment["workspace"]["members"] == [ - "packages/aws-durable-execution-sdk-python-otel" - ] diff --git a/packages/aws-durable-execution-sdk-python/README.md b/packages/aws-durable-execution-sdk-python/README.md index bfe54384f..08a70e06a 100644 --- a/packages/aws-durable-execution-sdk-python/README.md +++ b/packages/aws-durable-execution-sdk-python/README.md @@ -85,9 +85,8 @@ constraint applies to the combined `plugins=[...]` argument and Core 2.1+ with OTel 1.1+ rejects both at cold start with `PluginLoadError` naming the conflicting views; keep only one. Choose Invocation for work within each Lambda invocation or Execution for logical operations across the durable execution. -Unrelated instrumentation plugins can run alongside either view. Existing valid -registrations remain supported with OTel 1.1 on older core 2.0.x; those cores do -not implement the new exclusivity validation. +Unrelated instrumentation plugins can run alongside either view. OTel 1.1 requires core 2.1 or later for the invocation-worker lifecycle as well +as view-exclusivity validation. Release core 2.1 before releasing OTel 1.1. Plugin authors explicitly opt in by declaring `__durable_registration_api__ = 1` on a plugin class. That class and its subclasses @@ -112,22 +111,34 @@ in their hierarchy retain their existing attributes/helpers. The generic plugin neither attribute nor hook, and provider API version 1 and existing plugin lifecycle order are unchanged. -### Optional handler context scopes - -A plugin can provide `handler_context(info)` returning a context manager. The -updated core looks up this optional method normally, so inherited methods work -without extra declarations. A missing or non-callable attribute is ignored. -The core enters these scopes around the top-level handler on its worker thread, -in registration order, and closes them in reverse order. Cleanup receives no -handler exception and cannot suppress or replace its outcome. Invocation hooks -retain their original thread and order; failed lookup, setup or entry bindings -are discarded, and successful scope cleanup stays in the context that owns its -tokens. - -`handler_context` is an optional plugin API name: callable methods or attributes -with that name are invoked. The generic plugin base does not require a default -method. Older cores ignore this optional API and retain their existing behavior. -The provider API version and dependency requirements are unchanged. +### Invocation worker lifecycle + +The existing `on_invocation_start` and `on_invocation_end` hooks run on the +handler's invocation worker, both in registration order. Start hooks finish +before background checkpoint processing begins. The same worker runs the handler +(including its `finally` blocks), prepares the existing serialized output or +error, checkpoints large results when necessary, joins registered branches while +checkpointing is available, stops and waits for checkpoint processing, then calls +End before returning the outcome. The caller shuts down the handler executor. +There is no separate handler-context hook or context-manager plugin API. + +When plugins are registered, the worker begins with a copy of the caller's +context-variable bindings. Successful Start bindings flow into later Start hooks, +the handler, serialization and resource cleanup. A later successful Start may +replace an earlier binding. If a Start hook raises, its new bindings are discarded +for subsequent work; its End still runs in the original Context so its tokens can +be reset. End hooks retain forward registration order, not reverse stack order, +so they must not rely on observing a stack-like unwind of other plugins' contexts. +The worker's invocation context is discarded on return, including after plugin +cleanup failures, leaving the host's bindings unchanged. This isolates bindings, +not mutations to shared objects or external side effects. Without plugins, the +handler retains its existing fresh-worker context behavior. + +These lifecycle guarantees require core 2.1.0 or later. The OTel 1.1 plugin +requires that core version and uses the existing Start/End hooks. Upgrading the +core alone with OTel 1.0 isolates worker bindings, but does not add the newer +Invocation-view fallback to that older plugin; upgrade both packages for it. +Provider API version 1 and the independent registration API are unchanged. ## 🚀 Quick Start 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 732f6d6e8..e59820007 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 @@ -5,7 +5,7 @@ import functools import json import logging -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import Future, ThreadPoolExecutor, wait from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any @@ -192,7 +192,6 @@ def durable_execution( configured_plugins = load_configured_plugins(plugins) plugin_executor = PluginExecutor(configured_plugins) - @plugin_executor.handle_durable_output def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: invocation_input: DurableExecutionInvocationInput service_client: DurableServiceClient @@ -281,63 +280,13 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: ), ) - # Use ThreadPoolExecutor for concurrent execution of user code and background checkpoint processing - with ( - ThreadPoolExecutor( - max_workers=2, thread_name_prefix="dex-handler" - ) as executor, - contextlib.closing(execution_state) as execution_state, - ): - execution_operation = execution_state.get_execution_operation() - - # execute the plugins - plugin_executor.on_invocation_start( - execution_arn=invocation_input.durable_execution_arn, - lambda_context=context, - execution_start_time=( - execution_operation.start_timestamp - if execution_operation is not None - else None - ), - 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) - - # Thread 2: Execute user function + def invoke_handler() -> MutableMapping[str, Any]: logger.debug( "%s entering user-space...", invocation_input.durable_execution_arn ) - if configured_plugins: - # Invocation-start hooks can establish tracing and other contextvars. - # Context.run restores worker bindings on both return and failure. - user_future = executor.submit( - contextvars.copy_context().run, - plugin_executor.run_handler, - func, - input_event, - durable_context, - ) - else: - # Preserve the original fresh-worker context for uninstrumented - # handlers, including the absence of caller ContextVar bindings. - user_future = executor.submit(func, input_event, durable_context) - - logger.debug( - "%s waiting for user code completion...", - invocation_input.durable_execution_arn, - ) - try: # Background checkpointing errors will propagate through CompletionEvent.wait() as BackgroundThreadError - result = user_future.result() + result = func(input_event, durable_context) # done with userland logger.debug( @@ -473,6 +422,78 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: return result + def invoke_worker() -> MutableMapping[str, Any]: + # This task owns Start, user cleanup, output preparation, resource + # cleanup and End. The caller owns the executor that runs this task. + with plugin_executor.run(): + checkpoint_future: Future[None] | None = None + + def wait_for_checkpoint() -> None: + if checkpoint_future is not None: + # Match the former executor join: observed checkpoint + # errors retain their existing classification paths. + wait((checkpoint_future,)) + + try: + with contextlib.ExitStack() as resources: + # LIFO cleanup: branches join with checkpointing alive; + # then close stops checkpointing and the wait completes. + resources.callback(wait_for_checkpoint) + resources.callback( + plugin_executor._run_in_invocation_context, + execution_state.close, + ) + execution_operation = execution_state.get_execution_operation() + + # execute the plugins + plugin_executor.on_invocation_start( + execution_arn=invocation_input.durable_execution_arn, + lambda_context=context, + execution_start_time=( + execution_operation.start_timestamp + if execution_operation is not None + else None + ), + 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, + ) + # No checkpoint work starts until all Start hooks finish. + checkpoint_future = executor.submit( + execution_state.checkpoint_batches_forever + ) + output = plugin_executor._run_in_invocation_context( + invoke_handler + ) + plugin_executor.on_invocation_end( + DurableExecutionInvocationOutput.from_dict(output) + ) + return output + except Exception as error: + plugin_executor.on_invocation_end( + DurableExecutionInvocationOutput.create_retry( + ErrorObject.from_exception(error) + ) + ) + raise + + with ThreadPoolExecutor( + max_workers=2, thread_name_prefix="dex-handler" + ) as executor: + if configured_plugins: + invocation_future = executor.submit( + contextvars.copy_context().run, invoke_worker + ) + else: + # Preserve the original fresh-worker context without plugins. + invocation_future = executor.submit(invoke_worker) + return invocation_future.result() + return wrapper 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 7b3ccf0bc..856fc68c6 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 @@ -4,13 +4,12 @@ import contextvars import copy import datetime -import functools import logging -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from enum import Enum -from typing import Any, Callable, MutableMapping, cast +from typing import Any, Callable, cast from aws_durable_execution_sdk_python.identifier import OperationIdentifier from aws_durable_execution_sdk_python.lambda_service import ( @@ -467,9 +466,8 @@ def __init__(self, plugins: list[DurableInstrumentationPlugin] | None): self._executor: ThreadPoolExecutor | None = None self._invocation_status: InvocationStartInfo | None = None self._operations_provider: Callable[[], Mapping[str, Operation]] | None = None - # Non-None only after a start hook fails: the pre-hook snapshot excludes - # its partial bindings from later setup and the handler. Otherwise this - # stays None, and run_handler uses a fresh copy_context(). + # Non-None only after a Start hook fails: later setup and invocation + # work use its pre-hook snapshot. End still uses each hook's token owner. self._startup_context: contextvars.Context | None = None self._invocation_contexts: list[contextvars.Context | None] = [] @@ -557,56 +555,11 @@ 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) - @contextlib.contextmanager - def _safe_handler_context( - self, plugin: DurableInstrumentationPlugin, info: InvocationStartInfo - ) -> Iterator[bool]: - # Plugins may omit this optional scope method. - scope = None - succeeded = True - try: - factory = getattr(plugin, "handler_context", None) - if callable(factory): - scope = factory(info) - scope.__enter__() - except Exception: - scope = None - succeeded = False - logger.exception( - "Plugin %s handler context failed", plugin.__class__.__name__ - ) - try: - yield succeeded - finally: - if scope is not None: - try: - scope.__exit__(None, None, None) - except Exception: - logger.exception( - "Plugin %s handler context cleanup failed", - plugin.__class__.__name__, - ) - - def run_handler(self, handler: Callable[..., Any], *args: Any) -> Any: - """Run the user handler inside optional, balanced plugin context scopes.""" - if not self._plugins or self._invocation_status is None: - return handler(*args) - owner = ( - self._startup_context.copy() - if self._startup_context is not None - else contextvars.copy_context() - ) - with contextlib.ExitStack() as scopes: - for plugin in self._plugins: - before = owner.copy() - scope = self._safe_handler_context(plugin, self._invocation_status) - succeeded = owner.run(scope.__enter__) - # Keep finalizers with their entry Context: ContextVar tokens - # cannot be reset in a copy, even when its bindings are identical. - scopes.callback(owner.run, scope.__exit__, None, None, None) - if not succeeded: - owner = before - return owner.run(handler, *args) + def _run_in_invocation_context(self, invoke: Callable[[], Any]) -> Any: + """Continue in the pre-hook Context only after a failed Start hook.""" + if self._startup_context is not None: + return self._startup_context.run(invoke) + return invoke() def _snapshot_operation_infos( self, @@ -925,37 +878,3 @@ def _is_terminal_status(status): OperationStatus.CANCELLED, OperationStatus.STOPPED, ] - - @property - def handle_durable_output(self): - def decorator(func: Callable[[Any, LambdaContext], MutableMapping[str, Any]]): - def invoke(event: Any, context: LambdaContext): - with self.run(): - try: - output = func(event, context) - - self.on_invocation_end( - output=DurableExecutionInvocationOutput.from_dict(output), - ) - return output - except Exception as e: - self.on_invocation_end( - output=DurableExecutionInvocationOutput.create_retry( - ErrorObject.from_exception(e) - ), - ) - raise - - @functools.wraps(func) - def wrapper(event: Any, context: LambdaContext): - if not self._plugins: - return invoke(event, context) - # Keep hooks on their existing caller thread and in registration - # order, but isolate their context bindings from the host. Two - # plugins can otherwise restore a stale predecessor when their - # invocation-end hooks close scopes in the original order. - return contextvars.copy_context().run(invoke, event, context) - - return wrapper - - return decorator diff --git a/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py index c62377261..8cb081029 100644 --- a/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py @@ -1,513 +1,589 @@ -"""Focused handler-dispatch tests with a mock service and real worker thread.""" +"""Real public invocation-worker lifecycle and ContextVar ownership controls.""" from __future__ import annotations import contextvars -from collections.abc import Callable -from concurrent.futures import Future, ThreadPoolExecutor +import threading +from concurrent.futures import ThreadPoolExecutor from typing import Any from unittest.mock import Mock import pytest from aws_durable_execution_sdk_python.context import DurableContext -from aws_durable_execution_sdk_python.exceptions import InvocationError -from aws_durable_execution_sdk_python.execution import durable_execution +from aws_durable_execution_sdk_python.exceptions import ( + CheckpointError, + CheckpointErrorCategory, + InvocationError, + SuspendExecution, +) +from aws_durable_execution_sdk_python.execution import ( + DurableExecutionInvocationInputWithClient, + InitialExecutionState, + durable_execution, +) +from aws_durable_execution_sdk_python.lambda_service import ( + CheckpointOutput, + CheckpointUpdatedExecutionState, + DurableServiceClient, + ExecutionDetails, + Operation, + OperationStatus, + OperationType, + OperationUpdate, +) from aws_durable_execution_sdk_python.plugin import ( DurableInstrumentationPlugin, InvocationEndInfo, InvocationStartInfo, InvocationStatus, ) +from aws_durable_execution_sdk_python.state import ExecutionState -@pytest.mark.parametrize("outcome", ["success", "failure", "retry"]) -@pytest.mark.parametrize("plugin_mode", ["none", "healthy", "partial-failure"]) -def test_handler_worker_preserves_context_and_restores_its_caller( - monkeypatch: pytest.MonkeyPatch, outcome: str, plugin_mode: str +def invocation(client: Any = None) -> tuple[Any, Any]: + client = client or Mock(spec=DurableServiceClient) + client.checkpoint.return_value = CheckpointOutput( + checkpoint_token="next", new_execution_state=CheckpointUpdatedExecutionState() + ) + event = DurableExecutionInvocationInputWithClient( + durable_execution_arn="test-arn/worker-lifecycle", + checkpoint_token="initial", + initial_execution_state=InitialExecutionState( + operations=[ + Operation( + operation_id="execution", + operation_type=OperationType.EXECUTION, + status=OperationStatus.STARTED, + execution_details=ExecutionDetails(input_payload="{}"), + ) + ], + next_marker="", + ), + service_client=client, + ) + context = Mock() + context.aws_request_id = "worker-request" + context.client_context = context.identity = context.invoked_function_arn = None + context._epoch_deadline_time_in_ms = 0 + context.tenant_id = None + return event, context + + +@pytest.mark.parametrize("outcome", ["success", "failure", "retry", "pending"]) +@pytest.mark.parametrize("mode", ["none", "healthy", "failed-start", "failed-end"]) +def test_public_worker_lifecycle_owns_context_and_outcome( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + outcome: str, + mode: str, ) -> None: - marker = contextvars.ContextVar("handler-worker-context", default="worker-empty") - seen: list[str] = [] - worker_boundaries: list[tuple[str, str]] = [] + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + marker = contextvars.ContextVar("worker-owner", default="fresh-worker") + phases: list[tuple[str, int, str]] = [] + boundary: list[tuple[str, int, str, str]] = [] + retry = InvocationError("original retry") + fail = ValueError("original body failure") statuses: list[InvocationStatus] = [] + caller_thread = threading.get_ident() - class ClaimPlugin(DurableInstrumentationPlugin): - token: contextvars.Token[str] | None = None - - def on_invocation_start(self, info: InvocationStartInfo) -> None: - self.token = marker.set("invocation-start") - if plugin_mode == "partial-failure": - raise ValueError("partial plugin setup") - - def on_invocation_end(self, info: InvocationEndInfo) -> None: - statuses.append(info.status) - assert self.token is not None - marker.reset(self.token) - self.token = None - - def body(_event: Any, _context: DurableContext) -> str: - seen.append(marker.get()) - marker.set("worker-mutation") - if outcome == "failure": - raise ValueError("handler failure") - if outcome == "retry": - raise InvocationError("handler retry") - return "ok" - - class ObservingExecutor(ThreadPoolExecutor): - def submit( - self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any - ) -> Future[Any]: - if fn is not body and body not in args: - return super().submit(fn, *args, **kwargs) + class Observer(ThreadPoolExecutor): + def submit(self, fn: Any, /, *args: Any, **kwargs: Any) -> Any: + role = ( + "checkpoint" + if getattr(fn, "__name__", "") == "checkpoint_batches_forever" + else "invocation" + ) - def observe() -> Any: + def observed() -> Any: before = marker.get() try: return fn(*args, **kwargs) finally: - # This runs outside Context.run, on the actual SDK worker. - worker_boundaries.append((before, marker.get())) + boundary.append((role, threading.get_ident(), before, marker.get())) - return super().submit(observe) + return super().submit(observed) monkeypatch.setattr( - "aws_durable_execution_sdk_python.execution.ThreadPoolExecutor", - ObservingExecutor, + "aws_durable_execution_sdk_python.execution.ThreadPoolExecutor", Observer ) - monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) - client = Mock() - handler = durable_execution( - body, - boto3_client=client, - plugins=[ClaimPlugin()] if plugin_mode != "none" else [], - ) - event = { - "DurableExecutionArn": "test-arn/handler-context", - "CheckpointToken": "test-token", - "InitialExecutionState": { - "Operations": [ - { - "Id": "handler-context", - "Type": "EXECUTION", - "Status": "STARTED", - "ExecutionDetails": {"InputPayload": "{}"}, - } - ], - "NextMarker": "", - }, - } - lambda_context = Mock() - lambda_context.aws_request_id = "context-request" - lambda_context.client_context = None - lambda_context.identity = None - lambda_context._epoch_deadline_time_in_ms = 0 - lambda_context.invoked_function_arn = "test-arn" - lambda_context.tenant_id = None - token = marker.set("caller") - try: - if outcome == "retry": - with pytest.raises(InvocationError, match="handler retry"): - handler(event, lambda_context) - else: - result = handler(event, lambda_context) - assert result["Status"] == ( - "SUCCEEDED" if outcome == "success" else "FAILED" - ) - assert marker.get() == "caller" - finally: - marker.reset(token) - assert len(worker_boundaries) == 1 - worker_before, worker_after = worker_boundaries[0] - if plugin_mode != "none": - assert seen == [ - "caller" if plugin_mode == "partial-failure" else "invocation-start" - ] - assert worker_after == worker_before - assert worker_after != "worker-mutation" - assert statuses == [ - { - "success": InvocationStatus.SUCCEEDED, - "failure": InvocationStatus.FAILED, - "retry": InvocationStatus.RETRY, - }[outcome] - ] - else: - # No plugin means the original direct worker call: caller bindings are - # absent, and mutations belong to the worker's own context. - assert seen == [worker_before] - assert seen != ["caller"] - assert worker_after == "worker-mutation" - assert statuses == [] - client.checkpoint_durable_execution.assert_not_called() - - -@pytest.mark.parametrize("failure", [None, "body", "enter", "exit"]) -def test_optional_handler_scopes_are_balanced_and_cannot_change_outcome( - failure: str | None, -) -> None: - from contextlib import contextmanager - from datetime import UTC, datetime - from collections.abc import Iterator - from aws_durable_execution_sdk_python.plugin import PluginExecutor - marker = contextvars.ContextVar("handler-scope", default="caller") - events: list[str] = [] - - class ScopePlugin(DurableInstrumentationPlugin): - def __init__(self, name: str): - self.name = name + class Plugin(DurableInstrumentationPlugin): + token: Any = None - @contextmanager - def handler_context(self, info: InvocationStartInfo) -> Iterator[None]: - assert info.execution_arn == "handler-scope" - events.append("enter-" + self.name) - if self.name == "inner" and failure == "enter": - raise ValueError("plugin entry failure") - token = marker.set(self.name) - try: - yield - finally: - marker.reset(token) - events.append("exit-" + self.name) - if self.name == "inner" and failure == "exit": - raise ValueError("plugin cleanup failure") - - executor = PluginExecutor([ScopePlugin("outer"), ScopePlugin("inner")]) - - def body() -> str: - assert marker.get() == ("outer" if failure == "enter" else "inner") - events.append("body") - if failure in ("body", "exit"): - raise RuntimeError("original handler failure") - return "ok" - - with executor.run(): - executor.on_invocation_start("handler-scope", True, datetime.now(UTC), None) - if failure in ("body", "exit"): - with pytest.raises(RuntimeError, match="original handler failure"): - executor.run_handler(body) - else: - assert executor.run_handler(body) == "ok" - assert marker.get() == "caller" - assert events == ["enter-outer", "enter-inner", "body"] + ( - ["exit-outer"] if failure == "enter" else ["exit-inner", "exit-outer"] - ) - - -def test_handler_accepts_plugin_without_optional_scope() -> None: - from datetime import UTC, datetime - from types import SimpleNamespace - from typing import cast - from aws_durable_execution_sdk_python.plugin import PluginExecutor - - legacy = cast( - DurableInstrumentationPlugin, - SimpleNamespace(on_invocation_start=lambda info: None), - ) - executor = PluginExecutor([legacy]) - with executor.run(): - executor.on_invocation_start("legacy", True, datetime.now(UTC), None) - assert executor.run_handler(lambda: "unchanged") == "unchanged" - - -@pytest.mark.parametrize("outcome", ["SUCCEEDED", "PENDING", "FAILED", "retry"]) -@pytest.mark.parametrize("reverse", [False, True]) -@pytest.mark.parametrize("hook_failure", [None, "start", "end"]) -def test_invocation_context_scopes_do_not_escape_to_host( - outcome: str, reverse: bool, hook_failure: str | None -) -> None: - """Legacy hook order must not leave an already-ended plugin scope current.""" - from datetime import UTC, datetime - import threading - from aws_durable_execution_sdk_python.plugin import PluginExecutor - - marker = contextvars.ContextVar("invocation-scope", default="host") - events: list[tuple[str, str, str, int]] = [] - caller_thread = threading.get_ident() - - class ScopePlugin(DurableInstrumentationPlugin): - def __init__(self, name: str) -> None: - self.name = name - self.token: contextvars.Token[str] | None = None - - def on_invocation_start(self, _info: InvocationStartInfo) -> None: - events.append(("start", self.name, marker.get(), threading.get_ident())) - self.token = marker.set(self.name) - if self.name == names[0] and hook_failure == "start": - raise ValueError("plugin initialization failed") + def on_invocation_start(self, info: InvocationStartInfo) -> None: + phases.append(("start", threading.get_ident(), marker.get())) + self.token = marker.set("start-binding") + if mode == "failed-start": + raise ValueError("startup failure") - def on_invocation_end(self, _info: InvocationEndInfo) -> None: - events.append(("end", self.name, marker.get(), threading.get_ident())) - if self.name == names[0] and hook_failure == "end": - raise ValueError("plugin finalization failed") + def on_invocation_end(self, info: InvocationEndInfo) -> None: + phases.append(("end", threading.get_ident(), marker.get())) + statuses.append(info.status) assert self.token is not None marker.reset(self.token) self.token = None + if mode == "failed-end": + raise ValueError("end failure") - names = ["first", "second"] - if reverse: - names.reverse() - executor = PluginExecutor([ScopePlugin(name) for name in names]) - handler_failure = InvocationError("retry") - expected_output = {"Status": outcome} - - @executor.handle_durable_output - def invoke(_event: Any, _context: Any) -> dict[str, str]: - executor.on_invocation_start("test", True, datetime.now(UTC), None) - assert executor.run_handler(marker.get) == names[-1] - if outcome == "retry": - raise handler_failure - return expected_output - - token = marker.set("incoming") + def body(_event: Any, _context: DurableContext) -> str: + phases.append(("body", threading.get_ident(), marker.get())) + marker.set("body-binding") + try: + if outcome == "failure": + raise fail + if outcome == "retry": + raise retry + if outcome == "pending": + raise SuspendExecution("test suspension") + return "ok" + finally: + phases.append(("finally", threading.get_ident(), marker.get())) + + handler = durable_execution(body, plugins=[] if mode == "none" else [Plugin()]) + event, context = invocation() + token = marker.set("host") try: for _ in range(2): if outcome == "retry": - with pytest.raises(InvocationError, match="retry") as caught: - invoke({}, None) - assert caught.value is handler_failure + with pytest.raises(InvocationError) as caught: + handler(event, context) + assert caught.value is retry else: - assert invoke({}, None) is expected_output - assert marker.get() == "incoming" - finally: - marker.reset(token) - assert [(kind, name) for kind, name, _, _ in events] == [ - (kind, name) for _ in range(2) for kind in ("start", "end") for name in names - ] - assert all(thread == caller_thread for _, _, _, thread in events) - assert [ - value for kind, name, value, _ in events if kind == "start" and name == names[0] - ] == ["incoming", "incoming"] - - -def test_no_plugin_invocation_keeps_original_caller_context_semantics() -> None: - from aws_durable_execution_sdk_python.plugin import PluginExecutor - - marker = contextvars.ContextVar("no-plugin-caller", default="host") - executor = PluginExecutor([]) - - @executor.handle_durable_output - def invoke(_event: Any, _context: Any) -> dict[str, str]: - marker.set("caller-side-change") - return {"Status": "SUCCEEDED"} - - token = marker.set("incoming") - try: - invoke({}, None) - assert marker.get() == "caller-side-change" + result = handler(event, context) + assert ( + result["Status"] + == { + "success": "SUCCEEDED", + "failure": "FAILED", + "pending": "PENDING", + }[outcome] + ) + if outcome == "failure": + assert result["Error"]["ErrorMessage"] == "original body failure" + assert marker.get() == "host" finally: marker.reset(token) + per_call = 2 if mode == "none" else 4 + for offset in range(0, len(phases), per_call): + batch = phases[offset : offset + per_call] + assert [p[0] for p in batch] == ( + ["body", "finally"] + if mode == "none" + else ["start", "body", "finally", "end"] + ) + assert len({p[1] for p in batch}) == 1 + assert batch[0][1] != caller_thread + bodies = [p for p in phases if p[0] == "body"] + expected = ( + "fresh-worker" + if mode == "none" + else "host" + if mode == "failed-start" + else "start-binding" + ) + assert [p[2] for p in bodies] == [expected, expected] + for role, tid, before, after in boundary: + if role == "invocation": + assert after == ("body-binding" if mode == "none" else before) + if mode != "none": + assert ( + statuses + == [ + { + "success": InvocationStatus.SUCCEEDED, + "failure": InvocationStatus.FAILED, + "retry": InvocationStatus.RETRY, + "pending": InvocationStatus.PENDING, + }[outcome] + ] + * 2 + ) + assert "different Context" not in caplog.text -@pytest.mark.parametrize("stage", ["start", "factory", "enter"]) @pytest.mark.parametrize("bad_first", [False, True]) -@pytest.mark.parametrize("outcome", ["SUCCEEDED", "PENDING", "FAILED", "retry"]) -def test_failed_plugin_setup_discards_partial_context_bindings( - stage: str, bad_first: bool, outcome: str, caplog: pytest.LogCaptureFixture +@pytest.mark.parametrize("outcome", ["success", "failure", "retry", "pending"]) +def test_failed_start_preserves_clean_bindings_and_original_token_owners( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + bad_first: bool, + outcome: str, ) -> None: - """Keep healthy bindings, unset bindings, hook order, and token ownership.""" - from contextlib import nullcontext - from datetime import UTC, datetime - from typing import ContextManager - from aws_durable_execution_sdk_python.plugin import PluginExecutor - - marker = contextvars.ContextVar("partial-setup", default="default") - new_binding = contextvars.ContextVar[str]("partial-setup-no-default") - events: list[tuple[str, str]] = [] - cleanup: list[str] = [] - healthy_inputs: list[tuple[str, str | None]] = [] - - class Scope: - def __init__(self, plugin: SetupPlugin) -> None: - self.plugin = plugin - - def __enter__(self) -> None: - self.plugin.bind("enter") + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + marker = contextvars.ContextVar("failed-start-owner", default="host") + added = contextvars.ContextVar[str]("failed-start-added") + events: list[tuple[str, str, int]] = [] + cleanup_contexts: list[str] = [] + original_close = ExecutionState.close + + def close(state: ExecutionState) -> None: + cleanup_contexts.append(marker.get()) + with pytest.raises(LookupError): + added.get() + original_close(state) - def __exit__(self, *_args: Any) -> None: - self.plugin.reset("exit") + monkeypatch.setattr(ExecutionState, "close", close) - class SetupPlugin(DurableInstrumentationPlugin): + class Plugin(DurableInstrumentationPlugin): def __init__(self, name: str) -> None: self.name = name self.token: contextvars.Token[str] | None = None + self.added: contextvars.Token[str] | None = None - def bind(self, where: str) -> None: - events.append((where, self.name)) - if self.name == "healthy": - healthy_inputs.append((marker.get(), new_binding.get(None))) + def on_invocation_start(self, _info: InvocationStartInfo) -> None: + events.append(("start", self.name, threading.get_ident())) self.token = marker.set(self.name) if self.name == "bad": - new_binding.set("partial") - raise ValueError("partial plugin setup") + self.added = added.set("partial") + raise ValueError("partial startup") - def reset(self, where: str) -> None: - events.append((where, self.name)) + def on_invocation_end(self, _info: InvocationEndInfo) -> None: + events.append(("end", self.name, threading.get_ident())) assert self.token is not None marker.reset(self.token) - self.token = None - cleanup.append(self.name) - - def on_invocation_start(self, _info: InvocationStartInfo) -> None: - if stage == "start": - self.bind("start") - - def on_invocation_end(self, _info: InvocationEndInfo) -> None: - if stage == "start": - self.reset("end") + if self.added is not None: + added.reset(self.added) - def handler_context(self, _info: InvocationStartInfo) -> ContextManager[None]: - if stage == "start": - return nullcontext() - if self.name == "bad" and stage == "factory": - self.bind("factory") - return Scope(self) + failure = InvocationError("retry") - names = ["bad", "healthy"] if bad_first else ["healthy", "bad"] - executor = PluginExecutor([SetupPlugin(name) for name in names]) - output = {"Status": outcome} - failure = InvocationError("original retry") - handler_failure = RuntimeError("original handler error") - - def body() -> dict[str, str]: + def body(_event: Any, _context: DurableContext) -> str: + events.append(("body", "handler", threading.get_ident())) assert marker.get() == "healthy" with pytest.raises(LookupError): - new_binding.get() + added.get() if outcome == "retry": raise failure - if outcome == "FAILED": - raise handler_failure - return output - - @executor.handle_durable_output - def invoke(_event: Any, _context: Any) -> dict[str, str]: - executor.on_invocation_start("partial-setup", True, datetime.now(UTC), None) - try: - return executor.run_handler(body) - except RuntimeError as error: - assert error is handler_failure - return output + if outcome == "failure": + raise ValueError("body") + if outcome == "pending": + raise SuspendExecution("test suspension") + return "ok" - token = marker.set("incoming") - try: - for _ in range(2): - if outcome == "retry": - with pytest.raises(InvocationError) as caught: - invoke({}, None) - assert caught.value is failure - else: - assert invoke({}, None) is output - assert marker.get() == "incoming" - with pytest.raises(LookupError): - new_binding.get() - finally: - marker.reset(token) - assert healthy_inputs == [("incoming", None)] * 2 - assert cleanup == (names if stage == "start" else ["healthy"]) * 2 - setup = "start" if stage == "start" else "enter" - expected_setup = [ - ("factory" if name == "bad" and stage == "factory" else setup, name) - for name in names - ] - expected_cleanup = ( - [("end", name) for name in names] if stage == "start" else [("exit", "healthy")] - ) - assert events == (expected_setup + expected_cleanup) * 2 - # A token reset in the wrong Context is caught by the SDK, so explicitly - # check diagnostics and successful cleanup rather than relying on raises. + names = ["bad", "healthy"] if bad_first else ["healthy", "bad"] + handler = durable_execution(body, plugins=[Plugin(name) for name in names]) + event, context = invocation() + for _ in range(2): + if outcome == "retry": + with pytest.raises(InvocationError) as caught: + handler(event, context) + assert caught.value is failure + else: + assert ( + handler(event, context)["Status"] + == {"success": "SUCCEEDED", "failure": "FAILED", "pending": "PENDING"}[ + outcome + ] + ) + assert marker.get() == "host" + with pytest.raises(LookupError): + added.get() + assert cleanup_contexts == ["healthy"] * 2 + for offset in (0, 5): + batch = events[offset : offset + 5] + assert [(x[0], x[1]) for x in batch] == [("start", n) for n in names] + [ + ("body", "handler") + ] + [("end", n) for n in names] + assert len({x[2] for x in batch}) == 1 errors = [ - record.exc_info for record in caplog.records if record.exc_info is not None + r.exc_info + for r in caplog.records + if r.exc_info and r.name == "aws_durable_execution_sdk_python.plugin" ] - assert len(errors) == 2 - assert all(str(error[1]) == "partial plugin setup" for error in errors) + assert len(errors) == 2 and all(str(e[1]) == "partial startup" for e in errors) -@pytest.mark.parametrize("inherited", [False, True]) -@pytest.mark.parametrize("fail_handler", [False, True]) -def test_handler_scope_uses_direct_or_inherited_optional_method( - inherited: bool, fail_handler: bool +@pytest.mark.parametrize("result_kind", ["normal", "bad-json", "large", "large-error"]) +def test_worker_prepares_output_joins_branches_waits_checkpoint_then_ends( + monkeypatch: pytest.MonkeyPatch, + result_kind: str, ) -> None: - from contextlib import contextmanager - from collections.abc import Iterator - from datetime import UTC, datetime - from aws_durable_execution_sdk_python.plugin import PluginExecutor - - marker = contextvars.ContextVar("inherited-scope", default="outside") - events: list[str] = [] - error = RuntimeError("original handler failure") - - class ScopePlugin(DurableInstrumentationPlugin): - @contextmanager - def handler_context(self, _info: InvocationStartInfo) -> Iterator[None]: - events.append("enter") - token = marker.set("inside") - try: - yield - finally: - marker.reset(token) - events.append("exit") - - class InheritedPlugin(ScopePlugin): - pass - - plugin = InheritedPlugin() if inherited else ScopePlugin() - executor = PluginExecutor([plugin]) - - def handler() -> str: - assert marker.get() == "inside" - events.append("body") - if fail_handler: - raise error - return "ok" + import aws_durable_execution_sdk_python.execution as execution - with executor.run(): - for _ in range(2): - executor.on_invocation_start("inherited", True, datetime.now(UTC), None) - if fail_handler: - with pytest.raises(RuntimeError) as caught: - executor.run_handler(handler) - assert caught.value is error - else: - assert executor.run_handler(handler) == "ok" - assert marker.get() == "outside" - assert events == ["enter", "body", "exit"] * 2 + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + events: list[tuple[str, int]] = [] + started = threading.Event() + checkpoint_started = threading.Event() + closing = threading.Event() + branch_done = threading.Event() + real_checkpoint = ExecutionState.checkpoint_batches_forever + real_close = ExecutionState.close + real_stop = ExecutionState.stop_checkpointing + real_dumps = execution.json.dumps + result: Any = { + "normal": {"ok": True}, + "bad-json": object(), + "large": "x" * 256, + "large-error": None, + }[result_kind] + + def record(name: str) -> None: + events.append((name, threading.get_ident())) + + def checkpoint(state: ExecutionState) -> None: + assert started.is_set() + record("checkpoint-start") + checkpoint_started.set() + try: + real_checkpoint(state) + finally: + record("checkpoint-end") + + def stop(state: ExecutionState) -> None: + assert branch_done.is_set() + record("checkpoint-stop") + real_stop(state) + + def close(state: ExecutionState) -> None: + record("close") + closing.set() + real_close(state) + + def dumps(value: Any, *args: Any, **kwargs: Any) -> str: + if value is result or ( + isinstance(value, dict) and value.get("Status") == "FAILED" + ): + record("serialize") + return real_dumps(value, *args, **kwargs) + + monkeypatch.setattr(ExecutionState, "checkpoint_batches_forever", checkpoint) + monkeypatch.setattr(ExecutionState, "stop_checkpointing", stop) + monkeypatch.setattr(ExecutionState, "close", close) + monkeypatch.setattr(execution.json, "dumps", dumps) + if result_kind in ("large", "large-error"): + monkeypatch.setattr(execution, "LAMBDA_RESPONSE_SIZE_LIMIT", 64) + + class Plugin(DurableInstrumentationPlugin): + def on_invocation_start(self, _info: InvocationStartInfo) -> None: + record("start") + assert not checkpoint_started.is_set() + started.set() + def on_invocation_end(self, info: InvocationEndInfo) -> None: + record("end") + assert branch_done.is_set() + assert any(x[0] == "checkpoint-end" for x in events) + assert info.status is ( + InvocationStatus.FAILED + if result_kind in ("bad-json", "large-error") + else InvocationStatus.SUCCEEDED + ) -@pytest.mark.parametrize("value", [None, "not a method", 42]) -def test_noncallable_optional_handler_scope_is_ignored(value: Any) -> None: - from datetime import UTC, datetime - from types import SimpleNamespace - from typing import cast - from aws_durable_execution_sdk_python.plugin import PluginExecutor + def body(_event: Any, ctx: DurableContext) -> Any: + assert checkpoint_started.wait(5) + record("body") + pool = ThreadPoolExecutor(max_workers=1) + ctx.state.register_branch_pool(pool) + + def late_branch() -> None: + assert closing.wait(5) + assert not ctx.state._checkpointing_stopped.is_set() + ctx.state.create_checkpoint( + OperationUpdate.create_execution_succeed(payload='"branch"'), + is_sync=True, + ) + record("branch-done") + branch_done.set() - plugin = cast( - DurableInstrumentationPlugin, - SimpleNamespace(on_invocation_start=lambda info: None, handler_context=value), + pool.submit(late_branch) + try: + if result_kind == "large-error": + raise ValueError("x" * 256) + return result + finally: + record("finally") + + handler = durable_execution(body, plugins=[Plugin()]) + event, context = invocation() + output = handler(event, context) + record("caller") + names = [x[0] for x in events] + assert ( + names.index("start") + < names.index("checkpoint-start") + < names.index("body") + < names.index("finally") + < names.index("serialize") + < names.index("close") + ) + assert ( + names.index("branch-done") + < names.index("checkpoint-stop") + < names.index("checkpoint-end") + < names.index("end") + < names.index("caller") + ) + worker_events = { + tid + for name, tid in events + if name + in {"start", "body", "finally", "serialize", "close", "checkpoint-stop", "end"} + } + assert len(worker_events) == 1 and threading.get_ident() not in worker_events + assert output["Status"] == ( + "FAILED" if result_kind in ("bad-json", "large-error") else "SUCCEEDED" ) - executor = PluginExecutor([plugin]) - with executor.run(): - executor.on_invocation_start("noncallable", True, datetime.now(UTC), None) - assert executor.run_handler(lambda: "unchanged") == "unchanged" -def test_handler_scope_lookup_failure_discards_partial_bindings( +@pytest.mark.parametrize("shape", ["method", "property", "dynamic"]) +def test_removed_handler_scope_api_is_not_inspected( + monkeypatch: pytest.MonkeyPatch, shape: str +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + calls = [] + + def legacy(*_args: Any) -> Any: + calls.append("unexpected") + raise AssertionError("Removed API called") + + choices: dict[str, dict[str, Any]] = { + "method": {"handler_context": legacy}, + "property": {"handler_context": property(legacy)}, + "dynamic": { + "__getattr__": lambda self, name: legacy() + if name == "handler_context" + else (_ for _ in ()).throw(AttributeError(name)) + }, + } + plugin = type("LegacyPlugin", (DurableInstrumentationPlugin,), choices[shape])() + handler = durable_execution(lambda _e, _c: "ok", plugins=[plugin]) + event, context = invocation() + assert handler(event, context)["Status"] == "SUCCEEDED" + assert calls == [] + + +@pytest.mark.parametrize( + "checkpoint_path", ["step-start", "large-result", "large-error"] +) +@pytest.mark.parametrize("retryable", [False, True]) +def test_checkpoint_failure_reports_prepared_outcome_after_worker_cleanup( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, + checkpoint_path: str, + retryable: bool, ) -> None: - from datetime import UTC, datetime - from aws_durable_execution_sdk_python.plugin import PluginExecutor - - marker = contextvars.ContextVar("lookup-failure", default="outside") - - class BrokenPlugin(DurableInstrumentationPlugin): - @property - def handler_context(self) -> Any: - marker.set("partial") - raise ValueError("scope lookup failed") - - executor = PluginExecutor([BrokenPlugin()]) - with executor.run(): - executor.on_invocation_start("lookup", True, datetime.now(UTC), None) - assert executor.run_handler(marker.get) == "outside" - assert marker.get() == "outside" - assert "scope lookup failed" in caplog.text + import aws_durable_execution_sdk_python.execution as execution + + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + if checkpoint_path != "step-start": + monkeypatch.setattr(execution, "LAMBDA_RESPONSE_SIZE_LIMIT", 128) + marker = contextvars.ContextVar("checkpoint-worker", default="host") + failure = CheckpointError( + "actual checkpoint failure", + error_category=( + CheckpointErrorCategory.INVOCATION + if retryable + else CheckpointErrorCategory.EXECUTION + ), + ) + events: list[tuple[str, int]] = [] + ends: list[tuple[InvocationEndInfo, str]] = [] + real_checkpoint = ExecutionState.checkpoint_batches_forever + real_close = ExecutionState.close + + def record(name: str) -> None: + events.append((name, threading.get_ident())) + + def checkpoint(state: ExecutionState) -> None: + record("checkpoint-start") + try: + real_checkpoint(state) + finally: + record("checkpoint-end") + + def close(state: ExecutionState) -> None: + record("close") + assert marker.get() == "plugin" + real_close(state) + record("closed") + + def service_checkpoint(*_args: Any, **_kwargs: Any) -> Any: + record("service-failure") + raise failure + + monkeypatch.setattr(ExecutionState, "checkpoint_batches_forever", checkpoint) + monkeypatch.setattr(ExecutionState, "close", close) + + class Plugin(DurableInstrumentationPlugin): + token: contextvars.Token[str] | None = None + + def on_invocation_start(self, _info: InvocationStartInfo) -> None: + record("start") + self.token = marker.set("plugin") + + def on_invocation_end(self, info: InvocationEndInfo) -> None: + record("end") + ends.append((info, marker.get())) + assert self.token is not None + marker.reset(self.token) + self.token = None + + def body(_event: Any, ctx: DurableContext) -> str: + record("body") + try: + if checkpoint_path == "step-start": + ctx.step(lambda _step: "ok", name="failed-checkpoint") + if checkpoint_path == "large-error": + raise ValueError("x" * 256) + return "x" * 256 + finally: + record("finally") + + plugin = Plugin() + handler = durable_execution(body, plugins=[plugin]) + client = Mock(spec=DurableServiceClient) + client.checkpoint.side_effect = service_checkpoint + event, context = invocation(client) + for _ in range(2): + events.clear() + ends.clear() + if retryable: + with pytest.raises(CheckpointError) as caught: + handler(event, context) + assert caught.value is failure + else: + output = handler(event, context) + assert output["Status"] == "FAILED" + assert output["Error"]["ErrorMessage"] == str(failure) + assert output["Error"]["ErrorType"].endswith(".CheckpointError") + record("caller") + assert marker.get() == "host" and plugin.token is None + assert len(ends) == 1 + info, active = ends[0] + assert active == "plugin" + assert info.status is ( + InvocationStatus.RETRY if retryable else InvocationStatus.FAILED + ) + assert info.error is not None and info.error.message == str(failure) + assert info.error.type is not None and info.error.type.endswith( + ".CheckpointError" + ) + names = [name for name, _tid in events] + assert ( + names.index("start") + < names.index("checkpoint-start") + < names.index("service-failure") + ) + assert ( + names.index("body") + < names.index("finally") + < names.index("close") + < names.index("closed") + < names.index("end") + < names.index("caller") + ) + assert ( + names.index("service-failure") + < names.index("checkpoint-end") + < names.index("end") + ) + workers = { + tid + for name, tid in events + if name in {"start", "body", "finally", "close", "closed", "end"} + } + assert len(workers) == 1 and threading.get_ident() not in workers + assert not any( + record.exc_info and record.name == "aws_durable_execution_sdk_python.plugin" + for record in caplog.records + ) diff --git a/pyproject.toml b/pyproject.toml index bca7cb9e2..80eb21e0c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -131,66 +131,48 @@ dependencies = [ [tool.hatch.envs.dev-examples.scripts] test = "pytest packages/aws-durable-execution-sdk-python-examples/test {args}" -[tool.hatch.envs.test-pypi-otel] -# Test new exclusivity capability against an installed capable core. -# Legacy valid-registration combinations are verified separately against 2.0.x. -# Override inherited workspace membership: core must come from an installed distribution. -workspace.members = ["packages/aws-durable-execution-sdk-python-otel"] -dependencies = [ - "aws-durable-execution-sdk-python-testing>=1.2.1", - "aws-durable-execution-sdk-python>=2.1.0", - "opentelemetry-sdk>=1.20.0", - "opentelemetry-propagator-aws-xray", - "pytest", - "pytest-cov", - "coverage[toml]", -] -pre-install-commands = [ - "pip install -e packages/aws-durable-execution-sdk-python-otel", -] - -[tool.hatch.envs.test-pypi-otel.scripts] -test = "pytest packages/aws-durable-execution-sdk-python-otel/tests {args}" - -[tool.hatch.envs.test-pypi-otel-legacy] -template = "test-pypi-otel" -workspace.members = ["packages/aws-durable-execution-sdk-python-otel"] +# Install immutable artifacts built from this checkout, including before the +# first compatible core is published. No editable workspace packages in this lane. +[tool.hatch.envs.test-wheel-otel] +detached = true +workspace.members = [] dependencies = [ - "aws-durable-execution-sdk-python>=2.0.0,<2.1.0", "aws-durable-execution-sdk-python-testing>=1.2.1", "opentelemetry-sdk>=1.20.0", "opentelemetry-propagator-aws-xray", "pytest", - "pytest-cov", - "coverage[toml]", + "packaging", ] -[tool.hatch.envs.test-pypi-otel-legacy.scripts] +[tool.hatch.envs.test-wheel-otel.scripts] test = [ - "python -c 'from pathlib import Path; from importlib.metadata import version; from packaging.version import Version; import aws_durable_execution_sdk_python.execution as core; assert Version(version(\"aws-durable-execution-sdk-python\")).release[:2] == (2, 0); assert \"site-packages\" in Path(core.__file__).resolve().parts; print(version(\"aws-durable-execution-sdk-python\"), core.__file__)' ", - "pytest packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_view_registration.py -k 'one_view_and_unrelated_plugin_suspend_resume or no_otel_plugin_remains_valid or execution_constructor_retains_ambient_log_correlation' {args}" + "python .github/scripts/install_otel_test_wheels.py", + "pytest packages/aws-durable-execution-sdk-python-otel/tests {args}", + "pytest .github/tests/otel_lifecycle_compatibility_test.py {args}", ] -# Keep every pre-existing OTel test on the minimum core; the newly added view -# registration file has capability-only cases plus the legacy-valid cases below. -[tool.hatch.envs.test-pypi-otel-minimum] -template = "test-pypi-otel" -workspace.members = ["packages/aws-durable-execution-sdk-python-otel"] +# A core-only upgrade retains the released plugin's more limited Invocation +# fallback. Validate the real published plugin separately from the new pair. +[tool.hatch.envs.test-wheel-otel-legacy] +template = "test-wheel-otel" +detached = true +workspace.members = [] dependencies = [ - "aws-durable-execution-sdk-python==2.0.0", + "aws-durable-execution-sdk-python-otel==1.0.0", "aws-durable-execution-sdk-python-testing>=1.2.1", "opentelemetry-sdk>=1.20.0", "opentelemetry-propagator-aws-xray", "pytest", - "pytest-cov", - "coverage[toml]", + "packaging", ] -[tool.hatch.envs.test-pypi-otel-minimum.scripts] +[tool.hatch.envs.test-wheel-otel-legacy.env-vars] +OTEL_COMPAT_LEGACY = "1" + +[tool.hatch.envs.test-wheel-otel-legacy.scripts] test = [ - "python -c 'from pathlib import Path; from importlib.metadata import version; from packaging.version import Version; import aws_durable_execution_sdk_python.execution as core; assert version(\"aws-durable-execution-sdk-python\") == \"2.0.0\"; assert \"site-packages\" in Path(core.__file__).resolve().parts; print(version(\"aws-durable-execution-sdk-python\"), core.__file__)' ", - "pytest packages/aws-durable-execution-sdk-python-otel/tests --ignore=packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_view_registration.py {args}", - "pytest packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_view_registration.py -k 'one_view_and_unrelated_plugin_suspend_resume or no_otel_plugin_remains_valid or execution_constructor_retains_ambient_log_correlation' {args}" + "python .github/scripts/install_otel_test_wheels.py --legacy-plugin", + "pytest .github/tests/otel_lifecycle_compatibility_test.py {args}", ] [tool.hatch.envs.test-pypi-examples] From 29847f372255c5ac4ed0f43d3b67dcbd59b37173 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 9 Oct 2026 00:47:16 +0000 Subject: [PATCH 23/24] fix: propagate plugin context into concurrent branches --- .../README.md | 6 +- .../tests/e2e/test_branch_context_int.py | 179 ++++++++++++++++++ .../README.md | 5 + .../concurrency/executor.py | 27 ++- .../tests/e2e/branch_worker_context_test.py | 129 +++++++++++++ .../tests/operation/map_test.py | 29 +++ .../tests/operation/parallel_test.py | 29 +++ 7 files changed, 396 insertions(+), 8 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_branch_context_int.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/e2e/branch_worker_context_test.py diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index f42b5986c..2a095bf9c 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -200,7 +200,11 @@ Start hooks run in registration order. A later successful plugin that deliberate sets or clears the active span wins; OTel does not apply a second correction pass. Place OTel after a span-replacing plugin when OTel's view-specific context is desired. Baggage-only plugins that extend the current context can appear on either -side. End hooks also retain registration order and reset tokens in their owning +side. The SDK also copies the coordinator's current bindings for each `map` or +`parallel` branch admission/resume when plugins are registered. Baggage and +other successful bindings reach these SDK-managed callbacks, while a branch's +changes cannot leak into siblings, the coordinator or a reused worker. +End hooks also retain registration order and reset tokens in their owning Context; they do not promise a reverse-stack observation of other plugins' spans. Existing registration, factory lifetime and checkpoint formats are unchanged. diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_branch_context_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_branch_context_int.py new file mode 100644 index 000000000..c2ae779af --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_branch_context_int.py @@ -0,0 +1,179 @@ +"""Public SDK branch workers inherit instrumented bindings without sharing them.""" + +from __future__ import annotations + +import contextvars +import json +import threading +from typing import Any + +import pytest +from aws_durable_execution_sdk_python import DurableContext, durable_execution +from aws_durable_execution_sdk_python.config import Duration, MapConfig, ParallelConfig +from aws_durable_execution_sdk_python.concurrency.models import BatchResult +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationStatus, +) +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry import baggage, context, trace +from opentelemetry.sdk.trace import TracerProvider + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +@pytest.mark.parametrize( + ("view", "failed_start"), + [ + (None, False), + (InvocationOtelPlugin, False), + (ExecutionOtelPlugin, False), + (InvocationOtelPlugin, True), + (ExecutionOtelPlugin, True), + ], +) +@pytest.mark.parametrize("kind", ["map", "parallel"]) +@pytest.mark.parametrize("concurrency", [1, 2]) +@pytest.mark.parametrize("resume", [False, True]) +def test_public_branch_context_propagation_and_isolation( + monkeypatch: pytest.MonkeyPatch, + view: type[InvocationOtelPlugin] | type[ExecutionOtelPlugin] | None, + failed_start: bool, + caplog: pytest.LogCaptureFixture, + kind: str, + concurrency: int, + resume: bool, +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + marker = contextvars.ContextVar("branch-invocation-marker", default="unset") + observations: list[tuple[int, str, Any, bool]] = [] + bodies: list[int] = [] + visits = {0: 0, 1: 0} + statuses: list[InvocationStatus] = [] + starts: list[int] = [] + completed_ends: list[int] = [] + lock = threading.Lock() + barrier = threading.Barrier(2) if concurrency == 2 else None + provider = TracerProvider() + host = context.get_current() + + class Bind(DurableInstrumentationPlugin): + def on_invocation_start(self, info: Any) -> None: + starts.append(threading.get_ident()) + self.token = marker.set("plugin") + self.bag = context.attach(baggage.set_baggage("tenant", "present")) + + def on_invocation_end(self, info: Any) -> None: + context.detach(self.bag) + marker.reset(self.token) + # Record completion only after both tokens were reset in their owner. + completed_ends.append(threading.get_ident()) + statuses.append(info.status) + + class Broken(DurableInstrumentationPlugin): + def on_invocation_start(self, _info: Any) -> None: + marker.set("partial") + context.attach(baggage.set_baggage("tenant", "partial")) + raise ValueError("expected failed Start") + + plugins: list[DurableInstrumentationPlugin] = [] + otel_plugin: InvocationOtelPlugin | ExecutionOtelPlugin | None = None + if view is not None: + plugins.append(Bind()) + if failed_start: + plugins.append(Broken()) + otel_plugin = view( + OtelPluginConfig( + tracer_provider=provider, + enrich_logger=False, + context_extractor=lambda _: None, + ) + ) + plugins.append(otel_plugin) + + def branch(child: DurableContext, index: int) -> str: + with lock: + visits[index] += 1 + first_visit = visits[index] == 1 + observations.append( + ( + index, + marker.get(), + baggage.get_baggage("tenant"), + trace.get_current_span().get_span_context().is_valid, + ) + ) + if view is not None: + # Deliberately leave a binding behind: another logical branch or + # resume on this pool thread must still start with its parent's copy. + marker.set(f"branch-{index}") + if barrier is not None and first_visit: + barrier.wait(timeout=10) + + def step(_step: Any) -> str: + with lock: + bodies.append(index) + return f"saved-{index}" + + saved = child.step(step, name="save") + return saved + + def handler(_event: Any, durable: DurableContext) -> list[str]: + expected = "unset" if view is None else "plugin" + assert marker.get() == expected + token = marker.set("handler" if view is None else "plugin") + try: + result: BatchResult[str] = ( + durable.map( + [0, 1], + lambda child, item, index, items: branch(child, index), + name="mapped", + config=MapConfig(max_concurrency=concurrency), + ) + if kind == "map" + else durable.parallel( + [lambda child: branch(child, 0), lambda child: branch(child, 1)], + name="parallel", + config=ParallelConfig(max_concurrency=concurrency), + ) + ) + assert marker.get() == ("handler" if view is None else "plugin") + if resume: + durable.wait(Duration.from_seconds(1), name="resume") + return result.get_results() + finally: + marker.reset(token) + + wrapped = durable_execution(handler, plugins=plugins) + try: + with DurableFunctionTestRunner(handler=wrapped) as runner: + result = runner.run(input="{}", timeout=30) + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == ["saved-0", "saved-1"] + assert sorted(bodies) == [0, 1] + assert {item[0] for item in observations} == {0, 1} + expected = "unset" if view is None else "plugin" + assert all( + item[1:] + == (expected, None if view is None else "present", view is not None) + for item in observations + ) + # The final wait replays the completed batch without rerunning branches. + assert visits == {0: 1, 1: 1} + if view is not None and resume: + assert InvocationStatus.PENDING in statuses + assert starts == completed_ends + if view is not None: + assert starts and statuses[-1] is InvocationStatus.SUCCEEDED + assert otel_plugin is not None and otel_plugin._context_tokens == {} + assert not any( + r.exc_info and r.name == "opentelemetry.context" for r in caplog.records + ) + assert marker.get() == "unset" + assert context.get_current() == host + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python/README.md b/packages/aws-durable-execution-sdk-python/README.md index 08a70e06a..4b2463e52 100644 --- a/packages/aws-durable-execution-sdk-python/README.md +++ b/packages/aws-durable-execution-sdk-python/README.md @@ -129,6 +129,11 @@ replace an earlier binding. If a Start hook raises, its new bindings are discard for subsequent work; its End still runs in the original Context so its tokens can be reset. End hooks retain forward registration order, not reverse stack order, so they must not rely on observing a stack-like unwind of other plugins' contexts. +Each SDK-managed `map` or `parallel` branch admission, including an in-process +resume, receives a fresh copy of the coordinator's bindings when plugins are +registered. Branch changes remain local even when pool threads are reused. +User-created threads retain normal Python context-variable behavior; without +registered plugins, SDK branch submission keeps its existing behavior. The worker's invocation context is discarded on return, including after plugin cleanup failures, leaving the host's bindings unchanged. This isolates bindings, not mutations to shared objects or external side effects. Without plugins, the diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py index 2a39c4eff..bafa43b19 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py @@ -2,6 +2,7 @@ from __future__ import annotations +import contextvars import heapq import logging import queue @@ -291,13 +292,25 @@ def execute( def submit(branch: Branch[CallableType, ResultType]) -> None: branch.start() - pool.submit( - self._branch_worker, - execution_state, - executor_context, - events, - branch.executable, - ) + if execution_state._plugin_executor._plugins: # noqa: SLF001 + # Every admission/resume gets its own Context. Branch bindings + # cannot leak into siblings, the coordinator, or reused workers. + pool.submit( + contextvars.copy_context().run, + self._branch_worker, + execution_state, + executor_context, + events, + branch.executable, + ) + else: + pool.submit( + self._branch_worker, + execution_state, + executor_context, + events, + branch.executable, + ) try: # Only rebuild the items snapshot after a terminal event changes diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/branch_worker_context_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/branch_worker_context_test.py new file mode 100644 index 000000000..a9104a745 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/branch_worker_context_test.py @@ -0,0 +1,129 @@ +"""Registered Start bindings reach independently resumed public branch workers.""" + +from __future__ import annotations + +import contextvars +import json +import threading +from typing import Any + +import pytest +from aws_durable_execution_sdk_python import DurableContext, durable_execution +from aws_durable_execution_sdk_python.config import Duration, MapConfig, ParallelConfig +from aws_durable_execution_sdk_python.concurrency.models import BatchResult +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationEndInfo, + InvocationStartInfo, + InvocationStatus, +) +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner + + +@pytest.mark.parametrize("kind", ["map", "parallel"]) +@pytest.mark.parametrize("concurrency", [1, 2]) +@pytest.mark.parametrize("failed_start", [False, True]) +def test_partial_branch_resume_uses_fresh_successful_start_bindings( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + kind: str, + concurrency: int, + failed_start: bool, +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + marker = contextvars.ContextVar("resumed-branch", default="host") + partial = contextvars.ContextVar[str]("partial-start") + observations: list[tuple[int, str]] = [] + visits = {0: 0, 1: 0} + bodies: list[int] = [] + statuses: list[InvocationStatus] = [] + lock = threading.Lock() + barrier = threading.Barrier(2) if concurrency == 2 else None + + class Bind(DurableInstrumentationPlugin): + token: contextvars.Token[str] | None = None + + def on_invocation_start(self, _info: InvocationStartInfo) -> None: + self.token = marker.set("plugin") + + def on_invocation_end(self, info: InvocationEndInfo) -> None: + statuses.append(info.status) + assert self.token is not None + marker.reset(self.token) + self.token = None + + class Broken(DurableInstrumentationPlugin): + def on_invocation_start(self, _info: InvocationStartInfo) -> None: + marker.set("failed") + partial.set("failed") + raise ValueError("expected failed Start") + + plugin = Bind() + plugins: list[DurableInstrumentationPlugin] = [plugin] + if failed_start: + plugins.append(Broken()) + + def branch(child: DurableContext, index: int) -> str: + with lock: + visits[index] += 1 + first_visit = visits[index] == 1 + observations.append((index, marker.get())) + with pytest.raises(LookupError): + partial.get() + marker.set(f"branch-{index}") + if barrier is not None and first_visit: + barrier.wait(timeout=10) + + def step(_step: Any) -> str: + with lock: + bodies.append(index) + return f"saved-{index}" + + saved = child.step(step, name="save") + child.wait(Duration.from_seconds(index + 1), name="branch-wait") + return saved + + def handler(_event: Any, durable: DurableContext) -> list[str]: + assert marker.get() == "plugin" + result: BatchResult[str] + if kind == "map": + result = durable.map( + [0, 1], + lambda child, item, index, items: branch(child, index), + name="mapped", + config=MapConfig(max_concurrency=concurrency), + ) + else: + result = durable.parallel( + [lambda child: branch(child, 0), lambda child: branch(child, 1)], + name="parallel", + config=ParallelConfig(max_concurrency=concurrency), + ) + assert marker.get() == "plugin" + durable.wait(Duration.from_seconds(1), name="after-batch") + return result.get_results() + + wrapped = durable_execution(handler, plugins=plugins) + with DurableFunctionTestRunner(handler=wrapped) as runner: + result = runner.run(input="{}", timeout=30) + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == ["saved-0", "saved-1"] + assert sorted(bodies) == [0, 1] + assert all(count >= 2 for count in visits.values()) + assert all(binding == "plugin" for _, binding in observations) + assert InvocationStatus.PENDING in statuses + assert statuses[-1] is InvocationStatus.SUCCEEDED + assert marker.get() == "host" and plugin.token is None + with pytest.raises(LookupError): + partial.get() + errors = [ + record + for record in caplog.records + if record.exc_info and record.name == "aws_durable_execution_sdk_python.plugin" + ] + assert len(errors) == (len(statuses) if failed_start else 0) + assert all( + record.exc_info is not None + and str(record.exc_info[1]) == "expected failed Start" + for record in errors + ) diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/map_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/map_test.py index 2670e6747..926b1dd0b 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/map_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/map_test.py @@ -7,6 +7,7 @@ import pytest # Mock the executor.execute method +from aws_durable_execution_sdk_python.plugin import PluginExecutor from aws_durable_execution_sdk_python.concurrency.models import ( BatchItem, BatchItemStatus, @@ -184,6 +185,8 @@ def mock_run_in_child_context(func, name, config): # Create a minimal ExecutionState mock class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -222,6 +225,8 @@ def mock_run_in_child_context(func, name, config): return func("mock_context") class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -354,6 +359,8 @@ def callable_func(ctx, item, idx, items): ) as mock_execute: class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -407,6 +414,8 @@ def callable_func(ctx, item, idx, items): executor_context.create_child_context = lambda *args, **kwargs: Mock() class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -461,6 +470,8 @@ def callable_func(ctx, item, idx, items): executor_context.create_child_context = lambda *args, **kwargs: child_context class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -510,6 +521,8 @@ def mock_summary_generator(result): executor_context.create_child_context = Mock(return_value=_child_ctx) class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -550,6 +563,8 @@ def callable_func(ctx, item, idx, items): executor_context.create_child_context = Mock(return_value=Mock()) # SLF001 class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -589,6 +604,8 @@ def func(ctx, item, index, array): config = MapConfig(summary_generator=None) class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -634,6 +651,8 @@ def callable_func(ctx, item, idx, items): # Mock execution state that indicates operation already succeeded class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -703,6 +722,8 @@ def callable_func(ctx, item, idx, items): # Mock execution state that indicates operation succeeded but children need replay class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -802,6 +823,8 @@ def test_func(ctx, item, idx, items): execution_count = 0 class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -1032,6 +1055,8 @@ def func(ctx, item, idx, items): return {"item": item.upper(), "index": idx} class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -1481,6 +1506,8 @@ def test_map_handler_defaults_summary_generator_for_user_config(): """ class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -1556,6 +1583,8 @@ def predicate(s: CompletionStatus) -> CompletionDecision: ) as mock_execute: class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/parallel_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/parallel_test.py index 12231cfab..87d4d5a89 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/parallel_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/parallel_test.py @@ -8,6 +8,7 @@ import pytest +from aws_durable_execution_sdk_python.plugin import PluginExecutor from aws_durable_execution_sdk_python.concurrency.executor import ConcurrentExecutor from aws_durable_execution_sdk_python.identifier import OperationIdNamespace @@ -211,6 +212,8 @@ def func2(ctx): config = ParallelConfig(max_concurrency=2) class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -255,6 +258,8 @@ def func1(ctx): callables = [func1] class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -299,6 +304,8 @@ def func1(ctx): config = ParallelConfig(max_concurrency=5) class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -351,6 +358,8 @@ def func1(ctx): callables = [func1] class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -466,6 +475,8 @@ def func1(ctx): callables = [func1] class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -513,6 +524,8 @@ def mock_summary_generator(result): config = ParallelConfig(summary_generator=mock_summary_generator) class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -556,6 +569,8 @@ def func2(ctx): callables = [func1, func2] class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -606,6 +621,8 @@ def func3(ctx): config = ParallelConfig(summary_generator=None) class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -654,6 +671,8 @@ def func2(ctx): # Mock execution state that indicates operation already succeeded class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -723,6 +742,8 @@ def func1(ctx): # Mock execution state that indicates operation succeeded but children need replay class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -824,6 +845,8 @@ def task2(ctx): execution_count = 0 class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -1052,6 +1075,8 @@ def func3(ctx): callables = [func1, func2, func3] class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -1474,6 +1499,8 @@ def test_parallel_handler_defaults_summary_generator_for_user_config(): """A user config without a summary generator gets the default (JS parity).""" class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -1551,6 +1578,8 @@ def predicate(s: CompletionStatus) -> CompletionDecision: ) as mock_execute: class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass From 24563253734554cdd325aaf9035682d5066b2a3b Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 9 Oct 2026 01:31:15 +0000 Subject: [PATCH 24/24] test: cover instrumented branch context handoff --- .../tests/concurrency_test.py | 47 ++++++++++++++++++- 1 file changed, 46 insertions(+), 1 deletion(-) diff --git a/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py b/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py index a93135663..c30980c73 100644 --- a/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py @@ -1,5 +1,6 @@ """Tests for the concurrency module.""" +import contextvars import hashlib import json import queue @@ -77,7 +78,10 @@ from aws_durable_execution_sdk_python.operation.parallel import ( ParallelExecutor, ) -from aws_durable_execution_sdk_python.plugin import PluginExecutor +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + PluginExecutor, +) from aws_durable_execution_sdk_python.state import ( CheckpointedResult, ExecutionState, @@ -5275,3 +5279,44 @@ def predicate(s: CompletionStatus) -> CompletionDecision: # endregion Custom completion predicate (should_complete) integration tests + + +def test_instrumented_executor_isolates_bindings_on_a_reused_worker() -> None: + """Exercise actual admission/worker execution independently of durable I/O.""" + binding = contextvars.ContextVar("executor-unit-binding", default="empty") + seen: list[tuple[str, int]] = [] + + class RecordingExecutor(ConcurrentExecutor[Callable[[], str], str]): + def _execute_item_in_child_context( + self, + executor_context: DurableContext, + executable: Executable[Callable[[], str]], + ) -> str: + seen.append((binding.get(), threading.get_ident())) + binding.set(f"branch-{executable.index}") + return executable.func() + + executor = RecordingExecutor( + executables=[Executable(index, lambda: "ok") for index in range(2)], + max_concurrency=1, + completion_config=CompletionConfig(min_successful=2), + sub_type_top=OperationSubType.PARALLEL, + sub_type_iteration=OperationSubType.PARALLEL_BRANCH, + name_prefix="branch-", + serdes=None, + operation_id_namespace=_StubNamespace(), + ) + state = Mock(spec=ExecutionState) + state._plugin_executor = PluginExecutor([DurableInstrumentationPlugin()]) + token = binding.set("coordinator") + try: + result = executor.execute(state, Mock(spec=DurableContext)) + assert result.get_results() == ["ok", "ok"] + assert binding.get() == "coordinator" + assert [value for value, _thread in seen] == ["coordinator", "coordinator"] + assert seen[0][1] == seen[1][1] != threading.get_ident() + finally: + for call in state.register_branch_pool.call_args_list: + call.args[0].shutdown(wait=True) + binding.reset(token) + assert binding.get() == "empty"