fix(openai): preserve async parsing for streaming responses (#607)

resolves
https://linear.app/braintrustdata/issue/SDK-86/async-streaming-response-parsing-not-preserved

## Summary

- await asynchronous raw-response parsing before wrapping OpenAI
Responses and Chat Completions streams
- preserve the awaitable `parse()` contract exposed by
`with_streaming_response`
- extend existing cassette-backed streaming tests to await and consume
the traced streams without mocks or new cassette churn

## Root cause

`openai.AsyncAPIResponse.parse()` is asynchronous, but the tracing proxy
replaced it with a synchronous method that returned the traced stream
directly. The wrappers also attempted to instrument the un-awaited
parser coroutine rather than the parsed iterator.

This broke callers such as OpenAI Agents that enter
`with_streaming_response` and then call `await api_response.parse()`.
Chat Completions had the same issue on its async raw streaming path.

The wrappers now detect and await asynchronous parsing before
constructing the traced iterator, remember the original parse contract,
and expose an awaitable `parse()` on the raw-response proxy. Synchronous
raw-response behavior remains unchanged.

## Validation

- reproduced the Chat Completions failure with existing VCR coverage:
`parse()` was not awaitable and the underlying parser coroutine was
never awaited
- full `test_openai(latest)` session with cassette playback: 70 tests
passed, plus attachment and OpenRouter gateway suites
- focused cassette playback for OpenAI 1.92.0, 1.77.0, and 1.71.0
- pre-commit hooks: Ruff format/check, cassette consistency, pytest-pin
synchronization, codespell, and whitespace checks

---------

Co-authored-by: Abhijeet Prasad <abhijeet@braintrustdata.com>

Author: cjgalione

py/src/braintrust/integrations/openai/test_openai.py

@@ -1066,16 +1066,23 @@ async def test_openai_chat_streaming_async(memory_logger):
     for client, is_wrapped in clients:
         start = time.time()
 
-        stream = await client.chat.completions.create(
-            model=TEST_MODEL,
-            messages=[{"role": "user", "content": TEST_PROMPT}],
-            stream=True,
-            stream_options={"include_usage": True},
-        )
-
-        chunks = []
-        async for chunk in stream:
-            chunks.append(chunk)
+        create_kwargs = {
+            "model": TEST_MODEL,
+            "messages": [{"role": "user", "content": TEST_PROMPT}],
+            "stream": True,
+            "stream_options": {"include_usage": True},
+        }
+        if is_wrapped:
+            async with client.chat.completions.with_streaming_response.create(**create_kwargs) as raw_response:
+                assert raw_response.headers
+                parse_result = raw_response.parse()
+                assert inspect.isawaitable(parse_result)
+                stream = await parse_result
+                assert stream.response
+                chunks = [chunk async for chunk in stream]
+        else:
+            stream = await client.chat.completions.create(**create_kwargs)
+            chunks = [chunk async for chunk in stream]
         end = time.time()
 
         assert chunks
@@ -1978,7 +1985,7 @@ async def test_openai_responses_with_raw_response_async(memory_logger):
 @pytest.mark.asyncio
 @pytest.mark.vcr
 async def test_openai_responses_with_raw_response_create_stream_async(memory_logger):
-    """Async version of test_openai_responses_with_raw_response_create_stream."""
+    """Async raw-response variants preserve headers, parsing, streams, and tracing."""
     assert not memory_logger.pop()
 
     # Unwrapped client: headers accessible, stream iterable via parse(), no spans.
@@ -1996,21 +2003,23 @@ async def test_openai_responses_with_raw_response_create_stream_async(memory_log
     assert "24" in "".join(chunks) or "twenty-four" in "".join(chunks).lower()
     assert not memory_logger.pop()
 
-    # Wrapped client: headers still accessible, parse() yields traced stream, span generated.
+    # Wrapped client: the newer streaming-response wrapper keeps async parse() and tracing.
     client = wrap_openai(AsyncOpenAI())
     start = time.time()
-    raw = await client.responses.with_raw_response.create(
+    async with client.responses.with_streaming_response.create(
         model=TEST_MODEL,
         input=TEST_PROMPT,
         stream=True,
-    )
-    assert raw.headers
-    stream = raw.parse()
-    assert stream.response  # SDK-specific attribute preserved
-    chunks = []
-    async for chunk in stream:
-        if chunk.type == "response.output_text.delta":
-            chunks.append(chunk.delta)
+    ) as raw:
+        assert raw.headers
+        parse_result = raw.parse()
+        assert inspect.isawaitable(parse_result)
+        stream = await parse_result
+        assert stream.response  # SDK-specific attribute preserved
+        chunks = []
+        async for chunk in stream:
+            if chunk.type == "response.output_text.delta":
+                chunks.append(chunk.delta)
     end = time.time()
     assert "24" in "".join(chunks) or "twenty-four" in "".join(chunks).lower()
 

py/src/braintrust/integrations/openai/tracing.py

@@ -499,8 +499,12 @@ async def acreate(self, *args: Any, **kwargs: Any) -> Any:
             start = time.time()
             create_response = await self.acreate_fn(*args, **kwargs)
 
+            parse_was_awaitable = False
             if hasattr(create_response, "parse"):
                 raw_response = create_response.parse()
+                if inspect.isawaitable(raw_response):
+                    parse_was_awaitable = True
+                    raw_response = await raw_response
                 log_headers(create_response, span)
             else:
                 raw_response = create_response
@@ -529,7 +533,11 @@ async def gen():
                 should_end = False
                 streamer = gen()
                 if raw_requested and hasattr(create_response, "parse"):
-                    return _RawResponseWithTracedStream(create_response, _AsyncTracedStream(raw_response, streamer))
+                    return _RawResponseWithTracedStream(
+                        create_response,
+                        _AsyncTracedStream(raw_response, streamer),
+                        async_parse=parse_was_awaitable,
+                    )
                 return _AsyncTracedStream(raw_response, streamer)
             else:
                 log_response = _try_to_dict(raw_response)
@@ -844,11 +852,17 @@ class _RawResponseWithTracedStream(NamedWrapper):
     """Proxy for LegacyAPIResponse that replaces parse() with a traced stream,
     so that with_raw_response + stream=True preserves both headers and tracing."""
 
-    def __init__(self, raw_response: Any, traced_stream: Any) -> None:
+    def __init__(self, raw_response: Any, traced_stream: Any, *, async_parse: bool = False) -> None:
         self._traced_stream = traced_stream
+        self._async_parse = async_parse
         super().__init__(raw_response)
 
     def parse(self, *args: Any, **kwargs: Any) -> Any:
+        if self._async_parse:
+            return self._aparse()
+        return self._traced_stream
+
+    async def _aparse(self) -> Any:
         return self._traced_stream
 
 
@@ -1066,8 +1080,12 @@ async def acreate(self, *args: Any, **kwargs: Any) -> Any:
         try:
             start = time.time()
             create_response = await self.acreate_fn(*args, **kwargs)
+            parse_was_awaitable = False
             if hasattr(create_response, "parse"):
                 raw_response = create_response.parse()
+                if inspect.isawaitable(raw_response):
+                    parse_was_awaitable = True
+                    raw_response = await raw_response
                 log_headers(create_response, span)
             else:
                 raw_response = create_response
@@ -1097,7 +1115,11 @@ async def gen():
                 should_end = False
                 streamer = gen()
                 if raw_requested and hasattr(create_response, "parse"):
-                    return _RawResponseWithTracedStream(create_response, _AsyncTracedStream(raw_response, streamer))
+                    return _RawResponseWithTracedStream(
+                        create_response,
+                        _AsyncTracedStream(raw_response, streamer),
+                        async_parse=parse_was_awaitable,
+                    )
                 return _AsyncTracedStream(raw_response, streamer)
             else:
                 log_response = _try_to_dict(raw_response)