tool-result-cache 0.1.0

//! # tool-result-cache
//!
//! Content-addressable LRU cache for LLM agent tool calls.
//!
//! When an agent calls `search_web(q="anthropic prompt cache")` three times
//! in a row, it is wasting tokens, money, and time. `ToolCache` is a small
//! LRU cache keyed on a stable hash of `(tool_name, args)` with optional TTL
//! per entry and optional max size.
//!
//! ## Quick example
//!
//! ```
//! use std::time::Duration;
//! use serde_json::json;
//! use tool_result_cache::ToolCache;
//!
//! let mut cache: ToolCache<String> = ToolCache::new()
//!     .with_capacity(128)
//!     .with_ttl(Duration::from_secs(300));
//!
//! let args = json!({"q": "anthropic prompt cache"});
//! let value = cache
//!     .get_or_set("search_web", &args, || "result-a".to_string())
//!     .clone();
//! assert_eq!(value, "result-a");
//!
//! // Second call hits the cache and never invokes the closure.
//! let again = cache
//!     .get_or_set("search_web", &args, || panic!("would not run"))
//!     .clone();
//! assert_eq!(again, "result-a");
//! ```
//!
//! ## Key canonicalization
//!
//! Arguments are JSON values (via [`serde_json::Value`]). Object keys are
//! sorted recursively before hashing, so `{"a": 1, "b": 2}` and
//! `{"b": 2, "a": 1}` hit the same entry.
//!
//! ## LRU implementation
//!
//! The cache stores entries in a `HashMap<String, Entry>` and tracks recency
//! in a `Vec<String>` ordered from oldest at index 0 to most-recently used at
//! the tail. Capacity-based eviction pops index 0; `_touch` removes the key
//! from its current position and pushes it onto the tail. The internal
//! ordering vector keeps this dependency-light (only `serde_json` is pulled
//! in), at the cost of O(n) `Vec::remove` on touch where `n` is the number
//! of cached entries. For the typical agent-side cache size (a few hundred
//! to a few thousand entries) this is fine.

#![deny(missing_docs)]

use std::collections::HashMap;
use std::time::{Duration, Instant};

use serde_json::Value;

mod sha2;
use sha2::SimpleSha256;

/// Build a stable cache key from a tool name and JSON arguments.
///
/// The result is a 64-character lowercase hex SHA-256 of
/// `tool_name + "\0" + canonical_json(args)`. JSON object keys are sorted
/// recursively before hashing so that two semantically equal argument bags
/// hash to the same key regardless of key insertion order.
pub fn make_key(tool_name: &str, args: &Value) -> String {
    let mut buf = String::new();
    canonicalize_into(args, &mut buf);
    let mut hasher = SimpleSha256::new();
    hasher.update(tool_name.as_bytes());
    hasher.update(b"\0");
    hasher.update(buf.as_bytes());
    hex_lower(&hasher.finalize())
}

/// Snapshot of [`ToolCache`] counters at a point in time.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CacheStats {
    /// Number of `get` / `get_or_set` calls that found a live entry.
    pub hits: u64,
    /// Number of `get` / `get_or_set` calls that did not find a live entry.
    pub misses: u64,
    /// Number of entries dropped because the cache exceeded its capacity.
    pub evictions: u64,
    /// Number of entries dropped because their TTL elapsed.
    pub expirations: u64,
}

impl CacheStats {
    fn zero() -> Self {
        Self {
            hits: 0,
            misses: 0,
            evictions: 0,
            expirations: 0,
        }
    }
}

type ClockFn = Box<dyn Fn() -> Instant + Send + Sync + 'static>;

struct Entry<V> {
    value: V,
    expires_at: Option<Instant>,
}

/// LRU cache for tool results, keyed by `(tool_name, args)`.
///
/// `V` is the cached value type. Use [`ToolCache::new`], optionally chaining
/// [`ToolCache::with_capacity`] and [`ToolCache::with_ttl`] to configure the
/// cache:
///
/// ```
/// use std::time::Duration;
/// use tool_result_cache::ToolCache;
///
/// let _cache: ToolCache<String> = ToolCache::new()
///     .with_capacity(256)
///     .with_ttl(Duration::from_secs(60));
/// ```
pub struct ToolCache<V> {
    max_size: usize,
    default_ttl: Option<Duration>,
    clock: ClockFn,
    data: HashMap<String, Entry<V>>,
    /// LRU order: index 0 is least-recently used, last index is most-recent.
    order: Vec<String>,
    stats: CacheStats,
}

