scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

fix(integrations): cap Anthropic cache_control injection at 4 blocks (#30480)

Browse source 
* fix(integrations): cap Anthropic cache_control injection at 4 blocks

Respect Anthropic's 4 cache_control breakpoint limit by counting client-supplied blocks, skipping messages that already carry cache_control, and stopping further auto-injection once the limit is reached.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(integrations): reserve cache slot for tool_config and short-circuit cap

Address review feedback on the cache_control cap: break out of the injection loop before resolving target indices once the limit is reached, and reserve one of the four breakpoint slots when a tool_config injection point is present so the cachePoint appended by the Bedrock transform does not push the total past Anthropic's limit.

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
(cherry picked from commit fc9d789d24)
This commit is contained in:
Shivam Rawat 2026-06-15 20:50:38 -07:00 • committed by Yuneng Jiang
parent f05280f71b
commit 71cc179eca
No known key found for this signature in database
2 changed files with 476 additions and 33 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

155
litellm/integrations/anthropic_cache_control_hook.py
View file

@ -27,6 +27,11 @@ else:
LiteLLMLoggingObj = Any
# Anthropic (and Bedrock Claude) reject requests with more than 4 cache_control
# breakpoints: "A maximum of 4 blocks with cache_control may be provided."
MAX_CACHE_CONTROL_BLOCKS = 4
class AnthropicCacheControlHook(CustomPromptManagement):
def get_chat_completion_prompt(
self,
@ -61,16 +66,30 @@ class AnthropicCacheControlHook(CustomPromptManagement):
processed_messages = copy.deepcopy(messages)
# Separate message-level and non-message-level injection points
remaining_points = []
message_points: List[CacheControlMessageInjectionPoint] = []
remaining_points: List[CacheControlInjectionPoint] = []
for point in injection_points:
if point.get("location") == "message":
point = cast(CacheControlMessageInjectionPoint, point)
processed_messages = self._process_message_injection(
point=point, messages=processed_messages
)
message_points.append(cast(CacheControlMessageInjectionPoint, point))
else:
remaining_points.append(point)
# Non-message points (currently Bedrock tool_config) are handled in the
# provider transform, where each tool_config point appends at most one
# cachePoint to the tools. That block also counts toward Anthropic's
# limit, so reserve a slot for it here to leave room.
reserved_blocks = (
1
if any(p.get("location") == "tool_config" for p in remaining_points)
else 0
)
processed_messages = self._apply_message_injections(
points=message_points,
messages=processed_messages,
max_blocks=MAX_CACHE_CONTROL_BLOCKS - reserved_blocks,
)
# Pass through non-message injection points for provider-specific handling
if remaining_points:
non_default_params["cache_control_injection_points"] = remaining_points
@ -78,14 +97,71 @@ class AnthropicCacheControlHook(CustomPromptManagement):
return model, processed_messages, non_default_params
@staticmethod
def _process_message_injection(
point: CacheControlMessageInjectionPoint, messages: List[AllMessageValues]
def _apply_message_injections(
points: List[CacheControlMessageInjectionPoint],
messages: List[AllMessageValues],
max_blocks: int,
) -> List[AllMessageValues]:
"""Process message-level cache control injection."""
control: ChatCompletionCachedContent = point.get(
"control", None
) or ChatCompletionCachedContent(type="ephemeral")
"""Apply message-level cache control injection points in order.
Anthropic allows at most ``MAX_CACHE_CONTROL_BLOCKS`` cache_control
breakpoints per request. Client-supplied breakpoints count toward that
limit, so we never inject onto a message that already carries
cache_control (preserving the client's TTL) and we stop injecting once
``max_blocks`` is reached. Injection points are honored in config order,
so earlier points win when slots are scarce.
"""
used_blocks = sum(
AnthropicCacheControlHook._count_cache_control_blocks(msg)
for msg in messages
)
limit_reached = False
for point in points:
if used_blocks >= max_blocks:
limit_reached = True
break
control: ChatCompletionCachedContent = point.get(
"control", None
) or ChatCompletionCachedContent(type="ephemeral")
for target_index in AnthropicCacheControlHook._resolve_target_indices(
point=point, messages=messages
):
if used_blocks >= max_blocks:
limit_reached = True
break
if AnthropicCacheControlHook._message_has_cache_control(
messages[target_index]
):
# Client already marked this message; don't overwrite it.
continue
messages[target_index] = (
AnthropicCacheControlHook._safe_insert_cache_control_in_message(
messages[target_index], control
)
)
used_blocks += 1
if limit_reached:
break
if limit_reached:
verbose_logger.warning(
f"AnthropicCacheControlHook: Reached the Anthropic limit of "
f"{MAX_CACHE_CONTROL_BLOCKS} cache_control blocks. Skipping further injection."
)
return messages
@staticmethod
def _resolve_target_indices(
point: CacheControlMessageInjectionPoint, messages: List[AllMessageValues]
) -> List[int]:
"""Resolve which message indices an injection point targets."""
_targetted_index: Optional[Union[int, str]] = point.get("index", None)
targetted_index: Optional[int] = None
if isinstance(_targetted_index, str):
@ -96,36 +172,49 @@ class AnthropicCacheControlHook(CustomPromptManagement):
else:
targetted_index = _targetted_index
targetted_role = point.get("role", None)
# Case 1: Target by specific index
if targetted_index is not None:
original_index = targetted_index
# Handle negative indices (convert to positive)
if targetted_index < 0:
targetted_index += len(messages)
if 0 <= targetted_index < len(messages):
messages[targetted_index] = (
AnthropicCacheControlHook._safe_insert_cache_control_in_message(
messages[targetted_index], control
)
)
else:
verbose_logger.warning(
f"AnthropicCacheControlHook: Provided index {original_index} is out of bounds for message list of length {len(messages)}. "
f"Targeted index was {targetted_index}. Skipping cache control injection for this point."
)
return [targetted_index]
verbose_logger.warning(
f"AnthropicCacheControlHook: Provided index {original_index} is out of bounds for message list of length {len(messages)}. "
f"Targeted index was {targetted_index}. Skipping cache control injection for this point."
)
return []
# Case 2: Target by role
elif targetted_role is not None:
for msg in messages:
if msg.get("role") == targetted_role:
msg = (
AnthropicCacheControlHook._safe_insert_cache_control_in_message(
message=msg, control=control
)
)
return messages
targetted_role = point.get("role", None)
if targetted_role is not None:
return [
idx
for idx, msg in enumerate(messages)
if msg.get("role") == targetted_role
]
return []
@staticmethod
def _count_cache_control_blocks(message: AllMessageValues) -> int:
"""Count cache_control breakpoints on a message (message + content level)."""
count = 0
if message.get("cache_control") is not None:
count += 1
content = message.get("content")
if isinstance(content, list):
for block in content:
if isinstance(block, dict) and block.get("cache_control") is not None:
count += 1
return count
@staticmethod
def _message_has_cache_control(message: AllMessageValues) -> bool:
"""Return True if the message already carries any cache_control."""
return AnthropicCacheControlHook._count_cache_control_blocks(message) > 0
@staticmethod
def _safe_insert_cache_control_in_message(

354
tests/test_litellm/integrations/test_anthropic_cache_control_hook.py
View file

@ -1087,3 +1087,357 @@ async def test_anthropic_cache_control_hook_string_negative_index():
f"Expected cachePoint in last message content, got: {last_message_content}. "
"String index '-1' was not parsed correctly (str.isdigit() returns False for negative strings)."
)
def _count_cache_control(messages: List[AllMessageValues]) -> int:
"""Count cache_control breakpoints across messages (message + content level)."""
count = 0
for message in messages:
if message.get("cache_control") is not None:
count += 1
content = message.get("content")
if isinstance(content, list):
for block in content:
if isinstance(block, dict) and block.get("cache_control") is not None:
count += 1
return count
def _build_injection_points():
return [
{
"location": "message",
"role": "system",
"control": {"type": "ephemeral", "ttl": "1h"},
},
{
"location": "message",
"index": -1,
"control": {"type": "ephemeral", "ttl": "5m"},
},
]
def test_cache_control_hook_caps_at_four_blocks_with_client_cache_control():
"""Regression for LIT-3667 / Anthropic 'A maximum of 4 blocks ... Found 5'.
A Hermes-style request already carries 4 client cache_control breakpoints on
its system messages. With both auto-inject points configured the hook must
NOT add a 5th breakpoint, and must NOT overwrite the client's existing
breakpoints (TTL must be preserved).
"""
hook = AnthropicCacheControlHook()
messages: List[AllMessageValues] = [
{
"role": "system",
"content": [
{
"type": "text",
"text": f"System block {i}",
"cache_control": {"type": "ephemeral", "ttl": "1h"},
}
],
}
for i in range(4)
]
messages.append({"role": "user", "content": "hello"})
_, processed, _ = hook.get_chat_completion_prompt(
model="bedrock/us.anthropic.claude-opus-4-6-v1:0",
messages=messages,
non_default_params={
"cache_control_injection_points": _build_injection_points()
},
prompt_id=None,
prompt_variables=None,
dynamic_callback_params={},
)
assert (
_count_cache_control(processed) == 4
), "Hook must cap cache_control at Anthropic's limit of 4 blocks"
# Client TTL on system blocks must be preserved (not overwritten by config).
for i in range(4):
assert processed[i]["content"][-1]["cache_control"] == {
"type": "ephemeral",
"ttl": "1h",
}
# The last (user) message must not receive a 5th breakpoint.
user_message = processed[-1]
assert user_message.get("cache_control") is None
user_content = user_message.get("content")
if isinstance(user_content, list):
assert all(
block.get("cache_control") is None
for block in user_content
if isinstance(block, dict)
)
def test_cache_control_hook_caps_at_four_blocks_without_client_cache_control():
"""Four plain system messages + role:system + index:-1 must stay at 4 blocks.
role:system fills all four slots, so the index:-1 point is skipped.
"""
hook = AnthropicCacheControlHook()
messages: List[AllMessageValues] = [
{"role": "system", "content": f"System {i}"} for i in range(4)
]
messages.append({"role": "user", "content": "hello"})
_, processed, _ = hook.get_chat_completion_prompt(
model="bedrock/us.anthropic.claude-opus-4-6-v1:0",
messages=messages,
non_default_params={
"cache_control_injection_points": _build_injection_points()
},
prompt_id=None,
prompt_variables=None,
dynamic_callback_params={},
)
assert _count_cache_control(processed) == 4
# All four system messages cached; user message skipped (limit reached).
assert all(processed[i].get("cache_control") is not None for i in range(4))
assert processed[-1].get("cache_control") is None
def test_cache_control_hook_does_not_overwrite_existing_cache_control():
"""If a targeted message already has client cache_control, do not inject."""
hook = AnthropicCacheControlHook()
messages: List[AllMessageValues] = [
{
"role": "system",
"content": [
{
"type": "text",
"text": "Cached by client",
"cache_control": {"type": "ephemeral", "ttl": "1h"},
}
],
},
{"role": "user", "content": "hello"},
]
_, processed, _ = hook.get_chat_completion_prompt(
model="bedrock/us.anthropic.claude-opus-4-6-v1:0",
messages=messages,
# Target the already-cached system message with a different TTL.
non_default_params={
"cache_control_injection_points": [
{
"location": "message",
"index": 0,
"control": {"type": "ephemeral", "ttl": "5m"},
}
]
},
prompt_id=None,
prompt_variables=None,
dynamic_callback_params={},
)
# Client's 1h TTL must be preserved, not replaced by the config's 5m.
assert processed[0]["content"][-1]["cache_control"] == {
"type": "ephemeral",
"ttl": "1h",
}
assert _count_cache_control(processed) == 1
@pytest.mark.asyncio
async def test_cache_control_hook_bedrock_payload_caps_cachepoints_at_four():
"""End-to-end: outgoing Bedrock payload must not exceed 4 cachePoint blocks.
Reproduces the customer report where 4 client cache_control system blocks
plus auto-inject produced 5 cachePoint blocks and Bedrock returned 400.
"""
with patch.dict(
os.environ,
{
"AWS_ACCESS_KEY_ID": "fake_access_key_id",
"AWS_SECRET_ACCESS_KEY": "fake_secret_access_key",
"AWS_REGION_NAME": "us-east-1",
},
):
litellm.callbacks = [AnthropicCacheControlHook()]
mock_response = MagicMock()
mock_response.json.return_value = {
"output": {"message": {"role": "assistant", "content": "ok"}},
"stopReason": "end_turn",
"usage": {"inputTokens": 100, "outputTokens": 4, "totalTokens": 104},
}
mock_response.status_code = 200
client = AsyncHTTPHandler()
with patch.object(client, "post", return_value=mock_response) as mock_post:
messages = [
{
"role": "system",
"content": [
{
"type": "text",
"text": f"System block {i}",
"cache_control": {"type": "ephemeral", "ttl": "1h"},
}
],
}
for i in range(4)
]
messages.append({"role": "user", "content": "hello"})
await litellm.acompletion(
model="bedrock/us.anthropic.claude-opus-4-6-v1:0",
messages=messages,
max_tokens=32,
cache_control_injection_points=_build_injection_points(),
client=client,
)
request_body = json.loads(mock_post.call_args.kwargs["data"])
cache_points = sum(
1
for block in request_body.get("system", [])
if isinstance(block, dict) and "cachePoint" in block
)
for msg in request_body.get("messages", []):
content = msg.get("content", [])
if isinstance(content, list):
cache_points += sum(
1
for block in content
if isinstance(block, dict) and "cachePoint" in block
)
assert cache_points <= 4, (
f"Bedrock payload exceeded Anthropic's 4 cache_control block limit: "
f"found {cache_points} cachePoint blocks"
)
def test_cache_control_hook_reserves_slot_for_tool_config_point():
"""A tool_config injection point consumes one of the 4 slots downstream.
With role:system targeting 4 system messages plus a tool_config point, the
hook must inject at most 3 message-level blocks so the tool_config cachePoint
appended by the Bedrock transform keeps the total at 4, not 5.
"""
hook = AnthropicCacheControlHook()
messages: List[AllMessageValues] = [
{"role": "system", "content": f"System {i}"} for i in range(4)
]
messages.append({"role": "user", "content": "hello"})
_, processed, non_default_params = hook.get_chat_completion_prompt(
model="bedrock/us.anthropic.claude-opus-4-6-v1:0",
messages=messages,
non_default_params={
"cache_control_injection_points": [
{
"location": "message",
"role": "system",
"control": {"type": "ephemeral", "ttl": "1h"},
},
{"location": "tool_config"},
]
},
prompt_id=None,
prompt_variables=None,
dynamic_callback_params={},
)
assert _count_cache_control(processed) == 3
# The tool_config point is passed through for the provider transform.
assert non_default_params["cache_control_injection_points"] == [
{"location": "tool_config"}
]
@pytest.mark.asyncio
async def test_cache_control_hook_bedrock_payload_caps_with_tool_config_point():
"""End-to-end: message + tool_config injection must not exceed 4 cachePoints."""
with patch.dict(
os.environ,
{
"AWS_ACCESS_KEY_ID": "fake_access_key_id",
"AWS_SECRET_ACCESS_KEY": "fake_secret_access_key",
"AWS_REGION_NAME": "us-east-1",
},
):
litellm.callbacks = [AnthropicCacheControlHook()]
mock_response = MagicMock()
mock_response.json.return_value = {
"output": {"message": {"role": "assistant", "content": "ok"}},
"stopReason": "end_turn",
"usage": {"inputTokens": 100, "outputTokens": 4, "totalTokens": 104},
}
mock_response.status_code = 200
client = AsyncHTTPHandler()
with patch.object(client, "post", return_value=mock_response) as mock_post:
messages = [
{"role": "system", "content": f"System block {i}"} for i in range(4)
]
messages.append({"role": "user", "content": "What is the weather?"})
await litellm.acompletion(
model="bedrock/us.anthropic.claude-opus-4-6-v1:0",
messages=messages,
max_tokens=32,
tools=[
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather for a location",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string"}},
"required": ["location"],
},
},
}
],
cache_control_injection_points=[
{
"location": "message",
"role": "system",
"control": {"type": "ephemeral", "ttl": "1h"},
},
{"location": "tool_config"},
],
client=client,
)
request_body = json.loads(mock_post.call_args.kwargs["data"])
cache_points = sum(
1
for block in request_body.get("system", [])
if isinstance(block, dict) and "cachePoint" in block
)
for msg in request_body.get("messages", []):
content = msg.get("content", [])
if isinstance(content, list):
cache_points += sum(
1
for block in content
if isinstance(block, dict) and "cachePoint" in block
)
for tool in request_body.get("toolConfig", {}).get("tools", []):
if isinstance(tool, dict) and "cachePoint" in tool:
cache_points += 1
assert cache_points <= 4, (
f"Bedrock payload exceeded Anthropic's 4 cache_control block limit "
f"when mixing message and tool_config injection: found {cache_points}"
)