mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
9b53ee878f
* perf(ci): distribute unit cases and reuse security sweep connections * revert(ci): reject first speed trial after new failures * perf(ci): avoid repeated test discovery and parsing * perf(ci): isolate parallel security workers * revert(ci): restore measured security execution isolation * perf(tests): reuse source scans and deduplicate cache cleanup * perf(ci): restore parallel security with safe log batching * test(ci): isolate fixture state and schema comparisons * test(ci): make schema and scheduler fixtures deterministic * test(ci): drain queued logging before fixture teardown * test(ci): preserve logging delivery without idle worker restarts * test(ci): keep guardrail sinks available during parallel bursts * test(ci): isolate full proxy lifecycle from gateway imports * test(ci): preserve first failure responses and sweep evidence * test(ci): restore documentation URL settings between cases * test(ci): isolate open MCP tools in private gateways * test(ci): group private MCP cases and handle exited workers * test(ci): retain proxy exits and MCP failure details * test(ci): repair merge interactions with main Key the schema shard's full clone on artifact-name, since main removed the unit-flag input. Give each owned Redis its own log file again. Hand the realtime transcription tests main added a logging object whose dispatch handlers are awaitable, so the per-test logging drain no longer waits on a MagicMock task forever * perf(ci): shard owned integration cases and close lifecycle leaks * fix(ci): qualify shard execution and harden interrupted cleanup * fix(ci): bound worker supervision and own fault traffic * perf(ci): reuse unparametrized fixtures across bounded case groups * fix(ci): preserve boot fault allowances and complete exit records * test(ci): tolerate startup retries and retain readiness details * perf(ci): balance integration shards by recorded duration * fix(ci): load migration database helpers only for workers
232 lines
8.5 KiB
Python
232 lines
8.5 KiB
Python
import asyncio
|
|
import socket
|
|
import threading
|
|
import time
|
|
from collections.abc import Generator
|
|
from contextlib import contextmanager
|
|
from typing import Final
|
|
from urllib.parse import urlsplit, urlunsplit
|
|
|
|
from pydantic import TypeAdapter
|
|
|
|
from tests.integration._support.async_server import LoopbackServer, close_writer
|
|
|
|
PORT: Final = TypeAdapter(int)
|
|
OUTAGE_SECONDS: Final = 10.0
|
|
|
|
|
|
def _free_port() -> int:
|
|
with socket.socket() as reserve:
|
|
reserve.bind(("127.0.0.1", 0))
|
|
return PORT.validate_python(reserve.getsockname()[1])
|
|
|
|
|
|
class DatabaseRelay:
|
|
def __init__(self, upstream_host: str, upstream_port: int, trigger: bytes) -> None:
|
|
self.port: Final = _free_port()
|
|
self._upstream_host: Final = upstream_host
|
|
self._upstream_port: Final = upstream_port
|
|
self._trigger: Final = trigger
|
|
self._server: Final = LoopbackServer(self.port, self._serve)
|
|
self._loop: Final = self._server.loop
|
|
self._armed: Final = threading.Event()
|
|
self.tripped: Final = threading.Event()
|
|
self.refused = 0
|
|
self.reconnected: Final = threading.Event()
|
|
self._tripped_at = 0.0
|
|
self._writers: tuple[asyncio.StreamWriter, ...] = ()
|
|
|
|
def arm(self) -> None:
|
|
self._armed.set()
|
|
|
|
def start(self) -> None:
|
|
self._server.start()
|
|
|
|
def stop(self) -> None:
|
|
self._server.stop()
|
|
|
|
def _drop_all(self) -> None:
|
|
for writer in self._writers:
|
|
writer.close()
|
|
self._writers = ()
|
|
|
|
async def _serve(self, client_reader: asyncio.StreamReader, client_writer: asyncio.StreamWriter) -> None:
|
|
if self.tripped.is_set() and time.monotonic() - self._tripped_at < OUTAGE_SECONDS:
|
|
self.refused += 1
|
|
client_writer.close()
|
|
return
|
|
if self.tripped.is_set():
|
|
self.reconnected.set()
|
|
server_reader, server_writer = await asyncio.open_connection(self._upstream_host, self._upstream_port)
|
|
self._writers = (*self._writers, client_writer, server_writer)
|
|
|
|
async def forward(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, inspect: bool) -> None:
|
|
try:
|
|
while chunk := await reader.read(65536):
|
|
if inspect and self._armed.is_set() and not self.tripped.is_set() and self._trigger in chunk:
|
|
self._tripped_at = time.monotonic()
|
|
self.tripped.set()
|
|
self._drop_all()
|
|
return
|
|
writer.write(chunk)
|
|
await writer.drain()
|
|
except (ConnectionError, asyncio.IncompleteReadError):
|
|
return
|
|
finally:
|
|
await close_writer(writer)
|
|
|
|
await asyncio.gather(
|
|
forward(client_reader, server_writer, True),
|
|
forward(server_reader, client_writer, False),
|
|
)
|
|
|
|
|
|
class HeldStatementRelay:
|
|
def __init__(self, upstream_host: str, upstream_port: int, trigger: bytes) -> None:
|
|
self.port: Final = _free_port()
|
|
self._upstream_host: Final = upstream_host
|
|
self._upstream_port: Final = upstream_port
|
|
self._trigger: Final = trigger
|
|
self._server: Final = LoopbackServer(self.port, self._serve)
|
|
self._loop: Final = self._server.loop
|
|
self._released: Final = asyncio.Event()
|
|
self.held: Final = threading.Event()
|
|
|
|
def release(self) -> None:
|
|
if not self._loop.is_closed():
|
|
self._loop.call_soon_threadsafe(self._released.set)
|
|
|
|
def start(self) -> None:
|
|
self._server.start()
|
|
|
|
def stop(self) -> None:
|
|
self.release()
|
|
self._server.stop()
|
|
|
|
def _holds(self, window: bytes) -> bool:
|
|
return not self.held.is_set() and self._trigger in window
|
|
|
|
async def _serve(self, client_reader: asyncio.StreamReader, client_writer: asyncio.StreamWriter) -> None:
|
|
server_reader, server_writer = await asyncio.open_connection(self._upstream_host, self._upstream_port)
|
|
|
|
async def forward(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, inspect: bool) -> None:
|
|
tail = b"" # rebind-ok: carries the previous read's end so a trigger split across reads still matches
|
|
try:
|
|
while chunk := await reader.read(65536):
|
|
window: Final = tail + chunk
|
|
if inspect and self._holds(window):
|
|
self.held.set()
|
|
await self._released.wait()
|
|
tail = window[-(len(self._trigger) - 1) :]
|
|
writer.write(chunk)
|
|
await writer.drain()
|
|
except (ConnectionError, asyncio.IncompleteReadError):
|
|
return
|
|
finally:
|
|
await close_writer(writer)
|
|
|
|
await asyncio.gather(
|
|
forward(client_reader, server_writer, True),
|
|
forward(server_reader, client_writer, False),
|
|
)
|
|
|
|
|
|
class TriggerScanner:
|
|
def __init__(self, trigger: bytes) -> None:
|
|
self._trigger: Final = trigger
|
|
self._tail: bytes = b""
|
|
|
|
def feed(self, chunk: bytes) -> bool:
|
|
window: Final = self._tail + chunk
|
|
self._tail = window[-(len(self._trigger) - 1) :]
|
|
return self._trigger in window
|
|
|
|
|
|
class DroppedConnectionRelay:
|
|
def __init__(self, upstream_host: str, upstream_port: int, trigger: bytes) -> None:
|
|
self.port: Final = _free_port()
|
|
self._upstream_host: Final = upstream_host
|
|
self._upstream_port: Final = upstream_port
|
|
self._trigger: Final = trigger
|
|
self._server: Final = LoopbackServer(self.port, self._serve)
|
|
self._loop: Final = self._server.loop
|
|
self._armed: Final = threading.Event()
|
|
self.dropped: Final = threading.Event()
|
|
|
|
def arm(self) -> None:
|
|
self._armed.set()
|
|
|
|
def disarm(self) -> None:
|
|
self._armed.clear()
|
|
|
|
def start(self) -> None:
|
|
self._server.start()
|
|
|
|
def stop(self) -> None:
|
|
self._server.stop()
|
|
|
|
async def _serve(self, client_reader: asyncio.StreamReader, client_writer: asyncio.StreamWriter) -> None:
|
|
server_reader, server_writer = await asyncio.open_connection(self._upstream_host, self._upstream_port)
|
|
|
|
async def forward(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, inspect: bool) -> None:
|
|
scanner: Final = TriggerScanner(self._trigger)
|
|
try:
|
|
while chunk := await reader.read(65536):
|
|
matched: Final = scanner.feed(chunk)
|
|
if inspect and self._armed.is_set() and matched:
|
|
self.dropped.set()
|
|
client_writer.close()
|
|
return
|
|
writer.write(chunk)
|
|
await writer.drain()
|
|
except (ConnectionError, asyncio.IncompleteReadError):
|
|
return
|
|
finally:
|
|
await close_writer(writer)
|
|
|
|
await asyncio.gather(
|
|
forward(client_reader, server_writer, True),
|
|
forward(server_reader, client_writer, False),
|
|
)
|
|
|
|
|
|
def _relayed_url(database_url: str, port: int) -> str:
|
|
parts: Final = urlsplit(database_url)
|
|
credentials: Final = f"{parts.username}:{parts.password}@" if parts.username else ""
|
|
return urlunsplit(parts._replace(netloc=f"{credentials}127.0.0.1:{port}"))
|
|
|
|
|
|
@contextmanager
|
|
def database_relay(database_url: str, trigger: bytes) -> Generator[tuple[DatabaseRelay, str]]:
|
|
parts: Final = urlsplit(database_url)
|
|
assert parts.hostname is not None and parts.port is not None, database_url
|
|
relay: Final = DatabaseRelay(parts.hostname, parts.port, trigger)
|
|
try:
|
|
relay.start()
|
|
yield relay, _relayed_url(database_url, relay.port)
|
|
finally:
|
|
relay.stop()
|
|
|
|
|
|
@contextmanager
|
|
def held_statement_relay(database_url: str, trigger: bytes) -> Generator[tuple[HeldStatementRelay, str]]:
|
|
parts: Final = urlsplit(database_url)
|
|
assert parts.hostname is not None and parts.port is not None, database_url
|
|
relay: Final = HeldStatementRelay(parts.hostname, parts.port, trigger)
|
|
try:
|
|
relay.start()
|
|
yield relay, _relayed_url(database_url, relay.port)
|
|
finally:
|
|
relay.stop()
|
|
|
|
|
|
@contextmanager
|
|
def dropped_connection_relay(database_url: str, trigger: bytes) -> Generator[tuple[DroppedConnectionRelay, str]]:
|
|
parts: Final = urlsplit(database_url)
|
|
assert parts.hostname is not None and parts.port is not None, database_url
|
|
relay: Final = DroppedConnectionRelay(parts.hostname, parts.port, trigger)
|
|
try:
|
|
relay.start()
|
|
yield relay, _relayed_url(database_url, relay.port)
|
|
finally:
|
|
relay.stop()
|