1 commit

Author	SHA1	Message	Date
Sudhindra Desai	99f07e7bbd	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>	2026-08-11 09:55:24 -07:00