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

import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient

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,
    _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.asyncio
async def test_otel_span_debug_privacy_preserves_api_result(
    monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
    from typing import Final

    import litellm
    from opentelemetry.sdk.trace import TracerProvider
    from opentelemetry.sdk.trace.export import SimpleSpanProcessor
    from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter

    from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
    from litellm.proxy.common_utils.debug_utils import get_otel_spans

    monkeypatch.setattr(litellm, "redact_messages_in_debug_logs", True)
    exporter: Final = InMemorySpanExporter()
    provider: Final = TracerProvider()
    provider.add_span_processor(SimpleSpanProcessor(exporter))
    logger: Final = OpenTelemetry(config=OpenTelemetryConfig(exporter=exporter), tracer_provider=provider)
    monkeypatch.setattr(proxy_server, "open_telemetry_logger", logger)
    try:
        with provider.get_tracer("privacy-test").start_as_current_span("PRIVATE_SPAN_42896") as span:
            span.set_attribute("gen_ai.input.messages", "PRIVATE_MESSAGE_42896")
        result: Final = await get_otel_spans()
        output: Final = capsys.readouterr().out
        assert result["otel_spans"] == ["PRIVATE_SPAN_42896"]
        assert "redacted-by-litellm" in output
        assert "PRIVATE" not in output
    finally:
        provider.shutdown()