litellm/tests/proxy_migration_tests/test_image_admin_mcp.py

import os
import shutil
import subprocess
from typing import Final

import pytest
import test_offline_image_migration

offline_postgres: Final = test_offline_image_migration.offline_postgres

IMAGE: Final = os.getenv("LITELLM_IMAGE")
SCHEMA_IMAGE: Final = os.getenv("LITELLM_ADMIN_MCP_SCHEMA_IMAGE")
COMPONENT: Final = os.getenv("LITELLM_IMAGE_COMPONENT", "unified")
PROBE: Final = """
import asyncio
import base64
import importlib
import json
import sys
from typing import Final

import httpx2
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding, rsa
from fastapi import HTTPException
from litellm_admin_mcp.server import create_http_app
from litellm.proxy import proxy_server

module_name: Final = {
    "unified": "litellm.proxy.proxy_server",
    "backend": "backend.main",
    "gateway": "gateway.main",
}[sys.argv[2]]
app: Final = importlib.import_module(module_name).app

if sys.argv[1] == "base":
    private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    message: Final = json.dumps({"expiration_date": "2999-01-01", "user_id": "image-test"}).encode()
    signature: Final = private_key.sign(
        message,
        padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH),
        hashes.SHA256(),
    )
    proxy_server._license_check.public_key = private_key.public_key()
    proxy_server._license_check.license_str = base64.b64encode(message + b"." + signature).decode()

async def verify_admin_tools(client: httpx2.AsyncClient) -> None:
    user: Final = await client.post(
        "/user/new",
        headers={"Authorization": "Bearer sk-0123456789abcdef0123456789abcdef"},
        json={"user_id": "image-admin", "user_role": "proxy_admin", "auto_create_key": True},
    )
    assert user.status_code == 200, "Admin provisioning failed: " + str(user.status_code)
    headers: Final = {
        "Authorization": "Bearer " + user.json()["key"],
        "Accept": "application/json, text/event-stream",
    }
    discovery: Final = await client.post(
        "/admin/mcp", headers=headers, json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"}
    )
    assert discovery.status_code == 200, discovery.text
    assert {"create_team", "get_team", "delete_teams"} <= {
        tool["name"] for tool in discovery.json()["result"]["tools"]
    }

    async def call_tool(name: str, arguments: dict[str, object]) -> str:
        response: Final = await client.post(
            "/admin/mcp", headers=headers,
            json={"jsonrpc": "2.0", "id": 2, "method": "tools/call",
                  "params": {"name": name, "arguments": arguments}},
        )
        assert response.status_code == 200, response.text
        result: Final = response.json()["result"]
        assert not result.get("isError"), response.text
        return result["content"][0]["text"]

    team_id: Final = "image-admin-mcp-team"
    created: Final = json.loads(await call_tool(
        "create_team", {"body": {"team_id": team_id, "team_alias": team_id, "max_budget": 25}}
    ))
    assert created["team_id"] == team_id and created["max_budget"] == 25, created
    read: Final = json.loads(await call_tool("get_team", {"query": {"team_id": team_id}}))
    assert read["team_info"]["team_id"] == team_id and read["team_info"]["max_budget"] == 25, read
    await call_tool("delete_teams", {"body": {"team_ids": [team_id]}})
    deleted: Final = await client.get("/team/info", headers=headers, params={"team_id": team_id})
    assert deleted.status_code == 404, deleted.text
    print("admin-mcp-tools-ok")

async def probe() -> None:
    try:
        async with app.router.lifespan_context(app):
            async with httpx2.AsyncClient(
                transport=httpx2.ASGITransport(app=app), base_url="http://localhost:4000"
            ) as client:
                response: Final = await client.post(
                    "/admin/mcp",
                    json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"},
                    headers={"Accept": "application/json, text/event-stream"},
                )
                print("admin-mcp-status=" + str(response.status_code))
                if sys.argv[3] == "tools":
                    assert response.status_code == 401, response.text
                    await verify_admin_tools(client)
    except HTTPException as exc:
        assert exc.status_code == 403 and "LITELLM_LICENSE" in str(exc.detail)
        assert sys.argv[1] == "none"
        print("admin-mcp-license-required")

asyncio.run(probe())
"""

pytestmark = [
    pytest.mark.skipif(IMAGE is None, reason="requires a built image (set LITELLM_IMAGE)"),
    pytest.mark.skipif(shutil.which("docker") is None, reason="requires the docker CLI"),
]


def _run_probe(
    enabled: str | None,
    license_mode: str,
    network: str = "none",
    database_url: str | None = None,
) -> subprocess.CompletedProcess[str]:
    assert IMAGE is not None
    enabled_args: Final = () if enabled is None else ("--env", "LITELLM_ENABLE_ADMIN_MCP=" + enabled)
    database_args: Final = () if database_url is None else ("--env", "DATABASE_URL=" + database_url)
    return subprocess.run(
        [
            "docker",
            "run",
            "--rm",
            "--network",
            network,
            "--user",
            "12345:0",
            *enabled_args,
            *database_args,
            "--env",
            "LITELLM_LOCAL_MODEL_COST_MAP=true",
            "--env",
            "LITELLM_MASTER_KEY=sk-0123456789abcdef0123456789abcdef",
            "--entrypoint",
            "python",
            IMAGE,
            "-c",
            PROBE,
            license_mode,
            COMPONENT,
            "tools" if database_url is not None else "visibility",
        ],
        capture_output=True,
        text=True,
        timeout=120,
        check=False,
    )


@pytest.mark.parametrize(
    "enabled,license_mode,expected",
    [
        (None, "none", "status=404"),
        ("false", "base", "status=404"),
        ("true", "none", "license-required"),
        ("true", "base", "status=404" if COMPONENT == "gateway" else "status=401"),
    ],
)
def test_image_admin_mcp_requires_opt_in_license_and_management_component(
    enabled: str | None, license_mode: str, expected: str
) -> None:
    result: Final = _run_probe(enabled, license_mode)
    assert result.returncode == 0 and f"admin-mcp-{expected}" in result.stdout, (
        f"Admin MCP image probe failed with component={COMPONENT}, enabled={enabled}, license={license_mode}\n"
        f"{result.stdout}\n{result.stderr}"
    )


@pytest.mark.skipif(COMPONENT == "gateway", reason="the gateway excludes management endpoints")
def test_image_admin_mcp_personal_admin_manages_team(offline_postgres: tuple[str, str]) -> None:
    assert SCHEMA_IMAGE is not None, "set LITELLM_ADMIN_MCP_SCHEMA_IMAGE to the matching builder image"
    network, postgres = offline_postgres
    database_url: Final = f"postgresql://postgres:pw@{postgres}:5432/litellm"
    schema: Final = subprocess.run(
        [
            "docker", "run", "--rm", "--network", network,
            "--env", "DATABASE_URL=" + database_url,
            "--env", "HOME=/opt/prisma", "--env", "XDG_CACHE_HOME=/opt/prisma/.cache",
            "--env", "PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries",
            "--entrypoint", "prisma", SCHEMA_IMAGE,
            "db", "push", "--schema", "/app/schema.prisma", "--skip-generate", "--accept-data-loss",
        ],
        capture_output=True, text=True, timeout=180, check=False,
    )
    assert schema.returncode == 0, f"Schema provisioning failed\n{schema.stdout}\n{schema.stderr}"
    result: Final = _run_probe("true", "base", network, database_url)
    assert result.returncode == 0 and "admin-mcp-tools-ok" in result.stdout, (
        f"Admin MCP management failed in {COMPONENT}\n{result.stdout}\n{result.stderr}"
    )