test: verify database transactions and persisted accounting contracts

.circleci/config.yml

@@ -2953,7 +2953,7 @@ workflows:
           name: integration-<< matrix.suite >>
           matrix:
             parameters:
-              suite: [management, accounting, providers]
+              suite: [management, accounting, database, providers]
           filters:
             branches:
               only:

tests/e2e/coverage_registry/other.yaml

@@ -56,3 +56,7 @@
 - {id: other.auth.jwt.wrong_audience_denied, module: other, tier: P0, area: auth, assertions: [wrong_audience_denied], source: "auth/handle_jwt.py", rationale: "A signed token from the trusted issuer intended for another app is rejected"}
 - {id: other.provider_wire.internal_parameters_filtered, module: other, tier: P0, area: provider_wire, assertions: [internal_parameters_filtered], source: "main.py", rationale: "A real provider request preserves content and public parameters without internal limiter fields"}
 - {id: other.provider_wire.validator_rejects_corruption, module: other, tier: P0, area: provider_wire, assertions: [validator_rejects_corruption], source: "tests/integration/_support/upstream.py", rationale: "The controlled transport rejects missing messages and internal fields while accepting supported metadata"}
+- {id: other.database.partitions.lock_wait_outlives_transaction_default, module: other, tier: P0, area: database, assertions: [lock_wait_outlives_transaction_default], source: "spend_logs_partition_manager.py", rationale: "Actual partition DDL succeeds after a witnessed permitted lock wait beyond five seconds"}
+- {id: other.database.partitions.repeat_preserves_rows, module: other, tier: P0, area: database, assertions: [repeat_preserves_rows], source: "spend_logs_partition_manager.py", rationale: "Repeated real partition maintenance preserves existing rows and one partition"}
+- {id: other.database.regeneration.writer_updates_dependent_grants, module: other, tier: P0, area: database, assertions: [writer_updates_dependent_grants], source: "access_group_key_sync.py", rationale: "Regeneration preserves dependent grants when the configured reader rejects writes"}
+- {id: other.database.access_group.failed_second_write_rolls_back_first, module: other, tier: P0, area: database, assertions: [failed_second_write_rolls_back_first], source: "access_group_endpoints.py", rationale: "A real constraint failure leaves neither group nor partial key grants"}

tests/e2e/coverage_registry/quota_management.yaml

@@ -66,3 +66,9 @@
 - {id: quota_management.spend_tracking.custom_price.matches_input_rates, module: quota_management, tier: P0, behavior: spend_tracking, variant: custom_price, assertions: [matches_input_rates], exercised_on: [chat_completions], source: "router.py", rationale: "Configured deployment prices reach the response cost"}
 - {id: quota_management.spend_tracking.default_prices.survive_nullable_sibling_reload, module: quota_management, tier: P0, behavior: spend_tracking, variant: default_prices, assertions: [survive_nullable_sibling_reload], exercised_on: [chat_completions], source: "router.py", rationale: "Omitted and null prices retain defaults across sibling loading and reload, with persisted request charges"}
 - {id: quota_management.spend_tracking.default_prices.loaded_router_preserves_cached_defaults, module: quota_management, tier: P0, behavior: spend_tracking, variant: default_prices, assertions: [loaded_router_preserves_cached_defaults], exercised_on: [chat_completions], source: "router.py", rationale: "YAML-loaded omitted and null model-info prices preserve cached defaults across real SDK requests and reload order"}
+- {id: quota_management.spend_tracking.price_precedence.zero_and_default_rates, module: quota_management, tier: P0, behavior: spend_tracking, variant: controlled_contract, exercised_on: [chat_completions], assertions: [zero_and_default_rates], source: "tests/integration", rationale: "Controlled requests with independent arithmetic, database state, and transport observations"}
+- {id: quota_management.spend_tracking.alias_prices.remain_independent_on_reload, module: quota_management, tier: P0, behavior: spend_tracking, variant: controlled_contract, exercised_on: [chat_completions], assertions: [remain_independent_on_reload], source: "tests/integration", rationale: "Controlled requests with independent arithmetic, database state, and transport observations"}
+- {id: quota_management.response_cache.generated_sequences_preserve_content_and_accounting, module: quota_management, tier: P0, behavior: spend_tracking, variant: controlled_contract, exercised_on: [chat_completions], assertions: [generated_sequences_preserve_content_and_accounting], source: "tests/integration", rationale: "Controlled requests with independent arithmetic, database state, and transport observations"}
+- {id: quota_management.budget.key.boundary_blocks_before_provider_and_reset_restores, module: quota_management, tier: P0, behavior: spend_tracking, variant: controlled_contract, exercised_on: [chat_completions], assertions: [boundary_blocks_before_provider_and_reset_restores], source: "tests/integration", rationale: "Controlled requests with independent arithmetic, database state, and transport observations"}
+- {id: quota_management.response_cache.system_messages_partition_cache_identity, module: quota_management, tier: P0, behavior: spend_tracking, variant: controlled_contract, exercised_on: [chat_completions], assertions: [system_messages_partition_cache_identity], source: "tests/integration", rationale: "Controlled requests with independent arithmetic, database state, and transport observations"}
+- {id: quota_management.response_cache.repeated_hits_preserve_identity_and_single_charge, module: quota_management, tier: P0, behavior: spend_tracking, variant: repeated_cache_hit, exercised_on: [chat_completions], assertions: [single_charge], source: "spend_tracking_utils.py", rationale: "Repeated cached responses retain their identity while distinct spend rows bill only the original request"}

