litellm/tests/unit/tracing/test_remote.py

import asyncio
import json
from collections.abc import Mapping
from typing import Final

import httpx
import pytest
from pydantic import JsonValue

from litellm.tracing.errors import TraceChanged
from litellm.tracing.generated.types import AllQueryScope, TraceScope
from litellm.tracing.remote import LensConnection, RemoteTraceStore, bounded_response


@pytest.mark.parametrize(
    "url", ("", "ftp://lens", "http://", "http://user:secret@lens", "https://lens?q=1", "https://lens/#x")
)
def test_service_url_rejects_unsupported_or_credential_bearing_destinations(url: str) -> None:
    with pytest.raises(ValueError, match="URL"):
        LensConnection.from_env({"LITELLM_LENS_URL": url, "LITELLM_LENS_SERVICE_TOKEN": "x" * 32})


def test_connection_requires_a_strong_secret_and_preserves_the_configured_prefix() -> None:
    with pytest.raises(ValueError, match="secret"):
        LensConnection.from_env({"LITELLM_LENS_URL": "https://lens", "LITELLM_LENS_SERVICE_TOKEN": "short"})
    connection: Final = LensConnection.from_env(
        {"LITELLM_LENS_URL": "https://lens/prefix/", "LITELLM_LENS_SERVICE_TOKEN": "x" * 32}
    )
    assert connection.url == "https://lens/prefix"
    assert "x" * 32 not in repr(connection)


async def _read_case(
    store: RemoteTraceStore, operation: str, scope: TraceScope, query_scope: AllQueryScope
) -> tuple[JsonValue, Mapping[str, object]]:
    match operation:
        case "list":
            return (
                await store.list_traces(scope, 10, 20, "next", 17),
                {
                    "operation": operation,
                    "scope": scope,
                    "start_ms": 10,
                    "end_ms": 20,
                    "cursor": "next",
                    "limit": 17,
                },
            )
        case "trace":
            return (
                await store.get_trace("trace", scope, "ref", "next", 17),
                {
                    "operation": operation,
                    "scope": scope,
                    "trace_id": "trace",
                    "trace_ref": "ref",
                    "cursor": "next",
                    "page_size": 17,
                },
            )
        case "span":
            return (
                await store.get_span("trace", "span", scope, "ref"),
                {
                    "operation": operation,
                    "scope": scope,
                    "trace_id": "trace",
                    "trace_ref": "ref",
                    "span_id": "span",
                },
            )
        case "span_error":
            return (
                await store.get_span_error("trace", "span", scope, "ref", "next"),
                {
                    "operation": operation,
                    "scope": scope,
                    "trace_id": "trace",
                    "trace_ref": "ref",
                    "span_id": "span",
                    "cursor": "next",
                },
            )
        case "sql":
            return (
                json.loads(await store.query_sql("SELECT 1", query_scope, "unused-local-secret")),
                {"operation": operation, "scope": query_scope, "sql": "SELECT 1"},
            )
        case "help":
            return (
                await store.query_help(query_scope, "unused-local-secret"),
                {"operation": operation, "scope": query_scope},
            )
        case _:
            return (
                json.loads(await store.query("trace_agents", {"start_ms": 1})),
                {"operation": operation, "name": "trace_agents", "parameters": {"start_ms": 1}},
            )


