litellm/tests/unit/proxy/test_tracing_endpoints.py

"""
Tests for the agent tracing endpoints (litellm/proxy/tracing_endpoints.py).
"""

from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from typing import Final
from unittest.mock import AsyncMock, MagicMock

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

from litellm.proxy import tracing_endpoints
from litellm.proxy._types import LitellmUserRoles, ProxyLifespanState, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.tracing_runtime import manage_tracing, provide_storage
from litellm.rust_bridge.traces import ClickHouseStorage
from litellm.tracing import TraceReceiver, TracingPayloadTooLargeError
from litellm.tracing.store import TraceStore
from litellm.tracing.types import TraceScope

TEAM_KEY = UserAPIKeyAuth(
    token="hashed-key", team_id="team-research", org_id="org-1", user_role=LitellmUserRoles.INTERNAL_USER
)


@pytest.mark.parametrize(
    ("auth", "scope", "can_write"),
    (
        pytest.param(
            UserAPIKeyAuth(token="admin-key", team_id="team-a", user_role=LitellmUserRoles.PROXY_ADMIN),
            TraceScope(team_ids=(), api_key_hash=""),
            True,
            id="admin",
        ),
        pytest.param(
            UserAPIKeyAuth(token="view-key", team_id="team-a", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY),
            TraceScope(team_ids=(), api_key_hash=""),
            False,
            id="view-only-admin",
        ),
        pytest.param(
            TEAM_KEY,
            TraceScope(team_ids=("team-research",), api_key_hash=""),
            True,
            id="team-key",
        ),
        pytest.param(
            UserAPIKeyAuth(token="hashed-key", user_role=LitellmUserRoles.INTERNAL_USER),
            TraceScope(team_ids=("",), api_key_hash="hashed-key"),
            True,
            id="teamless-key",
        ),
    ),
)
def test_trace_read_and_write_permissions(
    client: TestClient, receiver: MagicMock, auth: UserAPIKeyAuth, scope: TraceScope, can_write: bool
) -> None:
    client.app.dependency_overrides[user_api_key_auth] = lambda: auth

    read: Final = client.get("/v1/traces?start_ms=1&end_ms=2")
    assert read.status_code == 200, read.text
    receiver.list_traces.assert_awaited_once_with(scope=scope, start_ms=1, end_ms=2, cursor=None)

    write: Final = client.post("/v1/traces", json={})
    assert write.status_code == (200 if can_write else 403), write.text
    if not can_write:
        receiver.ingest.assert_not_awaited()
        return
    receiver.ingest.assert_awaited_once()
    tenant: Final = receiver.ingest.await_args.kwargs["tenant"]
    assert (tenant.team_id, tenant.api_key_hash, tenant.org_id) == (
        auth.team_id or "",
        auth.token or "",
        auth.org_id or "",
    )


@pytest.fixture
def receiver(client) -> MagicMock:
    fake = MagicMock()
    fake.ingest = AsyncMock(return_value=1)
    fake.list_traces = AsyncMock(return_value={"data": [], "next_cursor": None})
    fake.get_trace = AsyncMock(return_value=None)
    fake.get_span = AsyncMock(return_value=None)
    client.app.dependency_overrides[tracing_endpoints.provide_receiver] = lambda: fake
    return fake


@pytest.fixture
def client() -> TestClient:
    app = FastAPI()
    app.include_router(tracing_endpoints.router)
    app.dependency_overrides[user_api_key_auth] = lambda: TEAM_KEY
    return TestClient(app)


@pytest.mark.parametrize("native_available", [True, False])
def test_501_when_tracing_not_enabled(
    client: TestClient, native_available: bool, monkeypatch: pytest.MonkeyPatch
) -> None:
    from google.rpc.status_pb2 import Status

    from litellm.rust_bridge import loader

    if not native_available:
        monkeypatch.setattr(loader, "_cached_bridge", None)
    response: Final = client.post("/v1/traces", content=b"")
    assert response.status_code == 501
    assert response.headers["content-type"] == "application/x-protobuf"
    assert Status.FromString(response.content).message == (
        "Agent tracing is not enabled. Set `tracing:` in general_settings and CLICKHOUSE_URL." if native_available else ""
    )
    assert client.get("/v1/traces").status_code == 501


