From 58b339f2cdb4b2cdeecf565626aef7f96bc91a2a Mon Sep 17 00:00:00 2001
From: buduoqiu <yaodong-shen@users.noreply.github.com>
Date: Mon, 27 Jul 2026 01:23:35 +0800
Subject: [PATCH] fix(sdk): restore agent context after decorator exits

Scope the agent name to the span context token so cleanup restores the prior OpenTelemetry context. Add regressions for standalone work and sibling tasks within a workflow.\n\nFixes #4378
---
 .../tests/test_agent_workflow_context.py      | 49 +++++++++++++++++++
 .../traceloop/sdk/decorators/base.py          | 11 +++--
 2 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/packages/traceloop-sdk/tests/test_agent_workflow_context.py b/packages/traceloop-sdk/tests/test_agent_workflow_context.py
index 8c0f92d02e..242f7254a7 100644
--- a/packages/traceloop-sdk/tests/test_agent_workflow_context.py
+++ b/packages/traceloop-sdk/tests/test_agent_workflow_context.py
@@ -10,6 +10,8 @@
 - @agent nested inside @workflow inherits workflow_name from the workflow.
 - The same agent name running under two different workflows stays distinct.
 - A bare @agent (no enclosing @workflow) leaves workflow_name unset.
+- A completed @agent does not tag later independent work.
+- A completed @agent restores the enclosing workflow context for sibling tasks.
 """
 
 from opentelemetry import trace
@@ -126,3 +128,50 @@ def solo_agent():
 
     assert SpanAttributes.TRACELOOP_WORKFLOW_NAME not in child_span.attributes
     assert child_span.attributes[GEN_AI_AGENT_NAME] == "solo"
+
+
+def test_completed_agent_does_not_tag_following_task(exporter):
+    """Agent metadata must not leak into a later independent task."""
+
+    @agent(name="completed")
+    def completed_agent():
+        pass
+
+    @task(name="standalone")
+    def standalone_task():
+        pass
+
+    completed_agent()
+    standalone_task()
+
+    spans = {span.name: span for span in exporter.get_finished_spans()}
+    standalone_span = spans["standalone.task"]
+
+    assert GEN_AI_AGENT_NAME not in standalone_span.attributes
+    assert standalone_span.parent is None
+
+
+def test_completed_agent_restores_enclosing_workflow_context(exporter):
+    """A task after an agent keeps its workflow but not the completed agent."""
+
+    @agent(name="planner")
+    def planner_agent():
+        pass
+
+    @task(name="after_agent")
+    def task_after_agent():
+        pass
+
+    @workflow(name="rag")
+    def rag_workflow():
+        planner_agent()
+        task_after_agent()
+
+    rag_workflow()
+
+    spans = {span.name: span for span in exporter.get_finished_spans()}
+    task_span = spans["after_agent.task"]
+
+    assert task_span.attributes[SpanAttributes.TRACELOOP_WORKFLOW_NAME] == "rag"
+    assert GEN_AI_AGENT_NAME not in task_span.attributes
+    assert task_span.parent.span_id == spans["rag.workflow"].context.span_id
diff --git a/packages/traceloop-sdk/traceloop/sdk/decorators/base.py b/packages/traceloop-sdk/traceloop/sdk/decorators/base.py
index 04f00e0ad9..5a1df4be55 100644
--- a/packages/traceloop-sdk/traceloop/sdk/decorators/base.py
+++ b/packages/traceloop-sdk/traceloop/sdk/decorators/base.py
@@ -17,10 +17,11 @@
 from opentelemetry import context as context_api
 from opentelemetry.semconv_ai import SpanAttributes, TraceloopSpanKindValues
 from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import (
+    GEN_AI_AGENT_NAME,
     GEN_AI_TOOL_NAME,
 )
 
-from traceloop.sdk.tracing import get_tracer, set_workflow_name, set_agent_name
+from traceloop.sdk.tracing import get_tracer, set_workflow_name
 from traceloop.sdk.tracing.tracing import (
     TracerWrapper,
     set_entity_path,
@@ -139,14 +140,16 @@ def _setup_span(entity_name, tlp_span_kind, version):
     """Sets up the OpenTelemetry span and context"""
     if tlp_span_kind == TraceloopSpanKindValues.WORKFLOW:
         set_workflow_name(entity_name)
-    elif tlp_span_kind == TraceloopSpanKindValues.AGENT:
-        set_agent_name(entity_name)
 
     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)
+        if tlp_span_kind == TraceloopSpanKindValues.AGENT:
+            # Keep the agent name in the same context scope as its span. Detaching
+            # ctx_token then restores the context that existed before the agent.
+            ctx = context_api.set_value("agent_name", entity_name, ctx)
         ctx_token = context_api.attach(ctx)
 
         if tlp_span_kind in [
@@ -158,6 +161,8 @@ def _setup_span(entity_name, tlp_span_kind, version):
 
         span.set_attribute(SpanAttributes.TRACELOOP_SPAN_KIND, tlp_span_kind.value)
         span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_NAME, entity_name)
+        if tlp_span_kind == TraceloopSpanKindValues.AGENT:
+            span.set_attribute(GEN_AI_AGENT_NAME, entity_name)
         if tlp_span_kind == TraceloopSpanKindValues.TOOL:
             span.set_attribute(GEN_AI_TOOL_NAME, entity_name)
         if version: