Fix spans not correctly nested when using streaming

python/packages/core/agent_framework/_types.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import base64
+import contextlib
 import json
 import logging
 import re
@@ -2890,6 +2891,7 @@ class ResponseStream(AsyncIterable[UpdateT], Generic[UpdateT, FinalT]):
         self._inner_stream_source: ResponseStream[Any, Any] | Awaitable[ResponseStream[Any, Any]] | None = None
         self._wrap_inner: bool = False
         self._map_update: Callable[[Any], UpdateT | Awaitable[UpdateT]] | None = None
+        self._pull_context_manager_factories: list[Callable[[], contextlib.AbstractContextManager[Any]]] = []
 
     def map(
         self,
@@ -3008,11 +3010,18 @@ class ResponseStream(AsyncIterable[UpdateT], Generic[UpdateT, FinalT]):
         return self
 
     async def __anext__(self) -> UpdateT:
-        if self._iterator is None:
-            stream = await self._get_stream()
-            self._iterator = stream.__aiter__()
         try:
-            update: UpdateT = await self._iterator.__anext__()
+            with contextlib.ExitStack() as stack:
+                for factory in self._pull_context_manager_factories:
+                    stack.enter_context(factory())
+                # Resolve the underlying stream inside the pull contexts so that any
+                # spans/contexts created during stream resolution (e.g. inner chat
+                # completion spans created on the first pull of a wrapped agent stream)
+                # inherit the active context (e.g. an outer agent invoke span).
+                if self._iterator is None:
+                    stream = await self._get_stream()
+                    self._iterator = stream.__aiter__()
+                update: UpdateT = await self._iterator.__anext__()
         except StopAsyncIteration:
             self._consumed = True
             await self._run_cleanup_hooks()
@@ -3177,6 +3186,25 @@ class ResponseStream(AsyncIterable[UpdateT], Generic[UpdateT, FinalT]):
         self._cleanup_hooks.append(hook)
         return self
 
+    def with_pull_context_manager(
+        self,
+        cm_factory: Callable[[], contextlib.AbstractContextManager[Any]],
+    ) -> ResponseStream[UpdateT, FinalT]:
+        """Register a context manager factory invoked around each underlying iterator pull.
+
+        The factory is called once per ``__anext__`` and the returned context manager wraps
+        the await of the underlying iterator. This is useful for state that needs to be
+        active while the inner async work runs - for example, attaching an OpenTelemetry
+        span to the current context so child spans created by inner code (HTTP clients,
+        tool execution) are correctly parented.
+
+        Because the context manager is entered and exited within the same ``__anext__``
+        invocation, attach/detach style operations remain symmetric in the same async
+        context regardless of where the stream is iterated.
+        """
+        self._pull_context_manager_factories.append(cm_factory)
+        return self
+
     async def _run_cleanup_hooks(self) -> None:
         if self._cleanup_run:
             return

python/packages/core/agent_framework/observability.py

@@ -26,6 +26,7 @@ from time import perf_counter, time_ns
 from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, TypedDict, cast, overload
 
 from dotenv import load_dotenv
+from opentelemetry import context as otel_context
 from opentelemetry import metrics, trace
 
 from . import __version__ as version_info
@@ -1277,27 +1278,8 @@ class ChatTelemetryLayer(Generic[OptionsCoT]):
         )
 
         if stream:
-            result_stream = cast(
-                ResponseStream[ChatResponseUpdate, ChatResponse[Any]],
-                super_get_response(
-                    messages=messages,
-                    stream=True,
-                    options=opts,
-                    compaction_strategy=compaction_strategy,
-                    tokenizer=tokenizer,
-                    function_invocation_kwargs=function_invocation_kwargs,
-                    client_kwargs=merged_client_kwargs,
-                ),
-            )
+            span = _start_streaming_span(attributes, OtelAttr.REQUEST_MODEL)
 
