swagger: stub-inject unloaded lazy features and warm on dropdown expand

This commit is contained in:
Michael Riad Zaky 2026-04-28 13:36:00 -07:00 • committed by Michael Riad Zaky
parent 0f8dd28542
commit adcab435a4
2 changed files with 73 additions and 2 deletions

View file

@ -8,8 +8,9 @@ omits each feature's routes until the feature is warmed.
import asyncio
import importlib
import sys
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Callable, Tuple
from typing import TYPE_CHECKING, Callable, Dict, Tuple
from starlette.types import Receive, Scope, Send
@ -46,6 +47,9 @@ class LazyFeature:
# For routes whose path has a leading parameter (e.g. /{server}/authorize)
# — startswith can't match those, so the matcher also checks endswith.
path_suffixes: Tuple[str, ...] = ()
# Keep the stub injected even after load — for mounted ASGI sub-apps
# whose routes don't appear in the parent app's openapi spec.
persistent_swagger_stub: bool = False
LAZY_FEATURES: Tuple[LazyFeature, ...] = (
@ -158,6 +162,7 @@ LAZY_FEATURES: Tuple[LazyFeature, ...] = (
module_path="litellm.proxy._experimental.mcp_server.server",
path_prefixes=("/mcp",),
register_fn=_mount_app("/mcp", attr_name="app"),
persistent_swagger_stub=True,
),
LazyFeature(
name="config_overrides",
@ -305,3 +310,33 @@ class LazyFeatureMiddleware:
def attach_lazy_features(app: "FastAPI") -> None:
app.add_middleware(LazyFeatureMiddleware, fastapi_app=app)
def inject_lazy_stubs(schema: Dict) -> Dict:
"""Stub openapi entries for unloaded features so Swagger renders sections."""
paths = schema.setdefault("paths", {})
for feat in LAZY_FEATURES:
if feat.module_path in sys.modules and not feat.persistent_swagger_stub:
continue
prefix = feat.path_prefixes[0]
if prefix in paths:
continue
paths[prefix] = {
"get": {
"tags": [feat.name],
"summary": feat.name,
"responses": {"200": {"description": "OK"}},
}
}
return schema
def lazy_tag_to_prefix() -> Dict[str, str]:
"""feature.name -> first prefix, used by the Swagger warmup JS plugin.
Excludes persistent-stub features (mounted sub-apps) — warming them
triggers a streaming hit and no useful new routes appear."""
return {
feat.name: feat.path_prefixes[0]
for feat in LAZY_FEATURES
if not feat.persistent_swagger_stub
}

View file

@ -1033,6 +1033,11 @@ def get_openapi_schema():
openapi_schema = CustomOpenAPISpec.add_llm_api_request_schema_body(openapi_schema)
# Stub unloaded lazy features so they appear as Swagger sections.
from litellm.proxy._lazy_features import inject_lazy_stubs
openapi_schema = inject_lazy_stubs(openapi_schema)
# Fix Swagger UI execute path error when server_root_path is set
if server_root_path:
openapi_schema["servers"] = [{"url": "/" + server_root_path.strip("/")}]
@ -1059,6 +1064,11 @@ def custom_openapi():
openapi_schema = CustomOpenAPISpec.add_llm_api_request_schema_body(openapi_schema)
# Stub unloaded lazy features so they appear as Swagger sections.
from litellm.proxy._lazy_features import inject_lazy_stubs
openapi_schema = inject_lazy_stubs(openapi_schema)
# Fix Swagger UI execute path error when server_root_path is set
if server_root_path:
openapi_schema["servers"] = [{"url": "/" + server_root_path.strip("/")}]
@ -1486,14 +1496,40 @@ def mount_swagger_ui():
app.mount("/swagger", StaticFiles(directory=swagger_directory), name="swagger")
# On dropdown expand: one-time fetch to the prefix (triggers lazy load),
# then spec re-download so real routes replace the stub. Raw JS (no
# <script> tag) since it's injected inside the existing inline script.
from fastapi.responses import HTMLResponse
from litellm.proxy._lazy_features import lazy_tag_to_prefix
_lazy_plugin_js = (
"const TAG_TO_PREFIX = "
+ json.dumps(lazy_tag_to_prefix())
+ ";const warmedTags = new Set();const LazyLoadPlugin = () => ({"
"statePlugins:{layout:{wrapActions:{show:(ori,sys)=>(...args)=>{"
"const thing=args[0];let tag=null;"
"if(Array.isArray(thing)){for(const t of thing)if(TAG_TO_PREFIX[t])tag=t;}"
"if(tag&&!warmedTags.has(tag)){warmedTags.add(tag);"
"fetch(TAG_TO_PREFIX[tag]).finally(()=>setTimeout(()=>sys.specActions.download(),800));}"
"return ori(...args);}}}}});"
)
def swagger_monkey_patch(*args, **kwargs):
return get_swagger_ui_html(
response = get_swagger_ui_html(
*args,
**kwargs,
swagger_js_url=f"{custom_root_path_swagger_path}/swagger-ui-bundle.js",
swagger_css_url=f"{custom_root_path_swagger_path}/swagger-ui.css",
swagger_favicon_url=f"{custom_root_path_swagger_path}/favicon.png",
)
body = response.body.decode("utf-8")
body = body.replace(
"const ui = SwaggerUIBundle({",
_lazy_plugin_js + "const ui = SwaggerUIBundle({plugins:[LazyLoadPlugin],",
1,
)
return HTMLResponse(content=body)
applications.get_swagger_ui_html = swagger_monkey_patch