litellm/tests/unit/proxy/test_tracing_endpoints.py

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

import asyncio
import json
from collections.abc import AsyncGenerator, Mapping
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from types import ModuleType
from typing import Final, Literal, TypedDict
from unittest.mock import AsyncMock, MagicMock, call

import httpx
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from httpx import Response
from pydantic import JsonValue, TypeAdapter
from typing_extensions import ReadOnly

from litellm.constants import (
    AGENT_TRACING_AGENT_LIST_LIMIT,
    DEFAULT_AGENT_TRACING_RETENTION_DAYS,
    TRACE_READ_RETRY_AFTER_SECONDS,
)
from litellm.proxy import tracing_endpoints
from litellm.proxy._types import LitellmUserRoles, ProxyLifespanState, UserAPIKeyAuth
from litellm.proxy.auth.authorization import OwnedRows, ReadScope
from litellm.proxy.auth.authorization_dependencies import get_log_team_lookup
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 import loader
from litellm.tracing.errors import TraceChanged
from litellm.tracing.generated.models import TraceQueryHelp
from litellm.tracing.generated.responses import TraceSQLResponse
from litellm.tracing.generated.types import AllQueryScope, TraceScope
from litellm.tracing.storage import LensTraceStorage
from litellm.tracing import TraceReceiver
from litellm.tracing.remote import RemoteTraceStore
from litellm.tracing.types import TraceAgent, TraceAgentList

SQL_ROWS: Final[tuple[Mapping[str, JsonValue], ...]] = (
    {
        "value": "9007199254740993",
        "count": 42,
        "fraction": 2.5,
        "nested": {"values": [True, None, "text"]},
    },
)
SQL_RESPONSE: Final = TraceSQLResponse(data=SQL_ROWS)
QUERY_HELP: Final[Mapping[str, object]] = {
    "dialect": "test SQL",
    "access": "authenticated scope",
    "response": (
        'JSON object {"data": [rows]}; each row maps selected columns to values; 64-bit integers may be strings'
    ),
    "tables": [{"name": "otel_traces", "columns": [{"name": "value", "type": "String", "comment": "label"}]}],
    "normalized_fields": [],
    "metadata": {
        "table": "spend_logs",
        "column": "metadata",
        "fields": [],
        "sampled_rows": 0,
        "invalid_json_rows": 0,
        "truncated": True,
        "sample_sql": "SELECT metadata FROM traces",
        "scope": "bounded sample",
        "error": "discovery unavailable",
    },
    "attributes": [],
    "relationships": [],
    "examples": [{"name": "recent", "sql": "SELECT * FROM traces LIMIT 1"}],
    "gotchas": ["Keep queries bounded"],
    "guide": "scoped",
}


TEAM_KEY = UserAPIKeyAuth(
    user_id="user",
    token="hashed-key",
    team_id="team-research",
    org_id="org-1",
    user_role=LitellmUserRoles.INTERNAL_USER,
)
TRACE_RESPONSE: Final = {
    "summary": {
        "trace_id": "t1",
        "name": "trace",
        "service": "test",
        "input_preview": "",
        "start_time": "2026-01-01T00:00:00Z",
        "duration_ms": 0,
        "status": "ok",
        "span_count": 0,
        "agent_count": 0,
        "agent_invocations": 0,
        "llm_calls": 0,
        "priced_calls": 0,
        "tool_calls": 0,
        "error_count": 0,
        "input_tokens": 0,
        "output_tokens": 0,
        "models": [],
        "spend": None,
    },
    "agents": [],
    "spans": [],
}
SPAN_DETAIL_RESPONSE: Final = {
    "span_id": "s1",
    "input": "",
    "output": "",
    "input_ui": {"kind": "text", "text": ""},
    "output_ui": {"kind": "text", "text": ""},
    "attributes": {},
}
SPAN_ERROR_RESPONSE: Final = {
    "span_id": "s1",
    "message": "span error",
    "total_chars": 10,
    "next_cursor": None,
}
NOW_MS: Final = 1_800_000_000_000


class RequestValidationError(TypedDict):
    type: ReadOnly[str]
    loc: ReadOnly[list[str | int]]


def _validation_errors(response: Response) -> tuple[RequestValidationError, ...]:
    return tuple(TypeAdapter(list[RequestValidationError]).validate_python(response.json()["detail"]))


def _assert_validation_error(response: Response, error_type: str, location: tuple[str | int, ...]) -> None:
    assert response.status_code == 422, response.text
    assert any(
        error["type"] == error_type and tuple(error["loc"]) == location for error in _validation_errors(response)
    )