-            # Create span directly without trace.use_span() context attachment.
-            # Streaming spans are closed asynchronously in cleanup hooks, which run
-            # in a different async context than creation — using use_span() would
-            # cause "Failed to detach context" errors from OpenTelemetry.
-            operation = attributes.get(OtelAttr.OPERATION, "operation")
-            span_name = attributes.get(OtelAttr.REQUEST_MODEL, "unknown")
-            span = get_tracer().start_span(f"{operation} {span_name}")
-            span.set_attributes(attributes)
             if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and messages:
                 _capture_messages(
                     span=span,
@@ -1319,6 +1301,19 @@ class ChatTelemetryLayer(Generic[OptionsCoT]):
             def _record_duration() -> None:
                 duration_state["duration"] = perf_counter() - start_time
 
+            result_stream = cast(
+                ResponseStream[ChatResponseUpdate, ChatResponse[Any]],
+                super_get_response(
+                    messages=messages,
+                    stream=True,
+                    options=opts,
+                    compaction_strategy=compaction_strategy,
+                    tokenizer=tokenizer,
+                    function_invocation_kwargs=function_invocation_kwargs,
+                    client_kwargs=merged_client_kwargs,
+                ),
+            )
+
             async def _finalize_stream() -> None:
                 from ._types import ChatResponse
 
@@ -1357,11 +1352,18 @@ class ChatTelemetryLayer(Generic[OptionsCoT]):
                 finally:
                     _close_span()
 
-            # Register a weak reference callback to close the span if stream is garbage collected
-            # without being consumed. This ensures spans don't leak if users don't consume streams.
-            wrapped_stream: ResponseStream[ChatResponseUpdate, ChatResponse[Any]] = result_stream.with_cleanup_hook(
-                _record_duration
-            ).with_cleanup_hook(_finalize_stream)
+            # The pull context manager attaches the span around each underlying iterator pull so
+            # that child spans created during the pull (e.g. HTTP requests, inner tool execution)
+            # are parented under this chat span. Attach and detach happen in the same async
+            # context as the pull, avoiding cross-context cleanup issues. The weakref finalizer
+            # ensures the span is closed even if the stream is garbage collected without being
+            # consumed.
+            wrapped_stream: ResponseStream[ChatResponseUpdate, ChatResponse[Any]] = (
+                result_stream
+                .with_cleanup_hook(_record_duration)
+                .with_cleanup_hook(_finalize_stream)
+                .with_pull_context_manager(lambda: _activate_span(span))
+            )
             weakref.finalize(wrapped_stream, _close_span)
             return wrapped_stream
 
@@ -1543,23 +1545,8 @@ class AgentTelemetryLayer:
         inner_accumulated_usage_token = INNER_ACCUMULATED_USAGE.set({})
 
         if stream:
-            try:
-                run_result: object = execute()
-                if isinstance(run_result, ResponseStream):
-                    result_stream: ResponseStream[AgentResponseUpdate, AgentResponse[Any]] = run_result  # pyright: ignore[reportUnknownVariableType]
-                elif isinstance(run_result, Awaitable):
-                    result_stream = ResponseStream.from_awaitable(run_result)  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType]
-                else:
-                    raise RuntimeError("Streaming telemetry requires a ResponseStream result.")
-            except Exception:
-                INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.reset(inner_response_telemetry_captured_fields_token)
-                INNER_ACCUMULATED_USAGE.reset(inner_accumulated_usage_token)
-                raise
+            span = _start_streaming_span(attributes, OtelAttr.AGENT_NAME)
 
