From bd9593b74d4a05e265225a43afa0fd6c8d837484 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 16:41:31 -0700
Subject: [PATCH] fix(proxy): share the Bedrock count-tokens predicate between
 auth and the passthrough handler

---
 litellm/proxy/auth/auth_utils.py                       |  7 ++++---
 .../llm_passthrough_endpoints.py                       |  7 +++++--
 tests/test_litellm/proxy/auth/test_auth_utils.py       | 10 ++++++++++
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py
index 12e6f75889e..1e4836654a1 100644
--- a/litellm/proxy/auth/auth_utils.py
+++ b/litellm/proxy/auth/auth_utils.py
@@ -1989,13 +1989,14 @@ def get_model_from_request(
 
 
 def _model_from_bedrock_route(route: str) -> str | None:
-    bedrock_endpoint: Final = re.sub(r"^/bedrock/", "", route, flags=re.IGNORECASE)
-    if "count_tokens" in bedrock_endpoint.lower() or "count-tokens" in bedrock_endpoint.lower():
-        return 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:
diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
index 29f216fd450..820312ac0fd 100644
--- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
+++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
@@ -703,6 +703,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.
@@ -977,8 +981,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,
diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py
index 61d1a308b9d..a996de4d40c 100644
--- a/tests/test_litellm/proxy/auth/test_auth_utils.py
+++ b/tests/test_litellm/proxy/auth/test_auth_utils.py
@@ -508,6 +508,16 @@ def test_get_model_from_request_bedrock_count_tokens_uses_body_model():
     )
 
 
+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(