fix(cost_calculator): keep custom-priced router ids when resolving slash aliases

litellm/cost_calculator.py

@@ -799,50 +799,27 @@ def _select_model_name_for_cost_calc(
         else:
             return_model = f"{custom_llm_provider}/{return_model}"
 
-        # A router-facing model_name alias may itself contain a "/" (e.g. "vertex/claude-opus-5")
-        # whose leading segment is NOT a registered provider, so the alias was re-prefixed above
-        # into a non-existent key like "vertex_ai/vertex/claude-opus-5". Strip unregistered
-        # leading segments until the assembled name is a known model_cost key (or nothing
-        # strippable remains), so cost lookup resolves to the real model instead of pricing
-        # the request at $0. See #38069.
-        stripped = _strip_unregistered_leading_segments(return_model, region_name)
-        if stripped is not None:
-            return_model = stripped
+        return_model = _strip_unregistered_leading_segments(return_model, region_name)
 
     return return_model
 
 
-def _strip_unregistered_leading_segments(model: str, region_name: str | None) -> str | None:
-    """Find the real cost-map key hidden inside a re-prefixed alias.
-
-    After the provider (and optional region) was prepended to a router-facing alias that
-    itself contains "/" segments (e.g. "vertex_ai/vertex/claude-opus-5"), walk the
-    remaining segments: keep the provider/region head, then try successively shorter
-    tails — "vertex_ai/vertex/claude-opus-5", "vertex_ai/claude-opus-5" — returning the
-    first that exists in `litellm.model_cost`. Segments that are themselves providers or
-    regions are never dropped from the head, so "vertex_ai/us-east-1/model" keeps its region.
-
-    Returns None if no candidate resolves; the caller keeps its original assembly.
-    """
-
-    segments = model.split("/")
-    # head = leading provider (+ optional region) that must be preserved
-    head_len = 1
-    if region_name is not None and len(segments) > 1 and segments[1] == region_name:
-        head_len = 2
-
-    head = segments[:head_len]
-    tail = segments[head_len:]
-    while tail:
-        candidate = "/".join(head + tail)
-        if candidate in litellm.model_cost:
-            return candidate
-        # only strip a leading tail segment that is NOT itself meaningful:
-        # a provider/region segment inside the tail (e.g. "azure_ai/o3") stays
-        if tail[0] in LlmProvidersSet or tail[0] == region_name:
-            return None
-        tail = tail[1:]
-    return None
+def _strip_unregistered_leading_segments(model: str, region_name: str | None) -> str:
+    """Resolve a provider-prefixed slash alias like "vertex_ai/vertex/claude-opus-5" to the
+    registered cost key ("vertex_ai/claude-opus-5"), keeping the model unchanged when it already
+    resolves downstream (custom-priced router ids) or no stripped candidate is registered (#38069)."""
+    segments: Final = model.split("/")
+    if "/".join(segments[1:]) in litellm.model_cost:
+        return model
+    head_len: Final = 2 if region_name is not None and len(segments) > 2 and segments[1] == region_name else 1
+    head: Final = "/".join(segments[:head_len])
+    tail: Final = segments[head_len:]
+    strippable: Final = next(
+        (index for index, segment in enumerate(tail) if segment in LlmProvidersSet or segment == region_name),
+        len(tail),
+    )
+    candidates: Final = tuple(f"{head}/{'/'.join(tail[start:])}" for start in range(min(strippable, len(tail) - 1) + 1))
+    return next((candidate for candidate in candidates if candidate in litellm.model_cost), model)
 
 
 @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)

tests/test_litellm/test_cost_calculator.py

@@ -3866,3 +3866,51 @@ def test_select_model_name_unresolvable_alias_unchanged(_local_model_cost_map):
     )
 
     assert selected == "vertex_ai/team/nonsense-model"
+
+
+def test_completion_cost_keeps_custom_priced_slash_router_id(_local_model_cost_map):
+    """A custom-priced router id containing "/" keeps its custom pricing instead of being
+    rewritten to the built-in key its suffix happens to match."""
+
+    from litellm.cost_calculator import _select_model_name_for_cost_calc
+
+    litellm.register_model(
+        model_cost={
+            "vertex/claude-opus-5": {
+                "input_cost_per_token": 7e-6,
+                "output_cost_per_token": 8e-6,
+                "litellm_provider": "vertex_ai",
+            }
+        }
+    )
+
+    selected = _select_model_name_for_cost_calc(
+        model="vertex_ai/claude-opus-5",
+        completion_response=None,
+        custom_pricing=True,
+        custom_llm_provider="vertex_ai",
+        router_model_id="vertex/claude-opus-5",
+    )
+    assert selected == "vertex_ai/vertex/claude-opus-5"
+
+    response = litellm.ModelResponse(
+        id="x",
+        choices=[
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": "hi"},
+                "finish_reason": "stop",
+            }
+        ],
+        model="vertex/claude-opus-5",
+    )
+    response._hidden_params = {"custom_llm_provider": "vertex_ai"}
+    response.usage = litellm.Usage(prompt_tokens=100, completion_tokens=50)
+
+    cost = litellm.completion_cost(
+        completion_response=response,
+        custom_llm_provider="vertex_ai",
+        custom_pricing=True,
+        router_model_id="vertex/claude-opus-5",
+    )
+    assert cost == pytest.approx(100 * 7e-6 + 50 * 8e-6, rel=1e-9)