-            operation = attributes.get(OtelAttr.OPERATION, "operation")
-            span_name = attributes.get(OtelAttr.AGENT_NAME, "unknown")
-            span = get_tracer().start_span(f"{operation} {span_name}")
-            span.set_attributes(attributes)
             if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and messages:
                 _capture_messages(
                     span=span,
@@ -1581,6 +1568,20 @@ class AgentTelemetryLayer:
             def _record_duration() -> None:
                 duration_state["duration"] = perf_counter() - start_time
 
+            try:
+                run_result: object = execute()
+                if isinstance(run_result, ResponseStream):
+                    result_stream: ResponseStream[AgentResponseUpdate, AgentResponse[Any]] = run_result  # pyright: ignore[reportUnknownVariableType]
+                elif isinstance(run_result, Awaitable):
+                    result_stream = ResponseStream.from_awaitable(run_result)  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType]
+                else:
+                    raise RuntimeError("Streaming telemetry requires a ResponseStream result.")
+            except Exception:
+                INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.reset(inner_response_telemetry_captured_fields_token)
+                INNER_ACCUMULATED_USAGE.reset(inner_accumulated_usage_token)
+                _close_span()
+                raise
+
             async def _finalize_stream() -> None:
                 from ._types import AgentResponse
 
@@ -1620,9 +1621,18 @@ class AgentTelemetryLayer:
                     INNER_ACCUMULATED_USAGE.reset(inner_accumulated_usage_token)
                     _close_span()
 
-            wrapped_stream: ResponseStream[AgentResponseUpdate, AgentResponse[Any]] = result_stream.with_cleanup_hook(
-                _record_duration
-            ).with_cleanup_hook(_finalize_stream)
+            # The pull context manager attaches the span around each underlying iterator pull so
+            # that child spans created during the pull (e.g. inner chat completion spans from the
+            # underlying ChatTelemetryLayer) are parented under this agent invoke span. Attach and
+            # detach happen in the same async context as the pull, avoiding cross-context cleanup
+            # issues. The weakref finalizer ensures the span is closed even if the stream is
+            # garbage collected without being consumed.
+            wrapped_stream: ResponseStream[AgentResponseUpdate, AgentResponse[Any]] = (
+                result_stream
+                .with_cleanup_hook(_record_duration)
+                .with_cleanup_hook(_finalize_stream)
+                .with_pull_context_manager(lambda: _activate_span(span))
+            )
             weakref.finalize(wrapped_stream, _close_span)
             return wrapped_stream
 
@@ -1809,6 +1819,27 @@ def get_function_span(
     )
 
 
