From 7d89a4e513041e51b7e242129fed22d5a6a0c08f Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 2 Oct 2026 23:22:03 +0000 Subject: [PATCH 01/54] 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/54] 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 e9cda8c36e317c791636fce0f9c69bf8889d4ed6 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Sat, 3 Oct 2026 00:13:42 +0000 Subject: [PATCH 03/54] ci: validate the stacked conformance branch --- .github/workflows/ci.yml | 3 +- .../opentelemetry-conformance-tests.yml | 9 +- .../README.md | 25 ++- .../src/Makefile | 6 + .../src/otel_21_completed_step_replay.py | 36 ++++ .../src/otel_22_user_function_context.py | 113 ++++++++++++ .../src/otel_23_callback_function_context.py | 136 ++++++++++++++ .../src/otel_24_invocation_retry_status.py | 33 ++++ .../template.yaml | 124 +++++++++++++ .../tests/test_callback_function_context.py | 122 +++++++++++++ .../tests/test_completed_step_replay.py | 100 ++++++++++ .../tests/test_invocation_retry_status.py | 110 +++++++++++ .../tests/test_otel_examples.py | 12 ++ .../README.md | 9 + .../execution_plugin.py | 7 +- .../invocation_plugin.py | 8 +- .../tests/test_execution_plugin.py | 4 +- .../tests/test_invocation_plugin.py | 6 +- .../tests/test_status_mapping.py | 172 ++++++++++++++++++ 19 files changed, 1023 insertions(+), 12 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_21_completed_step_replay.py create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_22_user_function_context.py create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_23_callback_function_context.py create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_24_invocation_retry_status.py create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_function_context.py create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_completed_step_replay.py create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_invocation_retry_status.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/test_status_mapping.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 63c4b44c6..b2d66fd44 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -10,7 +10,8 @@ on: branches: [ main ] pull_request: - branches: [ main ] + # Run checks while the conformance PR is stacked on the focused context fix. + branches: ["main", "fix/otel-handler-context-428"] jobs: lint-commits: diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index b31fff281..3a6387dae 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -8,7 +8,8 @@ name: OpenTelemetry Conformance Tests on: pull_request: - branches: ["main"] + # Run checks while the conformance PR is stacked on the focused context fix. + branches: ["main", "fix/otel-handler-context-428"] paths: - "packages/aws-durable-execution-sdk-python/**" - "packages/aws-durable-execution-sdk-python-otel/**" @@ -45,7 +46,7 @@ on: conformance_test_ref: description: Conformance test commit SHA or branch name required: true - default: main + default: 922dfafa575f771a6d5b117d9f1e631d149e01e4 type: string permissions: {} @@ -59,13 +60,13 @@ jobs: actions: write contents: read id-token: write - uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@a628f5589bbf067a441696c792ab3023c0d0899b + uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@922dfafa575f771a6d5b117d9f1e631d149e01e4 with: language: python resource_prefix: p sdk_repository: aws/aws-durable-execution-sdk-python sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }} - conformance_test_ref: ${{ inputs.conformance_test_ref || 'main' }} + conformance_test_ref: ${{ inputs.conformance_test_ref || '922dfafa575f771a6d5b117d9f1e631d149e01e4' }} # Check the SDK out so the handlers and templates below are on disk. The handlers # themselves are installed from sdk_ref by src/requirements.txt during the SAM build. checkout_sdk: true diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md index 5ff978c88..b2871c83d 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md @@ -33,7 +33,7 @@ template-long-running.yaml # otel-long-running suite tests/ # contract tests for the templates and handlers ``` -The 20 invocation and 20 execution requirements reuse the same scenario +The invocation and execution requirements reuse the same scenario handlers; the view is selected per function through the `OTEL_PLUGIN_MODE` environment variable, which `common.otel_plugin()` reads to pick `InvocationOtelPlugin` or `ExecutionOtelPlugin`. `template.yaml` deploys only the @@ -63,6 +63,10 @@ view named by its `OtelSuite` parameter. | `otel-invocation-18` | `otel_18_chained_invoke_failure.handler` | Verifies failed chained-invoke telemetry. | | `otel-invocation-19` | `otel_19_execution_failure.handler` | Verifies telemetry for a direct handler failure. | | `otel-invocation-20` | `otel_20_virtual_context.handler` | Verifies a virtual child-context span without context checkpoints. | +| `otel-invocation-21` | `otel_21_completed_step_replay.handler` | Replays a completed step after a successful wait/resume; each step body runs once. | +| `otel-invocation-22` | `otel_22_user_function_context.handler` | Creates ordinary user spans under the active handler, attempt, child, branch, and iteration contexts. | +| `otel-invocation-23` | `otel_23_callback_function_context.handler` | Verifies retry/check attempts, callback submitter, wrapped retry helper, and virtual-child callback parents. | +| `otel-invocation-24` | `otel_24_invocation_retry_status.handler` | Raises a retryable invocation error after a completed step, then resumes with its saved result. | | `otel-execution-1` | `otel_1_success.handler` | Verifies the execution-view workflow, step, and attempt hierarchy. | | `otel-execution-2` | `otel_2_wait_resume.handler` | Verifies the execution view across a resumed invocation. | | `otel-execution-3` | `otel_3_retry.handler` | Verifies the execution view across retry attempts. | @@ -83,6 +87,10 @@ view named by its `OtelSuite` parameter. | `otel-execution-18` | `otel_18_chained_invoke_failure.handler` | Verifies source and target failed workflow roots. | | `otel-execution-19` | `otel_19_execution_failure.handler` | Verifies a failed invocation without a completed workflow. | | `otel-execution-20` | `otel_20_virtual_context.handler` | Verifies a virtual child-context span under the workflow root. | +| `otel-execution-21` | `otel_21_completed_step_replay.handler` | Verifies completed-operation spans are exported once across normal successful replay. | +| `otel-execution-22` | `otel_22_user_function_context.handler` | Observes active execution-view callback contexts without supplying or repairing parents. | +| `otel-execution-23` | `otel_23_callback_function_context.handler` | Verifies the same SDK-owned callback lifecycle parents in the execution view. | +| `otel-execution-24` | `otel_24_invocation_retry_status.handler` | Verifies invocation retry status independently of step retry and recovery re-exports. | | `otel-long-running-1` | `otel_long_running_1_wait.handler` | Verifies wait and resume telemetry across a long durable suspension. | | `otel-long-running-2` | `otel_long_running_2_retry.handler` | Verifies retry telemetry across a long durable backoff. | | `otel-long-running-3` | `otel_long_running_3_callback.handler` | Verifies callback telemetry when completion arrives after a long delay. | @@ -91,6 +99,21 @@ view named by its `OtelSuite` parameter. The runner discovers each mapping from `TestingMetadata.TestDescription` on the functions in the templates. +## Callback coverage boundary + +Cases 22 and 23 create normal user spans from the active context, without an +explicit parent or a copied execution ARN. They cover handler context, step and +condition attempts, child and branch bodies, callback submission, wrapped +`with_retry` body/strategy callbacks, and a virtual child. Case 23 places the +retry helper and virtual child after callback completion so ordinary successful +replay does not repeat their probes. + +This does not promise an operation/attempt parent for every arbitrary callback. +General retry/wait policies, serializers, summary generators, and item naming +have phase-specific caller scopes outside the wrapped user-function lifecycle; +these cases do not impose a new ownership policy on them. Instrumentation +extensions and user-created threads are outside this business-callback contract. + ## How a handler maps to a requirement ```yaml diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/Makefile b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/Makefile index 320577e7f..41e731b57 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/Makefile +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/Makefile @@ -5,6 +5,7 @@ build-Otel14MapFailure build-Otel15WaitInterrupted build-Otel16WaitForConditionFailure \ build-Otel17WaitForCallbackFailure build-Otel18ChainedInvokeFailure \ build-Otel18InvokeTarget build-Otel19ExecutionFailure build-Otel20VirtualContext \ + build-Otel21CompletedStepReplay build-Otel22UserFunctionContext build-Otel23CallbackFunctionContext build-Otel24InvocationRetryStatus \ build-OtelExecution1Success \ build-OtelExecution2WaitResume build-OtelExecution3Retry \ build-OtelExecution4TerminalFailure build-OtelExecution5ChildContext \ @@ -17,6 +18,8 @@ build-OtelExecution17WaitForCallbackFailure build-OtelExecution18ChainedInvokeFailure \ build-OtelExecution18InvokeTarget build-OtelExecution19ExecutionFailure \ build-OtelExecution20VirtualContext \ + build-OtelExecution21CompletedStepReplay build-OtelExecution22UserFunctionContext build-OtelExecution23CallbackFunctionContext \ + build-OtelExecution24InvocationRetryStatus \ build-OtelLongRunning1Wait \ build-OtelLongRunning2Retry build-OtelLongRunning3Callback \ build-OtelLongRunning4ChainedInvoke build-OtelLongRunning4InvokeTarget @@ -28,6 +31,7 @@ build-Otel11InvokeTarget build-Otel12ChildContextFailure build-Otel13ParallelFai build-Otel14MapFailure build-Otel15WaitInterrupted build-Otel16WaitForConditionFailure \ build-Otel17WaitForCallbackFailure build-Otel18ChainedInvokeFailure \ build-Otel18InvokeTarget build-Otel19ExecutionFailure build-Otel20VirtualContext \ +build-Otel21CompletedStepReplay build-Otel22UserFunctionContext build-Otel23CallbackFunctionContext build-Otel24InvocationRetryStatus \ build-OtelExecution1Success \ build-OtelExecution2WaitResume build-OtelExecution3Retry \ build-OtelExecution4TerminalFailure build-OtelExecution5ChildContext \ @@ -40,6 +44,8 @@ build-OtelExecution16WaitForConditionFailure \ build-OtelExecution17WaitForCallbackFailure build-OtelExecution18ChainedInvokeFailure \ build-OtelExecution18InvokeTarget build-OtelExecution19ExecutionFailure \ build-OtelExecution20VirtualContext \ +build-OtelExecution21CompletedStepReplay build-OtelExecution22UserFunctionContext build-OtelExecution23CallbackFunctionContext \ +build-OtelExecution24InvocationRetryStatus \ build-OtelLongRunning1Wait \ build-OtelLongRunning2Retry build-OtelLongRunning3Callback \ build-OtelLongRunning4ChainedInvoke build-OtelLongRunning4InvokeTarget: diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_21_completed_step_replay.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_21_completed_step_replay.py new file mode 100644 index 000000000..2ea9ef4a0 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_21_completed_step_replay.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Normal successful replay of a completed step for OTel case 21.""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python import ( + DurableContext, + StepContext, + durable_execution, + durable_step, +) +from aws_durable_execution_sdk_python.config import Duration +from common import otel_plugin, require_scenario + + +@durable_step +def before_wait(_step_context: StepContext) -> str: + return "before" + + +@durable_step +def after_wait(_step_context: StepContext) -> str: + return "after" + + +@durable_execution(plugins=[otel_plugin()]) +def handler(event: dict[str, Any], context: DurableContext) -> str: + require_scenario(event, "completed-step-replay") + before = context.step(before_wait(), name="otel-before-wait") + context.wait(Duration.from_seconds(1), name="otel-replay-wait") + after = context.step(after_wait(), name="otel-after-wait") + return f"{before}-{after}" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_22_user_function_context.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_22_user_function_context.py new file mode 100644 index 000000000..4f103ba4e --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_22_user_function_context.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Observe the real active SDK context inside public user callbacks.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +from aws_durable_execution_sdk_python import ( + DurableContext, + StepContext, + durable_execution, +) +from aws_durable_execution_sdk_python.config import ( + Duration, + MapConfig, + ParallelBranch, + ParallelConfig, +) +from common import otel_plugin, require_scenario +from opentelemetry import trace + + +def probe(label: str) -> None: + if not trace.get_current_span().get_span_context().is_valid: + raise RuntimeError(f"No active span for conformance.{label}") + span = trace.get_tracer("aws-durable-execution-conformance").start_span( + f"conformance.{label}", attributes={"conformance.callback": label} + ) + span.end() + + +def step_body(_context: StepContext) -> str: + probe("step") + return "step" + + +def child_step(_context: StepContext) -> str: + probe("child-step") + return "child-step" + + +def child_body(context: DurableContext) -> str: + probe("child") + context.step(child_step, name="otel-context-child-step") + probe("child-restored") + return "child" + + +def parallel_step_a(_context: StepContext) -> str: + probe("parallel-step-a") + return "a" + + +def parallel_step_b(_context: StepContext) -> str: + probe("parallel-step-b") + return "b" + + +def parallel_a(context: DurableContext) -> str: + probe("parallel-a") + return context.step(parallel_step_a, name="otel-context-branch-step-a") + + +def parallel_b(context: DurableContext) -> str: + probe("parallel-b") + return context.step(parallel_step_b, name="otel-context-branch-step-b") + + +def iteration_name(_item: int, index: int) -> str: + return ("otel-context-iteration-0", "otel-context-iteration-1")[index] + + +def mapper( + context: DurableContext, item: int, index: int, _items: Sequence[int] +) -> int: + probe(("map-0", "map-1")[index]) + + def map_step(_step_context: StepContext) -> int: + probe(("map-step-0", "map-step-1")[index]) + return item + + return context.step( + map_step, name=("otel-context-map-step-0", "otel-context-map-step-1")[index] + ) + + +@durable_execution(plugins=[otel_plugin()]) +def handler(event: dict[str, Any], context: DurableContext) -> str: + require_scenario(event, "user-function-context") + probe("handler") + context.step(step_body, name="otel-context-step") + context.run_in_child_context(child_body, name="otel-context-child") + context.parallel( + [ + ParallelBranch(parallel_a, name="otel-context-branch-a"), + ParallelBranch(parallel_b, name="otel-context-branch-b"), + ], + name="otel-context-parallel", + config=ParallelConfig(max_concurrency=2), + ) + context.map( + [0, 1], + mapper, + name="otel-context-map", + config=MapConfig(max_concurrency=2, item_namer=iteration_name), + ) + probe("handler-restored") + context.wait(Duration.from_seconds(1), name="otel-context-resume") + probe("handler-after-resume") + return "context-complete" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_23_callback_function_context.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_23_callback_function_context.py new file mode 100644 index 000000000..c7fae30ec --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_23_callback_function_context.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Observe callback contexts at SDK-owned user-function lifecycle boundaries.""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python import ( + DurableContext, + StepContext, + durable_execution, +) +from aws_durable_execution_sdk_python.config import ( + ChildConfig, + Duration, + JitterStrategy, + StepConfig, +) +from aws_durable_execution_sdk_python.exceptions import ChildContextError +from aws_durable_execution_sdk_python.retries import ( + RetryDecision, + RetryStrategyConfig, + WithRetryConfig, + create_retry_strategy, + with_retry, +) +from aws_durable_execution_sdk_python.types import ( + WaitForCallbackContext, + WaitForConditionCheckContext, +) +from aws_durable_execution_sdk_python.waits import ( + WaitForConditionConfig, + WaitForConditionDecision, +) +from common import otel_plugin, require_scenario +from opentelemetry import trace + + +HELPER_FAILURE = "intentional-helper-failure" + + +def probe(label: str) -> None: + if not trace.get_current_span().get_span_context().is_valid: + raise RuntimeError(f"No active span for conformance.{label}") + span = trace.get_tracer("aws-durable-execution-conformance").start_span( + f"conformance.{label}", attributes={"conformance.callback": label} + ) + span.end() + + +def retry_step(context: StepContext) -> str: + probe(f"retry-attempt-{context.attempt}") + if context.attempt == 1: + raise RuntimeError("intentional-step-retry") + return "retried" + + +def check_condition(state: int, _context: WaitForConditionCheckContext) -> int: + next_state = state + 1 + probe(f"condition-check-{next_state}") + return next_state + + +def wait_strategy(state: int, _attempt: int) -> WaitForConditionDecision: + if state >= 2: + return WaitForConditionDecision.stop_polling() + return WaitForConditionDecision.continue_waiting(Duration.from_seconds(1)) + + +def submit_callback(_callback_id: str, _context: WaitForCallbackContext) -> None: + probe("callback-submitter") + + +def helper_body(_context: DurableContext, _attempt: int) -> str: + probe("with-retry-body") + raise RuntimeError(HELPER_FAILURE) + + +def helper_retry_strategy(_error: Exception, _attempt: int) -> RetryDecision: + probe("with-retry-strategy") + return RetryDecision.no_retry() + + +def virtual_child(_context: DurableContext) -> str: + probe("virtual-child") + return "virtual" + + +@durable_execution(plugins=[otel_plugin()]) +def handler(event: dict[str, Any], context: DurableContext) -> str: + require_scenario(event, "callback-function-context") + context.step( + retry_step, + name="otel-context-retry-step", + config=StepConfig( + retry_strategy=create_retry_strategy( + RetryStrategyConfig( + max_attempts=2, + initial_delay=Duration.from_seconds(1), + max_delay=Duration.from_seconds(1), + backoff_rate=1.0, + jitter_strategy=JitterStrategy.NONE, + retryable_error_types=[RuntimeError], + ) + ) + ), + ) + context.wait_for_condition( + check_condition, + name="otel-context-condition", + config=WaitForConditionConfig(initial_state=0, wait_strategy=wait_strategy), + ) + context.wait_for_callback(submit_callback, name="otel-context-callback") + + # These callbacks run after the last asynchronous wait, so normal replay + # does not repeat the helper body or the checkpointless virtual child. + try: + with_retry( + context, + helper_body, + WithRetryConfig(retry_strategy=helper_retry_strategy), + name="otel-context-with-retry", + ) + except ChildContextError as error: + if error.message != HELPER_FAILURE: + raise + else: + raise RuntimeError("Expected the intentional helper failure") + context.run_in_child_context( + virtual_child, + name="otel-context-virtual", + config=ChildConfig(is_virtual=True), + ) + return "callback-context-complete" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_24_invocation_retry_status.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_24_invocation_retry_status.py new file mode 100644 index 000000000..568622517 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_24_invocation_retry_status.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""A real invocation retry for OTel status-mapping case 24.""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python import ( + DurableContext, + StepContext, + durable_execution, + durable_step, +) +from aws_durable_execution_sdk_python.exceptions import InvocationError +from common import otel_plugin, require_scenario + + +@durable_step +def before_invocation_retry(_step_context: StepContext) -> str: + return "saved" + + +@durable_execution(plugins=[otel_plugin()]) +def handler(event: dict[str, Any], context: DurableContext) -> str: + require_scenario(event, "invocation-retry-status") + # Capture this at entry: consuming the completed step can end replay mode. + entered_replay = context.is_replaying() + context.step(before_invocation_retry(), name="otel-before-invocation-retry") + if not entered_replay: + raise InvocationError("Conformance invocation retry after a completed step") + return "retry-complete" diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/template.yaml b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/template.yaml index be70e122d..40b1dbd39 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/template.yaml +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/template.yaml @@ -409,6 +409,62 @@ Resources: FunctionName: !Sub "${AWS::StackName}-otel-invocation-20" Role: !Ref LambdaExecutionRoleArn + Otel21CompletedStepReplay: + Type: AWS::Serverless::Function + Condition: DeployInvocationView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-invocation-21 + Properties: + CodeUri: src/ + Handler: otel_21_completed_step_replay.handler + FunctionName: !Sub "${AWS::StackName}-otel-invocation-21" + Role: !Ref LambdaExecutionRoleArn + + Otel22UserFunctionContext: + Type: AWS::Serverless::Function + Condition: DeployInvocationView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-invocation-22 + Properties: + CodeUri: src/ + Handler: otel_22_user_function_context.handler + FunctionName: !Sub "${AWS::StackName}-otel-invocation-22" + Role: !Ref LambdaExecutionRoleArn + + Otel23CallbackFunctionContext: + Type: AWS::Serverless::Function + Condition: DeployInvocationView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-invocation-23 + Properties: + CodeUri: src/ + Handler: otel_23_callback_function_context.handler + FunctionName: !Sub "${AWS::StackName}-otel-invocation-23" + Role: !Ref LambdaExecutionRoleArn + + Otel24InvocationRetryStatus: + Type: AWS::Serverless::Function + Condition: DeployInvocationView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-invocation-24 + Properties: + CodeUri: src/ + Handler: otel_24_invocation_retry_status.handler + FunctionName: !Sub "${AWS::StackName}-otel-invocation-24" + Role: !Ref LambdaExecutionRoleArn + OtelExecution1Success: Type: AWS::Serverless::Function Condition: DeployExecutionView @@ -781,3 +837,71 @@ Resources: Environment: Variables: OTEL_PLUGIN_MODE: execution + + OtelExecution21CompletedStepReplay: + Type: AWS::Serverless::Function + Condition: DeployExecutionView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-execution-21 + Properties: + CodeUri: src/ + Handler: otel_21_completed_step_replay.handler + FunctionName: !Sub "${AWS::StackName}-otel-execution-21" + Role: !Ref LambdaExecutionRoleArn + Environment: + Variables: + OTEL_PLUGIN_MODE: execution + + OtelExecution22UserFunctionContext: + Type: AWS::Serverless::Function + Condition: DeployExecutionView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-execution-22 + Properties: + CodeUri: src/ + Handler: otel_22_user_function_context.handler + FunctionName: !Sub "${AWS::StackName}-otel-execution-22" + Role: !Ref LambdaExecutionRoleArn + Environment: + Variables: + OTEL_PLUGIN_MODE: execution + + OtelExecution23CallbackFunctionContext: + Type: AWS::Serverless::Function + Condition: DeployExecutionView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-execution-23 + Properties: + CodeUri: src/ + Handler: otel_23_callback_function_context.handler + FunctionName: !Sub "${AWS::StackName}-otel-execution-23" + Role: !Ref LambdaExecutionRoleArn + Environment: + Variables: + OTEL_PLUGIN_MODE: execution + + OtelExecution24InvocationRetryStatus: + Type: AWS::Serverless::Function + Condition: DeployExecutionView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-execution-24 + Properties: + CodeUri: src/ + Handler: otel_24_invocation_retry_status.handler + FunctionName: !Sub "${AWS::StackName}-otel-execution-24" + Role: !Ref LambdaExecutionRoleArn + Environment: + Variables: + OTEL_PLUGIN_MODE: execution diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_function_context.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_function_context.py new file mode 100644 index 000000000..2aecd8cc6 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_function_context.py @@ -0,0 +1,122 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Verify actual SDK callback parents across retries and external completion.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +import time +from collections import Counter +from pathlib import Path + +import pytest +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry import context, trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +SRC_DIR = Path(__file__).resolve().parents[1] / "src" +EXPECTED_PARENTS = { + "retry-attempt-1": "otel-context-retry-step attempt 1", + "retry-attempt-2": "otel-context-retry-step attempt 2", + "condition-check-1": "otel-context-condition attempt 1", + "condition-check-2": "otel-context-condition attempt 2", + "callback-submitter": "otel-context-callback submitter attempt 1", + "with-retry-body": "otel-context-with-retry", + "with-retry-strategy": "otel-context-with-retry", + "virtual-child": "otel-context-virtual", +} + + +@pytest.mark.parametrize("plugin_type", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_callback_probes_keep_their_actual_sdk_parent( + monkeypatch: pytest.MonkeyPatch, + plugin_type: type[ExecutionOtelPlugin] | type[InvocationOtelPlugin], +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + monkeypatch.setattr(trace, "get_tracer_provider", lambda: provider) + plugin = plugin_type( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + before_context = context.get_current() + common_spec = importlib.util.spec_from_file_location( + "common", SRC_DIR / "common.py" + ) + assert common_spec is not None and common_spec.loader is not None + common = importlib.util.module_from_spec(common_spec) + monkeypatch.setitem(sys.modules, "common", common) + common_spec.loader.exec_module(common) + monkeypatch.setattr(common, "otel_plugin", lambda: plugin) + spec = importlib.util.spec_from_file_location( + "otel_23_callback_function_context", + SRC_DIR / "otel_23_callback_function_context.py", + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + try: + with DurableFunctionTestRunner(handler=module.handler) as runner: + arn = runner.run_async( + input=json.dumps({"scenario": "callback-function-context"}), timeout=30 + ) + callback_id = runner.wait_for_callback( + arn, name="otel-context-callback create callback id", timeout=10 + ) + # External delivery delay, outside the durable handler. + time.sleep(1) + runner.send_callback_success(callback_id, result=b"callback-complete") + result = runner.wait_for_result(arn, timeout=15) + assert result.status.value == "SUCCEEDED" + assert result.result is not None + assert json.loads(result.result) == "callback-context-complete" + spans = exporter.get_finished_spans() + probes = [span for span in spans if span.name.startswith("conformance.")] + assert Counter(span.name for span in probes) == Counter( + f"conformance.{label}" for label in EXPECTED_PARENTS + ) + spans_by_id = { + (span.context.trace_id, span.context.span_id): span + for span in spans + if span.context is not None + } + sdk_traces = { + span.context.trace_id + for span in spans + if span.context is not None + and span.attributes is not None + and span.attributes.get("durable.execution.arn") == arn + } + for probe_span in probes: + assert probe_span.context is not None + assert probe_span.parent is not None + assert probe_span.attributes is not None + label = str(probe_span.attributes["conformance.callback"]) + assert "durable.execution.arn" not in probe_span.attributes + assert probe_span.context.trace_id in sdk_traces + parent = spans_by_id[ + (probe_span.context.trace_id, probe_span.parent.span_id) + ] + assert parent.name == EXPECTED_PARENTS[label] + assert parent.attributes is not None + assert parent.attributes["durable.execution.arn"] == arn + assert context.get_current() == before_context + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_completed_step_replay.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_completed_step_replay.py new file mode 100644 index 000000000..d331decca --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_completed_step_replay.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Exercise case 21 through the public decorator and local durable runner.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from pathlib import Path +from typing import Any + +import pytest +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry import context +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +SRC_DIR = Path(__file__).resolve().parents[1] / "src" + + +@pytest.mark.parametrize("plugin_type", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_completed_step_replay_uses_saved_result( + monkeypatch: pytest.MonkeyPatch, + plugin_type: type[ExecutionOtelPlugin] | type[InvocationOtelPlugin], +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = plugin_type( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + before_context = context.get_current() + common_spec = importlib.util.spec_from_file_location( + "common", SRC_DIR / "common.py" + ) + assert common_spec is not None and common_spec.loader is not None + common = importlib.util.module_from_spec(common_spec) + monkeypatch.setitem(sys.modules, "common", common) + common_spec.loader.exec_module(common) + monkeypatch.setattr(common, "otel_plugin", lambda: plugin) + spec = importlib.util.spec_from_file_location( + "otel_21_completed_step_replay", SRC_DIR / "otel_21_completed_step_replay.py" + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + calls: list[str] = [] + + def track(factory: Any, label: str) -> Any: + def make_step() -> Any: + step = factory() + + def counted(step_context: Any) -> str: + calls.append(label) + return step(step_context) + + return counted + + return make_step + + monkeypatch.setattr(module, "before_wait", track(module.before_wait, "before")) + monkeypatch.setattr(module, "after_wait", track(module.after_wait, "after")) + try: + with DurableFunctionTestRunner(handler=module.handler) as runner: + result = runner.run( + input=json.dumps({"scenario": "completed-step-replay"}), timeout=15 + ) + assert result.status.value == "SUCCEEDED" + assert result.result is not None + assert json.loads(result.result) == "before-after" + assert calls == ["before", "after"] + spans = exporter.get_finished_spans() + invocations = [span for span in spans if span.name == "Invocation"] + assert [ + span.attributes["durable.invocation.status"] + for span in invocations + if span.attributes is not None + ] == ["PENDING", "SUCCEEDED"] + if plugin_type is ExecutionOtelPlugin: + # Count raw exports, including duplicates that reuse a span ID. + for name in ("otel-before-wait", "otel-replay-wait", "otel-after-wait"): + assert len([span for span in spans if span.name == name]) == 1 + assert context.get_current() == before_context + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_invocation_retry_status.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_invocation_retry_status.py new file mode 100644 index 000000000..d2851a709 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_invocation_retry_status.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Exercise a real invocation retry through the public decorator and local runner.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from pathlib import Path +from typing import Any + +import pytest +from aws_durable_execution_sdk_python.plugin import InvocationEndInfo, InvocationStatus +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry import context +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +SRC_DIR = Path(__file__).resolve().parents[1] / "src" + + +@pytest.mark.parametrize("plugin_type", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_invocation_retry_preserves_completed_step( + monkeypatch: pytest.MonkeyPatch, + plugin_type: type[ExecutionOtelPlugin] | type[InvocationOtelPlugin], +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = plugin_type( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + statuses: list[InvocationStatus] = [] + original_end = plugin.on_invocation_end + + def observe_end(info: InvocationEndInfo) -> None: + statuses.append(info.status) + original_end(info) + + monkeypatch.setattr(plugin, "on_invocation_end", observe_end) + before_context = context.get_current() + common_spec = importlib.util.spec_from_file_location( + "common", SRC_DIR / "common.py" + ) + assert common_spec is not None and common_spec.loader is not None + common = importlib.util.module_from_spec(common_spec) + monkeypatch.setitem(sys.modules, "common", common) + common_spec.loader.exec_module(common) + monkeypatch.setattr(common, "otel_plugin", lambda: plugin) + spec = importlib.util.spec_from_file_location( + "otel_24_invocation_retry_status", + SRC_DIR / "otel_24_invocation_retry_status.py", + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + calls: list[str] = [] + + def track(factory: Any, label: str) -> Any: + def make_step() -> Any: + step = factory() + + def counted(step_context: Any) -> str: + calls.append(label) + return step(step_context) + + return counted + + return make_step + + monkeypatch.setattr( + module, + "before_invocation_retry", + track(module.before_invocation_retry, "saved"), + ) + try: + with DurableFunctionTestRunner(handler=module.handler) as runner: + result = runner.run( + input=json.dumps({"scenario": "invocation-retry-status"}), timeout=15 + ) + assert result.status.value == "SUCCEEDED" + assert result.result is not None + assert json.loads(result.result) == "retry-complete" + assert calls == ["saved"] + assert statuses == [InvocationStatus.RETRY, InvocationStatus.SUCCEEDED] + spans = exporter.get_finished_spans() + invocations = [span for span in spans if span.name == "Invocation"] + assert len(invocations) == 2 + assert [span.status.status_code.name for span in invocations] == ["UNSET", "OK"] + workflows = [span for span in spans if span.name == "Workflow"] + assert len(workflows) == 1 + assert workflows[0].status.status_code.name == "OK" + assert context.get_current() == before_context + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py index 5908dd7d9..2745f1611 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py @@ -50,6 +50,10 @@ ("Otel18ChainedInvokeFailure", "otel-invocation-18"), ("Otel19ExecutionFailure", "otel-invocation-19"), ("Otel20VirtualContext", "otel-invocation-20"), + ("Otel21CompletedStepReplay", "otel-invocation-21"), + ("Otel22UserFunctionContext", "otel-invocation-22"), + ("Otel23CallbackFunctionContext", "otel-invocation-23"), + ("Otel24InvocationRetryStatus", "otel-invocation-24"), ("OtelExecution1Success", "otel-execution-1"), ("OtelExecution2WaitResume", "otel-execution-2"), ("OtelExecution3Retry", "otel-execution-3"), @@ -70,6 +74,10 @@ ("OtelExecution18ChainedInvokeFailure", "otel-execution-18"), ("OtelExecution19ExecutionFailure", "otel-execution-19"), ("OtelExecution20VirtualContext", "otel-execution-20"), + ("OtelExecution21CompletedStepReplay", "otel-execution-21"), + ("OtelExecution22UserFunctionContext", "otel-execution-22"), + ("OtelExecution23CallbackFunctionContext", "otel-execution-23"), + ("OtelExecution24InvocationRetryStatus", "otel-execution-24"), ] EXPECTED_LONG_RUNNING_MAPPINGS: list[tuple[str, str]] = [ ("OtelLongRunning1Wait", "otel-long-running-1"), @@ -152,6 +160,10 @@ "otel_18_chained_invoke_failure", "otel_19_execution_failure", "otel_20_virtual_context", + "otel_21_completed_step_replay", + "otel_22_user_function_context", + "otel_23_callback_function_context", + "otel_24_invocation_retry_status", "otel_long_running_1_wait", "otel_long_running_2_retry", "otel_long_running_3_callback", diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 5e78f2a9c..bf30b2733 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -297,6 +297,15 @@ 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. +### Status attributes + +`durable.invocation.status` uses `RETRYING` when the core plugin hook reports +`InvocationStatus.RETRY`. The core enum remains unchanged. Operation spans use +OTel `OK` only for `SUCCEEDED`, `ERROR` when error details are delivered, and +`UNSET` for other outcomes without error details, including `FAILED`, +`CANCELLED`, `TIMED_OUT`, and `STOPPED`. The original durable operation status +remains in `durable.operation.status`. + ### 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/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..dba9bda80 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 @@ -43,6 +43,7 @@ InvocationStartInfo, OperationEndInfo, OperationStartInfo, + OperationStatus, OperationType, UserFunctionEndInfo, UserFunctionOutcome, @@ -607,7 +608,9 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: if self._invocation_span is not None: self._invocation_span.set_attribute( "durable.invocation.status", - info.status.value if info.status else "", + "RETRYING" + if info.status is InvocationStatus.RETRY + else info.status.value, ) if info.status in (InvocationStatus.SUCCEEDED, InvocationStatus.PENDING): self._invocation_span.set_status(StatusCode.OK) @@ -726,7 +729,7 @@ def on_operation_end(self, info: OperationEndInfo) -> None: span.record_exception( Exception(info.error.message or info.error.type or "Unknown error") ) - else: + elif info.status is OperationStatus.SUCCEEDED: span.set_status(StatusCode.OK) self._note_parent_end(info.parent_id, end_time) 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..1c4f74a34 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 @@ -14,6 +14,7 @@ InvocationStartInfo, OperationEndInfo, OperationStartInfo, + OperationStatus, OperationType, UserFunctionEndInfo, UserFunctionOutcome, @@ -676,7 +677,10 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: invocation_span = self._get_span(None) if invocation_span: invocation_span.set_attribute( - "durable.invocation.status", info.status.value + "durable.invocation.status", + "RETRYING" + if info.status is InvocationStatus.RETRY + else info.status.value, ) # Span status mapping: SUCCEEDED/PENDING -> OK, FAILED -> ERROR, # RETRY -> UNSET. RETRY is left UNSET because the plugin interface @@ -780,7 +784,7 @@ def on_operation_end(self, info: OperationEndInfo) -> None: span.record_exception( Exception(info.error.message or info.error.type or "Unknown error") ) - else: + elif info.status is OperationStatus.SUCCEEDED: span.set_status(StatusCode.OK) self._end_span(info.operation_id, info.end_time) 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..12edcc0ae 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 @@ -1123,7 +1123,9 @@ def test_invocation_span_status_kind_and_attributes(status, expected_code): invocation = {s.name: s for s in exporter.get_finished_spans()}["Invocation"] assert invocation.kind is trace.SpanKind.INTERNAL assert invocation.attributes is not None - assert invocation.attributes["durable.invocation.status"] == status.value + assert invocation.attributes["durable.invocation.status"] == ( + "RETRYING" if status is InvocationStatus.RETRY else status.value + ) assert invocation.attributes["durable.invocation.first"] is True assert invocation.status.status_code is expected_code 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..1894d2b48 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 @@ -563,7 +563,11 @@ def test_invocation_span_status_reflects_execution_status( invocation = next(s for s in spans if s.name == "Invocation") attributes = invocation.attributes assert attributes is not None - assert attributes["durable.invocation.status"] == invocation_status.value + assert attributes["durable.invocation.status"] == ( + "RETRYING" + if invocation_status is InvocationStatus.RETRY + else invocation_status.value + ) assert invocation.status.status_code is expected_span_status diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_status_mapping.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_status_mapping.py new file mode 100644 index 000000000..d1224a305 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_status_mapping.py @@ -0,0 +1,172 @@ +"""Shared status regressions through both public plugin lifecycles.""" + +from collections.abc import Iterator +from datetime import UTC, datetime, timedelta +from typing import Any + +import pytest +from aws_durable_execution_sdk_python.lambda_service import ( + ErrorObject, + OperationStatus, + OperationSubType, +) +from aws_durable_execution_sdk_python.plugin import ( + InvocationEndInfo, + InvocationStartInfo, + InvocationStatus, + OperationEndInfo, + OperationStartInfo, + OperationType, +) +from opentelemetry import context +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +START = datetime(2026, 10, 1, tzinfo=UTC) +END = START + timedelta(seconds=1) +ARN = ( + "arn:aws:lambda:us-west-2:123456789012:function:status:1/durable-execution/test/id" +) + + +InstrumentedView = tuple[ + ExecutionOtelPlugin | InvocationOtelPlugin, InMemorySpanExporter +] + + +@pytest.fixture(params=[ExecutionOtelPlugin, InvocationOtelPlugin]) +def instrumented_view(request: pytest.FixtureRequest) -> Iterator[InstrumentedView]: + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = request.param( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + original = context.get_current() + yield plugin, exporter + assert context.get_current() == original + provider.shutdown() + + +@pytest.mark.parametrize( + "status", + [ + OperationStatus.SUCCEEDED, + OperationStatus.FAILED, + OperationStatus.CANCELLED, + OperationStatus.TIMED_OUT, + OperationStatus.STOPPED, + ], +) +@pytest.mark.parametrize("with_error", [False, True]) +@pytest.mark.parametrize("started_here", [False, True]) +def test_operation_status_and_exception( + instrumented_view: InstrumentedView, + status: OperationStatus, + with_error: bool, + started_here: bool, +) -> None: + """A delivered terminal hook must not label an unsuccessful operation OK.""" + plugin, exporter = instrumented_view + invocation: dict[str, Any] = dict( + request_id="first", + execution_arn=ARN, + execution_start_time=START, + is_first_invocation=True, + ) + plugin.on_invocation_start(InvocationStartInfo(**invocation)) + operation: dict[str, Any] = dict( + operation_id="wait", + name="wait", + parent_id=None, + operation_type=OperationType.WAIT, + sub_type=OperationSubType.WAIT, + start_time=START, + is_replayed=False, + ) + if started_here: + plugin.on_operation_start( + OperationStartInfo(**operation, status=OperationStatus.STARTED) + ) + error = ( + ErrorObject( + message="deadline exceeded", + type="TimeoutError", + data=None, + stack_trace=None, + ) + if with_error + else None + ) + plugin.on_operation_end( + OperationEndInfo(**operation, end_time=END, status=status, error=error) + ) + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + span = next(s for s in exporter.get_finished_spans() if s.name == "wait") + assert span.attributes is not None + assert span.attributes["durable.operation.status"] == status.value + expected = ( + StatusCode.ERROR + if with_error + else ( + StatusCode.OK if status is OperationStatus.SUCCEEDED else StatusCode.UNSET + ) + ) + assert span.status.status_code is expected + if with_error: + assert span.status.description == "deadline exceeded" + assert len(span.events) == 1 + assert span.events[0].name == "exception" + assert span.events[0].attributes is not None + assert span.events[0].attributes["exception.message"] == "deadline exceeded" + else: + assert not span.events + + +def test_retry_then_resume_exports_retrying( + instrumented_view: InstrumentedView, +) -> None: + """Telemetry normalizes RETRY without changing the core enum or next invocation.""" + plugin, exporter = instrumented_view + original = context.get_current() + for index, status in enumerate( + [InvocationStatus.RETRY, InvocationStatus.SUCCEEDED] + ): + invocation: dict[str, Any] = dict( + request_id=f"request-{index}", + execution_arn=ARN, + execution_start_time=START, + is_first_invocation=index == 0, + ) + plugin.on_invocation_start(InvocationStartInfo(**invocation)) + plugin.on_invocation_end(InvocationEndInfo(**invocation, status=status)) + assert context.get_current() == original + spans = [s for s in exporter.get_finished_spans() if s.name == "Invocation"] + assert [ + s.attributes["durable.invocation.status"] + for s in spans + if s.attributes is not None + ] == [ + "RETRYING", + "SUCCEEDED", + ] + assert [s.status.status_code for s in spans] == [StatusCode.UNSET, StatusCode.OK] + assert InvocationStatus.RETRY.value == "RETRY" + assert spans[0].context is not None + assert spans[1].context is not None + assert spans[0].context.trace_id == spans[1].context.trace_id From 049099ef35abfd2834b146b03a605a99eef32d0f Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Sat, 3 Oct 2026 00:56:51 +0000 Subject: [PATCH 04/54] ci: pin revised OTel conformance requirements --- .../tests/test_opentelemetry_conformance_workflow.py | 8 ++++++-- .github/workflows/opentelemetry-conformance-tests.yml | 6 +++--- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py index 891cfd1be..dd2879e2d 100644 --- a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py +++ b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py @@ -1,3 +1,4 @@ +import re from pathlib import Path @@ -18,7 +19,10 @@ def test_opentelemetry_conformance_caller_uses_current_workflow_contract() -> No "uses: aws/aws-durable-execution-conformance-tests/.github/workflows/" "opentelemetry-orchestrator.yml@" ) - assert orchestrator in workflow + pinned_ref = re.search(re.escape(orchestrator) + r"([0-9a-f]{40})", workflow) + assert pinned_ref is not None + shared_ref = pinned_ref.group(1) + assert f"default: {shared_ref}" in workflow assert "python-opentelemetry.yml@" not in workflow assert "\n otlp_endpoint:" not in workflow @@ -27,7 +31,7 @@ def test_opentelemetry_conformance_caller_uses_current_workflow_contract() -> No "resource_prefix: p", "sdk_repository: aws/aws-durable-execution-sdk-python", "sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }}", - "conformance_test_ref: ${{ inputs.conformance_test_ref || 'main' }}", + f"conformance_test_ref: ${{{{ inputs.conformance_test_ref || '{shared_ref}' }}}}", "checkout_sdk: true", f"examples_dir: {EXAMPLES_DIR}", "adot_release_repository: aws-observability/aws-otel-python-instrumentation", diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index 3a6387dae..6d2ea2ef8 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -46,7 +46,7 @@ on: conformance_test_ref: description: Conformance test commit SHA or branch name required: true - default: 922dfafa575f771a6d5b117d9f1e631d149e01e4 + default: 98724ee08b7bd08bf334322e97c4b95e8f61d8bb type: string permissions: {} @@ -60,13 +60,13 @@ jobs: actions: write contents: read id-token: write - uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@922dfafa575f771a6d5b117d9f1e631d149e01e4 + uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@98724ee08b7bd08bf334322e97c4b95e8f61d8bb with: language: python resource_prefix: p sdk_repository: aws/aws-durable-execution-sdk-python sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }} - conformance_test_ref: ${{ inputs.conformance_test_ref || '922dfafa575f771a6d5b117d9f1e631d149e01e4' }} + conformance_test_ref: ${{ inputs.conformance_test_ref || '98724ee08b7bd08bf334322e97c4b95e8f61d8bb' }} # Check the SDK out so the handlers and templates below are on disk. The handlers # themselves are installed from sdk_ref by src/requirements.txt during the SAM build. checkout_sdk: true 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 05/54] 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 06/54] 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 6cb34302645965b2d31d83d22887c67ab44a9077 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Sat, 3 Oct 2026 02:37:01 +0000 Subject: [PATCH 07/54] ci: validate OTel conformance against main --- .github/workflows/ci.yml | 3 +-- .github/workflows/opentelemetry-conformance-tests.yml | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b2d66fd44..8e9d6c15e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -10,8 +10,7 @@ on: branches: [ main ] pull_request: - # Run checks while the conformance PR is stacked on the focused context fix. - branches: ["main", "fix/otel-handler-context-428"] + branches: ["main"] jobs: lint-commits: diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index 36e00a43f..c428d1ae4 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -8,8 +8,7 @@ name: OpenTelemetry Conformance Tests on: pull_request: - # Run checks while the conformance PR is stacked on the focused context fix. - branches: ["main", "fix/otel-handler-context-428"] + branches: ["main"] paths: - "packages/aws-durable-execution-sdk-python/**" - "packages/aws-durable-execution-sdk-python-otel/**" From 370295fbec57b748928781b68744726e367f714c Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Sat, 3 Oct 2026 02:55:43 +0000 Subject: [PATCH 08/54] 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 09/54] 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 10/54] 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 97f78adc2c5e4afc16836eda2925adcdca65329b Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Fri, 2 Oct 2026 22:20:29 -0700 Subject: [PATCH 11/54] 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 8e9d6c15e..50fa5e4d5 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 dd69e05728c60647aae76202321f5d34b605c00e Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Sat, 3 Oct 2026 04:51:04 +0000 Subject: [PATCH 12/54] 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 1894d2b48..723ad41ef 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 @@ -2202,3 +2202,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 146f2313801b8bce2f8402d5e8c8f1cd064e4279 Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Sat, 3 Oct 2026 18:38:56 -0700 Subject: [PATCH 13/54] test: pin reviewed OTel conformance revisions --- .../opentelemetry-conformance-tests.yml | 6 +- .../README.md | 13 +++ .../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 +++++++++++++++++++ 6 files changed, 209 insertions(+), 6 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python/tests/e2e/handler_invocation_context_int_test.py diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index c428d1ae4..626f4a2e3 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -45,7 +45,7 @@ on: conformance_test_ref: description: Conformance test commit SHA or branch name required: true - default: 98b802cdb172614f98e217f7464784d47b9bb484 + default: 743dc5080145160dc1932c88f858ca6fe90d91eb type: string permissions: {} @@ -59,13 +59,13 @@ jobs: actions: write contents: read id-token: write - uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@98b802cdb172614f98e217f7464784d47b9bb484 + uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@743dc5080145160dc1932c88f858ca6fe90d91eb with: language: python resource_prefix: p sdk_repository: aws/aws-durable-execution-sdk-python sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }} - conformance_test_ref: ${{ inputs.conformance_test_ref || '98b802cdb172614f98e217f7464784d47b9bb484' }} + conformance_test_ref: ${{ inputs.conformance_test_ref || '743dc5080145160dc1932c88f858ca6fe90d91eb' }} # Check the SDK out so the handlers and templates below are on disk. The handlers # themselves are installed from sdk_ref by src/requirements.txt during the SAM build. checkout_sdk: true diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 07dcd1fa8..195303ba6 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -318,6 +318,19 @@ OTel `OK` only for `SUCCEEDED`, `ERROR` when error details are delivered, and `CANCELLED`, `TIMED_OUT`, and `STOPPED`. The original durable operation status remains in `durable.operation.status`. +### Invocation context isolation + +Invocation hooks retain their caller thread and registration order. With the +updated core, invocation-local context-variable bindings are isolated from the +host: hooks see the incoming context and the handler receives their resulting +context, while invocation exit restores the host's original bindings even if a +plugin fails during setup or cleanup. Plugins must not use invocation context +bindings to mutate the host context after the invocation has returned. Older +supported cores retain their existing lifecycle behavior, including the +execution-view limitation when later plugins open invocation context scopes. +The new isolation applies only when plugins are registered. + + ### 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 ed3c6ba9ac89700ab9d8185112361e7941fdecbe Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Sat, 3 Oct 2026 19:13:43 -0700 Subject: [PATCH 14/54] 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 9055bdaaa93e8e0a9cdc8b387b76846db98c7b45 Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Mon, 5 Oct 2026 21:11:20 -0700 Subject: [PATCH 15/54] fix: bind execution view and stabilize CI --- .github/workflows/ci.yml | 1 + .github/workflows/cloud-tests.yml | 4 +- .../opentelemetry-conformance-tests.yml | 6 +- .../plugin/plugin_wait_replay_flag.py | 9 +- .../tests/test_plugin_wait_replay_flag.py | 72 +++++ .../README.md | 15 +- .../execution_plugin.py | 22 ++ .../invocation_plugin.py | 2 + .../e2e/test_invocation_wait_resume_int.py | 41 ++- .../tests/test_execution_plugin.py | 71 +++++ .../README.md | 19 ++ .../plugin.py | 108 +++++-- .../tests/handler_worker_context_test.py | 280 +++++++++++++++++- 13 files changed, 606 insertions(+), 44 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests/tests/test_plugin_wait_replay_flag.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 50fa5e4d5..a1d35c08f 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/.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/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index 12067e032..7782d2530 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -45,7 +45,7 @@ on: conformance_test_ref: description: Conformance test commit SHA or branch name required: true - default: f5dfb8cbfa5583f0ea73aec3fafd80542846d270 + default: 6182cee00ed33ce01e0579066fb31096f988817f type: string permissions: {} @@ -59,13 +59,13 @@ jobs: actions: write contents: read id-token: write - uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@f5dfb8cbfa5583f0ea73aec3fafd80542846d270 + uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@6182cee00ed33ce01e0579066fb31096f988817f with: language: python resource_prefix: p sdk_repository: aws/aws-durable-execution-sdk-python sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }} - conformance_test_ref: ${{ inputs.conformance_test_ref || 'f5dfb8cbfa5583f0ea73aec3fafd80542846d270' }} + conformance_test_ref: ${{ inputs.conformance_test_ref || '6182cee00ed33ce01e0579066fb31096f988817f' }} # Check the SDK out so the handlers and templates below are on disk. The handlers # themselves are installed from sdk_ref by src/requirements.txt during the SAM build. checkout_sdk: true 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 195303ba6..668169094 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -166,6 +166,15 @@ 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. +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) @@ -328,7 +337,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-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 dba9bda80..c4134da40 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 ( @@ -124,6 +126,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 = ( @@ -488,6 +492,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/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 ce1304d24..4fe006675 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 @@ -107,6 +107,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..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 @@ -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,11 @@ def handler_impl(_event: Any, context: DurableContext) -> str: # the legacy lane continues checking its supported combinations above. + ( [(ExecutionOtelPlugin, True)] - if hasattr(DurableInstrumentationPlugin, "handler_context") + + [(ExecutionOtelPlugin, kind) for kind in ("same", "unrelated", "absent")] + if getattr( + core_plugin_api, "DURABLE_INSTRUMENTATION_HANDLER_CONTEXT_API_VERSION", None + ) + == 1 else [] ), ) @@ -284,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) @@ -292,7 +297,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)) @@ -335,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) @@ -425,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 12edcc0ae..23c7394ba 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 ( @@ -1622,3 +1623,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 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 5ae5beb51..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,18 +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. 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 @@ -473,12 +462,26 @@ 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 [] 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 +495,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 +525,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,21 +570,29 @@ 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) + # 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__() 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 +605,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..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 @@ -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 == [ @@ -149,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 @@ -249,7 +257,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 +300,263 @@ 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): + __durable_handler_context_api__ = 1 + + 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) + + +@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 3fcf3778695d9f13090845196c62107120818722 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 7 Oct 2026 00:03:24 +0000 Subject: [PATCH 16/54] ci: route OTel conformance to CodeBuild --- .github/scripts/tests/test_opentelemetry_conformance_workflow.py | 1 + .github/workflows/opentelemetry-conformance-tests.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py index dd2879e2d..f1016456f 100644 --- a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py +++ b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py @@ -28,6 +28,7 @@ def test_opentelemetry_conformance_caller_uses_current_workflow_contract() -> No for configuration in ( "language: python", + "runs_on: codebuild-github-actions-runner-${{ github.run_id }}-${{ github.run_attempt }}", "resource_prefix: p", "sdk_repository: aws/aws-durable-execution-sdk-python", "sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }}", diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index 7782d2530..a0141a1a9 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -62,6 +62,7 @@ jobs: uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@6182cee00ed33ce01e0579066fb31096f988817f with: language: python + runs_on: codebuild-github-actions-runner-${{ github.run_id }}-${{ github.run_attempt }} resource_prefix: p sdk_repository: aws/aws-durable-execution-sdk-python sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }} From 680232b3ddb4435eafca8382b0272694f9eacaab Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 7 Oct 2026 00:36:36 +0000 Subject: [PATCH 17/54] ci: preserve queued conformance runs --- ...test_opentelemetry_conformance_workflow.py | 41 ++++++++++++++++--- .github/workflows/cloud-tests.yml | 1 + .github/workflows/conformance-tests.yml | 1 + .../opentelemetry-conformance-tests.yml | 7 ++++ 4 files changed, 45 insertions(+), 5 deletions(-) diff --git a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py index f1016456f..8cef0c104 100644 --- a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py +++ b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py @@ -1,13 +1,13 @@ import re from pathlib import Path +import yaml + WORKFLOW_PATH = ( Path(__file__).parents[2] / "workflows" / "opentelemetry-conformance-tests.yml" ) -EXAMPLES_DIR = ( - ".build/durable-sdk/packages/aws-durable-execution-sdk-python-conformance-tests-otel" -) +EXAMPLES_DIR = ".build/durable-sdk/packages/aws-durable-execution-sdk-python-conformance-tests-otel" def test_opentelemetry_conformance_caller_uses_current_workflow_contract() -> None: @@ -71,8 +71,39 @@ def test_opentelemetry_conformance_runs_when_the_handlers_change() -> None: workflow = WORKFLOW_PATH.read_text() trigger_path = ( - " - " - '"packages/aws-durable-execution-sdk-python-conformance-tests-otel/**"' + ' - "packages/aws-durable-execution-sdk-python-conformance-tests-otel/**"' ) # Once for pull_request, once for push. assert workflow.count(trigger_path) == 2 + + +def test_opentelemetry_conformance_queues_complete_runs() -> None: + workflow = yaml.safe_load(WORKFLOW_PATH.read_text()) + + assert workflow["concurrency"] == { + "group": "otel-conformance-tests", + "cancel-in-progress": False, + "queue": "max", + } + + +def test_cloud_tests_queue_each_shared_runtime_stack() -> None: + workflow = yaml.safe_load(WORKFLOW_PATH.with_name("cloud-tests.yml").read_text()) + + assert workflow["jobs"]["example-tests"]["concurrency"] == { + "group": "cloud-tests-${{ matrix.python-prefix }}", + "cancel-in-progress": False, + "queue": "max", + } + + +def test_general_conformance_keeps_all_pending_runs() -> None: + workflow = yaml.safe_load( + WORKFLOW_PATH.with_name("conformance-tests.yml").read_text() + ) + + assert workflow["concurrency"] == { + "group": "conformance-tests-global", + "cancel-in-progress": False, + "queue": "max", + } diff --git a/.github/workflows/cloud-tests.yml b/.github/workflows/cloud-tests.yml index cd3a5cb96..da80a43e2 100644 --- a/.github/workflows/cloud-tests.yml +++ b/.github/workflows/cloud-tests.yml @@ -46,6 +46,7 @@ jobs: concurrency: group: cloud-tests-${{ matrix.python-prefix }} cancel-in-progress: false + queue: max strategy: fail-fast: false matrix: diff --git a/.github/workflows/conformance-tests.yml b/.github/workflows/conformance-tests.yml index 2e9b73d7b..2d660987e 100644 --- a/.github/workflows/conformance-tests.yml +++ b/.github/workflows/conformance-tests.yml @@ -32,6 +32,7 @@ concurrency: # run against the shared stack. group: conformance-tests-global cancel-in-progress: false + queue: max permissions: contents: read diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index a0141a1a9..f76a3aab8 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -48,6 +48,13 @@ on: default: 6182cee00ed33ce01e0579066fb31096f988817f type: string +# Backend stacks are shared across PRs. Queue whole runs so reusable +# backend jobs do not displace another PR from their pending slots. +concurrency: + group: otel-conformance-tests + cancel-in-progress: false + queue: max + permissions: {} jobs: From 3373a3f6afade458ccc3d5924b6481424c611b8a Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 7 Oct 2026 01:26:25 +0000 Subject: [PATCH 18/54] ci: adopt shared backend queue preservation --- .../tests/test_opentelemetry_conformance_workflow.py | 9 ++++++--- .github/workflows/opentelemetry-conformance-tests.yml | 2 +- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py index 8cef0c104..741981058 100644 --- a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py +++ b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py @@ -4,6 +4,9 @@ import yaml +SHARED_WORKFLOW_REF = "f5855f2d0f60be996973173cf479c3567f83e30f" +CONFORMANCE_TEST_REF = "6182cee00ed33ce01e0579066fb31096f988817f" + WORKFLOW_PATH = ( Path(__file__).parents[2] / "workflows" / "opentelemetry-conformance-tests.yml" ) @@ -21,8 +24,8 @@ def test_opentelemetry_conformance_caller_uses_current_workflow_contract() -> No ) pinned_ref = re.search(re.escape(orchestrator) + r"([0-9a-f]{40})", workflow) assert pinned_ref is not None - shared_ref = pinned_ref.group(1) - assert f"default: {shared_ref}" in workflow + assert pinned_ref.group(1) == SHARED_WORKFLOW_REF + assert f"default: {CONFORMANCE_TEST_REF}" in workflow assert "python-opentelemetry.yml@" not in workflow assert "\n otlp_endpoint:" not in workflow @@ -32,7 +35,7 @@ def test_opentelemetry_conformance_caller_uses_current_workflow_contract() -> No "resource_prefix: p", "sdk_repository: aws/aws-durable-execution-sdk-python", "sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }}", - f"conformance_test_ref: ${{{{ inputs.conformance_test_ref || '{shared_ref}' }}}}", + f"conformance_test_ref: ${{{{ inputs.conformance_test_ref || '{CONFORMANCE_TEST_REF}' }}}}", "checkout_sdk: true", f"examples_dir: {EXAMPLES_DIR}", "adot_release_repository: aws-observability/aws-otel-python-instrumentation", diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index f76a3aab8..5f7cd354e 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -66,7 +66,7 @@ jobs: actions: write contents: read id-token: write - uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@6182cee00ed33ce01e0579066fb31096f988817f + uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@f5855f2d0f60be996973173cf479c3567f83e30f with: language: python runs_on: codebuild-github-actions-runner-${{ github.run_id }}-${{ github.run_attempt }} From a663841dd9ebecf58729707cd90ff3b20028b6b3 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 7 Oct 2026 02:23:05 +0000 Subject: [PATCH 19/54] fix: build new js example workspace dependencies --- .github/scripts/js_examples/run.py | 16 ++++++- .github/scripts/tests/test_js_examples.py | 52 +++++++++++++++++++++++ 2 files changed, 67 insertions(+), 1 deletion(-) diff --git a/.github/scripts/js_examples/run.py b/.github/scripts/js_examples/run.py index e2fa359a8..7dac7e829 100755 --- a/.github/scripts/js_examples/run.py +++ b/.github/scripts/js_examples/run.py @@ -61,15 +61,24 @@ TESTING_SOURCE = REPO_ROOT / "packages/aws-durable-execution-sdk-python-testing/src" JS_SDK_URL = "https://github.com/aws/aws-durable-execution-sdk-js.git" EXAMPLES_REL = Path("packages/aws-durable-execution-sdk-js-examples") -# The examples depend on these workspaces by "*". The root "npm run build" +# The examples depend on these local workspaces. The root "npm run build" # also builds the insight tools and the VS Code extension, which the harness # does not need, so the script builds only these, in dependency order. JS_BUILD_WORKSPACES = ( "packages/aws-durable-execution-sdk-js", "packages/aws-durable-execution-sdk-js-testing", "packages/aws-durable-execution-sdk-js-otel", + "packages/aws-durable-execution-sdk-js-extras", + "packages/aws-durable-execution-sdk-js-microvm-worker", "packages/aws-durable-execution-sdk-js-examples", ) +# Older JS refs predate these packages; their examples do not need them. +OPTIONAL_JS_BUILD_WORKSPACES = frozenset( + { + "packages/aws-durable-execution-sdk-js-extras", + "packages/aws-durable-execution-sdk-js-microvm-worker", + } +) # otel examples export spans to an OpenTelemetry collector. The harness does # not run one, so these examples cannot pass here. They are not selected. EXCLUDED_DIRS = ("/otel/",) @@ -202,6 +211,11 @@ def build_js_sdk(js_dir: Path, *, force: bool) -> None: [npm, "ci", "--no-audit", "--no-fund"], cwd=js_dir, env=env, check=True ) for workspace in JS_BUILD_WORKSPACES: + if ( + workspace in OPTIONAL_JS_BUILD_WORKSPACES + and not (js_dir / workspace / "package.json").is_file() + ): + continue log(f"npm run build -w {workspace}") subprocess.run( [npm, "run", "build", "-w", workspace], cwd=js_dir, env=env, check=True diff --git a/.github/scripts/tests/test_js_examples.py b/.github/scripts/tests/test_js_examples.py index 15522209b..dd56d143e 100644 --- a/.github/scripts/tests/test_js_examples.py +++ b/.github/scripts/tests/test_js_examples.py @@ -623,6 +623,58 @@ def test_source_fingerprint_changes_with_local_edits(tmp_path: Path) -> None: assert run.source_fingerprint(tmp_path) != untracked +@pytest.mark.parametrize("has_microvm_packages", [False, True]) +def test_build_produces_example_workspace_dependencies( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + has_microvm_packages: bool, +) -> None: + core = "packages/aws-durable-execution-sdk-js" + testing = "packages/aws-durable-execution-sdk-js-testing" + otel = "packages/aws-durable-execution-sdk-js-otel" + extras = "packages/aws-durable-execution-sdk-js-extras" + worker = "packages/aws-durable-execution-sdk-js-microvm-worker" + examples_package = "packages/aws-durable-execution-sdk-js-examples" + packages = [core, testing, otel, examples_package] + if has_microvm_packages: + packages.extend([extras, worker]) + for package in packages: + directory = tmp_path / package + directory.mkdir(parents=True) + (directory / "package.json").write_text("{}") + + built: list[str] = [] + + def execute(args: list[str], **_kwargs: Any) -> subprocess.CompletedProcess[str]: + if args[1] == "ci": + return subprocess.CompletedProcess(args, 0) + workspace = args[-1] + assert (tmp_path / workspace / "package.json").is_file() + if workspace == worker: + assert testing in built + if workspace == examples_package and has_microvm_packages: + # The generator imports extras/microvm and typechecks the worker. + # A successful examples build requires both artifacts to exist. + assert extras in built + assert worker in built + built.append(workspace) + return subprocess.CompletedProcess(args, 0) + + stamp = tmp_path / "built-stamp" + monkeypatch.setattr(run, "_head", lambda _: "test-head") + monkeypatch.setattr(run, "state_file", lambda *_: stamp) + monkeypatch.setattr(run, "source_fingerprint", lambda _: "fingerprint") + monkeypatch.setattr(run.shutil, "which", lambda _: "/fake/npm") + monkeypatch.setattr(run.subprocess, "run", execute) + + run.build_js_sdk(tmp_path, force=True) + + assert built[-1] == examples_package + assert stamp.read_text() == "fingerprint\n" + if not has_microvm_packages: + assert built == [core, testing, otel, examples_package] + + def test_proxy_creates_the_dump_directory(tmp_path: Path) -> None: dump_dir = tmp_path / "new" / "dump" proxy = invoke_proxy.ProxyServer(0, "http://127.0.0.1:1", dump_dir) From 746e38862e4989b78610827bd54ffea6d0fb120d Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 00:20:32 +0000 Subject: [PATCH 20/54] fix: deliver external completion hooks once Record local callback success, failure, and timeout updates for the next invocation. Deduplicate actual terminal update notifications within an invocation while preserving first delivery and clearing state between invocations. Keep checkpoint versions and error payload handling intact. Exercise public callback completion through both memory and file stores, including timeout, stored results, and two subsequent replays. Cover duplicate checkpoints, first notification, reset, and concurrent updates. --- .../execution.py | 3 + .../e2e/callback_updated_operations_test.py | 161 ++++++++++++++++++ .../tests/execution_test.py | 32 ++++ .../plugin.py | 16 ++ .../tests/plugin_test.py | 60 +++++++ 5 files changed, 272 insertions(+) create mode 100644 packages/aws-durable-execution-sdk-python-testing/tests/e2e/callback_updated_operations_test.py diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py index 82880a1ba..7548631bc 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py @@ -647,6 +647,7 @@ def complete_callback_success( end_timestamp=now if now is not None else real_now(), callback_details=updated_callback_details, ) + self._record_updated_operation(operation.operation_id) return self.operations[index] def complete_callback_failure( @@ -673,6 +674,7 @@ def complete_callback_failure( end_timestamp=now if now is not None else real_now(), callback_details=updated_callback_details, ) + self._record_updated_operation(operation.operation_id) return self.operations[index] def complete_callback_timeout( @@ -699,6 +701,7 @@ def complete_callback_timeout( end_timestamp=now if now is not None else real_now(), callback_details=updated_callback_details, ) + self._record_updated_operation(operation.operation_id) return self.operations[index] def complete_chained_invoke( diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/e2e/callback_updated_operations_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/callback_updated_operations_test.py new file mode 100644 index 000000000..b58118894 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/callback_updated_operations_test.py @@ -0,0 +1,161 @@ +"""Real callback completions are delivered once before resumed user code.""" + +from __future__ import annotations + +import json +import time +from pathlib import Path +from typing import Any + +import pytest +from aws_durable_execution_sdk_python import ( + DurableContext, + StepContext, + durable_execution, + durable_step, +) +from aws_durable_execution_sdk_python.config import ( + CallbackConfig, + Duration, + WaitForCallbackConfig, +) +from aws_durable_execution_sdk_python.exceptions import CallbackError +from aws_durable_execution_sdk_python.lambda_service import ErrorObject +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationStartInfo, + OperationEndInfo, + UserFunctionStartInfo, +) +from aws_durable_execution_sdk_python.serdes import JsonSerDes +from aws_durable_execution_sdk_python.types import WaitForCallbackContext + +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from aws_durable_execution_sdk_python_testing.stores.filesystem import ( + FileSystemExecutionStore, +) + + +class CallbackObserver(DurableInstrumentationPlugin): + def __init__(self) -> None: + self.invocations: list[InvocationStartInfo] = [] + self.ends: list[OperationEndInfo] = [] + self.order: list[str] = [] + + def on_invocation_start(self, info: InvocationStartInfo) -> None: + self.invocations.append(info) + + def on_operation_end(self, info: OperationEndInfo) -> None: + if info.name == "target": + self.ends.append(info) + self.order.append("target-end") + + def on_user_function_start(self, info: UserFunctionStartInfo) -> None: + if info.name == "observed": + self.order.append("observed-start") + + +def _submit(_callback_id: str, _context: WaitForCallbackContext) -> None: + return None + + +def _suspended_callback(runner: DurableFunctionTestRunner, arn: str, name: str) -> str: + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + events = runner.get_execution_history(arn, include_execution_data=True).events + starts = [ + event + for event in events + if event.event_type == "CallbackStarted" and event.name == name + ] + if starts and any( + event.event_type == "InvocationCompleted" + and event.event_id > starts[0].event_id + for event in events + ): + details = starts[0].callback_started_details + assert details is not None + assert details.callback_id is not None + return details.callback_id + time.sleep(0.01) + raise AssertionError(f"Callback {name} did not suspend") + + +@pytest.mark.parametrize("outcome", ["success", "failure", "timeout"]) +@pytest.mark.parametrize("filesystem", [False, True]) +def test_callback_update_is_consumed_before_user_code( + outcome: str, filesystem: bool, tmp_path: Path +) -> None: + observer = CallbackObserver() + marker_calls: list[str] = [] + + @durable_step + def observed(_context: StepContext, value: str) -> str: + marker_calls.append(value) + return value + + def handler(_event: Any, context: DurableContext) -> str: + config = CallbackConfig( + timeout=Duration.from_seconds(1 if outcome == "timeout" else 30), + serdes=JsonSerDes(), + ) + try: + target = context.create_callback(name="target", config=config).result() + except CallbackError: + target = outcome + assert isinstance(target, str) + saved = context.step(observed(target), name="observed") + callback_config = WaitForCallbackConfig(serdes=JsonSerDes()) + one = context.wait_for_callback(_submit, name="one", config=callback_config) + two = context.wait_for_callback(_submit, name="two", config=callback_config) + return "/".join((saved, one, two)) + + wrapped = durable_execution(handler, plugins=[observer]) + store = FileSystemExecutionStore.create(tmp_path) if filesystem else None + with DurableFunctionTestRunner( + handler=wrapped, + store=store, + skip_time=False, + poll_interval=0.01, + execution_timeout=25, + ) as runner: + arn = runner.run_async(input="{}") + callback_id = _suspended_callback(runner, arn, "target") + if outcome == "success": + runner.send_callback_success( + callback_id, result=json.dumps("target").encode() + ) + elif outcome == "failure": + runner.send_callback_failure( + callback_id, error=ErrorObject.from_message("explicit callback failure") + ) + # The timeout case uses the runner's real scheduled callback deadline. + for name in ["one", "two"]: + callback_id = _suspended_callback(runner, arn, name + " create callback id") + runner.send_callback_success(callback_id, result=json.dumps(name).encode()) + result = runner.wait_for_result(arn, timeout=10) + + expected = "target" if outcome == "success" else outcome + assert result.status.value == "SUCCEEDED" + assert result.result is not None + assert json.loads(result.result) == expected + "/one/two" + assert marker_calls == [expected] + assert len(observer.invocations) == 4 + assert len(observer.ends) == 1 + terminal = observer.ends[0] + assert ( + terminal.status.value + == {"success": "SUCCEEDED", "failure": "FAILED", "timeout": "TIMED_OUT"}[ + outcome + ] + ) + assert terminal.is_replayed is False + assert set(observer.invocations[1].updated_operations) == {terminal.operation_id} + assert all( + terminal.operation_id not in invocation.updated_operations + for invocation in observer.invocations[2:] + ) + assert observer.order == ["target-end", "observed-start"] + if outcome == "failure": + assert terminal.error is not None + assert terminal.error.message == "explicit callback failure" diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/execution_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/execution_test.py index bde5e3d4b..f3ac97185 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/execution_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/execution_test.py @@ -955,6 +955,38 @@ def test_from_dict_with_none_result(): # region callback +@pytest.mark.parametrize("outcome", ["success", "failure", "timeout"]) +def test_callback_completion_records_only_successful_state_changes(outcome): + """Preserve payloads, token versions and consumed metadata on rejection.""" + operation = Operation( + operation_id="callback", + operation_type=OperationType.CALLBACK, + status=OperationStatus.STARTED, + callback_details=CallbackDetails(callback_id="callback-id"), + ) + execution = Execution("test-arn", _make_start_input(), [operation]) + complete = getattr(execution, f"complete_callback_{outcome}") + payload = b'"result"' if outcome == "success" else ErrorObject.from_message("error") + token_version = execution.token_sequence + result = complete("callback-id", payload) + + assert execution.updated_operation_ids == ["callback"] + assert execution.token_sequence == token_version + assert execution.seq_counter == 1 + assert result.callback_details.result == ( + '"result"' if outcome == "success" else None + ) + assert result.callback_details.error == (None if outcome == "success" else payload) + restored = Execution.from_json_dict(execution.to_json_dict()) + assert restored.updated_operation_ids == ["callback"] + execution.mark_state_delivered() + with pytest.raises(IllegalStateException, match="not in STARTED state"): + complete("callback-id", payload) + assert execution.updated_operation_ids == [] + assert execution.seq_counter == 1 + assert execution.token_sequence == token_version + + def test_find_callback_operation_not_found(): """Test find_callback_operation raises exception when callback not found.""" execution = Execution("test-arn", Mock(), []) diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py index f9e202017..f10b6bf5d 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,6 +10,7 @@ from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from enum import Enum +from threading import Lock from typing import Any, Callable, MutableMapping, cast from aws_durable_execution_sdk_python.identifier import OperationIdentifier @@ -482,9 +483,13 @@ def __init__(self, plugins: list[DurableInstrumentationPlugin] | None): self._operations_provider: Callable[[], Mapping[str, Operation]] | None = None self._startup_context: contextvars.Context | None = None self._invocation_contexts: list[contextvars.Context | None] = [] + self._reported_terminal_updates: set[tuple[str, OperationStatus]] = set() + self._terminal_updates_lock = Lock() @contextlib.contextmanager def run(self): + with self._terminal_updates_lock: + self._reported_terminal_updates.clear() if self._plugins: self._executor = ThreadPoolExecutor( max_workers=1, @@ -500,6 +505,8 @@ def run(self): # Shut down the thread pool, waiting for pending tasks to complete. if self._executor: self._executor.shutdown(wait=True) + with self._terminal_updates_lock: + self._reported_terminal_updates.clear() @staticmethod def _dispatch_plugin(plugin: DurableInstrumentationPlugin, info) -> bool: @@ -878,6 +885,15 @@ def on_operation_update( ) for operation in updated_operations: if self._is_terminal_status(operation.status): + # Replay delivery and checkpoint responses can report the same + # completion. Deduplicate actual notifications, not state that + # may have arrived without an UpdatedOperationIds notification. + if self._plugins: + key = (operation.operation_id, operation.status) + with self._terminal_updates_lock: + if key in self._reported_terminal_updates: + continue + self._reported_terminal_updates.add(key) self.execute_plugins( OperationEndInfo( operation_id=operation.operation_id, diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py index 69cdfc502..0777fce2a 100644 --- a/packages/aws-durable-execution-sdk-python/tests/plugin_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_test.py @@ -2,6 +2,7 @@ import logging import pickle import unittest +from concurrent.futures import ThreadPoolExecutor from copy import deepcopy from dataclasses import asdict, fields from unittest.mock import MagicMock, patch @@ -1537,6 +1538,65 @@ def test_terminal_status_without_step_details_fires_operation_only(self): self.assertIn("operation_end:op-1", self.plugin.calls) + def test_checkpoint_does_not_repeat_an_observed_terminal_update(self): + """A resumed external result can also appear in the next checkpoint.""" + for status in ( + OperationStatus.SUCCEEDED, + OperationStatus.FAILED, + OperationStatus.CANCELLED, + OperationStatus.TIMED_OUT, + OperationStatus.STOPPED, + ): + with self.subTest(status=status): + plugin = _TrackingPlugin() + executor = PluginExecutor(plugins=[plugin]) + operation = self._make_operation(status=status) + with executor.run(): + # First completion delivered through UpdatedOperationIds. + executor.on_operation_update(operation) + # The next response carries the already observed state. + executor.on_operation_update( + [operation], + operations={operation.operation_id: operation}, + previous_operations={operation.operation_id: operation}, + ) + self.assertEqual(plugin.calls.count("operation_end:op-1"), 1) + + def test_checkpoint_terminal_transition_still_emits(self): + previous = self._make_operation(status=OperationStatus.STARTED) + operation = self._make_operation(status=OperationStatus.SUCCEEDED) + with self.executor.run(): + self.executor.on_operation_update( + [operation], + operations={operation.operation_id: operation}, + previous_operations={previous.operation_id: previous}, + ) + self.assertEqual(self.plugin.calls.count("operation_end:op-1"), 1) + + def test_checkpoint_preserves_first_terminal_notification(self): + """State can predate notification, e.g. without UpdatedOperationIds.""" + operation = self._make_operation(status=OperationStatus.SUCCEEDED) + with self.executor.run(): + self.executor.on_operation_update( + [operation], + operations={operation.operation_id: operation}, + previous_operations={operation.operation_id: operation}, + ) + self.assertEqual(self.plugin.calls.count("operation_end:op-1"), 1) + + def test_terminal_notification_resets_between_invocations(self): + operation = self._make_operation(status=OperationStatus.SUCCEEDED) + for _ in range(2): + with self.executor.run(): + self.executor.on_operation_update(operation) + self.assertEqual(self.plugin.calls.count("operation_end:op-1"), 2) + + def test_concurrent_completion_notifications_emit_once(self): + operation = self._make_operation(status=OperationStatus.SUCCEEDED) + with self.executor.run(), ThreadPoolExecutor(max_workers=8) as workers: + list(workers.map(self.executor.on_operation_update, [operation] * 16)) + self.assertEqual(self.plugin.calls.count("operation_end:op-1"), 1) + def test_non_terminal_status_without_step_details_fires_nothing(self): op = self._make_operation(status=OperationStatus.STARTED, step_details=None) From 12e760cf72dc14055ae0dec48bf0a20ad73451b2 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 00:41:32 +0000 Subject: [PATCH 21/54] test: add external callback conformance Add the root callback completion/replay handler and both view mappings, plus aliases of the existing callback-failure handler for case 25. Exercise four real invocations, stored results, first-completion timing, raw terminal export counts, parentage, and two subsequent replays. Document the tester/core completion fixes and the remaining cloud no-error-details precondition. Keep the workflow and test references unchanged until the shared 26-case definitions are published together. --- .../README.md | 31 +++ .../src/Makefile | 4 + ..._26_external_callback_completion_replay.py | 53 +++++ .../template.yaml | 62 ++++++ ...est_external_callback_completion_replay.py | 195 ++++++++++++++++++ .../tests/test_otel_examples.py | 5 + 6 files changed, 350 insertions(+) create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_26_external_callback_completion_replay.py create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_external_callback_completion_replay.py diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md index b2871c83d..0b0ea8b2a 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md @@ -67,6 +67,8 @@ view named by its `OtelSuite` parameter. | `otel-invocation-22` | `otel_22_user_function_context.handler` | Creates ordinary user spans under the active handler, attempt, child, branch, and iteration contexts. | | `otel-invocation-23` | `otel_23_callback_function_context.handler` | Verifies retry/check attempts, callback submitter, wrapped retry helper, and virtual-child callback parents. | | `otel-invocation-24` | `otel_24_invocation_retry_status.handler` | Raises a retryable invocation error after a completed step, then resumes with its saved result. | +| `otel-invocation-25` | `otel_17_wait_for_callback_failure.handler` | Targets a failed callback without error details; service history must satisfy the explicit no-error-details precondition. | +| `otel-invocation-26` | `otel_26_external_callback_completion_replay.handler` | Completes a root callback after suspension, saves its result in a step, and replays it through two callback barriers. | | `otel-execution-1` | `otel_1_success.handler` | Verifies the execution-view workflow, step, and attempt hierarchy. | | `otel-execution-2` | `otel_2_wait_resume.handler` | Verifies the execution view across a resumed invocation. | | `otel-execution-3` | `otel_3_retry.handler` | Verifies the execution view across retry attempts. | @@ -91,6 +93,8 @@ view named by its `OtelSuite` parameter. | `otel-execution-22` | `otel_22_user_function_context.handler` | Observes active execution-view callback contexts without supplying or repairing parents. | | `otel-execution-23` | `otel_23_callback_function_context.handler` | Verifies the same SDK-owned callback lifecycle parents in the execution view. | | `otel-execution-24` | `otel_24_invocation_retry_status.handler` | Verifies invocation retry status independently of step retry and recovery re-exports. | +| `otel-execution-25` | `otel_17_wait_for_callback_failure.handler` | Targets the same errorless failed callback and its `UNSET` leaf in the execution view. | +| `otel-execution-26` | `otel_26_external_callback_completion_replay.handler` | Requires one terminal root-callback export at first completion and no duplicate exports on two later replays. | | `otel-long-running-1` | `otel_long_running_1_wait.handler` | Verifies wait and resume telemetry across a long durable suspension. | | `otel-long-running-2` | `otel_long_running_2_retry.handler` | Verifies retry telemetry across a long durable backoff. | | `otel-long-running-3` | `otel_long_running_3_callback.handler` | Verifies callback telemetry when completion arrives after a long delay. | @@ -99,6 +103,33 @@ view named by its `OtelSuite` parameter. The runner discovers each mapping from `TestingMetadata.TestDescription` on the functions in the templates. +## External completion and status coverage + +The local testing library includes callback success, failure and timeout in the +next invocation's `UpdatedOperationIds`. The core delivers each terminal update +notification once per invocation, including when a resumed operation and a later +checkpoint response carry the same completion. It tracks actual notifications, +preserves first delivery when the update-ID metadata is absent, and clears that +tracking at invocation boundaries. Public runner regressions cover memory and +file stores, stored step results, failure payloads and two subsequent replays. + +Case 24 covers invocation `RETRY` becoming `RETRYING`/`UNSET`. Case 25 targets the +separate rule that a failed operation without error details remains `UNSET`. +Its cloud coverage requires both views to pass the raw service-history no-error-details +precondition and the telemetry assertions. Local file-store results establish +SDK behavior; that AWS service precondition still needs independent validation. +`CANCELLED`, `TIMED_OUT` and `STOPPED` without error details have explicit unit +coverage in both views; their corresponding cloud paths are not established. +The existing success/`OK` and detailed-error/`ERROR` controls remain in the suites. + +Case 26 revisits a root-level public callback on every replay. The driver waits +for `InvocationCompleted` before completing the target and each barrier. The +target's terminal span must precede the observed step and must not be exported +again during the two later resumes. Invocation view also retains its initial +pending callback segment. Cloud validation must count the raw S3 export records +without deduplicating equal span IDs; local runner checks alone do not establish +cloud coverage. + ## Callback coverage boundary Cases 22 and 23 create normal user spans from the active context, without an diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/Makefile b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/Makefile index 41e731b57..8382a49cb 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/Makefile +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/Makefile @@ -6,6 +6,7 @@ build-Otel17WaitForCallbackFailure build-Otel18ChainedInvokeFailure \ build-Otel18InvokeTarget build-Otel19ExecutionFailure build-Otel20VirtualContext \ build-Otel21CompletedStepReplay build-Otel22UserFunctionContext build-Otel23CallbackFunctionContext build-Otel24InvocationRetryStatus \ + build-Otel25CallbackFailureWithoutError build-Otel26ExternalCallbackReplay \ build-OtelExecution1Success \ build-OtelExecution2WaitResume build-OtelExecution3Retry \ build-OtelExecution4TerminalFailure build-OtelExecution5ChildContext \ @@ -20,6 +21,7 @@ build-OtelExecution20VirtualContext \ build-OtelExecution21CompletedStepReplay build-OtelExecution22UserFunctionContext build-OtelExecution23CallbackFunctionContext \ build-OtelExecution24InvocationRetryStatus \ + build-OtelExecution25CallbackFailureWithoutError build-OtelExecution26ExternalCallbackReplay \ build-OtelLongRunning1Wait \ build-OtelLongRunning2Retry build-OtelLongRunning3Callback \ build-OtelLongRunning4ChainedInvoke build-OtelLongRunning4InvokeTarget @@ -32,6 +34,7 @@ build-Otel14MapFailure build-Otel15WaitInterrupted build-Otel16WaitForConditionF build-Otel17WaitForCallbackFailure build-Otel18ChainedInvokeFailure \ build-Otel18InvokeTarget build-Otel19ExecutionFailure build-Otel20VirtualContext \ build-Otel21CompletedStepReplay build-Otel22UserFunctionContext build-Otel23CallbackFunctionContext build-Otel24InvocationRetryStatus \ +build-Otel25CallbackFailureWithoutError build-Otel26ExternalCallbackReplay \ build-OtelExecution1Success \ build-OtelExecution2WaitResume build-OtelExecution3Retry \ build-OtelExecution4TerminalFailure build-OtelExecution5ChildContext \ @@ -46,6 +49,7 @@ build-OtelExecution18InvokeTarget build-OtelExecution19ExecutionFailure \ build-OtelExecution20VirtualContext \ build-OtelExecution21CompletedStepReplay build-OtelExecution22UserFunctionContext build-OtelExecution23CallbackFunctionContext \ build-OtelExecution24InvocationRetryStatus \ +build-OtelExecution25CallbackFailureWithoutError build-OtelExecution26ExternalCallbackReplay \ build-OtelLongRunning1Wait \ build-OtelLongRunning2Retry build-OtelLongRunning3Callback \ build-OtelLongRunning4ChainedInvoke build-OtelLongRunning4InvokeTarget: diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_26_external_callback_completion_replay.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_26_external_callback_completion_replay.py new file mode 100644 index 000000000..4e562e2f5 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_26_external_callback_completion_replay.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Public callback completion followed by two controlled replay barriers.""" + +from __future__ import annotations + +from typing import Any, cast + +from aws_durable_execution_sdk_python import ( + DurableContext, + StepContext, + durable_execution, + durable_step, +) +from aws_durable_execution_sdk_python.config import ( + CallbackConfig, + WaitForCallbackConfig, +) +from aws_durable_execution_sdk_python.serdes import JsonSerDes +from aws_durable_execution_sdk_python.types import WaitForCallbackContext +from common import otel_plugin, require_scenario + + +def submit_callback(_callback_id: str, _context: WaitForCallbackContext) -> None: + return None + + +@durable_step +def observe_target(_context: StepContext, result: str) -> str: + return result + + +@durable_execution(plugins=[otel_plugin()]) +def handler(event: dict[str, Any], context: DurableContext) -> str: + require_scenario(event, "external-callback-completion-replay") + config = WaitForCallbackConfig(serdes=JsonSerDes()) + target = cast( + str, + context.create_callback( + name="otel-external-target", config=CallbackConfig(serdes=JsonSerDes()) + ).result(), + ) + observed = context.step( + observe_target(target), name="otel-external-target-observed" + ) + one = context.wait_for_callback( + submit_callback, name="otel-external-barrier-one", config=config + ) + two = context.wait_for_callback( + submit_callback, name="otel-external-barrier-two", config=config + ) + return "/".join((observed, one, two)) diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/template.yaml b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/template.yaml index 40b1dbd39..8d898ba46 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/template.yaml +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/template.yaml @@ -465,6 +465,34 @@ Resources: FunctionName: !Sub "${AWS::StackName}-otel-invocation-24" Role: !Ref LambdaExecutionRoleArn + Otel25CallbackFailureWithoutError: + Type: AWS::Serverless::Function + Condition: DeployInvocationView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-invocation-25 + Properties: + CodeUri: src/ + Handler: otel_17_wait_for_callback_failure.handler + FunctionName: !Sub "${AWS::StackName}-otel-invocation-25" + Role: !Ref LambdaExecutionRoleArn + + Otel26ExternalCallbackReplay: + Type: AWS::Serverless::Function + Condition: DeployInvocationView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-invocation-26 + Properties: + CodeUri: src/ + Handler: otel_26_external_callback_completion_replay.handler + FunctionName: !Sub "${AWS::StackName}-otel-invocation-26" + Role: !Ref LambdaExecutionRoleArn + OtelExecution1Success: Type: AWS::Serverless::Function Condition: DeployExecutionView @@ -905,3 +933,37 @@ Resources: Environment: Variables: OTEL_PLUGIN_MODE: execution + + OtelExecution25CallbackFailureWithoutError: + Type: AWS::Serverless::Function + Condition: DeployExecutionView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-execution-25 + Properties: + CodeUri: src/ + Handler: otel_17_wait_for_callback_failure.handler + FunctionName: !Sub "${AWS::StackName}-otel-execution-25" + Role: !Ref LambdaExecutionRoleArn + Environment: + Variables: + OTEL_PLUGIN_MODE: execution + + OtelExecution26ExternalCallbackReplay: + Type: AWS::Serverless::Function + Condition: DeployExecutionView + Metadata: + BuildMethod: makefile + TestingMetadata: + TestDescription: + - otel-execution-26 + Properties: + CodeUri: src/ + Handler: otel_26_external_callback_completion_replay.handler + FunctionName: !Sub "${AWS::StackName}-otel-execution-26" + Role: !Ref LambdaExecutionRoleArn + Environment: + Variables: + OTEL_PLUGIN_MODE: execution diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_external_callback_completion_replay.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_external_callback_completion_replay.py new file mode 100644 index 000000000..2ec2f34fa --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_external_callback_completion_replay.py @@ -0,0 +1,195 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""A real callback completion is exported before two later public replays.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +import time +from pathlib import Path +from typing import Any + +import pytest +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from aws_durable_execution_sdk_python_testing.stores.filesystem import ( + FileSystemExecutionStore, +) +from opentelemetry import context +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +SRC_DIR = Path(__file__).resolve().parents[1] / "src" + + +def _suspended_callback( + runner: DurableFunctionTestRunner, arn: str, name: str +) -> tuple[str, tuple[int, int]]: + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + events = runner.get_execution_history(arn, include_execution_data=True).events + starts = [ + event + for event in events + if event.event_type == "CallbackStarted" and event.name == name + ] + completions = [ + event + for event in events + if event.event_type == "InvocationCompleted" + and starts + and event.event_id > starts[0].event_id + ] + if starts and completions: + details = starts[0].callback_started_details + assert details is not None and details.callback_id is not None + return details.callback_id, (starts[0].event_id, completions[0].event_id) + time.sleep(0.01) + raise AssertionError(f"Callback {name} did not suspend") + + +@pytest.mark.parametrize("plugin_type", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_external_callback_completion_precedes_later_replays( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + plugin_type: type[ExecutionOtelPlugin] | type[InvocationOtelPlugin], +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = plugin_type( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + before_context = context.get_current() + common_spec = importlib.util.spec_from_file_location( + "common", SRC_DIR / "common.py" + ) + assert common_spec is not None and common_spec.loader is not None + common = importlib.util.module_from_spec(common_spec) + monkeypatch.setitem(sys.modules, "common", common) + common_spec.loader.exec_module(common) + monkeypatch.setattr(common, "otel_plugin", lambda: plugin) + spec = importlib.util.spec_from_file_location( + "otel_26_external_callback_completion_replay", + SRC_DIR / "otel_26_external_callback_completion_replay.py", + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + calls: list[str] = [] + original_observe = module.observe_target + + def counted_observe(value: str) -> Any: + step = original_observe(value) + + def counted(step_context: Any) -> str: + calls.append(value) + return step(step_context) + + return counted + + monkeypatch.setattr(module, "observe_target", counted_observe) + gates: list[tuple[int, int]] = [] + try: + with DurableFunctionTestRunner( + handler=module.handler, + store=FileSystemExecutionStore.create(tmp_path), + poll_interval=0.01, + execution_timeout=25, + skip_time=False, + ) as runner: + arn = runner.run_async( + input=json.dumps({"scenario": "external-callback-completion-replay"}) + ) + for name, payload in ( + ("otel-external-target", "target"), + ("otel-external-barrier-one create callback id", "one"), + ("otel-external-barrier-two create callback id", "two"), + ): + callback_id, gate = _suspended_callback(runner, arn, name) + gates.append(gate) + runner.send_callback_success( + callback_id, result=json.dumps(payload).encode("utf-8") + ) + result = runner.wait_for_result(arn, timeout=10) + history = runner.get_execution_history(arn, include_execution_data=True) + + assert result.status.value == "SUCCEEDED" + assert result.result is not None + assert json.loads(result.result) == "target/one/two" + assert calls == ["target"] + assert gates == [(2, 3), (8, 11), (15, 18)] + assert [ + event.event_id + for event in history.events + if event.event_type == "InvocationCompleted" + ] == [3, 11, 18, 21] + assert history.events[-1].event_type == "ExecutionSucceeded" + assert history.events[-1].event_id == 22 + + spans = exporter.get_finished_spans() + assert all(span.context is not None for span in spans) + assert len({span.context.trace_id for span in spans if span.context}) == 1 + invocations = sorted( + [span for span in spans if span.name == "Invocation"], + key=lambda span: span.start_time or 0, + ) + assert len(invocations) == 4 + assert [ + (span.attributes or {}).get("durable.invocation.status") + for span in invocations + ] == ["PENDING", "PENDING", "PENDING", "SUCCEEDED"] + targets = [ + span + for span in spans + if span.name == "otel-external-target" + and (span.attributes or {}).get("durable.operation.status") == "SUCCEEDED" + ] + # Count raw exports, including duplicate records with identical span IDs. + assert len(targets) == 1 + target = targets[0] + assert target.status.status_code.name == "OK" + observed = [ + span + for span in spans + if span.name == "otel-external-target-observed attempt 1" + ] + assert len(observed) == 1 + assert target.end_time is not None and observed[0].start_time is not None + assert target.end_time <= observed[0].start_time + assert invocations[2].start_time is not None + assert invocations[3].start_time is not None + assert target.end_time < invocations[2].start_time < invocations[3].start_time + parent = ( + invocations[1] + if plugin_type is InvocationOtelPlugin + else next(span for span in spans if span.name == "Workflow") + ) + assert target.parent is not None and parent.context is not None + assert target.parent.span_id == parent.context.span_id + # Invocation view also exports the legitimate first-invocation segment. + pending_target = [ + span + for span in spans + if span.name == "otel-external-target" + and (span.attributes or {}).get("durable.operation.status") == "STARTED" + ] + assert len(pending_target) == (1 if plugin_type is InvocationOtelPlugin else 0) + assert context.get_current() == before_context + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py index 2745f1611..61c0d24e1 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py @@ -54,6 +54,8 @@ ("Otel22UserFunctionContext", "otel-invocation-22"), ("Otel23CallbackFunctionContext", "otel-invocation-23"), ("Otel24InvocationRetryStatus", "otel-invocation-24"), + ("Otel25CallbackFailureWithoutError", "otel-invocation-25"), + ("Otel26ExternalCallbackReplay", "otel-invocation-26"), ("OtelExecution1Success", "otel-execution-1"), ("OtelExecution2WaitResume", "otel-execution-2"), ("OtelExecution3Retry", "otel-execution-3"), @@ -78,6 +80,8 @@ ("OtelExecution22UserFunctionContext", "otel-execution-22"), ("OtelExecution23CallbackFunctionContext", "otel-execution-23"), ("OtelExecution24InvocationRetryStatus", "otel-execution-24"), + ("OtelExecution25CallbackFailureWithoutError", "otel-execution-25"), + ("OtelExecution26ExternalCallbackReplay", "otel-execution-26"), ] EXPECTED_LONG_RUNNING_MAPPINGS: list[tuple[str, str]] = [ ("OtelLongRunning1Wait", "otel-long-running-1"), @@ -164,6 +168,7 @@ "otel_22_user_function_context", "otel_23_callback_function_context", "otel_24_invocation_retry_status", + "otel_26_external_callback_completion_replay", "otel_long_running_1_wait", "otel_long_running_2_retry", "otel_long_running_3_callback", From 246cad826359f3c7e96340ff115177bc8c89bd23 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 01:06:26 +0000 Subject: [PATCH 22/54] fix: consume delivered callback update IDs Prune update IDs covered by a successful checkpoint delivery while retaining later changes. Keep token and watermark values, pure reads, rejected requests, cached responses, and error payloads unchanged. Cover early success, failure and timeout during submission with real public APIs and two later replays on both stores. Exercise rejection, read, retry and concurrent update boundaries. Pin the reusable workflow and conformance definitions to the coordinated 26-case shared revision fb95e20acb0b9b10f5fc032ac521a4fd45efdd73. --- ...test_opentelemetry_conformance_workflow.py | 4 +- .../opentelemetry-conformance-tests.yml | 6 +- .../execution.py | 10 +- .../e2e/callback_updated_operations_test.py | 92 +++++++++++++++++++ .../tests/execution_test.py | 56 +++++++++++ .../tests/executor_checkpoint_test.py | 44 +++++++++ 6 files changed, 206 insertions(+), 6 deletions(-) diff --git a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py index 741981058..775348d07 100644 --- a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py +++ b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py @@ -4,8 +4,8 @@ import yaml -SHARED_WORKFLOW_REF = "f5855f2d0f60be996973173cf479c3567f83e30f" -CONFORMANCE_TEST_REF = "6182cee00ed33ce01e0579066fb31096f988817f" +SHARED_WORKFLOW_REF = "fb95e20acb0b9b10f5fc032ac521a4fd45efdd73" +CONFORMANCE_TEST_REF = "fb95e20acb0b9b10f5fc032ac521a4fd45efdd73" WORKFLOW_PATH = ( Path(__file__).parents[2] / "workflows" / "opentelemetry-conformance-tests.yml" diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index 5f7cd354e..1bee20cf6 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -45,7 +45,7 @@ on: conformance_test_ref: description: Conformance test commit SHA or branch name required: true - default: 6182cee00ed33ce01e0579066fb31096f988817f + default: fb95e20acb0b9b10f5fc032ac521a4fd45efdd73 type: string # Backend stacks are shared across PRs. Queue whole runs so reusable @@ -66,14 +66,14 @@ jobs: actions: write contents: read id-token: write - uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@f5855f2d0f60be996973173cf479c3567f83e30f + uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@fb95e20acb0b9b10f5fc032ac521a4fd45efdd73 with: language: python runs_on: codebuild-github-actions-runner-${{ github.run_id }}-${{ github.run_attempt }} resource_prefix: p sdk_repository: aws/aws-durable-execution-sdk-python sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }} - conformance_test_ref: ${{ inputs.conformance_test_ref || '6182cee00ed33ce01e0579066fb31096f988817f' }} + conformance_test_ref: ${{ inputs.conformance_test_ref || 'fb95e20acb0b9b10f5fc032ac521a4fd45efdd73' }} # Check the SDK out so the handlers and templates below are on disk. The handlers # themselves are installed from sdk_ref by src/requirements.txt during the SAM build. checkout_sdk: true diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py index 7548631bc..96138ee49 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py @@ -442,7 +442,8 @@ def mark_state_delivered(self) -> None: delivers state in the invocation input. So the list is reset when that input is built, not when the invocation completes: an operation that completes while the handler is still running is - reported on the next invocation. + reported on the next invocation unless a checkpoint response has + already delivered it through ``advance_handler_seen``. """ self.updated_operation_ids = [] @@ -883,6 +884,13 @@ def advance_handler_seen(self, seq: int) -> None: smaller or equal values are ignored.""" if seq > self.execution.handler_seen_seq: self.execution.handler_seen_seq = seq + # A checkpoint has delivered these updates to the running handler. + # Retain only changes newer than that response for the next input. + self.execution.updated_operation_ids = [ + operation_id + for operation_id in self.execution.updated_operation_ids + if self.execution.operation_last_touched_seq.get(operation_id, 0) > seq + ] # --- internals ------------------------------------------------- diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/e2e/callback_updated_operations_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/callback_updated_operations_test.py index b58118894..cae94c294 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/e2e/callback_updated_operations_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/callback_updated_operations_test.py @@ -5,6 +5,8 @@ import json import time from pathlib import Path +from queue import Queue +from threading import Event from typing import Any import pytest @@ -159,3 +161,93 @@ def handler(_event: Any, context: DurableContext) -> str: if outcome == "failure": assert terminal.error is not None assert terminal.error.message == "explicit callback failure" + + +@pytest.mark.parametrize("outcome", ["success", "failure", "timeout"]) +@pytest.mark.parametrize("filesystem", [False, True]) +def test_callback_delivered_by_checkpoint_is_not_updated_on_later_resume( + outcome: str, filesystem: bool, tmp_path: Path +) -> None: + observer = CallbackObserver() + submitter_entered = Event() + release_submitter = Event() + submission_calls: list[str] = [] + published_callbacks: Queue[str] = Queue() + + @durable_step + def submit(_context: StepContext, callback_id: str) -> str: + submission_calls.append(callback_id) + published_callbacks.put(callback_id) + submitter_entered.set() + assert release_submitter.wait(10) + return "submitted" + + def handler(_event: Any, context: DurableContext) -> str: + callback = context.create_callback( + name="target", + config=CallbackConfig( + timeout=Duration.from_seconds(1 if outcome == "timeout" else 30), + serdes=JsonSerDes(), + ), + ) + context.step(submit(callback.callback_id), name="submit") + try: + result = callback.result() + except CallbackError: + result = outcome + context.wait(Duration.from_seconds(1), name="first-replay") + context.wait(Duration.from_seconds(1), name="second-replay") + assert isinstance(result, str) + return result + + store = FileSystemExecutionStore.create(tmp_path) if filesystem else None + with DurableFunctionTestRunner( + handler=durable_execution(handler, plugins=[observer]), + store=store, + skip_time=False, + poll_interval=0.01, + execution_timeout=25, + ) as runner: + arn = runner.run_async(input="{}") + try: + assert submitter_entered.wait(5) + callback_id = published_callbacks.get(timeout=5) + if outcome == "success": + runner.send_callback_success(callback_id, result=b'"target"') + elif outcome == "failure": + runner.send_callback_failure( + callback_id, error=ErrorObject.from_message("early failure") + ) + expected_event = { + "success": "CallbackSucceeded", + "failure": "CallbackFailed", + "timeout": "CallbackTimedOut", + }[outcome] + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + events = runner.get_execution_history( + arn, include_execution_data=True + ).events + if any(event.event_type == expected_event for event in events): + break + time.sleep(0.01) + else: + raise AssertionError("Callback did not complete during submission") + assert not any( + event.event_type == "InvocationCompleted" for event in events + ) + finally: + release_submitter.set() + result = runner.wait_for_result(arn, timeout=15) + + assert result.status.value == "SUCCEEDED" + assert result.result is not None + assert json.loads(result.result) == ("target" if outcome == "success" else outcome) + assert len(submission_calls) == 1 + assert len(observer.invocations) >= 3 + assert len(observer.ends) == 1 + target_id = observer.ends[0].operation_id + assert all( + target_id not in invocation.updated_operations + for invocation in observer.invocations[1:] + ) diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/execution_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/execution_test.py index f3ac97185..353d320d4 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/execution_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/execution_test.py @@ -3,6 +3,7 @@ import json from dataclasses import replace from datetime import datetime, timezone +from threading import Event, Thread from unittest.mock import patch, Mock import pytest @@ -1722,6 +1723,61 @@ def test_record_invocation_completion_keeps_updated_operation_ids(): assert execution.updated_operation_ids == [] +@pytest.mark.parametrize("complete_before_advance", [False, True]) +def test_checkpoint_consumes_only_updates_covered_by_its_watermark( + complete_before_advance, +): + """A late update survives reads and retries of an older state delivery.""" + execution = Execution( + "test-arn", + _make_start_input(), + [ + Operation( + operation_id=name, + operation_type=OperationType.CALLBACK, + status=OperationStatus.STARTED, + callback_details=CallbackDetails(callback_id=name), + ) + for name in ["delivered", "later"] + ], + ) + execution.complete_callback_success("delivered", b"first") + response = OperationPaginatorState.pin(execution) + ready = Event() + finished = Event() + + def complete_later(): + assert ready.wait(5) + execution.complete_callback_failure("later", ErrorObject.from_message("later")) + finished.set() + + worker = Thread(target=complete_later) + worker.start() + if complete_before_advance: + ready.set() + assert finished.wait(5) + expected = ["delivered", "later"] if complete_before_advance else ["delivered"] + assert execution.updated_operation_ids == expected + response.page(None, max_size_bytes=1024 * 1024) + assert execution.updated_operation_ids == expected + + response.advance_handler_seen(1) + ready.set() + assert finished.wait(5) + worker.join(timeout=5) + assert not worker.is_alive() + assert execution.updated_operation_ids == ["later"] + assert execution.handler_seen_seq == 1 + assert execution.token_sequence == 0 + assert execution.seq_counter == 2 + # An idempotent/older delivery cannot consume the later completion. + response.advance_handler_seen(1) + response.advance_handler_seen(0) + assert execution.updated_operation_ids == ["later"] + OperationPaginatorState.pin(execution).advance_handler_seen(2) + assert execution.updated_operation_ids == [] + + def test_function_arn_is_qualified_with_the_executed_version(): execution = Execution.new(_make_start_input()) execution.region = "ap-southeast-2" diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/executor_checkpoint_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/executor_checkpoint_test.py index c789dcd74..fb21cf68f 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/executor_checkpoint_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/executor_checkpoint_test.py @@ -17,6 +17,7 @@ import pytest from aws_durable_execution_sdk_python.execution import InvocationStatus from aws_durable_execution_sdk_python.lambda_service import ( + CallbackDetails, ErrorObject, Operation as SvcOperation, OperationAction, @@ -126,6 +127,49 @@ def test_empty_poll_returns_empty_operations_and_advances_token(): assert CheckpointToken.from_str(response.checkpoint_token).token_sequence == 1 +@pytest.mark.parametrize("rejection", ["token", "operation"]) +def test_rejected_checkpoint_retains_undelivered_callback_update(rejection): + executor, store, execution, token = _make_executor_with_started_execution() + execution.operations.append( + SvcOperation( + operation_id="callback", + operation_type=OperationType.CALLBACK, + status=OperationStatus.STARTED, + callback_details=CallbackDetails(callback_id="callback-id"), + ) + ) + execution.complete_callback_success("callback-id", b"result") + store.save(execution) + before = ( + execution.token_sequence, + execution.handler_seen_seq, + execution.seq_counter, + ) + assert execution.updated_operation_ids == ["callback"] + with pytest.raises(InvalidParameterValueException): + executor.checkpoint_execution( + execution_arn=execution.durable_execution_arn, + checkpoint_token="invalid" if rejection == "token" else token, + updates=( + [ + OperationUpdate( + operation_id="callback", + operation_type=OperationType.CALLBACK, + action=OperationAction.SUCCEED, + ) + ] + if rejection == "operation" + else [] + ), + ) + assert execution.updated_operation_ids == ["callback"] + assert ( + execution.token_sequence, + execution.handler_seen_seq, + execution.seq_counter, + ) == before + + # endregion # region: Non-empty checkpoint returns only the delta From 3bacb932a6e5a3543082817d71b8c14b1a967f56 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 01:18:57 +0000 Subject: [PATCH 23/54] test: pin strict invocation parent conformance --- .../tests/test_opentelemetry_conformance_workflow.py | 4 ++-- .github/workflows/opentelemetry-conformance-tests.yml | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py index 775348d07..a11bcce6b 100644 --- a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py +++ b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py @@ -4,8 +4,8 @@ import yaml -SHARED_WORKFLOW_REF = "fb95e20acb0b9b10f5fc032ac521a4fd45efdd73" -CONFORMANCE_TEST_REF = "fb95e20acb0b9b10f5fc032ac521a4fd45efdd73" +SHARED_WORKFLOW_REF = "ad35bd36b67821d52a17db4f9d80d8696694a7f3" +CONFORMANCE_TEST_REF = "ad35bd36b67821d52a17db4f9d80d8696694a7f3" WORKFLOW_PATH = ( Path(__file__).parents[2] / "workflows" / "opentelemetry-conformance-tests.yml" diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index 1bee20cf6..919ae44e8 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -45,7 +45,7 @@ on: conformance_test_ref: description: Conformance test commit SHA or branch name required: true - default: fb95e20acb0b9b10f5fc032ac521a4fd45efdd73 + default: ad35bd36b67821d52a17db4f9d80d8696694a7f3 type: string # Backend stacks are shared across PRs. Queue whole runs so reusable @@ -66,14 +66,14 @@ jobs: actions: write contents: read id-token: write - uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@fb95e20acb0b9b10f5fc032ac521a4fd45efdd73 + uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@ad35bd36b67821d52a17db4f9d80d8696694a7f3 with: language: python runs_on: codebuild-github-actions-runner-${{ github.run_id }}-${{ github.run_attempt }} resource_prefix: p sdk_repository: aws/aws-durable-execution-sdk-python sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }} - conformance_test_ref: ${{ inputs.conformance_test_ref || 'fb95e20acb0b9b10f5fc032ac521a4fd45efdd73' }} + conformance_test_ref: ${{ inputs.conformance_test_ref || 'ad35bd36b67821d52a17db4f9d80d8696694a7f3' }} # Check the SDK out so the handlers and templates below are on disk. The handlers # themselves are installed from sdk_ref by src/requirements.txt during the SAM build. checkout_sdk: true From 22004eb5f496f0e8acd9912854dcd1756ee0d386 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 01:41:23 +0000 Subject: [PATCH 24/54] fix: align empty callback errors across stores Represent an exactly empty serialized callback error as absent, matching CallbackDetails.from_dict without requiring a file-store round trip. Preserve every present field, including empty strings and stack lists. Compare real public handler failures across memory/file stores and both OTel views, including caller error/status and replay behavior. Pin the coordinated shared S1c revision with the strict empty-payload, parent-occurrence and callback phase-gate assertions. --- ...test_opentelemetry_conformance_workflow.py | 4 +- .../opentelemetry-conformance-tests.yml | 6 +- .../README.md | 4 + .../test_callback_failure_without_error.py | 208 ++++++++++++++++++ .../execution.py | 5 +- 5 files changed, 221 insertions(+), 6 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_failure_without_error.py diff --git a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py index a11bcce6b..e32684716 100644 --- a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py +++ b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py @@ -4,8 +4,8 @@ import yaml -SHARED_WORKFLOW_REF = "ad35bd36b67821d52a17db4f9d80d8696694a7f3" -CONFORMANCE_TEST_REF = "ad35bd36b67821d52a17db4f9d80d8696694a7f3" +SHARED_WORKFLOW_REF = "bdb4f1cd0f9252c1aaa978bb8b341b71f2b9d9dc" +CONFORMANCE_TEST_REF = "bdb4f1cd0f9252c1aaa978bb8b341b71f2b9d9dc" WORKFLOW_PATH = ( Path(__file__).parents[2] / "workflows" / "opentelemetry-conformance-tests.yml" diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index 919ae44e8..ecf2e0f3f 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -45,7 +45,7 @@ on: conformance_test_ref: description: Conformance test commit SHA or branch name required: true - default: ad35bd36b67821d52a17db4f9d80d8696694a7f3 + default: bdb4f1cd0f9252c1aaa978bb8b341b71f2b9d9dc type: string # Backend stacks are shared across PRs. Queue whole runs so reusable @@ -66,14 +66,14 @@ jobs: actions: write contents: read id-token: write - uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@ad35bd36b67821d52a17db4f9d80d8696694a7f3 + uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@bdb4f1cd0f9252c1aaa978bb8b341b71f2b9d9dc with: language: python runs_on: codebuild-github-actions-runner-${{ github.run_id }}-${{ github.run_attempt }} resource_prefix: p sdk_repository: aws/aws-durable-execution-sdk-python sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }} - conformance_test_ref: ${{ inputs.conformance_test_ref || 'ad35bd36b67821d52a17db4f9d80d8696694a7f3' }} + conformance_test_ref: ${{ inputs.conformance_test_ref || 'bdb4f1cd0f9252c1aaa978bb8b341b71f2b9d9dc' }} # Check the SDK out so the handlers and templates below are on disk. The handlers # themselves are installed from sdk_ref by src/requirements.txt during the SAM build. checkout_sdk: true diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md index 0b0ea8b2a..22dbef8fb 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md @@ -112,6 +112,10 @@ checkpoint response carry the same completion. It tracks actual notifications, preserves first delivery when the update-ID metadata is absent, and clears that tracking at invocation boundaries. Public runner regressions cover memory and file stores, stored step results, failure payloads and two subsequent replays. +An exactly empty serialized callback error is represented as absent in both +stores, matching the SDK's service parser. Present fields remain intact, +including empty messages, types, data and stack lists; the enclosing failed +callback future still raises the same error as the file-store baseline. Case 24 covers invocation `RETRY` becoming `RETRYING`/`UNSET`. Case 25 targets the separate rule that a failed operation without error details remains `UNSET`. diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_failure_without_error.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_failure_without_error.py new file mode 100644 index 000000000..2eea016b3 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_failure_without_error.py @@ -0,0 +1,208 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Public callback failures preserve error-detail semantics across stores.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +import time +from pathlib import Path + +import pytest +from aws_durable_execution_sdk_python.lambda_service import ErrorObject +from aws_durable_execution_sdk_python.plugin import OperationEndInfo, OperationType +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from aws_durable_execution_sdk_python_testing.stores.filesystem import ( + FileSystemExecutionStore, +) +from opentelemetry import context +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +SRC_DIR = Path(__file__).resolve().parents[1] / "src" + + +def _run_public_callback_failure( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + plugin_type: type[ExecutionOtelPlugin] | type[InvocationOtelPlugin], + filesystem: bool, + error_kind: str, +) -> dict[str, object]: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = plugin_type( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + callback_errors: list[ErrorObject | None] = [] + original_end = plugin.on_operation_end + + def observe_error(info: OperationEndInfo) -> None: + if info.operation_type is OperationType.CALLBACK: + callback_errors.append(info.error) + original_end(info) + + monkeypatch.setattr(plugin, "on_operation_end", observe_error) + original_context = context.get_current() + common_spec = importlib.util.spec_from_file_location( + "common", SRC_DIR / "common.py" + ) + assert common_spec is not None and common_spec.loader is not None + common = importlib.util.module_from_spec(common_spec) + monkeypatch.setitem(sys.modules, "common", common) + common_spec.loader.exec_module(common) + monkeypatch.setattr(common, "otel_plugin", lambda: plugin) + spec = importlib.util.spec_from_file_location( + "otel_17_wait_for_callback_failure", + SRC_DIR / "otel_17_wait_for_callback_failure.py", + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + error = { + "omitted": None, + "empty": ErrorObject(message=None, type=None, data=None, stack_trace=None), + "rich": ErrorObject( + message="explicit failure", + type="CallbackFailure", + data=None, + stack_trace=None, + ), + "empty-message": ErrorObject( + message="", type=None, data=None, stack_trace=None + ), + "empty-stack": ErrorObject(message=None, type=None, data=None, stack_trace=[]), + "empty-type": ErrorObject(message=None, type="", data=None, stack_trace=None), + "empty-data": ErrorObject(message=None, type=None, data="", stack_trace=None), + }[error_kind] + try: + store = FileSystemExecutionStore.create(tmp_path) if filesystem else None + with DurableFunctionTestRunner( + handler=module.handler, + store=store, + poll_interval=0.01, + execution_timeout=20, + skip_time=False, + ) as runner: + arn = runner.run_async( + input=json.dumps({"scenario": "wait-for-callback-failure"}) + ) + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + events = runner.get_execution_history( + arn, include_execution_data=True + ).events + starts = [ + event for event in events if event.event_type == "CallbackStarted" + ] + if starts and any( + event.event_type == "InvocationCompleted" + and event.event_id > starts[0].event_id + for event in events + ): + details = starts[0].callback_started_details + assert details is not None and details.callback_id is not None + callback_id = details.callback_id + break + time.sleep(0.01) + else: + raise AssertionError("Callback did not suspend before failure") + if error_kind == "omitted": + runner.send_callback_failure(callback_id) + else: + runner.send_callback_failure(callback_id, error=error) + result = runner.wait_for_result(arn, timeout=10) + + assert result.status.value == "FAILED" + spans = exporter.get_finished_spans() + leaves = [ + span + for span in spans + if (span.attributes or {}).get("durable.operation.type") == "CALLBACK" + and (span.attributes or {}).get("durable.operation.status") == "FAILED" + ] + assert len(leaves) == 1 + has_details = error_kind not in {"omitted", "empty"} + assert [ + item.to_dict() if item is not None else None for item in callback_errors + ] == [error.to_dict() if has_details and error is not None else None] + if has_details and not filesystem: + assert callback_errors[0] is error + assert leaves[0].status.status_code.name == ( + "ERROR" if has_details else "UNSET" + ) + assert [event.name for event in leaves[0].events] == ( + ["exception"] if has_details else [] + ) + # The public failed future still raises, so its enclosing context fails. + parents = [ + span + for span in spans + if span.name == "otel-failed-callback" + and (span.attributes or {}).get("durable.operation.status") == "FAILED" + ] + assert len(parents) == 1 + assert parents[0].status.status_code.name == "ERROR" + invocations = sorted( + [span for span in spans if span.name == "Invocation"], + key=lambda span: span.start_time or 0, + ) + statuses = [ + (span.attributes or {}).get("durable.invocation.status") + for span in invocations + ] + assert statuses == ["PENDING", "FAILED"] + assert result.error is not None + assert context.get_current() == original_context + return { + "status": result.status.value, + "caller_error": dict(result.error.to_dict()), + "invocation_statuses": statuses, + } + finally: + provider.shutdown() + + +@pytest.mark.parametrize("plugin_type", [ExecutionOtelPlugin, InvocationOtelPlugin]) +@pytest.mark.parametrize( + "error_kind", + [ + "omitted", + "empty", + "rich", + "empty-message", + "empty-stack", + "empty-type", + "empty-data", + ], +) +def test_public_callback_failure_preserves_error_details( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + plugin_type: type[ExecutionOtelPlugin] | type[InvocationOtelPlugin], + error_kind: str, +) -> None: + memory = _run_public_callback_failure( + monkeypatch, tmp_path / "memory", plugin_type, False, error_kind + ) + filesystem = _run_public_callback_failure( + monkeypatch, tmp_path / "filesystem", plugin_type, True, error_kind + ) + assert memory == filesystem diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py index 96138ee49..2ab02fdce 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py @@ -665,8 +665,11 @@ def complete_callback_failure( self.touch_operation(operation.operation_id) updated_callback_details = None if operation.callback_details: + # Match CallbackDetails.from_dict without depending on a store + # serialization round trip: an empty wire Error has no details. updated_callback_details = replace( - operation.callback_details, error=error + operation.callback_details, + error=error if error is not None and error.to_dict() else None, ) self.operations[index] = replace( From cd81a3d6d3f63f7d083aa3df81dfb46500a90272 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 02:08:42 +0000 Subject: [PATCH 25/54] fix: preserve empty callback error history payload Match the observed service projection for a failed callback with no error details when IncludeExecutionData is true: keep Payload {} and mark it untruncated. SDK-facing state remains error=None, with an UNSET leaf and the same caller failure. Other statuses, nonempty errors and metadata-only history retain their existing representation. Verify the raw public history shape, state/telemetry separation and direct factory boundaries. The unchanged JS 19ee19f example suite passes all 148 tests without retries or assertion/selection changes. --- .../README.md | 3 + .../test_callback_failure_without_error.py | 31 ++++++++++ .../model.py | 8 +++ .../tests/event_factory_test.py | 61 +++++++++++++++++++ 4 files changed, 103 insertions(+) diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md index 22dbef8fb..421de8f64 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md @@ -116,6 +116,9 @@ An exactly empty serialized callback error is represented as absent in both stores, matching the SDK's service parser. Present fields remain intact, including empty messages, types, data and stack lists; the enclosing failed callback future still raises the same error as the file-store baseline. +Detailed callback-failure history retains the service's empty error `Payload` +object with `Truncated: false`, while SDK-facing state still has no error details. +Metadata-only history and nonempty errors retain their existing representation. Case 24 covers invocation `RETRY` becoming `RETRYING`/`UNSET`. Case 25 targets the separate rule that a failed operation without error details remains `UNSET`. diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_failure_without_error.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_failure_without_error.py index 2eea016b3..0da4ff4c6 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_failure_without_error.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_failure_without_error.py @@ -129,6 +129,10 @@ def observe_error(info: OperationEndInfo) -> None: else: runner.send_callback_failure(callback_id, error=error) result = runner.wait_for_result(arn, timeout=10) + with_data = runner.get_execution_history(arn, include_execution_data=True) + without_data = runner.get_execution_history( + arn, include_execution_data=False + ) assert result.status.value == "FAILED" spans = exporter.get_finished_spans() @@ -175,6 +179,17 @@ def observe_error(info: OperationEndInfo) -> None: "status": result.status.value, "caller_error": dict(result.error.to_dict()), "invocation_statuses": statuses, + "submitted_error": dict(error.to_dict()) if error is not None else {}, + "history_error": next( + event.to_dict()["CallbackFailedDetails"]["Error"] + for event in with_data.events + if event.event_type == "CallbackFailed" + ), + "metadata_error": next( + event.to_dict()["CallbackFailedDetails"]["Error"] + for event in without_data.events + if event.event_type == "CallbackFailed" + ), } finally: provider.shutdown() @@ -206,3 +221,19 @@ def test_public_callback_failure_preserves_error_details( monkeypatch, tmp_path / "filesystem", plugin_type, True, error_kind ) assert memory == filesystem + expected_payload = memory["submitted_error"] + for outcome in [memory, filesystem]: + # Actual AWS history retains an empty Payload object for this failure, + # independently of the SDK-facing absent callback error. + # Retain the existing flags for nonempty errors; only the observed + # no-details projection is changed here. + assert outcome["history_error"] == { + "Payload": expected_payload, + "Truncated": bool(expected_payload), + } + # Preserve this API's existing metadata-only projection. + assert outcome["metadata_error"] == ( + {"Payload": expected_payload, "Truncated": True} + if expected_payload + else {"Truncated": True} + ) diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/model.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/model.py index 43026ec60..cb99eeab1 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/model.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/model.py @@ -2239,6 +2239,14 @@ def create_callback_event_failed(cls, context: EventCreationContext) -> Event: event_error: EventError | None = ( EventError.from_details(callback_details) if callback_details else None ) + if ( + context.include_execution_data + and callback_details is not None + and callback_details.error is None + ): + # Detailed service history retains an empty Error.Payload object. + # This projection must not turn the SDK-facing absent error into one. + event_error = EventError(payload=ErrorObject.from_dict({}), truncated=False) return cls( event_type=EventType.CALLBACK_FAILED.value, event_timestamp=context.end_timestamp, diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/event_factory_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/event_factory_test.py index cd2fd8743..04334b6db 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/event_factory_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/event_factory_test.py @@ -9,8 +9,10 @@ import pytest from aws_durable_execution_sdk_python.lambda_service import ( + CallbackDetails, ChainedInvokeOptions, ErrorObject, + Operation as ServiceOperation, OperationStatus, OperationType, StepDetails, @@ -849,6 +851,65 @@ def test_create_callback_failed(): assert event.callback_failed_details.error.payload.message == "Callback failed" +@pytest.mark.parametrize("include_data", [False, True]) +@pytest.mark.parametrize("has_error", [False, True]) +def test_callback_failure_history_projection_preserves_sdk_state( + include_data, has_error +): + error = ErrorObject.from_message("details") if has_error else None + operation = ServiceOperation( + operation_id="callback", + operation_type=OperationType.CALLBACK, + status=OperationStatus.FAILED, + callback_details=CallbackDetails(callback_id="callback-id", error=error), + ) + context = EventCreationContext.create( + operation=operation, + event_id=3, + durable_execution_arn="arn:test", + start_input=StartDurableExecutionInput( + account_id="123", + function_name="test", + function_qualifier="$LATEST", + execution_name="test", + execution_timeout_seconds=300, + execution_retention_period_days=7, + ), + include_execution_data=include_data, + ) + event = Event.create_callback_event(context) + expected = {"Truncated": not (include_data and not has_error)} + if has_error or include_data: + expected["Payload"] = error.to_dict() if error else {} + assert event.callback_failed_details.error.to_dict() == expected + assert operation.callback_details.error is error + + +def test_callback_timeout_history_projection_is_unchanged(): + operation = ServiceOperation( + operation_id="callback", + operation_type=OperationType.CALLBACK, + status=OperationStatus.TIMED_OUT, + callback_details=CallbackDetails(callback_id="callback-id", error=None), + ) + context = EventCreationContext.create( + operation=operation, + event_id=3, + durable_execution_arn="arn:test", + start_input=StartDurableExecutionInput( + account_id="123", + function_name="test", + function_qualifier="$LATEST", + execution_name="test", + execution_timeout_seconds=300, + execution_retention_period_days=7, + ), + include_execution_data=True, + ) + event = Event.create_callback_event(context) + assert event.callback_timed_out_details.error.to_dict() == {"Truncated": True} + + def test_create_callback_timed_out(): operation = create_mock_operation("callback-1", status=OperationStatus.TIMED_OUT) error_obj = ErrorObject.from_message("Callback timed out") From 2a100531ed4a8b631dc6a5844b95927b3d2d4f08 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 02:44:35 +0000 Subject: [PATCH 26/54] fix(testing): preserve empty history payloads --- .../test_callback_failure_without_error.py | 31 ++++++++ .../model.py | 7 +- .../tests/model_test.py | 30 ++++++++ .../tests/runner_test.py | 76 +++++++++++++++++++ 4 files changed, 142 insertions(+), 2 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_failure_without_error.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_failure_without_error.py index 0da4ff4c6..a04a92c0d 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_failure_without_error.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_failure_without_error.py @@ -14,6 +14,10 @@ import pytest from aws_durable_execution_sdk_python.lambda_service import ErrorObject from aws_durable_execution_sdk_python.plugin import OperationEndInfo, OperationType +from aws_durable_execution_sdk_python_testing.model import ( + GetDurableExecutionHistoryResponse, + events_to_operations, +) from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner from aws_durable_execution_sdk_python_testing.stores.filesystem import ( FileSystemExecutionStore, @@ -175,6 +179,33 @@ def observe_error(info: OperationEndInfo) -> None: assert statuses == ["PENDING", "FAILED"] assert result.error is not None assert context.get_current() == original_context + for history in (with_data, without_data): + wire = history.to_dict() + decoded = GetDurableExecutionHistoryResponse.from_dict(wire) + original_error = next( + event["CallbackFailedDetails"]["Error"] + for event in wire["Events"] + if event["EventType"] == "CallbackFailed" + ) + for candidate in (history, decoded): + assert ( + next( + event.to_dict()["CallbackFailedDetails"]["Error"] + for event in candidate.events + if event.event_type == "CallbackFailed" + ) + == original_error + ) + callback = next( + operation + for operation in events_to_operations(candidate.events) + if operation.callback_details is not None + ) + details = callback.callback_details + assert details is not None + assert ( + details.error.to_dict() if details.error is not None else None + ) == (error.to_dict() if has_details and error is not None else None) return { "status": result.status.value, "caller_error": dict(result.error.to_dict()), diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/model.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/model.py index cb99eeab1..fdccb2cff 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/model.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/model.py @@ -677,7 +677,8 @@ class EventError: @classmethod def from_dict(cls, data: dict) -> EventError: payload = None - if payload_data := data.get("Payload"): + payload_data = data.get("Payload") + if payload_data is not None: payload = ErrorObject.from_dict(payload_data) return cls( @@ -2752,7 +2753,9 @@ def events_to_operations(events: list[Event]) -> list[Operation]: callback_details=CallbackDetails( callback_id=callback_id, result=result, - error=error, + # History preserves a present empty Payload object, while + # CallbackDetails uses None for an empty wire error. + error=error if error is not None and error.to_dict() else None, ), ) diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/model_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/model_test.py index 746f083d6..f3f8a9131 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/model_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/model_test.py @@ -1461,6 +1461,36 @@ def test_event_error_with_payload_only(): } +@pytest.mark.parametrize( + "payload", + [ + {}, + {"ErrorMessage": ""}, + {"ErrorType": ""}, + {"ErrorData": ""}, + {"StackTrace": []}, + ], +) +def test_event_error_roundtrip_preserves_present_payload(payload): + wire = {"Payload": payload, "Truncated": False} + assert EventError.from_dict(wire).to_dict() == wire + + +@pytest.mark.parametrize("truncated", [False, True]) +@pytest.mark.parametrize("payload_field", [{}, {"Payload": None}]) +def test_event_error_absent_or_null_payload_retains_existing_meaning( + payload_field, truncated +): + parsed = EventError.from_dict({**payload_field, "Truncated": truncated}) + assert parsed.payload is None + assert parsed.to_dict() == {"Truncated": truncated} + + +def test_event_error_empty_payload_retains_truncation_flag(): + wire = {"Payload": {}, "Truncated": True} + assert EventError.from_dict(wire).to_dict() == wire + + # Tests for RetryDetails def test_retry_details_serialization(): """Test RetryDetails from_dict/to_dict round-trip.""" diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/runner_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/runner_test.py index ab75f0d24..0662c07f1 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/runner_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/runner_test.py @@ -1116,6 +1116,82 @@ def test_durable_child_context_test_runner_init_with_args( # Tests for DurableFunctionCloudTestRunner and from_execution_history +@pytest.mark.parametrize("decode_wire", [False, True]) +@pytest.mark.parametrize( + "payload", + [ + {}, + {"ErrorMessage": ""}, + {"ErrorType": ""}, + {"ErrorData": ""}, + {"StackTrace": []}, + {"ErrorMessage": "failure"}, + ], +) +def test_callback_history_result_preserves_wire_and_canonical_error( + payload, decode_wire +): + from aws_durable_execution_sdk_python.lambda_service import ErrorObject + from aws_durable_execution_sdk_python_testing.model import ( + CallbackFailedDetails, + CallbackStartedDetails, + Event, + EventError, + GetDurableExecutionResponse, + ) + + timestamp = datetime.datetime(2026, 10, 7, tzinfo=datetime.UTC) + caller_error = ErrorObject.from_message("Callback failed") + execution = GetDurableExecutionResponse( + durable_execution_arn="arn:execution", + durable_execution_name="execution", + function_arn="arn:function", + status="FAILED", + start_timestamp=timestamp, + error=caller_error, + ) + history = GetDurableExecutionHistoryResponse( + events=[ + Event( + event_type="CallbackStarted", + event_timestamp=timestamp, + event_id=1, + operation_id="callback", + name="callback", + callback_started_details=CallbackStartedDetails( + callback_id="callback-id" + ), + ), + Event( + event_type="CallbackFailed", + event_timestamp=timestamp, + event_id=2, + operation_id="callback", + name="callback", + callback_failed_details=CallbackFailedDetails( + error=EventError( + payload=ErrorObject.from_dict(payload), truncated=False + ) + ), + ), + ] + ) + if decode_wire: + history = GetDurableExecutionHistoryResponse.from_dict(history.to_dict()) + assert history.events[-1].callback_failed_details.error.to_dict() == { + "Payload": payload, + "Truncated": False, + } + result = DurableFunctionTestResult.from_execution_history(execution, history) + callback = result.get_callback("callback") + assert callback.status is OperationStatus.FAILED + assert (callback.error.to_dict() if callback.error is not None else None) == ( + payload if payload else None + ) + assert result.status is InvocationStatus.FAILED + assert result.error is caller_error + + def test_durable_function_test_result_from_execution_history(): """Test DurableFunctionTestResult.from_execution_history factory method.""" import datetime From a6c8c904db8df57df8f9323dd7d9528440817559 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 20:41:01 +0000 Subject: [PATCH 27/54] refactor: simplify optional handler context scopes --- .github/workflows/ci.yml | 3 +- .../README.md | 10 +- .../execution_plugin.py | 1 - .../invocation_plugin.py | 2 - .../e2e/test_invocation_wait_resume_int.py | 23 +-- .../README.md | 30 ++- .../plugin.py | 25 +-- .../tests/handler_worker_context_test.py | 173 +++++++----------- 8 files changed, 97 insertions(+), 170 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 41b5ff5b2..cb945d8d0 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -95,8 +95,7 @@ jobs: assert version("aws-durable-execution-sdk-python") == "2.0.0" assert "site-packages" in Path(plugin.__file__).parts - assert not hasattr(plugin.DurableInstrumentationPlugin, "handler_context") - assert not hasattr(plugin, "DURABLE_INSTRUMENTATION_HANDLER_CONTEXT_API_VERSION") + assert not hasattr(plugin.PluginExecutor, "run_handler") PYTHON hatch run test-pypi-otel-minimum:test - name: Build distribution diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 2e69cd01d..fcc0f77e9 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -190,11 +190,11 @@ span current only while the handler runs. The scope closes on the same worker in reverse plugin order, including on failure and suspension, without changing the invocation-hook caller. Older cores ignore this optional scope and retain their existing behavior; install the updated core as well to get handler context propagation. -The bundled OTel classes explicitly opt in with `__durable_handler_context_api__ = 1`. -Custom subclasses must repeat that literal marker on their own concrete class to -use the scope; inherited or instance markers are ignored. Unopted legacy helpers -and properties with the same name are never inspected. Older cores ignore the -marker without importing any new core API. +The bundled OTel classes provide this optional method directly; custom subclasses +inherit it normally and can override it using standard Python method resolution. +The updated core invokes a callable `handler_context` attribute when present; +there is no separate plugin opt-in marker. Older cores ignore the method without +importing any new core API. Execution view similarly restores the Workflow span inside the handler scope if another invocation-start hook clears the active span or switches to an unrelated trace. Both views retain valid same-trace parents and baggage, and restore the diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py index b05a72a5e..9c384ed62 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py @@ -129,7 +129,6 @@ class ExecutionOtelPlugin(DurableInstrumentationPlugin): span). """ - __durable_handler_context_api__ = 1 __durable_registration_api__: ClassVar[int] = 1 exclusive_group: ClassVar[str | None] = "aws-durable-execution-otel-view" diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py index 73e4ac7c5..68c2893a4 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py @@ -110,8 +110,6 @@ class InvocationOtelPlugin(DurableInstrumentationPlugin): provider installed by the ADOT Lambda layer). """ - __durable_handler_context_api__ = 1 - DEFAULT_INSTRUMENT_NAME = "aws-durable-execution-sdk-python" __durable_registration_api__: ClassVar[int] = 1 diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py index 427854750..2e95429a7 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py @@ -261,6 +261,14 @@ def handler_impl(_event: Any, context: DurableContext) -> str: assert completed_wait_span.end_time <= after_resume.start_time +class InheritedInvocationPlugin(InvocationOtelPlugin): + pass + + +class InheritedExecutionPlugin(ExecutionOtelPlugin): + pass + + @pytest.mark.parametrize( ("plugin_type", "extra_context_plugin"), [ @@ -274,10 +282,8 @@ def handler_impl(_event: Any, context: DurableContext) -> str: + ( [(ExecutionOtelPlugin, True)] + [(ExecutionOtelPlugin, kind) for kind in ("same", "unrelated", "absent")] - if getattr( - core_plugin_api, "DURABLE_INSTRUMENTATION_HANDLER_CONTEXT_API_VERSION", None - ) - == 1 + + [(InheritedInvocationPlugin, False), (InheritedExecutionPlugin, "absent")] + if hasattr(core_plugin_api.PluginExecutor, "run_handler") else [] ), ) @@ -297,12 +303,7 @@ def test_handler_user_spans_inherit_context_across_resume_and_failure( # The documented PyPI compatibility environment deliberately uses an older # core. Keep exercising its supported operation tracing and lifecycle while # asserting the new handler contract only when that core exposes the scope. - supports_handler_context = ( - getattr( - core_plugin_api, "DURABLE_INSTRUMENTATION_HANDLER_CONTEXT_API_VERSION", None - ) - == 1 - ) + supports_handler_context = hasattr(core_plugin_api.PluginExecutor, "run_handler") exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) @@ -457,7 +458,7 @@ def on_invocation_end(self, _info: Any) -> None: else derive_workflow_span_id(EXECUTION_ARN) ] * 2 - if plugin_type is ExecutionOtelPlugin + if issubclass(plugin_type, ExecutionOtelPlugin) else ambient_ids if ambient_kind == "same" else [ diff --git a/packages/aws-durable-execution-sdk-python/README.md b/packages/aws-durable-execution-sdk-python/README.md index 8429e196e..bfe54384f 100644 --- a/packages/aws-durable-execution-sdk-python/README.md +++ b/packages/aws-durable-execution-sdk-python/README.md @@ -114,22 +114,20 @@ lifecycle order are unchanged. ### Optional handler context scopes -A plugin can declare `__durable_handler_context_api__ = 1` directly on its -concrete class and implement `handler_context(info)` returning a context manager. -The updated core enters these scopes around the top-level handler on its worker -thread, in registration order, and closes them in reverse order. Cleanup receives -no handler exception and cannot suppress or replace its outcome. Invocation hooks -retain their original thread and order; failed setup bindings are discarded and -successful scope cleanup stays in the context that owns its tokens. - -The marker must be the literal integer `1`; instance and inherited markers do not -opt in. A subclass must redeclare the marker to adopt this new hook. Unopted legacy -helpers, properties and dynamic attributes named `handler_context` are untouched. -The generic plugin base supplies neither a marker nor a default method. Core -support is advertised by module constant -`aws_durable_execution_sdk_python.plugin.DURABLE_INSTRUMENTATION_HANDLER_CONTEXT_API_VERSION`. -Older cores ignore this optional API. The provider API version and dependency -requirements are unchanged. +A plugin can provide `handler_context(info)` returning a context manager. The +updated core looks up this optional method normally, so inherited methods work +without extra declarations. A missing or non-callable attribute is ignored. +The core enters these scopes around the top-level handler on its worker thread, +in registration order, and closes them in reverse order. Cleanup receives no +handler exception and cannot suppress or replace its outcome. Invocation hooks +retain their original thread and order; failed lookup, setup or entry bindings +are discarded, and successful scope cleanup stays in the context that owns its +tokens. + +`handler_context` is an optional plugin API name: callable methods or attributes +with that name are invoked. The generic plugin base does not require a default +method. Older cores ignore this optional API and retain their existing behavior. +The provider API version and dependency requirements are unchanged. ## 🚀 Quick Start diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py index f10b6bf5d..a74a1cc86 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 @@ -31,7 +31,6 @@ logger = logging.getLogger(__name__) DURABLE_INSTRUMENTATION_PLUGIN_API_VERSION = 1 -DURABLE_INSTRUMENTATION_HANDLER_CONTEXT_API_VERSION = 1 class InvocationStatus(Enum): @@ -463,18 +462,6 @@ class DurableInstrumentationPluginProvider: plugin_api_version: int -def _handler_context_api_enabled( - plugin_type: type[DurableInstrumentationPlugin], -) -> bool: - """Read only the concrete class namespace, bypassing metaclass descriptors.""" - namespace = type.__dict__["__dict__"].__get__(plugin_type, type(plugin_type)) - version = namespace.get("__durable_handler_context_api__") - return ( - type(version) is int - and version == DURABLE_INSTRUMENTATION_HANDLER_CONTEXT_API_VERSION - ) - - class PluginExecutor: def __init__(self, plugins: list[DurableInstrumentationPlugin] | None): self._plugins = plugins or [] @@ -578,18 +565,12 @@ def execute_plugins(self, info, sync): def _safe_handler_context( self, plugin: DurableInstrumentationPlugin, info: InvocationStartInfo ) -> Iterator[bool]: - # Old plugin objects may not inherit this core's new optional method. + # Plugins may omit this optional scope method. scope = None succeeded = True try: - # Old plugins may have an unrelated helper/property with this name. - # Never even inspect it unless this concrete class explicitly opts in. - factory = ( - getattr(plugin, "handler_context", None) - if _handler_context_api_enabled(type(plugin)) - else None - ) - if factory is not None: + factory = getattr(plugin, "handler_context", None) + if callable(factory): scope = factory(info) scope.__enter__() except Exception: diff --git a/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py index 215624ea3..c62377261 100644 --- a/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py @@ -155,8 +155,6 @@ def test_optional_handler_scopes_are_balanced_and_cannot_change_outcome( events: list[str] = [] class ScopePlugin(DurableInstrumentationPlugin): - __durable_handler_context_api__ = 1 - def __init__(self, name: str): self.name = name @@ -331,8 +329,6 @@ def __exit__(self, *_args: Any) -> None: self.plugin.reset("exit") class SetupPlugin(DurableInstrumentationPlugin): - __durable_handler_context_api__ = 1 - def __init__(self, name: str) -> None: self.name = name self.token: contextvars.Token[str] | None = None @@ -427,85 +423,21 @@ def invoke(_event: Any, _context: Any) -> dict[str, str]: assert all(str(error[1]) == "partial plugin setup" for error in errors) -@pytest.mark.parametrize("shape", ["helper", "property", "dynamic"]) -def test_unopted_legacy_handler_context_is_never_inspected(shape: str) -> None: - from contextlib import contextmanager - from collections.abc import Iterator - from datetime import UTC, datetime - from aws_durable_execution_sdk_python.plugin import PluginExecutor - - calls: list[str] = [] - - @contextmanager - def helper(_info: InvocationStartInfo) -> Iterator[None]: - calls.append("helper") - yield - - def property_getter(_self: Any) -> Any: - calls.append("property") - return helper - - def dynamic_getter(_self: Any, name: str) -> Any: - if name == "handler_context": - calls.append("dynamic") - return "legacy-business-value" - raise AttributeError(name) - - members_by_shape: dict[str, dict[str, Any]] = { - "helper": {"handler_context": staticmethod(helper)}, - "property": {"handler_context": property(property_getter)}, - "dynamic": {"__getattr__": dynamic_getter}, - } - plugin_type = type( - "Legacy", (DurableInstrumentationPlugin,), members_by_shape[shape] - ) - plugin = plugin_type() - if shape == "dynamic": - assert plugin.handler_context == "legacy-business-value" - calls.clear() - executor = PluginExecutor([plugin]) - with executor.run(): - executor.on_invocation_start("legacy-helper", True, datetime.now(UTC), None) - assert executor.run_handler(lambda: list(calls)) == [] - assert calls == [] - - -@pytest.mark.parametrize("marker", [None, 0, 2, True, "1", property(lambda _: 1)]) -def test_handler_scope_requires_literal_class_local_version(marker: Any) -> None: - from contextlib import nullcontext - from datetime import UTC, datetime - from aws_durable_execution_sdk_python.plugin import PluginExecutor - - calls: list[str] = [] - - def helper(_self: Any, _info: InvocationStartInfo) -> Any: - calls.append("scope") - return nullcontext() - - plugin_type = type( - "Legacy", - (DurableInstrumentationPlugin,), - {"__durable_handler_context_api__": marker, "handler_context": helper}, - ) - executor = PluginExecutor([plugin_type()]) - with executor.run(): - executor.on_invocation_start("invalid-marker", True, datetime.now(UTC), None) - assert executor.run_handler(lambda: "ok") == "ok" - assert calls == [] - - -def test_handler_scope_opt_in_is_not_inherited_or_taken_from_instance() -> None: +@pytest.mark.parametrize("inherited", [False, True]) +@pytest.mark.parametrize("fail_handler", [False, True]) +def test_handler_scope_uses_direct_or_inherited_optional_method( + inherited: bool, fail_handler: bool +) -> None: from contextlib import contextmanager from collections.abc import Iterator from datetime import UTC, datetime from aws_durable_execution_sdk_python.plugin import PluginExecutor - marker = contextvars.ContextVar("explicit-scope", default="outside") + marker = contextvars.ContextVar("inherited-scope", default="outside") events: list[str] = [] + error = RuntimeError("original handler failure") - class OptedPlugin(DurableInstrumentationPlugin): - __durable_handler_context_api__ = 1 - + class ScopePlugin(DurableInstrumentationPlugin): @contextmanager def handler_context(self, _info: InvocationStartInfo) -> Iterator[None]: events.append("enter") @@ -516,47 +448,66 @@ def handler_context(self, _info: InvocationStartInfo) -> Iterator[None]: marker.reset(token) events.append("exit") - class LegacySubclass(OptedPlugin): + class InheritedPlugin(ScopePlugin): pass - class ExplicitSubclass(OptedPlugin): - __durable_handler_context_api__ = 1 - - for plugin, expected in [ - (OptedPlugin(), "inside"), - (LegacySubclass(), "outside"), - (ExplicitSubclass(), "inside"), - ]: - # Assigning a marker to an instance cannot accidentally enable the hook. - plugin.__durable_handler_context_api__ = 1 - executor = PluginExecutor([plugin]) - with executor.run(): - executor.on_invocation_start("subclass", True, datetime.now(UTC), None) - assert executor.run_handler(marker.get) == expected - assert marker.get() == "outside" - assert events == ["enter", "exit", "enter", "exit"] - - -def test_handler_opt_in_does_not_trigger_legacy_metaclass_descriptors() -> None: + plugin = InheritedPlugin() if inherited else ScopePlugin() + executor = PluginExecutor([plugin]) + + def handler() -> str: + assert marker.get() == "inside" + events.append("body") + if fail_handler: + raise error + return "ok" + + with executor.run(): + for _ in range(2): + executor.on_invocation_start("inherited", True, datetime.now(UTC), None) + if fail_handler: + with pytest.raises(RuntimeError) as caught: + executor.run_handler(handler) + assert caught.value is error + else: + assert executor.run_handler(handler) == "ok" + assert marker.get() == "outside" + assert events == ["enter", "body", "exit"] * 2 + + +@pytest.mark.parametrize("value", [None, "not a method", 42]) +def test_noncallable_optional_handler_scope_is_ignored(value: Any) -> None: from datetime import UTC, datetime + from types import SimpleNamespace + from typing import cast from aws_durable_execution_sdk_python.plugin import PluginExecutor - reads: list[str] = [] + plugin = cast( + DurableInstrumentationPlugin, + SimpleNamespace(on_invocation_start=lambda info: None, handler_context=value), + ) + executor = PluginExecutor([plugin]) + with executor.run(): + executor.on_invocation_start("noncallable", True, datetime.now(UTC), None) + assert executor.run_handler(lambda: "unchanged") == "unchanged" - def namespace(_cls: Any) -> Any: - reads.append("metaclass-dict") - raise RuntimeError("legacy namespace") - def helper(_self: Any, _info: Any) -> Any: - reads.append("legacy-helper") - raise RuntimeError("legacy helper") +def test_handler_scope_lookup_failure_discards_partial_bindings( + caplog: pytest.LogCaptureFixture, +) -> None: + from datetime import UTC, datetime + from aws_durable_execution_sdk_python.plugin import PluginExecutor - meta = type("LegacyMeta", (type,), {"__dict__": property(namespace)}) - plugin_type = meta( - "Legacy", (DurableInstrumentationPlugin,), {"handler_context": helper} - ) - executor = PluginExecutor([plugin_type()]) + marker = contextvars.ContextVar("lookup-failure", default="outside") + + class BrokenPlugin(DurableInstrumentationPlugin): + @property + def handler_context(self) -> Any: + marker.set("partial") + raise ValueError("scope lookup failed") + + executor = PluginExecutor([BrokenPlugin()]) with executor.run(): - executor.on_invocation_start("metaclass", True, datetime.now(UTC), None) - assert executor.run_handler(lambda: "ok") == "ok" - assert reads == [] + executor.on_invocation_start("lookup", True, datetime.now(UTC), None) + assert executor.run_handler(marker.get) == "outside" + assert marker.get() == "outside" + assert "scope lookup failed" in caplog.text From 21be89ee3cf5e32f99bec3e4eea3b15a1cf12263 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Thu, 8 Oct 2026 23:29:03 +0000 Subject: [PATCH 28/54] fix(otel): unify invocation lifecycle on handler worker --- .github/scripts/install_otel_test_wheels.py | 101 ++ .../otel_lifecycle_compatibility_test.py | 162 +++ .github/workflows/ci.yml | 18 +- CONTRIBUTING.md | 30 +- .../README.md | 74 +- .../pyproject.toml | 2 +- .../execution_plugin.py | 20 - .../invocation_plugin.py | 38 +- .../e2e/test_invocation_wait_resume_int.py | 93 +- .../tests/test_execution_plugin.py | 77 +- .../tests/test_invocation_plugin.py | 130 +-- .../tests/test_package_metadata.py | 69 +- .../README.md | 49 +- .../execution.py | 129 ++- .../plugin.py | 96 +- .../tests/handler_worker_context_test.py | 944 ++++++++++-------- pyproject.toml | 66 +- 17 files changed, 1156 insertions(+), 942 deletions(-) create mode 100644 .github/scripts/install_otel_test_wheels.py create mode 100644 .github/tests/otel_lifecycle_compatibility_test.py diff --git a/.github/scripts/install_otel_test_wheels.py b/.github/scripts/install_otel_test_wheels.py new file mode 100644 index 000000000..938996aca --- /dev/null +++ b/.github/scripts/install_otel_test_wheels.py @@ -0,0 +1,101 @@ +"""Install and verify built artifacts without silently using editable sources.""" + +from __future__ import annotations + +import argparse +import ast +import hashlib +import importlib.metadata +import json +import subprocess +import sys +import zipfile +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[2] +CORE = "aws-durable-execution-sdk-python" +OTEL = CORE + "-otel" + + +def built_wheel(package: str) -> Path: + directory = ROOT / "packages" / package + module = package.replace("-", "_") + about = ast.parse((directory / "src" / module / "__about__.py").read_text()) + version = next( + ast.literal_eval(node.value) + for node in about.body + if isinstance(node, ast.Assign) + and any( + isinstance(target, ast.Name) and target.id == "__version__" + for target in node.targets + ) + ) + wheels = list((directory / "dist").glob(f"{module}-{version}-*.whl")) + if len(wheels) != 1: + raise ValueError(f"Build exactly one {package} {version} wheel first: {wheels}") + return wheels[0] + + +def verify(wheel: Path, package: str) -> None: + installed = importlib.metadata.distribution(package) + direct = json.loads(installed.read_text("direct_url.json") or "{}") + assert not direct.get("dir_info", {}).get("editable"), direct + digest = hashlib.sha256(wheel.read_bytes()).hexdigest() + assert direct["archive_info"]["hashes"]["sha256"] == digest, direct + module = package.replace("-", "_") + with zipfile.ZipFile(wheel) as archive: + sources = [ + name + for name in archive.namelist() + if name.startswith(module + "/") and name.endswith(".py") + ] + assert sources + for name in sources: + path = Path(installed.locate_file(name)).resolve() + assert "site-packages" in path.parts, path + assert path.read_bytes() == archive.read(name), path + print( + json.dumps( + { + "package": package, + "version": installed.version, + "wheel": str(wheel), + "sha256": digest, + "verified_sources": len(sources), + } + ) + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--legacy-plugin", action="store_true") + args = parser.parse_args() + packages = [CORE] if args.legacy_plugin else [CORE, OTEL] + wheels = [built_wheel(package) for package in packages] + subprocess.run( + [ + sys.executable, + "-m", + "pip", + "install", + "--no-index", + "--no-deps", + "--force-reinstall", + *map(str, wheels), + ], + check=True, + ) + for wheel, package in zip(wheels, packages, strict=True): + verify(wheel, package) + if args.legacy_plugin: + assert importlib.metadata.version(OTEL) == "1.0.0" + import aws_durable_execution_sdk_python_otel as otel + + assert "site-packages" in Path(otel.__file__).resolve().parts + subprocess.run([sys.executable, "-m", "pip", "check"], check=True) + + +if __name__ == "__main__": + main() diff --git a/.github/tests/otel_lifecycle_compatibility_test.py b/.github/tests/otel_lifecycle_compatibility_test.py new file mode 100644 index 000000000..8e9ce857b --- /dev/null +++ b/.github/tests/otel_lifecycle_compatibility_test.py @@ -0,0 +1,162 @@ +"""Exercise real installed version pairs through the public durable handler.""" + +from __future__ import annotations + +import contextvars +from datetime import UTC, datetime +from importlib.metadata import version +import os +from pathlib import Path +import threading +from types import SimpleNamespace + +import pytest +from aws_durable_execution_sdk_python import durable_execution +from aws_durable_execution_sdk_python import execution as core_execution +from aws_durable_execution_sdk_python.lambda_service import ( + ExecutionDetails, + Operation, + OperationStatus, + OperationType, +) +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from aws_durable_execution_sdk_python_otel.execution_plugin import ExecutionOtelPlugin +from aws_durable_execution_sdk_python_otel.invocation_plugin import InvocationOtelPlugin +from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig +from opentelemetry import baggage, context, trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from packaging.version import Version + + +class NoNetworkClient: + def __getattr__(self, name: str): + raise AssertionError(f"Unexpected service access: {name}") + + +@pytest.mark.parametrize("view", [InvocationOtelPlugin, ExecutionOtelPlugin]) +@pytest.mark.parametrize("order", ["alone", "baggage-first", "baggage-last"]) +@pytest.mark.parametrize("ambient_present", [False, True]) +def test_installed_pair_preserves_context_and_documents_legacy_fallback( + view, order, ambient_present +): + legacy = os.environ.get("OTEL_COMPAT_LEGACY") == "1" + assert Version(version("aws-durable-execution-sdk-python")) >= Version("2.1.0") + assert version("aws-durable-execution-sdk-python-otel") == ( + "1.0.0" if legacy else "1.1.0" + ) + assert "site-packages" in Path(core_execution.__file__).resolve().parts + import aws_durable_execution_sdk_python_otel as installed_otel + + assert "site-packages" in Path(installed_otel.__file__).resolve().parts + + def run() -> None: + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + tracer = provider.get_tracer("installed-lifecycle") + phases = [] + worker_threads = [] + caller_thread = threading.get_ident() + + class BaggagePlugin(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + self.token = context.attach(baggage.set_baggage("customer", "present")) + phases.append("baggage-start") + worker_threads.append(threading.get_ident()) + + def on_invocation_end(self, info): + context.detach(self.token) + phases.append("baggage-end") + worker_threads.append(threading.get_ident()) + + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + enrich_logger=False, + context_extractor=lambda _: None, + ) + ) + plugins = { + "alone": [plugin], + "baggage-first": [BaggagePlugin(), plugin], + "baggage-last": [plugin, BaggagePlugin()], + }[order] + body_parents = [] + + @durable_execution(plugins=plugins, boto3_client=NoNetworkClient()) + def handler(event, durable_context): + phases.append("body") + worker_threads.append(threading.get_ident()) + assert baggage.get_baggage("incoming") == "keep" + assert baggage.get_baggage("customer") == ( + None if order == "alone" else "present" + ) + parent = trace.get_current_span().get_span_context() + body_parents.append(parent) + with tracer.start_as_current_span("customer-span"): + pass + return "ok" + + operation = Operation( + operation_id="installed", + operation_type=OperationType.EXECUTION, + status=OperationStatus.STARTED, + start_timestamp=datetime(2026, 10, 8, tzinfo=UTC), + execution_details=ExecutionDetails(input_payload="{}"), + ) + event = { + "DurableExecutionArn": "test-arn/installed", + "CheckpointToken": "token", + "InitialExecutionState": { + "Operations": [operation.to_json_dict()], + "NextMarker": "", + }, + } + lambda_context = SimpleNamespace( + aws_request_id="installed", + client_context=None, + identity=None, + _epoch_deadline_time_in_ms=0, + invoked_function_arn="test-arn", + tenant_id=None, + ) + ambient = tracer.start_span("host") if ambient_present else trace.INVALID_SPAN + host = baggage.set_baggage( + "incoming", "keep", trace.set_span_in_context(ambient) + ) + token = context.attach(host) + try: + for _ in range(2): + assert handler(event, lambda_context)["Status"] == "SUCCEEDED" + assert context.get_current() is host + assert caller_thread not in worker_threads + assert phases == ( + ["body"] * 2 + if order == "alone" + else ["baggage-start", "body", "baggage-end"] * 2 + ) + spans = exporter.get_finished_spans() + if legacy and view is InvocationOtelPlugin: + # The released plugin does not attach an Invocation fallback. + assert body_parents == [ambient.get_span_context()] * 2 + else: + name = "Invocation" if view is InvocationOtelPlugin else "Workflow" + contexts = [span.context for span in spans if span.name == name] + assert contexts + assert all( + parent.is_valid and parent in contexts for parent in body_parents + ) + users = [span for span in spans if span.name == "customer-span"] + assert len(users) == 2 + assert [span.parent for span in users] == [ + parent if parent.is_valid else None for parent in body_parents + ] + assert plugin._context_tokens == {} + finally: + context.detach(token) + ambient.end() + provider.shutdown() + + contextvars.Context().run(run) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index cb945d8d0..f5efea050 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -84,20 +84,6 @@ jobs: run: hatch run types:check - name: Run tests + coverage run: hatch run test:cov - - name: Verify supported legacy core compatibility - run: hatch run test-pypi-otel-legacy:test - - name: Test OTel with the minimum supported core - run: | - hatch run test-pypi-otel-minimum:python - <<'PYTHON' - from importlib.metadata import version - from pathlib import Path - import aws_durable_execution_sdk_python.plugin as plugin - - assert version("aws-durable-execution-sdk-python") == "2.0.0" - assert "site-packages" in Path(plugin.__file__).parts - assert not hasattr(plugin.PluginExecutor, "run_handler") - PYTHON - hatch run test-pypi-otel-minimum:test - name: Build distribution run: | for pkg in packages/*/; do @@ -108,6 +94,10 @@ jobs: cd "$GITHUB_WORKSPACE" fi done + - name: Test installed core and OTel wheels + run: hatch run test-wheel-otel:test + - name: Test released OTel with the new core wheel + run: hatch run test-wheel-otel-legacy:test - name: Verify OTel wheel dependency contract run: | OTEL_WHEEL=$(find packages/aws-durable-execution-sdk-python-otel/dist \ diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 15b865f78..a3b6bc2f7 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -70,23 +70,25 @@ hatch run dev-otel:typecheck # type check otel only hatch run dev-examples:test # run examples tests only ``` -### PyPI release testing +### Installed package compatibility testing -To verify packages work against the published PyPI version of the core SDK (rather than the local workspace): +Build the core and OTel distributions with `hatch build` in each package, then run +these commands from the repository root: ```bash -hatch run test-pypi-otel:test # test new OTel capabilities against capable installed core -hatch run test-pypi-otel-legacy:test # valid registrations/lifecycles on supported core 2.0.x -hatch run test-pypi-otel-minimum:test # all prior OTel tests and valid registrations on core 2.0.0 -hatch run test-pypi-examples:test # test examples against PyPI core SDK -``` - -The OTel minimum-core environment excludes the local core and pins 2.0.0, so -newer PyPI releases cannot remove legacy compatibility coverage. It retains all -pre-existing OTel tests plus valid registration/wait-resume cases. Exclusivity -validation requires the newer core and is exercised by the complete workspace -suite and the capable-core environment. Use `hatch run dev-otel:test` for the -current workspace core. +hatch run test-wheel-otel:test # full OTel suite on the two built wheels +hatch run test-wheel-otel-legacy:test # released OTel 1.0.0 with the built core +hatch run test-pypi-examples:test # examples against the published core +``` + +The wheel lanes have no editable workspace members. They verify installed source +bytes and artifact hashes before exercising public handlers. OTel 1.1 requires +the redesigned core 2.1.0 lifecycle; it no longer claims compatibility with core +2.0.x. Publish core first. Building both wheels lets CI verify the intended pair +before that minimum is available on PyPI. The legacy-plugin lane documents the +actual core-only upgrade: host isolation is provided by the new core, while old +Invocation OTel does not gain the new fallback. Workspace tests continue to cover +the complete current implementation with `hatch run dev-otel:test`. ### Package-level commands diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index fcc0f77e9..f42b5986c 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -55,9 +55,11 @@ DURABLE_EXECUTION_PLUGINS=otel-execution cold start, so the handler does not need to import or explicitly register the plugin. -Automatic mutual-exclusion validation requires core SDK 2.1.0 or later together -with OTel 1.1.0 or later. OTel 1.1 remains compatible with core 2.0.x for existing -valid registrations; those older cores do not enforce the new group metadata. +OTel 1.1.0 requires core SDK 2.1.0 or later for the invocation-worker lifecycle +and automatic mutual-exclusion validation. Publish the redesigned core first, +then the plugin. Core 2.0.x runs invocation hooks on the caller and cannot provide +the new plugin's handler propagation and host-context isolation; that version +pair is not supported by OTel 1.1.0. Configure only one OTel view on every core version. `InvocationOtelPlugin` shows work within each Lambda invocation; `ExecutionOtelPlugin` shows logical operations across the whole execution. They emit overlapping telemetry and @@ -182,24 +184,31 @@ lambda_.Function( ### Handler context propagation -A core SDK with handler-worker context propagation carries the context established -by invocation-start hooks into the handler. Invocation view preserves an active -ambient span on the canonical execution trace; when that context is absent or -belongs to a different trace, its optional `handler_context` scope makes the Invocation -span current only while the handler runs. The scope closes on the same worker in -reverse plugin order, including on failure and suspension, without changing the -invocation-hook caller. Older cores ignore this optional scope and retain their -existing behavior; install the updated core as well to get handler context propagation. -The bundled OTel classes provide this optional method directly; custom subclasses -inherit it normally and can override it using standard Python method resolution. -The updated core invokes a callable `handler_context` attribute when present; -there is no separate plugin opt-in marker. Older cores ignore the method without -importing any new core API. -Execution view similarly restores the Workflow span inside the handler scope if -another invocation-start hook clears the active span or switches to an unrelated -trace. Both views retain valid same-trace parents and baggage, and restore the -worker's previous context when the scope ends. -Existing plugin registration, factory lifetime, and checkpoint formats are unchanged. +The core runs the existing Start hooks, handler, output preparation, resource +cleanup and End hooks on one invocation worker. Start finishes before checkpoint +processing begins. End follows registered-branch joins and checkpoint shutdown, +and output serialization or checkpoint errors retain their normal classification. +There is no separate handler-context plugin API. + +During Start, Invocation view preserves a valid active span on the canonical +execution trace; when it is absent or unrelated, it attaches the Invocation span. +Execution view attaches the Workflow span, including when an incoming span is on +the same trace. The Invocation span's own ambient parenting is separate from the +active context supplied to handler instrumentation. Both views preserve baggage. + +Start hooks run in registration order. A later successful plugin that deliberately +sets or clears the active span wins; OTel does not apply a second correction pass. +Place OTel after a span-replacing plugin when OTel's view-specific context is +desired. Baggage-only plugins that extend the current context can appear on either +side. End hooks also retain registration order and reset tokens in their owning +Context; they do not promise a reverse-stack observation of other plugins' spans. +Existing registration, factory lifetime and checkpoint formats are unchanged. + +Upgrade both core to 2.1+ and OTel to 1.1+ for these guarantees. The released +OTel 1.0 plugin can run on the new core with worker/host isolation, but its +Invocation view does not attach the new fallback; a core-only upgrade does not +supply that plugin behavior. CI tests the new pair as installed wheels and tests +the actual released plugin separately, including before core 2.1 is on PyPI. ### 3. In your Lambda handler (index.py) @@ -353,19 +362,14 @@ remains in `durable.operation.status`. ### Invocation context isolation -Invocation hooks retain their caller thread and registration order. With the -updated core, invocation-local context-variable bindings are isolated from the -host: hooks see the incoming context and the handler receives their resulting -context, while invocation exit restores the host's original bindings even if a -plugin fails during setup or cleanup. Plugins must not use invocation context -bindings to mutate the host context after the invocation has returned. Older -supported cores retain their existing lifecycle behavior, including the -execution-view limitation when later plugins open invocation context scopes. -The new isolation applies only when plugins are registered. If an invocation-start -hook or handler-scope entry raises, subsequent setup and the handler retain the -bindings from before that hook. Successful scopes still clean up in their original -context, preserving token ownership. This isolates context-variable bindings; -it does not undo a plugin's mutations to shared objects or external side effects. +With core 2.1+, invocation hooks and the handler run on one worker in an +invocation-local Context initialized from the host's bindings. Successful Start +bindings are visible to later hooks, the handler and resource cleanup. The host's +bindings remain unchanged after return, even if a plugin fails during Start or +End. A failed Start's bindings are discarded for subsequent work, while End runs +in that Start's original Context to preserve token ownership. This isolates +context-variable bindings; it does not undo mutations to shared objects or +external side effects. The isolation applies when plugins are registered. ### Log Correlation @@ -467,7 +471,7 @@ setups. ## Requirements - Python >= 3.11 -- `aws-durable-execution-sdk-python` >= 2.0.0 (core >= 2.1.0 with OTel >= 1.1.0 for automatic view-exclusivity validation) +- `aws-durable-execution-sdk-python` >= 2.1.0 (release the redesigned core before OTel 1.1.0) - An ADOT/community OpenTelemetry Lambda layer, or the `standalone` extra ## License diff --git a/packages/aws-durable-execution-sdk-python-otel/pyproject.toml b/packages/aws-durable-execution-sdk-python-otel/pyproject.toml index cbc2f97a2..461718714 100644 --- a/packages/aws-durable-execution-sdk-python-otel/pyproject.toml +++ b/packages/aws-durable-execution-sdk-python-otel/pyproject.toml @@ -22,7 +22,7 @@ classifiers = [ "Programming Language :: Python :: Implementation :: PyPy", ] dependencies = [ - "aws-durable-execution-sdk-python>=2.0.0", + "aws-durable-execution-sdk-python>=2.1.0", ] [project.entry-points."aws_durable_execution.plugins"] diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py index 9c384ed62..40a62ca84 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py @@ -31,11 +31,9 @@ from __future__ import annotations -import contextlib import datetime import logging import threading -from collections.abc import Iterator from typing import Any, ClassVar from aws_durable_execution_sdk_python.plugin import ( @@ -511,24 +509,6 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: if self._config.enrich_logger: install_log_filter(self) - @contextlib.contextmanager - def handler_context(self, info: InvocationStartInfo) -> Iterator[None]: - """Keep handler instrumentation on this execution's trace in its worker.""" - ambient = trace.get_current_span().get_span_context() - workflow = self._workflow_span - token = None - if ( - self._tracing_enabled - and workflow is not None - and (not ambient.is_valid or ambient.trace_id != self._execution_trace_id) - ): - token = otel_context.attach(trace.set_span_in_context(workflow)) - try: - yield - finally: - if token is not None: - otel_context.detach(token) - def _start_workflow_span(self, info: InvocationStartInfo) -> None: """Install a non-recording placeholder for the execution-scoped Workflow span. diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py index 68c2893a4..0fb974fe1 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py @@ -2,11 +2,9 @@ from __future__ import annotations -import contextlib import datetime import logging import threading -from collections.abc import Iterator from typing import Any, ClassVar from aws_durable_execution_sdk_python.plugin import ( @@ -75,6 +73,7 @@ {InvocationStatus.SUCCEEDED, InvocationStatus.FAILED} ) _TIMESTAMP_STEP_NANOS = 1_000 +_INVOCATION_CONTEXT_KEY = "__invocation_context__" _SpanAttributes = dict[str, str | bool | int] @@ -302,7 +301,7 @@ def get_current_span_context(self) -> SpanContext | None: on_user_function_start). Unrelated ambient spans are ignored so logs stay correlated to the durable execution trace. 2. The invocation span from the plugin registry, including lifecycle - phases outside the optional handler-worker context scope. + phases where another plugin changed the active context. Returns: A valid SpanContext, or None if no span is active. @@ -591,27 +590,20 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: attributes=self._extract_attributes(info), ) - # Cover handlers installed after construction as well. - if self._enrich_logger: - install_log_filter(self) - - @contextlib.contextmanager - def handler_context(self, info: InvocationStartInfo) -> Iterator[None]: - """Bind the fallback only inside the SDK-owned handler worker scope.""" + # Start and End share the invocation worker and token-owning Context. + # Retain a valid same-trace ambient parent; otherwise bind this invocation. ambient = trace.get_current_span().get_span_context() invocation_span = self._get_span(None) - token = None - if ( - self._tracing_enabled - and invocation_span is not None - and (not ambient.is_valid or ambient.trace_id != self._execution_trace_id) + if invocation_span is not None and ( + not ambient.is_valid or ambient.trace_id != self._execution_trace_id ): - token = context.attach(trace.set_span_in_context(invocation_span)) - try: - yield - finally: - if token is not None: - context.detach(token) + self._attach_context( + _INVOCATION_CONTEXT_KEY, trace.set_span_in_context(invocation_span) + ) + + # Cover handlers installed after construction as well. + if self._enrich_logger: + install_log_filter(self) def _start_workflow_span(self, info: InvocationStartInfo) -> None: """Install a non-recording placeholder for the execution-scoped Workflow span. @@ -689,8 +681,8 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: self._reset_state() return - # User execution has finished; the worker has already closed its handler - # context scope without modifying the invocation-hook caller. + # User work and checkpoint cleanup have finished. Release Start bindings + # in the Context that created their tokens before closing the spans. self._detach_remaining_contexts() # Spans are registered parent-first, so close pending spans in reverse diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py index 2e95429a7..fa7036920 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py @@ -26,7 +26,6 @@ OperationType, StepDetails, ) -from aws_durable_execution_sdk_python import plugin as core_plugin_api from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( derive_workflow_span_id, @@ -272,20 +271,11 @@ class InheritedExecutionPlugin(ExecutionOtelPlugin): @pytest.mark.parametrize( ("plugin_type", "extra_context_plugin"), [ - (InvocationOtelPlugin, False), - (InvocationOtelPlugin, True), - (ExecutionOtelPlugin, False), + (view, extra) + for view in (InvocationOtelPlugin, ExecutionOtelPlugin) + for extra in (False, True, "same", "unrelated", "absent") ] - # Execution-view caller isolation requires the coordinated newer core. - # Released 2.0.x retains the pre-existing same-order teardown limitation; - # the legacy lane continues checking its supported combinations above. - + ( - [(ExecutionOtelPlugin, True)] - + [(ExecutionOtelPlugin, kind) for kind in ("same", "unrelated", "absent")] - + [(InheritedInvocationPlugin, False), (InheritedExecutionPlugin, "absent")] - if hasattr(core_plugin_api.PluginExecutor, "run_handler") - else [] - ), + + [(InheritedInvocationPlugin, False), (InheritedExecutionPlugin, "absent")], ) @pytest.mark.parametrize("reverse_plugins", [False, True]) @pytest.mark.parametrize("fail_after_resume", [False, True]) @@ -300,10 +290,6 @@ def test_handler_user_spans_inherit_context_across_resume_and_failure( ) -> None: monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) monkeypatch.setenv("_X_AMZN_TRACE_ID", XRAY_TRACE_HEADER) - # The documented PyPI compatibility environment deliberately uses an older - # core. Keep exercising its supported operation tracing and lifecycle while - # asserting the new handler contract only when that core exposes the scope. - supports_handler_context = hasattr(core_plugin_api.PluginExecutor, "run_handler") exporter = InMemorySpanExporter() provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) @@ -326,9 +312,7 @@ def step_body(_step_context: Any) -> str: def handler_body(_event: Any, context: DurableContext) -> str: if extra_context_plugin: - assert baggage.get_baggage("customer") == ( - "present" if supports_handler_context else None - ) + assert baggage.get_baggage("customer") == "present" user_span("handler-entry") saved = context.step(step_body, name="before-wait") user_span("handler-after-step") @@ -338,8 +322,8 @@ def handler_body(_event: Any, context: DurableContext) -> str: raise ValueError("handler failed after resume") return saved - # An unrelated plugin may own a caller-thread OTel baggage scope. The - # invocation-view fallback must never become part of its saved token. + # Start hooks share the invocation worker, in registration order. A later + # successful plugin may deliberately replace or clear the active span. from opentelemetry import baggage class BaggagePlugin(DurableInstrumentationPlugin): @@ -449,37 +433,54 @@ def on_invocation_end(self, _info: Any) -> None: assert calls == ["step"] assert otel_context.get_current() == before_context spans = exporter.get_finished_spans() - expected_parents = ( - [None, None] - if not supports_handler_context - else [ - 0xCAFE - if extra_context_plugin == "same" and not reverse_plugins - else derive_workflow_span_id(EXECUTION_ARN) - ] - * 2 - if issubclass(plugin_type, ExecutionOtelPlugin) - else ambient_ids - if ambient_kind == "same" - else [ - span.context.span_id - for span in spans - if span.name == "Invocation" and span.context is not None + expected_parents: list[int | None] + if isinstance(extra_context_plugin, str) and not reverse_plugins: + # No second OTel correction pass: the later successful Start wins. + expected_parents = [ + None if extra_context_plugin == "absent" else 0xCAFE + ] * 2 + expected_trace_id = {"same": XRAY_TRACE_ID, "unrelated": 1, "absent": None}[ + extra_context_plugin ] - ) + elif issubclass(plugin_type, ExecutionOtelPlugin): + expected_parents = [derive_workflow_span_id(EXECUTION_ARN)] * 2 + expected_trace_id = XRAY_TRACE_ID + elif reverse_plugins and isinstance(extra_context_plugin, str): + expected_parents = ( + [0xCAFE] * 2 + if extra_context_plugin == "same" + else [ + span.context.span_id + for span in spans + if span.name == "Invocation" and span.context is not None + ] + ) + expected_trace_id = XRAY_TRACE_ID + else: + expected_parents = ( + [*ambient_ids] + if ambient_kind == "same" + else [ + span.context.span_id + for span in spans + if span.name == "Invocation" and span.context is not None + ] + ) + expected_trace_id = XRAY_TRACE_ID for name in ("handler-entry", "handler-after-step"): users = [span for span in spans if span.name == name] assert len(users) == 2 assert [span.parent.span_id if span.parent else None for span in users] == ( expected_parents ) - assert all( - span.context is not None - and ( - (span.context.trace_id == XRAY_TRACE_ID) == supports_handler_context + assert all(span.context is not None for span in users) + if expected_trace_id is None: + assert all( + span.parent is None and span.context.trace_id != XRAY_TRACE_ID + for span in users ) - for span in users - ) + else: + assert all(span.context.trace_id == expected_trace_id for span in users) after_resume = next( span for span in spans if span.name == "handler-after-resume" ) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py index 23c7394ba..214a3a89d 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py @@ -1628,10 +1628,8 @@ def test_nested_suspension_unwinds_scopes_in_reverse_order(): @pytest.mark.parametrize("ambient_kind", ["same", "unrelated", "absent"]) @pytest.mark.parametrize("raises", [False, True]) @pytest.mark.parametrize("sampler", [ALWAYS_ON, ALWAYS_OFF]) -def test_handler_scope_keeps_execution_trace_and_restores_context( - ambient_kind: str, - raises: bool, - sampler: Sampler, +def test_invocation_hooks_bind_parent_and_restore_baggage( + ambient_kind: str, raises: bool, sampler: Sampler ) -> None: provider = TracerProvider(sampler=sampler) plugin = ExecutionOtelPlugin( @@ -1641,55 +1639,50 @@ def test_handler_scope_keeps_execution_trace_and_restores_context( enrich_logger=False, ) ) - info = _invocation_start_info() - plugin.on_invocation_start(info) - workflow = trace.get_current_span().get_span_context() - assert workflow.is_valid - assert workflow.span_id == derive_workflow_span_id(EXECUTION_ARN) - caller = baggage.set_baggage("customer", "preserved", Context()) - expected = workflow + caller = baggage.set_baggage("tenant", "hook-test", Context()) + ambient = None if ambient_kind != "absent": ambient = SpanContext( - trace_id=workflow.trace_id if ambient_kind == "same" else 1, - span_id=0xCAFE, + trace_id=_to_otel_trace_id(EXECUTION_ARN, START_TIME) + if ambient_kind == "same" + else 1, + span_id=0x42, is_remote=False, - trace_flags=workflow.trace_flags, + trace_flags=TraceFlags(1), ) caller = trace.set_span_in_context(NonRecordingSpan(ambient), caller) - if ambient_kind == "same": - expected = ambient token = otel_context.attach(caller) - error = ValueError("user failure") + error = ValueError("handler error") try: + plugin.on_invocation_start(_invocation_start_info()) + expected = trace.get_current_span().get_span_context() + assert expected.trace_id == _to_otel_trace_id(EXECUTION_ARN, START_TIME) + assert expected.span_id == derive_workflow_span_id(EXECUTION_ARN) def body() -> None: - with plugin.handler_context(info): - assert trace.get_current_span().get_span_context() == expected - assert baggage.get_baggage("customer") == "preserved" - if raises: - raise error - - if raises: - with pytest.raises(ValueError) as caught: + active = trace.get_current_span().get_span_context() + assert active == expected + assert active.is_valid + assert baggage.get_baggage("tenant") == "hook-test" + if raises: + raise error + + try: + if raises: + with pytest.raises(ValueError) as caught: + body() + assert caught.value is error + else: body() - assert caught.value is error - else: - body() + finally: + plugin.on_invocation_end( + _invocation_end_info( + InvocationStatus.FAILED if raises else InvocationStatus.SUCCEEDED + ) + ) + assert plugin._context_tokens == {} assert otel_context.get_current() is caller + assert baggage.get_baggage("tenant") == "hook-test" finally: otel_context.detach(token) - plugin.on_invocation_end(_invocation_end_info()) provider.shutdown() - - -@pytest.mark.parametrize("completed", [False, True]) -def test_handler_scope_without_live_workflow_is_noop(completed: bool) -> None: - plugin, _ = _create_plugin() - info = _invocation_start_info() - if completed: - plugin.on_invocation_start(info) - plugin.on_invocation_end(_invocation_end_info()) - caller = otel_context.get_current() - with plugin.handler_context(info): - assert otel_context.get_current() is caller - assert otel_context.get_current() is caller diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py index 723ad41ef..ba373ad8d 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py @@ -724,6 +724,7 @@ def test_operation_end_without_start_links_previous_logical_operation(): assert ( span.attributes["durable.operation.status"] == OperationStatus.SUCCEEDED.value ) + plugin.on_invocation_end(_invocation_end_info()) def test_continuation_span_uses_current_start_and_end_times(): @@ -753,6 +754,7 @@ def test_continuation_span_uses_current_start_and_end_times(): span = exporter.get_finished_spans()[0] assert invocation_span.start_time <= span.start_time assert before_callback <= span.start_time <= span.end_time <= after_callback + plugin.on_invocation_end(_invocation_end_info()) def test_resume_operation_timestamps_do_not_precede_current_invocation(): @@ -812,6 +814,7 @@ def test_resume_operation_timestamps_do_not_precede_current_invocation(): assert invocation_span.start_time <= after_resume_span.start_time assert after_resume_span.parent is not None assert after_resume_span.parent.span_id == invocation_span.context.span_id + plugin.on_invocation_end(_invocation_end_info()) def test_ordered_timestamps_are_thread_safe(): @@ -878,6 +881,7 @@ def test_retried_operation_uses_fresh_id_and_links_previous_logical_operation(): derive_workflow_span_id(EXECUTION_ARN), operation_id_to_span_id(EXECUTION_ARN, operation_id), } + plugin.on_invocation_end(_invocation_end_info()) def test_step_operation_span_parents_attempt_span(): @@ -1028,6 +1032,7 @@ def test_user_function_callbacks_emit_attempt_span_attributes(): == UserFunctionOutcome.SUCCEEDED.value ) assert "durable.operation.status" not in span.attributes + plugin.on_invocation_end(_invocation_end_info()) def test_step_attempt_span_name_includes_attempt_number(): @@ -1070,6 +1075,7 @@ def test_step_attempt_span_name_includes_attempt_number(): span = exporter.get_finished_spans()[0] assert span.name == "fetch-user attempt 2" + plugin.on_invocation_end(_invocation_end_info()) def test_step_attempt_span_name_defaults_to_first_attempt(): @@ -1112,6 +1118,7 @@ def test_step_attempt_span_name_defaults_to_first_attempt(): span = exporter.get_finished_spans()[0] assert span.name == "fetch-user attempt 1" + plugin.on_invocation_end(_invocation_end_info()) @pytest.mark.parametrize( @@ -1213,6 +1220,7 @@ def test_context_span_waits_for_terminal_status_and_omits_attempt_attributes( assert "durable.attempt.number" not in span.attributes assert "durable.attempt.outcome" not in span.attributes assert span.status.status_code is expected_span_status + plugin.on_invocation_end(_invocation_end_info()) def test_span_registry_helpers_can_be_called_from_multiple_threads(): @@ -1257,8 +1265,9 @@ def test_user_function_end_restores_enclosing_context(): # After the step, the enclosing context is restored and no scope is left # behind. Log correlation resolves the invocation span from the registry. assert otel_context.get_current() == enclosing_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} assert plugin.get_current_span_context().span_id == invocation_span_id + plugin.on_invocation_end(_invocation_end_info()) def test_user_function_start_preserves_baggage_in_current_context(): @@ -1292,7 +1301,8 @@ def test_user_function_end_restores_enclosing_context_on_failure(): ) assert otel_context.get_current() == enclosing_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} + plugin.on_invocation_end(_invocation_end_info()) def test_user_function_end_restores_enclosing_context_across_multiple_steps(): @@ -1309,8 +1319,9 @@ def test_user_function_end_restores_enclosing_context_across_multiple_steps(): # Between each step the context is back to where it started, and log # correlation still resolves the invocation span. assert otel_context.get_current() == enclosing_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} assert plugin.get_current_span_context().span_id == invocation_span_id + plugin.on_invocation_end(_invocation_end_info()) # ---------------------------------------------------------------------- @@ -1332,6 +1343,7 @@ def test_get_current_span_context_returns_invocation_span_at_top_level(): invocation_span = plugin._get_span(None) assert span_context is not None assert span_context.span_id == invocation_span.get_span_context().span_id + plugin.on_invocation_end(_invocation_end_info()) def test_get_current_span_context_returns_operation_span_inside_step(): @@ -1363,6 +1375,7 @@ def test_get_current_span_context_returns_invocation_span_between_steps(): invocation_span = plugin._get_span(None) assert span_context is not None assert span_context.span_id == invocation_span.get_span_context().span_id + plugin.on_invocation_end(_invocation_end_info()) # ---------------------------------------------------------------------- @@ -1441,11 +1454,12 @@ def test_top_level_step_end_falls_back_to_invocation_for_correlation(): plugin.on_user_function_start(_user_function_start_info(operation_id)) plugin.on_user_function_end(_user_function_end_info(operation_id)) - # No durable span is attached at the top level, so the registry fallback - # supplies the invocation span for log correlation. + # The invocation Start binding is restored after the step, including + # between-step user instrumentation and log correlation. assert otel_context.get_current() == enclosing_context - assert not trace.get_current_span().get_span_context().is_valid + assert trace.get_current_span().get_span_context().span_id == invocation_span_id assert plugin.get_current_span_context().span_id == invocation_span_id + plugin.on_invocation_end(_invocation_end_info()) def test_get_current_span_context_returns_context_span_between_nested_steps(): @@ -1709,6 +1723,7 @@ def test_replayed_context_span_links_previous_logical_operation(): derive_workflow_span_id(EXECUTION_ARN), operation_id_to_span_id(EXECUTION_ARN, operation_id), } + plugin.on_invocation_end(_invocation_end_info()) def test_checkpointed_context_first_span_uses_deterministic_id(): @@ -1867,7 +1882,8 @@ def test_child_context_end_restores_context_active_before_it(): ) assert otel_context.get_current() == enclosing_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} + plugin.on_invocation_end(_invocation_end_info()) def test_nested_scopes_are_released_without_accumulating(): @@ -1891,13 +1907,14 @@ def test_nested_scopes_are_released_without_accumulating(): ) # The inner step restored the child-context scope, not a copy of it. assert otel_context.get_current() == inside_context - assert set(plugin._context_tokens) == {context_id} + assert set(plugin._context_tokens) == {context_id, "__invocation_context__"} plugin.on_user_function_end( _user_function_end_info(context_id, operation_type=OperationType.CONTEXT) ) assert otel_context.get_current() == before_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} + plugin.on_invocation_end(_invocation_end_info()) def test_invocation_end_releases_scope_of_suspended_user_function(): @@ -1981,7 +1998,7 @@ def test_detach_ignores_token_attached_on_another_thread(): plugin._detach_context("step-1:attempt:1") - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} assert otel_context.get_current() == before_context plugin.on_invocation_end(_invocation_end_info()) @@ -2072,7 +2089,7 @@ def test_reentered_step_attempt_releases_the_previous_scope(): plugin.on_user_function_end(_user_function_end_info(operation_id)) assert otel_context.get_current() == before_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} plugin.on_invocation_end(_invocation_end_info()) @@ -2166,7 +2183,7 @@ def test_nested_suspension_unwinds_scopes_in_reverse_order(): ) ) assert otel_context.get_current() == before_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} # Neither span is ended: both operations are still in flight. assert not exporter.get_finished_spans() @@ -2199,7 +2216,7 @@ def test_nested_suspension_unwinds_scopes_in_reverse_order(): _user_function_end_info("ctx-outer", operation_type=OperationType.CONTEXT) ) assert otel_context.get_current() == before_context - assert plugin._context_tokens == {} + assert set(plugin._context_tokens) == {"__invocation_context__"} plugin.on_invocation_end(_invocation_end_info()) @@ -2207,66 +2224,61 @@ def test_nested_suspension_unwinds_scopes_in_reverse_order(): @pytest.mark.parametrize("ambient_kind", ["same", "unrelated", "absent"]) @pytest.mark.parametrize("raises", [False, True]) @pytest.mark.parametrize("sampler", [ALWAYS_ON, ALWAYS_OFF]) -def test_handler_context_preserves_or_replaces_parent_and_restores_baggage( +def test_invocation_hooks_bind_parent_and_restore_baggage( ambient_kind: str, raises: bool, sampler: Sampler ) -> None: - plugin, _ = _create_plugin_with_sampler(sampler) - info = _invocation_start_info() - plugin.on_invocation_start(info) - invocation = plugin.get_current_span_context() - assert invocation is not None and invocation.is_valid - caller = baggage.set_baggage("tenant", "scope-test", Context()) - expected = invocation + provider = TracerProvider(sampler=sampler) + plugin = InvocationOtelPlugin( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + caller = baggage.set_baggage("tenant", "hook-test", Context()) + ambient = None if ambient_kind != "absent": ambient = SpanContext( - trace_id=( - invocation.trace_id - if ambient_kind == "same" - else (1 if invocation.trace_id != 1 else 2) - ), + trace_id=_to_otel_trace_id(EXECUTION_ARN, START_TIME) + if ambient_kind == "same" + else 1, span_id=0x42, is_remote=False, - trace_flags=invocation.trace_flags, + trace_flags=TraceFlags(1), ) caller = trace.set_span_in_context(NonRecordingSpan(ambient), caller) - if ambient_kind == "same": - expected = ambient token = otel_context.attach(caller) error = ValueError("handler error") try: + plugin.on_invocation_start(_invocation_start_info()) + expected = ( + ambient if ambient_kind == "same" else plugin.get_current_span_context() + ) def body() -> None: - with plugin.handler_context(info): - active = trace.get_current_span().get_span_context() - assert active == expected - assert active.is_valid - assert baggage.get_baggage("tenant") == "scope-test" - if raises: - raise error - - if raises: - with pytest.raises(ValueError) as caught: + active = trace.get_current_span().get_span_context() + assert active == expected + assert active.is_valid + assert baggage.get_baggage("tenant") == "hook-test" + if raises: + raise error + + try: + if raises: + with pytest.raises(ValueError) as caught: + body() + assert caught.value is error + else: body() - assert caught.value is error - else: - body() + finally: + plugin.on_invocation_end( + _invocation_end_info( + InvocationStatus.FAILED if raises else InvocationStatus.SUCCEEDED + ) + ) + assert plugin._context_tokens == {} assert otel_context.get_current() is caller - assert baggage.get_baggage("tenant") == "scope-test" + assert baggage.get_baggage("tenant") == "hook-test" finally: otel_context.detach(token) - plugin.on_invocation_end(_invocation_end_info()) - - -@pytest.mark.parametrize("completed", [False, True]) -def test_handler_context_without_live_invocation_leaves_context_unchanged( - completed: bool, -) -> None: - plugin, _ = _create_plugin() - info = _invocation_start_info() - if completed: - plugin.on_invocation_start(info) - plugin.on_invocation_end(_invocation_end_info()) - caller = otel_context.get_current() - with plugin.handler_context(info): - assert otel_context.get_current() is caller - assert otel_context.get_current() is caller + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py index 28477a460..b7a34ebda 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_package_metadata.py @@ -6,8 +6,7 @@ PACKAGE_ROOT = Path(__file__).resolve().parents[1] REPOSITORY_ROOT = PACKAGE_ROOT.parents[1] -CORE_DEPENDENCY = "aws-durable-execution-sdk-python>=2.0.0" -EXCLUSIVITY_TEST_CORE = "aws-durable-execution-sdk-python>=2.1.0" +CORE_DEPENDENCY = "aws-durable-execution-sdk-python>=2.1.0" TEST_OTEL_DEPENDENCIES = { "opentelemetry-sdk>=1.20.0", "opentelemetry-propagator-aws-xray", @@ -68,7 +67,7 @@ def test_test_environments_install_layer_provided_dependencies() -> None: "test", "dev-otel", "dev-examples", - "test-pypi-otel", + "test-wheel-otel", "test-pypi-examples", ): assert TEST_OTEL_DEPENDENCIES <= set( @@ -77,22 +76,7 @@ def test_test_environments_install_layer_provided_dependencies() -> None: assert TEST_OTEL_DEPENDENCIES <= set(environments["types"]["extra-dependencies"]) -def test_pypi_compatibility_environment_requires_exclusivity_core() -> None: - dependencies = _load_pyproject(REPOSITORY_ROOT / "pyproject.toml")["tool"]["hatch"][ - "envs" - ]["test-pypi-otel"]["dependencies"] - - assert EXCLUSIVITY_TEST_CORE in dependencies - - -def test_pypi_otel_environment_installs_lifecycle_test_runner() -> None: - dependencies = _load_pyproject(REPOSITORY_ROOT / "pyproject.toml")["tool"]["hatch"][ - "envs" - ]["test-pypi-otel"]["dependencies"] - assert "aws-durable-execution-sdk-python-testing>=1.2.1" in dependencies - - -def test_core_dependency_preserves_previously_supported_releases() -> None: +def test_core_dependency_requires_worker_lifecycle_release() -> None: dependencies = _load_pyproject(PACKAGE_ROOT / "pyproject.toml")["project"][ "dependencies" ] @@ -101,42 +85,23 @@ def test_core_dependency_preserves_previously_supported_releases() -> None: for value in dependencies if Requirement(value).name == "aws-durable-execution-sdk-python" ) - assert requirement.specifier.contains("2.0.0") - assert requirement.specifier.contains("2.0.1") + assert not requirement.specifier.contains("2.0.0") + assert not requirement.specifier.contains("2.0.1") assert requirement.specifier.contains("2.1.0") -def test_pypi_otel_environment_does_not_shadow_installed_core() -> None: - environment = _load_pyproject(REPOSITORY_ROOT / "pyproject.toml")["tool"]["hatch"][ +def test_wheel_lanes_do_not_shadow_installed_artifacts() -> None: + environments = _load_pyproject(REPOSITORY_ROOT / "pyproject.toml")["tool"]["hatch"][ "envs" - ]["test-pypi-otel"] - assert environment["workspace"]["members"] == [ - "packages/aws-durable-execution-sdk-python-otel" ] - - -def test_legacy_lane_retains_supported_core_20_range() -> None: - environment = _load_pyproject(REPOSITORY_ROOT / "pyproject.toml")["tool"]["hatch"][ - "envs" - ]["test-pypi-otel-legacy"] - requirement = next( - Requirement(value) - for value in environment["dependencies"] - if Requirement(value).name == "aws-durable-execution-sdk-python" + for name in ("test-wheel-otel", "test-wheel-otel-legacy"): + assert environments[name]["workspace"]["members"] == [] + assert environments[name]["detached"] is True + assert ( + "aws-durable-execution-sdk-python-otel==1.0.0" + in environments["test-wheel-otel-legacy"]["dependencies"] + ) + assert ( + "aws-durable-execution-sdk-python-testing>=1.2.1" + in environments["test-wheel-otel"]["dependencies"] ) - assert requirement.specifier.contains("2.0.0") - assert requirement.specifier.contains("2.0.1") - assert not requirement.specifier.contains("2.1.0") - assert environment["workspace"]["members"] == [ - "packages/aws-durable-execution-sdk-python-otel" - ] - - -def test_minimum_core_lane_retains_exact_supported_floor() -> None: - environment = _load_pyproject(REPOSITORY_ROOT / "pyproject.toml")["tool"]["hatch"][ - "envs" - ]["test-pypi-otel-minimum"] - assert "aws-durable-execution-sdk-python==2.0.0" in environment["dependencies"] - assert environment["workspace"]["members"] == [ - "packages/aws-durable-execution-sdk-python-otel" - ] diff --git a/packages/aws-durable-execution-sdk-python/README.md b/packages/aws-durable-execution-sdk-python/README.md index bfe54384f..08a70e06a 100644 --- a/packages/aws-durable-execution-sdk-python/README.md +++ b/packages/aws-durable-execution-sdk-python/README.md @@ -85,9 +85,8 @@ constraint applies to the combined `plugins=[...]` argument and Core 2.1+ with OTel 1.1+ rejects both at cold start with `PluginLoadError` naming the conflicting views; keep only one. Choose Invocation for work within each Lambda invocation or Execution for logical operations across the durable execution. -Unrelated instrumentation plugins can run alongside either view. Existing valid -registrations remain supported with OTel 1.1 on older core 2.0.x; those cores do -not implement the new exclusivity validation. +Unrelated instrumentation plugins can run alongside either view. OTel 1.1 requires core 2.1 or later for the invocation-worker lifecycle as well +as view-exclusivity validation. Release core 2.1 before releasing OTel 1.1. Plugin authors explicitly opt in by declaring `__durable_registration_api__ = 1` on a plugin class. That class and its subclasses @@ -112,22 +111,34 @@ in their hierarchy retain their existing attributes/helpers. The generic plugin neither attribute nor hook, and provider API version 1 and existing plugin lifecycle order are unchanged. -### Optional handler context scopes - -A plugin can provide `handler_context(info)` returning a context manager. The -updated core looks up this optional method normally, so inherited methods work -without extra declarations. A missing or non-callable attribute is ignored. -The core enters these scopes around the top-level handler on its worker thread, -in registration order, and closes them in reverse order. Cleanup receives no -handler exception and cannot suppress or replace its outcome. Invocation hooks -retain their original thread and order; failed lookup, setup or entry bindings -are discarded, and successful scope cleanup stays in the context that owns its -tokens. - -`handler_context` is an optional plugin API name: callable methods or attributes -with that name are invoked. The generic plugin base does not require a default -method. Older cores ignore this optional API and retain their existing behavior. -The provider API version and dependency requirements are unchanged. +### Invocation worker lifecycle + +The existing `on_invocation_start` and `on_invocation_end` hooks run on the +handler's invocation worker, both in registration order. Start hooks finish +before background checkpoint processing begins. The same worker runs the handler +(including its `finally` blocks), prepares the existing serialized output or +error, checkpoints large results when necessary, joins registered branches while +checkpointing is available, stops and waits for checkpoint processing, then calls +End before returning the outcome. The caller shuts down the handler executor. +There is no separate handler-context hook or context-manager plugin API. + +When plugins are registered, the worker begins with a copy of the caller's +context-variable bindings. Successful Start bindings flow into later Start hooks, +the handler, serialization and resource cleanup. A later successful Start may +replace an earlier binding. If a Start hook raises, its new bindings are discarded +for subsequent work; its End still runs in the original Context so its tokens can +be reset. End hooks retain forward registration order, not reverse stack order, +so they must not rely on observing a stack-like unwind of other plugins' contexts. +The worker's invocation context is discarded on return, including after plugin +cleanup failures, leaving the host's bindings unchanged. This isolates bindings, +not mutations to shared objects or external side effects. Without plugins, the +handler retains its existing fresh-worker context behavior. + +These lifecycle guarantees require core 2.1.0 or later. The OTel 1.1 plugin +requires that core version and uses the existing Start/End hooks. Upgrading the +core alone with OTel 1.0 isolates worker bindings, but does not add the newer +Invocation-view fallback to that older plugin; upgrade both packages for it. +Provider API version 1 and the independent registration API are unchanged. ## 🚀 Quick Start diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py index 732f6d6e8..e59820007 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/execution.py @@ -5,7 +5,7 @@ import functools import json import logging -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import Future, ThreadPoolExecutor, wait from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any @@ -192,7 +192,6 @@ def durable_execution( configured_plugins = load_configured_plugins(plugins) plugin_executor = PluginExecutor(configured_plugins) - @plugin_executor.handle_durable_output def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: invocation_input: DurableExecutionInvocationInput service_client: DurableServiceClient @@ -281,63 +280,13 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: ), ) - # Use ThreadPoolExecutor for concurrent execution of user code and background checkpoint processing - with ( - ThreadPoolExecutor( - max_workers=2, thread_name_prefix="dex-handler" - ) as executor, - contextlib.closing(execution_state) as execution_state, - ): - execution_operation = execution_state.get_execution_operation() - - # execute the plugins - plugin_executor.on_invocation_start( - execution_arn=invocation_input.durable_execution_arn, - lambda_context=context, - execution_start_time=( - execution_operation.start_timestamp - if execution_operation is not None - else None - ), - is_first_invocation=not has_prior_operations, - execution_input=input_event, - # Read the map through a callable rather than snapshotting it - # here: the invocation-end hook needs the state as of the end of - # the invocation, and neither hook pays for the conversion until - # a plugin actually reads it. - operations_provider=lambda: execution_state.operations, - updated_operation_ids=invocation_input.updated_operation_ids, - ) - # Thread 1: Run background checkpoint processing - executor.submit(execution_state.checkpoint_batches_forever) - - # Thread 2: Execute user function + def invoke_handler() -> MutableMapping[str, Any]: logger.debug( "%s entering user-space...", invocation_input.durable_execution_arn ) - if configured_plugins: - # Invocation-start hooks can establish tracing and other contextvars. - # Context.run restores worker bindings on both return and failure. - user_future = executor.submit( - contextvars.copy_context().run, - plugin_executor.run_handler, - func, - input_event, - durable_context, - ) - else: - # Preserve the original fresh-worker context for uninstrumented - # handlers, including the absence of caller ContextVar bindings. - user_future = executor.submit(func, input_event, durable_context) - - logger.debug( - "%s waiting for user code completion...", - invocation_input.durable_execution_arn, - ) - try: # Background checkpointing errors will propagate through CompletionEvent.wait() as BackgroundThreadError - result = user_future.result() + result = func(input_event, durable_context) # done with userland logger.debug( @@ -473,6 +422,78 @@ def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: return result + def invoke_worker() -> MutableMapping[str, Any]: + # This task owns Start, user cleanup, output preparation, resource + # cleanup and End. The caller owns the executor that runs this task. + with plugin_executor.run(): + checkpoint_future: Future[None] | None = None + + def wait_for_checkpoint() -> None: + if checkpoint_future is not None: + # Match the former executor join: observed checkpoint + # errors retain their existing classification paths. + wait((checkpoint_future,)) + + try: + with contextlib.ExitStack() as resources: + # LIFO cleanup: branches join with checkpointing alive; + # then close stops checkpointing and the wait completes. + resources.callback(wait_for_checkpoint) + resources.callback( + plugin_executor._run_in_invocation_context, + execution_state.close, + ) + execution_operation = execution_state.get_execution_operation() + + # execute the plugins + plugin_executor.on_invocation_start( + execution_arn=invocation_input.durable_execution_arn, + lambda_context=context, + execution_start_time=( + execution_operation.start_timestamp + if execution_operation is not None + else None + ), + is_first_invocation=not has_prior_operations, + execution_input=input_event, + # Read the map through a callable rather than snapshotting it + # here: the invocation-end hook needs the state as of the end of + # the invocation, and neither hook pays for the conversion until + # a plugin actually reads it. + operations_provider=lambda: execution_state.operations, + updated_operation_ids=invocation_input.updated_operation_ids, + ) + # No checkpoint work starts until all Start hooks finish. + checkpoint_future = executor.submit( + execution_state.checkpoint_batches_forever + ) + output = plugin_executor._run_in_invocation_context( + invoke_handler + ) + plugin_executor.on_invocation_end( + DurableExecutionInvocationOutput.from_dict(output) + ) + return output + except Exception as error: + plugin_executor.on_invocation_end( + DurableExecutionInvocationOutput.create_retry( + ErrorObject.from_exception(error) + ) + ) + raise + + with ThreadPoolExecutor( + max_workers=2, thread_name_prefix="dex-handler" + ) as executor: + if configured_plugins: + invocation_future = executor.submit( + contextvars.copy_context().run, invoke_worker + ) + else: + # Preserve the original fresh-worker context without plugins. + invocation_future = executor.submit(invoke_worker) + return invocation_future.result() + return wrapper diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py index a74a1cc86..ed3792102 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/plugin.py @@ -4,14 +4,13 @@ import contextvars import copy import datetime -import functools import logging -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from enum import Enum from threading import Lock -from typing import Any, Callable, MutableMapping, cast +from typing import Any, Callable, cast from aws_durable_execution_sdk_python.identifier import OperationIdentifier from aws_durable_execution_sdk_python.lambda_service import ( @@ -468,6 +467,8 @@ def __init__(self, plugins: list[DurableInstrumentationPlugin] | None): self._executor: ThreadPoolExecutor | None = None self._invocation_status: InvocationStartInfo | None = None self._operations_provider: Callable[[], Mapping[str, Operation]] | None = None + # Non-None only after a Start hook fails: later setup and invocation + # work use its pre-hook snapshot. End still uses each hook's token owner. self._startup_context: contextvars.Context | None = None self._invocation_contexts: list[contextvars.Context | None] = [] self._reported_terminal_updates: set[tuple[str, OperationStatus]] = set() @@ -561,56 +562,11 @@ def execute_plugins(self, info, sync): # this is called asynchronously, so plugins cannot manipulate thread local objects self._executor.submit(self._dispatch_plugin, plugin, info) - @contextlib.contextmanager - def _safe_handler_context( - self, plugin: DurableInstrumentationPlugin, info: InvocationStartInfo - ) -> Iterator[bool]: - # Plugins may omit this optional scope method. - scope = None - succeeded = True - try: - factory = getattr(plugin, "handler_context", None) - if callable(factory): - scope = factory(info) - scope.__enter__() - except Exception: - scope = None - succeeded = False - logger.exception( - "Plugin %s handler context failed", plugin.__class__.__name__ - ) - try: - yield succeeded - finally: - if scope is not None: - try: - scope.__exit__(None, None, None) - except Exception: - logger.exception( - "Plugin %s handler context cleanup failed", - plugin.__class__.__name__, - ) - - def run_handler(self, handler: Callable[..., Any], *args: Any) -> Any: - """Run the user handler inside optional, balanced plugin context scopes.""" - if not self._plugins or self._invocation_status is None: - return handler(*args) - owner = ( - self._startup_context.copy() - if self._startup_context is not None - else contextvars.copy_context() - ) - with contextlib.ExitStack() as scopes: - for plugin in self._plugins: - before = owner.copy() - scope = self._safe_handler_context(plugin, self._invocation_status) - succeeded = owner.run(scope.__enter__) - # Keep finalizers with their entry Context: ContextVar tokens - # cannot be reset in a copy, even when its bindings are identical. - scopes.callback(owner.run, scope.__exit__, None, None, None) - if not succeeded: - owner = before - return owner.run(handler, *args) + def _run_in_invocation_context(self, invoke: Callable[[], Any]) -> Any: + """Continue in the pre-hook Context only after a failed Start hook.""" + if self._startup_context is not None: + return self._startup_context.run(invoke) + return invoke() def _snapshot_operation_infos( self, @@ -938,37 +894,3 @@ def _is_terminal_status(status): OperationStatus.CANCELLED, OperationStatus.STOPPED, ] - - @property - def handle_durable_output(self): - def decorator(func: Callable[[Any, LambdaContext], MutableMapping[str, Any]]): - def invoke(event: Any, context: LambdaContext): - with self.run(): - try: - output = func(event, context) - - self.on_invocation_end( - output=DurableExecutionInvocationOutput.from_dict(output), - ) - return output - except Exception as e: - self.on_invocation_end( - output=DurableExecutionInvocationOutput.create_retry( - ErrorObject.from_exception(e) - ), - ) - raise - - @functools.wraps(func) - def wrapper(event: Any, context: LambdaContext): - if not self._plugins: - return invoke(event, context) - # Keep hooks on their existing caller thread and in registration - # order, but isolate their context bindings from the host. Two - # plugins can otherwise restore a stale predecessor when their - # invocation-end hooks close scopes in the original order. - return contextvars.copy_context().run(invoke, event, context) - - return wrapper - - return decorator diff --git a/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py index c62377261..8cb081029 100644 --- a/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/handler_worker_context_test.py @@ -1,513 +1,589 @@ -"""Focused handler-dispatch tests with a mock service and real worker thread.""" +"""Real public invocation-worker lifecycle and ContextVar ownership controls.""" from __future__ import annotations import contextvars -from collections.abc import Callable -from concurrent.futures import Future, ThreadPoolExecutor +import threading +from concurrent.futures import ThreadPoolExecutor from typing import Any from unittest.mock import Mock import pytest from aws_durable_execution_sdk_python.context import DurableContext -from aws_durable_execution_sdk_python.exceptions import InvocationError -from aws_durable_execution_sdk_python.execution import durable_execution +from aws_durable_execution_sdk_python.exceptions import ( + CheckpointError, + CheckpointErrorCategory, + InvocationError, + SuspendExecution, +) +from aws_durable_execution_sdk_python.execution import ( + DurableExecutionInvocationInputWithClient, + InitialExecutionState, + durable_execution, +) +from aws_durable_execution_sdk_python.lambda_service import ( + CheckpointOutput, + CheckpointUpdatedExecutionState, + DurableServiceClient, + ExecutionDetails, + Operation, + OperationStatus, + OperationType, + OperationUpdate, +) from aws_durable_execution_sdk_python.plugin import ( DurableInstrumentationPlugin, InvocationEndInfo, InvocationStartInfo, InvocationStatus, ) +from aws_durable_execution_sdk_python.state import ExecutionState -@pytest.mark.parametrize("outcome", ["success", "failure", "retry"]) -@pytest.mark.parametrize("plugin_mode", ["none", "healthy", "partial-failure"]) -def test_handler_worker_preserves_context_and_restores_its_caller( - monkeypatch: pytest.MonkeyPatch, outcome: str, plugin_mode: str +def invocation(client: Any = None) -> tuple[Any, Any]: + client = client or Mock(spec=DurableServiceClient) + client.checkpoint.return_value = CheckpointOutput( + checkpoint_token="next", new_execution_state=CheckpointUpdatedExecutionState() + ) + event = DurableExecutionInvocationInputWithClient( + durable_execution_arn="test-arn/worker-lifecycle", + checkpoint_token="initial", + initial_execution_state=InitialExecutionState( + operations=[ + Operation( + operation_id="execution", + operation_type=OperationType.EXECUTION, + status=OperationStatus.STARTED, + execution_details=ExecutionDetails(input_payload="{}"), + ) + ], + next_marker="", + ), + service_client=client, + ) + context = Mock() + context.aws_request_id = "worker-request" + context.client_context = context.identity = context.invoked_function_arn = None + context._epoch_deadline_time_in_ms = 0 + context.tenant_id = None + return event, context + + +@pytest.mark.parametrize("outcome", ["success", "failure", "retry", "pending"]) +@pytest.mark.parametrize("mode", ["none", "healthy", "failed-start", "failed-end"]) +def test_public_worker_lifecycle_owns_context_and_outcome( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + outcome: str, + mode: str, ) -> None: - marker = contextvars.ContextVar("handler-worker-context", default="worker-empty") - seen: list[str] = [] - worker_boundaries: list[tuple[str, str]] = [] + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + marker = contextvars.ContextVar("worker-owner", default="fresh-worker") + phases: list[tuple[str, int, str]] = [] + boundary: list[tuple[str, int, str, str]] = [] + retry = InvocationError("original retry") + fail = ValueError("original body failure") statuses: list[InvocationStatus] = [] + caller_thread = threading.get_ident() - class ClaimPlugin(DurableInstrumentationPlugin): - token: contextvars.Token[str] | None = None - - def on_invocation_start(self, info: InvocationStartInfo) -> None: - self.token = marker.set("invocation-start") - if plugin_mode == "partial-failure": - raise ValueError("partial plugin setup") - - def on_invocation_end(self, info: InvocationEndInfo) -> None: - statuses.append(info.status) - assert self.token is not None - marker.reset(self.token) - self.token = None - - def body(_event: Any, _context: DurableContext) -> str: - seen.append(marker.get()) - marker.set("worker-mutation") - if outcome == "failure": - raise ValueError("handler failure") - if outcome == "retry": - raise InvocationError("handler retry") - return "ok" - - class ObservingExecutor(ThreadPoolExecutor): - def submit( - self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any - ) -> Future[Any]: - if fn is not body and body not in args: - return super().submit(fn, *args, **kwargs) + class Observer(ThreadPoolExecutor): + def submit(self, fn: Any, /, *args: Any, **kwargs: Any) -> Any: + role = ( + "checkpoint" + if getattr(fn, "__name__", "") == "checkpoint_batches_forever" + else "invocation" + ) - def observe() -> Any: + def observed() -> Any: before = marker.get() try: return fn(*args, **kwargs) finally: - # This runs outside Context.run, on the actual SDK worker. - worker_boundaries.append((before, marker.get())) + boundary.append((role, threading.get_ident(), before, marker.get())) - return super().submit(observe) + return super().submit(observed) monkeypatch.setattr( - "aws_durable_execution_sdk_python.execution.ThreadPoolExecutor", - ObservingExecutor, + "aws_durable_execution_sdk_python.execution.ThreadPoolExecutor", Observer ) - monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) - client = Mock() - handler = durable_execution( - body, - boto3_client=client, - plugins=[ClaimPlugin()] if plugin_mode != "none" else [], - ) - event = { - "DurableExecutionArn": "test-arn/handler-context", - "CheckpointToken": "test-token", - "InitialExecutionState": { - "Operations": [ - { - "Id": "handler-context", - "Type": "EXECUTION", - "Status": "STARTED", - "ExecutionDetails": {"InputPayload": "{}"}, - } - ], - "NextMarker": "", - }, - } - lambda_context = Mock() - lambda_context.aws_request_id = "context-request" - lambda_context.client_context = None - lambda_context.identity = None - lambda_context._epoch_deadline_time_in_ms = 0 - lambda_context.invoked_function_arn = "test-arn" - lambda_context.tenant_id = None - token = marker.set("caller") - try: - if outcome == "retry": - with pytest.raises(InvocationError, match="handler retry"): - handler(event, lambda_context) - else: - result = handler(event, lambda_context) - assert result["Status"] == ( - "SUCCEEDED" if outcome == "success" else "FAILED" - ) - assert marker.get() == "caller" - finally: - marker.reset(token) - assert len(worker_boundaries) == 1 - worker_before, worker_after = worker_boundaries[0] - if plugin_mode != "none": - assert seen == [ - "caller" if plugin_mode == "partial-failure" else "invocation-start" - ] - assert worker_after == worker_before - assert worker_after != "worker-mutation" - assert statuses == [ - { - "success": InvocationStatus.SUCCEEDED, - "failure": InvocationStatus.FAILED, - "retry": InvocationStatus.RETRY, - }[outcome] - ] - else: - # No plugin means the original direct worker call: caller bindings are - # absent, and mutations belong to the worker's own context. - assert seen == [worker_before] - assert seen != ["caller"] - assert worker_after == "worker-mutation" - assert statuses == [] - client.checkpoint_durable_execution.assert_not_called() - - -@pytest.mark.parametrize("failure", [None, "body", "enter", "exit"]) -def test_optional_handler_scopes_are_balanced_and_cannot_change_outcome( - failure: str | None, -) -> None: - from contextlib import contextmanager - from datetime import UTC, datetime - from collections.abc import Iterator - from aws_durable_execution_sdk_python.plugin import PluginExecutor - marker = contextvars.ContextVar("handler-scope", default="caller") - events: list[str] = [] - - class ScopePlugin(DurableInstrumentationPlugin): - def __init__(self, name: str): - self.name = name + class Plugin(DurableInstrumentationPlugin): + token: Any = None - @contextmanager - def handler_context(self, info: InvocationStartInfo) -> Iterator[None]: - assert info.execution_arn == "handler-scope" - events.append("enter-" + self.name) - if self.name == "inner" and failure == "enter": - raise ValueError("plugin entry failure") - token = marker.set(self.name) - try: - yield - finally: - marker.reset(token) - events.append("exit-" + self.name) - if self.name == "inner" and failure == "exit": - raise ValueError("plugin cleanup failure") - - executor = PluginExecutor([ScopePlugin("outer"), ScopePlugin("inner")]) - - def body() -> str: - assert marker.get() == ("outer" if failure == "enter" else "inner") - events.append("body") - if failure in ("body", "exit"): - raise RuntimeError("original handler failure") - return "ok" - - with executor.run(): - executor.on_invocation_start("handler-scope", True, datetime.now(UTC), None) - if failure in ("body", "exit"): - with pytest.raises(RuntimeError, match="original handler failure"): - executor.run_handler(body) - else: - assert executor.run_handler(body) == "ok" - assert marker.get() == "caller" - assert events == ["enter-outer", "enter-inner", "body"] + ( - ["exit-outer"] if failure == "enter" else ["exit-inner", "exit-outer"] - ) - - -def test_handler_accepts_plugin_without_optional_scope() -> None: - from datetime import UTC, datetime - from types import SimpleNamespace - from typing import cast - from aws_durable_execution_sdk_python.plugin import PluginExecutor - - legacy = cast( - DurableInstrumentationPlugin, - SimpleNamespace(on_invocation_start=lambda info: None), - ) - executor = PluginExecutor([legacy]) - with executor.run(): - executor.on_invocation_start("legacy", True, datetime.now(UTC), None) - assert executor.run_handler(lambda: "unchanged") == "unchanged" - - -@pytest.mark.parametrize("outcome", ["SUCCEEDED", "PENDING", "FAILED", "retry"]) -@pytest.mark.parametrize("reverse", [False, True]) -@pytest.mark.parametrize("hook_failure", [None, "start", "end"]) -def test_invocation_context_scopes_do_not_escape_to_host( - outcome: str, reverse: bool, hook_failure: str | None -) -> None: - """Legacy hook order must not leave an already-ended plugin scope current.""" - from datetime import UTC, datetime - import threading - from aws_durable_execution_sdk_python.plugin import PluginExecutor - - marker = contextvars.ContextVar("invocation-scope", default="host") - events: list[tuple[str, str, str, int]] = [] - caller_thread = threading.get_ident() - - class ScopePlugin(DurableInstrumentationPlugin): - def __init__(self, name: str) -> None: - self.name = name - self.token: contextvars.Token[str] | None = None - - def on_invocation_start(self, _info: InvocationStartInfo) -> None: - events.append(("start", self.name, marker.get(), threading.get_ident())) - self.token = marker.set(self.name) - if self.name == names[0] and hook_failure == "start": - raise ValueError("plugin initialization failed") + def on_invocation_start(self, info: InvocationStartInfo) -> None: + phases.append(("start", threading.get_ident(), marker.get())) + self.token = marker.set("start-binding") + if mode == "failed-start": + raise ValueError("startup failure") - def on_invocation_end(self, _info: InvocationEndInfo) -> None: - events.append(("end", self.name, marker.get(), threading.get_ident())) - if self.name == names[0] and hook_failure == "end": - raise ValueError("plugin finalization failed") + def on_invocation_end(self, info: InvocationEndInfo) -> None: + phases.append(("end", threading.get_ident(), marker.get())) + statuses.append(info.status) assert self.token is not None marker.reset(self.token) self.token = None + if mode == "failed-end": + raise ValueError("end failure") - names = ["first", "second"] - if reverse: - names.reverse() - executor = PluginExecutor([ScopePlugin(name) for name in names]) - handler_failure = InvocationError("retry") - expected_output = {"Status": outcome} - - @executor.handle_durable_output - def invoke(_event: Any, _context: Any) -> dict[str, str]: - executor.on_invocation_start("test", True, datetime.now(UTC), None) - assert executor.run_handler(marker.get) == names[-1] - if outcome == "retry": - raise handler_failure - return expected_output - - token = marker.set("incoming") + def body(_event: Any, _context: DurableContext) -> str: + phases.append(("body", threading.get_ident(), marker.get())) + marker.set("body-binding") + try: + if outcome == "failure": + raise fail + if outcome == "retry": + raise retry + if outcome == "pending": + raise SuspendExecution("test suspension") + return "ok" + finally: + phases.append(("finally", threading.get_ident(), marker.get())) + + handler = durable_execution(body, plugins=[] if mode == "none" else [Plugin()]) + event, context = invocation() + token = marker.set("host") try: for _ in range(2): if outcome == "retry": - with pytest.raises(InvocationError, match="retry") as caught: - invoke({}, None) - assert caught.value is handler_failure + with pytest.raises(InvocationError) as caught: + handler(event, context) + assert caught.value is retry else: - assert invoke({}, None) is expected_output - assert marker.get() == "incoming" - finally: - marker.reset(token) - assert [(kind, name) for kind, name, _, _ in events] == [ - (kind, name) for _ in range(2) for kind in ("start", "end") for name in names - ] - assert all(thread == caller_thread for _, _, _, thread in events) - assert [ - value for kind, name, value, _ in events if kind == "start" and name == names[0] - ] == ["incoming", "incoming"] - - -def test_no_plugin_invocation_keeps_original_caller_context_semantics() -> None: - from aws_durable_execution_sdk_python.plugin import PluginExecutor - - marker = contextvars.ContextVar("no-plugin-caller", default="host") - executor = PluginExecutor([]) - - @executor.handle_durable_output - def invoke(_event: Any, _context: Any) -> dict[str, str]: - marker.set("caller-side-change") - return {"Status": "SUCCEEDED"} - - token = marker.set("incoming") - try: - invoke({}, None) - assert marker.get() == "caller-side-change" + result = handler(event, context) + assert ( + result["Status"] + == { + "success": "SUCCEEDED", + "failure": "FAILED", + "pending": "PENDING", + }[outcome] + ) + if outcome == "failure": + assert result["Error"]["ErrorMessage"] == "original body failure" + assert marker.get() == "host" finally: marker.reset(token) + per_call = 2 if mode == "none" else 4 + for offset in range(0, len(phases), per_call): + batch = phases[offset : offset + per_call] + assert [p[0] for p in batch] == ( + ["body", "finally"] + if mode == "none" + else ["start", "body", "finally", "end"] + ) + assert len({p[1] for p in batch}) == 1 + assert batch[0][1] != caller_thread + bodies = [p for p in phases if p[0] == "body"] + expected = ( + "fresh-worker" + if mode == "none" + else "host" + if mode == "failed-start" + else "start-binding" + ) + assert [p[2] for p in bodies] == [expected, expected] + for role, tid, before, after in boundary: + if role == "invocation": + assert after == ("body-binding" if mode == "none" else before) + if mode != "none": + assert ( + statuses + == [ + { + "success": InvocationStatus.SUCCEEDED, + "failure": InvocationStatus.FAILED, + "retry": InvocationStatus.RETRY, + "pending": InvocationStatus.PENDING, + }[outcome] + ] + * 2 + ) + assert "different Context" not in caplog.text -@pytest.mark.parametrize("stage", ["start", "factory", "enter"]) @pytest.mark.parametrize("bad_first", [False, True]) -@pytest.mark.parametrize("outcome", ["SUCCEEDED", "PENDING", "FAILED", "retry"]) -def test_failed_plugin_setup_discards_partial_context_bindings( - stage: str, bad_first: bool, outcome: str, caplog: pytest.LogCaptureFixture +@pytest.mark.parametrize("outcome", ["success", "failure", "retry", "pending"]) +def test_failed_start_preserves_clean_bindings_and_original_token_owners( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + bad_first: bool, + outcome: str, ) -> None: - """Keep healthy bindings, unset bindings, hook order, and token ownership.""" - from contextlib import nullcontext - from datetime import UTC, datetime - from typing import ContextManager - from aws_durable_execution_sdk_python.plugin import PluginExecutor - - marker = contextvars.ContextVar("partial-setup", default="default") - new_binding = contextvars.ContextVar[str]("partial-setup-no-default") - events: list[tuple[str, str]] = [] - cleanup: list[str] = [] - healthy_inputs: list[tuple[str, str | None]] = [] - - class Scope: - def __init__(self, plugin: SetupPlugin) -> None: - self.plugin = plugin - - def __enter__(self) -> None: - self.plugin.bind("enter") + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + marker = contextvars.ContextVar("failed-start-owner", default="host") + added = contextvars.ContextVar[str]("failed-start-added") + events: list[tuple[str, str, int]] = [] + cleanup_contexts: list[str] = [] + original_close = ExecutionState.close + + def close(state: ExecutionState) -> None: + cleanup_contexts.append(marker.get()) + with pytest.raises(LookupError): + added.get() + original_close(state) - def __exit__(self, *_args: Any) -> None: - self.plugin.reset("exit") + monkeypatch.setattr(ExecutionState, "close", close) - class SetupPlugin(DurableInstrumentationPlugin): + class Plugin(DurableInstrumentationPlugin): def __init__(self, name: str) -> None: self.name = name self.token: contextvars.Token[str] | None = None + self.added: contextvars.Token[str] | None = None - def bind(self, where: str) -> None: - events.append((where, self.name)) - if self.name == "healthy": - healthy_inputs.append((marker.get(), new_binding.get(None))) + def on_invocation_start(self, _info: InvocationStartInfo) -> None: + events.append(("start", self.name, threading.get_ident())) self.token = marker.set(self.name) if self.name == "bad": - new_binding.set("partial") - raise ValueError("partial plugin setup") + self.added = added.set("partial") + raise ValueError("partial startup") - def reset(self, where: str) -> None: - events.append((where, self.name)) + def on_invocation_end(self, _info: InvocationEndInfo) -> None: + events.append(("end", self.name, threading.get_ident())) assert self.token is not None marker.reset(self.token) - self.token = None - cleanup.append(self.name) - - def on_invocation_start(self, _info: InvocationStartInfo) -> None: - if stage == "start": - self.bind("start") - - def on_invocation_end(self, _info: InvocationEndInfo) -> None: - if stage == "start": - self.reset("end") + if self.added is not None: + added.reset(self.added) - def handler_context(self, _info: InvocationStartInfo) -> ContextManager[None]: - if stage == "start": - return nullcontext() - if self.name == "bad" and stage == "factory": - self.bind("factory") - return Scope(self) + failure = InvocationError("retry") - names = ["bad", "healthy"] if bad_first else ["healthy", "bad"] - executor = PluginExecutor([SetupPlugin(name) for name in names]) - output = {"Status": outcome} - failure = InvocationError("original retry") - handler_failure = RuntimeError("original handler error") - - def body() -> dict[str, str]: + def body(_event: Any, _context: DurableContext) -> str: + events.append(("body", "handler", threading.get_ident())) assert marker.get() == "healthy" with pytest.raises(LookupError): - new_binding.get() + added.get() if outcome == "retry": raise failure - if outcome == "FAILED": - raise handler_failure - return output - - @executor.handle_durable_output - def invoke(_event: Any, _context: Any) -> dict[str, str]: - executor.on_invocation_start("partial-setup", True, datetime.now(UTC), None) - try: - return executor.run_handler(body) - except RuntimeError as error: - assert error is handler_failure - return output + if outcome == "failure": + raise ValueError("body") + if outcome == "pending": + raise SuspendExecution("test suspension") + return "ok" - token = marker.set("incoming") - try: - for _ in range(2): - if outcome == "retry": - with pytest.raises(InvocationError) as caught: - invoke({}, None) - assert caught.value is failure - else: - assert invoke({}, None) is output - assert marker.get() == "incoming" - with pytest.raises(LookupError): - new_binding.get() - finally: - marker.reset(token) - assert healthy_inputs == [("incoming", None)] * 2 - assert cleanup == (names if stage == "start" else ["healthy"]) * 2 - setup = "start" if stage == "start" else "enter" - expected_setup = [ - ("factory" if name == "bad" and stage == "factory" else setup, name) - for name in names - ] - expected_cleanup = ( - [("end", name) for name in names] if stage == "start" else [("exit", "healthy")] - ) - assert events == (expected_setup + expected_cleanup) * 2 - # A token reset in the wrong Context is caught by the SDK, so explicitly - # check diagnostics and successful cleanup rather than relying on raises. + names = ["bad", "healthy"] if bad_first else ["healthy", "bad"] + handler = durable_execution(body, plugins=[Plugin(name) for name in names]) + event, context = invocation() + for _ in range(2): + if outcome == "retry": + with pytest.raises(InvocationError) as caught: + handler(event, context) + assert caught.value is failure + else: + assert ( + handler(event, context)["Status"] + == {"success": "SUCCEEDED", "failure": "FAILED", "pending": "PENDING"}[ + outcome + ] + ) + assert marker.get() == "host" + with pytest.raises(LookupError): + added.get() + assert cleanup_contexts == ["healthy"] * 2 + for offset in (0, 5): + batch = events[offset : offset + 5] + assert [(x[0], x[1]) for x in batch] == [("start", n) for n in names] + [ + ("body", "handler") + ] + [("end", n) for n in names] + assert len({x[2] for x in batch}) == 1 errors = [ - record.exc_info for record in caplog.records if record.exc_info is not None + r.exc_info + for r in caplog.records + if r.exc_info and r.name == "aws_durable_execution_sdk_python.plugin" ] - assert len(errors) == 2 - assert all(str(error[1]) == "partial plugin setup" for error in errors) + assert len(errors) == 2 and all(str(e[1]) == "partial startup" for e in errors) -@pytest.mark.parametrize("inherited", [False, True]) -@pytest.mark.parametrize("fail_handler", [False, True]) -def test_handler_scope_uses_direct_or_inherited_optional_method( - inherited: bool, fail_handler: bool +@pytest.mark.parametrize("result_kind", ["normal", "bad-json", "large", "large-error"]) +def test_worker_prepares_output_joins_branches_waits_checkpoint_then_ends( + monkeypatch: pytest.MonkeyPatch, + result_kind: str, ) -> None: - from contextlib import contextmanager - from collections.abc import Iterator - from datetime import UTC, datetime - from aws_durable_execution_sdk_python.plugin import PluginExecutor - - marker = contextvars.ContextVar("inherited-scope", default="outside") - events: list[str] = [] - error = RuntimeError("original handler failure") - - class ScopePlugin(DurableInstrumentationPlugin): - @contextmanager - def handler_context(self, _info: InvocationStartInfo) -> Iterator[None]: - events.append("enter") - token = marker.set("inside") - try: - yield - finally: - marker.reset(token) - events.append("exit") - - class InheritedPlugin(ScopePlugin): - pass - - plugin = InheritedPlugin() if inherited else ScopePlugin() - executor = PluginExecutor([plugin]) - - def handler() -> str: - assert marker.get() == "inside" - events.append("body") - if fail_handler: - raise error - return "ok" + import aws_durable_execution_sdk_python.execution as execution - with executor.run(): - for _ in range(2): - executor.on_invocation_start("inherited", True, datetime.now(UTC), None) - if fail_handler: - with pytest.raises(RuntimeError) as caught: - executor.run_handler(handler) - assert caught.value is error - else: - assert executor.run_handler(handler) == "ok" - assert marker.get() == "outside" - assert events == ["enter", "body", "exit"] * 2 + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + events: list[tuple[str, int]] = [] + started = threading.Event() + checkpoint_started = threading.Event() + closing = threading.Event() + branch_done = threading.Event() + real_checkpoint = ExecutionState.checkpoint_batches_forever + real_close = ExecutionState.close + real_stop = ExecutionState.stop_checkpointing + real_dumps = execution.json.dumps + result: Any = { + "normal": {"ok": True}, + "bad-json": object(), + "large": "x" * 256, + "large-error": None, + }[result_kind] + + def record(name: str) -> None: + events.append((name, threading.get_ident())) + + def checkpoint(state: ExecutionState) -> None: + assert started.is_set() + record("checkpoint-start") + checkpoint_started.set() + try: + real_checkpoint(state) + finally: + record("checkpoint-end") + + def stop(state: ExecutionState) -> None: + assert branch_done.is_set() + record("checkpoint-stop") + real_stop(state) + + def close(state: ExecutionState) -> None: + record("close") + closing.set() + real_close(state) + + def dumps(value: Any, *args: Any, **kwargs: Any) -> str: + if value is result or ( + isinstance(value, dict) and value.get("Status") == "FAILED" + ): + record("serialize") + return real_dumps(value, *args, **kwargs) + + monkeypatch.setattr(ExecutionState, "checkpoint_batches_forever", checkpoint) + monkeypatch.setattr(ExecutionState, "stop_checkpointing", stop) + monkeypatch.setattr(ExecutionState, "close", close) + monkeypatch.setattr(execution.json, "dumps", dumps) + if result_kind in ("large", "large-error"): + monkeypatch.setattr(execution, "LAMBDA_RESPONSE_SIZE_LIMIT", 64) + + class Plugin(DurableInstrumentationPlugin): + def on_invocation_start(self, _info: InvocationStartInfo) -> None: + record("start") + assert not checkpoint_started.is_set() + started.set() + def on_invocation_end(self, info: InvocationEndInfo) -> None: + record("end") + assert branch_done.is_set() + assert any(x[0] == "checkpoint-end" for x in events) + assert info.status is ( + InvocationStatus.FAILED + if result_kind in ("bad-json", "large-error") + else InvocationStatus.SUCCEEDED + ) -@pytest.mark.parametrize("value", [None, "not a method", 42]) -def test_noncallable_optional_handler_scope_is_ignored(value: Any) -> None: - from datetime import UTC, datetime - from types import SimpleNamespace - from typing import cast - from aws_durable_execution_sdk_python.plugin import PluginExecutor + def body(_event: Any, ctx: DurableContext) -> Any: + assert checkpoint_started.wait(5) + record("body") + pool = ThreadPoolExecutor(max_workers=1) + ctx.state.register_branch_pool(pool) + + def late_branch() -> None: + assert closing.wait(5) + assert not ctx.state._checkpointing_stopped.is_set() + ctx.state.create_checkpoint( + OperationUpdate.create_execution_succeed(payload='"branch"'), + is_sync=True, + ) + record("branch-done") + branch_done.set() - plugin = cast( - DurableInstrumentationPlugin, - SimpleNamespace(on_invocation_start=lambda info: None, handler_context=value), + pool.submit(late_branch) + try: + if result_kind == "large-error": + raise ValueError("x" * 256) + return result + finally: + record("finally") + + handler = durable_execution(body, plugins=[Plugin()]) + event, context = invocation() + output = handler(event, context) + record("caller") + names = [x[0] for x in events] + assert ( + names.index("start") + < names.index("checkpoint-start") + < names.index("body") + < names.index("finally") + < names.index("serialize") + < names.index("close") + ) + assert ( + names.index("branch-done") + < names.index("checkpoint-stop") + < names.index("checkpoint-end") + < names.index("end") + < names.index("caller") + ) + worker_events = { + tid + for name, tid in events + if name + in {"start", "body", "finally", "serialize", "close", "checkpoint-stop", "end"} + } + assert len(worker_events) == 1 and threading.get_ident() not in worker_events + assert output["Status"] == ( + "FAILED" if result_kind in ("bad-json", "large-error") else "SUCCEEDED" ) - executor = PluginExecutor([plugin]) - with executor.run(): - executor.on_invocation_start("noncallable", True, datetime.now(UTC), None) - assert executor.run_handler(lambda: "unchanged") == "unchanged" -def test_handler_scope_lookup_failure_discards_partial_bindings( +@pytest.mark.parametrize("shape", ["method", "property", "dynamic"]) +def test_removed_handler_scope_api_is_not_inspected( + monkeypatch: pytest.MonkeyPatch, shape: str +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + calls = [] + + def legacy(*_args: Any) -> Any: + calls.append("unexpected") + raise AssertionError("Removed API called") + + choices: dict[str, dict[str, Any]] = { + "method": {"handler_context": legacy}, + "property": {"handler_context": property(legacy)}, + "dynamic": { + "__getattr__": lambda self, name: legacy() + if name == "handler_context" + else (_ for _ in ()).throw(AttributeError(name)) + }, + } + plugin = type("LegacyPlugin", (DurableInstrumentationPlugin,), choices[shape])() + handler = durable_execution(lambda _e, _c: "ok", plugins=[plugin]) + event, context = invocation() + assert handler(event, context)["Status"] == "SUCCEEDED" + assert calls == [] + + +@pytest.mark.parametrize( + "checkpoint_path", ["step-start", "large-result", "large-error"] +) +@pytest.mark.parametrize("retryable", [False, True]) +def test_checkpoint_failure_reports_prepared_outcome_after_worker_cleanup( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, + checkpoint_path: str, + retryable: bool, ) -> None: - from datetime import UTC, datetime - from aws_durable_execution_sdk_python.plugin import PluginExecutor - - marker = contextvars.ContextVar("lookup-failure", default="outside") - - class BrokenPlugin(DurableInstrumentationPlugin): - @property - def handler_context(self) -> Any: - marker.set("partial") - raise ValueError("scope lookup failed") - - executor = PluginExecutor([BrokenPlugin()]) - with executor.run(): - executor.on_invocation_start("lookup", True, datetime.now(UTC), None) - assert executor.run_handler(marker.get) == "outside" - assert marker.get() == "outside" - assert "scope lookup failed" in caplog.text + import aws_durable_execution_sdk_python.execution as execution + + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + if checkpoint_path != "step-start": + monkeypatch.setattr(execution, "LAMBDA_RESPONSE_SIZE_LIMIT", 128) + marker = contextvars.ContextVar("checkpoint-worker", default="host") + failure = CheckpointError( + "actual checkpoint failure", + error_category=( + CheckpointErrorCategory.INVOCATION + if retryable + else CheckpointErrorCategory.EXECUTION + ), + ) + events: list[tuple[str, int]] = [] + ends: list[tuple[InvocationEndInfo, str]] = [] + real_checkpoint = ExecutionState.checkpoint_batches_forever + real_close = ExecutionState.close + + def record(name: str) -> None: + events.append((name, threading.get_ident())) + + def checkpoint(state: ExecutionState) -> None: + record("checkpoint-start") + try: + real_checkpoint(state) + finally: + record("checkpoint-end") + + def close(state: ExecutionState) -> None: + record("close") + assert marker.get() == "plugin" + real_close(state) + record("closed") + + def service_checkpoint(*_args: Any, **_kwargs: Any) -> Any: + record("service-failure") + raise failure + + monkeypatch.setattr(ExecutionState, "checkpoint_batches_forever", checkpoint) + monkeypatch.setattr(ExecutionState, "close", close) + + class Plugin(DurableInstrumentationPlugin): + token: contextvars.Token[str] | None = None + + def on_invocation_start(self, _info: InvocationStartInfo) -> None: + record("start") + self.token = marker.set("plugin") + + def on_invocation_end(self, info: InvocationEndInfo) -> None: + record("end") + ends.append((info, marker.get())) + assert self.token is not None + marker.reset(self.token) + self.token = None + + def body(_event: Any, ctx: DurableContext) -> str: + record("body") + try: + if checkpoint_path == "step-start": + ctx.step(lambda _step: "ok", name="failed-checkpoint") + if checkpoint_path == "large-error": + raise ValueError("x" * 256) + return "x" * 256 + finally: + record("finally") + + plugin = Plugin() + handler = durable_execution(body, plugins=[plugin]) + client = Mock(spec=DurableServiceClient) + client.checkpoint.side_effect = service_checkpoint + event, context = invocation(client) + for _ in range(2): + events.clear() + ends.clear() + if retryable: + with pytest.raises(CheckpointError) as caught: + handler(event, context) + assert caught.value is failure + else: + output = handler(event, context) + assert output["Status"] == "FAILED" + assert output["Error"]["ErrorMessage"] == str(failure) + assert output["Error"]["ErrorType"].endswith(".CheckpointError") + record("caller") + assert marker.get() == "host" and plugin.token is None + assert len(ends) == 1 + info, active = ends[0] + assert active == "plugin" + assert info.status is ( + InvocationStatus.RETRY if retryable else InvocationStatus.FAILED + ) + assert info.error is not None and info.error.message == str(failure) + assert info.error.type is not None and info.error.type.endswith( + ".CheckpointError" + ) + names = [name for name, _tid in events] + assert ( + names.index("start") + < names.index("checkpoint-start") + < names.index("service-failure") + ) + assert ( + names.index("body") + < names.index("finally") + < names.index("close") + < names.index("closed") + < names.index("end") + < names.index("caller") + ) + assert ( + names.index("service-failure") + < names.index("checkpoint-end") + < names.index("end") + ) + workers = { + tid + for name, tid in events + if name in {"start", "body", "finally", "close", "closed", "end"} + } + assert len(workers) == 1 and threading.get_ident() not in workers + assert not any( + record.exc_info and record.name == "aws_durable_execution_sdk_python.plugin" + for record in caplog.records + ) diff --git a/pyproject.toml b/pyproject.toml index bca7cb9e2..80eb21e0c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -131,66 +131,48 @@ dependencies = [ [tool.hatch.envs.dev-examples.scripts] test = "pytest packages/aws-durable-execution-sdk-python-examples/test {args}" -[tool.hatch.envs.test-pypi-otel] -# Test new exclusivity capability against an installed capable core. -# Legacy valid-registration combinations are verified separately against 2.0.x. -# Override inherited workspace membership: core must come from an installed distribution. -workspace.members = ["packages/aws-durable-execution-sdk-python-otel"] -dependencies = [ - "aws-durable-execution-sdk-python-testing>=1.2.1", - "aws-durable-execution-sdk-python>=2.1.0", - "opentelemetry-sdk>=1.20.0", - "opentelemetry-propagator-aws-xray", - "pytest", - "pytest-cov", - "coverage[toml]", -] -pre-install-commands = [ - "pip install -e packages/aws-durable-execution-sdk-python-otel", -] - -[tool.hatch.envs.test-pypi-otel.scripts] -test = "pytest packages/aws-durable-execution-sdk-python-otel/tests {args}" - -[tool.hatch.envs.test-pypi-otel-legacy] -template = "test-pypi-otel" -workspace.members = ["packages/aws-durable-execution-sdk-python-otel"] +# Install immutable artifacts built from this checkout, including before the +# first compatible core is published. No editable workspace packages in this lane. +[tool.hatch.envs.test-wheel-otel] +detached = true +workspace.members = [] dependencies = [ - "aws-durable-execution-sdk-python>=2.0.0,<2.1.0", "aws-durable-execution-sdk-python-testing>=1.2.1", "opentelemetry-sdk>=1.20.0", "opentelemetry-propagator-aws-xray", "pytest", - "pytest-cov", - "coverage[toml]", + "packaging", ] -[tool.hatch.envs.test-pypi-otel-legacy.scripts] +[tool.hatch.envs.test-wheel-otel.scripts] test = [ - "python -c 'from pathlib import Path; from importlib.metadata import version; from packaging.version import Version; import aws_durable_execution_sdk_python.execution as core; assert Version(version(\"aws-durable-execution-sdk-python\")).release[:2] == (2, 0); assert \"site-packages\" in Path(core.__file__).resolve().parts; print(version(\"aws-durable-execution-sdk-python\"), core.__file__)' ", - "pytest packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_view_registration.py -k 'one_view_and_unrelated_plugin_suspend_resume or no_otel_plugin_remains_valid or execution_constructor_retains_ambient_log_correlation' {args}" + "python .github/scripts/install_otel_test_wheels.py", + "pytest packages/aws-durable-execution-sdk-python-otel/tests {args}", + "pytest .github/tests/otel_lifecycle_compatibility_test.py {args}", ] -# Keep every pre-existing OTel test on the minimum core; the newly added view -# registration file has capability-only cases plus the legacy-valid cases below. -[tool.hatch.envs.test-pypi-otel-minimum] -template = "test-pypi-otel" -workspace.members = ["packages/aws-durable-execution-sdk-python-otel"] +# A core-only upgrade retains the released plugin's more limited Invocation +# fallback. Validate the real published plugin separately from the new pair. +[tool.hatch.envs.test-wheel-otel-legacy] +template = "test-wheel-otel" +detached = true +workspace.members = [] dependencies = [ - "aws-durable-execution-sdk-python==2.0.0", + "aws-durable-execution-sdk-python-otel==1.0.0", "aws-durable-execution-sdk-python-testing>=1.2.1", "opentelemetry-sdk>=1.20.0", "opentelemetry-propagator-aws-xray", "pytest", - "pytest-cov", - "coverage[toml]", + "packaging", ] -[tool.hatch.envs.test-pypi-otel-minimum.scripts] +[tool.hatch.envs.test-wheel-otel-legacy.env-vars] +OTEL_COMPAT_LEGACY = "1" + +[tool.hatch.envs.test-wheel-otel-legacy.scripts] test = [ - "python -c 'from pathlib import Path; from importlib.metadata import version; from packaging.version import Version; import aws_durable_execution_sdk_python.execution as core; assert version(\"aws-durable-execution-sdk-python\") == \"2.0.0\"; assert \"site-packages\" in Path(core.__file__).resolve().parts; print(version(\"aws-durable-execution-sdk-python\"), core.__file__)' ", - "pytest packages/aws-durable-execution-sdk-python-otel/tests --ignore=packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_view_registration.py {args}", - "pytest packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_view_registration.py -k 'one_view_and_unrelated_plugin_suspend_resume or no_otel_plugin_remains_valid or execution_constructor_retains_ambient_log_correlation' {args}" + "python .github/scripts/install_otel_test_wheels.py --legacy-plugin", + "pytest .github/tests/otel_lifecycle_compatibility_test.py {args}", ] [tool.hatch.envs.test-pypi-examples] From d415329e80b36425f941f4af12833eec1a16922e Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 9 Oct 2026 00:47:16 +0000 Subject: [PATCH 29/54] fix: propagate plugin context into concurrent branches --- .../README.md | 6 +- .../tests/e2e/test_branch_context_int.py | 179 ++++++++++++++++++ .../README.md | 5 + .../concurrency/executor.py | 27 ++- .../tests/e2e/branch_worker_context_test.py | 129 +++++++++++++ .../tests/operation/map_test.py | 29 +++ .../tests/operation/parallel_test.py | 29 +++ 7 files changed, 396 insertions(+), 8 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_branch_context_int.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/e2e/branch_worker_context_test.py diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index f42b5986c..2a095bf9c 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -200,7 +200,11 @@ Start hooks run in registration order. A later successful plugin that deliberate sets or clears the active span wins; OTel does not apply a second correction pass. Place OTel after a span-replacing plugin when OTel's view-specific context is desired. Baggage-only plugins that extend the current context can appear on either -side. End hooks also retain registration order and reset tokens in their owning +side. The SDK also copies the coordinator's current bindings for each `map` or +`parallel` branch admission/resume when plugins are registered. Baggage and +other successful bindings reach these SDK-managed callbacks, while a branch's +changes cannot leak into siblings, the coordinator or a reused worker. +End hooks also retain registration order and reset tokens in their owning Context; they do not promise a reverse-stack observation of other plugins' spans. Existing registration, factory lifetime and checkpoint formats are unchanged. diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_branch_context_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_branch_context_int.py new file mode 100644 index 000000000..c2ae779af --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_branch_context_int.py @@ -0,0 +1,179 @@ +"""Public SDK branch workers inherit instrumented bindings without sharing them.""" + +from __future__ import annotations + +import contextvars +import json +import threading +from typing import Any + +import pytest +from aws_durable_execution_sdk_python import DurableContext, durable_execution +from aws_durable_execution_sdk_python.config import Duration, MapConfig, ParallelConfig +from aws_durable_execution_sdk_python.concurrency.models import BatchResult +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationStatus, +) +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry import baggage, context, trace +from opentelemetry.sdk.trace import TracerProvider + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +@pytest.mark.parametrize( + ("view", "failed_start"), + [ + (None, False), + (InvocationOtelPlugin, False), + (ExecutionOtelPlugin, False), + (InvocationOtelPlugin, True), + (ExecutionOtelPlugin, True), + ], +) +@pytest.mark.parametrize("kind", ["map", "parallel"]) +@pytest.mark.parametrize("concurrency", [1, 2]) +@pytest.mark.parametrize("resume", [False, True]) +def test_public_branch_context_propagation_and_isolation( + monkeypatch: pytest.MonkeyPatch, + view: type[InvocationOtelPlugin] | type[ExecutionOtelPlugin] | None, + failed_start: bool, + caplog: pytest.LogCaptureFixture, + kind: str, + concurrency: int, + resume: bool, +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + marker = contextvars.ContextVar("branch-invocation-marker", default="unset") + observations: list[tuple[int, str, Any, bool]] = [] + bodies: list[int] = [] + visits = {0: 0, 1: 0} + statuses: list[InvocationStatus] = [] + starts: list[int] = [] + completed_ends: list[int] = [] + lock = threading.Lock() + barrier = threading.Barrier(2) if concurrency == 2 else None + provider = TracerProvider() + host = context.get_current() + + class Bind(DurableInstrumentationPlugin): + def on_invocation_start(self, info: Any) -> None: + starts.append(threading.get_ident()) + self.token = marker.set("plugin") + self.bag = context.attach(baggage.set_baggage("tenant", "present")) + + def on_invocation_end(self, info: Any) -> None: + context.detach(self.bag) + marker.reset(self.token) + # Record completion only after both tokens were reset in their owner. + completed_ends.append(threading.get_ident()) + statuses.append(info.status) + + class Broken(DurableInstrumentationPlugin): + def on_invocation_start(self, _info: Any) -> None: + marker.set("partial") + context.attach(baggage.set_baggage("tenant", "partial")) + raise ValueError("expected failed Start") + + plugins: list[DurableInstrumentationPlugin] = [] + otel_plugin: InvocationOtelPlugin | ExecutionOtelPlugin | None = None + if view is not None: + plugins.append(Bind()) + if failed_start: + plugins.append(Broken()) + otel_plugin = view( + OtelPluginConfig( + tracer_provider=provider, + enrich_logger=False, + context_extractor=lambda _: None, + ) + ) + plugins.append(otel_plugin) + + def branch(child: DurableContext, index: int) -> str: + with lock: + visits[index] += 1 + first_visit = visits[index] == 1 + observations.append( + ( + index, + marker.get(), + baggage.get_baggage("tenant"), + trace.get_current_span().get_span_context().is_valid, + ) + ) + if view is not None: + # Deliberately leave a binding behind: another logical branch or + # resume on this pool thread must still start with its parent's copy. + marker.set(f"branch-{index}") + if barrier is not None and first_visit: + barrier.wait(timeout=10) + + def step(_step: Any) -> str: + with lock: + bodies.append(index) + return f"saved-{index}" + + saved = child.step(step, name="save") + return saved + + def handler(_event: Any, durable: DurableContext) -> list[str]: + expected = "unset" if view is None else "plugin" + assert marker.get() == expected + token = marker.set("handler" if view is None else "plugin") + try: + result: BatchResult[str] = ( + durable.map( + [0, 1], + lambda child, item, index, items: branch(child, index), + name="mapped", + config=MapConfig(max_concurrency=concurrency), + ) + if kind == "map" + else durable.parallel( + [lambda child: branch(child, 0), lambda child: branch(child, 1)], + name="parallel", + config=ParallelConfig(max_concurrency=concurrency), + ) + ) + assert marker.get() == ("handler" if view is None else "plugin") + if resume: + durable.wait(Duration.from_seconds(1), name="resume") + return result.get_results() + finally: + marker.reset(token) + + wrapped = durable_execution(handler, plugins=plugins) + try: + with DurableFunctionTestRunner(handler=wrapped) as runner: + result = runner.run(input="{}", timeout=30) + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == ["saved-0", "saved-1"] + assert sorted(bodies) == [0, 1] + assert {item[0] for item in observations} == {0, 1} + expected = "unset" if view is None else "plugin" + assert all( + item[1:] + == (expected, None if view is None else "present", view is not None) + for item in observations + ) + # The final wait replays the completed batch without rerunning branches. + assert visits == {0: 1, 1: 1} + if view is not None and resume: + assert InvocationStatus.PENDING in statuses + assert starts == completed_ends + if view is not None: + assert starts and statuses[-1] is InvocationStatus.SUCCEEDED + assert otel_plugin is not None and otel_plugin._context_tokens == {} + assert not any( + r.exc_info and r.name == "opentelemetry.context" for r in caplog.records + ) + assert marker.get() == "unset" + assert context.get_current() == host + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python/README.md b/packages/aws-durable-execution-sdk-python/README.md index 08a70e06a..4b2463e52 100644 --- a/packages/aws-durable-execution-sdk-python/README.md +++ b/packages/aws-durable-execution-sdk-python/README.md @@ -129,6 +129,11 @@ replace an earlier binding. If a Start hook raises, its new bindings are discard for subsequent work; its End still runs in the original Context so its tokens can be reset. End hooks retain forward registration order, not reverse stack order, so they must not rely on observing a stack-like unwind of other plugins' contexts. +Each SDK-managed `map` or `parallel` branch admission, including an in-process +resume, receives a fresh copy of the coordinator's bindings when plugins are +registered. Branch changes remain local even when pool threads are reused. +User-created threads retain normal Python context-variable behavior; without +registered plugins, SDK branch submission keeps its existing behavior. The worker's invocation context is discarded on return, including after plugin cleanup failures, leaving the host's bindings unchanged. This isolates bindings, not mutations to shared objects or external side effects. Without plugins, the diff --git a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py index 2a39c4eff..bafa43b19 100644 --- a/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py +++ b/packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py @@ -2,6 +2,7 @@ from __future__ import annotations +import contextvars import heapq import logging import queue @@ -291,13 +292,25 @@ def execute( def submit(branch: Branch[CallableType, ResultType]) -> None: branch.start() - pool.submit( - self._branch_worker, - execution_state, - executor_context, - events, - branch.executable, - ) + if execution_state._plugin_executor._plugins: # noqa: SLF001 + # Every admission/resume gets its own Context. Branch bindings + # cannot leak into siblings, the coordinator, or reused workers. + pool.submit( + contextvars.copy_context().run, + self._branch_worker, + execution_state, + executor_context, + events, + branch.executable, + ) + else: + pool.submit( + self._branch_worker, + execution_state, + executor_context, + events, + branch.executable, + ) try: # Only rebuild the items snapshot after a terminal event changes diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/branch_worker_context_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/branch_worker_context_test.py new file mode 100644 index 000000000..a9104a745 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/branch_worker_context_test.py @@ -0,0 +1,129 @@ +"""Registered Start bindings reach independently resumed public branch workers.""" + +from __future__ import annotations + +import contextvars +import json +import threading +from typing import Any + +import pytest +from aws_durable_execution_sdk_python import DurableContext, durable_execution +from aws_durable_execution_sdk_python.config import Duration, MapConfig, ParallelConfig +from aws_durable_execution_sdk_python.concurrency.models import BatchResult +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationEndInfo, + InvocationStartInfo, + InvocationStatus, +) +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner + + +@pytest.mark.parametrize("kind", ["map", "parallel"]) +@pytest.mark.parametrize("concurrency", [1, 2]) +@pytest.mark.parametrize("failed_start", [False, True]) +def test_partial_branch_resume_uses_fresh_successful_start_bindings( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + kind: str, + concurrency: int, + failed_start: bool, +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + marker = contextvars.ContextVar("resumed-branch", default="host") + partial = contextvars.ContextVar[str]("partial-start") + observations: list[tuple[int, str]] = [] + visits = {0: 0, 1: 0} + bodies: list[int] = [] + statuses: list[InvocationStatus] = [] + lock = threading.Lock() + barrier = threading.Barrier(2) if concurrency == 2 else None + + class Bind(DurableInstrumentationPlugin): + token: contextvars.Token[str] | None = None + + def on_invocation_start(self, _info: InvocationStartInfo) -> None: + self.token = marker.set("plugin") + + def on_invocation_end(self, info: InvocationEndInfo) -> None: + statuses.append(info.status) + assert self.token is not None + marker.reset(self.token) + self.token = None + + class Broken(DurableInstrumentationPlugin): + def on_invocation_start(self, _info: InvocationStartInfo) -> None: + marker.set("failed") + partial.set("failed") + raise ValueError("expected failed Start") + + plugin = Bind() + plugins: list[DurableInstrumentationPlugin] = [plugin] + if failed_start: + plugins.append(Broken()) + + def branch(child: DurableContext, index: int) -> str: + with lock: + visits[index] += 1 + first_visit = visits[index] == 1 + observations.append((index, marker.get())) + with pytest.raises(LookupError): + partial.get() + marker.set(f"branch-{index}") + if barrier is not None and first_visit: + barrier.wait(timeout=10) + + def step(_step: Any) -> str: + with lock: + bodies.append(index) + return f"saved-{index}" + + saved = child.step(step, name="save") + child.wait(Duration.from_seconds(index + 1), name="branch-wait") + return saved + + def handler(_event: Any, durable: DurableContext) -> list[str]: + assert marker.get() == "plugin" + result: BatchResult[str] + if kind == "map": + result = durable.map( + [0, 1], + lambda child, item, index, items: branch(child, index), + name="mapped", + config=MapConfig(max_concurrency=concurrency), + ) + else: + result = durable.parallel( + [lambda child: branch(child, 0), lambda child: branch(child, 1)], + name="parallel", + config=ParallelConfig(max_concurrency=concurrency), + ) + assert marker.get() == "plugin" + durable.wait(Duration.from_seconds(1), name="after-batch") + return result.get_results() + + wrapped = durable_execution(handler, plugins=plugins) + with DurableFunctionTestRunner(handler=wrapped) as runner: + result = runner.run(input="{}", timeout=30) + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == ["saved-0", "saved-1"] + assert sorted(bodies) == [0, 1] + assert all(count >= 2 for count in visits.values()) + assert all(binding == "plugin" for _, binding in observations) + assert InvocationStatus.PENDING in statuses + assert statuses[-1] is InvocationStatus.SUCCEEDED + assert marker.get() == "host" and plugin.token is None + with pytest.raises(LookupError): + partial.get() + errors = [ + record + for record in caplog.records + if record.exc_info and record.name == "aws_durable_execution_sdk_python.plugin" + ] + assert len(errors) == (len(statuses) if failed_start else 0) + assert all( + record.exc_info is not None + and str(record.exc_info[1]) == "expected failed Start" + for record in errors + ) diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/map_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/map_test.py index 2670e6747..926b1dd0b 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/map_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/map_test.py @@ -7,6 +7,7 @@ import pytest # Mock the executor.execute method +from aws_durable_execution_sdk_python.plugin import PluginExecutor from aws_durable_execution_sdk_python.concurrency.models import ( BatchItem, BatchItemStatus, @@ -184,6 +185,8 @@ def mock_run_in_child_context(func, name, config): # Create a minimal ExecutionState mock class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -222,6 +225,8 @@ def mock_run_in_child_context(func, name, config): return func("mock_context") class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -354,6 +359,8 @@ def callable_func(ctx, item, idx, items): ) as mock_execute: class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -407,6 +414,8 @@ def callable_func(ctx, item, idx, items): executor_context.create_child_context = lambda *args, **kwargs: Mock() class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -461,6 +470,8 @@ def callable_func(ctx, item, idx, items): executor_context.create_child_context = lambda *args, **kwargs: child_context class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -510,6 +521,8 @@ def mock_summary_generator(result): executor_context.create_child_context = Mock(return_value=_child_ctx) class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -550,6 +563,8 @@ def callable_func(ctx, item, idx, items): executor_context.create_child_context = Mock(return_value=Mock()) # SLF001 class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -589,6 +604,8 @@ def func(ctx, item, index, array): config = MapConfig(summary_generator=None) class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -634,6 +651,8 @@ def callable_func(ctx, item, idx, items): # Mock execution state that indicates operation already succeeded class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -703,6 +722,8 @@ def callable_func(ctx, item, idx, items): # Mock execution state that indicates operation succeeded but children need replay class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -802,6 +823,8 @@ def test_func(ctx, item, idx, items): execution_count = 0 class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -1032,6 +1055,8 @@ def func(ctx, item, idx, items): return {"item": item.upper(), "index": idx} class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -1481,6 +1506,8 @@ def test_map_handler_defaults_summary_generator_for_user_config(): """ class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -1556,6 +1583,8 @@ def predicate(s: CompletionStatus) -> CompletionDecision: ) as mock_execute: class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/parallel_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/parallel_test.py index 12231cfab..87d4d5a89 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/parallel_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/parallel_test.py @@ -8,6 +8,7 @@ import pytest +from aws_durable_execution_sdk_python.plugin import PluginExecutor from aws_durable_execution_sdk_python.concurrency.executor import ConcurrentExecutor from aws_durable_execution_sdk_python.identifier import OperationIdNamespace @@ -211,6 +212,8 @@ def func2(ctx): config = ParallelConfig(max_concurrency=2) class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -255,6 +258,8 @@ def func1(ctx): callables = [func1] class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -299,6 +304,8 @@ def func1(ctx): config = ParallelConfig(max_concurrency=5) class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -351,6 +358,8 @@ def func1(ctx): callables = [func1] class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -466,6 +475,8 @@ def func1(ctx): callables = [func1] class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -513,6 +524,8 @@ def mock_summary_generator(result): config = ParallelConfig(summary_generator=mock_summary_generator) class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -556,6 +569,8 @@ def func2(ctx): callables = [func1, func2] class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -606,6 +621,8 @@ def func3(ctx): config = ParallelConfig(summary_generator=None) class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -654,6 +671,8 @@ def func2(ctx): # Mock execution state that indicates operation already succeeded class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -723,6 +742,8 @@ def func1(ctx): # Mock execution state that indicates operation succeeded but children need replay class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -824,6 +845,8 @@ def task2(ctx): execution_count = 0 class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -1052,6 +1075,8 @@ def func3(ctx): callables = [func1, func2, func3] class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -1474,6 +1499,8 @@ def test_parallel_handler_defaults_summary_generator_for_user_config(): """A user config without a summary generator gets the default (JS parity).""" class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass @@ -1551,6 +1578,8 @@ def predicate(s: CompletionStatus) -> CompletionDecision: ) as mock_execute: class MockExecutionState: + _plugin_executor = PluginExecutor([]) + def register_branch_pool(self, pool): pass From 192ce1129c03699d5a59487776bc716203dd8a02 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 9 Oct 2026 01:31:15 +0000 Subject: [PATCH 30/54] test: cover instrumented branch context handoff --- .../tests/concurrency_test.py | 47 ++++++++++++++++++- 1 file changed, 46 insertions(+), 1 deletion(-) diff --git a/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py b/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py index a93135663..c30980c73 100644 --- a/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/concurrency_test.py @@ -1,5 +1,6 @@ """Tests for the concurrency module.""" +import contextvars import hashlib import json import queue @@ -77,7 +78,10 @@ from aws_durable_execution_sdk_python.operation.parallel import ( ParallelExecutor, ) -from aws_durable_execution_sdk_python.plugin import PluginExecutor +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + PluginExecutor, +) from aws_durable_execution_sdk_python.state import ( CheckpointedResult, ExecutionState, @@ -5275,3 +5279,44 @@ def predicate(s: CompletionStatus) -> CompletionDecision: # endregion Custom completion predicate (should_complete) integration tests + + +def test_instrumented_executor_isolates_bindings_on_a_reused_worker() -> None: + """Exercise actual admission/worker execution independently of durable I/O.""" + binding = contextvars.ContextVar("executor-unit-binding", default="empty") + seen: list[tuple[str, int]] = [] + + class RecordingExecutor(ConcurrentExecutor[Callable[[], str], str]): + def _execute_item_in_child_context( + self, + executor_context: DurableContext, + executable: Executable[Callable[[], str]], + ) -> str: + seen.append((binding.get(), threading.get_ident())) + binding.set(f"branch-{executable.index}") + return executable.func() + + executor = RecordingExecutor( + executables=[Executable(index, lambda: "ok") for index in range(2)], + max_concurrency=1, + completion_config=CompletionConfig(min_successful=2), + sub_type_top=OperationSubType.PARALLEL, + sub_type_iteration=OperationSubType.PARALLEL_BRANCH, + name_prefix="branch-", + serdes=None, + operation_id_namespace=_StubNamespace(), + ) + state = Mock(spec=ExecutionState) + state._plugin_executor = PluginExecutor([DurableInstrumentationPlugin()]) + token = binding.set("coordinator") + try: + result = executor.execute(state, Mock(spec=DurableContext)) + assert result.get_results() == ["ok", "ok"] + assert binding.get() == "coordinator" + assert [value for value, _thread in seen] == ["coordinator", "coordinator"] + assert seen[0][1] == seen[1][1] != threading.get_ident() + finally: + for call in state.register_branch_pool.call_args_list: + call.args[0].shutdown(wait=True) + binding.reset(token) + assert binding.get() == "empty" From eded23e8fc79c4118a837b1fa7c9756bbd8e8360 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 9 Oct 2026 01:38:48 +0000 Subject: [PATCH 31/54] fix(otel): retain completions until their parent span exists --- .github/scripts/install_otel_test_wheels.py | 5 +- .../opentelemetry-conformance-tests.yml | 4 +- CONTRIBUTING.md | 4 +- .../README.md | 7 + .../invocation_plugin.py | 30 ++- .../e2e/test_parent_creation_order_int.py | 193 ++++++++++++++++++ .../tests/test_invocation_plugin.py | 132 ++++++++++++ 7 files changed, 369 insertions(+), 6 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_parent_creation_order_int.py diff --git a/.github/scripts/install_otel_test_wheels.py b/.github/scripts/install_otel_test_wheels.py index 938996aca..b1332e5c2 100644 --- a/.github/scripts/install_otel_test_wheels.py +++ b/.github/scripts/install_otel_test_wheels.py @@ -16,6 +16,7 @@ ROOT = Path(__file__).resolve().parents[2] CORE = "aws-durable-execution-sdk-python" OTEL = CORE + "-otel" +TESTING = CORE + "-testing" def built_wheel(package: str) -> Path: @@ -72,7 +73,9 @@ def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--legacy-plugin", action="store_true") args = parser.parse_args() - packages = [CORE] if args.legacy_plugin else [CORE, OTEL] + # This PR also repairs the local service simulator. Validate its actual + # built artifact with the SDK/plugin, retaining every test assertion. + packages = [CORE] if args.legacy_plugin else [CORE, OTEL, TESTING] wheels = [built_wheel(package) for package in packages] subprocess.run( [ diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index ecf2e0f3f..cd35db05c 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -45,7 +45,7 @@ on: conformance_test_ref: description: Conformance test commit SHA or branch name required: true - default: bdb4f1cd0f9252c1aaa978bb8b341b71f2b9d9dc + default: cf325a89d5f1ba3c2d57dbdb71ee6deda778935a type: string # Backend stacks are shared across PRs. Queue whole runs so reusable @@ -73,7 +73,7 @@ jobs: resource_prefix: p sdk_repository: aws/aws-durable-execution-sdk-python sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }} - conformance_test_ref: ${{ inputs.conformance_test_ref || 'bdb4f1cd0f9252c1aaa978bb8b341b71f2b9d9dc' }} + conformance_test_ref: ${{ inputs.conformance_test_ref || 'cf325a89d5f1ba3c2d57dbdb71ee6deda778935a' }} # Check the SDK out so the handlers and templates below are on disk. The handlers # themselves are installed from sdk_ref by src/requirements.txt during the SAM build. checkout_sdk: true diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index a3b6bc2f7..b583d7461 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -72,11 +72,11 @@ hatch run dev-examples:test # run examples tests only ### Installed package compatibility testing -Build the core and OTel distributions with `hatch build` in each package, then run +Build the core, OTel and testing-library distributions with `hatch build` in each package, then run these commands from the repository root: ```bash -hatch run test-wheel-otel:test # full OTel suite on the two built wheels +hatch run test-wheel-otel:test # full OTel suite on the three built wheels hatch run test-wheel-otel-legacy:test # released OTel 1.0.0 with the built core hatch run test-pypi-examples:test # examples against the published core ``` diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 2a095bf9c..f1c4a0ff8 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -364,6 +364,13 @@ OTel `OK` only for `SUCCEEDED`, `ERROR` when error details are delivered, and `CANCELLED`, `TIMED_OUT`, and `STOPPED`. The original durable operation status remains in `durable.operation.status`. +When a concurrent branch's terminal completion arrives before replay has created +its parent span, Invocation view retains that completion until the actual parent +span is registered. The terminal segment is exported under that parent before +control returns to the branch's user code. No ancestor is invented and SDK +completion delivery is unchanged. If the parent never becomes active, normal +invocation cleanup and flushing still run before the missing parent is reported. + ### Invocation context isolation With core 2.1+, invocation hooks and the handler run on one worker in an 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 0fb974fe1..9712a81a7 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 @@ -156,6 +156,9 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._span_time_floor_ns: int | None = None # Maps operation ID (None for root) to the active span. self._operation_spans: dict[str | None, Span] = {} + # A sibling checkpoint can report completion before replay enters its + # parent context. Retain the real event until that parent span exists. + self._pending_operation_ends: dict[str, list[OperationEndInfo]] = {} # Replay state supplied by CONTEXT operation START hooks. Missing # entries identify checkpointless contexts such as FLAT branches. self._context_operation_replays: dict[str, bool] = {} @@ -491,10 +494,19 @@ def _start_span( links=links, ) self._operation_spans[registry_key] = span + pending_ends = ( + self._pending_operation_ends.pop(registry_key, []) + if registry_key is not None + else [] + ) if operation_id is None: self._span_time_floor_ns = span_start_time logger.debug("Started OTel span: %s", span) + # Registration and dequeue share the lock, so an arriving completion + # either sees its parent or is drained here. Do not hold it in callbacks. + for pending_end in pending_ends: + self.on_operation_end(pending_end) return span def _end_span( @@ -690,6 +702,7 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: with self._operation_spans_lock: operation_ids = list(reversed(self._operation_spans)) incomplete_attempt_span_keys = set(self._incomplete_attempt_span_keys) + unresolved_parents = tuple(self._pending_operation_ends) for operation_id in operation_ids: if operation_id: span = self._get_span(operation_id) @@ -738,6 +751,12 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: # Flush before Lambda freeze if hasattr(self._provider, "force_flush"): self._provider.force_flush() + if unresolved_parents: + # Keep the previous missing-parent failure observable, but finish + # normal token/span cleanup and flushing before reporting it. + raise ValueError( + f"No parent span found for deferred operation ends: {unresolved_parents}" + ) def _reset_state(self) -> None: """Clear per-invocation state for warm Lambda environment reuse.""" @@ -752,6 +771,7 @@ def _reset_state(self) -> None: self._span_time_floor_ns = None with self._operation_spans_lock: self._operation_spans = {} + self._pending_operation_ends = {} self._context_operation_replays = {} self._incomplete_attempt_span_keys = set() self._tracing_enabled = False @@ -795,7 +815,15 @@ def on_operation_end(self, info: OperationEndInfo) -> None: # The operation started in a prior invocation. Create a new # correlated segment and link it to the deterministic logical # operation context shared across invocations. - parent_span = self._resolve_parent_span(info.parent_id) + with self._operation_spans_lock: + parent_span = self._operation_spans.get(info.parent_id) + if parent_span is None and info.parent_id is not None: + self._pending_operation_ends.setdefault(info.parent_id, []).append( + info + ) + return + if parent_span is None: + parent_span = self._resolve_parent_span(info.parent_id) attributes = self._extract_attributes(info) span = self._start_span( operation_id=info.operation_id, diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_parent_creation_order_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_parent_creation_order_int.py new file mode 100644 index 000000000..b998de4d3 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_parent_creation_order_int.py @@ -0,0 +1,193 @@ +"""A real completion can precede its replayed sibling's parent span.""" + +from __future__ import annotations + +import json +import threading +import time +from typing import Any + +import pytest +from aws_durable_execution_sdk_python import DurableContext, durable_execution +from aws_durable_execution_sdk_python.config import ( + CallbackConfig, + Duration, + ParallelConfig, +) +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationStartInfo, + OperationEndInfo, + OperationType, + UserFunctionStartInfo, +) +from aws_durable_execution_sdk_python.serdes import JsonSerDes +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry import trace +from opentelemetry.sdk.trace import ReadableSpan, TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import InvocationOtelPlugin, OtelPluginConfig + + +def _span_start(span: ReadableSpan) -> int: + assert span.start_time is not None + return span.start_time + + +def test_early_sibling_callback_exports_under_real_parent_before_user_observation( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = InvocationOtelPlugin( + OtelPluginConfig( + tracer_provider=provider, + enrich_logger=False, + context_extractor=lambda _: None, + ) + ) + second_start = threading.Event() + both_completed = threading.Event() + end_observed = threading.Event() + parent_entered = threading.Event() + generations = [0] + observed: list[tuple[int, int]] = [] + hook_order: list[str] = [] + + class Gate(DurableInstrumentationPlugin): + def on_invocation_start(self, _info: InvocationStartInfo) -> None: + generations[0] += 1 + if generations[0] == 2: + second_start.set() + assert both_completed.wait(10) + + def on_user_function_start(self, info: UserFunctionStartInfo) -> None: + if generations[0] == 2 and info.name == "parallel-branch-1": + hook_order.append("parent-gated") + assert end_observed.wait(10) + parent_entered.set() + + class ObserveEnd(DurableInstrumentationPlugin): + def on_operation_end(self, info: OperationEndInfo) -> None: + if ( + generations[0] == 2 + and info.name == "target-1" + and info.operation_type is OperationType.CALLBACK + ): + # This observer follows OTel in the actual SDK dispatch order. + assert not parent_entered.is_set() + hook_order.append("early-completion-observed") + end_observed.set() + + def branch(child: DurableContext, index: int) -> Any: + value = child.create_callback( + name=f"target-{index}", config=CallbackConfig(serdes=JsonSerDes()) + ).result() + if index == 1: + finished = [ + span + for span in exporter.get_finished_spans() + if span.name == "target-1" + ] + observed.append( + (len(finished), trace.get_current_span().get_span_context().span_id) + ) + with provider.get_tracer("customer").start_as_current_span( + "after-target-1" + ): + pass + return value + + def handler(_event: Any, durable: DurableContext) -> list[Any]: + result = durable.parallel( + [lambda child: branch(child, 0), lambda child: branch(child, 1)], + name="parallel", + config=ParallelConfig(max_concurrency=2), + ) + durable.wait(Duration.from_seconds(1), name="later-one") + durable.wait(Duration.from_seconds(1), name="later-two") + return result.get_results() + + wrapped = durable_execution(handler, plugins=[Gate(), plugin, ObserveEnd()]) + try: + with DurableFunctionTestRunner( + handler=wrapped, poll_interval=0.01, skip_time=False, execution_timeout=25 + ) as runner: + arn = runner.run_async(input="{}") + deadline = time.monotonic() + 10 + callbacks = {} + while time.monotonic() < deadline: + history = runner.get_execution_history( + arn, include_execution_data=True + ).events + callbacks = { + event.name: event.callback_started_details.callback_id + for event in history + if event.event_type == "CallbackStarted" + } + if len(callbacks) == 2 and any( + event.event_type == "InvocationCompleted" for event in history + ): + break + time.sleep(0.01) + assert set(callbacks) == {"target-0", "target-1"} + runner.send_callback_success( + callbacks["target-0"], result=json.dumps("left").encode() + ) + assert second_start.wait(10) + # The second callback completes after the invocation input snapshot, + # so another branch's checkpoint returns this genuine new completion. + runner.send_callback_success( + callbacks["target-1"], result=json.dumps("right").encode() + ) + both_completed.set() + result = runner.wait_for_result(arn, timeout=15) + history = runner.get_execution_history( + arn, include_execution_data=True + ).events + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == ["left", "right"] + assert end_observed.is_set() and parent_entered.is_set() + assert len(observed) == 1 and observed[0][0] == 2 + spans = exporter.get_finished_spans() + # Keep every span: Invocation view exports a nonterminal segment in + # invocation one and a distinct terminal continuation in invocation two. + # Neither segment may stand in for the other or reappear on later replays. + for index in (0, 1): + segments = sorted( + [span for span in spans if span.name == f"target-{index}"], + key=_span_start, + ) + assert len(segments) == 2 + initial, terminal = segments + assert initial.attributes is not None and terminal.attributes is not None + assert initial.attributes["durable.operation.status"] == "STARTED" + assert terminal.attributes["durable.operation.status"] == "SUCCEEDED" + assert initial.context.span_id != terminal.context.span_id + parents = sorted( + [span for span in spans if span.name == f"parallel-branch-{index}"], + key=_span_start, + ) + assert len(parents) == 2 + assert initial.parent is not None + assert initial.parent.span_id == parents[0].context.span_id + assert terminal.parent is not None + assert terminal.parent.span_id == parents[1].context.span_id + assert initial.end_time is not None and terminal.start_time is not None + assert initial.end_time <= terminal.start_time + if index == 1: + assert terminal.parent.span_id == observed[0][1] + (marker,) = [span for span in spans if span.name == "after-target-1"] + assert terminal.end_time is not None and marker.start_time is not None + assert terminal.end_time <= marker.start_time + assert sum(event.event_type == "InvocationCompleted" for event in history) == 4 + assert not [record for record in caplog.records if record.exc_info] + assert plugin._context_tokens == {} + finally: + both_completed.set() + end_observed.set() + provider.shutdown() 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 ba373ad8d..b1f863c68 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 @@ -2282,3 +2282,135 @@ def body() -> None: finally: otel_context.detach(token) provider.shutdown() + + +@pytest.mark.parametrize( + ("status", "error", "expected"), + [ + (OperationStatus.SUCCEEDED, None, StatusCode.OK), + (OperationStatus.FAILED, None, StatusCode.UNSET), + (OperationStatus.CANCELLED, None, StatusCode.UNSET), + (OperationStatus.TIMED_OUT, None, StatusCode.UNSET), + (OperationStatus.STOPPED, None, StatusCode.UNSET), + ( + OperationStatus.FAILED, + ErrorObject(message="failure", type="Example", data=None, stack_trace=None), + StatusCode.ERROR, + ), + ], +) +def test_deferred_end_retains_real_parent_and_status_mapping( + status: OperationStatus, error: ErrorObject | None, expected: StatusCode +) -> None: + plugin, exporter = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + end = OperationEndInfo( + operation_id="child", + operation_type=OperationType.WAIT, + sub_type=None, + name="deferred-child", + parent_id="parent", + start_time=START_TIME, + is_replayed=False, + status=status, + end_time=END_TIME, + error=error, + ) + plugin.on_operation_end(end) + assert not exporter.get_finished_spans() + assert plugin._pending_operation_ends == {"parent": [end]} + plugin.on_user_function_start( + _user_function_start_info("parent", operation_type=OperationType.CONTEXT) + ) + parent = plugin._get_span("parent") + assert parent is not None + (child,) = exporter.get_finished_spans() + assert child.parent is not None and child.attributes is not None + assert child.parent.span_id == parent.get_span_context().span_id + assert child.status.status_code is expected + assert child.attributes["durable.operation.status"] == status.value + assert len(child.events) == (1 if error is not None else 0) + assert plugin._pending_operation_ends == {} + plugin.on_user_function_end( + _user_function_end_info("parent", operation_type=OperationType.CONTEXT) + ) + plugin.on_invocation_end(_invocation_end_info()) + + +def test_deferred_descendants_drain_from_real_parent_completion_in_child_first_order() -> ( + None +): + plugin, exporter = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + for operation_id, parent_id, kind in [ + ("child", "inner", OperationType.WAIT), + ("inner", "outer", OperationType.CONTEXT), + ]: + plugin.on_operation_end( + OperationEndInfo( + operation_id=operation_id, + operation_type=kind, + sub_type=None, + name=operation_id, + parent_id=parent_id, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=END_TIME, + ) + ) + assert not exporter.get_finished_spans() + plugin.on_user_function_start( + _user_function_start_info("outer", operation_type=OperationType.CONTEXT) + ) + child, inner = exporter.get_finished_spans() + assert [child.name, inner.name] == ["child", "inner"] + assert child.parent is not None and inner.parent is not None + assert child.parent.span_id == inner.context.span_id + outer = plugin._get_span("outer") + assert outer is not None + assert inner.parent.span_id == outer.get_span_context().span_id + assert child.end_time is not None and inner.end_time is not None + assert child.end_time <= inner.end_time + assert plugin._pending_operation_ends == {} + plugin.on_user_function_end( + _user_function_end_info("outer", operation_type=OperationType.CONTEXT) + ) + plugin.on_invocation_end(_invocation_end_info()) + + +def test_unresolved_parent_is_reported_after_cleanup_and_not_carried_into_reuse() -> ( + None +): + plugin, exporter = _create_plugin() + before = otel_context.get_current() + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_operation_end( + OperationEndInfo( + operation_id="missing-child", + operation_type=OperationType.WAIT, + sub_type=None, + name="missing-child", + parent_id="missing-parent", + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=END_TIME, + ) + ) + with pytest.raises(ValueError, match="No parent span found for deferred"): + plugin.on_invocation_end(_invocation_end_info()) + assert otel_context.get_current() == before + assert plugin._operation_spans == {} + assert plugin._pending_operation_ends == {} + assert plugin._context_tokens == {} + assert {span.name for span in exporter.get_finished_spans()} == { + "Invocation", + "Workflow", + } + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_invocation_end(_invocation_end_info(InvocationStatus.PENDING)) + assert not [ + span for span in exporter.get_finished_spans() if span.name == "missing-child" + ] + assert plugin._pending_operation_ends == {} From eb6c78fb8ad2d13797c2550459bb734e1d57cf6e Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 9 Oct 2026 01:46:23 +0000 Subject: [PATCH 32/54] test: align conformance runner reference assertion --- .../scripts/tests/test_opentelemetry_conformance_workflow.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py index e32684716..6ca1ff643 100644 --- a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py +++ b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py @@ -5,7 +5,7 @@ SHARED_WORKFLOW_REF = "bdb4f1cd0f9252c1aaa978bb8b341b71f2b9d9dc" -CONFORMANCE_TEST_REF = "bdb4f1cd0f9252c1aaa978bb8b341b71f2b9d9dc" +CONFORMANCE_TEST_REF = "cf325a89d5f1ba3c2d57dbdb71ee6deda778935a" WORKFLOW_PATH = ( Path(__file__).parents[2] / "workflows" / "opentelemetry-conformance-tests.yml" From d59af28c25fc73d7423cab456a0a8fa90641509e Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Fri, 9 Oct 2026 01:51:55 +0000 Subject: [PATCH 33/54] ci: use validated 26-case conformance runner --- .../scripts/tests/test_opentelemetry_conformance_workflow.py | 2 +- .github/workflows/opentelemetry-conformance-tests.yml | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py index 6ca1ff643..24b55be8c 100644 --- a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py +++ b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py @@ -5,7 +5,7 @@ SHARED_WORKFLOW_REF = "bdb4f1cd0f9252c1aaa978bb8b341b71f2b9d9dc" -CONFORMANCE_TEST_REF = "cf325a89d5f1ba3c2d57dbdb71ee6deda778935a" +CONFORMANCE_TEST_REF = "75987d46a915bc37409eed3ea9c3617a924c9756" WORKFLOW_PATH = ( Path(__file__).parents[2] / "workflows" / "opentelemetry-conformance-tests.yml" diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index cd35db05c..592b92603 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -45,7 +45,7 @@ on: conformance_test_ref: description: Conformance test commit SHA or branch name required: true - default: cf325a89d5f1ba3c2d57dbdb71ee6deda778935a + default: 75987d46a915bc37409eed3ea9c3617a924c9756 type: string # Backend stacks are shared across PRs. Queue whole runs so reusable @@ -73,7 +73,7 @@ jobs: resource_prefix: p sdk_repository: aws/aws-durable-execution-sdk-python sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }} - conformance_test_ref: ${{ inputs.conformance_test_ref || 'cf325a89d5f1ba3c2d57dbdb71ee6deda778935a' }} + conformance_test_ref: ${{ inputs.conformance_test_ref || '75987d46a915bc37409eed3ea9c3617a924c9756' }} # Check the SDK out so the handlers and templates below are on disk. The handlers # themselves are installed from sdk_ref by src/requirements.txt during the SAM build. checkout_sdk: true From c7d08fafd6b4860ec7a68ed659e0f19246fe74f5 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Tue, 13 Oct 2026 16:58:47 +0000 Subject: [PATCH 34/54] test(otel): narrow callback history fields before use --- .../tests/e2e/test_parent_creation_order_int.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_parent_creation_order_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_parent_creation_order_int.py index b998de4d3..16490c3f8 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_parent_creation_order_int.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_parent_creation_order_int.py @@ -119,16 +119,20 @@ def handler(_event: Any, durable: DurableContext) -> list[Any]: ) as runner: arn = runner.run_async(input="{}") deadline = time.monotonic() + 10 - callbacks = {} + callbacks: dict[str, str] = {} while time.monotonic() < deadline: history = runner.get_execution_history( arn, include_execution_data=True ).events - callbacks = { - event.name: event.callback_started_details.callback_id - for event in history - if event.event_type == "CallbackStarted" - } + callbacks = {} + for event in history: + if event.event_type == "CallbackStarted": + assert event.name is not None + assert event.callback_started_details is not None + assert event.callback_started_details.callback_id is not None + callbacks[event.name] = ( + event.callback_started_details.callback_id + ) if len(callbacks) == 2 and any( event.event_type == "InvocationCompleted" for event in history ): @@ -150,6 +154,7 @@ def handler(_event: Any, durable: DurableContext) -> list[Any]: arn, include_execution_data=True ).events assert result.status.value == "SUCCEEDED" + assert result.result is not None assert json.loads(result.result) == ["left", "right"] assert end_observed.is_set() and parent_entered.is_set() assert len(observed) == 1 and observed[0][0] == 2 From 70c4b8cc061b986250be092db27d8b17447094b9 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Tue, 13 Oct 2026 18:00:21 +0000 Subject: [PATCH 35/54] fix(testing): normalize empty errors in operation history --- .../model.py | 19 +++-- .../tests/model_test.py | 84 +++++++++++++++++++ 2 files changed, 97 insertions(+), 6 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/model.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/model.py index fdccb2cff..1de6b425c 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/model.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/model.py @@ -2601,6 +2601,11 @@ class HistoryEventTypeConfig: } +def _operation_error(error: ErrorObject | None) -> ErrorObject | None: + # History retains present Payload={}, while SDK state treats Error={} as absent. + return error if error is not None and error.to_dict() else None + + def events_to_operations(events: list[Event]) -> list[Operation]: """Convert a list of history events into operations. @@ -2753,9 +2758,7 @@ def events_to_operations(events: list[Event]) -> list[Operation]: callback_details=CallbackDetails( callback_id=callback_id, result=result, - # History preserves a present empty Payload object, while - # CallbackDetails uses None for an empty wire error. - error=error if error is not None and error.to_dict() else None, + error=_operation_error(error), ), ) @@ -2796,7 +2799,7 @@ def events_to_operations(events: list[Event]) -> list[Operation]: operation, step_details=StepDetails( result=result_val, - error=error_val, + error=_operation_error(error_val), attempt=attempt, next_attempt_timestamp=next_attempt_ts, ), @@ -2829,7 +2832,9 @@ def events_to_operations(events: list[Event]) -> list[Operation]: operation, context_details=ContextDetails( result=None, - error=event.context_failed_details.error.payload, + error=_operation_error( + event.context_failed_details.error.payload + ), ), ) @@ -2854,7 +2859,9 @@ def events_to_operations(events: list[Event]) -> list[Operation]: operation, chained_invoke_details=ChainedInvokeDetails( result=None, - error=event.chained_invoke_failed_details.error.payload, + error=_operation_error( + event.chained_invoke_failed_details.error.payload + ), ), ) diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/model_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/model_test.py index f3f8a9131..9885003a5 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/model_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/model_test.py @@ -7,8 +7,11 @@ import pytest from aws_durable_execution_sdk_python.lambda_service import ( + ChainedInvokeDetails as ServiceChainedInvokeDetails, + ContextDetails as ServiceContextDetails, OperationStatus, OperationType, + StepDetails as ServiceStepDetails, ) from aws_durable_execution_sdk_python_testing.exceptions import ( @@ -1491,6 +1494,87 @@ def test_event_error_empty_payload_retains_truncation_flag(): assert EventError.from_dict(wire).to_dict() == wire +@pytest.mark.parametrize( + "event_type,details_attribute,service_details", + [ + ("StepFailed", "step_details", ServiceStepDetails), + ("ContextFailed", "context_details", ServiceContextDetails), + ("ChainedInvokeFailed", "chained_invoke_details", ServiceChainedInvokeDetails), + ], +) +@pytest.mark.parametrize( + "wire_error", + [ + {"Truncated": False}, + {"Payload": None, "Truncated": False}, + {"Payload": {}, "Truncated": False}, + {"Payload": {}, "Truncated": True}, + {"Truncated": True}, + {"Payload": {"ErrorMessage": ""}, "Truncated": False}, + {"Payload": {"ErrorType": ""}, "Truncated": False}, + {"Payload": {"ErrorData": ""}, "Truncated": False}, + {"Payload": {"StackTrace": []}, "Truncated": False}, + { + "Payload": { + "ErrorMessage": "failure", + "ErrorType": "RuntimeError", + "ErrorData": "details", + "StackTrace": [], + }, + "Truncated": False, + }, + ], +) +def test_failed_history_preserves_payload_and_matches_service_error_semantics( + event_type, details_attribute, service_details, wire_error +): + from aws_durable_execution_sdk_python_testing.runner import ( + DurableFunctionTestResult, + ) + + history = GetDurableExecutionHistoryResponse.from_dict( + { + "Events": [ + { + "EventType": event_type, + "EventId": 1, + "Id": "operation", + "Name": "failed-operation", + "EventTimestamp": TIMESTAMP_2023_01_01_00_01, + event_type + "Details": {"Error": wire_error}, + } + ] + } + ) + history_wire = history.to_dict() + expected_history_error = dict(wire_error) + if expected_history_error.get("Payload") is None: + expected_history_error.pop("Payload", None) + assert history_wire["Events"][0][event_type + "Details"]["Error"] == ( + expected_history_error + ) + assert GetDurableExecutionHistoryResponse.from_dict(history_wire).to_dict() == ( + history_wire + ) + + expected = service_details.from_dict({"Error": wire_error.get("Payload")}).error + (operation,) = events_to_operations(history.events) + assert operation.status is OperationStatus.FAILED + + execution = GetDurableExecutionResponse( + durable_execution_arn="arn:aws:lambda:us-west-2:123456789012:function:test:execution:one", + durable_execution_name="one", + function_arn="arn:aws:lambda:us-west-2:123456789012:function:test", + status="FAILED", + start_timestamp=TIMESTAMP_2023_01_01_00_00, + ) + result = DurableFunctionTestResult.from_execution_history(execution, history) + assert result.get_operation_by_name("failed-operation").error == expected + assert getattr(operation, details_attribute).error == expected + # Reconstructing SDK-facing state must not mutate detailed history payloads. + assert history.to_dict() == history_wire + + # Tests for RetryDetails def test_retry_details_serialization(): """Test RetryDetails from_dict/to_dict round-trip.""" From ce1e99b0be90e5d2424762be438419540252524d Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 14 Oct 2026 01:10:47 +0000 Subject: [PATCH 36/54] ci: use validated shared conformance runner --- .../scripts/tests/test_opentelemetry_conformance_workflow.py | 2 +- .github/workflows/opentelemetry-conformance-tests.yml | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py index 24b55be8c..8c44cdfe5 100644 --- a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py +++ b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py @@ -5,7 +5,7 @@ SHARED_WORKFLOW_REF = "bdb4f1cd0f9252c1aaa978bb8b341b71f2b9d9dc" -CONFORMANCE_TEST_REF = "75987d46a915bc37409eed3ea9c3617a924c9756" +CONFORMANCE_TEST_REF = "a1ba6a96298a252eb85bef43866ef1d3ccab66be" WORKFLOW_PATH = ( Path(__file__).parents[2] / "workflows" / "opentelemetry-conformance-tests.yml" diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index 592b92603..981eb71a1 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -45,7 +45,7 @@ on: conformance_test_ref: description: Conformance test commit SHA or branch name required: true - default: 75987d46a915bc37409eed3ea9c3617a924c9756 + default: a1ba6a96298a252eb85bef43866ef1d3ccab66be type: string # Backend stacks are shared across PRs. Queue whole runs so reusable @@ -73,7 +73,7 @@ jobs: resource_prefix: p sdk_repository: aws/aws-durable-execution-sdk-python sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }} - conformance_test_ref: ${{ inputs.conformance_test_ref || '75987d46a915bc37409eed3ea9c3617a924c9756' }} + conformance_test_ref: ${{ inputs.conformance_test_ref || 'a1ba6a96298a252eb85bef43866ef1d3ccab66be' }} # Check the SDK out so the handlers and templates below are on disk. The handlers # themselves are installed from sdk_ref by src/requirements.txt during the SAM build. checkout_sdk: true From cff904febac3639747f913ed7e7c9dd8b009baa7 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 14 Oct 2026 03:12:30 +0000 Subject: [PATCH 37/54] fix(testing): acknowledge delivered input snapshots --- ...test_opentelemetry_conformance_workflow.py | 2 +- .../opentelemetry-conformance-tests.yml | 4 +- .../e2e/test_callback_input_delivery_int.py | 198 ++++++++++++++++++ .../docs/architecture.md | 16 +- .../execution.py | 19 +- .../executor.py | 10 +- .../e2e/invocation_input_delivery_test.py | 110 ++++++++++ .../tests/executor_test.py | 10 + .../tests/invocation_input_delivery_test.py | 98 +++++++++ 9 files changed, 447 insertions(+), 20 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_callback_input_delivery_int.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/tests/e2e/invocation_input_delivery_test.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/tests/invocation_input_delivery_test.py diff --git a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py index 8c44cdfe5..1479be395 100644 --- a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py +++ b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py @@ -5,7 +5,7 @@ SHARED_WORKFLOW_REF = "bdb4f1cd0f9252c1aaa978bb8b341b71f2b9d9dc" -CONFORMANCE_TEST_REF = "a1ba6a96298a252eb85bef43866ef1d3ccab66be" +CONFORMANCE_TEST_REF = "fb69b21ccc8805a74271554403707f5e649e5aae" WORKFLOW_PATH = ( Path(__file__).parents[2] / "workflows" / "opentelemetry-conformance-tests.yml" diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index 981eb71a1..ec8d40699 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -45,7 +45,7 @@ on: conformance_test_ref: description: Conformance test commit SHA or branch name required: true - default: a1ba6a96298a252eb85bef43866ef1d3ccab66be + default: fb69b21ccc8805a74271554403707f5e649e5aae type: string # Backend stacks are shared across PRs. Queue whole runs so reusable @@ -73,7 +73,7 @@ jobs: resource_prefix: p sdk_repository: aws/aws-durable-execution-sdk-python sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }} - conformance_test_ref: ${{ inputs.conformance_test_ref || 'a1ba6a96298a252eb85bef43866ef1d3ccab66be' }} + conformance_test_ref: ${{ inputs.conformance_test_ref || 'fb69b21ccc8805a74271554403707f5e649e5aae' }} # Check the SDK out so the handlers and templates below are on disk. The handlers # themselves are installed from sdk_ref by src/requirements.txt during the SAM build. checkout_sdk: true diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_callback_input_delivery_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_callback_input_delivery_int.py new file mode 100644 index 000000000..917716607 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_callback_input_delivery_int.py @@ -0,0 +1,198 @@ +"""Input-only callback resumes export each terminal outcome once.""" + +import json +import logging +import threading +import time + +import pytest +from aws_durable_execution_sdk_python import ( + DurableContext, + durable_execution, + durable_step, + durable_with_child_context, +) +from aws_durable_execution_sdk_python.config import CallbackConfig +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from aws_durable_execution_sdk_python.serdes import JsonSerDes +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +@pytest.mark.parametrize( + "view,nested,page_bytes,fresh_executor,concurrent_completion", + [ + (ExecutionOtelPlugin, False, None, False, False), + (InvocationOtelPlugin, False, None, False, False), + (ExecutionOtelPlugin, True, None, False, False), + (InvocationOtelPlugin, True, None, False, False), + (InvocationOtelPlugin, True, 1, False, False), + (ExecutionOtelPlugin, False, None, True, False), + (InvocationOtelPlugin, True, None, False, True), + ], + ids=[ + "execution", + "invocation", + "execution-nested", + "invocation-nested", + "paged", + "fresh-executor", + "concurrent", + ], +) +def test_two_callbacks_do_not_repeat_terminal_spans_after_input_only_resume( + view, nested, page_bytes, fresh_executor, concurrent_completion, caplog, monkeypatch +): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + starts = [] + ends = [] + marker_calls = [] + second_end = threading.Event() + release_end = threading.Event() + + class Observe(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + starts.append(info) + + def on_operation_end(self, info): + if info.name in ("a", "b"): + ends.append((len(starts), info.name, info.status.value)) + + def on_invocation_end(self, info): + if concurrent_completion and len(starts) == 2: + second_end.set() + assert release_end.wait(6) + + observer = Observe() + + def make_plugin(): + return view( + OtelPluginConfig( + tracer_provider=provider, + enrich_logger=False, + context_extractor=lambda _: None, + ) + ) + + @durable_step + def marker(context, value): + marker_calls.append(1) + return value + + def callbacks(context): + a = context.create_callback( + name="a", config=CallbackConfig(serdes=JsonSerDes()) + ) + b = context.create_callback( + name="b", config=CallbackConfig(serdes=JsonSerDes()) + ) + return [a.result(), b.result()] + + @durable_with_child_context + def child(context): + return callbacks(context) + + def body(event, context: DurableContext): + result = ( + context.run_in_child_context(child(), name="group") + if nested + else callbacks(context) + ) + return context.step(marker(result), name="after-b") + + wrapped = durable_execution(body, plugins=[make_plugin(), observer]) + + def entrypoint(event, context): + target = ( + durable_execution(body, plugins=[make_plugin(), observer]) + if fresh_executor + else wrapped + ) + return target(event, context) + + try: + with DurableFunctionTestRunner( + handler=entrypoint, + skip_time=False, + poll_interval=0.01, + execution_timeout=25, + max_invocation_page_bytes=page_bytes, + ) as runner: + arn = runner.run_async(input="{}") + a_id = runner.wait_for_callback(arn, name="a", timeout=5) + b_id = runner.wait_for_callback(arn, name="b", timeout=5) + runner.send_callback_success(a_id, result=b'"A"') + if concurrent_completion: + assert second_end.wait(5) + runner.send_callback_success(b_id, result=b'"B"') + release_end.set() + else: + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + events = runner.get_execution_history( + arn, include_execution_data=True + ).events + if ( + sum( + event.event_type == "InvocationCompleted" + for event in events + ) + >= 2 + ): + break + time.sleep(0.01) + else: + raise AssertionError("The second invocation did not suspend") + runner.send_callback_success(b_id, result=b'"B"') + result = runner.wait_for_result(arn, timeout=10) + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == ["A", "B"] + assert len(starts) == 3 + assert marker_calls == [1] + assert ends == [(2, "a", "SUCCEEDED"), (3, "b", "SUCCEEDED")] + spans = exporter.get_finished_spans() + callback_spans = [ + span + for span in spans + if (span.attributes or {}).get("durable.operation.type") == "CALLBACK" + ] + # Retain the nonterminal segments as well as the terminal segments; + # filtering down to successful spans would hide the invocation story. + expected = [("a", "SUCCEEDED"), ("b", "SUCCEEDED")] + if view is InvocationOtelPlugin: + expected = [ + ("b", "STARTED"), + ("a", "STARTED"), + ("a", "SUCCEEDED"), + ("b", "STARTED"), + ("b", "SUCCEEDED"), + ] + assert [ + (span.name, span.attributes["durable.operation.status"]) + for span in callback_spans + ] == expected + exported_ids = {span.context.span_id for span in spans} + assert all( + span.parent is not None and span.parent.span_id in exported_ids + for span in callback_spans + ) + assert not [ + record + for record in caplog.records + if record.levelno >= logging.ERROR + and ("plugin" in record.name or "otel" in record.name) + ] + finally: + release_end.set() + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-testing/docs/architecture.md b/packages/aws-durable-execution-sdk-python-testing/docs/architecture.md index c46f076a1..e82f81337 100644 --- a/packages/aws-durable-execution-sdk-python-testing/docs/architecture.md +++ b/packages/aws-durable-execution-sdk-python-testing/docs/architecture.md @@ -234,10 +234,18 @@ list in a checkpoint response. The delta for the next response is the operations whose `operation_last_touched_seq` is greater than `handler_seen_seq`. -The watermark advances only when a checkpoint response is built. It never -advances when invocation input is built. That is what makes retries -correct. A failed invocation never moves the watermark. So the retry -delivers the same state again. +The watermark advances when a checkpoint response is built or an invocation +response passes validation, including a `PENDING` response. Building invocation +input does not acknowledge it: input updates remain pending until delivery is +confirmed. A failed call or rejected response does not consume that input, so a +retry can receive it again. + +The invocation acknowledgment uses the sequence captured when its input was +built, on the execution's worker. It never moves a higher checkpoint watermark +backward or consumes operations that changed after the captured snapshot. This +also prevents an input-only callback resume from receiving the same terminal +update again at a later checkpoint. The acknowledgment concerns runner state +delivery, not telemetry export. ``` seq_counter handler_seen_seq delta next checkpoint diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py index 2ab02fdce..569b7f1f2 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py @@ -434,16 +434,11 @@ def record_invocation_completion( ) def mark_state_delivered(self) -> None: - """Reset the list of changed operations once an invocation's input - carries them. - - The service reports, on each invocation, the operations that - changed after the last state the handler observed. The runner - delivers state in the invocation input. So the list is reset when - that input is built, not when the invocation completes: an - operation that completes while the handler is still running is - reported on the next invocation unless a checkpoint response has - already delivered it through ``advance_handler_seen``. + """Explicitly clear all pending operation IDs after their delivery. + + The live executor uses versioned ``advance_handler_seen`` acknowledgments + so failed invocations retain their input updates and concurrent changes + newer than the acknowledged input snapshot remain pending. """ self.updated_operation_ids = [] @@ -887,8 +882,8 @@ def advance_handler_seen(self, seq: int) -> None: smaller or equal values are ignored.""" if seq > self.execution.handler_seen_seq: self.execution.handler_seen_seq = seq - # A checkpoint has delivered these updates to the running handler. - # Retain only changes newer than that response for the next input. + # A checkpoint response or completed invocation delivered this state. + # Retain changes newer than that acknowledged snapshot. self.execution.updated_operation_ids = [ operation_id for operation_id in self.execution.updated_operation_ids diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/executor.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/executor.py index 44a01bc01..b4033d976 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/executor.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/executor.py @@ -1526,7 +1526,8 @@ def _begin_invocation( self._set_invocation_gate(execution_arn, InvocationState.INVOKING) execution.begin_new_invocation() invocation_input = self._invoker.create_invocation_input(execution=execution) - execution.mark_state_delivered() + # Keep input updates until a valid response acknowledges this snapshot. + # A failed invocation must not consume their delivery before a retry. self._store.save(execution) return execution, invocation_input, execution.seq_counter @@ -1584,6 +1585,13 @@ def _finish_invocation( # A clean invocation resets the retry counter. reset_target = self._store.load(execution_arn) reset_target.consecutive_failed_invocation_attempts = 0 + if invocation_seq is not None: + # A valid response acknowledges the input snapshot, even when this + # invocation suspended without another checkpoint. Newer changes + # must remain available for the next invocation or checkpoint. + OperationPaginatorState.pin(reset_target).advance_handler_seen( + invocation_seq + ) self._store.save(reset_target) # Release the gate: terminal -> COMPLETED, otherwise PRE_INVOKE diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/e2e/invocation_input_delivery_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/invocation_input_delivery_test.py new file mode 100644 index 000000000..af69bc3dc --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/invocation_input_delivery_test.py @@ -0,0 +1,110 @@ +"""An input-only callback resume must not redeliver its terminal update later.""" + +import json +import time + +import pytest +from aws_durable_execution_sdk_python import ( + DurableContext, + durable_execution, + durable_step, +) +from aws_durable_execution_sdk_python.config import CallbackConfig +from aws_durable_execution_sdk_python.serdes import JsonSerDes +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from aws_durable_execution_sdk_python_testing.stores.filesystem import ( + FileSystemExecutionStore, +) +from aws_durable_execution_sdk_python_testing.stores.memory import ( + InMemoryExecutionStore, +) + + +@pytest.mark.parametrize("filesystem", [False, True]) +@pytest.mark.parametrize("failed_delivery", [False, True]) +def test_callback_input_acknowledgment_survives_checkpoint_free_resume( + filesystem, failed_delivery, tmp_path +): + store = ( + FileSystemExecutionStore.create(tmp_path / "store") + if filesystem + else InMemoryExecutionStore() + ) + starts = [] + invocation_ends = [] + operation_ends = [] + attempts = [] + marker_calls = [] + + class Observe(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + starts.append(info) + + def on_invocation_end(self, info): + invocation_ends.append(info) + + def on_operation_end(self, info): + if info.name in ("a", "b"): + operation_ends.append((len(starts), info.name, info.is_replayed)) + + @durable_step + def marker(context, value): + marker_calls.append(1) + return value + + @durable_execution(plugins=[Observe()]) + def handler(event, context: DurableContext): + a = context.create_callback( + name="a", config=CallbackConfig(serdes=JsonSerDes()) + ) + b = context.create_callback( + name="b", config=CallbackConfig(serdes=JsonSerDes()) + ) + result = [a.result(), b.result()] + return context.step(marker(result), name="after-b") + + def entrypoint(event, context): + attempts.append(1) + if failed_delivery and len(attempts) == 2: + raise RuntimeError("invocation failed before the SDK received its input") + return handler(event, context) + + with DurableFunctionTestRunner( + handler=entrypoint, + store=store, + skip_time=False, + poll_interval=0.01, + execution_timeout=25, + ) as runner: + arn = runner.run_async(input="{}") + a_id = runner.wait_for_callback(arn, name="a", timeout=5) + b_id = runner.wait_for_callback(arn, name="b", timeout=5) + runner.send_callback_success(a_id, result=b'"A"') + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + events = runner.get_execution_history( + arn, include_execution_data=True + ).events + if ( + len(invocation_ends) >= 2 + and sum(event.event_type == "InvocationCompleted" for event in events) + >= 2 + ): + break + time.sleep(0.01) + else: + raise AssertionError("The second SDK invocation did not suspend") + runner.send_callback_success(b_id, result=b'"B"') + result = runner.wait_for_result(arn, timeout=10) + + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == ["A", "B"] + assert [info.status.value for info in invocation_ends] == [ + "PENDING", + "PENDING", + "SUCCEEDED", + ] + assert operation_ends == [(2, "a", False), (3, "b", False)] + assert marker_calls == [1] + assert len(attempts) == (4 if failed_delivery else 3) diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/executor_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/executor_test.py index bb6b71f4f..8bd384aea 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/executor_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/executor_test.py @@ -698,6 +698,10 @@ def test_pending_response_is_valid_when_an_operation_completed_after_the_handler mock_execution.consecutive_failed_invocation_attempts = 0 mock_execution.seq_counter = 4 mock_execution.handler_seen_seq = 2 + mock_execution.token_sequence = 0 + mock_execution.operations = [] + mock_execution.updated_operation_ids = [] + mock_execution.operation_last_touched_seq = {} mock_execution.has_pending_operations.return_value = False mock_execution.has_changes_after.side_effect = lambda seq: seq == 4 mock_execution.has_unseen_changes.return_value = True @@ -1237,6 +1241,12 @@ def test_should_retry_invocation_when_under_limit_through_public_api( mock_execution.is_complete = False mock_execution.start_input = start_input mock_execution.consecutive_failed_invocation_attempts = 3 # Under limit (5 is max) + mock_execution.seq_counter = 0 + mock_execution.handler_seen_seq = 0 + mock_execution.token_sequence = 0 + mock_execution.operations = [] + mock_execution.updated_operation_ids = [] + mock_execution.operation_last_touched_seq = {} # Configure invoker to fail initially with validation error, then succeed on retry mock_invocation_input = Mock() diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/invocation_input_delivery_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/invocation_input_delivery_test.py new file mode 100644 index 000000000..42c64a35b --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/invocation_input_delivery_test.py @@ -0,0 +1,98 @@ +"""Invocation acknowledgments preserve newer checkpoint and completion state.""" + +from datetime import UTC, datetime +from unittest.mock import Mock, patch + +import pytest + +from aws_durable_execution_sdk_python.execution import ( + DurableExecutionInvocationOutput, + InvocationStatus, +) +from aws_durable_execution_sdk_python.lambda_service import ( + CallbackDetails, + Operation, + OperationStatus, + OperationType, +) +from aws_durable_execution_sdk_python_testing.execution import ( + Execution, + OperationPaginatorState, +) +from aws_durable_execution_sdk_python_testing.executor import Executor +from aws_durable_execution_sdk_python_testing.invoker import InvokeResponse +from aws_durable_execution_sdk_python_testing.model import StartDurableExecutionInput +from aws_durable_execution_sdk_python_testing.stores.memory import ( + InMemoryExecutionStore, +) +from aws_durable_execution_sdk_python_testing.token import CallbackToken + + +@pytest.mark.parametrize("checkpoint_during_invocation", [False, True]) +def test_input_ack_preserves_newer_completion_and_checkpoint_watermark( + checkpoint_during_invocation, +): + store = InMemoryExecutionStore() + execution = Execution.new( + StartDurableExecutionInput( + account_id="123456789012", + function_name="test", + function_qualifier="$LATEST", + execution_name="test", + execution_timeout_seconds=30, + execution_retention_period_days=1, + invocation_id="invocation", + ) + ) + execution.start() + arn = execution.durable_execution_arn + tokens = {} + for name in ("a", "b", "c"): + tokens[name] = CallbackToken(execution_arn=arn, operation_id=name).to_str() + execution.operations.append( + Operation( + operation_id=name, + operation_type=OperationType.CALLBACK, + status=OperationStatus.STARTED, + callback_details=CallbackDetails(callback_id=tokens[name]), + ) + ) + execution.touch_operation(name) + OperationPaginatorState.pin(execution).advance_handler_seen(execution.seq_counter) + execution.complete_callback_success(tokens["a"], b'"A"') + input_sequence = execution.seq_counter + if checkpoint_during_invocation: + execution.complete_callback_success(tokens["c"], b'"C"') + OperationPaginatorState.pin(execution).advance_handler_seen( + execution.seq_counter + ) + already_seen = execution.handler_seen_seq + execution.complete_callback_success(tokens["b"], b'"B"') + execution.needs_reinvoke = True + store.save(execution) + executor = Executor(store, Mock(), Mock(), Mock()) + try: + with patch.object(executor, "_invoke_execution") as reinvoke: + executor._finish_invocation( # noqa: SLF001 + arn, + InvokeResponse( + invocation_output=DurableExecutionInvocationOutput( + status=InvocationStatus.PENDING + ), + request_id="request", + ), + datetime.now(UTC), + datetime.now(UTC), + input_sequence, + ) + saved = store.load(arn) + assert saved.handler_seen_seq == max(already_seen, input_sequence) + assert saved.updated_operation_ids == ["b"] + assert [ + op.operation_id + for op in OperationPaginatorState.pin(saved).unseen_operations() + ] == ["b"] + reinvoke.assert_called_once_with(arn) + finally: + executor.shutdown() + executor._registry.shutdown() # noqa: SLF001 From 24b02665c3b90ac87a6aa8e86da6084d1e9c5ed6 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 14 Oct 2026 16:24:29 +0000 Subject: [PATCH 38/54] fix(otel): preserve explicit failure statuses --- .../test_callback_failure_without_error.py | 4 +- .../README.md | 12 +- .../execution_plugin.py | 6 + .../invocation_plugin.py | 6 + .../e2e/test_invocation_wait_resume_int.py | 125 +++++++++++++ .../tests/e2e/test_terminal_status_int.py | 169 ++++++++++++++++++ .../tests/test_invocation_plugin.py | 6 +- .../tests/test_status_mapping.py | 4 +- 8 files changed, 321 insertions(+), 11 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_terminal_status_int.py diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_failure_without_error.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_failure_without_error.py index a04a92c0d..2c7f41b80 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_failure_without_error.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_callback_failure_without_error.py @@ -153,9 +153,7 @@ def observe_error(info: OperationEndInfo) -> None: ] == [error.to_dict() if has_details and error is not None else None] if has_details and not filesystem: assert callback_errors[0] is error - assert leaves[0].status.status_code.name == ( - "ERROR" if has_details else "UNSET" - ) + assert leaves[0].status.status_code.name == "ERROR" assert [event.name for event in leaves[0].events] == ( ["exception"] if has_details else [] ) diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index d44c59506..3cba32506 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -395,10 +395,14 @@ each durable span in the same invocation. `durable.invocation.status` uses `RETRYING` when the core plugin hook reports `InvocationStatus.RETRY`. The core enum remains unchanged. Operation spans use -OTel `OK` only for `SUCCEEDED`, `ERROR` when error details are delivered, and -`UNSET` for other outcomes without error details, including `FAILED`, -`CANCELLED`, `TIMED_OUT`, and `STOPPED`. The original durable operation status -remains in `durable.operation.status`. +OTel `ERROR` for explicit `FAILED`, `TIMED_OUT`, or `STOPPED` outcomes even when +error details are absent. Supplied error details continue to set `ERROR` and +record an exception event; a failure status alone does not invent that event. +Without error details, `SUCCEEDED` is `OK` and other statuses, including +`CANCELLED`, remain `UNSET`. The original durable operation status remains in +`durable.operation.status`. A handler that catches an operation failure and +returns successfully still has a successful invocation/workflow outcome; the +failed operation retains its own error status on initial completion and replay. When a concurrent branch's terminal completion arrives before replay has created its parent span, Invocation view retains that completion until the actual parent 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 2008615e1..1fea5e665 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 @@ -760,6 +760,12 @@ def on_operation_end(self, info: OperationEndInfo) -> None: span.record_exception( Exception(info.error.message or info.error.type or "Unknown error") ) + elif info.status in ( + OperationStatus.FAILED, + OperationStatus.TIMED_OUT, + OperationStatus.STOPPED, + ): + span.set_status(StatusCode.ERROR) elif info.status is OperationStatus.SUCCEEDED: span.set_status(StatusCode.OK) 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 299e3803f..c053a8bc5 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 @@ -851,6 +851,12 @@ def on_operation_end(self, info: OperationEndInfo) -> None: span.record_exception( Exception(info.error.message or info.error.type or "Unknown error") ) + elif info.status in ( + OperationStatus.FAILED, + OperationStatus.TIMED_OUT, + OperationStatus.STOPPED, + ): + span.set_status(StatusCode.ERROR) elif info.status is OperationStatus.SUCCEEDED: span.set_status(StatusCode.OK) 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 fa7036920..7477a779d 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 @@ -15,10 +15,13 @@ InvocationStatus, durable_execution, ) +from aws_durable_execution_sdk_python.exceptions import InvokeError from aws_durable_execution_sdk_python.lambda_service import ( + ChainedInvokeDetails, CheckpointOutput, CheckpointUpdatedExecutionState, ExecutionDetails, + ErrorObject, Operation, OperationAction, OperationStatus, @@ -142,6 +145,128 @@ def checkpoint( return checkpoint, operations +@pytest.mark.parametrize("plugin_type", [ExecutionOtelPlugin, InvocationOtelPlugin]) +@pytest.mark.parametrize("with_error", [False, True]) +def test_caught_invoke_failure_without_details_is_error_on_live_and_replay( + plugin_type, with_error +) -> None: + """Use real SDK replay with service-shaped failed-invoke checkpoint responses. + + The local child runner synthesizes a diagnostic for FAILED children, so the + absent-error case belongs at the LambdaClient transport boundary. No plugin + lifecycle/status callback is fabricated here. + """ + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = plugin_type( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + error = ( + ErrorObject( + message="supplied invoke failure", + type="External", + data=None, + stack_trace=None, + ) + if with_error + else None + ) + caught: list[str] = [] + saved_calls: list[str] = [] + + @durable_step + def saved(_context) -> str: + saved_calls.append("saved") + return "caught" + + def body(_event, context: DurableContext) -> str: + try: + context.invoke("child", {}, name="failed-invoke") + except InvokeError as observed: + caught.append(type(observed).__name__) + value = context.step(saved(), name="saved-result") + context.wait(Duration.from_seconds(1), name="resume-wait") + return value + + handler = durable_execution(body, plugins=[plugin]) + + def failed_invoke_store(initial): + checkpoint, operations = _checkpoint_store(initial) + + def send(*args, **kwargs): + output = checkpoint(*args, **kwargs) + for key, operation in operations.items(): + if operation.name == "failed-invoke": + operations[key] = replace( + operation, + status=OperationStatus.FAILED, + sub_type=OperationSubType.CHAINED_INVOKE, + end_timestamp=datetime.now(UTC), + chained_invoke_details=ChainedInvokeDetails(error=error), + ) + return replace( + output, + new_execution_state=CheckpointUpdatedExecutionState( + operations=list(operations.values()) + ), + ) + + return send, operations + + initial = [_execution_operation()] + checkpoint, operations = failed_invoke_store(initial) + 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 + first = handler(_event(initial), _lambda_context()) + assert first["Status"] == "PENDING" + first_spans = [ + s for s in exporter.get_finished_spans() if s.name == "failed-invoke" + ] + resumed = [ + replace( + op, status=OperationStatus.SUCCEEDED, end_timestamp=datetime.now(UTC) + ) + if op.name == "resume-wait" + else op + for op in operations.values() + ] + wait_id = next(op.operation_id for op in resumed if op.name == "resume-wait") + checkpoint, _ = failed_invoke_store(resumed) + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as client_class: + client = Mock() + client.checkpoint = checkpoint + client_class.initialize_client.return_value = client + second = handler(_event(resumed, [wait_id]), _lambda_context()) + assert second == {"Status": "SUCCEEDED", "Result": '"caught"'} + assert caught == ["InvokeError", "InvokeError"] + assert saved_calls == ["saved"] + spans = [s for s in exporter.get_finished_spans() if s.name == "failed-invoke"] + assert spans == first_spans + (span,) = spans + assert span.attributes is not None + assert span.attributes["durable.operation.status"] == "FAILED" + assert span.status.status_code is trace.StatusCode.ERROR + assert len(span.events) == int(with_error) + workflow = next( + s for s in exporter.get_finished_spans() if s.name == "Workflow" + ) + assert workflow.status.status_code is trace.StatusCode.OK + finally: + provider.shutdown() + + @pytest.mark.parametrize( "plugin_type", [InvocationOtelPlugin, ExecutionOtelPlugin], diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_terminal_status_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_terminal_status_int.py new file mode 100644 index 000000000..6e0dbb692 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_terminal_status_int.py @@ -0,0 +1,169 @@ +"""A caught durable failure remains an error span without invented diagnostics.""" + +import json +import threading + +import pytest +from aws_durable_execution_sdk_python import durable_execution, durable_step +from aws_durable_execution_sdk_python.exceptions import CallbackError +from aws_durable_execution_sdk_python.lambda_service import ErrorObject +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationStatus, +) +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +@pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) +@pytest.mark.parametrize("completion", ["initial", "resume"]) +@pytest.mark.parametrize("outcome", ["empty-failure", "rich-failure", "success"]) +def test_caught_callback_status_survives_two_later_replays( + view, completion, outcome, monkeypatch +): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + created = threading.Event() + completed = threading.Event() + suspended = threading.Event() + finished = threading.Condition() + completed_invocations = [] + invocations = [] + target_ends = [] + side_effects = [] + first_exports = [] + caught = [] + callback_ids = [] + + class Observer(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + invocations.append(info) + + def on_operation_end(self, info): + if info.name == "target": + target_ends.append((len(invocations), info)) + + def on_invocation_end(self, info): + if info.status is InvocationStatus.PENDING: + suspended.set() + with finished: + completed_invocations.append(info) + finished.notify_all() + + @durable_step + def saved(_context, name, value): + side_effects.append(name) + if name == "observed": + first_exports.extend( + span for span in exporter.get_finished_spans() if span.name == "target" + ) + return value + + def body(_event, context): + callback = context.create_callback(name="target") + context.step(saved("before", "saved"), name="before") + if completion == "initial" and len(invocations) == 1: + callback_ids.append(callback.callback_id) + created.set() + assert completed.wait(5), "Driver did not complete the callback" + # A real checkpoint refreshes an external completion during the active + # invocation; on resume this completed step is consumed from history. + context.step(saved("refresh", "saved"), name="refresh") + try: + callback.result() + result = "success" + except CallbackError as error: + result = type(error).__name__ + caught.append(result) + result = context.step(saved("observed", result), name="observed") + for name in ("barrier-1", "barrier-2"): + context.create_callback(name=name).result() + return result + + handler = durable_execution(body, plugins=[plugin, Observer()]) + error = ( + ErrorObject( + message="supplied detail", type="External", data=None, stack_trace=None + ) + if outcome == "rich-failure" + else None + ) + try: + with DurableFunctionTestRunner( + handler=handler, skip_time=False, poll_interval=0.01, execution_timeout=25 + ) as runner: + arn = runner.run_async(input="{}") + if completion == "initial": + assert created.wait(5) + callback_id = callback_ids[0] + else: + callback_id = runner.wait_for_callback(arn, name="target", timeout=5) + assert suspended.wait(5) + if outcome == "success": + runner.send_callback_success(callback_id) + else: + runner.send_callback_failure(callback_id, error=error) + completed.set() + first_completion = 1 if completion == "initial" else 2 + for index, name in enumerate(("barrier-1", "barrier-2")): + callback_id = runner.wait_for_callback(arn, name=name, timeout=5) + with finished: + assert finished.wait_for( + lambda: len(completed_invocations) >= first_completion + index, + timeout=5, + ) + runner.send_callback_success(callback_id) + result = runner.wait_for_result(arn, timeout=10) + assert result.status.value == "SUCCEEDED" + expected_result = "success" if outcome == "success" else "CallbackExternalError" + assert json.loads(result.result) == expected_result + assert side_effects == ["before", "refresh", "observed"] + assert len(invocations) == (3 if completion == "initial" else 4) + assert caught == ([] if outcome == "success" else [expected_result] * 3) + assert len(target_ends) == 1 + invocation_number, info = target_ends[0] + assert invocation_number == (1 if completion == "initial" else 2) + assert info.status.value == ("SUCCEEDED" if outcome == "success" else "FAILED") + assert info.error == error + spans = [ + span for span in exporter.get_finished_spans() if span.name == "target" + ] + expected_statuses = [info.status.value] + if view is InvocationOtelPlugin and completion == "resume": + expected_statuses.insert(0, "STARTED") + assert [ + span.attributes["durable.operation.status"] for span in spans + ] == expected_statuses + assert first_exports == spans + terminal = spans[-1] + assert terminal.status.status_code is ( + StatusCode.OK if outcome == "success" else StatusCode.ERROR + ) + assert len(terminal.events) == (1 if error is not None else 0) + if error is not None: + assert terminal.events[0].attributes["exception.message"] == error.message + workflow = next( + s for s in exporter.get_finished_spans() if s.name == "Workflow" + ) + assert workflow.status.status_code is StatusCode.OK + finally: + completed.set() + provider.shutdown() 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 e7ae6cf67..773555a76 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 @@ -2306,10 +2306,10 @@ def body() -> None: ("status", "error", "expected"), [ (OperationStatus.SUCCEEDED, None, StatusCode.OK), - (OperationStatus.FAILED, None, StatusCode.UNSET), + (OperationStatus.FAILED, None, StatusCode.ERROR), (OperationStatus.CANCELLED, None, StatusCode.UNSET), - (OperationStatus.TIMED_OUT, None, StatusCode.UNSET), - (OperationStatus.STOPPED, None, StatusCode.UNSET), + (OperationStatus.TIMED_OUT, None, StatusCode.ERROR), + (OperationStatus.STOPPED, None, StatusCode.ERROR), ( OperationStatus.FAILED, ErrorObject(message="failure", type="Example", data=None, stack_trace=None), diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_status_mapping.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_status_mapping.py index d1224a305..b5423c726 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_status_mapping.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_status_mapping.py @@ -79,7 +79,7 @@ def test_operation_status_and_exception( with_error: bool, started_here: bool, ) -> None: - """A delivered terminal hook must not label an unsuccessful operation OK.""" + """Explicit failures are errors even when diagnostic details are absent.""" plugin, exporter = instrumented_view invocation: dict[str, Any] = dict( request_id="first", @@ -123,6 +123,8 @@ def test_operation_status_and_exception( expected = ( StatusCode.ERROR if with_error + or status + in (OperationStatus.FAILED, OperationStatus.TIMED_OUT, OperationStatus.STOPPED) else ( StatusCode.OK if status is OperationStatus.SUCCEEDED else StatusCode.UNSET ) From 4b8d6a40c7e4c4bc133e267647716ba812720e25 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 14 Oct 2026 17:01:03 +0000 Subject: [PATCH 39/54] ci: adopt corrected terminal status catalog --- .../scripts/tests/test_opentelemetry_conformance_workflow.py | 2 +- .github/workflows/opentelemetry-conformance-tests.yml | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py index 1479be395..c29b3f75e 100644 --- a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py +++ b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py @@ -5,7 +5,7 @@ SHARED_WORKFLOW_REF = "bdb4f1cd0f9252c1aaa978bb8b341b71f2b9d9dc" -CONFORMANCE_TEST_REF = "fb69b21ccc8805a74271554403707f5e649e5aae" +CONFORMANCE_TEST_REF = "f711bbfeb5464400d3aa1f8366e0bbb02087a67d" WORKFLOW_PATH = ( Path(__file__).parents[2] / "workflows" / "opentelemetry-conformance-tests.yml" diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index ec8d40699..d4177270c 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -45,7 +45,7 @@ on: conformance_test_ref: description: Conformance test commit SHA or branch name required: true - default: fb69b21ccc8805a74271554403707f5e649e5aae + default: f711bbfeb5464400d3aa1f8366e0bbb02087a67d type: string # Backend stacks are shared across PRs. Queue whole runs so reusable @@ -73,7 +73,7 @@ jobs: resource_prefix: p sdk_repository: aws/aws-durable-execution-sdk-python sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }} - conformance_test_ref: ${{ inputs.conformance_test_ref || 'fb69b21ccc8805a74271554403707f5e649e5aae' }} + conformance_test_ref: ${{ inputs.conformance_test_ref || 'f711bbfeb5464400d3aa1f8366e0bbb02087a67d' }} # Check the SDK out so the handlers and templates below are on disk. The handlers # themselves are installed from sdk_ref by src/requirements.txt during the SAM build. checkout_sdk: true From 81cf5262d2b10f179b56feb6713fc2a1b2b6bbb2 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 14 Oct 2026 17:18:19 +0000 Subject: [PATCH 40/54] fix(otel): keep exception events within spans --- .../README.md | 5 ++ .../execution_plugin.py | 6 +- .../invocation_plugin.py | 30 ++++--- .../e2e/test_exception_event_bounds_int.py | 78 +++++++++++++++++++ .../tests/test_execution_plugin.py | 8 +- .../tests/test_invocation_plugin.py | 8 +- .../tests/test_status_mapping.py | 40 ++++++++++ 7 files changed, 160 insertions(+), 15 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_exception_event_bounds_int.py diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 3cba32506..1c814a6b9 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -404,6 +404,11 @@ Without error details, `SUCCEEDED` is `OK` and other statuses, including returns successfully still has a successful invocation/workflow outcome; the failed operation retains its own error status on initial completion and replay. +Exception events use the same normalized end timestamp as their operation or +attempt span. This keeps events inside the exported span when a completion has +no durable start time or arrives later than its backend end time; the existing +view-specific span-boundary rules remain intact. + When a concurrent branch's terminal completion arrives before replay has created its parent span, Invocation view retains that completion until the actual parent span is registered. The terminal segment is exported under that parent before 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 1fea5e665..0a7a6e1ec 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 @@ -758,7 +758,8 @@ def on_operation_end(self, info: OperationEndInfo) -> None: if info.error: span.set_status(StatusCode.ERROR, info.error.message or "") span.record_exception( - Exception(info.error.message or info.error.type or "Unknown error") + Exception(info.error.message or info.error.type or "Unknown error"), + timestamp=_to_otel_timestamp(end_time), ) elif info.status in ( OperationStatus.FAILED, @@ -901,7 +902,8 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: (info.error.message or info.error.type) if info.error else "Unknown error" - ) + ), + timestamp=_to_otel_timestamp(end_time), ) else: span.set_status(StatusCode.OK) 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 c053a8bc5..f338f806b 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 @@ -512,7 +512,11 @@ def _start_span( return span def _end_span( - self, operation_id: str | None, end_timestamp: datetime.datetime | None = None + self, + operation_id: str | None, + end_timestamp: datetime.datetime | None = None, + *, + exception: Exception | None = None, ): """End and unregister the active span for an operation ID. @@ -521,6 +525,8 @@ def _end_span( span. end_timestamp: Optional durable end timestamp to use as the span end time. When omitted, OpenTelemetry uses the current time. + exception: Existing failure details to record at the same normalized + timestamp as the span end. """ logger.debug("Ending OTel span: operation_id=%s", operation_id) with self._operation_spans_lock: @@ -528,6 +534,8 @@ def _end_span( if span: # the span is not going to be populated if it has the same end_time and start_time end_time = self._next_ordered_timestamp(end_timestamp) + if exception is not None: + span.record_exception(exception, timestamp=end_time) span.end(end_time=end_time) logger.debug("Ended OTel span: %s", span) @@ -846,10 +854,11 @@ def on_operation_end(self, info: OperationEndInfo) -> None: else: span.set_attributes(self._extract_attributes(info)) + exception: Exception | None = None if info.error: span.set_status(StatusCode.ERROR, info.error.message or "") - span.record_exception( - Exception(info.error.message or info.error.type or "Unknown error") + exception = Exception( + info.error.message or info.error.type or "Unknown error" ) elif info.status in ( OperationStatus.FAILED, @@ -860,7 +869,7 @@ def on_operation_end(self, info: OperationEndInfo) -> None: elif info.status is OperationStatus.SUCCEEDED: span.set_status(StatusCode.OK) - self._end_span(info.operation_id, info.end_time) + self._end_span(info.operation_id, info.end_time, exception=exception) def on_user_function_start(self, info: UserFunctionStartInfo) -> None: """Called when a context or step operation starts user code. @@ -972,16 +981,15 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: and info.outcome is not UserFunctionOutcome.INCOMPLETE ): span.set_attributes(self._extract_attributes(info)) + exception: Exception | None = None if info.outcome is UserFunctionOutcome.FAILED: span.set_status( StatusCode.ERROR, info.error.message if info.error else "" ) - span.record_exception( - Exception( - (info.error.message or info.error.type) - if info.error - else "Unknown error" - ) + exception = Exception( + (info.error.message or info.error.type) + if info.error + else "Unknown error" ) else: span.set_status(StatusCode.OK) @@ -989,7 +997,7 @@ def on_user_function_end(self, info: UserFunctionEndInfo) -> None: end_timestamp = info.end_time if end_timestamp is not None and end_timestamp == info.start_time: end_timestamp += datetime.timedelta(microseconds=1) - self._end_span(span_key, end_timestamp) + self._end_span(span_key, end_timestamp, exception=exception) # Restore the enclosing context by releasing the scope this user # function attached. Code that runs after this operation (e.g. between # steps in a child context) correlates to its enclosing operation diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_exception_event_bounds_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_exception_event_bounds_int.py new file mode 100644 index 000000000..ed72ee651 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_exception_event_bounds_int.py @@ -0,0 +1,78 @@ +"""Real checkpointless child failures keep diagnostic events inside their span.""" + +import json + +import pytest +from aws_durable_execution_sdk_python import ( + durable_execution, + durable_with_child_context, +) +from aws_durable_execution_sdk_python.config import ChildConfig +from aws_durable_execution_sdk_python.exceptions import ChildContextError +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +@pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_caught_virtual_child_failure_has_bounded_event_time(view, monkeypatch): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + enrich_logger=False, + context_extractor=lambda _: None, + ) + ) + hooks = [] + + class Observe(DurableInstrumentationPlugin): + def on_operation_end(self, info): + if info.name == "virtual-error": + hooks.append(info) + + @durable_with_child_context + def child(_context): + raise ValueError("actual child failure") + + def body(_event, context): + try: + context.run_in_child_context( + child(), name="virtual-error", config=ChildConfig(is_virtual=True) + ) + except ChildContextError: + return "caught" + raise AssertionError("Expected the stored child failure") + + try: + with DurableFunctionTestRunner( + handler=durable_execution(body, plugins=[plugin, Observe()]), + poll_interval=0.01, + ) as runner: + result = runner.run(input="{}") + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == "caught" + (info,) = hooks + assert info.start_time is None + assert info.status.value == "FAILED" + assert info.error.message == "actual child failure" + (span,) = [ + s for s in exporter.get_finished_spans() if s.name == "virtual-error" + ] + assert span.status.status_code.name == "ERROR" + assert len(span.events) == 1 + assert span.events[0].attributes["exception.message"] == "actual child failure" + assert span.start_time <= span.events[0].timestamp <= span.end_time + finally: + provider.shutdown() 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 a3d79c8aa..9ef59c1d4 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 @@ -1300,7 +1300,7 @@ def test_invocation_end_restores_context_from_before_invocation_start(status): ) def test_step_scope_is_released_at_user_function_end(outcome): """Verify a finished step restores the context that enclosed it.""" - plugin, _ = _create_plugin() + plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) enclosing_context = otel_context.get_current() @@ -1320,6 +1320,12 @@ def test_step_scope_is_released_at_user_function_end(outcome): plugin.on_invocation_end(_invocation_end_info()) assert plugin._context_tokens == {} + events = [ + (span, event) for span in exporter.get_finished_spans() for event in span.events + ] + assert len(events) == (1 if outcome is UserFunctionOutcome.FAILED else 0) + for span, event in events: + assert span.start_time <= event.timestamp <= span.end_time def test_user_function_start_preserves_baggage_in_current_context(): 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 773555a76..ba9833f56 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 @@ -1304,7 +1304,7 @@ def test_user_function_start_preserves_baggage_in_current_context(): def test_user_function_end_restores_enclosing_context_on_failure(): """Verify the enclosing context is restored even when the step fails.""" - plugin, _ = _create_plugin() + plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) enclosing_context = otel_context.get_current() @@ -1317,6 +1317,12 @@ def test_user_function_end_restores_enclosing_context_on_failure(): assert otel_context.get_current() == enclosing_context assert set(plugin._context_tokens) == {"__invocation_context__"} plugin.on_invocation_end(_invocation_end_info()) + events = [ + (span, event) for span in exporter.get_finished_spans() for event in span.events + ] + assert len(events) == 1 + for span, event in events: + assert span.start_time <= event.timestamp <= span.end_time def test_user_function_end_restores_enclosing_context_across_multiple_steps(): diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_status_mapping.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_status_mapping.py index b5423c726..a97df2e1c 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_status_mapping.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_status_mapping.py @@ -136,10 +136,50 @@ def test_operation_status_and_exception( assert span.events[0].name == "exception" assert span.events[0].attributes is not None assert span.events[0].attributes["exception.message"] == "deadline exceeded" + assert span.start_time is not None and span.end_time is not None + assert span.start_time <= span.events[0].timestamp <= span.end_time else: assert not span.events +def test_error_event_is_inside_completion_without_a_start( + instrumented_view: InstrumentedView, +) -> None: + plugin, exporter = instrumented_view + invocation: dict[str, Any] = dict( + request_id="missing-start", + execution_arn=ARN, + execution_start_time=START, + is_first_invocation=True, + ) + plugin.on_invocation_start(InvocationStartInfo(**invocation)) + plugin.on_operation_end( + OperationEndInfo( + operation_id="missing-start", + name="missing-start", + parent_id=None, + operation_type=OperationType.WAIT, + sub_type=OperationSubType.WAIT, + start_time=None, + end_time=END, + is_replayed=False, + status=OperationStatus.FAILED, + error=ErrorObject( + message="real failure", type="Failure", data=None, stack_trace=None + ), + ) + ) + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + span = next(s for s in exporter.get_finished_spans() if s.name == "missing-start") + assert len(span.events) == 1 + assert span.events[0].attributes is not None + assert span.events[0].attributes["exception.message"] == "real failure" + assert span.start_time is not None and span.end_time is not None + assert span.start_time <= span.events[0].timestamp <= span.end_time + + def test_retry_then_resume_exports_retrying( instrumented_view: InstrumentedView, ) -> None: From b6ba1590e8f948c9a62cc357795bcff06e519d00 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 14 Oct 2026 18:00:50 +0000 Subject: [PATCH 41/54] fix(core): isolate invocation plugin executors --- .../execution.py | 2 +- .../e2e/plugin_invocation_overlap_int_test.py | 65 +++++++++++++++++ .../tests/execution_test.py | 73 +++++++++++++++++++ 3 files changed, 139 insertions(+), 1 deletion(-) create mode 100644 packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_overlap_int_test.py 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 e59820007..d9c0f157e 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 @@ -190,9 +190,9 @@ def durable_execution( logger.debug("Starting durable execution handler...") configured_plugins = load_configured_plugins(plugins) - plugin_executor = PluginExecutor(configured_plugins) def wrapper(event: Any, context: LambdaContext) -> MutableMapping[str, Any]: + plugin_executor = PluginExecutor(configured_plugins) invocation_input: DurableExecutionInvocationInput service_client: DurableServiceClient diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_overlap_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_overlap_int_test.py new file mode 100644 index 000000000..6073b9fce --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_invocation_overlap_int_test.py @@ -0,0 +1,65 @@ +"""Two local executions may overlap without sharing invocation-hook state.""" + +import contextvars +import json +import threading + +from aws_durable_execution_sdk_python import durable_execution +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner + + +def test_shared_handler_reports_both_invocation_ends(monkeypatch): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + ready = threading.Event() + release = threading.Event() + binding: contextvars.ContextVar[str | None] = contextvars.ContextVar( + "integration-plugin-owner", default=None + ) + tokens: contextvars.ContextVar[contextvars.Token[str | None] | None] = ( + contextvars.ContextVar("integration-plugin-token", default=None) + ) + starts = [] + ends = [] + + class SafePlugin(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + tokens.set(binding.set(info.execution_arn)) + starts.append(info.execution_arn) + + def on_invocation_end(self, info): + assert binding.get() == info.execution_arn + token = tokens.get() + assert token is not None + binding.reset(token) + ends.append(info.execution_arn) + + def body(event, _context): + if event["label"] == "a": + ready.set() + assert release.wait(5) + return event["label"] + + wrapped = durable_execution(body, plugins=[SafePlugin()]) + try: + with ( + DurableFunctionTestRunner(handler=wrapped, poll_interval=0.01) as first, + DurableFunctionTestRunner(handler=wrapped, poll_interval=0.01) as second, + ): + first_arn = first.run_async(input='{"label":"a"}') + assert ready.wait(5) + try: + second_result = second.run(input='{"label":"b"}') + finally: + release.set() + first_result = first.wait_for_result(first_arn, timeout=5) + assert first_result.status.value == second_result.status.value == "SUCCEEDED" + assert [json.loads(first_result.result), json.loads(second_result.result)] == [ + "a", + "b", + ] + assert len(starts) == 2 and len(set(starts)) == 2 + assert ends == list(reversed(starts)) + assert binding.get() is None + finally: + release.set() diff --git a/packages/aws-durable-execution-sdk-python/tests/execution_test.py b/packages/aws-durable-execution-sdk-python/tests/execution_test.py index 85ea72a97..3a6c9c1c1 100644 --- a/packages/aws-durable-execution-sdk-python/tests/execution_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/execution_test.py @@ -4,7 +4,11 @@ import json import time import warnings +import contextvars +import threading from collections.abc import Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import replace from copy import deepcopy from typing import Any from unittest.mock import Mock, patch @@ -3503,6 +3507,75 @@ def test_handler(event: Any, context: DurableContext) -> dict: assert callable(test_handler) +def test_overlapping_handler_calls_keep_each_plugin_end_and_payload(): + """One wrapper must not let another call replace its invocation-hook state.""" + entered = threading.Event() + release = threading.Event() + binding: contextvars.ContextVar[str | None] = contextvars.ContextVar( + "unit-plugin-owner", default=None + ) + tokens: contextvars.ContextVar[contextvars.Token[str | None] | None] = ( + contextvars.ContextVar("unit-plugin-token", default=None) + ) + starts = [] + ends = [] + + class SafePlugin(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + tokens.set(binding.set(info.request_id)) + starts.append(info.request_id) + + def on_invocation_end(self, info): + assert binding.get() == info.request_id + token = tokens.get() + assert token is not None + binding.reset(token) + ends.append( + ( + info.request_id, + info.execution_input["label"], + json.loads(info.execution_result), + ) + ) + + def body(event, _context): + if event["label"] == "a": + entered.set() + assert release.wait(5) + return event["label"] + + wrapped = durable_execution(body, plugins=[SafePlugin()]) + client = Mock(spec=DurableServiceClient) + client.checkpoint.return_value = CheckpointOutput( + checkpoint_token="new_token", + new_execution_state=CheckpointUpdatedExecutionState(), + ) + + def call(label): + incoming = replace( + _make_invocation_input(client, input_payload=json.dumps({"label": label})), + durable_execution_arn=f"arn:test:execution-{label}/exec1", + ) + context = _make_lambda_context() + context.aws_request_id = label + return wrapped(incoming, context) + + with ThreadPoolExecutor(max_workers=2) as callers: + first = callers.submit(call, "a") + try: + assert entered.wait(5) + second = callers.submit(call, "b") + second_result = second.result(timeout=5) + finally: + release.set() + first_result = first.result(timeout=5) + assert first_result == {"Status": "SUCCEEDED", "Result": '"a"'} + assert second_result == {"Status": "SUCCEEDED", "Result": '"b"'} + assert starts == ["a", "b"] + assert ends == [("b", "b", "b"), ("a", "a", "a")] + assert binding.get() is None + + def test_durable_execution_with_plugins_success(): """Test that plugins receive invocation start/end and execution end on success.""" mock_client = Mock(spec=DurableServiceClient) From 413c2b63b8e4b4e7cc6ebda9d258770814add7f7 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 14 Oct 2026 16:48:37 +0000 Subject: [PATCH 42/54] fix(otel): bound provider flush waiting --- .../README.md | 21 ++ .../_flush.py | 112 +++++++++ .../execution_plugin.py | 3 +- .../invocation_plugin.py | 3 +- .../tests/e2e/test_flush_lifecycle_int.py | 215 ++++++++++++++++++ .../tests/test_flush.py | 132 +++++++++++ 6 files changed, 484 insertions(+), 2 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_flush.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_flush_lifecycle_int.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/test_flush.py diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 1c814a6b9..1d02f9b1b 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -416,6 +416,27 @@ control returns to the branch's user code. No ancestor is invented and SDK completion delivery is unchanged. If the parent never becomes active, normal invocation cleanup and flushing still run before the missing parent is reported. +### Invocation flushing + +After span and ContextVar cleanup, the plugins wait up to the OpenTelemetry +provider's existing 30-second default budget for `force_flush()`. Only provider +I/O runs on a background daemon worker; invocation Start/End hooks and token +cleanup stay on their original worker and Context. The I/O receives a snapshot +of that post-cleanup Context. + +There is at most one outstanding provider flush in the process and no pending +queue. A timeout does not cancel the exporter or free that slot. Later requests +skip flushing with a warning until the earlier call actually finishes, including +requests using another provider. A failed or skipped flush is not an export +acknowledgment, and late exports may continue. Applications still own exporter +transport timeouts and provider shutdown. This bound does not cover synchronous +span-processor `on_end` work that runs before flushing. + +Ordinary provider exceptions that finish within the budget retain the existing +plugin exception isolation; fatal exceptions are propagated on the lifecycle +owner. Exceptions after the caller has stopped waiting are logged and cannot be +raised retroactively into the completed invocation. + ### Invocation context isolation With core 2.1+, invocation hooks and the handler run on one worker in an diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_flush.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_flush.py new file mode 100644 index 000000000..e288528ff --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_flush.py @@ -0,0 +1,112 @@ +"""Bound caller waiting for provider I/O without moving plugin lifecycle hooks.""" + +from __future__ import annotations + +import contextvars +import logging +import threading +import time +from collections.abc import Callable +from dataclasses import dataclass, field + + +logger = logging.getLogger(__name__) + +# Match the OTel provider's existing default budget. Some processors ignore +# that cooperative budget, so the SDK bounds its own wait independently. +_FLUSH_TIMEOUT_SECONDS = 30.0 +_lock = threading.Lock() + + +@dataclass +class _Attempt: + done: threading.Event = field(default_factory=threading.Event) + result: object = None + error: BaseException | None = None + timed_out: bool = False + + +_active: _Attempt | None = None + + +def _run(flush: Callable[[], object], attempt: _Attempt) -> None: + global _active # noqa: PLW0603 + try: + try: + attempt.result = flush() + except BaseException as error: + # Fatal exceptions completed within the wait are re-raised on the + # lifecycle owner. A late error cannot be thrown into a returned call. + attempt.error = error + with _lock: + late = attempt.timed_out + if not late and _active is attempt: + _active = None + attempt.done.set() + if late and attempt.error is not None: + logger.error( + "Provider flush failed after the caller stopped waiting", + exc_info=( + type(attempt.error), + attempt.error, + attempt.error.__traceback__, + ), + ) + finally: + # Retain the slot through actual completion, even after timeout. There + # is no queue and no new worker per invocation for a stuck exporter. + with _lock: + if _active is attempt: + _active = None + + +def flush_provider(flush: Callable[[], object]) -> bool: + """Wait for provider I/O, with one outstanding attempt process-wide. + + Start/End and ContextVar-token cleanup stay on their existing owner. Only + the supplied provider method runs elsewhere, using the post-cleanup Context + snapshot. Busy/expired attempts return False; they are never cancelled. + This does not bound synchronous span processors' earlier ``on_end`` calls, + application-owned provider shutdown, or exporter transport I/O itself. + """ + global _active # noqa: PLW0603 + deadline = time.monotonic() + _FLUSH_TIMEOUT_SECONDS + with _lock: + if _active is not None: + attempt = None + else: + attempt = _Attempt() + _active = attempt + if attempt is None: + logger.warning("Provider flush is still running; skipping another flush") + return False + + owner_context = contextvars.copy_context() + try: + threading.Thread( + target=owner_context.run, + args=(_run, flush, attempt), + name="durable-otel-flush", + daemon=True, + ).start() + except BaseException: + with _lock: + if _active is attempt: + _active = None + raise + + completed = attempt.done.wait(max(0.0, deadline - time.monotonic())) + if not completed: + with _lock: + completed = attempt.done.is_set() + if not completed: + attempt.timed_out = True + if not completed: + logger.warning("Provider flush exceeded the caller wait budget") + return False + if attempt.error is not None: + raise attempt.error + if attempt.result is False: + logger.warning("Provider reported an incomplete flush") + return False + return True 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 0a7a6e1ec..382ae4124 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 @@ -70,6 +70,7 @@ _ensure_extracted_context, xray_context_extractor, ) +from aws_durable_execution_sdk_python_otel._flush import flush_provider from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( DeterministicIdGenerator, derive_workflow_span_id, @@ -666,7 +667,7 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: if hasattr(self._provider, "force_flush"): try: - self._provider.force_flush() + flush_provider(self._provider.force_flush) except Exception: # noqa: BLE001 logger.exception("force_flush failed at invocation end") 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 f338f806b..5c019462c 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 @@ -40,6 +40,7 @@ _ensure_extracted_context, xray_context_extractor, ) +from aws_durable_execution_sdk_python_otel._flush import flush_provider from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( DeterministicIdGenerator, derive_workflow_span_id, @@ -768,7 +769,7 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: # Flush before Lambda freeze if hasattr(self._provider, "force_flush"): - self._provider.force_flush() + flush_provider(self._provider.force_flush) if unresolved_parents: # Keep the previous missing-parent failure observable, but finish # normal token/span cleanup and flushing before reporting it. diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_flush_lifecycle_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_flush_lifecycle_int.py new file mode 100644 index 000000000..6fa26b08a --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_flush_lifecycle_int.py @@ -0,0 +1,215 @@ +"""Only provider I/O may outlive the existing handler-worker End scope.""" + +import contextvars +import json +import threading +import time +from datetime import UTC, datetime +from unittest.mock import Mock, patch + +import pytest +from aws_durable_execution_sdk_python import durable_execution +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from aws_durable_execution_sdk_python.lambda_service import ( + ExecutionDetails, + Operation, + OperationStatus, + OperationType, +) +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, + _flush, +) + + +@pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_public_handler_finishes_while_one_provider_flush_is_blocked(view, monkeypatch): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + monkeypatch.setattr(_flush, "_FLUSH_TIMEOUT_SECONDS", 0.05) + ambient = contextvars.ContextVar("flush-owner", default="caller") + release = threading.Event() + exporting = threading.Event() + owners = [] + flush_calls = [] + exporter_contexts = [] + + class Exporter(InMemorySpanExporter): + def export(self, spans): + exporting.set() + exporter_contexts.append(ambient.get()) + assert release.wait(10) + return super().export(spans) + + class Provider(TracerProvider): + def force_flush(self, timeout_millis=30000): + flush_calls.append(threading.get_ident()) + return super().force_flush(timeout_millis) + + class Owner(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + self.token = ambient.set("handler") + self.thread = threading.get_ident() + + def on_invocation_end(self, info): + assert ambient.get() == "handler" + assert threading.get_ident() == self.thread + # A successful token reset proves ownership of the actual Context, + # rather than only comparing an OS thread ID that can be reused. + ambient.reset(self.token) + owners.append(self.thread) + + exporter = Exporter() + provider = Provider() + provider.add_span_processor( + BatchSpanProcessor(exporter, schedule_delay_millis=60000) + ) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + enrich_logger=False, + context_extractor=lambda _: None, + ) + ) + + def body(_event, _context): + assert ambient.get() == "handler" + return "done" + + handler = durable_execution(body, plugins=[Owner(), plugin]) + + def invoke(): + started = time.monotonic() + with DurableFunctionTestRunner(handler=handler, poll_interval=0.01) as runner: + result = runner.run(input="{}") + assert time.monotonic() - started < 2 + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == "done" + assert ambient.get() == "caller" + + try: + invoke() + assert exporting.wait(1) + invoke() + assert len(owners) == 2 + assert len(flush_calls) == 1 + assert all(worker != flush_calls[0] for worker in owners) + assert exporter_contexts == ["caller"] + finally: + release.set() + deadline = time.monotonic() + 2 + while time.monotonic() < deadline: + with _flush._lock: + idle = _flush._active is None + if idle: + break + time.sleep(0.001) + assert idle + try: + # Actual completion restores admission; no accumulated late error or + # stale invocation Context is delivered to this next request. + monkeypatch.setattr(_flush, "_FLUSH_TIMEOUT_SECONDS", 1.0) + invoke() + assert len(flush_calls) == 2 + assert len(owners) == 3 + assert all(value == "caller" for value in exporter_contexts) + finally: + provider.shutdown() + + +@pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) +@pytest.mark.parametrize("fatal", [False, True]) +def test_public_fast_flush_failure_preserves_owner_cleanup_and_reuse( + view, fatal, monkeypatch +): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + monkeypatch.setattr(_flush, "_FLUSH_TIMEOUT_SECONDS", 1.0) + binding = contextvars.ContextVar("flush-failure-owner", default="caller") + ended = [] + error = SystemExit("provider fatal") if fatal else RuntimeError("provider ordinary") + + class Owner(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + self.token = binding.set("handler") + self.thread = threading.get_ident() + + def on_invocation_end(self, info): + assert threading.get_ident() == self.thread + binding.reset(self.token) + ended.append(self.thread) + + class Provider(TracerProvider): + failure: BaseException | None = error + + def force_flush(self, timeout_millis=30000): + if self.failure is not None: + raise self.failure + return super().force_flush(timeout_millis) + + provider = Provider() + provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter())) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + enrich_logger=False, + context_extractor=lambda _: None, + ) + ) + + def body(_event, _context): + assert binding.get() == "handler" + return "done" + + handler = durable_execution(body, plugins=[Owner(), plugin]) + operation = Operation( + operation_id="flush-execution", + operation_type=OperationType.EXECUTION, + status=OperationStatus.STARTED, + start_timestamp=datetime.now(UTC), + execution_details=ExecutionDetails(input_payload="{}"), + ) + event = { + "DurableExecutionArn": "test-arn/flush-execution", + "CheckpointToken": "token", + "InitialExecutionState": { + "Operations": [operation.to_json_dict()], + "NextMarker": "", + }, + "LocalRunner": True, + } + context = Mock() + context.aws_request_id = "request" + context.invoked_function_arn = "test-arn" + context.client_context = None + context.identity = None + context.tenant_id = None + context._epoch_deadline_time_in_ms = 0 + try: + with patch("aws_durable_execution_sdk_python.execution.LambdaClient"): + if fatal: + with pytest.raises(SystemExit) as caught: + handler(event, context) + assert caught.value is error + else: + assert handler(event, context) == { + "Status": "SUCCEEDED", + "Result": '"done"', + } + assert len(ended) == 1 + assert binding.get() == "caller" + provider.failure = None + assert handler(event, context) == { + "Status": "SUCCEEDED", + "Result": '"done"', + } + assert len(ended) == 2 + assert binding.get() == "caller" + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_flush.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_flush.py new file mode 100644 index 000000000..2619cfc8d --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_flush.py @@ -0,0 +1,132 @@ +"""Bound provider I/O without spawning replacement workers for a stuck call.""" + +import contextvars +import threading +import time + +import pytest + +from aws_durable_execution_sdk_python_otel import _flush + + +def _wait_idle() -> None: + deadline = time.monotonic() + 2 + while time.monotonic() < deadline: + with _flush._lock: + if _flush._active is None: + return + time.sleep(0.001) + raise AssertionError("Provider worker did not actually complete") + + +@pytest.fixture(autouse=True) +def idle_flush(monkeypatch): + _wait_idle() + monkeypatch.setattr(_flush, "_FLUSH_TIMEOUT_SECONDS", 1.0) + yield + _wait_idle() + + +def test_timeout_retains_one_worker_and_rejects_new_work(monkeypatch): + monkeypatch.setattr(_flush, "_FLUSH_TIMEOUT_SECONDS", 0.02) + started = threading.Event() + release = threading.Event() + calls = [] + + def blocked(): + calls.append(threading.get_ident()) + started.set() + assert release.wait(2) + return True + + try: + began = time.monotonic() + assert _flush.flush_provider(blocked) is False + assert time.monotonic() - began < 1 + assert started.wait(1) + for _ in range(5): + assert _flush.flush_provider(lambda: calls.append("unexpected")) is False + assert len(calls) == 1 + workers = [t for t in threading.enumerate() if t.name == "durable-otel-flush"] + assert len(workers) == 1 + assert workers[0].daemon + finally: + release.set() + _wait_idle() + monkeypatch.setattr(_flush, "_FLUSH_TIMEOUT_SECONDS", 1.0) + assert _flush.flush_provider(lambda: True) is True + + +@pytest.mark.parametrize("result", [True, False, None]) +def test_completed_provider_result(result): + assert _flush.flush_provider(lambda: result) is (result is not False) + + +@pytest.mark.parametrize( + "error", [RuntimeError("ordinary"), SystemExit("fatal"), KeyboardInterrupt()] +) +def test_fast_error_preserves_identity_on_the_caller(error): + owner = threading.get_ident() + worker = [] + + def fail(): + worker.append(threading.get_ident()) + raise error + + with pytest.raises(type(error)) as caught: + _flush.flush_provider(fail) + assert caught.value is error + assert worker[0] != owner + assert threading.get_ident() == owner + + +@pytest.mark.parametrize( + "error", [RuntimeError("late ordinary"), SystemExit("late fatal")] +) +def test_late_error_is_reported_without_poisoning_next_call(error, monkeypatch, caplog): + monkeypatch.setattr(_flush, "_FLUSH_TIMEOUT_SECONDS", 0.02) + release = threading.Event() + + def fail_late(): + assert release.wait(2) + raise error + + try: + assert _flush.flush_provider(fail_late) is False + finally: + release.set() + _wait_idle() + messages = [ + r for r in caplog.records if "after the caller stopped waiting" in r.message + ] + assert len(messages) == 1 + assert messages[0].exc_info[1] is error + monkeypatch.setattr(_flush, "_FLUSH_TIMEOUT_SECONDS", 1.0) + assert _flush.flush_provider(lambda: True) is True + + +def test_start_failure_releases_capacity(monkeypatch): + def fail_start(_thread): + raise RuntimeError("start failed") + + with monkeypatch.context() as patch: + patch.setattr(threading.Thread, "start", fail_start) + with pytest.raises(RuntimeError, match="start failed"): + _flush.flush_provider(lambda: True) + assert _flush.flush_provider(lambda: True) is True + + +def test_only_provider_io_gets_post_cleanup_context_snapshot(): + value = contextvars.ContextVar("flush-test", default="unset") + token = value.set("owner-after-cleanup") + try: + + def flush(): + assert value.get() == "owner-after-cleanup" + value.set("provider-only") + return True + + assert _flush.flush_provider(flush) + assert value.get() == "owner-after-cleanup" + finally: + value.reset(token) From b179bf185b5dd2bba39e2a963cf9ab898bf1d85c Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 14 Oct 2026 18:19:28 +0000 Subject: [PATCH 43/54] fix(otel): wait for flush admission within budget --- .../README.md | 8 +- .../_flush.py | 21 ++-- .../tests/test_flush.py | 99 +++++++++++++++++++ 3 files changed, 119 insertions(+), 9 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 1d02f9b1b..315758a9f 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -425,9 +425,11 @@ cleanup stay on their original worker and Context. The I/O receives a snapshot of that post-cleanup Context. There is at most one outstanding provider flush in the process and no pending -queue. A timeout does not cancel the exporter or free that slot. Later requests -skip flushing with a warning until the earlier call actually finishes, including -requests using another provider. A failed or skipped flush is not an export +I/O queue. Concurrent requests, including requests using another provider, wait +for admission within the same absolute deadline used for completion. A timeout +does not cancel the exporter or free its slot. A request whose budget expires +before admission returns with a warning without starting more work. +A failed or skipped flush is not an export acknowledgment, and late exports may continue. Applications still own exporter transport timeouts and provider shutdown. This bound does not cover synchronous span-processor `on_end` work that runs before flushing. diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_flush.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_flush.py index e288528ff..26c1b46e8 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_flush.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_flush.py @@ -15,7 +15,7 @@ # Match the OTel provider's existing default budget. Some processors ignore # that cooperative budget, so the SDK bounds its own wait independently. _FLUSH_TIMEOUT_SECONDS = 30.0 -_lock = threading.Lock() +_lock = threading.Condition() @dataclass @@ -42,6 +42,7 @@ def _run(flush: Callable[[], object], attempt: _Attempt) -> None: late = attempt.timed_out if not late and _active is attempt: _active = None + _lock.notify_all() attempt.done.set() if late and attempt.error is not None: logger.error( @@ -58,6 +59,7 @@ def _run(flush: Callable[[], object], attempt: _Attempt) -> None: with _lock: if _active is attempt: _active = None + _lock.notify_all() def flush_provider(flush: Callable[[], object]) -> bool: @@ -65,20 +67,26 @@ def flush_provider(flush: Callable[[], object]) -> bool: Start/End and ContextVar-token cleanup stay on their existing owner. Only the supplied provider method runs elsewhere, using the post-cleanup Context - snapshot. Busy/expired attempts return False; they are never cancelled. + snapshot. Admission and completion share one absolute deadline. Expired + requests return False without queueing more work; active I/O is never cancelled. This does not bound synchronous span processors' earlier ``on_end`` calls, application-owned provider shutdown, or exporter transport I/O itself. """ global _active # noqa: PLW0603 deadline = time.monotonic() + _FLUSH_TIMEOUT_SECONDS with _lock: - if _active is not None: - attempt = None - else: + while _active is not None: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + _lock.wait(remaining) + if _active is None and time.monotonic() < deadline: attempt = _Attempt() _active = attempt + else: + attempt = None if attempt is None: - logger.warning("Provider flush is still running; skipping another flush") + logger.warning("Provider flush admission exceeded the caller wait budget") return False owner_context = contextvars.copy_context() @@ -93,6 +101,7 @@ def flush_provider(flush: Callable[[], object]) -> bool: with _lock: if _active is attempt: _active = None + _lock.notify_all() raise completed = attempt.done.wait(max(0.0, deadline - time.monotonic())) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_flush.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_flush.py index 2619cfc8d..ec67d7725 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_flush.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_flush.py @@ -3,8 +3,12 @@ import contextvars import threading import time +from concurrent.futures import ThreadPoolExecutor import pytest +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from aws_durable_execution_sdk_python_otel import _flush @@ -57,6 +61,101 @@ def blocked(): assert _flush.flush_provider(lambda: True) is True +def test_healthy_distinct_providers_wait_for_admission(): + exporting = threading.Event() + release = threading.Event() + second_called = threading.Event() + + class GatedExporter(InMemorySpanExporter): + def export(self, spans): + exporting.set() + assert release.wait(3) + return super().export(spans) + + first_exporter = GatedExporter() + second_exporter = InMemorySpanExporter() + first_provider = TracerProvider() + second_provider = TracerProvider() + first_provider.add_span_processor( + BatchSpanProcessor(first_exporter, schedule_delay_millis=60000) + ) + second_provider.add_span_processor( + BatchSpanProcessor(second_exporter, schedule_delay_millis=60000) + ) + with first_provider.get_tracer("first").start_as_current_span("first"): + pass + with second_provider.get_tracer("second").start_as_current_span("second"): + pass + + def second_flush(): + second_called.set() + return _flush.flush_provider(second_provider.force_flush) + + try: + with ThreadPoolExecutor(max_workers=2) as callers: + first = callers.submit(_flush.flush_provider, first_provider.force_flush) + try: + assert exporting.wait(1) + second = callers.submit(second_flush) + assert second_called.wait(1) + time.sleep(0.03) + assert not second.done(), ( + "A healthy caller was skipped instead of waiting" + ) + finally: + release.set() + assert first.result(timeout=2) is True + assert second.result(timeout=2) is True + assert [span.name for span in first_exporter.get_finished_spans()] == ["first"] + assert [span.name for span in second_exporter.get_finished_spans()] == [ + "second" + ] + finally: + release.set() + first_provider.shutdown() + second_provider.shutdown() + + +def test_admission_and_execution_share_the_same_deadline(): + first_entered = threading.Event() + first_release = threading.Event() + second_called = threading.Event() + second_entered = threading.Event() + second_release = threading.Event() + + def first_io(): + first_entered.set() + assert first_release.wait(3) + return True + + def second_io(): + second_entered.set() + assert second_release.wait(3) + return True + + def second_call(): + second_called.set() + return _flush.flush_provider(second_io) + + with ThreadPoolExecutor(max_workers=2) as callers: + first = callers.submit(_flush.flush_provider, first_io) + try: + assert first_entered.wait(1) + second = callers.submit(second_call) + assert second_called.wait(1) + # Spend most of the existing one-second budget on admission. + time.sleep(0.6) + first_release.set() + assert first.result(timeout=1) is True + assert second_entered.wait(1) + # A reset one-second I/O budget would still be blocked here. + assert second.result(timeout=0.65) is False + finally: + first_release.set() + second_release.set() + _wait_idle() + + @pytest.mark.parametrize("result", [True, False, None]) def test_completed_provider_result(result): assert _flush.flush_provider(lambda: result) is (result is not False) From 557de42600c7fae0cacec1ba26d747785d734507 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 14 Oct 2026 19:15:48 +0000 Subject: [PATCH 44/54] fix(otel): isolate shared invocation state --- .../README.md | 15 ++ .../_invocation_state.py | 103 +++++++++ .../execution_plugin.py | 198 ++++++++++++++--- .../invocation_plugin.py | 208 +++++++++++++++--- .../e2e/test_shared_invocation_state_int.py | 113 ++++++++++ .../tests/test_execution_plugin.py | 14 +- .../tests/test_invocation_owned_state.py | 118 ++++++++++ .../tests/test_invocation_plugin.py | 14 +- .../execution.py | 17 +- .../plugin.py | 4 +- 10 files changed, 724 insertions(+), 80 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_shared_invocation_state_int.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_owned_state.py diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 315758a9f..f083086ba 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -441,6 +441,21 @@ raised retroactively into the completed invocation. ### Invocation context isolation +Tracing state belongs to each invocation's Context, including its trace IDs, +span registry, deferred completions, timestamps and scope tokens. A shared plugin +instance can therefore serve overlapping invocations without one invocation's +cleanup clearing another's trace. Provider configuration and registration remain +instance resources. SDK-owned worker tasks capture the selected invocation +Context; Start/End hooks and token reset retain their existing owner and order. + +Completion metadata can be redelivered after a failed invocation or checkpoint +transport attempt that has no successful response acknowledgment. There is no +exporter acknowledgment in durable state, so this does not promise exactly-once +telemetry across failed/redelivered invocations. Successful replay preserves +checkpointed work; this limit is not a claim that already-checkpointed step +bodies run again. The plugin does not use a global exported-operation cache that +could discard a first delivery after a failed attempt. + With core 2.1+, invocation hooks and the handler run on one worker in an invocation-local Context initialized from the host's bindings. Successful Start bindings are visible to later hooks, the handler and resource cleanup. The host's diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py new file mode 100644 index 000000000..10dd88482 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py @@ -0,0 +1,103 @@ +"""Invocation-owned tracing data, propagated by the SDK's captured Contexts.""" + +from __future__ import annotations + +from collections.abc import Callable +from contextvars import ContextVar, Token +from dataclasses import dataclass, field +import datetime +import threading +from typing import Any, Generic, TypeVar + +from aws_durable_execution_sdk_python.plugin import OperationEndInfo +from opentelemetry.sdk.trace.sampling import Sampler +from opentelemetry.trace import Span, Tracer, TracerProvider + +from aws_durable_execution_sdk_python_otel.context_extractors import ExtractedContext +from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + DeterministicIdGenerator, +) +from aws_durable_execution_sdk_python_otel.durable_sampling import DurableSamplingIntent +from aws_durable_execution_sdk_python_otel.execution_trace_context import ( + ExecutionTraceContext, +) + + +_Key = TypeVar("_Key") + + +@dataclass +class _InvocationState(Generic[_Key]): + provider: TracerProvider + tracer: Tracer + id_generator: DeterministicIdGenerator + sampling_delegate: Sampler | None = None + execution_arn: str = "" + execution_trace_id: int | None = None + execution_start_time: datetime.datetime | None = None + extracted_context: ExtractedContext | None = None + execution_trace_context: ExecutionTraceContext | None = None + sampling_intent: DurableSamplingIntent | None = None + workflow_span: Span | None = None + operation_spans: dict[_Key, Span] = field(default_factory=dict) + context_tokens: dict[str, tuple[int, object]] = field(default_factory=dict) + tracing_enabled: bool = False + owner_token: Token[Any] | None = None + owner_thread: threading.Thread | None = None + + +@dataclass +class ExecutionState(_InvocationState[str]): + invocation_span: Span | None = None + checkpointed_context_ids: set[str] = field(default_factory=set) + ended_operation_ids: set[str] = field(default_factory=set) + + +@dataclass +class InvocationState(_InvocationState[str | None]): + span_time_floor_ns: int | None = None + pending_operation_ends: dict[str, list[OperationEndInfo]] = field( + default_factory=dict + ) + context_operation_replays: dict[str, bool] = field(default_factory=dict) + incomplete_attempt_span_keys: set[str] = field(default_factory=set) + + +_State = TypeVar("_State", bound=_InvocationState[Any]) + + +def begin_state( + variable: ContextVar[_State], + state: _State, + reset_previous: Callable[[], None], +) -> None: + """Start fresh data without clearing another invocation inherited in a copy. + + A direct warm-start caller can omit its previous End; retain that cleanup + behavior only when its token proves this is the actual owning Context. + Thread identity alone cannot distinguish a Context from a copy of it. + """ + previous = variable.get() + token = previous.owner_token + if token is not None and previous.owner_thread is threading.current_thread(): + try: + variable.reset(token) + except (ValueError, RuntimeError): + # An inherited Context (or an already-ended inherited token) does + # not own the previous invocation and must not mutate its data. + pass + else: + previous.owner_token = variable.set(previous) + reset_previous() + state.owner_thread = threading.current_thread() + state.owner_token = variable.set(state) + + +def finish_state(variable: ContextVar[_State]) -> None: + """Restore the binding on the Context that ran Start, after normal cleanup.""" + state = variable.get() + token = state.owner_token + if token is not None: + state.owner_token = None + state.owner_thread = None + variable.reset(token) 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 382ae4124..fb44a1f45 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,6 +31,7 @@ from __future__ import annotations +import contextvars import datetime import logging import threading @@ -64,6 +65,11 @@ TraceState, ) +from aws_durable_execution_sdk_python_otel._invocation_state import ( + ExecutionState, + begin_state, + finish_state, +) from aws_durable_execution_sdk_python_otel.context_extractors import ( ContextExtractor, ExtractedContext, @@ -141,43 +147,159 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._workflow_span_name = self._config.workflow_span_name result = create_tracer_provider(self._config) - self._provider = result.tracer_provider self._uses_global_provider = result.uses_global_provider - - self._tracer: Tracer = self._provider.get_tracer(self._config.instrument_name) - self._id_generator = DeterministicIdGenerator() - self._sampling_delegate: Sampler | None = None - self._bind_sdk_tracer() - - # Per-invocation state. - self._execution_arn = "" - self._execution_trace_id: int | None = None - self._execution_start_time: datetime.datetime | None = None - self._extracted_context: ExtractedContext | None = None - self._execution_trace_context: ExecutionTraceContext | None = None - self._sampling_intent: DurableSamplingIntent | None = None - self._workflow_span: Span | None = None - self._invocation_span: Span | None = None - self._operation_spans: dict[str, Span] = {} - # CONTEXT operations that emitted a durable START hook this invocation. - # A context absent from this set is checkpointless (for example, a FLAT - # map/parallel branch) and needs a fresh per-invocation span identity. - self._checkpointed_context_ids: set[str] = set() - # Operations whose span was already exported this invocation, so a - # repeated on_operation_end does not export it twice. - self._ended_operation_ids: set[str] = set() - # Tokens returned by context.attach(), keyed by the span registry key, - # paired with the thread that attached them. Every attach the plugin - # 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._template = ExecutionState( + provider=result.tracer_provider, + tracer=result.tracer_provider.get_tracer(self._config.instrument_name), + id_generator=DeterministicIdGenerator(), + ) + self._state: contextvars.ContextVar[ExecutionState] = contextvars.ContextVar( + "durable_execution_otel_invocation", default=self._template + ) self._lock = threading.RLock() - self._tracing_enabled = False + self._bind_sdk_tracer() self._registration_accepted = False if self._config.enrich_logger: install_log_filter(self) + # Private accessors route existing helpers through the current invocation. + @property + def _provider(self) -> trace.TracerProvider: + return self._state.get().provider + + @_provider.setter + def _provider(self, value: trace.TracerProvider) -> None: + self._state.get().provider = value + + @property + def _tracer(self) -> Tracer: + return self._state.get().tracer + + @_tracer.setter + def _tracer(self, value: Tracer) -> None: + self._state.get().tracer = value + + @property + def _id_generator(self) -> DeterministicIdGenerator: + return self._state.get().id_generator + + @_id_generator.setter + def _id_generator(self, value: DeterministicIdGenerator) -> None: + self._state.get().id_generator = value + + @property + def _sampling_delegate(self) -> Sampler | None: + return self._state.get().sampling_delegate + + @_sampling_delegate.setter + def _sampling_delegate(self, value: Sampler | None) -> None: + self._state.get().sampling_delegate = value + + @property + def _execution_arn(self) -> str: + return self._state.get().execution_arn + + @_execution_arn.setter + def _execution_arn(self, value: str) -> None: + self._state.get().execution_arn = value + + @property + def _execution_trace_id(self) -> int | None: + return self._state.get().execution_trace_id + + @_execution_trace_id.setter + def _execution_trace_id(self, value: int | None) -> None: + self._state.get().execution_trace_id = value + + @property + def _execution_start_time(self) -> datetime.datetime | None: + return self._state.get().execution_start_time + + @_execution_start_time.setter + def _execution_start_time(self, value: datetime.datetime | None) -> None: + self._state.get().execution_start_time = value + + @property + def _extracted_context(self) -> ExtractedContext | None: + return self._state.get().extracted_context + + @_extracted_context.setter + def _extracted_context(self, value: ExtractedContext | None) -> None: + self._state.get().extracted_context = value + + @property + def _execution_trace_context(self) -> ExecutionTraceContext | None: + return self._state.get().execution_trace_context + + @_execution_trace_context.setter + def _execution_trace_context(self, value: ExecutionTraceContext | None) -> None: + self._state.get().execution_trace_context = value + + @property + def _sampling_intent(self) -> DurableSamplingIntent | None: + return self._state.get().sampling_intent + + @_sampling_intent.setter + def _sampling_intent(self, value: DurableSamplingIntent | None) -> None: + self._state.get().sampling_intent = value + + @property + def _workflow_span(self) -> Span | None: + return self._state.get().workflow_span + + @_workflow_span.setter + def _workflow_span(self, value: Span | None) -> None: + self._state.get().workflow_span = value + + @property + def _context_tokens(self) -> dict[str, tuple[int, object]]: + return self._state.get().context_tokens + + @_context_tokens.setter + def _context_tokens(self, value: dict[str, tuple[int, object]]) -> None: + self._state.get().context_tokens = value + + @property + def _tracing_enabled(self) -> bool: + return self._state.get().tracing_enabled + + @_tracing_enabled.setter + def _tracing_enabled(self, value: bool) -> None: + self._state.get().tracing_enabled = value + + @property + def _invocation_span(self) -> Span | None: + return self._state.get().invocation_span + + @_invocation_span.setter + def _invocation_span(self, value: Span | None) -> None: + self._state.get().invocation_span = value + + @property + def _operation_spans(self) -> dict[str, Span]: + return self._state.get().operation_spans + + @_operation_spans.setter + def _operation_spans(self, value: dict[str, Span]) -> None: + self._state.get().operation_spans = value + + @property + def _checkpointed_context_ids(self) -> set[str]: + return self._state.get().checkpointed_context_ids + + @_checkpointed_context_ids.setter + def _checkpointed_context_ids(self, value: set[str]) -> None: + self._state.get().checkpointed_context_ids = value + + @property + def _ended_operation_ids(self) -> set[str]: + return self._state.get().ended_operation_ids + + @_ended_operation_ids.setter + def _ended_operation_ids(self, value: set[str]) -> None: + self._state.get().ended_operation_ids = value + def on_registration_result(self, registered: bool) -> None: """Preserve accepted resources, releasing only a discarded constructor.""" if registered: @@ -434,7 +556,15 @@ def _with_sampling(self, parent_context: Context) -> Context: def on_invocation_start(self, info: InvocationStartInfo) -> None: logger.debug("Durable invocation started: %s", info) self._registration_accepted = True - self._reset_state() + begin_state( + self._state, + ExecutionState( + provider=self._template.provider, + tracer=self._template.tracer, + id_generator=self._template.id_generator, + ), + self._reset_state, + ) if info.execution_start_time is None: logger.warning( "ExecutionOtelPlugin requires InvocationStartInfo.execution_start_time " @@ -663,11 +793,12 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: if info.status in _TERMINAL_INVOCATION_STATUSES: self._export_workflow_span(info) + provider = self._provider self._reset_state() - if hasattr(self._provider, "force_flush"): + if hasattr(provider, "force_flush"): try: - flush_provider(self._provider.force_flush) + flush_provider(provider.force_flush) except Exception: # noqa: BLE001 logger.exception("force_flush failed at invocation end") @@ -686,6 +817,7 @@ def _reset_state(self) -> None: self._checkpointed_context_ids = set() self._ended_operation_ids = set() self._tracing_enabled = False + finish_state(self._state) # ------------------------------------------------------------------ # Operation lifecycle 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 5c019462c..a3c6291d6 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,6 +2,7 @@ from __future__ import annotations +import contextvars import datetime import logging import threading @@ -34,6 +35,11 @@ TraceState, ) +from aws_durable_execution_sdk_python_otel._invocation_state import ( + InvocationState, + begin_state, + finish_state, +) from aws_durable_execution_sdk_python_otel.context_extractors import ( ContextExtractor, ExtractedContext, @@ -141,45 +147,167 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._enrich_logger = self._config.enrich_logger result = create_tracer_provider(self._config) - self._provider = result.tracer_provider self._uses_global_provider = result.uses_global_provider - self._tracer: Tracer = self._provider.get_tracer(self._config.instrument_name) - self._id_generator = DeterministicIdGenerator() - self._sampling_delegate: Sampler | None = None - self._bind_sdk_tracer() - - # per invocation status: - self._execution_arn = "" - self._execution_trace_id: int | None = None - self._execution_start_time: datetime.datetime | None = None - self._extracted_context: ExtractedContext | None = None - self._execution_trace_context: ExecutionTraceContext | None = None - self._sampling_intent: DurableSamplingIntent | None = None - self._workflow_span: Span | None = None - self._span_time_floor_ns: int | None = None - # Maps operation ID (None for root) to the active span. - self._operation_spans: dict[str | None, Span] = {} - # A sibling checkpoint can report completion before replay enters its - # parent context. Retain the real event until that parent span exists. - self._pending_operation_ends: dict[str, list[OperationEndInfo]] = {} - # Replay state supplied by CONTEXT operation START hooks. Missing - # entries identify checkpointless contexts such as FLAT branches. - self._context_operation_replays: dict[str, bool] = {} - # STEP attempt spans whose user function explicitly ended INCOMPLETE. - # Only these spans are marked truncated during invocation cleanup. - self._incomplete_attempt_span_keys: set[str] = set() - # Tokens returned by context.attach(), keyed by the span registry key, - # paired with the thread that attached them. Every attach the plugin - # 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._template = InvocationState( + provider=result.tracer_provider, + tracer=result.tracer_provider.get_tracer(self._config.instrument_name), + id_generator=DeterministicIdGenerator(), + ) + self._state: contextvars.ContextVar[InvocationState] = contextvars.ContextVar( + "durable_invocation_otel_invocation", default=self._template + ) self._operation_spans_lock = threading.RLock() - self._tracing_enabled = False + self._bind_sdk_tracer() self._registration_accepted = False if self._enrich_logger: install_log_filter(self) + # Private accessors route existing helpers through the current invocation. + @property + def _provider(self) -> trace.TracerProvider: + return self._state.get().provider + + @_provider.setter + def _provider(self, value: trace.TracerProvider) -> None: + self._state.get().provider = value + + @property + def _tracer(self) -> Tracer: + return self._state.get().tracer + + @_tracer.setter + def _tracer(self, value: Tracer) -> None: + self._state.get().tracer = value + + @property + def _id_generator(self) -> DeterministicIdGenerator: + return self._state.get().id_generator + + @_id_generator.setter + def _id_generator(self, value: DeterministicIdGenerator) -> None: + self._state.get().id_generator = value + + @property + def _sampling_delegate(self) -> Sampler | None: + return self._state.get().sampling_delegate + + @_sampling_delegate.setter + def _sampling_delegate(self, value: Sampler | None) -> None: + self._state.get().sampling_delegate = value + + @property + def _execution_arn(self) -> str: + return self._state.get().execution_arn + + @_execution_arn.setter + def _execution_arn(self, value: str) -> None: + self._state.get().execution_arn = value + + @property + def _execution_trace_id(self) -> int | None: + return self._state.get().execution_trace_id + + @_execution_trace_id.setter + def _execution_trace_id(self, value: int | None) -> None: + self._state.get().execution_trace_id = value + + @property + def _execution_start_time(self) -> datetime.datetime | None: + return self._state.get().execution_start_time + + @_execution_start_time.setter + def _execution_start_time(self, value: datetime.datetime | None) -> None: + self._state.get().execution_start_time = value + + @property + def _extracted_context(self) -> ExtractedContext | None: + return self._state.get().extracted_context + + @_extracted_context.setter + def _extracted_context(self, value: ExtractedContext | None) -> None: + self._state.get().extracted_context = value + + @property + def _execution_trace_context(self) -> ExecutionTraceContext | None: + return self._state.get().execution_trace_context + + @_execution_trace_context.setter + def _execution_trace_context(self, value: ExecutionTraceContext | None) -> None: + self._state.get().execution_trace_context = value + + @property + def _sampling_intent(self) -> DurableSamplingIntent | None: + return self._state.get().sampling_intent + + @_sampling_intent.setter + def _sampling_intent(self, value: DurableSamplingIntent | None) -> None: + self._state.get().sampling_intent = value + + @property + def _workflow_span(self) -> Span | None: + return self._state.get().workflow_span + + @_workflow_span.setter + def _workflow_span(self, value: Span | None) -> None: + self._state.get().workflow_span = value + + @property + def _context_tokens(self) -> dict[str, tuple[int, object]]: + return self._state.get().context_tokens + + @_context_tokens.setter + def _context_tokens(self, value: dict[str, tuple[int, object]]) -> None: + self._state.get().context_tokens = value + + @property + def _tracing_enabled(self) -> bool: + return self._state.get().tracing_enabled + + @_tracing_enabled.setter + def _tracing_enabled(self, value: bool) -> None: + self._state.get().tracing_enabled = value + + @property + def _operation_spans(self) -> dict[str | None, Span]: + return self._state.get().operation_spans + + @_operation_spans.setter + def _operation_spans(self, value: dict[str | None, Span]) -> None: + self._state.get().operation_spans = value + + @property + def _span_time_floor_ns(self) -> int | None: + return self._state.get().span_time_floor_ns + + @_span_time_floor_ns.setter + def _span_time_floor_ns(self, value: int | None) -> None: + self._state.get().span_time_floor_ns = value + + @property + def _pending_operation_ends(self) -> dict[str, list[OperationEndInfo]]: + return self._state.get().pending_operation_ends + + @_pending_operation_ends.setter + def _pending_operation_ends(self, value: dict[str, list[OperationEndInfo]]) -> None: + self._state.get().pending_operation_ends = value + + @property + def _context_operation_replays(self) -> dict[str, bool]: + return self._state.get().context_operation_replays + + @_context_operation_replays.setter + def _context_operation_replays(self, value: dict[str, bool]) -> None: + self._state.get().context_operation_replays = value + + @property + def _incomplete_attempt_span_keys(self) -> set[str]: + return self._state.get().incomplete_attempt_span_keys + + @_incomplete_attempt_span_keys.setter + def _incomplete_attempt_span_keys(self, value: set[str]) -> None: + self._state.get().incomplete_attempt_span_keys = value + def on_registration_result(self, registered: bool) -> None: """Preserve accepted resources, releasing only a discarded constructor.""" if registered: @@ -547,7 +675,15 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: """Called at the start of each invocation. Creates the invocation span.""" logger.debug("Durable invocation started: %s", info) self._registration_accepted = True - self._reset_state() + begin_state( + self._state, + InvocationState( + provider=self._template.provider, + tracer=self._template.tracer, + id_generator=self._template.id_generator, + ), + self._reset_state, + ) if info.execution_start_time is None: logger.warning( "InvocationOtelPlugin requires InvocationStartInfo.execution_start_time " @@ -765,11 +901,12 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: if info.status in _TERMINAL_INVOCATION_STATUSES: self._export_workflow_span(info) + provider = self._provider self._reset_state() # Flush before Lambda freeze - if hasattr(self._provider, "force_flush"): - flush_provider(self._provider.force_flush) + if hasattr(provider, "force_flush"): + flush_provider(provider.force_flush) if unresolved_parents: # Keep the previous missing-parent failure observable, but finish # normal token/span cleanup and flushing before reporting it. @@ -794,6 +931,7 @@ def _reset_state(self) -> None: self._context_operation_replays = {} self._incomplete_attempt_span_keys = set() self._tracing_enabled = False + finish_state(self._state) def on_operation_start(self, info: OperationStartInfo) -> None: """Called when an operation begins. Creates a span for the operation.""" diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_shared_invocation_state_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_shared_invocation_state_int.py new file mode 100644 index 000000000..1a0bbd625 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_shared_invocation_state_int.py @@ -0,0 +1,113 @@ +"""A shared plugin routes worker/checkpoint events to their real invocation.""" + +import json +import logging +import threading + +import pytest +from aws_durable_execution_sdk_python import durable_execution, durable_step +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +@pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_same_handler_shared_plugin_keeps_two_execution_traces( + view, monkeypatch, caplog +): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + ready = threading.Event() + release = threading.Event() + side_effects = [] + + @durable_step + def save(_context, label, stage): + side_effects.append((label, stage)) + return label + stage + + def body(event, context): + label = event["label"] + first = context.step(save(label, "first"), name="first") + if label == "a": + ready.set() + assert release.wait(5) + second = context.step(save(label, "second"), name="second") + return [first, second] + + handler = durable_execution(body, plugins=[plugin]) + try: + with ( + DurableFunctionTestRunner(handler=handler, poll_interval=0.01) as ra, + DurableFunctionTestRunner(handler=handler, poll_interval=0.01) as rb, + ): + a = ra.run_async(input='{"label":"a"}') + assert ready.wait(5) + try: + b = rb.run_async(input='{"label":"b"}') + result_b = rb.wait_for_result(b, timeout=5) + finally: + release.set() + result_a = ra.wait_for_result(a, timeout=5) + assert result_a.status.value == result_b.status.value == "SUCCEEDED" + assert json.loads(result_a.result) == ["afirst", "asecond"] + assert json.loads(result_b.result) == ["bfirst", "bsecond"] + assert sorted(side_effects) == [ + ("a", "first"), + ("a", "second"), + ("b", "first"), + ("b", "second"), + ] + spans = exporter.get_finished_spans() + workflows = [s for s in spans if s.name == "Workflow"] + assert len(workflows) == 2 + assert {s.attributes["durable.execution.arn"] for s in workflows} == {a, b} + trace_for = { + s.attributes["durable.execution.arn"]: s.context.trace_id for s in workflows + } + assert len(set(trace_for.values())) == 2 + for arn in (a, b): + own = [s for s in spans if s.attributes.get("durable.execution.arn") == arn] + assert all(s.context.trace_id == trace_for[arn] for s in own) + steps = [ + s for s in own if s.attributes.get("durable.operation.type") == "STEP" + ] + assert len(steps) == 4 + assert len({s.context.span_id for s in steps}) == 4 + terminals = [ + s + for s in steps + if s.attributes.get("durable.operation.status") == "SUCCEEDED" + ] + assert len(terminals) == 2 + assert {s.attributes["durable.operation.name"] for s in terminals} == { + "first", + "second", + } + assert all(s.parent is not None for s in terminals) + own_ids = {s.context.span_id for s in own} + assert all(s.parent.span_id in own_ids for s in steps) + assert not [ + r + for r in caplog.records + if r.levelno >= logging.ERROR and ("plugin" in r.name or "otel" in r.name) + ] + finally: + release.set() + provider.shutdown() 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 9ef59c1d4..d24232741 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 @@ -4,6 +4,7 @@ import threading from concurrent.futures import ThreadPoolExecutor +from contextvars import copy_context from datetime import UTC, datetime, timedelta from types import SimpleNamespace @@ -1452,7 +1453,9 @@ def test_detach_ignores_token_attached_on_another_thread(): with ThreadPoolExecutor(max_workers=1) as executor: executor.submit( - plugin.on_user_function_start, _step_start_info("step-1") + copy_context().run, + plugin.on_user_function_start, + _step_start_info("step-1"), ).result() plugin._detach_context("step-1:attempt:1") @@ -1545,20 +1548,23 @@ def test_suspension_releases_the_scope_on_the_originating_worker(): with ThreadPoolExecutor(max_workers=1) as worker: def suspend_on_worker() -> tuple[int, bool]: + worker_context = otel_context.get_current() plugin.on_user_function_start(_step_start_info("step-1")) attached_span_id = trace.get_current_span().get_span_context().span_id plugin.on_user_function_end(_step_incomplete_info("step-1")) return ( attached_span_id, - trace.get_current_span().get_span_context().is_valid, + otel_context.get_current() == worker_context, ) - attached_span_id, span_still_current = worker.submit(suspend_on_worker).result() + attached_span_id, context_restored = worker.submit( + copy_context().run, suspend_on_worker + ).result() suspended_span = plugin._get_span(span_key) # The scope was released on the worker, and its span is left open. assert attached_span_id != 0 - assert span_still_current is False + assert context_restored is True assert span_key not in plugin._context_tokens assert suspended_span is not None assert {s.name for s in exporter.get_finished_spans()} == { diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_owned_state.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_owned_state.py new file mode 100644 index 000000000..e56c52fea --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_owned_state.py @@ -0,0 +1,118 @@ +"""Invocation-owned state must not be reset through an inherited Context.""" + +import contextvars +from datetime import UTC, datetime + +import pytest +from aws_durable_execution_sdk_python.plugin import ( + InvocationEndInfo, + InvocationStartInfo, + InvocationStatus, +) +from opentelemetry import context as otel_context +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +@pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_copied_context_on_same_thread_does_not_clear_its_parent_invocation(view): + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + a = contextvars.Context() + info_a = dict( + request_id="a", + execution_arn="execution/a", + execution_start_time=datetime.now(UTC), + is_first_invocation=True, + ) + info_b = dict( + request_id="b", + execution_arn="execution/b", + execution_start_time=datetime.now(UTC), + is_first_invocation=True, + ) + try: + original = a.run(otel_context.get_current) + a.run(plugin.on_invocation_start, InvocationStartInfo(**info_a)) + state_a = a.run(plugin._state.get) + span_a = a.run(plugin.get_current_span_context) + b = a.copy() + b.run(plugin.on_invocation_start, InvocationStartInfo(**info_b)) + state_b = b.run(plugin._state.get) + assert state_a is not state_b + b.run( + plugin.on_invocation_end, + InvocationEndInfo(**info_b, status=InvocationStatus.SUCCEEDED), + ) + assert a.run(plugin._state.get) is state_a + assert a.run(plugin.get_current_span_context) == span_a + assert state_a.execution_arn == "execution/a" and state_a.tracing_enabled + a.run( + plugin.on_invocation_end, + InvocationEndInfo(**info_a, status=InvocationStatus.SUCCEEDED), + ) + assert a.run(otel_context.get_current) == original + assert plugin._state not in a + workflows = [s for s in exporter.get_finished_spans() if s.name == "Workflow"] + assert len(workflows) == 2 + assert {s.attributes["durable.execution.arn"] for s in workflows} == { + "execution/a", + "execution/b", + } + assert len({s.context.trace_id for s in workflows}) == 2 + finally: + provider.shutdown() + + +@pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_same_owner_warm_restart_releases_abandoned_scope(view): + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter())) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + before = otel_context.get_current() + first = dict( + request_id="first", + execution_arn="execution/first", + execution_start_time=datetime.now(UTC), + is_first_invocation=True, + ) + second = dict( + request_id="second", + execution_arn="execution/second", + execution_start_time=datetime.now(UTC), + is_first_invocation=True, + ) + try: + plugin.on_invocation_start(InvocationStartInfo(**first)) + abandoned = plugin._state.get() + plugin.on_invocation_start(InvocationStartInfo(**second)) + assert abandoned.context_tokens == {} + assert abandoned.owner_token is None + plugin.on_invocation_end( + InvocationEndInfo(**second, status=InvocationStatus.SUCCEEDED) + ) + assert otel_context.get_current() == before + assert plugin._state not in contextvars.copy_context() + finally: + provider.shutdown() 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 ba9833f56..dcab099b6 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 @@ -6,6 +6,7 @@ import threading import time from concurrent.futures import ThreadPoolExecutor +from contextvars import copy_context from datetime import UTC, datetime from types import SimpleNamespace @@ -2015,7 +2016,9 @@ def test_detach_ignores_token_attached_on_another_thread(): with ThreadPoolExecutor(max_workers=1) as executor: executor.submit( - plugin.on_user_function_start, _user_function_start_info("step-1") + copy_context().run, + plugin.on_user_function_start, + _user_function_start_info("step-1"), ).result() plugin._detach_context("step-1:attempt:1") @@ -2133,20 +2136,23 @@ def test_suspension_releases_the_scope_on_the_originating_worker(): with ThreadPoolExecutor(max_workers=1) as worker: def suspend_on_worker() -> tuple[int, bool]: + worker_context = otel_context.get_current() plugin.on_user_function_start(_user_function_start_info(operation_id)) attached_span_id = trace.get_current_span().get_span_context().span_id plugin.on_user_function_end(_user_function_incomplete_info(operation_id)) return ( attached_span_id, - trace.get_current_span().get_span_context().is_valid, + otel_context.get_current() == worker_context, ) - attached_span_id, span_still_current = worker.submit(suspend_on_worker).result() + attached_span_id, context_restored = worker.submit( + copy_context().run, suspend_on_worker + ).result() suspended_span = plugin._get_span(span_key) # The scope was released on the worker, and its span is left open. assert attached_span_id != 0 - assert span_still_current is False + assert context_restored is True assert span_key not in plugin._context_tokens assert suspended_span is not None assert {s.name for s in exporter.get_finished_spans()} == { 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 d9c0f157e..71d07605e 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 @@ -464,9 +464,20 @@ def wait_for_checkpoint() -> None: updated_operation_ids=invocation_input.updated_operation_ids, ) # No checkpoint work starts until all Start hooks finish. - checkpoint_future = executor.submit( - execution_state.checkpoint_batches_forever - ) + if configured_plugins: + checkpoint_context = ( + plugin_executor._run_in_invocation_context( + contextvars.copy_context + ) + ) + checkpoint_future = executor.submit( + checkpoint_context.run, + execution_state.checkpoint_batches_forever, + ) + else: + checkpoint_future = executor.submit( + execution_state.checkpoint_batches_forever + ) output = plugin_executor._run_in_invocation_context( invoke_handler ) 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 ed3792102..fbd1c53d5 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 @@ -560,7 +560,9 @@ def execute_plugins(self, info, sync): 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) + self._executor.submit( + contextvars.copy_context().run, self._dispatch_plugin, plugin, info + ) def _run_in_invocation_context(self, invoke: Callable[[], Any]) -> Any: """Continue in the pre-hook Context only after a failed Start hook.""" From accaec49f8dcb57e772c4dd1973b0b35ffaade49 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 14 Oct 2026 18:44:40 +0000 Subject: [PATCH 45/54] test(otel): replay a caught callback failure --- .../README.md | 30 +- .../otel_25_caught_callback_failure_replay.py | 60 +++ .../template.yaml | 4 +- .../test_caught_callback_failure_replay.py | 358 ++++++++++++++++++ .../tests/test_otel_examples.py | 20 + 5 files changed, 460 insertions(+), 12 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_25_caught_callback_failure_replay.py create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_caught_callback_failure_replay.py diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md index 421de8f64..f8d59871d 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md @@ -67,7 +67,7 @@ view named by its `OtelSuite` parameter. | `otel-invocation-22` | `otel_22_user_function_context.handler` | Creates ordinary user spans under the active handler, attempt, child, branch, and iteration contexts. | | `otel-invocation-23` | `otel_23_callback_function_context.handler` | Verifies retry/check attempts, callback submitter, wrapped retry helper, and virtual-child callback parents. | | `otel-invocation-24` | `otel_24_invocation_retry_status.handler` | Raises a retryable invocation error after a completed step, then resumes with its saved result. | -| `otel-invocation-25` | `otel_17_wait_for_callback_failure.handler` | Targets a failed callback without error details; service history must satisfy the explicit no-error-details precondition. | +| `otel-invocation-25` | `otel_25_caught_callback_failure_replay.handler` | Catches the failed callback child, records `caught`, replays through two typed barriers, and returns `caught/one/two`. | | `otel-invocation-26` | `otel_26_external_callback_completion_replay.handler` | Completes a root callback after suspension, saves its result in a step, and replays it through two callback barriers. | | `otel-execution-1` | `otel_1_success.handler` | Verifies the execution-view workflow, step, and attempt hierarchy. | | `otel-execution-2` | `otel_2_wait_resume.handler` | Verifies the execution view across a resumed invocation. | @@ -93,7 +93,7 @@ view named by its `OtelSuite` parameter. | `otel-execution-22` | `otel_22_user_function_context.handler` | Observes active execution-view callback contexts without supplying or repairing parents. | | `otel-execution-23` | `otel_23_callback_function_context.handler` | Verifies the same SDK-owned callback lifecycle parents in the execution view. | | `otel-execution-24` | `otel_24_invocation_retry_status.handler` | Verifies invocation retry status independently of step retry and recovery re-exports. | -| `otel-execution-25` | `otel_17_wait_for_callback_failure.handler` | Targets the same errorless failed callback and its `UNSET` leaf in the execution view. | +| `otel-execution-25` | `otel_25_caught_callback_failure_replay.handler` | Requires one `FAILED`/`ERROR` callback leaf without an invented exception event while the workflow succeeds after two replays. | | `otel-execution-26` | `otel_26_external_callback_completion_replay.handler` | Requires one terminal root-callback export at first completion and no duplicate exports on two later replays. | | `otel-long-running-1` | `otel_long_running_1_wait.handler` | Verifies wait and resume telemetry across a long durable suspension. | | `otel-long-running-2` | `otel_long_running_2_retry.handler` | Verifies retry telemetry across a long durable backoff. | @@ -120,14 +120,24 @@ Detailed callback-failure history retains the service's empty error `Payload` object with `Truncated: false`, while SDK-facing state still has no error details. Metadata-only history and nonempty errors retain their existing representation. -Case 24 covers invocation `RETRY` becoming `RETRYING`/`UNSET`. Case 25 targets the -separate rule that a failed operation without error details remains `UNSET`. -Its cloud coverage requires both views to pass the raw service-history no-error-details -precondition and the telemetry assertions. Local file-store results establish -SDK behavior; that AWS service precondition still needs independent validation. -`CANCELLED`, `TIMED_OUT` and `STOPPED` without error details have explicit unit -coverage in both views; their corresponding cloud paths are not established. -The existing success/`OK` and detailed-error/`ERROR` controls remain in the suites. +Case 24 covers invocation `RETRY` becoming `RETRYING`/`UNSET`. Case 25 uses +`caught-callback-failure-replay`: it catches only `CallbackExternalError`, which +`wait_for_callback` restores from the failed child on replay, then checkpoints `caught` and awaits +two root barriers. The final result is `caught/one/two`, with three PENDING +invocations followed by SUCCEEDED. The failed callback remains `FAILED`/`ERROR` +without inventing an exception event. The marker body and failed submitter run +once; replay reconstructs the failed child without revisiting its callback body. +Marker and barrier failures escape the narrow catch. Case17 retains its +unhandled supplied-error behavior. + +This fixture must be paired with the Shared case25 catalog that expects a +successful execution and the three phase-gated callback actions. The previous +case25 catalog expects a failed execution and is not compatible with this new +routing. Cloud coverage requires the typed empty, untruncated callback error +payload and actual trace evidence in both views. The initial-active-invocation, +completion-during-flush and pruned-child-history boundaries are separate claims. +Other rare terminal statuses without details retain local coverage; no new +cloud path is claimed for them. Case 26 revisits a root-level public callback on every replay. The driver waits for `InvocationCompleted` before completing the target and each barrier. The diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_25_caught_callback_failure_replay.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_25_caught_callback_failure_replay.py new file mode 100644 index 000000000..d9d1a2f45 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_25_caught_callback_failure_replay.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Catch a real callback failure, then replay it before a successful return.""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python import ( + DurableContext, + StepContext, + durable_execution, + durable_step, +) +from aws_durable_execution_sdk_python.config import WaitForCallbackConfig +from aws_durable_execution_sdk_python.exceptions import CallbackExternalError +from aws_durable_execution_sdk_python.serdes import JsonSerDes +from aws_durable_execution_sdk_python.types import LambdaContext, WaitForCallbackContext +from common import otel_plugin, require_scenario + + +def submit_callback(_callback_id: str, _context: WaitForCallbackContext) -> None: + return None + + +@durable_step +def observe_failure(_context: StepContext) -> str: + return "caught" + + +def _run(event: dict[str, Any], context: DurableContext) -> str: + require_scenario(event, "caught-callback-failure-replay") + config = WaitForCallbackConfig(serdes=JsonSerDes()) + try: + context.wait_for_callback( + submit_callback, name="otel-failed-callback", config=config + ) + except CallbackExternalError: + # wait_for_callback restores the callback-typed cause on both first + # consumption and replay of its persisted failed child context. + pass + else: + raise AssertionError("Expected the callback to fail") + + observed = context.step(observe_failure(), name="otel-failed-callback-observed") + one = context.wait_for_callback( + submit_callback, name="otel-failed-callback-barrier-one", config=config + ) + two = context.wait_for_callback( + submit_callback, name="otel-failed-callback-barrier-two", config=config + ) + return "/".join((observed, one, two)) + + +_durable_handler = durable_execution(_run, plugins=[otel_plugin()]) + + +def handler(event: Any, context: LambdaContext) -> Any: + return _durable_handler(event, context) diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/template.yaml b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/template.yaml index 8d898ba46..819b253b7 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/template.yaml +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/template.yaml @@ -475,7 +475,7 @@ Resources: - otel-invocation-25 Properties: CodeUri: src/ - Handler: otel_17_wait_for_callback_failure.handler + Handler: otel_25_caught_callback_failure_replay.handler FunctionName: !Sub "${AWS::StackName}-otel-invocation-25" Role: !Ref LambdaExecutionRoleArn @@ -944,7 +944,7 @@ Resources: - otel-execution-25 Properties: CodeUri: src/ - Handler: otel_17_wait_for_callback_failure.handler + Handler: otel_25_caught_callback_failure_replay.handler FunctionName: !Sub "${AWS::StackName}-otel-execution-25" Role: !Ref LambdaExecutionRoleArn Environment: diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_caught_callback_failure_replay.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_caught_callback_failure_replay.py new file mode 100644 index 000000000..676227def --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_caught_callback_failure_replay.py @@ -0,0 +1,358 @@ +"""The real case25 handler catches only its intended failure across replay.""" + +import importlib.util +import json +import os +from pathlib import Path +import sys +import time + +import pytest +from aws_durable_execution_sdk_python import DurableContext +from aws_durable_execution_sdk_python.exceptions import ( + CallbackExternalError, + ChildContextError, +) +from aws_durable_execution_sdk_python.lambda_service import ErrorObject +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from aws_durable_execution_sdk_python_testing.stores.filesystem import ( + FileSystemExecutionStore, +) +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +SRC = Path(__file__).resolve().parents[1] / "src" +FAILED = "otel-failed-callback" +OBSERVED = FAILED + "-observed" +BARRIERS = [FAILED + "-barrier-one", FAILED + "-barrier-two"] +EXPECTED_HISTORY = [ + "ExecutionStarted", + "ContextStarted", + "CallbackStarted", + "StepStarted", + "StepSucceeded", + "InvocationCompleted", + "CallbackFailed", + "ContextFailed", + "StepStarted", + "StepSucceeded", + "ContextStarted", + "CallbackStarted", + "StepStarted", + "StepSucceeded", + "InvocationCompleted", + "CallbackSucceeded", + "ContextSucceeded", + "ContextStarted", + "CallbackStarted", + "StepStarted", + "StepSucceeded", + "InvocationCompleted", + "CallbackSucceeded", + "ContextSucceeded", + "InvocationCompleted", + "ExecutionSucceeded", +] + + +def _callback_after_suspend(runner, arn, name): + deadline = time.monotonic() + 8 + while time.monotonic() < deadline: + events = runner.get_execution_history(arn, include_execution_data=True).events + starts = [ + e + for e in events + if e.event_type == "CallbackStarted" + and e.name == name + " create callback id" + ] + ends = [ + e + for e in events + if starts + and e.event_type == "InvocationCompleted" + and e.event_id > starts[0].event_id + ] + if starts and ends: + return starts[0].callback_started_details.callback_id, ( + starts[0].event_id, + ends[0].event_id, + ) + time.sleep(0.01) + raise AssertionError(f"No creating invocation completion for {name}") + + +def _load(monkeypatch, plugin): + spec = importlib.util.spec_from_file_location("common", SRC / "common.py") + common = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, "common", common) + spec.loader.exec_module(common) + monkeypatch.setattr(common, "otel_plugin", lambda: plugin) + spec = importlib.util.spec_from_file_location( + "otel_25_caught_callback_failure_replay", + SRC / "otel_25_caught_callback_failure_replay.py", + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.mark.parametrize( + "view,mode", + [ + (ExecutionOtelPlugin, "normal"), + (InvocationOtelPlugin, "normal"), + (ExecutionOtelPlugin, "unexpected-success"), + (ExecutionOtelPlugin, "barrier-failure"), + ], +) +def test_case25_public_caught_failure_and_replay(monkeypatch, tmp_path, view, mode): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + starts = [] + invocation_ends = [] + exported = [] + operations = [] + user_starts = [] + failures = [] + marker_calls = [] + + class Exporter(InMemorySpanExporter): + def export(self, spans): + exported.extend((s, len(starts)) for s in spans) + return super().export(spans) + + exporter = Exporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + + def record(method, collection): + original = getattr(plugin, method) + + def wrapped(info): + collection.append((len(starts), info)) + return original(info) + + monkeypatch.setattr(plugin, method, wrapped) + + original_start = plugin.on_invocation_start + + def on_start(info): + starts.append(info) + return original_start(info) + + monkeypatch.setattr(plugin, "on_invocation_start", on_start) + record("on_invocation_end", invocation_ends) + record("on_operation_end", operations) + record("on_user_function_start", user_starts) + original_wait = DurableContext.wait_for_callback + + def wait(context, *args, **kwargs): + try: + return original_wait(context, *args, **kwargs) + except (CallbackExternalError, ChildContextError) as error: + if kwargs.get("name") == FAILED: + failures.append( + { + "invocation": len(starts), + "class": type(error).__name__, + "error_type": error.error_type, + "message": error.message, + } + ) + raise + + monkeypatch.setattr(DurableContext, "wait_for_callback", wait) + module = _load(monkeypatch, plugin) + original_marker = module.observe_failure + + def counted_marker(): + step = original_marker() + + def run(context): + marker_calls.append(len(starts)) + return step(context) + + return run + + monkeypatch.setattr(module, "observe_failure", counted_marker) + gates = [] + try: + with DurableFunctionTestRunner( + handler=module.handler, + store=FileSystemExecutionStore.create(tmp_path / "store"), + skip_time=False, + poll_interval=0.01, + execution_timeout=30, + ) as runner: + arn = runner.run_async( + input=json.dumps({"scenario": "caught-callback-failure-replay"}) + ) + callback, gate = _callback_after_suspend(runner, arn, FAILED) + gates.append(gate) + if mode == "unexpected-success": + runner.send_callback_success(callback, result=b'"unexpected"') + else: + runner.send_callback_failure(callback) + for name, result in zip(BARRIERS, ("one", "two")): + callback, gate = _callback_after_suspend(runner, arn, name) + gates.append(gate) + if mode == "barrier-failure": + runner.send_callback_failure( + callback, error=ErrorObject.from_message("barrier failure") + ) + break + runner.send_callback_success( + callback, result=json.dumps(result).encode() + ) + result = runner.wait_for_result(arn, timeout=10) + history = runner.get_execution_history( + arn, include_execution_data=True + ).to_dict() + span_rows = [ + { + "name": s.name, + "invocation": n, + "trace_id": f"{s.context.trace_id:032x}", + "span_id": f"{s.context.span_id:016x}", + "parent_id": f"{s.parent.span_id:016x}" if s.parent else None, + "status": s.status.status_code.name, + "start_time": s.start_time, + "end_time": s.end_time, + "attributes": dict(s.attributes or {}), + "events": [ + { + "name": e.name, + "timestamp": e.timestamp, + "attributes": dict(e.attributes or {}), + } + for e in s.events + ], + "links": [ + { + "trace_id": f"{link.context.trace_id:032x}", + "span_id": f"{link.context.span_id:016x}", + "attributes": dict(link.attributes or {}), + } + for link in s.links + ], + } + for s, n in exported + ] + summary = { + "view": view.__name__, + "mode": mode, + "execution_status": result.status.value, + "result": result.result, + "error": result.error.to_dict() if result.error else None, + "invocation_statuses": [info.status.value for _, info in invocation_ends], + "gates": gates, + "caught_failures": failures, + "marker_body_invocations": marker_calls, + "operation_ends": [ + { + "invocation": n, + "id": i.operation_id, + "name": i.name, + "parent_id": i.parent_id, + "type": i.operation_type.value, + "status": i.status.value, + "error": i.error.to_dict() if i.error else None, + } + for n, i in operations + ], + "user_function_starts": [ + { + "invocation": n, + "id": i.operation_id, + "name": i.name, + "type": i.operation_type.value, + } + for n, i in user_starts + ], + } + proof_dir = Path(os.environ.get("PYTHON_CASE25_PROOF_DIR", str(tmp_path))) + proof_dir.mkdir(parents=True, exist_ok=True) + prefix = f"python-case25-{view.__name__}-{mode}" + for suffix, data in ( + ("history", history), + ("spans", span_rows), + ("summary", summary), + ): + (proof_dir / f"{prefix}-{suffix}.json").write_text( + json.dumps(data, indent=2, default=str) + ) + if mode != "normal": + assert result.status.value == "FAILED" + assert marker_calls == ([] if mode == "unexpected-success" else [2]) + assert not any(e.get("Name") == BARRIERS[1] for e in history["Events"]) + return + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == "caught/one/two" + assert summary["invocation_statuses"] == [ + "PENDING", + "PENDING", + "PENDING", + "SUCCEEDED", + ] + assert [e["EventType"] for e in history["Events"]] == EXPECTED_HISTORY + assert [e["EventId"] for e in history["Events"]] == list(range(1, 27)) + assert gates == [(3, 6), (12, 15), (19, 22)] + assert marker_calls == [2] + assert [f["invocation"] for f in failures] == [2, 3, 4] + assert {f["class"] for f in failures} == {"CallbackExternalError"} + assert len({f["error_type"] for f in failures}) == 1 + failed = history["Events"][6] + assert failed["CallbackFailedDetails"]["Error"] == { + "Payload": {}, + "Truncated": False, + } + target_id = failed["Id"] + leaves = [ + s + for s in span_rows + if s["attributes"].get("durable.operation.id") == target_id + and s["attributes"].get("durable.operation.status") == "FAILED" + ] + assert len(leaves) == 1 + leaf = leaves[0] + assert ( + leaf["invocation"] == 2 + and leaf["status"] == "ERROR" + and leaf["events"] == [] + ) + marker_spans = [ + s + for s in span_rows + if s["attributes"].get("durable.operation.name") == OBSERVED + ] + assert marker_spans and all(s["invocation"] == 2 for s in marker_spans) + assert leaf["end_time"] <= min(s["start_time"] for s in marker_spans) + workflow = [s for s in span_rows if s["name"] == "Workflow"] + assert len(workflow) == 1 and workflow[0]["status"] == "OK" + parent = next(s for s in span_rows if s["span_id"] == leaf["parent_id"]) + assert ( + parent["attributes"]["durable.operation.id"] == history["Events"][1]["Id"] + ) + if view is InvocationOtelPlugin: + invocation = next( + s for s in span_rows if s["span_id"] == parent["parent_id"] + ) + assert invocation["name"] == "Invocation" and invocation["invocation"] == 2 + submitter_id = history["Events"][3]["Id"] + assert sum(i.operation_id == submitter_id for _, i in user_starts) == 1 + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py index 61c0d24e1..fec2b7326 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py @@ -168,6 +168,7 @@ "otel_22_user_function_context", "otel_23_callback_function_context", "otel_24_invocation_retry_status", + "otel_25_caught_callback_failure_replay", "otel_26_external_callback_completion_replay", "otel_long_running_1_wait", "otel_long_running_2_retry", @@ -408,6 +409,25 @@ def test_every_registered_handler_exists( ) +def test_case25_uses_its_caught_replay_handler_without_retargeting_other_cases( + template: dict[str, Any], +) -> None: + resources = template["Resources"] + for prefix in ("Otel", "OtelExecution"): + assert resources[prefix + "25CallbackFailureWithoutError"]["Properties"][ + "Handler" + ] == ("otel_25_caught_callback_failure_replay.handler") + assert resources[prefix + "17WaitForCallbackFailure"]["Properties"][ + "Handler" + ] == ("otel_17_wait_for_callback_failure.handler") + assert resources[prefix + "24InvocationRetryStatus"]["Properties"][ + "Handler" + ] == ("otel_24_invocation_retry_status.handler") + assert resources[prefix + "26ExternalCallbackReplay"]["Properties"][ + "Handler" + ] == ("otel_26_external_callback_completion_replay.handler") + + def test_template_deploys_only_the_selected_view(template: dict[str, Any]) -> None: assert template["Parameters"]["OtelSuite"]["AllowedValues"] == [ "all", From 5ea2a4b431f3b85ddf2ac94c58ae135ee2d397fc Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 14 Oct 2026 19:31:07 +0000 Subject: [PATCH 46/54] fix(testing): acknowledge legacy updates --- .../execution.py | 22 ++++++------ .../e2e/invocation_input_delivery_test.py | 36 ++++++++++++++++++- .../tests/execution_test.py | 26 ++++++++++++++ 3 files changed, 73 insertions(+), 11 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py index 569b7f1f2..0a16afac1 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py @@ -877,18 +877,20 @@ def unseen_operations(self) -> list[Operation]: ] def advance_handler_seen(self, seq: int) -> None: - """Advance the wrapped execution's ``handler_seen_seq`` to - ``seq`` if that represents forward progress. Monotonic: - smaller or equal values are ignored.""" + """Acknowledge delivered updates without retreating the watermark. + + Equal acknowledgments also consume old-format, unversioned update IDs. + Only versions covered by this delivery may be removed. + """ if seq > self.execution.handler_seen_seq: self.execution.handler_seen_seq = seq - # A checkpoint response or completed invocation delivered this state. - # Retain changes newer than that acknowledged snapshot. - self.execution.updated_operation_ids = [ - operation_id - for operation_id in self.execution.updated_operation_ids - if self.execution.operation_last_touched_seq.get(operation_id, 0) > seq - ] + # A checkpoint response or completed invocation delivered this state. + # Retain changes newer than that acknowledged snapshot. + self.execution.updated_operation_ids = [ + operation_id + for operation_id in self.execution.updated_operation_ids + if self.execution.operation_last_touched_seq.get(operation_id, 0) > seq + ] # --- internals ------------------------------------------------- diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/e2e/invocation_input_delivery_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/invocation_input_delivery_test.py index af69bc3dc..dfc64e7aa 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/e2e/invocation_input_delivery_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/invocation_input_delivery_test.py @@ -12,6 +12,8 @@ from aws_durable_execution_sdk_python.config import CallbackConfig from aws_durable_execution_sdk_python.serdes import JsonSerDes from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from aws_durable_execution_sdk_python_testing.execution import Execution +from aws_durable_execution_sdk_python_testing.executor import Executor, InvocationState from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner from aws_durable_execution_sdk_python_testing.stores.filesystem import ( FileSystemExecutionStore, @@ -23,8 +25,9 @@ @pytest.mark.parametrize("filesystem", [False, True]) @pytest.mark.parametrize("failed_delivery", [False, True]) +@pytest.mark.parametrize("legacy_format", [False, True]) def test_callback_input_acknowledgment_survives_checkpoint_free_resume( - filesystem, failed_delivery, tmp_path + filesystem, failed_delivery, legacy_format, tmp_path, monkeypatch ): store = ( FileSystemExecutionStore.create(tmp_path / "store") @@ -36,6 +39,36 @@ def test_callback_input_acknowledgment_survives_checkpoint_free_resume( operation_ends = [] attempts = [] marker_calls = [] + legacy_restores = [] + + original_begin = Executor._begin_invocation # noqa: SLF001 + + def begin_with_old_store(self, arn): + # Rehydrate a real completed callback using the persisted format that + # predates sequence accounting. No SDK hook or invocation input is faked. + if ( + legacy_format + and not legacy_restores + and self._invocation_gate(arn) is InvocationState.PRE_INVOKE # noqa: SLF001 + ): + execution = store.load(arn) + if execution.updated_operation_ids: + snapshot = execution.to_json_dict() + for field in ( + "SeqCounter", + "HandlerSeenSeq", + "OperationLastTouchedSeq", + "OperationSizeBytes", + ): + snapshot.pop(field) + restored = Execution.from_json_dict(snapshot) + assert restored.seq_counter == restored.handler_seen_seq == 0 + assert restored.updated_operation_ids == execution.updated_operation_ids + store.save(restored) + legacy_restores.append(snapshot) + return original_begin(self, arn) + + monkeypatch.setattr(Executor, "_begin_invocation", begin_with_old_store) class Observe(DurableInstrumentationPlugin): def on_invocation_start(self, info): @@ -108,3 +141,4 @@ def entrypoint(event, context): assert operation_ends == [(2, "a", False), (3, "b", False)] assert marker_calls == [1] assert len(attempts) == (4 if failed_delivery else 3) + assert len(legacy_restores) == int(legacy_format) diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/execution_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/execution_test.py index 353d320d4..c9457a6dd 100644 --- a/packages/aws-durable-execution-sdk-python-testing/tests/execution_test.py +++ b/packages/aws-durable-execution-sdk-python-testing/tests/execution_test.py @@ -1778,6 +1778,32 @@ def complete_later(): assert execution.updated_operation_ids == [] +@pytest.mark.parametrize("seen,ack", [(0, 0), (5, 5), (5, 3)]) +def test_acknowledged_updates_are_pruned_without_watermark_progress(seen, ack): + execution = Execution("test-arn", _make_start_input(), []) + snapshot = execution.to_json_dict() + for field in ( + "SeqCounter", + "HandlerSeenSeq", + "OperationLastTouchedSeq", + "OperationSizeBytes", + ): + snapshot.pop(field) + snapshot["UpdatedOperationIds"] = ["legacy", "newer"] + restored = Execution.from_json_dict(snapshot) + restored.handler_seen_seq = seen + restored.operation_last_touched_seq["newer"] = ack + 1 + paginator = OperationPaginatorState.pin(restored) + + paginator.advance_handler_seen(ack) + + assert restored.updated_operation_ids == ["newer"] + assert restored.handler_seen_seq == seen + assert restored.token_sequence == restored.seq_counter == 0 + paginator.advance_handler_seen(ack) + assert restored.updated_operation_ids == ["newer"] + + def test_function_arn_is_qualified_with_the_executed_version(): execution = Execution.new(_make_start_input()) execution.region = "ap-southeast-2" From f8ccecb6c5c0fd56799c080799700b7631448413 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 14 Oct 2026 19:48:11 +0000 Subject: [PATCH 47/54] ci: pair caught replay fixture and catalog --- .../scripts/tests/test_opentelemetry_conformance_workflow.py | 2 +- .github/workflows/opentelemetry-conformance-tests.yml | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py index c29b3f75e..c27d0948d 100644 --- a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py +++ b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py @@ -5,7 +5,7 @@ SHARED_WORKFLOW_REF = "bdb4f1cd0f9252c1aaa978bb8b341b71f2b9d9dc" -CONFORMANCE_TEST_REF = "f711bbfeb5464400d3aa1f8366e0bbb02087a67d" +CONFORMANCE_TEST_REF = "05960beb6c81d7b24b53114f2581b3b6f9fef6d4" WORKFLOW_PATH = ( Path(__file__).parents[2] / "workflows" / "opentelemetry-conformance-tests.yml" diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index d4177270c..c45d2e583 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -45,7 +45,7 @@ on: conformance_test_ref: description: Conformance test commit SHA or branch name required: true - default: f711bbfeb5464400d3aa1f8366e0bbb02087a67d + default: 05960beb6c81d7b24b53114f2581b3b6f9fef6d4 type: string # Backend stacks are shared across PRs. Queue whole runs so reusable @@ -73,7 +73,7 @@ jobs: resource_prefix: p sdk_repository: aws/aws-durable-execution-sdk-python sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }} - conformance_test_ref: ${{ inputs.conformance_test_ref || 'f711bbfeb5464400d3aa1f8366e0bbb02087a67d' }} + conformance_test_ref: ${{ inputs.conformance_test_ref || '05960beb6c81d7b24b53114f2581b3b6f9fef6d4' }} # Check the SDK out so the handlers and templates below are on disk. The handlers # themselves are installed from sdk_ref by src/requirements.txt during the SAM build. checkout_sdk: true From 4d3df9d6a28239b47e7a05736be41b2b24448c27 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 14 Oct 2026 18:44:40 +0000 Subject: [PATCH 48/54] test(otel): replay a caught callback failure --- .../README.md | 30 +- .../otel_25_caught_callback_failure_replay.py | 60 +++ .../template.yaml | 4 +- .../test_caught_callback_failure_replay.py | 358 ++++++++++++++++++ .../tests/test_otel_examples.py | 20 + 5 files changed, 460 insertions(+), 12 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_25_caught_callback_failure_replay.py create mode 100644 packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_caught_callback_failure_replay.py diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md index 421de8f64..f8d59871d 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/README.md @@ -67,7 +67,7 @@ view named by its `OtelSuite` parameter. | `otel-invocation-22` | `otel_22_user_function_context.handler` | Creates ordinary user spans under the active handler, attempt, child, branch, and iteration contexts. | | `otel-invocation-23` | `otel_23_callback_function_context.handler` | Verifies retry/check attempts, callback submitter, wrapped retry helper, and virtual-child callback parents. | | `otel-invocation-24` | `otel_24_invocation_retry_status.handler` | Raises a retryable invocation error after a completed step, then resumes with its saved result. | -| `otel-invocation-25` | `otel_17_wait_for_callback_failure.handler` | Targets a failed callback without error details; service history must satisfy the explicit no-error-details precondition. | +| `otel-invocation-25` | `otel_25_caught_callback_failure_replay.handler` | Catches the failed callback child, records `caught`, replays through two typed barriers, and returns `caught/one/two`. | | `otel-invocation-26` | `otel_26_external_callback_completion_replay.handler` | Completes a root callback after suspension, saves its result in a step, and replays it through two callback barriers. | | `otel-execution-1` | `otel_1_success.handler` | Verifies the execution-view workflow, step, and attempt hierarchy. | | `otel-execution-2` | `otel_2_wait_resume.handler` | Verifies the execution view across a resumed invocation. | @@ -93,7 +93,7 @@ view named by its `OtelSuite` parameter. | `otel-execution-22` | `otel_22_user_function_context.handler` | Observes active execution-view callback contexts without supplying or repairing parents. | | `otel-execution-23` | `otel_23_callback_function_context.handler` | Verifies the same SDK-owned callback lifecycle parents in the execution view. | | `otel-execution-24` | `otel_24_invocation_retry_status.handler` | Verifies invocation retry status independently of step retry and recovery re-exports. | -| `otel-execution-25` | `otel_17_wait_for_callback_failure.handler` | Targets the same errorless failed callback and its `UNSET` leaf in the execution view. | +| `otel-execution-25` | `otel_25_caught_callback_failure_replay.handler` | Requires one `FAILED`/`ERROR` callback leaf without an invented exception event while the workflow succeeds after two replays. | | `otel-execution-26` | `otel_26_external_callback_completion_replay.handler` | Requires one terminal root-callback export at first completion and no duplicate exports on two later replays. | | `otel-long-running-1` | `otel_long_running_1_wait.handler` | Verifies wait and resume telemetry across a long durable suspension. | | `otel-long-running-2` | `otel_long_running_2_retry.handler` | Verifies retry telemetry across a long durable backoff. | @@ -120,14 +120,24 @@ Detailed callback-failure history retains the service's empty error `Payload` object with `Truncated: false`, while SDK-facing state still has no error details. Metadata-only history and nonempty errors retain their existing representation. -Case 24 covers invocation `RETRY` becoming `RETRYING`/`UNSET`. Case 25 targets the -separate rule that a failed operation without error details remains `UNSET`. -Its cloud coverage requires both views to pass the raw service-history no-error-details -precondition and the telemetry assertions. Local file-store results establish -SDK behavior; that AWS service precondition still needs independent validation. -`CANCELLED`, `TIMED_OUT` and `STOPPED` without error details have explicit unit -coverage in both views; their corresponding cloud paths are not established. -The existing success/`OK` and detailed-error/`ERROR` controls remain in the suites. +Case 24 covers invocation `RETRY` becoming `RETRYING`/`UNSET`. Case 25 uses +`caught-callback-failure-replay`: it catches only `CallbackExternalError`, which +`wait_for_callback` restores from the failed child on replay, then checkpoints `caught` and awaits +two root barriers. The final result is `caught/one/two`, with three PENDING +invocations followed by SUCCEEDED. The failed callback remains `FAILED`/`ERROR` +without inventing an exception event. The marker body and failed submitter run +once; replay reconstructs the failed child without revisiting its callback body. +Marker and barrier failures escape the narrow catch. Case17 retains its +unhandled supplied-error behavior. + +This fixture must be paired with the Shared case25 catalog that expects a +successful execution and the three phase-gated callback actions. The previous +case25 catalog expects a failed execution and is not compatible with this new +routing. Cloud coverage requires the typed empty, untruncated callback error +payload and actual trace evidence in both views. The initial-active-invocation, +completion-during-flush and pruned-child-history boundaries are separate claims. +Other rare terminal statuses without details retain local coverage; no new +cloud path is claimed for them. Case 26 revisits a root-level public callback on every replay. The driver waits for `InvocationCompleted` before completing the target and each barrier. The diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_25_caught_callback_failure_replay.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_25_caught_callback_failure_replay.py new file mode 100644 index 000000000..d9d1a2f45 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/src/otel_25_caught_callback_failure_replay.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Catch a real callback failure, then replay it before a successful return.""" + +from __future__ import annotations + +from typing import Any + +from aws_durable_execution_sdk_python import ( + DurableContext, + StepContext, + durable_execution, + durable_step, +) +from aws_durable_execution_sdk_python.config import WaitForCallbackConfig +from aws_durable_execution_sdk_python.exceptions import CallbackExternalError +from aws_durable_execution_sdk_python.serdes import JsonSerDes +from aws_durable_execution_sdk_python.types import LambdaContext, WaitForCallbackContext +from common import otel_plugin, require_scenario + + +def submit_callback(_callback_id: str, _context: WaitForCallbackContext) -> None: + return None + + +@durable_step +def observe_failure(_context: StepContext) -> str: + return "caught" + + +def _run(event: dict[str, Any], context: DurableContext) -> str: + require_scenario(event, "caught-callback-failure-replay") + config = WaitForCallbackConfig(serdes=JsonSerDes()) + try: + context.wait_for_callback( + submit_callback, name="otel-failed-callback", config=config + ) + except CallbackExternalError: + # wait_for_callback restores the callback-typed cause on both first + # consumption and replay of its persisted failed child context. + pass + else: + raise AssertionError("Expected the callback to fail") + + observed = context.step(observe_failure(), name="otel-failed-callback-observed") + one = context.wait_for_callback( + submit_callback, name="otel-failed-callback-barrier-one", config=config + ) + two = context.wait_for_callback( + submit_callback, name="otel-failed-callback-barrier-two", config=config + ) + return "/".join((observed, one, two)) + + +_durable_handler = durable_execution(_run, plugins=[otel_plugin()]) + + +def handler(event: Any, context: LambdaContext) -> Any: + return _durable_handler(event, context) diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/template.yaml b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/template.yaml index 8d898ba46..819b253b7 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/template.yaml +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/template.yaml @@ -475,7 +475,7 @@ Resources: - otel-invocation-25 Properties: CodeUri: src/ - Handler: otel_17_wait_for_callback_failure.handler + Handler: otel_25_caught_callback_failure_replay.handler FunctionName: !Sub "${AWS::StackName}-otel-invocation-25" Role: !Ref LambdaExecutionRoleArn @@ -944,7 +944,7 @@ Resources: - otel-execution-25 Properties: CodeUri: src/ - Handler: otel_17_wait_for_callback_failure.handler + Handler: otel_25_caught_callback_failure_replay.handler FunctionName: !Sub "${AWS::StackName}-otel-execution-25" Role: !Ref LambdaExecutionRoleArn Environment: diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_caught_callback_failure_replay.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_caught_callback_failure_replay.py new file mode 100644 index 000000000..676227def --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_caught_callback_failure_replay.py @@ -0,0 +1,358 @@ +"""The real case25 handler catches only its intended failure across replay.""" + +import importlib.util +import json +import os +from pathlib import Path +import sys +import time + +import pytest +from aws_durable_execution_sdk_python import DurableContext +from aws_durable_execution_sdk_python.exceptions import ( + CallbackExternalError, + ChildContextError, +) +from aws_durable_execution_sdk_python.lambda_service import ErrorObject +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from aws_durable_execution_sdk_python_testing.stores.filesystem import ( + FileSystemExecutionStore, +) +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +SRC = Path(__file__).resolve().parents[1] / "src" +FAILED = "otel-failed-callback" +OBSERVED = FAILED + "-observed" +BARRIERS = [FAILED + "-barrier-one", FAILED + "-barrier-two"] +EXPECTED_HISTORY = [ + "ExecutionStarted", + "ContextStarted", + "CallbackStarted", + "StepStarted", + "StepSucceeded", + "InvocationCompleted", + "CallbackFailed", + "ContextFailed", + "StepStarted", + "StepSucceeded", + "ContextStarted", + "CallbackStarted", + "StepStarted", + "StepSucceeded", + "InvocationCompleted", + "CallbackSucceeded", + "ContextSucceeded", + "ContextStarted", + "CallbackStarted", + "StepStarted", + "StepSucceeded", + "InvocationCompleted", + "CallbackSucceeded", + "ContextSucceeded", + "InvocationCompleted", + "ExecutionSucceeded", +] + + +def _callback_after_suspend(runner, arn, name): + deadline = time.monotonic() + 8 + while time.monotonic() < deadline: + events = runner.get_execution_history(arn, include_execution_data=True).events + starts = [ + e + for e in events + if e.event_type == "CallbackStarted" + and e.name == name + " create callback id" + ] + ends = [ + e + for e in events + if starts + and e.event_type == "InvocationCompleted" + and e.event_id > starts[0].event_id + ] + if starts and ends: + return starts[0].callback_started_details.callback_id, ( + starts[0].event_id, + ends[0].event_id, + ) + time.sleep(0.01) + raise AssertionError(f"No creating invocation completion for {name}") + + +def _load(monkeypatch, plugin): + spec = importlib.util.spec_from_file_location("common", SRC / "common.py") + common = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, "common", common) + spec.loader.exec_module(common) + monkeypatch.setattr(common, "otel_plugin", lambda: plugin) + spec = importlib.util.spec_from_file_location( + "otel_25_caught_callback_failure_replay", + SRC / "otel_25_caught_callback_failure_replay.py", + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.mark.parametrize( + "view,mode", + [ + (ExecutionOtelPlugin, "normal"), + (InvocationOtelPlugin, "normal"), + (ExecutionOtelPlugin, "unexpected-success"), + (ExecutionOtelPlugin, "barrier-failure"), + ], +) +def test_case25_public_caught_failure_and_replay(monkeypatch, tmp_path, view, mode): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + starts = [] + invocation_ends = [] + exported = [] + operations = [] + user_starts = [] + failures = [] + marker_calls = [] + + class Exporter(InMemorySpanExporter): + def export(self, spans): + exported.extend((s, len(starts)) for s in spans) + return super().export(spans) + + exporter = Exporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + + def record(method, collection): + original = getattr(plugin, method) + + def wrapped(info): + collection.append((len(starts), info)) + return original(info) + + monkeypatch.setattr(plugin, method, wrapped) + + original_start = plugin.on_invocation_start + + def on_start(info): + starts.append(info) + return original_start(info) + + monkeypatch.setattr(plugin, "on_invocation_start", on_start) + record("on_invocation_end", invocation_ends) + record("on_operation_end", operations) + record("on_user_function_start", user_starts) + original_wait = DurableContext.wait_for_callback + + def wait(context, *args, **kwargs): + try: + return original_wait(context, *args, **kwargs) + except (CallbackExternalError, ChildContextError) as error: + if kwargs.get("name") == FAILED: + failures.append( + { + "invocation": len(starts), + "class": type(error).__name__, + "error_type": error.error_type, + "message": error.message, + } + ) + raise + + monkeypatch.setattr(DurableContext, "wait_for_callback", wait) + module = _load(monkeypatch, plugin) + original_marker = module.observe_failure + + def counted_marker(): + step = original_marker() + + def run(context): + marker_calls.append(len(starts)) + return step(context) + + return run + + monkeypatch.setattr(module, "observe_failure", counted_marker) + gates = [] + try: + with DurableFunctionTestRunner( + handler=module.handler, + store=FileSystemExecutionStore.create(tmp_path / "store"), + skip_time=False, + poll_interval=0.01, + execution_timeout=30, + ) as runner: + arn = runner.run_async( + input=json.dumps({"scenario": "caught-callback-failure-replay"}) + ) + callback, gate = _callback_after_suspend(runner, arn, FAILED) + gates.append(gate) + if mode == "unexpected-success": + runner.send_callback_success(callback, result=b'"unexpected"') + else: + runner.send_callback_failure(callback) + for name, result in zip(BARRIERS, ("one", "two")): + callback, gate = _callback_after_suspend(runner, arn, name) + gates.append(gate) + if mode == "barrier-failure": + runner.send_callback_failure( + callback, error=ErrorObject.from_message("barrier failure") + ) + break + runner.send_callback_success( + callback, result=json.dumps(result).encode() + ) + result = runner.wait_for_result(arn, timeout=10) + history = runner.get_execution_history( + arn, include_execution_data=True + ).to_dict() + span_rows = [ + { + "name": s.name, + "invocation": n, + "trace_id": f"{s.context.trace_id:032x}", + "span_id": f"{s.context.span_id:016x}", + "parent_id": f"{s.parent.span_id:016x}" if s.parent else None, + "status": s.status.status_code.name, + "start_time": s.start_time, + "end_time": s.end_time, + "attributes": dict(s.attributes or {}), + "events": [ + { + "name": e.name, + "timestamp": e.timestamp, + "attributes": dict(e.attributes or {}), + } + for e in s.events + ], + "links": [ + { + "trace_id": f"{link.context.trace_id:032x}", + "span_id": f"{link.context.span_id:016x}", + "attributes": dict(link.attributes or {}), + } + for link in s.links + ], + } + for s, n in exported + ] + summary = { + "view": view.__name__, + "mode": mode, + "execution_status": result.status.value, + "result": result.result, + "error": result.error.to_dict() if result.error else None, + "invocation_statuses": [info.status.value for _, info in invocation_ends], + "gates": gates, + "caught_failures": failures, + "marker_body_invocations": marker_calls, + "operation_ends": [ + { + "invocation": n, + "id": i.operation_id, + "name": i.name, + "parent_id": i.parent_id, + "type": i.operation_type.value, + "status": i.status.value, + "error": i.error.to_dict() if i.error else None, + } + for n, i in operations + ], + "user_function_starts": [ + { + "invocation": n, + "id": i.operation_id, + "name": i.name, + "type": i.operation_type.value, + } + for n, i in user_starts + ], + } + proof_dir = Path(os.environ.get("PYTHON_CASE25_PROOF_DIR", str(tmp_path))) + proof_dir.mkdir(parents=True, exist_ok=True) + prefix = f"python-case25-{view.__name__}-{mode}" + for suffix, data in ( + ("history", history), + ("spans", span_rows), + ("summary", summary), + ): + (proof_dir / f"{prefix}-{suffix}.json").write_text( + json.dumps(data, indent=2, default=str) + ) + if mode != "normal": + assert result.status.value == "FAILED" + assert marker_calls == ([] if mode == "unexpected-success" else [2]) + assert not any(e.get("Name") == BARRIERS[1] for e in history["Events"]) + return + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == "caught/one/two" + assert summary["invocation_statuses"] == [ + "PENDING", + "PENDING", + "PENDING", + "SUCCEEDED", + ] + assert [e["EventType"] for e in history["Events"]] == EXPECTED_HISTORY + assert [e["EventId"] for e in history["Events"]] == list(range(1, 27)) + assert gates == [(3, 6), (12, 15), (19, 22)] + assert marker_calls == [2] + assert [f["invocation"] for f in failures] == [2, 3, 4] + assert {f["class"] for f in failures} == {"CallbackExternalError"} + assert len({f["error_type"] for f in failures}) == 1 + failed = history["Events"][6] + assert failed["CallbackFailedDetails"]["Error"] == { + "Payload": {}, + "Truncated": False, + } + target_id = failed["Id"] + leaves = [ + s + for s in span_rows + if s["attributes"].get("durable.operation.id") == target_id + and s["attributes"].get("durable.operation.status") == "FAILED" + ] + assert len(leaves) == 1 + leaf = leaves[0] + assert ( + leaf["invocation"] == 2 + and leaf["status"] == "ERROR" + and leaf["events"] == [] + ) + marker_spans = [ + s + for s in span_rows + if s["attributes"].get("durable.operation.name") == OBSERVED + ] + assert marker_spans and all(s["invocation"] == 2 for s in marker_spans) + assert leaf["end_time"] <= min(s["start_time"] for s in marker_spans) + workflow = [s for s in span_rows if s["name"] == "Workflow"] + assert len(workflow) == 1 and workflow[0]["status"] == "OK" + parent = next(s for s in span_rows if s["span_id"] == leaf["parent_id"]) + assert ( + parent["attributes"]["durable.operation.id"] == history["Events"][1]["Id"] + ) + if view is InvocationOtelPlugin: + invocation = next( + s for s in span_rows if s["span_id"] == parent["parent_id"] + ) + assert invocation["name"] == "Invocation" and invocation["invocation"] == 2 + submitter_id = history["Events"][3]["Id"] + assert sum(i.operation_id == submitter_id for _, i in user_starts) == 1 + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py index 61c0d24e1..fec2b7326 100644 --- a/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py +++ b/packages/aws-durable-execution-sdk-python-conformance-tests-otel/tests/test_otel_examples.py @@ -168,6 +168,7 @@ "otel_22_user_function_context", "otel_23_callback_function_context", "otel_24_invocation_retry_status", + "otel_25_caught_callback_failure_replay", "otel_26_external_callback_completion_replay", "otel_long_running_1_wait", "otel_long_running_2_retry", @@ -408,6 +409,25 @@ def test_every_registered_handler_exists( ) +def test_case25_uses_its_caught_replay_handler_without_retargeting_other_cases( + template: dict[str, Any], +) -> None: + resources = template["Resources"] + for prefix in ("Otel", "OtelExecution"): + assert resources[prefix + "25CallbackFailureWithoutError"]["Properties"][ + "Handler" + ] == ("otel_25_caught_callback_failure_replay.handler") + assert resources[prefix + "17WaitForCallbackFailure"]["Properties"][ + "Handler" + ] == ("otel_17_wait_for_callback_failure.handler") + assert resources[prefix + "24InvocationRetryStatus"]["Properties"][ + "Handler" + ] == ("otel_24_invocation_retry_status.handler") + assert resources[prefix + "26ExternalCallbackReplay"]["Properties"][ + "Handler" + ] == ("otel_26_external_callback_completion_replay.handler") + + def test_template_deploys_only_the_selected_view(template: dict[str, Any]) -> None: assert template["Parameters"]["OtelSuite"]["AllowedValues"] == [ "all", From d5738f3c344f05d8e8b9331564cfe49827a691e3 Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 14 Oct 2026 19:48:11 +0000 Subject: [PATCH 49/54] ci: pair caught replay fixture and catalog --- .../scripts/tests/test_opentelemetry_conformance_workflow.py | 2 +- .github/workflows/opentelemetry-conformance-tests.yml | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py index c29b3f75e..c27d0948d 100644 --- a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py +++ b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py @@ -5,7 +5,7 @@ SHARED_WORKFLOW_REF = "bdb4f1cd0f9252c1aaa978bb8b341b71f2b9d9dc" -CONFORMANCE_TEST_REF = "f711bbfeb5464400d3aa1f8366e0bbb02087a67d" +CONFORMANCE_TEST_REF = "05960beb6c81d7b24b53114f2581b3b6f9fef6d4" WORKFLOW_PATH = ( Path(__file__).parents[2] / "workflows" / "opentelemetry-conformance-tests.yml" diff --git a/.github/workflows/opentelemetry-conformance-tests.yml b/.github/workflows/opentelemetry-conformance-tests.yml index d4177270c..c45d2e583 100644 --- a/.github/workflows/opentelemetry-conformance-tests.yml +++ b/.github/workflows/opentelemetry-conformance-tests.yml @@ -45,7 +45,7 @@ on: conformance_test_ref: description: Conformance test commit SHA or branch name required: true - default: f711bbfeb5464400d3aa1f8366e0bbb02087a67d + default: 05960beb6c81d7b24b53114f2581b3b6f9fef6d4 type: string # Backend stacks are shared across PRs. Queue whole runs so reusable @@ -73,7 +73,7 @@ jobs: resource_prefix: p sdk_repository: aws/aws-durable-execution-sdk-python sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }} - conformance_test_ref: ${{ inputs.conformance_test_ref || 'f711bbfeb5464400d3aa1f8366e0bbb02087a67d' }} + conformance_test_ref: ${{ inputs.conformance_test_ref || '05960beb6c81d7b24b53114f2581b3b6f9fef6d4' }} # Check the SDK out so the handlers and templates below are on disk. The handlers # themselves are installed from sdk_ref by src/requirements.txt during the SAM build. checkout_sdk: true From 4028c491dd76d8ce776fadb8776f838685b8d42a Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 14 Oct 2026 19:51:32 +0000 Subject: [PATCH 50/54] docs(otel): state failed delivery limits --- packages/aws-durable-execution-sdk-python-otel/README.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 315758a9f..c5e8db1bd 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -441,6 +441,14 @@ raised retroactively into the completed invocation. ### Invocation context isolation +Completion metadata can be redelivered after a failed invocation or checkpoint +transport attempt that has no successful response acknowledgment. There is no +exporter acknowledgment in durable state, so this does not promise exactly-once +telemetry across failed/redelivered invocations. Successful replay preserves +checkpointed work; this limit is not a claim that already-checkpointed step +bodies run again. The plugin does not use a global exported-operation cache that +could discard a first delivery after a failed attempt. + With core 2.1+, invocation hooks and the handler run on one worker in an invocation-local Context initialized from the host's bindings. Successful Start bindings are visible to later hooks, the handler and resource cleanup. The host's From 15ab1abb8cf16b0242539747e63fd614ca2111f8 Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Wed, 14 Oct 2026 21:34:01 +0000 Subject: [PATCH 51/54] fix(otel): deliver cached child input updates --- .../README.md | 8 + .../_invocation_state.py | 34 ++- .../execution_plugin.py | 35 +++ .../invocation_plugin.py | 16 +- .../test_cached_child_input_delivery_int.py | 225 ++++++++++++++++++ .../tests/test_terminal_input_parent_refs.py | 181 ++++++++++++++ .../README.md | 8 + .../plugin.py | 71 +++++- .../tests/plugin_terminal_input_test.py | 132 ++++++++++ 9 files changed, 695 insertions(+), 15 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_cached_child_input_delivery_int.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/test_terminal_input_parent_refs.py create mode 100644 packages/aws-durable-execution-sdk-python/tests/plugin_terminal_input_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 f083086ba..cd4de725a 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -448,6 +448,14 @@ cleanup clearing another's trace. Provider configuration and registration remain instance resources. SDK-owned worker tasks capture the selected invocation Context; Start/End hooks and token reset retain their existing owner and order. +Terminal input updates can be observed even when replay skips their completed +child context. When the input contains that checkpointed parent's terminal +metadata, the views reference its original logical span identity without +re-exporting the parent. Missing or nonterminal parents are deferred until their +actual span is available; unresolved parentage is reported after cleanup instead +of inventing a root parent. Draining this private queue does not call the public +operation-end hook again. + Completion metadata can be redelivered after a failed invocation or checkpoint transport attempt that has no successful response acknowledgment. There is no exporter acknowledgment in durable state, so this does not promise exactly-once diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py index 10dd88482..cbece37fb 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py @@ -9,7 +9,13 @@ import threading from typing import Any, Generic, TypeVar -from aws_durable_execution_sdk_python.plugin import OperationEndInfo +from aws_durable_execution_sdk_python.plugin import ( + InvocationStartInfo, + OperationEndInfo, + OperationInfo, + OperationStatus, + OperationType, +) from opentelemetry.sdk.trace.sampling import Sampler from opentelemetry.trace import Span, Tracer, TracerProvider @@ -51,6 +57,9 @@ class ExecutionState(_InvocationState[str]): invocation_span: Span | None = None checkpointed_context_ids: set[str] = field(default_factory=set) ended_operation_ids: set[str] = field(default_factory=set) + pending_operation_ends: dict[str, list[OperationEndInfo]] = field( + default_factory=dict + ) @dataclass @@ -66,6 +75,29 @@ class InvocationState(_InvocationState[str | None]): _State = TypeVar("_State", bound=_InvocationState[Any]) +def completed_input_parents(info: InvocationStartInfo) -> list[OperationInfo]: + """Known checkpointed parents of input updates, without guessing missing data.""" + parents: dict[str, OperationInfo] = {} + for operation in info.updated_operations.values(): + parent = ( + info.operations.get(operation.parent_id) if operation.parent_id else None + ) + if ( + parent is not None + and parent.operation_type is OperationType.CONTEXT + and parent.status + in ( + OperationStatus.SUCCEEDED, + OperationStatus.FAILED, + OperationStatus.TIMED_OUT, + OperationStatus.STOPPED, + OperationStatus.CANCELLED, + ) + ): + parents[parent.operation_id] = parent + return list(parents.values()) + + def begin_state( variable: ContextVar[_State], state: _State, 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 fb44a1f45..f2208de8e 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 @@ -68,6 +68,7 @@ from aws_durable_execution_sdk_python_otel._invocation_state import ( ExecutionState, begin_state, + completed_input_parents, finish_state, ) from aws_durable_execution_sdk_python_otel.context_extractors import ( @@ -300,6 +301,14 @@ def _ended_operation_ids(self) -> set[str]: def _ended_operation_ids(self, value: set[str]) -> None: self._state.get().ended_operation_ids = value + @property + def _pending_operation_ends(self) -> dict[str, list[OperationEndInfo]]: + return self._state.get().pending_operation_ends + + @_pending_operation_ends.setter + def _pending_operation_ends(self, value: dict[str, list[OperationEndInfo]]) -> None: + self._state.get().pending_operation_ends = value + def on_registration_result(self, registered: bool) -> None: """Preserve accepted resources, releasing only a discarded constructor.""" if registered: @@ -333,6 +342,9 @@ def _bind_sdk_tracer(self) -> bool: def _set_span(self, key: str, span: Span) -> None: with self._lock: self._operation_spans[key] = span + pending = self._pending_operation_ends.pop(key, []) + for info in pending: + self._render_operation_end(info) def _get_span(self, key: str | None) -> Span | None: if key is None: @@ -633,6 +645,10 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: self._start_workflow_span(info) # Keep the invocation on the shared execution trace. self._start_invocation_span(info) + # A completed child will not run another user-function Start hook. + # Reference its original logical identity without exporting it again. + for parent in completed_input_parents(info): + self._register_operation_placeholder(parent.operation_id, parent.start_time) # Make the Workflow span the active span so auto-instrumented spans # created during the invocation become its children. The token is @@ -758,6 +774,9 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: self._reset_state() return + with self._lock: + unresolved_parents = tuple(self._pending_operation_ends) + # End the invocation span regardless of terminal status. Record the # invocation status and map it to a span status: # SUCCEEDED/PENDING -> OK (this invocation did its work, whether it @@ -801,6 +820,10 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: flush_provider(provider.force_flush) except Exception: # noqa: BLE001 logger.exception("force_flush failed at invocation end") + if unresolved_parents: + raise ValueError( + f"No parent span found for deferred operation ends: {unresolved_parents}" + ) def _reset_state(self) -> None: self._detach_remaining_contexts() @@ -816,6 +839,7 @@ def _reset_state(self) -> None: self._operation_spans = {} self._checkpointed_context_ids = set() self._ended_operation_ids = set() + self._pending_operation_ends = {} self._tracing_enabled = False finish_state(self._state) @@ -836,6 +860,10 @@ def on_operation_start(self, info: OperationStartInfo) -> None: self._note_parent_start(info.parent_id, info.start_time) def on_operation_end(self, info: OperationEndInfo) -> None: + self._render_operation_end(info) + + def _render_operation_end(self, info: OperationEndInfo) -> None: + """Render or defer one delivery without re-entering a public hook.""" logger.debug("Durable operation ended: %s", info) if not self._tracing_enabled: return @@ -844,6 +872,13 @@ def on_operation_end(self, info: OperationEndInfo) -> None: # replay-only and must not re-export the deterministic logical span. if info.is_replayed: return + with self._lock: + if ( + info.parent_id is not None + and info.parent_id not in self._operation_spans + ): + self._pending_operation_ends.setdefault(info.parent_id, []).append(info) + return # Export the span only on the first end for this operation. with self._lock: if info.operation_id in self._ended_operation_ids: 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 a3c6291d6..c23c60864 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 @@ -38,6 +38,7 @@ from aws_durable_execution_sdk_python_otel._invocation_state import ( InvocationState, begin_state, + completed_input_parents, finish_state, ) from aws_durable_execution_sdk_python_otel.context_extractors import ( @@ -637,7 +638,7 @@ def _start_span( # Registration and dequeue share the lock, so an arriving completion # either sees its parent or is drained here. Do not hold it in callbacks. for pending_end in pending_ends: - self.on_operation_end(pending_end) + self._render_operation_end(pending_end) return span def _end_span( @@ -756,6 +757,15 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: name="Invocation", attributes=self._extract_attributes(info), ) + # A cached completed child has no new segment. Keep its checkpointed + # logical parent identity, without creating or re-exporting a parent. + for parent in completed_input_parents(info): + parent_context = self._operation_link_context(parent.operation_id) + if parent_context is not None: + with self._operation_spans_lock: + self._operation_spans[parent.operation_id] = DurableParentSpan( + parent_context, start_time=parent.start_time + ) # Start and End share the invocation worker and token-owning Context. # Retain a valid same-trace ambient parent; otherwise bind this invocation. @@ -964,6 +974,10 @@ def on_operation_end(self, info: OperationEndInfo) -> None: invocation is completing an operation that began earlier, so a short continuation span with a fresh span ID is created and ended. """ + self._render_operation_end(info) + + def _render_operation_end(self, info: OperationEndInfo) -> None: + """Render or defer one delivery without re-entering a public hook.""" logger.debug("Durable operation ended: %s", info) if not self._tracing_enabled: return diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_cached_child_input_delivery_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_cached_child_input_delivery_int.py new file mode 100644 index 000000000..cf8272e2b --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_cached_child_input_delivery_int.py @@ -0,0 +1,225 @@ +"""A terminal input update remains observable when its completed child is skipped.""" + +import json +import logging +import os +from pathlib import Path +import time + +import pytest +from aws_durable_execution_sdk_python import durable_execution, durable_step +from aws_durable_execution_sdk_python.config import CallbackConfig +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from aws_durable_execution_sdk_python.serdes import JsonSerDes +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from aws_durable_execution_sdk_python_testing.stores.filesystem import ( + FileSystemExecutionStore, +) +from aws_durable_execution_sdk_python_testing.stores.memory import ( + InMemoryExecutionStore, +) +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +@pytest.mark.parametrize( + ("view", "filesystem", "failed_delivery"), + [ + (ExecutionOtelPlugin, False, False), + (InvocationOtelPlugin, False, False), + (ExecutionOtelPlugin, True, False), + (InvocationOtelPlugin, True, False), + (ExecutionOtelPlugin, False, True), + ], +) +def test_completed_child_callback_notifies_before_input_ack( + view, filesystem, failed_delivery, tmp_path, monkeypatch, caplog +): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + store = ( + FileSystemExecutionStore.create(tmp_path / "store") + if filesystem + else InMemoryExecutionStore() + ) + starts, ends, notifications, child_calls, marker_calls = [], [], [], [], [] + attempts = [] + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + + class Observe(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + starts.append(info) + + def on_invocation_end(self, info): + ends.append(info) + + def on_operation_end(self, info): + if info.name in ("a", "b"): + notifications.append((len(starts), info)) + + def child(context): + child_calls.append(1) + return context.create_callback( + name="b", config=CallbackConfig(serdes=JsonSerDes()) + ).callback_id + + @durable_step + def marker(_context, result): + marker_calls.append(1) + return result + + def handler(_event, context): + a = context.create_callback( + name="a", config=CallbackConfig(serdes=JsonSerDes()) + ) + b_id = context.run_in_child_context(child, name="cached-child") + result = context.step(marker(a.result()), name="after-a") + return {"a": result, "b_id": b_id} + + wrapped = durable_execution(handler, plugins=[Observe(), plugin]) + + def entrypoint(event, context): + attempts.append(1) + if failed_delivery and len(attempts) == 2: + raise RuntimeError("invocation failed before SDK input delivery") + return wrapped(event, context) + + def wait_for_completed(runner, arn, count): + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + events = runner.get_execution_history( + arn, include_execution_data=True + ).events + if ( + len(ends) >= count + and sum(e.event_type == "InvocationCompleted" for e in events) >= count + ): + return events + time.sleep(0.01) + raise AssertionError("Expected invocation did not finish") + + try: + with DurableFunctionTestRunner( + handler=entrypoint, + store=store, + skip_time=False, + poll_interval=0.01, + execution_timeout=25, + ) as runner: + arn = runner.run_async(input="{}") + a_id = runner.wait_for_callback(arn, name="a", timeout=5) + b_id = runner.wait_for_callback(arn, name="b", timeout=5) + wait_for_completed(runner, arn, 1) + runner.send_callback_success(b_id, result=b'"B"') + # A pre-SDK failure is a physical attempt, not an SDK End hook. + middle_history = wait_for_completed(runner, arn, 2) + middle_state = store.load(arn) + middle_updated_ids = list(middle_state.updated_operation_ids) + runner.send_callback_success(a_id, result=b'"A"') + result = runner.wait_for_result(arn, timeout=10) + final_history = runner.get_execution_history( + arn, include_execution_data=True + ).events + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == {"a": "A", "b_id": b_id} + assert [i.status.value for i in ends] == ["PENDING", "PENDING", "SUCCEEDED"] + assert child_calls == [1] and marker_calls == [1] + assert len(attempts) == (4 if failed_delivery else 3) + b_info = next( + info for info in starts[1].updated_operations.values() if info.name == "b" + ) + parent = starts[1].operations[b_info.parent_id] + assert parent.name == "cached-child" and parent.status.value == "SUCCEEDED" + b_ends = [ + (invocation, info) for invocation, info in notifications if info.name == "b" + ] + terminal_spans = [ + s + for s in exporter.get_finished_spans() + if s.attributes.get("durable.operation.name") == "b" + and s.attributes.get("durable.operation.status") == "SUCCEEDED" + ] + parent_spans = [ + s + for s in exporter.get_finished_spans() + if s.attributes.get("durable.operation.id") == b_info.parent_id + and s.attributes.get("durable.operation.type") == "CONTEXT" + and s.attributes.get("durable.operation.status") == "SUCCEEDED" + ] + hook_errors = [ + r.getMessage() + for r in caplog.records + if r.levelno >= logging.ERROR and ("plugin" in r.name or "otel" in r.name) + ] + proof_dir = os.environ.get("CACHED_CHILD_PROOF_DIR") + if proof_dir: + destination = Path(proof_dir) + destination.mkdir(parents=True, exist_ok=True) + ( + destination + / f"{view.__name__}-{'fs' if filesystem else 'memory'}-failed{failed_delivery}.json" + ).write_text( + json.dumps( + { + "business_succeeded": True, + "result_b_matches_issued": True, + "invocation_statuses": [i.status.value for i in ends], + "child_calls": len(child_calls), + "marker_calls": len(marker_calls), + "physical_attempts": len(attempts), + "failed_before_SDK": failed_delivery, + "input_updated_b": b_info.operation_id, + "input_parent_id": b_info.parent_id, + "parent_status": parent.status.value, + "middle_updated_ids": middle_updated_ids, + "b_notifications": [ + (n, i.operation_id, i.parent_id, i.is_replayed) + for n, i in b_ends + ], + "b_terminal_spans": len(terminal_spans), + "hook_errors": hook_errors, + "original_parent_span_ids": [ + s.context.span_id for s in parent_spans + ], + "b_terminal_parent_ids": [ + s.parent.span_id if s.parent else None + for s in terminal_spans + ], + "middle_history_events": [ + (e.event_id, e.event_type, e.name) for e in middle_history + ], + "final_history_events": [ + (e.event_id, e.event_type, e.name) for e in final_history + ], + }, + indent=2, + ) + + "\n" + ) + assert len(b_ends) == 1 + assert b_ends[0][0] == 2 + assert b_ends[0][1].operation_id == b_info.operation_id + assert b_ends[0][1].parent_id == b_info.parent_id + assert not b_ends[0][1].is_replayed + assert len(terminal_spans) == 1 + assert len(parent_spans) == 1 + assert terminal_spans[0].parent is not None + assert terminal_spans[0].parent.span_id == parent_spans[0].context.span_id + assert not hook_errors + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_terminal_input_parent_refs.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_terminal_input_parent_refs.py new file mode 100644 index 000000000..de922979d --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_terminal_input_parent_refs.py @@ -0,0 +1,181 @@ +"""Missing parent metadata is deferred rather than guessed or reparented.""" + +import contextvars +from datetime import UTC, datetime, timedelta + +import pytest +from aws_durable_execution_sdk_python.plugin import ( + InvocationEndInfo, + InvocationStartInfo, + InvocationStatus, + OperationEndInfo, + OperationInfo, + OperationStartInfo, + OperationStatus, + OperationType, + UserFunctionEndInfo, + UserFunctionStartInfo, +) +from opentelemetry import context as otel_context +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +def _info(cls, now, **changes): + fields = dict( + operation_id="b", + operation_type=OperationType.CALLBACK, + sub_type=None, + name="b", + parent_id="parent", + start_time=now, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=now + timedelta(seconds=1), + ) + fields.update(changes) + return cls(**fields) + + +@pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) +@pytest.mark.parametrize("parent_visible", [False, True]) +def test_nonterminal_or_missing_parent_waits_for_its_actual_span(view, parent_visible): + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + now = datetime.now(UTC) + invocation = dict( + request_id="request", + execution_arn="execution/parent", + is_first_invocation=False, + execution_start_time=now, + ) + b = _info(OperationInfo, now) + parent = _info( + OperationInfo, + now, + operation_id="parent", + operation_type=OperationType.CONTEXT, + name="parent", + parent_id=None, + status=OperationStatus.STARTED, + end_time=None, + ) + operations = {"b": b} + if parent_visible: + operations["parent"] = parent + + def invoke(): + before = otel_context.get_current() + plugin.on_invocation_start( + InvocationStartInfo( + **invocation, operations=operations, updated_operations={"b": b} + ) + ) + plugin.on_operation_end(_info(OperationEndInfo, now)) + assert not [ + s + for s in exporter.get_finished_spans() + if s.attributes.get("durable.operation.id") == "b" + ] + parent_fields = dict( + operation_id="parent", + operation_type=OperationType.CONTEXT, + name="parent", + parent_id=None, + status=OperationStatus.STARTED, + is_replayed=True, + end_time=None, + ) + plugin.on_operation_start(_info(OperationStartInfo, now, **parent_fields)) + start = _info(UserFunctionStartInfo, now, **parent_fields) + plugin.on_user_function_start(start) + actual_parent = plugin._get_span("parent").get_span_context() + plugin.on_user_function_end(UserFunctionEndInfo.from_start_info(start, None)) + plugin.on_operation_end( + _info( + OperationEndInfo, + now, + operation_id="parent", + operation_type=OperationType.CONTEXT, + name="parent", + parent_id=None, + ) + ) + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + spans = [ + s + for s in exporter.get_finished_spans() + if s.attributes.get("durable.operation.id") == "b" + ] + assert len(spans) == 1 and spans[0].parent.span_id == actual_parent.span_id + assert otel_context.get_current() == before + assert plugin._state not in contextvars.copy_context() + + try: + contextvars.Context().run(invoke) + finally: + provider.shutdown() + + +@pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_unavailable_parent_is_reported_after_context_cleanup(view): + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + now = datetime.now(UTC) + invocation = dict( + request_id="request", + execution_arn="execution/unknown-parent", + is_first_invocation=False, + execution_start_time=now, + ) + b = _info(OperationInfo, now) + + def invoke(): + before = otel_context.get_current() + plugin.on_invocation_start( + InvocationStartInfo( + **invocation, operations={"b": b}, updated_operations={"b": b} + ) + ) + plugin.on_operation_end(_info(OperationEndInfo, now)) + with pytest.raises(ValueError, match="No parent span found for deferred"): + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.PENDING) + ) + assert not [ + s + for s in exporter.get_finished_spans() + if s.attributes.get("durable.operation.id") == "b" + ] + assert otel_context.get_current() == before + assert plugin._state not in contextvars.copy_context() + + try: + contextvars.Context().run(invoke) + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python/README.md b/packages/aws-durable-execution-sdk-python/README.md index 4b2463e52..a2b9486ea 100644 --- a/packages/aws-durable-execution-sdk-python/README.md +++ b/packages/aws-durable-execution-sdk-python/README.md @@ -122,6 +122,14 @@ checkpointing is available, stops and waits for checkpoint processing, then call End before returning the outcome. The caller shuts down the handler executor. There is no separate handler-context hook or context-manager plugin API. +After all Start hooks, visible terminal operations named by the invocation +input's `UpdatedOperationIds` receive their existing `on_operation_end` callback. +This includes descendants that a cached completed child will not revisit. +Input delivery and later checkpoint/replay observations share the per-invocation +notification ledger; an operation/status is not reported twice in that invocation. +Missing operation metadata is not fabricated, and failed invocation delivery can +be redelivered on a later attempt. This is not an exporter acknowledgment. + When plugins are registered, the worker begins with a copy of the caller's context-variable bindings. Successful Start bindings flow into later Start hooks, the handler, serialization and resource cleanup. A later successful Start may 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 fbd1c53d5..195588296 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 @@ -416,9 +416,11 @@ def on_operation_start(self, info: OperationStartInfo) -> None: def on_operation_end(self, info: OperationEndInfo) -> None: """ - Called when an operation reaches a terminal status. Terminal operations - are not emitted again during replay. Child contexts without a terminal - checkpoint may emit this from the thread that runs the operation. + Called when an operation reaches a terminal status. Visible terminal + input updates are reported after all invocation-start hooks, even if + replay skips their completed child context. Duplicate deliveries of an + operation/status within one invocation are suppressed. Child contexts + without a terminal checkpoint may emit this from their operation thread. Args: info: Information about the operation. @@ -573,6 +575,8 @@ def _run_in_invocation_context(self, invoke: Callable[[], Any]) -> Any: def _snapshot_operation_infos( self, operations_provider: Callable[[], Mapping[str, Operation]] | None, + *, + raw_snapshot: dict[str, Operation] | None = None, ) -> dict[str, OperationInfo]: """Build the plugin ``OperationInfo`` view of the current operation map. @@ -593,7 +597,11 @@ def _snapshot_operation_infos( if not self._plugins or operations_provider is None: return {} try: - return _to_operation_info_map(operations_provider()) + source = dict(operations_provider()) + snapshot = _to_operation_info_map(source) + if raw_snapshot is not None: + raw_snapshot.update(source) + return snapshot except Exception: # A plugin-facing view must never break the execution. logger.exception("Failed to snapshot operations for plugin hook") @@ -624,7 +632,10 @@ def on_invocation_start( """ aws_request_id = lambda_context.aws_request_id if lambda_context else None self._operations_provider = operations_provider if self._plugins else None - operations = self._snapshot_operation_infos(operations_provider) + raw_operations: dict[str, Operation] = {} + operations = self._snapshot_operation_infos( + operations_provider, raw_snapshot=raw_operations + ) self._invocation_status = InvocationStartInfo( execution_arn=execution_arn, request_id=aws_request_id, @@ -638,7 +649,38 @@ def on_invocation_start( if operation_id in operations }, ) + # Capture delivery data before exposing the mutable snapshot maps to + # Start hooks. A cached child may never visit its updated descendant. + input_ends = [ + OperationEndInfo( + operation_id=operation.operation_id, + operation_type=_to_operation_type(operation.operation_type), + sub_type=operation.sub_type, + name=operation.name, + parent_id=operation.parent_id, + start_time=operation.start_timestamp, + end_time=operation.end_timestamp, + result=_extract_result(operation), + error=self._extract_error(operation), + attempt=( + operation.step_details.attempt if operation.step_details else None + ), + is_replayed=False, + status=operation.status, + ) + for operation_id in (updated_operation_ids or []) + if operation_id in raw_operations + for operation in (raw_operations[operation_id],) + if self._is_terminal_status(operation.status) + ] self.execute_plugins(self._invocation_status, sync=True) + if input_ends: + + def report_input_updates() -> None: + for operation_end in input_ends: + self._report_terminal_update(operation_end) + + self._run_in_invocation_context(report_input_updates) def _snapshot_execution_input(self, execution_input: Any) -> Any: """Deep-copy the execution input so the plugin view is isolated. @@ -827,13 +869,7 @@ def on_operation_update( # Replay delivery and checkpoint responses can report the same # completion. Deduplicate actual notifications, not state that # may have arrived without an UpdatedOperationIds notification. - if self._plugins: - key = (operation.operation_id, operation.status) - with self._terminal_updates_lock: - if key in self._reported_terminal_updates: - continue - self._reported_terminal_updates.add(key) - self.execute_plugins( + self._report_terminal_update( OperationEndInfo( operation_id=operation.operation_id, operation_type=_to_operation_type(operation.operation_type), @@ -852,7 +888,6 @@ def on_operation_update( ), is_replayed=False, ), - sync=True, ) if ( @@ -883,6 +918,16 @@ def on_operation_update( sync=True, ) + def _report_terminal_update(self, info: OperationEndInfo) -> None: + """Claim actual input/checkpoint delivery once in this invocation.""" + if self._plugins: + key = (info.operation_id, info.status) + with self._terminal_updates_lock: + if key in self._reported_terminal_updates: + return + self._reported_terminal_updates.add(key) + self.execute_plugins(info, sync=True) + @staticmethod def _extract_error(operation: Operation): return _extract_error(operation) diff --git a/packages/aws-durable-execution-sdk-python/tests/plugin_terminal_input_test.py b/packages/aws-durable-execution-sdk-python/tests/plugin_terminal_input_test.py new file mode 100644 index 000000000..363a984c4 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python/tests/plugin_terminal_input_test.py @@ -0,0 +1,132 @@ +"""Visible terminal input updates notify plugins before cached user paths skip them.""" + +import contextvars +from datetime import UTC, datetime + +from aws_durable_execution_sdk_python.execution import ( + DurableExecutionInvocationOutput, + InvocationStatus, +) +from aws_durable_execution_sdk_python.lambda_service import ( + CallbackDetails, + ErrorObject, + Operation, + OperationStatus, + OperationType, +) +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + PluginExecutor, +) + + +def _callback(operation_id, status=OperationStatus.SUCCEEDED, error=None): + return Operation( + operation_id=operation_id, + operation_type=OperationType.CALLBACK, + status=status, + name=operation_id, + callback_details=CallbackDetails( + callback_id=operation_id, result="saved", error=error + ), + ) + + +def _start(executor, operations, updated_ids): + executor.on_invocation_start( + execution_arn="execution/input-delivery", + is_first_invocation=False, + execution_start_time=datetime.now(UTC), + lambda_context=None, + operations_provider=lambda: operations, + updated_operation_ids=updated_ids, + ) + + +def test_input_delivery_follows_all_starts_and_shares_actual_notification_ledger(): + marker = contextvars.ContextVar("input-owner", default="host") + events = [] + + class Observe(DurableInstrumentationPlugin): + def __init__(self, name): + self.name = name + + def on_invocation_start(self, info): + marker.set(self.name) + events.append(("start", self.name)) + + def on_operation_end(self, info): + events.append( + ("end", self.name, info.operation_id, marker.get(), info.result) + ) + + b, active = _callback("b"), _callback("active", OperationStatus.STARTED) + operations = {"b": b, "active": active} + executor = PluginExecutor([Observe("first"), Observe("second")]) + + def invoke(): + with executor.run(): + _start(executor, operations, ["b", "active", "not-visible"]) + assert events == [ + ("start", "first"), + ("start", "second"), + ("end", "first", "b", "second", "saved"), + ("end", "second", "b", "second", "saved"), + ] + executor.on_operation_update(b) + assert len(events) == 4 + # Missing input metadata did not claim an event that arrives later. + executor.on_operation_update(_callback("not-visible")) + assert [event[2] for event in events if event[0] == "end"] == [ + "b", + "b", + "not-visible", + "not-visible", + ] + executor.on_invocation_end( + DurableExecutionInvocationOutput(status=InvocationStatus.PENDING) + ) + + contextvars.Context().run(invoke) + assert marker.get() == "host" + events.clear() + # A new invocation can receive the same explicit update after redelivery; + # the ledger is not a cross-invocation exporter acknowledgment. + contextvars.Context().run(invoke) + + +def test_failed_start_and_snapshot_mutation_cannot_lose_input_notification(): + marker = contextvars.ContextVar("input-failed-start", default="host") + error = ErrorObject( + message="saved error", type="TestError", data=None, stack_trace=["original"] + ) + operation = _callback("b", OperationStatus.FAILED, error) + seen = [] + + class MutatingFailure(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + marker.set("poison") + info.updated_operations["b"].error.stack_trace.append("mutated") + info.updated_operations.clear() + raise RuntimeError("intentional Start failure") + + class Observe(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + assert marker.get() == "host" + marker.set("survivor") + + def on_operation_end(self, info): + seen.append((info.operation_id, marker.get(), list(info.error.stack_trace))) + + def invoke(): + executor = PluginExecutor([MutatingFailure(), Observe()]) + with executor.run(): + _start(executor, {"b": operation}, ["b"]) + executor.on_invocation_end( + DurableExecutionInvocationOutput(status=InvocationStatus.PENDING) + ) + + contextvars.Context().run(invoke) + assert seen == [("b", "survivor", ["original"])] + assert operation.callback_details.error.stack_trace == ["original"] + assert marker.get() == "host" From 40e4fa6ec889c7a03e269079d8429cf2f9b7e629 Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:39:36 +0000 Subject: [PATCH 52/54] fix(otel): retain late callback parent context --- .../README.md | 8 +- .../_invocation_state.py | 1 + .../execution_plugin.py | 9 +- .../invocation_plugin.py | 32 ++- .../tests/e2e/test_late_child_parent_int.py | 202 ++++++++++++++++++ .../tests/test_terminal_input_parent_refs.py | 125 +++++++++++ 6 files changed, 365 insertions(+), 12 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_late_child_parent_int.py diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index cd4de725a..65177e2ca 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -449,8 +449,12 @@ instance resources. SDK-owned worker tasks capture the selected invocation Context; Start/End hooks and token reset retain their existing owner and order. Terminal input updates can be observed even when replay skips their completed -child context. When the input contains that checkpointed parent's terminal -metadata, the views reference its original logical span identity without +child context. A child that finishes during the current invocation retains its +actual parent context for later descendant checkpoint deliveries. In Invocation +view, input completions first wait for the handler to create any current replay +segment, including `ReplayChildren` for large results. At invocation End, only +completions still lacking that segment use a known completed input parent. When the input contains that checkpointed parent's +terminal metadata, the views reference its original logical span identity without re-exporting the parent. Missing or nonterminal parents are deferred until their actual span is available; unresolved parentage is reported after cleanup instead of inventing a root parent. Draining this private queue does not call the public diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py index cbece37fb..d8a5418ce 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py @@ -65,6 +65,7 @@ class ExecutionState(_InvocationState[str]): @dataclass class InvocationState(_InvocationState[str | None]): span_time_floor_ns: int | None = None + input_parent_refs: dict[str, OperationInfo] = field(default_factory=dict) pending_operation_ends: dict[str, list[OperationEndInfo]] = field( default_factory=dict ) 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 f2208de8e..ce58427b0 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 @@ -939,7 +939,14 @@ def _render_operation_end(self, info: OperationEndInfo) -> None: span.set_status(StatusCode.OK) self._note_parent_end(info.parent_id, end_time) - popped = self._pop_span(info.operation_id) + with self._lock: + popped = self._operation_spans.pop(info.operation_id, None) + if popped is not None and info.operation_type is OperationType.CONTEXT: + # A callback created by this child can finish in a later root + # checkpoint. Retain only its actual parent context until End. + self._operation_spans[info.operation_id] = DurableParentSpan( + popped.get_span_context(), start_time=start_time + ) if popped is not None: popped.end(end_time=_to_otel_timestamp(end_time)) 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 c23c60864..ad919106f 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 @@ -757,15 +757,11 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: name="Invocation", attributes=self._extract_attributes(info), ) - # A cached completed child has no new segment. Keep its checkpointed - # logical parent identity, without creating or re-exporting a parent. - for parent in completed_input_parents(info): - parent_context = self._operation_link_context(parent.operation_id) - if parent_context is not None: - with self._operation_spans_lock: - self._operation_spans[parent.operation_id] = DurableParentSpan( - parent_context, start_time=parent.start_time - ) + # A terminal checkpoint can still use ReplayChildren. Give the handler + # a chance to register its current segment before using an old identity. + self._state.get().input_parent_refs = { + parent.operation_id: parent for parent in completed_input_parents(info) + } # Start and End share the invocation worker and token-owning Context. # Retain a valid same-trace ambient parent; otherwise bind this invocation. @@ -858,6 +854,24 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: self._reset_state() return + # Actual replay parent segments drain their children on registration. + # Only now can an input completion with no such segment use the known + # completed checkpoint's identity. Missing/nonterminal parents still fail. + pending_ends = [] + with self._operation_spans_lock: + for parent_id in tuple(self._pending_operation_ends): + parent = self._state.get().input_parent_refs.get(parent_id) + if parent is None or parent_id in self._operation_spans: + continue + parent_context = self._operation_link_context(parent_id) + if parent_context is not None: + self._operation_spans[parent_id] = DurableParentSpan( + parent_context, start_time=parent.start_time + ) + pending_ends.extend(self._pending_operation_ends.pop(parent_id)) + for pending_end in pending_ends: + self._render_operation_end(pending_end) + # User work and checkpoint cleanup have finished. Release Start bindings # in the Context that created their tokens before closing the spans. self._detach_remaining_contexts() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_late_child_parent_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_late_child_parent_int.py new file mode 100644 index 000000000..f2d82dbbc --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_late_child_parent_int.py @@ -0,0 +1,202 @@ +"""Late callback completions keep the actual parent for their invocation.""" + +import json +import logging +import os +import queue +from pathlib import Path +import threading +import time + +import pytest +from aws_durable_execution_sdk_python import durable_execution +from aws_durable_execution_sdk_python.config import CallbackConfig +from aws_durable_execution_sdk_python.plugin import DurableInstrumentationPlugin +from aws_durable_execution_sdk_python.serdes import JsonSerDes +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +@pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) +@pytest.mark.parametrize("large_result", [False, True]) +def test_late_callback_keeps_completed_or_replayed_parent( + view, large_result, monkeypatch, caplog +): + monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = view( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + starts, ends, child_calls, saved_calls, root_calls = [], [], [], [], [] + root_step_started, release_root_step = threading.Event(), threading.Event() + callback_ids = queue.Queue(maxsize=1) + + class Observe(DurableInstrumentationPlugin): + def on_invocation_start(self, info): + starts.append(info) + + def on_invocation_end(self, info): + ends.append(info) + + def saved(_step): + saved_calls.append(1) + return "saved" + + def child(ctx): + child_calls.append(1) + cb = ctx.create_callback(name="b", config=CallbackConfig(serdes=JsonSerDes())) + assert ctx.step(saved, name="save-child") == "saved" + return dict( + b_id=cb.callback_id, large="x" * (300 * 1024) if large_result else "small" + ) + + def root_step(a_id, b_id): + callback_ids.put((a_id, b_id)) + root_calls.append(1) + root_step_started.set() + assert release_root_step.wait(10) + return "root-saved" + + def handler(_event, ctx): + a = ctx.create_callback(name="a", config=CallbackConfig(serdes=JsonSerDes())) + result = ctx.run_in_child_context(child, name="returned-child") + if not large_result: + assert ( + ctx.step( + lambda _: root_step(a.callback_id, result["b_id"]), + name="root-checkpoint", + ) + == "root-saved" + ) + return dict(a=a.result(), b_id=result["b_id"], length=len(result["large"])) + + wrapped = durable_execution(handler, plugins=[Observe(), plugin]) + + def completed(runner, arn, count): + deadline = time.monotonic() + 12 + while time.monotonic() < deadline: + history = runner.get_execution_history( + arn, include_execution_data=True + ).events + if ( + len(ends) >= count + and sum(e.event_type == "InvocationCompleted" for e in history) >= count + ): + return history + time.sleep(0.01) + raise AssertionError("invocation completion not observed") + + def spans(name, status=None): + return [ + s + for s in exporter.get_finished_spans() + if s.attributes.get("durable.operation.name") == name + and ( + status is None or s.attributes.get("durable.operation.status") == status + ) + ] + + try: + with DurableFunctionTestRunner( + handler=wrapped, skip_time=False, poll_interval=0.01, execution_timeout=35 + ) as runner: + arn = runner.run_async(input="{}") + if large_result: + a_id = runner.wait_for_callback(arn, name="a", timeout=5) + b_id = runner.wait_for_callback(arn, name="b", timeout=5) + completed(runner, arn, 1) + else: + # The runner wait_for_callback helper waits for suspension; + # this case deliberately completes during the active root step. + a_id, b_id = callback_ids.get(timeout=5) + assert root_step_started.wait(5) + runner.send_callback_success(b_id, result=b'"B"') + release_root_step.set() + completed(runner, arn, 2 if large_result else 1) + middle_callbacks = spans("b", "SUCCEEDED") + middle_parents = spans("returned-child", "SUCCEEDED") + invocation_spans = [ + s for s in exporter.get_finished_spans() if s.name == "Invocation" + ] + runner.send_callback_success(a_id, result=b'"A"') + result = runner.wait_for_result(arn, timeout=10) + history = runner.get_execution_history( + arn, include_execution_data=True + ).events + expected_invocations = 3 if large_result else 2 + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == dict( + a="A", b_id=b_id, length=300 * 1024 if large_result else 5 + ) + assert len(starts) == expected_invocations and [ + x.status.value for x in ends + ] == ["PENDING"] * (expected_invocations - 1) + ["SUCCEEDED"] + assert saved_calls == [1] and root_calls == ([] if large_result else [1]) + assert len(child_calls) == (3 if large_result else 1) + hook_errors = [ + r.getMessage() + for r in caplog.records + if r.levelno >= logging.ERROR and ("plugin" in r.name or "otel" in r.name) + ] + proof = dict( + view=view.__name__, + large_result=large_result, + statuses=[x.status.value for x in ends], + saved_step_calls=len(saved_calls), + root_step_calls=len(root_calls), + child_calls=len(child_calls), + middle_callback_spans=[ + dict(span_id=s.context.span_id, parent_id=s.parent.span_id) + for s in middle_callbacks + ], + middle_parent_spans=[ + dict(span_id=s.context.span_id, parent_id=s.parent.span_id) + for s in middle_parents + ], + middle_invocation_ids=[s.context.span_id for s in invocation_spans], + final_callback_spans=len(spans("b", "SUCCEEDED")), + hook_errors=hook_errors, + history=[(e.event_id, e.event_type, e.name) for e in history], + ) + if os.environ.get("LATE_PARENT_PROOF_DIR"): + out = Path(os.environ["LATE_PARENT_PROOF_DIR"]) + out.mkdir(parents=True, exist_ok=True) + (out / f"{view.__name__}-large{large_result}.json").write_text( + json.dumps(proof, indent=2) + "\n" + ) + assert len(middle_callbacks) == 1 + assert len(spans("b", "SUCCEEDED")) == 1 + if large_result and view is InvocationOtelPlugin: + assert len(middle_parents) == 2 + assert ( + middle_parents[-1].parent.span_id + == invocation_spans[-1].context.span_id + ) + assert ( + middle_callbacks[0].parent.span_id == middle_parents[-1].context.span_id + ) + assert ( + middle_callbacks[0].parent.span_id != middle_parents[0].context.span_id + ) + else: + assert ( + middle_callbacks[0].parent.span_id == middle_parents[0].context.span_id + ) + assert not hook_errors + finally: + release_root_step.set() + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_terminal_input_parent_refs.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_terminal_input_parent_refs.py index de922979d..a152cd571 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_terminal_input_parent_refs.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_terminal_input_parent_refs.py @@ -179,3 +179,128 @@ def invoke(): contextvars.Context().run(invoke) finally: provider.shutdown() + + +def test_completed_context_remains_parent_for_late_execution_checkpoint(): + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = ExecutionOtelPlugin( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + now = datetime.now(UTC) + invocation = dict( + request_id="request", + execution_arn="execution/late-parent", + is_first_invocation=True, + execution_start_time=now, + ) + parent_fields = dict( + operation_id="parent", + operation_type=OperationType.CONTEXT, + name="parent", + parent_id=None, + status=OperationStatus.STARTED, + end_time=None, + ) + before = otel_context.get_current() + try: + plugin.on_invocation_start(InvocationStartInfo(**invocation)) + plugin.on_operation_start(_info(OperationStartInfo, now, **parent_fields)) + start = _info(UserFunctionStartInfo, now, **parent_fields) + plugin.on_user_function_start(start) + plugin.on_operation_start( + _info( + OperationStartInfo, now, status=OperationStatus.STARTED, end_time=None + ) + ) + plugin.on_user_function_end(UserFunctionEndInfo.from_start_info(start, None)) + plugin.on_operation_end( + _info( + OperationEndInfo, + now, + operation_id="parent", + operation_type=OperationType.CONTEXT, + name="parent", + parent_id=None, + ) + ) + parents = [s for s in exporter.get_finished_spans() if s.name == "parent"] + assert len(parents) == 1 + plugin.on_operation_end(_info(OperationEndInfo, now)) + children = [s for s in exporter.get_finished_spans() if s.name == "b"] + assert ( + len(children) == 1 + and children[0].parent.span_id == parents[0].context.span_id + ) + assert ( + len([s for s in exporter.get_finished_spans() if s.name == "parent"]) == 1 + ) + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.SUCCEEDED) + ) + assert plugin._operation_spans == {} and otel_context.get_current() == before + finally: + provider.shutdown() + + +def test_completed_input_parent_waits_for_current_invocation_segment(): + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = InvocationOtelPlugin( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + now = datetime.now(UTC) + invocation = dict( + request_id="request", + execution_arn="execution/replayed-parent", + is_first_invocation=False, + execution_start_time=now, + ) + parent_fields = dict( + operation_id="parent", + operation_type=OperationType.CONTEXT, + name="parent", + parent_id=None, + is_replayed=True, + ) + b = _info(OperationInfo, now) + parent = _info(OperationInfo, now, **parent_fields) + before = otel_context.get_current() + try: + plugin.on_invocation_start( + InvocationStartInfo( + **invocation, + operations={"parent": parent, "b": b}, + updated_operations={"b": b}, + ) + ) + plugin.on_operation_end(_info(OperationEndInfo, now)) + assert not [s for s in exporter.get_finished_spans() if s.name == "b"] + start = _info( + UserFunctionStartInfo, now, **parent_fields, is_replay_children=True + ) + plugin.on_user_function_start(start) + actual_parent = plugin._get_span("parent").get_span_context() + children = [s for s in exporter.get_finished_spans() if s.name == "b"] + assert ( + len(children) == 1 and children[0].parent.span_id == actual_parent.span_id + ) + plugin.on_user_function_end(UserFunctionEndInfo.from_start_info(start, None)) + plugin.on_operation_end(_info(OperationEndInfo, now, **parent_fields)) + plugin.on_invocation_end( + InvocationEndInfo(**invocation, status=InvocationStatus.PENDING) + ) + assert plugin._operation_spans == {} and otel_context.get_current() == before + assert plugin._state.get().input_parent_refs == {} + finally: + provider.shutdown() From 7fd24a134ed83500400aa3bba1b86c4d0bea0805 Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Wed, 14 Oct 2026 23:06:48 +0000 Subject: [PATCH 53/54] fix(otel): retain input parent metadata for updates --- .../README.md | 4 +- .../_invocation_state.py | 32 +++++------- .../tests/e2e/test_late_child_parent_int.py | 44 +++++++++++++--- .../tests/test_terminal_input_parent_refs.py | 50 +++++++++++++++++++ 4 files changed, 103 insertions(+), 27 deletions(-) diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 65177e2ca..008e8766d 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -449,7 +449,9 @@ instance resources. SDK-owned worker tasks capture the selected invocation Context; Start/End hooks and token reset retain their existing owner and order. Terminal input updates can be observed even when replay skips their completed -child context. A child that finishes during the current invocation retains its +child context. Known completed context metadata is retained from the full input +snapshot, including parents whose descendants are still pending at Start and +complete in a later checkpoint. A child that finishes during the current invocation retains its actual parent context for later descendant checkpoint deliveries. In Invocation view, input completions first wait for the handler to create any current replay segment, including `ReplayChildren` for large results. At invocation End, only diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py index d8a5418ce..b4212541b 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/_invocation_state.py @@ -77,26 +77,20 @@ class InvocationState(_InvocationState[str | None]): def completed_input_parents(info: InvocationStartInfo) -> list[OperationInfo]: - """Known checkpointed parents of input updates, without guessing missing data.""" - parents: dict[str, OperationInfo] = {} - for operation in info.updated_operations.values(): - parent = ( - info.operations.get(operation.parent_id) if operation.parent_id else None + """Known terminal contexts that may parent input or later checkpoint updates.""" + return [ + operation + for operation in info.operations.values() + if operation.operation_type is OperationType.CONTEXT + and operation.status + in ( + OperationStatus.SUCCEEDED, + OperationStatus.FAILED, + OperationStatus.TIMED_OUT, + OperationStatus.STOPPED, + OperationStatus.CANCELLED, ) - if ( - parent is not None - and parent.operation_type is OperationType.CONTEXT - and parent.status - in ( - OperationStatus.SUCCEEDED, - OperationStatus.FAILED, - OperationStatus.TIMED_OUT, - OperationStatus.STOPPED, - OperationStatus.CANCELLED, - ) - ): - parents[parent.operation_id] = parent - return list(parents.values()) + ] def begin_state( diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_late_child_parent_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_late_child_parent_int.py index f2d82dbbc..fa7fec189 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_late_child_parent_int.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_late_child_parent_int.py @@ -26,9 +26,12 @@ @pytest.mark.parametrize("view", [ExecutionOtelPlugin, InvocationOtelPlugin]) -@pytest.mark.parametrize("large_result", [False, True]) +@pytest.mark.parametrize( + ("large_result", "resume_before_completion"), + [(False, False), (True, False), (False, True)], +) def test_late_callback_keeps_completed_or_replayed_parent( - view, large_result, monkeypatch, caplog + view, large_result, resume_before_completion, monkeypatch, caplog ): monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) provider = TracerProvider() @@ -74,6 +77,11 @@ def root_step(a_id, b_id): def handler(_event, ctx): a = ctx.create_callback(name="a", config=CallbackConfig(serdes=JsonSerDes())) result = ctx.run_in_child_context(child, name="returned-child") + if resume_before_completion: + gate = ctx.create_callback( + name="resume-gate", config=CallbackConfig(serdes=JsonSerDes()) + ) + assert gate.result() == "resume" if not large_result: assert ( ctx.step( @@ -120,13 +128,19 @@ def spans(name, status=None): b_id = runner.wait_for_callback(arn, name="b", timeout=5) completed(runner, arn, 1) else: + if resume_before_completion: + gate_id = runner.wait_for_callback( + arn, name="resume-gate", timeout=5 + ) + completed(runner, arn, 1) + runner.send_callback_success(gate_id, result=b'"resume"') # The runner wait_for_callback helper waits for suspension; # this case deliberately completes during the active root step. a_id, b_id = callback_ids.get(timeout=5) assert root_step_started.wait(5) runner.send_callback_success(b_id, result=b'"B"') release_root_step.set() - completed(runner, arn, 2 if large_result else 1) + completed(runner, arn, 2 if large_result or resume_before_completion else 1) middle_callbacks = spans("b", "SUCCEEDED") middle_parents = spans("returned-child", "SUCCEEDED") invocation_spans = [ @@ -137,7 +151,7 @@ def spans(name, status=None): history = runner.get_execution_history( arn, include_execution_data=True ).events - expected_invocations = 3 if large_result else 2 + expected_invocations = 3 if large_result or resume_before_completion else 2 assert result.status.value == "SUCCEEDED" assert json.loads(result.result) == dict( a="A", b_id=b_id, length=300 * 1024 if large_result else 5 @@ -147,6 +161,15 @@ def spans(name, status=None): ] == ["PENDING"] * (expected_invocations - 1) + ["SUCCEEDED"] assert saved_calls == [1] and root_calls == ([] if large_result else [1]) assert len(child_calls) == (3 if large_result else 1) + if resume_before_completion: + assert all( + info.name != "b" for info in starts[1].updated_operations.values() + ) + pending_b = next( + info for info in starts[1].operations.values() if info.name == "b" + ) + assert pending_b.status.value == "STARTED" + assert starts[1].operations[pending_b.parent_id].status.value == "SUCCEEDED" hook_errors = [ r.getMessage() for r in caplog.records @@ -155,10 +178,16 @@ def spans(name, status=None): proof = dict( view=view.__name__, large_result=large_result, + resume_before_completion=resume_before_completion, statuses=[x.status.value for x in ends], saved_step_calls=len(saved_calls), root_step_calls=len(root_calls), child_calls=len(child_calls), + middle_input_updated_names=[ + info.name for info in starts[1].updated_operations.values() + ] + if resume_before_completion + else None, middle_callback_spans=[ dict(span_id=s.context.span_id, parent_id=s.parent.span_id) for s in middle_callbacks @@ -175,9 +204,10 @@ def spans(name, status=None): if os.environ.get("LATE_PARENT_PROOF_DIR"): out = Path(os.environ["LATE_PARENT_PROOF_DIR"]) out.mkdir(parents=True, exist_ok=True) - (out / f"{view.__name__}-large{large_result}.json").write_text( - json.dumps(proof, indent=2) + "\n" - ) + ( + out + / f"{view.__name__}-large{large_result}-resume{resume_before_completion}.json" + ).write_text(json.dumps(proof, indent=2) + "\n") assert len(middle_callbacks) == 1 assert len(spans("b", "SUCCEEDED")) == 1 if large_result and view is InvocationOtelPlugin: diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_terminal_input_parent_refs.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_terminal_input_parent_refs.py index a152cd571..9531e6178 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_terminal_input_parent_refs.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_terminal_input_parent_refs.py @@ -304,3 +304,53 @@ def test_completed_input_parent_waits_for_current_invocation_segment(): assert plugin._state.get().input_parent_refs == {} finally: provider.shutdown() + + +def test_completed_parent_refs_cover_later_updates_without_guessing(): + from aws_durable_execution_sdk_python_otel._invocation_state import ( + completed_input_parents, + ) + + now = datetime.now(UTC) + parent = _info( + OperationInfo, + now, + operation_id="parent", + operation_type=OperationType.CONTEXT, + name="parent", + parent_id=None, + ) + unfinished = _info( + OperationInfo, + now, + operation_id="unfinished", + operation_type=OperationType.CONTEXT, + status=OperationStatus.STARTED, + ) + step = _info( + OperationInfo, now, operation_id="step", operation_type=OperationType.STEP + ) + invocation = dict( + request_id="request", + execution_arn="execution/full-parent", + is_first_invocation=False, + execution_start_time=now, + ) + assert completed_input_parents( + InvocationStartInfo( + **invocation, + operations={"parent": parent, "unfinished": unfinished, "step": step}, + updated_operations={}, + ) + ) == [parent] + missing_parent_child = _info(OperationInfo, now) + assert ( + completed_input_parents( + InvocationStartInfo( + **invocation, + operations={"b": missing_parent_child}, + updated_operations={"b": missing_parent_child}, + ) + ) + == [] + ) From 4654b348e83f0cf039f07acc28f16bf6ccd723f7 Mon Sep 17 00:00:00 2001 From: Frank Chen <65260095+zhongkechen@users.noreply.github.com> Date: Thu, 15 Oct 2026 00:33:10 +0000 Subject: [PATCH 54/54] fix(testing): match service replay-state pruning --- .../README.md | 30 +- .../test_cached_child_input_delivery_int.py | 43 ++- .../tests/e2e/test_late_child_parent_int.py | 12 +- .../README.md | 25 ++ .../checkpoint/processor.py | 5 +- .../checkpoint/processors/base.py | 8 +- .../execution.py | 21 +- .../invoker.py | 10 +- .../model.py | 7 +- .../operation_state.py | 64 +++++ ...replay_children_service_parity_int_test.py | 256 ++++++++++++++++++ .../replay_children_service_parity_test.py | 239 ++++++++++++++++ .../e2e/plugin_operation_type_int_test.py | 12 +- 13 files changed, 691 insertions(+), 41 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/operation_state.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/tests/e2e/replay_children_service_parity_int_test.py create mode 100644 packages/aws-durable-execution-sdk-python-testing/tests/replay_children_service_parity_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 008e8766d..a7a07a772 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -448,19 +448,23 @@ cleanup clearing another's trace. Provider configuration and registration remain instance resources. SDK-owned worker tasks capture the selected invocation Context; Start/End hooks and token reset retain their existing owner and order. -Terminal input updates can be observed even when replay skips their completed -child context. Known completed context metadata is retained from the full input -snapshot, including parents whose descendants are still pending at Start and -complete in a later checkpoint. A child that finishes during the current invocation retains its -actual parent context for later descendant checkpoint deliveries. In Invocation -view, input completions first wait for the handler to create any current replay -segment, including `ReplayChildren` for large results. At invocation End, only -completions still lacking that segment use a known completed input parent. When the input contains that checkpointed parent's -terminal metadata, the views reference its original logical span identity without -re-exporting the parent. Missing or nonterminal parents are deferred until their -actual span is available; unresolved parentage is reported after cleanup instead -of inventing a root parent. Draining this private queue does not call the public -operation-end hook again. +Visible terminal input updates are delivered after all Start hooks. The service +omits descendants of completed contexts unless their ancestors retain children +through `ReplayChildren`; an `UpdatedOperationIds` entry alone does not provide +missing operation metadata. The SDK does not synthesize a hook or a span from +that missing record. Public input-delivery tests use a real large child result +to retain metadata, while still asserting one saved-step execution and terminal +export with the correct parent in each view. + +Known completed context metadata is retained from the full visible input +snapshot for descendant completions delivered by later checkpoint responses. +A child that finishes during the invocation retains its actual non-recording +parent context until cleanup. In Invocation view, completions wait for any +current replay segment, including a large-result ReplayChildren parent. Only +at End does a completion still lacking a current segment use a known completed +input parent's identity. Missing or nonterminal parents remain deferred; +unresolved parentage is reported after cleanup. No root parent or missing +metadata is invented, and queue draining does not re-enter a public hook. Completion metadata can be redelivered after a failed invocation or checkpoint transport attempt that has no successful response acknowledgment. There is no diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_cached_child_input_delivery_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_cached_child_input_delivery_int.py index cf8272e2b..dfc1c0254 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_cached_child_input_delivery_int.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_cached_child_input_delivery_int.py @@ -1,4 +1,4 @@ -"""A terminal input update remains observable when its completed child is skipped.""" +"""Visible terminal input updates are delivered once with real retained children.""" import json import logging @@ -39,7 +39,7 @@ (ExecutionOtelPlugin, False, True), ], ) -def test_completed_child_callback_notifies_before_input_ack( +def test_retained_child_callback_notifies_before_input_ack( view, filesystem, failed_delivery, tmp_path, monkeypatch, caplog ): monkeypatch.delenv("DURABLE_EXECUTION_PLUGINS", raising=False) @@ -50,6 +50,7 @@ def test_completed_child_callback_notifies_before_input_ack( ) starts, ends, notifications, child_calls, marker_calls = [], [], [], [], [] attempts = [] + saved_child_calls = [] provider = TracerProvider() exporter = InMemorySpanExporter() provider.add_span_processor(SimpleSpanProcessor(exporter)) @@ -72,11 +73,20 @@ def on_operation_end(self, info): if info.name in ("a", "b"): notifications.append((len(starts), info)) + def saved_child(_step): + saved_child_calls.append(1) + return "saved" + def child(context): child_calls.append(1) - return context.create_callback( + callback_id = context.create_callback( name="b", config=CallbackConfig(serdes=JsonSerDes()) ).callback_id + assert context.step(saved_child, name="saved-child") == "saved" + # The real service hides descendants of a completed ordinary context. + # A real large result makes this context use ReplayChildren, retaining + # B's metadata without injecting input state or hook notifications. + return {"id": callback_id, "large": "x" * (300 * 1024)} @durable_step def marker(_context, result): @@ -87,7 +97,7 @@ def handler(_event, context): a = context.create_callback( name="a", config=CallbackConfig(serdes=JsonSerDes()) ) - b_id = context.run_in_child_context(child, name="cached-child") + b_id = context.run_in_child_context(child, name="cached-child")["id"] result = context.step(marker(a.result()), name="after-a") return {"a": result, "b_id": b_id} @@ -138,7 +148,8 @@ def wait_for_completed(runner, arn, count): assert result.status.value == "SUCCEEDED" assert json.loads(result.result) == {"a": "A", "b_id": b_id} assert [i.status.value for i in ends] == ["PENDING", "PENDING", "SUCCEEDED"] - assert child_calls == [1] and marker_calls == [1] + assert child_calls == [1, 1, 1] and saved_child_calls == [1] + assert marker_calls == [1] assert len(attempts) == (4 if failed_delivery else 3) b_info = next( info for info in starts[1].updated_operations.values() if info.name == "b" @@ -180,6 +191,8 @@ def wait_for_completed(runner, arn, count): "result_b_matches_issued": True, "invocation_statuses": [i.status.value for i in ends], "child_calls": len(child_calls), + "saved_child_calls": len(saved_child_calls), + "retention_via_public_large_result": True, "marker_calls": len(marker_calls), "physical_attempts": len(attempts), "failed_before_SDK": failed_delivery, @@ -193,7 +206,7 @@ def wait_for_completed(runner, arn, count): ], "b_terminal_spans": len(terminal_spans), "hook_errors": hook_errors, - "original_parent_span_ids": [ + "parent_segment_span_ids": [ s.context.span_id for s in parent_spans ], "b_terminal_parent_ids": [ @@ -217,9 +230,23 @@ def wait_for_completed(runner, arn, count): assert b_ends[0][1].parent_id == b_info.parent_id assert not b_ends[0][1].is_replayed assert len(terminal_spans) == 1 - assert len(parent_spans) == 1 assert terminal_spans[0].parent is not None - assert terminal_spans[0].parent.span_id == parent_spans[0].context.span_id + if view is ExecutionOtelPlugin: + assert len(parent_spans) == 1 + assert terminal_spans[0].parent.span_id == parent_spans[0].context.span_id + else: + assert len(parent_spans) == 3 + assert len({span.context.span_id for span in parent_spans}) == 3 + invocations = [ + span + for span in exporter.get_finished_spans() + if span.name == "Invocation" + ] + assert len(invocations) == 3 + assert [span.parent.span_id for span in parent_spans] == [ + span.context.span_id for span in invocations + ] + assert terminal_spans[0].parent.span_id == parent_spans[1].context.span_id assert not hook_errors finally: provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_late_child_parent_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_late_child_parent_int.py index fa7fec189..6845f2884 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_late_child_parent_int.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_late_child_parent_int.py @@ -165,11 +165,15 @@ def spans(name, status=None): assert all( info.name != "b" for info in starts[1].updated_operations.values() ) - pending_b = next( - info for info in starts[1].operations.values() if info.name == "b" + # Service input prunes B even while it is STARTED. Its later + # checkpoint response still carries the actual completion unchanged. + assert all(info.name != "b" for info in starts[1].operations.values()) + completed_parent = next( + info + for info in starts[1].operations.values() + if info.name == "returned-child" ) - assert pending_b.status.value == "STARTED" - assert starts[1].operations[pending_b.parent_id].status.value == "SUCCEEDED" + assert completed_parent.status.value == "SUCCEEDED" hook_errors = [ r.getMessage() for r in caplog.records diff --git a/packages/aws-durable-execution-sdk-python-testing/README.md b/packages/aws-durable-execution-sdk-python-testing/README.md index 58b21e793..fc31dbefc 100644 --- a/packages/aws-durable-execution-sdk-python-testing/README.md +++ b/packages/aws-durable-execution-sdk-python-testing/README.md @@ -170,6 +170,31 @@ flow, pagination, and a map of the code. ## Documentation +### Operation state on replay + +Invocation input and `GetDurableExecutionState` follow the service's retained +operation view. An operation is omitted when any ancestor is a `CONTEXT` that +has `SUCCEEDED` or `FAILED` without `ReplayChildren: true`. The walk includes +running intermediate contexts; a nested context's own retention flag cannot +override a completed ancestor that prunes it. Roots and operations whose parent +is absent are kept, and ordering does not change the rule. + +`ContextOptions.ReplayChildren` on START is ignored. A closing update records +true only when it explicitly supplies true; otherwise it records false. Testing +service responses include that closing false, while Core's public model and +generic serialization contract remain unchanged. + +Checkpoint responses still contain all updated operations. `UpdatedOperationIds` +can include an ID whose operation record is hidden from invocation input and +GetState. Raw stores, history, results and nested operation inspection retain the +complete execution. An updated ID alone does not restore missing metadata for a +plugin. A real large child result retains its children through ReplayChildren +when a test needs that public replay behavior. + +This corrects local service parity; it does not fix duplicate execution-view +spans for virtual FLAT map iterations while their map is still running. + + ### Error Handling The testing framework implements AWS-compliant error responses that match the exact format expected by boto3 and AWS services. For detailed information about error response formats, exception types, and troubleshooting, see: diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processor.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processor.py index 42ea9333b..f1b1d6da6 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processor.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processor.py @@ -21,6 +21,9 @@ CheckpointRequestDispatcher, ) from aws_durable_execution_sdk_python_testing.clock import RealClock +from aws_durable_execution_sdk_python_testing.operation_state import ( + prune_replay_children, +) from aws_durable_execution_sdk_python_testing.exceptions import ( InvalidParameterValueException, ) @@ -153,7 +156,7 @@ def get_execution_state( token = CheckpointToken.from_str(checkpoint_token) execution = self._store.load(token.execution_arn) return StateOutput( - operations=execution.get_navigable_operations(), + operations=prune_replay_children(execution.get_navigable_operations()), next_marker=None, ) diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processors/base.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processors/base.py index d2eaa4ec4..9ea3cb256 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processors/base.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/checkpoint/processors/base.py @@ -100,9 +100,11 @@ def _create_context_details(self, update: OperationUpdate) -> ContextDetails | N ContextDetails( result=update.payload, error=update.error, - replay_children=update.context_options.replay_children - if update.context_options - else False, + replay_children=( + update.action in (OperationAction.SUCCEED, OperationAction.FAIL) + and update.context_options is not None + and update.context_options.replay_children is True + ), ) if update.operation_type == OperationType.CONTEXT else None diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py index 0a16afac1..833081dcd 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/execution.py @@ -24,6 +24,9 @@ ) from aws_durable_execution_sdk_python_testing.clock import real_now +from aws_durable_execution_sdk_python_testing.operation_state import ( + prune_replay_children, +) from aws_durable_execution_sdk_python_testing.exceptions import ( IllegalStateException, InvalidParameterValueException, @@ -854,10 +857,9 @@ def page( ``max_size_bytes`` and, when given, ``max_items``. Second element of the tuple is a marker for the next page, or ``None`` when the page fits everything.""" - start_idx = self._resolve_marker(marker) - return self._walk_page( - self.snapshot_operations, start_idx, max_size_bytes, max_items - ) + operations = prune_replay_children(self.snapshot_operations) + start_idx = self._resolve_marker(marker, len(operations)) + return self._walk_page(operations, start_idx, max_size_bytes, max_items) def unseen_operations(self) -> list[Operation]: """Operations whose ``operation_last_touched_seq`` is strictly @@ -919,14 +921,21 @@ def _size_for(self, op: Operation) -> int: # floor at 1 to guarantee pagination always advances. return max(self.execution.operation_size_bytes.get(op.operation_id, 0), 1) - def _resolve_marker(self, marker: str | None) -> int: + def _resolve_marker( + self, marker: str | None, operation_count: int | None = None + ) -> int: if marker is None: return 0 seq, idx = self._decode_marker(marker) if seq != self.pinned_token_sequence: msg = "Invalid marker" raise InvalidParameterValueException(msg) - if idx < 0 or idx > len(self.snapshot_operations): + count = ( + len(self.snapshot_operations) + if operation_count is None + else operation_count + ) + if idx < 0 or idx > count: msg = "Invalid marker" raise InvalidParameterValueException(msg) return idx diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/invoker.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/invoker.py index 47e7ac5ed..69070cbb6 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/invoker.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/invoker.py @@ -25,6 +25,9 @@ ResourceNotFoundException, ) from aws_durable_execution_sdk_python_testing.execution import OperationPaginatorState +from aws_durable_execution_sdk_python_testing.operation_state import ( + operation_state_to_dict, +) from aws_durable_execution_sdk_python_testing.model import ( LambdaContext, executed_version, @@ -472,10 +475,15 @@ def invoke( input.durable_execution_arn, endpoint_url, region_name ) + invocation_payload = input.to_json_dict() + invocation_payload["InitialExecutionState"]["Operations"] = [ + operation_state_to_dict(operation, json_timestamps=True) + for operation in input.initial_execution_state.operations + ] invoke_kwargs: dict[str, Any] = { "FunctionName": function_name, "InvocationType": "RequestResponse", # Synchronous invocation - "Payload": json.dumps(input.to_json_dict()), + "Payload": json.dumps(invocation_payload), } if tenant_id is not None: invoke_kwargs["TenantId"] = tenant_id diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/model.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/model.py index 1de6b425c..617a40089 100644 --- a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/model.py +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/model.py @@ -1,6 +1,9 @@ """Model classes for the web API.""" from __future__ import annotations +from aws_durable_execution_sdk_python_testing.operation_state import ( + operation_state_to_dict, +) import datetime @@ -572,7 +575,7 @@ def from_dict(cls, data: dict) -> GetDurableExecutionStateResponse: def to_dict(self) -> dict[str, Any]: result: dict[str, Any] = { - "Operations": [op.to_dict() for op in self.operations] + "Operations": [operation_state_to_dict(op) for op in self.operations] } if self.next_marker is not None: result["NextMarker"] = self.next_marker @@ -3157,7 +3160,7 @@ def from_dict(cls, data: dict) -> CheckpointUpdatedExecutionState: def to_dict(self) -> dict[str, Any]: result: dict[str, Any] = { - "Operations": [op.to_dict() for op in self.operations] + "Operations": [operation_state_to_dict(op) for op in self.operations] } if self.next_marker is not None: result["NextMarker"] = self.next_marker diff --git a/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/operation_state.py b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/operation_state.py new file mode 100644 index 000000000..850eea00b --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/src/aws_durable_execution_sdk_python_testing/operation_state.py @@ -0,0 +1,64 @@ +"""Service views of retained operation state; raw execution storage stays complete.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +from aws_durable_execution_sdk_python.lambda_service import ( + Operation, + OperationStatus, + OperationType, +) + + +_CLOSED_CONTEXT = (OperationStatus.SUCCEEDED, OperationStatus.FAILED) + + +def prune_replay_children(operations: Sequence[Operation]) -> list[Operation]: + """Omit a descendant when any known completed context prunes its children.""" + by_id = {operation.operation_id: operation for operation in operations} + visible = [] + for operation in operations: + parent_id = operation.parent_id + visited = {operation.operation_id} + hidden = False + while parent_id is not None and parent_id not in visited: + parent = by_id.get(parent_id) + if parent is None: + break + if ( + parent.operation_type is OperationType.CONTEXT + and parent.status in _CLOSED_CONTEXT + and not ( + parent.context_details + and parent.context_details.replay_children is True + ) + ): + hidden = True + break + visited.add(parent_id) + parent_id = parent.parent_id + if not hidden: + visible.append(operation) + return visible + + +def operation_state_to_dict( + operation: Operation, *, json_timestamps: bool = False +) -> dict[str, Any]: + """Project service ReplayChildren presence without changing Core's model API.""" + data = dict(operation.to_json_dict() if json_timestamps else operation.to_dict()) + if ( + operation.operation_type is OperationType.CONTEXT + and operation.context_details is not None + ): + details = dict(data.get("ContextDetails", {})) + if operation.status is OperationStatus.STARTED: + details.pop("ReplayChildren", None) + elif operation.status in _CLOSED_CONTEXT: + details["ReplayChildren"] = ( + operation.context_details.replay_children is True + ) + data["ContextDetails"] = details + return data diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/e2e/replay_children_service_parity_int_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/replay_children_service_parity_int_test.py new file mode 100644 index 000000000..3c31a9412 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/e2e/replay_children_service_parity_int_test.py @@ -0,0 +1,256 @@ +"""Real public handlers replay through the service's retained-operation view.""" + +import json +import os +from pathlib import Path + +import pytest +from aws_durable_execution_sdk_python import durable_execution +from aws_durable_execution_sdk_python.config import CallbackConfig +from aws_durable_execution_sdk_python.exceptions import ChildContextError +from aws_durable_execution_sdk_python.serdes import JsonSerDes +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from aws_durable_execution_sdk_python_testing.stores.memory import ( + InMemoryExecutionStore, +) +from aws_durable_execution_sdk_python_testing.stores.filesystem import ( + FileSystemExecutionStore, +) + + +@pytest.mark.parametrize("filesystem", [False, True]) +@pytest.mark.parametrize("mode", ["completed", "retained", "nested-retained", "failed"]) +def test_public_resume_uses_service_state_but_inspection_remains_complete( + filesystem, mode, tmp_path +): + store = ( + FileSystemExecutionStore.create(tmp_path / "state") + if filesystem + else InMemoryExecutionStore() + ) + snapshots = [] + child_calls = [] + step_calls = [] + + def save(_step): + step_calls.append(1) + return "saved" + + def child(ctx): + child_calls.append(1) + saved = ctx.step(save, name="saved-child-step") + if mode == "failed": + raise ValueError("stored-child-failure") + return { + "saved": saved, + "large": "x" * (300 * 1024) + if mode in ["retained", "nested-retained"] + else "", + } + + def outer(ctx): + value = ctx.run_in_child_context(child, name="inner-retained") + return {"saved": value["saved"], "large": ""} + + def handler(_event, ctx): + try: + value = ctx.run_in_child_context( + outer if mode == "nested-retained" else child, name="outer" + ) + result = {"saved": value["saved"], "length": len(value["large"])} + except ChildContextError: + assert mode == "failed" + result = {"failed": True} + gate = ctx.create_callback( + name="resume", config=CallbackConfig(serdes=JsonSerDes()) + ) + assert gate.result() == "resume" + return result + + wrapped = durable_execution(handler) + + def entrypoint(event, context): + initial = event.initial_execution_state.operations + state = event.service_client.get_execution_state( + event.durable_execution_arn, event.checkpoint_token, "" + ) + snapshots.append( + { + "input_names": [op.name for op in initial], + "state_names": [op.name for op in state.operations], + "input_ids": [op.operation_id for op in initial], + "state_ids": [op.operation_id for op in state.operations], + "updated_ids": list(event.updated_operation_ids), + } + ) + return wrapped(event, context) + + with DurableFunctionTestRunner( + handler=entrypoint, + store=store, + skip_time=False, + poll_interval=0.01, + execution_timeout=20, + ) as runner: + arn = runner.run_async(input="{}") + callback_id = runner.wait_for_callback(arn, name="resume", timeout=5) + runner.send_callback_success(callback_id, result=b'"resume"') + result = runner.wait_for_result(arn, timeout=10) + history = runner.get_execution_history(arn, include_execution_data=True).events + assert result.status.value == "SUCCEEDED" + assert json.loads(result.result) == ( + {"failed": True} + if mode == "failed" + else {"saved": "saved", "length": 300 * 1024 if mode == "retained" else 0} + ) + assert ( + len(snapshots) == 2 and snapshots[0]["input_ids"] == snapshots[0]["state_ids"] + ) + assert snapshots[1]["input_ids"] == snapshots[1]["state_ids"] + visible = "saved-child-step" in snapshots[1]["input_names"] + assert visible is (mode == "retained") + if mode == "nested-retained": + assert "inner-retained" not in snapshots[1]["input_names"] + assert "outer" in snapshots[1]["input_names"] + assert step_calls == [1] and len(child_calls) == (2 if mode == "retained" else 1) + raw = store.load(arn) + assert any(op.name == "saved-child-step" for op in raw.operations) + inspected = result.get_operation_by_name("outer") + if mode == "nested-retained": + inspected = inspected.get_operation_by_name("inner-retained") + assert ( + inspected.get_operation_by_name("saved-child-step").name == "saved-child-step" + ) + assert ( + sum( + e.event_type == "StepSucceeded" and e.name == "saved-child-step" + for e in history + ) + == 1 + ) + assert any(e.name == "saved-child-step" for e in history) + if os.environ.get("SERVICE_PRUNING_PROOF_DIR"): + out = Path(os.environ["SERVICE_PRUNING_PROOF_DIR"]) + out.mkdir(parents=True, exist_ok=True) + (out / f"{mode}-fs{filesystem}.json").write_text( + json.dumps( + dict( + mode=mode, + filesystem=filesystem, + snapshots=snapshots, + result=json.loads(result.result), + step_calls=len(step_calls), + child_calls=len(child_calls), + history_names=[e.name for e in history], + ), + indent=2, + ) + + "\n" + ) + + +@pytest.mark.parametrize("filesystem", [False, True]) +def test_hidden_callback_updated_id_does_not_make_its_metadata_visible( + filesystem, tmp_path +): + import time + + store = ( + FileSystemExecutionStore.create(tmp_path / "hidden-state") + if filesystem + else InMemoryExecutionStore() + ) + snapshots = [] + child_calls = [] + saved_calls = [] + + def saved(_step): + saved_calls.append(1) + return "saved" + + def child(ctx): + child_calls.append(1) + callback = ctx.create_callback( + name="hidden-b", config=CallbackConfig(serdes=JsonSerDes()) + ) + assert ctx.step(saved, name="hidden-saved") == "saved" + return callback.callback_id + + def handler(_event, ctx): + a = ctx.create_callback( + name="root-a", config=CallbackConfig(serdes=JsonSerDes()) + ) + b_id = ctx.run_in_child_context(child, name="closed-parent") + return {"a": a.result(), "b_id": b_id} + + wrapped = durable_execution(handler) + + def entrypoint(event, context): + state = event.service_client.get_execution_state( + event.durable_execution_arn, event.checkpoint_token, "" + ) + snapshots.append( + { + "input_ids": [ + op.operation_id for op in event.initial_execution_state.operations + ], + "state_ids": [op.operation_id for op in state.operations], + "updated_ids": list(event.updated_operation_ids), + } + ) + return wrapped(event, context) + + def completed(runner, arn, count): + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + events = runner.get_execution_history( + arn, include_execution_data=True + ).events + if sum(e.event_type == "InvocationCompleted" for e in events) >= count: + return events + time.sleep(0.01) + raise AssertionError("invocation completion not observed") + + with DurableFunctionTestRunner( + handler=entrypoint, + store=store, + skip_time=False, + poll_interval=0.01, + execution_timeout=25, + ) as runner: + arn = runner.run_async(input="{}") + a_id = runner.wait_for_callback(arn, name="root-a", timeout=5) + b_id = runner.wait_for_callback(arn, name="hidden-b", timeout=5) + first = completed(runner, arn, 1) + b_operation_id = next( + e.operation_id + for e in first + if e.name == "hidden-b" and e.event_type == "CallbackStarted" + ) + runner.send_callback_success(b_id, result=b'"B"') + completed(runner, arn, 2) + runner.send_callback_success(a_id, result=b'"A"') + result = runner.wait_for_result(arn, timeout=10) + history = runner.get_execution_history(arn, include_execution_data=True).events + assert result.status.value == "SUCCEEDED" and json.loads(result.result) == { + "a": "A", + "b_id": b_id, + } + assert len(snapshots) == 3 and child_calls == [1] and saved_calls == [1] + assert b_operation_id in snapshots[1]["updated_ids"] + assert b_operation_id not in snapshots[1]["input_ids"] + assert b_operation_id not in snapshots[1]["state_ids"] + assert any(op.operation_id == b_operation_id for op in store.load(arn).operations) + assert ( + result.get_operation_by_name("closed-parent") + .get_operation_by_name("hidden-b") + .result + == '"B"' + ) + assert ( + sum( + e.event_type == "CallbackSucceeded" and e.operation_id == b_operation_id + for e in history + ) + == 1 + ) diff --git a/packages/aws-durable-execution-sdk-python-testing/tests/replay_children_service_parity_test.py b/packages/aws-durable-execution-sdk-python-testing/tests/replay_children_service_parity_test.py new file mode 100644 index 000000000..65302dbec --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-testing/tests/replay_children_service_parity_test.py @@ -0,0 +1,239 @@ +"""The service's retained-operation tree applies only to invocation/state reads.""" + +from datetime import UTC, datetime +from unittest.mock import Mock + +import pytest +from aws_durable_execution_sdk_python.lambda_service import ( + ContextDetails, + ContextOptions, + Operation, + OperationAction, + OperationStatus, + OperationType, + OperationUpdate, +) +from aws_durable_execution_sdk_python_testing.execution import ( + Execution, + OperationPaginatorState, +) +from aws_durable_execution_sdk_python_testing.invoker import ( + InProcessInvoker, + LambdaInvoker, +) +from aws_durable_execution_sdk_python_testing.model import ( + StartDurableExecutionInput, + GetDurableExecutionStateResponse, + CheckpointUpdatedExecutionState, +) +from aws_durable_execution_sdk_python_testing.checkpoint.processor import ( + CheckpointProcessor, +) +from aws_durable_execution_sdk_python_testing.checkpoint.processors.context import ( + ContextProcessor, +) +from aws_durable_execution_sdk_python_testing.stores.memory import ( + InMemoryExecutionStore, +) +from aws_durable_execution_sdk_python_testing.stores.filesystem import ( + FileSystemExecutionStore, +) + + +def step(identifier, parent=None, status=OperationStatus.SUCCEEDED): + return Operation( + operation_id=identifier, + parent_id=parent, + operation_type=OperationType.STEP, + status=status, + ) + + +def ctx(identifier, parent=None, status=OperationStatus.SUCCEEDED, replay=None): + return Operation( + operation_id=identifier, + parent_id=parent, + operation_type=OperationType.CONTEXT, + status=status, + context_details=ContextDetails.from_dict( + {} if replay is None else {"ReplayChildren": replay} + ), + ) + + +TREE = [ + step("1-step"), + ctx("2-context", status=OperationStatus.STARTED), + step("2.1-step", "2-context", OperationStatus.PENDING), + step("3-step", "missing"), + ctx("4-context", status=OperationStatus.FAILED), + ctx("4.1-context", "4-context", OperationStatus.FAILED, True), + step("4.1.1-step", "4.1-context", OperationStatus.FAILED), + step("4.1.2-step", "4.1-context", OperationStatus.FAILED), + ctx("5-context", replay=True), + ctx("5.1-context", "5-context", replay=False), + step("5.1.1-step", "5.1-context"), + step("5.1-step", "5-context"), + ctx("map"), + ctx("running-branch", "map", OperationStatus.STARTED), + step("pending-branch-step", "running-branch", OperationStatus.PENDING), + ctx("retained-inside-pruned", "running-branch", replay=True), + step("nested-step", "retained-inside-pruned"), +] +KEPT = [ + "1-step", + "2-context", + "2.1-step", + "3-step", + "4-context", + "5-context", + "5.1-context", + "5.1-step", + "map", +] + + +def execution_with(operations): + start = StartDurableExecutionInput( + account_id="123456789012", + function_name="parity", + function_qualifier="$LATEST", + execution_name="parity", + execution_timeout_seconds=60, + execution_retention_period_days=1, + invocation_id="invocation", + ) + execution = Execution( + "arn:aws:lambda:us-west-2:123456789012:function:parity:$LATEST/durable-execution/parity/id", + start, + list(operations), + ) + for op in operations: + execution.touch_operation(op.operation_id) + execution.updated_operation_ids = ["4.1.1-step", "pending-branch-step"] + return execution + + +@pytest.mark.parametrize("filesystem", [False, True]) +@pytest.mark.parametrize("reverse", [False, True]) +def test_service_tree_paging_inputs_and_get_state_keep_raw_inspection( + filesystem, reverse, tmp_path +): + operations = list(reversed(TREE)) if reverse else list(TREE) + expected = list(reversed(KEPT)) if reverse else KEPT + execution = execution_with(operations) + store = ( + FileSystemExecutionStore.create(tmp_path / "state") + if filesystem + else InMemoryExecutionStore() + ) + store.save(execution) + restored = store.load(execution.durable_execution_arn) + paginator = OperationPaginatorState.pin(restored) + paged = [] + marker = None + while True: + page, marker = paginator.page(marker, 100000, max_items=2) + paged.extend(op.operation_id for op in page) + if marker is None: + break + assert paged == expected + for invoker in [InProcessInvoker(Mock(), Mock()), LambdaInvoker(Mock())]: + invocation = invoker.create_invocation_input(restored) + assert [ + o.operation_id for o in invocation.initial_execution_state.operations + ] == expected + assert invocation.updated_operation_ids == ["4.1.1-step", "pending-branch-step"] + state = CheckpointProcessor(store, Mock()).get_execution_state( + restored.get_new_checkpoint_token(), "" + ) + assert [o.operation_id for o in state.operations] == expected + # Checkpoint deltas, raw store, inspection, and hidden updated IDs remain full. + assert paginator.unseen_operations() == operations + assert ( + restored.operations == operations + and restored.get_assertable_operations() == operations + ) + assert restored.updated_operation_ids == ["4.1.1-step", "pending-branch-step"] + + +@pytest.mark.parametrize("action", [OperationAction.SUCCEED, OperationAction.FAIL]) +@pytest.mark.parametrize( + "close_options", + [ + None, + ContextOptions(replay_children=True), + ContextOptions(replay_children=False), + ContextOptions(), + ], +) +def test_replay_children_is_recorded_only_by_closing_update(action, close_options): + processor = ContextProcessor() + now = datetime.now(UTC) + notifier = Mock() + start = OperationUpdate( + operation_id="ctx", + operation_type=OperationType.CONTEXT, + action=OperationAction.START, + context_options=ContextOptions(replay_children=True), + ) + opened = processor.process(start, None, notifier, "execution", now) + assert ( + opened.context_details is not None + and opened.context_details.replay_children is False + ) + assert "ReplayChildren" not in opened.to_json_dict()["ContextDetails"] + close = OperationUpdate( + operation_id="ctx", + operation_type=OperationType.CONTEXT, + action=action, + context_options=close_options, + payload="saved", + ) + completed = processor.process(close, opened, notifier, "execution", now) + expected = close_options is not None and close_options.replay_children is True + assert completed.context_details.replay_children is expected + assert ( + GetDurableExecutionStateResponse([completed]).to_dict()["Operations"][0][ + "ContextDetails" + ]["ReplayChildren"] + is expected + ) + assert ( + CheckpointUpdatedExecutionState([completed]).to_dict()["Operations"][0][ + "ContextDetails" + ]["ReplayChildren"] + is expected + ) + assert completed.context_details.result == "saved" + + +@pytest.mark.parametrize( + "operations", [[step("a", "a")], [step("a", "b"), step("b", "a")]] +) +def test_invalid_cycles_terminate_without_dropping_unrelated_operations(operations): + execution = execution_with(operations) + assert OperationPaginatorState.pin(execution).page(None, 100000)[0] == operations + + +def test_lambda_invocation_projects_close_false_and_start_absence(): + from io import BytesIO + import json + + execution = execution_with(TREE) + client = Mock() + client.invoke.return_value = { + "StatusCode": 200, + "Payload": BytesIO(b'{"Status":"SUCCEEDED","Result":"null"}'), + } + invoker = LambdaInvoker(client) + invocation = invoker.create_invocation_input(execution) + invoker.invoke("target", invocation) + payload = json.loads(client.invoke.call_args.kwargs["Payload"]) + operations = {op["Id"]: op for op in payload["InitialExecutionState"]["Operations"]} + assert list(operations) == KEPT + assert "ReplayChildren" not in operations["2-context"]["ContextDetails"] + assert operations["4-context"]["ContextDetails"]["ReplayChildren"] is False + assert operations["5-context"]["ContextDetails"]["ReplayChildren"] is True + assert operations["5.1-context"]["ContextDetails"]["ReplayChildren"] is False + assert payload["UpdatedOperationIds"] == ["4.1.1-step", "pending-branch-step"] diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_operation_type_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_operation_type_int_test.py index f4fca018c..c1ee85d63 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_operation_type_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/plugin_operation_type_int_test.py @@ -120,10 +120,16 @@ def handler(_event: Any, context: DurableContext) -> str: assert {info.sub_type for info in plugin.operation_starts} == set(expected) assert {info.sub_type for info in plugin.operation_ends} == set(expected) - # The final replay starts with every earlier operation checkpointed, - # including the execution root, whose service record has no subtype. + # Live hooks above cover every category. The service prunes descendants of + # the completed callback/map/parallel contexts from the replay input; the + # surviving records, including the subtype-less execution root, keep type. replay_operations = plugin.invocation_starts[-1].operations.values() - assert {info.sub_type for info in replay_operations} == {*expected, None} + retained_subtypes = set(expected) - { + OperationSubType.CALLBACK, + OperationSubType.MAP_ITERATION, + OperationSubType.PARALLEL_BRANCH, + } + assert {info.sub_type for info in replay_operations} == {*retained_subtypes, None} all_infos: list[OperationInfo] = [ *plugin.operation_starts, *plugin.operation_ends,