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

[P1] Clean up suspended scopes before branch workers are reused. These guards only remove scopes from an older invocation or the same operation key. A timed-out map/parallel branch skips on_user_function_end, and its pool thread can then resume a different branch in the same invocation. The new scope stacks above the abandoned sibling and detaching it restores that sibling's span, leaking log correlation, baggage, and suppression state across branches. Add a same-thread suspension cleanup hook in the core user-function lifecycle rather than relying on key/epoch heuristics, with an executor-level resubmission test.

[github-actions]

Codex AI review

[P2] Do not use thread-local depth as operation-context ownership. Both plugins now use this value to decide whether the current span is durable. When an operation's OTel context is propagated through asyncio.to_thread, copy_context, or an instrumented executor, the child thread has the correct operation span but depth zero, so logs incorrectly fall back to the invocation span and nested scopes discard propagated context. Keep detach tokens thread-local, but add an ownership marker to the attached OTel context or validate the current span against the plugin registry; cover propagated child-thread logging in both plugins.

[github-actions]

Codex AI review

The normal lifecycle is balanced, but two concurrency paths remain unsafe. Add executor-level tests for branch resubmission and propagated OTel contexts.

Reviewed commit 1e9fa5c0c3566e7b0f800bab59faec24d0be7b95. Workflow run

[github-actions]

Claude AI review

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

No blocking findings. This is a careful, well-tested fix that correctly addresses the trace-pollution bug (unbalanced context.attach()/re-attach leaving an ended span current on the reused Lambda handler thread) in both OTel plugins.

Verified correct:

• context_scope enforces strict LIFO detach (exit_scope/unwind/_discard_stale/_discard_reentered all unwind everything stacked above the target before deleting), so ContextVar.reset is never called out of order — matching the module docstring's rationale about reset reviving stale contexts.
• Scope keys mirror the span-registry keys in both plugins, so each enter_scope is popped by the matching exit_scope.
• The warm-reuse fix is sound: the Invocation span is created before any attach, so dropping the invocation-start attach does not change its parent, while GLOBAL-mode ambient-parent lookup and the X-Ray extractor no longer adopt the prior execution's ended Workflow span (confirmed by test_warm_reuse_does_not_share_a_trace_between_executions).
• Log correlation is preserved: get_current_span_context gating on depth(self) > 0 falls back to the registry between operations, and test_get_current_span_context_returns_invocation_span_between_steps passes untouched.
• The SuspendExecution premise checks out against the core SDK (state.py:1172-1173 re-raises without calling on_user_function_end), which the epoch and re-entry cleanup handle.
• _scope_base_context correctly distinguishes an intentionally-empty extracted Context (falsy, since Context subclasses dict) from None.

Test coverage is strong: the autouse fixtures were upgraded from hiding leaks (global reset) to asserting balance, and new tests cover suspension, re-entry, worker-thread confinement, FLAT branches, empty/populated extracted context, and both plugins on one thread.

Residual risk (documented in the PR, not a regression): a scope abandoned by a suspended operation on a per-invocation worker thread cannot be detached from the handler-thread unwind, and cannot be distinguished from a live sibling (parent_id is checkpoint hierarchy, not physical nesting). If a branch-pool worker that ran a suspended branch is reused for another branch within the same invocation, that sibling's ambient context (baggage / auto-instrumented parents) can be polluted. The durable span hierarchy itself is unaffected because parents are chosen explicitly in _start_span. This is strictly better than the base (which leaked every scope unconditionally) and, per the PR, needs an SDK-side suspended-user-function end hook to close fully. No test exercises this branch-worker-reuse path; worth adding once the SDK gap is addressed.

Reviewed commit 1e9fa5c0c3566e7b0f800bab59faec24d0be7b95. Workflow run