fix(bedrock): cache only probed context windows, memoise probe failures
`_resolve_bedrock_context_length` means to persist only probe-derived context
windows -- its comment says "Only persist probe-derived values (region present);
a pure table fallback must not poison the cache" -- but the code under it cannot
tell the two apart. `get_bedrock_context_length(model, region=region,
probe=bool(region))` returns the probed window when the probe succeeds and the
static table (or the 128K default) when it does not, in the same int, and the
guard `if ctx and region:` tests the region, not the probe. The region is never
empty: `resolve_bedrock_region()` ends in `or "us-east-1"`.
Observed on a Bedrock deployment whose SSO session expires every 8 hours: the
first turn after expiry probes without credentials, `probe_bedrock_context_length`
returns None, and
context_lengths:
'<opus-5 global inference profile id>@bedrock:': 128000
is written (the synthetic `bedrock://` key lands on disk as `@bedrock:`, since
`_context_cache_key` strips trailing slashes) -- `BEDROCK_DEFAULT_CONTEXT_LENGTH`, because that model has no
`BEDROCK_CONTEXT_LENGTHS` row (#74263, addressed by #75824). The resolver returns
any cached value before it considers probing, so the probe never runs again for
that model and the compressor works from a 128K window on a 1M model.
The resolver now asks the probe itself and persists only what the probe said:
probe returned a window persist under model@base_url (or model@bedrock://), return it
probe returned None return the static table / default, persist nothing,
memoise the failure in memory for 5 minutes
cached value present serve it, as before
The memo keeps this cheap. `probe_bedrock_context_length` pads prompts of 1.3M and
2.2M tokens and attempts up to two `converse` calls (`_BEDROCK_PROBE_TIERS`; the
second tier is sent only when the first yields no parseable limit, i.e. on the
failure path), and `_cached_client` is a bare `boto3.client` that never validates
credentials, so a model whose probe keeps returning None WITH working credentials
(un-enabled model, opaque InternalServerException, unparseable length error)
would otherwise re-send both on every resolution -- and resolution is not memoised
per instance (model switch, every turn containing '@', each fallback candidate,
each /models API call). `_BEDROCK_PROBE_FAILURE_CACHE` /
`_BEDROCK_PROBE_FAILURE_TTL_SECONDS` follow the file's existing convention for
this shape (`_ENDPOINT_PROBE_FAILURE_TTL_SECONDS`, `_LOCAL_CTX_PROBE_CACHE`):
negative results in memory only, bounded, so credentials that come back are
noticed. A successful probe drops the entry, and
`_invalidate_cached_context_length` clears it for the model alongside the
local-probe memos, since a dropped entry is the reason to probe again.
Prior report: #68049 (open, 2026-07-20) fixes the same persistence defect with
the same core replacement (probe directly, persist only a probed window, return
`get_bedrock_context_length(model, probe=False)` otherwise). Its hunk targets the
inline code in `get_model_context_length` that has since moved into
`_resolve_bedrock_context_length`, so it no longer applies to main. This commit
adds, on top of that mechanism, the in-memory failure memo, its clearing on
invalidation, and a test of the memo's TTL.
Deliberately unchanged: the cached-value path of `_resolve_bedrock_context_length`
(a persisted value is served as on main; the step-1 table floor is not applied
to the `bedrock://` key, because a probed window may legitimately be below the
table and flooring it would re-probe every second call);
`get_bedrock_context_length` (signature, `probe=` parameter, table lookup and its
tests); `BEDROCK_CONTEXT_LENGTHS` (the missing Opus 5 row is #75824);
`save_context_length` / `get_cached_context_length`; the step ordering in
`get_model_context_length`.
tests/agent/test_model_metadata.py::TestBedrockContextCachePersistence, 2 cases:
a failed probe for a model with no table row persists nothing, and a failed probe
is memoised so a second resolution inside the TTL costs no probe and writes
nothing, then once the memo lapses the probe runs again and its window is what
gets persisted. Both fail on main, which persists 128000 on the first call.
(cherry picked from commit 5bc156abbf55feea79be5d3c66ff66ea838c03a4)
This commit is contained in:
@@ -320,6 +320,16 @@ MINIMUM_CONTEXT_LENGTH = 64_000
|
||||
# startup resolves the same model several times (banner, /model, compressor). Never persisted.
|
||||
_LOCAL_CTX_PROBE_TTL_SECONDS = 30.0
|
||||
_LOCAL_CTX_PROBE_CACHE: Dict[tuple, tuple] = {}
|
||||
# In-process (model, region) -> monotonic_ts memo of a FAILED Bedrock context probe. That probe pads
|
||||
# prompts of 1.3M/2.2M tokens and attempts up to two converse calls (agent/bedrock_adapter.py
|
||||
# _BEDROCK_PROBE_TIERS; the second tier is sent only when the first yields no parseable limit, i.e. on
|
||||
# the failure path), and its failures are deliberately never persisted, so without a negative memo a
|
||||
# model whose probe keeps failing (un-enabled model, opaque InternalServerException, unparseable length
|
||||
# error) would re-send them on every resolution. Negative only, in memory only, and bounded — same
|
||||
# reasoning as _ENDPOINT_PROBE_FAILURE_TTL_SECONDS: a failure is usually transient (expired SSO
|
||||
# session, offline box), so it must expire rather than stick.
|
||||
_BEDROCK_PROBE_FAILURE_TTL_SECONDS = 300.0
|
||||
_BEDROCK_PROBE_FAILURE_CACHE: Dict[tuple, float] = {}
|
||||
# Family-pattern fallbacks, used only when provider-aware sources all miss.
|
||||
# Lookups are longest-key-first substring matches, so dict order is cosmetic
|
||||
# and a specific key must be STRICTLY longer than its catch-all.
|
||||
@@ -1145,6 +1155,11 @@ def _invalidate_cached_context_length(model: str, base_url: str) -> None:
|
||||
bare, stripped = _strip_provider_prefix(model), (base_url or "").rstrip("/")
|
||||
_LOCAL_CTX_PROBE_CACHE.pop((bare, stripped), None)
|
||||
_LOCAL_CTX_PROBE_CACHE.pop(("ollama_show", bare, stripped), None)
|
||||
# Same for a memoised Bedrock probe failure (keyed by region, which the caller does not know):
|
||||
# the entry being dropped is the reason to ask the probe again, not to wait out its TTL.
|
||||
for memo_key in list(_BEDROCK_PROBE_FAILURE_CACHE): # snapshot: another thread may be memoising
|
||||
if memo_key[0] in (model, bare):
|
||||
_BEDROCK_PROBE_FAILURE_CACHE.pop(memo_key, None)
|
||||
# Every key shape get_cached_context_length consults.
|
||||
stale_keys = {key, f"{model}@{base_url}", f"{key}/"}
|
||||
if not any(k in cache for k in stale_keys):
|
||||
@@ -1813,12 +1828,20 @@ def _validate_cached_context_length(model: str, base_url: str, cached: int, is_b
|
||||
return cached
|
||||
|
||||
|
||||
def _bedrock_probe_failed_recently(model: str, region: str) -> bool:
|
||||
"""True while a failed Bedrock context probe for *model* in *region* is still memoised
|
||||
(see _BEDROCK_PROBE_FAILURE_CACHE): answer from the static table without re-probing."""
|
||||
failed_at = _BEDROCK_PROBE_FAILURE_CACHE.get((model, region))
|
||||
return failed_at is not None and (time.monotonic() - failed_at) < _BEDROCK_PROBE_FAILURE_TTL_SECONDS
|
||||
|
||||
|
||||
def _resolve_bedrock_context_length(model: str, base_url: str) -> Optional[int]:
|
||||
"""Step 1b: Bedrock static table + one cached live probe (Bedrock exposes no context window via
|
||||
metadata APIs); None when boto3 is absent. Cached per model under base_url, else a synthetic
|
||||
bedrock:// key so display/offline paths share it."""
|
||||
metadata APIs); None when boto3 is absent. Only a PROBED window is cached (the table answers a
|
||||
call, never the cache), per model under base_url, else a synthetic bedrock:// key so
|
||||
display/offline paths share it."""
|
||||
try:
|
||||
from agent.bedrock_adapter import get_bedrock_context_length, resolve_bedrock_region
|
||||
from agent.bedrock_adapter import get_bedrock_context_length, probe_bedrock_context_length, resolve_bedrock_region
|
||||
except ImportError:
|
||||
return None # boto3 not installed — fall through to generic resolution
|
||||
cache_key_url = base_url or "bedrock://"
|
||||
@@ -1831,11 +1854,17 @@ def _resolve_bedrock_context_length(model: str, base_url: str) -> Optional[int]:
|
||||
if not region:
|
||||
with contextlib.suppress(Exception):
|
||||
region = resolve_bedrock_region()
|
||||
ctx = get_bedrock_context_length(model, region=region, probe=bool(region))
|
||||
# Only persist probe-derived values (region present); a pure table fallback must not poison the cache.
|
||||
if ctx and region:
|
||||
save_context_length(model, cache_key_url, ctx)
|
||||
return ctx
|
||||
if region and not _bedrock_probe_failed_recently(model, region):
|
||||
probed = probe_bedrock_context_length(model, region)
|
||||
if probed:
|
||||
# The probe is the only authoritative source, so it is the only thing worth persisting:
|
||||
# a table fallback written here would be served forever (this branch runs before it),
|
||||
# and the probe would never be consulted for the model again.
|
||||
save_context_length(model, cache_key_url, probed)
|
||||
_BEDROCK_PROBE_FAILURE_CACHE.pop((model, region), None) # success ends the failure window
|
||||
return probed
|
||||
_BEDROCK_PROBE_FAILURE_CACHE[(model, region)] = time.monotonic()
|
||||
return get_bedrock_context_length(model, probe=False) # static table / default: answers this call only
|
||||
|
||||
|
||||
def _resolve_custom_endpoint_context_length(model: str, base_url: str, api_key: str, provider: str) -> int:
|
||||
|
||||
@@ -1292,6 +1292,73 @@ class TestBedrockContextResolution:
|
||||
assert mock_fetch.called
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# Bedrock context cache persistence — only a probe result may be persisted
|
||||
# =========================================================================
|
||||
|
||||
class TestBedrockContextCachePersistence:
|
||||
"""``_resolve_bedrock_context_length`` persisted whatever
|
||||
``get_bedrock_context_length`` returned whenever a region was resolvable —
|
||||
and ``resolve_bedrock_region()`` always resolves one (it ends in
|
||||
``or "us-east-1"``). A probe that returned None (expired SSO session,
|
||||
offline, opaque server error) therefore froze the static table value — the
|
||||
128K default for a model with no table row — under ``model@<base_url>`` or
|
||||
``model@bedrock://`` (written as ``model@bedrock:``), and the probe, which
|
||||
the resolver treats as the only authoritative source, never ran for that
|
||||
model again.
|
||||
|
||||
Invariants: only a probe-derived window may be persisted, and a failed
|
||||
probe is memoised in memory for ``_BEDROCK_PROBE_FAILURE_TTL_SECONDS`` so
|
||||
it is not re-sent on every resolution, yet runs again once the memo lapses.
|
||||
"""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_bedrock_probe_memo(self):
|
||||
"""The failure memo is module-level state; it must not leak between tests."""
|
||||
from agent import model_metadata as mm
|
||||
mm._BEDROCK_PROBE_FAILURE_CACHE.clear()
|
||||
yield
|
||||
mm._BEDROCK_PROBE_FAILURE_CACHE.clear()
|
||||
|
||||
@patch("agent.bedrock_adapter.resolve_bedrock_region", return_value="us-east-1")
|
||||
@patch("agent.bedrock_adapter.probe_bedrock_context_length", return_value=None)
|
||||
def test_failed_probe_does_not_persist_static_fallback(self, mock_probe, mock_region, tmp_path):
|
||||
"""A failed probe answers from the table (the 128K default here) but writes nothing
|
||||
to disk. On main this persists ``amazon.future-model-v1:0@bedrock:: 128000``."""
|
||||
from agent.bedrock_adapter import BEDROCK_DEFAULT_CONTEXT_LENGTH
|
||||
model = "amazon.future-model-v1:0" # no BEDROCK_CONTEXT_LENGTHS row
|
||||
cache_file = tmp_path / "context_length_cache.yaml"
|
||||
with patch("agent.model_metadata._get_context_cache_path", return_value=cache_file):
|
||||
assert get_model_context_length(model, provider="bedrock") == BEDROCK_DEFAULT_CONTEXT_LENGTH
|
||||
assert get_cached_context_length(model, "bedrock://") is None
|
||||
assert not cache_file.exists()
|
||||
mock_probe.assert_called_once_with(model, "us-east-1")
|
||||
|
||||
@patch("agent.bedrock_adapter.resolve_bedrock_region", return_value="us-east-1")
|
||||
@patch("agent.bedrock_adapter.probe_bedrock_context_length", side_effect=[None, 1_000_000])
|
||||
def test_failed_probe_is_memoised_until_the_ttl_lapses(self, mock_probe, mock_region, tmp_path):
|
||||
"""Not persisting must not turn the probe into a per-resolution cost: inside the TTL a
|
||||
second resolution answers from the table without re-probing and still writes nothing;
|
||||
once the memo lapses the probe runs again and its window is what gets persisted. On
|
||||
main the first call persists 128K and every later call serves it."""
|
||||
from agent import model_metadata as mm
|
||||
from agent.bedrock_adapter import BEDROCK_DEFAULT_CONTEXT_LENGTH
|
||||
model = "amazon.future-model-v1:0"
|
||||
cache_file = tmp_path / "context_length_cache.yaml"
|
||||
with patch("agent.model_metadata._get_context_cache_path", return_value=cache_file):
|
||||
assert get_model_context_length(model, provider="bedrock") == BEDROCK_DEFAULT_CONTEXT_LENGTH
|
||||
assert get_model_context_length(model, provider="bedrock") == BEDROCK_DEFAULT_CONTEXT_LENGTH
|
||||
assert mock_probe.call_count == 1
|
||||
assert not cache_file.exists() # memoised in memory only
|
||||
# Age the entry past the failure TTL (as tests/agent/test_probe_cache_followups.py does).
|
||||
mm._BEDROCK_PROBE_FAILURE_CACHE[(model, "us-east-1")] = (
|
||||
time.monotonic() - mm._BEDROCK_PROBE_FAILURE_TTL_SECONDS - 1
|
||||
)
|
||||
assert get_model_context_length(model, provider="bedrock") == 1_000_000
|
||||
assert get_cached_context_length(model, "bedrock://") == 1_000_000
|
||||
assert mock_probe.call_count == 2
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# _strip_provider_prefix — Ollama model:tag vs provider:model
|
||||
# =========================================================================
|
||||
|
||||
Reference in New Issue
Block a user