import json from datetime import datetime, timezone from typing import Any from unittest.mock import AsyncMock, MagicMock import httpx import pytest from litellm.proxy.guardrails import usage_tracking from litellm.proxy.guardrails.usage_tracking import ( _MAX_PENDING_ROWS, PendingRollups, _capped, process_spend_logs_guardrail_usage, ) def _prisma() -> MagicMock: client = MagicMock() db = client.db db.litellm_dailyguardrailmetrics.upsert = AsyncMock() db.litellm_dailyguardrailusageunits.upsert = AsyncMock() db.litellm_spendlogguardrailindex.create_many = AsyncMock() return client def _payload( request_id: str, *, team_id: str | None = "team-a", api_key: str = "hashed-key-1", usage: dict[str, Any] | None = None, guardrail_status: str = "success", cost_by_unit: dict[str, Any] | None = None, cost_in_spend: bool | None = None, ) -> dict[str, Any]: entry: dict[str, Any] = { "guardrail_id": "bedrock-guard", "guardrail_status": guardrail_status, } if usage is not None: entry["guardrail_usage"] = usage if cost_by_unit is not None: entry["guardrail_cost_by_unit"] = cost_by_unit if cost_in_spend is not None: entry["guardrail_cost_in_spend"] = cost_in_spend return { "request_id": request_id, "startTime": datetime(2026, 8, 17, 12, 0, tzinfo=timezone.utc), "team_id": team_id, "api_key": api_key, "metadata": json.dumps({"guardrail_information": [entry]}), } def _units_upserts(prisma: MagicMock) -> dict[tuple, int]: calls = prisma.db.litellm_dailyguardrailusageunits.upsert.call_args_list out: dict[tuple, int] = {} for c in calls: where = c.kwargs["where"]["guardrail_id_date_team_id_api_key_usage_unit"] create = c.kwargs["data"]["create"] assert create["units"] == c.kwargs["data"]["update"]["units"]["increment"] assert {k: create[k] for k in where} == where out[tuple(where[k] for k in ("guardrail_id", "date", "team_id", "api_key", "usage_unit"))] = create["units"] return out def _cost_upserts(prisma: MagicMock) -> dict[str, tuple[float, int]]: """usage_unit -> (cost, untracked_units) written on create; the update path must increment by the same.""" calls = prisma.db.litellm_dailyguardrailusageunits.upsert.call_args_list out: dict[str, tuple[float, int]] = {} for c in calls: create = c.kwargs["data"]["create"] update = c.kwargs["data"]["update"] assert update["cost"] == {"increment": create["cost"]} assert update["untracked_units"] == {"increment": create["untracked_units"]} out[create["usage_unit"]] = (create["cost"], create["untracked_units"]) return out @pytest.mark.asyncio async def test_usage_units_rolled_up_by_guardrail_team_key_and_date(): """ LIT-5650: billable units must aggregate per (guardrail, date, team, key, counter): same-key payloads sum into one upsert, a team-less payload gets its own empty-string-team row, and blocked invocations (which Bedrock still bills for) count exactly like passed ones. """ prisma = _prisma() logs = [ _payload("r1", usage={"topicPolicyUnits": 1, "contentPolicyUnits": 1}), _payload( "r2", usage={"topicPolicyUnits": 1, "contentPolicyUnits": 2}, guardrail_status="guardrail_intervened", ), _payload("r3", team_id=None, api_key="hashed-key-2", usage={"topicPolicyUnits": 1}), ] await process_spend_logs_guardrail_usage(prisma, logs) assert _units_upserts(prisma) == { ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 2, ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 3, ("bedrock-guard", "2026-08-17", "", "hashed-key-2", "topicPolicyUnits"): 1, } @pytest.mark.asyncio async def test_flagged_status_counts_as_flagged_not_passed_or_blocked(): """LIT-6894: a custom code flag() verdict lands in flagged_count on the Monitor rollup.""" prisma = _prisma() logs = [ _payload("r1", guardrail_status="success"), _payload("r2", guardrail_status="guardrail_flagged"), _payload("r3", guardrail_status="guardrail_intervened"), ] await process_spend_logs_guardrail_usage(prisma, logs) create = prisma.db.litellm_dailyguardrailmetrics.upsert.call_args.kwargs["data"]["create"] assert (create["requests_evaluated"], create["passed_count"], create["flagged_count"], create["blocked_count"]) == ( 3, 1, 1, 1, ) def _fake_sleep() -> tuple[AsyncMock, list[float]]: delays: list[float] = [] sleep = AsyncMock(side_effect=lambda delay: delays.append(delay)) return sleep, delays @pytest.mark.asyncio async def test_one_failing_upsert_does_not_drop_remaining_writes(): """ A DB error on one daily-metrics or usage-unit upsert must not cancel the remaining upserts in the flushed batch, or the usage endpoints would permanently under-report billable counters. """ prisma = _prisma() prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = httpx.ConnectError("db down") prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [httpx.ConnectError("db down"), None, None] sleep, _ = _fake_sleep() logs = [ _payload("r1", usage={"topicPolicyUnits": 1}), _payload("r2", team_id=None, api_key="hashed-key-2", usage={"topicPolicyUnits": 1}), ] await process_spend_logs_guardrail_usage(prisma, logs, sleep=sleep, pending=PendingRollups()) assert _units_upserts(prisma) == { ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 1, ("bedrock-guard", "2026-08-17", "", "hashed-key-2", "topicPolicyUnits"): 1, } @pytest.mark.asyncio async def test_transient_upsert_failure_is_retried_with_backoff_for_failed_rows_only(): """ A connection error (the write provably never reached the database) must not permanently drop billed units from the aggregates: only the rows that failed are re-sent, after exponential backoff, and the batch ends once every row has landed. """ prisma = _prisma() prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [httpx.ConnectError("blip"), None, None] sleep, delays = _fake_sleep() logs = [ _payload("r1", usage={"topicPolicyUnits": 1}), _payload("r2", team_id=None, api_key="hashed-key-2", usage={"topicPolicyUnits": 1}), ] await process_spend_logs_guardrail_usage(prisma, logs, sleep=sleep) calls = prisma.db.litellm_dailyguardrailusageunits.upsert.call_args_list assert len(calls) == 3 assert calls[2].kwargs["where"] == calls[0].kwargs["where"] assert delays == [1] @pytest.mark.asyncio async def test_persistent_upsert_failure_stops_after_three_retries(): prisma = _prisma() prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = httpx.ConnectError("db down") sleep, delays = _fake_sleep() pending = PendingRollups() await process_spend_logs_guardrail_usage( prisma, [_payload("r1", usage={"topicPolicyUnits": 1})], sleep=sleep, pending=pending ) assert prisma.db.litellm_dailyguardrailmetrics.upsert.call_count == 4 assert delays == [1, 2, 4] assert prisma.db.litellm_dailyguardrailusageunits.upsert.call_count == 1 assert dict(pending.metrics) == { ("bedrock-guard", "2026-08-17"): { "requests_evaluated": 1, "passed_count": 1, "blocked_count": 0, "flagged_count": 0, } } @pytest.mark.asyncio async def test_retry_exhausted_rows_are_requeued_and_land_on_the_next_flush(): """ LIT-5761: rollup rows whose connection-error retries exhaust must not be silently lost. They are requeued and merged into the next flushed batch, so the aggregates catch up once the database is reachable again. """ pending = PendingRollups() down = _prisma() down.db.litellm_dailyguardrailmetrics.upsert.side_effect = httpx.ConnectError("db down") down.db.litellm_dailyguardrailusageunits.upsert.side_effect = httpx.ConnectError("db down") sleep, _ = _fake_sleep() await process_spend_logs_guardrail_usage( down, [_payload("r1", usage={"topicPolicyUnits": 2})], sleep=sleep, pending=pending ) assert dict(pending.units) == { ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): (2, 0.0, 2) } recovered = _prisma() await process_spend_logs_guardrail_usage( recovered, [_payload("r2", usage={"topicPolicyUnits": 3})], sleep=sleep, pending=pending ) assert _units_upserts(recovered) == { ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 5, } metrics_create = recovered.db.litellm_dailyguardrailmetrics.upsert.call_args.kwargs["data"]["create"] assert metrics_create["requests_evaluated"] == 2 assert not pending.units assert not pending.metrics @pytest.mark.asyncio async def test_ambiguous_failures_are_never_requeued(): """ A post-send failure (the increment may have committed) must stay dropped: requeueing it would re-send a possibly applied increment and double-count. """ pending = PendingRollups() prisma = _prisma() prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = httpx.ReadTimeout("maybe committed") sleep, delays = _fake_sleep() await process_spend_logs_guardrail_usage( prisma, [_payload("r1", usage={"topicPolicyUnits": 1})], sleep=sleep, pending=pending ) assert delays == [] assert not pending.units assert not pending.metrics def test_pending_requeue_is_capped_dropping_oldest_rows(): rows = {index: index for index in range(_MAX_PENDING_ROWS + 5)} capped = _capped(rows, "usage unit") assert len(capped) == _MAX_PENDING_ROWS assert 4 not in capped assert _MAX_PENDING_ROWS + 4 in capped def _units_upsert_wheres(prisma: MagicMock) -> list[tuple]: return [ tuple( c.kwargs["where"]["guardrail_id_date_team_id_api_key_usage_unit"][k] for k in ("guardrail_id", "date", "team_id", "api_key", "usage_unit") ) for c in prisma.db.litellm_dailyguardrailusageunits.upsert.call_args_list ] @pytest.mark.asyncio async def test_post_send_failure_is_never_retried_so_increments_cannot_double_count(): """ Follow-up to #37225: the units upsert is a non-idempotent increment, so an ambiguous post-send failure (read timeout after the statement may have committed) must be attempted exactly once. Re-sending it stacks a second increment and inflates billable unit totals. Only a connection error proves the write never reached the database and may be retried; the other rows in the batch still land either way. """ prisma = _prisma() prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [ httpx.ReadTimeout("read timed out"), httpx.ConnectError("refused"), None, ] sleep, delays = _fake_sleep() logs = [ _payload("r1", usage={"topicPolicyUnits": 1}), _payload("r2", team_id=None, api_key="hashed-key-2", usage={"topicPolicyUnits": 1}), ] await process_spend_logs_guardrail_usage(prisma, logs, sleep=sleep) timed_out_row = ("bedrock-guard", "2026-08-17", "", "hashed-key-2", "topicPolicyUnits") refused_row = ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits") assert _units_upsert_wheres(prisma) == [timed_out_row, refused_row, refused_row] assert delays == [1] @pytest.mark.asyncio async def test_generic_upsert_exception_is_terminal_for_that_row_only(): prisma = _prisma() prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = RuntimeError("constraint violation") sleep, delays = _fake_sleep() await process_spend_logs_guardrail_usage(prisma, [_payload("r1", usage={"topicPolicyUnits": 1})], sleep=sleep) assert prisma.db.litellm_dailyguardrailmetrics.upsert.call_count == 1 assert delays == [] assert _units_upserts(prisma) == { ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 1, } @pytest.mark.asyncio async def test_zero_and_non_int_usage_counters_are_skipped(): prisma = _prisma() logs = [ _payload( "r1", usage={ "topicPolicyUnits": 1, "wordPolicyUnits": 0, "contentPolicyImageUnits": 0, "oddball": "not-an-int", "boolish": True, }, ), _payload("r2", usage=None), ] await process_spend_logs_guardrail_usage(prisma, logs) assert _units_upserts(prisma) == { ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 1, } @pytest.mark.asyncio async def test_payload_without_request_id_is_skipped_like_the_metrics_path(): prisma = _prisma() logs = [ {**_payload("ignored", usage={"topicPolicyUnits": 5}), "request_id": None}, _payload("r2", usage={"topicPolicyUnits": 1}), ] await process_spend_logs_guardrail_usage(prisma, logs) assert _units_upserts(prisma) == { ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 1, } assert prisma.db.litellm_dailyguardrailmetrics.upsert.call_args.kwargs["data"]["create"]["requests_evaluated"] == 1 @pytest.mark.asyncio async def test_cost_rolled_up_per_counter_alongside_units(): """LIT-5652: the hook's per-counter cost lands on the same daily row as the units it priced, summed across payloads exactly like the units are, and the update path increments it so a second flush on the same day keeps adding.""" prisma = _prisma() logs = [ _payload( "r1", usage={"contentPolicyUnits": 1000, "wordPolicyUnits": 50}, cost_by_unit={"contentPolicyUnits": 0.15, "wordPolicyUnits": 0.0}, ), _payload( "r2", usage={"contentPolicyUnits": 2000, "wordPolicyUnits": 10}, cost_by_unit={"contentPolicyUnits": 0.3, "wordPolicyUnits": 0.0}, ), ] await process_spend_logs_guardrail_usage(prisma, logs) assert _units_upserts(prisma) == { ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 3000, ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "wordPolicyUnits"): 60, } costs = _cost_upserts(prisma) assert costs["contentPolicyUnits"] == (pytest.approx(0.45), 0) assert costs["wordPolicyUnits"] == (0.0, 0) @pytest.mark.asyncio async def test_counter_the_hook_could_not_price_is_stored_as_untracked_units_not_free(): """A counter the cost map does not list arrives stamped as None. Its units must land in untracked_units with no cost, so the row never reads as free, while the priced counter on the same request keeps its cost.""" prisma = _prisma() logs = [ _payload( "r1", usage={"contentPolicyUnits": 1000, "someFutureCounter": 3}, cost_by_unit={"contentPolicyUnits": 0.15, "someFutureCounter": None}, ) ] await process_spend_logs_guardrail_usage(prisma, logs) assert _units_upserts(prisma) == { ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 1000, ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "someFutureCounter"): 3, } costs = _cost_upserts(prisma) assert costs["contentPolicyUnits"] == (pytest.approx(0.15), 0) assert costs["someFutureCounter"] == (0.0, 3) @pytest.mark.asyncio async def test_mixed_priced_and_unpriced_increments_keep_the_subtotal_and_count_the_rest_untracked(): """Priced and unpriced increments on the same row (a hook without pricing, a pre-upgrade proxy in a mixed fleet) must keep the priced subtotal and count exactly the unpriced units as untracked. Nulling the cost would throw away a known number; keeping it alone would look exact while understating.""" prisma = _prisma() logs = [ _payload("r1", usage={"contentPolicyUnits": 1000}, cost_by_unit={"contentPolicyUnits": 0.15}), _payload("r2", usage={"contentPolicyUnits": 700}), _payload("r3", usage={"contentPolicyUnits": 300}, cost_by_unit={"contentPolicyUnits": None}), ] await process_spend_logs_guardrail_usage(prisma, logs) assert _units_upserts(prisma) == { ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 2000, } assert _cost_upserts(prisma) == {"contentPolicyUnits": (pytest.approx(0.15), 1000)} @pytest.mark.asyncio async def test_report_only_and_forged_costs_are_not_rolled_up_but_units_are(): """guardrail_cost_in_spend=False (Azure Prompt Shield) keeps its cost out of spend, so the rollup must not record it either or the dashboard would show a number the budget never charged. A negative or non-finite per-counter cost is treated the same way rather than subtracting from the day.""" prisma = _prisma() logs = [ _payload("r1", usage={"text_records": 3}, cost_by_unit={"text_records": 0.5}, cost_in_spend=False), _payload("r2", usage={"contentPolicyUnits": 10}, cost_by_unit={"contentPolicyUnits": -0.5}), _payload("r3", usage={"topicPolicyUnits": 10}, cost_by_unit={"topicPolicyUnits": float("inf")}), ] await process_spend_logs_guardrail_usage(prisma, logs) assert _units_upserts(prisma) == { ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "text_records"): 3, ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 10, ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 10, } assert _cost_upserts(prisma) == { "text_records": (0.0, 3), "contentPolicyUnits": (0.0, 10), "topicPolicyUnits": (0.0, 10), } @pytest.mark.asyncio async def test_requeued_cost_is_added_to_the_next_flush(): """Cost and untracked units must survive the connection-error requeue the same way units do, or a DB blip would silently drop dollars (or the record that some units had no price) while keeping the units themselves.""" pending = PendingRollups() down = _prisma() down.db.litellm_dailyguardrailmetrics.upsert.side_effect = httpx.ConnectError("db down") down.db.litellm_dailyguardrailusageunits.upsert.side_effect = httpx.ConnectError("db down") sleep, _ = _fake_sleep() await process_spend_logs_guardrail_usage( down, [ _payload( "r1", usage={"contentPolicyUnits": 1000, "someFutureCounter": 3}, cost_by_unit={"contentPolicyUnits": 0.15, "someFutureCounter": None}, ) ], sleep=sleep, pending=pending, ) recovered = _prisma() await process_spend_logs_guardrail_usage( recovered, [ _payload( "r2", usage={"contentPolicyUnits": 2000, "someFutureCounter": 4}, cost_by_unit={"contentPolicyUnits": 0.3, "someFutureCounter": None}, ) ], sleep=sleep, pending=pending, ) assert _units_upserts(recovered) == { ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 3000, ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "someFutureCounter"): 7, } costs = _cost_upserts(recovered) assert costs["contentPolicyUnits"] == (pytest.approx(0.45), 0) assert costs["someFutureCounter"] == (0.0, 7) def _fan_out_payload(request_id: str, guardrail_ids: tuple[str, ...]) -> dict[str, Any]: return { "request_id": request_id, "startTime": datetime(2026, 8, 17, 12, 0, tzinfo=timezone.utc), "team_id": "team-a", "api_key": "hashed-key-1", "metadata": json.dumps( {"guardrail_information": [{"guardrail_id": gid, "guardrail_status": "success"} for gid in guardrail_ids]} ), } def _index_rows_written(prisma: MagicMock) -> list[tuple[str, str]]: return [ (row["request_id"], row["guardrail_id"]) for call in prisma.db.litellm_spendlogguardrailindex.create_many.call_args_list for row in call.kwargs["data"] ] @pytest.mark.asyncio async def test_index_rows_are_written_in_row_bounded_statements(monkeypatch): """ LIT-5931: the drain caps logs, not logs x guardrails, so a fan-out must be split into statements the query engine can afford instead of one create_many. """ monkeypatch.setattr(usage_tracking, "SPEND_LOG_WRITE_BATCH_MAX_ROWS", 100) prisma = _prisma() guardrail_ids = tuple(f"guard-{i}" for i in range(50)) logs = [_fan_out_payload(f"r{i}", guardrail_ids) for i in range(5)] await process_spend_logs_guardrail_usage(prisma, logs, pending=PendingRollups()) statements = prisma.db.litellm_spendlogguardrailindex.create_many.call_args_list assert [len(call.kwargs["data"]) for call in statements] == [100, 100, 50] assert all(call.kwargs["skip_duplicates"] is True for call in statements) assert _index_rows_written(prisma) == [(f"r{i}", gid) for i in range(5) for gid in guardrail_ids] @pytest.mark.asyncio async def test_one_failing_index_statement_does_not_drop_the_others_or_the_rollup(monkeypatch): monkeypatch.setattr(usage_tracking, "SPEND_LOG_WRITE_BATCH_MAX_ROWS", 100) prisma = _prisma() prisma.db.litellm_spendlogguardrailindex.create_many.side_effect = [None, httpx.ReadTimeout("ambiguous"), None] guardrail_ids = tuple(f"guard-{i}" for i in range(50)) logs = [_fan_out_payload(f"r{i}", guardrail_ids) for i in range(5)] await process_spend_logs_guardrail_usage(prisma, logs, pending=PendingRollups()) assert prisma.db.litellm_spendlogguardrailindex.create_many.await_count == 3 assert prisma.db.litellm_dailyguardrailmetrics.upsert.await_count == len(guardrail_ids)