litellm/tests/unit/proxy/lens/test_adapter.py

from collections.abc import Mapping
from typing import Final

import httpx
import jwt
import pytest
from fastapi import HTTPException
from starlette.requests import Request

from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.lens.adapter import (
    Connection,
    Identity,
    delegated_identity,
    forward,
    lens_request,
    request_body,
    service_connection,
    service_status,
)
from litellm.tracing.remote import LensConnection

SECRET: Final = "gateway-delegated-identity-fixture-32"


@pytest.fixture
def connection() -> Connection:
    return Connection(LensConnection("http://lens.test", SECRET), SECRET)


@pytest.fixture
def identity() -> Identity:
    return Identity(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="user",
        team_id="team",
        org_id="org",
        token="token-hash",
        models=("allowed-model",),
        log_team_ids=("permitted-team",),
    )


@pytest.mark.parametrize(
    ("overrides", "available"),
    (
        ({}, True),
        ({"LENS_GATEWAY_SECRET": "short"}, False),
        ({"LENS_GATEWAY_SECRET": "x" * 513}, False),
        ({"LITELLM_LENS_URL": "file:///tmp/lens"}, False),
        ({"LITELLM_LENS_SERVICE_TOKEN": ""}, False),
    ),
    ids=("configured", "short_signing_secret", "long_signing_secret", "invalid_url", "missing_service_token"),
)
def test_connection_requires_valid_endpoint_and_both_secrets(overrides: Mapping[str, str], available: bool) -> None:
    configured: Final = Connection.from_env(
        {
            "LITELLM_LENS_URL": "http://lens.test",
            "LITELLM_LENS_SERVICE_TOKEN": SECRET,
            "LENS_GATEWAY_SECRET": SECRET,
            **overrides,
        }
    )
    assert (configured is not None) == available


def test_delegated_token_contains_only_authenticated_identity(connection: Connection, identity: Identity) -> None:
    token: Final = connection.identity_token(identity, 1000)
    assert token is not None
    decoded: Final = jwt.decode(
        token, SECRET, algorithms=["HS256"], audience="litellm-lens", issuer="litellm", options={"verify_exp": False}
    )
    assert decoded == {
        "iss": "litellm",
        "aud": "litellm-lens",
        "sub": "user",
        "iat": 1000,
        "exp": 1030,
        "identity": identity.model_dump(mode="json"),
    }


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "role",
    (
        LitellmUserRoles.PROXY_ADMIN,
        LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
        LitellmUserRoles.INTERNAL_USER,
        LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
        None,
    ),
    ids=("admin", "admin_viewer", "user", "user_viewer", "default_user"),
)
async def test_delegated_identity_preserves_authenticated_lens_scope_and_log_permissions(
    role: LitellmUserRoles | None,
) -> None:
    async def log_team_lookup(auth: UserAPIKeyAuth) -> tuple[str, ...]:
        assert auth.user_id == "user"
        assert role not in (LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY)
        return ("permitted-team",)

    identity: Final = await delegated_identity(
        UserAPIKeyAuth(
            user_role=role,
            user_id="user",
            team_id="credential-team",
            org_id="org",
            token="token-hash",
            models=["model"],
        ),
        log_team_lookup,
    )
    assert identity == Identity(
        user_role=role or LitellmUserRoles.INTERNAL_USER,
        user_id="user",
        team_id="credential-team",
        org_id="org",
        token="token-hash",
        models=("model",),
        log_team_ids=()
        if role in (LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY)
        else ("permitted-team",),
    )


@pytest.mark.asyncio
async def test_delegation_permission_failure_falls_back_to_own_user_scope() -> None:
    async def unavailable(auth: UserAPIKeyAuth) -> tuple[str, ...]:
        raise RuntimeError("Permission database unavailable")

    identity: Final = await delegated_identity(
        UserAPIKeyAuth(user_id="user", user_role=LitellmUserRoles.INTERNAL_USER), unavailable
    )
    assert identity.user_id == "user"
    assert identity.log_team_ids == ()


