Python: Extend streaming span parenting to synchronous stream construction and additional resolution paths #5559

Context

Follow-up to #5552 (which fixed #5528). #5552 correctly parents OpenTelemetry spans during async iterator pulls of streaming responses by introducing ResponseStream.with_pull_context_manager() and activating it from __anext__. Two human reviewers approved and the primary bug (spans not nested under streaming) is resolved.

Automated reviewers (github-actions[bot] and Copilot) flagged a narrower set of legitimate edge cases that were intentionally deferred to keep #5552 scoped to the regression fix and unblock a release. This issue tracks those items so the streaming-span story is fully closed out.

Items to address

1. Activate the span around synchronous stream construction in both telemetry wrappers

with_pull_context_manager() only makes the span current during async pulls. However:

• BaseChatClient.get_response() can return a direct ResponseStream from _inner_get_response without deferring construction (python/packages/core/agent_framework/_clients.py:509-515). Some clients do non-trivial synchronous setup before returning the stream (e.g. Bedrock, python/packages/bedrock/agent_framework_bedrock/_chat_client.py:341-363).
• The agent invocation path has the same shape: execute() may construct a ResponseStream directly and real agents do synchronous setup before returning (e.g. A2A, python/packages/a2a/agent_framework_a2a/_agent.py:279-307).

Any child spans emitted during that synchronous setup phase will still be unparented. Suggested shape (from the bot review):

try:
    with _activate_span(span):
        result_stream = cast(
            ResponseStream[ChatResponseUpdate, ChatResponse[Any]],
            super_get_response(...),
        )
except Exception as exception:
    capture_exception(span=span, exception=exception, timestamp=time_ns())
    _close_span()
    raise

Apply the equivalent change around execute() in _trace_agent_invocation. Keep with_pull_context_manager() for the async pulls.

Files:

• python/packages/core/agent_framework/observability.py (chat streaming wrapper around line 1306; agent streaming wrapper around line 1574)

2. Ensure get_final_response() resolves the inner stream under pull contexts

with_pull_context_manager() only wraps __anext__(). For streams created via ResponseStream.from_awaitable() (_wrap_inner=True), get_final_response() pre-resolves the inner stream via _get_stream() before any pulls happen, bypassing the registered pull context managers. Callers that await stream.get_final_response() without first iterating can still create inner telemetry spans outside the intended active context.

Either:

• wrap _get_stream() so the registered factories are also applied during inner stream resolution, or
• defer resolution entirely to the first pull.

File:

• python/packages/core/agent_framework/_types.py (around line 3223 for with_pull_context_manager; line ~3095 for the get_final_response() resolution)

3. Additional test coverage

The four integration tests added in #5552 cover the happy parent/child nesting paths. Three gaps remain:

• Stream raising mid-iteration: a test where _inner_get_response yields one update then raises, asserting the pull context manager was still properly exited (no leaked OTel context).
• execute() failure during streaming agent invocation: the new error-handling path in _trace_agent_invocation that calls _close_span() when execute() raises is not exercised. Add a test asserting the span is ended and context vars are reset.
• get_final_response() without prior iteration: a test that calls await stream.get_final_response() directly on a from_awaitable stream (no prior await stream or iteration) and asserts the pull context managers are active during resolution. This is a distinct entry point from the existing __await__ test.

File:

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

References

• PR: Python: Fix spans not correctly nested when using streaming #5552
• Original bug: Python: Streaming path: Chat span is sibling of invoke_agent span instead of child (broken trace context propagation) #5528
• Automated review threads on Python: Fix spans not correctly nested when using streaming #5552 (commit b9926c0) contain the inline diff suggestions for items 1 and 2.