From 24db29cd737b4736678846dce8c1d5350c01f320 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 4 Nov 2026 04:02:09 +0900 Subject: [PATCH] feat(proxy): supervise the opt-in Rust forwarding server Co-Authored-By: Claude Opus 5.5 --- .github/workflows/test-unit.yml | 1 + gateway/launch.py | 71 +- litellm/proxy/dev_reload.py | 7 +- litellm/proxy/proxy_cli.py | 107 ++- litellm/proxy/rust_forward/AGENTS.md | 18 + litellm/proxy/rust_forward/__init__.py | 0 litellm/proxy/rust_forward/mode.py | 177 +++++ litellm/proxy/rust_forward/peer.py | 57 ++ .../proxy/rust_forward/platform/__init__.py | 52 ++ litellm/proxy/rust_forward/platform/posix.py | 45 ++ .../proxy/rust_forward/platform/windows.py | 56 ++ litellm/proxy/rust_forward/server.py | 30 + litellm/proxy/rust_forward/supervisor.py | 332 +++++++++ .../proxy/shutdown/graceful_uvicorn_worker.py | 10 +- litellm/rust_bridge/configuration.py | 4 +- litellm_cli/main.py | 8 +- .../configuration/test_rust_forward.py | 180 +++++ .../test_rust_forward_traffic.py | 457 ++++++++++++ tests/unit/gateway/test_launch.py | 110 ++- tests/unit/proxy/rust_forward/__init__.py | 0 tests/unit/proxy/rust_forward/test_mode.py | 112 +++ tests/unit/proxy/rust_forward/test_peer.py | 143 ++++ .../proxy/rust_forward/test_supervisor.py | 241 ++++++ tests/unit/proxy/test_dev_reload.py | 41 +- tests/unit/proxy/test_proxy_cli.py | 692 +++++++----------- 25 files changed, 2485 insertions(+), 466 deletions(-) create mode 100644 litellm/proxy/rust_forward/AGENTS.md create mode 100644 litellm/proxy/rust_forward/__init__.py create mode 100644 litellm/proxy/rust_forward/mode.py create mode 100644 litellm/proxy/rust_forward/peer.py create mode 100644 litellm/proxy/rust_forward/platform/__init__.py create mode 100644 litellm/proxy/rust_forward/platform/posix.py create mode 100644 litellm/proxy/rust_forward/platform/windows.py create mode 100644 litellm/proxy/rust_forward/server.py create mode 100644 litellm/proxy/rust_forward/supervisor.py create mode 100644 tests/integration/configuration/test_rust_forward.py create mode 100644 tests/integration/configuration/test_rust_forward_traffic.py create mode 100644 tests/unit/proxy/rust_forward/__init__.py create mode 100644 tests/unit/proxy/rust_forward/test_mode.py create mode 100644 tests/unit/proxy/rust_forward/test_peer.py create mode 100644 tests/unit/proxy/rust_forward/test_supervisor.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 939c1a1c01f..adbf07f30a9 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -584,6 +584,7 @@ jobs: tests/unit/proxy/credential_endpoints tests/unit/proxy/discovery_endpoints tests/unit/proxy/shutdown + tests/unit/proxy/rust_forward tests/unit/proxy/public_endpoints tests/unit/proxy/prompts tests/unit/proxy/rag_endpoints diff --git a/gateway/launch.py b/gateway/launch.py index 59f7121e513..f7f98bd1156 100644 --- a/gateway/launch.py +++ b/gateway/launch.py @@ -22,15 +22,16 @@ import sys from collections.abc import Callable, Mapping, Sequence from typing import TYPE_CHECKING, Final +from typing_extensions import assert_never from uvicorn.main import main as uvicorn_main from litellm_cli.main import without_gc_pauses if TYPE_CHECKING: - from fastapi import FastAPI - from litellm.proxy.db.db_url_settings import DatabaseURLSettings from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings + from litellm.proxy.rust_forward.peer import ASGIApp + from litellm.proxy.rust_forward.server import PythonServer GATEWAY_APP: Final = "gateway.launch:worker_app" KEEPALIVE_FLAG: Final = "--timeout-keep-alive" @@ -65,17 +66,58 @@ def pool_database_url( return start_in_container_pgbouncer(pgbouncer, upstream_url, token_auth=settings.token_auth()) -def worker_app() -> "FastAPI": +def worker_app() -> "ASGIApp": with without_gc_pauses(): from gateway.main import app - return app + from litellm.proxy.rust_forward.mode import rust_forward_mode + from litellm.proxy.rust_forward.peer import wrap + return wrap(app, rust_forward_mode(os.environ)) def _serve(argv: Sequence[str]) -> None: uvicorn_main(tuple(argv), prog_name="uvicorn") -def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None: +def python_server(argv: Sequence[str]) -> "PythonServer": + """What uvicorn would serve the public port with, read from the same flags and environment it reads.""" + from pydantic import TypeAdapter + + from litellm.proxy.rust_forward.server import Listener, PythonServer, Server + from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager + + params: Final[Mapping[str, object]] = uvicorn_main.make_context( + "uvicorn", list(uvicorn_argv(argv, os.environ)) + ).params + tls: Final = ( + params["ssl_keyfile"] is not None + or params["ssl_certfile"] is not None + or any(arg.split("=")[0] == "--ssl-ciphers" for arg in argv) + ) + listener: Final = ( + Listener.TLS + if tls + else Listener.UNIX_SOCKET + if params["uds"] is not None + else Listener.INHERITED_FD + if params["fd"] is not None + else Listener.TCP + ) + graceful_timeout: Final = TypeAdapter[int | None](int | None).validate_python(params["timeout_graceful_shutdown"]) + return PythonServer( + server=Server.UVICORN, + listener=listener, + host=TypeAdapter(str).validate_python(params["host"]), + port=TypeAdapter(int).validate_python(params["port"]), + keepalive_seconds=TypeAdapter(int).validate_python(params["timeout_keep_alive"]), + graceful_timeout_seconds=GracefulShutdownManager.get_timeout() + if graceful_timeout is None + else graceful_timeout, + proxy_headers=TypeAdapter(bool).validate_python(params["proxy_headers"]), + trusted_proxies=TypeAdapter[str | None](str | None).validate_python(params["forwarded_allow_ips"]), + ) + + +def configure_database() -> None: from litellm.proxy.db.db_url_settings import DatabaseURLSettings from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings, export_pooled_database_url @@ -86,7 +128,24 @@ def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) - sys.exit(f"LiteLLM gateway: in-container pgbouncer could not start: {pooled_url.reason}") if pooled_url is not None: export_pooled_database_url(pooled_url) - serve(uvicorn_argv(argv, os.environ)) + + +def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None: + from litellm.proxy.rust_forward.mode import Disabled, Supervised, Supervisor, rust_forward_mode, uvicorn_flags + from litellm.proxy.rust_forward.supervisor import supervise + + mode: Final = rust_forward_mode(os.environ) + match mode: + case Disabled(): + configure_database() + serve(uvicorn_argv(argv, os.environ)) + case Supervisor(): + command: Final = (sys.executable, "-m", "gateway.launch", *argv) + sys.exit(supervise(python_server(argv), command, before_start=configure_database)) + case Supervised(): + serve((*uvicorn_argv(argv, os.environ), *uvicorn_flags(mode))) + case _: + assert_never(mode) if __name__ == "__main__": diff --git a/litellm/proxy/dev_reload.py b/litellm/proxy/dev_reload.py index 500f581b34d..7b057271cb6 100644 --- a/litellm/proxy/dev_reload.py +++ b/litellm/proxy/dev_reload.py @@ -13,7 +13,7 @@ from socket import socket from typing import Final, Literal from pydantic import ConfigDict, ImportString, JsonValue, TypeAdapter, with_config -from typing_extensions import ReadOnly, Required, TypedDict +from typing_extensions import NotRequired, ReadOnly, Required, TypedDict from uvicorn import Config, Server from uvicorn import config as uvicorn_config from uvicorn.supervisors import ChangeReload, Multiprocess @@ -42,6 +42,7 @@ class ServerOptions(TypedDict, total=False): @with_config(ConfigDict(extra="allow")) class ReloadOptions(ServerOptions): app: Required[ReadOnly[str]] + factory: NotRequired[ReadOnly[bool]] reload: ReadOnly[bool] reload_dirs: ReadOnly[tuple[str, ...]] reload_includes: ReadOnly[tuple[str, ...]] @@ -92,7 +93,9 @@ def load_worker_app(import_path: str, worker_app_factory: bool) -> Callable[..., def pool_config(options: ReloadOptions, workers: int, worker_overrides: Mapping[str, object] | None = None) -> Config: overrides: Final = worker_overrides or {} worker_app: Final = TypeAdapter(str).validate_python(overrides.get("app", options["app"])) - worker_app_factory: Final = TypeAdapter(bool).validate_python(overrides.get("factory", False)) + worker_app_factory: Final = TypeAdapter(bool).validate_python( + overrides.get("factory", options.get("factory", False)) + ) return Config( partial(load_worker_app, worker_app, worker_app_factory), factory=True, diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 0fbfa0b06a7..9bc4b8b5492 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -35,6 +35,9 @@ from litellm.types.llms.base import LiteLLMBaseModel if TYPE_CHECKING: from fastapi import FastAPI + + from litellm.proxy.rust_forward.peer import ASGIApp + from litellm.proxy.rust_forward.server import PythonServer else: FastAPI = Any @@ -674,12 +677,14 @@ class ProxyInitializationHelpers: def run_gunicorn_server( host: str, port: int, - app: FastAPI, + app: "ASGIApp", num_workers: int, ssl_certfile_path: str, ssl_keyfile_path: str, max_requests_before_restart: int | None = None, max_requests_before_restart_jitter: int | None = None, + bind: str | None = None, + proxy_headers: bool = True, ): """ Run litellm with `gunicorn` @@ -752,9 +757,11 @@ class ProxyInitializationHelpers: ) graceful_shutdown_timeout: Final = GracefulShutdownManager.get_timeout() gunicorn_options: Final = { - "bind": f"{host}:{port}", + "bind": f"{host}:{port}" if bind is None else bind, "workers": num_workers, # default is 1 - "worker_class": build_graceful_shutdown_uvicorn_worker(graceful_shutdown_timeout), + "worker_class": build_graceful_shutdown_uvicorn_worker( + graceful_shutdown_timeout, proxy_headers=proxy_headers + ), "graceful_timeout": math.ceil(graceful_shutdown_timeout) + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS, "preload": True, # Add the preload flag, "accesslog": "-", # Log to stdout @@ -815,7 +822,7 @@ class ProxyInitializationHelpers: _run_ollama_serve = run_ollama_serve @staticmethod - def is_port_in_use(port): + def is_port_in_use(port: int): import socket with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: @@ -823,6 +830,57 @@ class ProxyInitializationHelpers: _is_port_in_use = is_port_in_use + @staticmethod + def serving_port(port: int) -> int: + return random.randint(1024, 49152) if port == 4000 and ProxyInitializationHelpers.is_port_in_use(port) else port + + @staticmethod + def python_server( + host: str, + port: int, + keepalive_timeout: int | None, + run_gunicorn: bool, + run_hypercorn: bool, + run_granian: bool, + tls: bool, + ) -> "PythonServer": + """``run_gunicorn_server`` never hands ``--keepalive_timeout`` or forwarded IPs to gunicorn, so its workers keep + gunicorn's keep-alive and gunicorn's trusted proxies, which include ``::1`` where uvicorn's do not.""" + import uvicorn + from gunicorn.config import Config as GunicornConfig + + from litellm.proxy.rust_forward.server import Listener, PythonServer, Server + from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager + + server: Final = ( + Server.HYPERCORN + if run_hypercorn + else Server.GRANIAN + if run_granian + else Server.GUNICORN + if run_gunicorn + else Server.UVICORN + ) + keepalive: Final = ( + TypeAdapter(int).validate_python(GunicornConfig().keepalive) + if server is Server.GUNICORN + else uvicorn.Config(app="", log_config=None).timeout_keep_alive + if keepalive_timeout is None + else keepalive_timeout + ) + return PythonServer( + server=server, + listener=Listener.TLS if tls else Listener.TCP, + host=host, + port=ProxyInitializationHelpers.serving_port(port), + keepalive_seconds=keepalive, + graceful_timeout_seconds=GracefulShutdownManager.get_timeout(), + proxy_headers=True, + trusted_proxies=",".join(TypeAdapter(list[str]).validate_python(GunicornConfig().forwarded_allow_ips)) + if server is Server.GUNICORN + else None, + ) + @staticmethod def get_loop_type(): """Helper function to determine the event loop type based on platform""" @@ -1210,8 +1268,8 @@ class ProxyInitializationHelpers: ) def run_server( cli_args, - host, - port, + host: str, + port: int, api_base, api_version, model, @@ -1280,6 +1338,33 @@ def run_server( run_setup_wizard() return + from litellm.proxy.rust_forward.mode import ( + Supervisor, + gunicorn_bind, + rust_forward_mode, + server_proxy_headers, + uvicorn_options, + ) + + forward_mode: Final = rust_forward_mode(os.environ) + if isinstance(forward_mode, Supervisor) and not any( + (version, validate_config, health, test, skip_server_startup, check_migrations_against_existing_data) + ): + from litellm.proxy.rust_forward.supervisor import supervise + + server: Final = ProxyInitializationHelpers.python_server( + host, + port, + keepalive_timeout, + run_gunicorn=run_gunicorn, + run_hypercorn=run_hypercorn, + run_granian=run_granian, + tls=any((ssl_keyfile_path, ssl_certfile_path, ciphers)), + ) + raise SystemExit( + supervise(server, (sys.executable, "-m", "litellm.proxy.proxy_cli", *sys.argv[1:])) + ) + args: Final = locals() try: from litellm.proxy.proxy_server import ( @@ -1699,8 +1784,7 @@ def run_server( ) sys.exit(1) export_pooled_database_url(pooled_database_url) - if port == 4000 and ProxyInitializationHelpers.is_port_in_use(port): - port = random.randint(1024, 49152) + port = ProxyInitializationHelpers.serving_port(port) if prometheus_metrics_port == port: raise click.UsageError("--prometheus_metrics_port must differ from --port") @@ -1748,6 +1832,7 @@ def run_server( keepalive_timeout=keepalive_timeout, timeout_worker_healthcheck=(timeout_worker_healthcheck if running_uvicorn else None), ) + uvicorn_args.update(uvicorn_options(forward_mode, WORKER_APP)) # Optional: recycle uvicorn workers after N requests if max_requests_before_restart is not None: uvicorn_args["limit_max_requests"] = max_requests_before_restart @@ -1812,15 +1897,19 @@ def run_server( workers=num_workers, ) elif run_gunicorn is True: + from litellm.proxy.rust_forward.peer import wrap + ProxyInitializationHelpers.run_gunicorn_server( host=host, port=port, - app=app, + app=wrap(app, forward_mode), num_workers=num_workers, ssl_certfile_path=ssl_certfile_path, ssl_keyfile_path=ssl_keyfile_path, max_requests_before_restart=max_requests_before_restart, max_requests_before_restart_jitter=max_requests_before_restart_jitter, + bind=gunicorn_bind(forward_mode, f"{host}:{port}"), + proxy_headers=server_proxy_headers(forward_mode), ) elif run_hypercorn is True: ProxyInitializationHelpers.init_hypercorn_server( diff --git a/litellm/proxy/rust_forward/AGENTS.md b/litellm/proxy/rust_forward/AGENTS.md new file mode 100644 index 00000000000..294d94dd21b --- /dev/null +++ b/litellm/proxy/rust_forward/AGENTS.md @@ -0,0 +1,18 @@ +# Rust forwarding, Python half + +With `LITELLM_RUST` on, the supervisor here runs the Rust gateway on the public port and Python on a private loopback listener. This package is the Python half of forwarding and pairs with the crate `litellm-rust/crates/gateway-forward`. Once Rust serves every route, both get deleted together, along with one dispatch branch in each entry point + +Every forwarding decision lives here. Entry points (`litellm/proxy/proxy_cli.py`, `gateway/launch.py`, `litellm_cli/main.py`) only translate between their own flags and this package: + +- In the supervisor they describe what they would serve with as a `server.PythonServer` and hand it to `supervisor.supervise`, which decides whether forwarding can start, what Rust and the Python child receive, and how to report a refusal +- In the supervised child they merge `mode.uvicorn_options`, `mode.uvicorn_flags` or `mode.gunicorn_bind` and `mode.server_proxy_headers` into their server settings, and wrap the app with `peer.wrap` in their worker factory. These are identities when forwarding is off, so entry points never branch on it + +Every OS-specific operation lives behind `platform.Platform`, which `supervise` picks once with `platform.current_platform()`: the gateway binary name, how the private listener reaches Python, the spawn options, and how children are signalled and stopped. `PosixPlatform` hands Python an inherited socket (`mode.InheritedFd`). `WindowsPlatform` hands it an assigned loopback port (`mode.AssignedPort`), starts each child in its own process group and stops it with `CTRL_BREAK_EVENT`. The supervisor and `mode.py` never check the OS themselves + +Windows support exists only for running the proxy locally from the CLI. Nobody deploys the proxy on Windows, and the Windows design leans on that. A port picked by the supervisor can be taken by another process before Python binds it, `CTRL_BREAK_EVENT` needs the console a CLI run always has, and gunicorn is refused. Revisit those choices before treating Windows as a deployment target + +Each environment contract is written and read in one place. `mode.supervised_environment` and `mode.rust_forward_mode` are the two halves of the Python child's, and `supervisor.rust_spec` writes the `LITELLM_FORWARD_*` variables that `gateway-forward`'s `settings.rs` reads. Python owns every value Rust needs, so a new one gets a new variable in `rust_spec` + +The supervisor's whole lifecycle is `supervisor.step`, a pure function from state and event to the next state and the actions to run. Children start in `START_ORDER`, each once the one before it reports `Ready`, and on a stop every started child is signalled at once in reverse start order. Keep decisions in `step` and keep `_run` executing actions, so a new child is an entry in `START_ORDER` and the launch table rather than a new state + +Python must see each request as if the client had connected directly. `x-litellm-rust-peer` is the only Rust-specific input, and `peer.py` honors it only from the loopback hop. The proxy app itself never knows Rust is in front of it diff --git a/litellm/proxy/rust_forward/__init__.py b/litellm/proxy/rust_forward/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/proxy/rust_forward/mode.py b/litellm/proxy/rust_forward/mode.py new file mode 100644 index 00000000000..a95a66cf4bc --- /dev/null +++ b/litellm/proxy/rust_forward/mode.py @@ -0,0 +1,177 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Annotated, Final, TypeAlias + +from pydantic import BaseModel, BeforeValidator, ConfigDict, Field +from typing_extensions import assert_never + +from litellm.rust_bridge.configuration import ENV_NAME + +GATEWAY_FD: Final = "_LITELLM_GATEWAY_FD" +GATEWAY_PORT: Final = "_LITELLM_GATEWAY_PORT" +PROXY_HEADERS: Final = "_LITELLM_GATEWAY_PROXY_HEADERS" +TRUSTED_PROXIES: Final = "FORWARDED_ALLOW_IPS" +DEFAULT_TRUSTED_PROXIES: Final = "127.0.0.1" +RUST_HOP: Final = "127.0.0.1" + + +def _blank_as_unset(value: object) -> object: + stripped: Final = value.strip() if isinstance(value, str) else value + return None if stripped == "" else stripped + + +_Switch: TypeAlias = Annotated[bool | None, BeforeValidator(_blank_as_unset)] + + +class _ForwardEnvironment(BaseModel): + model_config = ConfigDict(frozen=True) + + rust: _Switch = Field(None, validation_alias=ENV_NAME) + listener_fd: int | None = Field(None, validation_alias=GATEWAY_FD) + listener_port: int | None = Field(None, validation_alias=GATEWAY_PORT) + proxy_headers: bool = Field(True, validation_alias=PROXY_HEADERS) + trusted_proxies: str = Field(DEFAULT_TRUSTED_PROXIES, validation_alias=TRUSTED_PROXIES) + + +@dataclass(frozen=True, slots=True) +class Disabled: + pass + + +@dataclass(frozen=True, slots=True) +class Supervisor: + pass + + +@dataclass(frozen=True, slots=True) +class InheritedFd: + fd: int + + +@dataclass(frozen=True, slots=True) +class AssignedPort: + port: int + + +ChildListener: TypeAlias = InheritedFd | AssignedPort + + +@dataclass(frozen=True, slots=True) +class Supervised: + listener: ChildListener + proxy_headers: bool + trusted_proxies: str + + +RustForwardMode: TypeAlias = Disabled | Supervisor | Supervised + + +def _listener_environment(listener: ChildListener) -> tuple[str, str]: + match listener: + case InheritedFd(fd): + return GATEWAY_FD, str(fd) + case AssignedPort(port): + return GATEWAY_PORT, str(port) + case _: + assert_never(listener) + + +def supervised_environment( + listener: ChildListener, proxy_headers: bool, trusted_proxies: str | None +) -> Mapping[str, str]: + name, value = _listener_environment(listener) + return MappingProxyType( + { + name: value, + PROXY_HEADERS: str(proxy_headers).lower(), + **({} if trusted_proxies is None else {TRUSTED_PROXIES: trusted_proxies}), + } + ) + + +def rust_forward_mode(environ: Mapping[str, str]) -> RustForwardMode: + parsed: Final = _ForwardEnvironment.model_validate(environ) + if parsed.listener_fd is not None: + return Supervised(InheritedFd(parsed.listener_fd), parsed.proxy_headers, parsed.trusted_proxies) + if parsed.listener_port is not None: + return Supervised(AssignedPort(parsed.listener_port), parsed.proxy_headers, parsed.trusted_proxies) + if parsed.rust: + return Supervisor() + return Disabled() + + +# A supervised server's TCP peer is always Rust on loopback, so server-level proxy headers would trust any +# client's forwarding headers. They stay off there and ``peer.wrap`` applies the operator's trust instead. + + +def _uvicorn_listener(listener: ChildListener) -> Mapping[str, object]: + match listener: + case InheritedFd(fd): + return MappingProxyType({"fd": fd}) + case AssignedPort(port): + return MappingProxyType({"host": RUST_HOP, "port": port}) + case _: + assert_never(listener) + + +def uvicorn_options(mode: RustForwardMode, factory: str) -> Mapping[str, object]: + match mode: + case Supervised(listener): + return MappingProxyType( + {"app": factory, "factory": True, **_uvicorn_listener(listener), "proxy_headers": False} + ) + case Disabled() | Supervisor(): + return MappingProxyType({}) + case _: + assert_never(mode) + + +def _uvicorn_listener_flags(listener: ChildListener) -> tuple[str, ...]: + match listener: + case InheritedFd(fd): + return ("--fd", str(fd)) + case AssignedPort(port): + return ("--host", RUST_HOP, "--port", str(port)) + case _: + assert_never(listener) + + +def uvicorn_flags(mode: RustForwardMode) -> tuple[str, ...]: + match mode: + case Supervised(listener): + return (*_uvicorn_listener_flags(listener), "--no-proxy-headers") + case Disabled() | Supervisor(): + return () + case _: + assert_never(mode) + + +def _gunicorn_listener(listener: ChildListener) -> str: + match listener: + case InheritedFd(fd): + return f"fd://{fd}" + case AssignedPort(port): + return f"{RUST_HOP}:{port}" + case _: + assert_never(listener) + + +def gunicorn_bind(mode: RustForwardMode, public: str) -> str: + match mode: + case Supervised(listener): + return _gunicorn_listener(listener) + case Disabled() | Supervisor(): + return public + case _: + assert_never(mode) + + +def server_proxy_headers(mode: RustForwardMode) -> bool: + match mode: + case Supervised(): + return False + case Disabled() | Supervisor(): + return True + case _: + assert_never(mode) diff --git a/litellm/proxy/rust_forward/peer.py b/litellm/proxy/rust_forward/peer.py new file mode 100644 index 00000000000..2d03d6ca35c --- /dev/null +++ b/litellm/proxy/rust_forward/peer.py @@ -0,0 +1,57 @@ +from collections.abc import Awaitable, Callable +from typing import TYPE_CHECKING, Final, TypeAlias + +from typing_extensions import assert_never + +from litellm.proxy.rust_forward.mode import RUST_HOP, Disabled, RustForwardMode, Supervised, Supervisor + +if TYPE_CHECKING: + from uvicorn._types import ASGIReceiveCallable, ASGISendCallable, Scope, WWWScope + +PEER_HEADER: Final = b"x-litellm-rust-peer" + +ASGIApp: TypeAlias = Callable[..., Awaitable[None]] + + +def _restore_peer(scope: "WWWScope") -> None: + headers: Final = tuple(scope["headers"]) + peer: Final = next((value for name, value in headers if name == PEER_HEADER), None) + if peer is None: + return + remaining: Final = tuple((name, value) for name, value in headers if name != PEER_HEADER) + scope["headers"] = remaining # rebind-ok: the peer header is private to the Rust hop + client: Final = scope["client"] + if client is None or client[0] != RUST_HOP: + return + host, _, port = peer.decode("latin-1").rpartition(":") + scope["client"] = (host.strip("[]"), int(port)) # rebind-ok: uvicorn's access log reads this same scope + + +class RustPeerMiddleware: + """Make Python see each request as if the client had connected to it directly. + + Rust owns the public port, so Python's peer is always Rust on loopback. Rust passes the real + TCP peer in ``x-litellm-rust-peer`` and leaves every forwarding header untouched. This puts the + real peer back into the scope, then runs uvicorn's own ``ProxyHeadersMiddleware`` with the trust + settings Python would have used on the public port. + """ + + def __init__(self, app: ASGIApp, mode: Supervised) -> None: + from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware + + self.app: Final = ProxyHeadersMiddleware(app, mode.trusted_proxies) if mode.proxy_headers else app + + async def __call__(self, scope: "Scope", receive: "ASGIReceiveCallable", send: "ASGISendCallable") -> None: + if scope["type"] != "lifespan": + _restore_peer(scope) + await self.app(scope, receive, send) + + +def wrap(app: ASGIApp, mode: RustForwardMode) -> ASGIApp: + match mode: + case Supervised(): + return RustPeerMiddleware(app, mode) + case Disabled() | Supervisor(): + return app + case _: + assert_never(mode) diff --git a/litellm/proxy/rust_forward/platform/__init__.py b/litellm/proxy/rust_forward/platform/__init__.py new file mode 100644 index 00000000000..b40ff5cb015 --- /dev/null +++ b/litellm/proxy/rust_forward/platform/__init__.py @@ -0,0 +1,52 @@ +import asyncio +import os +import signal +from collections.abc import Callable +from contextlib import AbstractContextManager +from dataclasses import dataclass +from typing import Protocol + +from litellm.proxy.rust_forward.mode import ChildListener +from litellm.proxy.rust_forward.server import Server + + +@dataclass(frozen=True, slots=True) +class PrivateListener: + listener: ChildListener + port: int + pass_fds: tuple[int, ...] + + +@dataclass(frozen=True, slots=True) +class SpawnOptions: + pass_fds: tuple[int, ...] + start_new_session: bool + creationflags: int + + +class Platform(Protocol): + @property + def name(self) -> str: ... + + @property + def gateway_name(self) -> str: ... + + @property + def servers(self) -> frozenset[Server]: ... + + def private_listener(self) -> AbstractContextManager[PrivateListener]: ... + + def spawn_options(self, pass_fds: tuple[int, ...]) -> SpawnOptions: ... + + def catch_stop(self, loop: asyncio.AbstractEventLoop, on_stop: Callable[[signal.Signals], None]) -> None: ... + + def signal_child(self, process: asyncio.subprocess.Process, signum: signal.Signals) -> None: ... + + def stop_group(self, process: asyncio.subprocess.Process) -> None: ... + + +def current_platform() -> Platform: + from litellm.proxy.rust_forward.platform.posix import PosixPlatform + from litellm.proxy.rust_forward.platform.windows import WindowsPlatform + + return WindowsPlatform() if os.name == "nt" else PosixPlatform() diff --git a/litellm/proxy/rust_forward/platform/posix.py b/litellm/proxy/rust_forward/platform/posix.py new file mode 100644 index 00000000000..b1b2e500657 --- /dev/null +++ b/litellm/proxy/rust_forward/platform/posix.py @@ -0,0 +1,45 @@ +import asyncio +import contextlib +import os +import signal +import socket +from collections.abc import Callable, Generator +from dataclasses import dataclass +from typing import Final + +from pydantic import TypeAdapter + +from litellm.proxy.rust_forward.mode import RUST_HOP, InheritedFd +from litellm.proxy.rust_forward.platform import PrivateListener, SpawnOptions +from litellm.proxy.rust_forward.server import Server + +_ADDRESS: Final = TypeAdapter(tuple[str, int]) + + +@dataclass(frozen=True, slots=True) +class PosixPlatform: + name: str = "POSIX" + gateway_name: str = "litellm-gateway" + servers: frozenset[Server] = frozenset({Server.UVICORN, Server.GUNICORN}) + + @contextlib.contextmanager + def private_listener(self) -> Generator[PrivateListener]: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as private: + private.bind((RUST_HOP, 0)) + _, port = _ADDRESS.validate_python(private.getsockname()) + yield PrivateListener(InheritedFd(private.fileno()), port, (private.fileno(),)) + + def spawn_options(self, pass_fds: tuple[int, ...]) -> SpawnOptions: + return SpawnOptions(pass_fds=pass_fds, start_new_session=True, creationflags=0) + + def catch_stop(self, loop: asyncio.AbstractEventLoop, on_stop: Callable[[signal.Signals], None]) -> None: + for signum in (signal.SIGTERM, signal.SIGINT): + loop.add_signal_handler(signum, on_stop, signum) + + def signal_child(self, process: asyncio.subprocess.Process, signum: signal.Signals) -> None: + with contextlib.suppress(ProcessLookupError): + process.send_signal(signum) + + def stop_group(self, process: asyncio.subprocess.Process) -> None: + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGTERM) diff --git a/litellm/proxy/rust_forward/platform/windows.py b/litellm/proxy/rust_forward/platform/windows.py new file mode 100644 index 00000000000..f32063428e3 --- /dev/null +++ b/litellm/proxy/rust_forward/platform/windows.py @@ -0,0 +1,56 @@ +import asyncio +import contextlib +import signal +import socket +from collections.abc import Callable, Generator +from dataclasses import dataclass +from types import FrameType +from typing import Final + +from pydantic import TypeAdapter + +from litellm.proxy.rust_forward.mode import RUST_HOP, AssignedPort +from litellm.proxy.rust_forward.platform import PrivateListener, SpawnOptions +from litellm.proxy.rust_forward.server import Server + +_ADDRESS: Final = TypeAdapter(tuple[str, int]) +# signal.CTRL_BREAK_EVENT, signal.SIGBREAK and subprocess.CREATE_NEW_PROCESS_GROUP exist only on Windows, so they are +# spelled out to keep this module importable and type checked everywhere. A Windows-only test pins them to the stdlib. +CTRL_BREAK_EVENT: Final = 1 +SIGBREAK: Final = 21 +CREATE_NEW_PROCESS_GROUP: Final = 0x00000200 + + +# Windows support exists only for running the proxy locally from the CLI, never for deploying it. That is why the +# supervisor picks a free loopback port and hands Python its number instead of an open socket, which Windows uvicorn +# cannot serve. Another process could take that port before Python binds it, and the readiness probe would then +# accept it as Python. Python's failed bind still ends the run, which is acceptable on a developer's machine. +@dataclass(frozen=True, slots=True) +class WindowsPlatform: + name: str = "Windows" + gateway_name: str = "litellm-gateway.exe" + servers: frozenset[Server] = frozenset({Server.UVICORN}) + + @contextlib.contextmanager + def private_listener(self) -> Generator[PrivateListener]: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.bind((RUST_HOP, 0)) + _, port = _ADDRESS.validate_python(probe.getsockname()) + yield PrivateListener(AssignedPort(port), port, ()) + + def spawn_options(self, pass_fds: tuple[int, ...]) -> SpawnOptions: + return SpawnOptions(pass_fds=pass_fds, start_new_session=False, creationflags=CREATE_NEW_PROCESS_GROUP) + + def catch_stop(self, loop: asyncio.AbstractEventLoop, on_stop: Callable[[signal.Signals], None]) -> None: + def handle(signum: int, _frame: FrameType | None) -> None: + loop.call_soon_threadsafe(on_stop, signal.Signals(signum)) + + for signum in (signal.SIGINT, SIGBREAK): + signal.signal(signum, handle) + + def signal_child(self, process: asyncio.subprocess.Process, signum: signal.Signals) -> None: + self.stop_group(process) + + def stop_group(self, process: asyncio.subprocess.Process) -> None: + with contextlib.suppress(ProcessLookupError): + process.send_signal(CTRL_BREAK_EVENT) diff --git a/litellm/proxy/rust_forward/server.py b/litellm/proxy/rust_forward/server.py new file mode 100644 index 00000000000..8398cbda533 --- /dev/null +++ b/litellm/proxy/rust_forward/server.py @@ -0,0 +1,30 @@ +from dataclasses import dataclass +from enum import Enum + + +class Server(Enum): + UVICORN = "uvicorn" + GUNICORN = "gunicorn" + HYPERCORN = "Hypercorn" + GRANIAN = "Granian" + + +class Listener(Enum): + TCP = "plain TCP" + TLS = "TLS" + UNIX_SOCKET = "a Unix socket" + INHERITED_FD = "an inherited file descriptor" + + +@dataclass(frozen=True, slots=True) +class PythonServer: + """What an entry point would serve with on the public port if Rust were not in front of it.""" + + server: Server + listener: Listener + host: str + port: int + keepalive_seconds: int + graceful_timeout_seconds: float + proxy_headers: bool + trusted_proxies: str | None diff --git a/litellm/proxy/rust_forward/supervisor.py b/litellm/proxy/rust_forward/supervisor.py new file mode 100644 index 00000000000..3331b3f0c0b --- /dev/null +++ b/litellm/proxy/rust_forward/supervisor.py @@ -0,0 +1,332 @@ +import asyncio +import contextlib +import os +import shutil +import signal +import sys +import sysconfig +from collections.abc import AsyncGenerator, Callable, Coroutine, Mapping, Sequence +from dataclasses import dataclass +from enum import Enum, IntEnum +from functools import partial +from types import MappingProxyType +from typing import Final, TypeAlias + +from typing_extensions import assert_never + +from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS +from litellm.proxy.rust_forward.mode import RUST_HOP, supervised_environment +from litellm.proxy.rust_forward.platform import Platform, PrivateListener, SpawnOptions, current_platform +from litellm.proxy.rust_forward.server import Listener, PythonServer, Server +from litellm.rust_bridge.configuration import ENV_NAME + +PROBE_INTERVAL_SECONDS: Final = 0.05 + + +def gateway_binary(platform: Platform, scripts_dir: str) -> str | None: + return shutil.which(platform.gateway_name, path=scripts_dir) + + +@dataclass(frozen=True, slots=True) +class StartupError: + reason: str + + +def startup_check(platform: Platform, server: PythonServer, binary: str | None) -> str | StartupError: + if server.server in (Server.HYPERCORN, Server.GRANIAN): + return StartupError(f"{ENV_NAME} does not support {server.server.value}") + if server.server not in platform.servers: + return StartupError(f"{ENV_NAME} does not support {server.server.value} on {platform.name}") + if server.listener is not Listener.TCP: + return StartupError(f"{ENV_NAME} does not support {server.listener.value} on the public port") + if binary is None: + return StartupError( + f"{ENV_NAME} requires the litellm-gateway binary. Install litellm[proxy] in this environment" + ) + return binary + + +class Child(Enum): + RUST = "rust" + PYTHON = "python" + + +# Each child starts once the one before it is ready. Python alone takes no connections until it has booted, so Rust +# only opens the public port once Python listens. +START_ORDER: Final = (Child.PYTHON, Child.RUST) + + +class ExitCode(IntEnum): + CLEAN = 0 + CHILD_FAILED = 1 + + +@dataclass(frozen=True, slots=True) +class Ready: + child: Child + + +@dataclass(frozen=True, slots=True) +class Stop: + signum: signal.Signals + + +@dataclass(frozen=True, slots=True) +class Exited: + child: Child + code: int + + +@dataclass(frozen=True, slots=True) +class Running: + up: tuple[Child, ...] + + +@dataclass(frozen=True, slots=True) +class Stopping: + up: frozenset[Child] + + +@dataclass(frozen=True, slots=True) +class Finished: + code: ExitCode + + +@dataclass(frozen=True, slots=True) +class Start: + child: Child + + +@dataclass(frozen=True, slots=True) +class Signal: + child: Child + signum: signal.Signals + + +Event: TypeAlias = Ready | Stop | Exited +State: TypeAlias = Running | Stopping | Finished +Action: TypeAlias = Start | Signal +Transition: TypeAlias = tuple[State, tuple[Action, ...]] + + +def _start_after(up: tuple[Child, ...]) -> Transition: + following: Final = START_ORDER[len(up) : len(up) + 1] + return Running((*up, *following)), tuple(Start(child) for child in following) + + +def boot() -> Transition: + return _start_after(()) + + +# Every child gets the signal at once, as Python alone stopped. Python ends its long streams within its graceful +# timeout, which Rust's drain already covers, so stopping Python only after Rust exits would hold Python's shutdown, +# spend flush included, behind streams only Python can end. Reverse start order lets Rust stop accepting first. +def _stop(up: tuple[Child, ...], signum: signal.Signals) -> Transition: + return Stopping(frozenset(up)), tuple(Signal(child, signum) for child in reversed(up)) + + +def step(state: State, event: Event) -> Transition: + match state: + case Running(up): + match event: + case Ready(child) if up[-1:] == (child,): + return _start_after(up) + case Ready(): + return state, () + case Stop(signum): + return _stop(up, signum) + case Exited(): + return Finished(ExitCode.CHILD_FAILED), () + case _: + assert_never(event) + case Stopping(up): + match event: + case Ready() | Stop(): + return state, () + case Exited(child, code) if child in up and code == 0: + remaining: Final = up - {child} + return (Stopping(remaining) if remaining else Finished(ExitCode.CLEAN)), () + case Exited(): + return Finished(ExitCode.CHILD_FAILED), () + case _: + assert_never(event) + case Finished(): + return state, () + case _: + assert_never(state) + + +@dataclass(frozen=True, slots=True) +class ChildSpec: + argv: tuple[str, ...] + env_overrides: Mapping[str, str] + spawn: SpawnOptions + + +def python_spec( + platform: Platform, command: Sequence[str], private: PrivateListener, server: PythonServer +) -> ChildSpec: + return ChildSpec( + argv=(*command, "--port", str(private.port)), + env_overrides=supervised_environment(private.listener, server.proxy_headers, server.trusted_proxies), + spawn=platform.spawn_options(private.pass_fds), + ) + + +def rust_spec(platform: Platform, binary: str, upstream_port: int, server: PythonServer) -> ChildSpec: + return ChildSpec( + argv=(binary,), + env_overrides=MappingProxyType( + { + "LITELLM_FORWARD_UPSTREAM": f"http://{RUST_HOP}:{upstream_port}", + "LITELLM_FORWARD_HOST": server.host, + "LITELLM_FORWARD_PORT": str(server.port), + "LITELLM_FORWARD_UPSTREAM_KEEPALIVE": str(server.keepalive_seconds), + "LITELLM_FORWARD_SHUTDOWN_DRAIN": str( + server.graceful_timeout_seconds + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS + ), + } + ), + spawn=platform.spawn_options(()), + ) + + +async def _spawn(spec: ChildSpec) -> asyncio.subprocess.Process: + return await asyncio.create_subprocess_exec( + *spec.argv, + env={**os.environ, **spec.env_overrides}, + pass_fds=spec.spawn.pass_fds, + start_new_session=spec.spawn.start_new_session, + creationflags=spec.spawn.creationflags, + ) + + +Launch: TypeAlias = Callable[[asyncio.Queue[Event]], contextlib.AbstractAsyncContextManager[asyncio.subprocess.Process]] + + +async def _report_exit(child: Child, process: asyncio.subprocess.Process, events: asyncio.Queue[Event]) -> None: + events.put_nowait(Exited(child, await process.wait())) + + +async def _accepts(port: int) -> bool: + try: + _, writer = await asyncio.wait_for(asyncio.open_connection(RUST_HOP, port), PROBE_INTERVAL_SECONDS) + except (OSError, asyncio.TimeoutError): + return False + writer.close() + return True + + +async def _report_ready_once_listening(child: Child, port: int, events: asyncio.Queue[Event]) -> None: + while not await _accepts(port): + await asyncio.sleep(PROBE_INTERVAL_SECONDS) + events.put_nowait(Ready(child)) + + +@contextlib.asynccontextmanager +async def _in_background(work: Coroutine[None, None, None]) -> AsyncGenerator[None]: + task: Final = asyncio.create_task(work) + try: + yield + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +@contextlib.asynccontextmanager +async def _launch( + platform: Platform, + child: Child, + spec: ChildSpec, + events: asyncio.Queue[Event], +) -> AsyncGenerator[asyncio.subprocess.Process]: + process: Final = await _spawn(spec) + waiter: Final = asyncio.create_task(_report_exit(child, process, events)) + try: + yield process + finally: + platform.stop_group(process) + await waiter + + +@contextlib.asynccontextmanager +async def _launch_until_listening( + platform: Platform, child: Child, spec: ChildSpec, port: int, events: asyncio.Queue[Event] +) -> AsyncGenerator[asyncio.subprocess.Process]: + async with ( + _launch(platform, child, spec, events) as process, + _in_background(_report_ready_once_listening(child, port, events)), + ): + yield process + + +async def _run( + platform: Platform, + transition: Transition, + events: asyncio.Queue[Event], + children: Mapping[Child, asyncio.subprocess.Process], + launches: Mapping[Child, Launch], +) -> ExitCode: + state, actions = transition + if not actions: + if isinstance(state, Finished): + return state.code + return await _run(platform, step(state, await events.get()), events, children, launches) + action: Final = actions[0] + rest: Final = (state, actions[1:]) + match action: + case Signal(child, signum): + platform.signal_child(children[child], signum) + return await _run(platform, rest, events, children, launches) + case Start(child): + async with launches[child](events) as process: + return await _run(platform, rest, events, MappingProxyType({**children, child: process}), launches) + case _: + assert_never(action) + + +async def _supervise(platform: Platform, binary: str, command: Sequence[str], server: PythonServer) -> ExitCode: + events: Final[asyncio.Queue[Event]] = asyncio.Queue() + platform.catch_stop(asyncio.get_running_loop(), lambda signum: events.put_nowait(Stop(signum))) + with platform.private_listener() as private: + launches: Final[Mapping[Child, Launch]] = MappingProxyType( + { + Child.PYTHON: partial( + _launch_until_listening, + platform, + Child.PYTHON, + python_spec(platform, command, private, server), + private.port, + ), + Child.RUST: partial(_launch, platform, Child.RUST, rust_spec(platform, binary, private.port, server)), + } + ) + return await _run(platform, boot(), events, MappingProxyType({}), launches) + + +def _nothing_to_prepare() -> None: + pass + + +def supervise( + server: PythonServer, + command: Sequence[str], + before_start: Callable[[], None] = _nothing_to_prepare, + scripts_dir: str = sysconfig.get_path("scripts"), + platform: Platform | None = None, +) -> int: + """Serve ``server``'s public port from Rust and run ``command`` behind it, or say why that cannot start. + + ``before_start`` runs only once forwarding is known to be possible, so a refused start leaves nothing behind. + """ + chosen: Final = current_platform() if platform is None else platform + checked: Final = startup_check(chosen, server, gateway_binary(chosen, scripts_dir)) + if isinstance(checked, StartupError): + sys.stderr.write(f"{checked.reason}\n") + return ExitCode.CHILD_FAILED + before_start() + try: + return asyncio.run(_supervise(chosen, checked, command, server)) + except OSError as error: + sys.stderr.write(f"Rust gateway startup failed: {error}\n") + return ExitCode.CHILD_FAILED diff --git a/litellm/proxy/shutdown/graceful_uvicorn_worker.py b/litellm/proxy/shutdown/graceful_uvicorn_worker.py index 6ae5ae60c3f..8bbd64190f8 100644 --- a/litellm/proxy/shutdown/graceful_uvicorn_worker.py +++ b/litellm/proxy/shutdown/graceful_uvicorn_worker.py @@ -1,10 +1,16 @@ from uvicorn.workers import UvicornWorker -def build_graceful_shutdown_uvicorn_worker(timeout_graceful_shutdown: float) -> type[UvicornWorker]: +def build_graceful_shutdown_uvicorn_worker( + timeout_graceful_shutdown: float, *, proxy_headers: bool = True +) -> type[UvicornWorker]: class GracefulShutdownUvicornWorker(UvicornWorker): def __init__(self, *args: object) -> None: - self.CONFIG_KWARGS = {**UvicornWorker.CONFIG_KWARGS, "timeout_graceful_shutdown": timeout_graceful_shutdown} + self.CONFIG_KWARGS = { + **UvicornWorker.CONFIG_KWARGS, + "timeout_graceful_shutdown": timeout_graceful_shutdown, + "proxy_headers": proxy_headers, + } super().__init__(*args) return GracefulShutdownUvicornWorker diff --git a/litellm/rust_bridge/configuration.py b/litellm/rust_bridge/configuration.py index 152ba632996..5b51f155df9 100644 --- a/litellm/rust_bridge/configuration.py +++ b/litellm/rust_bridge/configuration.py @@ -8,7 +8,7 @@ from typing import Final from pydantic import TypeAdapter, ValidationError from typing_extensions import assert_never -_GLOBAL_ENV_NAME: Final = "LITELLM_RUST" +ENV_NAME: Final = "LITELLM_RUST" _ENV_BOOL: Final = TypeAdapter(bool) @@ -72,7 +72,7 @@ def decision(rollout: Rollout) -> Decision: return decide( rollout, process_override=_CONFIGURATION.override, - environment_override=_parse_env_bool(os.getenv(_GLOBAL_ENV_NAME)), + environment_override=_parse_env_bool(os.getenv(ENV_NAME)), ) diff --git a/litellm_cli/main.py b/litellm_cli/main.py index a9dfa4be179..ea170603d08 100644 --- a/litellm_cli/main.py +++ b/litellm_cli/main.py @@ -33,7 +33,7 @@ from contextlib import contextmanager from typing import TYPE_CHECKING, Final if TYPE_CHECKING: - from fastapi import FastAPI + from litellm.proxy.rust_forward.peer import ASGIApp SDK_PACKAGE: Final = "litellm" CLI_PACKAGE: Final = "litellm.proxy.client" @@ -214,10 +214,12 @@ def without_gc_pauses() -> Generator[None, None, None]: gc.enable() -def proxy_worker_app() -> "FastAPI": +def proxy_worker_app() -> "ASGIApp": with without_gc_pauses(): from litellm.proxy.proxy_cli import ProxyInitializationHelpers ProxyInitializationHelpers.apply_worker_log_config(os.environ) from litellm.proxy.proxy_server import app - return app + from litellm.proxy.rust_forward.mode import rust_forward_mode + from litellm.proxy.rust_forward.peer import wrap + return wrap(app, rust_forward_mode(os.environ)) diff --git a/tests/integration/configuration/test_rust_forward.py b/tests/integration/configuration/test_rust_forward.py new file mode 100644 index 00000000000..b75da286bf7 --- /dev/null +++ b/tests/integration/configuration/test_rust_forward.py @@ -0,0 +1,180 @@ +import json +import os +import subprocess +import sys +from pathlib import Path +from typing import Final + +import httpx +import pytest + +from tests.integration._support.client import Gateway, eventually +from tests.integration._support.ports import allocate_port +from tests.integration._support.process import stop_root_process + +TRUSTED_BALANCER: Final = "127.0.0.3" + + +def ready(client: httpx.Client) -> bool: + try: + return client.get("/health/liveliness").status_code == 200 + except httpx.TransportError: + return False + + +def client_ip(base_url: str, local_address: str, key: str, forwarded_for: str) -> httpx.Response: + with httpx.Client( + base_url=base_url, + transport=httpx.HTTPTransport(local_address=local_address), + trust_env=False, + timeout=5, + ) as client: + return client.get( + "/v1/mcp/network/client-ip", + headers={"Authorization": f"Bearer {key}", "X-Forwarded-For": forwarded_for}, + ) + + +@pytest.mark.parametrize("server_arguments", [(), ("--run_gunicorn",)]) +def test_gateway_preserves_metrics_bind_and_mcp_client_ip( + gateway: Gateway, tmp_path: Path, server_arguments: tuple[str, ...] +) -> None: + config: Final = tmp_path / "config.yaml" + config.write_text( + json.dumps( + { + "model_list": [], + "general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY", "use_x_forwarded_for": False}, + } + ) + ) + metrics_port: Final = allocate_port() + public_port: Final = allocate_port() + environment: Final = { + **{key: value for key, value in os.environ.items() if key not in ("DATABASE_URL", "DATABASE_URL_READ_REPLICA")}, + "LITELLM_MASTER_KEY": gateway.key, + "STORE_MODEL_IN_DB": "False", + "PYTHON_DOTENV_DISABLED": "1", + "LITELLM_RUST": "true", + "FORWARDED_ALLOW_IPS": TRUSTED_BALANCER, + } + with (tmp_path / "proxy.log").open("wb") as output: + with subprocess.Popen( + ( + sys.executable, + "-I", + "-m", + "litellm.proxy.proxy_cli", + "--config", + str(config), + "--host", + "0.0.0.0", + "--port", + str(public_port), + "--num_workers", + "2", + "--timeout_worker_healthcheck", + "120", + "--prometheus_metrics_port", + str(metrics_port), + *server_arguments, + ), + env=environment, + stdout=output, + stderr=output, + ) as supervisor: + try: + with httpx.Client( + base_url=f"http://127.0.0.1:{public_port}", + transport=httpx.HTTPTransport(local_address="127.0.0.2"), + trust_env=False, + follow_redirects=True, + timeout=5, + ) as external: + assert eventually(lambda: ready(external), bool, seconds=180), (tmp_path / "proxy.log").read_text() + response: Final = external.get(f"http://127.0.0.2:{metrics_port}/metrics") + assert response.status_code == 200, response.text + public: Final = f"http://127.0.0.1:{public_port}" + spoofed: Final = client_ip(public, "127.0.0.2", gateway.key, "1.2.3.4") + assert spoofed.status_code == 200, spoofed.text + assert spoofed.json() == {"ip": "127.0.0.2"} + behind_balancer: Final = client_ip(public, TRUSTED_BALANCER, gateway.key, "203.0.113.9") + assert behind_balancer.status_code == 200, behind_balancer.text + assert behind_balancer.json() == {"ip": "203.0.113.9"} + borrowed: Final = client_ip(public, "127.0.0.2", gateway.key, f"203.0.113.9, {TRUSTED_BALANCER}") + assert borrowed.status_code == 200, borrowed.text + assert borrowed.json() == {"ip": "127.0.0.2"} + finally: + assert stop_root_process(supervisor), (tmp_path / "proxy.log").read_text() + assert supervisor.returncode == 0, (tmp_path / "proxy.log").read_text() + + +@pytest.mark.parametrize( + "server_arguments", + [("--no-proxy-headers",), ("--forwarded-allow-ips", "127.0.0.1")], +) +def test_gateway_launcher_private_hop_preserves_client_address( + tmp_path: Path, server_arguments: tuple[str, ...] +) -> None: + config: Final = tmp_path / "config.yaml" + config.write_text( + json.dumps({"model_list": [], "general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY"}}) + ) + public_port: Final = allocate_port() + environment: Final = { + **{key: value for key, value in os.environ.items() if not key.startswith("DATABASE_") and key != "DIRECT_URL"}, + "LITELLM_MASTER_KEY": "sk-rust-forwarding-integration-key-not-a-secret", + "LITELLM_RUST": "true", + "LITELLM_PGBOUNCER_ENABLED": "false", + "CONFIG_FILE_PATH": str(config), + "STORE_MODEL_IN_DB": "False", + "PYTHON_DOTENV_DISABLED": "1", + } + log: Final = tmp_path / "gateway.log" + with log.open("wb") as output: + with subprocess.Popen( + ( + sys.executable, + "-I", + "-m", + "gateway.launch", + "--host", + "0.0.0.0", + "--port", + str(public_port), + "--workers", + "2", + "--timeout-worker-healthcheck", + "120", + *server_arguments, + ), + env=environment, + stdout=output, + stderr=output, + ) as supervisor: + try: + with httpx.Client( + base_url=f"http://127.0.0.1:{public_port}", + transport=httpx.HTTPTransport(local_address="127.0.0.2"), + trust_env=False, + timeout=5, + ) as external: + assert eventually(lambda: ready(external), bool, seconds=180), log.read_text() + response: Final = external.get( + "/health/liveliness?client-ip-proof", headers={"X-Forwarded-For": "198.51.100.10"} + ) + assert response.status_code == 200, response.text + events: Final = eventually( + lambda: tuple( + line + for line in log.read_text().splitlines() + if "GET /health/liveliness?client-ip-proof " in line + ), + bool, + seconds=10, + ) + assert len(events) == 1, log.read_text() + assert "127.0.0.2:" in events[0], events + finally: + assert stop_root_process(supervisor), log.read_text() + assert supervisor.returncode == 0, log.read_text() diff --git a/tests/integration/configuration/test_rust_forward_traffic.py b/tests/integration/configuration/test_rust_forward_traffic.py new file mode 100644 index 00000000000..eeeeea366d8 --- /dev/null +++ b/tests/integration/configuration/test_rust_forward_traffic.py @@ -0,0 +1,457 @@ +import asyncio +import json +import os +import socket +import subprocess +import sys +import threading +import time +from collections.abc import AsyncIterator, Callable, Iterator, Mapping +from contextlib import AbstractContextManager, ExitStack +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final + +import anyio +import httpx +import openai +import pytest +from pydantic import BaseModel, JsonValue +from starlette.applications import Starlette +from starlette.requests import Request as StarletteRequest +from starlette.responses import StreamingResponse +from starlette.routing import Route, WebSocketRoute +from starlette.websockets import WebSocket, WebSocketDisconnect +from websockets.sync.client import connect + +from tests.integration._support.asgi import asgi_server +from tests.integration._support.client import eventually +from tests.integration._support.mcp import mcp_peer +from tests.integration._support.ports import allocate_port +from tests.integration._support.process import stop_root_process +from tests.integration._support.wire import Reply, Request, wire_server + +MASTER_KEY: Final = "sk-rust-forward-traffic-key-not-a-secret" +KEEPALIVE_SECONDS: Final = 1 +FILES_MODEL: Final = "rust-forward-files" +FILE_ID: Final = "file-rust-forward" +CHAT_MODEL: Final = "rust-forward-chat" +REALTIME_MODEL: Final = "rust-forward-realtime" +MCP_HEADERS: Final = MappingProxyType( + { + "Authorization": f"Bearer {MASTER_KEY}", + "Accept": "application/json, text/event-stream", + "mcp-protocol-version": "2025-06-18", + } +) + + +@dataclass(frozen=True, slots=True) +class ScriptedProvider: + url: str + released: threading.Event + hung_up: SimpleQueue[str] + + +@dataclass(frozen=True, slots=True) +class RustProxy: + url: str + port: int + log: Path + provider: ScriptedProvider + + +class _Message(BaseModel): + content: str + + +class _ChatRequest(BaseModel): + messages: tuple[_Message, ...] + + +class _InputText(BaseModel): + text: str + + +class _Item(BaseModel): + content: tuple[_InputText, ...] = () + + +class _RealtimeFrame(BaseModel): + type: str + item: _Item | None = None + + +def _sse(content: str | None, finish_reason: str | None = None) -> bytes: + chunk: Final = { + "id": "chatcmpl-rust-forward", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "delta": {} if content is None else {"content": content}, "finish_reason": finish_reason} + ], + } + return b"data: " + json.dumps(chunk).encode() + b"\n\n" + + +async def _gated_stream(released: threading.Event) -> AsyncIterator[bytes]: + yield _sse("before-gate") + await anyio.to_thread.run_sync(released.wait, 30) + yield _sse("after-gate") + yield _sse(None, "stop") + yield b"data: [DONE]\n\n" + + +async def _endless_stream(marker: str, hung_up: SimpleQueue[str]) -> AsyncIterator[bytes]: + try: + for index in range(600): + yield _sse(f"tick-{index} ") + await asyncio.sleep(0.1) + except (asyncio.CancelledError, GeneratorExit): + hung_up.put(marker) + raise + + +def _provider_app(released: threading.Event, hung_up: SimpleQueue[str]) -> Starlette: + async def chat(request: StarletteRequest) -> StreamingResponse: + prompt: Final = _ChatRequest.model_validate_json(await request.body()).messages[-1].content + return StreamingResponse( + _gated_stream(released) if prompt == "gated" else _endless_stream(prompt, hung_up), + media_type="text/event-stream", + ) + + async def realtime(websocket: WebSocket) -> None: + await websocket.accept() + await websocket.send_json({"type": "session.created", "session": {"id": "sess_rust_forward"}}) + try: + while True: + frame = _RealtimeFrame.model_validate_json(await websocket.receive_text()) # rebind-ok: per frame + if frame.item is not None: + text = "".join(part.text for part in frame.item.content) # rebind-ok: per frame + await websocket.send_json({"type": "response.output_text.delta", "delta": f"echo:{text}"}) + except WebSocketDisconnect: + return + + return Starlette( + routes=[ + Route("/v1/chat/completions", chat, methods=["POST"]), + WebSocketRoute("/v1/realtime", realtime), + ] + ) + + +def _routable_address() -> str: + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as probe: + probe.connect(("192.0.2.1", 9)) + address: Final[str] = probe.getsockname()[0] + assert not address.startswith("127."), address + return address + + +def _ready(client: httpx.Client) -> bool: + try: + return client.get("/health/liveliness").status_code == 200 + except httpx.TransportError: + return False + + +def _files_peer(content: bytes) -> Callable[[Request], Reply]: + file_object: Final = json.dumps( + { + "id": FILE_ID, + "object": "file", + "purpose": "batch", + "bytes": len(content), + "created_at": 1, + "filename": "upload.jsonl", + "status": "processed", + } + ).encode() + + def respond(request: Request) -> Reply: + if request.method == "POST" and request.target.startswith("/v1/files"): + return Reply(body=file_object) + if request.method == "GET" and request.target.startswith(f"/v1/files/{FILE_ID}/content"): + return Reply( + content_type="application/octet-stream", body=content, chunks=(content[:65536], content[65536:]) + ) + return Reply(status=404, body=b'{"error": {"message": "not found"}}') + + return respond + + +UPLOAD: Final = b"".join( + json.dumps( + { + "custom_id": f"row-{index}", + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": os.urandom(48).hex()}]}, + } + ).encode() + + b"\n" + for index in range(30_000) +) + + +@pytest.fixture(scope="module") +def files_upstream() -> Iterator[str]: + with wire_server(_files_peer(UPLOAD), keep_alive=True) as wire: + yield wire.url + + +@pytest.fixture(scope="module") +def rust_proxy(tmp_path_factory: pytest.TempPathFactory, files_upstream: str) -> Iterator[RustProxy]: + directory: Final = tmp_path_factory.mktemp("rust-forward-traffic") + log: Final = directory / "proxy.log" + port: Final = allocate_port() + released: Final = threading.Event() + hung_up: Final[SimpleQueue[str]] = SimpleQueue() + with ExitStack() as stack: + mcp: Final = stack.enter_context(mcp_peer()) + provider: Final = ScriptedProvider( + stack.enter_context(asgi_server(_provider_app(released, hung_up))), released, hung_up + ) + config: Final = directory / "config.yaml" + config.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": CHAT_MODEL, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{provider.url}/v1", + "api_key": "integration-provider-key", + }, + }, + { + "model_name": REALTIME_MODEL, + "litellm_params": { + "model": "openai/gpt-realtime", + "api_base": provider.url, + "api_key": "integration-provider-key", + }, + "model_info": {"mode": "realtime"}, + }, + { + "model_name": FILES_MODEL, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{files_upstream}/v1", + "api_key": "integration-provider-key", + }, + }, + ], + "mcp_servers": {"math": {**mcp.registration(), "available_on_public_internet": False}}, + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "mcp_internal_ip_ranges": ["127.0.0.0/8"], + }, + } + ) + ) + environment: Final = { + **{ + key: value + for key, value in os.environ.items() + if not key.startswith("DATABASE_") and key != "DIRECT_URL" + }, + "LITELLM_MASTER_KEY": MASTER_KEY, + "LITELLM_RUST": "true", + "STORE_MODEL_IN_DB": "False", + "PYTHON_DOTENV_DISABLED": "1", + } + output: Final = stack.enter_context(log.open("wb")) + supervisor: Final = stack.enter_context( + subprocess.Popen( + ( + sys.executable, + "-I", + "-m", + "litellm.proxy.proxy_cli", + "--config", + str(config), + "--host", + "0.0.0.0", + "--port", + str(port), + "--num_workers", + "2", + "--keepalive_timeout", + str(KEEPALIVE_SECONDS), + "--timeout_worker_healthcheck", + "120", + ), + env=environment, + stdout=output, + stderr=output, + ) + ) + try: + with httpx.Client(base_url=f"http://127.0.0.1:{port}", trust_env=False, timeout=5) as probe: + assert eventually(lambda: _ready(probe), bool, seconds=180), log.read_text() + yield RustProxy(f"http://127.0.0.1:{port}", port, log, provider) + finally: + released.set() + assert stop_root_process(supervisor), log.read_text() + assert supervisor.returncode == 0, log.read_text() + + +def _still_open(stream: httpx.Response) -> bool: + try: + next(stream.iter_raw()) + except httpx.ReadTimeout: + return True + except (StopIteration, httpx.TransportError): + return False + return True + + +def _initialize(client: httpx.Client, headers: Mapping[str, str]) -> httpx.Response: + return client.post( + "/mcp", + headers={**MCP_HEADERS, **headers}, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "rust-forward", "version": "1"}, + }, + }, + ) + + +def _mcp_post(client: httpx.Client, session: str, body: Mapping[str, JsonValue]) -> httpx.Response: + return client.post("/mcp", headers={**MCP_HEADERS, "mcp-session-id": session}, json=body) + + +def test_mcp_session_keeps_working_while_its_get_stream_sits_idle(rust_proxy: RustProxy) -> None: + idle_seconds: Final = 8 * KEEPALIVE_SECONDS + with httpx.Client(base_url=rust_proxy.url, trust_env=False, timeout=30) as client: + initialized: Final = _initialize(client, {}) + assert initialized.status_code == 200, initialized.text + session: Final = initialized.headers["mcp-session-id"] + assert _mcp_post(client, session, {"jsonrpc": "2.0", "method": "notifications/initialized"}).status_code == 202 + with httpx.Client(base_url=rust_proxy.url, trust_env=False, timeout=httpx.Timeout(5, read=2)) as listener: + with listener.stream( + "GET", "/mcp", headers={**MCP_HEADERS, "Accept": "text/event-stream", "mcp-session-id": session} + ) as notifications: + assert notifications.status_code == 200 + assert notifications.headers["content-type"].startswith("text/event-stream") + time.sleep(idle_seconds) + called: Final = _mcp_post( + client, + session, + { + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": {"name": "math-add", "arguments": {"a": 2, "b": 3}}, + }, + ) + assert called.status_code == 200, called.text + assert '"text":"5"' in called.text.replace(" ", ""), called.text + assert _still_open(notifications) + + +def test_uploaded_file_and_its_download_cross_byte_for_byte(rust_proxy: RustProxy) -> None: + with openai.OpenAI(api_key=MASTER_KEY, base_url=rust_proxy.url, max_retries=0, timeout=60) as client: + created: Final = client.files.create( + file=("upload.jsonl", UPLOAD), purpose="batch", extra_body={"model": FILES_MODEL} + ) + downloaded: Final = client.files.content(created.id, extra_query={"model": FILES_MODEL}).content + assert len(UPLOAD) > 4 * 1024 * 1024 + assert downloaded == UPLOAD, (len(downloaded), len(UPLOAD)) + + +def test_internal_only_mcp_server_judges_the_connecting_peer_not_its_forwarded_for_claim( + rust_proxy: RustProxy, +) -> None: + outsider: Final = _routable_address() + with httpx.Client( + base_url=f"http://{outsider}:{rust_proxy.port}", + transport=httpx.HTTPTransport(local_address=outsider), + trust_env=False, + timeout=10, + ) as spoofing: + spoofed: Final = _initialize(spoofing, {"X-Forwarded-For": "127.0.0.1"}) + with httpx.Client(base_url=rust_proxy.url, trust_env=False, timeout=10) as internal: + admitted: Final = _initialize(internal, {}) + assert admitted.status_code == 200, admitted.text + assert spoofed.status_code == 403, spoofed.text + assert "client IP" in spoofed.text, spoofed.text + + +def test_oauth_metadata_names_the_origin_the_client_used(rust_proxy: RustProxy) -> None: + with httpx.Client(base_url=rust_proxy.url, trust_env=False, timeout=10) as client: + direct: Final = client.get("/.well-known/oauth-protected-resource") + balanced: Final = client.get( + "/.well-known/oauth-protected-resource", + headers={"Host": "gateway.example.test", "X-Forwarded-Proto": "https"}, + ) + assert direct.status_code == 200, direct.text + assert direct.json()["resource"] == rust_proxy.url + assert direct.json()["authorization_servers"] == [f"{rust_proxy.url}/mcp"] + assert balanced.status_code == 200, balanced.text + assert balanced.json()["resource"] == "https://gateway.example.test" + + +def _chat_stream(client: httpx.Client, prompt: str) -> AbstractContextManager[httpx.Response]: + return client.stream( + "POST", + "/v1/chat/completions", + headers={"Authorization": f"Bearer {MASTER_KEY}"}, + json={"model": CHAT_MODEL, "stream": True, "messages": [{"role": "user", "content": prompt}]}, + ) + + +def _content(lines: Iterator[str]) -> str: + return next(line for line in lines if '"content"' in line) + + +def test_stream_chunk_reaches_the_client_before_the_provider_sends_the_next(rust_proxy: RustProxy) -> None: + with httpx.Client(base_url=rust_proxy.url, trust_env=False, timeout=httpx.Timeout(10, read=5)) as client: + with _chat_stream(client, "gated") as response: + assert response.status_code == 200, response.read() + lines: Final = response.iter_lines() + first: Final = _content(lines) + rust_proxy.provider.released.set() + rest: Final = tuple(line for line in lines if line) + assert "before-gate" in first, first + assert any("after-gate" in line for line in rest), rest + assert rest[-1] == "data: [DONE]", rest + + +def test_client_hang_up_mid_stream_closes_the_provider_stream(rust_proxy: RustProxy) -> None: + marker: Final = f"endless-{time.monotonic_ns()}" + with httpx.Client(base_url=rust_proxy.url, trust_env=False, timeout=httpx.Timeout(10, read=5)) as client: + with _chat_stream(client, marker) as response: + assert response.status_code == 200, response.read() + assert "tick-0" in _content(response.iter_lines()) + assert rust_proxy.provider.hung_up.get(timeout=15) == marker + + +def test_realtime_websocket_frames_cross_both_ways(rust_proxy: RustProxy) -> None: + marker: Final = f"realtime-{time.monotonic_ns()}" + with connect( + f"ws://127.0.0.1:{rust_proxy.port}/v1/realtime?model={REALTIME_MODEL}", + additional_headers={"Authorization": f"Bearer {MASTER_KEY}"}, + open_timeout=10, + ) as session: + session.send( + json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": marker}]}, + } + ) + ) + replies: Final = tuple(session.recv(timeout=10) for _ in range(2)) + assert any(f"echo:{marker}" in reply for reply in replies), replies diff --git a/tests/unit/gateway/test_launch.py b/tests/unit/gateway/test_launch.py index d682fcbad2d..50dab887688 100644 --- a/tests/unit/gateway/test_launch.py +++ b/tests/unit/gateway/test_launch.py @@ -12,13 +12,17 @@ from typing import Final, cast from unittest.mock import MagicMock, patch import pytest +import uvicorn from uvicorn.importer import import_from_string from uvicorn.main import main as uvicorn_main -from gateway.launch import GATEWAY_APP, main, pool_database_url, uvicorn_argv +from gateway.launch import GATEWAY_APP, main, pool_database_url, python_server, uvicorn_argv from litellm.proxy.db.db_url_settings import DatabaseURLSettings from litellm.proxy.db.pgbouncer import PGBOUNCER_POOLED_ENV_VAR, PgBouncerError, PgBouncerSettings from litellm.proxy.proxy_server import app as proxy_app +from litellm.proxy.rust_forward.mode import AssignedPort, InheritedFd, Supervised, uvicorn_flags +from litellm.proxy.rust_forward.server import Listener +from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager DB_ENV: Final = { "DATABASE_HOST": "db.internal", @@ -114,6 +118,100 @@ def _uvicorn_params(argv: tuple[str, ...]) -> dict[str, object]: class TestUvicornArgv: + @pytest.mark.parametrize( + ("arguments", "listener"), + [ + ((), Listener.TCP), + (("--uds", "/tmp/gateway.sock"), Listener.UNIX_SOCKET), + (("--fd", "0"), Listener.INHERITED_FD), + (("--fd=17",), Listener.INHERITED_FD), + (("--ssl-keyfile", "unused"), Listener.TLS), + (("--ssl-ciphers", "TLSv1.2"), Listener.TLS), + ], + ) + def test_forwarding_sees_the_public_listener_uvicorn_would_open( + self, arguments: tuple[str, ...], listener: Listener + ) -> None: + assert python_server(arguments).listener is listener + + def test_an_unsupported_listener_stops_forwarding_before_the_database_is_prepared( + self, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] + ) -> None: + monkeypatch.setenv("LITELLM_RUST", "true") + monkeypatch.delenv("_LITELLM_GATEWAY_FD", raising=False) + monkeypatch.setenv("LITELLM_PGBOUNCER_ENABLED", "true") + monkeypatch.setenv("LITELLM_PGBOUNCER_BINARY", "/nonexistent/pgbouncer") + with pytest.raises(SystemExit) as stopped: + main(("--uds", "/tmp/gateway.sock"), serve=lambda argv: None) + assert stopped.value.code == 1 + assert capsys.readouterr().err == "LITELLM_RUST does not support a Unix socket on the public port\n" + + @pytest.mark.parametrize( + ("argv", "trust"), + [ + ((), (True, None)), + (("--forwarded-allow-ips", "10.0.0.5"), (True, "10.0.0.5")), + (("--forwarded-allow-ips=*",), (True, "*")), + (("--no-proxy-headers",), (False, None)), + ], + ) + def test_forwarding_carries_operator_proxy_header_options_to_the_child( + self, argv: tuple[str, ...], trust: tuple[bool, str | None], monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv("FORWARDED_ALLOW_IPS", raising=False) + server: Final = python_server(("--host", "0.0.0.0", "--port", "4000", *argv)) + assert (server.host, server.port) == ("0.0.0.0", 4000) + assert (server.proxy_headers, server.trusted_proxies) == trust + + @pytest.mark.parametrize( + ("argv", "environ"), + [(("--timeout-keep-alive", "2"), {}), ((), {"KEEPALIVE_TIMEOUT": "7"}), ((), {})], + ) + def test_rust_retires_connections_by_the_keepalive_uvicorn_runs_with( + self, argv: tuple[str, ...], environ: dict[str, str], monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv("KEEPALIVE_TIMEOUT", raising=False) + for name, value in environ.items(): + monkeypatch.setenv(name, value) + served: Final = _uvicorn_params(uvicorn_argv(argv, os.environ))["timeout_keep_alive"] + assert python_server(argv).keepalive_seconds == served + + def test_without_a_keepalive_rust_follows_uvicorns_own_default(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("KEEPALIVE_TIMEOUT", raising=False) + assert python_server(()).keepalive_seconds == uvicorn.Config(app="", log_config=None).timeout_keep_alive + + @pytest.mark.parametrize("argv", [(), ("--timeout-graceful-shutdown", "12")]) + def test_rust_drains_for_pythons_graceful_timeout_plus_teardown( + self, argv: tuple[str, ...], monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "7.5") + flag: Final = _uvicorn_params(uvicorn_argv(argv, {}))["timeout_graceful_shutdown"] + graceful: Final = GracefulShutdownManager.get_timeout() if flag is None else float(str(flag)) + assert python_server(argv).graceful_timeout_seconds == graceful + + def test_supervised_listener_uses_the_inherited_loopback_socket(self): + supervised: Final = uvicorn_flags(Supervised(InheritedFd(17), True, "127.0.0.1")) + params: Final = _uvicorn_params((*uvicorn_argv(("--workers", "2"), {}), *supervised)) + assert params["fd"] == 17 + assert params["workers"] == 2 + + def test_assigned_port_listener_binds_only_loopback(self): + supervised: Final = uvicorn_flags(Supervised(AssignedPort(45556), True, "127.0.0.1")) + params: Final = _uvicorn_params((*uvicorn_argv(("--host", "0.0.0.0", "--port", "4000"), {}), *supervised)) + assert (params["host"], params["port"], params["fd"]) == ("127.0.0.1", 45556, None) + + @pytest.mark.parametrize("argv", [(), ("--proxy-headers",), ("--forwarded-allow-ips=*",)]) + def test_private_hop_leaves_proxy_headers_to_the_app(self, argv: tuple[str, ...]): + params: Final = _uvicorn_params( + (*uvicorn_argv(argv, {}), *uvicorn_flags(Supervised(InheritedFd(17), True, "*"))) + ) + assert params["proxy_headers"] is False + + def test_direct_listener_keeps_explicit_proxy_header_options(self): + params: Final = _uvicorn_params(uvicorn_argv(("--no-proxy-headers", "--forwarded-allow-ips=*"), {})) + assert params["proxy_headers"] is False + assert params["forwarded_allow_ips"] == "*" + def test_keepalive_env_reaches_uvicorn(self): params: Final = _uvicorn_params(uvicorn_argv(("--workers", "4"), {"KEEPALIVE_TIMEOUT": "75"})) assert params["app"] == GATEWAY_APP @@ -229,6 +327,16 @@ class TestMain: DatabaseURLSettings.from_env().apply_to_env() assert os.environ["DATABASE_URL"] == pooled + def test_forwarding_disabled_keeps_uvicorn_socket_options( + self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("LITELLM_RUST", "false") + monkeypatch.delenv("_LITELLM_GATEWAY_FD", raising=False) + monkeypatch.delenv("LITELLM_PGBOUNCER_ENABLED", raising=False) + served: Final[list[tuple[str, ...]]] = [] + main(("--uds", "/tmp/gateway.sock"), serve=lambda argv: served.append(tuple(argv))) + assert _uvicorn_params(served[0])["uds"] == "/tmp/gateway.sock" + def test_a_pooler_that_cannot_start_stops_the_gateway_before_uvicorn( self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path ): diff --git a/tests/unit/proxy/rust_forward/__init__.py b/tests/unit/proxy/rust_forward/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/rust_forward/test_mode.py b/tests/unit/proxy/rust_forward/test_mode.py new file mode 100644 index 00000000000..9d57d4806c3 --- /dev/null +++ b/tests/unit/proxy/rust_forward/test_mode.py @@ -0,0 +1,112 @@ +from typing import Final + +import pytest +import uvicorn +from pydantic import ValidationError + +from litellm.proxy.proxy_cli import WORKER_APP, ProxyInitializationHelpers +from litellm.proxy.rust_forward.mode import ( + AssignedPort, + Disabled, + InheritedFd, + RustForwardMode, + Supervised, + Supervisor, + gunicorn_bind, + rust_forward_mode, + server_proxy_headers, + supervised_environment, + uvicorn_options, +) + + +@pytest.mark.parametrize( + ("environ", "expected"), + [ + ({}, Disabled()), + ({"LITELLM_RUST": "false"}, Disabled()), + ({"LITELLM_RUST": "off"}, Disabled()), + ({"LITELLM_RUST": "0"}, Disabled()), + ({"LITELLM_RUST": "true"}, Supervisor()), + ({"LITELLM_RUST": "1"}, Supervisor()), + ({"LITELLM_RUST": "yes"}, Supervisor()), + ({"LITELLM_RUST": "on"}, Supervisor()), + ({"LITELLM_RUST": " 1 "}, Supervisor()), + ({"LITELLM_RUST": ""}, Disabled()), + ({"LITELLM_RUST": " "}, Disabled()), + ( + {"LITELLM_RUST": "true", "_LITELLM_GATEWAY_FD": "17", "FORWARDED_ALLOW_IPS": "127.0.0.1,10.0.0.5"}, + Supervised(InheritedFd(17), True, "127.0.0.1,10.0.0.5"), + ), + ], +) +def test_mode_is_resolved_from_environment(environ: dict[str, str], expected: RustForwardMode) -> None: + assert rust_forward_mode(environ) == expected + + +@pytest.mark.parametrize( + ("environ", "variable"), + [ + ({"LITELLM_RUST": "enabled"}, "LITELLM_RUST"), + ({"LITELLM_RUST": "2"}, "LITELLM_RUST"), + ({"_LITELLM_GATEWAY_FD": "seventeen"}, "_LITELLM_GATEWAY_FD"), + ({"_LITELLM_GATEWAY_FD": "17", "_LITELLM_GATEWAY_PROXY_HEADERS": "maybe"}, "_LITELLM_GATEWAY_PROXY_HEADERS"), + ], +) +def test_invalid_forwarding_value_stops_startup_naming_the_variable(environ: dict[str, str], variable: str) -> None: + with pytest.raises(ValidationError, match=variable): + rust_forward_mode(environ) + + +@pytest.mark.parametrize( + ("proxy_headers", "trusted_proxies", "inherited", "expected"), + [ + (True, None, {}, Supervised(InheritedFd(17), True, "127.0.0.1")), + (True, None, {"FORWARDED_ALLOW_IPS": "10.0.0.5"}, Supervised(InheritedFd(17), True, "10.0.0.5")), + (True, "*", {"FORWARDED_ALLOW_IPS": "10.0.0.5"}, Supervised(InheritedFd(17), True, "*")), + (False, None, {}, Supervised(InheritedFd(17), False, "127.0.0.1")), + (True, None, {}, Supervised(AssignedPort(45555), True, "127.0.0.1")), + (False, "*", {}, Supervised(AssignedPort(45555), False, "*")), + ], +) +def test_child_reads_back_the_listener_and_trust_the_supervisor_wrote( + proxy_headers: bool, trusted_proxies: str | None, inherited: dict[str, str], expected: Supervised +) -> None: + written: Final = supervised_environment(expected.listener, proxy_headers, trusted_proxies) + assert rust_forward_mode({**inherited, **written}) == expected + + +def _served_uvicorn_config(mode: RustForwardMode) -> uvicorn.Config: + options: Final = { + **ProxyInitializationHelpers.get_default_unvicorn_init_args("0.0.0.0", 45555), + **uvicorn_options(mode, WORKER_APP), + } + return uvicorn.Config(**options) + + +def test_supervised_uvicorn_serves_the_inherited_socket_without_trusting_the_rust_hop() -> None: + config: Final = _served_uvicorn_config(Supervised(InheritedFd(17), True, "*")) + assert (config.fd, config.proxy_headers) == (17, False) + assert (config.app, config.factory) == (WORKER_APP, True) + + +@pytest.mark.parametrize("mode", [Disabled(), Supervisor()]) +def test_unsupervised_servers_keep_their_own_socket_and_proxy_headers(mode: RustForwardMode) -> None: + config: Final = _served_uvicorn_config(mode) + assert (config.fd, config.proxy_headers, config.port) == (None, True, 45555) + assert gunicorn_bind(mode, "0.0.0.0:4000") == "0.0.0.0:4000" + assert server_proxy_headers(mode) is True + + +def test_supervised_gunicorn_binds_the_inherited_socket_without_trusting_the_rust_hop() -> None: + mode: Final = Supervised(InheritedFd(17), True, "*") + assert gunicorn_bind(mode, "0.0.0.0:4000") == "fd://17" + assert server_proxy_headers(mode) is False + + +def test_assigned_port_serves_only_the_loopback_hop_without_trusting_it() -> None: + mode: Final = Supervised(AssignedPort(45556), True, "*") + config: Final = _served_uvicorn_config(mode) + assert (config.host, config.port, config.fd, config.proxy_headers) == ("127.0.0.1", 45556, None, False) + assert gunicorn_bind(mode, "0.0.0.0:4000") == "127.0.0.1:45556" + assert server_proxy_headers(mode) is False diff --git a/tests/unit/proxy/rust_forward/test_peer.py b/tests/unit/proxy/rust_forward/test_peer.py new file mode 100644 index 00000000000..ee625c4a4be --- /dev/null +++ b/tests/unit/proxy/rust_forward/test_peer.py @@ -0,0 +1,143 @@ +import asyncio +import gc +from typing import TYPE_CHECKING, Final + +import httpx +import pytest +from fastapi import Request +from fastapi.routing import APIRoute +from uvicorn.importer import import_from_string +from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware + +if TYPE_CHECKING: + from uvicorn._types import ( + ASGIReceiveCallable, + ASGISendCallable, + ASGISendEvent, + HTTPDisconnectEvent, + HTTPScope, + Scope, + ) + +from litellm.proxy.proxy_server import app as proxy_app +from litellm.proxy.rust_forward.mode import InheritedFd, Supervised +from litellm.proxy.rust_forward.peer import wrap + + +Headers = tuple[tuple[bytes, bytes], ...] + + +async def _receive() -> "HTTPDisconnectEvent": + return {"type": "http.disconnect"} + + +async def _send(message: "ASGISendEvent") -> None: + return None + + +async def _seen_by_app( + behind_rust: bool, client: tuple[str, int], headers: Headers, mode: Supervised +) -> tuple[tuple[str, int] | None, str, Headers]: + seen: Final[asyncio.Future[tuple[tuple[str, int] | None, str, Headers]]] = ( + asyncio.get_running_loop().create_future() + ) + + async def probe(scope: "Scope", receive: "ASGIReceiveCallable", send: "ASGISendCallable") -> None: + assert scope["type"] == "http" + seen.set_result((scope["client"], scope["scheme"], tuple(scope["headers"]))) + + direct: Final = ProxyHeadersMiddleware(probe, mode.trusted_proxies) if mode.proxy_headers else probe + app: Final = wrap(probe, mode) if behind_rust else direct + scope: Final[HTTPScope] = { + "type": "http", + "asgi": {"version": "3.0", "spec_version": "2.3"}, + "http_version": "1.1", + "method": "GET", + "scheme": "http", + "path": "/", + "raw_path": b"/", + "query_string": b"", + "root_path": "", + "headers": headers, + "client": client, + "server": ("127.0.0.1", 4000), + } + await app(scope, _receive, _send) + return seen.result() + + +PEER: Final = ("203.0.113.7", 51234) +TRUSTED_PEER: Final = ("10.0.12.34", 40000) +FORWARDED: Final[Headers] = ( + (b"x-forwarded-for", b"198.51.100.1, 203.0.113.7"), + (b"x-forwarded-proto", b"https"), +) + + +@pytest.mark.parametrize("peer", [PEER, TRUSTED_PEER, ("2001:db8::1", 443)]) +@pytest.mark.parametrize("headers", [(), FORWARDED, ((b"x-forwarded-for", b"10.1.2.3"), (b"host", b"public.example"))]) +@pytest.mark.parametrize( + "mode", + [ + Supervised(InheritedFd(17), True, "127.0.0.1"), + Supervised(InheritedFd(17), True, "10.0.0.0/8"), + Supervised(InheritedFd(17), True, "*"), + Supervised(InheritedFd(17), False, "10.0.0.0/8"), + ], +) +async def test_python_behind_rust_sees_what_a_direct_connection_would( + peer: tuple[str, int], headers: Headers, mode: Supervised +) -> None: + rendered: Final = f"[{peer[0]}]:{peer[1]}" if ":" in peer[0] else f"{peer[0]}:{peer[1]}" + behind_rust: Final = await _seen_by_app( + True, ("127.0.0.1", 60000), (*headers, (b"x-litellm-rust-peer", rendered.encode())), mode + ) + direct: Final = await _seen_by_app(False, peer, headers, mode) + assert behind_rust == direct + + +async def test_untrusted_peer_cannot_set_the_scheme_behind_rust() -> None: + client, scheme, _ = await _seen_by_app( + True, + ("127.0.0.1", 60000), + ((b"x-forwarded-proto", b"https"), (b"x-litellm-rust-peer", b"203.0.113.7:51234")), + Supervised(InheritedFd(17), True, "127.0.0.1"), + ) + assert (client, scheme) == (("203.0.113.7", 51234), "http") + + +async def test_peer_header_from_anyone_but_the_rust_hop_is_dropped_and_ignored() -> None: + client, _, headers = await _seen_by_app( + True, + ("192.0.2.9", 1000), + ((b"x-litellm-rust-peer", b"10.0.12.34:1"),), + Supervised(InheritedFd(17), False, "127.0.0.1"), + ) + assert (client, headers) == (("192.0.2.9", 1000), ()) + + +async def _peer_probe(request: Request) -> dict[str, object]: + assert request.client is not None + return {"client": [request.client.host, request.client.port], "peer": request.headers.get("x-litellm-rust-peer")} + + +@pytest.mark.parametrize("entry_point", ["litellm_cli.main:proxy_worker_app", "gateway.launch:worker_app"]) +async def test_every_entry_point_hands_routes_the_peer_rust_saw( + entry_point: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("_LITELLM_GATEWAY_FD", "17") + monkeypatch.setattr(proxy_app.router, "lifespan_context", proxy_app.router.lifespan_context) + monkeypatch.setattr( + proxy_app.router, "routes", [APIRoute("/__rust_forward_peer_probe", _peer_probe), *proxy_app.router.routes] + ) + try: + served: Final = import_from_string(entry_point)() + finally: + gc.unfreeze() + transport: Final = httpx.ASGITransport(app=served, client=("127.0.0.1", 60000)) + async with httpx.AsyncClient(transport=transport, base_url="http://proxy") as client: + response: Final = await client.get( + "/__rust_forward_peer_probe", headers={"x-litellm-rust-peer": "203.0.113.7:51234"} + ) + assert response.status_code == 200, response.text + assert response.json() == {"client": ["203.0.113.7", 51234], "peer": None} diff --git a/tests/unit/proxy/rust_forward/test_supervisor.py b/tests/unit/proxy/rust_forward/test_supervisor.py new file mode 100644 index 00000000000..d0ba6872bbd --- /dev/null +++ b/tests/unit/proxy/rust_forward/test_supervisor.py @@ -0,0 +1,241 @@ +import functools +import signal +import socket +import subprocess +import sys +from pathlib import Path +from typing import Final + +import pytest + +from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS +from litellm.proxy.proxy_cli import run_server +from litellm.proxy.rust_forward.mode import AssignedPort, InheritedFd, Supervised, rust_forward_mode +from litellm.proxy.rust_forward.platform import Platform, PrivateListener, current_platform +from litellm.proxy.rust_forward.platform.posix import PosixPlatform +from litellm.proxy.rust_forward.platform.windows import ( + CREATE_NEW_PROCESS_GROUP, + CTRL_BREAK_EVENT, + SIGBREAK, + WindowsPlatform, +) +from litellm.proxy.rust_forward.server import Listener, PythonServer, Server +from litellm.proxy.rust_forward.supervisor import ( + Child, + Event, + ExitCode, + Exited, + Finished, + Ready, + Running, + Signal, + Start, + StartupError, + Stop, + Stopping, + Transition, + boot, + gateway_binary, + python_spec, + rust_spec, + startup_check, + step, + supervise, +) + +TERM: Final = signal.SIGTERM +INT: Final = signal.SIGINT +POSIX: Final = PosixPlatform() +WINDOWS: Final = WindowsPlatform() +POSIX_ONLY: Final = pytest.mark.skipif(sys.platform == "win32", reason="runs a POSIX shebang gateway") + + +def _server( + server: Server = Server.UVICORN, listener: Listener = Listener.TCP, trusted_proxies: str | None = None +) -> PythonServer: + return PythonServer(server, listener, "0.0.0.0", 4000, 2, 2.5, False, trusted_proxies) + + +def test_python_child_serves_the_private_socket_with_the_operators_trust() -> None: + private: Final = PrivateListener(InheritedFd(17), 45555, (17,)) + spec: Final = python_spec(POSIX, ("python",), private, _server(trusted_proxies="10.0.0.5")) + assert rust_forward_mode(spec.env_overrides) == Supervised(InheritedFd(17), False, "10.0.0.5") + assert (spec.spawn.pass_fds, spec.spawn.start_new_session) == ((17,), True) + + +def test_windows_python_child_binds_the_assigned_port_in_its_own_process_group() -> None: + with WINDOWS.private_listener() as private: + pass + spec: Final = python_spec(WINDOWS, ("python",), private, _server()) + assert rust_forward_mode(spec.env_overrides) == Supervised(AssignedPort(private.port), False, "127.0.0.1") + assert spec.argv[-2:] == ("--port", str(private.port)) + assert (spec.spawn.pass_fds, spec.spawn.creationflags) == ((), CREATE_NEW_PROCESS_GROUP) + with socket.create_server(("127.0.0.1", private.port)): + pass + + +@pytest.mark.skipif(sys.platform != "win32", reason="these names exist only on Windows") +def test_windows_constants_match_the_stdlib() -> None: + assert (CTRL_BREAK_EVENT, SIGBREAK) == (signal.CTRL_BREAK_EVENT, signal.SIGBREAK) + assert CREATE_NEW_PROCESS_GROUP == subprocess.CREATE_NEW_PROCESS_GROUP + + +def test_python_child_keeps_the_public_host_and_its_options() -> None: + arguments: Final = ("--host", "0.0.0.0", "--port", "4000", "--num_workers", "2", "--reload") + spec: Final = python_spec(POSIX, arguments, PrivateListener(InheritedFd(17), 45555, (17,)), _server()) + params: Final = run_server.make_context("litellm", list(spec.argv)).params + assert (params["host"], params["port"]) == ("0.0.0.0", 45555) + assert (params["num_workers"], params["reload"]) == (2, True) + + +def test_rust_child_serves_public_port_and_forwards_to_private_listener() -> None: + spec: Final = rust_spec(POSIX, "gateway", 45555, _server()) + assert spec.argv == ("gateway",) + assert spec.spawn.pass_fds == () + assert dict(spec.env_overrides) == { + "LITELLM_FORWARD_UPSTREAM": "http://127.0.0.1:45555", + "LITELLM_FORWARD_HOST": "0.0.0.0", + "LITELLM_FORWARD_PORT": "4000", + "LITELLM_FORWARD_UPSTREAM_KEEPALIVE": "2", + "LITELLM_FORWARD_SHUTDOWN_DRAIN": str(2.5 + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS), + } + + +@pytest.mark.parametrize( + ("server", "reason"), + [ + (_server(Server.HYPERCORN), "LITELLM_RUST does not support Hypercorn"), + (_server(Server.GRANIAN), "LITELLM_RUST does not support Granian"), + (_server(listener=Listener.TLS), "LITELLM_RUST does not support TLS on the public port"), + (_server(listener=Listener.UNIX_SOCKET), "LITELLM_RUST does not support a Unix socket on the public port"), + ( + _server(Server.GUNICORN, Listener.INHERITED_FD), + "LITELLM_RUST does not support an inherited file descriptor on the public port", + ), + ], +) +def test_unsupported_servers_are_refused_before_anything_starts( + server: PythonServer, reason: str, capsys: pytest.CaptureFixture[str] +) -> None: + assert startup_check(POSIX, server, "gateway") == StartupError(reason) + prepared: Final[list[bool]] = [] + refused: Final = supervise(server, ("python",), before_start=lambda: prepared.append(True), platform=POSIX) + assert refused == ExitCode.CHILD_FAILED + assert (capsys.readouterr().err, prepared) == (f"{reason}\n", []) + + +@pytest.mark.parametrize("server", [Server.UVICORN, Server.GUNICORN]) +def test_plain_tcp_uvicorn_and_gunicorn_can_be_forwarded(server: Server) -> None: + assert startup_check(POSIX, _server(server), "gateway") == "gateway" + + +def test_windows_forwards_uvicorn_and_refuses_gunicorn() -> None: + assert startup_check(WINDOWS, _server(), "gateway") == "gateway" + assert startup_check(WINDOWS, _server(Server.GUNICORN), "gateway") == StartupError( + "LITELLM_RUST does not support gunicorn on Windows" + ) + + +def test_missing_binary_gives_install_instruction() -> None: + assert startup_check(POSIX, _server(), None) == StartupError( + "LITELLM_RUST requires the litellm-gateway binary. Install litellm[proxy] in this environment" + ) + + +@POSIX_ONLY +def test_binary_is_found_only_in_the_environment_scripts_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + scripts: Final = tmp_path / "scripts" + elsewhere: Final = tmp_path / "elsewhere" + scripts.mkdir() + elsewhere.mkdir() + binary: Final = elsewhere / "litellm-gateway" + binary.write_text("#!/bin/sh\nexit 0\n") + binary.chmod(0o755) + monkeypatch.setenv("PATH", str(elsewhere)) + assert gateway_binary(POSIX, str(scripts)) is None + assert gateway_binary(POSIX, str(elsewhere)) == str(binary) + + +PYTHON_CHILD: Final = """ +import os, pathlib, socket, time +time.sleep(0.5) +pathlib.Path(os.environ["MARKERS"], "listening").touch() +fd = os.environ.get("_LITELLM_GATEWAY_FD") +port = os.environ.get("_LITELLM_GATEWAY_PORT") +listener = socket.socket(fileno=int(fd)) if fd else socket.create_server(("127.0.0.1", int(port))) +listener.listen() +time.sleep(10) +""" +RUST_CHILD: Final = """ +import os, pathlib +markers = pathlib.Path(os.environ["MARKERS"]) +(markers / "rust_saw_python_listening").write_text(str((markers / "listening").exists())) +""" + + +def _fake_gateway(scripts: Path) -> Platform: + if sys.platform != "win32": + gateway: Final = scripts / "litellm-gateway" + gateway.write_text(f"#!{sys.executable} -I\n{RUST_CHILD}") + gateway.chmod(0o755) + return current_platform() + script: Final = scripts / "rust_child.py" + script.write_text(RUST_CHILD) + (scripts / "litellm-gateway.cmd").write_text(f'@"{sys.executable}" -I "{script}"\r\n') + return WindowsPlatform(gateway_name="litellm-gateway.cmd") + + +def test_rust_opens_the_public_port_only_once_python_listens(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("MARKERS", str(tmp_path)) + command: Final = (sys.executable, "-I", "-c", PYTHON_CHILD) + supervise(_server(), command, scripts_dir=str(tmp_path), platform=_fake_gateway(tmp_path)) + assert (tmp_path / "rust_saw_python_listening").read_text() == "True" + + +def _fold(events: tuple[Event, ...]) -> Transition: + def advance(acc: Transition, event: Event) -> Transition: + state, done = acc + following, actions = step(state, event) + return following, (*done, *actions) + + return functools.reduce(advance, events, boot()) + + +RUST: Final = Child.RUST +PYTHON: Final = Child.PYTHON +PYTHON_READY: Final = Ready(PYTHON) +STARTED: Final = (Start(PYTHON), Start(RUST)) +BOTH_TERM: Final = (*STARTED, Signal(RUST, TERM), Signal(PYTHON, TERM)) +CLEAN: Final = Finished(ExitCode.CLEAN) +FAILED: Final = Finished(ExitCode.CHILD_FAILED) + + +@pytest.mark.parametrize( + ("events", "expected"), + [ + ((), (Running((PYTHON,)), (Start(PYTHON),))), + ((Ready(RUST),), (Running((PYTHON,)), (Start(PYTHON),))), + ((PYTHON_READY,), (Running((PYTHON, RUST)), STARTED)), + ((PYTHON_READY, PYTHON_READY, Ready(RUST)), (Running((PYTHON, RUST)), STARTED)), + ((PYTHON_READY, Stop(TERM)), (Stopping(frozenset({RUST, PYTHON})), BOTH_TERM)), + ((PYTHON_READY, Stop(TERM), Exited(RUST, 0), Exited(PYTHON, 0)), (CLEAN, BOTH_TERM)), + ((PYTHON_READY, Stop(TERM), Exited(PYTHON, 0), Exited(RUST, 0)), (CLEAN, BOTH_TERM)), + ( + (PYTHON_READY, Stop(INT), Stop(INT), Exited(RUST, 0), Exited(PYTHON, 0)), + (CLEAN, (*STARTED, Signal(RUST, INT), Signal(PYTHON, INT))), + ), + ((PYTHON_READY, Stop(TERM), Exited(RUST, 0), Exited(RUST, 0)), (FAILED, BOTH_TERM)), + ((PYTHON_READY, Stop(TERM), Exited(RUST, 1)), (FAILED, BOTH_TERM)), + ((PYTHON_READY, Stop(TERM), Exited(RUST, 0), Exited(PYTHON, 1)), (FAILED, BOTH_TERM)), + ((PYTHON_READY, Exited(PYTHON, 0)), (FAILED, STARTED)), + ((PYTHON_READY, Exited(RUST, 1)), (FAILED, STARTED)), + ((PYTHON_READY, Exited(RUST, 1), Stop(TERM), PYTHON_READY, Exited(PYTHON, 0)), (FAILED, STARTED)), + ((Exited(PYTHON, 1),), (FAILED, (Start(PYTHON),))), + ((Stop(TERM), PYTHON_READY, Exited(PYTHON, 0)), (CLEAN, (Start(PYTHON), Signal(PYTHON, TERM)))), + ((Stop(TERM), Exited(PYTHON, 1)), (FAILED, (Start(PYTHON), Signal(PYTHON, TERM)))), + ], +) +def test_children_start_in_order_once_ready_and_stop_together_in_reverse( + events: tuple[Event, ...], expected: Transition +) -> None: + assert _fold(events) == expected diff --git a/tests/unit/proxy/test_dev_reload.py b/tests/unit/proxy/test_dev_reload.py index 60a7c53820b..6650c1cd3e3 100644 --- a/tests/unit/proxy/test_dev_reload.py +++ b/tests/unit/proxy/test_dev_reload.py @@ -86,6 +86,39 @@ def test_a_broken_factory_exits_the_worker_as_a_startup_failure( assert exit_info.value.code == uvicorn_config.STARTUP_FAILURE +@pytest.mark.parametrize("worker_overrides", [None, {"factory": True}]) +def test_a_factory_flag_in_the_options_reaches_the_workers( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, worker_overrides: dict[str, object] | None +) -> None: + (tmp_path / "reload_factory_flag_app.py").write_text( + "async def served(scope, receive, send):\n" + " await send({'type': 'http.response.start', 'status': 204, 'headers': []})\n" + " await send({'type': 'http.response.body', 'body': b'', 'more_body': False})\n" + "\n" + "def factory():\n" + " return served\n" + ) + monkeypatch.syspath_prepend(str(tmp_path)) + config: Final = pool_config( + RELOAD_OPTIONS.validate_python({**_options("reload_factory_flag_app:factory"), "factory": True}), + workers=POOL_WORKERS, + worker_overrides=worker_overrides, + ) + config.load() + loaded_app: Final = config.loaded_app + assert loaded_app is not None + sent: Final[list[Message]] = [] + + async def receive() -> Message: + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(message: Message) -> None: + sent.append(message) + + asyncio.run(loaded_app({"type": "http", "method": "GET", "path": "/", "headers": []}, receive, send)) + assert sent[0] == {"type": "http.response.start", "status": 204, "headers": []} + + @REQUIRES_UVICORN_STARTUP_FAILURE @pytest.mark.parametrize( ("module", "source"), @@ -118,7 +151,9 @@ def test_the_reload_pool_defaults_the_worker_healthcheck_and_keeps_an_explicit_o explicit: Final = RELOAD_WORKER_HEALTHCHECK_SECONDS + 7 assert pool_config(_options("app:app"), workers=2).timeout_worker_healthcheck == RELOAD_WORKER_HEALTHCHECK_SECONDS assert ( - pool_config({**_options("app:app"), "timeout_worker_healthcheck": explicit}, workers=2).timeout_worker_healthcheck + pool_config( + {**_options("app:app"), "timeout_worker_healthcheck": explicit}, workers=2 + ).timeout_worker_healthcheck == explicit ) @@ -142,7 +177,9 @@ def test_the_reload_pool_accepts_and_keeps_every_uvicorn_arg_the_proxy_cli_build def test_the_reload_pool_passes_a_uvicorn_arg_it_does_not_list_through_to_the_workers() -> None: - config: Final = pool_config(RELOAD_OPTIONS.validate_python({**_options("app:app"), "backlog": 123}), workers=POOL_WORKERS) + config: Final = pool_config( + RELOAD_OPTIONS.validate_python({**_options("app:app"), "backlog": 123}), workers=POOL_WORKERS + ) assert (config.backlog, config.reload) == (123, False) diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index 643c01e3e7f..04ff27badaf 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -39,9 +39,11 @@ from litellm.proxy.shutdown.graceful_shutdown_manager import DEFAULT_GRACEFUL_SH @pytest.fixture(autouse=True) def fork_reservation(): """Reserving is irreversible: it would forbid native routes in this pytest worker for good""" - with patch( # test-quality-ok: process-global native state, a real reservation would poison every later test in the worker - "litellm.rust_bridge.fork_guard.reserve_process_for_forking" - ) as reserve: + with ( + patch( # test-quality-ok: process-global native state, a real reservation would poison every later test in the worker + "litellm.rust_bridge.fork_guard.reserve_process_for_forking" + ) as reserve + ): yield reserve @@ -98,42 +100,30 @@ class TestProxyInitializationHelpers: # Execute with pytest.raises(ValueError, match="Invalid test value"): - ProxyInitializationHelpers._run_test_chat_completion( - "localhost", 8000, "gpt-3.5-turbo", True - ) + ProxyInitializationHelpers._run_test_chat_completion("localhost", 8000, "gpt-3.5-turbo", True) # Test with valid string test value - ProxyInitializationHelpers._run_test_chat_completion( - "localhost", 8000, "gpt-3.5-turbo", "http://test-url" - ) + ProxyInitializationHelpers._run_test_chat_completion("localhost", 8000, "gpt-3.5-turbo", "http://test-url") # Assert - mock_openai.assert_called_once_with( - api_key="My API Key", base_url="http://test-url" - ) + mock_openai.assert_called_once_with(api_key="My API Key", base_url="http://test-url") mock_client.chat.completions.create.assert_called() def test_get_default_unvicorn_init_args(self): # Test without log_config - args = ProxyInitializationHelpers._get_default_unvicorn_init_args( - "localhost", 8000 - ) + args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000) assert args["app"] == "litellm.proxy.proxy_server:app" assert args["host"] == "localhost" assert args["port"] == 8000 assert args["server_header"] is False # Test with log_config - args = ProxyInitializationHelpers._get_default_unvicorn_init_args( - "localhost", 8000, "log_config.json" - ) + args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000, "log_config.json") assert args["log_config"] == "log_config.json" # Test with json_logs=True with patch("litellm.json_logs", True): - args = ProxyInitializationHelpers._get_default_unvicorn_init_args( - "localhost", 8000 - ) + args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000) # When json_logs is True, log_config should be set to the JSON log config dict assert args["log_config"] is not None assert isinstance(args["log_config"], dict) @@ -141,15 +131,11 @@ class TestProxyInitializationHelpers: assert "formatters" in args["log_config"] # Test with keepalive_timeout - args = ProxyInitializationHelpers._get_default_unvicorn_init_args( - "localhost", 8000, None, 60 - ) + args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000, None, 60) assert args["timeout_keep_alive"] == 60 # Test with both log_config and keepalive_timeout - args = ProxyInitializationHelpers._get_default_unvicorn_init_args( - "localhost", 8000, "log_config.json", 120 - ) + args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000, "log_config.json", 120) assert args["log_config"] == "log_config.json" assert args["timeout_keep_alive"] == 120 @@ -280,14 +266,10 @@ class TestProxyInitializationHelpers: py_file = tmp_path / "module.py" py_file.write_text("x = 1\n") - applied = ProxyInitializationHelpers._patch_statreload_extra_paths( - [str(config_file)] - ) + applied = ProxyInitializationHelpers._patch_statreload_extra_paths([str(config_file)]) assert applied is True - fake_self = types.SimpleNamespace( - config=types.SimpleNamespace(reload_dirs=[tmp_path]) - ) + fake_self = types.SimpleNamespace(config=types.SimpleNamespace(reload_dirs=[tmp_path])) yielded_paths = {Path(p).resolve() for p in StatReload.iter_py_files(fake_self)} assert config_file.resolve() in yielded_paths @@ -304,14 +286,10 @@ class TestProxyInitializationHelpers: env_file = tmp_path / ".env" env_file.write_text("FOO=bar\n") - applied = ProxyInitializationHelpers._patch_statreload_extra_paths( - [str(env_file)] - ) + applied = ProxyInitializationHelpers._patch_statreload_extra_paths([str(env_file)]) assert applied is True - fake_self = types.SimpleNamespace( - config=types.SimpleNamespace(reload_dirs=[tmp_path]) - ) + fake_self = types.SimpleNamespace(config=types.SimpleNamespace(reload_dirs=[tmp_path])) yielded_paths = {Path(p).resolve() for p in StatReload.iter_py_files(fake_self)} assert env_file.resolve() in yielded_paths @@ -323,10 +301,7 @@ class TestProxyInitializationHelpers: StatReload._litellm_patched_config_paths.clear() assert ProxyInitializationHelpers._patch_statreload_extra_paths([]) is False - assert ( - ProxyInitializationHelpers._patch_statreload_extra_paths([None, ""]) - is False - ) + assert ProxyInitializationHelpers._patch_statreload_extra_paths([None, ""]) is False def test_patch_statreload_extra_paths_is_idempotent(self, tmp_path): from pathlib import Path @@ -344,18 +319,14 @@ class TestProxyInitializationHelpers: for _ in range(3): ProxyInitializationHelpers._patch_statreload_extra_paths([str(config_file)]) - fake_self = types.SimpleNamespace( - config=types.SimpleNamespace(reload_dirs=[tmp_path]) - ) + fake_self = types.SimpleNamespace(config=types.SimpleNamespace(reload_dirs=[tmp_path])) yielded = list(StatReload.iter_py_files(fake_self)) assert len(yielded) == len(set(map(str, yielded))) yielded_paths = {Path(p).resolve() for p in yielded} assert config_file.resolve() in yielded_paths assert py_file.resolve() in yielded_paths - def test_configure_dev_reload_watches_env_and_sets_override_flag( - self, tmp_path, monkeypatch - ): + def test_configure_dev_reload_watches_env_and_sets_override_flag(self, tmp_path, monkeypatch): from pathlib import Path from uvicorn.supervisors.statreload import StatReload @@ -372,9 +343,7 @@ class TestProxyInitializationHelpers: uvicorn_args: dict = {} with patch("litellm._logging.verbose_proxy_logger.warning") as mock_warning: - ProxyInitializationHelpers._configure_dev_reload( - uvicorn_args, str(config_file) - ) + ProxyInitializationHelpers._configure_dev_reload(uvicorn_args, str(config_file)) assert os.environ["LITELLM_DEV_ENV_HOT_RELOAD"] == "True" assert uvicorn_args["reload"] is True @@ -385,9 +354,7 @@ class TestProxyInitializationHelpers: assert "override" in warning_text assert ".env" in warning_text - fake_self = types.SimpleNamespace( - config=types.SimpleNamespace(reload_dirs=[tmp_path]) - ) + fake_self = types.SimpleNamespace(config=types.SimpleNamespace(reload_dirs=[tmp_path])) yielded_paths = {Path(p).resolve() for p in StatReload.iter_py_files(fake_self)} assert env_file.resolve() in yielded_paths assert config_file.resolve() in yielded_paths @@ -411,17 +378,13 @@ class TestProxyInitializationHelpers: mock_app = MagicMock() # Execute - ProxyInitializationHelpers._init_hypercorn_server( - mock_app, "localhost", 8000, None, None, None - ) + ProxyInitializationHelpers._init_hypercorn_server(mock_app, "localhost", 8000, None, None, None) # Assert mock_asyncio_run.assert_called_once() # Test with SSL - ProxyInitializationHelpers._init_hypercorn_server( - mock_app, "localhost", 8000, "cert.pem", "key.pem", "ECDHE" - ) + ProxyInitializationHelpers._init_hypercorn_server(mock_app, "localhost", 8000, "cert.pem", "key.pem", "ECDHE") @patch("granian.Granian") @patch("builtins.print") @@ -590,10 +553,14 @@ class TestProxyInitializationHelpers: database_name_enc = urllib.parse.quote_plus(database_name) # Construct DATABASE_URL from the provided variables - database_url = f"postgresql://{database_username_enc}:{database_password_enc}@{database_host}/{database_name_enc}" + database_url = ( + f"postgresql://{database_username_enc}:{database_password_enc}@{database_host}/{database_name_enc}" + ) # Assert the correct URL was constructed with properly escaped characters - expected_url = "postgresql://user%40with%2Bspecial:test-password-special-chars@localhost:5432/db_name%2Ftest" + expected_url = ( + "postgresql://user%40with%2Bspecial:test-password-special-chars@localhost:5432/db_name%2Ftest" + ) assert database_url == expected_url # Test appending query parameters @@ -611,9 +578,7 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("atexit.register") # critical @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_skip_server_startup( self, mock_should_update, mock_setup_db, mock_atexit_register, mock_uvicorn_run, tmp_path: Path ): @@ -631,11 +596,7 @@ class TestProxyInitializationHelpers: ) # Remove DATABASE_URL/DIRECT_URL so the CLI doesn't attempt # real prisma operations when these are set in CI. - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} clean_env["PROMETHEUS_MULTIPROC_DIR"] = str(tmp_path) live_proxy_samples = tmp_path / "counter_123.db" live_proxy_samples.write_bytes(b"samples of a proxy that is still running") @@ -655,9 +616,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -668,18 +627,12 @@ class TestProxyInitializationHelpers: # --- skip startup --- result = runner.invoke(run_server, ["--local", "--skip_server_startup"]) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" assert "Skipping server startup" in result.output mock_uvicorn_run.assert_not_called() - result = runner.invoke( - run_server, ["--local", "--skip_server_startup", "--telemetry", "False"] - ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + result = runner.invoke(run_server, ["--local", "--skip_server_startup", "--telemetry", "False"]) + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" assert "Skipping server startup" in result.output assert "telemetry" not in runner.invoke(run_server, ["--help"]).output assert live_proxy_samples.exists() @@ -689,9 +642,7 @@ class TestProxyInitializationHelpers: result = runner.invoke(run_server, ["--local"]) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() assert not live_proxy_samples.exists() @@ -728,9 +679,7 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_limit_concurrency_passed_to_uvicorn( self, mock_should_update, mock_setup_db, mock_atexit_register, mock_uvicorn_run ): @@ -747,11 +696,7 @@ class TestProxyInitializationHelpers: KeyManagementSettings=MagicMock(), save_worker_config=MagicMock(), ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), patch.dict( @@ -761,9 +706,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.side_effect = lambda *a, **k: { "app": "litellm.proxy.proxy_server:app", @@ -771,20 +714,14 @@ class TestProxyInitializationHelpers: "port": 8000, } - result = runner.invoke( - run_server, ["--local", "--limit_concurrency", "250"] - ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + result = runner.invoke(run_server, ["--local", "--limit_concurrency", "250"]) + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() assert mock_uvicorn_run.call_args.kwargs.get("limit_concurrency") == 250 mock_uvicorn_run.reset_mock() result = runner.invoke(run_server, ["--local"]) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() assert "limit_concurrency" not in mock_uvicorn_run.call_args.kwargs @@ -802,7 +739,9 @@ class TestProxyInitializationHelpers: @patch("httpx.HTTPTransport.handle_request") @patch("atexit.register") @patch("subprocess.Popen") - @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch( + "litellm.proxy.db.prisma_client.PrismaManager.setup_database" + ) # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above @patch( # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False ) @@ -840,9 +779,7 @@ class TestProxyInitializationHelpers: save_worker_config=MagicMock(), ) clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL", "PROMETHEUS_METRICS_PORT") + k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL", "PROMETHEUS_METRICS_PORT") } clean_env["PROMETHEUS_MULTIPROC_DIR"] = str(tmp_path) with ( @@ -868,9 +805,7 @@ class TestProxyInitializationHelpers: run_server, ["--local", "--host", "127.0.0.1", "--port", "4000", "--prometheus_metrics_port", "4001"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_popen.assert_called_once() spawned = list(mock_popen.call_args.args[0]) assert spawned[1:3] == ["-m", "litellm.proxy.prometheus_metrics_server"] @@ -893,9 +828,7 @@ class TestProxyInitializationHelpers: mock_popen.reset_mock() mock_uvicorn_run.reset_mock() result = runner.invoke(run_server, ["--local"]) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_popen.assert_not_called() mock_uvicorn_run.assert_called_once() @@ -909,9 +842,7 @@ class TestProxyInitializationHelpers: mock_popen.assert_not_called() mock_uvicorn_run.assert_not_called() - result = runner.invoke( - run_server, ["--local", "--prometheus_metrics_port", "0"] - ) + result = runner.invoke(run_server, ["--local", "--prometheus_metrics_port", "0"]) assert result.exit_code == 2 assert "Invalid value for '--prometheus_metrics_port'" in result.output mock_popen.assert_not_called() @@ -933,9 +864,7 @@ class TestProxyInitializationHelpers: @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_db_timeout_settings_are_forwarded_to_pool_timeout( self, mock_should_update, @@ -968,11 +897,7 @@ class TestProxyInitializationHelpers: } ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), @@ -983,9 +908,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, patch( "litellm.proxy.proxy_cli.append_query_params", side_effect=lambda url, params: ( @@ -1004,9 +927,7 @@ class TestProxyInitializationHelpers: ["--local", "--config", "test-config.yaml", "--skip_server_startup"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_append_query_params.assert_called() appended_params = mock_append_query_params.call_args.args[1] assert appended_params["connection_limit"] == 5 @@ -1061,9 +982,7 @@ class TestProxyInitializationHelpers: @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_db_connection_extra_params_forwarded_to_url( self, mock_should_update, @@ -1098,11 +1017,7 @@ class TestProxyInitializationHelpers: } ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), @@ -1113,9 +1028,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, patch( "litellm.proxy.proxy_cli.append_query_params", side_effect=lambda url, params: str(url), @@ -1132,9 +1045,7 @@ class TestProxyInitializationHelpers: ["--local", "--config", "test-config.yaml", "--skip_server_startup"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_append_query_params.assert_called() appended_params = mock_append_query_params.call_args.args[1] assert appended_params["connect_timeout"] == 15 @@ -1185,9 +1096,7 @@ class TestProxyInitializationHelpers: @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_disable_prepared_statements_forwarded_to_url( self, mock_should_update, @@ -1219,11 +1128,7 @@ class TestProxyInitializationHelpers: } ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), @@ -1234,9 +1139,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, patch( "litellm.proxy.proxy_cli.append_query_params", side_effect=lambda url, params: str(url), @@ -1253,9 +1156,7 @@ class TestProxyInitializationHelpers: ["--local", "--config", "test-config.yaml", "--skip_server_startup"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_append_query_params.assert_called() appended_params = mock_append_query_params.call_args.args[1] if expect_pgbouncer: @@ -1277,9 +1178,7 @@ class TestProxyInitializationHelpers: @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_disable_prepared_statements_env_var_forwarded_to_url( self, mock_should_update, @@ -1327,9 +1226,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, patch( "litellm.proxy.proxy_cli.append_query_params", side_effect=lambda url, params: str(url), @@ -1346,9 +1243,7 @@ class TestProxyInitializationHelpers: ["--local", "--config", "test-config.yaml", "--skip_server_startup"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" appended_params = mock_append_query_params.call_args.args[1] if expect_pgbouncer: assert appended_params["pgbouncer"] == "true", appended_params @@ -1358,9 +1253,7 @@ class TestProxyInitializationHelpers: @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_malformed_disable_prepared_statements_env_var_is_rejected_even_when_config_enables_it( self, mock_should_update, @@ -1406,9 +1299,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, patch( "litellm.proxy.proxy_cli.append_query_params", side_effect=lambda url, params: str(url), @@ -1432,9 +1323,7 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_proxy_default_api_version_uses_azure_default( self, mock_should_update, mock_setup_db, mock_atexit_register, mock_uvicorn_run ): @@ -1451,11 +1340,7 @@ class TestProxyInitializationHelpers: KeyManagementSettings=MagicMock(), save_worker_config=MagicMock(), ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), patch.dict( @@ -1465,9 +1350,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -1475,9 +1358,7 @@ class TestProxyInitializationHelpers: "port": 8000, } result = runner.invoke(run_server, ["--local", "--skip_server_startup"]) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_proxy_module.save_worker_config.assert_called_once() call_kwargs = mock_proxy_module.save_worker_config.call_args[1] assert call_kwargs["api_version"] == litellm.AZURE_DEFAULT_API_VERSION @@ -1485,12 +1366,8 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("builtins.print") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) - def test_keepalive_timeout_flag( - self, mock_should_update, mock_setup_db, mock_print, mock_uvicorn_run - ): + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) + def test_keepalive_timeout_flag(self, mock_should_update, mock_setup_db, mock_print, mock_uvicorn_run): """Test that the keepalive_timeout flag is properly passed to uvicorn""" from click.testing import CliRunner @@ -1507,11 +1384,7 @@ class TestProxyInitializationHelpers: # DB-setup block (un-timeout'd `subprocess.run(["prisma"])` + # migrate-deploy retry loop) — same isolation every other run_server # test in this file uses. - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), @@ -1526,9 +1399,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, patch( "litellm.proxy.proxy_cli.ProxyInitializationHelpers.is_port_in_use", return_value=False, @@ -1560,12 +1431,8 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("builtins.print") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) - def test_timeout_worker_healthcheck_flag( - self, mock_should_update, mock_setup_db, mock_print, mock_uvicorn_run - ): + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) + def test_timeout_worker_healthcheck_flag(self, mock_should_update, mock_setup_db, mock_print, mock_uvicorn_run): """Test that the --timeout_worker_healthcheck flag is threaded through to the uvicorn init helper.""" from click.testing import CliRunner @@ -1582,11 +1449,7 @@ class TestProxyInitializationHelpers: # DB-setup block (un-timeout'd `subprocess.run(["prisma"])` + # migrate-deploy retry loop) — same isolation every other run_server # test in this file uses. - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), @@ -1601,9 +1464,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, patch( "litellm.proxy.proxy_cli.ProxyInitializationHelpers.is_port_in_use", return_value=False, @@ -1615,9 +1476,7 @@ class TestProxyInitializationHelpers: "port": 8000, } - result = runner.invoke( - run_server, ["--local", "--timeout_worker_healthcheck", "15"] - ) + result = runner.invoke(run_server, ["--local", "--timeout_worker_healthcheck", "15"]) assert result.exit_code == 0 mock_get_args.assert_called_once_with( @@ -1631,9 +1490,7 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("builtins.print") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - def test_max_requests_before_restart_flag( - self, mock_setup_db, mock_print, mock_uvicorn_run - ): + def test_max_requests_before_restart_flag(self, mock_setup_db, mock_print, mock_uvicorn_run): """Test that the max_requests_before_restart flag is passed to uvicorn as limit_max_requests""" from click.testing import CliRunner @@ -1646,11 +1503,7 @@ class TestProxyInitializationHelpers: mock_key_mgmt = MagicMock() mock_save_worker_config = MagicMock() - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict( os.environ, @@ -1668,9 +1521,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -1678,13 +1529,9 @@ class TestProxyInitializationHelpers: "port": 8000, } - result = runner.invoke( - run_server, ["--local", "--max_requests_before_restart", "123"] - ) + result = runner.invoke(run_server, ["--local", "--max_requests_before_restart", "123"]) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() # Check that uvicorn.run was called with limit_max_requests parameter @@ -1694,9 +1541,7 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("builtins.print") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - def test_max_requests_before_restart_jitter_flag( - self, mock_setup_db, mock_print, mock_uvicorn_run - ): + def test_max_requests_before_restart_jitter_flag(self, mock_setup_db, mock_print, mock_uvicorn_run): """--max_requests_before_restart_jitter maps to uvicorn limit_max_requests_jitter""" from click.testing import CliRunner @@ -1707,11 +1552,7 @@ class TestProxyInitializationHelpers: pass runner = CliRunner() - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), patch("uvicorn.Config", _NewUvicornConfig), @@ -1726,9 +1567,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -1747,9 +1586,7 @@ class TestProxyInitializationHelpers: ], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() call_args = mock_uvicorn_run.call_args assert call_args[1]["limit_max_requests"] == 1000 @@ -1768,11 +1605,7 @@ class TestProxyInitializationHelpers: from litellm.proxy.proxy_cli import run_server runner = CliRunner() - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), patch.dict( @@ -1786,9 +1619,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -1808,9 +1639,7 @@ class TestProxyInitializationHelpers: ], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_not_called() mock_run_gunicorn.assert_called_once() g_kwargs = mock_run_gunicorn.call_args[1] @@ -1856,20 +1685,14 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - def test_single_worker_max_requests_warning_printed_on_uvicorn( - self, mock_setup_db, mock_uvicorn_run - ): + def test_single_worker_max_requests_warning_printed_on_uvicorn(self, mock_setup_db, mock_uvicorn_run): """uvicorn single-worker startup prints the does-not-recycle warning""" from click.testing import CliRunner from litellm.proxy.proxy_cli import run_server runner = CliRunner() - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), patch.dict( @@ -1883,9 +1706,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -1893,13 +1714,9 @@ class TestProxyInitializationHelpers: "port": 8000, } - result = runner.invoke( - run_server, ["--local", "--max_requests_before_restart", "123"] - ) + result = runner.invoke(run_server, ["--local", "--max_requests_before_restart", "123"]) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" assert "--max_requests_before_restart 123 with --num_workers 1" in result.output assert "uvicorn exits the process" in result.output @@ -1916,11 +1733,7 @@ class TestProxyInitializationHelpers: from litellm.proxy.proxy_cli import run_server runner = CliRunner() - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), patch.dict( @@ -1934,9 +1747,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -1949,12 +1760,8 @@ class TestProxyInitializationHelpers: ["--local", "--num_workers", "2", "--max_requests_before_restart", "123"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" - assert not any( - "does not recycle a worker" in str(c) for c in mock_print.call_args_list - ) + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" + assert not any("does not recycle a worker" in str(c) for c in mock_print.call_args_list) @patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.run_gunicorn_server") @patch("uvicorn.run") @@ -1968,11 +1775,7 @@ class TestProxyInitializationHelpers: from litellm.proxy.proxy_cli import run_server runner = CliRunner() - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), patch.dict( @@ -1986,9 +1789,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -2001,9 +1802,7 @@ class TestProxyInitializationHelpers: ["--local", "--run_gunicorn", "--max_requests_before_restart", "123"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" assert "--max_requests_before_restart 123 with --num_workers 1" in result.output assert "gunicorn accepts no connections" in result.output @@ -2020,11 +1819,7 @@ class TestProxyInitializationHelpers: from litellm.proxy.proxy_cli import run_server runner = CliRunner() - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), patch.dict( @@ -2038,9 +1833,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -2053,13 +1846,9 @@ class TestProxyInitializationHelpers: ["--local", "--run_hypercorn", "--max_requests_before_restart", "5"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_init_hypercorn.assert_called_once() - assert not any( - "does not recycle a worker" in str(c) for c in mock_print.call_args_list - ) + assert not any("does not recycle a worker" in str(c) for c in mock_print.call_args_list) @pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows") def test_gunicorn_options_include_max_requests_jitter(self): @@ -2260,7 +2049,9 @@ class TestProxyInitializationHelpers: with patch.dict(os.environ, test_env_special): result = construct_database_url_from_env_vars() - expected_url = "postgresql://user%40with%2Bspecial:test-password-special-chars@localhost:5432/db_name%2Ftest" + expected_url = ( + "postgresql://user%40with%2Bspecial:test-password-special-chars@localhost:5432/db_name%2Ftest" + ) assert result == expected_url # Test without password (should still work) @@ -2388,11 +2179,7 @@ class TestQueryEngineReaperWiring: from litellm.proxy.proxy_cli import run_server runner = CliRunner() - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), patch.dict( @@ -2407,12 +2194,8 @@ class TestQueryEngineReaperWiring: }, ), patch("uvicorn.run") as mock_uvicorn_run, - patch( - "litellm.proxy.proxy_cli.start_query_engine_reaper" - ) as mock_start_reaper, - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.start_query_engine_reaper") as mock_start_reaper, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -2423,17 +2206,13 @@ class TestQueryEngineReaperWiring: return result, mock_uvicorn_run, mock_start_reaper def test_multi_worker_uvicorn_starts_reaper(self): - result, mock_uvicorn_run, mock_start_reaper = self._invoke_run_server( - ["--local", "--num_workers", "2"] - ) + result, mock_uvicorn_run, mock_start_reaper = self._invoke_run_server(["--local", "--num_workers", "2"]) assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() mock_start_reaper.assert_called_once() def test_single_worker_uvicorn_does_not_start_reaper(self): - result, mock_uvicorn_run, mock_start_reaper = self._invoke_run_server( - ["--local", "--num_workers", "1"] - ) + result, mock_uvicorn_run, mock_start_reaper = self._invoke_run_server(["--local", "--num_workers", "1"]) assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() mock_start_reaper.assert_not_called() @@ -2444,9 +2223,7 @@ class TestQueryEngineReaperWiring: with ( patch("gunicorn.app.base.BaseApplication.run"), - patch( - "litellm.proxy.proxy_cli.start_query_engine_reaper" - ) as mock_start_reaper, + patch("litellm.proxy.proxy_cli.start_query_engine_reaper") as mock_start_reaper, ): ProxyInitializationHelpers._run_gunicorn_server( host="127.0.0.1", @@ -2494,11 +2271,7 @@ class TestRunServerDbSetup: save_worker_config=MagicMock(), ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" with ( @@ -2510,9 +2283,7 @@ class TestRunServerDbSetup: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -2543,10 +2314,18 @@ class TestRunServerDbSetup: mock_setup_database.assert_called_with(use_migrate=False, use_v2_resolver=True) @patch("atexit.register") - @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above - @patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above - @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above - @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch( + "litellm.proxy.db.prisma_client.PrismaManager.setup_database" + ) # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch( + "litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes" + ) # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch( + "litellm.proxy.db.check_migration.check_prisma_schema_diff" + ) # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch( + "litellm.proxy.db.prisma_client.should_update_prisma_schema" + ) # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above def test_migrations_run_when_the_prisma_cli_is_not_on_path( self, mock_should_update_schema, @@ -2570,11 +2349,7 @@ class TestRunServerDbSetup: save_worker_config=MagicMock(), ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" clean_env["PATH"] = str(empty_bin) @@ -2734,11 +2509,7 @@ class TestRunServerDbSetup: save_worker_config=MagicMock(), ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" with ( @@ -2750,9 +2521,7 @@ class TestRunServerDbSetup: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -2795,11 +2564,7 @@ class TestRunServerDbSetup: save_worker_config=MagicMock(), ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} clean_env["DATABASE_URL"] = "sqlite:///data/litellm.db" with ( @@ -2813,9 +2578,7 @@ class TestRunServerDbSetup: ), ): with pytest.raises(SystemExit) as exc_info: - run_server.main( - ["--local", "--skip_server_startup"], standalone_mode=False - ) + run_server.main(["--local", "--skip_server_startup"], standalone_mode=False) assert exc_info.value.code == 1 mock_setup_database.assert_not_called() err: Final = capsys.readouterr().err @@ -2857,11 +2620,7 @@ class TestRunServerDbSetup: save_worker_config=MagicMock(), ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" clean_env["USE_V2_MIGRATION_RESOLVER"] = "true" @@ -2875,9 +2634,7 @@ class TestRunServerDbSetup: }, ), ): - run_server.main( - ["--local", "--skip_server_startup"], standalone_mode=False - ) + run_server.main(["--local", "--skip_server_startup"], standalone_mode=False) mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True) assert "--use_v2_migration_resolver is deprecated" not in capsys.readouterr().out @@ -2920,9 +2677,7 @@ class TestRunServerDbSetup: {"proxy_server": mock_proxy_module, "litellm.proxy.proxy_server": mock_proxy_module}, ), ): - run_server.main( - ["--local", "--skip_server_startup", "--use_v2_migration_resolver"], standalone_mode=False - ) + run_server.main(["--local", "--skip_server_startup", "--use_v2_migration_resolver"], standalone_mode=False) mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True) assert "--use_v2_migration_resolver is deprecated and has no effect" in capsys.readouterr().out @@ -3059,12 +2814,7 @@ class TestRunServerDbSetup: def test_resolve_v2_migration_resolver(self, use_legacy_flag, env_value, expected): from litellm.proxy.proxy_cli import resolve_v2_migration_resolver - assert ( - resolve_v2_migration_resolver( - use_legacy_flag=use_legacy_flag, env_value=env_value - ) - is expected - ) + assert resolve_v2_migration_resolver(use_legacy_flag=use_legacy_flag, env_value=env_value) is expected def test_deprecated_v2_flag_not_reported_outside_a_cli_invocation(self): from litellm.proxy.proxy_cli import deprecated_v2_flag_passed_on_cli @@ -3105,9 +2855,7 @@ class TestRunServerDbSetup: ) clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL", "USE_V2_MIGRATION_RESOLVER") + k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL", "USE_V2_MIGRATION_RESOLVER") } clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" @@ -3188,9 +2936,7 @@ class TestRunServerDbSetup: "sys.modules", {"proxy_server": mock_proxy_module, "litellm.proxy.proxy_server": mock_proxy_module}, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, outcome as exc_info, ): mock_get_args.return_value = {"app": "litellm.proxy.proxy_server:app", "host": "localhost", "port": 8000} @@ -3242,11 +2988,7 @@ class TestWorkerStartupHooks: "LITELLM_WORKER_STARTUP_HOOKS": "tests.unit.proxy.test_proxy_cli:_dummy_hook", } # Remove DATABASE_URL to avoid real DB setup - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} clean_env.update(env_overrides) with patch.dict(os.environ, clean_env, clear=True): @@ -3269,11 +3011,7 @@ class TestWorkerStartupHooks: env_overrides = { "LITELLM_WORKER_STARTUP_HOOKS": "tests.unit.proxy.test_proxy_cli:_dummy_async_hook", } - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} clean_env.update(env_overrides) with patch.dict(os.environ, clean_env, clear=True): @@ -3293,11 +3031,7 @@ class TestWorkerStartupHooks: env_overrides = { "LITELLM_WORKER_STARTUP_HOOKS": "tests.unit.proxy.test_proxy_cli:_failing_hook", } - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} clean_env.update(env_overrides) with patch.dict(os.environ, clean_env, clear=True): @@ -3326,18 +3060,11 @@ class TestWorkerStartupHooks: from litellm.proxy.proxy_server import proxy_startup_event - hooks = ( - "tests.unit.proxy.test_proxy_cli:_dummy_hook," - "tests.unit.proxy.test_proxy_cli:_dummy_async_hook" - ) + hooks = "tests.unit.proxy.test_proxy_cli:_dummy_hook,tests.unit.proxy.test_proxy_cli:_dummy_async_hook" env_overrides = { "LITELLM_WORKER_STARTUP_HOOKS": hooks, } - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} clean_env.update(env_overrides) with patch.dict(os.environ, clean_env, clear=True): @@ -3483,13 +3210,9 @@ class TestPostgresStatementTimeoutOptions: import yaml config_path = tmp_path / "config.yaml" - config_path.write_text( - yaml.dump({"model_list": [], "general_settings": {"database_statement_timeout": 60}}) - ) + config_path.write_text(yaml.dump({"model_list": [], "general_settings": {"database_statement_timeout": 60}})) - captured = _run_server_and_capture_urls( - str(config_path), direct_url="postgresql://t:t@localhost:5432/t" - ) + captured = _run_server_and_capture_urls(str(config_path), direct_url="postgresql://t:t@localhost:5432/t") assert "options" in urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL"]).query) assert "options" not in urlparse.parse_qs(urlparse.urlparse(captured["DIRECT_URL"]).query) @@ -3661,9 +3384,7 @@ class TestReadReplicaConnectionParams: import yaml config_path: Final = tmp_path / "config.yaml" - config_path.write_text( - yaml.dump({"model_list": [], "general_settings": {"database_connection_pool_limit": 3}}) - ) + config_path.write_text(yaml.dump({"model_list": [], "general_settings": {"database_connection_pool_limit": 3}})) _run_server_and_capture_urls( str(config_path), @@ -3683,9 +3404,7 @@ class TestReadReplicaConnectionParams: import yaml config_path: Final = tmp_path / "config.yaml" - config_path.write_text( - yaml.dump({"model_list": [], "general_settings": {"database_connection_pool_limit": 3}}) - ) + config_path.write_text(yaml.dump({"model_list": [], "general_settings": {"database_connection_pool_limit": 3}})) _run_server_and_capture_urls( str(config_path), @@ -3880,9 +3599,7 @@ class TestTokenAuthCliFlags: patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False), patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database"), patch("uvicorn.run"), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -3896,9 +3613,7 @@ class TestTokenAuthCliFlags: return result, database_url, toggle def test_azure_flag_assembles_a_token_bearing_database_url(self): - result, database_url, toggle = self._invoke_with_azure_host( - ["--local", "--azure_postgresql_auth"] - ) + result, database_url, toggle = self._invoke_with_azure_host(["--local", "--azure_postgresql_auth"]) assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" assert database_url is not None @@ -4126,7 +3841,11 @@ def probe_logger() -> Iterator[logging.Logger]: def _log_config_source(kind: str, tmp_path: Path) -> str | dict[str, object]: - config: Final[dict[str, object]] = {"version": 1, "incremental": True, "loggers": {_PROBE_LOGGER: {"level": "ERROR"}}} + config: Final[dict[str, object]] = { + "version": 1, + "incremental": True, + "loggers": {_PROBE_LOGGER: {"level": "ERROR"}}, + } if kind == "dict": return config written: Final = tmp_path / f"log_config.{kind}" @@ -4455,3 +4174,98 @@ def test_migration_data_check_requires_database_url(monkeypatch): with pytest.raises(click.UsageError, match="requires DATABASE_URL"): ProxyInitializationHelpers.run_migration_data_check(use_migrate=True) + + +class TestRustForwardServer: + """The public listener ``run_server`` reports to the Rust supervisor matches what it would serve itself.""" + + @staticmethod + def _server(argv: list[str]): + params: Final = run_server.make_context("litellm", ["--port", "45999", *argv]).params + return params, ProxyInitializationHelpers.python_server( + str(params["host"]), + int(str(params["port"])), + params["keepalive_timeout"], + run_gunicorn=params["run_gunicorn"] is True, + run_hypercorn=params["run_hypercorn"] is True, + run_granian=params["run_granian"] is True, + tls=any((params["ssl_keyfile_path"], params["ssl_certfile_path"], params["ciphers"])), + ) + + @pytest.mark.parametrize( + ("argv", "reason"), + [ + (["--ssl_keyfile_path", "unused"], "does not support TLS"), + (["--ssl_certfile_path", "unused"], "does not support TLS"), + (["--ciphers", "unused"], "does not support TLS"), + (["--run_hypercorn"], "does not support Hypercorn"), + (["--run_granian"], "does not support Granian"), + ], + ) + def test_unsupported_servers_fail_before_starting_children( + self, argv: list[str], reason: str, monkeypatch: pytest.MonkeyPatch + ) -> None: + from click.testing import CliRunner + + monkeypatch.setenv("LITELLM_RUST", "true") + monkeypatch.delenv("_LITELLM_GATEWAY_FD", raising=False) + result: Final = CliRunner().invoke(run_server, argv) + assert result.exit_code == 1 + assert f"LITELLM_RUST {reason}" in result.output + + @pytest.mark.parametrize( + ("argv", "environ"), + [(["--keepalive_timeout", "2"], {}), ([], {"KEEPALIVE_TIMEOUT": "7"}), ([], {})], + ) + def test_rust_retires_connections_by_the_keepalive_uvicorn_runs_with( + self, argv: list[str], environ: dict[str, str], monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv("KEEPALIVE_TIMEOUT", raising=False) + for name, value in environ.items(): + monkeypatch.setenv(name, value) + params, server = self._server(argv) + served: Final = uvicorn.Config( + **ProxyInitializationHelpers.get_default_unvicorn_init_args( + "127.0.0.1", 45555, keepalive_timeout=params["keepalive_timeout"] + ) + ) + assert server.keepalive_seconds == served.timeout_keep_alive + + def test_gunicorn_workers_keep_gunicorns_keepalive_and_rust_follows(self, monkeypatch: pytest.MonkeyPatch) -> None: + from gunicorn.config import Config as GunicornConfig + + monkeypatch.delenv("KEEPALIVE_TIMEOUT", raising=False) + _, server = self._server(["--run_gunicorn", "--keepalive_timeout", "9"]) + assert server.keepalive_seconds == GunicornConfig().keepalive + + @pytest.mark.parametrize("gunicorn", [True, False]) + def test_python_behind_rust_trusts_the_proxies_its_server_would_have( + self, gunicorn: bool, monkeypatch: pytest.MonkeyPatch + ) -> None: + from gunicorn.config import Config as GunicornConfig + + from litellm.proxy.rust_forward.mode import InheritedFd, Supervised, rust_forward_mode + from litellm.proxy.rust_forward.platform import PrivateListener + from litellm.proxy.rust_forward.platform.posix import PosixPlatform + from litellm.proxy.rust_forward.supervisor import python_spec + + monkeypatch.delenv("FORWARDED_ALLOW_IPS", raising=False) + _, server = self._server(["--run_gunicorn"] if gunicorn else []) + mode: Final = rust_forward_mode( + python_spec( + PosixPlatform(), ("litellm",), PrivateListener(InheritedFd(17), 45555, (17,)), server + ).env_overrides + ) + trusted: Final = ( + GunicornConfig().forwarded_allow_ips + if gunicorn + else [uvicorn.Config(app="", log_config=None).forwarded_allow_ips] + ) + assert isinstance(mode, Supervised) + assert mode.trusted_proxies.split(",") == trusted + + def test_rust_drains_for_the_graceful_timeout_uvicorn_runs_with(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "12.5") + _, server = self._server([]) + served: Final = ProxyInitializationHelpers.get_default_unvicorn_init_args("127.0.0.1", 45555) + assert server.graceful_timeout_seconds == served["timeout_graceful_shutdown"]