fix(otel): balance context attach and detach across plugin lifecycles

[wangyb-A]

Fixes #643

Problem

Both OTel plugins called opentelemetry.context.attach() without keeping the returned token, and "restored" the enclosing span by attaching another context rather than detaching. Every operation therefore pushed two context layers and popped none, leaving an ended span current after its scope had finished.

The worst consequence is not in the issue. The invocation-start attach (execution_plugin.py:219) ran on the Lambda handler thread, which is reused across warm invocations, so the next execution's context extractor (context_extractors.py:27) and GLOBAL-mode ambient-parent lookup (execution_plugin.py:252) adopted the previous execution's ended Workflow span. A valid parent overrides the deterministic ID generator's trace ID, so two unrelated durable executions merged into a single trace. Measured with a runtime probe before the fix:

execution 1 trace_id                 = 65937d25c4f40d011ac92bea7d425ca1
execution 2 Invocation span trace_id = 65937d25c4f40d011ac92bea7d425ca1   <- polluted

and after:

execution 1 trace_id                 = 00000000000000000000000000000000  (no span current)
execution 2 Invocation span trace_id = 65937d25006c2aae27b0d4a02a5926f2  <- own trace

Two runtime properties shaped the fix. First, the hooks run on several threads — the invocation hooks on the Lambda handler thread, the user-function hooks on the dex-handler worker that runs user code, and on a branch worker per map/parallel branch — and ContextVar.reset() only accepts a token created in the same contextvars.Context. Second, unlike Java's ScopeImpl.close(), which ignores a close that does not represent the current context, ContextVar.reset() writes back its captured value unconditionally, so detaching out of order revives a stale context instead of failing safe.

Changes

• New context_scope.py — a module-level, thread-confined LIFO stack of attach tokens. exit_scope unwinds downwards so the underlying ContextVar is always reset in order. The stack is module level rather than per plugin instance because both plugins ship as separate entry points and can be enabled together: hooks dispatch in registration order, so the second plugin's scope must come off while the first plugin's end hook runs. An epoch check discards scopes a suspended operation left behind, since the SDK re-raises SuspendExecution without calling on_user_function_end (state.py:1171).
• Paired every user-function attach with a detach on the same thread, replacing the re-attach that previously stood in for restoring the enclosing context. A nested operation now lands back on its parent's still-attached scope; a top-level one lands back on the thread's ambient context.
• Unwind any remaining scopes at invocation end, so the handler thread is left exactly as the plugin found it.
• Dropped the invocation-start attach entirely. User code runs on a separate worker and ThreadPoolExecutor does not copy contextvars, so that attach never reached the code it was meant to parent — it only leaked. The Workflow and Invocation spans are used as explicit parents instead, matching the Java plugins, which never make either span current.

Behaviour change

Ambient auto-instrumented spans emitted outside any operation — for example directly in the handler between two steps — are no longer parented to the Invocation span. In an ADOT deployment they attach to the ambient Lambda invocation span instead, which the previous invocation-start attach was shadowing. This matches ExecutionOtelPlugin/InvocationOtelPlugin in the Java SDK, which cover the same window with MDC rather than an attached span. Calls inside a step or child context are unaffected.

Log correlation is unchanged: the logging filter resolves through the plugin's span registry, so records emitted between operations still carry the invocation's traceId and spanId. One visible detail — handler-thread records now carry the Invocation span's spanId instead of the Workflow span's (same traceId), since the registry prefers the Invocation span; there is a test pinning this.

The durable span hierarchy itself is untouched: parents, links, and deterministic IDs are all chosen explicitly in _start_span, never taken from the ambient context.

Acceptance criteria

• Every context.attach() token owned by the plugin has a corresponding context.detach() — with one documented exception below.
• Invocation cleanup restores the context that was active before plugin invocation start — now true by construction on the handler thread, since the plugin attaches nothing there, plus an explicit unwind.
• User-function cleanup restores the exact enclosing context without accumulating stale scopes — 2 attaches / 0 detaches per operation becomes 1 / 1.
• Tests no longer require global context resets to hide plugin lifecycle leaks — the autouse fixtures now assert that each test leaves the context exactly as it found it.
• Tests cover nested child contexts, multiple sequential steps, failures, and warm invocation reuse.

Documented exception to the first criterion: scopes attached on a worker thread that suspends cannot be detached from the handler thread where invocation end runs, because a token is only resettable in the context that created it. Those threads are created per invocation and their ContextVar dies with them, and the epoch check discards any leftover if a thread is ever reused. The Java plugins have the same gap — their invocation-end sweep iterates an unordered map from the handler thread, so those closes hit ScopeImpl's guard and are ignored.

Testing

129 tests pass in the otel package, 3223 across the monorepo; hatch fmt --check and hatch run types:check clean.

