From ba3a2bcf7ffcaf3c1ac4e70b0b3175c362a1e698 Mon Sep 17 00:00:00 2001
From: Sam Brenner <106700075+sabrenner@users.noreply.github.com>
Date: Thu, 22 Aug 2024 09:14:31 -0400
Subject: [PATCH] fix(openai): propagate asyncio.TimeoutErrors correctly when
 openai operation canceled [backport 2.10] (#10320)

Backport 16cdd6aac0ab77a5e6caac1a7c7417d2e58eb12f from #10265 to 2.10.

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.

## Checklist
- [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))

## Reviewer Checklist
- [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               |  7 +++-
 ...async-timeout-errors-f9ccc1adbe4ed14e.yaml |  4 ++
 tests/contrib/openai/test_openai_v1.py        | 40 +++++++++++++++++++
 3 files changed, 50 insertions(+), 1 deletion(-)
 create mode 100644 releasenotes/notes/openai-async-timeout-errors-f9ccc1adbe4ed14e.yaml

diff --git a/ddtrace/contrib/openai/patch.py b/ddtrace/contrib/openai/patch.py
index 5e3bf2caead..a610cf55771 100644
--- a/ddtrace/contrib/openai/patch.py
+++ b/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`
diff --git a/releasenotes/notes/openai-async-timeout-errors-f9ccc1adbe4ed14e.yaml b/releasenotes/notes/openai-async-timeout-errors-f9ccc1adbe4ed14e.yaml
new file mode 100644
index 00000000000..17d85525559
--- /dev/null
+++ b/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.
\ No newline at end of file
diff --git a/tests/contrib/openai/test_openai_v1.py b/tests/contrib/openai/test_openai_v1.py
index e14d54bca8d..ec720dfb5d3 100644
--- a/tests/contrib/openai/test_openai_v1.py
+++ b/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"