car-inference 0.60.0

//! Prompt-cache separation for workers (docs/proposals/orchestrator-workers.md,
//! I1 and I5).
//!
//! A provider caches prompt prefixes per account, so two workers of different
//! labels on one CAR would share cache entries whenever their prefixes match,
//! and a cache hit (a cheaper, faster first token, or the cache counts in the
//! usage) would tell one label that another sent the same prefix. CAR puts a
//! short per-label value, derived from the per-CAR secret and the worker's full
//! label set with its epochs, at the start of every part of a worker's request
//! that a provider caches ahead of the conversation, and
//! [`cached_prefix_violation`] checks it there on the serialized request body
//! itself. Cache entries then never match across labels, and caching inside
//! one label still works.
//!
//! **What the guarantee requires.** A provider serves a cached prefix only when
//! it matches at least its minimum cacheable length: 1,024 tokens for OpenAI,
//! Anthropic and Gemini, more for some models. So two labels can never share a
//! cache entry if no run of request bytes long enough to be that many tokens
//! comes before the label's value. CAR holds each part a provider caches ahead
//! of the conversation to one of two shapes: the value starts within its first
//! [`CACHE_VALUE_WITHIN_BYTES`] bytes, or the whole part is no longer than
//! that (an empty `tools` list, a `{"type":"json_object"}` format: nothing
//! there to put the value in, and nothing a worker could fill). And at least
//! one of those parts must hold the value: in some formats the conversation
//! is not a checked part, so a request whose short parts all lack the value
//! could still put a long conversation in the cache. Whatever order
//! a provider caches the parts in, at most three parts' worth of those bytes,
//! 768, can come before the value, and 768 bytes cannot be 1,024 tokens. The
//! parts, per wire format, in the order the provider caches them:
//!
//! | Format | Parts checked |
//! |---|---|
//! | OpenAI chat, OpenRouter, Azure OpenAI | `tools`, `response_format`, `messages` (system message first) |
//! | OpenAI Responses | `tools`, `text` (the structured-output format), `instructions` |
//! | Anthropic | `tools`, `system` (`messages` come after both) |
//! | Google, Vertex | `tools`, `systemInstruction` |
//! | Bedrock | `toolConfig`, `system` |
//! | Parslee | `tools`, `input` (system item first) |
//!
//! Serialized JSON puts object keys in sorted order (serde_json's
//! `preserve_order` is off in the daemon), so within a part the first bytes
//! are its alphabetically first key. That decides where the value goes:
//!
//! - **System text** is a string, so the value simply starts it.
//! - **Tools**: the CAR-defined first tool's description starts with the value,
//!   and `description` is its first key (Anthropic) or its function's first key
//!   (OpenAI, Bedrock, Gemini).
//! - **Structured output**: OpenAI caches the schema ahead of the system
//!   message. Its first serialized field is `json_schema.name` (`format.name`
//!   on Responses), ahead of the schema, so CAR sets that name to the value. The
//!   worker's schema is left exactly as it was sent: `$defs`, `$ref`,
//!   `definitions`, a root `anyOf` and the like are all accepted, and none of
//!   them comes before the value. Gemini does not cache `generationConfig`
//!   (context caching covers system instruction, tools and contents), and
//!   Anthropic and Bedrock take no response schema.

use crate::tasks::generate::{GenerateRequest, ResponseFormat};
use car_inference_types::Message;

/// Put `partition` first in every cached position of `req`.
///
/// - **System text.** The request's `context`, and every system message, are
///   folded into one system text that starts with `partition`. A system
///   message is removed from `messages`, so no handler can place it ahead of
///   the value.
/// - **Tools.** A CAR-defined tool named `partition` goes first. Its
///   description tells the model it is not a tool to call.
/// - **Structured output.** A JSON schema's `name` is set to `partition`. The
///   schema itself is untouched, so the output does not change.
///
/// `context_stable_prefix` is cleared: a cache breakpoint is placed only after
/// the whole system text, never before the value.
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() {
        // Exhaustive on purpose: a new system-like variant must be decided
        // here, or it could sit ahead of the value.
        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());
    }
}

