[opentelemetry-instrumentation-genai-openai] Record response telemetry for async with_streaming_response - #610
Conversation
…try for async `with_streaming_response` `RawResponseStreamProxy.parse()` assumed the wrapped response returned a stream. That holds for `LegacyAPIResponse` and `APIResponse`, but the async client's `with_streaming_response` yields an `AsyncAPIResponse`, whose `parse()` is a coroutine. The isinstance check rejected it, so the caller awaited an uninstrumented stream and the span closed carrying only request-side attributes: no output messages, usage, finish reasons, response id or response model. This is the path the OpenAI Agents SDK takes for streamed runs, which resolves `responses.with_streaming_response.create` and then awaits `parse()`, so every streamed agent run produced an empty chat span. Await the coroutine and wrap what it resolves to. `parse()` stays awaitable, so callers are unaffected. Since `_raw_response` is shared, this covers both chat completions and responses.
There was a problem hiding this comment.
🟡 Changes recommended
The proxy’s async parse() path can return a non-awaitable memoized value on subsequent calls, which can break callers that consistently do await raw_response.parse().
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Fixes missing response telemetry for AsyncOpenAI.with_streaming_response by teaching the raw-response proxy to handle AsyncAPIResponse.parse() returning a coroutine, and adds regression tests to ensure streamed response attributes/usage are recorded.
Changes:
- Update
RawResponseStreamProxy.parse()to await an asyncparse()result and then wrap the resolved SDK stream. - Add integration tests covering
with_streaming_response.create(...).parse()for sync Responses plus async Responses and async Chat Completions. - Add unit tests for the “awaitable parse” branch, including close-fallback behavior when
parse()is never awaited.
File summaries
| File | Description |
|---|---|
| instrumentation/opentelemetry-instrumentation-genai-openai/src/opentelemetry/instrumentation/genai/openai/_raw_response.py | Await coroutine-returning parse() and wrap the resolved stream. |
| instrumentation/opentelemetry-instrumentation-genai-openai/tests/test_raw_response_proxy.py | Add unit coverage for awaitable parse() and close fallback cases. |
| instrumentation/opentelemetry-instrumentation-genai-openai/tests/test_responses.py | Add sync Responses regression test for with_streaming_response.parse(). |
| instrumentation/opentelemetry-instrumentation-genai-openai/tests/test_async_responses.py | Add async Responses regression test for coroutine parse() path. |
| instrumentation/opentelemetry-instrumentation-genai-openai/tests/test_async_chat_completions.py | Add async Chat Completions regression test for coroutine parse() path. |
| instrumentation/opentelemetry-instrumentation-genai-openai/.changelog/610.fixed | Add changelog entry for the telemetry fix. |
Review details
- Files reviewed: 6/6 changed files
- Comments generated: 1
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Pull request dashboard statusWaiting on maintainers · refreshed 2026-09-07 07:28 UTC Merge when ready. Status above doesn't look right?
|
The memo check at the top of parse() returned the wrapper directly, so a second call on an async response handed back a non-awaitable value and the caller's await raised TypeError. The wrapper type cannot tell us whether parse() is async: the async client pairs a synchronous LegacyAPIResponse.parse() with an AsyncStream. So record whether the wrapped parse() actually returned an awaitable, and keep memo hits awaitable when it did.
bc2d61e to
dfda651
Compare
|
great catch!! Just for my learning sake, how did you reproduce this? |
@eternalcuriouslearner thanks for the nudge! this should have had a tracking issue anyway, so I opened #624 with the symptom and reproduction steps. |
| stream = await pending | ||
| # parse() can hand out several coroutines before any of them resolves, | ||
| # so re-check the memo to keep one wrapper per response. | ||
| if self._self_parsed is not None: |
There was a problem hiding this comment.
Race: if the httpx response closes while this coroutine is still pending (e.g. the caller calls parse() without immediately awaiting it, as your own test_close_before_awaiting_parse_finalizes_once does), the close fallback already finalizes the span here, empty. When this coroutine then resolves anyway, _wrap_parsed still builds and returns a stream wrapper; draining it calls invocation.stop() again, which silently no-ops (span already ended), so whatever attributes it collected are dropped.
Fix: in _wrap_parsed, check whether the close fallback already fired before building the wrapper:
def _wrap_parsed(self, stream: Any) -> object:
if isinstance(stream, (Stream, AsyncStream)):
if self._self_finalize is None:
# Closed while parse() was pending; span is already gone,
# wrapping now would only drop attributes silently.
return stream
self._self_parsed = self._self_wrap_stream(stream)
return self._self_parsed
..._self_finalize is only ever cleared by the close fallback, so this is a reliable signal without adding new state.
Regression test (fails on this branch, passes with the fix above):
@pytest.mark.asyncio()
async def test_close_while_parse_pending_does_not_lose_the_stream():
resume = asyncio.Event()
stream = _FakeAsyncStream()
class _SlowAsyncRawResponse(_AsyncRawResponse):
async def parse(self, *, to=None):
await resume.wait()
return self._parse_result
raw = _SlowAsyncRawResponse(stream)
finalize_calls, wrap_calls = [], []
proxy = RawResponseStreamProxy(
raw,
wrap_stream=lambda s: wrap_calls.append(s) or ("wrapped", s),
finalize=lambda: finalize_calls.append(True),
)
pending = proxy.parse()
raw.http_response.close()
assert finalize_calls == [True] # close fallback already finalized
resume.set()
result = await pending
assert wrap_calls == [] # must not silently drop a live wrapper's data
assert result is streamThere was a problem hiding this comment.
Sorry for the deleted comment, I posted a suggestion and then went to verify it, and that turned out to be wrong. I've updated according to your suggestion, thanks!
When parse() is called but not awaited before the httpx response closes, the close fallback sees no wrapper and finalizes the span. If the coroutine then resolved to a stream, _wrap_parsed still built a wrapper: it would collect attributes and call stop() on an invocation that had already finished, so the response data was silently dropped. Hand the stream back uninstrumented in that case. _self_finalize is cleared only by the close fallback, so it already records whether the span is gone.
lmolkova
left a comment
lmolkova
left a comment
There was a problem hiding this comment.
Thanks for the update!
One more edge case and we should be good. I tend to ask AI to find edge cases and prove with failing test, it turns out to be relatively good at finding them around streaming.
| # Only ``AsyncAPIResponse`` (the async ``with_streaming_response``) | ||
| # has a coroutine ``parse()``. | ||
| self._self_parse_returns_awaitable = True | ||
| return self._parse_awaited(parsed) |
There was a problem hiding this comment.
This pulls the async path into #491: a stream the caller doesn't drain now emits no span at all, where before the close fallback still ended it (empty). Once _self_parsed is set the fallback steps aside and nothing else finalizes.
Fails here, passes on main:
@pytest.mark.asyncio()
async def test_abandoned_async_streaming_response_still_emits_span(
span_exporter, async_openai_client, instrument_with_content, vcr
):
with vcr.use_cassette("test_chat_completion_with_raw_response_streaming.yaml"):
async with async_openai_client.chat.completions.with_streaming_response.create(
messages=USER_ONLY_PROMPT, model=DEFAULT_MODEL, stream=True,
stream_options={"include_usage": True},
) as raw_response:
async for _chunk in await raw_response.parse():
break
assert len(span_exporter.get_finished_spans()) == 1#491 is hard in general because a plain stream gives no signal that the caller left. Here there is one - __aexit__ closes the http response - and the anthropic package already uses it: see _finalize_close_fallback / _afinalize_close_fallback in instrumentation/opentelemetry-instrumentation-genai-anthropic/src/opentelemetry/instrumentation/genai/anthropic/_raw_response.py, which close the stream wrapper instead of suppressing themselves. Both stream wrappers guard on _self_finalized, so a drained stream is unaffected.
There was a problem hiding this comment.
thanks! Rethinking this: instead of tracking the wrapper separately, I reused the existing _self_finalize callback. It starts as invocation.stop and is swapped for wrapper.close once parse() builds one, so whichever close hook fires just runs whatever is currently there. The field is cleared when the callback is taken, so finalization happens once and the re-entrant close finds nothing left to run.
Once parse() built a wrapper, the close fallback stood aside on the grounds that the wrapper owns finalization. But a wrapper only finalizes when it is drained, closed, or used as a context manager, so a caller that breaks out of the loop or raises inside the with block left nobody to end the span. On the async path this was a regression from wrapping the coroutine parse(): before, no wrapper was built and the fallback still ended the span, empty. The hooks now drive the wrapper instead of standing aside, so it finalizes with what it saw. Closing it closes the SDK stream, which closes this same httpx response and re-enters the hook. Rather than guard that separately, swap the existing take-once finalize callback for the wrapper close: taking it clears it, so finalization stays idempotent and the re-entrant call finds nothing to run.
b9e25c1 to
43cb3b5
Compare
Description
RawResponseStreamProxy.parse()assumed the wrapped response hands back astream. That holds for
LegacyAPIResponseandAPIResponse, but the asyncclient's
with_streaming_responseyields anAsyncAPIResponse, whoseparse()is a coroutine. The
isinstancecheck rejected it, so the caller awaited anuninstrumented stream: no
gen_ai.output.messages,gen_ai.usage.*,gen_ai.response.finish_reasons,gen_ai.response.idorgen_ai.response.model. The close fallback still ended the span with statusUNSET, so it looked healthy while carrying nothing.Of the four raw-response entry points, only this one is affected:
parse()OpenAIwith_raw_responseLegacyAPIResponseAsyncOpenAIwith_raw_responseLegacyAPIResponseOpenAIwith_streaming_responseAPIResponseAsyncOpenAIwith_streaming_responseAsyncAPIResponseIt is also the path the OpenAI Agents SDK takes for streamed runs
(
agents/models/openai_responses.pyresolvesresponses.with_streaming_response.create, then awaitsparse()), so everyRunner.run_streamedproduced an empty chat span.parse()now awaits the coroutine and wraps what it resolves to. It staysawaitable, so callers are unchanged.
_raw_responseis shared, so this coverschat completions and responses alike.
Type of change
How has this been tested?
New tests for
with_streaming_response.create()on responses (sync and async)and chat completions (async), plus three unit tests in
test_raw_response_proxy.pycovering the awaitable branch: wrapping andmemoizing the awaited stream, handing back a non-stream untouched, and the
close fallback still finalizing when the coroutine is never awaited.
The two async tests fail on
mainand pass here; the sync one passes on both,since that path was already correct.
openaisuite: 252 passed onpy310-latestandpy314-latestpy310-oldest: 113 passed; the two Responses tests skip, as the API doesnot exist at the declared floor
tox -e precommitandtox -e typecheckcleanChecklist
Fixes #624