def test_post_protobuf_returns_empty_protobuf(client, receiver):
    response = client.post(
        "/v1/traces",
        content=b"\x0a\x00",
        headers={"content-type": "application/x-protobuf", "content-encoding": "gzip"},
    )
    assert response.status_code == 200
    assert response.content == b""
    assert response.headers["content-type"] == "application/x-protobuf"
    kwargs = receiver.ingest.call_args.kwargs
    assert kwargs["body"] is not None
    assert kwargs["content_type"] == "application/x-protobuf"
    assert kwargs["content_encoding"] == "gzip"
    assert kwargs["tenant"].team_id == "team-research"


def test_post_json_returns_empty_json(client, receiver):
    response = client.post("/v1/traces", content=b"{}", headers={"content-type": "application/json"})
    assert response.status_code == 200
    assert response.json() == {}


def test_post_clickhouse_failure_is_503_with_retry_after(client, receiver):
    receiver.ingest.side_effect = RuntimeError("ClickHouse unavailable")
    response = client.post("/v1/traces", content=b"", headers={"content-type": "application/x-protobuf"})
    assert response.status_code == 503
    assert response.headers["retry-after"] == str(tracing_endpoints.OTLP_RETRY_AFTER_SECONDS)


def test_post_too_large_is_413(client, receiver):
    receiver.ingest.side_effect = TracingPayloadTooLargeError("OTLP body exceeds 10 bytes")
    response = client.post("/v1/traces", content=b"x" * 20)
    assert response.status_code == 413
    from google.rpc.status_pb2 import Status

    assert "exceeds" in Status.FromString(response.content).message


def test_list_traces_passes_scope_window_and_cursor(client, receiver):
    response = client.get("/v1/traces", params={"start_ms": 1, "end_ms": 2, "cursor": "abc"})
    assert response.status_code == 200
    assert response.json() == {"data": [], "next_cursor": None}
    receiver.list_traces.assert_awaited_once_with(
        scope={"team_ids": ("team-research",), "api_key_hash": ""}, start_ms=1, end_ms=2, cursor="abc"
    )


def test_list_traces_defaults_to_last_24h(client, receiver):
    client.get("/v1/traces")
    kwargs = receiver.list_traces.call_args.kwargs
    assert kwargs["end_ms"] - kwargs["start_ms"] == tracing_endpoints.MS_PER_DAY
    assert kwargs["cursor"] is None


def test_get_trace_404_and_200(client, receiver):
    assert client.get("/v1/traces/missing").status_code == 404
    trace = {"summary": {"trace_id": "t1"}, "agents": [], "spans": []}
    receiver.get_trace.return_value = trace
    response = client.get("/v1/traces/t1")
    assert response.status_code == 200
    assert response.json() == trace
    receiver.get_trace.assert_awaited_with("t1", {"team_ids": ("team-research",), "api_key_hash": ""}, "")


def test_get_span_404_and_200(client, receiver):
    assert client.get("/v1/traces/t1/spans/s1").status_code == 404
    receiver.get_span.return_value = {"span_id": "s1", "input": "", "output": "", "attributes": {}}
    response = client.get("/v1/traces/t1/spans/s1")
    assert response.status_code == 200
    assert response.json()["span_id"] == "s1"
    receiver.get_span.assert_awaited_with("t1", "s1", {"team_ids": ("team-research",), "api_key_hash": ""}, "")