New coverage: tests/test_context_scope.py (13 tests — LIFO nesting, unwind-above-target, unknown-key no-op, epoch discard, thread confinement, suspension, worker-thread hooks, both plugins on one thread), a warm-reuse test asserting two executions land in separate traces, a pre/post invocation context-restore test, and a log-filter test pinning the handler-thread spanId.

Both halves were mutation-tested. Restoring the invocation-start attach fails 11 tests including the warm-reuse and context-restore ones; skipping the detach in on_user_function_end fails 10 across the balance fixtures and the restore tests.

Four tests that asserted "the invocation span is current again after a step" were rewritten to assert exact context restore instead, since that behaviour came from the unbalanced re-attach. Three nested-context tests kept their assertions unchanged and needed only their lifecycle completed — they started a child context and never ended it, which the old reset fixture silently swallowed. Notably test_get_current_span_context_returns_invocation_span_between_steps passes untouched, which is the evidence that log correlation survived.

Follow-ups (not in this PR)

• wrap_user_function re-raises SuspendExecution without calling on_user_function_end (state.py:1171), so a suspended operation's hooks are structurally unpaired for every plugin, not just these two. Worth a core-side fix or an explicit suspension hook.
• No conformance requirement asserts that two executions in a warm container get distinct trace IDs. Per AGENTS.md that belongs upstream in aws-durable-execution-conformance-tests first.

[github-actions]

Codex AI review

[P2] Preserve existing scope owners when the other bundled plugin attaches. The second plugin rebases its context onto _extracted_context, so _owner_ids(context) contains only that plugin. With both entry points enabled, the first plugin no longer trusts the active attempt span and the installed log filter falls back to its Invocation span, losing operation-level correlation. Merge the current owner IDs into the produced context after cleanup, and test logging with both actual plugins enabled.

[github-actions]

Codex AI review

[P2] This unhandled case is reachable when a map/parallel branch suspends, another branch frees a concurrency slot, and the executor reuses the suspended branch's thread for a different branch. The different key and unchanged epoch bypass both cleanup checks, so the new scope inherits the suspended branch's context and restores it afterward, misattributing subsequent logs and instrumentation. Add a suspension lifecycle hook or finally path in core so plugins can detach on the same worker without reporting success, plus a worker-reuse test using different branch IDs.

[github-actions]

Codex AI review

Two P2 observability regressions remain in the new scope bookkeeping, affecting dual-plugin use and reused concurrent workers. Tests were not executed under the review constraints.

Reviewed commit 8ace0f5c977c949726e3a0a8600119987a60821f. Workflow run

[github-actions]

Claude AI review

Reviewed the OTel context attach/detach balancing change against the changed source and surrounding base-revision code in both plugins, log_filter.py, the core SDK's wrap_user_function suspension path (state.py:1156), and the worker-thread lifecycle (execution.py:284).

No actionable defects found. The fix is correct and well-tested. Confirmed:

• Every worker-thread enter_scope is paired with exit_scope in on_user_function_end; unwind sweeps the handler thread at invocation end. Dropping the invocation-start attach is sound — the dex-handler ThreadPoolExecutor does not copy contextvars, so that attach never reached user code and only leaked on the reused handler thread.
• exit_scope/unwind/_discard_* unwind LIFO from the target index upward, so ContextVar.reset() is never called out of order.
• Ownership is read from the context (_OWNER_KEY), correctly handling propagated operation contexts; the marker is a plain context value (not baggage), so it won't leak into propagator headers.
• _scope_base_context uses an is not None guard, correctly preserving an intentionally-empty extracted Context.
• get_current_span_context gating on owns_current preserves in-operation log correlation and falls back to the registry between operations and on the handler thread; the documented handler-thread Workflow→Invocation spanId shift (same traceId) is pinned by a test.
• No new lint/type/circular-import issues: context_scope imports opentelemetry lazily and nothing from the plugin package; Token[Context] is deferred via from __future__ import annotations.

Residual risk (non-blocking, already documented in the PR):

• A suspended operation on a reused branch-worker thread leaves its scope attached (no end hook fires for SuspendExecution); within the same invocation/epoch a subsequent branch on that thread stacks on it and inherits its ambient context/baggage. This is strictly better than the base, dies with the per-invocation branch pool, and is explicitly documented as requiring a core-side suspension hook — not fixable in this PR.
• The new autouse _assert_otel_context_balanced fixtures assert the balance invariant across the whole otel suite; any test not shown in the diff that drives an unpaired hook would newly fail. The PR reports the full suite passing (129 otel / 3223 monorepo), which could not be executed here to independently confirm.

Reviewed commit 8ace0f5c977c949726e3a0a8600119987a60821f. Workflow run