impl<V> Default for ToolCache<V> {
    fn default() -> Self {
        Self::new()
    }
}

impl<V> ToolCache<V> {
    /// Create a new cache with a default capacity of 1024 and no TTL.
    ///
    /// Capacity-based eviction only kicks in once the cache holds more than
    /// `max_size` entries. Use [`ToolCache::with_capacity`] or
    /// [`ToolCache::with_ttl`] to override.
    pub fn new() -> Self {
        Self {
            max_size: 1024,
            default_ttl: None,
            clock: Box::new(Instant::now),
            data: HashMap::new(),
            order: Vec::new(),
            stats: CacheStats::zero(),
        }
    }

    /// Set the maximum number of entries; `0` disables capacity-based eviction.
    pub fn with_capacity(mut self, max_size: usize) -> Self {
        self.max_size = max_size;
        self
    }

    /// Set the default TTL for new entries. Per-call TTLs override this value.
    pub fn with_ttl(mut self, ttl: Duration) -> Self {
        self.default_ttl = Some(ttl);
        self
    }

    /// Replace the cache clock. Default is [`Instant::now`]. Intended for tests.
    pub fn with_clock<F>(mut self, clock: F) -> Self
    where
        F: Fn() -> Instant + Send + Sync + 'static,
    {
        self.clock = Box::new(clock);
        self
    }

    /// Number of cached entries currently held.
    pub fn len(&self) -> usize {
        self.data.len()
    }

    /// `true` if the cache currently holds no entries.
    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    /// Number of `get` / `get_or_set` calls that found a live entry.
    pub fn hits(&self) -> u64 {
        self.stats.hits
    }

    /// Number of `get` / `get_or_set` calls that did not find a live entry.
    pub fn misses(&self) -> u64 {
        self.stats.misses
    }

    /// Number of entries dropped because the cache exceeded its capacity.
    pub fn evictions(&self) -> u64 {
        self.stats.evictions
    }

    /// Number of entries dropped because their TTL elapsed.
    pub fn expirations(&self) -> u64 {
        self.stats.expirations
    }

    /// Snapshot of all counters at this point in time. The snapshot is a copy
    /// and will not change as the cache continues to be used.
    pub fn stats(&self) -> CacheStats {
        self.stats
    }

    /// Drop all entries and reset statistics.
    pub fn clear(&mut self) {
        self.data.clear();
        self.order.clear();
        self.stats = CacheStats::zero();
    }

    /// Look up a cached value by tool name + args. Updates hit/miss stats.
    pub fn get(&mut self, tool_name: &str, args: &Value) -> Option<&V> {
        let key = make_key(tool_name, args);
        if self.touch(&key) {
            self.stats.hits += 1;
            self.data.get(&key).map(|e| &e.value)
        } else {
            self.stats.misses += 1;
            None
        }
    }

    /// Insert or replace a cached value using the cache's default TTL.
    pub fn set(&mut self, tool_name: &str, args: &Value, value: V) {
        self.insert(make_key(tool_name, args), value, self.default_ttl);
    }

    /// Insert or replace a cached value with a per-call TTL override.
    pub fn set_with_ttl(&mut self, tool_name: &str, args: &Value, value: V, ttl: Duration) {
        self.insert(make_key(tool_name, args), value, Some(ttl));
    }

    /// Return the cached value if present, else call `compute` and cache it.
    pub fn get_or_set<F>(&mut self, tool_name: &str, args: &Value, compute: F) -> &V
    where
        F: FnOnce() -> V,
    {
        self.get_or_set_with_ttl_opt(tool_name, args, compute, None)
    }

    /// Same as [`ToolCache::get_or_set`] but with a per-call TTL override.
    pub fn get_or_set_with_ttl<F>(
        &mut self,
        tool_name: &str,
        args: &Value,
        compute: F,
        ttl: Duration,
    ) -> &V
    where
        F: FnOnce() -> V,
    {
        self.get_or_set_with_ttl_opt(tool_name, args, compute, Some(ttl))
    }

    fn get_or_set_with_ttl_opt<F>(
        &mut self,
        tool_name: &str,
        args: &Value,
        compute: F,
        ttl: Option<Duration>,
    ) -> &V
    where
        F: FnOnce() -> V,
    {
        let key = make_key(tool_name, args);
        if self.touch(&key) {
            self.stats.hits += 1;
            return &self
                .data
                .get(&key)
                .expect("touch returned true so key must be present")
                .value;
        }
        self.stats.misses += 1;
        let value = compute();
        let effective_ttl = ttl.or(self.default_ttl);
        self.insert(key.clone(), value, effective_ttl);
        &self
            .data
            .get(&key)
            .expect("just inserted this key, must be present")
            .value
    }

