litellm/tests/unit/proxy/test_plugin_routes.py

"""Regression tests for UI-registered embed plugins.

Covers three bugs:
1. `general_settings.plugins` was not a field on ConfigGeneralSettings, so the
   admin UI's POST /config/field/update with field_name="plugins" was rejected
   with "Invalid field=plugins passed in."
2. The in-memory plugin registry only refreshed at startup, so a plugin added
   via the UI did not appear in /api/plugins until a restart.
3. Plugins persisted to DB general_settings were not loaded on startup (the
   registry only initialised from the YAML config), so UI-added plugins vanished
   after a restart.
"""

import asyncio
from unittest.mock import MagicMock

import pytest
import respx
from starlette.requests import Request

import litellm
from litellm.proxy._types import (
    ConfigGeneralSettings,
    LitellmUserRoles,
    PluginConfig,
    UserAPIKeyAuth,
)
from litellm.proxy.plugin_routes import list_plugins, register_plugins_from_config


def _admin() -> UserAPIKeyAuth:
    return UserAPIKeyAuth(api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN)


def _non_admin() -> UserAPIKeyAuth:
    return UserAPIKeyAuth(api_key="sk-user", user_role=LitellmUserRoles.INTERNAL_USER)


def test_plugins_is_a_valid_general_setting() -> None:
    """The config-update endpoint gates on this exact membership check."""
    assert "plugins" in ConfigGeneralSettings.model_fields


def test_config_general_settings_parses_plugin_list() -> None:
    """A list of plugin dicts (what the UI sends) coerces into PluginConfig."""
    settings = ConfigGeneralSettings.model_validate(
        {
            "plugins": [
                {
                    "name": "chat-ui",
                    "display_name": "Chat UI",
                    "url": "http://localhost:3300",
                },
                {
                    "name": "agent-builder",
                    "url": "http://127.0.0.1:4010",
                    "plugin_key": "sk-secret",
                },
            ]
        }
    )
    plugins = settings.plugins
    assert plugins is not None
    assert [p.name for p in plugins] == ["chat-ui", "agent-builder"]
    assert isinstance(plugins[0], PluginConfig)
    assert plugins[1].display_name is None
    assert plugins[1].plugin_key == "sk-secret"


def test_registered_plugins_appear_in_list_without_restart() -> None:
    """register_plugins_from_config makes UI-added plugins visible immediately,
    and replaces (not merges) so removed plugins disappear."""
    register_plugins_from_config(
        {
            "plugins": [
                {
                    "name": "chat-ui",
                    "display_name": "Chat UI",
                    "url": "http://localhost:3300",
                }
            ]
        }
    )
    names = [p["name"] for p in asyncio.run(list_plugins(user_api_key_dict=_admin()))]
    assert names == ["chat-ui"]

    register_plugins_from_config(
        {
            "plugins": [
                {
                    "name": "chat-ui",
                    "display_name": "Chat UI",
                    "url": "http://localhost:3300",
                },
                {
                    "name": "agent-builder",
                    "display_name": "Agent Builder",
                    "url": "http://127.0.0.1:4010",
                },
            ]
        }
    )
    names = sorted(
        p["name"] for p in asyncio.run(list_plugins(user_api_key_dict=_admin()))
    )
    assert names == ["agent-builder", "chat-ui"]

    # Removing a plugin from config drops it from the live list.
    register_plugins_from_config({})
    assert asyncio.run(list_plugins(user_api_key_dict=_admin())) == []


def test_plugin_key_is_never_returned_to_the_browser() -> None:
    """plugin_key is a credential the UI never needs; /api/plugins must omit it
    for every caller, admin included, so it never lands in browser state."""
    register_plugins_from_config(
        {
            "plugins": [
                {
                    "name": "p",
                    "display_name": "P",
                    "url": "http://localhost:9",
                    "plugin_key": "sk-secret",
                }
            ]
        }
    )

    admin_entry = asyncio.run(list_plugins(user_api_key_dict=_admin()))[0]
    user_entry = asyncio.run(list_plugins(user_api_key_dict=_non_admin()))[0]

    assert "plugin_key" not in admin_entry
    assert "plugin_key" not in user_entry
    assert admin_entry["url"] == "http://localhost:9"

    register_plugins_from_config({})


def test_db_persisted_plugins_load_on_startup(monkeypatch: pytest.MonkeyPatch) -> None:
    from litellm.proxy import proxy_server
    from litellm.proxy.proxy_server import ProxyConfig

    register_plugins_from_config({})
    monkeypatch.setattr(proxy_server, "general_settings", {})

    asyncio.run(
        ProxyConfig()._update_general_settings(
            {
                "plugins": [
                    {
                        "name": "db-plugin",
                        "display_name": "DB Plugin",
                        "url": "http://localhost:5000",
                    }
                ]
            }
        )
    )

    names = [p["name"] for p in asyncio.run(list_plugins(user_api_key_dict=_admin()))]
    assert names == ["db-plugin"]

    register_plugins_from_config({})


def test_safe_response_headers_sandbox_and_strips_wire_headers() -> None:
    """Proxied plugin responses must be inert and shed wire/cookie headers."""
    from litellm.proxy.plugin_routes import _safe_response_headers

    out = _safe_response_headers(
        {
            "content-type": "text/html",
            "content-encoding": "gzip",
            "content-length": "123",
            "set-cookie": "session=abc",
            "content-security-policy": "default-src *",
        }
    )

    assert out["content-security-policy"] == "sandbox"
    assert out["x-content-type-options"] == "nosniff"
    assert out["content-type"] == "text/html"
    for stripped in ("content-encoding", "content-length", "set-cookie"):
        assert stripped not in out


