use anyhow::Result;
use serde_json::{Value, json};
use crate::llmtrim::gate::{GateKind, PlanEntry, Transform};
use crate::llmtrim::ir::Request;
use crate::llmtrim::provider::Provider;
use crate::llmtrim::stages::tools::fnv1a;
pub struct CacheStage {
pub max_breakpoints: usize,
pub prompt_key: bool,
pub auto_ttl: String,
}
impl Transform for CacheStage {
fn name(&self) -> &str {
"cache"
}
fn gate_kind(&self) -> GateKind {
GateKind::Structural
}
fn scope(&self) -> crate::llmtrim::gate::Scope {
crate::llmtrim::gate::Scope::Tools
}
fn apply(
&self,
req: &mut Request,
provider: &dyn Provider,
_plan: &mut Vec<PlanEntry>,
) -> Result<()> {
if !has_client_breakpoint(req.raw()) {
sort_tools(req);
if self.prompt_key {
let key = format!("{:016x}", cache_prefix_hash(req));
provider.set_prompt_cache_key(req, &key);
}
set_router_cache_breakpoint(req, &self.auto_ttl);
}
provider.set_cache_breakpoints(req, self.max_breakpoints);
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum RouterCache {
Automatic,
SystemBlock,
Implicit,
}
fn set_router_cache_breakpoint(req: &mut Request, auto_ttl: &str) {
let kind = match crate::llmtrim::cache_zone::router_model(req.raw()) {
Some((vendor, model)) => classify_route(vendor, model),
None => return,
};
match kind {
RouterCache::Automatic => set_automatic_marker(req, auto_ttl),
RouterCache::SystemBlock => mark_system_message(req),
RouterCache::Implicit => {}
}
}
fn classify_route(vendor: &str, model: &str) -> RouterCache {
if vendor.eq_ignore_ascii_case("anthropic") {
return RouterCache::Automatic;
}
let gemini = vendor.eq_ignore_ascii_case("google")
&& model
.get(..6)
.is_some_and(|m| m.eq_ignore_ascii_case("gemini"));
if vendor.eq_ignore_ascii_case("qwen") || gemini {
return RouterCache::SystemBlock;
}
RouterCache::Implicit
}
fn set_automatic_marker(req: &mut Request, ttl: &str) {
let marker = if ttl.is_empty() {
json!({"type": "ephemeral"})
} else {
json!({"type": "ephemeral", "ttl": ttl})
};
if let Some(obj) = req.raw_mut().as_object_mut() {
obj.entry("cache_control").or_insert(marker);
}
}
fn mark_system_message(req: &mut Request) {
let Some(messages) = req
.raw_mut()
.get_mut("messages")
.and_then(Value::as_array_mut)
else {
return;
};
let Some(system) = messages
.iter_mut()
.find(|m| m.get("role").and_then(Value::as_str) == Some("system"))
else {
return;
};
match system.get_mut("content") {
Some(Value::String(text)) => {
let text = std::mem::take(text);
if let Some(obj) = system.as_object_mut() {
obj.insert(
"content".to_string(),
json!([{
"type": "text",
"text": text,
"cache_control": {"type": "ephemeral"},
}]),
);
}
}
Some(Value::Array(blocks)) => {
if let Some(last) = blocks.last_mut()
&& let Some(obj) = last.as_object_mut()
{
obj.insert("cache_control".to_string(), json!({"type": "ephemeral"}));
}
}
_ => {}
}
}
fn has_client_breakpoint(raw: &Value) -> bool {
["system", "messages", "tools"]
.iter()
.filter_map(|key| raw.get(*key))
.any(crate::llmtrim::cache_zone::has_cache_control)
}
fn sort_tools(req: &mut Request) {
let Some(Value::Array(tools)) = req.raw_mut().get_mut("tools") else {
return;
};
for tool in tools.iter_mut() {
sort_keys(tool);
sort_function_declarations(tool);
}
tools.sort_by(|a, b| tool_name(a).cmp(tool_name(b)));
}
fn sort_function_declarations(tool: &mut Value) {
for key in ["functionDeclarations", "function_declarations"] {
if let Some(Value::Array(decls)) = tool.get_mut(key) {
decls.sort_by(|a, b| tool_name(a).cmp(tool_name(b)));
}
}
}
fn tool_name(tool: &Value) -> &str {
tool.get("name")
.or_else(|| tool.get("function").and_then(|f| f.get("name")))
.and_then(Value::as_str)
.unwrap_or("")
}
fn sort_keys(v: &mut Value) {
match v {
Value::Object(map) => {
for child in map.values_mut() {
sort_keys(child);
}
let mut entries: Vec<(String, Value)> = std::mem::take(map).into_iter().collect();
entries.sort_by(|a, b| a.0.cmp(&b.0));
for (k, val) in entries {
map.insert(k, val);
}
}
Value::Array(a) => a.iter_mut().for_each(sort_keys),
_ => {}
}
}
pub fn cache_prefix_hash(req: &Request) -> u64 {
let raw = req.raw();
let mut buf = String::new();
let mut hashed_anything = false;
if let Some(sys) = raw.get("system") {
buf.push_str(&sys.to_string());
buf.push('\u{1f}'); hashed_anything = true;
}
if let Some(tools) = raw.get("tools") {
buf.push_str(&tools.to_string());
buf.push('\u{1f}');
hashed_anything = true;
}
if !hashed_anything && let Some(msgs) = raw.get("messages").and_then(Value::as_array) {
for m in msgs {
if m.get("role").and_then(Value::as_str) == Some("system") {
buf.push_str(&m.to_string());
buf.push('\u{1f}');
} else {
break;
}
}
}
fnv1a(buf.bytes())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::llmtrim::ir::ProviderKind;
use crate::llmtrim::provider::{AnthropicProvider, GoogleProvider, OpenAiProvider};
use serde_json::json;
fn anthropic(body: Value) -> Request {
Request::from_value(ProviderKind::Anthropic, body)
}
#[test]
fn anthropic_caches_system_string_as_block() {
let mut req = anthropic(json!({"system":"you are helpful","max_tokens":1,"messages":[]}));
AnthropicProvider.set_cache_breakpoints(&mut req, 4);
let sys = req.raw().get("system").unwrap();
assert_eq!(
sys.pointer("/0/cache_control/type").and_then(Value::as_str),
Some("ephemeral"),
"string system becomes a cached text block"
);
assert_eq!(
sys.pointer("/0/text").and_then(Value::as_str),
Some("you are helpful")
);
}
#[test]
fn anthropic_caches_last_tool() {
let mut req = anthropic(json!({
"max_tokens":1, "messages":[],
"tools":[{"name":"a","input_schema":{}},{"name":"b","input_schema":{}}]
}));
AnthropicProvider.set_cache_breakpoints(&mut req, 4);
assert_eq!(
req.raw()
.pointer("/tools/1/cache_control/type")
.and_then(Value::as_str),
Some("ephemeral")
);
assert!(req.raw().pointer("/tools/0/cache_control").is_none());
}
#[test]
fn respects_max_breakpoints() {
let mut req = anthropic(json!({
"system":"sys","max_tokens":1,"messages":[],
"tools":[{"name":"a","input_schema":{}}]
}));
AnthropicProvider.set_cache_breakpoints(&mut req, 1);
assert!(req.raw().pointer("/tools/0/cache_control").is_some());
assert!(
req.raw().get("system").unwrap().is_string(),
"system not converted (budget spent)"
);
}
#[test]
fn openai_is_noop() {
let body =
json!({"messages":[{"role":"system","content":"s"},{"role":"user","content":"hi"}]});
let mut req = Request::from_value(ProviderKind::OpenAi, body.clone());
OpenAiProvider.set_cache_breakpoints(&mut req, 4);
assert_eq!(
req.raw(),
&body,
"OpenAI request is unchanged (automatic caching)"
);
}
fn run_cache_stage(req: &mut Request, provider: &dyn Provider) {
run_cache_stage_with(req, provider, true, "1h");
}
fn run_cache_stage_with(
req: &mut Request,
provider: &dyn Provider,
prompt_key: bool,
auto_ttl: &str,
) {
let mut plan: Vec<PlanEntry> = Vec::new();
CacheStage {
max_breakpoints: 4,
prompt_key,
auto_ttl: auto_ttl.to_string(),
}
.apply(req, provider, &mut plan)
.unwrap();
}
#[test]
fn stabilize_sorts_tools_and_schema_keys() {
let mut req = Request::from_value(
ProviderKind::OpenAi,
json!({
"model": "gpt-4o",
"messages": [{"role": "user", "content": "hi"}],
"tools": [
{"type": "function", "function": {"name": "zebra", "parameters": {"b": 1, "a": 2}}},
{"type": "function", "function": {"name": "apple", "parameters": {}}},
]
}),
);
run_cache_stage(&mut req, &OpenAiProvider);
let tools = req.raw().get("tools").and_then(Value::as_array).unwrap();
assert_eq!(
tools[0].pointer("/function/name").unwrap(),
"apple",
"tools sorted by name"
);
assert_eq!(tools[1].pointer("/function/name").unwrap(), "zebra");
let keys: Vec<&str> = tools[1]
.pointer("/function/parameters")
.and_then(Value::as_object)
.unwrap()
.keys()
.map(String::as_str)
.collect();
assert_eq!(keys, ["a", "b"], "schema keys canonicalized");
}
fn routed(model: &str) -> Request {
Request::from_value(
ProviderKind::OpenAi,
json!({
"model": model,
"messages": [
{"role": "system", "content": "a long stable system prompt"},
{"role": "user", "content": "hi"},
],
"tools": [{"type": "function", "function": {"name": "read", "parameters": {}}}]
}),
)
}
#[test]
fn routed_anthropic_model_gets_automatic_caching() {
let mut req = routed("anthropic/claude-sonnet-4.5");
run_cache_stage(&mut req, &OpenAiProvider);
assert_eq!(
req.raw().pointer("/cache_control/ttl").unwrap(),
"1h",
"router asked to cache the Anthropic route"
);
assert!(
req.raw().get("prompt_cache_key").is_some(),
"prompt_cache_key still set"
);
}
#[test]
fn automatic_marker_honors_the_configured_ttl() {
let mut req = routed("anthropic/claude-sonnet-4.5");
run_cache_stage_with(&mut req, &OpenAiProvider, true, "");
assert_eq!(
req.raw().get("cache_control").unwrap(),
&json!({"type": "ephemeral"}),
"no ttl field for the 5m default"
);
let mut req = routed("anthropic/claude-sonnet-4.5");
run_cache_stage_with(&mut req, &OpenAiProvider, true, "1h");
assert_eq!(
req.raw().pointer("/cache_control/ttl").unwrap(),
"1h",
"1h is sent explicitly"
);
}
#[test]
fn prompt_cache_key_is_suppressed_for_strict_backends() {
let mut req = routed("gpt-5.2");
run_cache_stage_with(&mut req, &OpenAiProvider, false, "1h");
assert!(
req.raw().get("prompt_cache_key").is_none(),
"no key when the backend does not accept one"
);
assert_eq!(
req.raw().pointer("/tools/0/function/name").unwrap(),
"read",
"tools still canonicalized"
);
}
#[test]
fn qwen_and_gemini_routes_get_a_system_block_breakpoint() {
for model in ["qwen/qwen3-coder-plus", "google/gemini-3-pro"] {
let mut req = routed(model);
run_cache_stage(&mut req, &OpenAiProvider);
assert!(
req.raw().get("cache_control").is_none(),
"{model} must not get the top-level form"
);
assert_eq!(
req.raw()
.pointer("/messages/0/content/0/cache_control/type")
.unwrap(),
"ephemeral",
"{model} gets a system-block breakpoint"
);
assert_eq!(
req.raw().pointer("/messages/0/content/0/text").unwrap(),
"a long stable system prompt",
"{model} keeps its system text"
);
assert!(
req.raw()
.pointer("/messages/0/content/0/cache_control/ttl")
.is_none(),
"{model} must not carry a ttl"
);
}
}
#[test]
fn system_block_breakpoint_marks_the_last_block_of_an_array() {
let mut req = Request::from_value(
ProviderKind::OpenAi,
json!({
"model": "google/gemini-3-pro",
"messages": [{"role": "system", "content": [
{"type": "text", "text": "preamble"},
{"type": "text", "text": "the bulk"},
]}]
}),
);
run_cache_stage(&mut req, &OpenAiProvider);
assert!(
req.raw()
.pointer("/messages/0/content/0/cache_control")
.is_none(),
"earlier blocks stay unmarked"
);
assert_eq!(
req.raw()
.pointer("/messages/0/content/1/cache_control/type")
.unwrap(),
"ephemeral",
"the breakpoint ends the prefix"
);
}
#[test]
fn gemma_and_other_uncached_families_are_left_alone() {
for model in ["google/gemma-3-27b", "meta-llama/llama-4", "mistral/large"] {
let mut req = routed(model);
run_cache_stage(&mut req, &OpenAiProvider);
assert!(req.raw().get("cache_control").is_none(), "{model}");
assert!(
req.raw()
.pointer("/messages/0/content/0/cache_control")
.is_none(),
"{model} keeps its plain string system content"
);
}
}
#[test]
fn routed_variants_and_price_prefixes_are_recognized() {
for model in [
"~anthropic/claude-sonnet-latest",
"anthropic/claude-opus-4.5:nitro",
"Anthropic/claude-haiku-4.5",
] {
let mut req = routed(model);
run_cache_stage(&mut req, &OpenAiProvider);
assert!(
req.raw().get("cache_control").is_some(),
"{model} is an Anthropic route"
);
}
}
#[test]
fn implicitly_cached_routes_get_no_marker() {
for model in [
"openai/gpt-5.2",
"deepseek/deepseek-chat",
"google/gemini-3-pro",
"qwen/qwen3-coder-plus",
"openrouter/auto",
"gpt-5.2",
"claude-sonnet-4-5",
] {
let mut req = routed(model);
run_cache_stage(&mut req, &OpenAiProvider);
assert!(
req.raw().get("cache_control").is_none(),
"{model} must not get a top-level marker"
);
}
}
#[test]
fn routed_marker_defers_to_a_client_breakpoint() {
let mut req = Request::from_value(
ProviderKind::OpenAi,
json!({
"model": "anthropic/claude-sonnet-4.5",
"messages": [{"role": "system", "content": [
{"type": "text", "text": "s", "cache_control": {"type": "ephemeral"}}
]}]
}),
);
run_cache_stage(&mut req, &OpenAiProvider);
assert!(req.raw().get("cache_control").is_none());
}
#[test]
fn stabilize_sorts_gemini_function_declarations() {
let mut req = Request::from_value(
ProviderKind::Google,
json!({
"contents": [{"role": "user", "parts": [{"text": "hi"}]}],
"tools": [{"functionDeclarations": [
{"name": "zebra", "parameters": {"b": 1, "a": 2}},
{"name": "apple", "parameters": {}},
]}]
}),
);
run_cache_stage(&mut req, &GoogleProvider);
assert_eq!(
req.raw()
.pointer("/tools/0/functionDeclarations/0/name")
.unwrap(),
"apple",
"declarations sorted by name"
);
let keys: Vec<&str> = req
.raw()
.pointer("/tools/0/functionDeclarations/1/parameters")
.and_then(Value::as_object)
.unwrap()
.keys()
.map(String::as_str)
.collect();
assert_eq!(keys, ["a", "b"], "schema keys canonicalized");
}
#[test]
fn openai_gets_a_stable_prompt_cache_key() {
let mut req = Request::from_value(
ProviderKind::OpenAi,
json!({"model": "gpt-4o", "messages": [{"role": "system", "content": "s"}, {"role": "user", "content": "hi"}]}),
);
run_cache_stage(&mut req, &OpenAiProvider);
assert!(
req.raw()
.get("prompt_cache_key")
.and_then(Value::as_str)
.is_some(),
"prompt_cache_key injected for OpenAI"
);
}
#[test]
fn stabilize_defers_to_client_managed_caching() {
let mut req = anthropic(json!({
"max_tokens": 1, "messages": [],
"tools": [
{"name": "zebra", "input_schema": {}, "cache_control": {"type": "ephemeral"}},
{"name": "apple", "input_schema": {}},
]
}));
run_cache_stage(&mut req, &AnthropicProvider);
assert_eq!(
req.raw().pointer("/tools/0/name").unwrap(),
"zebra",
"tool order preserved when the client manages caching"
);
}
#[test]
fn top_level_automatic_caching_still_stabilizes_tools() {
let mut req = anthropic(json!({
"max_tokens": 1, "messages": [],
"cache_control": {"type": "ephemeral", "ttl": "1h"},
"tools": [
{"name": "zebra", "input_schema": {}},
{"name": "apple", "input_schema": {}},
]
}));
run_cache_stage(&mut req, &AnthropicProvider);
assert_eq!(
req.raw().pointer("/tools/0/name").unwrap(),
"apple",
"tools canonicalized despite the top-level automatic-caching marker"
);
assert!(
req.raw().pointer("/tools/1/cache_control").is_none(),
"no 5m breakpoints added alongside the top-level 1h marker"
);
}
#[test]
fn prefix_hash_is_stable_and_distinct() {
let a = anthropic(json!({"system":"SAME","messages":[{"role":"user","content":"q1"}]}));
let b = anthropic(
json!({"system":"SAME","messages":[{"role":"user","content":"q2 different"}]}),
);
let c = anthropic(json!({"system":"OTHER","messages":[{"role":"user","content":"q1"}]}));
assert_eq!(cache_prefix_hash(&a), cache_prefix_hash(&b));
assert_ne!(cache_prefix_hash(&a), cache_prefix_hash(&c));
}
}