litellm/tests/unit/proxy/telemetry/test_runtime.py

import asyncio
from typing import Final

import httpx
import pytest
from pydantic import HttpUrl

from litellm.proxy.telemetry.attempt_logger import TelemetryAttemptLogger
from litellm.proxy.telemetry.runtime import TelemetryRuntime, deployment_hasher
from litellm.proxy.telemetry.settings import TelemetryDestination, TelemetrySettings
from litellm.telemetry.consent import ConsentGatedSink, TelemetryConsent
from litellm.telemetry.records import AttemptRecord, InstanceInfo, RequestRecord, TelemetryGroup, UIEvent
from tests.unit.proxy.telemetry.fake_database import SettingsDatabase

_ENDPOINT: Final = HttpUrl("http://telemetry.invalid/v1/reports")


def _offline() -> httpx.AsyncClient:
    return httpx.AsyncClient(transport=httpx.MockTransport(lambda _request: httpx.Response(204)))


_HEARTBEAT: Final = TelemetryConsent(frozenset({TelemetryGroup.HEARTBEAT}))
_SUCCESS: Final = TelemetryConsent(frozenset({TelemetryGroup.HEARTBEAT, TelemetryGroup.REQUEST_SUCCESS}))


async def _started(settings: TelemetrySettings, db: SettingsDatabase | None) -> TelemetryRuntime:
    runtime: Final = TelemetryRuntime()
    await runtime.start(
        litellm_version="1.0.0", settings=settings, db=lambda: db, register=lambda _logger: None, http_client=_offline
    )
    return runtime


@pytest.mark.parametrize(
    "settings",
    [
        TelemetrySettings(),
        TelemetrySettings(groups="", endpoint=_ENDPOINT),
        TelemetrySettings(groups="heartbeat"),
        TelemetrySettings(groups="heartbeat", endpoint=_ENDPOINT, disabled=True),
    ],
)
@pytest.mark.asyncio
async def test_without_a_database_telemetry_needs_valid_pinned_groups_an_endpoint_and_no_veto(
    settings: TelemetrySettings,
) -> None:
    registered: Final[list[TelemetryAttemptLogger]] = []  # mutable-ok: captures the register callback
    runtime: Final = TelemetryRuntime()
    await runtime.start(
        litellm_version="1.0.0", settings=settings, db=lambda: None, register=registered.append, http_client=_offline
    )
    assert runtime.sink is None
    assert registered == []


@pytest.mark.asyncio
async def test_a_veto_never_opens_the_database() -> None:
    opened: Final[list[bool]] = []  # mutable-ok: records whether the db factory ran

    def _db() -> SettingsDatabase:
        opened.append(True)
        return SettingsDatabase(stored_groups='{"groups": ["heartbeat"]}')

    runtime: Final = TelemetryRuntime()
    await runtime.start(
        litellm_version="1.0.0", settings=TelemetrySettings(disabled=True), db=_db, register=lambda _logger: None
    )
    assert (runtime.sink, runtime.store, opened) == (None, None, [])


@pytest.mark.asyncio
async def test_pinned_groups_register_the_attempt_logger_and_stop_cleanly() -> None:
    registered: Final[list[TelemetryAttemptLogger]] = []  # mutable-ok: captures the register callback
    runtime: Final = TelemetryRuntime()
    await runtime.start(
        litellm_version="1.0.0",
        settings=TelemetrySettings(groups="HEARTBEAT", endpoint=_ENDPOINT, destination=TelemetryDestination.HTTP),
        db=lambda: None,
        register=registered.append,
        http_client=_offline,
    )
    assert runtime.sink is not None and runtime.sink.consent == _HEARTBEAT
    assert len(registered) == 1
    await runtime.stop()
    assert runtime.sink is None


@pytest.mark.asyncio
async def test_pinned_groups_win_over_the_stored_groups() -> None:
    db: Final = SettingsDatabase(
        stored_groups='{"groups": ["heartbeat", "request_success"], "destination": "local_table"}'
    )
    runtime: Final = await _started(TelemetrySettings(groups="heartbeat"), db)
    assert runtime.sink is not None and runtime.sink.consent == _HEARTBEAT
    await runtime.stop()


@pytest.mark.asyncio
async def test_stored_groups_are_picked_up_at_the_next_window_without_a_restart() -> None:
    db: Final = SettingsDatabase()
    runtime: Final = await _started(TelemetrySettings(), db)
    assert runtime.sink is None
    db.stored_groups = '{"groups": ["heartbeat", "request_success"], "destination": "local_table"}'
    await runtime.refresh()
    assert runtime.sink is not None and runtime.sink.consent == _SUCCESS
    db.stored_groups = '{"groups": [], "destination": "local_table"}'
    await runtime.refresh()
    assert runtime.sink is None
    await runtime.stop()


@pytest.mark.asyncio
async def test_a_proxy_that_never_opts_in_registers_no_logger_and_writes_nothing() -> None:
    db: Final = SettingsDatabase()
    registered: Final[list[TelemetryAttemptLogger]] = []  # mutable-ok: captures the register callback
    runtime: Final = TelemetryRuntime()
    await runtime.start(
        litellm_version="1.0.0",
        settings=TelemetrySettings(),
        db=lambda: db,
        register=registered.append,
        http_client=_offline,
    )
    await runtime.refresh()
    assert (runtime.sink, registered, db.inserted) == (None, [], ())

    for stored in ('{"groups": ["heartbeat"], "destination": "local_table"}', '{"groups": []}'):
        db.stored_groups = stored
        await runtime.refresh()
    db.stored_groups = '{"groups": ["heartbeat"], "destination": "local_table"}'
    await runtime.refresh()
    assert runtime.sink is not None
    assert len(registered) == 1
    assert db.inserted.count("telemetry_deployment_hash_secret") == 1
    await runtime.stop()


