Optimizing LLM Context Caching at the Gateway Level: Implementing Prefix-Based Prompt Caching in a Rust API Proxy — code snippets

snippet-2.txt

use sha2::{Sha256, Digest};
use serde::{Serialize, Deserialize};

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ChatMessage {
    pub role: String,
    pub content: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ChatCompletionRequest {
    pub model: String,
    pub messages: Vec<ChatMessage>,
    #[serde(default)]
    pub temperature: Option<f32>,
    #[serde(default)]
    pub system: Option<String>,
}

pub fn generate_normalized_prefix_hash(req: &ChatCompletionRequest, prefix_message_count: usize) -> Option<String> {
    if req.messages.is_empty() && req.system.is_none() {
        return None;
    }
    
    let mut hasher = Sha256::new();
    
    // 1. Normalize and hash the system prompt if present
    if let Some(ref system_prompt) = req.system {
        let normalized_system = system_prompt
            .lines()
            .map(|line| line.trim())
            .filter(|line| !line.is_empty())
            .collect::<Vec<&str>>()
            .join("\n");
        hasher.update(b"system:");
        hasher.update(normalized_system.as_bytes());
    }
    
    // 2. Normalize and hash the designated prefix messages
    let take_count = std::cmp::min(req.messages.len(), prefix_message_count);
    let prefix_slice = &req.messages[..take_count];
    
    for (index, msg) in prefix_slice.iter().enumerate() {
        let normalized_content = msg.content
            .lines()
            .map(|line| line.trim())
            .filter(|line| !line.is_empty())
            .collect::<Vec<&str>>()
            .join("\n");
            
        hasher.update(format!("msg_{}:", index).as_bytes());
        hasher.update(msg.role.as_bytes());
        hasher.update(normalized_content.as_bytes());
        if let Some(ref name) = msg.name {
            hasher.update(name.as_bytes());
        }
    }
    
    // 3. Caching is model-bound, so the model name must form part of the hash
    hasher.update(req.model.as_bytes());
    
    Some(hex::encode(hasher.finalize()))
}

snippet-3.txt

use tiktoken_rs::cl100k_base;

pub struct CachingDecisionEngine {
    min_token_threshold: usize,
    max_token_threshold: usize,
}

impl CachingDecisionEngine {
    pub fn new(min_tokens: usize, max_tokens: usize) -> Self {
        Self {
            min_token_threshold: min_tokens,
            max_token_threshold: max_tokens,
        }
    }

    pub fn should_cache_prefix(&self, system_prompt: Option<&str>, messages: &[ChatMessage]) -> (bool, usize) {
        // Fast path approximation: average 4 characters per token. 
        // If string length is significantly lower than threshold, bypass tokenizer entirely.
        let mut approximate_chars = system_prompt.map(|s| s.len()).unwrap_or(0);
        for m in messages {
            approximate_chars += m.content.len();
        }
        
        let estimated_tokens = approximate_chars / 4;
        if estimated_tokens < (self.min_token_threshold - 200) {
            return (false, estimated_tokens);
        }

        // Slow path: precise token calculation
        let bpe = match cl100k_base() {
            Ok(engine) => engine,
            Err(_) => return (false, 0),
        };
        
        let mut combined_text = system_prompt.unwrap_or("").to_string();
        for msg in messages {
            combined_text.push_str(&msg.content);
        }
        
        let token_count = bpe.encode_with_special_tokens(&combined_text).len();
        let eligible = token_count >= self.min_token_threshold && token_count <= self.max_token_threshold;
        
        (eligible, token_count)
    }
}

snippet-4.txt

use serde_json::{Value, json};

pub fn inject_anthropic_cache_control(payload: &mut Value) -> Result<(), &'static str> {
    let payload_obj = payload.as_object_mut().ok_or("Payload is not a valid JSON object")?;

    // Option A: Inject Cache Control into System Prompt
    if let Some(system) = payload_obj.get_mut("system") {
        match system {
            Value::String(text) => {
                *system = json!([
                    {
                        "type": "text",
                        "text": text,
                        "cache_control": { "type": "ephemeral" }
                    }
                ]);
                return Ok(());
            }
            Value::Array(blocks) => {
                if let Some(last_block) = blocks.last_mut() {
                    if let Some(obj) = last_block.as_object_mut() {
                        obj.insert("cache_control".to_string(), json!({ "type": "ephemeral" }));
                        return Ok(());
                    }
                }
            }
            _ => return Err("Invalid format for system block"),
        }
    }

    // Option B: Fallback to injecting Cache Control into Message Array (System is absent or empty)
    if let Some(messages) = payload_obj.get_mut("messages").and_then(|m| m.as_array_mut()) {
        if messages.len() >= 2 {
            // Target the last message in our prefix window (e.g. index 1 or 2)
            let target_index = messages.len() - 1;
            if let Some(msg) = messages.get_mut(target_index) {
                if let Some(content) = msg.get_mut("content") {
                    match content {
                        Value::String(text) => {
                            *content = json!([
                                {
                                    "type": "text",
                                    "text": text,
                                    "cache_control": { "type": "ephemeral" }
                                }
                            ]);
                            return Ok(());
                        }
                        Value::Array(blocks) => {
                            if let Some(last_block) = blocks.last_mut() {
                                if let Some(obj) = last_block.as_object_mut() {
                                    obj.insert("cache_control".to_string(), json!({ "type": "ephemeral" }));
                                    return Ok(());
                                }
                            }
                        }
                        _ => return Err("Invalid content structure in message"),
                    }
                }
            }
        }
    }

    Err("Could not find suitable caching breakpoint block")
}

