From feb92d3255ab44f645d17eb3afe988eb0b50ad0a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 14 Oct 2026 07:38:06 +0000 Subject: [PATCH] test(integration): cover orgs_users_customers endpoint contracts (#44725) --- .../test_access_group_model_listing.py | 354 +++++++++++ .../authorization/test_end_user_block.py | 367 ++++++++++++ .../test_user_self_update_guard.py | 287 +++++++++ .../management/test_organization_lifecycle.py | 150 ++++- .../test_user_delete_revokes_keys.py | 394 +++++++++++++ ..._user_updates_wedged_coordination_redis.py | 324 ++++++++++ .../spend/test_end_user_budget_enforcement.py | 200 +++++++ .../test_end_user_spend_without_proxy_user.py | 553 +++++++++++++++++- 8 files changed, 2626 insertions(+), 3 deletions(-) create mode 100644 tests/integration/authorization/test_end_user_block.py create mode 100644 tests/integration/authorization/test_user_self_update_guard.py create mode 100644 tests/integration/management/test_user_delete_revokes_keys.py diff --git a/tests/integration/authorization/test_access_group_model_listing.py b/tests/integration/authorization/test_access_group_model_listing.py index 9b37cc8f232..13439d19769 100644 --- a/tests/integration/authorization/test_access_group_model_listing.py +++ b/tests/integration/authorization/test_access_group_model_listing.py @@ -1,12 +1,15 @@ import uuid from collections.abc import Iterator from contextlib import contextmanager +from hashlib import sha256 from typing import Final import httpx import pytest +from pydantic import BaseModel, JsonValue from tests.integration._support.client import Gateway, eventually, object_value, string_value +from tests.integration._support.database import read_rows @contextmanager @@ -49,3 +52,354 @@ def test_team_key_lists_models_granted_through_team_access_group(gateway: Gatewa ) assert response.status_code == 200, response.text assert _listed_model_ids(response) == (model,), response.text + + +class _AccessGroupInfo(BaseModel): + access_group_id: str + assigned_team_ids: list[str] | None + assigned_key_ids: list[str] | None + + +def _denied_error(model: str, denied_type: str) -> dict[str, JsonValue]: + return { + "error": { + "message": f"The requested model '{model}' is not available for this API key, or the model name is invalid. Check the models available to you and try again.", + "type": denied_type, + "param": "model", + "code": "403", + } + } + + +def _observed(upstream: httpx.Client) -> list[dict[str, JsonValue]]: + response: Final = upstream.get("/__observations") + response.raise_for_status() + requests: Final = response.json()["requests"] + assert isinstance(requests, list) + return requests + + +def _chat(proxy: Gateway, model: str, key: str, text: str) -> httpx.Response: + return proxy.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": text}]}, + key=key, + ) + + +def _observation(text: str) -> dict[str, JsonValue]: + return { + "path": "/v1/chat/completions", + "authorization": "Bearer integration-provider-key", + "body": {"messages": [{"role": "user", "content": text}], "model": "gpt-4o-mini"}, + "method": "POST", + "api_key": "", + } + + +def _team_groups(team_id: str) -> list[dict[str, JsonValue]]: + return read_rows('SELECT access_group_ids FROM "LiteLLM_TeamTable" WHERE team_id=%s', (team_id,)) + + +def _key_groups(hashed: str) -> list[dict[str, JsonValue]]: + return read_rows('SELECT access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=%s', (hashed,)) + + +def _group_row(access_group_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT assigned_team_ids, assigned_key_ids FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', + (access_group_id,), + ) + + +def _assert_denied(response: httpx.Response, model: str, denied_type: str) -> None: + assert response.status_code == 403, response.text + assert response.json() == _denied_error(model, denied_type), response.text + + +def _assert_served(response: httpx.Response) -> None: + assert response.status_code == 200, response.text + + +@pytest.mark.parametrize("route_prefix", ("/v1/access_group", "/v1/unified_access_group")) +def test_access_group_assignment_grants_and_revokes_on_gateway_and_peer( + gateway: Gateway, peer: Gateway, route_prefix: str +) -> None: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): + model: Final = scenario.model() + team_id: Final = scenario.team(models=["no-default-models"]) + team_key: Final = scenario.key(team_id=team_id) + assigned_key: Final = scenario.key(models=["no-default-models"]) + hashed_key: Final = sha256(assigned_key.encode()).hexdigest() + _observed(upstream) + _assert_denied(_chat(gateway, model, team_key, "control team"), model, "team_model_access_denied") + _assert_denied(_chat(gateway, model, assigned_key, "control key"), model, "key_model_access_denied") + assert _observed(upstream) == [] + + created: Final = gateway.request( + "POST", + route_prefix, + { + "access_group_name": f"integration-{uuid.uuid4().hex}", + "access_model_names": [model], + "assigned_team_ids": [team_id], + "assigned_key_ids": [hashed_key], + }, + ) + assert created.status_code == 201, created.text + access_group_id: Final = string_value(created.json()["access_group_id"]) + scenario.cleanups.callback(gateway.request, "DELETE", f"{route_prefix}/{access_group_id}") + info: Final = _AccessGroupInfo.model_validate(gateway.get(f"{route_prefix}/{access_group_id}")) + assert info.assigned_team_ids == [team_id] + assert info.assigned_key_ids == [hashed_key] + assert _group_row(access_group_id) == [{"assigned_team_ids": [team_id], "assigned_key_ids": [hashed_key]}] + assert _team_groups(team_id) == [{"access_group_ids": [access_group_id]}] + assert _key_groups(hashed_key) == [{"access_group_ids": [access_group_id]}] + + expected: Final = [] + for proxy in (gateway, peer): + for key in (team_key, assigned_key): + text: Final = f"granted {uuid.uuid4().hex}" + _assert_served( + eventually( + lambda p=proxy, k=key, t=text: _chat(p, model, k, t), + lambda response: response.status_code == 200, + seconds=30, + ) + ) + expected.append(_observation(text)) + assert _observed(upstream) == expected + + removed_team: Final = gateway.request("PUT", f"{route_prefix}/{access_group_id}", {"assigned_team_ids": []}) + assert removed_team.status_code == 200, removed_team.text + assert _team_groups(team_id) == [{"access_group_ids": []}] + _assert_denied(_chat(gateway, model, team_key, "denied"), model, "team_model_access_denied") + denied_observations: Final = _observed(upstream) + assert denied_observations == [], denied_observations + still_granted: Final = f"still granted {uuid.uuid4().hex}" + _assert_served(_chat(gateway, model, assigned_key, still_granted)) + assert _observed(upstream) == [_observation(still_granted)] + + removed_key: Final = gateway.request("PUT", f"{route_prefix}/{access_group_id}", {"assigned_key_ids": []}) + assert removed_key.status_code == 200, removed_key.text + assert _key_groups(hashed_key) == [{"access_group_ids": []}] + _assert_denied(_chat(gateway, model, assigned_key, "denied"), model, "key_model_access_denied") + assert _observed(upstream) == [] + + restored: Final = gateway.request( + "PUT", + f"{route_prefix}/{access_group_id}", + {"assigned_team_ids": [team_id], "assigned_key_ids": [hashed_key]}, + ) + assert restored.status_code == 200, restored.text + assert _observed(upstream) == [] + regranted: Final = [] + for key in (team_key, assigned_key): + text = f"re-granted {uuid.uuid4().hex}" + _assert_served(_chat(gateway, model, key, text)) + regranted.append(_observation(text)) + assert _observed(upstream) == regranted + + deleted: Final = gateway.request("DELETE", f"{route_prefix}/{access_group_id}") + assert deleted.status_code == 204, deleted.text + assert _group_row(access_group_id) == [] + assert _team_groups(team_id) == [{"access_group_ids": []}] + assert _key_groups(hashed_key) == [{"access_group_ids": []}] + denied_pairs: Final = ( + (team_key, "team_model_access_denied"), + (assigned_key, "key_model_access_denied"), + ) + for key, denied_type in denied_pairs: + _assert_denied(_chat(gateway, model, key, "deleted"), model, denied_type) + deleted_observations: Final = _observed(upstream) + assert deleted_observations == [], deleted_observations + for key, denied_type in denied_pairs: + _assert_denied(_chat(gateway, model, key, "deleted"), model, denied_type) + assert _observed(upstream) == [] + + +@pytest.mark.parametrize("route_prefix", ("/v1/access_group", "/v1/unified_access_group")) +def test_deleting_an_access_group_revokes_a_still_assigned_team_and_key(gateway: Gateway, route_prefix: str) -> None: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): + model: Final = scenario.model() + team_id: Final = scenario.team(models=["no-default-models"]) + team_key: Final = scenario.key(team_id=team_id) + assigned_key: Final = scenario.key(models=["no-default-models"]) + hashed_key: Final = sha256(assigned_key.encode()).hexdigest() + created: Final = gateway.request( + "POST", + route_prefix, + { + "access_group_name": f"integration-{uuid.uuid4().hex}", + "access_model_names": [model], + "assigned_team_ids": [team_id], + "assigned_key_ids": [hashed_key], + }, + ) + assert created.status_code == 201, created.text + access_group_id: Final = string_value(created.json()["access_group_id"]) + scenario.cleanups.callback(gateway.request, "DELETE", f"{route_prefix}/{access_group_id}") + assert _group_row(access_group_id) == [{"assigned_team_ids": [team_id], "assigned_key_ids": [hashed_key]}] + assert _team_groups(team_id) == [{"access_group_ids": [access_group_id]}] + assert _key_groups(hashed_key) == [{"access_group_ids": [access_group_id]}] + assert _observed(upstream) == [] + + team_text: Final = f"delete warm team {uuid.uuid4().hex}" + team_response: Final = eventually( + lambda: _chat(gateway, model, team_key, team_text), + lambda response: response.status_code == 200, + seconds=30, + ) + _assert_served(team_response) + key_text: Final = f"delete warm key {uuid.uuid4().hex}" + key_response: Final = eventually( + lambda: _chat(gateway, model, assigned_key, key_text), + lambda response: response.status_code == 200, + seconds=30, + ) + _assert_served(key_response) + assert _observed(upstream) == [_observation(team_text), _observation(key_text)] + + deleted: Final = gateway.request("DELETE", f"{route_prefix}/{access_group_id}") + assert deleted.status_code == 204, deleted.text + assert _group_row(access_group_id) == [] + assert _team_groups(team_id) == [{"access_group_ids": []}] + assert _key_groups(hashed_key) == [{"access_group_ids": []}] + assert _observed(upstream) == [] + denied_pairs: Final = ( + (team_key, "team_model_access_denied"), + (assigned_key, "key_model_access_denied"), + ) + for key, denied_type in denied_pairs: + _assert_denied(_chat(gateway, model, key, "deleted"), model, denied_type) + assert _observed(upstream) == [] + + +@pytest.mark.parametrize("route_prefix", ("/v1/access_group", "/v1/unified_access_group")) +@pytest.mark.parametrize("revocation", ("team", "key", "delete")) +def test_access_group_revocation_reaches_a_warmed_peer_before_the_cache_ttl( + gateway: Gateway, peer: Gateway, route_prefix: str, revocation: str +) -> None: + pytest.skip("BUG: warmed peer returns 200 instead of 403 after access-group revocation") + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): + model: Final = scenario.model() + team_id: Final = scenario.team(models=["no-default-models"]) + team_key: Final = scenario.key(team_id=team_id) + assigned_key: Final = scenario.key(models=["no-default-models"]) + hashed_key: Final = sha256(assigned_key.encode()).hexdigest() + created: Final = gateway.request( + "POST", + route_prefix, + { + "access_group_name": f"integration-{uuid.uuid4().hex}", + "access_model_names": [model], + "assigned_team_ids": [team_id], + "assigned_key_ids": [hashed_key], + }, + ) + assert created.status_code == 201, created.text + access_group_id: Final = string_value(created.json()["access_group_id"]) + scenario.cleanups.callback(gateway.request, "DELETE", f"{route_prefix}/{access_group_id}") + _observed(upstream) + + if revocation == "team": + warmed: Final = _chat(peer, model, team_key, "warm team") + assert warmed.status_code == 200, warmed.text + assert _observed(upstream) == [_observation("warm team")] + team_update: Final = gateway.request("PUT", f"{route_prefix}/{access_group_id}", {"assigned_team_ids": []}) + assert team_update.status_code == 200, team_update.text + _assert_denied( + eventually( + lambda: _chat(peer, model, team_key, "revoked team"), + lambda response: response.status_code == 403, + seconds=10, + ), + model, + "team_model_access_denied", + ) + elif revocation == "key": + warmed: Final = _chat(peer, model, assigned_key, "warm key") + assert warmed.status_code == 200, warmed.text + assert _observed(upstream) == [_observation("warm key")] + key_update: Final = gateway.request("PUT", f"{route_prefix}/{access_group_id}", {"assigned_key_ids": []}) + assert key_update.status_code == 200, key_update.text + _assert_denied( + eventually( + lambda: _chat(peer, model, assigned_key, "revoked key"), + lambda response: response.status_code == 403, + seconds=10, + ), + model, + "key_model_access_denied", + ) + else: + warmed_team: Final = _chat(peer, model, team_key, "warm delete team") + warmed_key: Final = _chat(peer, model, assigned_key, "warm delete key") + assert warmed_team.status_code == 200, warmed_team.text + assert warmed_key.status_code == 200, warmed_key.text + assert _observed(upstream) == [ + _observation("warm delete team"), + _observation("warm delete key"), + ] + deleted: Final = gateway.request("DELETE", f"{route_prefix}/{access_group_id}") + assert deleted.status_code == 204, deleted.text + _assert_denied( + eventually( + lambda: _chat(peer, model, team_key, "revoked delete team"), + lambda response: response.status_code == 403, + seconds=10, + ), + model, + "team_model_access_denied", + ) + _assert_denied( + eventually( + lambda: _chat(peer, model, assigned_key, "revoked delete key"), + lambda response: response.status_code == 403, + seconds=10, + ), + model, + "key_model_access_denied", + ) + assert _observed(upstream) == [] + + +def test_raw_key_in_assigned_key_ids_grants_the_same_access_as_the_hashed_key(gateway: Gateway) -> None: + pytest.skip("BUG: raw assigned key does not populate access_group_ids for its hashed database key") + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): + model: Final = scenario.model() + key: Final = scenario.key(models=["no-default-models"]) + hashed_key: Final = sha256(key.encode()).hexdigest() + _observed(upstream) + _assert_denied(_chat(gateway, model, key, "control"), model, "key_model_access_denied") + assert _observed(upstream) == [] + + created: Final = gateway.request( + "POST", + "/v1/access_group", + { + "access_group_name": f"integration-{uuid.uuid4().hex}", + "access_model_names": [model], + "assigned_key_ids": [key], + }, + ) + assert created.status_code == 201, created.text + access_group_id: Final = string_value(created.json()["access_group_id"]) + scenario.cleanups.callback(gateway.request, "DELETE", f"/v1/access_group/{access_group_id}") + assert _key_groups(hashed_key) == [{"access_group_ids": [access_group_id]}] + text: Final = f"granted {uuid.uuid4().hex}" + _assert_served(_chat(gateway, model, key, text)) + assert _observed(upstream) == [_observation(text)] diff --git a/tests/integration/authorization/test_end_user_block.py b/tests/integration/authorization/test_end_user_block.py new file mode 100644 index 00000000000..9f4fdb7f8f0 --- /dev/null +++ b/tests/integration/authorization/test_end_user_block.py @@ -0,0 +1,367 @@ +import uuid +from contextlib import ExitStack +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, eventually, object_value +from integration._support.database import read_rows +from integration._support.process import owned_proxy +from pydantic import BaseModel, JsonValue + + +class _BlockedCustomerInfo(BaseModel): + user_id: str + blocked: bool + + +class _BlockUsersResponse(BaseModel): + blocked_users: list[_BlockedCustomerInfo] + + +def _blocked_row(user_id: str) -> list[dict[str, JsonValue]]: + return read_rows('SELECT blocked FROM "LiteLLM_EndUserTable" WHERE user_id=%s', (user_id,)) + + +def _observed(upstream: httpx.Client) -> list[dict[str, JsonValue]]: + response: Final = upstream.get("/__observations") + response.raise_for_status() + requests: Final = response.json()["requests"] + assert isinstance(requests, list) + return requests + + +def _chat(gateway: Gateway, model: str, key: str, end_user: str) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"end user block {uuid.uuid4().hex}"}], + "user": end_user, + }, + key=key, + ) + + +def _observation(text: str, end_user: str) -> dict[str, JsonValue]: + return { + "path": "/v1/chat/completions", + "authorization": "Bearer integration-provider-key", + "body": { + "messages": [{"role": "user", "content": text}], + "model": "gpt-4o-mini", + "user": end_user, + }, + "method": "POST", + "api_key": "", + } + + +def _assert_blocked_refusal(response: httpx.Response, end_user: str) -> None: + assert response.status_code == 400, response.text + assert response.json() == { + "error": { + "message": f"User blocked from making LLM API Calls. User={end_user}", + "type": "invalid_request_error", + "param": None, + "code": "400", + "provider_specific_fields": {"error": f"User blocked from making LLM API Calls. User={end_user}"}, + } + }, response.text + + +@dataclass(frozen=True, slots=True) +class _CallbackRig: + first: Gateway + second: Gateway + + +def _callback_config(directory: Path, blocked_user_list: list[str] | None = None) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + settings: Final = { + **object_value(config["litellm_settings"]), + "callbacks": ["blocked_user_check"], + } + if blocked_user_list is not None: + settings["blocked_user_list"] = blocked_user_list + config["litellm_settings"] = settings + path: Final = directory / "config.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _two_callback_proxies(gateway: Gateway, tmp_path: Path, stack: ExitStack) -> _CallbackRig: + config: Final = _callback_config(tmp_path) + first: Final = stack.enter_context(owned_proxy(gateway, tmp_path / "a", {}, config=config)) + second: Final = stack.enter_context(owned_proxy(gateway, tmp_path / "b", {}, config=config)) + return _CallbackRig(first, second) + + +def _one_callback_proxy(gateway: Gateway, tmp_path: Path, stack: ExitStack) -> Gateway: + config: Final = _callback_config(tmp_path) + return stack.enter_context(owned_proxy(gateway, tmp_path / "a", {}, config=config)) + + +@pytest.mark.parametrize("block_route", ("/customer/block", "/end_user/block")) +def test_block_route_marks_the_customer_blocked_and_update_unblocks(gateway: Gateway, block_route: str) -> None: + with gateway.scenario() as scenario: + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + created: Final = gateway.request("POST", "/customer/new", {"user_id": end_user, "max_budget": 100.0}) + assert created.status_code == 200, created.text + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + + gateway.request("POST", block_route, {"user_ids": [end_user]}) + assert eventually(lambda: _blocked_row(end_user), lambda value: value == [{"blocked": True}], seconds=15) == [ + {"blocked": True} + ] + info: Final = gateway.get("/customer/info", {"end_user_id": end_user}) + assert _BlockedCustomerInfo.model_validate(info).blocked is True + + updated: Final = gateway.request("POST", "/customer/update", {"user_id": end_user, "blocked": False}) + assert updated.status_code == 200, updated.text + assert eventually(lambda: _blocked_row(end_user), lambda value: value == [{"blocked": False}], seconds=15) == [ + {"blocked": False} + ] + assert ( + _BlockedCustomerInfo.model_validate(gateway.get("/customer/info", {"end_user_id": end_user})).blocked + is False + ) + + +def _assert_unseen_customer_is_blocked(gateway: Gateway, block_route: str) -> None: + with gateway.scenario() as scenario: + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + + blocked: Final = gateway.request("POST", block_route, {"user_ids": [end_user]}) + assert blocked.status_code == 200, blocked.text + blocked_response: Final = _BlockUsersResponse.model_validate_json(blocked.content) + assert blocked_response == _BlockUsersResponse( + blocked_users=[_BlockedCustomerInfo(user_id=end_user, blocked=True)] + ) + assert eventually(lambda: _blocked_row(end_user), lambda value: value == [{"blocked": True}], seconds=15) == [ + {"blocked": True} + ] + info: Final = gateway.request("GET", "/customer/info", params={"end_user_id": end_user}) + assert info.status_code == 200, info.text + assert _BlockedCustomerInfo.model_validate(info.json()).blocked is True + + updated: Final = gateway.request("POST", "/customer/update", {"user_id": end_user, "blocked": False}) + assert updated.status_code == 200, updated.text + assert eventually(lambda: _blocked_row(end_user), lambda value: value == [{"blocked": False}], seconds=15) == [ + {"blocked": False} + ] + info_after_update: Final = gateway.request("GET", "/customer/info", params={"end_user_id": end_user}) + assert info_after_update.status_code == 200, info_after_update.text + assert _BlockedCustomerInfo.model_validate(info_after_update.json()).blocked is False + + +def test_end_user_block_route_creates_and_blocks_an_unseen_customer(gateway: Gateway) -> None: + _assert_unseen_customer_is_blocked(gateway, "/end_user/block") + + +def test_customer_block_route_creates_and_blocks_an_unseen_customer(gateway: Gateway) -> None: + pytest.skip("BUG: POST /customer/block for an unseen customer returns 500 instead of 200 BlockUsersResponse") + _assert_unseen_customer_is_blocked(gateway, "/customer/block") + + +def test_listed_end_user_is_refused_and_unblock_restores_serving(gateway: Gateway, tmp_path: Path) -> None: + with ( + ExitStack() as stack, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + config: Final = _callback_config(tmp_path, blocked_user_list=[end_user]) + proxy: Final = stack.enter_context(owned_proxy(gateway, tmp_path / "listed", {}, config=config)) + with proxy.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + key: Final = scenario.key(models=[model]) + assert _observed(upstream) == [] + + refused: Final = _chat(proxy, model, key, end_user) + _assert_blocked_refusal(refused, end_user) + assert _observed(upstream) == [] + + unblocked: Final = proxy.request("POST", "/customer/unblock", {"user_ids": [end_user]}) + assert unblocked.status_code == 200, unblocked.text + assert unblocked.json() == {"blocked_users": []}, unblocked.text + assert _blocked_row(end_user) == [] + + text: Final = f"end user unblocked {uuid.uuid4().hex}" + served: Final = proxy.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": text}], "user": end_user}, + key=key, + ) + assert served.status_code == 200, served.text + assert _observed(upstream) == [_observation(text, end_user)] + + +@pytest.mark.parametrize("block_route", ("/customer/block", "/end_user/block")) +def test_blocked_end_user_is_refused_on_every_proxy_and_unblock_restores( + gateway: Gateway, tmp_path: Path, block_route: str +) -> None: + pytest.skip("BUG: blocked end-user chat requests return HTTP 200 instead of HTTP 400") + with ( + ExitStack() as stack, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): + rig: Final = _two_callback_proxies(gateway, tmp_path, stack) + with rig.first.scenario() as scenario: + _observed(upstream) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + created: Final = rig.first.request("POST", "/customer/new", {"user_id": end_user, "max_budget": 100.0}) + assert created.status_code == 200, created.text + scenario.cleanups.callback(rig.first.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + key: Final = scenario.key(models=[model]) + served: Final = _chat(rig.first, model, key, end_user) + assert served.status_code == 200, served.text + _observed(upstream) + + rig.first.request("POST", block_route, {"user_ids": [end_user]}) + assert eventually( + lambda: _blocked_row(end_user), lambda value: value == [{"blocked": True}], seconds=15 + ) == [{"blocked": True}] + info: Final = rig.first.get("/customer/info", {"end_user_id": end_user}) + assert _BlockedCustomerInfo.model_validate(info).blocked is True + + denied_first: Final = _chat(rig.first, model, key, end_user) + denied_second: Final = _chat(rig.second, model, key, end_user) + for denied in (denied_first, denied_second): + _assert_blocked_refusal(denied, end_user) + upstream_calls: Final = _observed(upstream) + assert upstream_calls == [], upstream_calls + + updated: Final = rig.first.request("POST", "/customer/update", {"user_id": end_user, "blocked": False}) + assert updated.status_code == 200, updated.text + assert eventually( + lambda: _blocked_row(end_user), lambda value: value == [{"blocked": False}], seconds=15 + ) == [{"blocked": False}] + assert ( + _BlockedCustomerInfo.model_validate(rig.first.get("/customer/info", {"end_user_id": end_user})).blocked + is False + ) + + text: Final = f"end user unblocked {uuid.uuid4().hex}" + + def call(candidate: Gateway) -> httpx.Response: + return candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": text}], "user": end_user}, + key=key, + ) + + assert ( + eventually( + lambda: call(rig.first), lambda response: response.status_code == 200, seconds=10 + ).status_code + == 200 + ) + assert ( + eventually( + lambda: call(rig.second), lambda response: response.status_code == 200, seconds=10 + ).status_code + == 200 + ) + expected_observations: Final = [ + { + "path": "/v1/chat/completions", + "authorization": "Bearer integration-provider-key", + "body": { + "messages": [{"role": "user", "content": text}], + "model": "gpt-4o-mini", + "user": end_user, + }, + "method": "POST", + "api_key": "", + } + ] * 2 + assert _observed(upstream) == expected_observations + + +def test_unblock_route_requires_the_blocked_user_check_callback(gateway: Gateway, tmp_path: Path) -> None: + with ExitStack() as stack: + callback_proxy: Final = _one_callback_proxy(gateway, tmp_path, stack) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + created: Final = callback_proxy.request( + "POST", "/customer/new", {"user_id": end_user, "max_budget": 100.0, "blocked": True} + ) + assert created.status_code == 200, created.text + stack.callback(callback_proxy.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + assert _blocked_row(end_user) == [{"blocked": True}] + + shared: Final = gateway.request("POST", "/customer/unblock", {"user_ids": [end_user]}) + assert shared.status_code == 400, shared.text + assert shared.json() == {"detail": {"error": "Blocked user check was never set. This call has no effect."}}, ( + shared.text + ) + assert _blocked_row(end_user) == [{"blocked": True}] + + callback_response: Final = callback_proxy.request("POST", "/customer/unblock", {"user_ids": [end_user]}) + assert callback_response.status_code == 400, callback_response.text + assert callback_response.json() == { + "detail": {"error": "Blocked user check was never set. This call has no effect."} + }, callback_response.text + assert _blocked_row(end_user) == [{"blocked": True}] + + +def test_block_route_returns_the_blocked_records(gateway: Gateway) -> None: + pytest.skip("BUG: POST /customer/block returns 500 instead of 200 BlockUsersResponse") + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + with gateway.scenario() as scenario: + created: Final = gateway.request("POST", "/customer/new", {"user_id": end_user}) + assert created.status_code == 200, created.text + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + blocked: Final = gateway.request("POST", "/customer/block", {"user_ids": [end_user]}) + assert blocked.status_code == 200, blocked.text + payload: Final = object_value(blocked.json()) + blocked_users: Final = payload["blocked_users"] + assert isinstance(blocked_users, list) and len(blocked_users) == 1 + assert object_value(blocked_users[0])["user_id"] == end_user + assert object_value(blocked_users[0])["blocked"] is True + + +def test_blocked_end_user_is_refused_without_the_callback(gateway: Gateway) -> None: + pytest.skip("BUG: blocked end-user request without callback returns 200 instead of 400") + with gateway.scenario() as scenario: + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + created: Final = gateway.request("POST", "/customer/new", {"user_id": end_user, "blocked": True}) + assert created.status_code == 200, created.text + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + key: Final = scenario.key(models=[model]) + refused: Final = _chat(gateway, model, key, end_user) + _assert_blocked_refusal(refused, end_user) + + +def test_blocked_end_user_header_spelling_is_refused_with_the_callback(gateway: Gateway, tmp_path: Path) -> None: + pytest.skip("BUG: x-litellm-end-user-id request for a blocked customer returns 200 instead of 400") + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + config: Final = _callback_config(tmp_path, blocked_user_list=[end_user]) + with ( + ExitStack() as stack, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): + proxy: Final = stack.enter_context(owned_proxy(gateway, tmp_path / "listed", {}, config=config)) + with proxy.scenario() as scenario: + created: Final = proxy.request("POST", "/customer/new", {"user_id": end_user, "max_budget": 100.0}) + assert created.status_code == 200, created.text + scenario.cleanups.callback(proxy.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + key: Final = scenario.key(models=[model]) + refused: Final = proxy.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"blocked {uuid.uuid4().hex}"}]}, + key=key, + headers={"x-litellm-end-user-id": end_user}, + ) + _assert_blocked_refusal(refused, end_user) + assert _observed(upstream) == [] diff --git a/tests/integration/authorization/test_user_self_update_guard.py b/tests/integration/authorization/test_user_self_update_guard.py new file mode 100644 index 00000000000..d83960e1e85 --- /dev/null +++ b/tests/integration/authorization/test_user_self_update_guard.py @@ -0,0 +1,287 @@ +import os +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway, object_value +from integration._support.database import read_rows +from pydantic import BaseModel, JsonValue + +from litellm.models.user import LiteLLM_UserTable +from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken + +_SELF_SERVICE_ROUTES: Final = ("/user/update", "/user/bulk_update", "/user/new", "/user/info") + + +class _UserInfo(BaseModel): + user_id: str + user_alias: str | None + user_role: str + max_budget: float | None + + +def _user_row(user_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + "SELECT user_id, user_alias, user_role, max_budget, model_max_budget, spend, object_permission_id, user_email" + ' FROM "LiteLLM_UserTable" WHERE user_id=%s', + (user_id,), + ) + + +def _user_info(gateway: Gateway, user_id: str) -> _UserInfo: + payload: Final = object_value(gateway.get("/user/info", {"user_id": user_id})) + return _UserInfo.model_validate(payload["user_info"]) + + +def _escalation_error(field: str) -> dict[str, JsonValue]: + return { + "error": { + "message": f"{{'error': \"Non-admin users cannot modify '{field}' on their own record. Contact your proxy admin.\"}}", + "type": "auth_error", + "param": "None", + "code": "403", + } + } + + +def _role_change_error() -> dict[str, JsonValue]: + return { + "error": { + "message": "Only proxy admins can modify user roles.", + "type": "auth_error", + "param": "None", + "code": "403", + } + } + + +def _ui_session_route_error(user_id: str) -> dict[str, JsonValue]: + masked_user_id: Final = f"{user_id[:6]}{'*' * (len(user_id) - 8)}{user_id[-2:]}" + return { + "error": { + "message": ( + "Authentication Error, Only proxy admin can be used to generate, delete, update info for new " + "keys/users/teams. Route=/user/update. Your role=internal_user. " + f"Your user_id={masked_user_id}" + ), + "type": "auth_error", + "param": "None", + "code": "401", + } + } + + +def test_internal_user_can_update_allowed_fields_on_own_record(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + user_id: Final = scenario.user(user_role="internal_user", max_budget=10.0) + key: Final = scenario.key(user_id=user_id, allowed_routes=list(_SELF_SERVICE_ROUTES)) + alias: Final = f"self-alias-{uuid.uuid4().hex}" + updated: Final = gateway.request("POST", "/user/update", {"user_id": user_id, "user_alias": alias}, key=key) + assert updated.status_code == 200, updated.text + assert object_value(object_value(updated.json())["data"])["user_alias"] == alias + assert _user_row(user_id)[0]["user_alias"] == alias + assert _user_info(gateway, user_id).user_alias == alias + + +@pytest.mark.parametrize("user_role", ("proxy_admin", "proxy_admin_viewer")) +def test_internal_user_cannot_promote_own_role(gateway: Gateway, user_role: str) -> None: + with gateway.scenario() as scenario: + user_id: Final = scenario.user(user_role="internal_user", max_budget=10.0) + key: Final = scenario.key(user_id=user_id, allowed_routes=list(_SELF_SERVICE_ROUTES)) + before_row: Final = _user_row(user_id) + before_info: Final = _user_info(gateway, user_id) + + refused: Final = gateway.request( + "POST", + "/user/update", + {"user_id": user_id, "user_role": user_role}, + key=key, + ) + assert refused.status_code == 403, refused.text + assert refused.json() == _role_change_error(), refused.text + assert _user_row(user_id) == before_row + assert _user_info(gateway, user_id) == before_info + + +def test_internal_user_can_update_allowed_fields_by_own_email(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + user_email: Final = f"integration-{uuid.uuid4().hex}@example.com" + user_id: Final = scenario.user( + user_email=user_email, + user_role="internal_user", + max_budget=10.0, + ) + key: Final = scenario.key(user_id=user_id, allowed_routes=list(_SELF_SERVICE_ROUTES)) + alias: Final = f"self-email-alias-{uuid.uuid4().hex}" + + updated: Final = gateway.request( + "POST", + "/user/update", + {"user_email": user_email, "user_alias": alias}, + key=key, + ) + assert updated.status_code == 200, updated.text + updated_data: Final = object_value(object_value(updated.json())["data"]) + assert updated_data["user_alias"] == alias + assert updated_data["user_id"] == user_id + assert _user_row(user_id)[0]["user_alias"] == alias + assert _user_info(gateway, user_id).user_alias == alias + + +def test_internal_user_cannot_escalate_protected_fields_by_own_email(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + user_email: Final = f"integration-{uuid.uuid4().hex}@example.com" + user_id: Final = scenario.user( + user_email=user_email, + user_role="internal_user", + max_budget=10.0, + ) + key: Final = scenario.key(user_id=user_id, allowed_routes=list(_SELF_SERVICE_ROUTES)) + before_row: Final = _user_row(user_id) + before_info: Final = _user_info(gateway, user_id) + + refused: Final = gateway.request( + "POST", + "/user/update", + {"user_email": user_email, "max_budget": 1000.0}, + key=key, + ) + assert refused.status_code == 403, refused.text + assert refused.json() == _escalation_error("max_budget"), refused.text + assert _user_row(user_id) == before_row + assert _user_info(gateway, user_id) == before_info + + +@pytest.mark.parametrize( + ("field", "value"), + (("max_budget", 1000), ("object_permission", {})), + ids=("max_budget", "object_permission"), +) +def test_ui_session_token_cannot_update_own_protected_fields( + gateway: Gateway, + monkeypatch: pytest.MonkeyPatch, + field: str, + value: JsonValue, +) -> None: + monkeypatch.setenv("LITELLM_SALT_KEY", os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt")) + with gateway.scenario() as scenario: + user_id: Final = scenario.user(user_role="internal_user", max_budget=10.0) + token: Final = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( + LiteLLM_UserTable(user_id=user_id, user_role="internal_user", models=[]) + ) + self_info: Final = gateway.request("GET", "/user/info", params={"user_id": user_id}, key=token) + assert self_info.status_code == 200, self_info.text + assert object_value(object_value(self_info.json())["user_info"])["user_id"] == user_id, self_info.text + before_row: Final = _user_row(user_id) + before_info: Final = _user_info(gateway, user_id) + + refused: Final = gateway.request("POST", "/user/update", {"user_id": user_id, field: value}, key=token) + assert refused.status_code == 401, refused.text + assert refused.json() == _ui_session_route_error(user_id), refused.text + assert _user_row(user_id) == before_row + assert _user_info(gateway, user_id) == before_info + + +@pytest.mark.parametrize( + ("field", "value"), + ( + ("max_budget", 1000), + ("max_budget", "1000"), + ("model_max_budget", {"gpt-4o": 5.0}), + ("spend", 0), + ("object_permission", {"vector_stores": ["vs-1"]}), + ("object_permission", {}), + ), + ids=( + "max_budget_number", + "max_budget_string", + "model_max_budget", + "spend", + "object_permission", + "object_permission_empty", + ), +) +def test_internal_user_cannot_escalate_protected_fields(gateway: Gateway, field: str, value: JsonValue) -> None: + with gateway.scenario() as scenario: + user_id: Final = scenario.user(user_role="internal_user", max_budget=10.0) + key: Final = scenario.key(user_id=user_id, allowed_routes=list(_SELF_SERVICE_ROUTES)) + before_row: Final = _user_row(user_id) + before_info: Final = _user_info(gateway, user_id) + + refused: Final = gateway.request("POST", "/user/update", {"user_id": user_id, field: value}, key=key) + assert refused.status_code == 403, refused.text + assert refused.json() == _escalation_error(field), refused.text + assert _user_row(user_id) == before_row + assert _user_info(gateway, user_id) == before_info + + +def test_bulk_update_self_escalation_fails_per_user(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + user_id: Final = scenario.user(user_role="internal_user", max_budget=10.0) + key: Final = scenario.key(user_id=user_id, allowed_routes=list(_SELF_SERVICE_ROUTES)) + before_row: Final = _user_row(user_id) + + refused: Final = gateway.request( + "POST", + "/user/bulk_update", + {"users": [{"user_id": user_id, "max_budget": 1000}]}, + key=key, + ) + assert refused.status_code == 200, refused.text + assert refused.json() == { + "results": [ + { + "user_id": user_id, + "user_email": None, + "success": False, + "error": "403: {'error': \"Non-admin users cannot modify 'max_budget' on their own record. Contact your proxy admin.\"}", + "updated_user": None, + } + ], + "total_requested": 1, + "successful_updates": 0, + "failed_updates": 1, + }, refused.text + assert _user_row(user_id) == before_row + + +def test_bulk_update_user_role_requires_proxy_admin(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + user_id: Final = scenario.user(user_role="internal_user", max_budget=10.0) + key: Final = scenario.key(user_id=user_id, allowed_routes=list(_SELF_SERVICE_ROUTES)) + before_row: Final = _user_row(user_id) + + refused: Final = gateway.request( + "POST", + "/user/bulk_update", + {"users": [{"user_id": user_id, "user_role": "proxy_admin"}]}, + key=key, + ) + assert refused.status_code == 403, refused.text + assert refused.json() == {"detail": "Only proxy admins can modify user roles."}, refused.text + assert _user_row(user_id) == before_row + + +@pytest.mark.parametrize("role", ("proxy_admin", "proxy_admin_viewer")) +def test_new_user_with_admin_role_requires_proxy_admin(gateway: Gateway, role: str) -> None: + with gateway.scenario() as scenario: + user_id: Final = scenario.user(user_role="internal_user", max_budget=10.0) + key: Final = scenario.key(user_id=user_id, allowed_routes=list(_SELF_SERVICE_ROUTES)) + attempted_id: Final = f"integration-{uuid.uuid4().hex}" + + refused: Final = gateway.request( + "POST", + "/user/new", + {"user_id": attempted_id, "user_role": role, "auto_create_key": False}, + key=key, + ) + assert refused.status_code == 403, refused.text + assert refused.json() == { + "error": { + "message": f"Only proxy admins can create administrative users (proxy_admin, proxy_admin_viewer). Attempted to create user with role: {role}. Your role: internal_user", + "type": "internal_server_error", + "param": None, + "code": "403", + } + }, refused.text + assert _user_row(attempted_id) == [] diff --git a/tests/integration/management/test_organization_lifecycle.py b/tests/integration/management/test_organization_lifecycle.py index 7b8d1d7f7e7..0193bf517b2 100644 --- a/tests/integration/management/test_organization_lifecycle.py +++ b/tests/integration/management/test_organization_lifecycle.py @@ -1,9 +1,10 @@ import uuid from concurrent.futures import ThreadPoolExecutor +from hashlib import sha256 from typing import Final import httpx -from integration._support.client import Gateway, object_value, string_value +from integration._support.client import Gateway, eventually, object_value, string_value from integration._support.database import read_rows from pydantic import JsonValue @@ -87,3 +88,150 @@ def test_member_role_update_and_removal_persist_and_read_back(gateway: Gateway) assert removed.status_code == 200, removed.text assert _membership_rows(organization_id) == [] assert _listed(gateway, organization_id)["members"] == [] + + +def _observed(upstream: httpx.Client) -> list[dict[str, JsonValue]]: + response: Final = upstream.get("/__observations") + response.raise_for_status() + requests: Final = response.json()["requests"] + assert isinstance(requests, list) + return requests + + +def _chat(proxy: Gateway, model: str, key: str, text: str) -> httpx.Response: + return proxy.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": text}]}, + key=key, + ) + + +def _observation(text: str) -> dict[str, JsonValue]: + return { + "path": "/v1/chat/completions", + "authorization": "Bearer integration-provider-key", + "body": {"messages": [{"role": "user", "content": text}], "model": "gpt-4o-mini"}, + "method": "POST", + "api_key": "", + } + + +def _token_not_found_error(key: str) -> dict[str, JsonValue]: + hashed: Final = sha256(key.encode()).hexdigest() + return { + "error": { + "message": ( + "Authentication Error, Invalid proxy server token passed. " + f"Received API Key = sk-...{key[-4:]}, Key Hash (Token) ={hashed}. " + "Unable to find token in cache or `LiteLLM_VerificationTokenTable`" + ), + "type": "token_not_found_in_db", + "param": "key", + "code": "401", + } + } + + +def _new_key(gateway: Gateway, **fields: JsonValue) -> str: + created: Final = gateway.post("/key/generate", fields) + return string_value(created["key"]) + + +def _key_rows(*hashes: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT token FROM "LiteLLM_VerificationToken" WHERE token = ANY(%s)', + (list(hashes),), + ) + + +def test_organization_delete_revokes_org_and_team_keys_on_gateway_and_peer(gateway: Gateway, peer: 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() + created_org: Final = gateway.post( + "/organization/new", {"organization_alias": f"integration-{uuid.uuid4().hex}"} + ) + scenario.cleanups.callback(scenario.delete_budget, string_value(created_org["budget_id"])) + organization_id: Final = string_value(created_org["organization_id"]) + other_organization_id: Final = scenario.organization() + member: Final = scenario.user() + membership: Final = gateway.request( + "POST", + "/organization/member_add", + {"organization_id": organization_id, "member": {"role": "internal_user", "user_id": member}}, + ) + assert membership.status_code == 200, membership.text + assert _membership_rows(organization_id) == [{"user_id": member, "user_role": "internal_user"}] + team_id: Final = string_value( + gateway.post( + "/team/new", + {"team_alias": f"integration-{uuid.uuid4().hex}", "organization_id": organization_id}, + )["team_id"] + ) + organization_key: Final = _new_key(gateway, organization_id=organization_id) + team_key: Final = _new_key(gateway, team_id=team_id) + other_key: Final = _new_key(gateway, organization_id=other_organization_id) + warmed_info: Final = peer.request("GET", "/organization/info", params={"organization_id": organization_id}) + assert warmed_info.status_code == 200, warmed_info.text + + assert _observed(upstream) == [] + for proxy in (gateway, peer): + for key in (organization_key, team_key, other_key): + text: Final = f"warm {uuid.uuid4().hex}" + response: Final = _chat(proxy, model, key, text) + assert response.status_code == 200, response.text + assert _observed(upstream) == [_observation(text)] + + deleted: Final = gateway.request("DELETE", "/organization/delete", {"organization_ids": [organization_id]}) + assert deleted.status_code == 200, deleted.text + + for key in (organization_key, team_key): + refused: Final = _chat(gateway, model, key, f"refused {uuid.uuid4().hex}") + assert refused.status_code == 401, refused.text + assert refused.json() == _token_not_found_error(key), refused.text + for key in (organization_key, team_key): + refused_peer: Final = eventually( + lambda k=key: _chat(peer, model, k, f"refused {uuid.uuid4().hex}"), + lambda response: response.status_code == 401, + seconds=10, + ) + assert refused_peer.json() == _token_not_found_error(key), refused_peer.text + refused_observations: Final = _observed(upstream) + assert refused_observations == [], refused_observations + for proxy in (gateway, peer): + for key in (organization_key, team_key): + refused_again: Final = _chat(proxy, model, key, f"refused {uuid.uuid4().hex}") + assert refused_again.status_code == 401, refused_again.text + assert _observed(upstream) == [] + + for proxy in (gateway, peer): + text: Final = f"served {uuid.uuid4().hex}" + response: Final = _chat(proxy, model, other_key, text) + assert response.status_code == 200, response.text + assert _observed(upstream) == [_observation(text)] + + hashed: Final = tuple(sha256(key.encode()).hexdigest() for key in (organization_key, team_key)) + assert _key_rows(*hashed) == [] + assert read_rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team_id,)) == [] + assert _membership_rows(organization_id) == [] + assert ( + read_rows( + 'SELECT organization_id FROM "LiteLLM_OrganizationTable" WHERE organization_id = %s', + (organization_id,), + ) + == [] + ) + info: Final = gateway.request("GET", "/organization/info", params={"organization_id": organization_id}) + assert info.status_code == 404, info.text + assert info.json() == {"detail": {"error": "Organization not found"}}, info.text + peer_info: Final = eventually( + lambda: peer.request("GET", "/organization/info", params={"organization_id": organization_id}), + lambda response: response.status_code == 404, + seconds=10, + ) + assert peer_info.json() == {"detail": {"error": "Organization not found"}}, peer_info.text + alive: Final = gateway.request("GET", "/organization/info", params={"organization_id": other_organization_id}) + assert alive.status_code == 200, alive.text diff --git a/tests/integration/management/test_user_delete_revokes_keys.py b/tests/integration/management/test_user_delete_revokes_keys.py new file mode 100644 index 00000000000..0cf2db9c3e6 --- /dev/null +++ b/tests/integration/management/test_user_delete_revokes_keys.py @@ -0,0 +1,394 @@ +import os +import uuid +from hashlib import sha256 +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, eventually, object_value, string_value +from integration._support.database import read_rows +from pydantic import JsonValue + +from litellm.models.user import LiteLLM_UserTable +from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken + + +def _observed(upstream: httpx.Client) -> list[dict[str, JsonValue]]: + response: Final = upstream.get("/__observations") + response.raise_for_status() + requests: Final = response.json()["requests"] + assert isinstance(requests, list) + return requests + + +def _chat(proxy: Gateway, model: str, key: str, text: str) -> httpx.Response: + return proxy.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": text}]}, + key=key, + ) + + +def _observation(text: str) -> dict[str, JsonValue]: + return { + "path": "/v1/chat/completions", + "authorization": "Bearer integration-provider-key", + "body": {"messages": [{"role": "user", "content": text}], "model": "gpt-4o-mini"}, + "method": "POST", + "api_key": "", + } + + +def _token_not_found_error(key: str) -> dict[str, JsonValue]: + hashed: Final = sha256(key.encode()).hexdigest() + return { + "error": { + "message": ( + "Authentication Error, Invalid proxy server token passed. " + f"Received API Key = sk-...{key[-4:]}, Key Hash (Token) ={hashed}. " + "Unable to find token in cache or `LiteLLM_VerificationTokenTable`" + ), + "type": "token_not_found_in_db", + "param": "key", + "code": "401", + } + } + + +def _blocked_key_error() -> dict[str, JsonValue]: + return { + "error": { + "message": "Authentication Error, Key is blocked. Update via `/key/unblock` if you're an admin.", + "type": "auth_error", + "param": "None", + "code": "401", + } + } + + +def _user_rows(user_id: str) -> dict[str, list[dict[str, JsonValue]]]: + return { + "user": read_rows('SELECT user_id FROM "LiteLLM_UserTable" WHERE user_id=%s', (user_id,)), + "keys": read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE user_id=%s', (user_id,)), + "team_memberships": read_rows('SELECT user_id FROM "LiteLLM_TeamMembership" WHERE user_id=%s', (user_id,)), + "org_memberships": read_rows( + 'SELECT user_id FROM "LiteLLM_OrganizationMembership" WHERE user_id=%s', (user_id,) + ), + "invitations": read_rows('SELECT id FROM "LiteLLM_InvitationLink" WHERE user_id=%s', (user_id,)), + } + + +def _new_user(gateway: Gateway, **fields: JsonValue) -> str: + created: Final = gateway.post( + "/user/new", {"user_id": f"integration-{uuid.uuid4().hex}", "auto_create_key": False, **fields} + ) + return string_value(created["user_id"]) + + +def _new_key(gateway: Gateway, **fields: JsonValue) -> str: + created: Final = gateway.post("/key/generate", fields) + return string_value(created["key"]) + + +def test_user_delete_revokes_keys_and_memberships(gateway: Gateway, peer: 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() + team_id: Final = scenario.team() + organization_id: Final = scenario.organization() + user_id: Final = _new_user(gateway) + gateway.post("/team/member_add", {"team_id": team_id, "member": {"role": "user", "user_id": user_id}}) + gateway.post( + "/organization/member_add", + {"organization_id": organization_id, "member": {"role": "internal_user", "user_id": user_id}}, + ) + keys: Final = ( + _new_key(gateway, user_id=user_id, models=[model]), + _new_key(gateway, user_id=user_id, models=[model]), + ) + invitation: Final = gateway.post("/invitation/new", {"user_id": user_id}) + assert string_value(object_value(invitation)["user_id"]) == user_id + + assert _observed(upstream) == [] + expected: Final = [] + for proxy in (gateway, peer): + for key in keys: + text: Final = f"warm {uuid.uuid4().hex}" + assert _chat(proxy, model, key, text).status_code == 200 + expected.append(_observation(text)) + assert _observed(upstream) == expected + + deleted: Final = gateway.request("POST", "/user/delete", {"user_ids": [user_id]}) + assert deleted.status_code == 200, deleted.text + assert deleted.json() == 1 + + for key in keys: + denied: Final = _chat(gateway, model, key, f"denied {uuid.uuid4().hex}") + assert denied.status_code == 401, denied.text + assert denied.json() == _token_not_found_error(key), denied.text + for key in keys: + denied_peer: Final = eventually( + lambda k=key: _chat(peer, model, k, f"denied {uuid.uuid4().hex}"), + lambda response: response.status_code == 401, + seconds=10, + ) + assert denied_peer.json() == _token_not_found_error(key), denied_peer.text + + denied_observations: Final = _observed(upstream) + assert denied_observations == [], denied_observations + for proxy in (gateway, peer): + for key in keys: + denied_again: Final = _chat(proxy, model, key, f"denied {uuid.uuid4().hex}") + assert denied_again.status_code == 401, denied_again.text + assert _observed(upstream) == [] + + assert _user_rows(user_id) == { + "user": [], + "keys": [], + "team_memberships": [], + "org_memberships": [], + "invitations": [], + } + team_row: Final = read_rows('SELECT members_with_roles FROM "LiteLLM_TeamTable" WHERE team_id=%s', (team_id,)) + assert team_row == [ + {"members_with_roles": [{"role": "admin", "user_email": None, "user_id": "default_user_id"}]} + ] + team_info: Final = object_value(object_value(gateway.get("/team/info", {"team_id": team_id}))["team_info"]) + assert team_info["members_with_roles"] == [ + {"role": "admin", "user_alias": None, "user_id": "default_user_id", "user_email": None} + ] + + info: Final = gateway.request("GET", "/user/info", params={"user_id": user_id}) + assert info.status_code == 404, info.text + assert info.json() == { + "error": { + "message": f"User {user_id} not found", + "type": "internal_server_error", + "param": None, + "code": "404", + } + }, info.text + + +def test_org_admin_can_only_delete_users_inside_their_orgs( + 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() + org_a: Final = scenario.organization() + org_b: Final = scenario.organization() + admin: Final = scenario.org_member(org_a, "org_admin") + admin_key: Final = scenario.key(user_id=admin) + + victim_b: Final = scenario.org_member(org_b, "internal_user") + victim_b_key: Final = scenario.key(user_id=victim_b, models=[model]) + victim_overlap: Final = scenario.org_member(org_a, "internal_user") + gateway.post( + "/organization/member_add", + {"organization_id": org_b, "member": {"role": "internal_user", "user_id": victim_overlap}}, + ) + victim_no_org: Final = scenario.user(user_role="internal_user") + victim_a: Final = _new_user(gateway) + gateway.post( + "/organization/member_add", + {"organization_id": org_a, "member": {"role": "internal_user", "user_id": victim_a}}, + ) + + assert _observed(upstream) == [] + refused: Final = gateway.request( + "POST", + "/user/delete", + {"user_ids": [victim_b], "organization_id": org_a}, + key=admin_key, + ) + assert refused.status_code == 403, refused.text + assert refused.json() == { + "detail": { + "error": ( + f"User {victim_b} is not within your admin scope. " + "Only PROXY_ADMIN may delete users outside your administered organizations." + ) + } + }, refused.text + assert _user_rows(victim_b)["user"] != [] + + refused_overlap: Final = gateway.request( + "POST", + "/user/delete", + {"user_ids": [victim_overlap], "organization_id": org_a}, + key=admin_key, + ) + assert refused_overlap.status_code == 403, refused_overlap.text + assert refused_overlap.json() == { + "detail": { + "error": ( + f"User {victim_overlap} is not within your admin scope. " + "Only PROXY_ADMIN may delete users outside your administered organizations." + ) + } + }, refused_overlap.text + assert _user_rows(victim_overlap)["user"] != [] + + refused_no_org: Final = gateway.request( + "POST", + "/user/delete", + {"user_ids": [victim_no_org], "organization_id": org_a}, + key=admin_key, + ) + assert refused_no_org.status_code == 403, refused_no_org.text + assert refused_no_org.json() == { + "detail": { + "error": ( + f"User {victim_no_org} is not within your admin scope. " + "Only PROXY_ADMIN may delete users outside your administered organizations." + ) + } + }, refused_no_org.text + assert _user_rows(victim_no_org)["user"] != [] + + alive: Final = f"still alive {uuid.uuid4().hex}" + served: Final = _chat(gateway, model, victim_b_key, alive) + assert served.status_code == 200, served.text + assert _observed(upstream) == [_observation(alive)] + + deleted: Final = gateway.request( + "POST", + "/user/delete", + {"user_ids": [victim_a], "organization_id": org_a}, + key=admin_key, + ) + assert deleted.status_code == 200, deleted.text + assert deleted.json() == 1 + assert _user_rows(victim_a)["user"] == [] + + +def test_scim_user_delete_revokes_keys(gateway: Gateway, peer: Gateway, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_SALT_KEY", os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt")) + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): + model: Final = scenario.model() + user_id: Final = _new_user(gateway) + remaining_user_id: Final = _new_user(gateway) + scenario.cleanups.callback(gateway.request, "POST", "/user/delete", {"user_ids": [remaining_user_id]}) + ui_token: Final = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( + LiteLLM_UserTable(user_id=user_id, user_role="internal_user", models=[]) + ) + team_id: Final = scenario.team( + members_with_roles=[ + {"role": "user", "user_id": user_id}, + {"role": "user", "user_id": remaining_user_id}, + ] + ) + initial_team_rows: Final = read_rows( + 'SELECT members_with_roles FROM "LiteLLM_TeamTable" WHERE team_id=%s', + (team_id,), + ) + assert len(initial_team_rows) == 1, initial_team_rows + initial_members: Final = initial_team_rows[0]["members_with_roles"] + assert isinstance(initial_members, list), initial_team_rows + initial_member_id_roles: Final = tuple( + ( + string_value(object_value(member)["user_id"]), + string_value(object_value(member)["role"]), + ) + for member in initial_members + ) + assert (user_id, "user") in initial_member_id_roles, initial_team_rows + assert (remaining_user_id, "user") in initial_member_id_roles, initial_team_rows + expected_remaining_members: Final = tuple(member for member in initial_member_id_roles if member[0] != user_id) + organization_id: Final = scenario.organization() + membership: Final = gateway.request( + "POST", + "/organization/member_add", + {"organization_id": organization_id, "member": {"role": "internal_user", "user_id": user_id}}, + ) + assert membership.status_code == 200, membership.text + invitation: Final = gateway.post("/invitation/new", {"user_id": user_id}) + assert string_value(object_value(invitation)["user_id"]) == user_id + key: Final = _new_key(gateway, user_id=user_id, models=[model]) + scenario.cleanups.callback(gateway.request, "POST", "/key/delete", {"keys": [key]}) + hashed: Final = sha256(key.encode()).hexdigest() + + gateway_info: Final = gateway.request("GET", "/user/info", params={"user_id": user_id}) + assert gateway_info.status_code == 200, gateway_info.text + peer_info: Final = peer.request("GET", "/user/info", params={"user_id": user_id}) + assert peer_info.status_code == 200, peer_info.text + for proxy in (gateway, peer): + session_info: Final = proxy.request("GET", "/user/info", params={"user_id": user_id}, key=ui_token) + assert session_info.status_code == 200, session_info.text + assert object_value(object_value(session_info.json())["user_info"])["user_id"] == user_id, session_info.text + assert _observed(upstream) == [] + + deleted: Final = gateway.request("DELETE", f"/scim/v2/Users/{user_id}") + assert deleted.status_code == 204, deleted.text + + expected_info: Final = { + "error": { + "message": f"User {user_id} not found", + "type": "internal_server_error", + "param": None, + "code": "404", + } + } + denied_info: Final = gateway.request("GET", "/user/info", params={"user_id": user_id}) + assert denied_info.status_code == 404, denied_info.text + assert denied_info.json() == expected_info, denied_info.text + denied_peer_info: Final = eventually( + lambda: peer.request("GET", "/user/info", params={"user_id": user_id}), + lambda response: response.status_code == 404, + seconds=10, + ) + assert denied_peer_info.json() == expected_info, denied_peer_info.text + for proxy in (gateway, peer): + expired_session: Final = proxy.request("GET", "/user/info", params={"user_id": user_id}, key=ui_token) + assert expired_session.status_code == 404, expired_session.text + assert expired_session.json() == expected_info, expired_session.text + + rows: Final = _user_rows(user_id) + assert rows["user"] == [] + assert rows["team_memberships"] == [] + assert rows["org_memberships"] == [] + assert rows["invitations"] == [] + assert rows["keys"] == [{"token": hashed}] + team_rows: Final = read_rows( + 'SELECT members_with_roles FROM "LiteLLM_TeamTable" WHERE team_id=%s', + (team_id,), + ) + assert len(team_rows) == 1, team_rows + members: Final = team_rows[0]["members_with_roles"] + assert isinstance(members, list), team_rows + member_id_roles: Final = tuple( + ( + string_value(object_value(member)["user_id"]), + string_value(object_value(member)["role"]), + ) + for member in members + ) + assert member_id_roles == expected_remaining_members, team_rows + assert ( + read_rows( + 'SELECT user_id FROM "LiteLLM_OrganizationMembership" WHERE organization_id=%s', + (organization_id,), + ) + == [] + ) + assert read_rows('SELECT blocked FROM "LiteLLM_VerificationToken" WHERE token=%s', (hashed,)) == [ + {"blocked": True} + ] + + denied: Final = _chat(gateway, model, key, f"denied {uuid.uuid4().hex}") + assert denied.status_code == 401, denied.text + assert denied.json() == _blocked_key_error(), denied.text + denied_peer: Final = _chat(peer, model, key, f"denied {uuid.uuid4().hex}") + assert denied_peer.status_code == 401, denied_peer.text + assert denied_peer.json() == _blocked_key_error(), denied_peer.text + assert _observed(upstream) == [] + assert _observed(upstream) == [] diff --git a/tests/integration/management/test_user_updates_wedged_coordination_redis.py b/tests/integration/management/test_user_updates_wedged_coordination_redis.py index 79b1f3166e3..41fefdcd3b2 100644 --- a/tests/integration/management/test_user_updates_wedged_coordination_redis.py +++ b/tests/integration/management/test_user_updates_wedged_coordination_redis.py @@ -1,4 +1,5 @@ import os +import re import time import uuid from collections.abc import Callable, Mapping @@ -7,6 +8,7 @@ from pathlib import Path from typing import Final from urllib.parse import urlsplit, urlunsplit +import httpx import psycopg import pytest from psycopg import sql @@ -16,6 +18,7 @@ from redis.client import PubSub from tests.integration._support.budget import case_timeout_seconds from tests.integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value +from tests.integration._support.database import read_rows from tests.integration._support.process import owned_proxy_process from tests.integration._support.redis_process import owned_redis @@ -345,3 +348,324 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( finally: admin.execute(sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(identity))) record_property("cell_elapsed_seconds", timings) + + +_USER_BUDGET_REFUSAL: Final = re.compile( + r"^ExceededBudget: User=(?P\S+) over budget\. Spend=(?P[0-9.eE+-]+), Budget=(?P[0-9.eE+-]+)$" +) + + +def _observed(upstream: httpx.Client) -> list[dict[str, JsonValue]]: + response: Final = upstream.get("/__observations") + response.raise_for_status() + requests: Final = response.json()["requests"] + assert isinstance(requests, list) + return requests + + +def _user_chat(proxy: Gateway, model: str, key: str, text: str) -> httpx.Response: + return proxy.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": text}]}, + key=key, + ) + + +def _chat_observation(text: str) -> dict[str, JsonValue]: + return { + "path": "/v1/chat/completions", + "authorization": "Bearer integration-provider-key", + "body": {"messages": [{"role": "user", "content": text}], "model": "gpt-4o-mini"}, + "method": "POST", + "api_key": "", + } + + +def _user_row(user_id: str) -> list[dict[str, JsonValue]]: + return read_rows('SELECT models, max_budget, metadata FROM "LiteLLM_UserTable" WHERE user_id=%s', (user_id,)) + + +def _partial_user_row(user_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + "SELECT models, max_budget, metadata, user_alias, rpm_limit, budget_duration, spend " + 'FROM "LiteLLM_UserTable" WHERE user_id=%s', + (user_id,), + ) + + +def _assert_user_budget_refusal(response: httpx.Response, user_id: str) -> None: + assert response.status_code == 422, response.text + assert response.json()["error"]["type"] == "budget_exceeded", response.text + match: Final = _USER_BUDGET_REFUSAL.match(response.json()["error"]["message"]) + assert match is not None, response.text + assert match.group("user") == user_id, response.text + assert float(match.group("spend")) > 0, response.text + + +def _user_model_denied_error(model: str) -> dict[str, JsonValue]: + return { + "error": { + "message": f"The requested model '{model}' is not available for this API key, or the model name is invalid. Check the models available to you and try again.", + "type": "key_model_access_denied", + "param": "model", + "code": "403", + } + } + + +def _model_user(gateway: Gateway, model: str, max_budget: float = 10.0) -> str: + created: Final = gateway.post( + "/user/new", + { + "user_id": f"integration-{uuid.uuid4().hex}", + "user_role": "internal_user", + "auto_create_key": False, + "max_budget": max_budget, + "models": [model], + "metadata": {"dept": "qa"}, + }, + ) + return string_value(created["user_id"]) + + +def _wait_for_spend(user_id: str, spend: float) -> None: + eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=%s', (user_id,)), + lambda rows: bool(rows) and float(rows[0]["spend"]) >= spend, + seconds=70, + ) + + +def test_user_budget_update_propagates_to_gateway_and_peer(gateway: Gateway, peer: 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) + user_id: Final = _model_user(gateway, model) + scenario.cleanups.callback(lambda: gateway.request("POST", "/user/delete", {"user_ids": [user_id]})) + key: Final = string_value(gateway.post("/key/generate", {"user_id": user_id})["key"]) + + assert _observed(upstream) == [] + for proxy in (gateway, peer): + warm: Final = f"warm {uuid.uuid4().hex}" + assert ( + eventually( + lambda p=proxy, t=warm: _user_chat(p, model, key, t), + lambda response: response.status_code == 200, + seconds=75, + ).status_code + == 200 + ) + assert _observed(upstream) == [_chat_observation(warm)] + _wait_for_spend(user_id, 0.11) + + shrink: Final = gateway.request("POST", "/user/update", {"user_id": user_id, "max_budget": 0.01}) + assert shrink.status_code == 200, shrink.text + _assert_user_budget_refusal(_user_chat(gateway, model, key, f"denied {uuid.uuid4().hex}"), user_id) + _assert_user_budget_refusal( + eventually( + lambda: _user_chat(peer, model, key, f"denied {uuid.uuid4().hex}"), + lambda response: response.status_code == 422, + seconds=10, + ), + user_id, + ) + denied_observations: Final = _observed(upstream) + assert denied_observations == [], denied_observations + for proxy in (gateway, peer): + _assert_user_budget_refusal(_user_chat(proxy, model, key, f"denied {uuid.uuid4().hex}"), user_id) + assert _observed(upstream) == [] + + restored: Final = gateway.request("POST", "/user/update", {"user_id": user_id, "max_budget": 100.0}) + assert restored.status_code == 200, restored.text + for proxy in (gateway, peer): + text: Final = f"restored {uuid.uuid4().hex}" + assert ( + eventually( + lambda p=proxy, t=text: _user_chat(p, model, key, t), + lambda response: response.status_code == 200, + seconds=30, + ).status_code + == 200 + ) + assert _observed(upstream) == [_chat_observation(text)] + + _wait_for_spend(user_id, 0.23) + email: Final = f"{user_id}@integration.local" + gateway.post("/user/update", {"user_id": user_id, "user_email": email}) + spelled: Final = gateway.request("POST", "/user/update", {"user_email": email, "max_budget": 0.01}) + assert spelled.status_code == 200, spelled.text + _assert_user_budget_refusal(_user_chat(gateway, model, key, f"denied {uuid.uuid4().hex}"), user_id) + assert _observed(upstream) == [] + + +def test_user_models_update_with_metadata_propagates_to_gateway_and_peer(gateway: Gateway, peer: 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) + user_id: Final = _model_user(gateway, model) + scenario.cleanups.callback(lambda: gateway.request("POST", "/user/delete", {"user_ids": [user_id]})) + key: Final = string_value(gateway.post("/key/generate", {"user_id": user_id})["key"]) + + assert _observed(upstream) == [] + for proxy in (gateway, peer): + warm: Final = f"warm {uuid.uuid4().hex}" + assert ( + eventually( + lambda p=proxy, t=warm: _user_chat(p, model, key, t), + lambda response: response.status_code == 200, + seconds=75, + ).status_code + == 200 + ) + assert _observed(upstream) == [_chat_observation(warm)] + + narrowed: Final = gateway.request( + "POST", + "/user/update", + {"user_id": user_id, "models": ["no-default-models"], "metadata": {"dept": "restricted"}}, + ) + assert narrowed.status_code == 200, narrowed.text + denied_gateway: Final = _user_chat(gateway, model, key, f"denied {uuid.uuid4().hex}") + assert denied_gateway.status_code == 403, denied_gateway.text + assert denied_gateway.json() == _user_model_denied_error(model), denied_gateway.text + denied_peer: Final = eventually( + lambda: _user_chat(peer, model, key, f"denied {uuid.uuid4().hex}"), + lambda response: response.status_code == 403, + seconds=30, + ) + assert denied_peer.json() == _user_model_denied_error(model), denied_peer.text + model_denied_observations: Final = _observed(upstream) + assert model_denied_observations == [], model_denied_observations + for proxy in (gateway, peer): + denied_again: Final = _user_chat(proxy, model, key, f"denied {uuid.uuid4().hex}") + assert denied_again.status_code == 403, denied_again.text + assert denied_again.json() == _user_model_denied_error(model), denied_again.text + assert _observed(upstream) == [] + + restored: Final = gateway.request( + "POST", + "/user/update", + {"user_id": user_id, "models": [model], "metadata": {"dept": "qa"}}, + ) + assert restored.status_code == 200, restored.text + for proxy in (gateway, peer): + text: Final = f"restored {uuid.uuid4().hex}" + assert ( + eventually( + lambda p=proxy, t=text: _user_chat(p, model, key, t), + lambda response: response.status_code == 200, + seconds=75, + ).status_code + == 200 + ) + assert _observed(upstream) == [_chat_observation(text)] + + +def test_partial_user_update_preserves_sibling_fields(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + user_id: Final = _model_user(gateway, model) + scenario.cleanups.callback(lambda: gateway.request("POST", "/user/delete", {"user_ids": [user_id]})) + + initial: Final = gateway.request( + "POST", + "/user/update", + { + "user_id": user_id, + "user_alias": "sibling-alias", + "rpm_limit": 123, + "budget_duration": "1d", + "spend": 0.25, + }, + ) + assert initial.status_code == 200, initial.text + before: Final = _partial_user_row(user_id) + assert float(str(before[0]["spend"])) == pytest.approx(0.25), before + updated: Final = gateway.request("POST", "/user/update", {"user_id": user_id, "tpm_limit": 1234}) + assert updated.status_code == 200, updated.text + after: Final = _partial_user_row(user_id) + assert after == before + info: Final = object_value(gateway.get("/user/info", {"user_id": user_id})["user_info"]) + assert info["models"] == [model], info + assert info["max_budget"] == 10.0, info + assert info["metadata"] == {"dept": "qa"}, info + assert info["user_alias"] == "sibling-alias", info + assert info["rpm_limit"] == 123, info + assert info["budget_duration"] == "1d", info + assert info["tpm_limit"] == 1234, info + assert float(str(info["spend"])) == pytest.approx(0.25), info + assert float(str(after[0]["spend"])) == pytest.approx(0.25), after + + +def test_user_models_only_update_is_enforced_on_the_gateway(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) + user_id: Final = _model_user(gateway, model, max_budget=100.0) + scenario.cleanups.callback(lambda: gateway.request("POST", "/user/delete", {"user_ids": [user_id]})) + key: Final = string_value(gateway.post("/key/generate", {"user_id": user_id})["key"]) + + assert _observed(upstream) == [] + warm: Final = f"warm {uuid.uuid4().hex}" + warmed: Final = _user_chat(gateway, model, key, warm) + assert warmed.status_code == 200, warmed.text + assert _observed(upstream) == [_chat_observation(warm)] + + narrowed: Final = gateway.request("POST", "/user/update", {"user_id": user_id, "models": ["no-default-models"]}) + assert narrowed.status_code == 200, narrowed.text + assert _observed(upstream) == [] + denied: Final = _user_chat(gateway, model, key, f"denied {uuid.uuid4().hex}") + assert denied.status_code == 403, denied.text + assert denied.json() == _user_model_denied_error(model), denied.text + assert _observed(upstream) == [] + + restored: Final = gateway.request("POST", "/user/update", {"user_id": user_id, "models": [model]}) + assert restored.status_code == 200, restored.text + text: Final = f"restored {uuid.uuid4().hex}" + served: Final = _user_chat(gateway, model, key, text) + assert served.status_code == 200, served.text + assert _observed(upstream) == [_chat_observation(text)] + + +def test_user_models_update_without_metadata_propagates_to_gateway_and_peer(gateway: Gateway, peer: Gateway) -> None: + pytest.skip("BUG: models-only /user/update leaves warmed peer serving the revoked model with 200") + 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) + user_id: Final = _model_user(gateway, model, max_budget=100.0) + scenario.cleanups.callback(lambda: gateway.request("POST", "/user/delete", {"user_ids": [user_id]})) + key: Final = string_value(gateway.post("/key/generate", {"user_id": user_id})["key"]) + + _observed(upstream) + for proxy in (gateway, peer): + warm: Final = f"warm {uuid.uuid4().hex}" + warmed: Final = _user_chat(proxy, model, key, warm) + assert warmed.status_code == 200, warmed.text + assert _observed(upstream) == [_chat_observation(warm)] + + narrowed: Final = gateway.request("POST", "/user/update", {"user_id": user_id, "models": ["no-default-models"]}) + assert narrowed.status_code == 200, narrowed.text + denied_gateway: Final = _user_chat(gateway, model, key, f"denied {uuid.uuid4().hex}") + assert denied_gateway.status_code == 403, denied_gateway.text + assert denied_gateway.json() == _user_model_denied_error(model), denied_gateway.text + denied_peer: Final = eventually( + lambda: _user_chat(peer, model, key, f"denied {uuid.uuid4().hex}"), + lambda response: response.status_code == 403, + seconds=10, + ) + assert denied_peer.json() == _user_model_denied_error(model), denied_peer.text + denied_observations: Final = _observed(upstream) + assert denied_observations == [], denied_observations + for proxy in (gateway, peer): + refused = _user_chat(proxy, model, key, f"denied {uuid.uuid4().hex}") + assert refused.json() == _user_model_denied_error(model), refused.text + assert _observed(upstream) == [] diff --git a/tests/integration/spend/test_end_user_budget_enforcement.py b/tests/integration/spend/test_end_user_budget_enforcement.py index de0bfcd5194..29b18f1f261 100644 --- a/tests/integration/spend/test_end_user_budget_enforcement.py +++ b/tests/integration/spend/test_end_user_budget_enforcement.py @@ -436,3 +436,203 @@ def test_key_end_user_budget_id_refuses_second_request(gateway: Gateway) -> None other_end_user_id, {"budget_id": None, "max_budget": None}, ) + + +class _Budget(BaseModel): + budget_id: str + max_budget: float | None + + +class _CustomerInfo(BaseModel): + user_id: str + blocked: bool + alias: str | None + spend: float + litellm_budget_table: _Budget | None + + +def _new_customer(gateway: Gateway, user_id: str, *, path: str = "/customer/new", **fields: JsonValue) -> None: + response: Final = gateway.request("POST", path, {"user_id": user_id, **fields}) + assert response.status_code == 200, response.text + + +def _customer_info(gateway: Gateway, user_id: str) -> _CustomerInfo: + response: Final = gateway.request("GET", "/customer/info", params={"end_user_id": user_id}) + assert response.status_code == 200, response.text + return _CustomerInfo.model_validate_json(response.content) + + +def _chat(gateway: Gateway, model: str, key: str, end_user: str, *, stream: bool = False) -> httpx.Response: + body: Final[dict[str, JsonValue]] = { + "model": model, + "messages": [{"role": "user", "content": f"end user budget {uuid.uuid4().hex}"}], + "user": end_user, + } + if stream: + body["stream"] = True + return gateway.request("POST", "/v1/chat/completions", body, key=key) + + +def _observed_requests(upstream: httpx.Client) -> list[dict[str, JsonValue]]: + response: Final = upstream.get("/__observations") + response.raise_for_status() + requests: Final = response.json()["requests"] + assert isinstance(requests, list) + return requests + + +def _expected_observation(body: dict[str, JsonValue]) -> dict[str, JsonValue]: + return { + "path": "/v1/chat/completions", + "authorization": "Bearer integration-provider-key", + "body": body, + "method": "POST", + "api_key": "", + } + + +def _sent_body(response: httpx.Response, end_user: str) -> dict[str, JsonValue]: + sent: Final = object_value(json.loads(response.request.content)) + expected: Final[dict[str, JsonValue]] = {"messages": sent["messages"], "model": "gpt-4o-mini", "user": end_user} + if sent.get("stream") is True: + expected["stream"] = True + expected["stream_options"] = {"include_usage": True} + return expected + + +def _end_user_row(user_id: str) -> list[dict[str, JsonValue]]: + return read_rows('SELECT spend, blocked FROM "LiteLLM_EndUserTable" WHERE user_id=%s', (user_id,)) + + +def _daily_rows(user_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT model, model_group, spend, api_requests, successful_requests FROM "LiteLLM_DailyEndUserSpend" WHERE end_user_id=%s', + (user_id,), + ) + + +_REFUSAL_MESSAGE: Final = re.compile( + r"^ExceededBudget: End User=(?P\S+) over budget\. Spend=(?P[0-9.eE+-]+), Budget=(?P[0-9.eE+-]+)$" +) + + +def _assert_budget_refusal(response: httpx.Response, end_user: str, budget: float) -> None: + error: Final = object_value(response.json()["error"]) + assert error["type"] == "budget_exceeded" + assert error["param"] is None + assert error["code"] == "422" + matched: Final = _REFUSAL_MESSAGE.match(str(error["message"])) + assert matched is not None, error["message"] + assert matched.group("user") == end_user + assert float(matched.group("spend")) == pytest.approx(0.06) + assert float(matched.group("budget")) == budget + + +def test_customer_budget_refuses_calls_and_writes_daily_spend(gateway: Gateway) -> None: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): + _observed_requests(upstream) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + _new_customer(gateway, end_user, max_budget=0.05) + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + info: Final = _customer_info(gateway, end_user) + assert info.blocked is False + assert info.litellm_budget_table is not None and info.litellm_budget_table.max_budget == 0.05, info + joined: Final = read_rows( + 'SELECT b.max_budget FROM "LiteLLM_EndUserTable" e JOIN "LiteLLM_BudgetTable" b' + " ON e.budget_id = b.budget_id WHERE e.user_id = %s", + (end_user,), + ) + assert len(joined) == 1 and float(str(joined[0]["max_budget"])) == 0.05, joined + + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + key: Final = scenario.key(models=[model]) + served: Final = _chat(gateway, model, key, end_user) + assert served.status_code == 200, served.text + assert _observed_requests(upstream) == [_expected_observation(_sent_body(served, end_user))] + + charged: Final = eventually( + lambda: _end_user_row(end_user), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= 0.05, + seconds=70, + ) + assert float(str(charged[0]["spend"])) == pytest.approx(20 * 0.001 + 20 * 0.002) + daily: Final = eventually(lambda: _daily_rows(end_user), lambda rows: len(rows) == 1, seconds=70) + assert daily[0]["model"] == "openai/gpt-4o-mini" + assert daily[0]["model_group"] == model + assert float(str(daily[0]["spend"])) == pytest.approx(0.06) + assert int(str(daily[0]["api_requests"])) == 1 + assert int(str(daily[0]["successful_requests"])) == 1 + + denied: Final = _chat(gateway, model, key, end_user) + assert denied.status_code == 422, denied.text + _assert_budget_refusal(denied, end_user, 0.05) + assert _observed_requests(upstream) == [] + + +def test_customer_budget_update_takes_effect_without_restart(gateway: Gateway) -> None: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): + _observed_requests(upstream) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + _new_customer(gateway, end_user, path="/end_user/new", max_budget=1.0) + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + key: Final = scenario.key(models=[model]) + served: Final = _chat(gateway, model, key, end_user) + assert served.status_code == 200, served.text + assert _observed_requests(upstream) == [_expected_observation(_sent_body(served, end_user))] + + updated: Final = gateway.request("POST", "/customer/update", {"user_id": end_user, "max_budget": 0.01}) + assert updated.status_code == 200, updated.text + info: Final = _customer_info(gateway, end_user) + assert info.litellm_budget_table is not None and info.litellm_budget_table.max_budget == 0.01, info + budget_row: Final = read_rows( + 'SELECT b.max_budget FROM "LiteLLM_EndUserTable" e JOIN "LiteLLM_BudgetTable" b' + " ON e.budget_id = b.budget_id WHERE e.user_id = %s", + (end_user,), + ) + assert len(budget_row) == 1 and float(str(budget_row[0]["max_budget"])) == 0.01, budget_row + eventually( + lambda: _end_user_row(end_user), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= 0.05, + seconds=70, + ) + denied: Final = _chat(gateway, model, key, end_user) + assert denied.status_code == 422, denied.text + _assert_budget_refusal(denied, end_user, 0.01) + assert _observed_requests(upstream) == [] + + +def test_streamed_call_is_refused_as_json_once_over_budget(gateway: Gateway) -> None: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): + _observed_requests(upstream) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + _new_customer(gateway, end_user, max_budget=0.05) + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + key: Final = scenario.key(models=[model]) + served: Final = _chat(gateway, model, key, end_user, stream=True) + assert served.status_code == 200, served.text + assert served.headers["content-type"].startswith("text/event-stream") + assert 'data: {"id":"chatcmpl-' in served.text and "data: [DONE]" in served.text, served.text + assert _observed_requests(upstream) == [_expected_observation(_sent_body(served, end_user))] + eventually( + lambda: _end_user_row(end_user), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= 0.05, + seconds=70, + ) + + denied: Final = _chat(gateway, model, key, end_user, stream=True) + assert denied.status_code == 422, denied.text + assert denied.headers["content-type"].startswith("application/json") + assert "data:" not in denied.text, denied.text + _assert_budget_refusal(denied, end_user, 0.05) + assert _observed_requests(upstream) == [] diff --git a/tests/integration/spend/test_end_user_spend_without_proxy_user.py b/tests/integration/spend/test_end_user_spend_without_proxy_user.py index 15d77f16db4..16210438e82 100644 --- a/tests/integration/spend/test_end_user_spend_without_proxy_user.py +++ b/tests/integration/spend/test_end_user_spend_without_proxy_user.py @@ -1,9 +1,40 @@ +import json import uuid -from typing import Final +from collections.abc import Mapping +from pathlib import Path +from typing import Final, Literal +import httpx import pytest -from integration._support.client import Gateway, eventually +import yaml +from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows +from integration._support.process import owned_proxy +from integration._support.upstream import JsonResponse, SseResponse, delete_scenario, register_scenario +from pydantic import BaseModel, JsonValue, TypeAdapter + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +class _AnthropicTextBlock(BaseModel): + type: Literal["text"] + text: str + + +class _AnthropicMessageUsage(BaseModel): + input_tokens: int + output_tokens: int + + +class _AnthropicMessageResponse(BaseModel): + id: str + type: Literal["message"] + role: Literal["assistant"] + model: str + content: list[_AnthropicTextBlock] + stop_reason: str | None + stop_sequence: str | None + usage: _AnthropicMessageUsage @pytest.mark.covers("spend.end_user.charged_when_key_has_no_user_id_and_auth_cache_is_redis") @@ -27,3 +58,521 @@ def test_end_user_spend_lands_for_key_without_user_id_when_auth_cache_is_redis(g seconds=70, ) assert float(charged[0]["spend"]) == pytest.approx(20 * 0.001 + 20 * 0.002) + + +def _observed(upstream: httpx.Client) -> list[dict[str, JsonValue]]: + response: Final = upstream.get("/__observations") + response.raise_for_status() + requests: Final = response.json()["requests"] + assert isinstance(requests, list) + return requests + + +def _observation(path: str, body: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return { + "path": path, + "authorization": "Bearer integration-provider-key", + "body": dict(body), + "method": "POST", + "api_key": "", + } + + +def _new_customer(gateway: Gateway, user_id: str, **fields: JsonValue) -> None: + response: Final = gateway.request("POST", "/customer/new", {"user_id": user_id, **fields}) + assert response.status_code == 200, response.text + + +def _charged_spend(user_id: str) -> float: + charged: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id=%s', (user_id,)), + lambda values: len(values) == 1 and float(str(values[0]["spend"])) >= 0.05, + seconds=70, + ) + return float(str(charged[0]["spend"])) + + +def _assert_daily_row(user_id: str, model: str, upstream_model: str = "openai/gpt-4o-mini") -> None: + daily: Final = eventually( + lambda: read_rows( + 'SELECT model, model_group, spend, api_requests FROM "LiteLLM_DailyEndUserSpend" WHERE end_user_id=%s', + (user_id,), + ), + lambda rows: len(rows) == 1, + seconds=70, + ) + assert daily[0]["model"] == upstream_model + assert daily[0]["model_group"] == model + assert float(str(daily[0]["spend"])) == pytest.approx(0.06) + assert int(str(daily[0]["api_requests"])) == 1 + + +def _assert_budget_refusal(response: httpx.Response, end_user: str) -> None: + assert response.status_code == 422, response.text + error: Final = object_value(response.json()["error"]) + assert error["type"] == "budget_exceeded" + assert error["param"] is None + assert error["code"] == "422" + assert str(error["message"]).startswith(f"ExceededBudget: End User={end_user} over budget. Spend="), error[ + "message" + ] + + +def _json_object(raw: bytes) -> dict[str, JsonValue]: + return _JSON_OBJECT.validate_json(raw) + + +_CHAT_SPELLINGS: Final = ( + pytest.param("x-litellm-end-user-id", id="header-x-litellm-end-user-id"), + pytest.param("x-litellm-customer-id", id="header-x-litellm-customer-id"), + pytest.param("litellm_metadata.user", id="body-litellm_metadata-user"), + pytest.param("metadata.user_id", id="body-metadata-user_id"), + pytest.param("litellm_metadata.user:json-string", id="body-litellm_metadata-json-string-user"), + pytest.param("metadata.user_id:json-string", id="body-metadata-json-string-user_id"), +) + + +def _chat_request( + gateway: Gateway, + model: str, + key: str, + spelling: str, + end_user: str, + text: str, + stream: bool = False, +) -> httpx.Response: + headers: Final[dict[str, str]] = {} + body: Final[dict[str, JsonValue]] = { + "model": model, + "messages": [{"role": "user", "content": text}], + } + if spelling.startswith("x-"): + headers[spelling] = end_user + elif spelling == "litellm_metadata.user": + body["litellm_metadata"] = {"user": end_user} + elif spelling == "litellm_metadata.user:json-string": + body["litellm_metadata"] = json.dumps({"user": end_user}) + elif spelling == "metadata.user_id:json-string": + body["metadata"] = json.dumps({"user_id": end_user}) + else: + body["metadata"] = {"user_id": end_user} + if stream: + body["stream"] = True + return gateway.request("POST", "/v1/chat/completions", body, key=key, headers=headers) + + +@pytest.mark.parametrize( + "spelling", + ( + pytest.param("x-litellm-end-user-id", id="header"), + pytest.param("metadata.user_id", id="metadata"), + ), +) +def test_streamed_end_user_spelling_is_charged_and_budget_enforced(gateway: Gateway, spelling: str) -> None: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): + _observed(upstream) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + _new_customer(gateway, end_user, max_budget=0.05) + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + key: Final = scenario.key(models=[model]) + + text: Final = f"streamed end user spelling {uuid.uuid4().hex}" + served: Final = _chat_request(gateway, model, key, spelling, end_user, text, stream=True) + assert served.status_code == 200, served.text + assert served.headers["content-type"].startswith("text/event-stream"), served.headers + chunks: Final = tuple( + _json_object(line.removeprefix("data: ").strip().encode()) + for line in served.text.splitlines() + if line.startswith("data: ") and line.removeprefix("data: ").strip() != "[DONE]" + ) + assert chunks + assert object_value(chunks[-1]["choices"][0])["finish_reason"] == "stop" + assert "data: [DONE]" in served.text + observed: Final = _observed(upstream) + assert observed == [ + _observation( + "/v1/chat/completions", + { + "messages": [{"role": "user", "content": text}], + "model": "gpt-4o-mini", + "stream": True, + "stream_options": {"include_usage": True}, + }, + ) + ] + spend: Final = _charged_spend(end_user) + assert spend == pytest.approx(20 * 0.001 + 20 * 0.002) + _assert_daily_row(end_user, model) + + denied: Final = _chat_request( + gateway, + model, + key, + spelling, + end_user, + f"again {uuid.uuid4().hex}", + stream=True, + ) + assert denied.headers["content-type"].startswith("application/json"), denied.headers + assert "data:" not in denied.text + _assert_budget_refusal(denied, end_user) + assert _observed(upstream) == [] + + +@pytest.mark.parametrize("stream", (False, True), ids=("non-streaming", "streaming")) +def test_anthropic_messages_metadata_user_id_is_charged_and_budget_enforced(gateway: Gateway, stream: bool) -> None: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): + _observed(upstream) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + _new_customer(gateway, end_user, max_budget=0.05) + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + scenario_id: Final = f"sc-{uuid.uuid4().hex}" + response: Final = ( + SseResponse( + content_type="text/event-stream", + frames=( + "event: message_start\ndata: " + '{"type":"message_start","message":{"id":"msg_$REQUEST_ID","type":"message",' + '"role":"assistant","model":"claude-sonnet-5-5","content":[],"stop_reason":null,' + '"stop_sequence":null,"usage":{"input_tokens":20,"output_tokens":0}}}', + "event: content_block_start\ndata: " + '{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}', + "event: content_block_delta\ndata: " + '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"scripted answer"}}', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}', + "event: message_delta\ndata: " + '{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},' + '"usage":{"output_tokens":20}}', + 'event: message_stop\ndata: {"type":"message_stop"}', + ), + ) + if stream + else JsonResponse( + content_type="application/json", + body={ + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "text", "text": "scripted answer"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 20, "output_tokens": 20}, + }, + ) + ) + handle: Final = register_scenario(scenario_id, response) + scenario.cleanups.callback(delete_scenario, handle) + model: Final = scenario.model( + model="anthropic/claude-sonnet-5-5", + api_base=handle.api_base(), + input_cost_per_token=0.001, + output_cost_per_token=0.002, + ) + key: Final = scenario.key(models=[model]) + text: Final = f"anthropic end user {uuid.uuid4().hex}" + body: Final[dict[str, JsonValue]] = { + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": text}], + "metadata": {"user_id": end_user}, + "stream": stream, + } + served: Final = gateway.client.request( + "POST", + "/v1/messages", + json=body, + headers={"x-api-key": key, "anthropic-version": "2023-06-01"}, + ) + assert served.status_code == 200, served.text + if stream: + assert served.headers["content-type"].startswith("text/event-stream"), served.headers + message_events: Final = tuple( + _json_object(line.removeprefix("data: ").strip().encode()) + for line in served.text.splitlines() + if line.startswith("data: ") + ) + assert tuple(str(event["type"]) for event in message_events) == ( + "message_start", + "content_block_start", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + ) + else: + message: Final = _AnthropicMessageResponse.model_validate_json(served.content) + assert message.type == "message" + assert message.role == "assistant" + assert message.content == [_AnthropicTextBlock(type="text", text="scripted answer")] + assert message.stop_reason == "end_turn" + assert message.stop_sequence is None + assert message.usage == _AnthropicMessageUsage(input_tokens=20, output_tokens=20) + observed: Final = _observed(upstream) + assert observed == [ + { + "path": f"/{scenario_id}/v1/messages", + "authorization": "", + "body": { + "model": "claude-sonnet-5-5", + "max_tokens": 64, + "messages": [{"role": "user", "content": text}], + "metadata": {"user_id": end_user}, + "stream": stream, + }, + "method": "POST", + "api_key": "", + } + ] + spend: Final = _charged_spend(end_user) + assert spend == pytest.approx(20 * 0.001 + 20 * 0.002) + _assert_daily_row(end_user, model, "anthropic/claude-sonnet-5-5") + + denied: Final = gateway.client.request( + "POST", + "/v1/messages", + json={**body, "messages": [{"role": "user", "content": f"again {uuid.uuid4().hex}"}]}, + headers={"x-api-key": key, "anthropic-version": "2023-06-01"}, + ) + _assert_budget_refusal(denied, end_user) + assert _observed(upstream) == [] + + +@pytest.mark.parametrize("spelling", _CHAT_SPELLINGS) +def test_end_user_id_spelling_is_charged_and_budget_enforced(gateway: Gateway, spelling: str) -> None: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): + _observed(upstream) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + _new_customer(gateway, end_user, max_budget=0.05) + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + key: Final = scenario.key(models=[model]) + + text: Final = f"end user spelling {uuid.uuid4().hex}" + served: Final = _chat_request(gateway, model, key, spelling, end_user, text) + assert served.status_code == 200, served.text + assert _observed(upstream) == [ + _observation( + "/v1/chat/completions", + {"messages": [{"role": "user", "content": text}], "model": "gpt-4o-mini"}, + ) + ] + spend: Final = _charged_spend(end_user) + assert spend == pytest.approx(20 * 0.001 + 20 * 0.002) + _assert_daily_row(end_user, model) + + denied: Final = _chat_request(gateway, model, key, spelling, end_user, f"again {uuid.uuid4().hex}") + _assert_budget_refusal(denied, end_user) + assert _observed(upstream) == [] + + +def test_safety_identifier_on_responses_is_charged_and_budget_enforced(gateway: Gateway) -> None: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): + _observed(upstream) + scenario_id: Final = f"sc-{uuid.uuid4().hex}" + handle: Final = register_scenario( + scenario_id, + JsonResponse( + content_type="application/json", + body={ + "id": f"resp_{scenario_id}", + "object": "response", + "created_at": 1759700000, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "scripted answer", "annotations": []}], + } + ], + "usage": { + "input_tokens": 20, + "output_tokens": 20, + "total_tokens": 40, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + }, + ), + ) + scenario.cleanups.callback( + lambda: httpx.delete(f"{handle.control_url}/__scenarios/{handle.scenario_id}", trust_env=False) + ) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + _new_customer(gateway, end_user, max_budget=0.05) + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + model: Final = scenario.model( + api_base=f"{gateway.upstream_url}/{handle.scenario_id}/v1", + input_cost_per_token=0.001, + output_cost_per_token=0.002, + ) + key: Final = scenario.key(models=[model]) + + input_message: Final[dict[str, JsonValue]] = { + "role": "user", + "content": [{"type": "input_text", "text": f"responses spelling {uuid.uuid4().hex}"}], + } + + def call() -> httpx.Response: + return gateway.request( + "POST", + "/v1/responses", + {"model": model, "input": [input_message], "safety_identifier": end_user}, + key=key, + ) + + served: Final = call() + assert served.status_code == 200, served.text + assert _observed(upstream) == [ + _observation( + f"/{handle.scenario_id}/v1/responses", + {"model": "gpt-4o-mini", "input": [input_message], "safety_identifier": end_user}, + ) + ] + spend: Final = _charged_spend(end_user) + assert spend == pytest.approx(0.06) + _assert_daily_row(end_user, model) + + denied: Final = call() + _assert_budget_refusal(denied, end_user) + assert _observed(upstream) == [] + + +def test_end_user_header_wins_over_body_user(gateway: Gateway) -> None: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): + _observed(upstream) + header_user: Final = f"integration-end-user-a-{uuid.uuid4().hex}" + body_user: Final = f"integration-end-user-b-{uuid.uuid4().hex}" + _new_customer(gateway, header_user, max_budget=0.05) + _new_customer(gateway, body_user, max_budget=1.0) + for end_user in (header_user, body_user): + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + key: Final = scenario.key(models=[model]) + + text: Final = f"header precedence {uuid.uuid4().hex}" + served: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": text}], "user": body_user}, + key=key, + headers={"x-litellm-end-user-id": header_user}, + ) + assert served.status_code == 200, served.text + assert _observed(upstream) == [ + _observation( + "/v1/chat/completions", + { + "messages": [{"role": "user", "content": text}], + "model": "gpt-4o-mini", + "user": body_user, + }, + ) + ] + assert _charged_spend(header_user) == pytest.approx(0.06) + _assert_daily_row(header_user, model) + body_rows: Final = read_rows('SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id=%s', (body_user,)) + assert all(float(str(row["spend"])) == 0.0 for row in body_rows), body_rows + assert ( + read_rows( + 'SELECT id FROM "LiteLLM_DailyEndUserSpend" WHERE end_user_id=%s', + (body_user,), + ) + == [] + ) + + denied: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"again {uuid.uuid4().hex}"}], "user": body_user}, + key=key, + headers={"x-litellm-end-user-id": header_user}, + ) + _assert_budget_refusal(denied, header_user) + assert _observed(upstream) == [] + + +def _config_with_general_settings(tmp_path: Path, general_settings: Mapping[str, JsonValue]) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["general_settings"] = {**object_value(config["general_settings"]), **dict(general_settings)} + path: Final = tmp_path / "config.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.mark.parametrize( + ("general_settings", "injects_user_body_field"), + ( + pytest.param( + {"user_header_mappings": [{"header_name": "x-integration-end-user", "litellm_user_role": "customer"}]}, + False, + id="config-user_header_mappings", + ), + pytest.param({"user_header_name": "x-integration-end-user"}, True, id="config-user_header_name"), + ), +) +def test_configured_end_user_header_is_charged_and_budget_enforced( + gateway: Gateway, tmp_path: Path, general_settings: Mapping[str, JsonValue], injects_user_body_field: bool +) -> None: + config: Final = _config_with_general_settings(tmp_path, general_settings) + with ( + owned_proxy(gateway, tmp_path, {}, config=config) as candidate, + candidate.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): + _observed(upstream) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + _new_customer(candidate, end_user, max_budget=0.05) + scenario.cleanups.callback(candidate.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + key: Final = scenario.key(models=[model]) + headers: Final = {"x-integration-end-user": end_user} + + text: Final = f"configured header {uuid.uuid4().hex}" + served: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": text}]}, + key=key, + headers=headers, + ) + assert served.status_code == 200, served.text + expected_body: Final[dict[str, JsonValue]] = { + "messages": [{"role": "user", "content": text}], + "model": "gpt-4o-mini", + } + if injects_user_body_field: + expected_body["user"] = end_user + assert _observed(upstream) == [_observation("/v1/chat/completions", expected_body)] + assert _charged_spend(end_user) == pytest.approx(0.06) + _assert_daily_row(end_user, model) + + denied: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"again {uuid.uuid4().hex}"}]}, + key=key, + headers=headers, + ) + _assert_budget_refusal(denied, end_user) + assert _observed(upstream) == []