[Feat] Lazy-load optional feature routers on first request

[Michael-RZ-Berri]

Relevant issues

Memory usage has gone up from prior versions due to added routes and unconditional imports, this instead lazy loads the new imports on first route usage to save on memory from unused modules.

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have Added testing in the tests/test_litellm/ directory, Adding at least 1 test is a hard requirement - see details
• My PR passes all unit tests on make test-unit
• My PR's scope is as isolated as possible, it only solves 1 specific problem
• I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

Screenshots / Proof of Fix

Down ~700mb on docker with two workers between lazy loading and baseline.

Guardrails UI usable after being lazy loaded.

Type

🆕 New Feature
🚄 Infrastructure
✅ Test

Changes

New lazy_features file and changes to imports in proxy_server.

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.

---

Michael Riad Zaky seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[greptile-apps]

Greptile Summary

This PR introduces LazyFeatureMiddleware — an ASGI middleware that defers importing ~30 optional feature routers until the first HTTP/WebSocket request matching their path prefix, cutting idle memory by ~700 MB. The implementation addresses previous review concerns: imports are off-loaded to a thread executor, per-feature locks allow independent features to load concurrently, and WebSocket scopes are handled alongside HTTP scopes.

Two minor path-matching issues remain:

• \"/.well-known/oauth-\" appears in both mcp_byok_oauth and mcp_discoverable prefixes, causing both routers to load for any /.well-known/oauth-* request and creating a route-ordering dependency for overlapping endpoints.
• policies uses the prefix \"/policy/\" (trailing slash), which won't match a bare GET /policy request if such a route exists in the router.

Confidence Score: 5/5

Safe to merge; only P2 style/edge-case concerns remain.

No P0 or P1 issues found. All blocking-import, global-lock, WebSocket bypass, and vector-store prefix concerns from prior reviews are addressed in this revision. Two remaining P2 findings affect uncommon edge cases and do not break any currently tested paths.

litellm/proxy/_lazy_features.py — review the mcp_byok_oauth/mcp_discoverable prefix overlap and the policies trailing-slash prefix before merging if those OAuth and policy endpoints are actively used.

Important Files Changed

Filename	Overview
litellm/proxy/_lazy_features.py	New file implementing lazy ASGI middleware for optional feature routers. Core logic is solid (per-feature locks, thread-executor import, websocket support, idempotent loading). Two P2 concerns: /.well-known/oauth- prefix listed in both mcp_byok_oauth and mcp_discoverable causes both to load together with potential route-ordering side-effects; policies trailing slash /policy/ may miss requests to a bare /policy route.
litellm/proxy/proxy_server.py	Removes ~30 eager module imports and app.include_router/app.mount calls, replacing them with attach_lazy_features(app). Google router kept eager intentionally due to /models/{name}:method path overlap. In-function lazy imports added for global_mcp_tool_registry, global_agent_registry, and append_agents_to_model_* helpers. No regressions observed.
tests/test_litellm/proxy/test_proxy_server.py	Adds three test classes covering registry shape, absence of heavy modules at startup (subprocess), and middleware behaviour (idempotent load, concurrent first-requests, failing import handling). Good coverage of the new mechanism.
tests/proxy_unit_tests/test_proxy_routes.py	Force-loads all lazy features synchronously before asserting route coverage. Continues on per-feature import failure so CI doesn't break when optional extras are absent.
tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py	Adds a guard that force-registers the vector_store_management lazy feature before asserting its routes exist. Logic is idempotent via the already_registered sentinel check.

Sequence Diagram

