fix(openai): propagate asyncio.TimeoutErrors correctly when openai …

…operation canceled (#10265)

Fixes #10191

In the case of an `asyncio` future cancellation, we want to propagate
that error up appropriately. What this looks like is the `response` from
the OpenAI function being `None`, which does not conform to the typing
from the OpenAI SDK (see original issue). If the result is ever `None`,
we don't pass that result to our shared generator to handle OpenAI
responses, and instead let the error bubble up.

- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

- [x] Reviewer has checked that all the criteria below are met
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

ddtrace/contrib/openai/patch.py

@@ -294,7 +294,12 @@ async def patched_endpoint(func, args, kwargs):
             raise
         finally:
             try:
-                g.send((resp, err))
+                if resp is not None:
+                    # openai responses cannot be None
+                    # if resp is None, it is likely because the context
+                    # of the request was cancelled, so we want that to propagate up properly
+                    # see: https://github.com/DataDog/dd-trace-py/issues/10191
+                    g.send((resp, err))
             except StopIteration as e:
                 if err is None:
                     # This return takes priority over `return resp`

releasenotes/notes/openai-async-timeout-errors-f9ccc1adbe4ed14e.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    openai: Fixes a bug where `asyncio.TimeoutError`s were not being propagated correctly from canceled OpenAI API requests.

tests/contrib/openai/test_openai_v1.py

@@ -1682,3 +1682,43 @@ def test_integration_service_name(openai_api_key, ddtrace_run_python_code_in_sub
         assert status == 0, err
         assert out == b""
         assert err == b""
+
+
+async def test_openai_asyncio_cancellation(openai):
+    import asyncio
+
+    import httpx
+
+    class DelayedTransport(httpx.AsyncBaseTransport):
+        def __init__(self, delay: float):
+            self.delay = delay
+            self._transport = httpx.AsyncHTTPTransport()
+
+        async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+            # Introduce a delay before making the actual request
+            await asyncio.sleep(self.delay)
+            return await self._transport.handle_async_request(request)
+
+    client = openai.AsyncOpenAI(http_client=httpx.AsyncClient(transport=DelayedTransport(delay=10)))
+    asyncio_timeout = False
+
+    try:
+        await asyncio.wait_for(
+            client.chat.completions.create(
+                model="gpt-3.5-turbo",
+                messages=[
+                    {
+                        "role": "user",
+                        "content": "Write a Python program that writes a Python program for a given task.",
+                    },
+                ],
+                user="ddtrace-test",
+            ),
+            timeout=1,
+        )
+    except asyncio.TimeoutError:
+        asyncio_timeout = True
+    except Exception as e:
+        assert False, f"Unexpected exception: {e}"
+
+    assert asyncio_timeout, "Expected asyncio.TimeoutError"