feat(proxy): supervise the opt-in Rust forwarding server (#49606)

* feat(proxy): supervise the opt-in Rust forwarding server

With LITELLM_RUST on, `litellm` and `python -m gateway.launch` start a
supervisor that runs the Rust gateway on the public port and Python on a
private loopback listener. Python still serves every route and sees each
request as if the client had connected to it directly.

Rust sends the client's TCP peer in `x-litellm-rust-peer`. The supervised
server turns server-level proxy headers off, and a middleware restores
the peer and then applies the operator's own trust, so client IPs,
trusted proxies and every IP-based setting behave the same with Rust on
or off.

The supervisor is transparent to Python. Signals pass through, Python's
exit status is the result, and both children shut down when the
supervisor's end of their stdin closes, so nothing outlives it. Rust
starts only once Python listens. When forwarding cannot start (no
gateway binary, TLS, a Unix socket, Hypercorn or Granian), the entry
point warns and serves from Python alone.

Every forwarding decision lives in litellm/proxy/rust_forward, paired
with the gateway-forward crate. OS differences sit behind a Platform
protocol, with Windows supported for local CLI runs. AGENTS.md files in
the touched folders state the boundary. The `--rust_gateway` flag is
gone in favor of the environment variable.

* refactor(proxy): switch Rust forwarding on with LITELLM_RUST_FIRST

Forwarding gets its own switch, independent of LITELLM_RUST, which keeps
selecting only the Rust paths inside the SDK. Neither turns the other on.

---------

Co-authored-by: narmstrong-berri <narmstrong@berri.ai>
This commit is contained in:
yujonglee 2026-11-12 04:40:00 +09:00 • committed by GitHub
parent b079bf5a51
commit 95ecfbbac6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
47 changed files with 3344 additions and 26 deletions

View file

@ -591,6 +591,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

16
gateway/AGENTS.md Normal file
View file

@ -0,0 +1,16 @@
# Gateway image entry point
`python -m gateway.launch` is how the gateway image starts. `launch.py` hands its arguments to uvicorn's own CLI and serves `main.py`, which trims the proxy app to the LLM data plane
- Importing `launch.py` stays cheap. Every uvicorn worker imports it again before it can answer its first health ping, so LiteLLM imports wait until `main` or `worker_app` runs
- PgBouncer starts exactly once per pod, before the workers start, and the workers inherit the loopback URL through `DATABASE_URL`. `configure_database` is the only place that does this
## Rust forwarding
`launch.py` is an entry point of Rust forwarding, and every forwarding decision lives in `litellm/proxy/rust_forward/`, whose `AGENTS.md` states the contract
- `main` dispatches on `rust_forward_mode` and does nothing else about forwarding. With forwarding off or refused it serves exactly as before
- `python_server` reads the flags and environment through uvicorn's own CLI parser, so Rust gets the values uvicorn would run with. Never reparse a uvicorn flag by hand
- Under forwarding the supervisor process runs `configure_database`, as `before_start`, once forwarding is known to be possible. The supervised child never runs it and inherits the exported URL
- The supervised child appends `mode.uvicorn_flags` after the operator's flags, so the private listener and `--no-proxy-headers` win
- `worker_app` wraps the app with `peer.wrap`, which is the identity when forwarding is off

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,16 +66,39 @@ 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 python_server(argv: Sequence[str]) -> "PythonServer":
"""What uvicorn would serve the public port with, read from the same flags and environment it reads."""
from litellm.proxy.rust_forward.server import PythonServer, Server, UvicornCliParams
from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
params: Final = UvicornCliParams.model_validate(
uvicorn_main.make_context("uvicorn", list(uvicorn_argv(argv, os.environ))).params
)
return PythonServer(
server=Server.UVICORN,
listener=params.listener,
host=params.host,
port=params.port,
keepalive_seconds=params.timeout_keep_alive,
graceful_timeout_seconds=GracefulShutdownManager.get_timeout()
if params.timeout_graceful_shutdown is None
else params.timeout_graceful_shutdown,
trusted_proxies=params.trusted_proxies,
)
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
@ -88,10 +112,57 @@ def configure_database() -> None:
export_pooled_database_url(pooled_url)
def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None:
def _serve_python_alone(argv: Sequence[str], serve: Callable[[Sequence[str]], None]) -> None:
configure_database()
serve(uvicorn_argv(argv, os.environ))
def _python_server_or_exit(argv: Sequence[str]) -> "PythonServer":
"""Report ``--help``, ``--version`` and a bad flag the way uvicorn's own CLI does, since it never gets to."""
import click
try:
return python_server(argv)
except click.exceptions.Exit as stopped:
sys.exit(stopped.exit_code)
except click.ClickException as error:
error.show()
sys.exit(error.exit_code)
def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None:
from pydantic import ValidationError
from litellm.proxy.rust_forward.mode import (
Mode,
Supervised,
environment_error,
rust_forward_mode,
stop_with_supervisor,
uvicorn_flags,
)
from litellm.proxy.rust_forward.supervisor import supervise
from litellm.proxy.rust_forward.supervisor.preflight import StartupError
try:
mode: Final = rust_forward_mode(os.environ)
except ValidationError as error:
sys.exit(environment_error(error))
match mode:
case Mode.DISABLED:
_serve_python_alone(argv, serve)
case Mode.SUPERVISOR:
supervised: Final = supervise(_python_server_or_exit(argv), before_start=configure_database)
if isinstance(supervised, StartupError):
_serve_python_alone(argv, serve)
else:
sys.exit(supervised)
case Supervised():
stop_with_supervisor(mode)
serve((*uvicorn_argv(argv, os.environ), *uvicorn_flags(mode)))
case _:
assert_never(mode)
if __name__ == "__main__":
main(sys.argv[1:])

View file

@ -13,6 +13,7 @@
- Pooled upstream connections are retired well before Python's keep-alive timeout, so a request never lands on a connection Python is closing
- Idle client connections close after Python's keep-alive timeout, as uvicorn closed them, so a client that vanishes cannot hold a socket forever. Only the wait for the next request is timed, never an exchange in progress
- WebSocket is the one protocol switch Python serves. Upgrades are relayed as raw bytes after the 101 and drained on shutdown
- Drain on EOF of stdin exactly as on SIGTERM. The supervisor holds the only write end, so EOF means it is gone and nobody is left to stop the gateway. `input_closed` reads on a plain thread, since a read blocked on tokio's blocking pool would hold up the runtime's exit
## Layout
@ -22,6 +23,8 @@ Every forwarding decision has one home. `rewrite.rs` owns what Python and the cl
`Settings` reads only these variables, and only when `LITELLM_FORWARD_UPSTREAM` is set. The Python supervisor in `litellm/proxy/rust_forward/` sets every one of them from the values Python actually runs with, so all of them are required and a missing or invalid one fails startup naming it. The crate never reads Python's own settings (`HOST`, `PORT`, `KEEPALIVE_TIMEOUT`, `GRACEFUL_SHUTDOWN_TIMEOUT`) and keeps no copy of Python's defaults, since Python can get those values from CLI flags the environment never shows
Only the Python supervisor starts the gateway in forwarding mode, from `litellm` and from `python -m gateway.launch`, on a laptop and in every image. `supervisor/children.py` is the only writer of these variables, so adding, renaming or removing one changes `settings.rs` and `children.rust_spec` in the same change. Without `LITELLM_FORWARD_UPSTREAM` the binary runs as the standalone Rust gateway and this crate does nothing
| Variable | Meaning |
|---|---|
| `LITELLM_FORWARD_UPSTREAM` | Python's loopback origin, e.g. `http://127.0.0.1:53211` |

View file

@ -13,7 +13,7 @@ hyper = { version = "1.12.0", features = ["http1", "server"] }
hyper-util = { workspace = true, features = ["client-legacy", "http1", "tokio"] }
serde.workspace = true
thiserror.workspace = true
tokio = { workspace = true, features = ["io-util"] }
tokio = { workspace = true, features = ["io-util", "sync"] }
tokio-util = { version = "0.7", features = ["rt"] }
tower-service = "0.3.3"
tracing.workspace = true

View file

@ -0,0 +1,29 @@
use std::io::{self, Read};
use tokio::sync::oneshot;
/// Resolves once `input` reaches end of file or fails to read. The Python supervisor holds the only
/// write end of the gateway's stdin, so EOF means it is gone and the gateway drains as on SIGTERM.
///
/// The read runs on a plain thread because a blocked read on tokio's blocking pool would hold up
/// runtime shutdown after a signal had already stopped the gateway.
pub fn input_closed(mut input: impl Read + Send + 'static) -> impl Future<Output = ()> + Send {
let (closed, on_closed) = oneshot::channel::<()>();
let reader = std::thread::Builder::new()
.name("stdin-eof".into())
.spawn(move || {
let _ = io::copy(&mut input, &mut io::sink());
let _ = closed.send(());
});
async move {
match reader {
Ok(_) => {
let _ = on_closed.await;
}
Err(error) => {
tracing::warn!(%error, "could not watch stdin, so only signals stop the gateway");
std::future::pending::<()>().await;
}
}
}
}

View file

@ -2,6 +2,7 @@ mod connections;
mod connector;
mod error;
mod forward;
mod input;
mod rewrite;
mod server;
mod settings;
@ -10,6 +11,7 @@ mod websocket;
pub use error::Error;
pub(crate) use error::Rejection;
pub use forward::Forward;
pub use input::input_closed;
pub use rewrite::PEER_HEADER;
pub use settings::{Settings, Upstream};

View file

@ -7,9 +7,11 @@ use std::{
use axum::{Router, body::Body, extract::Request, http::StatusCode, response::Response};
use http_body_util::BodyExt;
use litellm_gateway_forward::input_closed;
use rstest::rstest;
use support::{
DRAIN, Gateway, Raw, client, raw_upstream, recording, serve, unreachable_upstream, within,
DRAIN, Gateway, Raw, UVICORN_KEEPALIVE, client, raw_upstream, recording, serve, serve_until,
unreachable_upstream, within,
};
use tokio::{
net::{TcpListener, TcpStream},
@ -287,6 +289,55 @@ async fn shutdown_lets_an_in_flight_stream_finish() {
within(server).await.unwrap();
}
#[rstest]
#[tokio::test]
async fn the_supervisor_closing_stdin_drains_like_a_stop_signal() {
let release = Arc::new(Notify::new());
let released = release.clone();
let upstream = raw_upstream(move |mut python: Raw| {
let released = released.clone();
async move {
python.request_head().await.unwrap();
python
.send("HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n5\r\nfirst\r\n")
.await;
released.notified().await;
python.send("4\r\nlast\r\n0\r\n\r\n").await;
}
})
.await;
let (stdin, supervisor) = std::io::pipe().unwrap();
let (address, server) = serve_until(upstream, UVICORN_KEEPALIVE, input_closed(stdin)).await;
let response = within(
client().request(
Request::get(format!("http://{address}/"))
.body(Body::empty())
.unwrap(),
),
)
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(100)).await;
assert!(TcpStream::connect(address).await.is_ok());
drop(supervisor);
within(async {
while TcpStream::connect(address).await.is_ok() {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await;
assert!(!server.is_finished());
release.notify_one();
let body = within(response.into_body().collect())
.await
.unwrap()
.to_bytes();
assert_eq!(body, "firstlast");
within(server).await.unwrap();
}
#[rstest]
#[tokio::test]
async fn shutdown_stops_waiting_for_a_stream_still_open_after_the_drain() {

View file

@ -32,6 +32,21 @@ pub async fn within<T>(future: impl Future<Output = T>) -> T {
tokio::time::timeout(WAIT, future).await.expect("timed out")
}
/// A forwarder in front of `upstream` that drains once `shutdown` resolves, as `run` does.
pub async fn serve_until(
upstream: SocketAddr,
keepalive: Duration,
shutdown: impl Future<Output = ()> + Send + 'static,
) -> (SocketAddr, JoinHandle<()>) {
let forward = Forward::new(&format!("http://{upstream}").parse().unwrap(), keepalive);
let public = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = public.local_addr().unwrap();
(
address,
tokio::spawn(forward.serve(public, shutdown, DRAIN)),
)
}
/// A running forwarder in front of `upstream`, listening on a real socket like production.
pub struct Gateway {
pub address: SocketAddr,
@ -45,17 +60,11 @@ impl Gateway {
}
pub async fn with_keepalive(upstream: SocketAddr, keepalive: Duration) -> Self {
let forward = Forward::new(&format!("http://{upstream}").parse().unwrap(), keepalive);
let public = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = public.local_addr().unwrap();
let (stop, stopped) = oneshot::channel::<()>();
let server = tokio::spawn(forward.serve(
public,
async {
stopped.await.ok();
},
DRAIN,
));
let (address, server) = serve_until(upstream, keepalive, async {
stopped.await.ok();
})
.await;
Self {
address,
stop,

View file

@ -48,7 +48,11 @@ async fn main() -> Result<(), Box<dyn Error>> {
level: log.log_level(),
})
.install_global()?;
litellm_gateway_forward::run(forward, shutdown_signal()).await?;
let stdin_closed = litellm_gateway_forward::input_closed(std::io::stdin());
litellm_gateway_forward::run(forward, async {
tokio::select! { () = shutdown_signal() => (), () = stdin_closed => () }
})
.await?;
return Ok(());
}
let settings: Settings = settings::from_iter(std::env::vars_os())?;

5
litellm/proxy/AGENTS.md Normal file
View file

@ -0,0 +1,5 @@
# Proxy
## Rust forwarding
`proxy_server.py`, routes, auth and middleware never know whether Rust is in front of the proxy, and none of them may branch on it. `proxy_cli.py` only describes what it would serve with as a `rust_forward.server.PythonServer`, calls `rust_forward.supervisor.supervise`, and merges the options `rust_forward.mode` returns, which are identities when forwarding is off. Every forwarding decision lives in `rust_forward/`, and `rust_forward/AGENTS.md` states the contract

View file

@ -36,6 +36,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
@ -680,7 +683,7 @@ 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,
@ -836,6 +839,65 @@ class ProxyInitializationHelpers:
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 litellm.proxy.rust_forward.server import Listener, PythonServer, Server, uvicorn_trusted_proxies
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, trusted_proxies = (
ProxyInitializationHelpers.gunicorn_keepalive_and_trusted_proxies()
if server is Server.GUNICORN
else (
uvicorn.Config(app="", log_config=None).timeout_keep_alive
if keepalive_timeout is None
else keepalive_timeout,
uvicorn_trusted_proxies(None),
)
)
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(),
trusted_proxies=trusted_proxies,
)
@staticmethod
def gunicorn_keepalive_and_trusted_proxies() -> tuple[int, str]:
"""Loaded only when gunicorn is the server, because gunicorn imports ``fcntl``, which Windows lacks."""
from gunicorn.config import Config as GunicornConfig
config: Final = GunicornConfig()
keepalive: Final[object] = config.keepalive
forwarded_allow_ips: Final[object] = config.forwarded_allow_ips
return (
TypeAdapter(int).validate_python(keepalive),
",".join(TypeAdapter(list[str]).validate_python(forwarded_allow_ips)),
)
@staticmethod
def get_loop_type():
"""Helper function to determine the event loop type based on platform"""
@ -1276,7 +1338,7 @@ def run_server(
skip_server_startup,
validate_config: bool,
check_migrations_against_existing_data: bool,
keepalive_timeout,
keepalive_timeout: int | None,
timeout_worker_healthcheck,
max_request_header_size: int | None,
max_requests_before_restart: int | None,
@ -1306,6 +1368,49 @@ def run_server(
run_setup_wizard()
return
from pydantic import ValidationError
from litellm.proxy.rust_forward.mode import (
Mode,
environment_error,
private_listener,
rust_forward_mode,
stop_with_supervisor,
uvicorn_options,
)
try:
forward_mode: Final = rust_forward_mode(os.environ)
except ValidationError as error:
raise click.ClickException(environment_error(error)) from None
stop_with_supervisor(forward_mode)
runs_one_command: Final = (
version is True
or validate_config
or health is True
or test is True
or skip_server_startup is True
or check_migrations_against_existing_data
)
if forward_mode is Mode.SUPERVISOR and not runs_one_command:
from litellm.proxy.rust_forward.supervisor import supervise
from litellm.proxy.rust_forward.supervisor.preflight import StartupError
server: Final = ProxyInitializationHelpers.python_server(
host,
port,
keepalive_timeout,
run_gunicorn=run_gunicorn is True,
run_hypercorn=run_hypercorn is True,
run_granian=run_granian is True,
tls=None not in (ssl_keyfile_path, ssl_certfile_path),
)
if prometheus_metrics_port == server.port:
raise click.UsageError("--prometheus_metrics_port must differ from --port")
supervised: Final = supervise(server)
if not isinstance(supervised, StartupError):
raise SystemExit(supervised)
args: Final = locals()
try:
from litellm.proxy.proxy_server import (
@ -1773,6 +1878,8 @@ def run_server(
timeout_worker_healthcheck=(timeout_worker_healthcheck if running_uvicorn else None),
max_request_header_size=max_request_header_size,
)
for option, value in uvicorn_options(forward_mode, WORKER_APP).items():
uvicorn_args[option] = value
# Optional: recycle uvicorn workers after N requests
if max_requests_before_restart is not None:
uvicorn_args["limit_max_requests"] = max_requests_before_restart
@ -1837,16 +1944,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,
max_request_header_size=max_request_header_size,
private_listener=private_listener(forward_mode),
)
elif run_hypercorn is True:
ProxyInitializationHelpers.init_hypercorn_server(

View file

@ -0,0 +1,48 @@
# Rust forwarding, Python half
With `LITELLM_RUST_FIRST` on, a supervisor runs the Rust gateway on the public port and Python on a private loopback listener. Rust forwards every request to Python unchanged, so Python still serves every route. 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
The same path runs everywhere. `litellm` (the main images and a local CLI run) and `python -m gateway.launch` (the gateway image) both call `supervisor.supervise`, and nothing else starts the Rust gateway in forwarding mode. In a container the supervisor is PID 1, so its signal and exit rules are deployment behavior and not CLI polish
## Boundary
A supervised Python server differs from Python alone in exactly two ways, and nothing else in the proxy may change because Rust is in front
- It listens on the private loopback listener the supervisor gives it, never on the public address
- It restores the client address Rust saw, because its own TCP peer is always Rust
`proxy_server.py`, routes, auth and middleware never know Rust is there. A change that needs one of them to branch on forwarding is in the wrong place
## Modules
- `server.py` describes what an entry point would serve with on the public port, as a `PythonServer`. The supervisor reads nothing else about the entry point
- `mode.py` is the Python child's side. `Supervised` is the whole contract the supervisor hands the child, and the `uvicorn_*` and `gunicorn_*` functions translate it into each server's own options
- `peer.py` restores the client address
- `supervisor/` starts, watches and stops both children. Its rules are in `supervisor/AGENTS.md`
- `platform/` holds every OS-specific operation. Its rules are in `platform/AGENTS.md`
## Entry points
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 and what Rust and the Python child receive. A refusal (no gateway binary, or a server or listener Rust cannot front) is reported and returned, and the entry point then serves from Python alone exactly as with `LITELLM_RUST_FIRST` off
- In the supervised child they merge `mode.uvicorn_options` or `mode.uvicorn_flags` into their server settings, or hand `mode.private_listener` to `run_gunicorn_server`, 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
- A server is started three ways here, through uvicorn keyword arguments, uvicorn's own CLI and gunicorn, and each one states the same two facts, the private listener and that the server itself trusts no peer's forwarding headers. The uvicorn translations are in `mode.py` and turn proxy headers off. `run_gunicorn_server` applies `server_listener.gunicorn_overrides` for the `PrivateListener` it is given, which sets gunicorn's `bind` and an empty `forwarded_allow_ips`. A new way to start a server gets its own translation, never a branch in the entry point
## Client address
Python must see each request as if the client had connected directly. `x-litellm-rust-peer` is the only Rust-specific input, and the proxy app itself never knows Rust is in front of it
- Rust sends the client's TCP peer in `x-litellm-rust-peer` and leaves every forwarding header untouched. It never appends itself to `X-Forwarded-For`, because the proxy parses that header on its own (`use_x_forwarded_for`, `mcp_xff_num_trusted_hops`) and an extra hop would shift what those settings resolve
- Server-level proxy headers are off in a supervised server, through `proxy_headers: False` for uvicorn, `--no-proxy-headers` for its CLI and an empty `forwarded_allow_ips` for gunicorn. The server's TCP peer is always loopback, which uvicorn and gunicorn trust by default, so leaving them on would let any client set its own address and scheme
- `peer.RustPeerMiddleware` puts the real peer back into the scope and then runs uvicorn's own `ProxyHeadersMiddleware` with the operator's trust. It wraps the app from outside, so it runs before anything that reads the client address
- The operator's trust is one string, `trusted_proxies`. It is resolved by the server library the entry point uses (`server.uvicorn_trusted_proxies` or gunicorn's config), never defaulted in this package. `server.NO_TRUSTED_PROXIES` stands for proxy headers the operator turned off, since a trust list that matches nobody behaves the same
- `peer.py` honors the peer header only from the loopback hop and always removes it. Any process on the host can connect to the private listener and send that header, so loopback is the trust boundary of the private listener. That widens what Python alone allows only when the operator's trust list leaves loopback out
## Environment contracts
Each environment contract is written and read in one place
- The Python child's contract is the `mode.Supervised` type itself. `mode.supervised_environment` serializes it into one variable and `mode.rust_forward_mode` validates it back, so a new field needs no parsing code
- `supervisor.children.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 there and a matching field in `settings.rs` in the same change
- `mode.rust_forward_mode` reads `LITELLM_RUST_FIRST` strictly, so an invalid value stops the CLI at startup with one line naming it. It is independent of `LITELLM_RUST`, which only selects the Rust paths inside the SDK in `litellm/rust_bridge/configuration.py`, so neither turns the other on

View file

View file

@ -0,0 +1,142 @@
import contextlib
import os
import signal
import threading
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from enum import Enum, auto
from types import MappingProxyType
from typing import Final, TypeAlias
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
from typing_extensions import assert_never
from litellm.proxy.server_listener import LOOPBACK_HOST, AssignedPort, InheritedFd, PrivateListener
from litellm.rust_bridge.configuration import RustSwitch
SWITCH_ENV: Final = "LITELLM_RUST_FIRST"
SUPERVISED_ENV: Final = "_LITELLM_GATEWAY_SUPERVISED"
RUST_HOP: Final = LOOPBACK_HOST
class _Switch(BaseModel):
model_config = ConfigDict(frozen=True)
rust: RustSwitch = Field(None, validation_alias=SWITCH_ENV)
class Mode(Enum):
DISABLED = auto()
SUPERVISOR = auto()
@dataclass(frozen=True, slots=True)
class Supervised:
listener: PrivateListener
trusted_proxies: str
RustForwardMode: TypeAlias = Mode | Supervised
_SUPERVISED: Final = TypeAdapter(Supervised)
def supervised_environment(supervised: Supervised) -> Mapping[str, str]:
return MappingProxyType({SUPERVISED_ENV: _SUPERVISED.dump_json(supervised).decode()})
def rust_forward_mode(environ: Mapping[str, str]) -> RustForwardMode:
switch: Final = _Switch.model_validate(environ)
supervised: Final = environ.get(SUPERVISED_ENV)
if supervised is not None:
return _SUPERVISED.validate_json(supervised)
return Mode.SUPERVISOR if switch.rust else Mode.DISABLED
def _wait_for_eof(fd: int) -> None:
with contextlib.suppress(OSError):
while os.read(fd, 4096):
pass
def _stop_when_closed(fd: int, stop: Callable[[], None]) -> None:
_wait_for_eof(fd)
stop()
def _terminate_self() -> None:
signal.raise_signal(signal.SIGTERM)
def stop_with_supervisor(mode: RustForwardMode, fd: int = 0, stop: Callable[[], None] = _terminate_self) -> None:
"""Shut the supervised server down gracefully once the supervisor's end of its stdin closes, so it never outlives
a supervisor that was killed. Call it only in the process the supervisor started, not in its workers."""
match mode:
case Supervised():
threading.Thread(target=_stop_when_closed, args=(fd, stop), name="supervisor-stdin", daemon=True).start()
case Mode.DISABLED | Mode.SUPERVISOR:
pass
case _:
assert_never(mode)
def environment_error(error: ValidationError) -> str:
return "; ".join(
f"{'.'.join(str(part) for part in detail['loc'])} is invalid: {detail['msg']}"
for detail in error.errors(include_url=False)
)
# 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: PrivateListener) -> 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 Mode.DISABLED | Mode.SUPERVISOR:
return MappingProxyType({})
case _:
assert_never(mode)
def _uvicorn_listener_flags(listener: PrivateListener) -> 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 Mode.DISABLED | Mode.SUPERVISOR:
return ()
case _:
assert_never(mode)
def private_listener(mode: RustForwardMode) -> PrivateListener | None:
match mode:
case Supervised(listener):
return listener
case Mode.DISABLED | Mode.SUPERVISOR:
return None
case _:
assert_never(mode)

View file

@ -0,0 +1,69 @@
from typing import (
TYPE_CHECKING,
Final,
TypeAlias,
cast, # noqa: TID251 # the one Starlette-to-uvicorn ASGI conversion in served
)
from typing_extensions import assert_never
from litellm.proxy.rust_forward.mode import RUST_HOP, Mode, RustForwardMode, Supervised
if TYPE_CHECKING:
import starlette.types
from uvicorn._types import ASGI3Application, ASGIReceiveCallable, ASGISendCallable, Scope, WWWScope
PEER_HEADER: Final = b"x-litellm-rust-peer"
ASGIApp: TypeAlias = "ASGI3Application"
def served(app: "starlette.types.ASGIApp") -> ASGIApp:
return cast( # cast-ok: Starlette and uvicorn type the same ASGI 3 callable with unassignable scope/receive/send
"ASGI3Application", app
)
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)
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: "starlette.types.ASGIApp", mode: RustForwardMode) -> ASGIApp:
converted: Final = served(app)
match mode:
case Supervised():
return RustPeerMiddleware(converted, mode)
case Mode.DISABLED | Mode.SUPERVISOR:
return converted
case _:
assert_never(mode)

View file

@ -0,0 +1,19 @@
# Platforms
Every OS-specific operation sits behind `Platform`, which `supervise` picks once with `current_platform()`. The supervisor, `mode.py` and the entry points never check the OS themselves, so a new OS difference becomes a `Platform` member and never an `os.name` check somewhere else
A `Platform` answers these questions and nothing else. What the gateway binary is called, which servers it can supervise, how the private listener reaches Python, which spawn options the children get, which signals are relayed and which of them stop Rust, and how a child is signalled and stopped
## POSIX
`PosixPlatform` binds the private socket itself and hands Python the open descriptor (`server_listener.InheritedFd`). Nobody else can take the port between the supervisor picking it and Python serving it. The socket does not listen until Python does, so the readiness probe only succeeds against Python. Children run in their own session (`start_new_session`) only so that a terminal Ctrl-C reaches Python once, through the supervisor
## Windows
Windows is supported for running the proxy locally from the CLI. Nobody deploys the proxy on Windows, but `LITELLM_RUST_FIRST` is meant to become the default, and then this is the path every Windows laptop runs. Revisit the choices below before treating Windows as a deployment target
- Python gets an assigned loopback port (`server_listener.AssignedPort`) instead of an open socket, which uvicorn on Windows cannot serve. Another process can 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
- Each child runs in its own process group, which can only be sent `CTRL_BREAK_EVENT`. Both relayed signals arrive as that, so a second Ctrl-C does not force Python's exit the way it does on POSIX
- gunicorn is refused, since it does not run on Windows
- The Windows-only names are spelled out as constants in `windows.py` so the module imports and type checks on every OS. `tests/windows_tests/test_rust_forward_platform.py` pins them to the stdlib on the Windows runner
- No CI job runs the supervisor itself on Windows, because the `windows-sdk` job installs no proxy dependencies and no `litellm-gateway.exe`. Run a change to `windows.py` by hand on Windows until one does

View file

@ -0,0 +1,58 @@
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.server import Server
from litellm.proxy.server_listener import PrivateListener
@dataclass(frozen=True, slots=True)
class PlatformListener:
listener: PrivateListener
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[PlatformListener]: ...
def spawn_options(self, pass_fds: tuple[int, ...]) -> SpawnOptions: ...
@property
def relayed_signals(self) -> tuple[signal.Signals, ...]: ...
@property
def stop_signals(self) -> frozenset[signal.Signals]: ...
def catch(self, loop: asyncio.AbstractEventLoop, on_signal: 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()

View file

@ -0,0 +1,64 @@
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
from litellm.proxy.rust_forward.platform import PlatformListener, SpawnOptions
from litellm.proxy.rust_forward.server import Server
from litellm.proxy.server_listener import InheritedFd
_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[PlatformListener]:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as private:
private.bind((RUST_HOP, 0))
_, port = _ADDRESS.validate_python(private.getsockname())
yield PlatformListener(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)
@property
def relayed_signals(self) -> tuple[signal.Signals, ...]:
return (
signal.SIGTERM,
signal.SIGINT,
signal.SIGHUP,
signal.SIGQUIT,
signal.SIGUSR1,
signal.SIGUSR2,
signal.SIGTTIN,
signal.SIGTTOU,
signal.SIGWINCH,
)
@property
def stop_signals(self) -> frozenset[signal.Signals]:
return frozenset({signal.SIGTERM, signal.SIGINT, signal.SIGQUIT})
def catch(self, loop: asyncio.AbstractEventLoop, on_signal: Callable[[signal.Signals], None]) -> None:
for signum in self.relayed_signals:
loop.add_signal_handler(signum, on_signal, 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)

View file

@ -0,0 +1,68 @@
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
from litellm.proxy.rust_forward.platform import PlatformListener, SpawnOptions
from litellm.proxy.rust_forward.server import Server
from litellm.proxy.server_listener import AssignedPort
_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[PlatformListener]:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
probe.bind((RUST_HOP, 0))
_, port = _ADDRESS.validate_python(probe.getsockname())
yield PlatformListener(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)
@property
def relayed_signals(self) -> tuple[signal.Signals, ...]:
return (signal.SIGINT, signal.Signals(SIGBREAK))
@property
def stop_signals(self) -> frozenset[signal.Signals]:
return frozenset(self.relayed_signals)
def catch(self, loop: asyncio.AbstractEventLoop, on_signal: Callable[[signal.Signals], None]) -> None:
def handle(signum: int, _frame: FrameType | None) -> None:
loop.call_soon_threadsafe(on_signal, signal.Signals(signum))
for signum in self.relayed_signals:
signal.signal(signum, handle)
# A child in its own process group can only be sent CTRL_BREAK_EVENT, since Windows never delivers CTRL_C_EVENT
# to a new group. Both children shut down on it as on Ctrl-C, but uvicorn force-exits only on a repeated SIGINT,
# so a second Ctrl-C does not cut a Windows shutdown short the way it does on POSIX.
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)

View file

@ -0,0 +1,75 @@
from dataclasses import dataclass
from enum import Enum
from typing import Final
from pydantic import BaseModel, ConfigDict, TypeAdapter
from litellm.proxy.server_listener import NO_TRUSTED_PROXIES
_TRUSTED_PROXIES: Final = TypeAdapter(str)
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
trusted_proxies: str
class UvicornCliParams(BaseModel):
"""The options uvicorn's own CLI parsed that decide how it serves the public port."""
model_config = ConfigDict(frozen=True, extra="ignore")
host: str
port: int
uds: str | None
fd: int | None
ssl_keyfile: str | None
ssl_certfile: str | None
timeout_keep_alive: int
timeout_graceful_shutdown: int | None
proxy_headers: bool
forwarded_allow_ips: str | None
@property
def listener(self) -> Listener:
if self.ssl_keyfile is not None or self.ssl_certfile is not None:
return Listener.TLS
if self.uds is not None:
return Listener.UNIX_SOCKET
if self.fd is not None:
return Listener.INHERITED_FD
return Listener.TCP
@property
def trusted_proxies(self) -> str:
return uvicorn_trusted_proxies(self.forwarded_allow_ips) if self.proxy_headers else NO_TRUSTED_PROXIES
def uvicorn_trusted_proxies(configured: str | None) -> str:
import uvicorn
return _TRUSTED_PROXIES.validate_python(
uvicorn.Config(app="", log_config=None, forwarded_allow_ips=configured).forwarded_allow_ips
)

View file

@ -0,0 +1,21 @@
# Forwarding supervisor
`supervise` serves the public port from Rust and runs the entry point's own command behind it, or says why that cannot start. That command is the supervisor's own interpreter command line (`OWN_COMMAND`, from `sys.orig_argv`), so the child has the same import path and interpreter flags as the process the operator started. An entry point never builds a second way to start itself. It is transparent to Python, so a supervised proxy behaves on signals and exits exactly as Python alone. `tests/integration/configuration/test_rust_forward_lifecycle.py` checks that equivalence against Python alone, and a change that breaks one of the rules below breaks it
## Split by job
- `preflight` decides whether forwarding can start. It is the only place that refuses, and a refusal is a returned `StartupError`, never an exit, so the entry point can serve from Python alone
- `children` says what each child runs with, and it is the only writer of both environment contracts. `python_spec` writes the `mode.Supervised` variable and hands the Python child the entry point's own arguments with the private port appended. `rust_spec` writes the `LITELLM_FORWARD_*` variables. Each child gets only those variables on top of the supervisor's own environment
- `lifecycle` decides. `step` is a pure function from a `Lifecycle` of facts and an event to the next facts and the actions to run, and `exit_status` reads the result from those facts. It touches no process, socket or clock
- `runtime` acts. It only executes the actions `step` returns and feeds events back. The `contextlib.AsyncExitStack` it enters each child into owns their cleanup, last spawned first
A new behavior is a new fact, event or action in `lifecycle` with a property test in `test_lifecycle.py`, never a condition in `runtime`
## Rules
- Signals pass through. Every signal in `Platform.relayed_signals` goes to Python unchanged, and the ones in `Platform.stop_signals` go to Rust first. The supervisor never reads meaning into a signal, so a repeated SIGINT reaches Python again and forces its exit as it would alone
- Python's exit status is the result. When Python exits for any reason, Rust gets SIGTERM, and once Rust is gone the supervisor exits with Python's status, dying by the same signal when Python died by one. The one exception is Rust exiting while Python serves with no stop requested, which stops Python and exits 1
- A container's PID 1 cannot die by a signal it sends itself, since the kernel drops it. uvicorn alone survives its own re-raise there and exits 0, so `runtime.exit_like` exits 0 when a stop signal killed Python and with 128 plus the signal number for any other signal, which is how a kill of Python alone is reported
- Children cannot outlive the supervisor. Each child's stdin is a pipe only the supervisor holds open. The Python process the supervisor started (`mode.stop_with_supervisor`, never its workers) and the Rust gateway (`litellm_gateway_forward::input_closed`) shut down gracefully on EOF, so a supervisor killed by SIGKILL or a closed terminal leaves nothing behind
- Rust starts once Python listens, never after a stop was requested or Python exited. Python alone takes no connections until it has booted, so the public port refuses them until then too
- Only a preflight refusal falls back to Python alone. `before_start` runs once forwarding is known to be possible, so a refused start leaves nothing behind, and a child that fails to spawn after that exits 1

View file

@ -0,0 +1,55 @@
import asyncio
import sys
import sysconfig
from collections.abc import Callable, Sequence
from typing import Final
from typing_extensions import assert_never
from litellm.proxy.rust_forward.platform import Platform, current_platform
from litellm.proxy.rust_forward.server import PythonServer
from litellm.proxy.rust_forward.supervisor.lifecycle import FAILURE
from litellm.proxy.rust_forward.supervisor.preflight import StartupError, gateway_binary, startup_check
from litellm.proxy.rust_forward.supervisor.runtime import SpawnFailed, exit_like, run
OWN_COMMAND: Final = (sys.executable, *sys.orig_argv[1:])
def _nothing_to_prepare() -> None:
pass
def supervise(
server: PythonServer,
command: Sequence[str] = OWN_COMMAND,
before_start: Callable[[], None] = _nothing_to_prepare,
scripts_dir: str = sysconfig.get_path("scripts"),
platform: Platform | None = None,
) -> int | StartupError:
"""Serve ``server``'s public port from Rust and run ``command`` behind it, or say why that cannot start.
A refusal is returned for the caller to serve from Python alone, and ``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}, so Python serves the public port without Rust forwarding\n")
return checked
before_start()
try:
outcome: Final = asyncio.run(run(chosen, checked, command, server))
except OSError as error:
return _startup_failed(error)
match outcome:
case SpawnFailed(error):
return _startup_failed(error)
case int():
return exit_like(outcome, chosen.stop_signals)
case _:
assert_never(outcome)
def _startup_failed(error: OSError) -> int:
sys.stderr.write(f"Rust gateway startup failed: {error}\n")
return FAILURE

View file

@ -0,0 +1,43 @@
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS
from litellm.proxy.rust_forward.mode import RUST_HOP, Supervised, supervised_environment
from litellm.proxy.rust_forward.platform import Platform, PlatformListener, SpawnOptions
from litellm.proxy.rust_forward.server import PythonServer
@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: PlatformListener, server: PythonServer
) -> ChildSpec:
return ChildSpec(
argv=(*command, "--port", str(private.port)),
env_overrides=supervised_environment(Supervised(private.listener, 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(()),
)

View file

@ -0,0 +1,138 @@
import signal
from dataclasses import dataclass, replace
from enum import Enum, auto
from typing import Final, Literal, TypeAlias
from typing_extensions import assert_never
class Child(Enum):
RUST = "rust"
PYTHON = "python"
FAILURE: Final = 1
class Phase(Enum):
NOT_STARTED = auto()
SPAWNED = auto()
LISTENING = auto()
@dataclass(frozen=True, slots=True)
class Exited:
status: int
@dataclass(frozen=True, slots=True)
class Lost:
"""Rust exited while Python was serving and nobody had asked either to stop."""
status: int
PythonState: TypeAlias = Literal[Phase.SPAWNED, Phase.LISTENING] | Exited
RustState: TypeAlias = Literal[Phase.NOT_STARTED, Phase.SPAWNED] | Exited | Lost
@dataclass(frozen=True, slots=True)
class Lifecycle:
python: PythonState
rust: RustState
stop_requested: bool
@dataclass(frozen=True, slots=True)
class Caught:
signum: signal.Signals
class Probe(Enum):
PYTHON_LISTENING = auto()
@dataclass(frozen=True, slots=True)
class ChildExited:
child: Child
status: int
@dataclass(frozen=True, slots=True)
class Spawn:
child: Child
@dataclass(frozen=True, slots=True)
class Send:
child: Child
signum: signal.Signals
Event: TypeAlias = Caught | Literal[Probe.PYTHON_LISTENING] | ChildExited
Action: TypeAlias = Spawn | Send
Transition: TypeAlias = tuple[Lifecycle, tuple[Action, ...]]
def boot() -> Transition:
return Lifecycle(Phase.SPAWNED, Phase.NOT_STARTED, stop_requested=False), (Spawn(Child.PYTHON),)
def _relay(state: Lifecycle, signum: signal.Signals, stops: bool) -> Transition:
to_rust: Final = (Send(Child.RUST, signum),) if stops and state.rust is Phase.SPAWNED else ()
to_python: Final = (Send(Child.PYTHON, signum),) if not isinstance(state.python, Exited) else ()
return replace(state, stop_requested=state.stop_requested or stops), (*to_rust, *to_python)
# Python alone takes no connections until it has booted, so Rust only opens the public port once Python listens.
def _python_listening(state: Lifecycle) -> Transition:
if state.python is not Phase.SPAWNED:
return state, ()
if state.stop_requested:
return replace(state, python=Phase.LISTENING), ()
return replace(state, python=Phase.LISTENING, rust=Phase.SPAWNED), (Spawn(Child.RUST),)
def _python_exited(state: Lifecycle, status: int) -> Transition:
if isinstance(state.python, Exited):
return state, ()
exited: Final = replace(state, python=Exited(status))
return exited, ((Send(Child.RUST, signal.SIGTERM),) if state.rust is Phase.SPAWNED else ())
def _rust_exited(state: Lifecycle, status: int) -> Transition:
if state.rust is not Phase.SPAWNED:
return state, ()
if state.stop_requested or isinstance(state.python, Exited):
return replace(state, rust=Exited(status)), ()
return replace(state, rust=Lost(status)), (Send(Child.PYTHON, signal.SIGTERM),)
def step(state: Lifecycle, event: Event, stop_signals: frozenset[signal.Signals]) -> Transition:
"""The supervisor never gives a signal or an exit a meaning of its own. Every signal reaches Python as sent,
Python's exit status is the supervisor's, and Rust only follows Python, except that losing Rust stops Python."""
match event:
case Caught(signum):
return _relay(state, signum, signum in stop_signals)
case Probe.PYTHON_LISTENING:
return _python_listening(state)
case ChildExited(child, status):
match child:
case Child.PYTHON:
return _python_exited(state, status)
case Child.RUST:
return _rust_exited(state, status)
case _:
assert_never(child)
case _:
assert_never(event)
def exit_status(state: Lifecycle) -> int | None:
match state.python, state.rust:
case Exited(status), Phase.NOT_STARTED | Exited():
return status
case Exited(), Lost():
return FAILURE
case _:
return None

View file

@ -0,0 +1,28 @@
import os
import shutil
from dataclasses import dataclass
from litellm.proxy.rust_forward.mode import SWITCH_ENV
from litellm.proxy.rust_forward.platform import Platform
from litellm.proxy.rust_forward.server import Listener, PythonServer, Server
@dataclass(frozen=True, slots=True)
class StartupError:
reason: str
def gateway_binary(platform: Platform, scripts_dir: str) -> str | None:
return shutil.which(os.path.join(scripts_dir, platform.gateway_name))
def startup_check(platform: Platform, server: PythonServer, binary: str | None) -> str | StartupError:
if server.server in (Server.HYPERCORN, Server.GRANIAN):
return StartupError(f"{SWITCH_ENV} does not support {server.server.value}")
if server.server not in platform.servers:
return StartupError(f"{SWITCH_ENV} does not support {server.server.value} on {platform.name}")
if server.listener is not Listener.TCP:
return StartupError(f"{SWITCH_ENV} does not support {server.listener.value} on the public port")
if binary is None:
return StartupError(f"{SWITCH_ENV} requires the litellm-gateway binary from the litellm-proxy package")
return binary

View file

@ -0,0 +1,194 @@
import asyncio
import contextlib
import math
import os
import signal
from collections.abc import AsyncGenerator, Callable, Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, Protocol
import anyio
from anyio.abc import TaskGroup
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from typing_extensions import assert_never
from litellm.proxy.rust_forward.mode import RUST_HOP
from litellm.proxy.rust_forward.platform import Platform
from litellm.proxy.rust_forward.server import PythonServer
from litellm.proxy.rust_forward.supervisor.children import ChildSpec, python_spec, rust_spec
from litellm.proxy.rust_forward.supervisor.lifecycle import (
Caught,
Child,
ChildExited,
Event,
Probe,
Send,
Spawn,
boot,
exit_status,
step,
)
PROBE_INTERVAL_SECONDS: Final = 0.05
@dataclass(frozen=True, slots=True)
class SpawnFailed:
error: OSError
class Executor(Protocol):
async def spawn(self, child: Child) -> SpawnFailed | None: ...
def send(self, child: Child, signum: signal.Signals) -> None: ...
async def next_event(self) -> Event: ...
async def run(platform: Platform, binary: str, command: Sequence[str], server: PythonServer) -> int | SpawnFailed:
with platform.private_listener() as private:
specs: Final = MappingProxyType(
{
Child.PYTHON: python_spec(platform, command, private, server),
Child.RUST: rust_spec(platform, binary, private.port, server),
}
)
async with open_children(platform, specs, private.port) as children:
platform.catch(asyncio.get_running_loop(), children.caught)
return await drive(children, platform.stop_signals)
async def drive(executor: Executor, stop_signals: frozenset[signal.Signals]) -> int | SpawnFailed:
"""Carry out what ``lifecycle.step`` decides until it settles the exit status."""
state, actions = boot()
while True:
for action in actions:
match action:
case Spawn(child):
if (failed := await executor.spawn(child)) is not None:
return failed
case Send(child, signum):
executor.send(child, signum)
case _:
assert_never(action)
if (status := exit_status(state)) is not None:
return status
state, actions = step(state, await executor.next_event(), stop_signals)
def _die_by(signum: signal.Signals) -> None:
with contextlib.suppress(OSError):
signal.signal(signum, signal.SIG_DFL)
signal.raise_signal(signum)
def exit_like(
status: int, stop_signals: frozenset[signal.Signals], die_by: Callable[[signal.Signals], None] = _die_by
) -> int:
"""A child killed by a signal kills the supervisor with the same signal, as uvicorn re-raises the one it caught.
SIGKILL takes no handler and needs none to kill. The kernel drops the signal when the supervisor is a container's
PID 1. uvicorn alone survives its own re-raise there and exits 0, so a stop signal ends the supervisor with 0 too,
and any other signal with the status a shell reports for it."""
if status >= 0:
return status
signum: Final = signal.Signals(-status)
die_by(signum)
return 0 if signum in stop_signals else 128 + signum
@dataclass(frozen=True, slots=True)
class Children:
platform: Platform
specs: Mapping[Child, ChildSpec]
private_port: int
stack: contextlib.AsyncExitStack
tasks: TaskGroup
events_in: MemoryObjectSendStream[Event]
events_out: MemoryObjectReceiveStream[Event]
processes: dict[Child, asyncio.subprocess.Process] # mutable-ok: the running children the stack owns
def caught(self, signum: signal.Signals) -> None:
with contextlib.suppress(anyio.ClosedResourceError):
self.events_in.send_nowait(Caught(signum))
async def spawn(self, child: Child) -> SpawnFailed | None:
try:
process: Final = await self.stack.enter_async_context(_process(self.platform, self.specs[child]))
except OSError as error:
return SpawnFailed(error)
self.processes[child] = process
self.tasks.start_soon(_report_exit, child, process, self.events_in.clone())
if child is Child.PYTHON:
probe: Final = self.tasks.start_soon(_report_once_listening, self.private_port, self.events_in.clone())
self.stack.callback(probe.cancel)
return None
def send(self, child: Child, signum: signal.Signals) -> None:
self.platform.signal_child(self.processes[child], signum)
async def next_event(self) -> Event:
return await self.events_out.receive()
# The task group is entered before any child, so on the way out every child is stopped, last spawned first, before
# the task group waits for the tasks reporting their exits. Signal handlers cannot wait for room, hence the
# unbounded event buffer.
@contextlib.asynccontextmanager
async def open_children(
platform: Platform, specs: Mapping[Child, ChildSpec], private_port: int
) -> AsyncGenerator[Children]:
events_in, events_out = anyio.create_memory_object_stream[Event](math.inf)
with events_in, events_out:
async with contextlib.AsyncExitStack() as stack:
tasks: Final = await stack.enter_async_context(anyio.create_task_group())
yield Children(platform, specs, private_port, stack, tasks, events_in, events_out, {})
# A child's stdin is a pipe only the supervisor holds open, so the child sees EOF as soon as the supervisor is gone,
# however it went, and shuts down instead of outliving it. Stopping is shielded so a cancelled supervisor still waits
# for the child it is stopping.
@contextlib.asynccontextmanager
async def _process(platform: Platform, spec: ChildSpec) -> AsyncGenerator[asyncio.subprocess.Process]:
process: Final = await asyncio.create_subprocess_exec(
*spec.argv,
stdin=asyncio.subprocess.PIPE,
env={**os.environ, **spec.env_overrides},
pass_fds=spec.spawn.pass_fds,
start_new_session=spec.spawn.start_new_session,
creationflags=spec.spawn.creationflags,
)
try:
yield process
finally:
with anyio.CancelScope(shield=True):
platform.stop_group(process)
await process.wait()
if process.stdin is not None:
process.stdin.close()
async def _report_exit(
child: Child, process: asyncio.subprocess.Process, events: MemoryObjectSendStream[Event]
) -> None:
with events:
events.send_nowait(ChildExited(child, await process.wait()))
async def _report_once_listening(port: int, events: MemoryObjectSendStream[Event]) -> None:
with events:
while not await _accepts(port):
await anyio.sleep(PROBE_INTERVAL_SECONDS)
events.send_nowait(Probe.PYTHON_LISTENING)
async def _accepts(port: int) -> bool:
with anyio.move_on_after(PROBE_INTERVAL_SECONDS):
try:
stream: Final = await anyio.connect_tcp(RUST_HOP, port)
except OSError:
return False
await stream.aclose()
return True
return False

View file

@ -46,6 +46,11 @@
- A failure after execution starts is terminal and never triggers fallback
- Creating an unstarted native coroutine acquires no clients or credentials and captures no execution context
# The `LITELLM_RUST` switch
- `decision` treats a blank or invalid `LITELLM_RUST` as unset and caches the parse by raw value. It runs inside SDK calls, so it never fails one over a switch and never validates per call
- It does not put the Rust gateway in front of Python. That is `LITELLM_RUST_FIRST`, which `litellm/proxy/rust_forward/mode.py` reads strictly at startup with this same `RustSwitch` type, stopping the CLI with one line naming the variable
# Native extension
- `_native` ships in every wheel and the sdist builds it, so it is always importable

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

View file

@ -145,6 +145,11 @@ def signal_group(group: int, action: int) -> None:
pass
def send_signal(process: subprocess.Popen[bytes], signum: signal.Signals) -> None:
with suppress(ProcessLookupError):
process.send_signal(signum)
def graceful_stop_seconds() -> float:
return readiness_for(import_seconds(), 1)

View file

@ -0,0 +1,117 @@
import json
import os
import subprocess
import sys
from collections.abc import Iterator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from types import MappingProxyType
from typing import Final
import httpx
from pydantic import JsonValue
from tests.integration._support.client import eventually
from tests.integration._support.ports import allocate_port
from tests.integration._support.process import stop_root_process
MASTER_KEY: Final = "sk-rust-forward-integration-key-not-a-secret"
WORKERS: Final = "2"
HEALTHCHECK_SECONDS: Final = "120"
class EntryPoint(Enum):
LITELLM = "litellm.proxy.proxy_cli"
GATEWAY = "gateway.launch"
@dataclass(frozen=True, slots=True)
class ForwardedProxy:
port: int
log: Path
@property
def url(self) -> str:
return f"http://127.0.0.1:{self.port}"
def _serving(url: str) -> bool:
try:
return httpx.get(f"{url}/health/liveliness", timeout=5, trust_env=False).status_code == 200
except httpx.TransportError:
return False
def _launch(entry_point: EntryPoint, config: Path) -> tuple[tuple[str, ...], Mapping[str, str]]:
match entry_point:
case EntryPoint.LITELLM:
return (
(
"--config",
str(config),
"--num_workers",
WORKERS,
"--timeout_worker_healthcheck",
HEALTHCHECK_SECONDS,
),
MappingProxyType({}),
)
case EntryPoint.GATEWAY:
return (
("--workers", WORKERS, "--timeout-worker-healthcheck", HEALTHCHECK_SECONDS),
MappingProxyType({"CONFIG_FILE_PATH": str(config), "LITELLM_PGBOUNCER_ENABLED": "false"}),
)
@contextmanager
def forwarded_proxy(
entry_point: EntryPoint,
config: Mapping[str, JsonValue],
directory: Path,
arguments: tuple[str, ...] = (),
environment: Mapping[str, str] = MappingProxyType({}),
) -> Iterator[ForwardedProxy]:
"""``entry_point`` on 0.0.0.0 with two workers, ``LITELLM_RUST_FIRST=1`` and no database. It must exit 0 on SIGTERM."""
proxy: Final = ForwardedProxy(allocate_port(), directory / f"{entry_point.name.lower()}.log")
config_file: Final = directory / "config.yaml"
config_file.write_text(json.dumps(config))
launch_arguments, launch_environment = _launch(entry_point, config_file)
inherited: Final = {
name: value for name, value in os.environ.items() if not name.startswith("DATABASE_") and name != "DIRECT_URL"
}
with (
proxy.log.open("wb") as output,
subprocess.Popen(
(
sys.executable,
"-I",
"-m",
entry_point.value,
"--host",
"0.0.0.0",
"--port",
str(proxy.port),
*launch_arguments,
*arguments,
),
env={
**inherited,
"LITELLM_MASTER_KEY": MASTER_KEY,
"LITELLM_RUST_FIRST": "true",
"STORE_MODEL_IN_DB": "False",
"PYTHON_DOTENV_DISABLED": "1",
**launch_environment,
**environment,
},
stdout=output,
stderr=output,
) as process,
):
try:
assert eventually(lambda: _serving(proxy.url), bool, seconds=180), proxy.log.read_text()
yield proxy
finally:
assert stop_root_process(process), proxy.log.read_text()
assert process.returncode == 0, proxy.log.read_text()

View file

@ -0,0 +1,110 @@
from dataclasses import dataclass
from pathlib import Path
from typing import Final
import httpx
import psutil
import pytest
from tests.integration._support.client import eventually
from tests.integration._support.ports import allocate_port
from tests.integration._support.rust_forward import MASTER_KEY, EntryPoint, forwarded_proxy
TRUSTED_BALANCER: Final = "127.0.0.3"
CLIENT: Final = "127.0.0.2"
@dataclass(frozen=True, slots=True)
class Caller:
connects_from: str
forwarded_for: str
resolved_ip: str
CALLERS: Final = (
Caller(CLIENT, "1.2.3.4", resolved_ip=CLIENT),
Caller(TRUSTED_BALANCER, "203.0.113.9", resolved_ip="203.0.113.9"),
Caller(CLIENT, f"203.0.113.9, {TRUSTED_BALANCER}", resolved_ip=CLIENT),
)
def _client(url: str, local_address: str) -> httpx.Client:
return httpx.Client(
base_url=url, transport=httpx.HTTPTransport(local_address=local_address), trust_env=False, timeout=5
)
def _resolved_ip(url: str, caller: Caller) -> str:
with _client(url, caller.connects_from) as client:
response: Final = client.get(
"/v1/mcp/network/client-ip",
headers={"Authorization": f"Bearer {MASTER_KEY}", "X-Forwarded-For": caller.forwarded_for},
)
assert response.status_code == 200, response.text
return response.json()["ip"]
@pytest.mark.parametrize("server_arguments", [(), ("--run_gunicorn",)])
def test_only_the_trusted_balancer_can_name_the_client_ip_behind_rust(
tmp_path: Path, server_arguments: tuple[str, ...]
) -> None:
metrics_port: Final = allocate_port()
with forwarded_proxy(
EntryPoint.LITELLM,
{
"model_list": [],
"general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY", "use_x_forwarded_for": False},
},
tmp_path,
("--prometheus_metrics_port", str(metrics_port), *server_arguments),
{"FORWARDED_ALLOW_IPS": TRUSTED_BALANCER},
) as proxy:
with _client(f"http://{CLIENT}:{metrics_port}", CLIENT) as external:
metrics: Final = external.get("/metrics", follow_redirects=True)
assert metrics.status_code == 200, metrics.text
assert tuple(_resolved_ip(proxy.url, caller) for caller in CALLERS) == tuple(
caller.resolved_ip for caller in CALLERS
)
@pytest.mark.parametrize("server_arguments", [("--no-proxy-headers",), ("--forwarded-allow-ips", "127.0.0.1")])
def test_gateway_launcher_logs_the_real_client_address_behind_rust(
tmp_path: Path, server_arguments: tuple[str, ...]
) -> None:
with forwarded_proxy(
EntryPoint.GATEWAY,
{"model_list": [], "general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY"}},
tmp_path,
server_arguments,
) as proxy:
with _client(proxy.url, CLIENT) as external:
response: Final = external.get(
"/health/liveliness?client-ip-proof", headers={"X-Forwarded-For": "198.51.100.10"}
)
assert response.status_code == 200, response.text
logged: Final = eventually(
lambda: tuple(
line for line in proxy.log.read_text().splitlines() if "GET /health/liveliness?client-ip-proof " in line
),
bool,
seconds=10,
)
assert len(logged) == 1 and f"{CLIENT}:" in logged[0], logged
@pytest.mark.parametrize("entry_point", list(EntryPoint), ids=[entry_point.name for entry_point in EntryPoint])
def test_the_python_child_is_started_with_the_supervisors_own_command(tmp_path: Path, entry_point: EntryPoint) -> None:
with forwarded_proxy(
entry_point,
{"model_list": [], "general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY"}},
tmp_path,
):
supervisor: Final = next(
started for started in psutil.Process().children() if entry_point.value in started.cmdline()
)
own_command: Final = supervisor.cmdline()[1:]
python_children: Final = [
child.cmdline()[1:] for child in supervisor.children() if entry_point.value in child.cmdline()
]
assert "-I" in own_command
assert [child[: len(own_command)] for child in python_children] == [own_command]

View file

@ -0,0 +1,269 @@
import os
import signal
import socket
import subprocess
import sys
import time
from collections.abc import Callable
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import Final, TypeVar
import httpx
import psutil
import pytest
from tests.integration._support.client import eventually
from tests.integration._support.ports import allocate_port
from tests.integration._support.process import kill_stray_process, send_signal
pytestmark = pytest.mark.skipif(sys.platform == "win32", reason="drives POSIX signals")
APP: Final = """
import asyncio
async def app(scope, receive, send):
if scope["type"] != "http":
return
await send({"type": "http.response.start", "status": 200, "headers": [(b"content-type", b"text/plain")]})
if scope["path"] != "/stream":
await send({"type": "http.response.body", "body": b"ok"})
return
while True:
await send({"type": "http.response.body", "body": b"tick\\n", "more_body": True})
await asyncio.sleep(0.1)
"""
SERVE: Final = """
import os
import pathlib
import sys
import uvicorn
def listener():
from litellm.proxy.rust_forward.mode import Supervised, rust_forward_mode, stop_with_supervisor
mode = rust_forward_mode(os.environ)
if not isinstance(mode, Supervised):
return {"host": "127.0.0.1", "port": int(sys.argv[2])}
stop_with_supervisor(mode)
pathlib.Path(os.environ["MARKERS"], "python.pid").write_text(str(os.getpid()))
return {"fd": mode.listener.fd}
if __name__ == "__main__":
uvicorn.run("app:app", app_dir=os.path.dirname(__file__), workers=int(sys.argv[1]), log_level="warning", **listener())
"""
GATEWAY: Final = """
import asyncio
import os
import pathlib
import signal
import threading
def wait_for_eof(loop, stop):
while os.read(0, 4096):
pass
loop.call_soon_threadsafe(stop.set)
async def pipe(reader, writer):
try:
while data := await reader.read(65536):
writer.write(data)
await writer.drain()
except OSError:
pass
finally:
writer.close()
async def main():
pathlib.Path(os.environ["MARKERS"], "rust.pid").write_text(str(os.getpid()))
host, port = os.environ["LITELLM_FORWARD_UPSTREAM"].removeprefix("http://").rsplit(":", 1)
loop = asyncio.get_running_loop()
stop = asyncio.Event()
for signum in (signal.SIGTERM, signal.SIGINT):
loop.add_signal_handler(signum, stop.set)
threading.Thread(target=wait_for_eof, args=(loop, stop), daemon=True).start()
relays = set()
async def relay(client_reader, client_writer):
relays.add(asyncio.current_task())
try:
upstream_reader, upstream_writer = await asyncio.open_connection(host, int(port))
await asyncio.gather(pipe(client_reader, upstream_writer), pipe(upstream_reader, client_writer))
finally:
relays.discard(asyncio.current_task())
server = await asyncio.start_server(
relay, os.environ["LITELLM_FORWARD_HOST"], int(os.environ["LITELLM_FORWARD_PORT"])
)
await stop.wait()
server.close()
if relays:
await asyncio.wait(relays, timeout=float(os.environ["LITELLM_FORWARD_SHUTDOWN_DRAIN"]))
asyncio.run(main())
"""
SUPERVISE: Final = """
import sys
from pathlib import Path
from litellm.proxy.rust_forward.server import Listener, PythonServer, Server
from litellm.proxy.rust_forward.supervisor import supervise
workers, port, scripts = sys.argv[1], int(sys.argv[2]), Path(sys.argv[3])
server = PythonServer(Server.UVICORN, Listener.TCP, "127.0.0.1", port, 5, 30.0, "127.0.0.1")
result = supervise(server, (sys.executable, str(scripts / "serve.py"), workers, str(port)), scripts_dir=str(scripts))
raise SystemExit(result if isinstance(result, int) else 3)
"""
T = TypeVar("T")
EXIT_SECONDS: Final = 15
TIME_TOLERANCE_SECONDS: Final = 5
class Scenario(Enum):
TERM = "SIGTERM"
INT = "SIGINT"
INT_TWICE_WHILE_STREAMING = "SIGINT twice during an open stream"
HUP = "SIGHUP"
@dataclass(frozen=True, slots=True)
class Outcome:
status: int | None
survived_hup: bool
@dataclass(frozen=True, slots=True)
class Run:
outcome: Outcome
seconds: float
def _scripts(directory: Path) -> Path:
(directory / "app.py").write_text(APP)
(directory / "serve.py").write_text(SERVE)
(directory / "supervise.py").write_text(SUPERVISE)
gateway: Final = directory / "litellm-gateway"
gateway.write_text(f"#!{sys.executable}\n{GATEWAY}")
gateway.chmod(0o755)
return directory
def _serving(port: int) -> bool:
try:
return httpx.get(f"http://127.0.0.1:{port}/", timeout=1, trust_env=False).text == "ok"
except httpx.TransportError:
return False
def _exit_status(process: subprocess.Popen[bytes], seconds: float) -> int | None:
try:
return process.wait(timeout=seconds)
except subprocess.TimeoutExpired:
return None
def _streaming(port: int) -> socket.socket:
stream: Final = socket.create_connection(("127.0.0.1", port), timeout=10)
stream.sendall(b"GET /stream HTTP/1.1\r\nHost: test\r\n\r\n")
assert eventually(lambda: stream.recv(4096), lambda received: b"tick" in received, seconds=10)
return stream
def _signal_after_ready(process: subprocess.Popen[bytes], port: int, scenario: Scenario) -> Run:
assert eventually(lambda: _serving(port), bool, seconds=120)
started: Final = time.monotonic()
match scenario:
case Scenario.TERM | Scenario.INT:
send_signal(process, signal.SIGTERM if scenario is Scenario.TERM else signal.SIGINT)
return Run(Outcome(_exit_status(process, EXIT_SECONDS), False), time.monotonic() - started)
case Scenario.INT_TWICE_WHILE_STREAMING:
with _streaming(port):
send_signal(process, signal.SIGINT)
time.sleep(0.5)
send_signal(process, signal.SIGINT)
status: Final = _exit_status(process, EXIT_SECONDS)
return Run(Outcome(status, False), time.monotonic() - started)
case Scenario.HUP:
send_signal(process, signal.SIGHUP)
after_hup: Final = _exit_status(process, 3)
if after_hup is not None:
return Run(Outcome(after_hup, False), time.monotonic() - started)
assert eventually(lambda: _serving(port), bool, seconds=30)
send_signal(process, signal.SIGTERM)
return Run(Outcome(_exit_status(process, EXIT_SECONDS), True), time.monotonic() - started)
def _run(argv: tuple[str, ...], environment: dict[str, str], act: Callable[[subprocess.Popen[bytes]], T]) -> T:
with subprocess.Popen(argv, env=environment, start_new_session=True) as process:
try:
return act(process)
finally:
if process.poll() is None:
for leftover in (*psutil.Process(process.pid).children(recursive=True), psutil.Process(process.pid)):
kill_stray_process(leftover.pid)
def _environment(markers: Path) -> dict[str, str]:
return {**os.environ, "MARKERS": str(markers), "PYTHONUNBUFFERED": "1"}
@pytest.mark.parametrize("workers", [1, 2])
@pytest.mark.parametrize("scenario", list(Scenario), ids=[scenario.name for scenario in Scenario])
def test_the_supervisor_ends_as_python_alone_would(tmp_path: Path, workers: int, scenario: Scenario) -> None:
scripts: Final = _scripts(tmp_path)
alone_port: Final = allocate_port()
alone: Final = _run(
(sys.executable, str(scripts / "serve.py"), str(workers), str(alone_port)),
_environment(tmp_path),
lambda process: _signal_after_ready(process, alone_port, scenario),
)
supervised_port: Final = allocate_port()
supervised: Final = _run(
(sys.executable, str(scripts / "supervise.py"), str(workers), str(supervised_port), str(scripts)),
_environment(tmp_path),
lambda process: _signal_after_ready(process, supervised_port, scenario),
)
assert supervised.outcome == alone.outcome
assert abs(supervised.seconds - alone.seconds) < TIME_TOLERANCE_SECONDS, (alone, supervised)
def _alive(pid: int) -> bool:
try:
return psutil.Process(pid).status() != psutil.STATUS_ZOMBIE
except psutil.NoSuchProcess:
return False
@pytest.mark.parametrize("workers", [1, 2])
def test_children_exit_when_the_supervisor_is_killed(tmp_path: Path, workers: int) -> None:
scripts: Final = _scripts(tmp_path)
port: Final = allocate_port()
def survivors_of_killing_the_supervisor(process: subprocess.Popen[bytes]) -> tuple[int, ...]:
assert eventually(lambda: _serving(port), bool, seconds=120)
python: Final = int((tmp_path / "python.pid").read_text())
rust: Final = int((tmp_path / "rust.pid").read_text())
descendants: Final = (python, rust, *(child.pid for child in psutil.Process(python).children(recursive=True)))
kill_stray_process(process.pid)
return eventually(lambda: tuple(pid for pid in descendants if _alive(pid)), lambda alive: not alive)
survivors: Final = _run(
(sys.executable, str(scripts / "supervise.py"), str(workers), str(port), str(scripts)),
_environment(tmp_path),
survivors_of_killing_the_supervisor,
)
assert survivors == ()

View file

@ -0,0 +1,387 @@
import asyncio
import json
import os
import socket
import threading
import time
from collections.abc import AsyncIterator, Callable, Iterator, Mapping
from contextlib import AbstractContextManager, ExitStack
from dataclasses import dataclass
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.mcp import mcp_peer
from tests.integration._support.rust_forward import MASTER_KEY, EntryPoint, forwarded_proxy
from tests.integration._support.wire import Reply, Request, wire_server
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
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 _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]:
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
)
proxy: Final = stack.enter_context(
forwarded_proxy(
EntryPoint.LITELLM,
{
"model_list": [
_deployment(CHAT_MODEL, "openai/gpt-4o-mini", f"{provider.url}/v1"),
{
**_deployment(REALTIME_MODEL, "openai/gpt-realtime", provider.url),
"model_info": {"mode": "realtime"},
},
_deployment(FILES_MODEL, "openai/gpt-4o-mini", f"{files_upstream}/v1"),
],
"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"],
},
},
tmp_path_factory.mktemp("rust-forward-traffic"),
("--keepalive_timeout", str(KEEPALIVE_SECONDS)),
)
)
stack.callback(released.set)
yield RustProxy(proxy.url, proxy.port, provider)
def _deployment(name: str, model: str, api_base: str) -> dict[str, JsonValue]:
return {
"model_name": name,
"litellm_params": {"model": model, "api_base": api_base, "api_key": "integration-provider-key"},
}
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

View file

@ -15,10 +15,19 @@ import pytest
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 (
Mode,
RustForwardMode,
Supervised,
uvicorn_flags,
)
from litellm.proxy.rust_forward.server import NO_TRUSTED_PROXIES, Listener
from litellm.proxy.server_listener import AssignedPort, InheritedFd
from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
DB_ENV: Final = {
"DATABASE_HOST": "db.internal",
@ -114,6 +123,91 @@ 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-certfile", "unused"), Listener.TLS),
(("--ssl-ciphers", "TLSv1.2"), Listener.TCP),
],
)
def test_forwarding_sees_the_public_listener_uvicorn_would_open(
self, arguments: tuple[str, ...], listener: Listener
) -> None:
assert python_server(arguments).listener is listener
@pytest.mark.parametrize(
("argv", "environ", "trust"),
[
((), {}, "127.0.0.1"),
((), {"FORWARDED_ALLOW_IPS": "10.9.9.9"}, "10.9.9.9"),
(("--forwarded-allow-ips", "10.0.0.5"), {"FORWARDED_ALLOW_IPS": "10.9.9.9"}, "10.0.0.5"),
(("--forwarded-allow-ips=*",), {}, "*"),
(("--no-proxy-headers", "--forwarded-allow-ips=*"), {}, NO_TRUSTED_PROXIES),
],
)
def test_forwarding_carries_the_trust_uvicorn_resolves_to_the_child(
self,
argv: tuple[str, ...],
environ: dict[str, str],
trust: str,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.delenv("FORWARDED_ALLOW_IPS", raising=False)
for name, value in environ.items():
monkeypatch.setenv(name, value)
server: Final = python_server(("--host", "0.0.0.0", "--port", "4000", *argv))
assert (server.host, server.port) == ("0.0.0.0", 4000)
assert 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
@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
@pytest.mark.parametrize(
("mode", "served"),
[
(Mode.DISABLED, {"host": "0.0.0.0", "port": 4000, "fd": None, "proxy_headers": True}),
(
Supervised(InheritedFd(17), "*"),
{"host": "0.0.0.0", "port": 4000, "fd": 17, "proxy_headers": False},
),
(
Supervised(AssignedPort(45556), "*"),
{"host": "127.0.0.1", "port": 45556, "fd": None, "proxy_headers": False},
),
],
)
def test_supervised_uvicorn_serves_the_private_hop_and_never_trusts_it(
self, mode: RustForwardMode, served: dict[str, object]
) -> None:
argv: Final = ("--host", "0.0.0.0", "--port", "4000", "--workers", "2", "--proxy-headers")
params: Final = _uvicorn_params((*uvicorn_argv((*argv, "--forwarded-allow-ips=*"), {}), *uvicorn_flags(mode)))
assert {name: params[name] for name in served} == served
assert (params["workers"], params["forwarded_allow_ips"]) == (2, "*")
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 +323,70 @@ class TestMain:
DatabaseURLSettings.from_env().apply_to_env()
assert os.environ["DATABASE_URL"] == pooled
@pytest.mark.parametrize(
("rust", "warning"),
[
("false", ""),
(
"true",
"LITELLM_RUST_FIRST does not support a Unix socket on the public port, "
"so Python serves the public port without Rust forwarding\n",
),
],
)
def test_python_alone_keeps_uvicorn_socket_options(
self,
rust: str,
warning: str,
password_env: dict[str, str],
monkeypatch: pytest.MonkeyPatch,
capsys: pytest.CaptureFixture[str],
) -> None:
monkeypatch.setenv("LITELLM_RUST_FIRST", rust)
monkeypatch.delenv("_LITELLM_GATEWAY_SUPERVISED", 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"
assert capsys.readouterr().err == warning
def test_an_invalid_rust_switch_stops_the_gateway_naming_the_variable(
self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("LITELLM_RUST_FIRST", "maybe")
monkeypatch.delenv("_LITELLM_GATEWAY_SUPERVISED", raising=False)
served: Final[list[tuple[str, ...]]] = []
with pytest.raises(SystemExit) as stopped:
main(("--workers", "4"), serve=lambda argv: served.append(tuple(argv)))
assert str(stopped.value).startswith("LITELLM_RUST_FIRST is invalid: ")
assert served == []
@pytest.mark.parametrize(
("argv", "status", "said"),
[
pytest.param(("--nope",), 2, "--nope", id="a bad flag is a usage error"),
pytest.param(("--help",), 0, "--timeout-keep-alive", id="help is printed once and exits cleanly"),
],
)
def test_the_supervisor_reports_uvicorn_arguments_the_way_uvicorn_does(
self,
argv: tuple[str, ...],
status: int,
said: str,
password_env: dict[str, str],
monkeypatch: pytest.MonkeyPatch,
capsys: pytest.CaptureFixture[str],
) -> None:
monkeypatch.setenv("LITELLM_RUST_FIRST", "true")
monkeypatch.delenv("_LITELLM_GATEWAY_SUPERVISED", raising=False)
served: Final[list[tuple[str, ...]]] = []
with pytest.raises(SystemExit) as stopped:
main(argv, serve=lambda served_argv: served.append(tuple(served_argv)))
printed: Final = capsys.readouterr()
assert stopped.value.code == status
assert (printed.out + printed.err).count(said) == 1
assert served == []
def test_a_pooler_that_cannot_start_stops_the_gateway_before_uvicorn(
self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path
):

View file

@ -0,0 +1,7 @@
from litellm.proxy.rust_forward.server import Listener, PythonServer, Server
def python_server(
server: Server = Server.UVICORN, listener: Listener = Listener.TCP, trusted_proxies: str = "127.0.0.1"
) -> PythonServer:
return PythonServer(server, listener, "0.0.0.0", 4000, 2, 2.5, trusted_proxies)

View file

@ -0,0 +1,54 @@
import socket
from typing import Final
from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS
from litellm.proxy.proxy_cli import run_server
from litellm.proxy.rust_forward.mode import Supervised, rust_forward_mode
from litellm.proxy.rust_forward.platform import PlatformListener
from litellm.proxy.rust_forward.platform.posix import PosixPlatform
from litellm.proxy.rust_forward.platform.windows import CREATE_NEW_PROCESS_GROUP, WindowsPlatform
from litellm.proxy.rust_forward.supervisor.children import python_spec, rust_spec
from litellm.proxy.server_listener import AssignedPort, InheritedFd
from tests.unit.proxy.rust_forward.supervisor._support import python_server
POSIX: Final = PosixPlatform()
WINDOWS: Final = WindowsPlatform()
def test_python_child_serves_the_private_socket_with_the_operators_trust() -> None:
private: Final = PlatformListener(InheritedFd(17), 45555, (17,))
spec: Final = python_spec(POSIX, ("python",), private, python_server(trusted_proxies="10.0.0.5"))
assert rust_forward_mode(spec.env_overrides) == Supervised(InheritedFd(17), "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, python_server())
assert rust_forward_mode(spec.env_overrides) == Supervised(AssignedPort(private.port), "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
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, PlatformListener(InheritedFd(17), 45555, (17,)), python_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, python_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),
}

View file

@ -0,0 +1,158 @@
import functools
import itertools
import signal
from dataclasses import dataclass, replace
from typing import Final
import pytest
from hypothesis import given
from hypothesis import strategies as st
from litellm.proxy.rust_forward.platform.posix import PosixPlatform
from litellm.proxy.rust_forward.supervisor.lifecycle import (
FAILURE,
Action,
Caught,
Child,
ChildExited,
Event,
Lifecycle,
Probe,
Send,
Spawn,
boot,
exit_status,
step,
)
PY: Final = Child.PYTHON
RUST: Final = Child.RUST
LISTENING: Final = Probe.PYTHON_LISTENING
TERM: Final = signal.SIGTERM
INT: Final = signal.SIGINT
HUP: Final = signal.SIGHUP
STOPS: Final = PosixPlatform().stop_signals
def _replay(events: tuple[Event, ...]) -> tuple[tuple[tuple[Action, ...], ...], int | None]:
transitions: Final = tuple(
itertools.accumulate(events, lambda transition, event: step(transition[0], event, STOPS), initial=boot())
)
return tuple(actions for _, actions in transitions[1:]), exit_status(transitions[-1][0])
@pytest.mark.parametrize(
("trace", "status"),
[
pytest.param(
(
(LISTENING, (Spawn(RUST),)),
(Caught(TERM), (Send(RUST, TERM), Send(PY, TERM))),
(ChildExited(RUST, 0), ()),
(ChildExited(PY, -15), ()),
),
-15,
id="TERM reaches Rust first, then Python, and the supervisor dies as Python did",
),
pytest.param(
(
(LISTENING, (Spawn(RUST),)),
(Caught(INT), (Send(RUST, INT), Send(PY, INT))),
(Caught(INT), (Send(RUST, INT), Send(PY, INT))),
(ChildExited(PY, 1), (Send(RUST, TERM),)),
(ChildExited(RUST, 0), ()),
),
1,
id="a second Ctrl-C reaches Python again to force the exit",
),
pytest.param(
(
(LISTENING, (Spawn(RUST),)),
(Caught(HUP), (Send(PY, HUP),)),
(Caught(signal.SIGTTIN), (Send(PY, signal.SIGTTIN),)),
(ChildExited(PY, 0), (Send(RUST, TERM),)),
(ChildExited(RUST, 0), ()),
),
0,
id="gunicorn's HUP and TTIN reach only Python and Rust keeps serving",
),
pytest.param(
(
(LISTENING, (Spawn(RUST),)),
(ChildExited(RUST, 0), (Send(PY, TERM),)),
(ChildExited(PY, -15), ()),
),
FAILURE,
id="losing Rust while Python serves stops Python and fails",
),
pytest.param(
(
(Caught(TERM), (Send(PY, TERM),)),
(LISTENING, ()),
(ChildExited(PY, -15), ()),
),
-15,
id="a stop before Python listens never starts Rust",
),
pytest.param(((ChildExited(PY, 3), ()),), 3, id="Python failing to boot exits with its status"),
],
)
def test_supervisor_relays_signals_and_exits_as_python_did(
trace: tuple[tuple[Event, tuple[Action, ...]], ...], status: int
) -> None:
events: Final = tuple(event for event, _ in trace)
assert _replay(events) == (tuple(actions for _, actions in trace), status)
@dataclass(frozen=True, slots=True)
class Observed:
"""What an outside observer of the two processes knows, without looking inside ``Lifecycle``."""
state: Lifecycle
alive: frozenset[Child] = frozenset({PY})
rust_spawns: int = 0
stop_requested: bool = False
rust_lost: bool = False
python_status: int | None = None
def _observe(seen: Observed, event: Event) -> Observed:
state, actions = step(seen.state, event, STOPS)
death: Final = event if isinstance(event, ChildExited) and event.child in seen.alive else None
exited: Final = None if death is None else death.child
running: Final = seen.alive - {exited}
assert all(action.child in running for action in actions if isinstance(action, Send)), "signalled a dead child"
if isinstance(event, Caught):
to_rust: Final = (Send(RUST, event.signum),) if event.signum in STOPS and RUST in running else ()
assert actions == (*to_rust, *((Send(PY, event.signum),) if PY in running else ()))
spawned: Final = Spawn(RUST) in actions
assert not spawned or (event is LISTENING and not seen.stop_requested), "started Rust out of turn"
stop_requested: Final = seen.stop_requested or (isinstance(event, Caught) and event.signum in STOPS)
after: Final = replace(
seen,
state=state,
alive=running | ({RUST} if spawned else frozenset()),
rust_spawns=seen.rust_spawns + int(spawned),
stop_requested=stop_requested,
rust_lost=seen.rust_lost or (exited is RUST and PY in running and not stop_requested),
python_status=death.status if death is not None and exited is PY else seen.python_status,
)
settled: Final = None if after.alive or after.python_status is None else after.python_status
assert exit_status(state) == (FAILURE if after.rust_lost and settled is not None else settled)
return after
EVENTS: Final = st.lists(
st.one_of(
st.sampled_from(PosixPlatform().relayed_signals).map(Caught),
st.just(LISTENING),
st.builds(ChildExited, st.sampled_from(Child), st.integers(-15, 3)),
),
max_size=40,
)
@given(EVENTS)
def test_any_event_order_keeps_the_supervisor_transparent_to_python(events: list[Event]) -> None:
final: Final = functools.reduce(_observe, events, Observed(boot()[0]))
assert final.rust_spawns <= 1

View file

@ -0,0 +1,77 @@
import sys
from pathlib import Path
from typing import Final
import pytest
from litellm.proxy.rust_forward.platform import Platform
from litellm.proxy.rust_forward.platform.posix import PosixPlatform
from litellm.proxy.rust_forward.platform.windows import WindowsPlatform
from litellm.proxy.rust_forward.server import Listener, PythonServer, Server
from litellm.proxy.rust_forward.supervisor import supervise
from litellm.proxy.rust_forward.supervisor.preflight import StartupError, gateway_binary, startup_check
from tests.unit.proxy.rust_forward.supervisor._support import python_server
POSIX: Final = PosixPlatform()
WINDOWS: Final = WindowsPlatform()
POSIX_ONLY: Final = pytest.mark.skipif(sys.platform == "win32", reason="runs a POSIX shebang gateway")
@pytest.mark.parametrize(
("platform", "server", "refusal"),
[
(POSIX, python_server(), None),
(POSIX, python_server(Server.GUNICORN), None),
(WINDOWS, python_server(), None),
(WINDOWS, python_server(Server.GUNICORN), "does not support gunicorn on Windows"),
(POSIX, python_server(Server.HYPERCORN), "does not support Hypercorn"),
(POSIX, python_server(Server.GRANIAN), "does not support Granian"),
(POSIX, python_server(listener=Listener.TLS), "does not support TLS on the public port"),
(POSIX, python_server(listener=Listener.UNIX_SOCKET), "does not support a Unix socket on the public port"),
(
POSIX,
python_server(Server.GUNICORN, Listener.INHERITED_FD),
"does not support an inherited file descriptor on the public port",
),
],
)
def test_only_plain_tcp_on_a_supported_server_can_be_forwarded(
platform: Platform, server: PythonServer, refusal: str | None
) -> None:
expected: Final = "gateway" if refusal is None else StartupError(f"LITELLM_RUST_FIRST {refusal}")
assert startup_check(platform, server, "gateway") == expected
@pytest.mark.parametrize(
("server", "refusal"),
[
(python_server(Server.HYPERCORN), "does not support Hypercorn"),
(python_server(), "requires the litellm-gateway binary from the litellm-proxy package"),
],
)
def test_a_refused_start_says_why_in_one_line_and_starts_nothing(
server: PythonServer, refusal: str, tmp_path: Path, capsys: pytest.CaptureFixture[str]
) -> None:
prepared: Final[list[bool]] = []
refused: Final = supervise(
server, ("python",), before_start=lambda: prepared.append(True), scripts_dir=str(tmp_path), platform=POSIX
)
assert (refused, prepared) == (StartupError(f"LITELLM_RUST_FIRST {refusal}"), [])
assert (
capsys.readouterr().err
== f"LITELLM_RUST_FIRST {refusal}, so Python serves the public port without Rust forwarding\n"
)
@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)

View file

@ -0,0 +1,189 @@
import signal
import subprocess
import sys
from dataclasses import dataclass, field
from pathlib import Path
from types import MappingProxyType
from typing import Final
import anyio
import pytest
from litellm.proxy.rust_forward.platform import Platform, current_platform
from litellm.proxy.rust_forward.platform.windows import WindowsPlatform
from litellm.proxy.rust_forward.supervisor import supervise
from litellm.proxy.rust_forward.supervisor.children import ChildSpec
from litellm.proxy.rust_forward.supervisor.lifecycle import (
Action,
Caught,
Child,
ChildExited,
Event,
Probe,
Send,
Spawn,
)
from litellm.proxy.rust_forward.supervisor.runtime import SpawnFailed, drive, exit_like, open_children
from tests.unit.proxy.rust_forward.supervisor._support import python_server
PYTHON_CHILD: Final = """
import json, os, pathlib, socket, time
time.sleep(0.5)
pathlib.Path(os.environ["MARKERS"], "listening").touch()
private = json.loads(os.environ["_LITELLM_GATEWAY_SUPERVISED"])["listener"]
listener = socket.socket(fileno=private["fd"]) if "fd" in private else socket.create_server(("127.0.0.1", private["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(python_server(), command, scripts_dir=str(tmp_path), platform=_fake_gateway(tmp_path))
assert (tmp_path / "rust_saw_python_listening").read_text() == "True"
TERM: Final = signal.SIGTERM
POSIX_ONLY: Final = pytest.mark.skipif(sys.platform == "win32", reason="the child traps SIGTERM")
STOPS: Final = frozenset({signal.SIGTERM, signal.SIGINT, signal.SIGQUIT})
RUST_MISSING: Final = SpawnFailed(FileNotFoundError("litellm-gateway"))
@dataclass
class ScriptedExecutor:
"""Replays a fixed event script and records every action the driver carries out."""
events: list[Event]
failing: frozenset[Child]
done: list[Action] = field(default_factory=list)
async def spawn(self, child: Child) -> SpawnFailed | None:
self.done.append(Spawn(child))
return RUST_MISSING if child in self.failing else None
def send(self, child: Child, signum: signal.Signals) -> None:
self.done.append(Send(child, signum))
async def next_event(self) -> Event:
return self.events.pop(0)
@pytest.mark.parametrize(
("events", "failing", "done", "outcome", "unread"),
[
pytest.param(
[Probe.PYTHON_LISTENING, Caught(TERM), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, -15)],
frozenset(),
[Spawn(Child.PYTHON), Spawn(Child.RUST), Send(Child.RUST, TERM), Send(Child.PYTHON, TERM)],
-15,
[],
id="carries out each action in order and ends with Python's status",
),
pytest.param(
[ChildExited(Child.PYTHON, 3), Caught(TERM)],
frozenset(),
[Spawn(Child.PYTHON)],
3,
[Caught(TERM)],
id="reads no event once the status is settled",
),
pytest.param(
[Probe.PYTHON_LISTENING, Caught(TERM)],
frozenset({Child.RUST}),
[Spawn(Child.PYTHON), Spawn(Child.RUST)],
RUST_MISSING,
[Caught(TERM)],
id="returns a failed spawn without waiting for more events",
),
],
)
async def test_driver_executes_lifecycle_decisions(
events: list[Event], failing: frozenset[Child], done: list[Action], outcome: int | SpawnFailed, unread: list[Event]
) -> None:
executor: Final = ScriptedExecutor(events, failing)
assert await drive(executor, STOPS) == outcome
assert (executor.done, executor.events) == (done, unread)
def _children(argv: tuple[str, ...]) -> MappingProxyType[Child, ChildSpec]:
spec: Final = ChildSpec(argv, MappingProxyType({}), current_platform().spawn_options(()))
return MappingProxyType({Child.PYTHON: spec, Child.RUST: spec})
SLOW_TO_STOP: Final = (
sys.executable,
"-c",
"import signal, time; signal.signal(signal.SIGTERM, lambda *_: (time.sleep(0.5), exit(0))); time.sleep(30)",
)
@POSIX_ONLY
async def test_a_cancelled_supervisor_still_waits_for_the_child_it_stops() -> None:
with anyio.CancelScope() as scope:
async with open_children(current_platform(), _children(SLOW_TO_STOP), 1) as children:
assert await children.spawn(Child.RUST) is None
process: Final = children.processes[Child.RUST]
await anyio.sleep(0.3)
scope.cancel()
await anyio.sleep(1)
assert process.returncode == 0
async def test_exiting_children_are_reported_as_events() -> None:
async with open_children(current_platform(), _children((sys.executable, "-c", "exit(7)")), 1) as children:
assert await children.spawn(Child.RUST) is None
with anyio.fail_after(5):
assert await children.next_event() == ChildExited(Child.RUST, 7)
async def test_a_child_that_cannot_start_is_returned_as_a_failure(tmp_path: Path) -> None:
async with open_children(current_platform(), _children((str(tmp_path / "missing"),)), 1) as children:
outcome: Final = await children.spawn(Child.RUST)
assert isinstance(outcome, SpawnFailed)
assert isinstance(outcome.error, FileNotFoundError)
def _dropped(_signum: signal.Signals) -> None:
pass
@POSIX_ONLY
@pytest.mark.parametrize(
("status", "exits"),
[
pytest.param(3, 3, id="an exit status is passed on"),
pytest.param(-signal.SIGTERM, 0, id="a stop signal ends it the way uvicorn alone ends as PID 1"),
pytest.param(-signal.SIGKILL, 137, id="a kill is reported the way a shell reports it"),
],
)
def test_a_supervisor_the_kernel_will_not_signal_still_exits_as_python_alone_would(status: int, exits: int) -> None:
assert exit_like(status, STOPS, die_by=_dropped) == exits
@POSIX_ONLY
@pytest.mark.parametrize("signum", [signal.SIGTERM, signal.SIGKILL])
def test_the_supervisor_dies_by_the_signal_that_killed_python(signum: signal.Signals) -> None:
script: Final = (
"from litellm.proxy.rust_forward.supervisor.runtime import exit_like\n"
f"raise SystemExit(exit_like({-signum}, frozenset()))"
)
assert subprocess.run((sys.executable, "-I", "-c", script), check=False, timeout=60).returncode == -signum

View file

@ -0,0 +1,127 @@
import os
import threading
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 (
SUPERVISED_ENV,
Mode,
RustForwardMode,
Supervised,
rust_forward_mode,
private_listener,
stop_with_supervisor,
supervised_environment,
uvicorn_options,
)
from litellm.proxy.rust_forward.server import NO_TRUSTED_PROXIES
from litellm.proxy.server_listener import AssignedPort, InheritedFd, PrivateListener
@pytest.mark.parametrize(
("value", "expected"),
[
*((value, Mode.DISABLED) for value in (None, "", " ", "0", "false", "off")),
*((value, Mode.SUPERVISOR) for value in ("1", "true", "yes", "on", " 1 ")),
],
)
def test_mode_is_resolved_from_environment(value: str | None, expected: RustForwardMode) -> None:
assert rust_forward_mode({} if value is None else {"LITELLM_RUST_FIRST": value}) == expected
@pytest.mark.parametrize(
("environ", "expected"),
[
({"LITELLM_RUST": "true"}, Mode.DISABLED),
({"LITELLM_RUST": "enabled"}, Mode.DISABLED),
({"LITELLM_RUST": "false", "LITELLM_RUST_FIRST": "true"}, Mode.SUPERVISOR),
],
)
def test_forwarding_is_independent_of_the_sdk_rust_switch(environ: dict[str, str], expected: RustForwardMode) -> None:
assert rust_forward_mode(environ) == expected
@pytest.mark.parametrize(
("environ", "invalid"),
[
({"LITELLM_RUST_FIRST": "enabled"}, "LITELLM_RUST_FIRST"),
({"LITELLM_RUST_FIRST": "2"}, "LITELLM_RUST_FIRST"),
({SUPERVISED_ENV: '{"listener": {"fd": "seventeen"}, "trusted_proxies": "*"}'}, "fd"),
({SUPERVISED_ENV: '{"listener": {"fd": 17}, "trusted_proxies": 7}'}, "trusted_proxies"),
({SUPERVISED_ENV: '{"listener": {"fd": 17}}'}, "trusted_proxies"),
],
)
def test_invalid_forwarding_value_stops_startup_naming_it(environ: dict[str, str], invalid: str) -> None:
with pytest.raises(ValidationError, match=invalid):
rust_forward_mode(environ)
@pytest.mark.parametrize(
"written",
[
Supervised(InheritedFd(17), "127.0.0.1"),
Supervised(InheritedFd(17), "127.0.0.1,10.0.0.5"),
Supervised(InheritedFd(17), NO_TRUSTED_PROXIES),
Supervised(AssignedPort(45555), "*"),
],
)
def test_child_reads_back_exactly_what_the_supervisor_wrote(written: Supervised) -> None:
inherited: Final = {"LITELLM_RUST_FIRST": "true", "FORWARDED_ALLOW_IPS": "10.9.9.9"}
assert rust_forward_mode({**inherited, **supervised_environment(written)}) == written
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)
PUBLIC: Final = ("0.0.0.0", 45555, None)
@pytest.mark.parametrize(
("mode", "app", "uvicorn_listener", "gunicorn_listener", "proxy_headers"),
[
(Mode.DISABLED, "litellm.proxy.proxy_server:app", PUBLIC, None, True),
(Mode.SUPERVISOR, "litellm.proxy.proxy_server:app", PUBLIC, None, True),
(Supervised(InheritedFd(17), "*"), WORKER_APP, ("0.0.0.0", 45555, 17), InheritedFd(17), False),
(Supervised(AssignedPort(45556), "*"), WORKER_APP, ("127.0.0.1", 45556, None), AssignedPort(45556), False),
],
)
def test_supervised_server_serves_the_private_hop_and_never_trusts_it(
mode: RustForwardMode,
app: str,
uvicorn_listener: tuple[str, int, int | None],
gunicorn_listener: PrivateListener | None,
proxy_headers: bool,
) -> None:
config: Final = _served_uvicorn_config(mode)
assert (config.app, (config.host, config.port, config.fd), config.proxy_headers) == (
app,
uvicorn_listener,
proxy_headers,
)
assert config.factory is (app == WORKER_APP)
assert private_listener(mode) == gunicorn_listener
@pytest.mark.parametrize(
("mode", "stops"),
[(Supervised(InheritedFd(17), "127.0.0.1"), True), (Mode.DISABLED, False), (Mode.SUPERVISOR, False)],
)
def test_only_a_supervised_server_stops_once_its_supervisor_is_gone(mode: RustForwardMode, stops: bool) -> None:
read_end, write_end = os.pipe()
stopped: Final = threading.Event()
try:
stop_with_supervisor(mode, read_end, stopped.set)
assert not stopped.wait(0.2)
os.close(write_end)
assert stopped.wait(5 if stops else 0.5) is stops
finally:
os.close(read_end)

View file

@ -0,0 +1,135 @@
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 SUPERVISED_ENV, Supervised, supervised_environment
from litellm.proxy.rust_forward.peer import RustPeerMiddleware
from litellm.proxy.rust_forward.server import NO_TRUSTED_PROXIES
from litellm.proxy.server_listener import InheritedFd
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"])))
proxy_headers_off: Final = mode.trusted_proxies == NO_TRUSTED_PROXIES
direct: Final = probe if proxy_headers_off else ProxyHeadersMiddleware(probe, mode.trusted_proxies)
app: Final = RustPeerMiddleware(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), "127.0.0.1"),
Supervised(InheritedFd(17), "10.0.0.0/8"),
Supervised(InheritedFd(17), "*"),
Supervised(InheritedFd(17), NO_TRUSTED_PROXIES),
],
)
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_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), "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(SUPERVISED_ENV, supervised_environment(Supervised(InheritedFd(17), "127.0.0.1"))[SUPERVISED_ENV])
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}

View file

@ -26,6 +26,7 @@ import urllib.parse as urlparse
import uvicorn
import yaml
from gunicorn.config import Config as GunicornConfig
from uvicorn.config import LOOP_FACTORIES
from uvicorn.importer import import_from_string
@ -34,8 +35,9 @@ from litellm._logging import turn_on_json
from litellm.constants import DEFAULT_MAX_REQUEST_HEADER_SIZE, GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS
from litellm.proxy import proxy_cli
from litellm.proxy.proxy_cli import ProxyInitializationHelpers, run_server
from litellm.proxy.rust_forward.server import Listener, Server
from litellm.proxy.server_listener import AssignedPort, InheritedFd
from litellm.proxy.shutdown.graceful_shutdown_manager import DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT
from litellm.proxy.shutdown.graceful_shutdown_manager import DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT, GracefulShutdownManager
@pytest.fixture(autouse=True)
@ -4382,3 +4384,154 @@ 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=None not in (params["ssl_keyfile_path"], params["ssl_certfile_path"]),
)
@pytest.mark.parametrize(
("argv", "server", "listener"),
[
([], Server.UVICORN, Listener.TCP),
(["--run_gunicorn"], Server.GUNICORN, Listener.TCP),
(["--run_hypercorn"], Server.HYPERCORN, Listener.TCP),
(["--run_granian"], Server.GRANIAN, Listener.TCP),
(["--ssl_keyfile_path", "key.pem", "--ssl_certfile_path", "cert.pem"], Server.UVICORN, Listener.TLS),
],
)
def test_the_supervisor_sees_the_server_and_listener_python_would_use(
self, argv: list[str], server: Server, listener: Listener
) -> None:
_, described = self._server(argv)
assert (described.server, described.listener) == (server, listener)
def test_uvicorn_is_described_where_gunicorn_cannot_be_imported(self, monkeypatch: pytest.MonkeyPatch) -> None:
_, with_gunicorn = self._server([])
monkeypatch.setitem(sys.modules, "gunicorn.config", None)
_, without_gunicorn = self._server([])
assert without_gunicorn == with_gunicorn
@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)
@pytest.mark.parametrize(
("argv", "reason", "served_keyfile"),
[
(
["--ssl_keyfile_path", "key.pem", "--ssl_certfile_path", "cert.pem"],
"does not support TLS on the public port",
"key.pem",
),
(["--ssl_keyfile_path", "key.pem"], "requires the litellm-gateway binary", None),
(["--ssl_certfile_path", "cert.pem"], "requires the litellm-gateway binary", None),
(["--ciphers", "ECDHE"], "requires the litellm-gateway binary", None),
],
)
def test_a_refused_forwarding_start_serves_from_python_alone(
self,
mock_should_update,
mock_setup_db,
mock_atexit_register,
mock_uvicorn_run,
argv: list[str],
reason: str,
served_keyfile: str | None,
) -> None:
from click.testing import CliRunner
proxy_module: Final = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
KeyManagementSettings=MagicMock(),
save_worker_config=MagicMock(),
)
environ: Final = {
**{
k: v
for k, v in os.environ.items()
if k not in ("DATABASE_URL", "DIRECT_URL", "_LITELLM_GATEWAY_SUPERVISED")
},
"LITELLM_RUST_FIRST": "true",
}
with (
patch.dict(os.environ, environ, clear=True),
patch.dict("sys.modules", {"proxy_server": proxy_module, "litellm.proxy.proxy_server": proxy_module}),
patch("litellm.proxy.rust_forward.supervisor.gateway_binary", return_value=None),
):
result: Final = CliRunner().invoke(run_server, ["--port", "45999", *argv])
assert result.exit_code == 0, result.output
assert f"LITELLM_RUST_FIRST {reason}" in result.output
assert "so Python serves the public port without Rust forwarding" in result.output
mock_uvicorn_run.assert_called_once()
assert mock_uvicorn_run.call_args.kwargs.get("ssl_keyfile") == served_keyfile
assert mock_uvicorn_run.call_args.kwargs["port"] == 45999
def test_an_invalid_rust_switch_stops_startup_with_one_line_naming_it(self) -> None:
from click.testing import CliRunner
with patch.dict(os.environ, {"LITELLM_RUST_FIRST": "sometimes"}):
result: Final = CliRunner().invoke(run_server, ["--port", "45999"])
lines: Final = result.output.splitlines()
assert (result.exit_code, len(lines)) == (1, 1)
assert lines[0].startswith("Error: LITELLM_RUST_FIRST is invalid: ")
def test_a_metrics_port_on_the_public_port_is_refused_before_forwarding_starts(self) -> None:
from click.testing import CliRunner
with (
patch.dict(os.environ, {"LITELLM_RUST_FIRST": "true"}),
patch("litellm.proxy.rust_forward.supervisor.supervise") as supervise,
):
result: Final = CliRunner().invoke(run_server, ["--port", "45999", "--prometheus_metrics_port", "45999"])
assert result.exit_code == 2
assert "--prometheus_metrics_port must differ from --port" in result.output
supervise.assert_not_called()
@staticmethod
def _served_by_python(params: dict[str, object]) -> tuple[object, ...]:
if params["run_gunicorn"]:
gunicorn: Final = GunicornConfig()
return gunicorn.keepalive, GracefulShutdownManager.get_timeout(), ",".join(gunicorn.forwarded_allow_ips)
served: Final = uvicorn.Config(
**ProxyInitializationHelpers.get_default_unvicorn_init_args(
"127.0.0.1", 45555, keepalive_timeout=params["keepalive_timeout"]
)
)
return served.timeout_keep_alive, served.timeout_graceful_shutdown, served.forwarded_allow_ips
@pytest.mark.parametrize(
("argv", "environ"),
[
([], {}),
(["--keepalive_timeout", "2"], {"GRACEFUL_SHUTDOWN_TIMEOUT": "12.5"}),
([], {"KEEPALIVE_TIMEOUT": "7", "FORWARDED_ALLOW_IPS": "10.0.0.5"}),
(["--run_gunicorn", "--keepalive_timeout", "9"], {}),
],
)
def test_rust_follows_the_keepalive_drain_and_trust_python_serves_with(
self, argv: list[str], environ: dict[str, str], monkeypatch: pytest.MonkeyPatch
) -> None:
for name in ("KEEPALIVE_TIMEOUT", "GRACEFUL_SHUTDOWN_TIMEOUT", "FORWARDED_ALLOW_IPS"):
monkeypatch.delenv(name, raising=False)
for name, value in environ.items():
monkeypatch.setenv(name, value)
params, server = self._server(argv)
assert (
server.keepalive_seconds,
server.graceful_timeout_seconds,
server.trusted_proxies,
) == self._served_by_python(params)

View file

@ -0,0 +1,7 @@
# Windows tests
The `windows-sdk` job in `.github/workflows/test-merge-smoke.yml` is the only place these run, on a Windows runner with no proxy dependencies installed. They prove the SDK works on Windows without them, so a test here must not import uvicorn, fastapi, psutil or anything else outside the `dev` group, and must not need the `litellm-gateway.exe` binary
A test that cannot run off Windows carries a `skipif` on `sys.platform`, so collecting this folder on another OS never fails
Windows-only names such as `CTRL_BREAK_EVENT` come from `litellm/proxy/rust_forward/platform/windows.py`, which spells them out so the type checker accepts them on every OS. `test_rust_forward_platform.py` pins those constants to the stdlib

View file

@ -0,0 +1,35 @@
import signal
import subprocess
import sys
from pathlib import Path
from typing import Final
import pytest
from litellm.proxy.rust_forward.platform.windows import (
CREATE_NEW_PROCESS_GROUP,
CTRL_BREAK_EVENT,
SIGBREAK,
WindowsPlatform,
)
from litellm.proxy.rust_forward.supervisor.preflight import gateway_binary
@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
@pytest.mark.skipif(sys.platform != "win32", reason="only Windows searches the working directory for an executable")
def test_gateway_binary_in_the_working_directory_is_never_chosen(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
scripts: Final = tmp_path / "scripts"
scripts.mkdir()
(tmp_path / "litellm-gateway.exe").write_bytes(b"")
monkeypatch.chdir(tmp_path)
assert gateway_binary(WindowsPlatform(), str(scripts)) is None
installed: Final = scripts / "litellm-gateway.exe"
installed.write_bytes(b"")
assert gateway_binary(WindowsPlatform(), str(scripts)) == str(installed)