fix(bedrock): resolve application-inference-profile ARNs before the prompt-cache allowlist match

build_converse_kwargs gated cachePoint markers on the raw model id, so a profile ARN wrapping
Claude matched nothing in _CACHE_POINT_PATTERNS and silently lost Bedrock prompt caching.
_model_supports_prompt_cache now resolves the profile through the per-process-cached
_resolve_inference_profile_model_id first; the request still targets the profile ARN.

Part of #114476
This commit is contained in:
teknium1
2026-09-18 02:28:38 -07:00
committed by Teknium
parent 418a19cfc7
commit a51143fbbe
2 changed files with 16 additions and 0 deletions

View File

@@ -439,6 +439,9 @@ def _model_supports_tool_use(model_id: str) -> bool:
def _model_supports_prompt_cache(model_id: str) -> bool:
# An application-inference-profile ARN names no model: match on the wrapped model (cached lookup).
if _APPLICATION_PROFILE_ARN_RE.search(model_id):
model_id = _resolve_inference_profile_model_id(model_id)
return any(pattern in model_id.lower() for pattern in _CACHE_POINT_PATTERNS)

View File

@@ -1105,6 +1105,19 @@ class TestInferenceProfileContextLength:
assert len(warnings) == 1 and self.ARN in warnings[0].getMessage()
assert "GetInferenceProfile" in warnings[0].getMessage()
def test_profile_wrapping_claude_keeps_prompt_cache_markers(self):
# build_converse_kwargs gates cachePoint on the model id; the opaque profile ARN must be
# resolved to the wrapped Claude (cached lookup) or the profile silently loses prompt caching.
from agent.bedrock_adapter import build_converse_kwargs
client = MagicMock()
client.get_inference_profile.return_value = {"models": [
{"modelArn": "arn:aws:bedrock:us-west-2::foundation-model/anthropic.claude-sonnet-4-6"}]}
messages = [{"role": "system", "content": "Be helpful."}, {"role": "user", "content": "Hi"}]
with patch("agent.bedrock_adapter._get_bedrock_control_client", return_value=client):
kwargs = build_converse_kwargs(model=self.ARN, messages=messages)
assert kwargs["modelId"] == self.ARN # the request still targets the profile
assert kwargs["system"][-1] == {"cachePoint": {"type": "default"}}
class TestBedrockContextProbe:
"""Test the live context-window probe that reads the real window from