fix(sdk): restore agent context after decorator exits

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

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: