diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index 769bc0fed1e..6997afafd8d 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -18,6 +18,7 @@ from litellm.secret_managers.main import str_to_bool AIOHTTP_EXC_MAP: Dict = { # Order matters here, most specific exception first # Timeout related exceptions + asyncio.TimeoutError: httpx.TimeoutException, aiohttp.ServerTimeoutError: httpx.TimeoutException, aiohttp.ConnectionTimeoutError: httpx.ConnectTimeout, aiohttp.SocketTimeoutError: httpx.ReadTimeout, @@ -253,6 +254,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport): allow_redirects=False, auto_decompress=False, timeout=ClientTimeout( + total=timeout.get("read"), sock_connect=timeout.get("connect"), sock_read=timeout.get("read"), connect=timeout.get("pool"), diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py index 894fc45e361..1f1a36fd7ab 100644 --- a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py @@ -255,6 +255,48 @@ def _make_mock_response(should_fail=False, fail_count={"count": 0}): return MockResp() +@pytest.mark.asyncio +async def test_handle_async_request_total_timeout_triggers(): + """ + Ensure that LiteLLMAiohttpTransport raises httpx.TimeoutException + when the total timeout duration elapses. + """ + import asyncio + from aiohttp import web + + async def slow_handler(request): + await asyncio.sleep(0.3) + return web.Response(text="ok") + + app = web.Application() + app.router.add_get("/", slow_handler) + runner = web.AppRunner(app) + await runner.setup() + site = web.TCPSite(runner, "127.0.0.1", 0) + await site.start() + + port = site._server.sockets[0].getsockname()[1] + + def factory(): + return aiohttp.ClientSession() + + transport = LiteLLMAiohttpTransport(client=factory) # type: ignore + + request = httpx.Request("GET", f"http://127.0.0.1:{port}/") + + request.extensions["timeout"] = { + "connect": 0.1, + "read": 0.1, + "pool": 0.1, + "total": 0.1, + } + + try: + with pytest.raises(httpx.TimeoutException): + await transport.handle_async_request(request) + finally: + await transport.aclose() + await runner.cleanup() def _make_mock_session(closed=False): """Helper to create a mock aiohttp session"""