From 040e30aa722a64d9c4a9fc4f6316135af2290bab Mon Sep 17 00:00:00 2001 From: Kshitij Kapoor Date: Thu, 9 Jul 2026 13:04:31 +0530 Subject: [PATCH] perf(model_metadata): cache ollama /api/show probe + normalize context-cache keys MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up hunks completing the probe-cache cluster: 1. _query_ollama_api_show now goes through the existing _LOCAL_CTX_PROBE_CACHE (30s TTL, positive-only, namespaced key) — it was the one remaining per-resolution POST not covered by the #56431-era wrapper. Failures are never memoized so a server that comes up mid-startup is re-probed. Idea credit: #42081 (@Morad37), reworked to comply with the positive-only rule. 2. Persistent context-cache keys are normalized through _context_cache_key (trailing-slash strip) so http://host/v1 and http://host/v1/ share one entry; reads and invalidation honor legacy un-normalized rows. Idea credit: #37905 (@stevenau21). Tests: TTL hit collapses to one POST, failure-not-memoized (mutation-verified: unconditional caching makes it fail), namespace no-collision vs the sibling probe, slash-variant dedup, legacy-row read, dual-shape invalidation. --- agent/model_metadata.py | 57 +++++++++- tests/agent/test_probe_cache_followups.py | 132 ++++++++++++++++++++++ 2 files changed, 183 insertions(+), 6 deletions(-) create mode 100644 tests/agent/test_probe_cache_followups.py diff --git a/agent/model_metadata.py b/agent/model_metadata.py index 8396133c907..b911544811f 100644 --- a/agent/model_metadata.py +++ b/agent/model_metadata.py @@ -1052,13 +1052,23 @@ def _load_context_cache() -> Dict[str, int]: return {} +def _context_cache_key(model: str, base_url: str) -> str: + """Canonical ``model@base_url`` key for the persistent context cache. + + Trailing slashes are stripped so ``http://host/v1`` and + ``http://host/v1/`` share one entry instead of creating duplicates + that can go stale independently. + """ + return f"{model}@{(base_url or '').rstrip('/')}" + + def save_context_length(model: str, base_url: str, length: int) -> None: """Persist a discovered context length for a model+provider combo. Cache key is ``model@base_url`` so the same model name served from different providers can have different limits. """ - key = f"{model}@{base_url}" + key = _context_cache_key(model, base_url) cache = _load_context_cache() if cache.get(key) == length: return # already stored @@ -1075,18 +1085,28 @@ def save_context_length(model: str, base_url: str, length: int) -> None: def get_cached_context_length(model: str, base_url: str) -> Optional[int]: """Look up a previously discovered context length for model+provider.""" - key = f"{model}@{base_url}" + key = _context_cache_key(model, base_url) cache = _load_context_cache() - return cache.get(key) + hit = cache.get(key) + if hit is not None: + return hit + # Legacy rows written before key normalization may carry a trailing + # slash — honor them rather than re-probing. + if base_url and base_url != base_url.rstrip("/"): + return cache.get(f"{model}@{base_url}") + return None def _invalidate_cached_context_length(model: str, base_url: str) -> None: """Drop a stale cache entry so it gets re-resolved on the next lookup.""" - key = f"{model}@{base_url}" + key = _context_cache_key(model, base_url) cache = _load_context_cache() - if key not in cache: + # Also clear any legacy un-normalized row for the same pair. + legacy_key = f"{model}@{base_url}" + if key not in cache and legacy_key not in cache: return - del cache[key] + cache.pop(key, None) + cache.pop(legacy_key, None) path = _get_context_cache_path() try: path.parent.mkdir(parents=True, exist_ok=True) @@ -1473,6 +1493,12 @@ def _query_ollama_api_show(model: str, base_url: str, api_key: str = "") -> Opti hosting behind a reverse proxy, etc. For non-Ollama servers the POST returns 404/405 quickly; the function handles errors gracefully. + Results are cached in ``_LOCAL_CTX_PROBE_CACHE`` (same 30s TTL, + positive-only — see ``_query_local_context_length``) so back-to-back + resolutions during one startup issue a single POST instead of one per + call site. Failures are never memoized: a server that isn't up yet must + be re-probed once it comes up. + For hosted servers the GGUF ``model_info.*.context_length`` is the authoritative source: the user can't set their own ``num_ctx``, and the OpenAI-compat ``/v1/models`` endpoint correctly omits ``context_length`` @@ -1484,6 +1510,25 @@ def _query_ollama_api_show(model: str, base_url: str, api_key: str = "") -> Opti The order is flipped vs ``query_ollama_num_ctx()`` because local users control ``num_ctx`` themselves; hosted users can't. """ + import time as _time + + # Namespaced cache key: shares the TTL store with + # _query_local_context_length but never collides with its (model, url) + # keys — the two probes can return different values for the same pair. + cache_key = ("ollama_show", _strip_provider_prefix(model), base_url.rstrip("/")) + now = _time.monotonic() + cached = _LOCAL_CTX_PROBE_CACHE.get(cache_key) + if cached is not None and (now - cached[1]) < _LOCAL_CTX_PROBE_TTL_SECONDS: + return cached[0] + + result = _query_ollama_api_show_uncached(model, base_url, api_key=api_key) + if result: # positive-only — never memoize a failed probe + _LOCAL_CTX_PROBE_CACHE[cache_key] = (result, now) + return result + + +def _query_ollama_api_show_uncached(model: str, base_url: str, api_key: str = "") -> Optional[int]: + """Uncached body of ``_query_ollama_api_show`` — one POST to ``/api/show``.""" import httpx server_url = base_url.rstrip("/") diff --git a/tests/agent/test_probe_cache_followups.py b/tests/agent/test_probe_cache_followups.py new file mode 100644 index 00000000000..faf849b71a4 --- /dev/null +++ b/tests/agent/test_probe_cache_followups.py @@ -0,0 +1,132 @@ +"""Tests for probe-cache follow-ups on the #29988/#37595/#50572 salvage. + +Covers: +- _query_ollama_api_show TTL caching (positive-only, namespaced key) +- persistent context-cache key normalization (trailing-slash dedup) +""" + +from __future__ import annotations + +import os +import sys +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + + +@pytest.fixture(autouse=True) +def _clear_probe_cache(): + """Module-level TTL cache must not leak between tests.""" + from agent import model_metadata + model_metadata._LOCAL_CTX_PROBE_CACHE.clear() + yield + model_metadata._LOCAL_CTX_PROBE_CACHE.clear() + + +def _mock_show_response(ctx=131072): + resp = MagicMock() + resp.status_code = 200 + resp.json.return_value = { + "model_info": {"llama.context_length": ctx}, + "parameters": "", + } + return resp + + +def _client_mock(resp): + client = MagicMock() + client.__enter__ = lambda s: client + client.__exit__ = MagicMock(return_value=False) + client.post.return_value = resp + return client + + +class TestOllamaApiShowCaching: + def test_positive_result_cached_within_ttl(self): + from agent.model_metadata import _query_ollama_api_show + + client = _client_mock(_mock_show_response(131072)) + with patch("httpx.Client", return_value=client): + first = _query_ollama_api_show("llama3", "http://127.0.0.1:11434") + second = _query_ollama_api_show("llama3", "http://127.0.0.1:11434") + + assert first == second == 131072 + assert client.post.call_count == 1 # second call served from cache + + def test_failure_never_memoized(self): + """A down server must be re-probed on the next call (startup race).""" + from agent.model_metadata import _query_ollama_api_show + + bad = MagicMock() + bad.status_code = 404 + client = _client_mock(bad) + with patch("httpx.Client", return_value=client): + assert _query_ollama_api_show("llama3", "http://127.0.0.1:11434") is None + assert _query_ollama_api_show("llama3", "http://127.0.0.1:11434") is None + + assert client.post.call_count == 2 # None was NOT cached + + def test_cache_key_does_not_collide_with_local_ctx_probe(self): + """The ollama_show namespace must not read _query_local_context_length rows.""" + from agent import model_metadata + from agent.model_metadata import _query_ollama_api_show + import time as _time + + # Seed a same-(model,url) entry under the sibling probe's key shape. + model_metadata._LOCAL_CTX_PROBE_CACHE[("llama3", "http://127.0.0.1:11434")] = ( + 999, _time.monotonic(), + ) + + client = _client_mock(_mock_show_response(131072)) + with patch("httpx.Client", return_value=client): + result = _query_ollama_api_show("llama3", "http://127.0.0.1:11434") + + assert result == 131072 # probed for real, not the sibling's 999 + assert client.post.call_count == 1 + + +class TestContextCacheKeyNormalization: + def test_trailing_slash_variants_share_one_entry(self, tmp_path, monkeypatch): + from agent import model_metadata + + monkeypatch.setattr( + model_metadata, "_get_context_cache_path", + lambda: tmp_path / "context_lengths.yaml", + ) + + model_metadata.save_context_length("m1", "http://host/v1/", 200_000) + # Both slash variants resolve to the same row. + assert model_metadata.get_cached_context_length("m1", "http://host/v1") == 200_000 + assert model_metadata.get_cached_context_length("m1", "http://host/v1/") == 200_000 + + cache = model_metadata._load_context_cache() + assert list(cache.keys()) == ["m1@http://host/v1"] + + def test_legacy_unnormalized_row_still_honored(self, tmp_path, monkeypatch): + """Rows written pre-normalization (trailing slash in key) must not force a re-probe.""" + import yaml + from agent import model_metadata + + path = tmp_path / "context_lengths.yaml" + monkeypatch.setattr(model_metadata, "_get_context_cache_path", lambda: path) + path.write_text(yaml.dump({"context_lengths": {"m1@http://host/v1/": 128_000}})) + + assert model_metadata.get_cached_context_length("m1", "http://host/v1/") == 128_000 + + def test_invalidate_clears_both_key_shapes(self, tmp_path, monkeypatch): + import yaml + from agent import model_metadata + + path = tmp_path / "context_lengths.yaml" + monkeypatch.setattr(model_metadata, "_get_context_cache_path", lambda: path) + path.write_text(yaml.dump({"context_lengths": { + "m1@http://host/v1": 128_000, + "m1@http://host/v1/": 64_000, + }})) + + model_metadata._invalidate_cached_context_length("m1", "http://host/v1/") + cache = model_metadata._load_context_cache() + assert "m1@http://host/v1" not in cache + assert "m1@http://host/v1/" not in cache