import base64 import hashlib import json import re from datetime import datetime from itertools import chain from pathlib import Path from typing import Final from unittest.mock import AsyncMock import httpx import pytest from prisma import Json, Prisma from pydantic import InstanceOf, TypeAdapter from litellm.tracing.generated.responses import TraceSQLResponse from litellm.tracing.types import SpendLogRecord from scripts.seed_tracing_fixtures import ( JSON, TRACE_FIXTURES, fixture_capture, fixture_replays, managed_response, postgres_row, rebase, rebase_spend, response_ids, response_pattern, seed_arguments, seed_copy, seed_id, spend_fixtures, timestamps, ) from tests._master_key import MASTER_KEY CALL_KEYS: Final = TypeAdapter(tuple[str, ...]) DATETIMES: Final = TypeAdapter(tuple[datetime, datetime]) SPAN_IDENTITY: Final = TypeAdapter(tuple[str, str, str, int]) JSON_FIELDS: Final[TypeAdapter[tuple[Json, Json, Json]]] = TypeAdapter( tuple[InstanceOf[Json], InstanceOf[Json], InstanceOf[Json]] ) @pytest.mark.parametrize("path", sorted(TRACE_FIXTURES.glob("*.json")), ids=lambda path: path.stem) def test_fixture_replays_preserve_raw_payloads_except_time_and_identity(path: Path) -> None: export: Final = JSON.validate_json(path.read_bytes()) now_ms: Final = max(timestamps(export)) // 1_000_000 + 86_400_000 replays: Final = fixture_replays(TRACE_FIXTURES, now_ms, "all-fixtures", re.compile(r"(?!)")) replay: Final = next(item for item in replays if item.name == path.stem) group: Final = tuple(item for item in replays if item.namespace == replay.namespace) assert max(max(timestamps(item.export)) for item in group) // 1_000_000 == now_ms - 1000 assert len(frozenset(item.offset_ms for item in group)) == 1 assert tuple(timestamps(replay.export)) == tuple( timestamp + replay.offset_ms * 1_000_000 for timestamp in timestamps(export) ) assert replay.export != export def test_postgres_rows_preserve_clickhouse_cost_identity_and_payloads() -> None: spends: Final = dict(spend_fixtures())["deepagents_swarm"] for spend, postgres in ((spend, postgres_row(spend)) for spend in spends): start_time, end_time = DATETIMES.validate_python((postgres["startTime"], postgres["endTime"])) messages, response, proxy_request = JSON_FIELDS.validate_python( (postgres["messages"], postgres["response"], postgres["proxy_server_request"]) ) assert postgres["request_id"] == spend["response_id"] assert (postgres["api_key"], postgres["team_id"], postgres["user"], postgres["session_id"]) == ( spend["api_key"], spend["team_id"], spend["user"], spend["session_id"], ) assert postgres["spend"] == spend["spend"] assert postgres["total_tokens"] == spend["prompt_tokens"] + spend["completion_tokens"] assert round(start_time.timestamp() * 1000) == spend["start_time"] assert round(end_time.timestamp() * 1000) == spend["end_time"] assert postgres["request_duration_ms"] == spend["end_time"] - spend["start_time"] assert JSON.validate_python(getattr(messages, "data")) == JSON.validate_json(spend["messages"]) assert JSON.validate_python(getattr(response, "data")) == JSON.validate_json(spend["response"]) assert JSON.validate_python(getattr(proxy_request, "data")) is None @pytest.mark.parametrize("name,spends", spend_fixtures()) def test_captured_spend_replay_preserves_real_cost_and_call_identity( name: str, spends: tuple[SpendLogRecord, ...] ) -> None: export: Final = JSON.validate_json((TRACE_FIXTURES / f"{name}.json").read_bytes()) pattern: Final = response_pattern(spends) offset_ms: Final = 1123 namespace: Final = f"captured-{name}" shifted: Final = rebase(export, offset_ms * 1_000_000, namespace, pattern) replayed: Final = rebase_spend(spends, offset_ms, namespace, pattern) capture: Final = fixture_capture(name, replayed[0]) assert capture.trace_id == seed_id(fixture_capture(name, spends[0]).trace_id, namespace, 32) for before, after in zip(spends, replayed, strict=True): assert after["spend"] == before["spend"] assert (after["prompt_tokens"], after["completion_tokens"], after["total_tokens"]) == ( before["prompt_tokens"], before["completion_tokens"], before["total_tokens"], ) assert after["request_id"] != before["request_id"] assert after["start_time"] == before["start_time"] + offset_ms assert after["end_time"] == before["end_time"] + offset_ms if before["litellm_call_id"]: assert after["litellm_call_id"] != before["litellm_call_id"] @pytest.mark.parametrize("call_id", (None, "gateway")) def test_spend_fixture_loading_preserves_gateway_ids_and_defaults_legacy_rows( tmp_path: Path, call_id: str | None ) -> None: original: Final = dict(spend_fixtures())["deepagents_swarm"][0] fields: Final = {key: value for key, value in original.items() if key != "litellm_call_id"} supplied: Final = fields if call_id is None else {**fields, "litellm_call_id": call_id} (tmp_path / "example_spend_logs.jsonl").write_text(json.dumps(supplied) + "\n") loaded: Final = spend_fixtures(tmp_path) assert loaded == (("example", ({**original, "litellm_call_id": call_id or ""},)),) def test_seed_cli_rejects_nonpositive_copies() -> None: with pytest.raises(SystemExit) as error: seed_arguments(["--copies", "0"]) assert error.value.code == 2 assert seed_arguments(["--profile", "large", "--copies", "5"]).copies == 5 @pytest.mark.parametrize("timeout", ("0", "-1", "inf", "nan")) def test_seed_cli_rejects_invalid_http_timeouts(timeout: str) -> None: with pytest.raises(SystemExit) as error: seed_arguments(["--timeout-seconds", timeout]) assert error.value.code == 2 def test_replay_rebases_provider_request_evidence_with_the_spend_row() -> None: spend: Final = {**dict(spend_fixtures())["deepagents_swarm"][0], "provider_request_id": "req_replay"} pattern: Final = response_pattern((spend,)) replayed: Final = rebase_spend((spend,), 0, "another-run", pattern)[0] export: Final = rebase({"request_id": "req_replay"}, 0, "another-run", pattern) assert export == {"request_id": replayed["provider_request_id"]} assert replayed["provider_request_id"] != spend["provider_request_id"] assert replayed["provider_request_id"].startswith("req_") @pytest.mark.parametrize("upstream", ("msg_response", "req_response", "chatcmpl-response")) def test_replay_preserves_identity_between_managed_and_plain_response_ids(upstream: str) -> None: payload: Final = f"model:example;response_id:{upstream}" managed: Final = "resp_" + base64.b64encode(payload.encode()).decode() spend: Final = {**dict(spend_fixtures())["deepagents_swarm"][0], "response_id": managed} pattern: Final = response_pattern((spend,)) replayed: Final = rebase_spend((spend,), 0, "managed-replay", pattern)[0] plain: Final = rebase(upstream, 0, "managed-replay", pattern) assert plain != upstream assert managed_response(replayed["response_id"]) == f"model:example;response_id:{plain}" @pytest.mark.asyncio async def test_seed_copy_uses_lens_authenticated_tenant_for_both_spend_stores(monkeypatch: pytest.MonkeyPatch) -> None: import asyncio import httpx from litellm.tracing.remote import RemoteTraceStore monkeypatch.setenv("LITELLM_MASTER_KEY", "fixture-admin") fixture: Final = next(item for item in spend_fixtures() if item[0] == "deepagents_swarm") pattern: Final = response_pattern(fixture[1]) replay: Final = next( item for item in fixture_replays(TRACE_FIXTURES, 1_800_000_000_000, "remote-seed", pattern) if item.name == fixture[0] ) requests: Final = asyncio.Queue[httpx.Request]() def accept(request: httpx.Request) -> httpx.Response: requests.put_nowait(request) if request.url.path == "/internal/read": return httpx.Response( 200, json={ "meta": [], "data": [{"team_id": "lens-team", "api_key": "lens-key-hash", "user": "lens-user"}], "rows": 1, "statistics": {"elapsed": 0, "rows_read": 1, "bytes_read": 1}, }, ) return httpx.Response(204) database: Final = AsyncMock(spec=Prisma, litellm_spendlogs=AsyncMock()) async with httpx.AsyncClient(base_url="http://lens", transport=httpx.MockTransport(accept)) as client: captures: Final = await seed_copy(client, RemoteTraceStore(client), database, (replay,), (fixture,), pattern) upload: Final = requests.get_nowait() assert upload.url.path == "/v1/traces" assert json.loads(upload.content) == replay.export read: Final = requests.get_nowait() assert json.loads(read.content)["scope"] == {"kind": "all"} inserted: Final = requests.get_nowait() assert inserted.url.path == "/internal/spend" rows: Final = json.loads(inserted.content) assert len(rows) == len(fixture[1]) assert all( (row["team_id"], row["api_key"], row["user"]) == ("lens-team", "lens-key-hash", "lens-user") for row in rows ) assert tuple(row["request_id"] for row in captures[0][1]) == tuple(row["request_id"] for row in rows) saved: Final = database.litellm_spendlogs.create_many.call_args.kwargs["data"] assert tuple((row["request_id"], row["spend"]) for row in saved) == tuple( (row["request_id"], row["spend"]) for row in rows ) assert requests.empty() @pytest.mark.asyncio async def test_seed_copy_ids_come_from_lens_normalization(monkeypatch: pytest.MonkeyPatch) -> None: from litellm.tracing.remote import RemoteTraceStore from scripts.seed_tracing_fixtures import seeded_trace_ids monkeypatch.setenv("LITELLM_MASTER_KEY", "fixture-admin") def accept(request: httpx.Request) -> httpx.Response: assert json.loads(request.content) == { "operation": "sql", "scope": {"kind": "all"}, "sql": "SELECT DISTINCT TraceId AS trace_id FROM otel_traces WHERE ApiKeyHash = 'key-hash' ORDER BY trace_id", } return httpx.Response( 200, json={ "meta": [], "data": [{"trace_id": "normalized-session"}, {"trace_id": "original-trace"}], "rows": 2, "statistics": {"elapsed": 0, "rows_read": 2, "bytes_read": 1}, }, ) async with httpx.AsyncClient(base_url="http://lens", transport=httpx.MockTransport(accept)) as client: assert await seeded_trace_ids(RemoteTraceStore(client), "key-hash") == ("normalized-session", "original-trace")