774df367ee
build_prompt_cache_plan deep-copied the whole request history on every request just to strip stale markers and place at most four new ones. Take a shallow list copy, copy-on-write only the rows the marker strip may touch (rows carrying cache_control or list content), and deep-copy the rows that receive markers in the tool-cache layout. The planner's output is byte-identical; the caller's history is never mutated (#106061). (cherry picked from commit a74fe3ed8af9fd38b4e4a7221a231b7484db9334)
336 lines
17 KiB
Python
336 lines
17 KiB
Python
"""Anthropic prompt caching strategy — pure functions, no AIAgent dependency.
|
||
|
||
Default layout: 4 cache_control breakpoints — the static system prefix, the end of the
|
||
system prompt, and the last 2 non-system messages (without a static prefix: one system
|
||
breakpoint plus the last 3 messages). All markers share one TTL (5m or 1h).
|
||
"""
|
||
|
||
import copy
|
||
from dataclasses import dataclass
|
||
from typing import Any, Dict, List
|
||
|
||
from agent.prompt_cache_boundary import find_stable_prefix
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class PromptCachePlan:
|
||
"""Request-local message and tool sections with their cache markers."""
|
||
|
||
messages: List[Dict[str, Any]]
|
||
tools: List[Dict[str, Any]]
|
||
|
||
|
||
def envelope_tool_part_cache_markers_supported(provider: str | None, base_url: str | None) -> bool:
|
||
"""Whether the envelope-layout route honors part-level markers on role:tool.
|
||
|
||
OpenRouter/Nous Portal relocate a part-level ``cache_control`` onto the ``tool_result``
|
||
block; LiteLLM-style proxies copy parts verbatim, so it lands at ``tool_result.content[0]``
|
||
(non-retryable 400). There, role:tool carries no part markers and the budget reallocates.
|
||
"""
|
||
from agent.agent_runtime_helpers import _is_litellm_route
|
||
|
||
return not _is_litellm_route((provider or "").strip().lower(), base_url or "")
|
||
|
||
|
||
def _text_part(text: str, cache_marker: dict | None = None) -> dict:
|
||
part: dict = {"type": "text", "text": text}
|
||
if cache_marker is not None:
|
||
part["cache_control"] = cache_marker
|
||
return part
|
||
|
||
|
||
def _apply_cache_marker(msg: dict, cache_marker: dict, native_anthropic: bool = False,
|
||
tool_part_markers: bool = True) -> None:
|
||
"""Add cache_control to a single message, handling all format variations."""
|
||
role = msg.get("role", "")
|
||
content = msg.get("content")
|
||
|
||
if role == "tool" and not native_anthropic and not tool_part_markers:
|
||
# LiteLLM-style envelope: a part marker → tool_result.content[0] → non-retryable 400.
|
||
return
|
||
if (role == "tool" and native_anthropic) or content is None or content == "":
|
||
# Native role:tool: top-level marker, the adapter moves it inside tool_result. Empty
|
||
# content: no part can carry it, and OpenRouter rejects a top-level marker on role:tool
|
||
# (silent hang) and ignores it on empty assistant turns — skip those on the envelope.
|
||
if not (role in ("tool", "assistant") and not native_anthropic):
|
||
msg["cache_control"] = cache_marker
|
||
elif isinstance(content, str):
|
||
stable_prefix = find_stable_prefix(content) if role == "user" else None
|
||
if stable_prefix is not None and content[len(stable_prefix):].strip():
|
||
# Builder-declared boundary: the scaffold carries the breakpoint and the volatile
|
||
# tail rides unmarked. Request-local only — the stored message stays a string.
|
||
msg["content"] = [_text_part(stable_prefix, cache_marker), _text_part(content[len(stable_prefix):])]
|
||
else:
|
||
msg["content"] = [_text_part(content, cache_marker)]
|
||
elif isinstance(content, list) and content and isinstance(content[-1], dict):
|
||
content[-1]["cache_control"] = cache_marker
|
||
|
||
|
||
def _can_carry_marker(msg: dict, native_anthropic: bool, tool_part_markers: bool = True) -> bool:
|
||
"""True if a marker on this message is actually honored by the provider.
|
||
|
||
Native Anthropic honors every message; the envelope layout only honors markers inside
|
||
content parts (empty content wastes a breakpoint) and ``tool_part_markers=False`` excludes
|
||
role:tool too (400). Must agree with :func:`_apply_cache_marker` (marks the LAST part).
|
||
"""
|
||
if native_anthropic:
|
||
return True
|
||
if msg.get("role") == "tool" and not tool_part_markers:
|
||
return False
|
||
content = msg.get("content")
|
||
return isinstance(content[-1], dict) if isinstance(content, list) and content else isinstance(content, str) and content != ""
|
||
|
||
|
||
def _build_marker(ttl: str) -> Dict[str, str]:
|
||
"""Build a cache_control marker dict for the given TTL ('5m' or '1h')."""
|
||
return {"type": "ephemeral", "ttl": "1h"} if ttl == "1h" else {"type": "ephemeral"}
|
||
|
||
|
||
# Alibaba-family providers (Qwen routes): five-minute context cache, 1h tier rejected. Shared
|
||
# with agent_runtime_helpers.anthropic_prompt_cache_policy so the opt-in and the TTL clamp
|
||
# never desync. Do NOT narrow this set to extend a TTL — narrowing DISABLES caching.
|
||
ALIBABA_FAMILY_PROVIDERS = frozenset({"opencode", "opencode-go", "opencode-zen", "alibaba"})
|
||
|
||
# 1h-tier ALLOW-list: only routes wire-measured to retain a 1h marker. Other opencode routes
|
||
# are UNMEASURED, not known-bad (opencode-go's `ephemeral_5m_input_tokens` label is not
|
||
# evidence of the retention window).
|
||
MEASURED_1H_PROVIDERS = frozenset({"opencode-go"})
|
||
|
||
# Models measured to ignore the 1h tier on a MEASURED_1H_PROVIDERS route; consulted only
|
||
# there (the same model on its own endpoint is a separate route).
|
||
NO_1H_TIER_MODELS = frozenset({"minimax-m2.5"})
|
||
|
||
|
||
def _flat_model(model: str) -> str:
|
||
"""Bare model id, tolerating aggregator prefixes (``vendor/model``)."""
|
||
return (model or "").strip().rsplit("/", 1)[-1].lower()
|
||
|
||
|
||
def is_qwen_model(model: str) -> bool:
|
||
"""True when ``model`` names a Qwen-family model (shared with anthropic_prompt_cache_policy)."""
|
||
return "qwen" in (model or "").lower()
|
||
|
||
|
||
# ``prompt_caching.cache_ttl: auto`` picks the tier by who paces the session. The 1h tier
|
||
# writes at 2x base vs 1.25x for 5m and only pays off when turns are more than five minutes
|
||
# apart — a person stepping away and coming back. Machine-paced sessions call every few
|
||
# seconds until they finish and never collect the retention, so they stay on 5m.
|
||
# Measured on one install (2 days of per-call logs): 63% of interactive cache-write tokens
|
||
# were cold re-writes after a 5–60 min idle gap (1h saves ~42% of write cost there), while
|
||
# 1h on subagents/cron would have cost ~49% more.
|
||
AUTO_CACHE_TTL = "auto"
|
||
MACHINE_PACED_SOURCES = frozenset({"subagent", "cron", "oneshot", "webhook", "kanban", "api", "tool", "batch"})
|
||
|
||
|
||
def auto_cache_ttl_for_source(source: str | None) -> str:
|
||
"""The tier ``auto`` resolves to for a session source: ``5m`` for machine-paced sources,
|
||
``1h`` for everything a human types into (cli, tui, desktop, messaging platforms)."""
|
||
return "5m" if (source or "").strip().lower() in MACHINE_PACED_SOURCES else "1h"
|
||
|
||
|
||
def effective_cache_ttl(ttl: str | None, *, model: str = "", provider: str = "") -> str:
|
||
"""Clamp a requested cache TTL to what the destination route supports (``None`` → ``5m``).
|
||
|
||
Qwen/Alibaba routes drop ``1h`` (→ ``5m``) except on ``MEASURED_1H_PROVIDERS`` minus
|
||
``NO_1H_TIER_MODELS``; that check runs BEFORE the generic Qwen clamp, which would swallow it.
|
||
"""
|
||
if ttl != "1h":
|
||
return ttl or "5m"
|
||
provider_lower = (provider or "").lower()
|
||
if provider_lower in MEASURED_1H_PROVIDERS:
|
||
return "5m" if _flat_model(model) in NO_1H_TIER_MODELS else "1h"
|
||
return "5m" if is_qwen_model(model) or provider_lower in ALIBABA_FAMILY_PROVIDERS else "1h"
|
||
|
||
|
||
def _apply_system_cache_markers(
|
||
message: dict, cache_marker: dict, static_system_prefix: str | None, *,
|
||
native_anthropic: bool, mark_suffix: bool = True, fallback_to_whole: bool = True,
|
||
) -> int:
|
||
"""Mark the static system prefix (and optionally the full prompt); returns markers applied.
|
||
|
||
The stored system prompt stays one string, split only in the request. ``mark_suffix=False``
|
||
is the tool-cache-plan layout (suffix budget spent on the tools array); ``fallback_to_whole=
|
||
False`` marks nothing when the split is impossible. When the prompt IS the prefix the whole
|
||
message is one block — never an empty text block (400).
|
||
"""
|
||
content = message.get("content")
|
||
if isinstance(static_system_prefix, str) and static_system_prefix and isinstance(content, str) and content.startswith(static_system_prefix):
|
||
suffix = content[len(static_system_prefix):]
|
||
if suffix.strip():
|
||
message["content"] = [_text_part(static_system_prefix, cache_marker),
|
||
_text_part(suffix, cache_marker if mark_suffix else None)]
|
||
return 2 if mark_suffix else 1
|
||
elif not fallback_to_whole:
|
||
return 0
|
||
_apply_cache_marker(message, cache_marker, native_anthropic=native_anthropic)
|
||
return 1
|
||
|
||
|
||
def _has_part_marker(content: Any) -> bool:
|
||
return isinstance(content, list) and any(isinstance(part, dict) and "cache_control" in part for part in content)
|
||
|
||
|
||
def strip_anthropic_cache_control(api_messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||
"""Remove ``cache_control`` markers and undo decoration-produced list shapes (in place).
|
||
|
||
Used before re-decorating after a mid-turn failover. Flattening to a string is restricted
|
||
to the exact shapes :func:`apply_anthropic_cache_control` produces from string content
|
||
(single text part, two-part system split, two-part skill split) so the ``""``-join is
|
||
byte-exact. Marker removal is copy-on-write on part dicts: parts can alias caller-held
|
||
lists and stripping must never rewrite the stored transcript.
|
||
"""
|
||
for msg in api_messages:
|
||
if not isinstance(msg, dict):
|
||
continue
|
||
msg.pop("cache_control", None)
|
||
content = msg.get("content")
|
||
if not isinstance(content, list):
|
||
continue
|
||
role = msg.get("role")
|
||
# The skill split is the only decoration marking the FIRST part of a user message,
|
||
# so the shape alone identifies it even after the prefix registry evicted the entry.
|
||
skill_split_shape = (role == "user" and len(content) == 2 and all(isinstance(p, dict) for p in content)
|
||
and "cache_control" in content[0] and "cache_control" not in content[1])
|
||
if _has_part_marker(content):
|
||
content = msg["content"] = [
|
||
{k: v for k, v in part.items() if k != "cache_control"}
|
||
if isinstance(part, dict) and "cache_control" in part else part for part in content
|
||
]
|
||
plain_text_parts = content and all(
|
||
isinstance(part, dict) and part.get("type", "text") == "text"
|
||
and isinstance(part.get("text"), str) and set(part.keys()) <= {"type", "text"} for part in content
|
||
)
|
||
if plain_text_parts and (len(content) == 1 or (role == "system" and len(content) == 2) or skill_split_shape):
|
||
msg["content"] = "".join(part["text"] for part in content)
|
||
return api_messages
|
||
|
||
|
||
def strip_anthropic_tool_cache_control(tools: List[Dict[str, Any]] | None) -> List[Dict[str, Any]]:
|
||
"""Return copied tools without request-local Anthropic cache markers."""
|
||
cleaned = copy.deepcopy(tools or [])
|
||
for tool in cleaned:
|
||
if isinstance(tool, dict):
|
||
tool.pop("cache_control", None)
|
||
return cleaned
|
||
|
||
|
||
def _count_cache_markers(messages: List[Dict[str, Any]], tools: List[Dict[str, Any]]) -> int:
|
||
"""Count the wire-visible cache markers in a request-local plan."""
|
||
parts = [p for m in messages if isinstance(m, dict) and isinstance(m.get("content"), list) for p in m["content"]]
|
||
return sum(1 for item in [*messages, *parts, *tools] if isinstance(item, dict) and "cache_control" in item)
|
||
|
||
|
||
def _completed_transaction_endpoint_indexes(messages: List[Dict[str, Any]], *, native_anthropic: bool) -> List[int]:
|
||
"""Select legal ends of completed tool runs and ordinary turns."""
|
||
|
||
def _tool_run_end(start: int) -> int:
|
||
end = start
|
||
while end < len(messages) and isinstance(messages[end], dict) and messages[end].get("role") == "tool":
|
||
end += 1
|
||
return end
|
||
|
||
endpoints: List[int] = []
|
||
index = 0
|
||
while index < len(messages):
|
||
message = messages[index]
|
||
if not isinstance(message, dict) or message.get("role") == "system":
|
||
index += 1
|
||
continue
|
||
role = message.get("role")
|
||
|
||
if role == "assistant" and message.get("tool_calls"):
|
||
result_end = _tool_run_end(index + 1)
|
||
if result_end > index + 1 and _can_carry_marker(messages[result_end - 1], native_anthropic):
|
||
endpoints.append(result_end - 1)
|
||
index = result_end
|
||
continue
|
||
|
||
if role == "tool":
|
||
index = _tool_run_end(index)
|
||
continue
|
||
|
||
open_turn = (role == "user" and index + 1 < len(messages)) or (
|
||
role == "assistant" and message.get("content") in (None, ""))
|
||
if not open_turn and _can_carry_marker(message, native_anthropic):
|
||
endpoints.append(index)
|
||
index += 1
|
||
return endpoints
|
||
|
||
|
||
def build_prompt_cache_plan(
|
||
api_messages: List[Dict[str, Any]], tools: List[Dict[str, Any]] | None, *,
|
||
cache_ttl: str = "5m", native_anthropic: bool = False, static_system_prefix: str | None = None,
|
||
direct_native_tool_cache: bool = False, tool_part_markers: bool = True,
|
||
) -> PromptCachePlan:
|
||
"""Build copy-on-write cache sections for one resolved request destination
|
||
(``tool_part_markers=False`` keeps markers off role:tool parts on LiteLLM-style routes)."""
|
||
messages = list(api_messages or [])
|
||
for i, msg in enumerate(messages):
|
||
if isinstance(msg, dict) and (
|
||
"cache_control" in msg or isinstance(msg.get("content"), list)
|
||
):
|
||
messages[i] = strip_anthropic_cache_control([dict(msg)])[0]
|
||
planned_tools = strip_anthropic_tool_cache_control(tools)
|
||
|
||
if not direct_native_tool_cache or not planned_tools:
|
||
planned_messages = apply_anthropic_cache_control(
|
||
messages, cache_ttl=cache_ttl, native_anthropic=native_anthropic,
|
||
static_system_prefix=static_system_prefix, tool_part_markers=tool_part_markers)
|
||
return PromptCachePlan(messages=planned_messages, tools=planned_tools)
|
||
|
||
marker = _build_marker(cache_ttl)
|
||
if messages and isinstance(messages[0], dict) and messages[0].get("role") == "system":
|
||
# Tool-cache layout: only the static prefix carries a system-side marker; the
|
||
# volatile suffix's budget is spent on the tools array.
|
||
messages[0] = copy.deepcopy(messages[0])
|
||
_apply_system_cache_markers(messages[0], marker, static_system_prefix,
|
||
native_anthropic=True, mark_suffix=False, fallback_to_whole=False)
|
||
planned_tools[-1]["cache_control"] = dict(marker)
|
||
for endpoint in _completed_transaction_endpoint_indexes(messages, native_anthropic=True)[-2:]:
|
||
messages[endpoint] = copy.deepcopy(messages[endpoint])
|
||
_apply_cache_marker(messages[endpoint], marker, native_anthropic=True)
|
||
|
||
return PromptCachePlan(messages=messages, tools=planned_tools)
|
||
|
||
|
||
def apply_anthropic_cache_control(
|
||
api_messages: List[Dict[str, Any]], cache_ttl: str = "5m", native_anthropic: bool = False,
|
||
static_system_prefix: str | None = None, tool_part_markers: bool = True,
|
||
) -> List[Dict[str, Any]]:
|
||
"""Apply Anthropic cache-control markers to API messages.
|
||
|
||
With a matching ``static_system_prefix`` the prefix and full system prompt each get a
|
||
marker and the remaining two go to the latest cacheable non-system messages; otherwise
|
||
the legacy system-and-3 layout applies. Idempotent: pre-existing markers are stripped from
|
||
a per-message copy first. Returns a shallow list copy with deep copies of modified messages.
|
||
|
||
Idempotent: pre-existing ``cache_control`` markers are stripped from a per-message copy before new ones
|
||
are placed, so calling this twice (or handing it messages a prior call already marked) can never
|
||
accumulate past 4 markers. Only messages that already carry a marker pay the copy cost — a shallow
|
||
top-level copy suffices because :func:`strip_anthropic_cache_control` is copy-on-write on content parts
|
||
— and the rest of the copy-on-write contract is unchanged (#90971).
|
||
"""
|
||
if not api_messages:
|
||
return api_messages
|
||
|
||
messages = list(api_messages)
|
||
marker = _build_marker(cache_ttl)
|
||
|
||
for i, msg in enumerate(messages):
|
||
if isinstance(msg, dict) and ("cache_control" in msg or _has_part_marker(msg.get("content"))):
|
||
messages[i] = strip_anthropic_cache_control([dict(msg)])[0]
|
||
|
||
breakpoints_used = 0
|
||
if messages[0].get("role") == "system":
|
||
messages[0] = copy.deepcopy(messages[0])
|
||
breakpoints_used = _apply_system_cache_markers(messages[0], marker, static_system_prefix,
|
||
native_anthropic=native_anthropic)
|
||
|
||
non_sys = [i for i, m in enumerate(messages) if m.get("role") != "system"
|
||
and _can_carry_marker(m, native_anthropic=native_anthropic, tool_part_markers=tool_part_markers)]
|
||
for idx in non_sys[-(4 - breakpoints_used):]:
|
||
messages[idx] = copy.deepcopy(messages[idx])
|
||
_apply_cache_marker(messages[idx], marker, native_anthropic=native_anthropic, tool_part_markers=tool_part_markers)
|
||
|
||
return messages
|