def test_litellm_credential_header_names_covers_every_auth_header() -> None:
    """The canonical strip set must list every header user_api_key_auth accepts
    as a litellm key, so a new auth header can't silently start leaking."""
    from litellm.proxy._types import SpecialHeaders

    assert SpecialHeaders.litellm_credential_header_names() == {
        "authorization",
        "api-key",
        "x-api-key",
        "x-goog-api-key",
        "ocp-apim-subscription-key",
        "x-litellm-api-key",
    }


def test_every_litellm_auth_header_is_stripped_before_forwarding() -> None:
    """A plugin must never receive any header that authenticates against litellm,
    only the hop-by-hop set and benign headers are forwarded."""
    from litellm.proxy.plugin_routes import _request_strip_headers

    strip = _request_strip_headers()
    incoming = {
        "Authorization": "Bearer sk-litellm",
        "API-Key": "sk-litellm",
        "X-Api-Key": "sk-litellm",
        "X-Goog-Api-Key": "sk-litellm",
        "Ocp-Apim-Subscription-Key": "sk-litellm",
        "X-Litellm-Api-Key": "sk-litellm",
        "Cookie": "litellm_session=abc",
        "Accept": "application/json",
        "X-Trace-Id": "t-1",
    }
    forwarded = {k: v for k, v in incoming.items() if k.lower() not in strip}

    assert forwarded == {"Accept": "application/json", "X-Trace-Id": "t-1"}


def test_configured_custom_key_header_is_stripped() -> None:
    """A custom general_settings.litellm_key_header_name must also be stripped,
    read live so config changes are honoured without a restart."""
    from litellm.proxy import proxy_server
    from litellm.proxy.plugin_routes import _request_strip_headers

    original = getattr(proxy_server, "general_settings", None)
    proxy_server.general_settings = {"litellm_key_header_name": "X-My-Tenant-Key"}
    try:
        assert "x-my-tenant-key" in _request_strip_headers()
    finally:
        proxy_server.general_settings = original


def _plugin_proxy_request(headers: dict[str, str]) -> Request:
    async def receive() -> dict[str, object]:
        return {"type": "http.request", "body": b'{"q": 1}', "more_body": False}

    return Request(
        {
            "type": "http",
            "method": "POST",
            "path": "/plugin-proxy/p/items",
            "query_string": b"page=2",
            "headers": [(name.lower().encode(), value.encode()) for name, value in headers.items()],
        },
        receive,
    )


@pytest.mark.parametrize(
    ("user_id", "expected_identity_headers"),
    [
        ("admin-1", {"x-litellm-user-id": "admin-1", "x-litellm-user-role": "proxy_admin"}),
        (None, {"x-litellm-user-role": "proxy_admin"}),
        ("", {"x-litellm-user-role": "proxy_admin"}),
    ],
)
def test_plugin_proxy_forwards_caller_identity_and_replaces_the_caller_credential(
    monkeypatch: pytest.MonkeyPatch,
    respx_mock: respx.MockRouter,
    user_id: str | None,
    expected_identity_headers: dict[str, str],
) -> None:
    from litellm.proxy.plugin_routes import plugin_proxy

    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
    litellm.in_memory_llm_clients_cache.flush_cache()
    register_plugins_from_config(
        {"plugins": [{"name": "p", "url": "http://plugin.example", "plugin_key": "sk-plugin"}]}
    )
    upstream = respx_mock.post("http://plugin.example/items?page=2").respond(text="plugin says hi")

    response = asyncio.run(
        plugin_proxy(
            "p",
            "items",
            _plugin_proxy_request({"Authorization": "Bearer sk-caller", "X-Trace-Id": "t-1"}),
            UserAPIKeyAuth(api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN, user_id=user_id),
        )
    )
    register_plugins_from_config({})
    litellm.in_memory_llm_clients_cache.flush_cache()

    forwarded = upstream.calls[0].request
    assert response.status_code == 200
    assert response.body == b"plugin says hi"
    assert forwarded.content == b'{"q": 1}'
    assert forwarded.headers["authorization"] == "Bearer sk-plugin"
    assert forwarded.headers["x-trace-id"] == "t-1"
    assert {
        name: value for name, value in forwarded.headers.items() if name.startswith("x-litellm-user-")
    } == expected_identity_headers


def test_plugin_proxy_refuses_non_admin_callers_without_contacting_the_plugin(
    monkeypatch: pytest.MonkeyPatch,
    respx_mock: respx.MockRouter,
) -> None:
    from litellm.proxy.plugin_routes import plugin_proxy

    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
    litellm.in_memory_llm_clients_cache.flush_cache()
    register_plugins_from_config({"plugins": [{"name": "p", "url": "http://plugin.example"}]})
    upstream = respx_mock.post("http://plugin.example/items?page=2").respond(text="plugin says hi")

    response = asyncio.run(
        plugin_proxy(
            "p",
            "items",
            _plugin_proxy_request({"Authorization": "Bearer sk-caller"}),
            UserAPIKeyAuth(api_key="sk-user", user_role=LitellmUserRoles.INTERNAL_USER, user_id="user-1"),
        )
    )
    register_plugins_from_config({})

    assert response.status_code == 403
    assert response.body == b"Plugin proxy access requires proxy_admin role."
    assert not upstream.called