@pytest.mark.parametrize(
    ("auth", "scope"),
    (
        pytest.param(
            UserAPIKeyAuth(token="admin-key", team_id="team-a", user_role=LitellmUserRoles.PROXY_ADMIN),
            TraceScope(all_teams=1, user_id="", team_ids=()),
            id="admin",
        ),
        pytest.param(
            UserAPIKeyAuth(token="view-key", team_id="team-a", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY),
            TraceScope(all_teams=1, user_id="", team_ids=()),
            id="view-only-admin",
        ),
        pytest.param(
            TEAM_KEY,
            TraceScope(all_teams=0, user_id="user", team_ids=()),
            id="team-key",
        ),
        pytest.param(
            UserAPIKeyAuth(user_id="user", token="hashed-key", user_role=LitellmUserRoles.INTERNAL_USER),
            TraceScope(all_teams=0, user_id="user", team_ids=()),
            id="teamless-key",
        ),
        pytest.param(
            UserAPIKeyAuth(token="hashed-key", user_role=LitellmUserRoles.INTERNAL_USER),
            None,
            id="key-without-user-can-only-write",
        ),
    ),
)
def test_trace_read_permissions_with_retired_uploads(
    client: TestClient, receiver: MagicMock, auth: UserAPIKeyAuth, scope: TraceScope | None
) -> 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 == (403 if scope is None else 200), read.text
    if scope is None:
        receiver.list_traces.assert_not_awaited()
    else:
        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 == 410
    receiver.ingest.assert_not_called()


@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)
    fake.list_agents = AsyncMock(return_value=TraceAgentList(agents=()))
    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

    async def lookup(auth: UserAPIKeyAuth) -> tuple[str, ...]:
        return ()

    app.dependency_overrides[get_log_team_lookup] = lambda: lookup
    return TestClient(app)


@pytest.mark.parametrize("endpoint", ("/v1/traces", "/v1/logs"))
@pytest.mark.parametrize("media_type", ("application/json", "application/x-protobuf"))
def test_gateway_uploads_return_setup_guidance_without_reading_the_body(
    client: TestClient, receiver: MagicMock, endpoint: str, media_type: str
) -> None:
    from google.rpc.status_pb2 import Status

    response: Final = client.post(endpoint, content=b"invalid payload", headers={"content-type": media_type})
    assert response.status_code == 410
    assert response.headers["content-type"] == media_type
    message: Final = (
        response.json()["message"] if media_type == "application/json" else Status.FromString(response.content).message
    )
    assert message == "Send traces and logs directly to the Lens endpoint shown in Lens setup."
    receiver.ingest.assert_not_called()


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={"all_teams": 0, "user_id": "user", "team_ids": ()},
        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


@pytest.mark.parametrize(
    ("params", "expected_start_ms", "expected_end_ms"),
    (
        ({}, NOW_MS - tracing_endpoints.MS_PER_DAY, NOW_MS),
        ({"start_ms": 123}, 123, NOW_MS),
        ({"end_ms": -7}, NOW_MS - tracing_endpoints.MS_PER_DAY, -7),
    ),
)
def test_list_traces_resolves_default_bounds_from_injected_clock(
    client: TestClient,
    receiver: MagicMock,
    params: Mapping[str, int],
    expected_start_ms: int,
    expected_end_ms: int,
) -> None:
    client.app.dependency_overrides[tracing_endpoints.current_time_ms] = lambda: NOW_MS
    response: Final = client.get("/v1/traces", params=params)
    assert response.status_code == 200, response.text
    receiver.list_traces.assert_awaited_once_with(
        scope={"all_teams": 0, "user_id": "user", "team_ids": ()},
        start_ms=expected_start_ms,
        end_ms=expected_end_ms,
        cursor=None,
    )


@pytest.mark.parametrize(
    ("params", "expected_start_ms", "expected_end_ms"),
    (
        ({}, NOW_MS - DEFAULT_AGENT_TRACING_RETENTION_DAYS * tracing_endpoints.MS_PER_DAY, NOW_MS),
        ({"start_ms": 123, "end_ms": 456}, 123, 456),
    ),
)
def test_list_trace_agents_passes_reader_scope_and_window(
    client: TestClient,
    receiver: MagicMock,
    params: Mapping[str, int],
    expected_start_ms: int,
    expected_end_ms: int,
) -> None:
    client.app.dependency_overrides[tracing_endpoints.current_time_ms] = lambda: NOW_MS
    receiver.list_agents.return_value = TraceAgentList(
        agents=(
            TraceAgent(
                name="moyai",
                runs=3,
                failed_runs=1,
                last_seen=datetime(2026, 10, 7, 20, 31, tzinfo=timezone.utc),
                frameworks=("openai-agents",),
            ),
        )
    )
    response: Final = client.get("/v1/traces/agents", params=params)
    assert response.status_code == 200, response.text
    assert response.json() == {
        "agents": [
            {
                "name": "moyai",
                "runs": 3,
                "failed_runs": 1,
                "last_seen": "2026-10-07T20:31:00Z",
                "frameworks": ["openai-agents"],
            }
        ]
    }
    receiver.list_agents.assert_awaited_once_with(
        scope={"all_teams": 0, "user_id": "user", "team_ids": ()},
        start_ms=expected_start_ms,
        end_ms=expected_end_ms,
    )
    receiver.get_trace.assert_not_awaited()


def test_list_trace_agents_requires_read_access(client: TestClient, receiver: MagicMock) -> None:
    client.app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
        token="hashed-key", user_role=LitellmUserRoles.INTERNAL_USER
    )
    response: Final = client.get("/v1/traces/agents")
    assert response.status_code == 403, response.text
    receiver.list_agents.assert_not_awaited()


