Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_/migrate-simple-table-631c9b

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:

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

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))

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