def test_get_span_serves_ui_content_from_stored_payloads(client):
    storage = MagicMock()
    stored_output = '{"role": "ai", "content": "", "tool_calls": [{"name": "lookup", "args": {"id": 7}}]}'
    storage.query = AsyncMock(
        return_value=[{"span_id": "s1", "input": '{"city": "Paris"}', "output": stored_output, "attributes": {}}]
    )
    client.app.dependency_overrides[tracing_endpoints.provide_receiver] = lambda: TraceReceiver(TraceStore(storage))
    body = client.get("/v1/traces/t1/spans/s1").json()
    assert body["output"] == stored_output
    assert body["input_ui"] == {"kind": "fields", "fields": [{"key": "city", "value": "Paris"}]}
    assert body["output_ui"] == {
        "kind": "messages",
        "messages": [
            {"role": "assistant", "content": "", "tool_calls": [{"name": "lookup", "arguments": '{"id": 7}'}]}
        ],
    }


def test_trace_detail_passes_scoped_reference(client, receiver):
    receiver.get_trace.return_value = {"summary": {"trace_id": "t1"}, "agents": [], "spans": []}
    assert client.get("/v1/traces/t1?trace_ref=run-one").status_code == 200
    receiver.get_trace.assert_awaited_with("t1", {"team_ids": ("team-research",), "api_key_hash": ""}, "run-one")


def test_invalid_export_and_cursor_are_client_errors(client, receiver):
    from litellm.tracing.decode import InvalidOTLPPayloadError

    receiver.ingest.side_effect = InvalidOTLPPayloadError("invalid OTLP trace payload")
    assert client.post("/v1/traces", content=b"broken").status_code == 400
    receiver.list_traces.side_effect = ValueError("Invalid trace cursor")
    assert client.get("/v1/traces?cursor=broken").status_code == 400


def test_teamless_key_without_token_gets_403_on_reads(client, receiver):
    client.app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER
    )
    assert client.get("/v1/traces").status_code == 403
    receiver.list_traces.assert_not_called()


def test_view_only_admin_cannot_ingest_traces(client, receiver):
    client.app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
        token="admin-key", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
    )
    response = client.post("/v1/traces", content=b"{}")
    assert response.status_code == 403
    receiver.ingest.assert_not_called()


@pytest.mark.parametrize(
    "status_code, field, message",
    [(401, "detail", "Invalid API key"), (403, "message", "Not allowed to ingest agent traces")],
)
def test_auth_failure_precedes_disabled_receiver(
    client: TestClient, status_code: int, field: str, message: str
) -> None:
    def unavailable() -> None:
        return None

    def authenticate() -> UserAPIKeyAuth:
        if status_code == 401:
            raise HTTPException(status_code=401, detail="Invalid API key")
        return UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY)

    client.app.dependency_overrides[user_api_key_auth] = authenticate
    client.app.dependency_overrides[tracing_endpoints.provide_receiver] = unavailable
    response: Final = client.post("/v1/traces", content=b"{}", headers={"content-type": "application/json"})
    assert response.status_code == status_code
    assert response.json() == {field: message}


def test_disabled_receiver_precedes_read_scope_rejection(client: TestClient) -> None:
    client.app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER
    )
    response: Final = client.get("/v1/traces")
    assert response.status_code == 501
    assert response.json() == {
        "detail": "Agent tracing is not enabled. Set `tracing:` in general_settings and CLICKHOUSE_URL."
    }