def test_list_trace_agents_maps_storage_outage_to_503(client: TestClient, receiver: MagicMock) -> None:
    receiver.list_agents.side_effect = RuntimeError("private database details")
    response: Final = client.get("/v1/traces/agents")
    assert response.status_code == 503
    assert response.json()["detail"]["code"] == "unavailable"


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("role", "expected_scope"),
    (
        pytest.param(
            LitellmUserRoles.PROXY_ADMIN,
            TraceScope(all_teams=1, user_id="", team_ids=()),
            id="admin",
        ),
        pytest.param(
            LitellmUserRoles.INTERNAL_USER,
            TraceScope(all_teams=0, user_id="agent-owner", team_ids=("managed-team",)),
            id="owner-and-permitted-teams",
        ),
    ),
)
async def test_agent_picker_reads_through_worker_with_authenticated_scope(
    client: TestClient, role: LitellmUserRoles, expected_scope: TraceScope
) -> None:
    requests: Final = asyncio.Queue[httpx.Request]()
    secret: Final = "test-only-lens-service-secret-32-characters"

    def accept(request: httpx.Request) -> httpx.Response:
        requests.put_nowait(request)
        return httpx.Response(
            200,
            json={
                "data": [
                    {
                        "agent_name": "research-agent",
                        "runs": "3",
                        "failed_runs": "1",
                        "last_seen_ms": "1791405060000",
                        "frameworks": ["openai-agents"],
                    }
                ]
            },
        )

    async def lookup(auth: UserAPIKeyAuth) -> tuple[str, ...]:
        return ("managed-team",)

    client.app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
        user_id="agent-owner", token="user-key", team_id="unmanaged-team", user_role=role
    )
    client.app.dependency_overrides[get_log_team_lookup] = lambda: lookup
    async with httpx.AsyncClient(
        base_url="http://lens",
        headers={"Authorization": f"Bearer {secret}"},
        transport=httpx.MockTransport(accept),
    ) as worker:
        tracing: Final = TraceReceiver(storage=LensTraceStorage(RemoteTraceStore(worker)))
        client.app.dependency_overrides[tracing_endpoints.provide_receiver] = lambda: tracing
        async with httpx.AsyncClient(
            transport=httpx.ASGITransport(app=client.app), base_url="http://gateway"
        ) as gateway:
            response: Final = await gateway.get("/v1/traces/agents", params={"start_ms": 123, "end_ms": 456})
    assert response.status_code == 200, response.text
    assert response.json() == {
        "agents": [
            {
                "name": "research-agent",
                "runs": 3,
                "failed_runs": 1,
                "last_seen": "2026-10-07T20:31:00Z",
                "frameworks": ["openai-agents"],
            }
        ]
    }
    request: Final = requests.get_nowait()
    assert requests.empty()
    assert request.method == "POST"
    assert request.url.path == "/internal/read"
    assert request.headers["Authorization"] == f"Bearer {secret}"
    assert json.loads(request.content) == {
        "operation": "query",
        "name": "trace_agents",
        "parameters": {
            **expected_scope,
            "team_ids": list(expected_scope["team_ids"]),
            "start_ms": 123,
            "end_ms": 456,
            "limit": AGENT_TRACING_AGENT_LIST_LIMIT,
        },
    }


@pytest.mark.asyncio
@pytest.mark.parametrize(("status", "code"), ((503, "unavailable"), (413, "too_large")))
async def test_agent_picker_reports_worker_failures_without_leaking_details(
    client: TestClient, status: int, code: str
) -> None:
    async with httpx.AsyncClient(
        base_url="http://lens",
        transport=httpx.MockTransport(lambda request: httpx.Response(status, text="private storage details")),
    ) as worker:
        tracing: Final = TraceReceiver(storage=LensTraceStorage(RemoteTraceStore(worker)))
        client.app.dependency_overrides[tracing_endpoints.provide_receiver] = lambda: tracing
        async with httpx.AsyncClient(
            transport=httpx.ASGITransport(app=client.app), base_url="http://gateway"
        ) as gateway:
            response: Final = await gateway.get("/v1/traces/agents", params={"start_ms": 123, "end_ms": 456})
    assert response.status_code == status
    assert response.json()["detail"]["code"] == code
    assert "private storage details" not in response.text
    if status == 503:
        assert response.headers["Retry-After"] == str(TRACE_READ_RETRY_AFTER_SECONDS)


def test_list_traces_forwards_large_and_negative_bounds_unchanged(client: TestClient, receiver: MagicMock) -> None:
    response: Final = client.get("/v1/traces", params={"start_ms": 2**63, "end_ms": -1, "cursor": "next"})
    assert response.status_code == 200, response.text
    receiver.list_traces.assert_awaited_once_with(
        scope={"all_teams": 0, "user_id": "user", "team_ids": ()},
        start_ms=2**63,
        end_ms=-1,
        cursor="next",
    )


