fix(bedrock): probe real context window instead of stale static table

Bedrock models resolved their context window from a hardcoded table
(BEDROCK_CONTEXT_LENGTHS) keyed by longest-substring match. AWS ships
new model versions faster than the table tracks, so a new model like
claude-opus-4-8 (1M-token window) silently matched the older
"anthropic.claude-opus-4" entry and got capped at 200K — wasting 80%
of the available context.

Bedrock exposes the real window nowhere in metadata: get-foundation-model
omits it, Converse usage metrics omit it, CountTokens is unsupported on
several models. The only authoritative source is the ValidationException
raised when a prompt exceeds the window:

    "prompt is too long: 1300032 tokens > 1000000 maximum"

Length validation runs before inference, so an oversized request is
rejected immediately and cheaply (no tokens generated, no input
processed). This adds probe_bedrock_context_length(): pad a request just
past a tier, parse the reported maximum, return it. get_bedrock_context_length()
now probes first and falls back to the static table only when the probe
can't run (missing creds, network error, unparseable error). The static
table stays as a safety net.

get_model_context_length() caches the probe result per model+region, so
the network cost is paid once, not every turn. probe=False / empty region
disables probing for offline/display paths — backward compatible with the
single-arg callers.

Verified E2E against live Bedrock (eu-central-1): claude-opus-4-8 resolves
to 1000000. Unit tests cover error parsing, unparseable errors, missing
client, probe-beats-table, and table fallback.
This commit is contained in:
kubolko 2026-06-18 16:43:12 +02:00 committed by Teknium
parent 222772ad61
commit 6be4944bc0
3 changed files with 214 additions and 4 deletions

View file