@pytest.mark.requires_rust_extension
def test_injected_receiver_persists_authenticated_tenant(client: TestClient) -> None:
    storage: Final = MagicMock(spec=ClickHouseStorage)
    storage.insert_rows = AsyncMock()
    tracing: Final = TraceReceiver(TraceStore(storage))
    client.app.dependency_overrides[tracing_endpoints.provide_receiver] = lambda: tracing
    response: Final = client.post(
        "/v1/traces",
        json={
            "resourceSpans": [
                {
                    "resource": {
                        "attributes": [
                            {"key": "litellm.team_id", "value": {"stringValue": "spoofed-team"}},
                            {"key": "litellm.api_key_hash", "value": {"stringValue": "spoofed-key"}},
                            {"key": "litellm.org_id", "value": {"stringValue": "spoofed-org"}},
                        ]
                    },
                    "scopeSpans": [
                        {
                            "spans": [
                                {
                                    "traceId": "01" * 16,
                                    "spanId": "02" * 8,
                                    "name": "dependency-injection",
                                    "startTimeUnixNano": "1000000000",
                                    "endTimeUnixNano": "1000000001",
                                }
                            ]
                        }
                    ],
                }
            ],
        },
    )
    assert response.status_code == 200, response.text
    assert response.json() == {}
    storage.insert_rows.assert_awaited_once()
    table, rows = storage.insert_rows.await_args.args
    assert table == "otel_traces"
    assert len(rows) == 1
    assert rows[0]["TeamId"] == TEAM_KEY.team_id
    assert rows[0]["ApiKeyHash"] == TEAM_KEY.token
    assert rows[0]["ResourceAttributes"] == {
        "litellm.team_id": TEAM_KEY.team_id,
        "litellm.api_key_hash": TEAM_KEY.token,
        "litellm.org_id": TEAM_KEY.org_id,
    }


def test_lifespan_receivers_are_app_local() -> None:
    first_storage: Final = MagicMock(spec=ClickHouseStorage)
    first_storage.query = AsyncMock(
        return_value=[
            {
                "span_id": "first-span",
                "input": "first-input",
                "output": "",
                "attributes": {},
            }
        ]
    )
    second_storage: Final = MagicMock(spec=ClickHouseStorage)
    second_storage.query = AsyncMock(
        return_value=[
            {
                "span_id": "second-span",
                "input": "second-input",
                "output": "",
                "attributes": {},
            }
        ]
    )
    first_receiver: Final = TraceReceiver(TraceStore(first_storage))
    second_receiver: Final = TraceReceiver(TraceStore(second_storage))
    first_storage.ensure_schema = AsyncMock()
    second_storage.ensure_schema = AsyncMock()

    @asynccontextmanager
    async def first_lifespan(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]:
        async with manage_tracing(True, lambda: first_receiver) as receiver:
            state: Final[ProxyLifespanState] = {"tracing_receiver": receiver}
            yield state

    @asynccontextmanager
    async def second_lifespan(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]:
        async with manage_tracing(True, lambda: second_receiver) as receiver:
            state: Final[ProxyLifespanState] = {"tracing_receiver": receiver}
            yield state

    first_app: Final = FastAPI(lifespan=first_lifespan)
    second_app: Final = FastAPI(lifespan=second_lifespan)
    first_app.include_router(tracing_endpoints.router)
    second_app.include_router(tracing_endpoints.router)
    first_app.dependency_overrides[user_api_key_auth] = lambda: TEAM_KEY
    second_app.dependency_overrides[user_api_key_auth] = lambda: TEAM_KEY

    with TestClient(first_app) as first_client:
        with TestClient(second_app) as second_client:
            second_response: Final = second_client.get("/v1/traces/t1/spans/second-span?trace_ref=second-run")
            simultaneous: Final = first_client.get("/v1/traces/t1/spans/first-span?trace_ref=first-run")
        first_response: Final = first_client.get("/v1/traces/t1/spans/first-span?trace_ref=first-run")
        assert simultaneous.json() == first_response.json()
    first_storage.ensure_schema.assert_awaited_once()
    second_storage.ensure_schema.assert_awaited_once()

    assert first_response.status_code == second_response.status_code == 200
    assert first_response.json() == {
        "span_id": "first-span",
        "input": "first-input",
        "output": "",
        "attributes": {},
        "input_ui": {"kind": "text", "text": "first-input"},
        "output_ui": {"kind": "text", "text": ""},
    }
    assert second_response.json() == {
        "span_id": "second-span",
        "input": "second-input",
        "output": "",
        "attributes": {},
        "input_ui": {"kind": "text", "text": "second-input"},
        "output_ui": {"kind": "text", "text": ""},
    }
    assert first_storage.query.await_count == 2
    first_storage.query.assert_awaited_with(
        "span_detail",
        {
            "team_ids": (TEAM_KEY.team_id,),
            "api_key_hash": "",
            "trace_id": "t1",
            "span_id": "first-span",
            "trace_ref": "first-run",
        },
    )
    second_storage.query.assert_awaited_once_with(
        "span_detail",
        {
            "team_ids": (TEAM_KEY.team_id,),
            "api_key_hash": "",
            "trace_id": "t1",
            "span_id": "second-span",
            "trace_ref": "second-run",
        },
    )