snippet-5.txt

use redis::AsyncCommands;

pub struct RedisTracker {
    client: redis::Client,
    hit_threshold: i64,
    window_seconds: i64,
}

impl RedisTracker {
    pub fn new(connection_str: &str, hit_threshold: i64, window_seconds: i64) -> redis::RedisResult<Self> {
        let client = redis::Client::open(connection_str)?;
        Ok(Self {
            client,
            hit_threshold,
            window_seconds,
        })
    }

    pub async fn track_and_should_cache(&self, hash: &str) -> Result<bool, redis::RedisError> {
        let mut conn = self.client.get_multiplexed_async_connection().await?;
        let key = format!("prefix_hits:{}", hash);
        
        // Atomically increment counter and reset expiration
        let (hits, _): (i64, ()) = redis::pipe()
            .incr(&key, 1)
            .expire(&key, self.window_seconds)
            .query_async(&mut conn)
            .await?;
            
        Ok(hits >= self.hit_threshold)
    }
}

snippet-6.txt

use axum::{
    body::Body,
    http::{Request, Response, StatusCode},
    response::IntoResponse,
};
use http_body_util::BodyExt;
use std::convert::Infallible;
use std::sync::Arc;

pub async fn handle_proxy_request(
    redis_tracker: Arc<RedisTracker>,
    decision_engine: Arc<CachingDecisionEngine>,
    client: hyper_util::client::legacy::Client<
        hyper_rustls::HttpsConnector<hyper_util::client::legacy::connect::HttpConnector>, 
        Body
    >,
    req: Request<Body>,
) -> Result<impl IntoResponse, Infallible> {
    let (parts, body) = req.into_parts();
    
    // Buffer body to inspect the JSON payload
    let bytes = match body.collect().await {
        Ok(b) => b.to_bytes(),
        Err(_) => return Ok(Response::builder().status(StatusCode::BAD_REQUEST).body(Body::empty()).unwrap()),
    };

    let mut payload: Value = match serde_json::from_slice(&bytes) {
        Ok(val) => val,
        Err(_) => return Ok(Response::builder().status(StatusCode::BAD_REQUEST).body(Body::from(bytes)).unwrap()),
    };

    // Extract structure for analysis
    let req_struct: ChatCompletionRequest = match serde_json::from_value(payload.clone()) {
        Ok(r) => r,
        Err(_) => return Ok(Response::builder().status(StatusCode::UNPROCESSABLE_ENTITY).body(Body::from("Invalid request format")).unwrap()),
    };

    let (eligible, _tokens) = decision_engine.should_cache_prefix(req_struct.system.as_deref(), &req_struct.messages);
    
    if eligible {
        if let Some(hash) = generate_normalized_prefix_hash(&req_struct, 2) {
            match redis_tracker.track_and_should_cache(&hash).await {
                Ok(true) => {
                    // Inject ephemeral caching markers into the request payload
                    if inject_anthropic_cache_control(&mut payload).is_ok() {
                        metrics::counter!("proxy.prompt_cache.injected", "model" => req_struct.model.clone()).increment(1);
                    }
                }
                Ok(false) => {
                    metrics::counter!("proxy.prompt_cache.cold_prefix", "model" => req_struct.model.clone()).increment(1);
                }
                Err(e) => {
                    tracing::error!("Redis tracking failure: {:?}", e);
                }
            }
        }
    }

    let serialized = serde_json::to_vec(&payload).unwrap();
    let api_key = match parts.headers.get("x-api-key") {
        Some(k) => k.to_str().unwrap_or(""),
        None => return Ok(Response::builder().status(StatusCode::UNAUTHORIZED).body(Body::from("Missing API Key")).unwrap()),
    };

    let mut upstream_req = Request::post("https://api.anthropic.com/v1/messages")
        .header("content-type", "application/json")
        .header("x-api-key", api_key)
        .header("anthropic-version", "2023-06-01")
        .body(Body::from(serialized))
        .unwrap();

    let upstream_res = match client.request(upstream_req).await {
        Ok(res) => res,
        Err(err) => {
            tracing::error!("Upstream routing failed: {:?}", err);
            return Ok(Response::builder().status(StatusCode::BAD_GATEWAY).body(Body::from("Upstream error")).unwrap());
        }
    };

    // Copy status and headers from upstream response to client response
    let mut response_builder = Response::builder().status(upstream_res.status());
    for (name, val) in upstream_res.headers() {
        response_builder = response_builder.header(name, val);
    }

    Ok(response_builder.body(upstream_res.into_body()).unwrap())
}

snippet-7.txt

use tokio::time::{timeout, Duration};

pub async fn safe_check_redis_cache_eligibility(
    tracker: &RedisTracker,
    hash: &str
) -> bool {
    let fallback_timeout = Duration::from_millis(15);
    
    match timeout(fallback_timeout, tracker.track_and_should_cache(hash)).await {
        Ok(Ok(should_cache)) => should_cache,
        Ok(Err(redis_err)) => {
            tracing::error!(
                target: "llm_caching_proxy",
                "Redis lookup failed: {:?}. Falling back to default routing.",
                redis_err
            );
            metrics::counter!("proxy.redis.errors", "error_type" => "redis_error").increment(1);
            false // Fail open: do not modify request payload
        }
        Err(_elapsed) => {
            tracing::warn!(
                target: "llm_caching_proxy",
                "Redis connection timed out after {}ms. Failing open.",
                fallback_timeout.as_millis()
            );
            metrics::counter!("proxy.redis.errors", "error_type" => "timeout").increment(1);
            false // Fail open
        }
    }
}