From 1af27283d5cc274eb45bced497ccbfe659b9963b Mon Sep 17 00:00:00 2001
From: shivam <shivam@berri.ai>
Date: Fri, 24 Jul 2026 22:34:11 +0000
Subject: [PATCH] fix(cost_tracking): bill cache writes at the request's
 service-tier rate

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 litellm/utils.py                              | 11 +++++-
 .../llm_cost_calc/test_llm_cost_calc_utils.py | 38 +++++++++++++++++++
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/litellm/utils.py b/litellm/utils.py
index a11c5500503..3e4e53de98d 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -199,6 +199,7 @@ from litellm.types.utils import (
     SandboxProviders,
     SearchProviders,
     SelectTokenizerResponse,
+    ServiceTier,
     StreamingChoices,
     TextChoices,
     TextCompletionResponse,
@@ -5219,7 +5220,13 @@ def _is_potential_model_name_in_model_cost(
     )
 
 
-_ABOVE_THRESHOLD_COST_KEY = re.compile(r"_above_\d+k?_tokens$")
+# Cost keys the cost calculator composes at request time (context-length
+# thresholds and service tiers) instead of reading from a declared ModelInfo
+# field, e.g. "cache_creation_input_token_cost_priority" or
+# "input_cost_per_token_above_272k_tokens". They are copied onto ModelInfo
+# verbatim so a new pricing dimension in the cost map does not need a matching
+# field before it can be billed.
+_DYNAMIC_COST_KEY = re.compile(r"cost.*(?:_above_\d+k?_tokens|_(?:" + "|".join(st.value for st in ServiceTier) + r"))$")
 
 
 def _get_model_info_helper(
@@ -5535,7 +5542,7 @@ def _get_model_info_helper(
                 supports_image_size=_model_info.get("supports_image_size", None),
             )
             for cost_key, cost_value in _model_info.items():
-                if cost_key not in returned_model_info and _ABOVE_THRESHOLD_COST_KEY.search(cost_key) is not None:
+                if cost_key not in returned_model_info and _DYNAMIC_COST_KEY.search(cost_key) is not None:
                     returned_model_info[cost_key] = cost_value  # type: ignore[literal-required]
             return returned_model_info
     except Exception as e:
diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
index d282e656ce8..b1dffa38f43 100644
--- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
+++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
@@ -2399,3 +2399,41 @@ def test_generic_cost_per_token_gemini_35_flash_lite():
     )
     assert prompt_cost == pytest.approx(0.0003)
     assert completion_cost == pytest.approx(0.00125)
+
+
+@pytest.mark.parametrize(
+    "service_tier, expected_cache_write_rate, expected_input_rate, expected_output_rate",
+    [
+        (None, 6.25e-6, 5e-6, 3e-5),
+        ("priority", 1.25e-5, 1e-5, 6e-5),
+        ("flex", 3.125e-6, 2.5e-6, 1.5e-5),
+    ],
+)
+def test_service_tier_cache_write_pricing(
+    service_tier,
+    expected_cache_write_rate,
+    expected_input_rate,
+    expected_output_rate,
+    _local_model_cost_map,
+):
+    """Regression: cache-write tokens on a service-tier request must bill at that tier's
+    cache_creation_input_token_cost_<tier> rate. gpt-5.6-sol publishes $6.25/M standard,
+    $12.50/M priority and $3.125/M flex; billing every tier at the standard rate
+    under-charges priority by 50% and over-charges flex by 100%."""
+    usage = Usage(
+        prompt_tokens=4_020,
+        completion_tokens=4,
+        total_tokens=4_024,
+        prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=0, cache_write_tokens=4_017),
+    )
+
+    prompt_cost, completion_cost = generic_cost_per_token(
+        model="gpt-5.6-sol",
+        usage=usage,
+        custom_llm_provider="openai",
+        service_tier=service_tier,
+    )
+
+    expected_prompt = 4_017 * expected_cache_write_rate + 3 * expected_input_rate
+    assert prompt_cost == pytest.approx(expected_prompt, rel=1e-9)
+    assert completion_cost == pytest.approx(4 * expected_output_rate, rel=1e-9)