test(integration): regression tests for August cost tracking and budgeting bugs

tests/integration/_support/wire.py

@@ -2,11 +2,13 @@
 
 import ssl
 import threading
+import time
 from collections.abc import Callable, Generator, Mapping
 from contextlib import contextmanager
 from dataclasses import dataclass
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from queue import SimpleQueue
+from types import MappingProxyType
 from typing import Final
 
 
@@ -26,6 +28,8 @@ class Reply:
     chunks: tuple[bytes, ...] | None = None
     abort_after: int | None = None
     gate_after_first: threading.Event | None = None
+    pause_between_chunks: float = 0
+    headers: Mapping[str, str] = MappingProxyType({})
 
 
 @dataclass(frozen=True, slots=True)
@@ -64,6 +68,8 @@ def respond(self) -> None:
                 reply = Reply(status=500)
             self.send_response(reply.status)
             self.send_header("content-type", reply.content_type)
+            for name, value in reply.headers.items():
+                self.send_header(name, value)
             if reply.chunks is None:
                 self.send_header("content-length", str(len(reply.body)))
             else:
@@ -81,6 +87,8 @@ def respond(self) -> None:
                         self.wfile.flush()
                         if index == 0 and reply.gate_after_first is not None:
                             assert reply.gate_after_first.wait(timeout=5), "Stream barrier was never released"
+                        if reply.pause_between_chunks and index + 1 < len(reply.chunks):
+                            time.sleep(reply.pause_between_chunks)
                     else:
                         self.wfile.write(b"0\r\n\r\n")
                         self.wfile.flush()

tests/integration/contracts.json

@@ -1789,6 +1789,78 @@
     "tests/integration/mcp/test_mcp_lifecycle.py::test_same_url_server_grants_scope_discovery_and_direct_or_virtual_execution[bearer]": [
       "other.mcp.permissions.same_url_servers_enforce_discovery_and_execution"
     ],
