From 541b20290c306f93fae4abf0fa3221ef55873b3e Mon Sep 17 00:00:00 2001 From: Victor Kyriazakos Date: Mon, 14 Sep 2026 21:00:51 +0000 Subject: [PATCH] fix(bedrock): restore Grok context with provider-confirmed cache provenance --- agent/bedrock_adapter.py | 2 + agent/model_metadata.py | 146 +++++++---- agent/turn_overflow.py | 4 +- agent/turn_usage.py | 4 +- tests/agent/test_bedrock_context_cache.py | 246 ++++++++++++++++++ tests/agent/test_model_metadata.py | 7 +- .../context-compression-and-caching.md | 29 +++ 7 files changed, 383 insertions(+), 55 deletions(-) create mode 100644 tests/agent/test_bedrock_context_cache.py diff --git a/agent/bedrock_adapter.py b/agent/bedrock_adapter.py index c71e59d1f7..485626a31a 100644 --- a/agent/bedrock_adapter.py +++ b/agent/bedrock_adapter.py @@ -1049,6 +1049,8 @@ def _extract_provider_from_arn(arn: str) -> str: # substring, so versioned entries win over the generic "anthropic.claude-opus-4". BEDROCK_CONTEXT_LENGTHS: Dict[str, int] = { + # https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-6.html + "xai.grok-4.6": 500_000, # Anthropic Claude: 1M GA vs 200K. The 1M entries must match agent/model_metadata.py # DEFAULT_CONTEXT_LENGTHS or context compresses early. **dict.fromkeys(( diff --git a/agent/model_metadata.py b/agent/model_metadata.py index 940a5107ad..84ccceaabd 100644 --- a/agent/model_metadata.py +++ b/agent/model_metadata.py @@ -1097,23 +1097,37 @@ def _get_context_cache_path() -> Path: return get_hermes_home() / "context_length_cache.yaml" -def _load_context_cache() -> Dict[str, int]: - """Load the model+provider -> context_length cache from disk.""" +def _load_context_cache_document() -> dict: + """Read scalar lengths and their provenance from the same atomic document.""" path = _get_context_cache_path() if not path.exists(): return {} try: with open(path, encoding="utf-8") as f: - return (yaml.safe_load(f) or {}).get("context_lengths") or {} + data = yaml.safe_load(f) + if not isinstance(data, dict): + return {} + for section in ("context_lengths", "bedrock_confirmed_v1"): + if not isinstance(data.get(section), dict): + data[section] = {} + return data except Exception as e: logger.debug("Failed to load context length cache: %s", e) return {} -def _write_context_cache(cache: Dict[str, int]) -> None: +def _load_context_cache() -> Dict[str, int]: + """Load scalar lengths, preserving the legacy reader contract.""" + return _load_context_cache_document().get("context_lengths") or {} + + +def _write_context_cache(cache: Dict[str, int], bedrock_confirmed: dict | None = None) -> None: """Atomic write (a truncating write killed mid-dump leaves a partial file that _load_context_cache() swallows as {}, wiping EVERY cached length). Raises on failure.""" - atomic_yaml_write(_get_context_cache_path(), {"context_lengths": cache}) + document = {"context_lengths": cache} + if bedrock_confirmed: + document["bedrock_confirmed_v1"] = bedrock_confirmed + atomic_yaml_write(_get_context_cache_path(), document) def _context_cache_key(model: str, base_url: str) -> str: @@ -1121,53 +1135,88 @@ def _context_cache_key(model: str, base_url: str) -> str: return f"{model}@{(base_url or '').rstrip('/')}" -def save_context_length(model: str, base_url: str, length: int) -> None: +def save_context_length(model: str, base_url: str, length: int, *, source: str = "") -> None: """Persist a discovered context length under ``model@base_url`` (same model, different providers, different limits).""" # 0/negative is always a bug and would make get_model_context_length() return 0 (`0 is not None`). if length <= 0: logger.warning("Refusing to cache non-positive context length %s -> %s tokens", f"{model}@{base_url}", length) return key = _context_cache_key(model, base_url) - cache = _load_context_cache() - if cache.get(key) == length: + document = _load_context_cache_document() + cache = document.get("context_lengths") or {} + confirmed = document.get("bedrock_confirmed_v1") or {} + confirmed = confirmed if isinstance(confirmed, dict) else {} + old_confirmed = confirmed.copy() + # Generic writes revoke provenance even for the same number. The marker + # binds to the value, so an older writer cannot leave mismatched evidence. + for alias in (key, f"{model}@{base_url}", f"{key}/"): + confirmed.pop(alias, None) + if source == "bedrock-confirmed-v1": + confirmed[key] = length + if cache.get(key) == length and confirmed == old_confirmed: return # already stored cache[key] = length try: - _write_context_cache(cache) + _write_context_cache(cache, confirmed) logger.info("Cached context length %s -> %s tokens", key, f"{length:,}") except Exception as e: logger.debug("Failed to save context length cache: %s", e) -def get_cached_context_length(model: str, base_url: str) -> Optional[int]: +def save_provider_context_length(model: str, base_url: str, length: int, provider: str = "") -> None: + """Persist a provider-confirmed window, distinguishing it from legacy Bedrock fallbacks.""" + if _is_bedrock_context(base_url, provider): + save_context_length(model, base_url or "bedrock://", length, source="bedrock-confirmed-v1") + else: + save_context_length(model, base_url, length) + + +def get_cached_context_length(model: str, base_url: str, *, bedrock_confirmed: bool = False) -> Optional[int]: """Look up a previously discovered context length for model+provider.""" key = _context_cache_key(model, base_url) - cache = _load_context_cache() + document = _load_context_cache_document() + cache = document.get("context_lengths") or {} + if not isinstance(cache, dict): + return None # Legacy rows may carry a trailing slash, so probe the canonical key, the literal form and the slashed canonical form. - return next((hit for hit in map(cache.get, (key, f"{model}@{base_url}", f"{key}/")) if hit is not None), None) + matched_key = next((k for k in (key, f"{model}@{base_url}", f"{key}/") if cache.get(k) is not None), None) + length = cache.get(matched_key) + if type(length) is not int: + return None + if bedrock_confirmed: + confirmed = document.get("bedrock_confirmed_v1") + marker = confirmed.get(matched_key) if isinstance(confirmed, dict) else None + if type(marker) is not int or marker != length: + return None + return length 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 = _context_cache_key(model, base_url) - cache = _load_context_cache() + document = _load_context_cache_document() + cache = document.get("context_lengths") or {} + confirmed = document.get("bedrock_confirmed_v1") or {} + confirmed = confirmed if isinstance(confirmed, dict) else {} # Also drop the in-memory TTL probe entries, or the next resolution inside the TTL window reuses the stale value. 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. + from hermes_constants import hermes_home_key for memo_key in list(_BEDROCK_PROBE_FAILURE_CACHE): # snapshot: another thread may be memoising - if memo_key[0] in (model, bare): + if memo_key[:2] == (hermes_home_key(), stripped) and memo_key[2] 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): + if not any(k in cache or k in confirmed for k in stale_keys): return for k in stale_keys: cache.pop(k, None) + confirmed.pop(k, None) try: - _write_context_cache(cache) + _write_context_cache(cache, confirmed) except Exception as e: logger.debug("Failed to invalidate context length cache entry %s: %s", key, e) @@ -1783,7 +1832,7 @@ def _resolve_nous_context_length(model: str, base_url: str = "", api_key: str = return None, "" -def _validate_cached_context_length(model: str, base_url: str, cached: int, is_bedrock_context: bool, *, api_key: str = "") -> Optional[int]: +def _validate_cached_context_length(model: str, base_url: str, cached: int, *, api_key: str = "") -> Optional[int]: """Step 1 of get_model_context_length: accept, repair, or drop a persisted entry. Returns the value to use, or None to fall through to live resolution. Order matters: a value must be rejected as bogus before any provider-specific handling.""" @@ -1807,18 +1856,7 @@ def _validate_cached_context_length(model: str, base_url: str, cached: int, is_b if _infer_provider_from_url(base_url) == "nous": logger.debug("Bypassing persistent cache for %s@%s (Nous portal authoritative)", model, base_url) return None - if is_bedrock_context: - # Bedrock: the static table is a FLOOR — probe-derived entries may legitimately exceed it. - try: - from agent.bedrock_adapter import get_bedrock_context_length - bedrock_ctx = get_bedrock_context_length(model) - except ImportError: - return cached - if cached < bedrock_ctx: - logger.info("Dropping stale Bedrock cache entry %s@%s -> %s; using static Bedrock table value %s", model, base_url, f"{cached:,}", f"{bedrock_ctx:,}") - _invalidate_cached_context_length(model, base_url) - return bedrock_ctx - return cached + # For local endpoints, run the probe that respects configured Modelfile context values first. # _query_local_context_length prefers num_ctx from Modelfile, while _query_ollama_api_show returns the # GGUF training max first which can be larger and would create a false-safe window for compression @@ -1828,42 +1866,56 @@ 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 _bedrock_probe_failed_recently(key: tuple) -> bool: + """Prune expired failures and check this home/endpoint/model/region's cooldown.""" + now = time.monotonic() + for memo_key, failed_at in list(_BEDROCK_PROBE_FAILURE_CACHE.items()): + if now - failed_at >= _BEDROCK_PROBE_FAILURE_TTL_SECONDS: + _BEDROCK_PROBE_FAILURE_CACHE.pop(memo_key, None) + return key in _BEDROCK_PROBE_FAILURE_CACHE + + +def _is_bedrock_context(base_url: str, provider: str = "") -> bool: + return provider == "bedrock" or bool( + base_url and base_url_hostname(base_url).startswith("bedrock-runtime.") + and base_url_host_matches(base_url, "amazonaws.com") + ) 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. 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.""" + metadata APIs); None when boto3 is absent. Only provider-confirmed windows from a probe + or runtime error are reused. The table answers a call, never the cache. Keys use base_url, + or synthetic bedrock:// when absent, consistently with provider-error writers.""" try: 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://" - cached = get_cached_context_length(model, cache_key_url) - if cached is not None: + cached = get_cached_context_length(model, cache_key_url, bedrock_confirmed=True) + if cached is not None and cached > 0: return cached + # Legacy scalars have no trustworthy source. Ignore them until a successful + # probe replaces them; deleting here could erase a concurrent confirmed write + # or reset the failure cooldown repeatedly when the cache is read-only. # Region from the base_url host first, then the standard AWS chain. An empty region disables probing (table only). _m = re.search(r"bedrock-runtime\.([a-z0-9-]+)\.", base_url) if base_url else None region = _m.group(1) if _m else "" if not region: with contextlib.suppress(Exception): region = resolve_bedrock_region() - if region and not _bedrock_probe_failed_recently(model, region): + from hermes_constants import hermes_home_key + memo_key = (hermes_home_key(), cache_key_url.rstrip('/'), model, region) + if region and not _bedrock_probe_failed_recently(memo_key): 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 + save_provider_context_length(model, cache_key_url, probed, provider="bedrock") + _BEDROCK_PROBE_FAILURE_CACHE.pop(memo_key, None) # success ends the failure window return probed - _BEDROCK_PROBE_FAILURE_CACHE[(model, region)] = time.monotonic() + _BEDROCK_PROBE_FAILURE_CACHE[memo_key] = time.monotonic() return get_bedrock_context_length(model, probe=False) # static table / default: answers this call only @@ -2050,12 +2102,10 @@ def get_model_context_length( endpoint_context = _endpoint_scoped_context_length(model, base_url) if endpoint_context is not None: return endpoint_context - is_bedrock_context = provider == "bedrock" or ( - base_url and base_url_hostname(base_url).startswith("bedrock-runtime.") and base_url_host_matches(base_url, "amazonaws.com") - ) + is_bedrock_context = _is_bedrock_context(base_url, provider) # 1. Persistent cache (LM Studio / Codex OAuth excluded — see _skip_persistent_context_cache). - cached = get_cached_context_length(model, base_url) if base_url and not _skip_persistent_context_cache(base_url, provider) else None - validated = _validate_cached_context_length(model, base_url, cached, is_bedrock_context, api_key=api_key) if cached is not None else None + cached = get_cached_context_length(model, base_url) if base_url and not is_bedrock_context and not _skip_persistent_context_cache(base_url, provider) else None + validated = _validate_cached_context_length(model, base_url, cached, api_key=api_key) if cached is not None else None if validated is not None: return validated # 1b. AWS Bedrock. Must run BEFORE the custom-endpoint step: bedrock-runtime.* is not in diff --git a/agent/turn_overflow.py b/agent/turn_overflow.py index 8bfe095dc6..925a328ccb 100644 --- a/agent/turn_overflow.py +++ b/agent/turn_overflow.py @@ -316,7 +316,7 @@ def _adopt_provider_context_limit(st: _Recovery, error_msg: str, old_ctx: int) - """Shrink context_length only when the provider reports the real limit; else keep the window and compress. Guessed probe tiers can turn a configured 1M window into 256K/128K/64K. Returns the provider-reported limit, or ``None``.""" - from agent.model_metadata import save_context_length + from agent.model_metadata import save_provider_context_length agent = st.agent compressor = agent.context_compressor @@ -330,7 +330,7 @@ def _adopt_provider_context_limit(st: _Recovery, error_msg: str, old_ctx: int) - # Persist the provider-reported limit BEFORE compression/retry: rate limit, # missing usage, or restart must not lose confirmed metadata. Probe flags # remain a fallback if this write fails. - save_context_length(agent.model, agent.base_url, new_ctx) + save_provider_context_length(agent.model, agent.base_url, new_ctx, agent.provider) # Probe flags only on the built-in compressor (plugin engines manage their # own); provider-sourced value, so safe to cache. if hasattr(compressor, "_context_probed"): diff --git a/agent/turn_usage.py b/agent/turn_usage.py index 8d621fd6e5..b97159839b 100644 --- a/agent/turn_usage.py +++ b/agent/turn_usage.py @@ -163,9 +163,9 @@ def record_response_usage( if getattr(compressor, "_context_probed", False): ctx = compressor.context_length if getattr(compressor, "_context_probe_persistable", False): - from agent.model_metadata import save_context_length + from agent.model_metadata import save_provider_context_length - save_context_length(agent.model, agent.base_url, ctx) + save_provider_context_length(agent.model, agent.base_url, ctx, agent.provider) agent._safe_print(f"{agent.log_prefix}💾 Cached context length: {ctx:,} tokens for {agent.model}") compressor._context_probed = False compressor._context_probe_persistable = False diff --git a/tests/agent/test_bedrock_context_cache.py b/tests/agent/test_bedrock_context_cache.py new file mode 100644 index 0000000000..5321f482bd --- /dev/null +++ b/tests/agent/test_bedrock_context_cache.py @@ -0,0 +1,246 @@ +"""Bedrock cache provenance and compressor budgets across disk-backed restarts. + +AWS documents Grok 4.6's Bedrock context window as 500K, independently of +xAI's direct API catalog: https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-6.html +Only provider I/O is stubbed; cache/config readers and compressor are real. +""" + +import json +import os +import subprocess +import sys +import time +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import yaml + +from agent import bedrock_adapter as ba +from agent import model_metadata as mm +from agent.context_compressor import ContextCompressor + + +@pytest.fixture(autouse=True) +def isolated_home(tmp_path, monkeypatch): + monkeypatch.setenv("HERMES_HOME", str(tmp_path)) + mm._BEDROCK_PROBE_FAILURE_CACHE.clear() + monkeypatch.setattr(ba, "resolve_bedrock_region", lambda: "us-east-1") + yield tmp_path + mm._BEDROCK_PROBE_FAILURE_CACHE.clear() + + +@pytest.mark.parametrize("model", ["xai.grok-4.6", "global.xai.grok-4.6", "us.xai.grok-4.6"]) +@pytest.mark.parametrize("base_url", ["", "https://bedrock-runtime.us-east-1.amazonaws.com"]) +@pytest.mark.parametrize("legacy", [None, 128_000, 700_000]) +@pytest.mark.parametrize("probed", [None, 128_000, 800_000]) +def test_bedrock_resolution_migrates_ambiguous_cache_and_preserves_probe( + isolated_home, monkeypatch, model, base_url, legacy, probed, +): + """Neither an old small nor large scalar proves where it came from. + + A successful probe is authoritative, including below-table limits. Its + provenance must survive unrelated writes and an actual process restart. + """ + cache_url = base_url or "bedrock://" + key = mm._context_cache_key(model, cache_url) + cache_file = isolated_home / "context_length_cache.yaml" + if legacy is not None: + cache_file.write_text(yaml.safe_dump({"context_lengths": {key: legacy}})) + probe = Mock(return_value=probed) + monkeypatch.setattr(ba, "probe_bedrock_context_length", probe) + expected = probed if probed is not None else 500_000 + compressor = ContextCompressor(model, provider="bedrock", base_url=base_url, quiet_mode=True) + assert compressor.context_length == expected + # Preserve the existing raise-only 75% floor for windows below 512K. + expected_threshold = int(expected * (0.75 if expected < 512_000 else 0.5)) + assert compressor.threshold_tokens == expected_threshold + assert mm.get_model_context_length(model, provider="bedrock", base_url=base_url) == expected + probe.assert_called_once_with(model, "us-east-1") + assert mm.get_cached_context_length(model, cache_url, bedrock_confirmed=True) == probed + # Ordinary cache updates must not erase another entry's provenance. + mm.save_context_length("other-model", "https://other.example/v1", 64_000) + mm._invalidate_cached_context_length("other-model", "https://other.example/v1") + if probed is None: + assert yaml.safe_load(cache_file.read_text())["context_lengths"].get(key) == legacy + else: + script = ''' +import json, sys +from agent import bedrock_adapter as ba +from agent.context_compressor import ContextCompressor +ba.probe_bedrock_context_length = lambda *a, **k: (_ for _ in ()).throw(AssertionError("reprobed persisted success")) +c = ContextCompressor(sys.argv[1], provider="bedrock", base_url=sys.argv[2], quiet_mode=True) +print(json.dumps([c.context_length, c.threshold_tokens])) +''' + result = subprocess.run([sys.executable, "-c", script, model, base_url], + env=dict(os.environ), text=True, capture_output=True, check=True) + assert json.loads(result.stdout) == [expected, expected_threshold] + + +@pytest.mark.parametrize("base_url", ["", "https://bedrock-runtime.us-east-1.amazonaws.com"]) +@pytest.mark.parametrize("retry", ["ttl", "invalidate", "restart", "profile", "endpoint"]) +def test_failure_memo_retry_scope_and_expiry(isolated_home, monkeypatch, base_url, retry): + model = "global.xai.grok-4.6" + cache_url = base_url or "bedrock://" + probe = Mock(side_effect=[None, 96_000]) + monkeypatch.setattr(ba, "probe_bedrock_context_length", probe) + assert mm.get_model_context_length(model, provider="bedrock", base_url=base_url) == 500_000 + assert mm.get_model_context_length(model, provider="bedrock", base_url=base_url) == 500_000 + assert probe.call_count == 1 + assert not (isolated_home / "context_length_cache.yaml").exists() + if retry == "ttl": + expired = time.monotonic() - mm._BEDROCK_PROBE_FAILURE_TTL_SECONDS - 1 + for key in mm._BEDROCK_PROBE_FAILURE_CACHE: + mm._BEDROCK_PROBE_FAILURE_CACHE[key] = expired + elif retry == "invalidate": + mm._invalidate_cached_context_length(model, cache_url) + elif retry == "restart": + mm._BEDROCK_PROBE_FAILURE_CACHE.clear() + elif retry == "profile": + new_home = isolated_home / "other-profile" + new_home.mkdir() + monkeypatch.setenv("HERMES_HOME", str(new_home)) + else: + base_url = "https://bedrock-runtime.us-east-1.amazonaws.com/other" + assert mm.get_model_context_length(model, provider="bedrock", base_url=base_url) == 96_000 + assert probe.call_count == 2 + if retry == "ttl": + assert not mm._BEDROCK_PROBE_FAILURE_CACHE + + +@pytest.mark.parametrize("base_url", ["", "https://bedrock-runtime.us-east-1.amazonaws.com"]) +@pytest.mark.parametrize("override", ["argument", "config"]) +def test_explicit_context_override_and_compressor_caps_win(isolated_home, monkeypatch, base_url, override): + model = "us.xai.grok-4.6" + probe = Mock(side_effect=AssertionError("explicit override must not probe")) + monkeypatch.setattr(ba, "probe_bedrock_context_length", probe) + explicit_context = 80_000 if override == "argument" else None + if override == "config": + (isolated_home / "config.yaml").write_text(yaml.safe_dump({ + "model_overrides": {"bedrock": {model: {"context_window": 80_000}}}, + })) + compressor = ContextCompressor(model, provider="bedrock", base_url=base_url, + threshold_tokens_cap=30_000, max_tokens=10_000, + quiet_mode=True, config_context_length=explicit_context) + assert compressor.context_length == 80_000 + assert compressor.threshold_tokens == 30_000 + assert not (isolated_home / "context_length_cache.yaml").exists() + probe.assert_not_called() + + +def test_unknown_model_fallback_and_host_inference(monkeypatch): + probe = Mock(return_value=None) + monkeypatch.setattr(ba, "probe_bedrock_context_length", probe) + base_url = "https://bedrock-runtime.us-east-1.amazonaws.com" + assert mm.get_model_context_length("unknown.future-model", base_url=base_url) == ba.BEDROCK_DEFAULT_CONTEXT_LENGTH + assert mm.get_model_context_length("xai.grok-4.6", base_url=base_url) == 500_000 + assert mm.get_cached_context_length("unknown.future-model", base_url) is None + + +@pytest.mark.parametrize("base_url", ["", "https://bedrock-runtime.us-east-1.amazonaws.com"]) +@pytest.mark.parametrize("writer", ["overflow", "usage"]) +def test_provider_confirmed_writers_survive_restart(monkeypatch, base_url, writer): + from agent.turn_overflow import _adopt_provider_context_limit + from agent.turn_usage import record_response_usage + + model = "global.xai.grok-4.6" + compressor = ContextCompressor(model, base_url=base_url, provider="bedrock", quiet_mode=True) + compressor.context_length = 500_000 + agent = SimpleNamespace(model=model, provider="bedrock", api_mode="bedrock", base_url=base_url, + context_compressor=compressor, _buffer_vprint=lambda *a: None, + _safe_print=lambda *a: None, log_prefix="", client=None, + _session_db=None, verbose_logging=False, quiet_mode=True, + session_api_calls=0, session_estimated_cost_usd=0) + if writer == "overflow": + assert _adopt_provider_context_limit(SimpleNamespace(agent=agent), + "maximum context length is 96000 tokens", 500_000) == 96_000 + else: + compressor.context_length = 96_000 + compressor._context_probed = compressor._context_probe_persistable = True + for name in ("prompt", "completion", "total", "input", "output", "cache_read", "cache_write", "reasoning"): + setattr(agent, f"session_{name}_tokens", 0) + response = SimpleNamespace(usage={"input_tokens": 100, "output_tokens": 5}) + record_response_usage(agent, response, messages=[{"role": "user", "content": "hi"}], + api_call_count=1, api_duration=0.1, compression_attempts=0, max_compression_attempts=3) + assert mm.get_cached_context_length(model, base_url or "bedrock://") == 96_000 + script = ''' +import sys +from agent import model_metadata as mm, bedrock_adapter as ba +ba.probe_bedrock_context_length = lambda *a, **k: (_ for _ in ()).throw(AssertionError("lost provider limit")) +assert mm.get_model_context_length(sys.argv[1], base_url=sys.argv[2], provider="bedrock") == 96000 +''' + subprocess.run([sys.executable, "-c", script, model, base_url], check=True, capture_output=True, text=True) + + +@pytest.mark.parametrize("base_url", ["", "https://bedrock-runtime.us-east-1.amazonaws.com"]) +def test_readonly_legacy_cache_does_not_reset_probe_cooldown(isolated_home, monkeypatch, base_url): + model = "global.xai.grok-4.6" + key = mm._context_cache_key(model, base_url or "bedrock://") + cache_file = isolated_home / "context_length_cache.yaml" + cache_file.write_text(yaml.safe_dump({"context_lengths": {key: 128_000}})) + probe = Mock(return_value=None) + monkeypatch.setattr(ba, "probe_bedrock_context_length", probe) + monkeypatch.setattr(mm, "_write_context_cache", Mock(side_effect=OSError("read-only"))) + for _ in range(3): + assert mm.get_model_context_length(model, provider="bedrock", base_url=base_url) == 500_000 + assert probe.call_count == 1 + + +@pytest.mark.parametrize("rewrite", [None, 96_000, 120_000]) +def test_provenance_is_backward_readable_and_generic_writes_clear_it(isolated_home, monkeypatch, rewrite): + model, base_url = "xai.grok-4.6", "https://bedrock-runtime.us-east-1.amazonaws.com" + probe = Mock(side_effect=[96_000, 110_000]) + monkeypatch.setattr(ba, "probe_bedrock_context_length", probe) + assert mm.get_model_context_length(model, base_url=base_url) == 96_000 + raw = yaml.safe_load((isolated_home / "context_length_cache.yaml").read_text()) + # Old readers take this value directly into arithmetic. Metadata is additive. + assert raw["context_lengths"][mm._context_cache_key(model, base_url)] + 1 == 96_001 + if rewrite is not None: + mm.save_context_length(model, base_url, rewrite) + assert mm.get_model_context_length(model, base_url=base_url) == 110_000 + else: + assert mm.get_model_context_length(model, base_url=base_url) == 96_000 + + +@pytest.mark.parametrize("lengths", [True, [128_000], "bad"]) +def test_malformed_lengths_do_not_block_provider_persistence(isolated_home, monkeypatch, lengths): + (isolated_home / "context_length_cache.yaml").write_text(yaml.safe_dump({"context_lengths": lengths})) + monkeypatch.setattr(ba, "probe_bedrock_context_length", lambda *a: 96_000) + assert mm.get_model_context_length("xai.grok-4.6", provider="bedrock") == 96_000 + assert mm.get_cached_context_length("xai.grok-4.6", "bedrock://") == 96_000 + + +@pytest.mark.parametrize("marker", [True, "96000", 128_000, [96_000], {"source": "probe"}]) +def test_malformed_or_mismatched_provenance_requires_revalidation(isolated_home, monkeypatch, marker): + model, base_url = "xai.grok-4.6", "bedrock://" + key = mm._context_cache_key(model, base_url) + (isolated_home / "context_length_cache.yaml").write_text(yaml.safe_dump({ + "context_lengths": {key: 96_000}, "bedrock_confirmed_v1": {key: marker}, + })) + probe = Mock(return_value=100_000) + monkeypatch.setattr(ba, "probe_bedrock_context_length", probe) + assert mm.get_model_context_length(model, provider="bedrock") == 100_000 + probe.assert_called_once() + + +def test_context_local_profile_memos_do_not_cross_and_expired_rows_are_pruned(isolated_home, monkeypatch): + from hermes_constants import set_hermes_home_override, reset_hermes_home_override + + probe = Mock(side_effect=[None, None, 96_000]) + monkeypatch.setattr(ba, "probe_bedrock_context_length", probe) + model = "global.xai.grok-4.6" + assert mm.get_model_context_length(model, provider="bedrock") == 500_000 + # Leave a different model's expired row: lookup must prune it too. + assert mm.get_model_context_length("unknown.future", provider="bedrock") == ba.BEDROCK_DEFAULT_CONTEXT_LENGTH + for key in mm._BEDROCK_PROBE_FAILURE_CACHE: + if "unknown.future" in key: + mm._BEDROCK_PROBE_FAILURE_CACHE[key] = time.monotonic() - mm._BEDROCK_PROBE_FAILURE_TTL_SECONDS - 1 + token = set_hermes_home_override(isolated_home / "routed-profile") + try: + assert mm.get_model_context_length(model, provider="bedrock") == 96_000 + finally: + reset_hermes_home_override(token) + assert probe.call_count == 3 + assert all("unknown.future" not in key for key in mm._BEDROCK_PROBE_FAILURE_CACHE) + assert mm.get_model_context_length(model, provider="bedrock") == 500_000 + assert probe.call_count == 3 diff --git a/tests/agent/test_model_metadata.py b/tests/agent/test_model_metadata.py index a4dee9044f..9372cca14e 100644 --- a/tests/agent/test_model_metadata.py +++ b/tests/agent/test_model_metadata.py @@ -1351,9 +1351,10 @@ class TestBedrockContextCachePersistence: 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 - ) + for key in mm._BEDROCK_PROBE_FAILURE_CACHE: + mm._BEDROCK_PROBE_FAILURE_CACHE[key] = ( + 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 diff --git a/website/docs/developer-guide/context-compression-and-caching.md b/website/docs/developer-guide/context-compression-and-caching.md index 6150d5e522..7f237caa4e 100644 --- a/website/docs/developer-guide/context-compression-and-caching.md +++ b/website/docs/developer-guide/context-compression-and-caching.md @@ -7,6 +7,35 @@ Source files: `agent/context_engine.py` (ABC), `agent/context_compressor.py` (de `agent/prompt_caching.py`, `gateway/run_turn.py` (session hygiene), `agent/compression_facade.py` (search for `_compress_context`) +## Bedrock context window cache + +Bedrock context resolution in `agent/model_metadata.py` uses this precedence: + +- **Explicit overrides win.** Configured context lengths take priority over cache, + probes, and the static table. +- **Provider-confirmed limits persist.** A successful probe or a limit learned + from a provider error remains authoritative, even below the static table. + The compressor uses the same value after restart. +- **Legacy entries are revalidated.** Old scalar entries have no provenance and + may be either probe results or fallbacks. Their size does not establish which. +- **Failed probes use the current table without persisting it.** Failures have a + five-minute in-memory cooldown scoped to Hermes home, endpoint, model, and + region. Expiry or explicit cache invalidation permits another attempt. + +The cache remains at `context_length_cache.yaml` under the active Hermes home. +`context_lengths` retains scalar values for older readers. An additive +`bedrock_confirmed_v1` map binds each confirmed key to its exact value in the +same atomic write. Generic writes clear that key's provenance. Older writers +may drop the additive map, which causes revalidation after upgrading again. +Downgrading remains readable but restores the older runtime's resolution rules. + +The static fallback for `xai.grok-4.6` (including `global.` and `us.` inference +profiles) is 500,000 tokens, per the +[AWS model card](https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-6.html). +This is Bedrock-specific, not the direct xAI API window. Existing compression +rules still apply: without output reservation or an explicit token cap, the +small-window 75% threshold floor yields a 375,000-token trigger at this window. + ## Pluggable Context Engine Context management is built on the `ContextEngine` ABC (`agent/context_engine.py`). The built-in `ContextCompressor` is the default implementation, but plugins can replace it with alternative engines (e.g., Lossless Context Management).