Merge pull request #39660 from BerriAI/litellm_bedrock_passthrough_model_access

fix(proxy): enforce key and team model access on Bedrock passthrough routes (internal copy of #34244)

litellm/proxy/auth/auth_utils.py

@@ -1981,9 +1981,28 @@ def get_model_from_request(
         if vertex_match:
             model = vertex_match.group(1)
 
+    if route.lower().startswith("/bedrock"):
+        bedrock_model: Final = _model_from_bedrock_route(route)
+        return model if bedrock_model is None else bedrock_model
+
     return model
 
 
+def _model_from_bedrock_route(route: str) -> str | None:
+    from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
+        _extract_model_from_bedrock_endpoint,
+        is_bedrock_count_tokens_endpoint,
+    )
+
+    bedrock_endpoint: Final = re.sub(r"^/bedrock/", "", route, flags=re.IGNORECASE)
+    if is_bedrock_count_tokens_endpoint(bedrock_endpoint):
+        return None
+    try:
+        return _extract_model_from_bedrock_endpoint(bedrock_endpoint)
+    except ValueError:
+        return None
+
+
 def abbreviate_api_key(api_key: str) -> str:
     if len(api_key) < MINIMUM_CUSTOM_KEY_LENGTH:
         return "sk-..."

litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py

@@ -712,6 +712,10 @@ BEDROCK_ENDPOINT_ACTIONS: Final = {
 BEDROCK_STREAMING_ACTIONS: Final = {"invoke-with-response-stream", "converse-stream"}
 
 
+def is_bedrock_count_tokens_endpoint(endpoint: str) -> bool:
+    return "count_tokens" in endpoint or "count-tokens" in endpoint
+
+
 def _extract_model_from_bedrock_endpoint(endpoint: str) -> str:
     """
     Extract model name from Bedrock endpoint path.
@@ -986,8 +990,7 @@ async def bedrock_llm_proxy_route(
 
     request_body: Final = await _read_request_body(request=request)
 
-    # Special handling for count_tokens endpoints
-    if "count_tokens" in endpoint or "count-tokens" in endpoint:
+    if is_bedrock_count_tokens_endpoint(endpoint):
         return await handle_bedrock_count_tokens(
             endpoint=endpoint,
             request=request,

litellm/proxy/spend_tracking/budget_reservation.py

@@ -1389,12 +1389,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:
@@ -1442,11 +1445,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
@@ -1457,9 +1456,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:

tests/test_litellm/proxy/auth/test_auth_utils.py

@@ -428,6 +428,106 @@ def test_get_model_from_request_openai_deployment_route_still_works():
     )
 
 
+def test_get_model_from_request_bedrock_converse_passthrough():
+    assert (
+        get_model_from_request(
+            request_data={},
+            route="/bedrock/model/us.anthropic.claude-sonnet-4-6/converse",
+        )
+        == "us.anthropic.claude-sonnet-4-6"
+    )
+
+
+def test_get_model_from_request_bedrock_invoke_passthrough():
+    assert (
+        get_model_from_request(
+            request_data={},
+            route="/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke",
+        )
+        == "us.anthropic.claude-sonnet-4-6"
+    )
+
+
+def test_get_model_from_request_bedrock_v2_converse_stream_passthrough():
+    assert (
+        get_model_from_request(
+            request_data={},
+            route="/bedrock/v2/model/us.anthropic.claude-sonnet-4-6/converse-stream",
+        )
+        == "us.anthropic.claude-sonnet-4-6"
+    )
+
+
+def test_get_model_from_request_bedrock_model_id_with_slashes():
+    assert (
+        get_model_from_request(
+            request_data={},
+            route="/bedrock/model/aws/anthropic/model-name/invoke",
+        )
+        == "aws/anthropic/model-name"
+    )
+
+
+def test_get_model_from_request_bedrock_unparseable_endpoint_returns_none():
+    assert (
+        get_model_from_request(
+            request_data={},
+            route="/bedrock/agents/some-agent-route",
+        )
+        is None
+    )
+
+
+def test_get_model_from_request_bedrock_url_model_overrides_body_model():
+    assert (
+        get_model_from_request(
+            request_data={"model": "us.anthropic.claude-sonnet-4-6"},
+            route="/bedrock/model/us.anthropic.claude-opus-4-6-v1/converse",
+        )
+        == "us.anthropic.claude-opus-4-6-v1"
+    )
+
+
+def test_get_model_from_request_bedrock_invoke_url_model_overrides_body_model():
+    assert (
+        get_model_from_request(
+            request_data={"model": "us.anthropic.claude-sonnet-4-6"},
+            route="/bedrock/model/us.anthropic.claude-opus-4-6-v1/invoke",
+        )
+        == "us.anthropic.claude-opus-4-6-v1"
+    )
+
+
+def test_get_model_from_request_bedrock_count_tokens_uses_body_model():
+    assert (
+        get_model_from_request(
+            request_data={"model": "us.anthropic.claude-sonnet-4-6"},
+            route="/bedrock/v1/messages/count_tokens",
+        )
+        == "us.anthropic.claude-sonnet-4-6"
+    )
+
+
+def test_get_model_from_request_bedrock_uppercase_count_tokens_segment_is_not_count_tokens():
+    assert (
+        get_model_from_request(
+            request_data={"model": "us.anthropic.claude-haiku-4-5-20251001-v1:0"},
+            route="/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke/COUNT_TOKENS",
+        )
+        == "us.anthropic.claude-sonnet-4-6"
+    )
+
+
+def test_get_model_from_request_bedrock_unparseable_endpoint_keeps_body_model():
+    assert (
+        get_model_from_request(
+            request_data={"model": "us.anthropic.claude-sonnet-4-6"},
+            route="/bedrock/agents/some-agent-route",
+        )
+        == "us.anthropic.claude-sonnet-4-6"
+    )
+
+
 def test_get_model_from_request_includes_file_endpoint_header_model():
     assert (
         get_model_from_request(

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