/// How far into each cached part the label's value may start, in serialized
/// bytes. Well under any provider's minimum cacheable prefix (1,024 tokens).
pub const CACHE_VALUE_WITHIN_BYTES: usize = 256;

/// The parts of a request body a provider caches ahead of the conversation,
/// by wire format (see the module documentation).
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"],
    }
}

/// The parts of a Parslee gateway request body that are cached ahead of the
/// conversation.
pub const PARSLEE_CACHED_PARTS: &[&str] = &["tools", "input"];

/// Where the guarantee fails in `body` (see the module documentation), or
/// `None` when it holds: the first of `parts` that is longer than
/// [`CACHE_VALUE_WITHIN_BYTES`] serialized bytes without `partition` starting
/// within that many, or [`NO_PART_HOLDS_THE_VALUE`] when none of them holds
/// it there. Checked on the bytes the request is sent as.
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)
}

/// What [`cached_prefix_violation`] reports when no cached part holds the
/// value at all.
pub const NO_PART_HOLDS_THE_VALUE: &str = "no cached part holds the value";

tokio::task_local! {
    /// The per-label value of the worker whose call runs in this scope.
    static CACHE_PARTITION: Option<String>;
}

/// Run a model call with a worker's per-label value, so every request body
/// built for it is checked before it is sent ([`enforce_cached_prefix`]).
/// `None` (not a worker's call) checks nothing.
pub async fn scope_cache_partition<F: std::future::Future>(
    partition: Option<String>,
    future: F,
) -> F::Output {
    CACHE_PARTITION.scope(partition, future).await
}

/// Refuse to send a worker's request whose label value does not lead every
/// cached part (`parts`). Outside a worker's call this checks nothing. The
/// refusal is raised before dispatch, so it is never billed.
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()
        }
    }

    /// The body a handler sends, built the way `remote.rs` builds it.
    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,
        }
    }

    /// A large block of the worker's own bytes, far longer than the bound.
    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)
    }

    /// The schema shapes workers send, none of which may be refused.
    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),
            ),
        ]
    }

    /// The deterministic half of `worker_prompt_cache_does_not_cross_labels`,
    /// on the serialized request bytes: in every protocol CAR speaks, and for
    /// every schema shape a worker sends, the per-label value starts within the
    /// bound in every part the provider caches, the worker's own system text,
    /// tools, schema and messages all still arrive, and the worker's schema is
    /// sent exactly as it was given.
    #[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}"
                );
            }
            // The schema is the worker's, byte for byte.
            let Some(ResponseFormat::JsonSchema { schema: sent, .. }) = &req.response_format else {
                panic!("no schema");
            };
            assert_eq!(sent, &schema, "{variant}");
        }
    }

    /// The check sees what it is meant to: without the value, or with the
    /// worker's own name or schema ahead of it, a cached part is flagged.
    #[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());
        // The value is there, but past the bound.
        let late = json!({"system": format!("{}{PARTITION}", "W".repeat(300))});
        assert_eq!(
            cached_prefix_violation(cached_parts(ApiProtocol::Anthropic), &late, PARTITION),
            Some("system")
        );
        // The worker's own (long) schema name kept ahead of the value.
        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")
        );
        // A part with no value and no room for one is let through only while
        // it is no longer than the bound.
        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")
        );
    }

    /// T3: a request whose cached parts are all short and none holds the
    /// value is refused in every format, even where its long conversation is
    /// not a checked part.
    #[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}"
            );
        }
    }

    /// B1: a worker request with an empty tools list or a plain JSON-object
    /// format, parts with no slot for the value and nothing a worker can fill,
    /// passes the check in every format that sends them, and still carries the
    /// value in its system text.
    #[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}");
            }
        }
    }

    /// Enforcement: inside a worker's call a body that breaks the guarantee is
    /// refused before it is sent; outside one nothing is checked.
    #[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());
    }

    /// A system message never survives as a message a handler could place
    /// ahead of the value.
    #[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())
        );
    }
}