tests/integration/README.md

@@ -2,7 +2,7 @@
 
 These tests exercise a running gateway, PostgreSQL and Redis with an owned local upstream. CircleCI owns this suite. Tests are grouped by behavior, with no automatic test retries or fallback to paid provider calls
 
-Use `tests/integration/run.py management`, `accounting` or `providers` to run a selected group. Set `INTEGRATION_PROXY_URL`, `INTEGRATION_UPSTREAM_URL`, `INTEGRATION_MASTER_KEY` and `DATABASE_URL` to an isolated test deployment. The runner selects the new domain directories explicitly; the legacy OCI and sandbox selections remain separate
+Use `tests/integration/run.py management`, `accounting`, `database` or `providers` to run a selected group. Set `INTEGRATION_PROXY_URL`, `INTEGRATION_UPSTREAM_URL`, `INTEGRATION_MASTER_KEY` and `DATABASE_URL` to an isolated test deployment. The runner selects the new domain directories explicitly; the legacy OCI and sandbox selections remain separate
 
 Management also requires `INTEGRATION_PEER_URL`, `REDIS_HOST` and `REDIS_PORT`. CircleCI starts two directly addressed proxy processes sharing only that job's stores. The test-only CLI wrapper supplies enterprise route entitlement, following the existing behavior suite's convention. It does not qualify license validation; run it with one worker and no reload
 
@@ -17,3 +17,7 @@ Add contract definitions to the existing `tests/e2e/coverage_registry` and map c
 Provider sentinels currently use the controlled server, not live recordings. The provider shard also runs the existing strict replay controls for changed requests, exhausted interactions, leftover interactions and no provider connection. Future recorded scenarios must use that replay-only implementation; missing recordings cannot fall back to a real provider. The observation endpoint is destructive and the current selection runs serially against one owned upstream
 
 Fixtures must contain synthetic data only. Keep private incident records and source documents out of code, fixtures, logs and PR descriptions
+
+Database cases own their temporary schemas, roles, constraints and proxy processes. They prove reader-versus-writer execution with PostgreSQL lock observations, exercise real transaction wait limits and verify rollback after a reached database failure
+
+Accounting cases compare persisted input and output cost components against literal rates, including zero and default prices. Cache state models assert actual upstream calls, response identity and every persisted charge. Generated accounting tests have a 180-second test limit to accommodate the asynchronous spend writer; CircleCI keeps the whole shard capped at 11 minutes

tests/integration/_support/process.py

@@ -0,0 +1,92 @@
+import os
+import socket
+import signal
+import subprocess
+import sys
+import time
+import uuid
+from collections.abc import Iterator, Mapping
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Final
+
+import httpx
+import psutil
+
+from integration._support.client import Gateway
+
+
+def in_group(process: psutil.Process, group: int) -> bool:
+    try:
+        return os.getpgid(process.pid) == group
+    except ProcessLookupError:
+        return False
+
+
+def group_members(group: int) -> tuple[psutil.Process, ...]:
+    return tuple(process for process in psutil.process_iter() if in_group(process, group))
+
+
+def signal_group(group: int, action: int) -> None:
+    try:
+        os.killpg(group, action)
+    except ProcessLookupError:
+        pass
+
+
+@contextmanager
+def owned_proxy(gateway: Gateway, directory: Path, overrides: Mapping[str, str]) -> Iterator[Gateway]:
+    with socket.socket() as reserve:
+        reserve.bind(("127.0.0.1", 0))
+        port: Final = reserve.getsockname()[1]
+    root: Final = Path(__file__).resolve().parents[3]
+    environment: Final = {
+        **os.environ,
+        "LITELLM_MASTER_KEY": gateway.key,
+        "LITELLM_SALT_KEY": os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt"),
+        "STORE_MODEL_IN_DB": "True",
+        **overrides,
+    }
+    output: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", str(directory)))
+    output.mkdir(parents=True, exist_ok=True)
+    with (output / f"owned-proxy-{uuid.uuid4().hex}.log").open("w") as log:
+        process: Final = subprocess.Popen(
+            [sys.executable, "-m", "integration._support.proxy", "--config", "tests/integration/proxy_config.yaml",
+             "--host", "127.0.0.1", "--port", str(port), "--num_workers", "1", "--telemetry", "False",
+             "--use_prisma_db_push", "--enforce_prisma_migration_check"],
+            cwd=root, env=environment, stdout=log, stderr=subprocess.STDOUT, start_new_session=True,
+        )
+        try:
+            with httpx.Client(base_url=f"http://127.0.0.1:{port}", timeout=15, trust_env=False) as client:
+                deadline: Final = time.monotonic() + 70
+                while True:
+                    assert process.poll() is None, "Owned proxy exited before readiness"
+                    try:
+                        if client.get("/health/readiness", timeout=2).status_code == 200:
+                            break
+                    except httpx.TransportError:
+                        pass
+                    assert time.monotonic() < deadline, "Owned proxy readiness deadline exceeded"
+                    time.sleep(0.1)
+                yield Gateway(client, gateway.key, gateway.upstream_url)
+        finally:
+            forced = False
+            if process.poll() is None:
+                process.terminate()
+                try:
+                    process.wait(timeout=30)
+                except subprocess.TimeoutExpired:
+                    forced = True
+            residual: Final = group_members(process.pid)
+            if residual:
+                signal_group(process.pid, signal.SIGTERM)
+                psutil.wait_procs(residual, timeout=5)
+            remaining: Final = group_members(process.pid)
+            if remaining:
+                forced = True
+                signal_group(process.pid, signal.SIGKILL)
+                psutil.wait_procs(remaining, timeout=3)
+            process.wait(timeout=3)
+            survivors: Final = group_members(process.pid)
+            assert not survivors, "Owned proxy child survived cleanup"
+            assert not forced, "Owned proxy required forced cleanup"

