use crate::tasks::generate::{GenerateRequest, ResponseFormat};
use car_inference_types::Message;
pub fn apply_cache_partition(req: &mut GenerateRequest, partition: &str) {
let mut system = vec![partition.to_string()];
if let Some(messages) = req.messages.as_mut() {
messages.retain(|message| match message {
Message::System { content } => {
system.push(content.clone());
false
}
Message::User { .. }
| Message::UserMultimodal { .. }
| Message::Assistant { .. }
| Message::ToolResult { .. }
| Message::ProviderOutputItems { .. } => true,
});
}
if let Some(context) = req.context.take().filter(|context| !context.is_empty()) {
system.push(context);
}
req.context = Some(system.join("\n\n"));
req.context_stable_prefix = None;
if let Some(tools) = req.tools.as_mut().filter(|tools| !tools.is_empty()) {
tools.insert(
0,
serde_json::json!({
"name": partition,
"description": format!(
"{partition}: reserved by CAR to keep prompt caches apart. Never call it."
),
"parameters": { "type": "object", "properties": {} },
}),
);
}
if let Some(ResponseFormat::JsonSchema { name, .. }) = req.response_format.as_mut() {
*name = Some(partition.to_string());
}
}
pub const CACHE_VALUE_WITHIN_BYTES: usize = 256;
pub fn cached_parts(protocol: crate::schema::ApiProtocol) -> &'static [&'static str] {
use crate::schema::ApiProtocol;
match protocol {
ApiProtocol::OpenAiCompat | ApiProtocol::OpenRouter | ApiProtocol::AzureOpenAi => {
&["tools", "response_format", "messages"]
}
ApiProtocol::OpenAiResponses => &["tools", "text", "instructions"],
ApiProtocol::Anthropic => &["tools", "system"],
ApiProtocol::Google | ApiProtocol::VertexAi => &["tools", "systemInstruction"],
ApiProtocol::Bedrock => &["toolConfig", "system"],
}
}
pub const PARSLEE_CACHED_PARTS: &[&str] = &["tools", "input"];
pub fn cached_prefix_violation(
parts: &[&'static str],
body: &serde_json::Value,
partition: &str,
) -> Option<&'static str> {
let mut holds_value = false;
for part in parts.iter().copied() {
let Some(value) = body.get(part).filter(|value| !value.is_null()) else {
continue;
};
let bytes = serde_json::to_string(value).unwrap_or_default();
let leads = bytes
.find(partition)
.is_some_and(|at| at <= CACHE_VALUE_WITHIN_BYTES);
if !leads && bytes.len() > CACHE_VALUE_WITHIN_BYTES {
return Some(part);
}
holds_value |= leads;
}
(!holds_value).then_some(NO_PART_HOLDS_THE_VALUE)
}
pub const NO_PART_HOLDS_THE_VALUE: &str = "no cached part holds the value";
tokio::task_local! {
static CACHE_PARTITION: Option<String>;
}
pub async fn scope_cache_partition<F: std::future::Future>(
partition: Option<String>,
future: F,
) -> F::Output {
CACHE_PARTITION.scope(partition, future).await
}
pub(crate) fn enforce_cached_prefix(
parts: &[&'static str],
body: &serde_json::Value,
) -> Result<(), crate::InferenceError> {
let violation = CACHE_PARTITION
.try_with(|partition| {
partition
.as_deref()
.and_then(|partition| cached_prefix_violation(parts, body, partition))
})
.ok()
.flatten();
match violation {
None => Ok(()),
Some(_) => Err(crate::InferenceError::UnsupportedMode {
mode: "worker request",
backend: "prompt-cache isolation",
reason: "the label's cache value would not lead every cached part of the request",
}),
}
}
#[cfg(test)]
mod tests {
use super::{
apply_cache_partition, cached_parts, cached_prefix_violation, enforce_cached_prefix,
scope_cache_partition, CACHE_VALUE_WITHIN_BYTES,
};
use crate::protocol::{
AnthropicHandler, ApiRequest, AzureOpenAiHandler, BedrockHandler, GoogleHandler,
OpenAiHandler, OpenAiResponsesHandler, OpenRouterHandler, ProtocolHandler, VertexHandler,
};
use crate::tasks::generate::{GenerateRequest, ResponseFormat};
use car_inference_types::Message;
use serde_json::{json, Value};
const PARTITION: &str = "car_label_0123456789abcdef";
fn worker_request() -> GenerateRequest {
GenerateRequest {
prompt: "WORKER_PROMPT".into(),
context: Some("WORKER_CTX".into()),
messages: Some(vec![
Message::System {
content: "WORKER_SYS".into(),
},
Message::User {
content: "WORKER_USER".into(),
},
]),
tools: Some(vec![json!({
"name": "worker_tool",
"description": "WORKER_TOOL",
"parameters": {"type": "object", "properties": {"q": {"type": "string"}}},
})]),
response_format: Some(ResponseFormat::JsonSchema {
schema: json!({
"type": "object",
"description": "WORKER_SCHEMA",
"properties": {"WORKER_PROP": {"type": "string"}},
}),
strict: false,
name: Some("answer".into()),
}),
cache_control: true,
..Default::default()
}
}
fn body(handler: &dyn ProtocolHandler, req: &GenerateRequest) -> Value {
let (messages, system) = handler.build_messages(
req.messages.as_deref().unwrap_or(&[]),
&req.prompt,
req.context.as_deref(),
None,
);
let tools = req
.tools
.as_deref()
.map(|tools| handler.build_tools(tools))
.filter(|tools| !tools.is_empty());
handler.build_request_body(&ApiRequest {
model: "model".into(),
messages,
system,
system_stable_prefix: None,
temperature: 0.0,
max_tokens: 64,
tools,
tool_choice: None,
parallel_tool_calls: None,
stream: false,
budget_tokens: 0,
cache_control: req.cache_control,
cache_ttl: Default::default(),
response_format: req.response_format.clone(),
stop: Vec::new(),
reasoning_effort: None,
enable_thinking: None,
template_reasoning_effort: None,
})
}
fn contains(value: &Value, needle: &str) -> bool {
match value {
Value::String(text) => text.contains(needle),
Value::Array(items) => items.iter().any(|item| contains(item, needle)),
Value::Object(object) => object
.iter()
.any(|(key, item)| key.contains(needle) || contains(item, needle)),
_ => false,
}
}
fn big_defs() -> Value {
let mut defs = serde_json::Map::new();
for i in 0..40 {
defs.insert(
format!("WORKER_DEF_{i}"),
json!({"type": "object", "description": "WORKER_DEFINITION_TEXT ".repeat(4),
"properties": {"value": {"type": "string"}}}),
);
}
Value::Object(defs)
}
fn schema_variants() -> Vec<(&'static str, Value)> {
vec![
(
"pydantic v2 ($defs and $ref)",
json!({
"$defs": big_defs(),
"$ref": "#/$defs/WORKER_DEF_0",
"title": "WORKER_TITLE",
}),
),
(
"zod / draft-07 (definitions and top-level $ref)",
json!({
"$schema": "http://json-schema.org/draft-07/schema#",
"definitions": big_defs(),
"$ref": "#/definitions/WORKER_DEF_1",
}),
),
(
"root anyOf",
json!({"anyOf": [{"type": "string", "description": "WORKER_SCHEMA"}, {"type": "number"}]}),
),
(
"root oneOf",
json!({"oneOf": [{"type": "string"}, {"type": "null"}], "description": "WORKER_SCHEMA"}),
),
(
"root allOf",
json!({"allOf": [{"type": "object"}, {"properties": {"WORKER_PROP": {"type": "string"}}}]}),
),
(
"additionalProperties as a schema",
json!({
"type": "object",
"additionalProperties": {"type": "object", "properties": {"WORKER_PROP": {"type": "string"}}},
}),
),
(
"description last",
json!({"type": "object", "properties": {"WORKER_PROP": {"type": "string"}}, "description": "WORKER_SCHEMA"}),
),
]
}
fn handlers() -> Vec<(
&'static str,
crate::schema::ApiProtocol,
Box<dyn ProtocolHandler>,
)> {
use crate::schema::ApiProtocol;
vec![
("openai", ApiProtocol::OpenAiCompat, Box::new(OpenAiHandler)),
(
"openai-responses",
ApiProtocol::OpenAiResponses,
Box::new(OpenAiResponsesHandler),
),
(
"openrouter",
ApiProtocol::OpenRouter,
Box::new(OpenRouterHandler),
),
(
"anthropic",
ApiProtocol::Anthropic,
Box::new(AnthropicHandler),
),
("google", ApiProtocol::Google, Box::new(GoogleHandler)),
("vertex", ApiProtocol::VertexAi, Box::new(VertexHandler)),
("bedrock", ApiProtocol::Bedrock, Box::new(BedrockHandler)),
(
"azure",
ApiProtocol::AzureOpenAi,
Box::new(AzureOpenAiHandler),
),
]
}
#[test]
fn per_label_value_precedes_every_worker_byte() {
for (variant, schema) in schema_variants() {
let mut req = worker_request();
req.response_format = Some(ResponseFormat::JsonSchema {
schema: schema.clone(),
strict: false,
name: Some("WORKER_CHOSEN_NAME_".repeat(20)),
});
apply_cache_partition(&mut req, PARTITION);
for (name, protocol, handler) in handlers() {
let body = body(handler.as_ref(), &req);
assert!(contains(&body, PARTITION), "{name}: no value in {body}");
for marker in ["WORKER_SYS", "WORKER_CTX", "WORKER_TOOL", "WORKER_USER"] {
assert!(contains(&body, marker), "{name}: {marker} was lost: {body}");
}
assert_eq!(
cached_prefix_violation(cached_parts(protocol), &body, PARTITION),
None,
"{name}, {variant}: {body}"
);
}
let Some(ResponseFormat::JsonSchema { schema: sent, .. }) = &req.response_format else {
panic!("no schema");
};
assert_eq!(sent, &schema, "{variant}");
}
}
#[test]
fn the_check_flags_worker_bytes_ahead_of_the_value() {
use crate::schema::ApiProtocol;
let unpartitioned = body(&AnthropicHandler, &worker_request());
assert!(cached_prefix_violation(
cached_parts(ApiProtocol::Anthropic),
&unpartitioned,
PARTITION
)
.is_some());
let late = json!({"system": format!("{}{PARTITION}", "W".repeat(300))});
assert_eq!(
cached_prefix_violation(cached_parts(ApiProtocol::Anthropic), &late, PARTITION),
Some("system")
);
let mut req = worker_request();
req.response_format = Some(ResponseFormat::JsonSchema {
schema: json!({"$defs": big_defs()}),
strict: false,
name: Some("WORKER_CHOSEN_NAME_".repeat(20)),
});
apply_cache_partition(&mut req, PARTITION);
if let Some(ResponseFormat::JsonSchema { name, .. }) = req.response_format.as_mut() {
*name = Some("WORKER_CHOSEN_NAME_".repeat(20));
}
let sent = body(&OpenAiHandler, &req);
assert_eq!(
cached_prefix_violation(cached_parts(ApiProtocol::OpenAiCompat), &sent, PARTITION),
Some("response_format")
);
let tiny = json!({"tools": [], "system": format!("{PARTITION} WORKER")});
assert_eq!(
cached_prefix_violation(cached_parts(ApiProtocol::Anthropic), &tiny, PARTITION),
None
);
let over = json!({"tools": ["W".repeat(CACHE_VALUE_WITHIN_BYTES)], "system": PARTITION});
assert_eq!(
cached_prefix_violation(cached_parts(ApiProtocol::Anthropic), &over, PARTITION),
Some("tools")
);
}
#[test]
fn a_request_with_no_value_is_flagged_in_every_format() {
let req = GenerateRequest {
prompt: "WORKER_PROMPT".into(),
context: Some("WORKER_SYS".into()),
messages: Some(vec![Message::User {
content: "WORKER_CONVERSATION ".repeat(2_000),
}]),
..Default::default()
};
for (name, protocol, handler) in handlers() {
let sent = body(handler.as_ref(), &req);
assert!(
cached_prefix_violation(cached_parts(protocol), &sent, PARTITION).is_some(),
"{name}: {sent}"
);
}
}
#[test]
fn empty_tools_and_a_json_object_format_pass_in_every_format() {
for (tools, format) in [
(Some(vec![]), None),
(None, Some(ResponseFormat::JsonObject)),
(Some(vec![]), Some(ResponseFormat::JsonObject)),
] {
let mut req = worker_request();
req.tools = tools.clone();
req.response_format = format.clone();
apply_cache_partition(&mut req, PARTITION);
for (name, protocol, handler) in handlers() {
let sent = body(handler.as_ref(), &req);
assert_eq!(
cached_prefix_violation(cached_parts(protocol), &sent, PARTITION),
None,
"{name}, tools {tools:?}, format {format:?}: {sent}"
);
assert!(contains(&sent, PARTITION), "{name}: {sent}");
}
}
}
#[tokio::test]
async fn a_body_that_breaks_the_guarantee_is_refused_before_it_is_sent() {
use crate::schema::ApiProtocol;
let unpartitioned = body(&AnthropicHandler, &worker_request());
let parts = cached_parts(ApiProtocol::Anthropic);
assert!(enforce_cached_prefix(parts, &unpartitioned).is_ok());
let refused = scope_cache_partition(Some(PARTITION.to_string()), async {
enforce_cached_prefix(parts, &unpartitioned)
})
.await;
assert!(matches!(
refused,
Err(crate::InferenceError::UnsupportedMode { .. })
));
let mut req = worker_request();
apply_cache_partition(&mut req, PARTITION);
let partitioned = body(&AnthropicHandler, &req);
let sent = scope_cache_partition(Some(PARTITION.to_string()), async {
enforce_cached_prefix(parts, &partitioned)
})
.await;
assert!(sent.is_ok());
}
#[test]
fn system_messages_fold_into_the_system_text_after_the_value() {
let mut req = worker_request();
apply_cache_partition(&mut req, PARTITION);
assert!(req
.messages
.as_ref()
.unwrap()
.iter()
.all(|message| !matches!(message, Message::System { .. })));
assert_eq!(
req.context.as_deref(),
Some(format!("{PARTITION}\n\nWORKER_SYS\n\nWORKER_CTX").as_str())
);
}
}