diff --git a/trpc_agent_sdk/agents/_base_agent.py b/trpc_agent_sdk/agents/_base_agent.py index 930cd896..7b16cc16 100644 --- a/trpc_agent_sdk/agents/_base_agent.py +++ b/trpc_agent_sdk/agents/_base_agent.py @@ -260,9 +260,18 @@ async def run_async( from trpc_agent_sdk.telemetry._trace import tracer from trpc_agent_sdk.telemetry._trace import trace_agent - # Avoid start_as_current_span in async generators; cancellation may close - # the generator from another context and trigger detach token errors. + # Manually propagate span context using attach/detach instead of + # start_as_current_span. This ensures child spans (call_llm, execute_tool, + # etc.) can correctly resolve their parent. + # We use start_span + attach/detach rather than start_as_current_span + # because __aexit__ of the context manager is not guaranteed to run when + # an async generator is cancelled, but try/finally always executes + # even under CancelledError (PEP 492). + from opentelemetry import context as context_api + from opentelemetry.trace import set_span_in_context + span = tracer.start_span(f"agent_run [{self.name}]") + _ctx_token = context_api.attach(set_span_in_context(span, context_api.get_current())) try: ctx = self._create_invocation_context(parent_context) if ctx.agent_context is None: @@ -325,6 +334,7 @@ async def run_async( # avoid memory leak reset_invocation_ctx(token) finally: + context_api.detach(_ctx_token) span.end() @abstractmethod diff --git a/trpc_agent_sdk/runners.py b/trpc_agent_sdk/runners.py index 9c292a84..7173886d 100644 --- a/trpc_agent_sdk/runners.py +++ b/trpc_agent_sdk/runners.py @@ -380,9 +380,18 @@ async def run_async( Yields: The events generated by the agent. """ - # Avoid start_as_current_span in async generators; cancellation may close - # the generator from another context and trigger detach token errors. + # Manually propagate span context using attach/detach instead of + # start_as_current_span. This ensures child spans (agent_run, call_llm, + # execute_tool, etc.) can correctly resolve their parent. + # We use start_span + attach/detach rather than start_as_current_span + # because __aexit__ of the context manager is not guaranteed to run when + # an async generator is cancelled, but try/finally always executes + # even under CancelledError (PEP 492). + from opentelemetry import context as context_api + from opentelemetry.trace import set_span_in_context + span = tracer.start_span("invocation") + _ctx_token = context_api.attach(set_span_in_context(span, context_api.get_current())) try: # Create default agent context if not provided if agent_context is None: @@ -621,6 +630,7 @@ async def run_async( session_id=session_id, ) finally: + context_api.detach(_ctx_token) span.end() async def _append_new_message_to_session( diff --git a/trpc_agent_sdk/telemetry/_trace.py b/trpc_agent_sdk/telemetry/_trace.py index aac6657b..f6936afe 100644 --- a/trpc_agent_sdk/telemetry/_trace.py +++ b/trpc_agent_sdk/telemetry/_trace.py @@ -26,6 +26,7 @@ from __future__ import annotations import json +from collections.abc import Sequence from typing import Any from typing import Optional @@ -231,7 +232,23 @@ def trace_agent( span.set_attribute(f"{_trpc_agent_span_name}.agent.user_id", invocation_context.session.user_id) input_str = "" - if invocation_context.user_content and invocation_context.user_content.parts: + # When override_messages is set (e.g., member agent delegated by TeamAgent), + # use override_messages as the actual input instead of user_content, + # because user_content still holds the original input to the leader agent. + # Use getattr + Sequence check to avoid false positives when the context + # is a MagicMock (e.g., in unit tests) where accessing an unset attribute + # returns a truthy MagicMock object (isinstance(m, Sequence) is False). + override_messages = getattr(invocation_context, "override_messages", None) + if (isinstance(override_messages, Sequence) and not isinstance(override_messages, (str, bytes)) + and override_messages): + text_parts = [] + for content in override_messages: + if content and content.parts: + for part in content.parts: + if part.text and not part.thought: + text_parts.append(part.text) + input_str = "\n".join(text_parts) + elif invocation_context.user_content and invocation_context.user_content.parts: input_str = "\n".join([part.text or "" for part in invocation_context.user_content.parts]) span.set_attribute(f"{_trpc_agent_span_name}.agent.input", input_str)