@pytest.mark.asyncio
async def test_invalid_stored_groups_are_ignored_and_telemetry_stays_off() -> None:
    db: Final = SettingsDatabase(stored_groups='{"groups": ["token_info"]}')
    runtime: Final = await _started(TelemetrySettings(), db)
    assert runtime.sink is None
    await runtime.stop()


@pytest.mark.asyncio
async def test_a_failing_settings_read_keeps_the_current_groups() -> None:
    db: Final = SettingsDatabase(stored_groups='{"groups": ["heartbeat"], "destination": "local_table"}')
    runtime: Final = await _started(TelemetrySettings(), db)
    db.fail_reads = True
    await runtime.refresh()
    assert runtime.sink is not None and runtime.sink.consent == _HEARTBEAT
    await runtime.stop()


@pytest.mark.asyncio
async def test_groups_without_a_chosen_destination_record_nothing_until_one_is_saved() -> None:
    db: Final = SettingsDatabase(stored_groups='{"groups": ["heartbeat"]}')
    runtime: Final = await _started(TelemetrySettings(groups="heartbeat"), db)
    assert (runtime.sink, runtime.destination) == (None, None)
    db.stored_groups = '{"groups": [], "destination": "local_table"}'
    await runtime.refresh()
    assert runtime.sink is not None and runtime.destination is TelemetryDestination.LOCAL_TABLE
    await runtime.stop()


@pytest.mark.asyncio
async def test_the_env_destination_wins_over_the_stored_one() -> None:
    db: Final = SettingsDatabase(stored_groups='{"groups": ["heartbeat"], "destination": "local_table"}')
    runtime: Final = await _started(TelemetrySettings(endpoint=_ENDPOINT, destination=TelemetryDestination.HTTP), db)
    assert runtime.sink is not None and runtime.destination is TelemetryDestination.HTTP
    await runtime.stop()


@pytest.mark.asyncio
async def test_a_stored_http_destination_without_an_endpoint_records_nothing() -> None:
    db: Final = SettingsDatabase(stored_groups='{"groups": ["heartbeat"], "destination": "http"}')
    runtime: Final = await _started(TelemetrySettings(), db)
    assert runtime.sink is None
    await runtime.stop()


def test_deployment_hashes_are_stable_per_secret_and_differ_across_secrets() -> None:
    assert deployment_hasher(b"secret-a")("model-1") == deployment_hasher(b"secret-a")("model-1")
    assert deployment_hasher(b"secret-a")("model-1") != deployment_hasher(b"secret-b")("model-1")
    assert "model-1" not in deployment_hasher(b"secret-a")("model-1")


@pytest.mark.asyncio
async def test_stop_waits_for_in_flight_request_finalizers_before_the_last_flush() -> None:
    runtime: Final = await _started(
        TelemetrySettings(groups="heartbeat", endpoint=_ENDPOINT, destination=TelemetryDestination.HTTP), None
    )
    finished: Final[list[bool]] = []  # mutable-ok: records that the finalizer ran to completion

    async def _finalizer() -> None:
        await asyncio.sleep(0.01)
        finished.append(True)

    runtime.spawn(_finalizer())
    await runtime.stop()
    assert finished == [True]


@pytest.mark.asyncio
async def test_stopping_a_runtime_that_never_started_is_a_no_op() -> None:
    runtime: Final = TelemetryRuntime()
    await runtime.stop()
    assert runtime.sink is None


class _HeldFlushSink:
    def __init__(self) -> None:
        self.started: Final = asyncio.Event()
        self.release: Final = asyncio.Event()
        self.flushes: Final[list[bool]] = []  # mutable-ok: records each flush that ran to completion

    def set_instance(self, info: InstanceInfo) -> None: ...

    def record_request(self, record: RequestRecord) -> None: ...

    def record_attempt(self, record: AttemptRecord) -> None: ...

    def record_ui_event(self, event: UIEvent) -> None: ...

    async def flush(self) -> None:
        self.started.set()
        await self.release.wait()
        self.flushes.append(True)


@pytest.mark.asyncio
async def test_stop_during_an_export_lets_that_export_finish() -> None:
    runtime: Final = TelemetryRuntime()
    await runtime.start(
        litellm_version="1.0.0",
        settings=TelemetrySettings(groups="heartbeat", endpoint=_ENDPOINT, destination=TelemetryDestination.HTTP),
        flush_interval_s=0.001,
        db=lambda: None,
        register=lambda _logger: None,
        http_client=_offline,
    )
    held: Final = _HeldFlushSink()
    runtime.sink = ConsentGatedSink(held, _HEARTBEAT)
    await held.started.wait()
    stopping: Final = asyncio.create_task(runtime.stop())
    await asyncio.sleep(0)
    held.release.set()
    await stopping
    assert len(held.flushes) == 2


async def test_settle_timeout_reads_the_configured_value_after_start_and_the_default_before() -> None:
    runtime: Final = TelemetryRuntime()
    assert runtime.settings.settle_timeout_seconds == 2.0
    await runtime.start(
        litellm_version="1.0.0",
        settings=TelemetrySettings(groups="heartbeat", endpoint=_ENDPOINT, settle_timeout_seconds=7.5),
        db=lambda: None,
        register=lambda _logger: None,
        http_client=_offline,
    )
    assert runtime.settings.settle_timeout_seconds == 7.5
    await runtime.stop()