litellm/tests/unit/proxy/common_utils/test_debug_utils.py

import json
import os
import socket
from collections.abc import Callable, Iterator, Mapping
from dataclasses import asdict
from pathlib import Path
from types import SimpleNamespace
from typing import Final
from unittest.mock import MagicMock, patch

import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from redis import Redis
from redis.connection import Connection

from litellm.caching.redis_cache import RedisCache
from litellm.proxy import proxy_server
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.bug_report_config import build_proxy_bug_report
from litellm.proxy.common_utils.debug_utils import (
    PSUTIL_MISSING_ERROR,
    _ProcFilesystemProcess,
    _redis_memory_stats,
    _summary_process_memory,
    get_memory_summary,
)
from litellm.proxy.common_utils.debug_utils import router as debug_router
from litellm.router import Router
from litellm.router_utils.prompt_caching_cache import PromptCachingCache

PAGE_SIZE = 4096
STATM_SIZE_PAGES = 100_000
STATM_RESIDENT_PAGES = 30_000
MEMINFO_TOTAL_KB = 1_000_000


@pytest.fixture
def proc_process(tmp_path: Path) -> _ProcFilesystemProcess:
    statm = tmp_path / "statm"
    statm.write_text(f"{STATM_SIZE_PAGES} {STATM_RESIDENT_PAGES} 5000 1 0 20000 0\n")
    meminfo = tmp_path / "meminfo"
    meminfo.write_text(
        f"MemTotal:       {MEMINFO_TOTAL_KB} kB\nMemFree:         400000 kB\nMemAvailable:    600000 kB\n"
    )
    return _ProcFilesystemProcess(statm_path=str(statm), meminfo_path=str(meminfo), page_size=PAGE_SIZE)


def test_proc_filesystem_process_reads_resident_and_virtual_bytes_from_statm(
    proc_process: _ProcFilesystemProcess,
) -> None:
    memory_info = proc_process.memory_info()

    assert memory_info.rss == STATM_RESIDENT_PAGES * PAGE_SIZE
    assert memory_info.vms == STATM_SIZE_PAGES * PAGE_SIZE


def test_proc_filesystem_process_reports_share_of_meminfo_total(proc_process: _ProcFilesystemProcess) -> None:
    expected_percent = STATM_RESIDENT_PAGES * PAGE_SIZE / (MEMINFO_TOTAL_KB * 1024) * 100

    assert proc_process.memory_percent() == pytest.approx(expected_percent)


def test_summary_reports_rss_from_the_proc_filesystem(proc_process: _ProcFilesystemProcess) -> None:
    memory, health_status = _summary_process_memory(proc_process)

    assert memory["ram_usage_mb"] == round(STATM_RESIDENT_PAGES * PAGE_SIZE / (1024 * 1024), 2)
    assert memory["system_memory_percent"] == pytest.approx(12.0)
    assert health_status == "healthy"
    assert "error" not in memory


def test_summary_without_any_memory_source_names_psutil_and_reports_no_rss() -> None:
    memory, health_status = _summary_process_memory(None)

    assert memory == {"error": PSUTIL_MISSING_ERROR}
    assert health_status == "healthy"


@pytest.mark.asyncio
async def test_memory_summary_names_the_host_and_worker_that_answered() -> None:
    summary = await get_memory_summary(UserAPIKeyAuth())

    assert summary["hostname"] == socket.gethostname()
    assert summary["worker_pid"] == os.getpid()
    assert summary["memory"]["ram_usage_mb"] > 0


HOSTILE_CONFIG: Mapping[str, object] = {
    "model_list": [
        {
            "model_name": "acme-prod-gpt4",
            "litellm_params": {
                "model": "azure/acme-gpt4o-deployment",
                "api_base": "https://acme-eastus.openai.azure.com",
                "api_key": "sk-live-secret-1",
            },
        }
    ],
    "litellm_settings": {"drop_params": True, "callbacks": ["langfuse", "acme_hooks.audit_logger"]},
}

HOSTILE_GENERAL_SETTINGS: Mapping[str, object] = {
    "master_key": "sk-live-secret-master",
    "database_url": "postgres://user:hunter2@10.0.0.7/litellm",
    "store_model_in_db": True,
}

HOSTILE_STRINGS = ("acme", "sk-live-secret", "hunter2", "10.0.0.7", "azure.com")


