diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index fa28f471c77..13dc644f0af 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -630,7 +630,7 @@ class AsyncHTTPHandler: files=files, content=request_content, ) - response = await self._send(req, stream=stream) + response = await self.client.send(req, stream=stream) response.raise_for_status() return response except httpx.TimeoutException as e: @@ -681,7 +681,7 @@ class AsyncHTTPHandler: timeout=timeout, content=request_content, # type: ignore ) - response = await self._send(req) + response = await self.client.send(req) response.raise_for_status() return response except httpx.TimeoutException as e: @@ -730,7 +730,7 @@ class AsyncHTTPHandler: timeout=timeout, content=request_content, # type: ignore ) - response = await self._send(req) + response = await self.client.send(req) response.raise_for_status() return response except httpx.TimeoutException as e: @@ -779,7 +779,7 @@ class AsyncHTTPHandler: timeout=timeout, content=request_content, # type: ignore ) - response = await self._send(req, stream=stream) + response = await self.client.send(req, stream=stream) response.raise_for_status() return response except httpx.HTTPStatusError as e: @@ -787,24 +787,6 @@ class AsyncHTTPHandler: except Exception as e: raise e - async def _send(self, req: httpx.Request, stream: bool = False) -> httpx.Response: - """ - Send `req`, retrying once if the connection died before a response arrived. - - A pooled keep-alive connection the upstream has already dropped fails this - way. The transport evicts the dead connection, so the retry is served over - a fresh one from the same client: no single-use client to own, no lifetime - coupling between a client and a streaming body, and a faithful replay of - the original request (same method, same body, same client configuration). - - `raise_for_status()` stays with the caller so a retry-path 4xx/5xx still - reaches its `except httpx.HTTPStatusError` clause and gets masked. - """ - try: - return await self.client.send(req, stream=stream) - except (httpx.RemoteProtocolError, httpx.ConnectError): - return await self.client.send(req, stream=stream) - def __del__(self) -> None: try: asyncio.get_running_loop().create_task(self.close()) diff --git a/tests/local_testing/conftest.py b/tests/local_testing/conftest.py index d134a7439a8..48d50ea534c 100644 --- a/tests/local_testing/conftest.py +++ b/tests/local_testing/conftest.py @@ -77,6 +77,9 @@ _VCR_INCOMPATIBLE_FILES = frozenset( "test_router_caching.py", # Hits the local fake OpenAI endpoint on 127.0.0.1; nothing to record. "test_fake_openai_endpoint.py", + # Drives a loopback server that poisons its own keep-alive connections; + # vcrpy patches the transport that behaviour lives in. + "test_handler_no_replay_on_dead_connection.py", } ) diff --git a/tests/local_testing/test_handler_no_replay_on_dead_connection.py b/tests/local_testing/test_handler_no_replay_on_dead_connection.py new file mode 100644 index 00000000000..e96aa1f8c3c --- /dev/null +++ b/tests/local_testing/test_handler_no_replay_on_dead_connection.py @@ -0,0 +1,341 @@ +"""``AsyncHTTPHandler`` must not replay a request whose connection died. + +``post``, ``put``, ``patch`` and ``delete`` each used to catch +``(httpx.RemoteProtocolError, httpx.ConnectError)`` and re-send the request. A +dead pooled keep-alive produces those errors, but so does an upstream that read +the request in full, ran the work, and then died before answering. The two are +the same exception, so a replay could re-run billed work just as readily as it +could rescue a connection that never carried the request, and nothing available +to the client separates the two. The replay carried no header delta, no log +line and no counter, so neither the operator nor the upstream could tell it +from a genuine second call. + +RFC 9110 section 9.2.2 describes exactly that heuristic, calls it riskier, and +then says a proxy MUST NOT automatically retry non-idempotent requests. aiohttp, +which is litellm's default transport, implements the retry one layer below httpx +and deliberately scopes POST and PATCH out of it +(``IDEMPOTENT_METHODS`` in ``aiohttp/client.py``). Above this layer the Router +already retries at its default ``num_retries``, where the retry is counted, +configurable, and free to pick a different deployment. + +So the specification these tests pin is one line: **one call into a handler +method issues exactly one httpx send.** Counting sends rather than server-side +requests is what makes that checkable, because aiohttp retries PUT and DELETE +underneath httpx and the server therefore sees a second request either way. + +The tests drive a real loopback server (``_StaleKeepAliveServer``) rather than +monkeypatching ``client.send``, because a mocked transport has no connection +pool of its own to poison, and the pool is the thing under test. The server +binds ``127.0.0.1:0``, so it reaches no network and needs no credentials. +""" + +import asyncio +import os +import socket +import struct +import sys +from typing import List, NamedTuple, Optional, Tuple + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + +# The two transports map a dead pooled connection onto different httpx errors, +# so each is driven with the drop shape that reaches the code path at issue: +# +# httpcore FIN after the request is read -> httpx.RemoteProtocolError +# RST -> httpx.ReadError +# aiohttp RST -> aiohttp.ClientOSError -> httpx.ConnectError +# FIN -> ServerDisconnectedError -> httpx.ReadError +# +# The deleted retry caught the first error of each pair, which is why those are +# the shapes that would show a replay if it came back. +TRANSPORT_PARAMS = [ + pytest.param("aiohttp", "rst", httpx.ConnectError, id="aiohttp"), + pytest.param("httpcore", "fin", httpx.RemoteProtocolError, id="httpcore"), +] + +# aiohttp treats these as idempotent and retries them itself, below httpx. +AIOHTTP_RETRIES_THESE = ["put", "delete"] +# aiohttp declines to retry these, so nothing but litellm could re-send them. +AIOHTTP_DECLINES_THESE = ["post", "patch"] + +# Distinct from the warm-up body, so a replay is countable even when the method +# under test is the POST that warmed the pool. +BODY = b'{"n": 1}' + + +class _SeenRequest(NamedTuple): + method: str + target: str + body: bytes + # 1 for the first request on a TCP connection, 2 for the one that gets + # dropped. A request that arrives as 1 went out on a fresh connection. + position: int = 0 + + +class _StaleKeepAliveServer: + """A loopback HTTP server that poisons its own keep-alive connections. + + The first request on any TCP connection is answered normally. A *second* + request on the same connection is read in full, so a test can see exactly + what the client sent, and then the connection is dropped without a + response, either with a FIN (``drop="fin"``) or an RST (``drop="rst"``). + + Reading the second request in full before dropping it is the point: it is + the shape where a retry replays work the upstream already did. + + Every request the server reads, answered or not, lands in ``requests``. + """ + + def __init__(self, drop: str, stream_frames: int = 0): + self.drop = drop + self.stream_frames = stream_frames + self.requests: List[_SeenRequest] = [] + self._server: Optional[asyncio.AbstractServer] = None + self.url = "" + + async def __aenter__(self) -> "_StaleKeepAliveServer": + self._server = await asyncio.start_server(self._handle, "127.0.0.1", 0) + port = self._server.sockets[0].getsockname()[1] + self.url = f"http://127.0.0.1:{port}/v1/things" + return self + + async def __aexit__(self, *exc_info) -> None: + assert self._server is not None + self._server.close() + await self._server.wait_closed() + + async def _handle(self, reader, writer) -> None: + served = 0 + try: + while True: + request = await self._read_request(reader) + if request is None: + return + served += 1 + self.requests.append(request._replace(position=served)) + if served == 1: + await self._respond(writer) + else: + self._drop(writer) + return + except (ConnectionError, asyncio.IncompleteReadError): + return + finally: + writer.close() + + @staticmethod + async def _read_request(reader) -> Optional[_SeenRequest]: + head = b"" + while b"\r\n\r\n" not in head: + byte = await reader.read(1) + if not byte: + return None + head += byte + lines = head.split(b"\r\n") + method, target, _ = lines[0].split(b" ", 2) + headers = {} + for line in lines[1:]: + if b":" in line: + name, _, value = line.partition(b":") + headers[name.strip().lower()] = value.strip() + + if headers.get(b"transfer-encoding", b"").lower() == b"chunked": + body = b"" + while True: + size = int((await reader.readuntil(b"\r\n")).strip(), 16) + if size == 0: + await reader.readuntil(b"\r\n") + break + body += await reader.readexactly(size) + await reader.readexactly(2) + else: + length = int(headers.get(b"content-length", b"0")) + body = await reader.readexactly(length) if length else b"" + + return _SeenRequest(method.decode(), target.decode(), body) + + async def _respond(self, writer) -> None: + if self.stream_frames: + writer.write(b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nTransfer-Encoding: chunked\r\n\r\n") + await writer.drain() + for index in range(self.stream_frames): + frame = b"data: frame-%d\n\n" % index + writer.write(b"%x\r\n" % len(frame) + frame + b"\r\n") + await writer.drain() + await asyncio.sleep(0.01) + writer.write(b"0\r\n\r\n") + else: + writer.write(b"HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nContent-Length: 2\r\n\r\nok") + await writer.drain() + + def _drop(self, writer) -> None: + if self.drop == "rst": + sock = writer.get_extra_info("socket") + # "ii" is the Linux and macOS `struct linger` (two C ints); Windows + # wants two shorts. CI is ubuntu-only, so this stays as it is. + sock.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0)) + writer.close() + + +@pytest.fixture +def restore_transport_globals(): + """litellm picks its transport from module globals; put them back afterwards.""" + saved = (litellm.disable_aiohttp_transport, litellm.force_ipv4) + yield + litellm.disable_aiohttp_transport, litellm.force_ipv4 = saved + + +def _make_handler(transport: str) -> Tuple[AsyncHTTPHandler, List[str]]: + """Build a handler that records the method of every request httpx sends. + + Counting sends is the only way to tell a litellm retry from aiohttp's own. + aiohttp retries a dropped keep-alive one layer below httpx for the methods + it treats as idempotent, so the server sees a second request either way and + a server-side count cannot separate them. Only a litellm retry is a second + ``send``, and httpx runs its ``request`` event hook once per ``send``. + + The hook goes in at construction rather than onto ``handler.client`` + afterwards, so it covers every client the handler uses. A re-send issued on + a second, short-lived client, which is how the retry was written before it + was removed, therefore still shows up in this count. + """ + litellm.disable_aiohttp_transport = transport == "httpcore" + litellm.force_ipv4 = False + + sends: List[str] = [] + + async def record(request: httpx.Request) -> None: + sends.append(request.method) + + handler = AsyncHTTPHandler( + timeout=httpx.Timeout(10.0, connect=5.0), + event_hooks={"request": [record]}, + ) + return handler, sends + + +async def _warm_pool(handler: AsyncHTTPHandler, url: str) -> None: + """Leave exactly one idle keep-alive connection in the pool. + + Wall-clock sensitive, and nothing here extends the connection's life. httpx + defaults ``Limits.keepalive_expiry`` to 5.0s and ``create_client`` does not + override it, so on httpcore a pause longer than that between this warm-up + and the request under test drops the idle connection: the request then goes + out on a fresh connection, succeeds, and the test no longer exercises a + poisoned one. ``_assert_pool_was_poisoned`` catches that rather than letting + the case pass empty. Keep the code between the warm-up and the request under + test short, and on an unexplained failure suspect a stalled runner first. + (aiohttp holds its connector open for ``AIOHTTP_KEEPALIVE_TIMEOUT``, 120s by + default, and tolerates far more.) + """ + response = await handler.post(url, data=b"{}") + assert response.status_code == 200 + + +def _assert_pool_was_poisoned(server: _StaleKeepAliveServer) -> None: + """Fail unless the request under test actually landed on the reused connection. + + A request count cannot show this: a fresh connection produces a second + server-side request just as a reused one does. Only the request's position + on its own connection separates them, and it is position 2 that gets + dropped. Without this check a case could go green having exercised a + connection that was never at risk of being replayed. + """ + assert any(r.position == 2 for r in server.requests), ( + "no request reached the poisoned keep-alive connection: positions " + f"{[r.position for r in server.requests]}. The request under test went out on a " + "fresh connection, so nothing was exercised." + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", AIOHTTP_DECLINES_THESE) +@pytest.mark.parametrize("transport,drop,expected", TRANSPORT_PARAMS) +async def test_a_dead_connection_is_not_replayed(method, transport, drop, expected, restore_transport_globals): + """The upstream reads the request, dies, and litellm lets the failure through. + + The server has already consumed the body at this point, so a replay would + be a second execution of work the upstream may have completed. One send in, + one request at the server, and the transport error reaches the caller. + """ + async with _StaleKeepAliveServer(drop=drop) as server: + handler, sends = _make_handler(transport) + try: + await _warm_pool(handler, server.url) + + with pytest.raises(expected): + await getattr(handler, method)(server.url, data=BODY) + + _assert_pool_was_poisoned(server) + assert len(sends) == 2, f"litellm re-sent the request: {sends}" + # The body the upstream may already have acted on reached it once. + assert sum(r.body == BODY for r in server.requests) == 1 + assert server.requests[-1].method == method.upper() + finally: + await handler.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transport,drop,expected", TRANSPORT_PARAMS) +async def test_a_dead_connection_on_a_streaming_post_is_not_replayed( + transport, drop, expected, restore_transport_globals +): + """Streaming takes the same path, and used to hide the replay best. + + The retry returned a stream over a client it then closed, so the body died + partway through and the second billed generation was invisible to the + caller. With no retry the connection failure is what the caller sees. + """ + async with _StaleKeepAliveServer(drop=drop, stream_frames=6) as server: + handler, sends = _make_handler(transport) + try: + await _warm_pool(handler, server.url) + + with pytest.raises(expected): + response = await handler.post(server.url, data=BODY, stream=True) + async for _ in response.aiter_bytes(): + pass + + _assert_pool_was_poisoned(server) + assert len(sends) == 2, f"litellm re-sent the request: {sends}" + finally: + await handler.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", AIOHTTP_RETRIES_THESE) +@pytest.mark.parametrize("transport,drop,expected", TRANSPORT_PARAMS) +async def test_litellm_adds_no_retry_where_the_transport_has_its_own( + method, transport, drop, expected, restore_transport_globals +): + """PUT and DELETE are aiohttp's to retry, and litellm does not add a second one. + + aiohttp retries these below httpx, so on that transport the call can still + succeed and the server can still see a second request. Neither is litellm's + doing, and the send count says so on both transports. + + The two aiohttp cases are a control rather than a guard: aiohttp recovers + before any exception reaches litellm, so they hold whether or not litellm + has a retry of its own. The httpcore cases are the ones that fail if it + comes back. + """ + async with _StaleKeepAliveServer(drop=drop) as server: + handler, sends = _make_handler(transport) + try: + await _warm_pool(handler, server.url) + + try: + await getattr(handler, method)(server.url, data=BODY) + except httpx.TransportError: + pass + + _assert_pool_was_poisoned(server) + assert len(sends) == 2, f"litellm re-sent the request: {sends}" + finally: + await handler.close() diff --git a/tests/test_litellm/llms/custom_httpx/test_handler_retry_preserves_request.py b/tests/test_litellm/llms/custom_httpx/test_handler_retry_preserves_request.py deleted file mode 100644 index 3737aa4cfed..00000000000 --- a/tests/test_litellm/llms/custom_httpx/test_handler_retry_preserves_request.py +++ /dev/null @@ -1,546 +0,0 @@ -"""Regression tests for ``AsyncHTTPHandler``'s connection-error retry. - -The retry fires when a pooled keep-alive connection turns out to be dead: the -client picks a connection that still looks healthy, writes a request onto it, -and never hears back. These tests drive that from a real loopback server -(``_StaleKeepAliveServer``) rather than by monkeypatching ``client.send``, so -they exercise the transport, the connection pool, and the retry together. - -What they pin down: - -* a streaming response returned by the retry stays readable to the last byte; -* the retry replays the *original* request -- same method, same body, whether - the body came from ``content=`` or ``files=``; -* a non-2xx served on the retry path is still routed through litellm's - credential masking; -* handler-level TLS settings survive the retry. - -Both transports are covered. litellm's default is aiohttp; httpcore is the -opt-in (``litellm.disable_aiohttp_transport = True``). - -``tests/test_litellm/readme.md`` says this directory holds mocked tests, and a -mock is what this bug cannot be caught with: it needs a live connection pool -holding a keep-alive connection the upstream has already dropped, and a mocked -transport has no pool of its own to poison. The in-tree precedent for a real -loopback server here is -``tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py``. The -server below binds ``127.0.0.1:0``, so it reaches no network, needs no API key -or credentials, and dies with the test that started it. -""" - -import asyncio -import datetime -import io -import ipaddress -import os -import socket -import ssl -import struct -import sys -from typing import List, NamedTuple, Optional, Tuple - -import httpx -import pytest - -sys.path.insert(0, os.path.abspath("../../../..")) - -import litellm -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler - -# litellm retries on ``(httpx.RemoteProtocolError, httpx.ConnectError)``, and the -# two transports map a dead pooled connection onto different httpx errors: -# -# httpcore FIN after the request is read -> httpx.RemoteProtocolError (caught) -# RST -> httpx.ReadError (not caught) -# aiohttp RST -> aiohttp.ClientOSError -> httpx.ConnectError (caught) -# FIN -> ServerDisconnectedError-> httpx.ReadError (not caught) -# -# So each transport is driven with the drop shape that actually reaches the -# retry. (aiohttp additionally retries dropped keep-alives itself, but only for -# the methods it considers idempotent -- PUT and DELETE, not POST or PATCH.) -TRANSPORT_PARAMS = [ - pytest.param("aiohttp", "rst", id="aiohttp"), - pytest.param("httpcore", "fin", id="httpcore"), -] - - -class _SeenRequest(NamedTuple): - method: str - target: str - body: bytes - - -class _StaleKeepAliveServer: - """A loopback HTTP server that poisons its own keep-alive connections. - - The first request on any TCP connection is answered normally. A *second* - request on the same connection is read in full -- so a test can inspect - exactly what the client re-sent -- and then the connection is dropped - without a response, either with a FIN (``drop="fin"``) or an RST - (``drop="rst"``). - - Every request the server reads, answered or not, lands in ``requests``. - """ - - def __init__( - self, - drop: str, - stream_frames: int = 0, - tls: Optional[Tuple[str, str]] = None, - ): - self.drop = drop - self.stream_frames = stream_frames - self.tls = tls - self.status = 200 - self.requests: List[_SeenRequest] = [] - self._server: Optional[asyncio.AbstractServer] = None - self.url = "" - - async def __aenter__(self) -> "_StaleKeepAliveServer": - ssl_ctx = None - if self.tls is not None: - ssl_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) - ssl_ctx.load_cert_chain(certfile=self.tls[0], keyfile=self.tls[1]) - self._server = await asyncio.start_server(self._handle, "127.0.0.1", 0, ssl=ssl_ctx) - port = self._server.sockets[0].getsockname()[1] - self.url = f"{'https' if self.tls else 'http'}://127.0.0.1:{port}/v1/things" - return self - - async def __aexit__(self, *exc_info) -> None: - assert self._server is not None - self._server.close() - await self._server.wait_closed() - - async def _handle(self, reader, writer) -> None: - served = 0 - try: - while True: - request = await self._read_request(reader) - if request is None: - return - self.requests.append(request) - served += 1 - if served == 1: - await self._respond(writer) - else: - self._drop(writer) - return - except (ConnectionError, asyncio.IncompleteReadError, ssl.SSLError): - return - finally: - writer.close() - - @staticmethod - async def _read_request(reader) -> Optional[_SeenRequest]: - head = b"" - while b"\r\n\r\n" not in head: - byte = await reader.read(1) - if not byte: - return None - head += byte - lines = head.split(b"\r\n") - method, target, _ = lines[0].split(b" ", 2) - headers = {} - for line in lines[1:]: - if b":" in line: - name, _, value = line.partition(b":") - headers[name.strip().lower()] = value.strip() - - if headers.get(b"transfer-encoding", b"").lower() == b"chunked": - body = b"" - while True: - size = int((await reader.readuntil(b"\r\n")).strip(), 16) - if size == 0: - await reader.readuntil(b"\r\n") - break - body += await reader.readexactly(size) - await reader.readexactly(2) - else: - length = int(headers.get(b"content-length", b"0")) - body = await reader.readexactly(length) if length else b"" - - return _SeenRequest(method.decode(), target.decode(), body) - - async def _respond(self, writer) -> None: - if self.stream_frames: - writer.write( - b"HTTP/1.1 %d OK\r\nContent-Type: text/event-stream\r\nTransfer-Encoding: chunked\r\n\r\n" % self.status - ) - await writer.drain() - for index in range(self.stream_frames): - frame = b"data: frame-%d\n\n" % index - writer.write(b"%x\r\n" % len(frame) + frame + b"\r\n") - await writer.drain() - await asyncio.sleep(0.01) - writer.write(b"0\r\n\r\n") - else: - body = b"ok" if self.status == 200 else b"unauthorized" - writer.write( - b"HTTP/1.1 %d STATUS\r\nContent-Type: text/plain\r\n" - b"Content-Length: %d\r\n\r\n%s" % (self.status, len(body), body) - ) - await writer.drain() - - def _drop(self, writer) -> None: - if self.drop == "rst": - sock = writer.get_extra_info("socket") - # "ii" is the Linux and macOS `struct linger` (two C ints); Windows - # wants two shorts. CI is ubuntu-only, so this stays as it is. - sock.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0)) - writer.close() - - -class _NonSeekableReader: - """A read-once file-like object, as a pipe or a socket wrapper would be.""" - - def __init__(self, data: bytes): - self._buffer = io.BytesIO(data) - - def read(self, size: int = -1) -> bytes: - return self._buffer.read(size) - - def seek(self, *args, **kwargs) -> int: - raise io.UnsupportedOperation("seek") - - -@pytest.fixture -def restore_transport_globals(): - """litellm picks its transport from module globals; put them back afterwards.""" - saved = (litellm.disable_aiohttp_transport, litellm.force_ipv4) - yield - litellm.disable_aiohttp_transport, litellm.force_ipv4 = saved - - -def _make_handler(transport: str, **kwargs) -> AsyncHTTPHandler: - litellm.disable_aiohttp_transport = transport == "httpcore" - litellm.force_ipv4 = False - return AsyncHTTPHandler(timeout=httpx.Timeout(10.0, connect=5.0), **kwargs) - - -def _track_sends(handler: AsyncHTTPHandler) -> List[str]: - """Record the method of every request handed to ``httpx.AsyncClient.send``. - - That is the layer litellm retries at, and counting it is the only way to - tell litellm's retry from aiohttp's. aiohttp retries a dropped keep-alive - itself, one layer below httpx, for the methods it treats as idempotent: the - server sees a third request either way, so a server-side count cannot - separate them, but only litellm's retry is a second ``send``. httpx runs its - ``request`` event hook once per ``send``, so this list counts sends. - """ - sends: List[str] = [] - - async def record(request: httpx.Request) -> None: - sends.append(request.method) - - handler.client.event_hooks["request"].append(record) - return sends - - -def _assert_retried(server: _StaleKeepAliveServer, sends: List[str]) -> None: - """Fail unless litellm's retry ran, so a green test cannot be an empty one. - - Most of these cases still pass when the retry never fires -- the request - simply succeeds on its first attempt and there is nothing left to get wrong. - An httpx, httpcore or aiohttp upgrade that maps a dead pooled connection - onto an error the retry clause does not catch would quietly turn them into - that, so each one says out loud that the retry happened. - """ - assert len(sends) == 3, ( - f"litellm's retry did not fire: httpx sent {len(sends)} request(s), " - f"expected 3 (warm-up, poisoned attempt, retry)" - ) - assert len(server.requests) == 3, ( - f"the server saw {len(server.requests)} request(s), expected 3 (warm-up, poisoned attempt, retry)" - ) - - -async def _warm_pool(handler: AsyncHTTPHandler, url: str) -> None: - """Leave exactly one idle keep-alive connection in the pool. - - Wall-clock sensitive, and nothing here extends the connection's life. httpx - defaults ``Limits.keepalive_expiry`` to 5.0s and ``create_client`` does not - override it, so on httpcore a pause longer than that between this warm-up - and the request under test drops the idle connection: no poisoned - connection is picked, litellm's retry never fires, and ``_assert_retried`` - fails. (aiohttp holds its own connector open for - ``AIOHTTP_KEEPALIVE_TIMEOUT``, 120s by default, and tolerates far more.) So - keep the code between the warm-up and the request under test short, and on - an unexplained failure suspect a stalled runner before the retry itself. - """ - response = await handler.post(url, data=b"{}") - assert response.status_code == 200 - - -def _self_signed_cert(directory) -> Tuple[str, str]: - from cryptography import x509 - from cryptography.hazmat.primitives import hashes, serialization - from cryptography.hazmat.primitives.asymmetric import rsa - from cryptography.x509.oid import NameOID - - key = rsa.generate_private_key(public_exponent=65537, key_size=2048) - name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "127.0.0.1")]) - now = datetime.datetime.now(datetime.timezone.utc) - certificate = ( - x509.CertificateBuilder() - .subject_name(name) - .issuer_name(name) - .public_key(key.public_key()) - .serial_number(x509.random_serial_number()) - .not_valid_before(now - datetime.timedelta(days=1)) - .not_valid_after(now + datetime.timedelta(days=1)) - .add_extension( - x509.SubjectAlternativeName([x509.IPAddress(ipaddress.ip_address("127.0.0.1"))]), - critical=False, - ) - .sign(key, hashes.SHA256()) - ) - cert_path = str(directory / "cert.pem") - key_path = str(directory / "key.pem") - with open(cert_path, "wb") as handle: - handle.write(certificate.public_bytes(serialization.Encoding.PEM)) - with open(key_path, "wb") as handle: - handle.write( - key.private_bytes( - encoding=serialization.Encoding.PEM, - format=serialization.PrivateFormat.TraditionalOpenSSL, - encryption_algorithm=serialization.NoEncryption(), - ) - ) - return cert_path, key_path - - -@pytest.mark.asyncio -@pytest.mark.parametrize("transport,drop", TRANSPORT_PARAMS) -async def test_retry_keeps_the_streaming_body_readable(transport, drop, restore_transport_globals): - """The response the retry hands back must still stream to the last byte. - - Closing a single-use retry client in a ``finally`` returns the streaming - response with its transport already torn down, so the body dies almost at - once: on either transport at most one of the six frames arrives before - ``aiter_bytes()`` raises ``httpx.ReadError``, and sometimes none does. - """ - async with _StaleKeepAliveServer(drop=drop, stream_frames=6) as server: - handler = _make_handler(transport) - sends = _track_sends(handler) - try: - await _warm_pool(handler, server.url) - - response = await handler.post(server.url, data=b"{}", stream=True) - frames = 0 - async for chunk in response.aiter_bytes(): - frames += chunk.count(b"data: frame-") - - assert frames == 6 - _assert_retried(server, sends) - finally: - await handler.close() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("method", ["post", "put", "patch", "delete"]) -@pytest.mark.parametrize("transport,drop", TRANSPORT_PARAMS) -async def test_retry_preserves_the_http_method(method, transport, drop, restore_transport_globals): - """A retried PUT must reach the upstream as a PUT, not as a POST.""" - if transport == "aiohttp" and method in ("put", "delete"): - pytest.skip( - "aiohttp retries a dropped keep-alive itself for the methods it treats as " - "idempotent, one layer below httpx, so litellm's retry never runs here and " - "the case would pass with or without this change" - ) - - async with _StaleKeepAliveServer(drop=drop) as server: - handler = _make_handler(transport) - sends = _track_sends(handler) - try: - await _warm_pool(handler, server.url) - - response = await getattr(handler, method)(server.url, data=b'{"n": 1}') - - assert response.status_code == 200 - _assert_retried(server, sends) - dropped, retried = server.requests[-2:] - assert (dropped.method, retried.method) == (method.upper(), method.upper()) - finally: - await handler.close() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("transport,drop", TRANSPORT_PARAMS) -async def test_retry_preserves_a_content_body(transport, drop, restore_transport_globals): - """``content=`` must survive the retry rather than being dropped.""" - payload = b'{"prompt": "' + b"x" * 256 + b'"}' - async with _StaleKeepAliveServer(drop=drop) as server: - handler = _make_handler(transport) - sends = _track_sends(handler) - try: - await _warm_pool(handler, server.url) - - response = await handler.post(server.url, content=payload) - - assert response.status_code == 200 - _assert_retried(server, sends) - dropped, retried = server.requests[-2:] - assert dropped.body == payload - assert retried.body == payload - finally: - await handler.close() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("transport,drop", TRANSPORT_PARAMS) -async def test_retry_of_a_one_shot_content_iterator_fails_loudly(transport, drop, restore_transport_globals): - """An unreplayable ``content=`` iterator must raise, not send nothing. - - The retry replays the original ``httpx.Request``, and a request whose body - is a one-shot async iterator has already spent it, so httpx refuses with - ``StreamConsumed``. Failing is the honest outcome: the old retry built a - fresh request that was never passed ``content=`` at all, so it sent a - zero-byte body and returned the upstream's 200 as if the payload had landed. - - The two transports surface it differently. httpcore lets httpx's own - ``StreamConsumed`` out; aiohttp catches it while writing the body and - re-raises it as a connection error, which httpx maps to ``NetworkError`` - carrying the same text. - """ - - async def one_shot(): - yield b'{"prompt": "sent exactly once"}' - - expected = httpx.StreamConsumed if transport == "httpcore" else httpx.NetworkError - async with _StaleKeepAliveServer(drop=drop) as server: - handler = _make_handler(transport) - sends = _track_sends(handler) - try: - await _warm_pool(handler, server.url) - - with pytest.raises(expected) as raised: - await handler.post(server.url, content=one_shot()) - - assert "the content has already been streamed" in str(raised.value) - assert len(sends) == 3, f"litellm's retry did not fire: httpx sent {len(sends)} request(s)" - # The retry raised before anything reached the wire, so the upstream - # saw only the warm-up and the attempt that hit the poisoned - # connection: no empty-bodied third request quietly answered 200. - assert len(server.requests) == 2 - finally: - await handler.close() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("transport,drop", TRANSPORT_PARAMS) -async def test_retry_preserves_a_multipart_upload(transport, drop, restore_transport_globals): - """A retried ``files=`` upload must carry the file, not an empty body. - - Silently uploading zero bytes and returning 200 is worse than failing: the - caller has no way to tell the upload did not happen. - """ - files = {"file": ("audio.wav", b"RIFF" + b"\x00" * 2048, "audio/wav")} - async with _StaleKeepAliveServer(drop=drop) as server: - handler = _make_handler(transport) - sends = _track_sends(handler) - try: - await _warm_pool(handler, server.url) - - response = await handler.post(server.url, files=files) - - assert response.status_code == 200 - _assert_retried(server, sends) - dropped, retried = server.requests[-2:] - assert len(dropped.body) > 2048 - assert len(retried.body) == len(dropped.body) - assert retried.body == dropped.body - finally: - await handler.close() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("transport,drop", TRANSPORT_PARAMS) -async def test_retry_of_a_non_seekable_upload_replays_an_empty_part(transport, drop, restore_transport_globals): - """A ``files=`` part read from a non-seekable handle replays empty. - - A limitation of httpx, not of this change, and the reason the test above - uses bytes. ``FileField.render_data`` rewinds with ``seek(0)`` and swallows - ``io.UnsupportedOperation``, so on the replay it reads an exhausted handle, - yields nothing, and the retry sends the multipart envelope around a - zero-length part, which the upstream answers 200. Bytes, str and seekable - handles all replay faithfully. Pinned so a change in httpx is noticed. - """ - payload = b"RIFF" + b"\x00" * 2048 - files = {"file": ("audio.wav", _NonSeekableReader(payload), "audio/wav")} - async with _StaleKeepAliveServer(drop=drop) as server: - handler = _make_handler(transport) - sends = _track_sends(handler) - try: - await _warm_pool(handler, server.url) - - response = await handler.post(server.url, files=files) - - assert response.status_code == 200 - _assert_retried(server, sends) - dropped, retried = server.requests[-2:] - assert payload in dropped.body - assert payload not in retried.body - assert len(retried.body) == len(dropped.body) - len(payload) - finally: - await handler.close() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("transport,drop", TRANSPORT_PARAMS) -async def test_retry_path_error_still_masks_the_api_key(transport, drop, restore_transport_globals): - """A non-2xx served on the retry must go through litellm's masking. - - ``raise_for_status()`` has to stay in the caller's outer ``try`` for that: - an ``HTTPStatusError`` raised from inside the retry's own ``except`` clause - cannot reach the sibling ``except httpx.HTTPStatusError`` handler, so the - full URL -- API key and all -- ends up in the exception text. - """ - secret = "sk-not-a-real-key-0123456789" - async with _StaleKeepAliveServer(drop=drop) as server: - url = f"{server.url}?key={secret}" - handler = _make_handler(transport) - sends = _track_sends(handler) - try: - await _warm_pool(handler, url) - server.status = 401 - - with pytest.raises(httpx.HTTPStatusError) as raised: - await handler.post(url, data=b"{}") - - message = str(raised.value) - assert "key=[REDACTED_API_KEY]" in message - assert secret not in message - _assert_retried(server, sends) - finally: - await handler.close() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("transport,drop", [pytest.param("httpcore", "fin", id="httpcore")]) -async def test_retry_preserves_handler_ssl_verify(transport, drop, tmp_path, restore_transport_globals): - """The retry must keep the handler's own TLS settings. - - Building a second client for the retry without forwarding ``ssl_verify`` - silently swaps the handler's configuration for litellm's global default, so - a handler that deliberately trusts a private (here, self-signed) certificate - fails the moment it retries. - - httpcore only: over TLS, aiohttp reports a dropped pooled connection as - ``ServerDisconnectedError`` whichever way the peer goes away (the SSL layer - absorbs the reset), and that maps to ``httpx.ReadError``, which the retry - clause does not catch. The retry is unreachable there, so there is nothing - to assert. - """ - cert, key = _self_signed_cert(tmp_path) - async with _StaleKeepAliveServer(drop=drop, tls=(cert, key)) as server: - handler = _make_handler(transport, ssl_verify=False) - sends = _track_sends(handler) - try: - await _warm_pool(handler, server.url) - - response = await handler.post(server.url, data=b"{}") - - assert response.status_code == 200 - _assert_retried(server, sends) - finally: - await handler.close()