From cc4bce9cbed53a643b72f23a688f9c20ebb91f2f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 6 Nov 2026 13:34:34 +0900 Subject: [PATCH] refactor(proxy): lean on the stdlib, uvicorn and pydantic in Rust forwarding The supervisor runs its children through one contextlib.AsyncExitStack loop instead of recursive async with blocks, and payload-free states and modes are Enum members The Python child's env contract is the Supervised type itself, written and read through one pydantic TypeAdapter, so the trusted proxies it carries are always set. Each entry point resolves them from its server library, uvicorn's Config or gunicorn's, instead of a default copied into this package gateway.launch validates uvicorn's parsed CLI options with one pydantic model. TLS now means what each server serves TLS with: a key or cert for uvicorn's CLI, both for proxy_cli. A cipher list alone, or one half of the pair under proxy_cli, no longer refuses forwarding, since Python alone serves plain HTTP there --- gateway/launch.py | 46 ++----- litellm/proxy/proxy_cli.py | 10 +- litellm/proxy/rust_forward/AGENTS.md | 4 +- litellm/proxy/rust_forward/mode.py | 75 +++-------- litellm/proxy/rust_forward/peer.py | 4 +- litellm/proxy/rust_forward/server.py | 42 +++++- litellm/proxy/rust_forward/supervisor.py | 123 ++++++------------ .../test_rust_forward_lifecycle.py | 14 +- tests/unit/gateway/test_launch.py | 26 ++-- tests/unit/proxy/rust_forward/test_mode.py | 69 +++++----- tests/unit/proxy/rust_forward/test_peer.py | 6 +- .../proxy/rust_forward/test_supervisor.py | 45 ++++--- tests/unit/proxy/test_proxy_cli.py | 47 +++++-- 13 files changed, 252 insertions(+), 259 deletions(-) diff --git a/gateway/launch.py b/gateway/launch.py index 979c40984d9..1adb2ef1034 100644 --- a/gateway/launch.py +++ b/gateway/launch.py @@ -80,40 +80,23 @@ def _serve(argv: Sequence[str]) -> None: def python_server(argv: Sequence[str]) -> "PythonServer": """What uvicorn would serve the public port with, read from the same flags and environment it reads.""" - from pydantic import TypeAdapter - - from litellm.proxy.rust_forward.server import Listener, PythonServer, Server + from litellm.proxy.rust_forward.server import PythonServer, Server, UvicornCliParams, uvicorn_trusted_proxies from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager - params: Final[Mapping[str, object]] = uvicorn_main.make_context( - "uvicorn", list(uvicorn_argv(argv, os.environ)) - ).params - tls: Final = ( - params["ssl_keyfile"] is not None - or params["ssl_certfile"] is not None - or any(arg.split("=")[0] == "--ssl-ciphers" for arg in argv) + params: Final = UvicornCliParams.model_validate( + uvicorn_main.make_context("uvicorn", list(uvicorn_argv(argv, os.environ))).params ) - listener: Final = ( - Listener.TLS - if tls - else Listener.UNIX_SOCKET - if params["uds"] is not None - else Listener.INHERITED_FD - if params["fd"] is not None - else Listener.TCP - ) - graceful_timeout: Final = TypeAdapter[int | None](int | None).validate_python(params["timeout_graceful_shutdown"]) return PythonServer( server=Server.UVICORN, - listener=listener, - host=TypeAdapter(str).validate_python(params["host"]), - port=TypeAdapter(int).validate_python(params["port"]), - keepalive_seconds=TypeAdapter(int).validate_python(params["timeout_keep_alive"]), + listener=params.listener, + host=params.host, + port=params.port, + keepalive_seconds=params.timeout_keep_alive, graceful_timeout_seconds=GracefulShutdownManager.get_timeout() - if graceful_timeout is None - else graceful_timeout, - proxy_headers=TypeAdapter(bool).validate_python(params["proxy_headers"]), - trusted_proxies=TypeAdapter[str | None](str | None).validate_python(params["forwarded_allow_ips"]), + if params.timeout_graceful_shutdown is None + else params.timeout_graceful_shutdown, + proxy_headers=params.proxy_headers, + trusted_proxies=uvicorn_trusted_proxies(params.forwarded_allow_ips), ) @@ -132,9 +115,8 @@ def configure_database() -> None: def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None: from litellm.proxy.rust_forward.mode import ( - Disabled, + Mode, Supervised, - Supervisor, rust_forward_mode, stop_with_supervisor, uvicorn_flags, @@ -143,14 +125,14 @@ def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) - mode: Final = rust_forward_mode(os.environ) match mode: - case Supervisor() if not isinstance( + case Mode.SUPERVISOR if not isinstance( exit_code := supervise( python_server(argv), (sys.executable, "-m", "gateway.launch", *argv), before_start=configure_database ), StartupError, ): sys.exit(exit_code) - case Disabled() | Supervisor(): + case Mode.DISABLED | Mode.SUPERVISOR: configure_database() serve(uvicorn_argv(argv, os.environ)) case Supervised(): diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 0338b7261b9..632a71ec5c8 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -854,7 +854,7 @@ class ProxyInitializationHelpers: import uvicorn from gunicorn.config import Config as GunicornConfig - from litellm.proxy.rust_forward.server import Listener, PythonServer, Server + from litellm.proxy.rust_forward.server import Listener, PythonServer, Server, uvicorn_trusted_proxies from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager server: Final = ( @@ -883,7 +883,7 @@ class ProxyInitializationHelpers: proxy_headers=True, trusted_proxies=",".join(TypeAdapter(list[str]).validate_python(GunicornConfig().forwarded_allow_ips)) if server is Server.GUNICORN - else None, + else uvicorn_trusted_proxies(None), ) @staticmethod @@ -1359,7 +1359,7 @@ def run_server( from pydantic import ValidationError from litellm.proxy.rust_forward.mode import ( - Supervisor, + Mode, environment_error, gunicorn_bind, rust_forward_mode, @@ -1373,7 +1373,7 @@ def run_server( except ValidationError as error: raise click.ClickException(environment_error(error)) from None stop_with_supervisor(forward_mode) - if isinstance(forward_mode, Supervisor) and not any( + if forward_mode is Mode.SUPERVISOR and not any( (version, validate_config, health, test, skip_server_startup, check_migrations_against_existing_data) ): from litellm.proxy.rust_forward.supervisor import StartupError, supervise @@ -1385,7 +1385,7 @@ def run_server( run_gunicorn=run_gunicorn, run_hypercorn=run_hypercorn, run_granian=run_granian, - tls=any((ssl_keyfile_path, ssl_certfile_path, ciphers)), + tls=None not in (ssl_keyfile_path, ssl_certfile_path), ) supervised: Final = supervise(server, (sys.executable, "-m", "litellm.proxy.proxy_cli", *sys.argv[1:])) if not isinstance(supervised, StartupError): diff --git a/litellm/proxy/rust_forward/AGENTS.md b/litellm/proxy/rust_forward/AGENTS.md index 2c76bf73674..907b7dc1a3c 100644 --- a/litellm/proxy/rust_forward/AGENTS.md +++ b/litellm/proxy/rust_forward/AGENTS.md @@ -11,9 +11,9 @@ Every OS-specific operation lives behind `platform.Platform`, which `supervise` Windows support exists only for running the proxy locally from the CLI. Nobody deploys the proxy on Windows, and the Windows design leans on that. A port picked by the supervisor can be taken by another process before Python binds it, `CTRL_BREAK_EVENT` needs the console a CLI run always has, and gunicorn is refused. Revisit those choices before treating Windows as a deployment target -Each environment contract is written and read in one place. `mode.supervised_environment` and `mode.rust_forward_mode` are the two halves of the Python child's, and `supervisor.rust_spec` writes the `LITELLM_FORWARD_*` variables that `gateway-forward`'s `settings.rs` reads. Python owns every value Rust needs, so a new one gets a new variable in `rust_spec` +Each environment contract is written and read in one place. The Python child's 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. The trust it carries is resolved by the server library the entry point uses (`server.uvicorn_trusted_proxies` or gunicorn's config), never defaulted in this package. `supervisor.rust_spec` writes the `LITELLM_FORWARD_*` variables that `gateway-forward`'s `settings.rs` reads. Python owns every value Rust needs, so a new one gets a new variable in `rust_spec` -The supervisor is transparent to Python, so `python` with `LITELLM_RUST` behaves on signals and exits exactly as Python alone. Its whole lifecycle is `supervisor.step`, 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. `_run` only executes actions, and the `async with` around each child owns its cleanup. Three rules decide everything there, and a change that breaks one of them breaks the equivalence `tests/integration/configuration/test_rust_forward_lifecycle.py` checks against Python alone +The supervisor is transparent to Python, so `python` with `LITELLM_RUST` behaves on signals and exits exactly as Python alone. Its whole lifecycle is `supervisor.step`, 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. `_run` only executes actions, and the `contextlib.AsyncExitStack` it enters each child into owns their cleanup, last spawned first. Three rules decide everything there, and a change that breaks one of them breaks the equivalence `tests/integration/configuration/test_rust_forward_lifecycle.py` checks against Python alone - 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. Children run in their own session (`start_new_session`) only so that a terminal Ctrl-C reaches Python once, through the supervisor - 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 diff --git a/litellm/proxy/rust_forward/mode.py b/litellm/proxy/rust_forward/mode.py index 4f995d4a602..84c0a24d0fc 100644 --- a/litellm/proxy/rust_forward/mode.py +++ b/litellm/proxy/rust_forward/mode.py @@ -4,40 +4,28 @@ 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, ValidationError +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError from typing_extensions import assert_never from litellm.rust_bridge.configuration import ENV_NAME, RustSwitch -GATEWAY_FD: Final = "_LITELLM_GATEWAY_FD" -GATEWAY_PORT: Final = "_LITELLM_GATEWAY_PORT" -PROXY_HEADERS: Final = "_LITELLM_GATEWAY_PROXY_HEADERS" -TRUSTED_PROXIES: Final = "FORWARDED_ALLOW_IPS" -DEFAULT_TRUSTED_PROXIES: Final = "127.0.0.1" +SUPERVISED_ENV: Final = "_LITELLM_GATEWAY_SUPERVISED" RUST_HOP: Final = "127.0.0.1" -class _ForwardEnvironment(BaseModel): +class _Switch(BaseModel): model_config = ConfigDict(frozen=True) rust: RustSwitch = Field(None, validation_alias=ENV_NAME) - listener_fd: int | None = Field(None, validation_alias=GATEWAY_FD) - listener_port: int | None = Field(None, validation_alias=GATEWAY_PORT) - proxy_headers: bool = Field(True, validation_alias=PROXY_HEADERS) - trusted_proxies: str = Field(DEFAULT_TRUSTED_PROXIES, validation_alias=TRUSTED_PROXIES) -@dataclass(frozen=True, slots=True) -class Disabled: - pass - - -@dataclass(frozen=True, slots=True) -class Supervisor: - pass +class Mode(Enum): + DISABLED = auto() + SUPERVISOR = auto() @dataclass(frozen=True, slots=True) @@ -60,41 +48,20 @@ class Supervised: trusted_proxies: str -RustForwardMode: TypeAlias = Disabled | Supervisor | Supervised +RustForwardMode: TypeAlias = Mode | Supervised +_SUPERVISED: Final = TypeAdapter(Supervised) -def _listener_environment(listener: ChildListener) -> tuple[str, str]: - match listener: - case InheritedFd(fd): - return GATEWAY_FD, str(fd) - case AssignedPort(port): - return GATEWAY_PORT, str(port) - case _: - assert_never(listener) - - -def supervised_environment( - listener: ChildListener, proxy_headers: bool, trusted_proxies: str | None -) -> Mapping[str, str]: - name, value = _listener_environment(listener) - return MappingProxyType( - { - name: value, - PROXY_HEADERS: str(proxy_headers).lower(), - **({} if trusted_proxies is None else {TRUSTED_PROXIES: trusted_proxies}), - } - ) +def 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: - parsed: Final = _ForwardEnvironment.model_validate(environ) - if parsed.listener_fd is not None: - return Supervised(InheritedFd(parsed.listener_fd), parsed.proxy_headers, parsed.trusted_proxies) - if parsed.listener_port is not None: - return Supervised(AssignedPort(parsed.listener_port), parsed.proxy_headers, parsed.trusted_proxies) - if parsed.rust: - return Supervisor() - return Disabled() + 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: @@ -118,7 +85,7 @@ def stop_with_supervisor(mode: RustForwardMode, fd: int = 0, stop: Callable[[], match mode: case Supervised(): threading.Thread(target=_stop_when_closed, args=(fd, stop), name="supervisor-stdin", daemon=True).start() - case Disabled() | Supervisor(): + case Mode.DISABLED | Mode.SUPERVISOR: pass case _: assert_never(mode) @@ -151,7 +118,7 @@ def uvicorn_options(mode: RustForwardMode, factory: str) -> Mapping[str, object] return MappingProxyType( {"app": factory, "factory": True, **_uvicorn_listener(listener), "proxy_headers": False} ) - case Disabled() | Supervisor(): + case Mode.DISABLED | Mode.SUPERVISOR: return MappingProxyType({}) case _: assert_never(mode) @@ -171,7 +138,7 @@ def uvicorn_flags(mode: RustForwardMode) -> tuple[str, ...]: match mode: case Supervised(listener): return (*_uvicorn_listener_flags(listener), "--no-proxy-headers") - case Disabled() | Supervisor(): + case Mode.DISABLED | Mode.SUPERVISOR: return () case _: assert_never(mode) @@ -191,7 +158,7 @@ def gunicorn_bind(mode: RustForwardMode, public: str) -> str: match mode: case Supervised(listener): return _gunicorn_listener(listener) - case Disabled() | Supervisor(): + case Mode.DISABLED | Mode.SUPERVISOR: return public case _: assert_never(mode) @@ -201,7 +168,7 @@ def server_proxy_headers(mode: RustForwardMode) -> bool: match mode: case Supervised(): return False - case Disabled() | Supervisor(): + case Mode.DISABLED | Mode.SUPERVISOR: return True case _: assert_never(mode) diff --git a/litellm/proxy/rust_forward/peer.py b/litellm/proxy/rust_forward/peer.py index cf891495c2b..78445aad4e4 100644 --- a/litellm/proxy/rust_forward/peer.py +++ b/litellm/proxy/rust_forward/peer.py @@ -7,7 +7,7 @@ from typing import ( from typing_extensions import assert_never -from litellm.proxy.rust_forward.mode import RUST_HOP, Disabled, RustForwardMode, Supervised, Supervisor +from litellm.proxy.rust_forward.mode import RUST_HOP, Mode, RustForwardMode, Supervised if TYPE_CHECKING: import starlette.types @@ -63,7 +63,7 @@ def wrap(app: "starlette.types.ASGIApp", mode: RustForwardMode) -> ASGIApp: match mode: case Supervised(): return RustPeerMiddleware(converted, mode) - case Disabled() | Supervisor(): + case Mode.DISABLED | Mode.SUPERVISOR: return converted case _: assert_never(mode) diff --git a/litellm/proxy/rust_forward/server.py b/litellm/proxy/rust_forward/server.py index 8398cbda533..58d0ca1c2ec 100644 --- a/litellm/proxy/rust_forward/server.py +++ b/litellm/proxy/rust_forward/server.py @@ -1,5 +1,10 @@ from dataclasses import dataclass from enum import Enum +from typing import Final + +from pydantic import BaseModel, ConfigDict, TypeAdapter + +_TRUSTED_PROXIES: Final = TypeAdapter(str) class Server(Enum): @@ -27,4 +32,39 @@ class PythonServer: keepalive_seconds: int graceful_timeout_seconds: float proxy_headers: bool - trusted_proxies: str | None + 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 + + +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 + ) diff --git a/litellm/proxy/rust_forward/supervisor.py b/litellm/proxy/rust_forward/supervisor.py index f02d6d4e1cd..d29db08dba0 100644 --- a/litellm/proxy/rust_forward/supervisor.py +++ b/litellm/proxy/rust_forward/supervisor.py @@ -7,15 +7,15 @@ import sys import sysconfig from collections.abc import AsyncGenerator, Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass, replace -from enum import Enum +from enum import Enum, auto from functools import partial from types import MappingProxyType -from typing import Final, TypeAlias +from typing import Final, Literal, TypeAlias from typing_extensions import assert_never from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS -from litellm.proxy.rust_forward.mode import RUST_HOP, supervised_environment +from litellm.proxy.rust_forward.mode import RUST_HOP, Supervised, supervised_environment from litellm.proxy.rust_forward.platform import Platform, PrivateListener, SpawnOptions, current_platform from litellm.proxy.rust_forward.server import Listener, PythonServer, Server from litellm.rust_bridge.configuration import ENV_NAME @@ -52,19 +52,10 @@ class Child(Enum): FAILURE: Final = 1 -@dataclass(frozen=True, slots=True) -class NotStarted: - pass - - -@dataclass(frozen=True, slots=True) -class Spawned: - pass - - -@dataclass(frozen=True, slots=True) -class Listening: - pass +class Phase(Enum): + NOT_STARTED = auto() + SPAWNED = auto() + LISTENING = auto() @dataclass(frozen=True, slots=True) @@ -79,8 +70,8 @@ class Lost: status: int -PythonState: TypeAlias = Spawned | Listening | Exited -RustState: TypeAlias = NotStarted | Spawned | Exited | Lost +PythonState: TypeAlias = Literal[Phase.SPAWNED, Phase.LISTENING] | Exited +RustState: TypeAlias = Literal[Phase.NOT_STARTED, Phase.SPAWNED] | Exited | Lost @dataclass(frozen=True, slots=True) @@ -95,9 +86,8 @@ class Caught: signum: signal.Signals -@dataclass(frozen=True, slots=True) -class PythonListening: - pass +class Probe(Enum): + PYTHON_LISTENING = auto() @dataclass(frozen=True, slots=True) @@ -117,39 +107,39 @@ class Send: signum: signal.Signals -Event: TypeAlias = Caught | PythonListening | ChildExited +Event: TypeAlias = Caught | Literal[Probe.PYTHON_LISTENING] | ChildExited Action: TypeAlias = Spawn | Send Transition: TypeAlias = tuple[Lifecycle, tuple[Action, ...]] def boot() -> Transition: - return Lifecycle(Spawned(), NotStarted(), stop_requested=False), (Spawn(Child.PYTHON),) + 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 isinstance(state.rust, Spawned) else () + 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 not isinstance(state.python, Spawned): + if state.python is not Phase.SPAWNED: return state, () if state.stop_requested: - return replace(state, python=Listening()), () - return replace(state, python=Listening(), rust=Spawned()), (Spawn(Child.RUST),) + 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 isinstance(state.rust, Spawned) else ()) + return exited, ((Send(Child.RUST, signal.SIGTERM),) if state.rust is Phase.SPAWNED else ()) def _rust_exited(state: Lifecycle, status: int) -> Transition: - if not isinstance(state.rust, Spawned): + if state.rust is not Phase.SPAWNED: return state, () if state.stop_requested or isinstance(state.python, Exited): return replace(state, rust=Exited(status)), () @@ -162,7 +152,7 @@ def step(state: Lifecycle, event: Event, stop_signals: frozenset[signal.Signals] match event: case Caught(signum): return _relay(state, signum, signum in stop_signals) - case PythonListening(): + case Probe.PYTHON_LISTENING: return _python_listening(state) case ChildExited(child, status): match child: @@ -178,7 +168,7 @@ def step(state: Lifecycle, event: Event, stop_signals: frozenset[signal.Signals] def exit_status(state: Lifecycle) -> int | None: match state.python, state.rust: - case Exited(status), NotStarted() | Exited(): + case Exited(status), Phase.NOT_STARTED | Exited(): return status case Exited(), Lost(): return FAILURE @@ -198,7 +188,9 @@ def python_spec( ) -> ChildSpec: return ChildSpec( argv=(*command, "--port", str(private.port)), - env_overrides=supervised_environment(private.listener, server.proxy_headers, server.trusted_proxies), + env_overrides=supervised_environment( + Supervised(private.listener, server.proxy_headers, server.trusted_proxies) + ), spawn=platform.spawn_options(private.pass_fds), ) @@ -251,7 +243,7 @@ async def _accepts(port: int) -> bool: async def _report_once_listening(port: int, events: asyncio.Queue[Event]) -> None: while not await _accepts(port): await asyncio.sleep(PROBE_INTERVAL_SECONDS) - events.put_nowait(PythonListening()) + events.put_nowait(Probe.PYTHON_LISTENING) @contextlib.asynccontextmanager @@ -295,53 +287,22 @@ async def _launch_python( yield process -@dataclass(frozen=True, slots=True) -class _Launching: - state: Lifecycle - child: Child - rest: tuple[Action, ...] - - -async def _advance( - platform: Platform, - transition: Transition, - events: asyncio.Queue[Event], - children: Mapping[Child, asyncio.subprocess.Process], -) -> int | _Launching: - state, actions = transition - while True: - for index, action in enumerate(actions): - match action: - case Send(child, signum): - platform.signal_child(children[child], signum) - case Spawn(child): - return _Launching(state, child, actions[index + 1 :]) - case _: - assert_never(action) - status = exit_status(state) - if status is not None: - return status - state, actions = step(state, await events.get(), platform.stop_signals) - - -async def _run( - platform: Platform, - transition: Transition, - events: asyncio.Queue[Event], - children: Mapping[Child, asyncio.subprocess.Process], - launches: Mapping[Child, Launch], -) -> int: - advanced: Final = await _advance(platform, transition, events, children) - match advanced: - case int(): - return advanced - case _Launching(state, child, rest): - async with launches[child](events) as process: - return await _run( - platform, (state, rest), events, MappingProxyType({**children, child: process}), launches - ) - case _: - assert_never(advanced) +async def _run(platform: Platform, events: asyncio.Queue[Event], launches: Mapping[Child, Launch]) -> int: + async with contextlib.AsyncExitStack() as stack: + children: Final[dict[Child, asyncio.subprocess.Process]] = {} # mutable-ok: indexes the stack's children + state, actions = boot() + while True: + for action in actions: + match action: + case Spawn(child): + children[child] = await stack.enter_async_context(launches[child](events)) + case Send(child, signum): + platform.signal_child(children[child], signum) + case _: + assert_never(action) + if (status := exit_status(state)) is not None: + return status + state, actions = step(state, await events.get(), platform.stop_signals) async def _supervise(platform: Platform, binary: str, command: Sequence[str], server: PythonServer) -> int: @@ -356,7 +317,7 @@ async def _supervise(platform: Platform, binary: str, command: Sequence[str], se Child.RUST: partial(_launch, platform, Child.RUST, rust_spec(platform, binary, private.port, server)), } ) - return await _run(platform, boot(), events, MappingProxyType({}), launches) + return await _run(platform, events, launches) def _exit_like(status: int) -> int: diff --git a/tests/integration/configuration/test_rust_forward_lifecycle.py b/tests/integration/configuration/test_rust_forward_lifecycle.py index 3c639ce6d9a..7abafc6c6cd 100644 --- a/tests/integration/configuration/test_rust_forward_lifecycle.py +++ b/tests/integration/configuration/test_rust_forward_lifecycle.py @@ -46,14 +46,14 @@ import uvicorn def listener(): - fd = os.environ.get("_LITELLM_GATEWAY_FD") - if fd is None: - return {"host": "127.0.0.1", "port": int(sys.argv[2])} - from litellm.proxy.rust_forward.mode import rust_forward_mode, stop_with_supervisor + from litellm.proxy.rust_forward.mode import Supervised, rust_forward_mode, stop_with_supervisor - stop_with_supervisor(rust_forward_mode(os.environ)) + 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": int(fd)} + return {"fd": mode.listener.fd} if __name__ == "__main__": @@ -123,7 +123,7 @@ 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, True, None) +server = PythonServer(Server.UVICORN, Listener.TCP, "127.0.0.1", port, 5, 30.0, True, "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) """ diff --git a/tests/unit/gateway/test_launch.py b/tests/unit/gateway/test_launch.py index 8a916560c1b..becafe07d5d 100644 --- a/tests/unit/gateway/test_launch.py +++ b/tests/unit/gateway/test_launch.py @@ -126,7 +126,8 @@ class TestUvicornArgv: (("--fd", "0"), Listener.INHERITED_FD), (("--fd=17",), Listener.INHERITED_FD), (("--ssl-keyfile", "unused"), Listener.TLS), - (("--ssl-ciphers", "TLSv1.2"), Listener.TLS), + (("--ssl-certfile", "unused"), Listener.TLS), + (("--ssl-ciphers", "TLSv1.2"), Listener.TCP), ], ) def test_forwarding_sees_the_public_listener_uvicorn_would_open( @@ -135,18 +136,25 @@ class TestUvicornArgv: assert python_server(arguments).listener is listener @pytest.mark.parametrize( - ("argv", "trust"), + ("argv", "environ", "trust"), [ - ((), (True, None)), - (("--forwarded-allow-ips", "10.0.0.5"), (True, "10.0.0.5")), - (("--forwarded-allow-ips=*",), (True, "*")), - (("--no-proxy-headers",), (False, None)), + ((), {}, (True, "127.0.0.1")), + ((), {"FORWARDED_ALLOW_IPS": "10.9.9.9"}, (True, "10.9.9.9")), + (("--forwarded-allow-ips", "10.0.0.5"), {"FORWARDED_ALLOW_IPS": "10.9.9.9"}, (True, "10.0.0.5")), + (("--forwarded-allow-ips=*",), {}, (True, "*")), + (("--no-proxy-headers",), {}, (False, "127.0.0.1")), ], ) - def test_forwarding_carries_operator_proxy_header_options_to_the_child( - self, argv: tuple[str, ...], trust: tuple[bool, str | None], monkeypatch: pytest.MonkeyPatch + def test_forwarding_carries_the_trust_uvicorn_resolves_to_the_child( + self, + argv: tuple[str, ...], + environ: dict[str, str], + trust: tuple[bool, 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.proxy_headers, server.trusted_proxies) == trust @@ -335,7 +343,7 @@ class TestMain: capsys: pytest.CaptureFixture[str], ) -> None: monkeypatch.setenv("LITELLM_RUST", rust) - monkeypatch.delenv("_LITELLM_GATEWAY_FD", raising=False) + 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))) diff --git a/tests/unit/proxy/rust_forward/test_mode.py b/tests/unit/proxy/rust_forward/test_mode.py index 34329bacbb2..5acd76860fb 100644 --- a/tests/unit/proxy/rust_forward/test_mode.py +++ b/tests/unit/proxy/rust_forward/test_mode.py @@ -8,12 +8,12 @@ from pydantic import ValidationError from litellm.proxy.proxy_cli import WORKER_APP, ProxyInitializationHelpers from litellm.proxy.rust_forward.mode import ( + SUPERVISED_ENV, AssignedPort, - Disabled, InheritedFd, + Mode, RustForwardMode, Supervised, - Supervisor, gunicorn_bind, rust_forward_mode, server_proxy_headers, @@ -26,21 +26,17 @@ from litellm.proxy.rust_forward.mode import ( @pytest.mark.parametrize( ("environ", "expected"), [ - ({}, Disabled()), - ({"LITELLM_RUST": "false"}, Disabled()), - ({"LITELLM_RUST": "off"}, Disabled()), - ({"LITELLM_RUST": "0"}, Disabled()), - ({"LITELLM_RUST": "true"}, Supervisor()), - ({"LITELLM_RUST": "1"}, Supervisor()), - ({"LITELLM_RUST": "yes"}, Supervisor()), - ({"LITELLM_RUST": "on"}, Supervisor()), - ({"LITELLM_RUST": " 1 "}, Supervisor()), - ({"LITELLM_RUST": ""}, Disabled()), - ({"LITELLM_RUST": " "}, Disabled()), - ( - {"LITELLM_RUST": "true", "_LITELLM_GATEWAY_FD": "17", "FORWARDED_ALLOW_IPS": "127.0.0.1,10.0.0.5"}, - Supervised(InheritedFd(17), True, "127.0.0.1,10.0.0.5"), - ), + ({}, Mode.DISABLED), + ({"LITELLM_RUST": "false"}, Mode.DISABLED), + ({"LITELLM_RUST": "off"}, Mode.DISABLED), + ({"LITELLM_RUST": "0"}, Mode.DISABLED), + ({"LITELLM_RUST": "true"}, Mode.SUPERVISOR), + ({"LITELLM_RUST": "1"}, Mode.SUPERVISOR), + ({"LITELLM_RUST": "yes"}, Mode.SUPERVISOR), + ({"LITELLM_RUST": "on"}, Mode.SUPERVISOR), + ({"LITELLM_RUST": " 1 "}, Mode.SUPERVISOR), + ({"LITELLM_RUST": ""}, Mode.DISABLED), + ({"LITELLM_RUST": " "}, Mode.DISABLED), ], ) def test_mode_is_resolved_from_environment(environ: dict[str, str], expected: RustForwardMode) -> None: @@ -48,35 +44,36 @@ def test_mode_is_resolved_from_environment(environ: dict[str, str], expected: Ru @pytest.mark.parametrize( - ("environ", "variable"), + ("environ", "invalid"), [ ({"LITELLM_RUST": "enabled"}, "LITELLM_RUST"), ({"LITELLM_RUST": "2"}, "LITELLM_RUST"), - ({"_LITELLM_GATEWAY_FD": "seventeen"}, "_LITELLM_GATEWAY_FD"), - ({"_LITELLM_GATEWAY_FD": "17", "_LITELLM_GATEWAY_PROXY_HEADERS": "maybe"}, "_LITELLM_GATEWAY_PROXY_HEADERS"), + ({SUPERVISED_ENV: '{"listener": {"fd": "seventeen"}, "proxy_headers": true, "trusted_proxies": "*"}'}, "fd"), + ( + {SUPERVISED_ENV: '{"listener": {"fd": 17}, "proxy_headers": "maybe", "trusted_proxies": "*"}'}, + "proxy_headers", + ), + ({SUPERVISED_ENV: '{"listener": {"fd": 17}, "proxy_headers": true}'}, "trusted_proxies"), ], ) -def test_invalid_forwarding_value_stops_startup_naming_the_variable(environ: dict[str, str], variable: str) -> None: - with pytest.raises(ValidationError, match=variable): +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( - ("proxy_headers", "trusted_proxies", "inherited", "expected"), + "written", [ - (True, None, {}, Supervised(InheritedFd(17), True, "127.0.0.1")), - (True, None, {"FORWARDED_ALLOW_IPS": "10.0.0.5"}, Supervised(InheritedFd(17), True, "10.0.0.5")), - (True, "*", {"FORWARDED_ALLOW_IPS": "10.0.0.5"}, Supervised(InheritedFd(17), True, "*")), - (False, None, {}, Supervised(InheritedFd(17), False, "127.0.0.1")), - (True, None, {}, Supervised(AssignedPort(45555), True, "127.0.0.1")), - (False, "*", {}, Supervised(AssignedPort(45555), False, "*")), + Supervised(InheritedFd(17), True, "127.0.0.1"), + Supervised(InheritedFd(17), True, "127.0.0.1,10.0.0.5"), + Supervised(InheritedFd(17), False, "127.0.0.1"), + Supervised(AssignedPort(45555), True, "*"), + Supervised(AssignedPort(45555), False, "*"), ], ) -def test_child_reads_back_the_listener_and_trust_the_supervisor_wrote( - proxy_headers: bool, trusted_proxies: str | None, inherited: dict[str, str], expected: Supervised -) -> None: - written: Final = supervised_environment(expected.listener, proxy_headers, trusted_proxies) - assert rust_forward_mode({**inherited, **written}) == expected +def test_child_reads_back_exactly_what_the_supervisor_wrote(written: Supervised) -> None: + inherited: Final = {"LITELLM_RUST": "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: @@ -93,7 +90,7 @@ def test_supervised_uvicorn_serves_the_inherited_socket_without_trusting_the_rus assert (config.app, config.factory) == (WORKER_APP, True) -@pytest.mark.parametrize("mode", [Disabled(), Supervisor()]) +@pytest.mark.parametrize("mode", list(Mode)) def test_unsupervised_servers_keep_their_own_socket_and_proxy_headers(mode: RustForwardMode) -> None: config: Final = _served_uvicorn_config(mode) assert (config.fd, config.proxy_headers, config.port) == (None, True, 45555) @@ -117,7 +114,7 @@ def test_assigned_port_serves_only_the_loopback_hop_without_trusting_it() -> Non @pytest.mark.parametrize( ("mode", "stops"), - [(Supervised(InheritedFd(17), False, "127.0.0.1"), True), (Disabled(), False), (Supervisor(), False)], + [(Supervised(InheritedFd(17), False, "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() diff --git a/tests/unit/proxy/rust_forward/test_peer.py b/tests/unit/proxy/rust_forward/test_peer.py index b0c59641e7f..01331e2c69e 100644 --- a/tests/unit/proxy/rust_forward/test_peer.py +++ b/tests/unit/proxy/rust_forward/test_peer.py @@ -20,7 +20,7 @@ if TYPE_CHECKING: ) from litellm.proxy.proxy_server import app as proxy_app -from litellm.proxy.rust_forward.mode import InheritedFd, Supervised +from litellm.proxy.rust_forward.mode import SUPERVISED_ENV, InheritedFd, Supervised, supervised_environment from litellm.proxy.rust_forward.peer import RustPeerMiddleware Headers = tuple[tuple[bytes, bytes], ...] @@ -124,7 +124,9 @@ async def _peer_probe(request: Request) -> dict[str, object]: async def test_every_entry_point_hands_routes_the_peer_rust_saw( entry_point: str, monkeypatch: pytest.MonkeyPatch ) -> None: - monkeypatch.setenv("_LITELLM_GATEWAY_FD", "17") + monkeypatch.setenv( + SUPERVISED_ENV, supervised_environment(Supervised(InheritedFd(17), True, "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] diff --git a/tests/unit/proxy/rust_forward/test_supervisor.py b/tests/unit/proxy/rust_forward/test_supervisor.py index 1931b11c7a9..bb14443c96e 100644 --- a/tests/unit/proxy/rust_forward/test_supervisor.py +++ b/tests/unit/proxy/rust_forward/test_supervisor.py @@ -31,10 +31,10 @@ from litellm.proxy.rust_forward.supervisor import ( ChildExited, Event, Lifecycle, - PythonListening, + Phase, + Probe, Send, Spawn, - Spawned, StartupError, boot, exit_status, @@ -54,7 +54,7 @@ POSIX_ONLY: Final = pytest.mark.skipif(sys.platform == "win32", reason="runs a P def _server( - server: Server = Server.UVICORN, listener: Listener = Listener.TCP, trusted_proxies: str | None = None + 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, False, trusted_proxies) @@ -168,12 +168,11 @@ def test_binary_is_found_only_in_the_environment_scripts_dir(tmp_path: Path, mon PYTHON_CHILD: Final = """ -import os, pathlib, socket, time +import json, os, pathlib, socket, time time.sleep(0.5) pathlib.Path(os.environ["MARKERS"], "listening").touch() -fd = os.environ.get("_LITELLM_GATEWAY_FD") -port = os.environ.get("_LITELLM_GATEWAY_PORT") -listener = socket.socket(fileno=int(fd)) if fd else socket.create_server(("127.0.0.1", int(port))) +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) """ @@ -240,7 +239,7 @@ def _after(world: World, event: Event) -> World: match event: case Caught(signum): return replace(world, stop_requested=world.stop_requested or signum in STOPS) - case PythonListening(): + case Probe.PYTHON_LISTENING: return replace(world, reported_listening=world.reported_listening or Child.PYTHON in world.alive) case ChildExited(child, status) if child in world.alive: lost: Final = child is Child.RUST and Child.PYTHON in world.alive and not world.stop_requested @@ -264,7 +263,7 @@ STOPS: Final = POSIX.stop_signals EVENTS: Final = st.lists( st.one_of( st.sampled_from(SIGNALS).map(Caught), - st.just(PythonListening()), + st.just(Probe.PYTHON_LISTENING), st.builds(ChildExited, st.sampled_from(Child), st.integers(-15, 3)), ), max_size=40, @@ -272,8 +271,8 @@ EVENTS: Final = st.lists( @given(EVENTS) -@example([Caught(TERM), Caught(signal.SIGHUP), PythonListening()]) -@example([PythonListening(), Caught(TERM), Caught(signal.SIGWINCH), ChildExited(Child.RUST, 0)]) +@example([Caught(TERM), Caught(signal.SIGHUP), Probe.PYTHON_LISTENING]) +@example([Probe.PYTHON_LISTENING, Caught(TERM), Caught(signal.SIGWINCH), ChildExited(Child.RUST, 0)]) def test_lifecycle_relays_every_signal_and_ends_with_pythons_status(events: list[Event]) -> None: state, actions = boot() world = World() @@ -284,7 +283,7 @@ def test_lifecycle_relays_every_signal_and_ends_with_pythons_status(events: list match event: case Caught(signum): assert actions == _expected_relay(before, signum) - case PythonListening(): + case Probe.PYTHON_LISTENING: starts_rust = Child.PYTHON in world.alive and not world.reported_listening and not world.stop_requested assert actions == ((Spawn(Child.RUST),) if starts_rust else ()) case ChildExited(child) if child in world.alive: @@ -308,7 +307,7 @@ def test_rust_starts_at_most_once_and_only_while_python_listens_unstopped(events spawns: Final = tuple((event, state) for event, state, action in history if isinstance(action, Spawn)) assert len(spawns) <= 1 assert all( - isinstance(event, PythonListening) and isinstance(state.python, Spawned) and not state.stop_requested + event is Probe.PYTHON_LISTENING and state.python is Phase.SPAWNED and not state.stop_requested for event, state in spawns ) @@ -316,10 +315,22 @@ def test_rust_starts_at_most_once_and_only_while_python_listens_unstopped(events @pytest.mark.parametrize( ("events", "expected"), [ - ((PythonListening(), Caught(TERM), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, -15)), -15), - ((PythonListening(), Caught(INT), Caught(INT), ChildExited(Child.PYTHON, 1), ChildExited(Child.RUST, 0)), 1), - ((PythonListening(), Caught(signal.SIGHUP), ChildExited(Child.PYTHON, -1), ChildExited(Child.RUST, 0)), -1), - ((PythonListening(), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, 0)), FAILURE), + ((Probe.PYTHON_LISTENING, Caught(TERM), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, -15)), -15), + ( + ( + Probe.PYTHON_LISTENING, + Caught(INT), + Caught(INT), + ChildExited(Child.PYTHON, 1), + ChildExited(Child.RUST, 0), + ), + 1, + ), + ( + (Probe.PYTHON_LISTENING, Caught(signal.SIGHUP), ChildExited(Child.PYTHON, -1), ChildExited(Child.RUST, 0)), + -1, + ), + ((Probe.PYTHON_LISTENING, ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, 0)), FAILURE), ((ChildExited(Child.PYTHON, 3),), 3), ], ) diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index 7a30d78fee4..5ad21572605 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -4297,15 +4297,16 @@ class TestRustForwardServer: run_gunicorn=params["run_gunicorn"] is True, run_hypercorn=params["run_hypercorn"] is True, run_granian=params["run_granian"] is True, - tls=any((params["ssl_keyfile_path"], params["ssl_certfile_path"], params["ciphers"])), + tls=None not in (params["ssl_keyfile_path"], params["ssl_certfile_path"]), ) @pytest.mark.parametrize( ("argv", "reason"), [ - (["--ssl_keyfile_path", "unused"], "does not support TLS on the public port"), - (["--ssl_certfile_path", "unused"], "does not support TLS on the public port"), - (["--ciphers", "unused"], "does not support TLS on the public port"), + ( + ["--ssl_keyfile_path", "key.pem", "--ssl_certfile_path", "cert.pem"], + "does not support TLS on the public port", + ), (["--run_hypercorn"], "does not support Hypercorn"), (["--run_granian"], "does not support Granian"), ], @@ -4321,8 +4322,28 @@ class TestRustForwardServer: @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 + 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 @@ -4333,20 +4354,24 @@ class TestRustForwardServer: 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_FD")}, + **{ + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL", "_LITELLM_GATEWAY_SUPERVISED") + }, "LITELLM_RUST": "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", "--ssl_keyfile_path", "key.pem", "--ssl_certfile_path", "cert.pem"] - ) + result: Final = CliRunner().invoke(run_server, ["--port", "45999", *argv]) assert result.exit_code == 0, result.output - assert "LITELLM_RUST does not support TLS on the public port, so Python serves the public port" in result.output + assert f"LITELLM_RUST {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["ssl_keyfile"] == "key.pem" + 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: