From 7d89a4e513041e51b7e242129fed22d5a6a0c08f Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 2 Oct 2026 23:22:03 +0000 Subject: [PATCH 01/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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