diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 844d3c2ed26..43921a847a9 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3320,7 +3320,24 @@ class PrismaClient: elif query_type == "find_all" and reset_at is not None: response = await UserRepository(self).table.find_many( where={ # type: ignore - "budget_reset_at": {"lt": reset_at}, + # A user seeded from default_internal_user_params + # (or created via /user/new without an explicit + # budget_reset_at) has budget_duration set but + # budget_reset_at = NULL. `{"lt": reset_at}` never + # matches NULL, so such users would never be reset + # and their spend would accumulate for the lifetime + # of the row, silently exceeding max_budget. Treat a + # NULL budget_reset_at with a non-NULL budget_duration + # as due, matching the budget-table query below. + "OR": [ + { + "AND": [ + {"budget_reset_at": None}, + {"NOT": {"budget_duration": None}}, + ] + }, + {"budget_reset_at": {"lt": reset_at}}, + ], } ) elif query_type == "find_all" and user_id_list is not None: @@ -3406,7 +3423,18 @@ class PrismaClient: elif query_type == "find_all" and reset_at is not None: response = await TeamRepository(self).table.find_many( where={ # type: ignore - "budget_reset_at": {"lt": reset_at}, + # Same NULL budget_reset_at gap as the user query + # above: a team with a budget_duration but no + # initialized budget_reset_at would never be reset. + "OR": [ + { + "AND": [ + {"budget_reset_at": None}, + {"NOT": {"budget_duration": None}}, + ] + }, + {"budget_reset_at": {"lt": reset_at}}, + ], } ) elif query_type == "find_all" and user_id is not None: diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 9eba7c1277c..086a78b64be 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -94,7 +94,17 @@ class SafeAttributeModel: """ def __delattr__(self, name): + # Dropping an unset optional field stored in __dict__ goes straight to + # object.__delattr__, skipping pydantic's __delattr__ whose per-call + # class getattr lookup and _check_frozen dominate response construction. try: + if ( + name in type(self).__pydantic_fields__ + and name in self.__dict__ + and not type(self).model_config.get("frozen") + ): + object.__delattr__(self, name) + return super().__delattr__(name) except AttributeError: # noop if attribute does not exist diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 0b683745369..5e348b1bb7e 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -1803,3 +1803,71 @@ def test_reset_budget_for_tags_linked_to_budgets_management_cache_delete_failure asyncio.run(job.reset_budget_for_tags_linked_to_budgets([expired_budget])) prisma_client.db.litellm_tagtable.update_many.assert_awaited_once() + + +def _extract_reset_where(find_many_mock): + """Return the ``where`` dict passed to a mocked repository ``find_many``.""" + assert find_many_mock.await_count == 1 + _, kwargs = find_many_mock.await_args + return kwargs["where"] + + +def _asserts_null_reset_is_due(where): + """A budget-reset ``find_many`` filter must select rows whose + ``budget_reset_at`` is NULL but which have a ``budget_duration`` set, in + addition to rows whose ``budget_reset_at`` is already in the past. + + Regression guard: a user/team seeded from ``default_internal_user_params`` + (or created via ``/user/new`` without an explicit ``budget_reset_at``) has + ``budget_duration`` set but ``budget_reset_at = NULL``. A plain + ``{"budget_reset_at": {"lt": now}}`` filter never matches NULL, so such rows + would never be reset and their spend would accumulate for the lifetime of + the row, silently exceeding ``max_budget``. + """ + branches = where.get("OR") + assert isinstance(branches, list), f"expected an OR filter, got {where!r}" + + has_null_branch = any( + b.get("AND") + == [ + {"budget_reset_at": None}, + {"NOT": {"budget_duration": None}}, + ] + for b in branches + if isinstance(b, dict) + ) + has_expired_branch = any( + isinstance(b, dict) + and "budget_reset_at" in b + and b["budget_reset_at"] is not None + for b in branches + ) + assert has_null_branch, f"missing NULL-reset_at branch in {where!r}" + assert has_expired_branch, f"missing expired-reset_at branch in {where!r}" + + +@pytest.mark.parametrize("table_name", ["user", "team"]) +def test_get_data_reset_query_selects_null_budget_reset_at(table_name): + """``PrismaClient.get_data(..., reset_at=...)`` for the user and team tables + must select rows with a NULL ``budget_reset_at`` (and a non-NULL + ``budget_duration``), matching the budget-table query. Without this, users + auto-created from ``default_internal_user_params`` are never reset.""" + from litellm.proxy.utils import PrismaClient + + # Build a PrismaClient without running its heavy __init__; only .db is used. + client = PrismaClient.__new__(PrismaClient) + client.db = MagicMock() + + find_many = AsyncMock(return_value=[]) + table_attr = { + "user": "litellm_usertable", + "team": "litellm_teamtable", + }[table_name] + setattr(getattr(client.db, table_attr), "find_many", find_many) + + now = datetime.now(timezone.utc) + asyncio.run( + client.get_data(table_name=table_name, query_type="find_all", reset_at=now) + ) + + _asserts_null_reset_is_due(_extract_reset_where(find_many)) diff --git a/tests/test_litellm/types/test_types_utils.py b/tests/test_litellm/types/test_types_utils.py index 98820d657b5..21f28f54b8b 100644 --- a/tests/test_litellm/types/test_types_utils.py +++ b/tests/test_litellm/types/test_types_utils.py @@ -526,3 +526,103 @@ def test_delta_serialization_contract(): keys = list(extra_dump.keys()) assert extra_dump["custom_field"] == "v" assert keys.index("custom_field") < keys.index("content") + + +def test_safe_attribute_model_delattr(): + """ + SafeAttributeModel.__delattr__ must remove a field from the instance so it + is omitted from model_dump (OpenAI spec), whether the field is a declared + model field or an extra, and deleting a missing attribute must be a no-op. + """ + from litellm.types.utils import Message + + # Unset optional declared fields are dropped during __init__ -> absent from dump + msg = Message(content="hi", role="assistant") + assert not hasattr(msg, "audio") + assert not hasattr(msg, "reasoning_content") + assert "audio" not in msg.model_dump() + assert "reasoning_content" not in msg.model_dump() + + # Explicitly deleting a present declared field removes it from the dump + msg2 = Message(content="hi", role="assistant", reasoning_content="because") + assert msg2.reasoning_content == "because" + del msg2.reasoning_content + assert not hasattr(msg2, "reasoning_content") + assert "reasoning_content" not in msg2.model_dump() + + # Extra fields (extra='allow') are still deletable via the fallback path + msg3 = Message(content="hi", role="assistant", custom_field=123) + assert msg3.custom_field == 123 + del msg3.custom_field + assert not hasattr(msg3, "custom_field") + assert "custom_field" not in msg3.model_dump() + + # Deleting a non-existent attribute is a silent no-op + msg4 = Message(content="hi", role="assistant") + del msg4.does_not_exist + + +def test_delattr_fast_path_matches_pydantic_exactly(): + """ + The fast path must be observationally identical to pydantic's own + __delattr__ for a declared field, including model_fields_set membership and + the exclude_unset dump, both of which the fast path never touches. Deleting + the same field through the fast path and through pydantic's __delattr__ + (reached by skipping SafeAttributeModel in the MRO) must leave identical + state, so if a future pydantic release makes __delattr__ mutate + __pydantic_fields_set__ the two diverge and this fails rather than silently + shifting the serialization contract. + """ + from litellm.types.utils import Message, SafeAttributeModel + + def observe(m: Message) -> tuple: + return ( + hasattr(m, "reasoning_content"), + "reasoning_content" in m.model_fields_set, + "reasoning_content" in m.model_dump(), + "reasoning_content" in m.model_dump(exclude_unset=True), + ) + + fast = Message(content="hi", role="assistant", reasoning_content="x") + del fast.reasoning_content + + control = Message(content="hi", role="assistant", reasoning_content="x") + super(SafeAttributeModel, control).__delattr__("reasoning_content") + + assert observe(fast) == observe(control) + # A deleted field is gone from __dict__ (so absent from both dumps) yet + # stays in model_fields_set, since neither delete path clears fields_set. + assert observe(fast) == (False, True, False, False) + + +def test_delattr_fast_path_missing_attribute_is_noop(): + """ + The declared-field fast path must stay a silent no-op when the object delete + fails: the field passes the __dict__ membership guard but is already gone by + the time object.__delattr__ runs. This models a concurrent removal of the same + field on a shared response object. Previously the fast-path delete ran outside + the AttributeError handler, so the error leaked onto the Message/Delta/Choices/ + Usage construction hot path instead of being swallowed like the slow path. + + _VanishingDict reports every key as present (passing the guard) while storing + nothing, so the real object.__delattr__ still raises AttributeError. + """ + from litellm.types.utils import SafeAttributeModel + + class _VanishingDict(dict): + def __contains__(self, key: object) -> bool: + return True + + class _RacyModel(SafeAttributeModel): + __pydantic_fields__ = {"x": object()} + model_config: dict = {} + + def __init__(self) -> None: + self.__dict__ = _VanishingDict() + + racy = _RacyModel() + assert "x" in racy.__dict__ + assert "x" not in dict.keys(racy.__dict__) + + del racy.x + del racy.x