fix(proxy): cache litellm model resolution to stop repeated Provider List spam

## Description

The proxy repeatedly prints LiteLLM's `Provider List:
https://docs.litellm.ai/docs/providers` banner during normal operation,
with no explanation or way to suppress it (#2851).

Root cause: `_resolve_litellm_model()` in
`headroom/proxy/savings_tracker.py` runs on every savings-tracking
update (i.e. every request). For any model LiteLLM can't price (a
custom/local/gateway model name — e.g. the reporter's local oMLX setup),
the uncached fallback path calls `litellm.cost_per_token(...)` purely to
probe resolvability. When that probe fails, LiteLLM prints the banner as
an internal side effect before raising, and since the probe was never
cached, it re-fires on every single request for the same unresolvable
model.

**Update:** review flagged that the first version of this fix cached
into a plain, unbounded `dict` keyed by the (client-controlled) model
name — a memory-retention path on a request-facing proxy, since a caller
can grow it without limit by sending a new model string on every
request. Replaced with a bounded `functools.lru_cache`; see Changes Made
below.

Closes #2851

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- `headroom/proxy/savings_tracker.py`: `_resolve_litellm_model()` is now
decorated with `@lru_cache(maxsize=256)` instead of backing onto a
hand-rolled unbounded `dict`. An evicted model name simply re-probes
LiteLLM on next use — never a correctness issue, only whether the noisy
failure banner reruns for that specific name.
- `tests/conftest.py`: added a global `autouse` fixture,
`_reset_litellm_model_resolution_cache`, that clears the cache before
and after every test. It's process-lifetime and module-global, and
several existing tests monkeypatch `savings_tracker.litellm` with
different behavior per test while reusing common model names like
`"gpt-4o"` — without a reset, whichever test resolves a name first
silently wins that cache slot for the rest of the run and later tests
stop exercising their own fake.
- `tests/test_savings_tracker_litellm_resolution_cache.py` (new):
regression tests for the three properties that actually matter —
repeated resolution of one unknown model only probes LiteLLM once,
resolving far more distinct names than the bound never grows the cache
past it, and an evicted name is transparently re-probed rather than
reusing a slot it no longer owns.
- No behavior change for models LiteLLM can already price (fast path via
`model_cost` lookup) — only the noisy uncached probe path is memoized,
same as before.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [ ] Type checking passes (`mypy headroom`) — not run; `mypy` isn't
installed in this environment
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
$ python3 -m pytest tests/test_proxy_savings_history.py tests/test_savings_tracker_zero_price.py \
    tests/test_savings_tracker_litellm_resolution_cache.py -q
tests/test_proxy_savings_history.py .................................... [ 73%]
...                                                                       [ 79%]
tests/test_savings_tracker_zero_price.py .......                         [ 93%]
tests/test_savings_tracker_litellm_resolution_cache.py ...               [100%]
49 passed, 1 warning in 1.26s

# Re-run in reversed file order to check for the exact order-dependence the
# review flagged — same 49 passed, no failures either direction:
$ python3 -m pytest tests/test_savings_tracker_litellm_resolution_cache.py \
    tests/test_savings_tracker_zero_price.py tests/test_proxy_savings_history.py -q
49 passed, 1 warning in 1.11s

$ python3 -m ruff check headroom/proxy/savings_tracker.py tests/conftest.py \
    tests/test_savings_tracker_litellm_resolution_cache.py
All checks passed!
```

## Real Behavior Proof

- Environment: macOS, Python 3.12.3, this repo checked out locally.
- What changed since the last review pass: I got the compiled
`headroom._core` Rust extension in hand (by installing the published
`headroom-ai[all]` wheel into a separate venv and copying its
`_core.abi3.so` next to this local source tree — same Python ABI,
pure-Python edits in `savings_tracker.py` don't touch the compiled
boundary). That unblocked the full test files this fix touches,
including `tests/test_proxy_savings_history.py`, which was previously
reported as untestable here.
- Exact command / steps: three properties asserted directly against the
real (now-bounded) cache in
`tests/test_savings_tracker_litellm_resolution_cache.py`:
1. Resolve the same unresolvable model 5 times → assert the underlying
`litellm.cost_per_token` probe fired exactly once.
2. Resolve `_MODEL_RESOLUTION_CACHE_MAXSIZE + 50` distinct model names →
assert `_resolve_litellm_model.cache_info().currsize` stays at exactly
`_MODEL_RESOLUTION_CACHE_MAXSIZE` (256), never higher — this is the
actual memory-retention fix the review asked for.
3. Resolve one model, push exactly `maxsize` other distinct names
through to evict it via LRU, then resolve it again → assert it re-probed
(call count went 1 → 2), proving eviction is real and not just an
untested cache_info number.
- Observed result: all three pass; full affected-file suite (49 tests)
passes in both forward and reversed run order, confirming the new
`conftest.py` fixture actually fixes the cross-test leakage risk
(verified by literally reordering the files, not just by inspection).
- Not tested: a live HTTP request against a running `headroom proxy`
process specifically re-exercising this bounded-cache commit — the
earlier "20 simulated requests" proof against the previous
(unbounded-dict) version of this fix was via a standalone script, not a
real server; I have not repeated that specific live-server pass against
this commit. The unit-level proof above exercises the exact same
function (`_resolve_litellm_model`) the real proxy calls per-request
from `headroom/proxy/server.py`, so I'm confident it generalizes, but
flagging the gap rather than implying I re-ran it live.

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
— the bound/eviction rationale is commented above
`_resolve_litellm_model`, and the cross-test leakage rationale is
commented above the new `conftest.py` fixture
- [ ] I have made corresponding changes to the documentation — N/A,
internal implementation detail with no user-facing API/doc surface
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I did **not** edit `CHANGELOG.md`

## Additional Notes

- `mypy` still hasn't been run — not installed in this sandbox, and I
didn't want to widen the PR further by installing/configuring it just
for this. Flagging rather than silently skipping.
- The earlier "Additional Notes" gap about
`test_proxy_savings_history.py` being untestable in this environment is
resolved (see Real Behavior Proof) — it now runs and passes, including
the pre-existing
`test_litellm_resolution_and_savings_estimation_fallbacks` test that
exercises `_resolve_litellm_model` with a mutated `model_cost` dict
across several assertions in one test.
- Deliberately did not also bound
`headroom/pricing/litellm_pricing.py`'s sibling `_resolved_model_cache`
— same shape of cache, arguably the same exposure — since it's outside
this PR's diff and touching it wasn't asked for. Flagging in case a
maintainer wants it as a fast follow-up rather than silently leaving it
unmentioned.

---------

Co-authored-by: connectsudhindra-gif <connectsudhindra-gif@users.noreply.github.com>
Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
Sudhindra Desai
2026-08-11 11:55:24 -05:00
committed by GitHub
parent 4bd8ecd1e3
commit 99f07e7bbd
3 changed files with 135 additions and 0 deletions
+24
View File
@@ -16,6 +16,7 @@ import tempfile
import threading
from csv import DictWriter
from datetime import datetime, timedelta, timezone
from functools import lru_cache
from io import StringIO
from pathlib import Path
from typing import Any
@@ -164,6 +165,23 @@ def _normalize_model(value: Any) -> str:
return cleaned or MODEL_UNKNOWN
# `_resolve_litellm_model` is called on every savings-tracking update (i.e.
# every request), and `model` is client-controlled — it comes straight off
# the request body. For a model LiteLLM can't price (a custom / local /
# gateway name), the uncached fallback below calls `litellm.cost_per_token`
# purely to probe resolvability, which prints LiteLLM's noisy "Provider
# List: https://docs.litellm.ai/docs/providers" banner on every failed probe
# (#2851). Cache the resolution per model name so that probe runs at most
# once per distinct model — bounded, not a plain dict: a request-facing
# proxy must not let a caller grow an unbounded cache for free by sending a
# fresh model string on every request. `maxsize` caps memory; LRU eviction
# means a model that stops being sent eventually falls out and simply
# re-probes if it's ever sent again — never a correctness issue, only
# whether the probe (and its noisy failure banner) reruns.
_MODEL_RESOLUTION_CACHE_MAXSIZE = 256
@lru_cache(maxsize=_MODEL_RESOLUTION_CACHE_MAXSIZE)
def _resolve_litellm_model(model: str) -> str:
"""Resolve model name to one LiteLLM recognizes.
@@ -173,6 +191,12 @@ def _resolve_litellm_model(model: str) -> str:
"claude-opus" identically to the live /stats path. Uses the shared result
only when it maps to a priced model_cost key; otherwise falls through to the
bare-prefix logic below. Fail-soft: pricing never breaks bookkeeping.
Bounded LRU cache, keyed by model name — see
``_MODEL_RESOLUTION_CACHE_MAXSIZE`` above. Tests that mock the LiteLLM
module across calls with the same model name must call
``_resolve_litellm_model.cache_clear()`` between cases, or results from
an earlier case leak in.
"""
litellm = _get_litellm_module()
if litellm is None:
+23
View File
@@ -91,6 +91,29 @@ def _reset_copilot_routing_flag():
reset_request_routed_to_copilot()
# `savings_tracker._resolve_litellm_model` is an `lru_cache`d, module-global,
# process-lifetime cache keyed by model name (bounded — see #2860). Many test
# files monkeypatch `savings_tracker.litellm` to a fake with different
# `model_cost`/`cost_per_token` behavior per test, but reuse common model
# names like "gpt-4o" across them. Without a reset, whichever test resolves
# "gpt-4o" first "wins" the cache entry for the rest of the run, and later
# tests silently stop exercising their own fake — a real-not-hypothetical
# order-dependence bug once the cache is process-lifetime instead of per-call.
# Clear before AND after so a test's own within-test resolutions never leak
# in from, or leak out to, a neighboring test either.
@pytest.fixture(autouse=True)
def _reset_litellm_model_resolution_cache():
try:
from headroom.proxy.savings_tracker import _resolve_litellm_model
except ModuleNotFoundError:
yield
return
_resolve_litellm_model.cache_clear()
yield
_resolve_litellm_model.cache_clear()
# =============================================================================
# Global test hooks
# =============================================================================
@@ -0,0 +1,88 @@
"""Regression: `_resolve_litellm_model`'s cache must be bounded (PR #2860 review).
A plain unbounded dict cache keyed by a client-controlled model string is a
memory-retention path on a request-facing proxy: a caller can grow it without
limit by sending a new model name on every request. The fix uses a bounded
`functools.lru_cache`. These tests pin the three properties that actually
matter, independent of the litellm pricing behavior covered elsewhere:
- repeated resolution of the same unresolvable model only probes litellm once
- the cache never grows past its bound, no matter how many distinct model
names get resolved
- an evicted name is transparently re-probed (never silently wrong or stuck)
rather than growing the cache further
"""
from __future__ import annotations
import types
from headroom.proxy import savings_tracker as st
def _fake_litellm_always_unresolvable(probe_calls: dict[str, int]) -> types.SimpleNamespace:
"""A fake litellm where every model is unpriced and unresolvable.
`cost_per_token` always raises — exactly what a real custom/local model
litellm has never heard of does — which is the call this cache exists to
memoize (see the comment above `_resolve_litellm_model` in
savings_tracker.py: that raise is also where real litellm prints its
noisy "Provider List" banner, #2851).
"""
def cost_per_token(*, model, prompt_tokens, completion_tokens):
probe_calls[model] = probe_calls.get(model, 0) + 1
raise RuntimeError("unknown model")
return types.SimpleNamespace(model_cost={}, cost_per_token=cost_per_token)
def test_resolve_litellm_model_probes_unknown_model_once(monkeypatch):
probe_calls: dict[str, int] = {}
monkeypatch.setattr(
st, "_get_litellm_module", lambda: _fake_litellm_always_unresolvable(probe_calls)
)
for _ in range(5):
resolved = st._resolve_litellm_model("widget-local-model")
assert resolved == "widget-local-model"
assert probe_calls == {"widget-local-model": 1}
def test_resolve_litellm_model_cache_is_bounded(monkeypatch):
probe_calls: dict[str, int] = {}
monkeypatch.setattr(
st, "_get_litellm_module", lambda: _fake_litellm_always_unresolvable(probe_calls)
)
extra_beyond_bound = 50
for i in range(st._MODEL_RESOLUTION_CACHE_MAXSIZE + extra_beyond_bound):
st._resolve_litellm_model(f"widget-local-model-{i}")
info = st._resolve_litellm_model.cache_info()
assert info.maxsize == st._MODEL_RESOLUTION_CACHE_MAXSIZE
# However many distinct names were resolved, the cache itself never
# grows past its bound -- this is the actual memory-retention fix.
assert info.currsize == st._MODEL_RESOLUTION_CACHE_MAXSIZE
def test_resolve_litellm_model_evicted_name_reprobes(monkeypatch):
probe_calls: dict[str, int] = {}
monkeypatch.setattr(
st, "_get_litellm_module", lambda: _fake_litellm_always_unresolvable(probe_calls)
)
st._resolve_litellm_model("seed-model")
assert probe_calls["seed-model"] == 1
# Push exactly `maxsize` new distinct names through without ever touching
# "seed-model" again -- LRU eviction must push it out to make room.
for i in range(st._MODEL_RESOLUTION_CACHE_MAXSIZE):
st._resolve_litellm_model(f"filler-model-{i}")
# A resolvable name being evicted is not a correctness bug (it just
# re-probes) -- the assertion that matters is that it *does* re-probe
# rather than silently reusing a slot it no longer legitimately owns.
st._resolve_litellm_model("seed-model")
assert probe_calls["seed-model"] == 2