Conversation
…me leak across nested entities
📝 WalkthroughWalkthroughThe change adds token-based context cleanup for OpenAI agent spans and SDK decorators. It covers synchronous, asynchronous, generator, and async-generator paths, including span failures, task or thread boundaries, and conversation ID scoping. ChangesTracing context cleanup
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant DecoratedEntity
participant DecoratorWrapper
participant ContextSetters
participant OpenTelemetrySpan
DecoratedEntity->>DecoratorWrapper: invoke sync, async, or generator function
DecoratorWrapper->>ContextSetters: attach scoped context values
ContextSetters-->>DecoratorWrapper: return detach tokens
DecoratorWrapper->>OpenTelemetrySpan: create and run span
OpenTelemetrySpan-->>DecoratorWrapper: complete or close execution
DecoratorWrapper->>ContextSetters: safely detach tokens
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (5)
packages/traceloop-sdk/tests/test_conversation_id.py (2)
201-219: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThe patched
detachleaves the token attached.
raising_detachnever detaches, so theconv-crashtoken stays on the context that the coroutine ran in. Today this does not affecttest_conversation_id_does_not_leak_to_later_sibling, because pytest-asyncio runs the coroutine in aTaskwith its own copy of the context. That isolation is incidental. If the test later becomes synchronous, or the runner changes how it creates the context,conv-crashleaks into the following tests and they fail.Make the cleanup explicit: capture the real
detach, raiseValueError, then call the realdetachso state is restored.♻️ Proposed test hardening
+ real_detach = context_api.detach + def raising_detach(token): + real_detach(token) raise ValueError("was created in a different Context")Also add
assert context_api.get_value("conversation_id") is Noneafter the call once the real detach runs.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/traceloop-sdk/tests/test_conversation_id.py` around lines 201 - 219, Update test_conversation_decorator_async_detach_does_not_crash to preserve the original context_api.detach before monkeypatching, have the replacement raise the ValueError and then invoke the saved real detach with the token, and assert context_api.get_value("conversation_id") is None after handler() completes.
156-198: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDefine
inner_taskbefore the generator that calls it.
streaming_chatreferencesinner_taskat line 162, andinner_taskis defined at line 165. This works, because the generator body runs at iteration time and resolves the closure then. Reordering the definitions makes the dependency clear and removes the reliance on deferred name resolution.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/traceloop-sdk/tests/test_conversation_id.py` around lines 156 - 198, In both test_conversation_decorator_sync_generator and test_conversation_decorator_async_generator, move the inner_task definition before the streaming_chat generator definition that invokes it. Preserve the existing decorators, yielded values, and assertions unchanged.packages/traceloop-sdk/traceloop/sdk/decorators/__init__.py (1)
54-55: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider a module-level import for
_safe_detach.
_safe_detachandset_conversation_idare imported insideconversation. If the import guards against a circular import, add a short comment stating that. Otherwise move both imports to module level.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/traceloop-sdk/traceloop/sdk/decorators/__init__.py` around lines 54 - 55, Update the imports used by conversation to be module-level for _safe_detach and set_conversation_id; if _safe_detach must remain inside conversation to avoid a circular import, keep that placement and add a brief explanatory comment.packages/opentelemetry-instrumentation-openai-agents/tests/test_tracing_processor.py (1)
1337-1363: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the "no OTel span created" detach branch.
on_span_startalso detaches the name token when_start_agent_spanreturns a falsy span (theelif name_token is not Nonebranch in_hooks.py). The failure test covers the raise path only. A test that makesstart_spanreturnNonewould cover the remaining branch.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/opentelemetry-instrumentation-openai-agents/tests/test_tracing_processor.py` around lines 1337 - 1363, Extend test_name_detached_when_span_creation_fails to cover the falsy-span path by configuring processor.tracer.start_span to return None instead of raising. Invoke processor.on_span_start with the same named AgentSpanData and assert get_value("agent_name") is None, covering the elif name_token is not None branch in on_span_start.packages/traceloop-sdk/traceloop/sdk/decorators/base.py (1)
330-346: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueTokens stay attached after
sync_wrapreturns a generator.If
fnreturns a generator,sync_wrapreturns_handle_generator(...)without detaching. The name, span, and path tokens stay attached on the caller's context until the returned generator runs itsfinally. If the caller never iterates the generator, the entity name leaks for the rest of the trace. This matches the earlier behavior for the span token, so it is not a regression, but the new name and path tokens widen the leaked surface. Consider documenting this alongside theasync_gen_wrap"KNOWN LIMITATION" note.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/traceloop-sdk/traceloop/sdk/decorators/base.py` around lines 330 - 346, Document this known limitation in the synchronous generator path around sync_wrap and _handle_generator: when fn returns an unconsumed generator, the entity name, span, and path tokens remain attached until generator cleanup runs, so they may leak if iteration never begins. Align the documentation with the existing async_gen_wrap KNOWN LIMITATION note without changing the generator behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@packages/opentelemetry-instrumentation-openai-agents/opentelemetry/instrumentation/openai_agents/_hooks.py`:
- Around line 670-672: Update shutdown() to clear _agent_name_tokens alongside
_otel_spans, _span_contexts, _root_spans, and _reverse_handoffs_dict, ensuring
shutdown releases all tracked agent-name tokens and related references.
In `@packages/traceloop-sdk/traceloop/sdk/decorators/__init__.py`:
- Around line 93-101: Update sync_wrapper in the conversation decorator to
detect when fn returns a generator and keep the conversation token attached
until iteration completes, mirroring base._handle_generator; retain immediate
cleanup for non-generator results and ensure cleanup also occurs if iteration
raises. Add a test covering `@conversation` stacked above `@task` on a generator
function, verifying the conversation ID remains available while yielding.
---
Nitpick comments:
In
`@packages/opentelemetry-instrumentation-openai-agents/tests/test_tracing_processor.py`:
- Around line 1337-1363: Extend test_name_detached_when_span_creation_fails to
cover the falsy-span path by configuring processor.tracer.start_span to return
None instead of raising. Invoke processor.on_span_start with the same named
AgentSpanData and assert get_value("agent_name") is None, covering the elif
name_token is not None branch in on_span_start.
In `@packages/traceloop-sdk/tests/test_conversation_id.py`:
- Around line 201-219: Update
test_conversation_decorator_async_detach_does_not_crash to preserve the original
context_api.detach before monkeypatching, have the replacement raise the
ValueError and then invoke the saved real detach with the token, and assert
context_api.get_value("conversation_id") is None after handler() completes.
- Around line 156-198: In both test_conversation_decorator_sync_generator and
test_conversation_decorator_async_generator, move the inner_task definition
before the streaming_chat generator definition that invokes it. Preserve the
existing decorators, yielded values, and assertions unchanged.
In `@packages/traceloop-sdk/traceloop/sdk/decorators/__init__.py`:
- Around line 54-55: Update the imports used by conversation to be module-level
for _safe_detach and set_conversation_id; if _safe_detach must remain inside
conversation to avoid a circular import, keep that placement and add a brief
explanatory comment.
In `@packages/traceloop-sdk/traceloop/sdk/decorators/base.py`:
- Around line 330-346: Document this known limitation in the synchronous
generator path around sync_wrap and _handle_generator: when fn returns an
unconsumed generator, the entity name, span, and path tokens remain attached
until generator cleanup runs, so they may leak if iteration never begins. Align
the documentation with the existing async_gen_wrap KNOWN LIMITATION note without
changing the generator behavior.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 0c8599f6-45be-4f3e-96b2-15ae4357a8f2
⛔ Files ignored due to path filters (3)
packages/opentelemetry-instrumentation-openai-agents/uv.lockis excluded by!**/*.lockpackages/sample-app/uv.lockis excluded by!**/*.lockpackages/traceloop-sdk/uv.lockis excluded by!**/*.lock
📒 Files selected for processing (7)
packages/opentelemetry-instrumentation-openai-agents/opentelemetry/instrumentation/openai_agents/_hooks.pypackages/opentelemetry-instrumentation-openai-agents/tests/test_tracing_processor.pypackages/traceloop-sdk/tests/test_agent_name_scope.pypackages/traceloop-sdk/tests/test_conversation_id.pypackages/traceloop-sdk/traceloop/sdk/decorators/__init__.pypackages/traceloop-sdk/traceloop/sdk/decorators/base.pypackages/traceloop-sdk/traceloop/sdk/tracing/tracing.py
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@packages/traceloop-sdk/traceloop/sdk/decorators/__init__.py`:
- Around line 109-125: Update the generator-return path in the wrapper around fn
and _consume_with_token so the invocation token is always detached before
returning the generator. Have _consume_with_token attach a fresh conversation
token when iteration begins and detach it when consumption finishes, while
preserving conversation propagation during iteration. Add coverage for storing
an unstarted returned generator before unrelated tracing work executes.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 08a6efac-07cc-4e7e-beea-35bf88b2dbc9
📒 Files selected for processing (2)
packages/traceloop-sdk/tests/test_conversation_id.pypackages/traceloop-sdk/traceloop/sdk/decorators/__init__.py
2 participants
Summary by CodeRabbit
Bug Fixes
Documentation