def test_get_trace_404_and_200(client, receiver):
    assert client.get("/v1/traces/missing").status_code == 404
    receiver.get_trace.return_value = TRACE_RESPONSE
    response = client.get("/v1/traces/t1")
    assert response.status_code == 200
    assert response.json() == TRACE_RESPONSE
    receiver.get_trace.assert_awaited_with("t1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "", None, None)


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_DETAIL_RESPONSE
    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", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "")


@pytest.mark.parametrize("suffix,cursor,page_size", [("", None, None), ("&cursor=next&page_size=200", "next", 200)])
def test_trace_detail_passes_scoped_reference(client, receiver, suffix, cursor, page_size):
    receiver.get_trace.return_value = TRACE_RESPONSE
    assert client.get(f"/v1/traces/t1?trace_ref=run-one{suffix}").status_code == 200
    receiver.get_trace.assert_awaited_with(
        "t1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "run-one", cursor, page_size
    )


@pytest.mark.parametrize("page_size", (1, 500))
def test_trace_detail_accepts_page_size_bounds(client: TestClient, receiver: MagicMock, page_size: int) -> None:
    receiver.get_trace.return_value = TRACE_RESPONSE
    response: Final = client.get("/v1/traces/t1", params={"trace_ref": "run-one", "page_size": page_size})
    assert response.status_code == 200, response.text
    receiver.get_trace.assert_awaited_once_with(
        "t1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "run-one", None, page_size
    )


@pytest.mark.parametrize(
    ("value", "error_type"),
    (("0", "greater_than_equal"), ("501", "less_than_equal"), ("abc", "int_parsing")),
)
def test_trace_detail_reports_page_size_validation(
    client: TestClient, receiver: MagicMock, value: str, error_type: str
) -> None:
    response: Final = client.get("/v1/traces/t1", params={"page_size": value})
    _assert_validation_error(response, error_type, ("query", "page_size"))
    receiver.get_trace.assert_not_awaited()


def test_trace_read_routes_accept_and_forward_512_character_cursors(client: TestClient, receiver: MagicMock) -> None:
    cursor: Final = "x" * 512
    receiver.get_trace.return_value = TRACE_RESPONSE
    receiver.get_span_error = AsyncMock(return_value=SPAN_ERROR_RESPONSE)
    client.app.dependency_overrides[tracing_endpoints.current_time_ms] = lambda: NOW_MS

    list_response: Final = client.get("/v1/traces", params={"cursor": cursor})
    detail_response: Final = client.get("/v1/traces/t1", params={"trace_ref": "run-one", "cursor": cursor})
    error_response: Final = client.get(
        "/v1/traces/t1/spans/s1/error", params={"trace_ref": "run-one", "cursor": cursor}
    )

    assert list_response.status_code == 200, list_response.text
    assert detail_response.status_code == 200, detail_response.text
    assert error_response.status_code == 200, error_response.text
    receiver.list_traces.assert_awaited_once_with(
        scope={"all_teams": 0, "user_id": "user", "team_ids": ()},
        start_ms=NOW_MS - tracing_endpoints.MS_PER_DAY,
        end_ms=NOW_MS,
        cursor=cursor,
    )
    receiver.get_trace.assert_awaited_once_with(
        "t1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "run-one", cursor, None
    )
    receiver.get_span_error.assert_awaited_once_with(
        "t1", "s1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "run-one", cursor
    )


@pytest.mark.parametrize(
    "path",
    ("/v1/traces", "/v1/traces/t1", "/v1/traces/t1/spans/s1/error"),
)
def test_trace_read_routes_reject_513_character_cursors(client: TestClient, receiver: MagicMock, path: str) -> None:
    response: Final = client.get(path, params={"cursor": "x" * 513})
    _assert_validation_error(response, "string_too_long", ("query", "cursor"))


def test_trace_read_routes_ignore_unknown_query_parameters(client: TestClient, receiver: MagicMock) -> None:
    receiver.get_trace.return_value = TRACE_RESPONSE
    receiver.get_span.return_value = SPAN_DETAIL_RESPONSE
    receiver.get_span_error = AsyncMock(return_value=SPAN_ERROR_RESPONSE)

    list_params: Final = {"start_ms": 1, "end_ms": 2, "cursor": "list-cursor"}
    list_response: Final = client.get("/v1/traces", params=list_params)
    list_unknown_response: Final = client.get("/v1/traces", params={**list_params, "foo": "bar"})
    detail_params: Final = {"trace_ref": "run-one", "cursor": "detail-cursor", "page_size": 10}
    detail_response: Final = client.get("/v1/traces/t1", params=detail_params)
    detail_unknown_response: Final = client.get("/v1/traces/t1", params={**detail_params, "foo": "bar"})
    span_response: Final = client.get("/v1/traces/t1/spans/s1", params={"trace_ref": "run-one"})
    span_unknown_response: Final = client.get("/v1/traces/t1/spans/s1", params={"trace_ref": "run-one", "foo": "bar"})
    error_params: Final = {"trace_ref": "run-one", "cursor": "error-cursor"}
    error_response: Final = client.get("/v1/traces/t1/spans/s1/error", params=error_params)
    error_unknown_response: Final = client.get("/v1/traces/t1/spans/s1/error", params={**error_params, "foo": "bar"})

    assert list_response.status_code == 200, list_response.text
    assert list_unknown_response.status_code == 200, list_unknown_response.text
    assert detail_response.status_code == 200, detail_response.text
    assert detail_unknown_response.status_code == 200, detail_unknown_response.text
    assert span_response.status_code == 200, span_response.text
    assert span_unknown_response.status_code == 200, span_unknown_response.text
    assert error_response.status_code == 200, error_response.text
    assert error_unknown_response.status_code == 200, error_unknown_response.text
    receiver.list_traces.assert_has_awaits(
        (
            call(scope={"all_teams": 0, "user_id": "user", "team_ids": ()}, start_ms=1, end_ms=2, cursor="list-cursor"),
            call(scope={"all_teams": 0, "user_id": "user", "team_ids": ()}, start_ms=1, end_ms=2, cursor="list-cursor"),
        )
    )
    receiver.get_trace.assert_has_awaits(
        (
            call("t1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "run-one", "detail-cursor", 10),
            call("t1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "run-one", "detail-cursor", 10),
        )
    )
    receiver.get_span.assert_has_awaits(
        (
            call("t1", "s1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "run-one"),
            call("t1", "s1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "run-one"),
        )
    )
    receiver.get_span_error.assert_has_awaits(
        (
            call("t1", "s1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "run-one", "error-cursor"),
            call("t1", "s1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "run-one", "error-cursor"),
        )
    )


@pytest.mark.parametrize(
    "path,method",
    (
        ("/v1/traces", "list_traces"),
        ("/v1/traces/t1", "get_trace"),
        ("/v1/traces/t1/spans/s1", "get_span"),
        ("/v1/traces/t1/spans/s1/error", "get_span_error"),
    ),
)
@pytest.mark.parametrize(
    "error,status,code,message",
    (
        (
            RuntimeError("private database details"),
            503,
            "unavailable",
            "Traces are temporarily unavailable. Please try again.",
        ),
        (
            OverflowError("private query details"),
            413,
            "too_large",
            "Trace is too large for this view. Use a filtered trace query.",
        ),
        (
            TraceChanged("Trace changed while paging; refresh the trace to continue"),
            409,
            "trace_changed",
            "Trace changed while paging; refresh the trace to continue",
        ),
        (ValueError("Invalid span cursor"), 400, "invalid_request", "Invalid span cursor"),
    ),
)
def test_read_failures_carry_a_code_per_kind_without_exposing_database_details(
    client: TestClient,
    receiver: MagicMock,
    path: str,
    method: str,
    error: Exception,
    status: int,
    code: str,
    message: str,
) -> None:
    getattr(receiver, method).side_effect = error
    response: Final = client.get(path)
    assert response.status_code == status
    assert response.json() == {"detail": {"code": code, "message": message}}
    retry_after: Final = response.headers.get("Retry-After")
    assert (retry_after == str(TRACE_READ_RETRY_AFTER_SECONDS)) == (status == 503), retry_after


@pytest.mark.parametrize(
    "query",
    ("page_size=0", "page_size=501", "cursor=" + "x" * 513),
    ids=("zero-page-size", "oversized-page-size", "oversized-cursor"),
)
def test_trace_page_rejects_unbounded_parameters(client: TestClient, receiver: MagicMock, query: str) -> None:
    response: Final = client.get(f"/v1/traces/t1?{query}")
    assert response.status_code == 422
    receiver.get_trace.assert_not_awaited()


def test_invalid_cursor_is_a_client_error(client: TestClient, receiver: MagicMock) -> None:
    receiver.list_traces.side_effect = ValueError("Invalid trace cursor")
    assert client.get("/v1/traces?cursor=broken").status_code == 400


@pytest.mark.parametrize(
    "auth",
    (
        UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER),
        UserAPIKeyAuth(token="key"),
        UserAPIKeyAuth(token="key", team_id="unpermitted"),
        UserAPIKeyAuth(user_id="", token="key"),
    ),
)
def test_key_without_user_cannot_read_traces(client: TestClient, auth: UserAPIKeyAuth) -> None:
    storage: Final = MagicMock(spec=LensTraceStorage)
    client.app.dependency_overrides[user_api_key_auth] = lambda: auth
    client.app.dependency_overrides[tracing_endpoints.provide_receiver] = lambda: TraceReceiver(storage)
    client.app.dependency_overrides[tracing_endpoints.provide_trace_query_secret] = lambda: "test-secret"
    for path in (
        "/v1/traces",
        "/v1/traces/t1",
        "/v1/traces/t1/spans/s1",
        "/v1/traces/t1/spans/s1/error",
        "/v1/traces/query/help",
    ):
        response: Final = client.get(path)
        assert response.status_code == 403, response.text
    query: Final = client.post("/v1/traces/query", json={"sql": "SELECT * FROM otel_traces"})
    assert query.status_code == 403, query.text
    for read in (storage.list_traces, storage.get_trace, storage.get_span, storage.get_span_error):
        read.assert_not_called()
    storage.query_sql.assert_not_called()
    storage.query_help.assert_not_called()


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. Configure the Lens service and LITELLM_LENS_URL."
    }


def test_lifespan_receivers_are_app_local(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("LITELLM_LENS_URL", "http://lens.test")
    monkeypatch.setenv("LITELLM_LENS_SERVICE_TOKEN", "test-service-token-with-32-characters")
    first_storage: Final = MagicMock(spec=LensTraceStorage)
    first_storage.get_span = AsyncMock(return_value={**SPAN_DETAIL_RESPONSE, "span_id": "first-span"})
    second_storage: Final = MagicMock(spec=LensTraceStorage)
    second_storage.get_span = AsyncMock(return_value={**SPAN_DETAIL_RESPONSE, "span_id": "second-span"})
    first_receiver: Final = TraceReceiver(first_storage)
    second_receiver: Final = TraceReceiver(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_not_awaited()
    second_storage.ensure_schema.assert_not_awaited()

    assert first_response.status_code == second_response.status_code == 200
    assert first_response.json()["span_id"] == "first-span"
    assert second_response.json()["span_id"] == "second-span"
    scope: Final = TraceScope(all_teams=0, user_id=TEAM_KEY.user_id or "", team_ids=())
    assert first_storage.get_span.await_count == 2
    first_storage.get_span.assert_awaited_with("t1", "first-span", scope, "first-run")
    second_storage.get_span.assert_awaited_once_with("t1", "second-span", scope, "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_unconfigured_lifespan_receiver_returns_501(enabled: bool, monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.delenv("LITELLM_LENS_URL", raising=False)
    storage: Final = MagicMock(spec=LensTraceStorage)
    storage.ensure_schema = AsyncMock(side_effect=RuntimeError("storage unavailable"))
    tracing: Final = TraceReceiver(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
    storage.ensure_schema.assert_not_awaited()
    storage.list_traces.assert_not_called()


@pytest.mark.parametrize(
    ("auth", "expected_scope"),
    (
        (UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), {"kind": "all"}),
        (UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY), {"kind": "all"}),
        (TEAM_KEY, {"kind": "owned", "user_id": "user", "team_ids": ()}),
        (
            UserAPIKeyAuth(user_id="user", token="project-key", team_id="team-a", project_id="project-a"),
            {"kind": "owned", "user_id": "user", "team_ids": ()},
        ),
        (
            UserAPIKeyAuth(user_id="user", token="solo-key"),
            {"kind": "owned", "user_id": "user", "team_ids": ()},
        ),
    ),
)
def test_sql_and_help_use_authenticated_scope(
    client: TestClient, receiver: MagicMock, auth: UserAPIKeyAuth, expected_scope: dict[str, str]
) -> None:
    client.app.dependency_overrides[user_api_key_auth] = lambda: auth
    client.app.dependency_overrides[tracing_endpoints.provide_trace_query_secret] = lambda: "test-secret"
    receiver.storage.query_sql = AsyncMock(return_value=SQL_RESPONSE)
    receiver.storage.query_help = AsyncMock(return_value=TraceQueryHelp.model_validate(QUERY_HELP))
    result: Final = client.post("/v1/traces/query", json={"sql": "SELECT * FROM otel_traces"})
    assert result.status_code == 200, result.text
    assert result.json() == {"data": list(SQL_ROWS)}
    assert type(result.json()["data"][0]["value"]) is str
    assert type(result.json()["data"][0]["count"]) is int
    assert type(result.json()["data"][0]["fraction"]) is float
    receiver.storage.query_sql.assert_awaited_once_with("SELECT * FROM otel_traces", expected_scope, "test-secret")
    help_result: Final = client.get("/v1/traces/query/help")
    assert help_result.status_code == 200, help_result.text
    assert help_result.json() == QUERY_HELP
    receiver.storage.query_help.assert_awaited_once_with(expected_scope, "test-secret")
    forged: Final = client.post("/v1/traces/query", json={"sql": "SELECT 1", "scope": {"kind": "all"}})
    assert forged.status_code == 422, forged.text
    assert receiver.storage.query_sql.await_count == 1


def test_sql_query_returns_empty_data(client: TestClient, receiver: MagicMock) -> None:
    client.app.dependency_overrides[tracing_endpoints.provide_trace_query_secret] = lambda: "test-secret"
    receiver.storage.query_sql = AsyncMock(return_value=TraceSQLResponse(data=()))

    result: Final = client.post("/v1/traces/query", json={"sql": "SELECT * FROM otel_traces"})

    assert result.status_code == 200, result.text
    assert result.json() == {"data": []}


def test_sql_query_openapi_declares_a_closed_response_object(client: TestClient) -> None:
    openapi: Final = client.app.openapi()
    response: Final = openapi["paths"]["/v1/traces/query"]["post"]["responses"]["200"]["content"]["application/json"][
        "schema"
    ]
    component_name: Final = response["$ref"].rsplit("/", 1)[-1]
    component: Final = openapi["components"]["schemas"][component_name]

    assert set(component["properties"]) == {"data"}
    assert component["additionalProperties"] is False


@pytest.mark.parametrize(
    ("body", "error_type", "location"),
    (
        (b"{}", "missing", ("body", "sql")),
        (b'{"sql": null}', "string_type", ("body", "sql")),
        (b'{"sql": 1}', "string_type", ("body", "sql")),
        (b'{"sql": "SELECT 1", "extra": true}', "extra_forbidden", ("body", "extra")),
        (b"{", "json_invalid", ("body", 1)),
        (b"[]", "model_attributes_type", ("body",)),
    ),
)
def test_sql_query_rejects_invalid_request_bodies(
    client: TestClient, receiver: MagicMock, body: bytes, error_type: str, location: tuple[str | int, ...]
) -> None:
    client.app.dependency_overrides[tracing_endpoints.provide_trace_query_secret] = lambda: "test-secret"
    receiver.storage.query_sql = AsyncMock()
    response: Final = client.post("/v1/traces/query", content=body, headers={"content-type": "application/json"})
    _assert_validation_error(response, error_type, location)
    receiver.storage.query_sql.assert_not_awaited()


@pytest.mark.parametrize("auth", (UserAPIKeyAuth(), UserAPIKeyAuth(team_id="a", project_id="p")))
def test_sql_rejects_missing_identity_without_querying(
    client: TestClient, receiver: MagicMock, auth: UserAPIKeyAuth
) -> None:
    client.app.dependency_overrides[user_api_key_auth] = lambda: auth
    client.app.dependency_overrides[tracing_endpoints.provide_trace_query_secret] = lambda: "test-secret"
    result: Final = client.post("/v1/traces/query", json={"sql": "SELECT * FROM otel_traces"})
    assert result.status_code == 403, result.text
    assert client.get("/v1/traces/query/help").status_code == 403
    receiver.storage.query_sql.assert_not_called()
    receiver.storage.query_help.assert_not_called()


@pytest.mark.parametrize(
    ("error", "status"), ((ValueError("invalid SQL"), 400), (RuntimeError("reader unavailable"), 503))
)
def test_sql_reports_rejected_queries_and_unavailable_readers(
    client: TestClient, receiver: MagicMock, error: Exception, status: int
) -> None:
    client.app.dependency_overrides[tracing_endpoints.provide_trace_query_secret] = lambda: "test-secret"
    receiver.storage.query_sql = AsyncMock(side_effect=error)
    result: Final = client.post("/v1/traces/query", json={"sql": "SELECT 1"})
    assert result.status_code == status, result.text
    receiver.storage.query_sql.assert_awaited_once_with(
        "SELECT 1", {"kind": "owned", "user_id": "user", "team_ids": ()}, "test-secret"
    )


def test_query_help_does_not_fall_back_when_reader_provisioning_fails(client: TestClient, receiver: MagicMock) -> None:
    client.app.dependency_overrides[tracing_endpoints.provide_trace_query_secret] = lambda: "test-secret"
    receiver.storage.query_help = AsyncMock(side_effect=RuntimeError("reader provisioning failed"))
    result: Final = client.get("/v1/traces/query/help")
    assert result.status_code == 503, result.text
    receiver.storage.query_help.assert_awaited_once_with(
        {"kind": "owned", "user_id": "user", "team_ids": ()}, "test-secret"
    )


@pytest.mark.parametrize("secret", (None, "configured-master-key"))
def test_queries_require_a_proxy_secret(
    client: TestClient, receiver: MagicMock, monkeypatch: pytest.MonkeyPatch, secret: str | None
) -> None:
    from litellm.proxy import proxy_server

    monkeypatch.setattr(proxy_server, "master_key", secret)
    receiver.storage.query_sql = AsyncMock(return_value=SQL_RESPONSE)
    result: Final = client.post("/v1/traces/query", json={"sql": "SELECT 1"})
    if secret is None:
        assert result.status_code == 503, result.text
        assert "master key" in result.json()["detail"]
        receiver.storage.query_sql.assert_not_awaited()
        return
    assert result.status_code == 200, result.text
    receiver.storage.query_sql.assert_awaited_once_with(
        "SELECT 1", {"kind": "owned", "user_id": "user", "team_ids": ()}, secret
    )


@pytest.mark.parametrize(
    ("auth", "teams", "expected"),
    (
        (UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), ("team-a",), (1, "", ())),
        (UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY), ("team-a",), (1, "", ())),
        (UserAPIKeyAuth(user_id="user", token="key", team_id="unpermitted"), ("a", "b"), (0, "user", ("a", "b"))),
        (UserAPIKeyAuth(user_id="user", token="key"), (), (0, "user", ())),
        (UserAPIKeyAuth(user_id="user"), ("a",), (0, "user", ("a",))),
    ),
)
def test_shared_trace_permissions_reach_read_and_sql_boundaries(
    client: TestClient,
    auth: UserAPIKeyAuth,
    teams: tuple[str, ...],
    expected: tuple[Literal[0, 1], str, tuple[str, ...]],
) -> None:
    async def lookup(caller: UserAPIKeyAuth) -> tuple[str, ...]:
        assert caller is auth
        return teams

    team_lookup: Final = AsyncMock(side_effect=lookup)
    storage: Final = MagicMock(spec=LensTraceStorage)
    storage.get_span = AsyncMock(return_value=SPAN_DETAIL_RESPONSE)
    storage.query_sql = AsyncMock(return_value=SQL_RESPONSE)
    storage.query_help = AsyncMock(return_value=TraceQueryHelp.model_validate(QUERY_HELP))
    client.app.dependency_overrides[user_api_key_auth] = lambda: auth
    client.app.dependency_overrides[get_log_team_lookup] = lambda: team_lookup
    client.app.dependency_overrides[tracing_endpoints.provide_receiver] = lambda: TraceReceiver(storage)
    client.app.dependency_overrides[tracing_endpoints.provide_trace_query_secret] = lambda: "test-secret"

    response: Final = client.get("/v1/traces/t1/spans/s1?trace_ref=run-one")
    assert response.status_code == 200, response.text
    assert response.json()["span_id"] == "s1"
    storage.get_span.assert_awaited_once_with(
        "t1", "s1", TraceScope(all_teams=expected[0], user_id=expected[1], team_ids=expected[2]), "run-one"
    )
    sql_response: Final = client.post("/v1/traces/query", json={"sql": "SELECT * FROM otel_traces"})
    assert sql_response.status_code == 200, sql_response.text
    assert sql_response.json() == {"data": list(SQL_ROWS)}
    assert client.get("/v1/traces/query/help").json() == QUERY_HELP
    query_scope: Final = (
        {"kind": "all"}
        if expected[0]
        else {
            "kind": "owned",
            "user_id": expected[1],
            "team_ids": expected[2],
        }
    )
    storage.query_sql.assert_awaited_once_with("SELECT * FROM otel_traces", query_scope, "test-secret")
    storage.query_help.assert_awaited_once_with(query_scope, "test-secret")
    assert team_lookup.await_count == (
        3
        if auth.user_id and auth.user_role not in (LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY)
        else 0
    )


@pytest.mark.parametrize(
    ("scope", "expected"),
    (
        (OwnedRows(None), ("", ())),
        (OwnedRows("user"), ("user", ())),
        (OwnedRows("user", ("a", "b")), ("user", ("a", "b"))),
    ),
)
def test_trace_storage_permissions_map_owned_rows(
    scope: ReadScope,
    expected: tuple[str, tuple[str, ...]],
) -> None:
    assert tracing_endpoints._trace_scope(scope) == TraceScope(all_teams=0, user_id=expected[0], team_ids=expected[1])
    assert tracing_endpoints.trace_query_scope(scope) == {
        "kind": "owned",
        "user_id": expected[0],
        "team_ids": expected[1],
    }


@pytest.mark.parametrize("cursor,page_size", ((None, None), ("next", 200)))
async def test_storage_preserves_page_cursor_and_normalizes_remote_trace_data(
    cursor: str | None, page_size: int | None
) -> None:
    scope: Final[TraceScope] = {"all_teams": 0, "user_id": "owner", "team_ids": ()}

    def accept(request: httpx.Request) -> httpx.Response:
        assert json.loads(request.content) == {
            "operation": "trace", "scope": {**scope, "team_ids": []}, "trace_id": "t1", "trace_ref": "run",
            "cursor": cursor, "page_size": page_size,
        }
        return httpx.Response(200, json={**TRACE_RESPONSE, "next_cursor": "more"})

    async with httpx.AsyncClient(base_url="http://lens", transport=httpx.MockTransport(accept)) as client:
        trace: Final = await LensTraceStorage(RemoteTraceStore(client)).get_trace("t1", scope, "run", cursor, page_size)
    assert trace is not None
    assert trace["next_cursor"] == "more"
    assert trace["spans"] == ()
    assert trace["summary"]["span_count"] == 0


async def test_storage_validates_the_remote_query_help_value() -> None:
    async with httpx.AsyncClient(
        base_url="http://lens", transport=httpx.MockTransport(lambda request: httpx.Response(200, json=QUERY_HELP))
    ) as client:
        assert await LensTraceStorage(RemoteTraceStore(client)).query_help({"kind": "all"}, "secret") == TraceQueryHelp.model_validate(QUERY_HELP)


@pytest.mark.parametrize(
    "drift",
    (
        {
            "metadata": {
                "table": "spend_logs",
                "column": "metadata",
                "fields": [{"path": ["a"], "types": ["boolen"], "expression": "a"}],
                "sampled_rows": 1,
                "invalid_json_rows": 0,
                "truncated": False,
                "sample_sql": "SELECT metadata FROM spend_logs",
                "scope": "bounded sample",
            }
        },
        {"tables": [{"name": "traces", "columns": [{"name": "value", "type": "String"}]}]},
        {"unexpected": True},
    ),
)
async def test_storage_rejects_remote_query_help_that_drifts_from_the_contract(drift: Mapping[str, object]) -> None:
    async with httpx.AsyncClient(
        base_url="http://lens", transport=httpx.MockTransport(lambda request: httpx.Response(200, json={**QUERY_HELP, **drift}))
    ) as client:
        with pytest.raises(RuntimeError, match="invalid response"):
            await LensTraceStorage(RemoteTraceStore(client)).query_help({"kind": "all"}, "secret")