+@contextlib.contextmanager
+def _activate_span(span: trace.Span) -> Generator[None, None, None]:
+    """Attach ``span`` as the current span in the OpenTelemetry context.
+
+    Designed to be used as a per-pull context manager registered on a
+    ``ResponseStream`` via ``with_pull_context_manager``: it attaches the span
+    before each underlying iterator pull and detaches immediately after, so
+    child spans created during the pull (HTTP clients, inner chat completions,
+    tool execution) are correctly parented under ``span``.
+
+    Because attach and detach happen within the same ``__anext__`` invocation
+    (and therefore the same async task / contextvars context), there is no risk
+    of "Failed to detach context" warnings from cross-context cleanup.
+    """
+    token = otel_context.attach(trace.set_span_in_context(span))
+    try:
+        yield
+    finally:
+        otel_context.detach(token)
+
+
 @contextlib.contextmanager
 def _get_span(
     attributes: dict[str, Any],
@@ -1831,6 +1862,29 @@ def _get_span(
         yield current_span
 
 
+def _start_streaming_span(attributes: dict[str, Any], span_name_attribute: str) -> trace.Span:
+    """Start a non-current span for a streaming operation.
+
+    Unlike :func:`_get_span`, the returned span is not attached to the current
+    OpenTelemetry context. The caller is responsible for:
+
+    - Ending the span via cleanup hooks on the wrapped
+      :class:`~agent_framework._types.ResponseStream`.
+    - Activating the span around each iterator pull via
+      :func:`_activate_span` registered with ``with_pull_context_manager`` so
+      that child spans created during stream production inherit it as parent.
+
+    Streaming spans are closed asynchronously in cleanup hooks that run in a
+    different async context than creation, so attaching the span at creation
+    time would cause "Failed to detach context" errors from OpenTelemetry.
+    """
+    operation = attributes.get(OtelAttr.OPERATION, "operation")
+    span_name = attributes.get(span_name_attribute, "unknown")
+    span = get_tracer().start_span(f"{operation} {span_name}")
+    span.set_attributes(attributes)
+    return span
+
+
 def _get_instructions_from_options(options: Any) -> str | list[str] | None:
     """Extract instructions from options dict."""
     if options is None:

python/packages/core/tests/core/test_observability.py

@@ -3313,3 +3313,293 @@ async def test_agent_invoke_span_aggregates_usage_on_max_iterations_exhaustion(s
     # The invoke_agent span must aggregate usage from the in-loop call and the final exhaustion call
     assert agent_span.attributes.get(OtelAttr.INPUT_TOKENS) == 500
     assert agent_span.attributes.get(OtelAttr.OUTPUT_TOKENS) == 100
+
+
+# region Test span nesting (parent-child relationships)
+
+
+@pytest.mark.parametrize("stream", [False, True])
+async def test_chat_span_nested_under_agent_span(span_exporter: InMemorySpanExporter, stream: bool):
+    """The inner chat span must be a child of the outer agent invoke span."""
+
+    class NestedChatClient(ChatTelemetryLayer, BaseChatClient[Any]):
+        def service_url(self):
+            return "https://test.example.com"
+
+        def _inner_get_response(
+            self, *, messages: MutableSequence[Message], stream: bool, options: dict[str, Any], **kwargs: Any
+        ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
+            if stream:
+
+                async def _stream() -> AsyncIterable[ChatResponseUpdate]:
+                    yield ChatResponseUpdate(contents=[Content.from_text("Hello")], role="assistant")
+                    yield ChatResponseUpdate(
+                        contents=[Content.from_text(" world")], role="assistant", finish_reason="stop"
+                    )
+
+                def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
+                    return ChatResponse(
+                        messages=[Message(role="assistant", contents=["Hello world"])],
+                        response_id="resp_1",
+                        usage_details=UsageDetails(input_token_count=3, output_token_count=4),
+                        finish_reason="stop",
+                    )
+
+                return ResponseStream(_stream(), finalizer=_finalize)
+
+            async def _get() -> ChatResponse:
+                return ChatResponse(
+                    messages=[Message(role="assistant", contents=["Hello world"])],
+                    response_id="resp_1",
+                    usage_details=UsageDetails(input_token_count=3, output_token_count=4),
+                    finish_reason="stop",
+                )
+
+            return _get()
+
+    agent = Agent(
+        client=NestedChatClient(),
+        id="nested_agent_id",
+        name="nested_agent",
+        default_options={"model": "NestedModel"},
+    )
+
+    span_exporter.clear()
+    if stream:
+        result_stream = agent.run("Test message", stream=True)
+        async for _ in result_stream:
+            pass
+        await result_stream.get_final_response()
+    else:
+        await agent.run("Test message")
+
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 2
+
+    span_by_op = {s.attributes[OtelAttr.OPERATION.value]: s for s in spans}
+    agent_span = span_by_op[OtelAttr.AGENT_INVOKE_OPERATION]
+    chat_span = span_by_op[OtelAttr.CHAT_COMPLETION_OPERATION]
+
+    # Agent span has no parent (it is the root)
+    assert agent_span.parent is None
+
+    # Chat span's parent must be the agent span
+    assert chat_span.parent is not None
+    assert chat_span.parent.span_id == agent_span.context.span_id
+    assert chat_span.parent.trace_id == agent_span.context.trace_id
+
+    # Both spans must share the same trace
+    assert chat_span.context.trace_id == agent_span.context.trace_id
+
+
+@pytest.mark.parametrize("stream", [False, True])
+async def test_function_call_spans_nested_under_agent_span(span_exporter: InMemorySpanExporter, stream: bool):
+    """All inner spans (chat completions and execute_tool) must be children of the agent span."""
+    from agent_framework import Content
+    from agent_framework._tools import FunctionInvocationLayer
+
+    @tool(name="get_weather", description="Get the weather for a location")
+    def get_weather(location: str) -> str:
+        return f"The weather in {location} is sunny."
+
+    class NestedToolChatClient(FunctionInvocationLayer, ChatTelemetryLayer, BaseChatClient[Any]):
+        def __init__(self) -> None:
+            super().__init__()
+            self.call_count = 0
+
+        def service_url(self):
+            return "https://test.example.com"
+
+        def _inner_get_response(
+            self, *, messages: MutableSequence[Message], stream: bool, options: dict[str, Any], **kwargs: Any
+        ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
+            self.call_count += 1
+            is_first = self.call_count == 1
+
+            if stream:
+
+                async def _stream() -> AsyncIterable[ChatResponseUpdate]:
+                    if is_first:
+                        yield ChatResponseUpdate(
+                            contents=[
+                                Content.from_function_call(
+                                    call_id="call_123",
+                                    name="get_weather",
+                                    arguments='{"location": "Seattle"}',
+                                )
+                            ],
+                            role="assistant",
+                        )
+                    else:
+                        yield ChatResponseUpdate(
+                            contents=[Content.from_text("The weather in Seattle is sunny!")],
+                            role="assistant",
+                            finish_reason="stop",
+                        )
+
+                def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
+                    return ChatResponse.from_updates(updates)
+
+                return ResponseStream(_stream(), finalizer=_finalize)
+
+            async def _get() -> ChatResponse:
+                if is_first:
+                    return ChatResponse(
+                        messages=[
+                            Message(
+                                role="assistant",
+                                contents=[
+                                    Content.from_function_call(
+                                        call_id="call_123",
+                                        name="get_weather",
+                                        arguments='{"location": "Seattle"}',
+                                    )
+                                ],
+                            )
+                        ],
+                    )
+                return ChatResponse(
+                    messages=[Message(role="assistant", contents=["The weather in Seattle is sunny!"])],
+                    finish_reason="stop",
+                )
+
+            return _get()
+
+    agent = Agent(
+        client=NestedToolChatClient(),
+        id="tool_agent_id",
+        name="tool_agent",
+        default_options={"model": "ToolModel", "tools": [get_weather], "tool_choice": "auto"},
+    )
+
+    span_exporter.clear()
+    if stream:
+        result_stream = agent.run("What's the weather in Seattle?", stream=True)
+        async for _ in result_stream:
+            pass
+        await result_stream.get_final_response()
+    else:
+        await agent.run("What's the weather in Seattle?")
+
+    spans = span_exporter.get_finished_spans()
+
+    invoke_spans = [s for s in spans if s.attributes.get(OtelAttr.OPERATION.value) == OtelAttr.AGENT_INVOKE_OPERATION]
+    chat_spans = [s for s in spans if s.attributes.get(OtelAttr.OPERATION.value) == OtelAttr.CHAT_COMPLETION_OPERATION]
+    tool_spans = [s for s in spans if s.attributes.get(OtelAttr.OPERATION.value) == OtelAttr.TOOL_EXECUTION_OPERATION]
+
+    assert len(invoke_spans) == 1, f"Expected 1 invoke_agent span, got {len(invoke_spans)}"
+    assert len(chat_spans) == 2, f"Expected 2 chat spans, got {len(chat_spans)}"
+    assert len(tool_spans) == 1, f"Expected 1 execute_tool span, got {len(tool_spans)}"
+
+    agent_span = invoke_spans[0]
+    assert agent_span.parent is None
+
+    # All inner spans must be parented under the agent invoke span
+    for inner in (*chat_spans, *tool_spans):
+        assert inner.parent is not None, f"Span {inner.name} has no parent"
+        assert inner.parent.span_id == agent_span.context.span_id, (
+            f"Span {inner.name} parent={inner.parent.span_id} != agent={agent_span.context.span_id}"
+        )
+        assert inner.context.trace_id == agent_span.context.trace_id
+
+
+@pytest.mark.parametrize("stream", [False, True])
+async def test_chat_span_nested_under_explicit_outer_span(
+    span_exporter: InMemorySpanExporter, mock_chat_client, stream: bool
+):
+    """Chat telemetry spans (including streaming) must inherit a user-provided outer span as parent."""
+    from agent_framework.observability import get_tracer
+
+    client = mock_chat_client()
+    span_exporter.clear()
+
+    tracer = get_tracer()
+    with tracer.start_as_current_span("outer") as outer_span:
+        outer_ctx = outer_span.get_span_context()
+        if stream:
+            stream_obj = client.get_response(
+                stream=True, messages=[Message(role="user", contents=["Test"])], options={"model": "Test"}
+            )
+            async for _ in stream_obj:
+                pass
+            await stream_obj.get_final_response()
+        else:
+            await client.get_response(messages=[Message(role="user", contents=["Test"])], options={"model": "Test"})
+
+    spans = span_exporter.get_finished_spans()
+    chat_spans = [s for s in spans if s.attributes.get(OtelAttr.OPERATION.value) == OtelAttr.CHAT_COMPLETION_OPERATION]
+    assert len(chat_spans) == 1
+    chat_span = chat_spans[0]
+
+    assert chat_span.parent is not None
+    assert chat_span.parent.span_id == outer_ctx.span_id
+    assert chat_span.context.trace_id == outer_ctx.trace_id
+
+
+@pytest.mark.parametrize("stream", [False, True])
+async def test_http_span_nested_under_chat_span(span_exporter: InMemorySpanExporter, stream: bool):
+    """A span created inside ``_inner_get_response`` (e.g. an HTTP client call to the LLM provider)
+    must be parented under the chat completion span.
+
+    This validates that the chat span context is active while the inner client implementation
+    runs, both for non-streaming responses and while streaming updates are being pulled.
+    """
+    from agent_framework.observability import get_tracer
+
+    tracer = get_tracer()
+
+    class HttpEmittingClient(ChatTelemetryLayer, BaseChatClient[Any]):
+        def service_url(self):
+            return "https://test.example.com"
+
+        def _inner_get_response(
+            self, *, messages: MutableSequence[Message], stream: bool, options: dict[str, Any], **kwargs: Any
+        ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
+            if stream:
+
+                async def _stream() -> AsyncIterable[ChatResponseUpdate]:
+                    # Simulate an HTTP request to the model provider while producing the stream.
+                    with tracer.start_as_current_span("HTTP POST"):
+                        pass
+                    yield ChatResponseUpdate(contents=[Content.from_text("hi")], role="assistant", finish_reason="stop")
+
+                def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
+                    return ChatResponse.from_updates(updates)
+
+                return ResponseStream(_stream(), finalizer=_finalize)
+
+            async def _get() -> ChatResponse:
+                # Simulate an HTTP request to the model provider during the call.
+                with tracer.start_as_current_span("HTTP POST"):
+                    pass
+                return ChatResponse(
+                    messages=[Message(role="assistant", contents=["done"])],
+                    usage_details=UsageDetails(input_token_count=1, output_token_count=1),
+                )
+
+            return _get()
+
+    span_exporter.clear()
+    client = HttpEmittingClient()
+    if stream:
+        result_stream = client.get_response(
+            stream=True, messages=[Message(role="user", contents=["Test"])], options={"model": "Test"}
+        )
+        async for _ in result_stream:
+            pass
+        await result_stream.get_final_response()
+    else:
+        await client.get_response(messages=[Message(role="user", contents=["Test"])], options={"model": "Test"})
+
+    spans = span_exporter.get_finished_spans()
+    chat_spans = [s for s in spans if s.attributes.get(OtelAttr.OPERATION.value) == OtelAttr.CHAT_COMPLETION_OPERATION]
+    http_spans = [s for s in spans if s.name == "HTTP POST"]
+    assert len(chat_spans) == 1
+    assert len(http_spans) == 1
+
+    chat_span = chat_spans[0]
+    http_span = http_spans[0]
+
+    assert http_span.parent is not None
+    assert http_span.parent.span_id == chat_span.context.span_id
+    assert http_span.context.trace_id == chat_span.context.trace_id