tests/integration/contracts.json

@@ -75,6 +75,34 @@
     ],
     "tests/integration/authorization/test_warmed_policy.py::test_expiry_and_explicit_clear_reach_both_warmed_workers": [
       "mgmt.key.update.expiry_changes_reach_warmed_workers"
+    ],
+    "tests/integration/database/test_partition_transactions.py::test_real_partition_ddl_survives_witnessed_lock_and_is_idempotent": [
+      "other.database.partitions.lock_wait_outlives_transaction_default",
+      "other.database.partitions.repeat_preserves_rows"
+    ],
+    "tests/integration/database/test_reader_writer_regeneration.py::test_key_regeneration_uses_writer_with_a_real_readonly_reader": [
+      "other.database.regeneration.writer_updates_dependent_grants"
+    ],
+    "tests/integration/pricing/test_price_precedence.py::test_generated_zero_null_and_omitted_prices_follow_independent_arithmetic": [
+      "quota_management.spend_tracking.price_precedence.zero_and_default_rates"
+    ],
+    "tests/integration/pricing/test_price_precedence.py::test_same_upstream_aliases_keep_distinct_prices_after_reload": [
+      "quota_management.spend_tracking.alias_prices.remain_independent_on_reload"
+    ],
+    "tests/integration/spend/test_cache_and_quota.py::test_generated_cache_sequences_preserve_content_usage_and_zero_hit_cost": [
+      "quota_management.response_cache.generated_sequences_preserve_content_and_accounting"
+    ],
+    "tests/integration/spend/test_cache_and_quota.py::test_key_budget_at_boundary_blocks_provider_then_explicit_reset_restores": [
+      "quota_management.budget.key.boundary_blocks_before_provider_and_reset_restores"
+    ],
+    "tests/integration/spend/test_cache_and_quota.py::test_different_system_messages_do_not_share_a_cached_response": [
+      "quota_management.response_cache.system_messages_partition_cache_identity"
+    ],
+    "tests/integration/database/test_transaction_atomicity.py::test_access_group_second_key_constraint_failure_rolls_back_all_writes": [
+      "other.database.access_group.failed_second_write_rolls_back_first"
+    ],
+    "tests/integration/spend/test_cache_and_quota.py::test_repeated_hits_keep_response_identity_and_create_distinct_zero_cost_rows": [
+      "quota_management.response_cache.repeated_hits_preserve_identity_and_single_charge"
     ]
   }
 }

tests/integration/database/test_partition_transactions.py

