fix(proxy): scope the search free-tier counter to the request's provider (#47646)

* fix(proxy): scope the search free-tier counter to the request's provider

* fix(proxy): read the provider with get so partial spend payloads still get the free tier

* test(proxy): give the grounded spend payload fixture its provider

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/hooks/search_free_tier.py

@@ -34,8 +34,9 @@ def freed_units(allowance: int, units: int, total_after: int) -> int:
     return max(0, min(units, allowance - (total_after - units)))
 
 
-def search_free_tier_key(allowance: SearchFreeRequests, now: datetime) -> str:
-    return f"litellm:search_free_tier:{allowance['shared_group']}:{now.strftime(_PERIOD_STAMP[allowance['period']])}"
+def search_free_tier_key(provider: str, allowance: SearchFreeRequests, now: datetime) -> str:
+    period_stamp: Final = now.strftime(_PERIOD_STAMP[allowance["period"]])
+    return f"litellm:search_free_tier:{provider}:{allowance['shared_group']}:{period_stamp}"
 
 
 def search_free_tier_ttl(allowance: SearchFreeRequests, now: datetime) -> int:
@@ -107,7 +108,8 @@ async def apply_search_free_tier(
             "apply_search_free_tier is on but the proxy has no Redis, grounding is billed at full price"
         )
         return response_cost
-    key: Final = search_free_tier_key(allowance, now)
+    provider: Final = sl_object["custom_llm_provider"] or model_info["litellm_provider"]
+    key: Final = search_free_tier_key(provider, allowance, now)
     try:
         total_after: Final = await counter.async_increment(key, units, ttl=search_free_tier_ttl(allowance, now))
     except Exception:  # noqa: BLE001  # any counter failure means full price, never a lost spend row

tests/unit/proxy/hooks/test_proxy_track_cost_callback.py

@@ -3202,6 +3202,7 @@ def _grounded_gemini_kwargs(response_cost: float) -> dict:
             "response_cost": response_cost,
             "request_tags": None,
             "cache_hit": None,
