fix(tracing): release the trace scope when a generator is closed

[adityasingh2400]

Closing an async generator that has a with trace(...) block open behaves differently depending on whether tracing is enabled and on which task runs the close. Measured on main, advancing the generator once and then closing it:

trace type	close	aclose() result	caller's current trace afterwards
NoOpTrace	same task	ok	restored
NoOpTrace	other task	raises ValueError	still the closed trace
TraceImpl	same task	ok	still the closed trace
TraceImpl	other task	ok	still the closed trace

Two separate defects sit in that table.

The first is the crash. NoOpTrace.finish calls Scope.reset_current_trace unconditionally, and a Token is only valid in the Context that created it. asyncio finalizes an abandoned async generator from whichever task happens to run its aclose, so the body resumes in a context that never set the token and ContextVar.reset raises ValueError: <Token ...> was created in a different Context. Disabling tracing should not turn a working teardown into an exception.

The second is a context leak. TraceImpl.__exit__ and ReattachedTrace.__exit__ passed reset_current=exc_type is not GeneratorExit, so they skipped the reset for every GeneratorExit, including the ordinary same-task close() where the token is perfectly valid. The ended trace stays current for the rest of the caller's scope and later traces attach under it.

This replaces the blanket skip with a _reset_current_trace helper that resets and swallows only the cross-context ValueError. Same-task closes now restore the caller's trace for all three trace types, and cross-task closes no longer raise. What this cannot do is repair the caller's context from the finalizing task, because that task cannot rewrite another Context. The remaining leak on the cross-task rows above is therefore unchanged, and the helper's docstring says so rather than pretending the reset succeeded.

Tests in tests/tracing/test_traces_impl.py are parametrized over NoOpTrace and TraceImpl and cover both closes. On main the same-task case fails for TraceImpl and the cross-task case fails for NoOpTrace, which is exactly the two-defect split above. make lint, make typecheck, and the 180 tests across tests/tracing/, tests/test_tracing.py, tests/test_trace_processor.py, tests/test_agent_tracing.py, tests/test_tracing_errors.py and tests/test_tracing_errors_streamed.py are green.

Note that SpanImpl.__exit__ and NoOpSpan.__exit__ still carry the same GeneratorExit skip and so still have the second defect at span level. That is left out of this change deliberately to keep the diff to one lifecycle type, and it is easy to follow up on.

This is a refile of #3225, which was closed in an inactivity sweep with an invitation to revisit.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 731abaa2bb

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 83506d4a00

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 01493d715a

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[seratch]

The two generator-close defects are valid, and the trace-only scope is reasonable. Before merging, please scope the ValueError suppression to GeneratorExit-driven __exit__ cleanup. The current helper is used by all three finish(reset_current=True) implementations, so an explicit finish from the wrong context now silently discards the saved token instead of exposing the context-ownership violation.

Please add a regression showing that a non-GeneratorExit foreign-context finish still raises, include ReattachedTrace in the same-task and cross-task test matrix (or stop changing it), and clear the caller context after the cross-task test. The unavoidable caller-context residue can remain documented, and the analogous span behavior can be handled separately. After these focused changes, this should be ready for another review.

[adityasingh2400]

All four points addressed in e245662.

Scoping the suppression. You are right that the shared helper was too broad. It sat inside finish, so all three finish(reset_current=True) implementations swallowed the error and an explicit finish from the wrong context silently dropped the saved token.

finish now calls Scope.reset_current_trace directly again, so it is strict. The tolerance moved to __exit__ and only applies when GeneratorExit is unwinding the block:

def __exit__(self, exc_type, exc_val, exc_tb):
    if exc_type is GeneratorExit:
        _finish_on_generator_exit(self)
    else:
        self.finish(reset_current=True)

_finish_on_generator_exit catches the ValueError, logs at debug, and clears the token. That is the one path where the reset genuinely cannot succeed, since the finalizing task cannot rewrite the caller's context.

Foreign-context regression. test_explicit_finish_from_another_context_still_raises starts a trace as current, then calls finish(reset_current=True) from a separate task and asserts ValueError. On the previous commit it fails with DID NOT RAISE <class 'ValueError'> for all three trace types, which is exactly the silent discard you described.

ReattachedTrace. Added to the matrix rather than dropped. Both generator tests and the new one now run over NoOpTrace, TraceImpl and ReattachedTrace through a shared _TRACE_FACTORIES list.

Cross-task residue. The cross-task test now clears the caller context at the end, with a comment saying the residue is the documented consequence of finalizing from another task rather than something the test is asserting away.

Verification: tests/tracing is 48 passed, with ruff, ruff format and mypy clean. I left the analogous span behavior alone as you suggested.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: e245662446

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[adityasingh2400]

Codex posted one more P2 here a few minutes after this merged, and it is correct, so flagging it since the defect is now on main.

My _finish_on_generator_exit wraps the whole finish() call, but TraceImpl.finish runs the processor before resetting the scope:

def finish(self, reset_current: bool = False):
    if not self._started:
        return
    self._processor.on_trace_end(self)          # can raise
    if reset_current and self._prev_context_token is not None:
        Scope.reset_current_trace(self._prev_context_token)
        self._prev_context_token = None

So a processor whose on_trace_end raises ValueError is mistaken for a foreign context token. The error is swallowed, the saved token is dropped, and the reset never runs, which leaves the finished trace current for everything after the close. That is the same over-broad catch you had me narrow, one level further out, and I should have seen it when I moved the handler.

I have the fix on fix-generator-exit-catch-only-reset, branched from current main. It runs finish(reset_current=False) so processor failures surface normally, then resets the scope separately and tolerates only that failure:

trace.finish(reset_current=False)

token = trace._prev_context_token
if token is None:
    return
trace._prev_context_token = None
try:
    Scope.reset_current_trace(token)
except ValueError:
    logger.debug("Skipping trace context reset, token belongs to another context")

test_generator_close_surfaces_processor_failure covers it. On current main it fails with DID NOT RAISE <class 'ValueError'>, and with the change the processor error propagates and the token is left in place rather than silently dropped. tests/tracing is 49 passed with ruff, ruff format and mypy clean.

Happy to open that as a PR, or if you would rather fold it in differently just say so.

[seratch]

Yeah, fixing it as well would be appreciated!

[adityasingh2400]

Thanks, will do. The branch is fix-generator-exit-catch-only-reset on my fork, currently one commit ahead of main and zero behind, touching src/agents/tracing/traces.py and tests/tracing/test_traces_impl.py only. I will put the PR up shortly.

[adityasingh2400]

Opened as #4232.