From a51143fbbe6ddbc0c7f403d0579c4d75504c6793 Mon Sep 17 00:00:00 2001
From: teknium1 <127238744+teknium1@users.noreply.github.com>
Date: Fri, 18 Sep 2026 02:28:38 -0700
Subject: [PATCH] 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
---
 agent/bedrock_adapter.py            |  3 +++
 tests/agent/test_bedrock_adapter.py | 13 +++++++++++++
 2 files changed, 16 insertions(+)

diff --git a/agent/bedrock_adapter.py b/agent/bedrock_adapter.py
index 279e74dac9..c58abe0c86 100644
--- a/agent/bedrock_adapter.py
+++ b/agent/bedrock_adapter.py
@@ -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)
 
 
diff --git a/tests/agent/test_bedrock_adapter.py b/tests/agent/test_bedrock_adapter.py
index 51c817d48f..5d47079f4c 100644
--- a/tests/agent/test_bedrock_adapter.py
+++ b/tests/agent/test_bedrock_adapter.py
@@ -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