braintrustdata · Abhijeet Prasad (AbhiPrasad) · Mar 23, 2026
diff --git a/py/noxfile.py b/py/noxfile.py
@@ -210,6 +210,16 @@ def test_openai(session, version):
     _run_core_tests(session)
 
 
+@nox.session()
+def test_openai_http2_streaming(session):
+    _install_test_deps(session)
+    _install(session, "openai")
+    # h2 is isolated to this session because it's only needed to force the
+    # HTTP/2 LegacyAPIResponse streaming path used by the regression test.
+    session.install("h2")
+    _run_tests(session, f"{WRAPPER_DIR}/test_openai_http2.py")
+
+
 @nox.session()
 def test_openrouter(session):
     """Test wrap_openai with OpenRouter. Requires OPENROUTER_API_KEY env var."""

diff --git a/py/src/braintrust/oai.py b/py/src/braintrust/oai.py
@@ -19,10 +19,13 @@
 
 class NamedWrapper:
     def __init__(self, wrapped: Any):
+        self._wrapped = wrapped
+        # Keep the legacy mangled attribute for existing wrapped-client checks
+        # that introspect `_NamedWrapper__wrapped` directly.
         self.__wrapped = wrapped
 
     def __getattr__(self, name: str) -> Any:
-        return getattr(self.__wrapped, name)
+        return getattr(self._wrapped, name)
 
 
 class AsyncResponseWrapper:
@@ -188,7 +191,7 @@ def gen():
                         span.end()
 
                 should_end = False
-                return gen()
+                return _TracedStream(raw_response, gen())
             else:
                 log_response = _try_to_dict(raw_response)
                 metrics = _parse_metrics_from_usage(log_response.get("usage", {}))
@@ -244,7 +247,7 @@ async def gen():
 
                 should_end = False
                 streamer = gen()
-                return AsyncResponseWrapper(streamer)
+                return _AsyncTracedStream(raw_response, streamer)
             else:
                 log_response = _try_to_dict(raw_response)
                 metrics = _parse_metrics_from_usage(log_response.get("usage"))
@@ -365,6 +368,16 @@ def __iter__(self) -> Any:
     def __next__(self) -> Any:
         return next(self._traced_generator)
 
+    def __enter__(self) -> Any:
+        if hasattr(self._wrapped, "__enter__"):
+            self._wrapped.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> Any:
+        if hasattr(self._wrapped, "__exit__"):
+            return self._wrapped.__exit__(exc_type, exc_val, exc_tb)
+        return None
+
 
 class _AsyncTracedStream(NamedWrapper):
     """Traced async stream. Iterates via the traced generator while delegating
@@ -380,6 +393,16 @@ def __aiter__(self) -> Any:
     async def __anext__(self) -> Any:
         return await self._traced_generator.__anext__()
 
+    async def __aenter__(self) -> Any:
+        if hasattr(self._wrapped, "__aenter__"):
+            await self._wrapped.__aenter__()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> Any:
+        if hasattr(self._wrapped, "__aexit__"):
+            return await self._wrapped.__aexit__(exc_type, exc_val, exc_tb)
+        return None
+
 
 class _RawResponseWithTracedStream(NamedWrapper):
     """Proxy for LegacyAPIResponse that replaces parse() with a traced stream,
@@ -445,7 +468,7 @@ def gen():
                 should_end = False
                 if self.return_raw and hasattr(create_response, "parse"):
                     return _RawResponseWithTracedStream(create_response, _TracedStream(raw_response, gen()))
-                return gen()
+                return _TracedStream(raw_response, gen())
             else:
                 log_response = _try_to_dict(raw_response)
                 event_data = self._parse_event_from_result(log_response)
@@ -498,7 +521,7 @@ async def gen():
                 streamer = gen()
                 if self.return_raw and hasattr(create_response, "parse"):
                     return _RawResponseWithTracedStream(create_response, _AsyncTracedStream(raw_response, streamer))
-                return AsyncResponseWrapper(streamer)
+                return _AsyncTracedStream(raw_response, streamer)
             else:
                 log_response = _try_to_dict(raw_response)
                 event_data = self._parse_event_from_result(log_response)

