add openapi snapshotting to CI pipeline

.github/workflows/check-lazy-openapi-snapshot.yml

@@ -0,0 +1,75 @@
+name: Check Lazy OpenAPI Snapshot
+
+on:
+  pull_request:
+    branches:
+      - main
+      - litellm_internal_staging
+      - "litellm_**"
+
+permissions:
+  contents: read
+  checks: write
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  verify:
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    steps:
+      - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
+        with:
+          persist-credentials: false
+
+      - name: Set up Python
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.12"
+
+      - name: Set up uv
+        uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
+        with:
+          version: "0.10.9"
+
+      - name: Cache uv dependencies
+        uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
+        with:
+          path: |
+            ~/.cache/uv
+            .venv
+          key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
+          restore-keys: |
+            ${{ runner.os }}-uv-
+
+      - name: Install dependencies
+        run: uv sync --frozen --all-groups --all-extras
+
+      - name: Regenerate snapshot to /tmp
+        id: regen
+        run: |
+          cp litellm/proxy/_lazy_openapi_snapshot.json /tmp/snapshot.committed.json
+          uv run --no-sync python -m litellm.proxy._lazy_openapi_snapshot
+          mv litellm/proxy/_lazy_openapi_snapshot.json /tmp/snapshot.fresh.json
+          mv /tmp/snapshot.committed.json litellm/proxy/_lazy_openapi_snapshot.json
+
+      - name: Compare
+        id: diff
+        continue-on-error: true
+        run: |
+          diff -q /tmp/snapshot.fresh.json litellm/proxy/_lazy_openapi_snapshot.json
+
+      - name: Mark neutral if drift
+        if: steps.diff.outcome == 'failure'
+        uses: LouisBrunner/checks-action@v2.0.0
+        with:
+          token: ${{ secrets.GITHUB_TOKEN }}
+          name: lazy-openapi-snapshot
+          conclusion: neutral
+          output: |
+            {
+              "title": "Lazy openapi snapshot is stale",
+              "summary": "Run `python -m litellm.proxy._lazy_openapi_snapshot` and commit the regenerated `litellm/proxy/_lazy_openapi_snapshot.json`. Not blocking — the snapshot will regenerate at release if not committed."
+            }

litellm/proxy/_lazy_features.py

@@ -309,15 +309,81 @@ class LazyFeatureMiddleware:
 
 
 def attach_lazy_features(app: "FastAPI") -> None:
+    app.include_router(_make_warmup_router(app))
     app.add_middleware(LazyFeatureMiddleware, fastapi_app=app)
 
 