@@ -0,0 +1,86 @@
+import asyncio
+import os
+import time
+import uuid
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Final
+from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
+
+import psycopg
+import pytest
+from psycopg import sql
+from prisma import Prisma
+
+from integration._support.database import read_rows
+from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import SpendLogsPartitionManager
+
+
+@dataclass(frozen=True)
+class PartitionConnection:
+    db: Prisma
+
+
+@pytest.mark.covers("other.database.partitions.lock_wait_outlives_transaction_default", "other.database.partitions.repeat_preserves_rows")
+async def test_real_partition_ddl_survives_witnessed_lock_and_is_idempotent() -> None:
+    schema: Final = f"integration_{uuid.uuid4().hex}"
+    url: Final = os.environ["DATABASE_URL"]
+    parsed: Final = urlsplit(url)
+    scoped_url: Final = urlunsplit(parsed._replace(query=urlencode({**dict(parse_qsl(parsed.query)), "schema": schema})))
+    parent: Final = sql.Identifier(schema, "LiteLLM_SpendLogs")
+    with psycopg.connect(url, autocommit=True) as setup:
+        setup.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema)))
+        try:
+            setup.execute(sql.SQL('CREATE TABLE {} (request_id text, "startTime" timestamp NOT NULL) PARTITION BY RANGE ("startTime")').format(parent))
+            database: Final = Prisma(datasource={"url": scoped_url})
+            await database.connect()
+            try:
+                manager: Final = SpendLogsPartitionManager(interval="day", precreate_ahead=0)
+                with psycopg.connect(url) as blocker:
+                    blocker.execute(sql.SQL("LOCK TABLE {} IN ACCESS SHARE MODE").format(parent))
+                    blocker_pid: Final = blocker.info.backend_pid
+                    operation: Final = asyncio.create_task(manager.ensure_partitions(PartitionConnection(database), lambda: 7000))
+                    wait_deadline: Final = time.monotonic() + 3
+                    try:
+                        while True:
+                            witnesses: Final = read_rows(
+                                "SELECT a.pid, extract(epoch FROM clock_timestamp()-a.query_start)::double precision AS age "
+                                "FROM pg_stat_activity a WHERE %s = ANY(pg_blocking_pids(a.pid)) "
+                                "AND a.wait_event_type = 'Lock' AND a.query LIKE 'CREATE TABLE IF NOT EXISTS%%'",
+                                (blocker_pid,),
+                            )
+                            if witnesses:
+                                break
+                            assert time.monotonic() < wait_deadline, "Partition DDL never reached the held lock"
+                            await asyncio.sleep(0.02)
+                        assert len(witnesses) == 1
+                        held_at: Final = time.monotonic()
+                        age: Final = float(witnesses[0]["age"])
+                        assert age < 1, "DDL lock witness arrived too late for the qualification window"
+                        await asyncio.sleep(5.6 - age)
+                        held_seconds: Final = age + time.monotonic() - held_at
+                        assert 5.5 <= held_seconds < 6.5, f"Lock qualification timing outside window: {held_seconds}"
+                        assert not operation.done(), "DDL completed while its required lock was held"
+                    except BaseException:
+                        operation.cancel()
+                        await asyncio.gather(operation, return_exceptions=True)
+                        raise
+                    finally:
+                        blocker.rollback()
+                    ensured: Final = await asyncio.wait_for(operation, timeout=5)
+                assert len(ensured) == 1, "Partition DDL failed after the permitted lock wait"
+                catalog: Final = read_rows(
+                    "SELECT child.relname FROM pg_inherits i JOIN pg_class child ON child.oid=i.inhrelid "
+                    "JOIN pg_class parent ON parent.oid=i.inhparent JOIN pg_namespace n ON n.oid=parent.relnamespace "
+                    "WHERE n.nspname=%s AND parent.relname='LiteLLM_SpendLogs'", (schema,),
+                )
+                assert catalog == [{"relname": ensured[0]}]
+                now: Final = datetime.now(timezone.utc).replace(tzinfo=None)
+                setup.execute(sql.SQL('INSERT INTO {} VALUES (%s, %s)').format(parent), ("retained", now))
+                assert await manager.ensure_partitions(PartitionConnection(database), lambda: 7000) == ensured
+                assert setup.execute(sql.SQL("SELECT request_id FROM {}").format(parent)).fetchall() == [("retained",)]
+            finally:
+                await database.disconnect()
+        finally:
+            setup.execute(sql.SQL("DROP SCHEMA {} CASCADE").format(sql.Identifier(schema)))
+    assert read_rows("SELECT nspname FROM pg_namespace WHERE nspname=%s", (schema,)) == []

tests/integration/database/test_reader_writer_regeneration.py

