litellm/tests/integration/_support/wire.py

from __future__ import annotations

import ssl
import threading
import time
from collections.abc import Callable, Generator, Iterator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from functools import partial
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from queue import SimpleQueue
from types import MappingProxyType
from typing import BinaryIO, Final


@dataclass(frozen=True, slots=True)
class Request:
    method: str
    target: str
    headers: Mapping[str, str]
    body: bytes


class AbortedBody(Exception):
    """The client closed the connection before the body it announced was complete."""


def exactly(stream: BinaryIO, size: int) -> bytes:
    data: Final = stream.read(size)
    if len(data) < size:
        raise AbortedBody
    return data


def chunked_body(stream: BinaryIO) -> Iterator[bytes]:
    while True:
        size_line: Final = stream.readline()
        if not size_line:
            raise AbortedBody
        size: Final = int(size_line.split(b";")[0].strip(), 16)
        if size == 0:
            while stream.readline().strip():
                pass
            return
        yield exactly(stream, size)
        stream.readline()


def read_body(headers: Mapping[str, str], stream: BinaryIO) -> bytes:
    if headers.get("transfer-encoding", "").lower() == "chunked":
        return b"".join(chunked_body(stream))
    return exactly(stream, int(headers.get("content-length", "0")))


@dataclass(frozen=True, slots=True)
class Reply:
    status: int = 200
    body: bytes = b"{}"
    content_type: str = "application/json"
    chunks: tuple[bytes, ...] | None = None
    abort_after: int | None = None
    gate_after_first: threading.Event | None = None
    pause_between_chunks: float = 0
    heartbeat: bytes | None = None
    headers: Mapping[str, str] = MappingProxyType({})
    drop_connection: bool = False


@dataclass(frozen=True, slots=True)
class Wire:
    url: str
    received: SimpleQueue[Request]
    disconnected: SimpleQueue[str]
    connected: SimpleQueue[str]
    closed: SimpleQueue[str]

    def drain(self) -> tuple[Request, ...]:
        return tuple(self.received.get_nowait() for _ in range(self.received.qsize()))

    def connections(self) -> int:
        return self.connected.qsize()

    def open_connections(self) -> int:
        return self.connected.qsize() - self.closed.qsize()


def _await_release(
    gate: threading.Event, closing: threading.Event, beat: Callable[[], None] | None = None, every: float = 0.05
) -> bool:
    while not closing.is_set():
        if gate.wait(timeout=every):
            return True
        if beat is not None:
            beat()
    return gate.is_set()


@contextmanager
def wire_server(
    respond: Callable[[Request], Reply],
    tls: ssl.SSLContext | None = None,
    port: int = 0,
    keep_alive: bool = False,
) -> Generator[Wire, None, None]:
    """Owned TCP peer; requests traverse the real HTTP client and serialization. With `keep_alive` the
    peer honours HTTP/1.1 persistent connections so `connections()` counts the client's TCP sessions."""
    received: Final[SimpleQueue[Request]] = SimpleQueue()
    errors: Final[SimpleQueue[Exception]] = SimpleQueue()
    disconnected: Final[SimpleQueue[str]] = SimpleQueue()
    connected: Final[SimpleQueue[str]] = SimpleQueue()
    closed: Final[SimpleQueue[str]] = SimpleQueue()
    closing: Final = threading.Event()

    class Handler(BaseHTTPRequestHandler):
        protocol_version = "HTTP/1.1"
        timeout = 5

        def setup(self) -> None:
            super().setup()
            connected.put(f"{self.client_address[0]}:{self.client_address[1]}")

        def finish(self) -> None:
            super().finish()
            closed.put(f"{self.client_address[0]}:{self.client_address[1]}")

        def write_chunk(self, chunk: bytes) -> None:
            self.wfile.write(b"%x\r\n%s\r\n" % (len(chunk), chunk))
            self.wfile.flush()

        def respond(self) -> None:
            headers: Final = {name.lower(): value for name, value in self.headers.items()}
            try:
                body: Final = read_body(headers, self.rfile)
            except AbortedBody:
                self.close_connection = True
                disconnected.put(self.path)
                return
            request: Final = Request(self.command, self.path, headers, body)
            received.put(request)
            try:
                reply = respond(request)
            except Exception as error:
                errors.put(error)
                reply = Reply(status=500)
            if reply.drop_connection:
                self.close_connection = True
                disconnected.put(request.target)
                return
            self.send_response(reply.status)
            self.send_header("content-type", reply.content_type)
            for name, value in reply.headers.items():
                self.send_header(name, value)
            if reply.chunks is None:
                self.send_header("content-length", str(len(reply.body)))
            else:
                self.send_header("transfer-encoding", "chunked")
            if not keep_alive:
                self.send_header("connection", "close")
            self.end_headers()
            try:
                if self.command == "HEAD":
                    self.wfile.flush()
                elif reply.chunks is None:
                    self.wfile.write(reply.body)
                else:
                    for index, chunk in enumerate(reply.chunks):
                        if reply.abort_after == index:
                            break
                        self.write_chunk(chunk)
                        if (
                            index == 0
                            and reply.gate_after_first is not None
                            and not (
                                _await_release(reply.gate_after_first, closing)
                                if reply.heartbeat is None
                                else _await_release(
                                    reply.gate_after_first,
                                    closing,
                                    partial(self.write_chunk, reply.heartbeat),
                                    reply.pause_between_chunks,
                                )
                            )
                        ):
                            break
                        if reply.pause_between_chunks and index + 1 < len(reply.chunks):
                            time.sleep(reply.pause_between_chunks)
                    else:
                        self.wfile.write(b"0\r\n\r\n")
                        self.wfile.flush()
            except (BrokenPipeError, ConnectionResetError):
                disconnected.put(request.target)
            except Exception as error:
                errors.put(error)
            self.close_connection = not keep_alive

        do_POST = respond
        do_PUT = respond
        do_GET = respond
        do_DELETE = respond
        do_PATCH = respond
        do_HEAD = respond

        def log_message(self, format: str, *args: object) -> None:
            pass

    class OwnedHTTPServer(ThreadingHTTPServer):
        daemon_threads = False
        request_queue_size = 128

        def server_bind(self) -> None:
            super().server_bind()
            if tls is not None:
                self.socket = tls.wrap_socket(self.socket, server_side=True)

    with OwnedHTTPServer(("127.0.0.1", port), Handler) as server:
        thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05})
        thread.start()
        try:
            yield Wire(
                f"{'https' if tls is not None else 'http'}://127.0.0.1:{server.server_port}",
                received,
                disconnected,
                connected,
                closed,
            )
        finally:
            closing.set()
            server.shutdown()
            thread.join(timeout=6)
            assert not thread.is_alive(), "Owned HTTP server survived cleanup"
            server.server_close()
            failure: Final = None if errors.empty() else errors.get_nowait()
            assert failure is None, f"Owned HTTP peer failed: {failure!r}"