From fbda977c5ee632d0fcb475511ad6e88b6f099d1f Mon Sep 17 00:00:00 2001
From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com>
Date: Sun, 11 Oct 2026 22:31:07 -0500
Subject: [PATCH] fix(proxy): keep token-budget code within the basedpyright
 ceiling

Type the token lookups instead of reading untyped callback data:
- _extract_total_tokens takes a StandardLoggingPayload and returns its typed
  total_tokens, the same source the batched tokens_used DB write uses
- read max_tokens_budget off the typed LiteLLM_TeamTable, not getattr
- read total_tokens off the typed SpendLogsPayload, not .get
- drop Final on the loop-assigned team_update_data (illegal, and the gate
  exempts loop vars from the Final rule anyway)
The one remaining untyped boundary (the CustomLogger kwargs dict) is cast once
with a reason.
---
 litellm/proxy/auth/auth_checks.py             |  6 +++--
 litellm/proxy/db/db_spend_update_writer.py    |  4 +--
 .../proxy/hooks/proxy_track_cost_callback.py  | 27 ++++++++-----------
 3 files changed, 17 insertions(+), 20 deletions(-)

diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py
index cfeb13f3b3f..54d7f7de155 100644
--- a/litellm/proxy/auth/auth_checks.py
+++ b/litellm/proxy/auth/auth_checks.py
@@ -6537,8 +6537,10 @@ async def team_max_tokens_check(
     back to the ``tokens_used`` column loaded in auth. No budget set on the team
     means the check is skipped, so existing teams are unaffected.
     """
-    max_tokens_budget: Final = getattr(team_object, "max_tokens_budget", None)
-    if team_object is None or max_tokens_budget is None:
+    if team_object is None:
+        return
+    max_tokens_budget: Final = team_object.max_tokens_budget
+    if max_tokens_budget is None:
         return
 
     from litellm.proxy.proxy_server import get_current_tokens_used
diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py
index 657456f37a9..98c788b6fc4 100644
--- a/litellm/proxy/db/db_spend_update_writer.py
+++ b/litellm/proxy/db/db_spend_update_writer.py
@@ -1066,7 +1066,7 @@ class DBSpendUpdateWriter:
                 team_id=team_id,
                 user_id=user_id,
                 prisma_client=prisma_client,
-                tokens=int(payload.get("total_tokens") or 0),
+                tokens=payload["total_tokens"],
             )
         except Exception:
             verbose_proxy_logger.debug(
@@ -2187,7 +2187,7 @@ class DBSpendUpdateWriter:
                                     response_cost,
                                     token_increment,
                                 )
-                                team_update_data: Final = (
+                                team_update_data = (
                                     {
                                         "spend": {"increment": response_cost},
                                         "tokens_used": {"increment": token_increment},
diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py
index 3c3a05353e3..6144dc37e7a 100644
--- a/litellm/proxy/hooks/proxy_track_cost_callback.py
+++ b/litellm/proxy/hooks/proxy_track_cost_callback.py
@@ -733,24 +733,15 @@ class _IncrementSpendCounters(Protocol):
     ) -> None: ...
 
 
-def _extract_total_tokens(kwargs: Mapping[str, object], completion_response: object) -> int:
+def _extract_total_tokens(sl_object: StandardLoggingPayload | None) -> int:
     """Total tokens for this request, for the team token-budget counter.
 
-    Prefers the StandardLoggingPayload (always populated for a billed call),
-    falling back to the response usage. Returns 0 when unknown, which leaves the
-    token counter untouched.
+    Read from the StandardLoggingPayload, the same source the batched tokens_used
+    DB write uses, so the cross-pod counter and the column stay in agreement.
     """
-    sl_object: Final = kwargs.get("standard_logging_object") or {}
-    sl_total: Final = sl_object.get("total_tokens")
-    if sl_total:
-        return int(sl_total)
-    usage: Final = (
-        completion_response.get("usage")
-        if isinstance(completion_response, dict)
-        else getattr(completion_response, "usage", None)
-    )
-    usage_total: Final = usage.get("total_tokens") if isinstance(usage, dict) else getattr(usage, "total_tokens", None)
-    return int(usage_total) if usage_total else 0
+    if sl_object is None:
+        return 0
+    return sl_object["total_tokens"]
 
 
 async def _update_database_and_spend_counters(
@@ -901,7 +892,11 @@ async def _update_database_and_spend_counters_in_batch(
     # Awaited like the spend counter above; best-effort so a counter hiccup
     # never fails a request that was already billed.
     if team_id is not None:
-        total_tokens: Final = _extract_total_tokens(kwargs=kwargs, completion_response=completion_response)
+        sl_object: Final = cast(  # cast-ok: typed shape of the untyped CustomLogger callback payload
+            "StandardLoggingPayload | None",
+            kwargs.get("standard_logging_object"),  # pyright: ignore[reportUnknownMemberType]  # untyped callback kwargs
+        )
+        total_tokens: Final = _extract_total_tokens(sl_object)
         if total_tokens:
             from litellm.proxy.proxy_server import increment_spend_counter_cache