diff --git a/...s/cassettes/test_openai_chat_streaming_async_http2_context_manager_preserves_wrapper.yaml b/...s/cassettes/test_openai_chat_streaming_async_http2_context_manager_preserves_wrapper.yaml
@@ -0,0 +1,271 @@
+interactions:
+  - request:
+      body: '{"messages":[{"role":"user","content":"What''s 12 + 12?"}],"model":"gpt-4o-mini","stream":true,"stream_options":{"include_usage":true}}'
+      headers:
+        accept:
+          - application/json
+        accept-encoding:
+          - gzip, deflate
+        connection:
+          - keep-alive
+        content-length:
+          - "134"
+        content-type:
+          - application/json
+        host:
+          - api.openai.com
+        user-agent:
+          - AsyncOpenAI/Python 1.82.0
+        x-stainless-arch:
+          - arm64
+        x-stainless-async:
+          - async:asyncio
+        x-stainless-lang:
+          - python
+        x-stainless-os:
+          - MacOS
+        x-stainless-package-version:
+          - 1.82.0
+        x-stainless-read-timeout:
+          - "600"
+        x-stainless-retry-count:
+          - "0"
+        x-stainless-runtime:
+          - CPython
+        x-stainless-runtime-version:
+          - 3.13.3
+      method: POST
+      uri: https://api.openai.com/v1/chat/completions
+    response:
+      body:
+        string:
+          'data: {"id":"chatcmpl-BcNw6be6DuhpHqDlc8nKm7x3iMAL6","object":"chat.completion.chunk","created":1748488174,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[{"index":0,"delta":{"role":"assistant","content":"","refusal":null},"logprobs":null,"finish_reason":null}],"usage":null}
+
+
+          data: {"id":"chatcmpl-BcNw6be6DuhpHqDlc8nKm7x3iMAL6","object":"chat.completion.chunk","created":1748488174,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[{"index":0,"delta":{"content":"12"},"logprobs":null,"finish_reason":null}],"usage":null}
+
+
+          data: {"id":"chatcmpl-BcNw6be6DuhpHqDlc8nKm7x3iMAL6","object":"chat.completion.chunk","created":1748488174,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[{"index":0,"delta":{"content":"
+          +"},"logprobs":null,"finish_reason":null}],"usage":null}
+
+
+          data: {"id":"chatcmpl-BcNw6be6DuhpHqDlc8nKm7x3iMAL6","object":"chat.completion.chunk","created":1748488174,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[{"index":0,"delta":{"content":"
+          "},"logprobs":null,"finish_reason":null}],"usage":null}
+
+
+          data: {"id":"chatcmpl-BcNw6be6DuhpHqDlc8nKm7x3iMAL6","object":"chat.completion.chunk","created":1748488174,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[{"index":0,"delta":{"content":"12"},"logprobs":null,"finish_reason":null}],"usage":null}
+
+
+          data: {"id":"chatcmpl-BcNw6be6DuhpHqDlc8nKm7x3iMAL6","object":"chat.completion.chunk","created":1748488174,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[{"index":0,"delta":{"content":"
+          equals"},"logprobs":null,"finish_reason":null}],"usage":null}
+
+
+          data: {"id":"chatcmpl-BcNw6be6DuhpHqDlc8nKm7x3iMAL6","object":"chat.completion.chunk","created":1748488174,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[{"index":0,"delta":{"content":"
+          "},"logprobs":null,"finish_reason":null}],"usage":null}
+
+
+          data: {"id":"chatcmpl-BcNw6be6DuhpHqDlc8nKm7x3iMAL6","object":"chat.completion.chunk","created":1748488174,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[{"index":0,"delta":{"content":"24"},"logprobs":null,"finish_reason":null}],"usage":null}
+
+
+          data: {"id":"chatcmpl-BcNw6be6DuhpHqDlc8nKm7x3iMAL6","object":"chat.completion.chunk","created":1748488174,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[{"index":0,"delta":{"content":"."},"logprobs":null,"finish_reason":null}],"usage":null}
+
+
+          data: {"id":"chatcmpl-BcNw6be6DuhpHqDlc8nKm7x3iMAL6","object":"chat.completion.chunk","created":1748488174,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[{"index":0,"delta":{},"logprobs":null,"finish_reason":"stop"}],"usage":null}
+
+
+          data: {"id":"chatcmpl-BcNw6be6DuhpHqDlc8nKm7x3iMAL6","object":"chat.completion.chunk","created":1748488174,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[],"usage":{"prompt_tokens":14,"completion_tokens":8,"total_tokens":22,"prompt_tokens_details":{"cached_tokens":0,"audio_tokens":0},"completion_tokens_details":{"reasoning_tokens":0,"audio_tokens":0,"accepted_prediction_tokens":0,"rejected_prediction_tokens":0}}}
+
+
+          data: [DONE]
+
+
+          '
+      headers:
+        CF-RAY:
+          - 9472cb30e8d5e56c-EWR
+        Connection:
+          - keep-alive
+        Content-Type:
+          - text/event-stream; charset=utf-8
+        Date:
+          - Thu, 29 May 2025 03:09:34 GMT
+        Server:
+          - cloudflare
+        Set-Cookie:
+          - __cf_bm=AIq4yjX3ru.J9gwNR6zSYLQNorConUZa5qtJ6wXxuvE-1748488174-1.0.1.1-KquMaoYitsL5z76ow2IPzasSn98mtC1_QEt9VOT1pvvQt_obPUDugNtsEGJCc_wP50_X4wP.kC7nYuf98KX8dCPpiq2ZqY5vwVCdgocqRxU;
+            path=/; expires=Thu, 29-May-25 03:39:34 GMT; domain=.api.openai.com; HttpOnly;
+            Secure; SameSite=None
+          - _cfuvid=o0LrLIiV.VvLFX1H1bbtbV01AjzSfXrfrVn0fU7pANY-1748488174648-0.0.1.1-604800000;
+            path=/; domain=.api.openai.com; HttpOnly; Secure; SameSite=None
+        Transfer-Encoding:
+          - chunked
+        X-Content-Type-Options:
+          - nosniff
+        access-control-expose-headers:
+          - X-Request-ID
+        alt-svc:
+          - h3=":443"; ma=86400
+        cf-cache-status:
+          - DYNAMIC
+        openai-organization:
+          - braintrust-data
+        openai-processing-ms:
+          - "313"
+        openai-version:
+          - "2020-10-01"
+        strict-transport-security:
+          - max-age=31536000; includeSubDomains; preload
+        x-envoy-upstream-service-time:
+          - "317"
+        x-ratelimit-limit-requests:
+          - "30000"
+        x-ratelimit-limit-tokens:
+          - "150000000"
+        x-ratelimit-remaining-requests:
+          - "29999"
+        x-ratelimit-remaining-tokens:
+          - "149999993"
+        x-ratelimit-reset-requests:
+          - 2ms
+        x-ratelimit-reset-tokens:
+          - 0s
+        x-request-id:
+          - req_7718454117290f001d635e2ea50bf0b5
+      status:
+        code: 200
+        message: OK
+  - request:
+      body: '{"messages":[{"role":"user","content":"What''s 12 + 12?"}],"model":"gpt-4o-mini","stream":true,"stream_options":{"include_usage":true}}'
+      headers:
+        accept:
+          - application/json
+        accept-encoding:
+          - gzip, deflate
+        connection:
+          - keep-alive
+        content-length:
+          - "134"
+        content-type:
+          - application/json
+        cookie:
+          - __cf_bm=AIq4yjX3ru.J9gwNR6zSYLQNorConUZa5qtJ6wXxuvE-1748488174-1.0.1.1-KquMaoYitsL5z76ow2IPzasSn98mtC1_QEt9VOT1pvvQt_obPUDugNtsEGJCc_wP50_X4wP.kC7nYuf98KX8dCPpiq2ZqY5vwVCdgocqRxU;
+            _cfuvid=o0LrLIiV.VvLFX1H1bbtbV01AjzSfXrfrVn0fU7pANY-1748488174648-0.0.1.1-604800000
+        host:
+          - api.openai.com
+        user-agent:
+          - AsyncOpenAI/Python 1.82.0
+        x-stainless-arch:
+          - arm64
+        x-stainless-async:
+          - async:asyncio
+        x-stainless-lang:
+          - python
+        x-stainless-os:
+          - MacOS
+        x-stainless-package-version:
+          - 1.82.0
+        x-stainless-raw-response:
+          - "true"
+        x-stainless-read-timeout:
+          - "600"
+        x-stainless-retry-count:
+          - "0"
+        x-stainless-runtime:
+          - CPython
+        x-stainless-runtime-version:
+          - 3.13.3
+      method: POST
+      uri: https://api.openai.com/v1/chat/completions
+    response:
+      body:
+        string:
+          'data: {"id":"chatcmpl-BcNw7TPkf33tHhtv5BBYd6zunIaPW","object":"chat.completion.chunk","created":1748488175,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[{"index":0,"delta":{"role":"assistant","content":"","refusal":null},"logprobs":null,"finish_reason":null}],"usage":null}
+
+
+          data: {"id":"chatcmpl-BcNw7TPkf33tHhtv5BBYd6zunIaPW","object":"chat.completion.chunk","created":1748488175,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[{"index":0,"delta":{"content":"12"},"logprobs":null,"finish_reason":null}],"usage":null}
+
+
+          data: {"id":"chatcmpl-BcNw7TPkf33tHhtv5BBYd6zunIaPW","object":"chat.completion.chunk","created":1748488175,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[{"index":0,"delta":{"content":"
+          +"},"logprobs":null,"finish_reason":null}],"usage":null}
+
+
+          data: {"id":"chatcmpl-BcNw7TPkf33tHhtv5BBYd6zunIaPW","object":"chat.completion.chunk","created":1748488175,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[{"index":0,"delta":{"content":"
+          "},"logprobs":null,"finish_reason":null}],"usage":null}
+
+
+          data: {"id":"chatcmpl-BcNw7TPkf33tHhtv5BBYd6zunIaPW","object":"chat.completion.chunk","created":1748488175,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[{"index":0,"delta":{"content":"12"},"logprobs":null,"finish_reason":null}],"usage":null}
+
+
+          data: {"id":"chatcmpl-BcNw7TPkf33tHhtv5BBYd6zunIaPW","object":"chat.completion.chunk","created":1748488175,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[{"index":0,"delta":{"content":"
+          equals"},"logprobs":null,"finish_reason":null}],"usage":null}
+
+
+          data: {"id":"chatcmpl-BcNw7TPkf33tHhtv5BBYd6zunIaPW","object":"chat.completion.chunk","created":1748488175,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[{"index":0,"delta":{"content":"
+          "},"logprobs":null,"finish_reason":null}],"usage":null}
+
+
+          data: {"id":"chatcmpl-BcNw7TPkf33tHhtv5BBYd6zunIaPW","object":"chat.completion.chunk","created":1748488175,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[{"index":0,"delta":{"content":"24"},"logprobs":null,"finish_reason":null}],"usage":null}
+
+
+          data: {"id":"chatcmpl-BcNw7TPkf33tHhtv5BBYd6zunIaPW","object":"chat.completion.chunk","created":1748488175,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[{"index":0,"delta":{"content":"."},"logprobs":null,"finish_reason":null}],"usage":null}
+
+
+          data: {"id":"chatcmpl-BcNw7TPkf33tHhtv5BBYd6zunIaPW","object":"chat.completion.chunk","created":1748488175,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[{"index":0,"delta":{},"logprobs":null,"finish_reason":"stop"}],"usage":null}
+
+
+          data: {"id":"chatcmpl-BcNw7TPkf33tHhtv5BBYd6zunIaPW","object":"chat.completion.chunk","created":1748488175,"model":"gpt-4o-mini-2024-07-18","service_tier":"default","system_fingerprint":"fp_54eb4bd693","choices":[],"usage":{"prompt_tokens":14,"completion_tokens":8,"total_tokens":22,"prompt_tokens_details":{"cached_tokens":0,"audio_tokens":0},"completion_tokens_details":{"reasoning_tokens":0,"audio_tokens":0,"accepted_prediction_tokens":0,"rejected_prediction_tokens":0}}}
+
+
+          data: [DONE]
+
+
+          '
+      headers:
+        CF-RAY:
+          - 9472cb34d8b0e56c-EWR
+        Connection:
+          - keep-alive
+        Content-Type:
+          - text/event-stream; charset=utf-8
+        Date:
+          - Thu, 29 May 2025 03:09:35 GMT
+        Server:
+          - cloudflare
+        Transfer-Encoding:
+          - chunked
+        X-Content-Type-Options:
+          - nosniff
+        access-control-expose-headers:
+          - X-Request-ID
+        alt-svc:
+          - h3=":443"; ma=86400
+        cf-cache-status:
+          - DYNAMIC
+        openai-organization:
+          - braintrust-data
+        openai-processing-ms:
+          - "324"
+        openai-version:
+          - "2020-10-01"
+        strict-transport-security:
+          - max-age=31536000; includeSubDomains; preload
+        x-envoy-upstream-service-time:
+          - "328"
+        x-ratelimit-limit-requests:
+          - "30000"
+        x-ratelimit-limit-tokens:
+          - "150000000"
+        x-ratelimit-remaining-requests:
+          - "29999"
+        x-ratelimit-remaining-tokens:
+          - "149999993"
+        x-ratelimit-reset-requests:
+          - 2ms
+        x-ratelimit-reset-tokens:
+          - 0s
+        x-request-id:
+          - req_8e14ac90b3fc7df08ab71458200c0b80
+      status:
+        code: 200
+        message: OK
+version: 1