    /// Drop a single entry. Returns `true` if it was present.
    pub fn invalidate(&mut self, tool_name: &str, args: &Value) -> bool {
        let key = make_key(tool_name, args);
        if self.data.remove(&key).is_some() {
            self.remove_from_order(&key);
            true
        } else {
            false
        }
    }

    // ---- internals -------------------------------------------------------

    fn insert(&mut self, key: String, value: V, ttl: Option<Duration>) {
        let expires_at = ttl.map(|d| (self.clock)() + d);
        // Remove any prior position in the order vec; HashMap::insert replaces
        // in place, but recency needs to reset to the tail.
        if self.data.contains_key(&key) {
            self.remove_from_order(&key);
        }
        self.data.insert(key.clone(), Entry { value, expires_at });
        self.order.push(key);
        self.evict_if_needed();
    }

    /// Returns `true` if the key is present and live (not expired). On a
    /// live hit, the key is moved to the tail of the LRU order. On an
    /// expired hit, the entry is dropped and the expirations counter is
    /// incremented; the caller still treats this as a miss.
    fn touch(&mut self, key: &str) -> bool {
        let Some(entry) = self.data.get(key) else {
            return false;
        };
        if let Some(expires_at) = entry.expires_at {
            if (self.clock)() >= expires_at {
                self.data.remove(key);
                self.remove_from_order(key);
                self.stats.expirations += 1;
                return false;
            }
        }
        // Promote to tail.
        self.remove_from_order(key);
        self.order.push(key.to_string());
        true
    }

    fn remove_from_order(&mut self, key: &str) {
        if let Some(pos) = self.order.iter().position(|k| k == key) {
            self.order.remove(pos);
        }
    }

    fn evict_if_needed(&mut self) {
        if self.max_size == 0 {
            return;
        }
        while self.data.len() > self.max_size {
            // Pop oldest. The `order` vec mirrors `data` and must be drained
            // in lock-step.
            if self.order.is_empty() {
                break;
            }
            let oldest = self.order.remove(0);
            if self.data.remove(&oldest).is_some() {
                self.stats.evictions += 1;
            }
        }
    }
}

/// Convenience: look up `(tool_name, args)` in `cache` or compute and store
/// `compute()`. Returns an owned clone of the cached value.
///
/// This is the Rust equivalent of the Python `cache.wrap(...)` decorator;
/// since Rust does not decorate functions, callers wrap a call site instead.
pub fn cached_call<V, F>(cache: &mut ToolCache<V>, tool_name: &str, args: &Value, compute: F) -> V
where
    V: Clone,
    F: FnOnce() -> V,
{
    cache.get_or_set(tool_name, args, compute).clone()
}

// ---- canonical JSON --------------------------------------------------------

fn canonicalize_into(value: &Value, out: &mut String) {
    match value {
        Value::Null => out.push_str("null"),
        Value::Bool(true) => out.push_str("true"),
        Value::Bool(false) => out.push_str("false"),
        Value::Number(n) => out.push_str(&n.to_string()),
        Value::String(s) => {
            // Reuse serde_json string escaping for correctness.
            let escaped = serde_json::to_string(s).expect("string serialization cannot fail");
            out.push_str(&escaped);
        }
        Value::Array(items) => {
            out.push('[');
            for (i, item) in items.iter().enumerate() {
                if i > 0 {
                    out.push(',');
                }
                canonicalize_into(item, out);
            }
            out.push(']');
        }
        Value::Object(map) => {
            // Sort object keys lexicographically. `serde_json::Map` iteration
            // order depends on the `preserve_order` feature, so we must sort
            // explicitly to make keys stable.
            let mut keys: Vec<&String> = map.keys().collect();
            keys.sort();
            out.push('{');
            for (i, k) in keys.iter().enumerate() {
                if i > 0 {
                    out.push(',');
                }
                let escaped =
                    serde_json::to_string(k.as_str()).expect("string serialization cannot fail");
                out.push_str(&escaped);
                out.push(':');
                canonicalize_into(&map[*k], out);
            }
            out.push('}');
        }
    }
}

fn hex_lower(bytes: &[u8]) -> String {
    const HEX: &[u8; 16] = b"0123456789abcdef";
    let mut s = String::with_capacity(bytes.len() * 2);
    for &b in bytes {
        s.push(HEX[(b >> 4) as usize] as char);
        s.push(HEX[(b & 0x0f) as usize] as char);
    }
    s
}