From 571927705a7866dc358d01b2bdef68e65f253b5a Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 03:19:27 -0700
Subject: [PATCH] fix(passthrough): name Google failures from the URL on
 untracked methods too

A Gemini or Vertex pass-through failure on a method the success handler does
not track (countTokens, for one) still logged model "" and no provider,
although the model sits in the URL. The failure payload now falls back to the
Google AI host for the provider and reads /models/<name> from the URL, so
countTokens 400s name their model and provider like generateContent ones do.
Non-Google hosts keep the tracked-route rule and still log nothing extra.
---
 .../vertex_passthrough_logging_handler.py     |  9 +++++-
 .../pass_through_endpoints.py                 |  4 ++-
 .../test_pass_through_endpoints.py            | 32 +++++++++++++++++++
 3 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py
index b667a7916fb..960e4d5882b 100644
--- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py
+++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py
@@ -35,6 +35,8 @@ from litellm.types.utils import (
 )
 
 vertex_search_api_config: Final = VertexSearchAPIVectorStoreConfig()
+GEMINI_API_HOST_SUFFIX: Final = "generativelanguage.googleapis.com"
+VERTEX_API_HOST_SUFFIX: Final = "aiplatform.googleapis.com"
 if TYPE_CHECKING:
     from litellm.types.utils import LiteLLMBatch
 
@@ -560,10 +562,15 @@ class VertexPassthroughLoggingHandler:
             return "openapi"
         return None
 
+    @staticmethod
+    def is_google_ai_host(url: str) -> bool:
+        hostname: Final = urlparse(url).hostname
+        return hostname is not None and hostname.endswith((GEMINI_API_HOST_SUFFIX, VERTEX_API_HOST_SUFFIX))
+
     @staticmethod
     def custom_llm_provider_from_url(url: str) -> str:
         parsed_url: Final = urlparse(url)
-        if parsed_url.hostname and parsed_url.hostname.endswith("generativelanguage.googleapis.com"):
+        if parsed_url.hostname and parsed_url.hostname.endswith(GEMINI_API_HOST_SUFFIX):
             return litellm.LlmProviders.GEMINI.value
         return litellm.LlmProviders.VERTEX_AI.value
 
diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
index e59907bfa4c..eb81a02871e 100644
--- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
+++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
@@ -735,7 +735,9 @@ def _carry_guardrail_logging_info(request_data: dict, guardrail_data: dict | Non
 def _passthrough_google_provider(url_route: str, custom_llm_provider: str | None) -> str | None:
     if pass_through_endpoint_logging.is_gemini_route(url_route, custom_llm_provider):
         return "gemini"
-    if pass_through_endpoint_logging.is_vertex_route(url_route):
+    if pass_through_endpoint_logging.is_vertex_route(url_route) or VertexPassthroughLoggingHandler.is_google_ai_host(
+        url_route
+    ):
         return VertexPassthroughLoggingHandler.custom_llm_provider_from_url(url_route)
     return None
 
diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
index 79c2e7fb3e9..f79b755f84e 100644
--- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
+++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
@@ -5897,6 +5897,8 @@ def _upstream_400_as_event_stream(target: str) -> httpx.Response:
         (_GEMINI_FLASH_TARGET, "gemini", True, "gemini"),
         (_VERTEX_FLASH_TARGET, None, False, "vertex_ai"),
         (_VERTEX_FLASH_TARGET, None, True, "vertex_ai"),
+        (_GEMINI_FLASH_TARGET.replace(":generateContent", ":countTokens"), "gemini", False, "gemini"),
+        (_VERTEX_FLASH_TARGET.replace(":generateContent", ":countTokens"), None, False, "vertex_ai"),
     ],
 )
 async def test_passthrough_upstream_error_logs_model_from_url(
@@ -5986,6 +5988,36 @@ async def test_passthrough_mid_stream_drop_logs_model_from_url(
             "",
             "vertex_ai",
         ),
+        (
+            {"contents": []},
+            _GEMINI_FLASH_TARGET.replace(":generateContent", ":countTokens"),
+            "gemini",
+            "gemini-3.8-flash",
+            "gemini",
+        ),
+        (
+            {"contents": []},
+            _VERTEX_FLASH_TARGET.replace(":generateContent", ":countTokens"),
+            None,
+            "gemini-3.8-flash",
+            "vertex_ai",
+        ),
+        (
+            {"contents": []},
+            "https://aiplatform.googleapis.com/v1/projects/p/locations/global"
+            "/publishers/google/models/gemini-3.8-flash:countTokens",
+            None,
+            "gemini-3.8-flash",
+            "vertex_ai",
+        ),
+        (
+            {},
+            "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1/cachedContents",
+            None,
+            "",
+            "vertex_ai",
+        ),
+        ({"contents": []}, "https://api.example.com/v1/models/gpt-x:countTokens", None, "", None),
         ({"contents": []}, "https://api.example.com/v1/models/gpt-x:thing", None, "", None),
         (None, None, None, "", None),
     ],