@pytest.mark.asyncio
@pytest.mark.parametrize("operation", ("list", "trace", "span", "span_error", "sql", "help", "query"))
async def test_remote_reads_preserve_scope_and_pagination(operation: str) -> None:
    scope: Final = TraceScope(all_teams=0, user_id="owner", team_ids=("team",))
    query_scope: Final = AllQueryScope(kind="all")
    requests: Final = asyncio.Queue[httpx.Request]()

    def accept(request: httpx.Request) -> httpx.Response:
        requests.put_nowait(request)
        return httpx.Response(200, json={"data": [{"value": "safe"}]})

    async with httpx.AsyncClient(base_url="http://lens/prefix/", transport=httpx.MockTransport(accept)) as client:
        store: Final = RemoteTraceStore(client)
        result, expected = await _read_case(store, operation, scope, query_scope)
        request: Final = requests.get_nowait()
        assert request.url.path == "/prefix/internal/read"
        assert json.loads(request.content) == json.loads(json.dumps(expected))
        assert result == {"data": [{"value": "safe"}]}
        assert b"unused-local-secret" not in request.content


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "status,error",
    ((400, ValueError), (409, TraceChanged), (413, OverflowError), (503, RuntimeError), (302, RuntimeError)),
)
async def test_remote_failures_preserve_public_error_categories_without_leaking_storage_details(
    status: int, error: type[Exception]
) -> None:
    async with httpx.AsyncClient(
        base_url="http://lens",
        transport=httpx.MockTransport(lambda request: httpx.Response(status, text="private storage credentials")),
    ) as client:
        with pytest.raises(error) as failure:
            await RemoteTraceStore(client).get_trace("trace", TraceScope(all_teams=1, user_id="", team_ids=()), "ref")
        assert "private storage credentials" not in str(failure.value)


@pytest.mark.asyncio
async def test_network_failure_is_retryable_without_exposing_the_remote_url() -> None:
    def fail(request: httpx.Request) -> httpx.Response:
        raise httpx.ConnectError("private storage credentials", request=request)

    async with httpx.AsyncClient(base_url="http://lens", transport=httpx.MockTransport(fail)) as client:
        with pytest.raises(RuntimeError, match="Lens trace storage is unavailable"):
            await RemoteTraceStore(client).query_help(AllQueryScope(kind="all"), "secret")


@pytest.mark.asyncio
async def test_reads_reject_oversized_responses_and_invalid_json() -> None:
    with pytest.raises(RuntimeError, match="size limit"):
        await bounded_response(httpx.Response(200, content=b"abcd"), 3)
    assert await bounded_response(httpx.Response(200, content=b"abcd"), 4) == b"abcd"
    async with httpx.AsyncClient(
        base_url="http://lens", transport=httpx.MockTransport(lambda request: httpx.Response(200, content=b"{"))
    ) as client:
        with pytest.raises(ValueError, match="Invalid Lens response"):
            await RemoteTraceStore(client).query_help(AllQueryScope(kind="all"), "secret")


@pytest.mark.asyncio
async def test_gateway_cannot_write_arbitrary_tables() -> None:
    def fail(request: httpx.Request) -> httpx.Response:
        raise AssertionError("No network access is allowed for refused writes")

    async with httpx.AsyncClient(base_url="http://lens", transport=httpx.MockTransport(fail)) as client:
        store: Final = RemoteTraceStore(client)
        with pytest.raises(ValueError, match="request records and feedback"):
            await store.insert_rows("otel_traces", ())


@pytest.mark.asyncio
async def test_request_records_use_the_internal_service_endpoint() -> None:
    requests: Final = asyncio.Queue[httpx.Request]()

    def accept(request: httpx.Request) -> httpx.Response:
        requests.put_nowait(request)
        return httpx.Response(204)

    async with httpx.AsyncClient(base_url="http://lens", transport=httpx.MockTransport(accept)) as client:
        await RemoteTraceStore(client).insert_rows("spend_logs", ({"request_id": "r"},))
    request: Final = requests.get_nowait()
    assert request.url.path == "/internal/spend"
    assert json.loads(request.content) == [{"request_id": "r"}]


@pytest.mark.asyncio
async def test_feedback_rows_use_the_internal_feedback_endpoint() -> None:
    requests: Final = asyncio.Queue[httpx.Request]()

    def accept(request: httpx.Request) -> httpx.Response:
        requests.put_nowait(request)
        return httpx.Response(204)

    async with httpx.AsyncClient(base_url="http://lens", transport=httpx.MockTransport(accept)) as client:
        await RemoteTraceStore(client).insert_rows("lens_feedback", ({"TraceId": "t", "Score": 2},))
    request: Final = requests.get_nowait()
    assert request.url.path == "/internal/feedback"
    assert json.loads(request.content) == [{"TraceId": "t", "Score": 2}]