sequenceDiagram
    participant Client
    participant LazyFeatureMiddleware
    participant ThreadExecutor
    participant FeatureModule
    participant FastAPI

    Client->>LazyFeatureMiddleware: HTTP/WS request (path="/guardrails/...")
    activate LazyFeatureMiddleware
    LazyFeatureMiddleware->>LazyFeatureMiddleware: check path prefixes → match "guardrails"
    LazyFeatureMiddleware->>LazyFeatureMiddleware: acquire per-feature asyncio.Lock
    LazyFeatureMiddleware->>ThreadExecutor: run_in_executor(importlib.import_module)
    ThreadExecutor->>FeatureModule: import litellm.proxy.guardrails.guardrail_endpoints
    FeatureModule-->>ThreadExecutor: module object
    ThreadExecutor-->>LazyFeatureMiddleware: module
    LazyFeatureMiddleware->>FastAPI: register_fn(app, module) → app.include_router(router)
    LazyFeatureMiddleware->>LazyFeatureMiddleware: mark module_path loaded, clear openapi_schema
    LazyFeatureMiddleware->>LazyFeatureMiddleware: release lock
    LazyFeatureMiddleware->>FastAPI: forward request
    FastAPI-->>Client: response (newly registered route handles it)
    deactivate LazyFeatureMiddleware

    Note over LazyFeatureMiddleware: Subsequent requests skip load (module_path in _loaded)

Loading

_{Reviews (12): Last reviewed commit: "lazy-load optional feature routers on fi..." | Re-trigger Greptile}

[greptile-apps]

importlib.import_module blocks the event loop

importlib.import_module is a synchronous, blocking call. For heavy modules the docstring itself estimates 1–3 s (e.g. MCP). Because this runs on the asyncio event loop thread with no await, the entire event loop is frozen for that duration — every other in-flight request is stalled until the import finishes.

The fix is to off-load the import to a thread executor and keep only the synchronous register_fn call on the event loop:

loop = asyncio.get_event_loop()
module = await loop.run_in_executor(None, importlib.import_module, feat.module_path)
feat.register_fn(self._fastapi_app, module)

Note that register_fn must remain on the event loop thread since it mutates self._fastapi_app.router.routes (not thread-safe to call from the executor).

[Michael-RZ-Berri]

Addressed, moved to await

[greptile-apps]

Global lock serialises all concurrent first-loads

A single asyncio.Lock is used for every feature. If two unrelated features need to be loaded concurrently (e.g. a request to /guardrails and a request to /mcp arrive simultaneously), the second load is blocked waiting for the first, even though they touch completely different modules and routes. Once import_module is moved to a thread executor (see the blocking-call comment above), this serialisation becomes a real latency bottleneck.

Consider one lock-per-feature — a dict[str, asyncio.Lock] keyed by module_path — so independent features load in parallel.

[Michael-RZ-Berri]

Addressed, locks set

[greptile-apps]

/models prefix is too broad and triggers unintended lazy-loads

/models matches any request whose path starts with that string. The eagerly-registered model_management_router likely has routes under /model (singular), but if any eagerly-registered router exposes paths under /models, every such request will trigger the import of google_endpoints.endpoints as a side-effect — defeating the memory savings for that module and adding unnecessary import overhead to unrelated requests.

Consider a more specific prefix (e.g., /models/ with a trailing slash) to narrow the trigger to actual Google passthrough calls like GET /models/{name}:generateContent.

[Michael-RZ-Berri]

Addressed, this isn't lazy loaded

[greptile-apps]

Silent failure on import error surfaces as opaque 404

When importlib.import_module or register_fn raises, the module path is marked as loaded and subsequent requests to that prefix will permanently receive a 404 with no explanation. If a required optional dependency is missing (e.g. mcp extras not installed), users will see silent 404s with no error message indicating the feature is unavailable.

Consider returning a 503/501 response for paths belonging to a failed feature, or at minimum including the feature name in the warning so operators can connect the 404 to the startup warning log.

[Michael-RZ-Berri]

Addressed, included warning

[greptile-apps]

policy_engine and policy_resolve prefixes are subsets of policies

policy_engine is triggered by /policies and policy_resolve by /policies/resolve. Since /policies/resolve.startswith(/policies) is true, every request to /policies/resolve/... loads both modules. If this is intentional, a comment would help — but if policy_resolve only handles the more specific sub-paths, it should use only those prefixes and drop the overlap.

[Michael-RZ-Berri]

Left in, this is intentional

[veria-ai]