@pytest.mark.asyncio
async def test_forwarding_replaces_caller_headers_and_preserves_wire_response(
    connection: Connection,
    identity: Identity,
) -> None:
    async def receive() -> Mapping[str, object]:
        return {"type": "http.request", "body": b'{"case":"value"}', "more_body": False}

    def transport(request: httpx.Request) -> httpx.Response:
        assert request.url == "http://lens.test/lens/evals/runs?offset=2&tag=a&tag=b"
        assert request.content == b'{"case":"value"}'
        assert request.headers["x-lens-contract"] == "1"
        assert request.headers["idempotency-key"] == "dedupe"
        assert "cookie" not in request.headers
        assert "x-lens-internal" not in request.headers
        token: Final = request.headers["authorization"].removeprefix("Bearer ")
        assert jwt.decode(token, SECRET, algorithms=["HS256"], audience="litellm-lens", options={"verify_exp": False})[
            "identity"
        ] == identity.model_dump(mode="json")
        return httpx.Response(
            409,
            content=b'{"detail":"contract mismatch"}',
            headers={"content-type": "application/json", "retry-after": "2", "set-cookie": "unsafe"},
        )

    request: Final = Request(
        {
            "type": "http",
            "method": "POST",
            "path": "/lens/evals/runs",
            "query_string": b"offset=2&tag=a&tag=b",
            "headers": [
                (b"authorization", b"Bearer raw-key"),
                (b"x-lens-contract", b"1"),
                (b"idempotency-key", b"dedupe"),
                (b"cookie", b"private"),
                (b"x-lens-internal", b"forged"),
            ],
            "scheme": "http",
            "server": ("gateway.test", 80),
        },
        receive,
    )
    async with httpx.AsyncClient(transport=httpx.MockTransport(transport)) as client:
        response: Final = await forward(request, identity, "evals/runs", connection, client, 1000)
    assert response.status_code == 409
    assert response.body == b'{"detail":"contract mismatch"}'
    assert response.headers["retry-after"] == "2"
    assert "set-cookie" not in response.headers


@pytest.mark.asyncio
@pytest.mark.parametrize("path", ("../internal/read", "foo/../bar", "foo\\bar"))
async def test_forwarding_cannot_escape_lens_namespace(connection: Connection, identity: Identity, path: str) -> None:
    request: Final = Request({"type": "http", "method": "GET", "headers": []})
    async with httpx.AsyncClient() as client:
        with pytest.raises(HTTPException) as error:
            await forward(request, identity, path, connection, client, 1000)
    assert error.value.status_code == 400


@pytest.mark.asyncio
@pytest.mark.parametrize("status", (401, 503))
async def test_unavailable_status_stays_disconnected(connection: Connection, status: int) -> None:
    async with httpx.AsyncClient(transport=httpx.MockTransport(lambda _: httpx.Response(status))) as client:
        result: Final = await service_status(connection, client)
    assert result.protocol_version == 0
    assert not result.storage_ready


@pytest.mark.asyncio
async def test_status_preserves_public_contract_separately_from_worker_protocol(connection: Connection) -> None:
    payload: Final = {
        "storage_ready": True,
        "credentials_ready": True,
        "release": "test-release",
        "protocol_version": 7,
        "public_contract": 1,
    }
    async with httpx.AsyncClient(transport=httpx.MockTransport(lambda _: httpx.Response(200, json=payload))) as client:
        result: Final = await service_status(connection, client)
    assert result.model_dump() == payload


@pytest.mark.asyncio
async def test_trailing_slash_redirect_retains_gateway_prefix_and_query(
    connection: Connection, identity: Identity
) -> None:
    request: Final = Request(
        {
            "type": "http",
            "method": "GET",
            "path": "/gateway/lens/datasets/",
            "query_string": b"offset=2",
            "headers": [],
            "scheme": "https",
            "server": ("gateway.test", 443),
        }
    )
    async with httpx.AsyncClient() as client:
        response: Final = await forward(request, identity, "datasets/", connection, client, 1000)
    assert response.status_code == 307
    assert response.headers["location"] == "https://gateway.test/gateway/lens/datasets?offset=2"


@pytest.mark.asyncio
async def test_duplicate_contract_headers_cannot_be_collapsed_to_a_supported_version(
    connection: Connection, identity: Identity
) -> None:
    request: Final = Request(
        {"type": "http", "method": "GET", "headers": [(b"x-lens-contract", b"2"), (b"x-lens-contract", b"1")]}
    )
    async with httpx.AsyncClient() as client:
        response: Final = await forward(request, identity, "datasets", connection, client, 1000)
    assert response.status_code == 409
    assert response.body == b'{"detail":"contract_version","code":"contract_version"}'