+    "tests/integration/management/test_budget_updates.py::test_shortening_budget_duration_moves_reset_at_onto_the_new_schedule": [
+      "mgmt.budget.update.duration_change_recomputes_reset_at"
+    ],
+    "tests/integration/management/test_organization_budget_clear.py::test_patch_organization_update_with_null_tpm_limit_clears_it_and_keeps_sibling_limits": [
+      "mgmt.organization.update.null_clears_budget_limit"
+    ],
+    "tests/integration/management/test_team_budget_duration_defaults.py::test_team_new_explicit_null_budget_duration_is_not_replaced_by_default": [
+      "mgmt.team.new.explicit_null_budget_duration_overrides_default"
+    ],
+    "tests/integration/management/test_team_member_budget_cache.py::test_team_member_default_budget_lands_in_redis_after_first_member_call": [
+      "mgmt.team_member_budget.default_budget_is_cached_in_redis_as_json"
+    ],
+    "tests/integration/observability/test_callback_delivery.py::test_streamed_responses_success_callback_carries_provider_apim_request_id": [
+      "other.observability.callbacks.streamed_responses_events_carry_provider_response_headers"
+    ],
+    "tests/integration/observability/test_guardrail_effects.py::test_bedrock_passthrough_converse_guardrail_ignores_denied_term_in_tool_definition": [
+      "other.observability.guardrails.bedrock_passthrough_converse_scans_only_caller_content"
+    ],
+    "tests/integration/pricing/test_configured_prices.py::test_cost_estimate_reports_configured_prices_for_model_absent_from_cost_map": [
+      "quota_management.cost_estimate.configured_price.reported_for_model_absent_from_cost_map"
+    ],
+    "tests/integration/pricing/test_configured_prices.py::test_saving_echoed_model_info_does_not_freeze_cost_map_price_into_deployment": [
+      "pricing.model_update.echoed_cost_map_price_is_not_persisted_as_override"
+    ],
+    "tests/integration/pricing/test_databricks_cache_pricing.py::test_databricks_cached_prompt_tokens_bill_at_cache_rates_not_input_rate": [
+      "pricing.databricks.cached_prompt_tokens_bill_at_cache_rates"
+    ],
+    "tests/integration/pricing/test_ocr_page_pricing.py::test_ocr_annotation_pages_are_billed_at_annotation_cost_per_page": [
+      "pricing.ocr.annotation_pages_billed_at_annotation_rate"
+    ],
+    "tests/integration/pricing/test_service_tier_pricing.py::test_ultrafast_service_tier_bills_ultrafast_rates_and_keeps_pricing_off_the_wire": [
+      "quota_management.spend_tracking.service_tier_pricing.ultrafast_bills_ultrafast_rates"
+    ],
+    "tests/integration/spend/test_batch_completion_accounting.py::test_completed_batch_spend_row_records_reasoning_tokens_and_error_file_failures": [
+      "quota_management.spend_tracking.batch_costs.reasoning_tokens_and_error_file_failures_recorded"
+    ],
+    "tests/integration/spend/test_batch_observability.py::test_batch_retrieval_row_sums_reasoning_tokens_and_counts_output_and_error_file_failures": [
+      "spend.batches.retrieval_row_aggregates_reasoning_tokens_and_per_request_counts"
+    ],
+    "tests/integration/spend/test_batch_poll_starvation.py::test_batches_gone_at_provider_do_not_starve_a_newer_batch_out_of_cost_polling": [
+      "quota_management.spend_tracking.batch_costs.uncostable_rows_retire_so_newer_batches_are_costed"
+    ],
+    "tests/integration/spend/test_cache_and_quota.py::test_in_flight_count_tokens_does_not_reserve_key_budget_away_from_a_completion": [
+      "quota_management.budget.key.in_flight_count_tokens_reserves_nothing_so_completion_reaches_provider"
+    ],
+    "tests/integration/spend/test_cache_and_quota.py::test_repeated_count_tokens_on_budgeted_key_does_not_reserve_budget_or_block_later_completion": [
+      "quota_management.budget.key.count_tokens_reserves_nothing_so_completion_within_budget_succeeds"
+    ],
+    "tests/integration/spend/test_daily_rollup_retry.py::test_failed_daily_user_rollup_commit_is_retried_so_spend_report_and_daily_activity_agree": [
+      "spend.daily_rollup.failed_user_commit_is_retried_until_report_and_daily_activity_agree"
+    ],
+    "tests/integration/spend/test_disconnected_bedrock_messages_stream_billing.py::test_client_disconnect_mid_bedrock_messages_stream_still_bills_terminal_usage": [
+      "spend.anthropic_messages_stream.client_disconnect_bills_terminal_bedrock_usage"
+    ],
+    "tests/integration/spend/test_failed_dispatch_tokens.py::test_provider_500_after_dispatch_records_estimated_prompt_tokens_on_failure_row": [
+      "spend.failed_dispatch.failure_row_records_estimated_input_tokens"
+    ],
+    "tests/integration/spend/test_model_router_selected_model.py::test_model_router_alias_without_router_in_name_keeps_selected_model_in_response_and_spend_log": [
+      "spend.model_router.selected_model_is_returned_and_persisted_for_plain_alias"
+    ],
+    "tests/integration/spend/test_org_budget_cli_session_token.py::test_cli_session_token_without_org_id_charges_and_caps_the_team_organization": [
+      "quota_management.organization_budget.cli_session_token_without_org_id_charges_team_organization"
+    ],
+    "tests/integration/spend/test_passthrough_budget_reservation.py::test_repeated_gemini_passthrough_calls_stay_served_while_key_spend_is_below_max_budget": [
+      "spend.budget_reservation.gemini_passthrough_success_releases_reservation_from_spend_counter"
+    ],
+    "tests/integration/spend/test_team_daily_activity_aggregated.py::test_aggregated_team_activity_reports_the_whole_range_team_spend_in_one_page": [
+      "quota_management.spend_tracking.team_daily_activity_aggregated_reports_whole_range_team_spend"
+    ],
+    "tests/integration/spend/test_team_member_spend.py::test_member_added_without_any_budget_is_charged_on_its_membership_row": [
+      "spend.team_member.member_without_budget_gets_membership_row_and_spend"
+    ],
     "tests/integration/management/test_vector_store_config_ownership.py::test_config_store_is_listed_beside_db_store_and_survives_listing": [
       "mgmt.vector_store.list.keeps_config_store_beside_db_stores"
     ],

tests/integration/management/test_budget_updates.py

