From 92c2e69ac534fc3588b2bff7f544003c0a7f25ae Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Mon, 9 Nov 2026 15:38:59 -0800 Subject: [PATCH] fix(proxy): fail-closed budget enforcement returns 403 naming /team/member_add for a team user with no membership row (#50603) * fix(proxy): read an absent team membership row as zero spend under fail-closed budget enforcement A team user under a default member budget has no LiteLLM_TeamMembership row until spend is persisted. SpendCounterReseed.from_db returned None for that absent row, the same value as a failed DB read, so fail_closed_budget_enforcement rejected every request from that user with 503. The DB answering that the row is absent is a verified zero, so from_db now returns 0.0 for a team-member counter without a row and still returns None when the read fails. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): only read an absent membership row as zero when the team member key parses unambiguously A team_id or user_id containing ":" makes the counter key split onto the wrong membership row, so that miss keeps returning None instead of a verified zero. Type the new unit tests. * fix(proxy): refuse a fail-closed team member with no membership row with a 403 that names /team/member_add Replaces the earlier reading of an absent LiteLLM_TeamMembership row as verified zero spend. A key minted with team_id and user_id but no /team/member_add has no row and the spend writer never creates one, so admitting it at zero would leave its member spend tracked only in Redis. With fail_closed_budget_enforcement on, _check_team_member_budget now rejects that key with a 403 that names the user, the team and the member_add call, before any spend read. With the flag off nothing changes. The membership load raises on a database failure and returns None only for a genuinely absent row, so the 403 never fires during an outage. spend_counter_reseed.py and the counter prefix constant go back to main. * test(proxy): drive the no-membership 403 tests through the real membership and spend readers Fake Postgres rows and the FakeSpendCounters double stand in for the external stores instead of patching get_team_membership and get_current_spend. Strip CR/LF from the ids in the fail-closed warning log. * fix(proxy): read fail_closed_budget_enforcement from general_settings in the membership check Importing proxy_server's private helper grew reportPrivateUsage past the type gate ceiling. Type the integration helper's return as JsonValue. * fix(proxy): expose fail_closed_budget_enforcement_enabled for the membership check Reading general_settings directly from auth_checks grew the unknown-type counts past the type gate ceiling; a small public wrapper in proxy_server keeps the gate flat. * test(proxy): drop docstrings that restate the no-membership tests --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 28 +++++++- litellm/proxy/proxy_server.py | 4 ++ .../test_fail_closed_team_member_budget.py | 66 ++++++++++++++++++ ...st_auth_checks_object_access_and_lookup.py | 69 ++++++++++++++++++- 4 files changed, 164 insertions(+), 3 deletions(-) create mode 100644 tests/integration/spend/test_fail_closed_team_member_budget.py diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e546c093473..f0c9632837d 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -7076,6 +7076,27 @@ def team_member_max_budget_alert_check( _team_member_max_budget_alert_check: Final = team_member_max_budget_alert_check +def _raise_team_member_without_membership(user_id: str, team_id: str) -> None: + from litellm.proxy.litellm_pre_call_utils import sanitize_for_log + + verbose_proxy_logger.warning( + "fail_closed_budget_enforcement: rejecting request, user_id=%s has no membership row in team_id=%s " + "so spend against the team member budget cannot be verified", + sanitize_for_log(user_id), + sanitize_for_log(team_id), + ) + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": ( + f"User={user_id} has no membership in Team={team_id}, so their spend against the team member " + "budget cannot be verified and fail_closed_budget_enforcement is enabled. Add them to the team " + "with POST /team/member_add." + ) + }, + ) + + async def _check_team_member_budget( team_object: LiteLLM_TeamTable | None, user_object: LiteLLM_UserTable | None, @@ -7125,11 +7146,14 @@ async def _check_team_member_budget( ) if team_member_budget is not None: + from litellm.proxy.proxy_server import fail_closed_budget_enforcement_enabled, get_current_spend + + if loaded_membership is None and fail_closed_budget_enforcement_enabled(): + _raise_team_member_without_membership(user_id=valid_token.user_id, team_id=team_object.team_id) + team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0 # Read from cross-pod counter (Redis-first) if available - from litellm.proxy.proxy_server import get_current_spend - team_member_spend = await get_current_spend( counter_key=team_member_counter(valid_token.user_id, team_object.team_id), fallback_spend=team_member_spend, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 046510d026d..37dc135d4d7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3389,6 +3389,10 @@ def _fail_closed_budget_enforcement() -> bool: return general_settings.get("fail_closed_budget_enforcement") is True +def fail_closed_budget_enforcement_enabled() -> bool: + return _fail_closed_budget_enforcement() + + def _raise_budget_unverifiable(counter_key: str) -> None: verbose_proxy_logger.warning( "fail_closed_budget_enforcement: rejecting request — spend for %s could " diff --git a/tests/integration/spend/test_fail_closed_team_member_budget.py b/tests/integration/spend/test_fail_closed_team_member_budget.py new file mode 100644 index 00000000000..52b2e8dd0f3 --- /dev/null +++ b/tests/integration/spend/test_fail_closed_team_member_budget.py @@ -0,0 +1,66 @@ +import uuid +from collections.abc import Iterator +from contextlib import ExitStack +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, gateway_from_environment, object_value, string_value +from integration._support.process import owned_proxy_process +from pydantic import JsonValue + + +def _fail_closed_config(directory: Path) -> Path: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = { + **base, + "general_settings": {**object_value(base["general_settings"]), "fail_closed_budget_enforcement": True}, + } + path: Final = directory / f"fail-closed-{uuid.uuid4().hex[:8]}.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.fixture(scope="module") +def fail_closed_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("fail-closed-team-member") + with ExitStack() as stack: + rig_gateway: Final = stack.enter_context(gateway_from_environment()) + owned: Final = stack.enter_context( + owned_proxy_process(rig_gateway, directory, {}, config=_fail_closed_config(directory)) + ) + yield owned.gateway + + +def _chat(gateway: Gateway, model: str, key: str) -> tuple[int, dict[str, JsonValue]]: + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "integration control"}]}, + key=key, + ) + return response.status_code, JSON_OBJECT.validate_json(response.content) + + +def test_team_member_without_membership_row_is_told_to_be_added_then_admitted(fail_closed_gateway: Gateway) -> None: + with fail_closed_gateway.scenario() as scenario: + model: Final = scenario.model() + team_id: Final = scenario.team(models=[model], team_member_budget=100.0) + user_id: Final = scenario.user(max_budget=100.0) + key: Final = scenario.key(team_id=team_id, user_id=user_id, models=[model]) + + status, body = _chat(fail_closed_gateway, model, key) + + assert status == 403, body + message: Final = string_value(object_value(body["error"])["message"]) + assert user_id in message and team_id in message and "/team/member_add" in message + + fail_closed_gateway.post( + "/team/member_add", {"team_id": team_id, "member": {"role": "user", "user_id": user_id}} + ) + + admitted_status, admitted = _chat(fail_closed_gateway, model, key) + + assert admitted_status == 200, admitted + assert object_value(admitted["usage"])["total_tokens"] diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 78ce9728f3b..e7bcc696c0d 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -6080,6 +6080,71 @@ class TestGuardrailModificationCheck: self._call({"metadata": '"just a string"'}) +def _member_without_row_prisma(max_budget: float) -> MagicMock: + budget_row: Final = MagicMock() + budget_row.max_budget = max_budget + budget_row.dict = MagicMock(return_value={"budget_id": "budget-default", "max_budget": max_budget}) + prisma_client: Final = MagicMock() + prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=budget_row) + prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + return prisma_client + + +async def _check_member_without_membership_row(prisma_client: MagicMock) -> None: + from litellm.caching.dual_cache import DualCache + from litellm.proxy.utils import ProxyLogging + + await _check_team_member_budget( + team_object=LiteLLM_TeamTable(team_id="team-1", metadata={"team_member_budget_id": "budget-default"}), + user_object=LiteLLM_UserTable(user_id="user-1"), + valid_token=UserAPIKeyAuth(token="tok-1", user_id="user-1", team_id="team-1"), + prisma_client=prisma_client, + user_api_key_cache=DualCache(), + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + ) + + +@pytest.mark.asyncio +async def test_fail_closed_rejects_member_without_membership_row_with_a_403_naming_member_add( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from fastapi import HTTPException + + import litellm.proxy.proxy_server as ps + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters(redis={}) + prisma_client: Final = _member_without_row_prisma(100.0) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) + monkeypatch.setattr(ps, "spend_counter_store", spend.store()) + monkeypatch.setattr(ps, "prisma_client", prisma_client) + + with pytest.raises(HTTPException) as exc_info: + await _check_member_without_membership_row(prisma_client) + + assert exc_info.value.status_code == 403 + message: Final = str(exc_info.value.detail) + assert "user-1" in message and "team-1" in message and "/team/member_add" in message + prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() + assert spend.commands == [] + + +@pytest.mark.asyncio +async def test_fail_closed_off_still_admits_member_without_membership_row(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm.proxy.proxy_server as ps + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters(redis={}) + prisma_client: Final = _member_without_row_prisma(100.0) + monkeypatch.setattr(ps, "general_settings", {}) + monkeypatch.setattr(ps, "spend_counter_store", spend.store()) + monkeypatch.setattr(ps, "prisma_client", prisma_client) + + await _check_member_without_membership_row(prisma_client) + + assert spend.commands, "the member budget check must still read the spend counter with the flag off" + + @pytest.mark.asyncio async def test_team_member_budget_check_falls_back_to_team_default_budget_id(): """When a member's TeamMembership has no linked budget row, the check @@ -11850,7 +11915,9 @@ async def test_baggage_mapping_rejects_exactly_the_client_tags_rejected_without_ unmapped_request = _client_tag_request(client_tags, baggage=False) unmapped = _client_tag_body(client_tags, client_slot) LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(unmapped, route) - LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth(unmapped_request, unmapped, UserAPIKeyAuth(api_key="qa-key")) + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( + unmapped_request, unmapped, UserAPIKeyAuth(api_key="qa-key") + ) mapped_request = _client_tag_request(client_tags, baggage=baggage) mapped = _client_tag_body(client_tags, client_slot)