@@ -0,0 +1,91 @@
+import os
+import uuid
+from hashlib import sha256
+from pathlib import Path
+from concurrent.futures import ThreadPoolExecutor
+from typing import Final
+from urllib.parse import urlsplit, urlunsplit
+
+import psycopg
+import pytest
+from psycopg import sql
+
+from integration._support.client import Gateway, eventually, string_value
+from integration._support.database import read_rows
+from integration._support.process import owned_proxy
+
+
+def delete_if_present(candidate: Gateway, key: str) -> None:
+    digest: Final = sha256(key.encode()).hexdigest()
+    if read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)):
+        candidate.post("/key/delete", {"keys": [key]})
+    assert read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)) == []
+
+
+@pytest.mark.covers("other.database.regeneration.writer_updates_dependent_grants")
+def test_key_regeneration_uses_writer_with_a_real_readonly_reader(gateway: Gateway, tmp_path: Path) -> None:
+    role: Final = f"integration_reader_{uuid.uuid4().hex}"
+    url: Final = os.environ["DATABASE_URL"]
+    parsed: Final = urlsplit(url)
+    reader_url: Final = urlunsplit(parsed._replace(netloc=f"{role}:integration-reader-password@{parsed.hostname}:{parsed.port}"))
+    with psycopg.connect(url, autocommit=True) as admin:
+        admin.execute(sql.SQL("CREATE ROLE {} LOGIN PASSWORD 'integration-reader-password' NOSUPERUSER NOINHERIT").format(sql.Identifier(role)))
+        try:
+            admin.execute(sql.SQL("GRANT USAGE ON SCHEMA public TO {}").format(sql.Identifier(role)))
+            admin.execute(sql.SQL("GRANT SELECT ON ALL TABLES IN SCHEMA public TO {}").format(sql.Identifier(role)))
+            admin.execute(sql.SQL("ALTER ROLE {} SET default_transaction_read_only = on").format(sql.Identifier(role)))
+            with psycopg.connect(reader_url, autocommit=True) as reader:
+                assert reader.execute("SHOW transaction_read_only").fetchone() == ("on",)
+                with pytest.raises(psycopg.errors.ReadOnlySqlTransaction):
+                    reader.execute('UPDATE "LiteLLM_VerificationToken" SET blocked = true WHERE false')
+            with owned_proxy(gateway, tmp_path, {"DATABASE_URL_READ_REPLICA": reader_url}) as candidate:
+                assert read_rows("SELECT pid FROM pg_stat_activity WHERE usename=%s", (role,)), "Candidate reader was never connected"
+                with gateway.scenario() as scenario:
+                    model: Final = scenario.model()
+                    outside: Final = scenario.model()
+                    old: Final = string_value(candidate.post("/key/generate", {"models": [outside]})["key"])
+                    new: Final = f"sk-integration-{uuid.uuid4().hex}"
+                    scenario.cleanups.callback(delete_if_present, gateway, old)
+                    scenario.cleanups.callback(delete_if_present, gateway, new)
+                    old_hash: Final = sha256(old.encode()).hexdigest()
+                    before: Final = candidate.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "no grant yet"}]}, key=old)
+                    assert before.status_code == 403 and before.json()["error"]["type"] == "key_model_access_denied", before.text
+                    response: Final = candidate.request("POST", "/v1/access_group", {
+                        "access_group_name": f"integration-{uuid.uuid4().hex}",
+                        "access_model_names": [model], "assigned_key_ids": [old_hash],
+                    })
+                    assert response.status_code == 201, response.text
+                    group: Final = string_value(response.json()["access_group_id"])
+                    try:
+                        with psycopg.connect(url) as blocker, ThreadPoolExecutor(max_workers=1) as executor:
+                            blocker.execute('LOCK TABLE "LiteLLM_AccessGroupTable" IN ACCESS EXCLUSIVE MODE')
+                            pending: Final = executor.submit(candidate.request, "GET", f"/v1/access_group/{group}")
+                            try:
+                                reached: Final = eventually(lambda: read_rows(
+                                    "SELECT usename FROM pg_stat_activity WHERE %s=ANY(pg_blocking_pids(pid)) "
+                                    "AND usename=%s AND query LIKE 'SELECT%%'", (blocker.info.backend_pid, role),
+                                ), bool, seconds=3)
+                                assert reached == [{"usename": role}]
+                            finally:
+                                blocker.rollback()
+                            selected: Final = pending.result(timeout=5)
+                            assert selected.status_code == 200 and selected.json()["access_group_id"] == group, selected.text
+                        assert candidate.chat(model, key=old)["usage"]["total_tokens"] == 40
+                        regenerated: Final = candidate.post("/key/regenerate", {"key": old, "new_key": new, "grace_period": "0s"})
+                        assert regenerated["key"] == new
+                        new_hash: Final = sha256(new.encode()).hexdigest()
+                        assert new != old
+                        assert read_rows('SELECT assigned_key_ids FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', (group,)) == [{"assigned_key_ids": [new_hash]}]
+                        assert read_rows('SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s)', ([old_hash, new_hash],)) == [{"token": new_hash, "access_group_ids": [group]}]
+                        assert candidate.chat(model, key=new)["usage"]["total_tokens"] == 40
+                        assert candidate.chat(outside, key=new)["usage"]["total_tokens"] == 40
+                        denied: Final = candidate.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "rotated key"}]}, key=old)
+                        assert denied.status_code == 401 and denied.json()["error"]["type"] == "token_not_found_in_db", denied.text
+                    finally:
+                        deleted: Final = gateway.request("DELETE", f"/v1/access_group/{group}")
+                        assert deleted.status_code == 204, deleted.text
+                        assert read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', (group,)) == []
+        finally:
+            admin.execute(sql.SQL("DROP OWNED BY {}").format(sql.Identifier(role)))
+            admin.execute(sql.SQL("DROP ROLE {}").format(sql.Identifier(role)))
+    assert read_rows("SELECT rolname FROM pg_roles WHERE rolname=%s", (role,)) == []

tests/integration/database/test_transaction_atomicity.py

