scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

Merge 24db29cd73 into edd3542297

Browse source
This commit is contained in:
yujonglee 2026-11-03 19:02:44 +00:00 • committed by GitHub
commit 938ff6c90f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
25 changed files with 2485 additions and 466 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.github/workflows/test-unit.yml vendored
View file

@ -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

71
gateway/launch.py
View file

@ -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__":

7
litellm/proxy/dev_reload.py
View file

@ -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,

107
litellm/proxy/proxy_cli.py
View file

@ -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(

18
litellm/proxy/rust_forward/AGENTS.md Normal file
View file

@ -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

0
litellm/proxy/rust_forward/__init__.py Normal file
View file

177
litellm/proxy/rust_forward/mode.py Normal file
View file

@ -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)

57
litellm/proxy/rust_forward/peer.py Normal file
View file

@ -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)

52
litellm/proxy/rust_forward/platform/__init__.py Normal file
View file

@ -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()

45
litellm/proxy/rust_forward/platform/posix.py Normal file
View file

@ -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)

56
litellm/proxy/rust_forward/platform/windows.py Normal file
View file

@ -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)

30
litellm/proxy/rust_forward/server.py Normal file
View file

@ -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

332
litellm/proxy/rust_forward/supervisor.py Normal file
View file

@ -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

10
litellm/proxy/shutdown/graceful_uvicorn_worker.py
View file

@ -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

4
litellm/rust_bridge/configuration.py
View file

@ -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)),
)

8
litellm_cli/main.py
View file

@ -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))

180
tests/integration/configuration/test_rust_forward.py Normal file
View file

@ -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()

457
tests/integration/configuration/test_rust_forward_traffic.py Normal file
View file

@ -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

110
tests/unit/gateway/test_launch.py
View file

@ -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
):

0
tests/unit/proxy/rust_forward/__init__.py Normal file
View file

112
tests/unit/proxy/rust_forward/test_mode.py Normal file
View file

@ -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

143
tests/unit/proxy/rust_forward/test_peer.py Normal file
View file

@ -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}

241
tests/unit/proxy/rust_forward/test_supervisor.py Normal file
View file

@ -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

41
tests/unit/proxy/test_dev_reload.py
View file

@ -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)

692
tests/unit/proxy/test_proxy_cli.py
View file

File diff suppressed because it is too large Load diff