+            "custom_llm_provider": "vertex_ai",
             "cost_breakdown": {
                 "input_cost": 0.0,
                 "output_cost": 0.0,

tests/unit/proxy/hooks/test_search_free_tier.py

@@ -31,9 +31,9 @@ class _BrokenCounter:
         raise ConnectionError("redis is down")
 
 
-def _model_info(billing_unit: str, count: int, period: str = "month") -> dict:
+def _model_info(billing_unit: str, count: int, period: str = "month", provider: str = "vertex_ai") -> dict:
     return {
-        "litellm_provider": "vertex_ai",
+        "litellm_provider": provider,
         "mode": "chat",
         "web_search_billing_unit": billing_unit,
         "search_context_cost_per_query": {"search_context_size_medium": RATE},
@@ -52,6 +52,7 @@ def _call(model_info: dict, response_cost: float, cache_hit: bool | None = None)
     sl_object = {
         "response_cost": response_cost,
         "cache_hit": cache_hit,
+        "custom_llm_provider": model_info["litellm_provider"],
         "cost_breakdown": {"input_cost": 0.1, "output_cost": 0.2, "tool_usage_cost": 0.5, "total_cost": response_cost},
         "model_map_information": {"model_map_key": "gemini-test", "model_map_value": model_info},
     }
@@ -73,12 +74,12 @@ def test_freed_units_covers_only_what_landed_under_the_allowance(allowance, unit
     assert freed_units(allowance, units, total_after) == expected
 
 
-def test_counter_key_is_per_group_and_utc_period():
-    assert search_free_tier_key({"count": 1, "period": "month", "shared_group": "g3"}, NOW) == (
-        "litellm:search_free_tier:g3:2026-10"
+def test_counter_key_is_per_provider_group_and_utc_period():
+    assert search_free_tier_key("gemini", {"count": 1, "period": "month", "shared_group": "g3"}, NOW) == (
+        "litellm:search_free_tier:gemini:g3:2026-10"
     )
-    assert search_free_tier_key({"count": 1, "period": "day", "shared_group": "g2"}, NOW) == (
-        "litellm:search_free_tier:g2:2026-10-20"
+    assert search_free_tier_key("vertex_ai", {"count": 1, "period": "day", "shared_group": "g2"}, NOW) == (
+        "litellm:search_free_tier:vertex_ai:g2:2026-10-20"
     )
 
 
@@ -107,8 +108,23 @@ async def test_per_query_allowance_is_shared_across_calls_and_split_on_the_cross
     assert c2 == k2 == s2 == pytest.approx(1.0 - 2 * RATE)
     assert b2["discount_amount"] == pytest.approx(2 * RATE)
     assert c3 == k3 == s3 == 1.0 and "discount_amount" not in b3
-    assert counter.totals == {"litellm:search_free_tier:test-group:2026-10": 15}
-    assert counter.ttls["litellm:search_free_tier:test-group:2026-10"] > 0
+    assert counter.totals == {"litellm:search_free_tier:vertex_ai:test-group:2026-10": 15}
+    assert counter.ttls["litellm:search_free_tier:vertex_ai:test-group:2026-10"] > 0
+
+
+@pytest.mark.asyncio
+async def test_gemini_api_and_vertex_ai_grounding_draw_on_separate_allowances():
+    counter = _Counter()
+    gemini_api = _model_info("per_query", count=3, provider="gemini")
+    vertex_ai = _model_info("per_query", count=3, provider="vertex_ai")
+    kwargs, sl_object = _call(gemini_api, response_cost=1.0)
+    await apply_search_free_tier(kwargs, sl_object, 1.0, _response(2), counter, NOW)
+    kwargs, sl_object = _call(vertex_ai, response_cost=1.0)
+    vertex_cost = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(4), counter, NOW)
+    kwargs, sl_object = _call(gemini_api, response_cost=1.0)
+    gemini_cost = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(1), counter, NOW)
+    assert vertex_cost == pytest.approx(1.0 - 3 * RATE), "vertex ai had its whole allowance left"
+    assert gemini_cost == pytest.approx(1.0 - RATE), "gemini api had one free query left"
 
 
 @pytest.mark.asyncio
@@ -121,7 +137,7 @@ async def test_per_prompt_allowance_counts_a_grounded_prompt_once_whatever_the_q
     second = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW)
     assert first == pytest.approx(1.0 - RATE)
     assert second == 1.0
-    assert counter.totals == {"litellm:search_free_tier:test-group:2026-10-20": 2}
+    assert counter.totals == {"litellm:search_free_tier:vertex_ai:test-group:2026-10-20": 2}
 
 
 @pytest.mark.asyncio
@@ -172,7 +188,7 @@ async def test_a_second_callback_run_on_the_same_request_counts_it_once():
     first = await apply_search_free_tier(kwargs, sl_object, 1.0, _response(5), counter, NOW)
     second = await apply_search_free_tier(kwargs, sl_object, first, _response(5), counter, NOW)
     assert first == second == pytest.approx(1.0 - 5 * RATE)
-    assert counter.totals == {"litellm:search_free_tier:test-group:2026-10": 5}
+    assert counter.totals == {"litellm:search_free_tier:vertex_ai:test-group:2026-10": 5}
 
 
 @pytest.mark.asyncio
@@ -191,4 +207,4 @@ async def test_freed_cost_never_exceeds_the_billed_cost():
     kwargs, sl_object = _call(_model_info("per_query", count=100), response_cost=billed)
     assert await apply_search_free_tier(kwargs, sl_object, billed, _response(5), counter, NOW) == 0.0
     assert sl_object["cost_breakdown"]["discount_amount"] == pytest.approx(billed)
-    assert counter.totals == {"litellm:search_free_tier:test-group:2026-10": 5}
+    assert counter.totals == {"litellm:search_free_tier:vertex_ai:test-group:2026-10": 5}