Lazy-load optional feature routers on first request

This PR introduces a middleware that defers importing and registering ~30 optional feature routers until the first HTTP request matching their path prefix arrives. The module paths are hardcoded in a static tuple — no user input flows into importlib.import_module. Route authentication (via Depends(user_api_key_auth)) is preserved on each lazily-registered router. The path_suffixes matcher for MCP discoverable endpoints (/authorize, /token, /register) may trigger imports more broadly than intended, but this only affects import timing, not route accessibility or auth enforcement.

---

Status: 0 open
Risk: 2/10

[yuneng-berri]

If this is true, I am concerned that this will break the UI. Is this recoverable from the UI? Can we get more confidence in how these changes interact with the UI before we merge this?

[Michael-RZ-Berri]

Yeah I went and tested on the UI, the only place the openAPI json directly gets used in the UI is the optional features dropdown in the add key modal which gets a little latency added. I also went and tested on the guardrails page that it still works with lazy loading, seems fine.

[greptile-apps]

vector_store_files missing /vector_stores prefix causes 404 on first non-versioned file request

The vector_store_files LazyFeature only lists /v1/vector_stores as a path prefix. The vector_store_files_endpoints/endpoints.py router exposes routes under both /v1/vector_stores/{id}/files and /vector_stores/{id}/files (i.e. without the /v1/ prefix). A first request to e.g. POST /vector_stores/{id}/files matches the vector_stores feature (which loads vector_store_endpoints.endpoints) but does not match vector_store_files, so vector_store_files_endpoints.endpoints is never registered. Starlette then returns 404 — the route doesn't exist until a subsequent /v1/vector_stores/... request triggers the load.

LazyFeature(
    name="vector_store_files",
    module_path="litellm.proxy.vector_store_files_endpoints.endpoints",
-    path_prefixes=("/v1/vector_stores",),
+    path_prefixes=("/v1/vector_stores", "/vector_stores"),
),

[Michael-RZ-Berri]

Addressed, fixed as suggested

[greptile-apps]

Non-deterministic MCP route registration order can break /mcp/oauth and /mcp/discoverable

When the first MCP request targets a path that only matches mcp_app (e.g. /mcp or /mcp/sse), the middleware loads mcp_app and appends a Mount("/mcp", ...) to app.router.routes. A subsequent request to /mcp/oauth then loads mcp_byok_oauth, but its explicit routes are appended after the already-registered mount. Starlette iterates app.router.routes in order and the mount at /mcp returns Match.FULL for /mcp/oauth, so the ASGI sub-app handles the request instead of mcp_byok_oauth's router — silently returning the wrong response.

The same applies to mcp_discoverable with /mcp/discoverable.

The safe fix is to have mcp_app's load trigger also force-load mcp_byok_oauth and mcp_discoverable at the same time, or to add /mcp/oauth and /mcp/discoverable as prefixes of mcp_app so all three are always loaded together (guaranteeing the sub-routers are registered first). Alternatively, simply drop the non-versioned /mcp/oauth and /mcp/discoverable path prefixes (since /v1/mcp/... variants already exist) to eliminate the overlap entirely.

[Michael-RZ-Berri]

Addressed by dropping so there's no overlap

[greptile-apps]

WebSocket connections bypass lazy loading for mcp_app

LazyFeatureMiddleware.__call__ only triggers lazy loading when scope["type"] == "http". WebSocket upgrade requests arrive as scope["type"] == "websocket" per the ASGI spec, so they pass through the guard unmodified. If the very first request to /mcp is a WebSocket connection (common for many MCP clients), mcp_app won't be mounted yet and the connection will be rejected with a 404/500.

SSE-based MCP clients (which use regular HTTP) are fine, but WebSocket-first clients are silently broken on cold-start. Consider extending the check to also handle "websocket" scopes:

if scope["type"] in ("http", "websocket") and len(self._loaded) < len(self._features):

[Michael-RZ-Berri]

This got added in, it's addressed

[Michael-RZ-Berri]

@greptile-ai