@ -1393,9 +1393,22 @@ BEDROCK_CONTEXT_LENGTHS: Dict[str, int] = {
# Default for unknown Bedrock models
BEDROCK_DEFAULT_CONTEXT_LENGTH = 128_000
# Probe tiers (in tokens). We send a request padded just past each tier and
# read the real window from Bedrock's length-validation error. Two reasons
# this is tiered rather than one giant request:
# 1. A wildly oversized payload (e.g. 5M tokens) makes Bedrock return an
# opaque InternalServerException after retries instead of a clean
# ValidationException — so we must stay within a sane overage.
# 2. Stepping up lets us discover larger windows (2M+) without over-padding
# smaller ones.
# Each tier value is the *padding target*; the error reports the true maximum,
# which is what we actually return.
_BEDROCK_PROBE_TIERS = (1_300_000, 2_200_000)
_WORDS_PER_TOKEN = 0.9 # conservative: ensures the padded prompt clears the tier
def get_bedrock_context_length(model_id: str) -> int:
"""Look up the context window size for a Bedrock model.
def _static_bedrock_context_length(model_id: str) -> int:
"""Longest-substring-match lookup against the static fallback table.
Uses substring matching so versioned IDs like
``anthropic.claude-sonnet-4-6-20250514-v1:0`` resolve correctly.
@ -1408,3 +1421,103 @@ def get_bedrock_context_length(model_id: str) -> int:
best_key = key
best_val = val
return best_val
def probe_bedrock_context_length(model_id: str, region: str) -> Optional[int]:
"""Discover a Bedrock model's real context window by provoking a length error.
Bedrock does not expose the context window via any metadata API
(``get-foundation-model`` omits it, ``Converse`` metrics omit it,
``CountTokens`` is unsupported on several models). The only authoritative
source is the ``ValidationException`` raised when a prompt exceeds the
window:
"The model returned the following errors: prompt is too long:
1300032 tokens > 1000000 maximum"
Length validation happens *before* inference, so an oversized request is
rejected immediately and cheaply no tokens are generated and no input is
actually processed. We pad a request just past each tier in
``_BEDROCK_PROBE_TIERS`` and parse the reported ``maximum``. Tiers exist
because (a) a *wildly* oversized payload makes Bedrock fail with an opaque
InternalServerException instead of a clean length error, and (b) stepping
up discovers larger windows without over-padding smaller ones.
Returns the detected window, or ``None`` if the probe could not run
(missing credentials, network error, or no parseable limit) so the caller
can fall back to the static table.
"""
try:
from agent.model_metadata import parse_context_limit_from_error
except ImportError: # pragma: no cover — same package
return None
try:
client = _get_bedrock_runtime_client(region)
except Exception as exc: # boto3 missing / credential resolution failure
logger.debug("Bedrock context probe skipped for %s: %s", model_id, exc)
return None
last_error = ""
for tier_tokens in _BEDROCK_PROBE_TIERS:
pad_words = int(tier_tokens / _WORDS_PER_TOKEN)
oversized = "data " * pad_words
try:
client.converse(
modelId=model_id,
messages=[{"role": "user", "content": [{"text": oversized}]}],
inferenceConfig={"maxTokens": 8},
)
# Accepted a prompt this large → the window is at least this tier.
# Returning the tier as a lower bound is safe and avoids inventing
# a number we can't confirm.
logger.debug(
"Bedrock context probe for %s accepted ~%s-token prompt; "
"window is at least that", model_id, f"{tier_tokens:,}",
)
return tier_tokens
except Exception as exc:
msg = str(exc)
last_error = msg
limit = parse_context_limit_from_error(msg)
if limit and limit >= 1024:
logger.info(
"Probed Bedrock context window for %s: %s tokens",
model_id, f"{limit:,}",
)
return limit
# No parseable limit at this tier (opaque server error, auth,
# throttle). Try the next, smaller-overage strategy is N/A here —
# tiers ascend — so just continue; if all fail we return None.
continue
logger.debug(
"Bedrock context probe for %s returned no parseable limit: %s",
model_id, last_error[:200],
)
return None
def get_bedrock_context_length(model_id: str, region: str = "", probe: bool = True) -> int:
"""Resolve the context window for a Bedrock model.
Resolution order:
1. Live probe against Bedrock (authoritative; cached by the caller).
2. Static fallback table (longest-substring match).
3. Conservative default.
The static table is intentionally a *fallback*, not the primary source:
AWS ships new model versions (opus-4-7, opus-4-8, ...) faster than the
table can track, and a stale entry silently caps the window (e.g. a
1M-token Opus pinned to 200K via an ``opus-4`` substring match). The
probe asks Bedrock directly so every model current or future gets its
real window with no table maintenance.
``probe=False`` (or an empty ``region``) skips the network call and uses
the static table only used by pure-offline/display code paths.
"""
if probe and region:
probed = probe_bedrock_context_length(model_id, region)
if probed:
return probed
return _static_bedrock_context_length(model_id)

View file

@ -2287,10 +2287,42 @@ def get_model_context_length(
# back to the 128K default before reaching the original step 4b branch.
if is_bedrock_context:
try:
from agent.bedrock_adapter import get_bedrock_context_length
return get_bedrock_context_length(model)
from agent.bedrock_adapter import (
get_bedrock_context_length,
resolve_bedrock_region,
)
except ImportError:
pass # boto3 not installed — fall through to generic resolution
else:
# Bedrock does not expose the context window via any metadata API,
# so get_bedrock_context_length() probes the live endpoint (one
# fast, pre-inference length rejection) to read the real window.
# Cache the probe result per model so we pay that cost once, not
# every turn — keyed by base_url when present, else a synthetic
# bedrock:// key so display/offline paths share the entry.
cache_key_url = base_url or "bedrock://"
cached = get_cached_context_length(model, cache_key_url)
if cached is not None:
return cached
# Resolve region from the base_url host first, then the standard
# AWS region chain. An empty region disables probing (table only).
region = ""
if base_url:
_m = re.search(r"bedrock-runtime\.([a-z0-9-]+)\.", base_url)
if _m:
region = _m.group(1)
if not region:
try:
region = resolve_bedrock_region()
except Exception:
region = ""
ctx = get_bedrock_context_length(model, region=region, probe=bool(region))
if ctx and region:
# Only persist probe-derived values (region present); a pure
# table fallback shouldn't poison the cache against a later
# successful probe.
save_context_length(model, cache_key_url, ctx)
return ctx
if provider == "novita" or (base_url and base_url_host_matches(base_url, "api.novita.ai")):
ctx = _resolve_endpoint_context_length(model, base_url or "https://api.novita.ai/openai/v1", api_key=api_key)

View file

@ -1238,6 +1238,71 @@ class TestBedrockContextLength:
# "anthropic.claude-3-5-sonnet" should match before "anthropic.claude-3"
assert get_bedrock_context_length("anthropic.claude-3-5-sonnet-20240620-v1:0") == 200_000
def test_no_region_skips_probe_uses_table(self):
# Default call (no region) must NOT hit the network — returns the
# static table value. Guards backward compatibility for callers that
# still invoke get_bedrock_context_length(model_id) with one arg.
from agent.bedrock_adapter import get_bedrock_context_length
with patch("agent.bedrock_adapter.probe_bedrock_context_length") as mock_probe:
assert get_bedrock_context_length("anthropic.claude-opus-4-6") == 200_000
mock_probe.assert_not_called()
class TestBedrockContextProbe:
"""Test the live context-window probe that reads the real window from
Bedrock's 'prompt is too long' validation error."""
def _client_raising(self, message):
client = MagicMock()
client.converse.side_effect = Exception(message)
return client
def test_probe_parses_real_window_from_error(self):
from agent.bedrock_adapter import probe_bedrock_context_length
err = (
"An error occurred (ValidationException) when calling the Converse "
"operation: The model returned the following errors: prompt is too "
"long: 5000032 tokens > 1000000 maximum"
)
with patch("agent.bedrock_adapter._get_bedrock_runtime_client",
return_value=self._client_raising(err)):
assert probe_bedrock_context_length(
"eu.anthropic.claude-opus-4-8", "eu-central-1") == 1_000_000
def test_probe_returns_none_on_unparseable_error(self):
from agent.bedrock_adapter import probe_bedrock_context_length
err = "An error occurred (AccessDeniedException): not authorized"
with patch("agent.bedrock_adapter._get_bedrock_runtime_client",
return_value=self._client_raising(err)):
assert probe_bedrock_context_length(
"eu.anthropic.claude-opus-4-8", "eu-central-1") is None
def test_probe_returns_none_when_client_unavailable(self):
from agent.bedrock_adapter import probe_bedrock_context_length
with patch("agent.bedrock_adapter._get_bedrock_runtime_client",
side_effect=RuntimeError("boto3 missing")):
assert probe_bedrock_context_length("any.model", "eu-central-1") is None
def test_probe_result_beats_static_table(self):
# A successful probe (1M) must override the stale table value (200K
# via the 'anthropic.claude-opus-4' substring match).
from agent.bedrock_adapter import get_bedrock_context_length
err = "prompt is too long: 5000032 tokens > 1000000 maximum"
with patch("agent.bedrock_adapter._get_bedrock_runtime_client",
return_value=self._client_raising(err)):
assert get_bedrock_context_length(
"eu.anthropic.claude-opus-4-8",
region="eu-central-1") == 1_000_000
def test_probe_failure_falls_back_to_table(self):
from agent.bedrock_adapter import get_bedrock_context_length
err = "AccessDeniedException: nope"
with patch("agent.bedrock_adapter._get_bedrock_runtime_client",
return_value=self._client_raising(err)):
# opus-4-6 is in the table at 200K; probe fails → table wins.
assert get_bedrock_context_length(
"anthropic.claude-opus-4-6", region="eu-central-1") == 200_000
# ---------------------------------------------------------------------------
# Tool-calling capability detection