litellm/tests/unit/proxy/test__lazy_features.py

import sys
from collections.abc import AsyncGenerator, Mapping
from contextlib import asynccontextmanager
from dataclasses import replace
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 (
    LAZY_FEATURES,
    LazyFeature,
    LazyFeatureMiddleware,
    attach_lazy_features,
    lazy_tag_to_prefix,
    loaded_lazy_modules,
    reserve_lazy_slot,
)
from litellm.types.passthrough_endpoints.pass_through_endpoints import LITELLM_PASS_THROUGH_ENDPOINT_MARKER

FLAG: Final = "LITELLM_DISABLE_LAZY_ROUTES"
WARMUP_PATH: Final = "/lazy/warm/{name}"


def test_cimd_metadata_is_available_before_any_oauth_request(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv(FLAG, "false")
    monkeypatch.setenv("PROXY_BASE_URL", "https://gateway.example.com")
    app: Final = FastAPI()
    attach_lazy_features(app)

    with TestClient(app) as client:
        response: Final = client.get("/oauth/client-metadata.json")

    assert response.status_code == 200
    assert response.json()["client_id"] == "https://gateway.example.com/oauth/client-metadata.json"
    assert response.json()["redirect_uris"] == ["https://gateway.example.com/callback"]


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)


def test_a2a_feature_matches_v1_message_stream() -> None:
    feature: Final = next(feature for feature in LAZY_FEATURES if feature.name == "a2a")
    assert feature.matches("/v1/a2a/x/message/stream")


@pytest.mark.parametrize("disable_lazy_routes", ("false", "true"))
def test_retired_roi_calculator_cannot_be_loaded_or_called(
    monkeypatch: pytest.MonkeyPatch, disable_lazy_routes: str
) -> None:
    monkeypatch.setenv(FLAG, disable_lazy_routes)
    app: Final = FastAPI()
    attach_lazy_features(app)

    with TestClient(app) as client:
        for method, path in (
            ("POST", "/lazy/warm/roi_calculator"),
            ("GET", "/roi-calculator/settings"),
            ("GET", "/roi-calculator/report?mode=demo"),
            ("POST", "/roi-calculator/sync"),
            ("GET", "/roi-calculator/observed/report"),
            ("GET", "/roi-calculator/observed/apps"),
            ("POST", "/roi-calculator/observed/sync"),
            ("POST", "/roi-calculator/observed/oauth/github/start"),
            ("GET", "/roi-calculator/observed/oauth/github/callback"),
            ("GET", "/roi-calculator/observed/oauth/github/installed"),
            ("POST", "/roi-calculator/observed/oauth/gitlab/start"),
            ("GET", "/roi-calculator/observed/oauth/gitlab/callback"),
        ):
            response: Final = client.request(method, path)
            assert response.status_code == 404, f"{method} {path}: {response.text}"


@pytest.mark.parametrize("disable_lazy_routes", ("false", "true"))
@pytest.mark.parametrize("path", ("/lazy/warm/usage_ai", "/usage/ai/chat"))
def test_retired_usage_chat_cannot_be_loaded_or_called(
    monkeypatch: pytest.MonkeyPatch, disable_lazy_routes: str, path: str
) -> None:
    monkeypatch.setenv(FLAG, disable_lazy_routes)
    app: Final = FastAPI()
    attach_lazy_features(app)

    with TestClient(app) as client:
        response: Final = client.post(path)
        assert response.status_code == 404, f"{path}: {response.text}"