@pytest.mark.parametrize("auth", [TEAM_KEY, UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER)])
def test_query_validation_precedes_trace_access_checks(client: TestClient, auth: UserAPIKeyAuth) -> None:
    client.app.dependency_overrides[user_api_key_auth] = lambda: auth
    response: Final = client.get("/v1/traces", params={"start_ms": "invalid"})
    assert response.status_code == 422
    assert response.json()["detail"][0]["loc"] == ["query", "start_ms"]


@pytest.mark.parametrize("enabled", [True, False])
def test_unavailable_lifespan_receiver_returns_501(enabled: bool) -> None:
    storage: Final = MagicMock(spec=ClickHouseStorage)
    storage.ensure_schema = AsyncMock(side_effect=RuntimeError("storage unavailable"))
    tracing: Final = TraceReceiver(TraceStore(storage))

    @asynccontextmanager
    async def lifespan(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]:
        async with manage_tracing(enabled, lambda: tracing) as receiver:
            state: Final[ProxyLifespanState] = {"tracing_receiver": receiver}
            yield state

    app: Final = FastAPI(lifespan=lifespan)
    app.include_router(tracing_endpoints.router)
    app.dependency_overrides[user_api_key_auth] = lambda: TEAM_KEY
    with TestClient(app) as client:
        response: Final = client.get("/v1/traces")
    assert response.status_code == 501
    assert storage.ensure_schema.await_count == int(enabled)
    storage.query.assert_not_called()


def test_lens_reads_from_the_lifespan_storage() -> None:
    from litellm.proxy.lens.endpoints import router as lens_router

    storage: Final = MagicMock(spec=ClickHouseStorage)
    storage.ensure_schema = AsyncMock()
    storage.lens_sample = AsyncMock(return_value=[])
    tracing: Final = TraceReceiver(TraceStore(storage))

    @asynccontextmanager
    async def lifespan(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]:
        async with manage_tracing(True, lambda: tracing) as receiver:
            state: Final[ProxyLifespanState] = {"tracing_receiver": receiver}
            yield state

    app: Final = FastAPI(lifespan=lifespan)
    app.include_router(lens_router)
    app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
    with TestClient(app) as client:
        response: Final = client.post(
            "/lens/preview/sample",
            json={"settings": {"name": "Review", "model": "analysis", "context": "Find failed executions"}},
        )
    assert response.status_code == 200, response.text
    assert response.json()["executions"] == []
    storage.lens_sample.assert_awaited_once()
    assert storage.lens_sample.await_args.args[0]["all_teams"] == 1


def test_lens_reads_from_injected_storage_without_receiver() -> None:
    from litellm.proxy.lens.endpoints import router as lens_router
    from litellm.proxy.lens.sources import Storage

    storage: Final = MagicMock(spec=Storage)
    storage.lens_sample = AsyncMock(return_value=[])
    app: Final = FastAPI()
    app.include_router(lens_router)
    app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
    app.dependency_overrides[provide_storage] = lambda: storage

    with TestClient(app) as client:
        response: Final = client.post(
            "/lens/preview/sample",
            json={"settings": {"name": "Review", "model": "analysis", "context": "Find failed executions"}},
        )

    assert response.status_code == 200, response.text
    assert response.json()["executions"] == []
    storage.lens_sample.assert_awaited_once()