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

Author: Michael Riad Zaky

litellm/proxy/_lazy_features.py

@@ -8,8 +8,9 @@
 
 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 @@ class 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 @@ async def _load(self, feat: LazyFeature) -> None:
 
 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
+    }

litellm/proxy/proxy_server.py

@@ -1019,6 +1019,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("/")}]
@@ -1045,6 +1050,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("/")}]
@@ -1472,14 +1482,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