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: