Python: fix: establish correct span parenting during streaming agent invocation#5098
Python: fix: establish correct span parenting during streaming agent invocation#5098JasonOA888 wants to merge 1 commit intomicrosoft:mainfrom
Conversation
When stream=True, agent-invoke spans were not set as the active OTel context, causing child spans (chat completion, execute_tool) to appear as siblings rather than children of invoke_agent. The original code avoided trace.use_span() for streaming spans due to "Failed to detach context" errors when cleanup hooks run in a different async context. Fix: introduce _STREAMING_AGENT_INVOKE_CONTEXT ContextVar that holds the OTel context with the agent-invoke span. Child spans read from this to establish correct parent-child relationships without needing use_span()/attach(). Closes microsoft#5089
There was a problem hiding this comment.
Pull request overview
Fixes OpenTelemetry span hierarchy in Python streaming agent runs so chat and execute_tool spans correctly appear as children of the streaming invoke_agent span (closing issue #5089).
Changes:
- Introduces a ContextVar (
_STREAMING_AGENT_INVOKE_CONTEXT) to hold an OTelContextwith the streaminginvoke_agentspan set as current. - Uses that stored context as the explicit parent when starting streaming
chatspans. - Uses that stored context as the explicit parent when starting
execute_toolspans viaget_function_span().
| # Activate the agent-invoke span as the current OTel context so that | ||
| # child spans (chat completion, execute_tool) created during streaming | ||
| # inherit it as their parent. | ||
| agent_ctx = trace.set_span_in_context(span) | ||
| streaming_agent_ctx_token = _STREAMING_AGENT_INVOKE_CONTEXT.set(agent_ctx) | ||
|
|
There was a problem hiding this comment.
In the streaming path, _STREAMING_AGENT_INVOKE_CONTEXT is set before any potential exceptions from _capture_messages(...)/attribute serialization. If an exception is raised here, the ContextVar token is never reset and the span is never ended, which can leak the streaming parent context into later operations in the same task. Consider wrapping the post-set(...) setup in a try/except/finally that resets the ContextVar (and ends the span) on early failures before the cleanup hooks are registered.
Problem
Closes #5089
When
stream=True, theinvoke_agentspan is not activated as the current OTel context. Child spans (chat,execute_tool) are created without a parent reference and appear as siblings instead of children.Root Cause
The streaming path in
_trace_agent_invocationuses bareget_tracer().start_span()without context activation. The original code deliberately avoidedtrace.use_span()because streaming spans are closed in cleanup hooks that run in a different async context, causing "Failed to detach context" errors.Fix
Introduce
_STREAMING_AGENT_INVOKE_CONTEXT— acontextvars.ContextVarthat holds the OTel context with the agent-invoke span set as current.Three changes:
_trace_agent_invocation(streaming path): After creating theinvoke_agentspan, storetrace.set_span_in_context(span)in the ContextVar. Reset in_close_span._trace_chat_response(streaming path): Read the ContextVar and pass ascontext=tostart_span(), so chat spans become children ofinvoke_agent.get_function_span: Read the ContextVar and pass ascontext=tostart_as_current_span(), soexecute_toolspans become children ofinvoke_agent.This avoids
use_span()/attach()entirely — no detach errors — while propagating parent context via Python contextvars.Testing
Files Changed
python/packages/core/agent_framework/observability.py(+32 lines)