@@ -0,0 +1,29 @@
+from datetime import datetime, timedelta, timezone
+from typing import Final
+
+import pytest
+
+from tests.integration._support.client import Gateway, string_value
+from tests.integration._support.database import read_rows
+
+
+def _persisted_reset_at(budget_id: str) -> datetime:
+    rows: Final = read_rows(
+        'SELECT budget_reset_at::text AS reset_at FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (budget_id,)
+    )
+    assert len(rows) == 1, rows
+    reset_at: Final = datetime.fromisoformat(string_value(rows[0]["reset_at"]))
+    return reset_at if reset_at.tzinfo is not None else reset_at.replace(tzinfo=timezone.utc)
+
+
+@pytest.mark.covers("mgmt.budget.update.duration_change_recomputes_reset_at")
+def test_shortening_budget_duration_moves_reset_at_onto_the_new_schedule(gateway: Gateway) -> None:
+    with gateway.scenario() as scenario:
+        budget_id: Final = scenario.budget(max_budget=10.0, budget_duration="10d")
+        ten_day_reset_at: Final = _persisted_reset_at(budget_id)
+        before: Final = datetime.now(timezone.utc)
+        response: Final = gateway.request("POST", "/budget/update", {"budget_id": budget_id, "budget_duration": "1d"})
+        assert response.status_code == 200, response.text
+        updated: Final = _persisted_reset_at(budget_id)
+        assert updated < ten_day_reset_at, f"{updated} not before {ten_day_reset_at}"
+        assert before < updated <= before + timedelta(days=1, minutes=5), f"{updated} not within 1d of {before}"

tests/integration/management/test_organization_budget_clear.py

@@ -0,0 +1,59 @@
+import uuid
+from typing import Final
+
+import pytest
+
+from tests.integration._support.client import Gateway, object_value, string_value
+from tests.integration._support.database import read_rows
+
+
+def _budget_rows(budget_id: str) -> list[dict[str, object]]:
+    return read_rows(
+        'SELECT tpm_limit, rpm_limit, max_budget FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (budget_id,)
+    )
+
+
+@pytest.mark.covers("mgmt.organization.update.null_clears_budget_limit")
+def test_patch_organization_update_with_null_tpm_limit_clears_it_and_keeps_sibling_limits(gateway: Gateway) -> None:
+    created: Final = gateway.post(
+        "/organization/new",
+        {
+            "organization_alias": f"integration-{uuid.uuid4().hex}",
+            "tpm_limit": 4000,
+            "rpm_limit": 40,
+            "max_budget": 12.5,
+        },
+    )
+    organization_id: Final = string_value(created["organization_id"])
+    budget_id: Final = string_value(created["budget_id"])
+    try:
+        assert _budget_rows(budget_id) == [{"tpm_limit": 4000, "rpm_limit": 40, "max_budget": 12.5}]
+        updated: Final = gateway.request(
+            "PATCH", "/organization/update", {"organization_id": organization_id, "tpm_limit": None}
+        )
+        assert updated.status_code == 200, updated.text
+        updated_budget: Final = object_value(object_value(updated.json())["litellm_budget_table"])
+        assert (updated_budget["tpm_limit"], updated_budget["rpm_limit"], updated_budget["max_budget"]) == (
+            None,
+            40,
+            12.5,
+        ), updated.text
+        assert _budget_rows(budget_id) == [{"tpm_limit": None, "rpm_limit": 40, "max_budget": 12.5}]
+        info: Final = gateway.request("GET", "/organization/info", params={"organization_id": organization_id})
+        assert info.status_code == 200, info.text
+        info_budget: Final = object_value(object_value(info.json())["litellm_budget_table"])
+        assert (info_budget["tpm_limit"], info_budget["rpm_limit"], info_budget["max_budget"]) == (
+            None,
+            40,
+            12.5,
+        ), info.text
+    finally:
+        deleted: Final = gateway.request("DELETE", "/organization/delete", {"organization_ids": [organization_id]})
+        assert deleted.status_code == 200, deleted.text
+        gateway.post("/budget/delete", {"id": budget_id})
+        assert (
+            read_rows(
+                'SELECT organization_id FROM "LiteLLM_OrganizationTable" WHERE organization_id = %s', (organization_id,)
+            )
+            == []
+        )

tests/integration/management/test_team_budget_duration_defaults.py