@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_configured_pass_through_added_after_attach_beats_a_reserved_slot_feature(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.delenv(FLAG, raising=False)
    zeta: Final = _feature_module(monkeypatch, "zeta", "/zeta/{endpoint:path}")
    features: Final = (LazyFeature(name="zeta", module_path=zeta.module_path, path_prefixes=("/zeta",)),)

    async def early() -> dict[str, str]:
        return {"feature": "early"}

    async def late() -> dict[str, str]:
        return {"feature": "late"}

    async def configured() -> dict[str, str]:
        return {"feature": "configured"}

    setattr(configured, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, True)

    app: Final = FastAPI()
    app.add_api_route("/early", early, methods=["GET"])
    reserve_lazy_slot(app, "zeta", features=features)
    app.add_api_route("/late", late, methods=["GET"])
    attach_lazy_features(app, features)
    app.add_api_route("/zeta/{subpath:path}", configured, methods=["GET"])

    with TestClient(app) as client:
        assert client.get("/zeta/health").json() == {"feature": "configured"}, (
            "a deployment's own pass-through endpoint must outrank the built-in provider route"
        )


def test_flag_configured_pass_through_added_during_startup_beats_a_reserved_slot_feature(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setenv(FLAG, "true")
    features: Final = (_feature_module(monkeypatch, "zeta", "/zeta/{endpoint:path}"),)

    async def early() -> dict[str, str]:
        return {"feature": "early"}

    async def late() -> dict[str, str]:
        return {"feature": "late"}

    async def configured() -> dict[str, str]:
        return {"feature": "configured"}

    setattr(configured, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, True)

    @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)
    app.add_api_route("/early", early, methods=["GET"])
    reserve_lazy_slot(app, "zeta", features=features)
    app.add_api_route("/late", late, methods=["GET"])
    attach_lazy_features(app, features)

    with TestClient(app) as client:
        assert client.get("/zeta/health").json() == {"feature": "configured"}, (
            "a deployment's own pass-through endpoint must outrank the built-in provider route"
        )


def test_builtin_routes_first_flag_keeps_a_reserved_slot_feature_ahead_of_a_configured_route(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.delenv(FLAG, raising=False)
    monkeypatch.setenv("LITELLM_BUILTIN_PASS_THROUGH_ROUTES_FIRST", "true")
    zeta: Final = _feature_module(monkeypatch, "zeta", "/zeta/{endpoint:path}")
    features: Final = (LazyFeature(name="zeta", module_path=zeta.module_path, path_prefixes=("/zeta",)),)

    async def early() -> dict[str, str]:
        return {"feature": "early"}

    async def configured() -> dict[str, str]:
        return {"feature": "configured"}

    setattr(configured, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, True)

    app: Final = FastAPI()
    app.add_api_route("/early", early, methods=["GET"])
    reserve_lazy_slot(app, "zeta", features=features)
    attach_lazy_features(app, features)
    app.add_api_route("/zeta/{subpath:path}", configured, methods=["GET"])

    with TestClient(app) as client:
        assert client.get("/zeta/health").json() == {"feature": "zeta"}, (
            "LITELLM_BUILTIN_PASS_THROUGH_ROUTES_FIRST restores the old built-in-first order"
        )


def test_late_route_without_the_marker_still_loses_to_a_reserved_slot_feature(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.delenv(FLAG, raising=False)
    zeta: Final = _feature_module(monkeypatch, "zeta", "/zeta/{endpoint:path}")
    features: Final = (LazyFeature(name="zeta", module_path=zeta.module_path, path_prefixes=("/zeta",)),)

    async def early() -> dict[str, str]:
        return {"feature": "early"}

    async def late() -> dict[str, str]:
        return {"feature": "late"}

    app: Final = FastAPI()
    app.add_api_route("/early", early, methods=["GET"])
    reserve_lazy_slot(app, "zeta", features=features)
    app.add_api_route("/late", late, methods=["GET"])
    attach_lazy_features(app, features)
    app.add_api_route("/zeta/{subpath:path}", late, methods=["GET"])

    with TestClient(app) as client:
        assert client.get("/zeta/health").json() == {"feature": "zeta"}, (
            "only configured pass-through routes are promoted; ordinary late routes keep losing"
        )


def test_configured_pass_through_under_a_different_segment_is_not_moved(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.delenv(FLAG, raising=False)
    zeta: Final = _feature_module(monkeypatch, "zeta", "/zeta/{endpoint:path}")
    features: Final = (LazyFeature(name="zeta", module_path=zeta.module_path, path_prefixes=("/zeta",)),)

    async def configured() -> dict[str, str]:
        return {"feature": "configured"}

    setattr(configured, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, True)

    app: Final = FastAPI()
    reserve_lazy_slot(app, "zeta", features=features)
    attach_lazy_features(app, features)
    app.add_api_route("/omega/{subpath:path}", configured, methods=["GET"])

    with TestClient(app) as client:
        assert client.get("/zeta/health").json() == {"feature": "zeta"}

    routes: Final = app.router.routes
    omega_index: Final = next(
        i for i, route in enumerate(routes) if getattr(route, "path", "") == "/omega/{subpath:path}"
    )
    zeta_index: Final = next(
        i for i, route in enumerate(routes) if getattr(route, "path", "") == "/zeta/{endpoint:path}"
    )
    assert omega_index > zeta_index, "an unrelated configured route keeps its position after the feature routes"


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)


def test_root_authorization_server_metadata_is_the_gateways_not_the_byok_connector_document(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Two features claim the root ``/.well-known/oauth-authorization-server``. The one a client
    gets must be the gateway's own authorization server metadata, whose ``/token`` serves the
    gateway's grants, not the BYOK connector form flow whose token endpoint is
    ``/v1/mcp/oauth/token``; the connector keeps its path-appended document."""
    monkeypatch.setenv(FLAG, "false")
    app: Final = FastAPI()
    attach_lazy_features(app)

    with TestClient(app, base_url="https://gateway.example.com") as client:
        root: Final = client.get("/.well-known/oauth-authorization-server")
        connector: Final = client.get("/.well-known/oauth-authorization-server/v1/mcp/oauth")

    assert root.status_code == 200
    assert root.json()["issuer"] == "https://gateway.example.com"
    assert root.json()["token_endpoint"] == "https://gateway.example.com/token"
    assert connector.status_code == 200
    assert connector.json()["token_endpoint"] == "https://gateway.example.com/v1/mcp/oauth/token"


@pytest.mark.parametrize(
    "load_order",
    (("/realtime/unrouted", "/mcp/unrouted"), ("/mcp/unrouted", "/realtime/unrouted")),
    ids=("realtime_first", "realtime_last"),
)
def test_realtime_and_mcp_features_keep_every_route_once_in_either_load_order(
    monkeypatch: pytest.MonkeyPatch, load_order: tuple[str, str]
) -> None:
    monkeypatch.delenv(FLAG, raising=False)

    async def mcp_mount() -> Mapping[str, str]:
        return {"route": "mcp_mount"}

    mounted_app: Final = FastAPI()
    mounted_app.add_api_route("/{path:path}", mcp_mount, methods=["GET"])
    mcp_module: Final = ModuleType("tests.unit.proxy.lazy_fixture_mcp_app")
    mcp_module.app = mounted_app  # pyright: ignore[reportAttributeAccessIssue]  # fixture module built at test time
    monkeypatch.setitem(sys.modules, mcp_module.__name__, mcp_module)
    features: Final = tuple(
        replace(feature, module_path=mcp_module.__name__) if feature.name == "mcp_app" else feature
        for feature in LAZY_FEATURES
        if feature.name in ("mcp_app", "realtime")
    )

    async def mcp_proxy() -> Mapping[str, str]:
        return {"route": "mcp_proxy"}

    app: Final = FastAPI()
    attach_lazy_features(app, features)
    app.add_api_route("/mcp/proxy", mcp_proxy, methods=["GET"])

    with TestClient(app) as client:
        for path in load_order:
            client.get(path)
        response: Final = client.get("/mcp/proxy")

    assert loaded_lazy_modules(app) == {feature.module_path for feature in features}
    assert response.json() == {"route": "mcp_proxy"}
    route_ids: Final = tuple(id(route) for route in app.routes)
    assert len(route_ids) == len(frozenset(route_ids))