+def _make_warmup_router(app: "FastAPI") -> "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."""
+    from fastapi import APIRouter, HTTPException
+    from fastapi.openapi.utils import get_openapi
+
+    router = APIRouter()
+
+    @router.post("/lazy/warm/{name}", include_in_schema=False)
+    async def warm(name: str):
+        feat = next((f for f in LAZY_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": {}}}
+
+        already = any(
+            any(getattr(r, "path", "").startswith(p) for p in feat.path_prefixes)
+            for r in app.routes
+        )
+        if not already:
+            loop = asyncio.get_running_loop()
+            module = await loop.run_in_executor(
+                None, importlib.import_module, feat.module_path
+            )
+            feat.register_fn(app, module)
+            app.openapi_schema = None
+
+        feat_routes = [
+            r
+            for r in app.routes
+            if any(getattr(r, "path", "").startswith(p) for p in feat.path_prefixes)
+        ]
+        full = get_openapi(title=app.title, version=app.version, routes=feat_routes)
+        # Force all operations under one tag so they group under a single Swagger
+        # section — many lazy modules tag routes inconsistently.
+        for path_ops in full.get("paths", {}).values():
+            for op in path_ops.values():
+                if isinstance(op, dict):
+                    op["tags"] = [feat.name]
+        return {
+            "stub_path": feat.path_prefixes[0],
+            "paths": full.get("paths", {}),
+            "components": {"schemas": full.get("components", {}).get("schemas", {})},
+        }
+
+    return router
+
+
 def inject_lazy_stubs(schema: Dict) -> Dict:
-    """Stub openapi entries for unloaded features so Swagger renders sections."""
+    """Inject openapi entries for unloaded features. Uses the snapshot file
+    when available (full route info), otherwise falls back to a single
+    placeholder per feature."""
+    from litellm.proxy._lazy_openapi_snapshot import load_snapshot
+
+    snapshot = load_snapshot()
     paths = schema.setdefault("paths", {})
+    schemas = schema.setdefault("components", {}).setdefault("schemas", {})
+
     for feat in LAZY_FEATURES:
         if feat.module_path in sys.modules and not feat.persistent_swagger_stub:
             continue
+
+        fragment = (snapshot or {}).get(feat.name)
+        if fragment:
+            for p, ops in fragment.get("paths", {}).items():
+                paths.setdefault(p, ops)
+            for name, sch in fragment.get("components", {}).get("schemas", {}).items():
+                schemas.setdefault(name, sch)
+            continue
+
         prefix = feat.path_prefixes[0]
         if prefix in paths:
             continue
@@ -333,8 +399,12 @@ def inject_lazy_stubs(schema: Dict) -> Dict:
 
 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."""
+    Returns empty when the snapshot is loaded — the plugin is unnecessary
+    because /openapi.json already has full route info."""
+    from litellm.proxy._lazy_openapi_snapshot import load_snapshot
+
+    if load_snapshot():
+        return {}
     return {
         feat.name: feat.path_prefixes[0]
         for feat in LAZY_FEATURES

litellm/proxy/_lazy_openapi_snapshot.py

@@ -0,0 +1,70 @@
+"""
+Per-feature OpenAPI snapshot for lazy-loaded routers.
+
+The committed JSON is generated by `python -m litellm.proxy._lazy_openapi_snapshot`
+and consumed at runtime so /openapi.json can show full route info for unloaded
+features without importing them. CI verifies the file is current and surfaces
+any drift as a neutral check.
+"""
+
+import json
+import sys
+from pathlib import Path
+from typing import Dict, Optional
+
+SNAPSHOT_FILE = Path(__file__).parent / "_lazy_openapi_snapshot.json"
+
+
+def load_snapshot() -> Optional[Dict[str, Dict]]:
+    if not SNAPSHOT_FILE.exists():
+        return None
+    try:
+        with SNAPSHOT_FILE.open() as f:
+            return json.load(f)
+    except (json.JSONDecodeError, OSError):
+        return None
+
+
+def generate_snapshot() -> Dict[str, Dict]:
+    import importlib
+
+    from fastapi.openapi.utils import get_openapi
+
+    from litellm.proxy._lazy_features import LAZY_FEATURES
+    from litellm.proxy.proxy_server import app
+
+    for feat in LAZY_FEATURES:
+        if feat.module_path in sys.modules:
+            continue
+        try:
+            module = importlib.import_module(feat.module_path)
+            feat.register_fn(app, module)
+        except Exception as exc:
+            print(f"warning: skip {feat.name}: {exc}", file=sys.stderr)
+
+    fragments: Dict[str, Dict] = {}
+    for feat in LAZY_FEATURES:
+        feat_routes = [
+            r
+            for r in app.routes
+            if any(getattr(r, "path", "").startswith(p) for p in feat.path_prefixes)
+        ]
+        if not feat_routes:
+            continue
+        full = get_openapi(title=app.title, version=app.version, routes=feat_routes)
+        # Group all of a feature's routes under one tag.
+        for path_ops in full.get("paths", {}).values():
+            for op in path_ops.values():
+                if isinstance(op, dict):
+                    op["tags"] = [feat.name]
+        fragments[feat.name] = {
+            "paths": full.get("paths", {}),
+            "components": {"schemas": full.get("components", {}).get("schemas", {})},
+        }
+    return fragments
+
+
+if __name__ == "__main__":
+    fragments = generate_snapshot()
+    SNAPSHOT_FILE.write_text(json.dumps(fragments, indent=2, sort_keys=True) + "\n")
+    print(f"wrote {len(fragments)} feature fragments to {SNAPSHOT_FILE}")

litellm/proxy/proxy_server.py

@@ -1504,14 +1504,51 @@ def mount_swagger_ui():
     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 = () => ({"
+        "const TAG_TO_PREFIX = " + json.dumps(lazy_tag_to_prefix()) + ";"
+        "const warmedTags = new Set();"
+        "const LAZY_TAGS = new Set(Object.keys(TAG_TO_PREFIX));"
+        "const hideStubRows = () => {"
+        "document.querySelectorAll('.opblock').forEach(op => {"
+        "const d = op.querySelector('.opblock-summary-description');"
+        "if (d && LAZY_TAGS.has(d.textContent.trim())) op.style.display = 'none';"
+        "});};"
+        "const annotateLazyHeaders = () => {"
+        "document.querySelectorAll('.opblock-tag').forEach(tagEl => {"
+        "const m = (tagEl.id || '').match(/^operations-tag-(.+)$/);"
+        "if (!m || !LAZY_TAGS.has(m[1])) return;"
+        "const existing = tagEl.querySelector('.lazy-load-hint');"
+        "if (warmedTags.has(m[1])) { if (existing) existing.remove(); return; }"
+        "if (existing) return;"
+        "const hint = document.createElement('small');"
+        "hint.className = 'lazy-load-hint';"
+        "hint.textContent = ' (expand to load routes)';"
+        "hint.style.opacity = '0.6';"
+        "hint.style.marginLeft = '6px';"
+        "const target = tagEl.querySelector('a span') || tagEl.querySelector('span') || tagEl;"
+        "target.appendChild(hint);"
+        "});};"
+        "setInterval(() => { hideStubRows(); annotateLazyHeaders(); }, 200);"
+        "const LazyLoadPlugin = () => ({"
+        "afterLoad:function(system){setTimeout(()=>{"
+        "for(const tag of LAZY_TAGS)system.layoutActions.show(['operations-tag',tag],false);"
+        "},200);},"
         "statePlugins:{layout:{wrapActions:{show:(ori,sys)=>(...args)=>{"
-        "const thing=args[0];let tag=null;"
+        "const thing=args[0];const shown=args[1];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));}"
+        "if(shown!==false&&tag&&!warmedTags.has(tag)){warmedTags.add(tag);"
+        "fetch('/lazy/warm/'+tag,{method:'POST'}).then(r=>r.json()).then(d=>{"
+        "if(!d.paths||Object.keys(d.paths).length===0)return;"
+        "const cur=sys.specSelectors.specJson().toJS();"
+        "const merged={};let inserted=false;"
+        "for(const k in (cur.paths||{})){"
+        "if(k===d.stub_path){for(const nk in d.paths)merged[nk]=d.paths[nk];inserted=true;}"
+        "else{merged[k]=cur.paths[k];}}"
+        "if(!inserted)Object.assign(merged,d.paths);"
+        "cur.paths=merged;"
+        "cur.components=cur.components||{};"
+        "cur.components.schemas=Object.assign(cur.components.schemas||{},(d.components||{}).schemas||{});"
+        "sys.specActions.updateSpec(JSON.stringify(cur));"
+        "}).catch(()=>{});}"
         "return ori(...args);}}}}});"
     )
 
@@ -1526,7 +1563,8 @@ def mount_swagger_ui():
         body = response.body.decode("utf-8")
         body = body.replace(
             "const ui = SwaggerUIBundle({",
-            _lazy_plugin_js + "const ui = SwaggerUIBundle({plugins:[LazyLoadPlugin],",
+            _lazy_plugin_js
+            + 'const ui = SwaggerUIBundle({plugins:[LazyLoadPlugin],tagsSorter:"alpha",',
             1,
         )
         return HTMLResponse(content=body)