@@ -0,0 +1,61 @@
+import uuid
+from pathlib import Path
+from typing import Final
+
+import pytest
+import yaml
+from pydantic import JsonValue
+
+from tests.integration._support.client import Gateway, string_value
+from tests.integration._support.database import read_rows
+from tests.integration._support.process import owned_proxy
+
+
+def _budget_row(team_id: str) -> dict[str, JsonValue]:
+    rows: Final = read_rows(
+        'SELECT max_budget, budget_duration, budget_reset_at::text FROM "LiteLLM_TeamTable" WHERE team_id = %s',
+        (team_id,),
+    )
+    assert len(rows) == 1, rows
+    return rows[0]
+
+
+@pytest.mark.covers("mgmt.team.new.explicit_null_budget_duration_overrides_default")
+def test_team_new_explicit_null_budget_duration_is_not_replaced_by_default(gateway: Gateway, tmp_path: Path) -> None:
+    config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
+    config["litellm_settings"]["default_team_params"] = {"budget_duration": "30d"}
+    path: Final = tmp_path / "team-defaults.yaml"
+    path.write_text(yaml.safe_dump(config))
+    with (
+        owned_proxy(gateway, tmp_path, {"STORE_MODEL_IN_DB": "False"}, config=path) as candidate,
+        candidate.scenario() as scenario,
+    ):
+        never_resetting: Final = candidate.request(
+            "POST",
+            "/team/new",
+            {"team_alias": f"integration-{uuid.uuid4().hex}", "max_budget": 500, "budget_duration": None},
+        )
+        assert never_resetting.status_code == 200, never_resetting.text
+        never_resetting_id: Final = string_value(never_resetting.json()["team_id"])
+        scenario.cleanups.callback(scenario.delete_team, never_resetting_id)
+        assert never_resetting.json()["max_budget"] == 500.0, never_resetting.text
+        assert never_resetting.json()["budget_duration"] is None, never_resetting.text
+        assert never_resetting.json()["budget_reset_at"] is None, never_resetting.text
+        assert _budget_row(never_resetting_id) == {
+            "max_budget": 500.0,
+            "budget_duration": None,
+            "budget_reset_at": None,
+        }
+
+        inheriting: Final = candidate.request(
+            "POST", "/team/new", {"team_alias": f"integration-{uuid.uuid4().hex}", "max_budget": 500}
+        )
+        assert inheriting.status_code == 200, inheriting.text
+        inheriting_id: Final = string_value(inheriting.json()["team_id"])
+        scenario.cleanups.callback(scenario.delete_team, inheriting_id)
+        assert inheriting.json()["budget_duration"] == "30d", inheriting.text
+        assert inheriting.json()["budget_reset_at"] is not None, inheriting.text
+        inheriting_row: Final = _budget_row(inheriting_id)
+        assert inheriting_row["max_budget"] == 500.0, inheriting_row
+        assert inheriting_row["budget_duration"] == "30d", inheriting_row
+        assert inheriting_row["budget_reset_at"] is not None, inheriting_row

tests/integration/management/test_team_member_budget_cache.py

@@ -0,0 +1,40 @@
+import os
+from typing import Final
+
+import pytest
+from pydantic import JsonValue, TypeAdapter
+from redis import Redis
+
+from tests.integration._support.client import Gateway, eventually, object_value, string_value
+from tests.integration._support.database import read_rows
+
+_CACHED_BUDGET: Final = TypeAdapter(dict[str, JsonValue])
+
+
+@pytest.mark.covers("mgmt.team_member_budget.default_budget_is_cached_in_redis_as_json")
+def test_team_member_default_budget_lands_in_redis_after_first_member_call(gateway: Gateway) -> None:
+    with gateway.scenario() as scenario:
+        model: Final = scenario.model()
+        user: Final = scenario.user()
+        team: Final = scenario.team(team_member_budget=25)
+        key: Final = scenario.key(team_id=team, user_id=user, models=[model])
+        teams: Final = read_rows('SELECT metadata FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team,))
+        assert len(teams) == 1, teams
+        budget_id: Final = string_value(object_value(teams[0]["metadata"])["team_member_budget_id"])
+        response: Final = gateway.request(
+            "POST",
+            "/v1/chat/completions",
+            {"model": model, "messages": [{"role": "user", "content": "member budget cache"}]},
+            key=key,
+        )
+        assert response.status_code == 200, response.text
+        with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache:
+            cached: Final = eventually(
+                lambda: cache.get(f"team_member_default_budget:{budget_id}"),
+                lambda value: value is not None,
+                seconds=10,
+            )
+        assert isinstance(cached, bytes), cached
+        budget: Final = _CACHED_BUDGET.validate_json(cached)
+        assert budget["budget_id"] == budget_id, cached
+        assert budget["max_budget"] == 25, cached