@pytest.fixture
def hostile_proxy_config(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
    previous_config = proxy_server.proxy_config.get_config_state()
    proxy_server.proxy_config.update_config_state(config=HOSTILE_CONFIG)
    monkeypatch.setattr(proxy_server, "general_settings", dict(HOSTILE_GENERAL_SETTINGS))
    yield
    proxy_server.proxy_config.update_config_state(config=previous_config)


def _debug_client(caller: UserAPIKeyAuth) -> TestClient:
    app = FastAPI()
    app.include_router(debug_router)
    app.dependency_overrides[user_api_key_auth] = lambda: caller
    return TestClient(app)


@pytest.mark.parametrize(
    "caller",
    [
        UserAPIKeyAuth(),
        UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER),
        UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY),
    ],
)
@pytest.mark.usefixtures("hostile_proxy_config")
def test_debug_report_refuses_everyone_but_proxy_admins(caller: UserAPIKeyAuth) -> None:
    response = _debug_client(caller).get("/debug/report")

    assert response.status_code == 403, response.text
    assert "litellm_version" not in response.text


@pytest.mark.usefixtures("hostile_proxy_config")
def test_debug_report_returns_what_the_bug_report_link_carries_and_nothing_from_the_operator() -> None:
    response = _debug_client(UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)).get("/debug/report")

    assert response.status_code == 200, response.text
    assert response.json() == json.loads(json.dumps(asdict(build_proxy_bug_report(RuntimeError("boom")).environment)))
    assert response.json()["config_lines"] == [
        "general_settings.store_model_in_db = true",
        "litellm_settings.drop_params = true",
        "litellm_settings.callbacks = [langfuse]",
        "model_list[*].provider = [azure]",
    ]
    assert not any(hostile in response.text for hostile in HOSTILE_STRINGS), response.text


PINNED_PREFIXES = 3


@pytest.fixture
def router_with_pins(monkeypatch: pytest.MonkeyPatch) -> None:
    router = Router(model_list=[{"model_name": "claude", "litellm_params": {"model": "anthropic/claude-sonnet-4-5"}}])
    pins = PromptCachingCache(cache=router.prompt_caching_pin_cache)
    for prefix in range(PINNED_PREFIXES):
        pins.add_model_id(
            model_id=f"deployment-{prefix}",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": f"prefix {prefix} " * 300, "cache_control": {"type": "ephemeral"}}
                    ],
                }
            ],
            tools=None,
        )
    monkeypatch.setattr(proxy_server, "llm_router", router)


@pytest.mark.parametrize(
    ("path", "pin_count", "expected"),
    [
        (
            "/memory-usage-in-mem-cache",
            lambda body: body["num_items_in_prompt_caching_pin_cache"],
            2 * PINNED_PREFIXES,
        ),
        ("/memory-usage-in-mem-cache-items", lambda body: len(body["prompt_caching_pin_cache"]), PINNED_PREFIXES),
        ("/debug/memory/summary", lambda body: body["caches"]["breakdown"]["prompt_caching_pins"]["count"], PINNED_PREFIXES),
        (
            "/debug/memory/details?include_process_info=false",
            lambda body: body["cache_memory"]["prompt_caching_pin_cache"]["num_items"],
            PINNED_PREFIXES,
        ),
    ],
)
@pytest.mark.usefixtures("router_with_pins")
def test_memory_routes_count_prompt_caching_pins(
    path: str, pin_count: Callable[[dict[str, object]], int], expected: int
) -> None:
    response = _debug_client(UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)).get(path)

    assert response.status_code == 200
    assert pin_count(response.json()) == expected


@pytest.mark.parametrize(
    ("pool", "expected"),
    [
        pytest.param(
            SimpleNamespace(
                max_connections=23,
                connection_class=Connection,
                _available_connections=(object(),),
                _in_use_connections=(object(), object()),
            ),
            {
                "enabled": True,
                "cache_type": "RedisCache",
                "connection_pool": {"max_connections": 23, "connection_class": Connection.__name__},
            },
            id="pool-metadata-preserves-diagnostic-shape",
        ),
        pytest.param(
            SimpleNamespace(),
            {
                "enabled": True,
                "cache_type": "RedisCache",
                "connection_pool": {"max_connections": None, "connection_class": None},
            },
            id="unavailable-pool-fields",
        ),
        pytest.param(None, {"enabled": True, "cache_type": "RedisCache"}, id="no-pool"),
        pytest.param(
            SimpleNamespace(max_connections=object()),
            {"enabled": True, "cache_type": "RedisCache"},
            id="invalid-metadata-remains-best-effort",
        ),
    ],
)
def test_redis_memory_diagnostics_preserve_pool_shape(pool: object, expected: Mapping[str, object]) -> None:
    client: Final = MagicMock(spec=Redis)
    client.connection_pool = pool
    with patch("redis.Redis", return_value=client):
        cache: Final = RedisCache(host="diagnostic-pool", port=6379)

    with pytest.warns(FutureWarning, match="get_connection_pool_info is experimental and may change"):
        result: Final = _redis_memory_stats(cache)

    assert result == expected


def test_redis_memory_diagnostics_without_cache_remain_disabled() -> None:
    assert _redis_memory_stats(None) == {"enabled": False}