@@ -0,0 +1,67 @@
+import os
+import uuid
+from contextlib import ExitStack
+from hashlib import sha256
+from typing import Final
+
+import psycopg
+import pytest
+from psycopg import sql
+
+from integration._support.client import Gateway
+from integration._support.database import read_rows
+
+
+@pytest.mark.covers("other.database.access_group.failed_second_write_rolls_back_first")
+def test_access_group_second_key_constraint_failure_rolls_back_all_writes(gateway: Gateway) -> None:
+    with gateway.scenario() as scenario:
+        model: Final = scenario.model()
+        outside: Final = scenario.model()
+        keys: Final = (scenario.key(models=[outside]), scenario.key(models=[outside]))
+        tokens: Final = [sha256(key.encode()).hexdigest() for key in keys]
+        name: Final = f"integration-{uuid.uuid4().hex}"
+        constraint: Final = f"integration_reject_{uuid.uuid4().hex}"
+        witness: Final = constraint + "_seq"
+        check_function: Final = constraint + "_check"
+        body: Final = {"access_group_name": name, "access_model_names": [model], "assigned_key_ids": tokens}
+        def remove_partial_group() -> None:
+            for row in read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_name=%s', (name,)):
+                response: Final = gateway.request("DELETE", f"/v1/access_group/{row['access_group_id']}")
+                assert response.status_code == 204, response.text
+            assert read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_name=%s', (name,)) == []
+
+        scenario.cleanups.callback(remove_partial_group)
+        before: Final = read_rows('SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s) ORDER BY token', (tokens,))
+        with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection, ExitStack() as cleanup:
+            connection.execute(sql.SQL("CREATE SEQUENCE {}").format(sql.Identifier(witness)))
+            cleanup.callback(connection.execute, sql.SQL("DROP SEQUENCE {}").format(sql.Identifier(witness)))
+            connection.execute(sql.SQL("CREATE FUNCTION {}(text[]) RETURNS boolean LANGUAGE plpgsql AS $$ BEGIN IF cardinality($1)>0 THEN PERFORM nextval({}); RETURN false; END IF; RETURN true; END $$").format(sql.Identifier(check_function), sql.Literal(witness)))
+            cleanup.callback(connection.execute, sql.SQL("DROP FUNCTION {}(text[])").format(sql.Identifier(check_function)))
+            connection.execute(sql.SQL('ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT {} CHECK (token <> {} OR {}(access_group_ids))').format(sql.Identifier(constraint), sql.Literal(tokens[1]), sql.Identifier(check_function)))
+            cleanup.callback(connection.execute, sql.SQL('ALTER TABLE "LiteLLM_VerificationToken" DROP CONSTRAINT {}').format(sql.Identifier(constraint)))
+            try:
+                assert connection.execute(sql.SQL("SELECT is_called FROM {}").format(sql.Identifier(witness))).fetchone() == (False,)
+                failed: Final = gateway.request("POST", "/v1/access_group", body)
+                assert failed.status_code == 500, failed.text
+                # Sequence advancement survives rollback and proves the rejecting
+                # constraint actually evaluated the second key's nonempty grant.
+                assert connection.execute(sql.SQL("SELECT is_called FROM {}").format(sql.Identifier(witness))).fetchone() == (True,)
+                assert read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_name=%s', (name,)) == []
+                assert read_rows('SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s) ORDER BY token', (tokens,)) == before
+                for key in keys:
+                    denied: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "rolled back grant"}]}, key=key)
+                    assert denied.status_code == 403 and denied.json()["error"]["type"] == "key_model_access_denied", denied.text
+            finally:
+                cleanup.close()
+        created: Final = gateway.request("POST", "/v1/access_group", body)
+        assert created.status_code == 201, created.text
+        identity: Final = created.json()["access_group_id"]
+        try:
+            for key in keys:
+                assert gateway.chat(model, key=key)["usage"]["total_tokens"] == 40
+        finally:
+            deleted: Final = gateway.request("DELETE", f"/v1/access_group/{identity}")
+            assert deleted.status_code == 204, deleted.text
+            assert read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', (identity,)) == []
+        assert read_rows('SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s) ORDER BY token', (tokens,)) == before
+        assert read_rows('SELECT conname FROM pg_constraint WHERE conname=%s', (constraint,)) == []

tests/integration/pricing/test_price_precedence.py

