From 33c7f6cd3ae14857effcbd110e9cd0668cf1ac1b Mon Sep 17 00:00:00 2001 From: Kenneth Wong Date: Fri, 10 Jul 2026 05:56:23 +0800 Subject: [PATCH] fix(bedrock): resolve global.* inference profiles + pin per-user app-profile ARNs (#1795) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What & why Bedrock model resolution failed on accounts whose inference profiles use the newer `global.` cross-region prefix and undated version suffixes. On such an account, `list_inference_profiles` returns current-gen models as `global.anthropic.claude-opus-4-8`, `global.anthropic.claude-sonnet-5`, `global.anthropic.claude-opus-4-6-v1`, `global.anthropic.claude-fable-5`, etc. `_normalize_bedrock_profile_id` only stripped `us.`/`eu.`/`apac.`/`au.` and only matched a dated `-vN:M` suffix, so every `global.`-prefixed profile was silently dropped from the discovered model map. Requests then fell through to the fabricated fallback id and Bedrock rejected them: ``` litellm.BadRequestError: BedrockException - {"message":"The provided model identifier is invalid."} ``` On the affected account the discovered-profile count went from 5 → 14 after the fix. ## Changes 1. **`_normalize_bedrock_profile_id`** — strip the `global.` prefix in addition to the region prefixes, and match undated version suffixes (`-v1`, or none at all) alongside the legacy dated `-vN:M`. 2. **`HEADROOM_BEDROCK_MODEL_MAP` operator override** (read from the process environment). AWS discovery keys the model map by normalized model name, so it cannot disambiguate application inference profiles that share one underlying model — e.g. a team where `claude-sonnet-5-alice` and `claude-sonnet-5-bob` both resolve to `claude-sonnet-5`. When you need requests billed to a *specific* application profile (per-user cost attribution), pin it explicitly: ``` HEADROOM_BEDROCK_MODEL_MAP="claude-sonnet-5=arn:aws:bedrock:REGION:ACCT:application-inference-profile/abc123,claude-opus-4-8=arn:...:application-inference-profile/def456" ``` The plain model name (kept plain so a client's tool-search deferral stays on) resolves to the pinned ARN, routed via the converse endpoint. The override wins over discovery; when unset, discovery-only behaviour is unchanged. ## Tests `tests/test_bedrock_region.py` (43 passing): `global.`-prefixed normalization across dated / bare-`-v1` / no-suffix shapes; override-map parsing (empty, single, multi, whitespace, malformed-skip); and `map_model_id` override routing (pinned name → app-profile ARN via converse, wins over discovery; unpinned name falls through). No behavioural change for accounts already on system-defined region-prefixed profiles. --------- Co-authored-by: JerrettDavis --- headroom/backends/litellm.py | 74 +++++++++++++++++++++++++++-- tests/test_bedrock_region.py | 92 ++++++++++++++++++++++++++++++++++++ 2 files changed, 162 insertions(+), 4 deletions(-) diff --git a/headroom/backends/litellm.py b/headroom/backends/litellm.py index 6860fed58..1d597b544 100644 --- a/headroom/backends/litellm.py +++ b/headroom/backends/litellm.py @@ -226,6 +226,39 @@ def _fetch_bedrock_inference_profiles( return model_map +def _parse_bedrock_model_overrides(raw: str | None) -> dict[str, str]: + """Parse the ``HEADROOM_BEDROCK_MODEL_MAP`` operator override. + + AWS discovery keys the model map by the *normalized model name*, so it + cannot disambiguate application inference profiles that share one + underlying model — e.g. a team where ``claude-sonnet-5-kenneth`` and + ``claude-sonnet-5-jeremy`` both resolve to ``claude-sonnet-5``. When you + need requests billed to a *specific* application profile (per-user cost + attribution), pin the mapping explicitly here. The plain name Claude Code + sends (kept plain so tool-search deferral stays on) resolves to your ARN. + + Format: comma-separated ``name=target`` pairs, where ``target`` is an + application-inference-profile ARN (routed via the converse endpoint) or + any LiteLLM model string. Whitespace around pairs is ignored; blank + entries are skipped. + + HEADROOM_BEDROCK_MODEL_MAP="claude-sonnet-5=arn:aws:bedrock:...:application-inference-profile/x57j1esjrt66,claude-opus-4-8=arn:aws:bedrock:...:application-inference-profile/3dy9ytxuq2ci" + """ + overrides: dict[str, str] = {} + if not raw: + return overrides + for pair in raw.split(","): + pair = pair.strip() + if not pair or "=" not in pair: + continue + name, _, target = pair.partition("=") + name = name.strip() + target = target.strip() + if name and target: + overrides[name] = target + return overrides + + def _normalize_bedrock_profile_id(profile_id: str) -> str | None: """Extract standard Anthropic model name from Bedrock profile ID. @@ -248,8 +281,10 @@ def _normalize_bedrock_profile_id(profile_id: str) -> str | None: if profile_id.startswith("bedrock/"): profile_id = profile_id[8:] - # Strip region prefix (us., eu., apac., au.) - for prefix in ["us.", "eu.", "apac.", "au."]: + # Strip region prefix (us., eu., apac., au.) or the newer "global." + # cross-region prefix used by current-gen profiles (e.g. + # "global.anthropic.claude-sonnet-4-6"). + for prefix in ["us.", "eu.", "apac.", "au.", "global."]: if profile_id.startswith(prefix): profile_id = profile_id[len(prefix) :] break @@ -262,8 +297,12 @@ def _normalize_bedrock_profile_id(profile_id: str) -> str | None: if not profile_id.startswith("claude"): return None - # Strip version suffix (-v1:0, -v2:0, etc.) - normalized = re.sub(r"-v\d+:\d+$", "", profile_id) + # Strip version suffix. Legacy dated profiles use "-v1:0" / "-v2:0"; + # newer undated profiles use a bare "-v1" (no colon/revision) or carry + # no version suffix at all (e.g. "claude-opus-4-8"). Match all three + # shapes so undated current-gen profiles normalize instead of + # silently falling out of the resolvable model map. + normalized = re.sub(r"-v\d+(?::\d+)?$", "", profile_id) return normalized if normalized else None @@ -491,6 +530,20 @@ class LiteLLMBackend(Backend): else: self._model_map = self._config.model_map + # Operator override map (all providers; only meaningful for Bedrock + # today). Lets you pin a plain model name to a specific target the + # AWS discovery can't disambiguate — e.g. a per-user application + # inference profile ARN for cost attribution. See + # `_parse_bedrock_model_overrides`. + self._model_overrides = _parse_bedrock_model_overrides( + os.environ.get("HEADROOM_BEDROCK_MODEL_MAP") + ) + if self._model_overrides: + logger.info( + f"Loaded {len(self._model_overrides)} Bedrock model override(s) " + f"from HEADROOM_BEDROCK_MODEL_MAP: {sorted(self._model_overrides)}" + ) + logger.info(f"LiteLLM backend initialized (provider={provider}, region={region})") @property @@ -507,6 +560,19 @@ class LiteLLMBackend(Backend): - "bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0" (LiteLLM format) - "arn:aws:bedrock:...:application-inference-profile/..." (application inference profile) """ + # Operator override wins over everything — an explicit pin the AWS + # discovery cannot express (e.g. a per-user application inference + # profile). Keyed by the plain name Claude Code sends. + override = self._model_overrides.get(anthropic_model) + if override: + if override.startswith("arn:aws:"): + # Application inference profile ARNs must use the converse + # route — the invoke route rejects ARNs with HTTP 400. + return f"bedrock/converse/{override}" + if override.startswith(f"{self.provider}/"): + return override + return f"{self.provider}/{override}" + # Check direct mapping first if anthropic_model in self._model_map: return self._model_map[anthropic_model] diff --git a/tests/test_bedrock_region.py b/tests/test_bedrock_region.py index 46df34b2d..f96cff1a6 100644 --- a/tests/test_bedrock_region.py +++ b/tests/test_bedrock_region.py @@ -347,6 +347,33 @@ class TestBedrockModelMapping: result = backend.map_model_id("claude-sonnet-4-5-20250929") assert result == "bedrock/au.anthropic.claude-sonnet-4-5-20250929-v1:0" + def test_override_pins_plain_name_to_app_profile_arn(self, monkeypatch): + """HEADROOM_BEDROCK_MODEL_MAP pins a plain name to a specific app + profile ARN and routes it via the converse endpoint, winning over + the discovered map.""" + arn = "arn:aws:bedrock:ap-southeast-1:1:application-inference-profile/x57j1esjrt66" + monkeypatch.setenv("HEADROOM_BEDROCK_MODEL_MAP", f"claude-sonnet-5={arn}") + with patch( + "headroom.backends.litellm._fetch_bedrock_inference_profiles", + # Discovery also has a system-defined sonnet-5; the override must win. + return_value={"claude-sonnet-5": "bedrock/global.anthropic.claude-sonnet-5"}, + ): + backend = LiteLLMBackend(provider="bedrock", region="ap-southeast-1") + assert backend.map_model_id("claude-sonnet-5") == f"bedrock/converse/{arn}" + + def test_override_absent_falls_through_to_discovery(self, monkeypatch): + """A model not pinned in the override map resolves via discovery.""" + arn = "arn:aws:bedrock:ap-southeast-1:1:application-inference-profile/x57j1esjrt66" + monkeypatch.setenv("HEADROOM_BEDROCK_MODEL_MAP", f"claude-sonnet-5={arn}") + with patch( + "headroom.backends.litellm._fetch_bedrock_inference_profiles", + return_value={"claude-opus-4-8": "bedrock/global.anthropic.claude-opus-4-8"}, + ): + backend = LiteLLMBackend(provider="bedrock", region="ap-southeast-1") + assert backend.map_model_id("claude-opus-4-8") == ( + "bedrock/global.anthropic.claude-opus-4-8" + ) + # ============================================================================= # Normalize Bedrock Profile ID (edge cases) @@ -390,6 +417,71 @@ class TestNormalizeBedrockProfileId: "claude-sonnet-4-20250514" ) + def test_global_prefix_no_version_suffix(self): + # Current-gen cross-region profile: "global." prefix, no version + # suffix at all. Must strip the prefix and keep the bare name. + assert _normalize_bedrock_profile_id("global.anthropic.claude-opus-4-8") == ( + "claude-opus-4-8" + ) + + def test_global_prefix_bare_v_suffix(self): + # "global." prefix with an undated "-v1" (no ":revision") suffix. + assert _normalize_bedrock_profile_id("global.anthropic.claude-opus-4-6-v1") == ( + "claude-opus-4-6" + ) + + def test_global_prefix_dated_full_suffix(self): + # "global." prefix with the legacy dated "-vN:M" suffix. + assert ( + _normalize_bedrock_profile_id("global.anthropic.claude-haiku-4-5-20251001-v1:0") + == "claude-haiku-4-5-20251001" + ) + + def test_global_prefix_next_gen_names(self): + assert ( + _normalize_bedrock_profile_id("global.anthropic.claude-sonnet-5") == "claude-sonnet-5" + ) + assert _normalize_bedrock_profile_id("global.anthropic.claude-fable-5") == "claude-fable-5" + + +# ============================================================================= +# HEADROOM_BEDROCK_MODEL_MAP operator override parsing +# ============================================================================= + + +class TestParseBedrockModelOverrides: + """Test the HEADROOM_BEDROCK_MODEL_MAP override parser.""" + + def test_none_and_empty_yield_empty(self): + from headroom.backends.litellm import _parse_bedrock_model_overrides + + assert _parse_bedrock_model_overrides(None) == {} + assert _parse_bedrock_model_overrides("") == {} + assert _parse_bedrock_model_overrides(" ") == {} + + def test_single_pair(self): + from headroom.backends.litellm import _parse_bedrock_model_overrides + + arn = "arn:aws:bedrock:ap-southeast-1:1:application-inference-profile/x57j1esjrt66" + assert _parse_bedrock_model_overrides(f"claude-sonnet-5={arn}") == {"claude-sonnet-5": arn} + + def test_multiple_pairs_and_whitespace(self): + from headroom.backends.litellm import _parse_bedrock_model_overrides + + raw = " claude-sonnet-5=arn:a , claude-opus-4-8=arn:b " + assert _parse_bedrock_model_overrides(raw) == { + "claude-sonnet-5": "arn:a", + "claude-opus-4-8": "arn:b", + } + + def test_skips_malformed_entries(self): + from headroom.backends.litellm import _parse_bedrock_model_overrides + + # Missing "=" and blank segments are skipped, valid pairs survive. + assert _parse_bedrock_model_overrides("garbage,,claude-sonnet-5=arn:a,=noname") == { + "claude-sonnet-5": "arn:a", + } + # ============================================================================= # Named profile forwarded to acompletion kwargs