@pytest.mark.asyncio
@pytest.mark.parametrize("limit", (4, 5), ids=("overflow", "exact_limit"))
async def test_streamed_body_bound_is_enforced_without_content_length(limit: int) -> None:
    chunks: Final = iter((b"ab", b"cde"))

    async def receive() -> Mapping[str, object]:
        chunk: Final = next(chunks)
        return {"type": "http.request", "body": chunk, "more_body": chunk == b"ab"}

    request: Final = Request({"type": "http", "method": "POST", "headers": []}, receive)
    if limit == 5:
        assert await request_body(request, limit) == b"abcde"
        return
    with pytest.raises(HTTPException) as error:
        await request_body(request, limit)
    assert error.value.status_code == 413


@pytest.mark.asyncio
@pytest.mark.parametrize("configured", (False, True), ids=("unconfigured", "missing_signing_secret"))
async def test_missing_connection_keeps_setup_readable_and_rejects_lens_requests(
    monkeypatch: pytest.MonkeyPatch, identity: Identity, configured: bool
) -> None:
    monkeypatch.delenv("LENS_GATEWAY_SECRET", raising=False)
    monkeypatch.setenv("LITELLM_LENS_SERVICE_TOKEN", SECRET)
    monkeypatch.setenv("LITELLM_LENS_URL", "http://lens.test" if configured else "")
    monkeypatch.setenv("LITELLM_LENS_PUBLIC_URL", "https://lens.example.test/")
    status: Final = await service_connection(identity)
    assert status.configured == configured
    assert not status.connected
    assert not status.status.storage_ready
    assert status.url == "https://lens.example.test"
    request: Final = Request({"type": "http", "method": "GET", "headers": []})
    with pytest.raises(HTTPException) as routed:
        await lens_request(request, identity, "datasets")
    assert routed.value.status_code == 503
    async with httpx.AsyncClient() as client:
        with pytest.raises(HTTPException) as forwarded:
            await forward(request, identity, "datasets", None, client, 1000)
    assert forwarded.value.status_code == 503


@pytest.mark.asyncio
async def test_forwarding_requires_a_user_or_authenticated_key(connection: Connection, identity: Identity) -> None:
    anonymous: Final = identity.model_copy(update={"user_id": None, "token": None})
    request: Final = Request({"type": "http", "method": "GET", "headers": []})
    async with httpx.AsyncClient() as client:
        with pytest.raises(HTTPException) as error:
            await forward(request, anonymous, "datasets", connection, client, 1000)
    assert error.value.status_code == 403


@pytest.mark.asyncio
@pytest.mark.parametrize("failure", ("malformed", "oversized", "transport"))
async def test_failed_status_response_never_reports_a_ready_service(connection: Connection, failure: str) -> None:
    def transport(request: httpx.Request) -> httpx.Response:
        if failure == "transport":
            raise httpx.ConnectError("Fixture connection refused", request=request)
        return httpx.Response(200, content=b"x" * (16 * 1024 + 1) if failure == "oversized" else b"invalid-json")

    async with httpx.AsyncClient(transport=httpx.MockTransport(transport)) as client:
        result: Final = await service_status(connection, client)
    assert not result.storage_ready
    assert not result.credentials_ready
    assert result.public_contract == 0


@pytest.mark.asyncio
async def test_failed_forward_reports_service_unavailability(connection: Connection, identity: Identity) -> None:
    async def receive() -> Mapping[str, object]:
        return {"type": "http.request", "body": b"", "more_body": False}

    def transport(request: httpx.Request) -> httpx.Response:
        raise httpx.ReadTimeout("Fixture response unavailable", request=request)

    request: Final = Request(
        {"type": "http", "method": "GET", "path": "/lens", "query_string": b"", "headers": []}, receive
    )
    async with httpx.AsyncClient(transport=httpx.MockTransport(transport)) as client:
        with pytest.raises(HTTPException) as error:
            await forward(request, identity, "", connection, client, 1000)
    assert error.value.status_code == 503
    assert error.value.detail == "Lens service is unavailable"