mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-02 02:11:58 +00:00
feat(proxy): add LITELLM_DISABLE_LAZY_ROUTES to register optional routers at startup (#43911)
* feat(proxy): add LITELLM_DISABLE_LAZY_ROUTES to register optional routers at startup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): register eager lazy routes at startup so late eager routes keep precedence Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): share one optional-feature install path between lazy and eager registration Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * docs(proxy): say eager lazy routes register at worker startup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): name the import callable passed to _install Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): prove a startup hook can drop eager lazy routes for good Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): strip the lazy-routes flag from the lazy-mode control proxies Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): cover the lazy warmup route registering a feature Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci(unit): run tests/unit/proxy/test__lazy_features.py in the proxy-server-core shard Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): audit cells for LITELLM_DISABLE_LAZY_ROUTES Flag spellings, /openapi.json at boot, the warm-up route in both modes, /mcp/proxy ahead of the /mcp mount, route table and OpenAPI parity with a fully warmed lazy proxy, a broken optional import in both modes, every client SDK against the completion endpoints, a boot burst with a killed worker, and a restart * fix(proxy): keep config pass-through routes ahead of eagerly registered features With LITELLM_DISABLE_LAZY_ROUTES set, features registered before the proxy lifespan added config pass-through endpoints, so a pass-through overlapping a feature path (e.g. a self-hosted /langfuse) lost to the built-in route. Restore lazy mode's registry order once startup finishes, without bringing back routes a startup hook removed --------- Co-authored-by: ryan <ryan@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
46196339d2
commit
03743ae020
4 changed files with 945 additions and 53 deletions
|
|
@ -146,6 +146,7 @@ legacy_paths() {
|
|||
echo tests/unit/proxy/test_proxy_token_counter.py
|
||||
echo tests/unit/proxy/test_server_root_path.py ;;
|
||||
proxy-db-proxy-server-core)
|
||||
echo tests/unit/proxy/test__lazy_features.py
|
||||
echo tests/unit/proxy/test_aproxy_startup.py
|
||||
echo tests/unit/proxy/test_proxy_server.py ;;
|
||||
proxy-db-proxy-utils) echo tests/unit/proxy/test_proxy_utils.py ;;
|
||||
|
|
|
|||
|
|
@ -3,19 +3,24 @@ Lazy registration for optional feature routers. Each LAZY_FEATURES entry
|
|||
imports its module only on the first request matching its path prefix,
|
||||
saving ~700 MB at idle for deployments that don't use these features.
|
||||
First hit pays the import cost (1-3 s for heavy modules); /openapi.json
|
||||
omits each feature's routes until the feature is warmed.
|
||||
omits each feature's routes until the feature is warmed. Setting
|
||||
LITELLM_DISABLE_LAZY_ROUTES registers every feature at worker startup
|
||||
instead, so the route table is complete before the first request.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import importlib
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
import os
|
||||
from collections.abc import AsyncGenerator, Callable, Mapping, Sequence
|
||||
from collections.abc import Set as AbstractSet
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass, field
|
||||
from functools import partial
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from starlette.routing import BaseRoute, Match
|
||||
from starlette.types import ASGIApp, Receive, Scope, Send
|
||||
from starlette.types import ASGIApp, Lifespan, Receive, Scope, Send
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy.route_priority import hot_routes_first
|
||||
|
|
@ -428,57 +433,127 @@ def _in_registry_order(
|
|||
async def _force_load(app: "FastAPI", feat: LazyFeature, features: tuple[LazyFeature, ...] = LAZY_FEATURES) -> bool:
|
||||
"""Import + register a lazy feature exactly once per (app, module).
|
||||
Shared by the middleware and the /lazy/warm endpoint."""
|
||||
async with _lazy_lock(app, feat.module_path):
|
||||
if feat.module_path in _lazy_loaded(app):
|
||||
return False
|
||||
# Import on a thread (heavy modules take 1-3 s). register_fn
|
||||
# mutates app.router.routes, so it stays on the loop thread.
|
||||
imported: Final = asyncio.get_running_loop().run_in_executor(None, importlib.import_module, feat.module_path)
|
||||
await asyncio.wait((imported,))
|
||||
return _install(app, feat, imported.result, features)
|
||||
|
||||
|
||||
def _install(
|
||||
app: "FastAPI", feat: LazyFeature, import_module: Callable[[], object], features: tuple[LazyFeature, ...]
|
||||
) -> bool:
|
||||
try:
|
||||
_register_feature(app, feat, import_module(), features)
|
||||
return True
|
||||
except Exception as exc:
|
||||
_mark_failed(app, feat, exc)
|
||||
return False
|
||||
|
||||
|
||||
def _lazy_loaded(app: "FastAPI") -> set[str]:
|
||||
if not hasattr(app.state, "lazy_loaded"):
|
||||
app.state.lazy_loaded = set()
|
||||
app.state.lazy_locks = {}
|
||||
lock: Final = app.state.lazy_locks.setdefault(feat.module_path, asyncio.Lock())
|
||||
async with lock:
|
||||
if feat.module_path in app.state.lazy_loaded:
|
||||
return False
|
||||
try:
|
||||
# Import on a thread (heavy modules take 1-3 s). register_fn
|
||||
# mutates app.router.routes, so it stays on the loop thread.
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
module: Final = await loop.run_in_executor(None, importlib.import_module, feat.module_path)
|
||||
before: Final = len(app.router.routes)
|
||||
feat.register_fn(app, module)
|
||||
previous: Final[Mapping[str, tuple[BaseRoute, ...]]] = (
|
||||
app.state.lazy_routes if hasattr(app.state, "lazy_routes") else MappingProxyType({})
|
||||
)
|
||||
lazy_routes: Final[Mapping[str, tuple[BaseRoute, ...]]] = MappingProxyType(
|
||||
{**previous, feat.module_path: tuple(app.router.routes[before:])}
|
||||
)
|
||||
app.state.lazy_routes = lazy_routes # rebind-ok: the app owns the record of which routes each feature added
|
||||
app.router.routes[:] = hot_routes_first( # rebind-ok: the app owns its route table
|
||||
_in_registry_order(app.router.routes, lazy_routes, features, _lazy_slots(app))
|
||||
)
|
||||
app.state.lazy_loaded.add(feat.module_path)
|
||||
app.openapi_schema = None
|
||||
verbose_proxy_logger.info(
|
||||
"Lazy-loaded optional feature %r (module: %s)",
|
||||
feat.name,
|
||||
feat.module_path,
|
||||
)
|
||||
return True
|
||||
except Exception as exc:
|
||||
# Mark loaded anyway so we don't retry on every request.
|
||||
app.state.lazy_loaded.add(feat.module_path)
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to lazy-load optional feature %r (module: %s): %s. "
|
||||
"This feature's endpoints will return 404 until restart.",
|
||||
feat.name,
|
||||
feat.module_path,
|
||||
exc,
|
||||
)
|
||||
return False
|
||||
app.state.lazy_loaded = set[str]()
|
||||
app.state.lazy_locks = dict[str, asyncio.Lock]()
|
||||
loaded: Final[set[str]] = app.state.lazy_loaded
|
||||
return loaded
|
||||
|
||||
|
||||
def attach_lazy_features(app: "FastAPI") -> None:
|
||||
app.include_router(_make_warmup_router(app))
|
||||
app.add_middleware(LazyFeatureMiddleware, fastapi_app=app)
|
||||
def _lazy_lock(app: "FastAPI", module_path: str) -> asyncio.Lock:
|
||||
if not hasattr(app.state, "lazy_locks"):
|
||||
app.state.lazy_locks = dict[str, asyncio.Lock]()
|
||||
locks: Final[dict[str, asyncio.Lock]] = app.state.lazy_locks
|
||||
return locks.setdefault(module_path, asyncio.Lock())
|
||||
|
||||
|
||||
def _make_warmup_router(app: "FastAPI") -> "APIRouter":
|
||||
def _register_feature(app: "FastAPI", feat: LazyFeature, module: object, features: tuple[LazyFeature, ...]) -> None:
|
||||
before: Final = len(app.router.routes)
|
||||
feat.register_fn(app, module)
|
||||
previous: Final[Mapping[str, tuple[BaseRoute, ...]]] = (
|
||||
app.state.lazy_routes if hasattr(app.state, "lazy_routes") else MappingProxyType({})
|
||||
)
|
||||
lazy_routes: Final[Mapping[str, tuple[BaseRoute, ...]]] = MappingProxyType(
|
||||
{**previous, feat.module_path: tuple(app.router.routes[before:])}
|
||||
)
|
||||
app.state.lazy_routes = lazy_routes # rebind-ok: the app owns the record of which routes each feature added
|
||||
app.router.routes[:] = hot_routes_first( # rebind-ok: the app owns its route table
|
||||
_in_registry_order(app.router.routes, lazy_routes, features, _lazy_slots(app))
|
||||
)
|
||||
_lazy_loaded(app).add(feat.module_path)
|
||||
app.openapi_schema = None
|
||||
verbose_proxy_logger.info(
|
||||
"Lazy-loaded optional feature %r (module: %s)",
|
||||
feat.name,
|
||||
feat.module_path,
|
||||
)
|
||||
|
||||
|
||||
def _mark_failed(app: "FastAPI", feat: LazyFeature, exc: Exception) -> None:
|
||||
# Mark loaded anyway so we don't retry on every request.
|
||||
_lazy_loaded(app).add(feat.module_path)
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to lazy-load optional feature %r (module: %s): %s. "
|
||||
"This feature's endpoints will return 404 until restart.",
|
||||
feat.name,
|
||||
feat.module_path,
|
||||
exc,
|
||||
)
|
||||
|
||||
|
||||
def lazy_routes_disabled() -> bool:
|
||||
return os.getenv("LITELLM_DISABLE_LAZY_ROUTES", "").lower() in ("1", "true", "yes", "on")
|
||||
|
||||
|
||||
def register_all_features(app: "FastAPI", features: tuple[LazyFeature, ...] = LAZY_FEATURES) -> None:
|
||||
"""Register every feature router now, in registry order, so app.routes is
|
||||
complete before the app serves its first request."""
|
||||
for feat in features:
|
||||
_install(app, feat, partial(importlib.import_module, feat.module_path), features)
|
||||
|
||||
|
||||
def attach_lazy_features(app: "FastAPI", features: tuple[LazyFeature, ...] = LAZY_FEATURES) -> None:
|
||||
if lazy_routes_disabled():
|
||||
app.router.lifespan_context = _register_all_on_startup(app.router.lifespan_context, features)
|
||||
return
|
||||
app.include_router(_make_warmup_router(app, features))
|
||||
app.add_middleware(LazyFeatureMiddleware, fastapi_app=app, features=features)
|
||||
|
||||
|
||||
def _register_all_on_startup(inner: "Lifespan[FastAPI]", features: tuple[LazyFeature, ...]) -> "Lifespan[FastAPI]":
|
||||
"""Registering at startup, once every route the app defines exists, lands the features
|
||||
where lazy mode splices them: after every eager route (so /mcp/proxy, defined after
|
||||
attach_lazy_features(), still beats the /mcp mount) and before LITELLM_WORKER_STARTUP_HOOKS
|
||||
or an outer lifespan can filter the table. The inner lifespan then adds routes of its own
|
||||
(config pass-through endpoints), so the table is put back in lazy mode's order once it is up."""
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: "FastAPI") -> AsyncGenerator[None]:
|
||||
register_all_features(app, features)
|
||||
async with inner(app):
|
||||
_restore_registry_order(app, features)
|
||||
yield
|
||||
|
||||
return lifespan
|
||||
|
||||
|
||||
def _restore_registry_order(app: "FastAPI", features: tuple[LazyFeature, ...]) -> None:
|
||||
present: Final = frozenset(id(route) for route in app.router.routes)
|
||||
registered: Final[Mapping[str, tuple[BaseRoute, ...]]] = (
|
||||
app.state.lazy_routes if hasattr(app.state, "lazy_routes") else MappingProxyType({})
|
||||
)
|
||||
still_routed: Final = MappingProxyType(
|
||||
{module_path: tuple(r for r in routes if id(r) in present) for module_path, routes in registered.items()}
|
||||
)
|
||||
app.router.routes[:] = hot_routes_first( # rebind-ok: the app owns its route table
|
||||
_in_registry_order(app.router.routes, still_routed, features, _lazy_slots(app))
|
||||
)
|
||||
app.openapi_schema = None
|
||||
|
||||
|
||||
def _make_warmup_router(app: "FastAPI", features: tuple[LazyFeature, ...] = LAZY_FEATURES) -> "APIRouter":
|
||||
"""POST /lazy/warm/{name}: load a feature and return its partial openapi
|
||||
so the Swagger plugin can merge in-place without a full /openapi.json refetch.
|
||||
Requires auth — anyone who can hit the proxy can already trigger the same
|
||||
|
|
@ -497,13 +572,13 @@ def _make_warmup_router(app: "FastAPI") -> "APIRouter":
|
|||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
async def warm(name: str):
|
||||
feat: Final = next((f for f in LAZY_FEATURES if f.name == name), None)
|
||||
feat: Final = next((f for f in features if f.name == name), None)
|
||||
if feat is None:
|
||||
raise HTTPException(404, f"unknown lazy feature: {name}")
|
||||
if feat.persistent_swagger_stub:
|
||||
return {"stub_path": None, "paths": {}, "components": {"schemas": {}}}
|
||||
|
||||
await _force_load(app, feat)
|
||||
await _force_load(app, feat, features)
|
||||
|
||||
feat_routes: Final = [r for r in app.routes if feat.matches(getattr(r, "path", ""))]
|
||||
full: Final = get_openapi(title=app.title, version=app.version, routes=feat_routes)
|
||||
|
|
@ -524,7 +599,7 @@ def _make_warmup_router(app: "FastAPI") -> "APIRouter":
|
|||
|
||||
def loaded_lazy_modules(app: "FastAPI") -> frozenset[str]:
|
||||
"""The set of lazy feature modules whose routers are actually registered
|
||||
on this app (tracked by _force_load), empty before the middleware ever ran.
|
||||
on this app (tracked by _install), empty until a feature loads or eager startup runs.
|
||||
sys.modules is the wrong signal: boot code imports several feature modules
|
||||
(mcp_management, cloudzero, vantage, config_overrides) without mounting
|
||||
their routers, and their stubs must still be injected."""
|
||||
|
|
@ -583,6 +658,6 @@ def lazy_tag_to_prefix() -> dict[str, str]:
|
|||
because /openapi.json already has full route info."""
|
||||
from litellm.proxy._lazy_openapi_snapshot import load_snapshot
|
||||
|
||||
if load_snapshot():
|
||||
if lazy_routes_disabled() or load_snapshot():
|
||||
return {}
|
||||
return {feat.name: feat.path_prefixes[0] for feat in LAZY_FEATURES if not feat.persistent_swagger_stub}
|
||||
|
|
|
|||
605
tests/integration/configuration/test_lazy_routes_flag.py
Normal file
605
tests/integration/configuration/test_lazy_routes_flag.py
Normal file
|
|
@ -0,0 +1,605 @@
|
|||
"""Route table contract for the LITELLM_DISABLE_LAZY_ROUTES startup flag.
|
||||
|
||||
By default optional feature routers (``LAZY_FEATURES``) are registered on the first
|
||||
request to their path prefix, so an operator inspecting the route table right after
|
||||
boot cannot see or gate them. With the flag set every feature is registered at worker
|
||||
startup, so ``GET /routes`` lists them before any feature request is served and the
|
||||
first feature request changes nothing.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import openai
|
||||
import psutil
|
||||
import pytest
|
||||
import yaml
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
from litellm.proxy._lazy_features import LAZY_FEATURES, LazyFeature
|
||||
from tests.integration._support.client import Gateway, eventually, object_value, string_value
|
||||
from tests.integration._support.mcp import McpPeer, call_tool, echo_tool, scripted_peer, tool_calls, tool_names
|
||||
from tests.integration._support.process import OwnedProxy, owned_proxy_process
|
||||
from tests.integration._support.wire import Reply, Request, Wire, wire_server
|
||||
|
||||
TICKET_FEATURES: Final = ("mcp_management", "mcp_byok_oauth")
|
||||
FLAG: Final = "LITELLM_DISABLE_LAZY_ROUTES"
|
||||
WARMUP_ROUTE: Final = "/lazy/warm/{name}"
|
||||
MCP_WARM_PATH: Final = "/mcp/enabled"
|
||||
MARKER: Final = re.compile(rb"lazyroutes-[0-9a-f]{32}")
|
||||
FAILED_FEATURE: Final = re.compile(r"Failed to lazy-load optional feature '([a-z_]+)'")
|
||||
JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
HOOK_MODULE: Final = "lazy_routes_route_filter_hook"
|
||||
HOOK_SOURCE: Final = """from litellm.proxy.proxy_server import app
|
||||
|
||||
|
||||
def drop_mcp_routes() -> None:
|
||||
app.router.routes[:] = [
|
||||
route for route in app.router.routes if not getattr(route, "path", "").startswith(("/mcp", "/v1/mcp"))
|
||||
]
|
||||
"""
|
||||
|
||||
|
||||
def _paths(candidate: Gateway) -> tuple[str, ...]:
|
||||
routes: Final = candidate.get("/routes")["routes"]
|
||||
assert isinstance(routes, list), routes
|
||||
return tuple(string_value(object_value(route)["path"]) for route in routes)
|
||||
|
||||
|
||||
def _routed_features(candidate: Gateway) -> Mapping[str, tuple[str, ...]]:
|
||||
paths: Final = _paths(candidate)
|
||||
return {feature.name: tuple(path for path in paths if feature.matches(path)) for feature in LAZY_FEATURES}
|
||||
|
||||
|
||||
def _mcp_paths(candidate: Gateway) -> tuple[str, ...]:
|
||||
return tuple(path for path in _paths(candidate) if path.startswith(("/mcp", "/v1/mcp")))
|
||||
|
||||
|
||||
def _route_filter_hook(directory: Path) -> Mapping[str, str]:
|
||||
(directory / f"{HOOK_MODULE}.py").write_text(HOOK_SOURCE)
|
||||
search_path: Final = (str(directory), os.environ.get("PYTHONPATH", ""))
|
||||
return {
|
||||
"PYTHONPATH": os.pathsep.join(entry for entry in search_path if entry),
|
||||
"LITELLM_WORKER_STARTUP_HOOKS": f"{HOOK_MODULE}:drop_mcp_routes",
|
||||
}
|
||||
|
||||
|
||||
def test_lazy_routes_are_absent_from_the_route_table_until_first_request(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with owned_proxy_process(gateway, tmp_path, {}, remove_environment=(FLAG,)) as owned:
|
||||
at_boot: Final = _routed_features(owned.gateway)
|
||||
assert {name: at_boot[name] for name in TICKET_FEATURES} == {name: () for name in TICKET_FEATURES}, at_boot
|
||||
listing: Final = owned.gateway.request("GET", "/v1/mcp/server")
|
||||
assert listing.status_code == 200, listing.text
|
||||
after_first_request: Final = _routed_features(owned.gateway)
|
||||
assert after_first_request["mcp_management"] != (), "first request did not register the router"
|
||||
assert after_first_request["mcp_byok_oauth"] == (), "only the requested feature is mounted"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("workers", (1, 4))
|
||||
def test_disable_lazy_routes_flag_registers_every_feature_at_startup(
|
||||
gateway: Gateway, tmp_path: Path, workers: int
|
||||
) -> None:
|
||||
with owned_proxy_process(gateway, tmp_path, {FLAG: "true"}, workers=workers) as owned:
|
||||
at_boot: Final = tuple(_routed_features(owned.gateway) for _ in range(2 * workers))
|
||||
unregistered: Final = sorted(name for name, paths in at_boot[0].items() if not paths)
|
||||
assert unregistered == [], f"features still missing from /routes at startup: {unregistered}"
|
||||
assert all(table == at_boot[0] for table in at_boot), "workers disagree on the route table"
|
||||
listing: Final = owned.gateway.request("GET", "/v1/mcp/server")
|
||||
assert listing.status_code == 200, listing.text
|
||||
assert _routed_features(owned.gateway) == at_boot[0], "first feature request changed the route table"
|
||||
|
||||
|
||||
def test_startup_hook_cannot_remove_lazy_routes_that_register_after_it_ran(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with owned_proxy_process(gateway, tmp_path, _route_filter_hook(tmp_path), remove_environment=(FLAG,)) as owned:
|
||||
assert _mcp_paths(owned.gateway) == (), "hook should have removed the routes registered before it ran"
|
||||
listing: Final = owned.gateway.request("GET", "/v1/mcp/server")
|
||||
assert listing.status_code == 200, listing.text
|
||||
assert _mcp_paths(owned.gateway) != (), "first request should have registered the routes the hook never saw"
|
||||
|
||||
|
||||
def test_disable_lazy_routes_flag_lets_a_startup_hook_remove_optional_routes_for_good(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
overrides: Final = {**_route_filter_hook(tmp_path), FLAG: "true"}
|
||||
with owned_proxy_process(gateway, tmp_path, overrides) as owned:
|
||||
assert _mcp_paths(owned.gateway) == (), "hook should have seen and removed every MCP route"
|
||||
listing: Final = owned.gateway.request("GET", "/v1/mcp/server")
|
||||
assert listing.status_code == 404, listing.text
|
||||
mounted: Final = owned.gateway.request("POST", "/mcp", {"jsonrpc": "2.0", "id": 1, "method": "tools/list"})
|
||||
assert mounted.status_code == 404, mounted.text
|
||||
guardrails: Final = owned.gateway.request("GET", "/guardrails/list")
|
||||
assert guardrails.status_code == 200, guardrails.text
|
||||
assert _mcp_paths(owned.gateway) == (), "a feature request re-registered routes the hook removed"
|
||||
|
||||
|
||||
def _openapi_paths(candidate: Gateway) -> Mapping[str, tuple[str, ...]]:
|
||||
paths: Final = object_value(candidate.get("/openapi.json")["paths"])
|
||||
return {path: tuple(sorted(object_value(operations))) for path, operations in paths.items()}
|
||||
|
||||
|
||||
def _published(feature: LazyFeature, paths: Mapping[str, tuple[str, ...]]) -> bool:
|
||||
return any(feature.matches(path) for path in paths)
|
||||
|
||||
|
||||
def _warm_every_feature(candidate: Gateway) -> None:
|
||||
warmed: Final = tuple(
|
||||
(feature.name, candidate.request("POST", f"/lazy/warm/{feature.name}")) for feature in LAZY_FEATURES
|
||||
)
|
||||
cold: Final = [
|
||||
(name, response.status_code, response.text) for name, response in warmed if response.status_code != 200
|
||||
]
|
||||
assert cold == [], cold
|
||||
enabled: Final = candidate.request("GET", MCP_WARM_PATH)
|
||||
assert enabled.status_code == 200, enabled.text
|
||||
unregistered: Final = sorted(name for name, paths in _routed_features(candidate).items() if not paths)
|
||||
assert unregistered == [], f"features still missing after warming every one of them: {unregistered}"
|
||||
|
||||
|
||||
def _shadowed_dependency(directory: Path) -> Mapping[str, str]:
|
||||
package: Final = directory / "shadow" / "RestrictedPython"
|
||||
package.mkdir(parents=True)
|
||||
(package / "__init__.py").write_text('raise ImportError("shadowed by the lazy routes audit")\n')
|
||||
search_path: Final = (str(package.parent), os.environ.get("PYTHONPATH", ""))
|
||||
return {"PYTHONPATH": os.pathsep.join(entry for entry in search_path if entry)}
|
||||
|
||||
|
||||
def _failed_features(owned: OwnedProxy) -> frozenset[str]:
|
||||
return frozenset(FAILED_FEATURE.findall(owned.log.read_text()))
|
||||
|
||||
|
||||
def _marker() -> str:
|
||||
return "lazyroutes-" + uuid.uuid4().hex
|
||||
|
||||
|
||||
def _chat_reply(identity: str, stream: bool) -> Reply:
|
||||
if not stream:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [
|
||||
{"index": 0, "message": {"role": "assistant", "content": "lazy ok"}, "finish_reason": "stop"}
|
||||
],
|
||||
"usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
chunk: Final[dict[str, JsonValue]] = {
|
||||
"id": identity,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
}
|
||||
deltas: Final[tuple[dict[str, JsonValue], ...]] = (
|
||||
{**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "lazy"}}]},
|
||||
{**chunk, "choices": [{"index": 0, "delta": {"content": " ok"}, "finish_reason": "stop"}]},
|
||||
{**chunk, "choices": [], "usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9}},
|
||||
)
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=(*(b"data: " + json.dumps(delta).encode() + b"\n\n" for delta in deltas), b"data: [DONE]\n\n"),
|
||||
)
|
||||
|
||||
|
||||
def _responses_reply(identity: str, stream: bool) -> Reply:
|
||||
response: Final[dict[str, JsonValue]] = {
|
||||
"id": identity,
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": "gpt-4o-mini",
|
||||
"output": [
|
||||
{
|
||||
"id": "msg_" + identity,
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [{"type": "output_text", "text": "lazy ok", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {"input_tokens": 7, "output_tokens": 2, "total_tokens": 9},
|
||||
}
|
||||
if not stream:
|
||||
return Reply(body=json.dumps(response).encode())
|
||||
events: Final[tuple[dict[str, JsonValue], ...]] = (
|
||||
{"type": "response.created", "sequence_number": 0, "response": {**response, "status": "in_progress"}},
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"sequence_number": 1,
|
||||
"item_id": "msg_" + identity,
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": "lazy ok",
|
||||
},
|
||||
{"type": "response.completed", "sequence_number": 2, "response": response},
|
||||
)
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events),
|
||||
)
|
||||
|
||||
|
||||
def _upstream(request: Request) -> Reply:
|
||||
found: Final = MARKER.search(request.body)
|
||||
if found is None:
|
||||
return Reply(status=404, body=b'{"error":"no marker"}')
|
||||
marker: Final = found.group(0).decode()
|
||||
stream: Final = object_value(JSON.validate_json(request.body)).get("stream") is True
|
||||
if request.target.endswith("/responses"):
|
||||
return _responses_reply(f"resp_{marker}", stream)
|
||||
return _chat_reply(f"chatcmpl-{marker}", stream)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def provider() -> Iterator[Wire]:
|
||||
with wire_server(_upstream) as wire:
|
||||
yield wire
|
||||
|
||||
|
||||
async def _stream_chat(base_url: str, key: str, model: str, marker: str) -> tuple[frozenset[str], str]:
|
||||
client: Final = openai.AsyncOpenAI(base_url=base_url + "/v1", api_key=key, max_retries=0)
|
||||
stream: Final = await client.chat.completions.create(
|
||||
model=model, messages=[{"role": "user", "content": marker}], stream=True
|
||||
)
|
||||
chunks: Final = [chunk async for chunk in stream]
|
||||
text: Final = "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices)
|
||||
return frozenset(chunk.id for chunk in chunks), text
|
||||
|
||||
|
||||
async def _stream_message(base_url: str, key: str, model: str, marker: str) -> str:
|
||||
client: Final = anthropic.AsyncAnthropic(base_url=base_url, api_key=key, max_retries=0)
|
||||
async with client.messages.stream(
|
||||
model=model, max_tokens=16, messages=[{"role": "user", "content": marker}]
|
||||
) as stream:
|
||||
return "".join([text async for text in stream.text_stream])
|
||||
|
||||
|
||||
def _status(candidate: Gateway, path: str) -> int:
|
||||
return candidate.request("GET", path).status_code
|
||||
|
||||
|
||||
def _workers(owned: OwnedProxy) -> tuple[psutil.Process, ...]:
|
||||
return tuple(child for child in psutil.Process(owned.process.pid).children() if _is_worker(child))
|
||||
|
||||
|
||||
def _is_worker(child: psutil.Process) -> bool:
|
||||
try:
|
||||
return "spawn_main" in " ".join(child.cmdline()) and child.status() != psutil.STATUS_ZOMBIE
|
||||
except psutil.Error:
|
||||
return False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("spelling", ("1", "Yes", "ON"))
|
||||
def test_every_truthy_spelling_of_the_flag_registers_the_ticket_features_at_startup(
|
||||
gateway: Gateway, tmp_path: Path, spelling: str
|
||||
) -> None:
|
||||
with owned_proxy_process(gateway, tmp_path, {FLAG: spelling}) as owned:
|
||||
at_boot: Final = _routed_features(owned.gateway)
|
||||
assert all(at_boot[name] for name in TICKET_FEATURES), {name: at_boot[name] for name in TICKET_FEATURES}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"spelling", ("", "0", "off", "maybe", "x" * 5000), ids=("empty", "zero", "off", "unknown-word", "five-kilobytes")
|
||||
)
|
||||
def test_a_falsey_or_unknown_flag_value_keeps_the_default_lazy_registration(
|
||||
gateway: Gateway, tmp_path: Path, spelling: str
|
||||
) -> None:
|
||||
with owned_proxy_process(gateway, tmp_path, {FLAG: spelling}) as owned:
|
||||
at_boot: Final = _routed_features(owned.gateway)
|
||||
assert {name: at_boot[name] for name in TICKET_FEATURES} == {name: () for name in TICKET_FEATURES}, at_boot
|
||||
listing: Final = owned.gateway.request("GET", "/v1/mcp/server")
|
||||
assert listing.status_code == 200, listing.text
|
||||
assert _routed_features(owned.gateway)["mcp_management"] != (), "first request did not register the router"
|
||||
|
||||
|
||||
def test_disable_lazy_routes_flag_publishes_the_live_route_table_in_openapi_before_any_feature_request(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
with owned_proxy_process(gateway, tmp_path, {FLAG: "true"}) as owned:
|
||||
at_boot: Final = _openapi_paths(owned.gateway)
|
||||
unpublished: Final = sorted(
|
||||
feature.name
|
||||
for feature in LAZY_FEATURES
|
||||
if feature.name in TICKET_FEATURES and not _published(feature, at_boot)
|
||||
)
|
||||
assert unpublished == [], f"ticket features missing from /openapi.json at startup: {unpublished}"
|
||||
assert "get" in at_boot["/v1/mcp/server"], at_boot["/v1/mcp/server"]
|
||||
assert WARMUP_ROUTE not in at_boot
|
||||
stranger: Final = owned.gateway.request("GET", "/v1/mcp/server", key="sk-not-a-real-key")
|
||||
assert stranger.status_code == 401, stranger.text
|
||||
listing: Final = owned.gateway.request("GET", "/v1/mcp/server")
|
||||
assert listing.status_code == 200, listing.text
|
||||
assert _openapi_paths(owned.gateway) == at_boot, "first feature request changed /openapi.json"
|
||||
|
||||
|
||||
def test_disable_lazy_routes_flag_removes_the_warmup_route(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with owned_proxy_process(gateway, tmp_path, {FLAG: "true"}) as owned:
|
||||
assert WARMUP_ROUTE not in _paths(owned.gateway)
|
||||
warmed: Final = owned.gateway.request("POST", "/lazy/warm/mcp_management")
|
||||
assert warmed.status_code == 404, warmed.text
|
||||
listing: Final = owned.gateway.request("GET", "/v1/mcp/server")
|
||||
assert listing.status_code == 200, listing.text
|
||||
|
||||
|
||||
def test_the_warmup_route_registers_a_feature_on_demand_by_default(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with owned_proxy_process(gateway, tmp_path, {}, remove_environment=(FLAG,)) as owned:
|
||||
assert WARMUP_ROUTE in _paths(owned.gateway)
|
||||
warmed: Final = owned.gateway.request("POST", "/lazy/warm/mcp_management")
|
||||
assert warmed.status_code == 200, warmed.text
|
||||
assert "/v1/mcp/server" in object_value(object_value(JSON.validate_json(warmed.content))["paths"])
|
||||
assert _routed_features(owned.gateway)["mcp_management"] != (), "warmup did not register the router"
|
||||
|
||||
|
||||
def test_disable_lazy_routes_flag_keeps_the_fixed_mcp_proxy_route_ahead_of_the_mcp_mount(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
with owned_proxy_process(gateway, tmp_path, {}, remove_environment=(FLAG,)) as lazy:
|
||||
control: Final = lazy.gateway.request("POST", "/mcp/proxy", {})
|
||||
assert control.status_code == 400, control.text
|
||||
warmed: Final = _paths(lazy.gateway)
|
||||
assert warmed.count("/mcp") == 2, "expected the fixed /mcp route and the /mcp mount"
|
||||
with owned_proxy_process(gateway, tmp_path, {FLAG: "true"}) as eager:
|
||||
at_boot: Final = _paths(eager.gateway)
|
||||
assert at_boot.count("/mcp") == 2, "expected the fixed /mcp route and the /mcp mount at startup"
|
||||
mount: Final = max(index for index, path in enumerate(at_boot) if path == "/mcp")
|
||||
assert at_boot.index("/mcp/proxy") < mount, "the /mcp mount shadows /mcp/proxy"
|
||||
proxied: Final = eager.gateway.request("POST", "/mcp/proxy", {})
|
||||
assert (proxied.status_code, proxied.text) == (control.status_code, control.text)
|
||||
|
||||
|
||||
def test_disable_lazy_routes_flag_matches_the_fully_warmed_lazy_route_table_and_openapi(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
with owned_proxy_process(gateway, tmp_path, {}, remove_environment=(FLAG,)) as lazy:
|
||||
_warm_every_feature(lazy.gateway)
|
||||
warmed_paths: Final = tuple(path for path in _paths(lazy.gateway) if path != WARMUP_ROUTE)
|
||||
warmed_openapi: Final = _openapi_paths(lazy.gateway)
|
||||
with owned_proxy_process(gateway, tmp_path, {FLAG: "true"}) as eager:
|
||||
assert _paths(eager.gateway) == warmed_paths
|
||||
assert _openapi_paths(eager.gateway) == warmed_openapi
|
||||
|
||||
|
||||
def test_disable_lazy_routes_flag_keeps_registering_after_an_optional_dependency_fails_to_import(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
with owned_proxy_process(gateway, tmp_path, {**_shadowed_dependency(tmp_path), FLAG: "true"}) as owned:
|
||||
at_boot: Final = _routed_features(owned.gateway)
|
||||
unregistered: Final = frozenset(name for name, paths in at_boot.items() if not paths)
|
||||
failed: Final = _failed_features(owned)
|
||||
assert "guardrails" in failed, owned.log.read_text()
|
||||
assert unregistered == failed, (sorted(unregistered), sorted(failed))
|
||||
guardrails: Final = owned.gateway.request("GET", "/guardrails/list")
|
||||
assert guardrails.status_code == 404, guardrails.text
|
||||
listing: Final = owned.gateway.request("GET", "/v1/mcp/server")
|
||||
assert listing.status_code == 200, listing.text
|
||||
stores: Final = owned.gateway.request("GET", "/vector_store/list")
|
||||
assert stores.status_code == 200, stores.text
|
||||
assert _routed_features(owned.gateway) == at_boot, "feature requests changed the route table"
|
||||
|
||||
|
||||
def test_a_broken_optional_dependency_only_404s_its_own_feature_by_default(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with owned_proxy_process(gateway, tmp_path, _shadowed_dependency(tmp_path), remove_environment=(FLAG,)) as owned:
|
||||
guardrails: Final = owned.gateway.request("GET", "/guardrails/list")
|
||||
assert guardrails.status_code == 404, guardrails.text
|
||||
assert "guardrails" in _failed_features(owned), owned.log.read_text()
|
||||
listing: Final = owned.gateway.request("GET", "/v1/mcp/server")
|
||||
assert listing.status_code == 200, listing.text
|
||||
assert _routed_features(owned.gateway)["mcp_management"] != ()
|
||||
|
||||
|
||||
def test_disable_lazy_routes_flag_leaves_the_completion_endpoints_serving_every_client(
|
||||
gateway: Gateway, tmp_path: Path, provider: Wire
|
||||
) -> None:
|
||||
with owned_proxy_process(gateway, tmp_path, {FLAG: "true"}) as owned, owned.gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1")
|
||||
base_url: Final = str(owned.gateway.client.base_url)
|
||||
key: Final = owned.gateway.key
|
||||
markers: Final = tuple(_marker() for _ in range(6))
|
||||
|
||||
completion: Final = openai.OpenAI(
|
||||
base_url=base_url + "/v1", api_key=key, max_retries=0
|
||||
).chat.completions.create(model=model, messages=[{"role": "user", "content": markers[0]}])
|
||||
assert (completion.id, completion.choices[0].message.content) == (f"chatcmpl-{markers[0]}", "lazy ok")
|
||||
|
||||
assert asyncio.run(_stream_chat(base_url, key, model, markers[1])) == (
|
||||
frozenset({f"chatcmpl-{markers[1]}"}),
|
||||
"lazy ok",
|
||||
)
|
||||
|
||||
message: Final = anthropic.Anthropic(base_url=base_url, api_key=key, max_retries=0).messages.create(
|
||||
model=model, max_tokens=16, messages=[{"role": "user", "content": markers[2]}]
|
||||
)
|
||||
assert [block.text for block in message.content if block.type == "text"] == ["lazy ok"]
|
||||
|
||||
assert asyncio.run(_stream_message(base_url, key, model, markers[3])) == "lazy ok"
|
||||
|
||||
responded: Final = owned.gateway.request("POST", "/v1/responses", {"model": model, "input": markers[4]})
|
||||
assert responded.status_code == 200, responded.text
|
||||
response: Final = object_value(JSON.validate_json(responded.content))
|
||||
assert (response["status"], response["object"]) == ("completed", "response"), responded.text
|
||||
assert "lazy ok" in responded.text, responded.text
|
||||
|
||||
streamed: Final = owned.gateway.request(
|
||||
"POST", "/v1/responses", {"model": model, "input": markers[5], "stream": True}
|
||||
)
|
||||
assert streamed.status_code == 200, streamed.text
|
||||
assert "response.completed" in streamed.text and "lazy ok" in streamed.text, streamed.text
|
||||
|
||||
reached: Final = tuple(request.target for request in provider.drain() if MARKER.search(request.body))
|
||||
assert len(reached) == 6, reached
|
||||
|
||||
|
||||
def test_disable_lazy_routes_flag_route_table_survives_a_boot_burst_and_a_killed_worker(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
probes: Final = ("/routes", "/v1/mcp/server", "/openapi.json", "/guardrails/list", "/vector_store/list") * 8
|
||||
with owned_proxy_process(gateway, tmp_path, {FLAG: "true"}, workers=2) as owned:
|
||||
at_boot: Final = _routed_features(owned.gateway)
|
||||
unregistered: Final = sorted(name for name, paths in at_boot.items() if not paths)
|
||||
assert unregistered == [], f"features still missing from /routes at startup: {unregistered}"
|
||||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||||
statuses: Final = tuple(pool.map(partial(_status, owned.gateway), probes))
|
||||
assert statuses == (200,) * len(probes), statuses
|
||||
assert _routed_features(owned.gateway) == at_boot, "the boot burst changed the route table"
|
||||
|
||||
victim: Final = eventually(lambda: _workers(owned), lambda workers: len(workers) == 2)[0]
|
||||
victim.kill()
|
||||
with httpx.Client(base_url=owned.gateway.client.base_url, timeout=15, trust_env=False) as fresh:
|
||||
survivor: Final = Gateway(fresh, owned.gateway.key, owned.gateway.upstream_url)
|
||||
during: Final = tuple(survivor.request("GET", "/v1/mcp/server").status_code for _ in range(10))
|
||||
assert during == (200,) * 10, during
|
||||
respawned: Final = eventually(
|
||||
lambda: frozenset(worker.pid for worker in _workers(owned)),
|
||||
lambda pids: len(pids) == 2 and victim.pid not in pids,
|
||||
seconds=30,
|
||||
)
|
||||
assert f"Child process [{victim.pid}] died" in owned.log.read_text(), respawned
|
||||
tables: Final = tuple(_routed_features(owned.gateway) for _ in range(4))
|
||||
assert all(table == at_boot for table in tables), "the respawned worker disagrees on the route table"
|
||||
|
||||
|
||||
def test_disable_lazy_routes_flag_yields_the_same_route_table_after_a_restart(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with owned_proxy_process(gateway, tmp_path, {FLAG: "true"}) as first:
|
||||
table: Final = _paths(first.gateway)
|
||||
with owned_proxy_process(gateway, tmp_path, {FLAG: "true"}) as second:
|
||||
assert _paths(second.gateway) == table
|
||||
assert all(_routed_features(second.gateway).values()), "a feature is missing after restart"
|
||||
|
||||
|
||||
SELF_HOSTED_LANGFUSE: Final = "/self-hosted-langfuse"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ConfiguredFeatures:
|
||||
alias: str
|
||||
config: Path
|
||||
policy: Wire
|
||||
langfuse: Wire
|
||||
peer: McpPeer
|
||||
|
||||
|
||||
def _allow(request: Request) -> Reply:
|
||||
return Reply(body=json.dumps({"action": "NONE"}).encode())
|
||||
|
||||
|
||||
def _langfuse_health(request: Request) -> Reply:
|
||||
return Reply(body=json.dumps({"status": "OK"}).encode())
|
||||
|
||||
|
||||
def _config_declaring(directory: Path, alias: str, policy: Wire, langfuse: Wire, peer: McpPeer) -> Path:
|
||||
base: Final = object_value(
|
||||
JSON.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
|
||||
)
|
||||
config: Final = {
|
||||
**base,
|
||||
"guardrails": [
|
||||
{
|
||||
"guardrail_name": alias,
|
||||
"litellm_params": {
|
||||
"guardrail": "generic_guardrail_api",
|
||||
"mode": "pre_call",
|
||||
"default_on": True,
|
||||
"api_base": policy.url,
|
||||
"api_key": "synthetic-guardrail-key",
|
||||
},
|
||||
}
|
||||
],
|
||||
"mcp_servers": {alias: peer.registration()},
|
||||
"general_settings": {
|
||||
**object_value(base["general_settings"]),
|
||||
"pass_through_endpoints": [
|
||||
{
|
||||
"path": "/langfuse",
|
||||
"target": langfuse.url + SELF_HOSTED_LANGFUSE,
|
||||
"include_subpath": True,
|
||||
"auth": True,
|
||||
}
|
||||
],
|
||||
},
|
||||
}
|
||||
path: Final = directory / "configured-features.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _configured_features(directory: Path) -> Iterator[_ConfiguredFeatures]:
|
||||
alias: Final = "lazyroutes" + uuid.uuid4().hex[:8]
|
||||
with (
|
||||
wire_server(_allow) as policy,
|
||||
wire_server(_langfuse_health) as langfuse,
|
||||
scripted_peer(echo_tool("add")) as peer,
|
||||
):
|
||||
config: Final = _config_declaring(directory, alias, policy, langfuse, peer)
|
||||
yield _ConfiguredFeatures(alias, config, policy, langfuse, peer)
|
||||
|
||||
|
||||
def _config_server_id(candidate: Gateway, alias: str) -> str:
|
||||
servers: Final = JSON.validate_json(candidate.request("GET", "/v1/mcp/server").content)
|
||||
assert isinstance(servers, list), servers
|
||||
return next(
|
||||
string_value(object_value(server)["server_id"])
|
||||
for server in servers
|
||||
if object_value(server)["server_name"] == alias
|
||||
)
|
||||
|
||||
|
||||
def _assert_config_declared_features_serve(owned: OwnedProxy, features: _ConfiguredFeatures, provider: Wire) -> None:
|
||||
marker: Final = _marker()
|
||||
with owned.gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=provider.url + "/v1")
|
||||
completion: Final = owned.gateway.request(
|
||||
"POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": marker}]}
|
||||
)
|
||||
assert completion.status_code == 200, completion.text
|
||||
screened: Final = [request for request in features.policy.drain() if marker.encode() in request.body]
|
||||
assert len(screened) == 1, "the config-declared guardrail did not screen the completion"
|
||||
assert len([request for request in provider.drain() if marker.encode() in request.body]) == 1
|
||||
|
||||
identity: Final = _config_server_id(owned.gateway, features.alias)
|
||||
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
||||
tool: Final = tool_names(owned.gateway, key, identity)["add"]
|
||||
features.peer.drain()
|
||||
called: Final = call_tool(owned.gateway, key, identity, tool, {"marker": marker})
|
||||
assert called.status_code == 200, called.text
|
||||
reached_peer: Final = [
|
||||
object_value(object_value(call["body"])["params"]) for call in tool_calls(features.peer.drain())
|
||||
]
|
||||
assert [(params["name"], params["arguments"]) for params in reached_peer] == [("add", {"marker": marker})]
|
||||
|
||||
forwarded: Final = owned.gateway.request("GET", "/langfuse/api/public/health")
|
||||
assert forwarded.status_code == 200, forwarded.text
|
||||
reached_langfuse: Final = tuple(request.target for request in features.langfuse.drain())
|
||||
assert reached_langfuse == (SELF_HOSTED_LANGFUSE + "/api/public/health",), (
|
||||
f"the config pass-through for /langfuse lost to the built-in Langfuse route: {reached_langfuse}"
|
||||
)
|
||||
|
||||
|
||||
def test_disable_lazy_routes_flag_serves_config_declared_features_like_the_warmed_lazy_proxy(
|
||||
gateway: Gateway, tmp_path: Path, provider: Wire
|
||||
) -> None:
|
||||
with _configured_features(tmp_path) as features:
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=features.config, remove_environment=(FLAG,)) as lazy:
|
||||
_assert_config_declared_features_serve(lazy, features, provider)
|
||||
_warm_every_feature(lazy.gateway)
|
||||
warmed_paths: Final = tuple(path for path in _paths(lazy.gateway) if path != WARMUP_ROUTE)
|
||||
warmed_openapi: Final = _openapi_paths(lazy.gateway)
|
||||
with owned_proxy_process(gateway, tmp_path, {FLAG: "true"}, config=features.config) as eager:
|
||||
_assert_config_declared_features_serve(eager, features, provider)
|
||||
assert _paths(eager.gateway) == warmed_paths
|
||||
assert _openapi_paths(eager.gateway) == warmed_openapi
|
||||
211
tests/unit/proxy/test__lazy_features.py
Normal file
211
tests/unit/proxy/test__lazy_features.py
Normal file
|
|
@ -0,0 +1,211 @@
|
|||
import sys
|
||||
from collections.abc import AsyncGenerator, Mapping
|
||||
from contextlib import asynccontextmanager
|
||||
from types import ModuleType
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from fastapi import APIRouter, FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
from pydantic import BaseModel
|
||||
|
||||
from litellm.proxy._lazy_features import (
|
||||
LazyFeature,
|
||||
LazyFeatureMiddleware,
|
||||
attach_lazy_features,
|
||||
lazy_tag_to_prefix,
|
||||
loaded_lazy_modules,
|
||||
)
|
||||
|
||||
FLAG: Final = "LITELLM_DISABLE_LAZY_ROUTES"
|
||||
WARMUP_PATH: Final = "/lazy/warm/{name}"
|
||||
|
||||
|
||||
class _Operation(BaseModel):
|
||||
tags: tuple[str, ...]
|
||||
|
||||
|
||||
class _WarmupBody(BaseModel):
|
||||
stub_path: str
|
||||
paths: Mapping[str, Mapping[str, _Operation]]
|
||||
|
||||
|
||||
def _feature_module(monkeypatch: pytest.MonkeyPatch, name: str, path: str) -> LazyFeature:
|
||||
async def served() -> dict[str, str]:
|
||||
return {"feature": name}
|
||||
|
||||
router: Final = APIRouter()
|
||||
router.add_api_route(path, served, methods=["GET"])
|
||||
module: Final = ModuleType(f"tests.unit.proxy.lazy_fixture_{name}")
|
||||
module.router = router # pyright: ignore[reportAttributeAccessIssue] # fixture module built at test time
|
||||
monkeypatch.setitem(sys.modules, module.__name__, module)
|
||||
return LazyFeature(name=name, module_path=module.__name__, path_prefixes=(path,))
|
||||
|
||||
|
||||
def _paths(app: FastAPI) -> tuple[str, ...]:
|
||||
return tuple(str(getattr(route, "path", "")) for route in app.routes)
|
||||
|
||||
|
||||
def _has_lazy_middleware(app: FastAPI) -> bool:
|
||||
return any(middleware.cls is LazyFeatureMiddleware for middleware in app.user_middleware)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ("1", "true", "TRUE", "yes", "on"))
|
||||
def test_flag_registers_every_feature_at_startup(monkeypatch: pytest.MonkeyPatch, value: str) -> None:
|
||||
monkeypatch.setenv(FLAG, value)
|
||||
features: Final = (
|
||||
_feature_module(monkeypatch, "alpha", "/alpha/list"),
|
||||
_feature_module(monkeypatch, "beta", "/beta/list"),
|
||||
)
|
||||
app: Final = FastAPI()
|
||||
|
||||
attach_lazy_features(app, features)
|
||||
|
||||
assert WARMUP_PATH not in _paths(app)
|
||||
assert not _has_lazy_middleware(app)
|
||||
assert loaded_lazy_modules(app) == set()
|
||||
with TestClient(app) as client:
|
||||
at_startup: Final = _paths(app)
|
||||
assert {"/alpha/list", "/beta/list"} <= set(at_startup)
|
||||
assert loaded_lazy_modules(app) == {features[0].module_path, features[1].module_path}
|
||||
assert client.get("/beta/list").json() == {"feature": "beta"}
|
||||
assert client.post("/lazy/warm/alpha").status_code == 404
|
||||
assert _paths(app) == at_startup, "first feature request changed the table"
|
||||
|
||||
|
||||
def test_flag_registers_before_the_inner_lifespan_and_after_late_routes(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv(FLAG, "true")
|
||||
features: Final = (_feature_module(monkeypatch, "epsilon", "/epsilon/{name}"),)
|
||||
seen_by_inner_lifespan: Final[list[tuple[str, ...]]] = [] # mutable-ok: captured from inside the lifespan
|
||||
|
||||
@asynccontextmanager
|
||||
async def inner_lifespan(app_: FastAPI) -> AsyncGenerator[None]:
|
||||
seen_by_inner_lifespan.append(_paths(app_))
|
||||
yield
|
||||
|
||||
async def late() -> dict[str, str]:
|
||||
return {"feature": "late"}
|
||||
|
||||
app: Final = FastAPI(lifespan=inner_lifespan)
|
||||
attach_lazy_features(app, features)
|
||||
app.add_api_route("/epsilon/list", late, methods=["GET"])
|
||||
|
||||
with TestClient(app) as client:
|
||||
assert client.get("/epsilon/list").json() == {"feature": "late"}, "late eager route must win, as in lazy mode"
|
||||
assert client.get("/epsilon/x").json() == {"feature": "epsilon"}
|
||||
assert seen_by_inner_lifespan == [_paths(app)], "startup hooks inside the proxy lifespan must see the full table"
|
||||
|
||||
|
||||
def test_flag_lets_a_route_added_during_startup_beat_an_overlapping_feature_route(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv(FLAG, "true")
|
||||
features: Final = (_feature_module(monkeypatch, "zeta", "/zeta/{endpoint:path}"),)
|
||||
|
||||
async def configured() -> dict[str, str]:
|
||||
return {"feature": "configured"}
|
||||
|
||||
@asynccontextmanager
|
||||
async def adds_a_pass_through(app_: FastAPI) -> AsyncGenerator[None]:
|
||||
app_.add_api_route("/zeta/{subpath:path}", configured, methods=["GET"])
|
||||
yield
|
||||
|
||||
app: Final = FastAPI(lifespan=adds_a_pass_through)
|
||||
attach_lazy_features(app, features)
|
||||
|
||||
with TestClient(app) as client:
|
||||
assert client.get("/zeta/health").json() == {"feature": "configured"}, (
|
||||
"lazy mode routes this to startup's route"
|
||||
)
|
||||
|
||||
|
||||
def test_flag_does_not_bring_back_a_feature_route_removed_during_startup(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv(FLAG, "true")
|
||||
features: Final = (
|
||||
_feature_module(monkeypatch, "eta", "/eta/list"),
|
||||
_feature_module(monkeypatch, "theta", "/theta/list"),
|
||||
)
|
||||
|
||||
@asynccontextmanager
|
||||
async def drops_eta(app_: FastAPI) -> AsyncGenerator[None]:
|
||||
app_.router.routes[:] = [route for route in app_.router.routes if getattr(route, "path", "") != "/eta/list"]
|
||||
yield
|
||||
|
||||
app: Final = FastAPI(lifespan=drops_eta)
|
||||
attach_lazy_features(app, features)
|
||||
|
||||
with TestClient(app) as client:
|
||||
assert client.get("/eta/list").status_code == 404
|
||||
assert client.get("/theta/list").json() == {"feature": "theta"}
|
||||
assert "/eta/list" not in _paths(app)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", (None, "", "0", "false", "off"))
|
||||
def test_without_the_flag_features_still_mount_on_first_request(
|
||||
monkeypatch: pytest.MonkeyPatch, value: str | None
|
||||
) -> None:
|
||||
if value is None:
|
||||
monkeypatch.delenv(FLAG, raising=False)
|
||||
else:
|
||||
monkeypatch.setenv(FLAG, value)
|
||||
features: Final = (_feature_module(monkeypatch, "gamma", "/gamma/list"),)
|
||||
app: Final = FastAPI()
|
||||
|
||||
attach_lazy_features(app, features)
|
||||
|
||||
assert "/gamma/list" not in _paths(app)
|
||||
assert WARMUP_PATH in _paths(app)
|
||||
assert _has_lazy_middleware(app)
|
||||
with TestClient(app) as client:
|
||||
assert client.get("/gamma/list").json() == {"feature": "gamma"}
|
||||
assert "/gamma/list" in _paths(app)
|
||||
|
||||
|
||||
def test_flag_keeps_registering_after_one_feature_fails_to_import(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv(FLAG, "true")
|
||||
broken: Final = LazyFeature(
|
||||
name="broken", module_path="tests.unit.proxy.lazy_fixture_does_not_exist", path_prefixes=("/broken",)
|
||||
)
|
||||
healthy: Final = _feature_module(monkeypatch, "delta", "/delta/list")
|
||||
app: Final = FastAPI()
|
||||
|
||||
attach_lazy_features(app, (broken, healthy))
|
||||
|
||||
with TestClient(app) as client:
|
||||
assert "/delta/list" in _paths(app)
|
||||
assert loaded_lazy_modules(app) == {broken.module_path, healthy.module_path}
|
||||
assert client.get("/delta/list").json() == {"feature": "delta"}
|
||||
assert client.get("/broken").status_code == 404
|
||||
|
||||
|
||||
def test_flag_hides_the_swagger_warmup_plugin(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
import litellm.proxy._lazy_openapi_snapshot as snapshot
|
||||
|
||||
monkeypatch.setattr(snapshot, "SNAPSHOT_FILE", snapshot.SNAPSHOT_FILE.with_name("missing-snapshot.json"))
|
||||
monkeypatch.setenv(FLAG, "false")
|
||||
assert lazy_tag_to_prefix() != {}, "control: without the flag and without a snapshot the plugin has tags"
|
||||
monkeypatch.setenv(FLAG, "true")
|
||||
assert lazy_tag_to_prefix() == {}
|
||||
|
||||
|
||||
def test_without_the_flag_the_warmup_route_registers_a_feature_and_returns_its_paths(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.delenv(FLAG, raising=False)
|
||||
features: Final = (
|
||||
_feature_module(monkeypatch, "alpha", "/alpha/list"),
|
||||
_feature_module(monkeypatch, "beta", "/beta/list"),
|
||||
)
|
||||
app: Final = FastAPI()
|
||||
attach_lazy_features(app, features)
|
||||
|
||||
with TestClient(app) as client:
|
||||
assert client.post("/lazy/warm/zeta").status_code == 404
|
||||
warmed: Final = client.post("/lazy/warm/alpha")
|
||||
assert warmed.status_code == 200, warmed.text
|
||||
body: Final = _WarmupBody.model_validate_json(warmed.text)
|
||||
assert body.stub_path == "/alpha/list"
|
||||
assert set(body.paths) == {"/alpha/list"}
|
||||
assert body.paths["/alpha/list"]["get"].tags == ("alpha",)
|
||||
assert loaded_lazy_modules(app) == {features[0].module_path}
|
||||
assert "/alpha/list" in _paths(app) and "/beta/list" not in _paths(app)
|
||||
Loading…
Add table
Reference in a new issue