fix(sdk): restore the previous context when an entity span ends (#4378) - #4395
Open
Anai-Guo wants to merge 1 commit into
Open
fix(sdk): restore the previous context when an entity span ends (#4378)#4395Anai-Guo wants to merge 1 commit into
Anai-Guo wants to merge 1 commit into
Conversation
…eloop#4378) `_setup_span` called `set_workflow_name`/`set_agent_name`, which attach a new OTel context and drop the returned token. Nothing ever detaches it, so `_cleanup_span`'s single detach of the span token restores a context that still carries the entity name. Spans created after a completed `@agent` therefore inherit `gen_ai.agent.name`, and the same applies to `traceloop.workflow.name` after a `@workflow` returns. Fold the entity name into the context that `_setup_span` attaches for the span, so the existing cleanup token restores it in one operation. The name is still active while `tracer.start_span()` runs -- the span processor's `on_start` reads it from the active context -- so it is attached briefly around the `start_span` call and then folded into the long-lived context.
|
|
📝 WalkthroughWalkthroughThe decorator span setup now scopes workflow and agent names to each invocation. Regression tests cover standalone tasks, workflow siblings, and cleanup after workflow completion. ChangesContext scoping
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 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: 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/base.py`:
- Around line 140-161: Update the synchronous generator path that calls
_handle_generator so it receives the context attachment token created for the
entity context. Ensure _handle_generator detaches that token and invokes
_cleanup_span in its finally block when generator execution completes,
preventing workflow or agent names from leaking to later spans. Add a regression
test covering a synchronous generator followed by a decorated `@task`.
🪄 Autofix (Beta)
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: 308f3f01-a5e3-4510-8687-a74a69eb8ab6
📒 Files selected for processing (2)
packages/traceloop-sdk/tests/test_agent_workflow_context.pypackages/traceloop-sdk/traceloop/sdk/decorators/base.py
Comment on lines
+140
to
+161
| # Build the entity name into a context value instead of attaching it on its | ||
| # own. `set_workflow_name`/`set_agent_name` attach and drop the token, so | ||
| # nothing ever restores it and the name leaks onto sibling spans created | ||
| # after this entity returns. The name still has to be in the *active* | ||
| # context while the span starts, because the span processor's `on_start` | ||
| # reads it from there -- so attach it, start the span, then fold it into the | ||
| # single long-lived context that `_cleanup_span` detaches. | ||
| entity_ctx = context_api.get_current() | ||
| if tlp_span_kind == TraceloopSpanKindValues.WORKFLOW: | ||
| set_workflow_name(entity_name) | ||
| entity_ctx = context_api.set_value("workflow_name", entity_name, entity_ctx) | ||
| elif tlp_span_kind == TraceloopSpanKindValues.AGENT: | ||
| set_agent_name(entity_name) | ||
| entity_ctx = context_api.set_value("agent_name", entity_name, entity_ctx) | ||
|
|
||
| span_name = f"{entity_name}.{tlp_span_kind.value}" | ||
|
|
||
| with get_tracer() as tracer: | ||
| span = tracer.start_span(span_name) | ||
| ctx = trace.set_span_in_context(span) | ||
| entity_token = context_api.attach(entity_ctx) | ||
| try: | ||
| span = tracer.start_span(span_name) | ||
| finally: | ||
| context_api.detach(entity_token) | ||
| ctx = trace.set_span_in_context(span, entity_ctx) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Clean up context for synchronous generator entities.
Line 161 creates a long-lived context that includes entity_ctx. The synchronous generator branch returns _handle_generator(span, res) without ctx_token. It cannot detach the attachment created at Line 162 when the generator completes.
A decorated synchronous generator can still leak its agent or workflow name to later spans. Pass ctx_token to the synchronous generator handler and call _cleanup_span in its finally block. Add a regression test for a synchronous generator followed by a @task.
🤖 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 140 -
161, Update the synchronous generator path that calls _handle_generator so it
receives the context attachment token created for the entity context. Ensure
_handle_generator detaches that token and invokes _cleanup_span in its finally
block when generator execution completes, preventing workflow or agent names
from leaking to later spans. Add a regression test covering a synchronous
generator followed by a decorated `@task`.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #4378.
Problem
_setup_span()sets the entity name withset_workflow_name()/set_agent_name(). Both do:— they attach a new context and throw the token away. Nothing ever detaches it, so
_cleanup_span()'s singledetach(ctx_token)restores a context that still carries the entity name. The name then leaks onto every span created after the entity returns:@taskrun after a completed bare@agentis taggedgen_ai.agent.name = "planner"@taskinside the same@workflow, after a nested@agentfinished, is tagged with the completed agent's nametraceloop.workflow.nameafter a@workflowreturnsSibling work therefore appears to belong to the preceding agent in the backend.
Fix
Fold the entity name into the same context that
_setup_span()attaches for the span, so the existing cleanup token restores both in one operation — no second, un-detached attach.One subtlety worth calling out: the name has to be in the active context while
tracer.start_span()runs, because the span processor'son_start(default_span_processor_on_start) readsget_value("workflow_name")/get_value("agent_name")off the current context, not off the span's parent context. Passingcontext=tostart_spanwould not be enough. So the entity context is attached just around thestart_span()call, detached immediately (properly nested), and then folded into the long-lived context that is attached for the span's children.set_workflow_name()/set_agent_name()are public API and are left untouched.Tests
Three regression tests added to
tests/test_agent_workflow_context.py, which already covers this area:test_task_after_bare_agent_is_not_tagged_with_agent_nametest_sibling_task_in_workflow_is_not_tagged_with_agent_name— also asserts the sibling keepstraceloop.workflow.name = "rag", per the issue's expected behaviortest_workflow_name_does_not_leak_after_workflow_returnsVerified against
traceloop-sdk==0.62.1, whosedecorators/base.pyandtracing/tracing.pyare byte-identical tomain@93429cf:test_agent_workflow_context.pytest_class_tasks,test_conversation_id,test_nested_tasks,test_manual,test_workflows,test_privacy_no_promptsThe 5 remaining failures are pre-existing in my environment and identical before and after the change (VCR cassettes needing live OpenAI access, plus a timestamp assertion in
test_manual).ruff checkclean under the package's config (line-length = 120,select = ["E", "F", "W"]).🤖 Generated with Claude Code
Summary by CodeRabbit
Bug Fixes
Tests