diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 5d6e744d504..974ff18d673 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -72,6 +72,11 @@ mod _native { request_auth_team_override, request_auth_unpriced_model, }; #[pymodule_export] + use crate::request_auth::{ + request_auth_budget_lookup_emails, request_auth_budget_lookup_member, + request_auth_budget_lookup_organization, request_auth_budget_lookup_soft, + }; + #[pymodule_export] use crate::router::{RouterCall, RouterCoreHandle}; #[pymodule_export] use crate::routes::audio_transcription::{atranscription, transcription}; diff --git a/litellm-rust/crates/python-bridge/src/request_auth.rs b/litellm-rust/crates/python-bridge/src/request_auth.rs index 8cea1f00e6f..9d07c977f29 100644 --- a/litellm-rust/crates/python-bridge/src/request_auth.rs +++ b/litellm-rust/crates/python-bridge/src/request_auth.rs @@ -140,6 +140,45 @@ pub(crate) fn request_auth_access(py: Python<'_>, facts: &Bound<'_, PyAny>) -> P to_py(py, &py.detach(move || core::access::decide(facts))) } +#[pyfunction] +pub(crate) fn request_auth_budget_lookup_soft( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult { + let facts: core::budget_lookup::SoftFacts = from_py(facts)?; + Ok(py.detach(move || core::budget_lookup::read_soft_emails(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_lookup_organization( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_lookup::OrganizationFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_lookup::organization(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_lookup_member( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_lookup::MemberFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::budget_lookup::member(&facts))) +} + +#[pyfunction] +pub(crate) fn request_auth_budget_lookup_emails( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_lookup::EmailFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::budget_lookup::emails(&facts))) +} + #[pyfunction] pub(crate) fn request_auth_budget_window_start( py: Python<'_>, diff --git a/litellm-rust/crates/request-auth-core/src/budget_lookup.rs b/litellm-rust/crates/request-auth-core/src/budget_lookup.rs new file mode 100644 index 00000000000..c5c150d2bdb --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/budget_lookup.rs @@ -0,0 +1,284 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Default, Deserialize)] +pub struct OrganizationFacts { + pub token_present: bool, + pub client_present: bool, + pub team_present: bool, + pub token_id: Option, + pub team_id: Option, + pub row_present: Option, + pub lookup_failed: bool, +} + +#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum OrganizationAction { + Skip, + ReadTokenId, + ReadTeamId, + ReadOrganization, + UseOrganization, +} + +pub fn organization(facts: &OrganizationFacts) -> OrganizationAction { + use OrganizationAction::*; + if !facts.token_present || !facts.client_present { + return Skip; + } + match facts.token_id { + None => return ReadTokenId, + Some(false) if !facts.team_present => return Skip, + Some(false) => match facts.team_id { + None => return ReadTeamId, + Some(false) => return Skip, + Some(true) => {} + }, + Some(true) => {} + } + if facts.lookup_failed { + return Skip; + } + match facts.row_present { + None => ReadOrganization, + Some(false) => Skip, + Some(true) => UseOrganization, + } +} + +#[derive(Clone, Debug, Default, Deserialize)] +pub struct MemberFacts { + pub team_present: bool, + pub token_present: bool, + pub team_id: Option, + pub user_id: Option, + pub membership_loaded: bool, + pub limits_read: bool, + pub maximum_present: bool, + pub default_id_string: bool, + pub default_read: bool, +} + +#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum MemberAction { + Skip, + ReadTeamId, + ReadUserId, + ReadMembership, + ReadLimits, + ReadDefault, + UseLimits, +} + +pub fn member(facts: &MemberFacts) -> MemberAction { + use MemberAction::*; + if !facts.team_present { + return Skip; + } + match facts.team_id { + None => return ReadTeamId, + Some(false) => return Skip, + Some(true) => {} + } + if !facts.token_present { + return Skip; + } + match facts.user_id { + None => return ReadUserId, + Some(false) => return Skip, + Some(true) => {} + } + if !facts.membership_loaded { + return ReadMembership; + } + if !facts.limits_read { + return ReadLimits; + } + if !facts.maximum_present && facts.default_id_string && !facts.default_read { + return ReadDefault; + } + UseLimits +} + +#[derive(Clone, Debug, Deserialize)] +pub enum EmailFacts { + Text(String), + List(Vec>), + Other, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct SoftFacts { + pub limit: Option, + pub spend: Option, + pub token_present: bool, +} + +pub fn read_soft_emails(facts: &SoftFacts) -> bool { + facts.token_present + && facts + .limit + .zip(facts.spend) + .is_some_and(|(limit, spend)| spend >= limit) +} + +fn stripped(value: &str) -> &str { + value.trim_matches(|character: char| { + character.is_whitespace() || ('\u{1c}'..='\u{1f}').contains(&character) + }) +} + +pub fn emails(facts: &EmailFacts) -> Vec { + match facts { + EmailFacts::Text(value) => value + .split(',') + .map(stripped) + .filter(|value| !value.is_empty()) + .map(str::to_owned) + .collect(), + EmailFacts::List(values) => values + .iter() + .flatten() + .filter(|value| !stripped(value).is_empty()) + .cloned() + .collect(), + EmailFacts::Other => Vec::new(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::below(Some(1.0), Some(0.0), true, false)] + #[case::at_limit(Some(1.0), Some(1.0), true, true)] + #[case::no_token(Some(1.0), Some(2.0), false, false)] + #[case::no_limit(None, Some(2.0), true, false)] + #[case::nan(Some(f64::NAN), Some(2.0), true, false)] + fn lazy_soft_metadata( + #[case] limit: Option, + #[case] spend: Option, + #[case] token_present: bool, + #[case] expected: bool, + ) { + assert_eq!( + read_soft_emails(&SoftFacts { + limit, + spend, + token_present + }), + expected + ); + } + + #[rstest] + fn failed_organization_lookup_is_skipped() { + assert_eq!( + organization(&OrganizationFacts { + token_present: true, + client_present: true, + token_id: Some(true), + lookup_failed: true, + ..Default::default() + }), + OrganizationAction::Skip + ); + } + + #[rstest] + #[case::missing_token(false, true, None, None, OrganizationAction::Skip)] + #[case::missing_client(true, false, None, None, OrganizationAction::Skip)] + #[case::token_first(true, true, None, None, OrganizationAction::ReadTokenId)] + #[case::token_selected(true, true, Some(true), None, OrganizationAction::ReadOrganization)] + #[case::missing_row(true, true, Some(true), Some(false), OrganizationAction::Skip)] + #[case::loaded_row( + true, + true, + Some(true), + Some(true), + OrganizationAction::UseOrganization + )] + fn organization_stages( + #[case] token_present: bool, + #[case] client_present: bool, + #[case] token_id: Option, + #[case] row_present: Option, + #[case] expected: OrganizationAction, + ) { + assert_eq!( + organization(&OrganizationFacts { + token_present, + client_present, + token_id, + row_present, + ..Default::default() + }), + expected + ); + } + + #[rstest] + #[case::no_team(false, None, OrganizationAction::Skip)] + #[case::read_team(true, None, OrganizationAction::ReadTeamId)] + #[case::no_team_id(true, Some(false), OrganizationAction::Skip)] + #[case::team_selected(true, Some(true), OrganizationAction::ReadOrganization)] + fn organization_fallback( + #[case] team_present: bool, + #[case] team_id: Option, + #[case] expected: OrganizationAction, + ) { + assert_eq!( + organization(&OrganizationFacts { + token_present: true, + client_present: true, + token_id: Some(false), + team_present, + team_id, + row_present: None, + lookup_failed: false + }), + expected + ); + } + + #[rstest] + #[case::load_membership(false, false, false, false, false, MemberAction::ReadMembership)] + #[case::read_limits(true, false, false, false, false, MemberAction::ReadLimits)] + #[case::own_maximum(true, true, true, true, false, MemberAction::UseLimits)] + #[case::default(true, true, false, true, false, MemberAction::ReadDefault)] + #[case::invalid_default(true, true, false, false, false, MemberAction::UseLimits)] + #[case::loaded_default(true, true, false, true, true, MemberAction::UseLimits)] + fn member_stages( + #[case] membership_loaded: bool, + #[case] limits_read: bool, + #[case] maximum_present: bool, + #[case] default_id_string: bool, + #[case] default_read: bool, + #[case] expected: MemberAction, + ) { + assert_eq!( + member(&MemberFacts { + team_present: true, + token_present: true, + team_id: Some(true), + user_id: Some(true), + membership_loaded, + limits_read, + maximum_present, + default_id_string, + default_read + }), + expected + ); + } + + #[rstest] + #[case::comma_list(EmailFacts::Text(" a, ,b,a ".into()), vec!["a", "b", "a"])] + #[case::list_keeps_space(EmailFacts::List(vec![Some(" a ".into()), None, Some("\u{1c}".into())]), vec![" a "])] + #[case::python_whitespace(EmailFacts::Text("\u{1c}a\u{1f},\u{85}b\u{a0}".into()), vec!["a", "b"])] + #[case::other(EmailFacts::Other, vec![])] + fn email_normalization(#[case] facts: EmailFacts, #[case] expected: Vec<&str>) { + assert_eq!(emails(&facts), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/lib.rs b/litellm-rust/crates/request-auth-core/src/lib.rs index c52896cfa10..592e7268c52 100644 --- a/litellm-rust/crates/request-auth-core/src/lib.rs +++ b/litellm-rust/crates/request-auth-core/src/lib.rs @@ -8,6 +8,7 @@ pub mod agent_grants; pub mod auth_context; pub mod budget_alerts; pub mod budget_context; +pub mod budget_lookup; pub mod budget_pricing; pub mod common_budgets; pub mod common_flow; diff --git a/litellm/_auth/budget_lookup.py b/litellm/_auth/budget_lookup.py new file mode 100644 index 00000000000..a980a1dee26 --- /dev/null +++ b/litellm/_auth/budget_lookup.py @@ -0,0 +1,90 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class OrganizationFacts: + token_present: bool + client_present: bool + team_present: bool + token_id: bool | None = None + team_id: bool | None = None + row_present: bool | None = None + lookup_failed: bool = False + + +OrganizationAction: TypeAlias = Literal["Skip", "ReadTokenId", "ReadTeamId", "ReadOrganization", "UseOrganization"] + + +@dataclass(frozen=True, slots=True) +class MemberFacts: + team_present: bool + token_present: bool + membership_loaded: bool + team_id: bool | None = None + user_id: bool | None = None + limits_read: bool = False + maximum_present: bool = False + default_id_string: bool = False + default_read: bool = False + + +MemberAction: TypeAlias = Literal[ + "Skip", "ReadTeamId", "ReadUserId", "ReadMembership", "ReadLimits", "ReadDefault", "UseLimits" +] +_ORGANIZATION: Final = TypeAdapter[OrganizationAction](OrganizationAction) +_MEMBER: Final = TypeAdapter[MemberAction](MemberAction) +_ORGANIZATION_FACTS: Final = TypeAdapter(OrganizationFacts) +_MEMBER_FACTS: Final = TypeAdapter(MemberFacts) +_ITEMS: Final = TypeAdapter(tuple[object, ...]) +_EMAILS: Final = TypeAdapter(tuple[str, ...]) +_METADATA: Final = TypeAdapter(Mapping[object, object]) + + +def soft_email_value(raw: object) -> object: + return _METADATA.validate_python(raw).get("soft_budget_alerting_emails") if isinstance(raw, dict) else None + + +@dataclass(frozen=True, slots=True) +class SoftFacts: + limit: float | None + spend: float | None + token_present: bool + + +_SOFT: Final = TypeAdapter(SoftFacts) + + +@native_boundary +def read_soft_emails(facts: SoftFacts) -> bool: + return _native.request_auth_budget_lookup_soft(_SOFT.dump_python(facts)) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def organization(facts: OrganizationFacts) -> OrganizationAction: + return _ORGANIZATION.validate_python( + _native.request_auth_budget_lookup_organization(_ORGANIZATION_FACTS.dump_python(facts)) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + ) + + +@native_boundary +def member(facts: MemberFacts) -> MemberAction: + return _MEMBER.validate_python(_native.request_auth_budget_lookup_member(_MEMBER_FACTS.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def emails(raw: object) -> tuple[str, ...]: + facts: Final = ( + {"Text": raw} + if isinstance(raw, str) + else {"List": tuple(value if isinstance(value, str) else None for value in _ITEMS.validate_python(raw))} + if isinstance(raw, list) + else "Other" + ) + return _EMAILS.validate_python(_native.request_auth_budget_lookup_emails(facts)) diff --git a/litellm/proxy/auth/budget_lookup_adapter.py b/litellm/proxy/auth/budget_lookup_adapter.py new file mode 100644 index 00000000000..1f12b0eb053 --- /dev/null +++ b/litellm/proxy/auth/budget_lookup_adapter.py @@ -0,0 +1,144 @@ +from collections.abc import Mapping +from dataclasses import dataclass, replace +from datetime import datetime +from typing import Final, assert_never + +from litellm._auth import budget_lookup as core +from litellm._auth.failure import NativeAuthFailure +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.proxy._types import ( + LiteLLM_OrganizationTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth.common_budget_adapter import metadata +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging + + +@dataclass(frozen=True, slots=True) +class LookupIO: + token: UserAPIKeyAuth | None + team: LiteLLM_TeamTable | None + prisma: PrismaClient | None + cache: UserApiKeyCache + logging: ProxyLogging + + +@dataclass(frozen=True, slots=True) +class OrganizationReading: + org_id: str + row: LiteLLM_OrganizationTable + + +async def organization( + facts: core.OrganizationFacts, io: LookupIO, org_id: str | None = None +) -> OrganizationReading | None: + from litellm.proxy.auth.auth_checks import get_org_object + + action: Final = core.organization(facts) + match action: + case "Skip": + return None + case "ReadTokenId": + assert io.token is not None + token_id: Final = io.token.org_id + return await organization(replace(facts, token_id=token_id is not None), io, token_id) + case "ReadTeamId": + assert io.team is not None + team_id: Final = io.team.organization_id + return await organization(replace(facts, team_id=team_id is not None), io, team_id) + case "ReadOrganization": + assert org_id is not None and io.prisma is not None + try: + row: Final = await get_org_object( + org_id=org_id, + prisma_client=io.prisma, + user_api_key_cache=io.cache, + proxy_logging_obj=io.logging, + include_budget_table=True, + ) + except NativeAuthFailure: + raise + except Exception: # noqa: BLE001 # pass the raw lookup failure to native Phase 1 policy + return await organization(replace(facts, lookup_failed=True), io, org_id) + result: Final = core.organization(replace(facts, row_present=row is not None)) + assert result in ("Skip", "UseOrganization") + if result == "Skip": + return None + assert row is not None + return OrganizationReading(org_id, row) + case "UseOrganization": + raise AssertionError("organization row was not retained") + case _: + assert_never(action) + + +@dataclass(frozen=True, slots=True) +class MemberReading: + membership: LiteLLM_TeamMembership | None + row: LiteLLM_BudgetTable | None + now: datetime + metadata: Mapping[str, object] + default_id: object + default: LiteLLM_BudgetTable | None = None + + +async def member( + facts: core.MemberFacts, + io: LookupIO, + membership: LiteLLM_TeamMembership | None, + reading: MemberReading | None = None, +) -> MemberReading | None: + from litellm.proxy.auth.auth_checks import get_team_member_default_budget, get_team_membership, get_utc_datetime + + action: Final = core.member(facts) + match action: + case "Skip": + return None + case "ReadTeamId": + assert io.team is not None + return await member(replace(facts, team_id=io.team.team_id is not None), io, membership) # pyright: ignore[reportUnnecessaryComparison] # Phase 1 accepts constructed cache rows with missing IDs + case "ReadUserId": + assert io.token is not None + return await member(replace(facts, user_id=io.token.user_id is not None), io, membership) + case "ReadMembership": + assert io.token is not None and io.token.user_id is not None and io.team is not None + loaded: Final = await get_team_membership( + user_id=io.token.user_id, + team_id=io.team.team_id, + prisma_client=io.prisma, + user_api_key_cache=io.cache, + proxy_logging_obj=io.logging, + ) + return await member(replace(facts, membership_loaded=True), io, loaded) + case "ReadLimits": + assert io.team is not None + row: Final = membership.litellm_budget_table if membership is not None else None + now: Final = get_utc_datetime() + values: Final = metadata(io.team.metadata) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # raw legacy metadata boundary + default_id: Final = values.get("team_member_budget_id") + snapshot: Final = MemberReading(membership, row, now, values, default_id) + return await member( + replace( + facts, + limits_read=True, + maximum_present=row is not None and row.max_budget is not None, + default_id_string=isinstance(default_id, str), + ), + io, + membership, + snapshot, + ) + case "ReadDefault": + assert reading is not None and isinstance(reading.default_id, str) + default: Final = await get_team_member_default_budget( + budget_id=reading.default_id, prisma_client=io.prisma, user_api_key_cache=io.cache + ) + return await member(replace(facts, default_read=True), io, membership, replace(reading, default=default)) + case "UseLimits": + assert reading is not None + return reading + case _: + assert_never(action) diff --git a/litellm/proxy/auth/common_budget_adapter.py b/litellm/proxy/auth/common_budget_adapter.py index 30420fc279e..0ca63ca0444 100644 --- a/litellm/proxy/auth/common_budget_adapter.py +++ b/litellm/proxy/auth/common_budget_adapter.py @@ -8,7 +8,7 @@ from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict import litellm -from litellm._auth import common_budgets +from litellm._auth import budget_lookup, common_budgets from litellm._auth.common_budgets import AlertFacts, BudgetDecision, BudgetFacts, MemberLimitFacts from litellm.constants import EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE from litellm.models.budget import LiteLLM_BudgetTable @@ -24,7 +24,6 @@ from litellm.proxy._types import ( from litellm.proxy.utils import ProxyLogging _METADATA: Final = TypeAdapter(Mapping[str, object]) -_ITEMS: Final = TypeAdapter(tuple[object, ...]) AlertKind: TypeAlias = Literal[ "token_budget", "user_budget", @@ -302,11 +301,7 @@ async def virtual_key_soft_budget_check( def _soft_emails(raw: object) -> tuple[str, ...]: - if isinstance(raw, str): - return tuple(email.strip() for email in raw.split(",") if email.strip()) - if isinstance(raw, list): - return tuple(email for email in _ITEMS.validate_python(raw) if isinstance(email, str) and email.strip()) - return () + return budget_lookup.emails(raw) async def team_soft_budget_check( @@ -316,7 +311,11 @@ async def team_soft_budget_check( ) -> None: if team_object is None: return - emails: Final = _soft_emails(metadata(team_object.metadata).get("soft_budget_alerting_emails")) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # validated legacy metadata boundary + if not budget_lookup.read_soft_emails( + budget_lookup.SoftFacts(team_object.soft_budget, team_object.spend, valid_token is not None) + ): + return + emails: Final = _soft_emails(budget_lookup.soft_email_value(team_object.metadata)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # raw legacy metadata projection result: Final = common_budgets.check( BudgetFacts( "TeamSoft", diff --git a/litellm/proxy/auth/common_budget_lookups.py b/litellm/proxy/auth/common_budget_lookups.py index 7c9e473f5b4..cf93ce95a33 100644 --- a/litellm/proxy/auth/common_budget_lookups.py +++ b/litellm/proxy/auth/common_budget_lookups.py @@ -2,9 +2,8 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import Final -from litellm._auth import common_budgets +from litellm._auth import budget_lookup, common_budgets from litellm._auth.common_budgets import BudgetFacts -from litellm._auth.failure import NativeAuthFailure from litellm.proxy._types import ( Litellm_EntityType, LiteLLM_TagTable, @@ -13,11 +12,11 @@ from litellm.proxy._types import ( LiteLLM_UserTable, UserAPIKeyAuth, ) +from litellm.proxy.auth import budget_lookup_adapter from litellm.proxy.auth.common_budget_adapter import ( CounterRead, emit_alert, key_alert_info, - metadata, raise_budget, read_budget, resolve_effective_team_member_budget, @@ -35,31 +34,15 @@ async def organization_max_budget_check( user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ) -> None: - from litellm.proxy.auth.auth_checks import get_org_object - - if valid_token is None or prisma_client is None: - return - org_id: Final = ( - valid_token.org_id - if valid_token.org_id is not None - else (team_object.organization_id if team_object is not None else None) + reading: Final = await budget_lookup_adapter.organization( + budget_lookup.OrganizationFacts(valid_token is not None, prisma_client is not None, team_object is not None), + budget_lookup_adapter.LookupIO(valid_token, team_object, prisma_client, user_api_key_cache, proxy_logging_obj), ) - if org_id is None: - return - try: - org: Final = await get_org_object( - org_id=org_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - include_budget_table=True, - ) - except NativeAuthFailure: - raise - except Exception: # noqa: BLE001 # preserve organization lookup failure handling - return - if org is None: + if reading is None: return + assert valid_token is not None + org_id: Final = reading.org_id + org: Final = reading.row carry_organization_budget_state(valid_token=valid_token, org_table=org) limit: Final = org.litellm_budget_table.max_budget if org.litellm_budget_table is not None else None facts, result = await read_budget( @@ -164,38 +147,19 @@ async def check_team_member_budget( team_membership: LiteLLM_TeamMembership | None = None, team_membership_loaded: bool = False, ) -> None: - from litellm.proxy.auth.auth_checks import ( - get_team_member_default_budget, - get_team_membership, - get_utc_datetime, - ) from litellm.proxy.auth.budget_alert_adapter import team_member_max_budget_alert_check - if team_object is None or team_object.team_id is None or valid_token is None or valid_token.user_id is None: # pyright: ignore[reportUnnecessaryComparison] # retain the legacy None guard for constructed cache rows + reading: Final = await budget_lookup_adapter.member( + budget_lookup.MemberFacts(team_object is not None, valid_token is not None, team_membership_loaded), + budget_lookup_adapter.LookupIO(valid_token, team_object, prisma_client, user_api_key_cache, proxy_logging_obj), + team_membership, + ) + if reading is None: return - membership: Final = ( - team_membership - if team_membership_loaded - else await get_team_membership( - user_id=valid_token.user_id, - team_id=team_object.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - ) - row: Final = membership.litellm_budget_table if membership is not None else None - now: Final = get_utc_datetime() - team_metadata: Final = metadata(team_object.metadata) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # typed projection of legacy metadata - default_id: Final = team_metadata.get("team_member_budget_id") - default: Final = ( - await get_team_member_default_budget( - budget_id=default_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache - ) - if (row is None or row.max_budget is None) and isinstance(default_id, str) - else None - ) - limit: Final = resolve_effective_team_member_budget(row, default, now) + assert team_object is not None and valid_token is not None and valid_token.user_id is not None + membership: Final = reading.membership + team_metadata: Final = reading.metadata + limit: Final = resolve_effective_team_member_budget(reading.row, reading.default, reading.now) facts, result = await read_budget( BudgetFacts("TeamMember", limit), CounterRead( diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 3b155741ed1..c3dcc77475c 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -884,6 +884,11 @@ def request_auth_end_user_lookup(facts: object) -> object: ... def request_auth_end_user_validation(facts: object) -> object: ... +def request_auth_budget_lookup_organization(facts: object) -> object: ... +def request_auth_budget_lookup_member(facts: object) -> object: ... +def request_auth_budget_lookup_emails(facts: object) -> object: ... +def request_auth_budget_lookup_soft(facts: object) -> bool: ... + def request_auth_route_role(facts: object) -> object: ... def request_auth_route_role_api(facts: object) -> object: ... diff --git a/tests/unit/proxy/auth/phase1_budget_lookup_reference.py b/tests/unit/proxy/auth/phase1_budget_lookup_reference.py new file mode 100644 index 00000000000..6d11b7f398c --- /dev/null +++ b/tests/unit/proxy/auth/phase1_budget_lookup_reference.py @@ -0,0 +1,341 @@ +"""Unchanged Phase 1 budget lookup policy with raw IO supplied by tests.""" + +from __future__ import annotations + +import asyncio +import math +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ( + CallInfo, + Litellm_EntityType, + LiteLLM_OrganizationTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.spend_tracking.carried_budget_state import carry_organization_budget_state +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.utils import get_utc_datetime + +from .phase1_budget_auth_reference import ( + get_team_member_default_budget, + get_team_membership, + resolve_effective_team_member_budget, +) + +TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY: Final = "team_member_max_budget_alert_emails" +_TEAM_MEMBER_ALERT_CONFIG_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +async def get_org_object( + org_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + include_budget_table: bool, +) -> LiteLLM_OrganizationTable | None: + raise AssertionError("Frozen reference raw IO must be supplied by the test: get_org_object") + + +def _parse_email_list(raw: str | Sequence[object] | None) -> list[str]: # mutable-ok: unchanged Phase 1 collection contract in independent oracle + """Parse emails from a list or comma-separated string.""" + if isinstance(raw, list): + return [e.strip() for e in raw if isinstance(e, str) and e.strip()] + elif isinstance(raw, str): + return [e.strip() for e in raw.split(",") if e.strip()] + return [] + + +def _merge_budget_alert_email_configs( + global_cfg: Mapping[str, str | Sequence[object] | None] | None, + per_key_cfg: Mapping[str, str | Sequence[object] | None] | None, +) -> dict[str, list[str]] | None: # mutable-ok: unchanged Phase 1 collection contract in independent oracle + """ + Per-threshold additive merge: each threshold's recipient list is the union + of global + per-key entries (deduped, global-first ordering). Missing + thresholds on one side are inherited from the other. + """ + global_cfg_normalized: Final = _normalize_alert_emails(global_cfg) + per_key_cfg_normalized: Final = _normalize_alert_emails(per_key_cfg) + if not global_cfg_normalized and (not per_key_cfg_normalized): + return None + thresholds: Final = set(global_cfg_normalized) | set(per_key_cfg_normalized) + return { + t: list(dict.fromkeys(global_cfg_normalized.get(t, []) + per_key_cfg_normalized.get(t, []))) for t in thresholds + } + + +def _is_valid_alert_threshold_pct(pct: str) -> bool: + return pct.isdigit() and len(pct) <= 3 and (1 <= int(pct) <= 100) + + +def _alert_recipients(raw: object) -> Sequence[str] | None: + if isinstance(raw, (str, Sequence)): + return _parse_email_list(raw) # pyright: ignore[reportUnknownArgumentType] # preserve unvalidated Phase 1 oracle values and call shape + return None + + +def _valid_alert_threshold_config(raw_config: object) -> Mapping[str, str | Sequence[object] | None] | None: + try: + config: Final = _TEAM_MEMBER_ALERT_CONFIG_ADAPTER.validate_python(raw_config) + except ValidationError: + return None + return MappingProxyType( + {pct: _alert_recipients(emails) for pct, emails in config.items() if _is_valid_alert_threshold_pct(pct)} + ) + + +def team_member_max_budget_alert_check( + team_id: str, + team_alias: str | None, + team_metadata: Mapping[str, object] | None, + organization_id: str | None, + user_id: str, + user_email: str | None, + proxy_logging_obj: ProxyLogging, + spend: float, + max_budget: float, +) -> None: + raw_config: Final = team_metadata.get(TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY) if team_metadata else None + alert_email_config: Final = _merge_budget_alert_email_configs( + global_cfg=None, per_key_cfg=_valid_alert_threshold_config(raw_config) + ) + if not alert_email_config or spend <= 0: + return + min_pct: Final = min(int(pct) for pct in alert_email_config) + if spend < max_budget * (min_pct / 100.0): + return + call_info: Final = CallInfo( + spend=spend, + max_budget=max_budget, + user_id=user_id, + team_id=team_id, + team_alias=team_alias, + organization_id=organization_id, + user_email=user_email, + event_group=Litellm_EntityType.TEAM_MEMBER, + max_budget_alert_emails=alert_email_config, + ) + asyncio.create_task(proxy_logging_obj.budget_alerts(type="max_budget_alert", user_info=call_info)) + + +async def _check_team_member_budget( # pyright: ignore[reportUnusedFunction] # preserve unvalidated Phase 1 oracle values and call shape + team_object: LiteLLM_TeamTable | None, + user_object: LiteLLM_UserTable | None, + valid_token: UserAPIKeyAuth | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + team_membership: LiteLLM_TeamMembership | None = None, + team_membership_loaded: bool = False, +): + """Check if team member is over their max budget within the team.""" + if ( + team_object is not None + and team_object.team_id is not None # pyright: ignore[reportUnnecessaryComparison] # preserve unvalidated Phase 1 oracle values and call shape + and (valid_token is not None) + and (valid_token.user_id is not None) + ): + if not team_membership_loaded: + team_membership = await get_team_membership( # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + user_id=valid_token.user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + loaded_membership = team_membership # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + member_budget_row: Final = loaded_membership.litellm_budget_table if loaded_membership is not None else None + now: Final = get_utc_datetime() + default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # preserve unvalidated Phase 1 oracle values and call shape + team_default_budget: Final = ( + await get_team_member_default_budget( + budget_id=default_budget_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache + ) + if (member_budget_row is None or member_budget_row.max_budget is None) + and isinstance(default_budget_id, str) + else None + ) + team_member_budget: Final = resolve_effective_team_member_budget( + member_budget=member_budget_row, team_default_budget=team_default_budget, now=now + ) + if team_member_budget is not None: + team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0 # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + from litellm.proxy.proxy_server import get_current_spend + + team_member_spend = await get_current_spend( # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + counter_key=f"spend:team_member:{valid_token.user_id}:{team_object.team_id}", + fallback_spend=team_member_spend, + max_budget=team_member_budget, + ) + if not math.isfinite(team_member_budget): + return + team_member_max_budget_alert_check( + team_id=team_object.team_id, + team_alias=team_object.team_alias, + team_metadata=team_object.metadata, # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # preserve unvalidated Phase 1 oracle values and call shape + organization_id=team_object.organization_id, + user_id=valid_token.user_id, + user_email=user_object.user_email if user_object is not None else None, + proxy_logging_obj=proxy_logging_obj, + spend=team_member_spend, + max_budget=team_member_budget, + ) + if team_member_spend >= team_member_budget: + raise litellm.BudgetExceededError( + current_cost=team_member_spend, + max_budget=team_member_budget, + message=f"Budget has been exceeded! User={valid_token.user_id} in Team={team_object.team_id} Current cost: {team_member_spend}, Max budget: {team_member_budget}", + entity_type=Litellm_EntityType.TEAM_MEMBER.value, + entity_id=f"{valid_token.user_id}:{team_object.team_id}", + ) + + +async def _team_soft_budget_check( # pyright: ignore[reportUnusedFunction] # preserve unvalidated Phase 1 oracle values and call shape + team_object: LiteLLM_TeamTable | None, valid_token: UserAPIKeyAuth | None, proxy_logging_obj: ProxyLogging +): + """ + Triggers a budget alert if the team is over it's soft budget. + """ + if ( + team_object is not None + and team_object.soft_budget is not None + and (team_object.spend is not None) + and (team_object.spend >= team_object.soft_budget) + ): + verbose_proxy_logger.debug( + "Crossed Soft Budget for team %s, spend %s, soft_budget %s", + team_object.team_id, + team_object.spend, + team_object.soft_budget, + ) + if valid_token: + alert_emails: list[str] | None = None # mutable-ok: unchanged Phase 1 collection contract in independent oracle # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + if team_object.metadata is not None and isinstance(team_object.metadata, dict): # pyright: ignore[reportUnknownMemberType, reportUnnecessaryIsInstance] # preserve unvalidated Phase 1 oracle values and call shape + soft_budget_alert_emails: Final = team_object.metadata.get("soft_budget_alerting_emails") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # preserve unvalidated Phase 1 oracle values and call shape + if soft_budget_alert_emails is not None: + if isinstance(soft_budget_alert_emails, list): + alert_emails = [ # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + email for email in soft_budget_alert_emails if isinstance(email, str) and email.strip() # pyright: ignore[reportUnknownVariableType] # preserve unvalidated Phase 1 oracle values and call shape + ] + elif isinstance(soft_budget_alert_emails, str): + alert_emails = [email.strip() for email in soft_budget_alert_emails.split(",") if email.strip()] # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + if alert_emails: + alert_emails = [email for email in alert_emails if email] # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + else: + alert_emails = None # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + if alert_emails is None or len(alert_emails) == 0: + verbose_proxy_logger.debug( + "Skipping team soft budget alert for team %s: no alert_emails configured in metadata.soft_budget_alerting_emails", + team_object.team_id, + ) + return + call_info: Final = CallInfo( + token=valid_token.token, + spend=team_object.spend, + max_budget=team_object.max_budget, + soft_budget=team_object.soft_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + user_email=None, + key_alias=valid_token.key_alias, + event_group=Litellm_EntityType.TEAM, + alert_emails=alert_emails, + ) + asyncio.create_task(proxy_logging_obj.budget_alerts(type="soft_budget", user_info=call_info)) + + +async def organization_max_budget_check( + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> None: + """ + Check if the organization is over its max budget. + + This function checks the organization budget using: + 1. First, tries to use valid_token.org_id (if key has organization_id set) + 2. Falls back to team_object.organization_id (if key doesn't have org_id but team does) + + This ensures organization budget checks work even when keys don't have organization_id + set directly, as long as their team belongs to an organization. + + Raises: + BudgetExceededError if the organization is over its max budget. + Triggers a budget alert if the organization is over its max budget. + """ + if valid_token is None or prisma_client is None: + return + org_id: str | None = None # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + if valid_token.org_id is not None: + org_id = valid_token.org_id # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + elif team_object is not None and team_object.organization_id is not None: + org_id = team_object.organization_id # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + if org_id is None: + return + try: + org_table: Final = await get_org_object( + org_id=org_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + include_budget_table=True, + ) + except Exception: + return + if org_table is None: + return + carry_organization_budget_state(valid_token=valid_token, org_table=org_table) + org_max_budget: float | None = None # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + if org_table.litellm_budget_table is not None: + org_max_budget = org_table.litellm_budget_table.max_budget # rebind-ok: preserve exact Phase 1 evaluation and mutation sequence + if org_max_budget is None: + return + from litellm.proxy.proxy_server import get_current_spend + + org_spend: Final = await get_current_spend( + counter_key=f"spend:org:{org_id}", fallback_spend=org_table.spend or 0.0, max_budget=org_max_budget + ) + if math.isfinite(org_max_budget) and org_spend >= org_max_budget: + call_info: Final = CallInfo( + token=valid_token.token, + spend=org_spend, + max_budget=org_max_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=org_id, + event_group=Litellm_EntityType.ORGANIZATION, + ) + asyncio.create_task(proxy_logging_obj.budget_alerts(type="organization_budget", user_info=call_info)) + raise litellm.BudgetExceededError( + current_cost=org_spend, + max_budget=org_max_budget, + message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_spend}, Max budget: {org_max_budget}", + entity_type=Litellm_EntityType.ORGANIZATION.value, + entity_id=org_id, + ) + + +def _normalize_alert_emails(cfg: Mapping[str, str | Sequence[object] | None] | None) -> dict[str, list[str]]: # mutable-ok: unchanged Phase 1 collection contract in independent oracle + """Coerce user-supplied threshold→recipients mapping to Dict[str, List[str]]. + + Values may legitimately arrive as list, comma-separated string, or None + from YAML/metadata; _parse_email_list tolerates all three. + """ + if not cfg: + return {} + return {k: _parse_email_list(v) for k, v in cfg.items()} diff --git a/tests/unit/proxy/auth/test_budget_lookup_adapter.py b/tests/unit/proxy/auth/test_budget_lookup_adapter.py new file mode 100644 index 00000000000..38707a68983 --- /dev/null +++ b/tests/unit/proxy/auth/test_budget_lookup_adapter.py @@ -0,0 +1,276 @@ +from collections.abc import Mapping +from datetime import datetime, timezone +from math import inf, isnan, nan +from types import MappingProxyType +from typing import Final +from unittest.mock import AsyncMock, Mock + +import pytest +from pydantic import TypeAdapter + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.proxy import proxy_server +from litellm.proxy._types import LiteLLM_OrganizationTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.auth import auth_checks, common_budget_adapter, common_budget_lookups +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging + +from . import phase1_budget_lookup_reference as reference +from .test_common_budget_adapter import observe + +_NOW: Final = datetime(2025, 1, 1, tzinfo=timezone.utc) +_CALL: Final = TypeAdapter(tuple[str, tuple[object, ...], Mapping[str, object]]) + + +def normalized_calls(events: Mock) -> tuple[tuple[str, tuple[object, ...], Mapping[str, object]], ...]: + return tuple( + ( + name, + args, + MappingProxyType( + {key: "NaN" if isinstance(value, float) and isnan(value) else value for key, value in kwargs.items()} + ), + ) + for name, args, kwargs in (_CALL.validate_python(value) for value in events.mock_calls) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("token_id", "team_id"), ((None, None), (None, "team-org"), ("", "team-org"), ("key-org", "team-org")) +) +@pytest.mark.parametrize(("token_present", "client_present"), ((True, True), (False, True), (True, False))) +@pytest.mark.parametrize("row_state", ("present", "missing", "failed")) +@pytest.mark.parametrize("limit", (None, 1.0, inf)) +async def test_organization_source_and_raw_reads_match_phase1( + monkeypatch: pytest.MonkeyPatch, + token_id: str | None, + team_id: str | None, + token_present: bool, + client_present: bool, + row_state: str, + limit: float | None, +) -> None: + row: Final = LiteLLM_OrganizationTable.model_construct( + organization_id="row", spend=2.0, budget_id="budget", litellm_budget_table=LiteLLM_BudgetTable(max_budget=limit) + ) + reader: Final = AsyncMock( + return_value=row if row_state == "present" else None, + side_effect=RuntimeError("unavailable") if row_state == "failed" else None, + ) + counter: Final = AsyncMock(return_value=2.0) + monkeypatch.setattr(reference, "get_org_object", reader) + monkeypatch.setattr(auth_checks, "get_org_object", reader) + monkeypatch.setattr(proxy_server, "get_current_spend", counter) + old: Final = UserAPIKeyAuth(org_id=token_id, token="token", key_alias="alias") + new: Final = old.model_copy(deep=True) + team: Final = LiteLLM_TeamTable(team_id="team", organization_id=team_id) + prisma: Final = Mock(spec=PrismaClient) if client_present else None + cache: Final = UserApiKeyCache() + alerts: Final = AsyncMock() + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=alerts) + new_log: Final = old_log + before: Final = await observe( + lambda: reference.organization_max_budget_check(old if token_present else None, team, prisma, cache, old_log), + old_log, + old, + ) + old_reads: Final = (tuple(reader.await_args_list), tuple(counter.await_args_list)) + reader.reset_mock() + counter.reset_mock() + alerts.reset_mock() + after: Final = await observe( + lambda: common_budget_lookups.organization_max_budget_check( + new if token_present else None, team, prisma, cache, new_log + ), + new_log, + new, + ) + assert after == before + assert (tuple(reader.await_args_list), tuple(counter.await_args_list)) == old_reads + assert new.model_dump() == old.model_dump() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("loaded", (False, True)) +@pytest.mark.parametrize("member_limit", (None, 0.0, 1.0, inf, nan)) +@pytest.mark.parametrize("default_id", (None, "", "default", 1)) +@pytest.mark.parametrize("missing_member", (False, True)) +async def test_member_default_and_counter_order_match_phase1( + monkeypatch: pytest.MonkeyPatch, + loaded: bool, + member_limit: float | None, + default_id: object, + missing_member: bool, +) -> None: + row: Final = LiteLLM_TeamMembership( + team_id="team", user_id="user", spend=2.0, litellm_budget_table=LiteLLM_BudgetTable(max_budget=member_limit) + ) + membership: Final = None if missing_member else row + member_read: Final = AsyncMock(return_value=membership) + default_read: Final = AsyncMock(return_value=LiteLLM_BudgetTable(max_budget=1)) + counter: Final = AsyncMock(return_value=2.0) + clock: Final = Mock(return_value=_NOW) + events: Final = Mock() + events.attach_mock(member_read, "membership") + events.attach_mock(default_read, "default") + events.attach_mock(counter, "counter") + events.attach_mock(clock, "clock") + for module in (reference, auth_checks): + monkeypatch.setattr(module, "get_team_membership", member_read) + monkeypatch.setattr(module, "get_team_member_default_budget", default_read) + monkeypatch.setattr(module, "get_utc_datetime", clock) + monkeypatch.setattr(proxy_server, "get_current_spend", counter) + team: Final = LiteLLM_TeamTable( + team_id="team", + metadata={ + "team_member_budget_id": default_id, + "team_member_max_budget_alert_emails": {"50": ["owner"]}, + }, + ) + key: Final = UserAPIKeyAuth(user_id="user") + prisma: Final = Mock(spec=PrismaClient) + cache: Final = UserApiKeyCache() + alerts: Final = AsyncMock() + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=alerts) + new_log: Final = old_log + before: Final = await observe( + lambda: reference._check_team_member_budget(team, None, key, prisma, cache, old_log, membership, loaded), # pyright: ignore[reportPrivateUsage] # exact Phase 1 reference + old_log, + key, + ) + old_reads: Final = normalized_calls(events) + events.reset_mock() + alerts.reset_mock() + after: Final = await observe( + lambda: common_budget_lookups.check_team_member_budget( + team, None, key, prisma, cache, new_log, membership, loaded + ), + new_log, + key, + ) + assert after == before + assert normalized_calls(events) == old_reads + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "raw", (None, "", " a,b , ,a", "\x1ca\x1f,\u0085b\u00a0", [" a ", "", "\x1c", 4, "a"], ("a",), 4, {"a": 1}) +) +@pytest.mark.parametrize( + ("limit", "spend", "token"), + ((1, 2, True), (1, 0, True), (0, 0, True), (1, 2, False), (None, 2, True), (nan, 2, True)), +) +async def test_soft_email_normalization_matches_phase1( + raw: object, + limit: float | None, + spend: float, + token: bool, +) -> None: + team: Final = LiteLLM_TeamTable( + team_id="team", soft_budget=limit, spend=spend, metadata={"soft_budget_alerting_emails": raw} + ) + key: Final = UserAPIKeyAuth() + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + new_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + before: Final = await observe( + lambda: reference._team_soft_budget_check(team, key if token else None, old_log), # pyright: ignore[reportPrivateUsage] # exact Phase 1 reference + old_log, + key, + ) + after: Final = await observe( + lambda: common_budget_adapter.team_soft_budget_check(team, key if token else None, new_log), + new_log, + key, + ) + assert after == before + + +@pytest.mark.asyncio +async def test_organization_native_failure_from_lookup_is_terminal(monkeypatch: pytest.MonkeyPatch) -> None: + failure: Final = NativeAuthRuntimeError("native failed") + reader: Final = AsyncMock(side_effect=failure) + monkeypatch.setattr(auth_checks, "get_org_object", reader) + with pytest.raises(NativeAuthRuntimeError) as raised: + await common_budget_lookups.organization_max_budget_check( + UserAPIKeyAuth(org_id="org"), None, Mock(spec=PrismaClient), UserApiKeyCache(), Mock(spec=ProxyLogging) + ) + assert raised.value is failure + assert reader.await_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_present", "team_id", "token_present", "user_id"), + ( + (False, "team", True, "user"), + (True, None, True, "user"), + (True, "team", False, "user"), + (True, "team", True, None), + (True, "", True, ""), + ), +) +async def test_member_identity_stages_match_phase1( + monkeypatch: pytest.MonkeyPatch, + team_present: bool, + team_id: str | None, + token_present: bool, + user_id: str | None, +) -> None: + team: Final = LiteLLM_TeamTable.model_construct(team_id=team_id) if team_present else None + key: Final = UserAPIKeyAuth(user_id=user_id) + reader: Final = AsyncMock(return_value=None) + clock: Final = Mock(return_value=_NOW) + for module in (reference, auth_checks): + monkeypatch.setattr(module, "get_team_membership", reader) + monkeypatch.setattr(module, "get_utc_datetime", clock) + logging: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + cache: Final = UserApiKeyCache() + before: Final = await observe( + lambda: reference._check_team_member_budget(team, None, key if token_present else None, None, cache, logging), # pyright: ignore[reportPrivateUsage] # exact Phase 1 reference + logging, + key, + ) + reads: Final = (tuple(reader.await_args_list), clock.call_count) + reader.reset_mock() + clock.reset_mock() + after: Final = await observe( + lambda: common_budget_lookups.check_team_member_budget( + team, None, key if token_present else None, None, cache, logging + ), + logging, + key, + ) + assert after == before + assert (tuple(reader.await_args_list), clock.call_count) == reads + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native", (False, True)) +async def test_soft_budget_skip_does_not_read_metadata(native: bool) -> None: + team: Final = LiteLLM_TeamTable(team_id="team", soft_budget=10, spend=0) + object.__setattr__(team, "metadata", object()) + alerts: Final = AsyncMock() + logging: Final = Mock(spec=ProxyLogging, budget_alerts=alerts) + policy: Final = common_budget_adapter.team_soft_budget_check if native else reference._team_soft_budget_check # pyright: ignore[reportPrivateUsage] # exact Phase 1 reference + result: Final = await policy(team, UserAPIKeyAuth(), logging) + assert result is None + alerts.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("raw", (None, [], 1, "metadata", {1: "ignored", "soft_budget_alerting_emails": " a "})) +async def test_soft_metadata_shape_preserves_phase1_projection(raw: object) -> None: + team: Final = LiteLLM_TeamTable(team_id="team", soft_budget=1, spend=2) + object.__setattr__(team, "metadata", raw) + key: Final = UserAPIKeyAuth() + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + new_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + before: Final = await observe( + lambda: reference._team_soft_budget_check(team, key, old_log), # pyright: ignore[reportPrivateUsage] # exact Phase 1 reference + old_log, + key, + ) + after: Final = await observe(lambda: common_budget_adapter.team_soft_budget_check(team, key, new_log), new_log, key) + assert after == before