From 533a1c959c21f659a78fb130b17c700a3d6d1a06 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 17:42:46 -0700
Subject: [PATCH] fix(proxy): reserve budget for the Bedrock Converse prompt,
 not the context window

Resolving the model from the /bedrock path sends passthrough calls through
optimistic budget reservation, whose tokenizer cannot walk Converse content
blocks and so fell back to the model's max_input_tokens. Count those messages
as text and read inferenceConfig.maxTokens so a budgeted key reserves the
request's cost.
---
 .../spend_tracking/budget_reservation.py      | 37 ++++++++++++-------
 .../spend_tracking/test_budget_reservation.py | 31 +++++++++++++++-
 2 files changed, 53 insertions(+), 15 deletions(-)

diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py
index 91d2ece7a51..5ed52a327d4 100644
--- a/litellm/proxy/spend_tracking/budget_reservation.py
+++ b/litellm/proxy/spend_tracking/budget_reservation.py
@@ -1362,12 +1362,15 @@ def _approximate_input_size(request_body: Mapping[str, object]) -> int:
 def _count_input_tokens(request_body: dict, model: str) -> int | None:
     try:
         if "messages" in request_body:
-            return litellm.token_counter(
-                model=model,
-                messages=request_body.get("messages") or [],
-                tools=request_body.get("tools"),
-                tool_choice=request_body.get("tool_choice"),
-            )
+            try:
+                return litellm.token_counter(
+                    model=model,
+                    messages=request_body.get("messages") or (),
+                    tools=request_body.get("tools"),
+                    tool_choice=request_body.get("tool_choice"),
+                )
+            except ValueError:
+                return _count_text_tokens(model=model, text=request_body.get("messages"))
         if "prompt" in request_body:
             return _count_text_tokens(model=model, text=request_body.get("prompt"))
         if "input" in request_body:
@@ -1415,11 +1418,7 @@ def _estimate_output_tokens(
     if _is_input_only_route(route=route):
         return 0
 
-    requested: int | None = None
-    for key in ("max_completion_tokens", "max_tokens", "max_output_tokens"):
-        requested = _to_int(request_body.get(key))
-        if requested is not None:
-            break
+    requested: Final = _requested_output_tokens(request_body)
 
     # Clamp at min(requested-or-default, model_max-or-default). Two purposes:
     # (1) Without an explicit cap we still need a finite reservation so the
@@ -1430,9 +1429,19 @@ def _estimate_output_tokens(
     #     at the cap — the model can only physically emit max_output_tokens
     #     anyway, so reserving more is both wasteful and a DoS surface.
     model_ceiling: Final = _to_int(model_info.get("max_output_tokens")) or DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK
-    if requested is None:
-        requested = DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK
-    return min(requested, model_ceiling)
+    return min(DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK if requested is None else requested, model_ceiling)
+
+
+_OUTPUT_TOKEN_FIELDS: Final = ("max_completion_tokens", "max_tokens", "max_output_tokens")
+
+
+def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None:
+    inference_config: Final = request_body.get("inferenceConfig")
+    candidates: Final = (
+        *(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS),
+        inference_config.get("maxTokens") if isinstance(inference_config, Mapping) else None,
+    )
+    return next((tokens for tokens in map(_to_int, candidates) if tokens is not None), None)
 
 
 def _count_text_tokens(model: str, text: object) -> int:
diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py
index f65f68812a2..c7de8c943ad 100644
--- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py
+++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py
@@ -5,7 +5,7 @@ import pytest
 from litellm.caching import DualCache
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
-from litellm.proxy.spend_tracking.budget_reservation import reserve_budget_for_request
+from litellm.proxy.spend_tracking.budget_reservation import estimate_request_max_cost, reserve_budget_for_request
 from litellm.proxy.utils import ProxyLogging
 
 TOKEN_COUNTING_ROUTES: Final = (
@@ -46,3 +46,32 @@ async def test_non_exempt_llm_route_still_reserves_budget():
 
     assert reservation is not None
     assert reservation["reserved_cost"] > 0
+
+
+BEDROCK_SONNET: Final = "us.anthropic.claude-sonnet-4-6"
+CONVERSE_BODY: Final = {
+    "messages": [{"role": "user", "content": [{"text": "Reply with one word: pong"}]}],
+    "inferenceConfig": {"maxTokens": 5},
+}
+INVOKE_BODY: Final = {
+    "anthropic_version": "bedrock-2023-05-31",
+    "max_tokens": 5,
+    "messages": [{"role": "user", "content": "Reply with one word: pong"}],
+}
+
+
+def test_bedrock_converse_body_reserves_the_prompt_not_the_context_window():
+    converse_cost: Final = estimate_request_max_cost(
+        request_body=CONVERSE_BODY,
+        route=f"/bedrock/model/{BEDROCK_SONNET}/converse",
+        llm_router=None,
+        input_token_counts={},
+    )
+    invoke_cost: Final = estimate_request_max_cost(
+        request_body=INVOKE_BODY,
+        route=f"/bedrock/model/{BEDROCK_SONNET}/invoke",
+        llm_router=None,
+        input_token_counts={},
+    )
+    assert converse_cost is not None and invoke_cost is not None
+    assert invoke_cost < converse_cost < 2 * invoke_cost