don't share httpx.AsyncHTTPTransport between event loops

Author: ewjoachim

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -6,12 +6,14 @@
 
 from __future__ import annotations as _annotations
 
+import asyncio
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator, Iterator
 from contextlib import asynccontextmanager, contextmanager
 from dataclasses import dataclass, field
 from datetime import datetime
-from functools import cache
+from functools import cache, lru_cache
+from types import TracebackType
 
 import httpx
 from typing_extensions import Literal, TypeAliasType
@@ -506,17 +508,47 @@ def cached_async_http_client(*, provider: str | None = None, timeout: int = 600,
 @cache
 def _cached_async_http_client(provider: str | None, timeout: int = 600, connect: int = 5) -> httpx.AsyncClient:
     return httpx.AsyncClient(
-        transport=_cached_async_http_transport(),
+        transport=_PerLoopTransport(),
         timeout=httpx.Timeout(timeout=timeout, connect=connect),
         headers={'User-Agent': get_user_agent()},
     )
 
 
-@cache
-def _cached_async_http_transport() -> httpx.AsyncHTTPTransport:
+@lru_cache(maxsize=32)
+def _get_transport_for_loop(loop: asyncio.AbstractEventLoop) -> httpx.AsyncHTTPTransport:
     return httpx.AsyncHTTPTransport()
 
 
+class _PerLoopTransport(httpx.AsyncBaseTransport):
+    def get_transport(self) -> httpx.AsyncHTTPTransport:
+        return _get_transport_for_loop(asyncio.get_running_loop())
+
+    def __init__(self):
+        # We need to ensure that if we call __aenter__ on a transport, it won't be removed
+        # from the lru cache until __aexit__ is called, so ins
+        self._currently_opened: dict[asyncio.AbstractEventLoop, httpx.AsyncHTTPTransport] = {}
+
+    async def __aenter__(self):
+        loop = asyncio.get_running_loop()
+        transport = self._currently_opened[loop] = self.get_transport()
+        return transport
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None = None,
+        exc_value: BaseException | None = None,
+        traceback: TracebackType | None = None,
+    ) -> None:
+        if transport := self._currently_opened.pop(asyncio.get_running_loop(), None):
+            await transport.__aexit__(exc_type, exc_value, traceback)
+
+    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+        return await self.get_transport().handle_async_request(request)
+
+    async def aclose(self) -> None:
+        await self.get_transport().aclose()
+
+
 @cache
 def get_user_agent() -> str:
     """Get the user agent string for the HTTP client."""