@@ -0,0 +1,80 @@
+import json
+import uuid
+from typing import Final
+
+import pytest
+from hypothesis import Phase, example, given, settings, strategies as st
+
+from integration._support.client import Gateway, eventually, object_value
+from integration._support.database import read_rows
+
+
+@pytest.mark.covers("quota_management.spend_tracking.price_precedence.zero_and_default_rates")
+@pytest.mark.timeout(180)
+def test_generated_zero_null_and_omitted_prices_follow_independent_arithmetic(gateway: Gateway) -> None:
+    @settings(max_examples=20, deadline=None, database=None, phases=(Phase.explicit, Phase.generate, Phase.shrink))
+    @example(rates=(0, 0))
+    @example(rates=(1, 2))
+    @example(rates=("null", "null"))
+    @given(rates=st.one_of(st.sampled_from((("omitted", "omitted"), ("null", "null"))), st.tuples(st.integers(0, 25), st.integers(0, 25))))
+    def check(rates: tuple[str | int, str | int]) -> None:
+        if rates[0] in ("omitted", "null"):
+            parameters: Final = {} if rates[0] == "omitted" else {"input_cost_per_token": None, "output_cost_per_token": None}
+            input_rate, output_rate = 0.00000015, 0.0000006
+        else:
+            assert isinstance(rates[0], int) and isinstance(rates[1], int)
+            input_rate, output_rate = rates[0] / 1_000_000, rates[1] / 1_000_000
+            parameters = {"input_cost_per_token": input_rate, "output_cost_per_token": output_rate}
+        with gateway.scenario() as scenario:
+            model: Final = scenario.model(**parameters)
+            response: Final = gateway.request("POST", "/v1/chat/completions", {
+                "model": model, "messages": [{"role": "user", "content": f"independent price {uuid.uuid4().hex}"}],
+            })
+            assert response.status_code == 200, response.text
+            assert response.json()["usage"] == {"prompt_tokens": 20, "completion_tokens": 20, "total_tokens": 40}
+            expected: Final = 20 * input_rate + 20 * output_rate
+            if expected:
+                assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected, rel=1e-6)
+            else:
+                assert response.headers.get("x-litellm-response-cost") in (None, "0", "0.0")
+            rows: Final = eventually(lambda: read_rows(
+                'SELECT spend, metadata, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
+                (response.json()["id"],),
+            ), lambda values: len(values) == 1, seconds=70)
+            assert rows[0]["prompt_tokens"] == 20 and rows[0]["completion_tokens"] == 20
+            assert float(rows[0]["spend"]) == pytest.approx(expected, rel=1e-6)
+            metadata: Final = rows[0]["metadata"]
+            parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata)
+            breakdown: Final = object_value(parsed["cost_breakdown"])
+            assert float(breakdown["input_cost"]) == pytest.approx(20 * input_rate, rel=1e-6)
+            assert float(breakdown["output_cost"]) == pytest.approx(20 * output_rate, rel=1e-6)
+
+    check()
+
+
+@pytest.mark.covers("quota_management.spend_tracking.alias_prices.remain_independent_on_reload")
+def test_same_upstream_aliases_keep_distinct_prices_after_reload(gateway: Gateway) -> None:
+    for order in (("free", "paid"), ("paid", "free")):
+        with gateway.scenario() as scenario:
+            rates: Final = {"free": {"input_cost_per_token": 0, "output_cost_per_token": 0}, "paid": {"input_cost_per_token": 0.001, "output_cost_per_token": 0.003}}
+            aliases: Final = {kind: scenario.model(**rates[kind]) for kind in order}
+            for generation in range(2):
+                for kind in order if generation == 0 else reversed(order):
+                    model: Final = aliases[kind]
+                    cost: Final = 0.08 if kind == "paid" else 0.0
+                    response: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": f"alias price {model} {generation}"}]})
+                    assert response.status_code == 200, response.text
+                    assert response.json()["usage"]["total_tokens"] == 40
+                    if cost:
+                        assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(cost)
+                    rows: Final = eventually(lambda response=response: read_rows('SELECT spend, metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (response.json()["id"],)), lambda values: len(values) == 1, seconds=70)
+                    assert float(rows[0]["spend"]) == pytest.approx(cost)
+                    metadata: Final = rows[0]["metadata"]
+                    parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata)
+                    breakdown: Final = object_value(parsed["cost_breakdown"])
+                    assert float(breakdown["input_cost"]) == pytest.approx(20 * rates[kind]["input_cost_per_token"])
+                    assert float(breakdown["output_cost"]) == pytest.approx(20 * rates[kind]["output_cost_per_token"])
+                entries: Final = gateway.get("/model/info")["data"]
+                target: Final = next(entry for entry in entries if entry["model_name"] == aliases["paid"])
+                changed: Final = gateway.request("PATCH", f"/model/{target['model_info']['id']}/update", {"model_info": {"description": "price reload"}})
+                assert changed.status_code == 200, changed.text

tests/integration/spend/test_cache_and_quota.py

@@ -0,0 +1,166 @@
+import uuid
+from contextlib import ExitStack
+from hashlib import sha256
+from typing import Final
+
+import httpx
+import pytest
+from hypothesis import strategies as st
+from hypothesis.stateful import RuleBasedStateMachine, rule, run_state_machine_as_test
+
+from integration._support.client import Gateway, eventually
+from integration._support.database import read_rows
+from integration._support.generation import LIFECYCLE_SETTINGS, bounded_http_requests
+
+
+@pytest.mark.covers("quota_management.response_cache.generated_sequences_preserve_content_and_accounting")
+@pytest.mark.timeout(180)
+def test_generated_cache_sequences_preserve_content_usage_and_zero_hit_cost(gateway: Gateway) -> None:
+    class CacheRequests(RuleBasedStateMachine):
+        def __init__(self) -> None:
+            super().__init__()
+            self.resources = ExitStack()
+            try:
+                self.scenario = self.resources.enter_context(gateway.scenario())
+                self.upstream = self.resources.enter_context(httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False))
+                self.model = self.scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
+                self.key = self.scenario.key(models=[self.model])
+                self.prefix = uuid.uuid4().hex
+                self.seen: frozenset[int] = frozenset()
+                self.requests = 0
+                self.paid = 0
+                self.failed = False
+                self.identities: dict[int, str] = {}
+            except BaseException:
+                with budget.cleanup():
+                    self.resources.close()
+                raise
+
+        @rule(marker=st.integers(min_value=0, max_value=2))
+        def request(self, marker: int) -> None:
+            try:
+                self.perform_request(marker)
+            except BaseException:
+                self.failed = True
+                raise
+
+        def perform_request(self, marker: int) -> None:
+            self.upstream.get("/__observations").raise_for_status()
+            response: Final = gateway.request("POST", "/v1/chat/completions", {
+                "model": self.model, "messages": [{"role": "user", "content": f"{self.prefix}-{marker}"}],
+            }, key=self.key)
+            assert response.status_code == 200, response.text
+            self.requests += 1
+            body: Final = response.json()
+            assert body["choices"][0]["message"]["content"] == "Hello! This is a mock response from the fake OpenAI endpoint."
+            assert body["usage"]["total_tokens"] == 40
+            observed: Final = self.upstream.get("/__observations").json()["requests"]
+            expected_calls: Final = 0 if marker in self.seen else 1
+            assert len(observed) == expected_calls, observed
+            # The response can retain its original cost header on a cache hit.
+            # Per-request billed cost is checked against fresh spend rows below.
+            if marker not in self.seen:
+                assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(0.06)
+            if marker in self.identities:
+                assert body["id"] == self.identities[marker]
+            else:
+                assert body["id"] not in self.identities.values()
+                self.identities = {**self.identities, marker: body["id"]}
+            self.paid += expected_calls
+            self.seen = self.seen.union((marker,))
+
+        def teardown(self) -> None:
+            try:
+                if self.requests and not self.failed:
+                    rows: Final = eventually(lambda: read_rows(
+                        'SELECT request_id, spend, cache_hit, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s',
+                        (sha256(self.key.encode()).hexdigest(),),
+                    ), lambda values: len(values) == self.requests, seconds=70)
+                    assert len({row["request_id"] for row in rows}) == self.requests
+                    assert sum(float(row["spend"]) for row in rows) == pytest.approx(self.paid * 0.06)
+                    assert sum(row["cache_hit"] == "True" for row in rows) == self.requests - self.paid
+                    for row in rows:
+                        assert row["prompt_tokens"] == 20 and row["completion_tokens"] == 20
+                        if row["cache_hit"] == "True":
+                            assert float(row["spend"]) == 0 and "_cache_hit" in row["request_id"]
+                            assert any(row["request_id"].startswith(identity + "_cache_hit") for identity in self.identities.values())
+                        else:
+                            assert row["request_id"] in self.identities.values()
+                            assert float(row["spend"]) == pytest.approx(0.06)
+            finally:
+                with budget.cleanup():
+                    self.resources.close()
+
+    with bounded_http_requests((gateway,), limit=2000) as budget:
+        run_state_machine_as_test(CacheRequests, settings=LIFECYCLE_SETTINGS)
+
+
+@pytest.mark.covers("quota_management.response_cache.repeated_hits_preserve_identity_and_single_charge")
+def test_repeated_hits_keep_response_identity_and_create_distinct_zero_cost_rows(gateway: Gateway) -> None:
+    with gateway.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream:
+        model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
+        key: Final = scenario.key(models=[model])
+        prompt: Final = f"repeated cache {uuid.uuid4().hex}"
+        upstream.get("/__observations").raise_for_status()
+        results: Final = tuple(gateway.chat(model, key=key, text=prompt) for _ in range(3))
+        assert len(upstream.get("/__observations").json()["requests"]) == 1
+        assert len({result["id"] for result in results}) == 1
+        for result in results:
+            assert result["choices"][0]["message"]["content"] == "Hello! This is a mock response from the fake OpenAI endpoint."
+            assert result["usage"]["total_tokens"] == 40
+        rows: Final = eventually(lambda: read_rows('SELECT request_id, spend, cache_hit FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (sha256(key.encode()).hexdigest(),)), lambda values: len(values) == 3, seconds=70)
+        assert len({row["request_id"] for row in rows}) == 3
+        assert sorted(float(row["spend"]) for row in rows) == [0, 0, 0.06]
+        for row in rows:
+            if row["cache_hit"] == "True":
+                assert float(row["spend"]) == 0
+                assert row["request_id"].startswith(results[0]["id"] + "_cache_hit")
+            else:
+                assert row["request_id"] == results[0]["id"] and float(row["spend"]) == pytest.approx(0.06)
+
+
+@pytest.mark.covers("quota_management.budget.key.boundary_blocks_before_provider_and_reset_restores")
+def test_key_budget_at_boundary_blocks_provider_then_explicit_reset_restores(gateway: Gateway) -> None:
+    with gateway.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream:
+        model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
+        key: Final = scenario.key(models=[model], max_budget=0.06)
+        control: Final = scenario.key(models=[model])
+        first: Final = gateway.chat(model, key=key, text=f"budget {uuid.uuid4().hex}")
+        assert first["usage"]["total_tokens"] == 40
+        digest: Final = sha256(key.encode()).hexdigest()
+        spent: Final = eventually(lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), lambda values: len(values) == 1 and float(values[0]["spend"]) >= 0.06, seconds=70)
+        assert float(spent[0]["spend"]) == pytest.approx(0.06)
+        upstream.get("/__observations").raise_for_status()
+        denied: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": f"over budget {uuid.uuid4().hex}"}]}, key=key)
+        assert denied.status_code == 429 and denied.json()["error"]["type"] == "budget_exceeded", denied.text
+        assert upstream.get("/__observations").json()["requests"] == []
+        assert gateway.chat(model, key=control, text=f"control {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40
+        gateway.post("/key/update", {"key": key, "spend": 0})
+        assert read_rows('SELECT spend, max_budget FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)) == [{"spend": 0.0, "max_budget": 0.06}]
+        assert gateway.chat(model, key=key, text=f"reset {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40
+        eventually(lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), lambda values: len(values) == 1 and float(values[0]["spend"]) >= 0.06, seconds=70)
+        upstream.get("/__observations").raise_for_status()
+        denied_again: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": f"boundary again {uuid.uuid4().hex}"}]}, key=key)
+        assert denied_again.status_code == 429 and denied_again.json()["error"]["type"] == "budget_exceeded", denied_again.text
+        assert upstream.get("/__observations").json()["requests"] == []
+
+
+@pytest.mark.covers("quota_management.response_cache.system_messages_partition_cache_identity")
+def test_different_system_messages_do_not_share_a_cached_response(gateway: Gateway) -> None:
+    with gateway.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream:
+        model: Final = scenario.model()
+        prompt: Final = uuid.uuid4().hex
+        identities: dict[str, str] = {}
+        for system, expected_calls in (("first policy", 1), ("second policy", 1), ("first policy", 0)):
+            upstream.get("/__observations").raise_for_status()
+            response: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "system", "content": system}, {"role": "user", "content": prompt}]})
+            assert response.status_code == 200 and response.json()["usage"]["total_tokens"] == 40, response.text
+            calls: Final = upstream.get("/__observations").json()["requests"]
+            assert len(calls) == expected_calls
+            if system in identities:
+                assert response.json()["id"] == identities[system]
+            else:
+                assert response.json()["id"] not in identities.values()
+                identities = {**identities, system: response.json()["id"]}
+            if calls:
+                assert calls[0]["body"]["messages"] == [{"role": "system", "content": system}, {"role": "user", "content": prompt}]