From 36c87909d0c1652feb4efc544246199634bcc50e Mon Sep 17 00:00:00 2001 From: Madan Date: Mon, 17 Aug 2026 16:19:12 -0700 Subject: [PATCH 1/3] Python: fix streaming when tracing replaces the raw response (#7461) Setting AZURE_EXPERIMENTAL_ENABLE_GENAI_TRACING=true made every streaming request fail with: AttributeError: 'AsyncStreamWrapper' object has no attribute 'parse' The Azure GenAI instrumentor replaces the SDK's raw-response wrapper with an object that *is* the event stream and exposes neither .parse() nor .headers. Both streaming paths already read .headers defensively via getattr, with a comment explaining that instrumentors wrap the response -- but then called .parse() unconditionally on that same object. Read .parse defensively too: _open_event_stream() uses .parse() when present and otherwise iterates the object directly, letting the instrumentor own the stream's lifetime. Behavior is unchanged for the normal SDK wrapper. The non-streaming .parse() call sites are left alone: the instrumentor's wrapper is stream-specific and those paths were not reported as failing. Co-Authored-By: Claude Opus 5 (1M context) --- .../agent_framework_openai/_chat_client.py | 35 ++++++++- .../tests/openai/test_openai_chat_client.py | 74 +++++++++++++++++++ 2 files changed, 107 insertions(+), 2 deletions(-) diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index 2e4c06dff70..9c3712e4c03 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -7,6 +7,7 @@ import shlex import sys from collections.abc import ( + AsyncGenerator, AsyncIterable, Awaitable, Callable, @@ -14,6 +15,7 @@ MutableMapping, Sequence, ) +from contextlib import asynccontextmanager from datetime import datetime, timezone from itertools import chain from typing import ( @@ -295,6 +297,35 @@ class OpenAIChatOptions(ChatOptions[ResponseFormatT], Generic[ResponseFormatT], # region Helpers +@asynccontextmanager +async def _open_event_stream(raw_response: Any) -> AsyncGenerator[Any]: + """Yield the event stream for a raw streaming response. + + Normally ``raw_response`` is the SDK's raw-response wrapper, whose ``.parse()`` + returns the event stream as an async context manager so the underlying socket is + closed deterministically. + + Telemetry instrumentors replace that wrapper with an object that *is* the event + stream and exposes neither ``.parse()`` nor ``.headers`` -- for example the + ``AsyncStreamWrapper`` installed when ``AZURE_EXPERIMENTAL_ENABLE_GENAI_TRACING`` + is enabled. ``.headers`` is already read defensively at the call sites, so read + ``.parse`` defensively too and iterate such an object directly, letting the + instrumentor own the stream's lifetime. + + Args: + raw_response: The object returned by a ``with_raw_response`` streaming call. + + Yields: + The event stream to iterate. + """ + parse = getattr(raw_response, "parse", None) + if parse is None: + yield raw_response + return + async with parse() as stream: + yield stream + + def _annotations_to_output_text(annotations: Sequence[Annotation] | None) -> list[dict[str, Any]]: """Convert framework `Annotation` objects to Responses API `output_text` annotation dicts. @@ -694,7 +725,7 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: # proxy ``.headers``. Degrade gracefully so the served-model surfacing is # best-effort instead of crashing the whole call. served_model = self._extract_served_model(getattr(raw_stream_response, "headers", None)) - async with raw_stream_response.parse() as stream_response: + async with _open_event_stream(raw_stream_response) as stream_response: async for chunk in stream_response: update = self._parse_chunk_from_openai( chunk, @@ -738,7 +769,7 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: ) # See note above on ``raw_stream_response.headers``. served_model = self._extract_served_model(getattr(raw_create_response, "headers", None)) - async with raw_create_response.parse() as stream_response: + async with _open_event_stream(raw_create_response) as stream_response: async for chunk in stream_response: update = self._parse_chunk_from_openai( chunk, diff --git a/python/packages/openai/tests/openai/test_openai_chat_client.py b/python/packages/openai/tests/openai/test_openai_chat_client.py index cf3363174b7..85a608cf87e 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -791,6 +791,80 @@ async def test_served_model_header_propagated_to_streaming_updates() -> None: assert update.model == "gpt-4o-2024-08-06" +class _FakeInstrumentedStream: + """A raw streaming response as replaced by a telemetry instrumentor. + + Stands in for the ``AsyncStreamWrapper`` installed when + ``AZURE_EXPERIMENTAL_ENABLE_GENAI_TRACING`` is enabled: the object *is* the event + stream and deliberately exposes neither ``parse`` nor ``headers``. + """ + + def __init__(self, events: Sequence[object]) -> None: + self._events = list(events) + self._iterator: Iterator[object] = iter(()) + + def __aiter__(self) -> "_FakeInstrumentedStream": + self._iterator = iter(self._events) + return self + + async def __anext__(self) -> object: + try: + return next(self._iterator) + except StopIteration as exc: + raise StopAsyncIteration from exc + + +async def test_streaming_survives_instrumented_response_without_parse() -> None: + """Streaming should work when tracing replaces the raw response with a bare stream. + + Regression test for #7461: the client read ``.headers`` defensively but called + ``.parse()`` unconditionally, so enabling Azure GenAI tracing raised + ``AttributeError: 'AsyncStreamWrapper' object has no attribute 'parse'`` and every + streaming request failed. + """ + client = OpenAIChatClient(model="test-model", api_key="test-key") + + events = [ + ResponseTextDeltaEvent( + type="response.output_text.delta", + content_index=0, + item_id="text_item", + output_index=0, + sequence_number=1, + logprobs=[], + delta="Hello", + ), + ResponseTextDeltaEvent( + type="response.output_text.delta", + content_index=0, + item_id="text_item", + output_index=0, + sequence_number=2, + logprobs=[], + delta=" world", + ), + ] + + instrumented = _FakeInstrumentedStream(events) + assert not hasattr(instrumented, "parse") + assert not hasattr(instrumented, "headers") + + with ( + patch.object(client, "_prepare_request", new=AsyncMock(return_value=(client.client, {}, {}))), + patch.object(client.client.responses.with_raw_response, "create", new=AsyncMock(return_value=instrumented)), + patch.object(client, "_get_metadata_from_response", return_value={}), + ): + stream = _as_chat_response_stream( + client._inner_get_response(messages=[Message(role="user", contents=["Hi"])], options={}, stream=True) + ) + updates = [update async for update in stream] + + assert "".join(update.text or "" for update in updates) == "Hello world" + # No served-model header is available on an instrumented stream, so updates keep + # the deployment alias rather than failing. + assert all(update.model == "test-model" for update in updates) + + async def test_served_model_header_aggregates_into_final_streaming_response() -> None: """Aggregating updates via to_chat_response() should preserve the served-model value.""" client = OpenAIChatClient(model="test-model", api_key="test-key") From 4407ab0f996433ebf738047d5632c6f639012d35 Mon Sep 17 00:00:00 2001 From: Madan Date: Mon, 17 Aug 2026 20:55:52 -0700 Subject: [PATCH 2/3] Parse the raw response the telemetry wrapper holds Review feedback: yielding the telemetry wrapper as-is only moved the AttributeError. Verified against azure-ai-projects==2.3.0 with openai==2.53.0: with_raw_response.create() routes through the instrumented AsyncResponses.create, so AsyncStreamWrapper.stream_async_iter is the still-unparsed LegacyAPIResponse, which is not an async iterator. Iterating the wrapper fails on the first __anext__ and traced streaming stays broken. Parse that inner raw response and hand it back to the wrapper instead, so the wrapper stays in the iteration path and keeps recording telemetry while real events flow through it. The previous test patched with_raw_response.create, i.e. above the layer that does the wrapping, so it could not catch this. The test now models the observed object graph -- a wrapper with no parse/headers whose stream_async_iter is an unparsed raw response -- and fails against the previous fix. A second test covers a bare event stream with nothing to parse. Also guard with callable() rather than an is-None check, per review. Co-Authored-By: Claude Opus 5 (1M context) --- .../agent_framework_openai/_chat_client.py | 38 ++++-- .../tests/openai/test_openai_chat_client.py | 114 +++++++++++++++--- 2 files changed, 121 insertions(+), 31 deletions(-) diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index 9c3712e4c03..55ff20950ad 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -305,25 +305,41 @@ async def _open_event_stream(raw_response: Any) -> AsyncGenerator[Any]: returns the event stream as an async context manager so the underlying socket is closed deterministically. - Telemetry instrumentors replace that wrapper with an object that *is* the event - stream and exposes neither ``.parse()`` nor ``.headers`` -- for example the - ``AsyncStreamWrapper`` installed when ``AZURE_EXPERIMENTAL_ENABLE_GENAI_TRACING`` - is enabled. ``.headers`` is already read defensively at the call sites, so read - ``.parse`` defensively too and iterate such an object directly, letting the - instrumentor own the stream's lifetime. + A telemetry instrumentor can replace that wrapper with one of its own that is + itself the async iterator and exposes neither ``.parse()`` nor ``.headers`` -- for + example the ``AsyncStreamWrapper`` installed by ``azure-ai-projects`` when + ``AZURE_EXPERIMENTAL_ENABLE_GENAI_TRACING`` is enabled. That wrapper still holds + the unparsed raw response (its ``stream_async_iter``), because + ``with_raw_response.create()`` routes through the instrumented ``create``. + + Parse that inner raw response and hand it back to the wrapper, so the wrapper + stays in the iteration path and keeps recording telemetry while we iterate real + events. Iterating the wrapper as-is would fail, since the unparsed raw response + is not an async iterator. Args: raw_response: The object returned by a ``with_raw_response`` streaming call. Yields: - The event stream to iterate. + The object to iterate for streaming events. """ parse = getattr(raw_response, "parse", None) - if parse is None: - yield raw_response + if callable(parse): + async with cast("Any", parse()) as stream: + yield stream return - async with parse() as stream: - yield stream + + # Telemetry wrapper: parse the raw response it wraps, in place. + inner = getattr(raw_response, "stream_async_iter", None) + inner_parse = getattr(inner, "parse", None) + if callable(inner_parse): + async with cast("Any", inner_parse()) as stream: + raw_response.stream_async_iter = stream + yield raw_response + return + + # Already an event stream (or an unrecognized wrapper): iterate it directly. + yield raw_response def _annotations_to_output_text(annotations: Sequence[Annotation] | None) -> list[dict[str, Any]]: diff --git a/python/packages/openai/tests/openai/test_openai_chat_client.py b/python/packages/openai/tests/openai/test_openai_chat_client.py index 85a608cf87e..121227c12fb 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -791,36 +791,56 @@ async def test_served_model_header_propagated_to_streaming_updates() -> None: assert update.model == "gpt-4o-2024-08-06" -class _FakeInstrumentedStream: - """A raw streaming response as replaced by a telemetry instrumentor. +class _UnparsedRawResponse: + """The still-unparsed raw response a telemetry wrapper captures. - Stands in for the ``AsyncStreamWrapper`` installed when - ``AZURE_EXPERIMENTAL_ENABLE_GENAI_TRACING`` is enabled: the object *is* the event - stream and deliberately exposes neither ``parse`` nor ``headers``. + Like ``LegacyAPIResponse``: it exposes ``parse()`` but is not an async iterator. """ - def __init__(self, events: Sequence[object]) -> None: - self._events = list(events) - self._iterator: Iterator[object] = iter(()) + def __init__(self, parsed: object) -> None: + self._parsed = parsed - def __aiter__(self) -> "_FakeInstrumentedStream": - self._iterator = iter(self._events) + def parse(self) -> object: + return self._parsed + + +class _FakeTelemetryStreamWrapper: + """The wrapper a telemetry instrumentor substitutes for the raw-response wrapper. + + Mirrors the ``AsyncStreamWrapper`` that ``azure-ai-projects`` installs when + ``AZURE_EXPERIMENTAL_ENABLE_GENAI_TRACING`` is enabled, as observed against + azure-ai-projects==2.3.0: + + * it exposes neither ``parse`` nor ``headers``; + * it is the async iterator, and ``__anext__`` delegates to ``stream_async_iter``; + * because ``with_raw_response.create()`` routes through the instrumented + ``create``, ``stream_async_iter`` is the still-unparsed raw response + (a ``LegacyAPIResponse``), which is *not* itself an async iterator. + + So iterating this wrapper as handed over raises ``AttributeError`` until the + inner raw response is parsed and handed back. + """ + + def __init__(self, stream_async_iter: object) -> None: + self.stream_async_iter = stream_async_iter + + def __aiter__(self) -> "_FakeTelemetryStreamWrapper": + self.stream_async_iter = self.stream_async_iter.__aiter__() # type: ignore[attr-defined] return self async def __anext__(self) -> object: - try: - return next(self._iterator) - except StopIteration as exc: - raise StopAsyncIteration from exc + return await self.stream_async_iter.__anext__() # type: ignore[attr-defined] -async def test_streaming_survives_instrumented_response_without_parse() -> None: - """Streaming should work when tracing replaces the raw response with a bare stream. +async def test_streaming_survives_telemetry_wrapped_raw_response() -> None: + """Streaming should work when tracing replaces the raw-response wrapper. - Regression test for #7461: the client read ``.headers`` defensively but called + Regression test for #7461. The client read ``.headers`` defensively but called ``.parse()`` unconditionally, so enabling Azure GenAI tracing raised - ``AttributeError: 'AsyncStreamWrapper' object has no attribute 'parse'`` and every - streaming request failed. + ``AttributeError: 'AsyncStreamWrapper' object has no attribute 'parse'``. + + The telemetry wrapper must stay in the iteration path so it still records + telemetry, while the raw response it wraps gets parsed into real events. """ client = OpenAIChatClient(model="test-model", api_key="test-key") @@ -845,7 +865,12 @@ async def test_streaming_survives_instrumented_response_without_parse() -> None: ), ] - instrumented = _FakeInstrumentedStream(events) + # The unparsed raw response the instrumentor captured: it has .parse() but is not + # an async iterator, exactly like LegacyAPIResponse. + unparsed_raw = _UnparsedRawResponse(_FakeAsyncEventStream(events)) + assert not hasattr(unparsed_raw, "__anext__") + + instrumented = _FakeTelemetryStreamWrapper(unparsed_raw) assert not hasattr(instrumented, "parse") assert not hasattr(instrumented, "headers") @@ -860,11 +885,60 @@ async def test_streaming_survives_instrumented_response_without_parse() -> None: updates = [update async for update in stream] assert "".join(update.text or "" for update in updates) == "Hello world" + # The telemetry wrapper stays in the iteration path rather than being bypassed. + assert isinstance(instrumented.stream_async_iter, _FakeAsyncEventStream) # No served-model header is available on an instrumented stream, so updates keep # the deployment alias rather than failing. assert all(update.model == "test-model" for update in updates) +async def test_streaming_accepts_raw_response_that_is_already_an_event_stream() -> None: + """An object with no ``parse`` and no wrapped raw response is iterated directly.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + + events = [ + ResponseTextDeltaEvent( + type="response.output_text.delta", + content_index=0, + item_id="text_item", + output_index=0, + sequence_number=1, + logprobs=[], + delta="Hello", + ), + ] + + class _BareEventStream: + def __init__(self, items: Sequence[object]) -> None: + self._items = list(items) + self._iterator: Iterator[object] = iter(()) + + def __aiter__(self) -> "_BareEventStream": + self._iterator = iter(self._items) + return self + + async def __anext__(self) -> object: + try: + return next(self._iterator) + except StopIteration as exc: + raise StopAsyncIteration from exc + + bare = _BareEventStream(events) + assert not hasattr(bare, "parse") + + with ( + patch.object(client, "_prepare_request", new=AsyncMock(return_value=(client.client, {}, {}))), + patch.object(client.client.responses.with_raw_response, "create", new=AsyncMock(return_value=bare)), + patch.object(client, "_get_metadata_from_response", return_value={}), + ): + stream = _as_chat_response_stream( + client._inner_get_response(messages=[Message(role="user", contents=["Hi"])], options={}, stream=True) + ) + updates = [update async for update in stream] + + assert "".join(update.text or "" for update in updates) == "Hello" + + async def test_served_model_header_aggregates_into_final_streaming_response() -> None: """Aggregating updates via to_chat_response() should preserve the served-model value.""" client = OpenAIChatClient(model="test-model", api_key="test-key") From c85069cdc244aff849d9af936e412249f348c04d Mon Sep 17 00:00:00 2001 From: Madan Date: Tue, 18 Aug 2026 16:02:59 -0700 Subject: [PATCH 3/3] Fix test typing failures (ty, zuban) Test Typing Checks caught two problems in the new test helpers: - ty: the wrapper's stream_async_iter was annotated `object`, so delegating to __aiter__/__anext__ was an attribute error. Annotate it `Any`, which also makes the two `type: ignore` comments unnecessary. - zuban: `_BareEventStream` was declared inside the test function, and its own forward-referenced return annotation does not resolve there. Move it to module scope alongside the other stream fakes. Verified with the task CI runs, `poe test-typing -P openai`: mypy, pyrefly, ty, zuban and pyright all pass. Co-Authored-By: Claude Opus 5 (1M context) --- .../tests/openai/test_openai_chat_client.py | 41 ++++++++++--------- 1 file changed, 22 insertions(+), 19 deletions(-) diff --git a/python/packages/openai/tests/openai/test_openai_chat_client.py b/python/packages/openai/tests/openai/test_openai_chat_client.py index 121227c12fb..aab0fc1fe24 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -804,6 +804,24 @@ def parse(self) -> object: return self._parsed +class _BareEventStream: + """An object that is already the event stream: no ``parse``, nothing to unwrap.""" + + def __init__(self, items: Sequence[object]) -> None: + self._items = list(items) + self._iterator: Iterator[object] = iter(()) + + def __aiter__(self) -> "_BareEventStream": + self._iterator = iter(self._items) + return self + + async def __anext__(self) -> object: + try: + return next(self._iterator) + except StopIteration as exc: + raise StopAsyncIteration from exc + + class _FakeTelemetryStreamWrapper: """The wrapper a telemetry instrumentor substitutes for the raw-response wrapper. @@ -821,15 +839,15 @@ class _FakeTelemetryStreamWrapper: inner raw response is parsed and handed back. """ - def __init__(self, stream_async_iter: object) -> None: - self.stream_async_iter = stream_async_iter + def __init__(self, stream_async_iter: Any) -> None: + self.stream_async_iter: Any = stream_async_iter def __aiter__(self) -> "_FakeTelemetryStreamWrapper": - self.stream_async_iter = self.stream_async_iter.__aiter__() # type: ignore[attr-defined] + self.stream_async_iter = self.stream_async_iter.__aiter__() return self async def __anext__(self) -> object: - return await self.stream_async_iter.__anext__() # type: ignore[attr-defined] + return await self.stream_async_iter.__anext__() async def test_streaming_survives_telemetry_wrapped_raw_response() -> None: @@ -908,21 +926,6 @@ async def test_streaming_accepts_raw_response_that_is_already_an_event_stream() ), ] - class _BareEventStream: - def __init__(self, items: Sequence[object]) -> None: - self._items = list(items) - self._iterator: Iterator[object] = iter(()) - - def __aiter__(self) -> "_BareEventStream": - self._iterator = iter(self._items) - return self - - async def __anext__(self) -> object: - try: - return next(self._iterator) - except StopIteration as exc: - raise StopAsyncIteration from exc - bare = _BareEventStream(events) assert not hasattr(bare, "parse")