Skip to content

Commit 7971826

Browse files
committed
fix(telemetry): propagate span context in async generators and fix member agent input tracing
- Use start_span + attach/detach instead of start_as_current_span in runners.py and _base_agent.py to properly propagate span context in async generators (CancelledError safe per PEP 492) - Fix trace_agent to prefer override_messages over user_content when tracing member agents delegated by TeamAgent
1 parent 240e3c9 commit 7971826

3 files changed

Lines changed: 42 additions & 5 deletions

File tree

trpc_agent_sdk/agents/_base_agent.py

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -260,9 +260,18 @@ async def run_async(
260260
from trpc_agent_sdk.telemetry._trace import tracer
261261
from trpc_agent_sdk.telemetry._trace import trace_agent
262262

263-
# Avoid start_as_current_span in async generators; cancellation may close
264-
# the generator from another context and trigger detach token errors.
263+
# Manually propagate span context using attach/detach instead of
264+
# start_as_current_span. This ensures child spans (call_llm, execute_tool,
265+
# etc.) can correctly resolve their parent.
266+
# We use start_span + attach/detach rather than start_as_current_span
267+
# because __aexit__ of the context manager is not guaranteed to run when
268+
# an async generator is cancelled, but try/finally always executes
269+
# even under CancelledError (PEP 492).
270+
from opentelemetry import context as context_api
271+
from opentelemetry.trace import set_span_in_context
272+
265273
span = tracer.start_span(f"agent_run [{self.name}]")
274+
_ctx_token = context_api.attach(set_span_in_context(span, context_api.get_current()))
266275
try:
267276
ctx = self._create_invocation_context(parent_context)
268277
if ctx.agent_context is None:
@@ -325,6 +334,7 @@ async def run_async(
325334
# avoid memory leak
326335
reset_invocation_ctx(token)
327336
finally:
337+
context_api.detach(_ctx_token)
328338
span.end()
329339

330340
@abstractmethod

trpc_agent_sdk/runners.py

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -380,9 +380,18 @@ async def run_async(
380380
Yields:
381381
The events generated by the agent.
382382
"""
383-
# Avoid start_as_current_span in async generators; cancellation may close
384-
# the generator from another context and trigger detach token errors.
383+
# Manually propagate span context using attach/detach instead of
384+
# start_as_current_span. This ensures child spans (agent_run, call_llm,
385+
# execute_tool, etc.) can correctly resolve their parent.
386+
# We use start_span + attach/detach rather than start_as_current_span
387+
# because __aexit__ of the context manager is not guaranteed to run when
388+
# an async generator is cancelled, but try/finally always executes
389+
# even under CancelledError (PEP 492).
390+
from opentelemetry import context as context_api
391+
from opentelemetry.trace import set_span_in_context
392+
385393
span = tracer.start_span("invocation")
394+
_ctx_token = context_api.attach(set_span_in_context(span, context_api.get_current()))
386395
try:
387396
# Create default agent context if not provided
388397
if agent_context is None:
@@ -621,6 +630,7 @@ async def run_async(
621630
session_id=session_id,
622631
)
623632
finally:
633+
context_api.detach(_ctx_token)
624634
span.end()
625635

626636
async def _append_new_message_to_session(

trpc_agent_sdk/telemetry/_trace.py

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
from __future__ import annotations
2727

2828
import json
29+
from collections.abc import Sequence
2930
from typing import Any
3031
from typing import Optional
3132

@@ -231,7 +232,23 @@ def trace_agent(
231232
span.set_attribute(f"{_trpc_agent_span_name}.agent.user_id", invocation_context.session.user_id)
232233

233234
input_str = ""
234-
if invocation_context.user_content and invocation_context.user_content.parts:
235+
# When override_messages is set (e.g., member agent delegated by TeamAgent),
236+
# use override_messages as the actual input instead of user_content,
237+
# because user_content still holds the original input to the leader agent.
238+
# Use getattr + Sequence check to avoid false positives when the context
239+
# is a MagicMock (e.g., in unit tests) where accessing an unset attribute
240+
# returns a truthy MagicMock object (isinstance(m, Sequence) is False).
241+
override_messages = getattr(invocation_context, "override_messages", None)
242+
if (isinstance(override_messages, Sequence) and not isinstance(override_messages, (str, bytes))
243+
and override_messages):
244+
text_parts = []
245+
for content in override_messages:
246+
if content and content.parts:
247+
for part in content.parts:
248+
if part.text and not part.thought:
249+
text_parts.append(part.text)
250+
input_str = "\n".join(text_parts)
251+
elif invocation_context.user_content and invocation_context.user_content.parts:
235252
input_str = "\n".join([part.text or "" for part in invocation_context.user_content.parts])
236253
span.set_attribute(f"{_trpc_agent_span_name}.agent.input", input_str)
237254

0 commit comments

Comments
 (0)