use std::collections::{BTreeMap, HashMap, HashSet};
use std::sync::{Arc, Mutex, PoisonError};
use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::value::RawValue;
use sha2::{Digest, Sha256};
use crate::completion::CacheCost;
pub type Clock = Arc<dyn Fn() -> u64 + Send + Sync>;
fn system_now() -> u64 {
#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
{
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_or(0, |elapsed| elapsed.as_secs())
}
#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
{
0
}
}
const MAX_TTL_SECS: u64 = 81 * 60;
const EXPIRY_MARGIN_SECS: u64 = 5;
const MIN_IDLE_SECS: u64 = 60;
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct AutoCache {
pub ttl: Duration,
pub implicit_floor: u64,
pub implicit_ceiling: u64,
pub min_tokens: u64,
pub min_gain: u64,
pub cached_ratio: f64,
pub storage_ratio_per_hour: f64,
}
impl Default for AutoCache {
fn default() -> Self {
Self {
ttl: Duration::from_secs(60 * 60),
implicit_floor: 4_096,
implicit_ceiling: 16_000,
min_tokens: 1_024,
min_gain: 256,
cached_ratio: 0.1,
storage_ratio_per_hour: 0.5 / 0.75,
}
}
}
impl AutoCache {
pub fn for_model(model: &str) -> Self {
Self::priced(
crate::catalog::lookup(super::PROVIDER_NAME, model)
.and_then(|spec| spec.pricing.as_ref()),
)
}
fn priced(pricing: Option<&crate::catalog::Pricing>) -> Self {
let ratio = pricing
.and_then(|pricing| Some((pricing.cache_read?, pricing.input)))
.filter(|(read, input)| *read >= 0.0 && *input > 0.0)
.map(|(read, input)| read / input);
let default = Self::default();
Self {
cached_ratio: ratio.unwrap_or(default.cached_ratio),
..default
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Lease {
pub name: String,
pub digest: String,
pub covers: usize,
pub tokens: u64,
pub expires_at: u64,
pub model: String,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum CacheEvent {
Created {
name: String,
tokens: u64,
estimated: u64,
covers: usize,
ttl_secs: u64,
at: u64,
},
Retired {
name: String,
at: u64,
},
Lost {
name: String,
at: u64,
},
Extended {
name: String,
expires_at: u64,
at: u64,
},
CreateFailed {
status: Option<u16>,
message: String,
at: u64,
},
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CreatedCache {
pub name: String,
pub tokens: u64,
pub at: u64,
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct CacheReport {
pub created: Vec<CreatedCache>,
pub reads: BTreeMap<String, u64>,
pub retired: Vec<String>,
pub lost: Vec<String>,
pub token_hours: f64,
pub live: Vec<Lease>,
}
impl From<&CacheReport> for CacheCost {
fn from(report: &CacheReport) -> Self {
Self {
cache_writes: report.created.iter().map(|created| created.tokens).sum(),
storage_token_hours: report.token_hours,
..Self::default()
}
}
}
#[derive(Clone, Debug, Default)]
struct Line {
premium: f64,
implicit: Option<f64>,
rolled_at: u64,
last_call: Option<u64>,
longest_gap: u64,
failed_at: Option<u64>,
}
#[derive(Clone, Debug)]
struct Life {
tokens: u64,
created: u64,
expires_at: u64,
ended: Option<u64>,
last_read: u64,
line_gap: u64,
}
#[derive(Default)]
pub(crate) struct Book {
pub(crate) leases: HashMap<String, Lease>,
lives: BTreeMap<String, Life>,
lineages: HashMap<String, HashSet<String>>,
lines: HashMap<String, Line>,
events: Vec<CacheEvent>,
report: CacheReport,
calibration: Option<f64>,
}
#[derive(Clone)]
pub struct CacheBook {
inner: Arc<Mutex<Book>>,
pub(crate) create: Arc<futures::lock::Mutex<()>>,
policy: AutoCache,
clock: Clock,
pub(crate) display_prefix: Arc<str>,
}
impl std::fmt::Debug for CacheBook {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("CacheBook")
.field("policy", &self.policy)
.field("display_prefix", &self.display_prefix)
.finish_non_exhaustive()
}
}
impl CacheBook {
pub fn new(policy: AutoCache) -> Self {
Self {
inner: Arc::default(),
create: Arc::default(),
policy,
clock: Arc::new(system_now),
display_prefix: Arc::from("rig-cache-"),
}
}
pub fn with_clock(mut self, clock: impl Fn() -> u64 + Send + Sync + 'static) -> Self {
self.clock = Arc::new(clock);
self
}
pub fn with_display_prefix(mut self, prefix: &str) -> Self {
self.display_prefix = Arc::from(prefix);
self
}
pub fn policy(&self) -> AutoCache {
self.policy
}
pub(crate) fn book(&self) -> std::sync::MutexGuard<'_, Book> {
self.inner.lock().unwrap_or_else(PoisonError::into_inner)
}
pub(crate) fn now(&self) -> u64 {
(self.clock)()
}
pub fn events(&self) -> Vec<CacheEvent> {
self.book().events.clone()
}
pub fn leases(&self) -> Vec<Lease> {
let mut leases: Vec<Lease> = self.book().leases.values().cloned().collect();
leases.sort_by(|a, b| a.name.cmp(&b.name));
leases
}
pub fn restore(&self, leases: Vec<Lease>) {
let now = self.now();
let mut book = self.book();
for lease in leases {
book.lives.entry(lease.name.clone()).or_insert(Life {
tokens: lease.tokens,
created: now,
expires_at: lease.expires_at,
ended: None,
last_read: now,
line_gap: 0,
});
book.leases.insert(lease.digest.clone(), lease);
}
}
pub fn report(&self) -> CacheReport {
let now = self.now();
let book = self.book();
let mut report = book.report.clone();
report.token_hours = book
.lives
.values()
.map(|life| {
let end = life.ended.unwrap_or(now).min(life.expires_at);
life.tokens as f64 * end.saturating_sub(life.created) as f64 / 3600.0
})
.sum();
report.live = book.leases.values().cloned().collect();
report.live.sort_by(|a, b| a.name.cmp(&b.name));
report
}
pub(crate) fn retired(&self, lease: &Lease) {
let now = self.now();
let mut book = self.book();
book.leases.remove(&lease.digest);
end_life(&mut book, &lease.name, now);
book.events.push(CacheEvent::Retired {
name: lease.name.clone(),
at: now,
});
book.report.retired.push(lease.name.clone());
tracing::info!(target: "gemini.cache.retired", name = %lease.name);
}
pub(crate) fn lost(&self, lease: &Lease) {
let now = self.now();
let mut book = self.book();
book.leases.remove(&lease.digest);
end_life(&mut book, &lease.name, now);
book.events.push(CacheEvent::Lost {
name: lease.name.clone(),
at: now,
});
book.report.lost.push(lease.name.clone());
tracing::info!(target: "gemini.cache.lost", name = %lease.name);
}
}
pub(crate) fn end_life(book: &mut Book, name: &str, now: u64) {
if let Some(life) = book.lives.get_mut(name)
&& life.ended.is_none()
{
life.ended = Some(now);
}
}
pub(crate) fn short_digest(digest: &str) -> &str {
digest.get(..40).unwrap_or(digest)
}
const PREFIX_KEYS: [&str; 3] = ["systemInstruction", "tools", "toolConfig"];
pub(crate) struct Parsed {
pub(crate) rest: Vec<(String, Box<RawValue>)>,
pub(crate) prefix: [Option<Box<RawValue>>; 3],
pub(crate) contents: Vec<Box<RawValue>>,
pub(crate) has_cached_content: bool,
}
pub(crate) fn parse(bytes: &[u8]) -> Option<Parsed> {
let fields: serde_json::Map<String, serde_json::Value> = serde_json::from_slice(bytes).ok()?;
let raw: HashMap<String, Box<RawValue>> = serde_json::from_slice(bytes).ok()?;
let mut parsed = Parsed {
rest: Vec::new(),
prefix: [None, None, None],
contents: Vec::new(),
has_cached_content: false,
};
for key in fields.keys() {
let value = raw.get(key)?.clone();
if key == "contents" {
parsed.contents = serde_json::from_str(value.get()).ok()?;
} else if let Some(slot) = PREFIX_KEYS.iter().position(|prefix| prefix == key) {
if value.get() != "null"
&& let Some(field) = parsed.prefix.get_mut(slot)
{
*field = Some(value);
}
} else {
if key == "cachedContent" && value.get() != "null" {
parsed.has_cached_content = true;
}
parsed.rest.push((key.clone(), value));
}
}
Some(parsed)
}
fn hex(bytes: &[u8]) -> String {
use std::fmt::Write;
bytes.iter().fold(String::new(), |mut out, byte| {
let _ = write!(out, "{byte:02x}");
out
})
}
fn canonical(fragment: &str) -> std::borrow::Cow<'_, str> {
serde_json::from_str::<serde_json::Value>(fragment)
.map_or(std::borrow::Cow::Borrowed(fragment), |value| {
std::borrow::Cow::Owned(crate::json_utils::to_canonical_string(&value))
})
}
pub(crate) fn digests(model: &str, parsed: &Parsed) -> Vec<String> {
let mut hasher = Sha256::new();
hasher.update(model.as_bytes());
for field in &parsed.prefix {
hasher.update([0u8]);
if let Some(raw) = field {
hasher.update(canonical(raw.get()).as_bytes());
}
}
let mut state = hasher.finalize().to_vec();
let mut out = vec![hex(&state)];
for content in &parsed.contents {
let mut hasher = Sha256::new();
hasher.update(&state);
hasher.update(canonical(content.get()).as_bytes());
state = hasher.finalize().to_vec();
out.push(hex(&state));
}
out
}
fn estimate(json: &str) -> u64 {
let Ok(mut value) = serde_json::from_str::<serde_json::Value>(json) else {
return (json.len() / 4) as u64;
};
fn scrub(value: &mut serde_json::Value) {
match value {
serde_json::Value::Object(map) => {
map.shift_remove("thoughtSignature");
map.values_mut().for_each(scrub);
}
serde_json::Value::Array(items) => items.iter_mut().for_each(scrub),
_ => {}
}
}
scrub(&mut value);
(value.to_string().len() / 4) as u64
}
pub(crate) fn is_user_text(raw: &RawValue) -> bool {
#[derive(Deserialize)]
struct Content {
role: Option<String>,
#[serde(default)]
parts: Vec<serde_json::Map<String, serde_json::Value>>,
}
serde_json::from_str::<Content>(raw.get()).is_ok_and(|content| {
content.role.as_deref() == Some("user")
&& content.parts.iter().any(|part| part.contains_key("text"))
&& !content
.parts
.iter()
.any(|part| part.contains_key("functionResponse"))
})
}
pub(crate) fn stripped(parsed: &Parsed, lease: &Lease) -> Option<Vec<u8>> {
let mut out: Vec<(String, &RawValue)> = Vec::new();
let name = serde_json::value::to_raw_value(&lease.name).ok()?;
let contents = serde_json::value::to_raw_value(&parsed.contents.get(lease.covers..)?).ok()?;
out.push(("cachedContent".to_owned(), &name));
out.push(("contents".to_owned(), &contents));
for (key, value) in &parsed.rest {
if key != "cachedContent" {
out.push((key.clone(), value));
}
}
let mut bytes = b"{".to_vec();
for (index, (key, value)) in out.iter().enumerate() {
if index > 0 {
bytes.push(b',');
}
bytes.extend(serde_json::to_vec(key).ok()?);
bytes.push(b':');
bytes.extend(value.get().as_bytes());
}
bytes.push(b'}');
Some(bytes)
}
pub(crate) fn cache_body(
model: &str,
parsed: &Parsed,
covers: usize,
display_name: &str,
ttl_secs: u64,
) -> Option<Vec<u8>> {
let mut bytes = b"{".to_vec();
let mut field = |key: &str, value: &str| {
if bytes.len() > 1 {
bytes.push(b',');
}
bytes.extend(format!("\"{key}\":").as_bytes());
bytes.extend(value.as_bytes());
};
field(
"model",
&serde_json::to_string(&format!("models/{model}")).ok()?,
);
field("displayName", &serde_json::to_string(display_name).ok()?);
field("ttl", &format!("\"{ttl_secs}s\""));
for (key, value) in PREFIX_KEYS.iter().zip(&parsed.prefix) {
if let Some(value) = value {
field(key, value.get());
}
}
if covers > 0 {
let contents = serde_json::value::to_raw_value(&parsed.contents.get(..covers)?).ok()?;
field("contents", contents.get());
}
bytes.push(b'}');
Some(bytes)
}
pub(crate) struct Plan {
pub(crate) read: Option<Lease>,
pub(crate) create: Option<Create>,
pub(crate) retire: Vec<Lease>,
pub(crate) extend: Option<(Lease, u64)>,
pub(crate) line: String,
pub(crate) coverable: u64,
}
pub(crate) struct Create {
pub(crate) covers: usize,
pub(crate) digest: String,
pub(crate) ttl_secs: u64,
estimate: u64,
expected: u64,
pub(crate) replaces: Option<Lease>,
}
impl CacheBook {
pub(crate) fn plan(
&self,
model: &str,
d: &[String],
parsed: &Parsed,
roll_allowed: bool,
) -> Plan {
let p = self.policy;
let now = self.now();
let n = parsed.contents.len();
let at = |k: usize| d.get(k).cloned().unwrap_or_default();
let mut book = self.book();
let calibration = book.calibration.unwrap_or(1.0);
let prefix_estimate: u64 = parsed
.prefix
.iter()
.flatten()
.map(|raw| estimate(raw.get()))
.sum();
let content_estimates: Vec<u64> = parsed
.contents
.iter()
.map(|raw| estimate(raw.get()))
.collect();
if n >= 1 {
book.lineages.entry(at(0)).or_default().insert(at(1));
}
let shared = book
.lineages
.get(&at(0))
.is_some_and(|seen| seen.len() >= 2);
let idle: Vec<Lease> = book
.leases
.values()
.filter(|lease| lease.covers > 0 && lease.model == model)
.filter(|lease| {
book.lives.get(&lease.name).is_some_and(|life| {
now.saturating_sub(life.last_read)
> MIN_IDLE_SECS.max(life.line_gap.saturating_mul(3))
})
})
.cloned()
.collect();
let read = (0..n.max(1))
.rev()
.find_map(|k| {
book.leases
.get(&at(k))
.filter(|lease| lease.model == model)
.filter(|lease| lease.expires_at > now + EXPIRY_MARGIN_SECS)
.cloned()
})
.filter(|lease| !idle.iter().any(|gone| gone.name == lease.name));
let covered = read.as_ref().map_or(0, |lease| lease.covers);
let tail: u64 = content_estimates
.get(covered..n.saturating_sub(1).max(covered))
.unwrap_or_default()
.iter()
.sum();
let coverable = match &read {
Some(lease) => lease.tokens + (tail as f64 * calibration) as u64,
None => ((prefix_estimate + tail) as f64 * calibration) as u64,
};
let line_key = read
.as_ref()
.filter(|lease| lease.covers > 0)
.map_or_else(|| at(1.min(n)), |lease| lease.digest.clone());
let prefix_leased = book.leases.contains_key(&at(0));
let line = book.lines.entry(line_key.clone()).or_insert_with(|| Line {
rolled_at: now,
..Line::default()
});
if let Some(last) = line.last_call {
line.longest_gap = line.longest_gap.max(now.saturating_sub(last));
}
line.last_call = Some(now);
let longest_gap = line.longest_gap;
let ttl_secs = p
.ttl
.as_secs()
.max(longest_gap.saturating_mul(2).min(MAX_TTL_SECS))
.max(1);
let mut plan = Plan {
read: read.clone(),
create: None,
retire: idle,
extend: None,
line: line_key,
coverable,
};
let implicit_serving = line.implicit.is_none_or(|ratio| ratio >= 0.5);
if coverable >= p.implicit_ceiling && implicit_serving {
plan.read = None;
return plan;
}
let failed_below = line.failed_at.is_some_and(|at| coverable < at + p.min_gain);
if read.is_none()
&& shared
&& !prefix_leased
&& (prefix_estimate as f64 * calibration) as u64 >= p.min_tokens
&& !failed_below
{
plan.create = Some(Create {
covers: 0,
digest: at(0),
ttl_secs,
estimate: prefix_estimate,
expected: (prefix_estimate as f64 * calibration) as u64,
replaces: None,
});
} else if n >= 2 && roll_allowed && !failed_below {
let cached = read.as_ref().map_or(0, |lease| lease.tokens);
let gain = coverable.saturating_sub(cached);
let held_hours = now.saturating_sub(line.rolled_at).max(60) as f64 / 3600.0;
let cost = coverable as f64 * (1.0 + p.storage_ratio_per_hour * held_hours);
if gain >= p.min_gain && coverable >= p.min_tokens && line.premium >= cost {
plan.create = Some(Create {
covers: n - 1,
digest: at(n - 1),
ttl_secs,
estimate: prefix_estimate
+ content_estimates
.get(..n - 1)
.unwrap_or_default()
.iter()
.sum::<u64>(),
expected: coverable,
replaces: read.clone().filter(|lease| lease.covers > 0),
});
}
}
if plan.create.is_none()
&& let Some(lease) = &read
&& longest_gap > 0
&& lease.expires_at.saturating_sub(now) < longest_gap.saturating_mul(3) / 2
{
plan.extend = Some((lease.clone(), ttl_secs));
}
plan
}
pub(crate) fn observe(&self, line: &str, read: Option<&str>, coverable: u64, cached: u64) {
let p = self.policy;
let mut book = self.book();
if let Some(name) = read {
*book.report.reads.entry(name.to_owned()).or_default() += 1;
}
let entry = book.lines.entry(line.to_owned()).or_default();
let mut missed = coverable.saturating_sub(cached) as f64;
if read.is_none() && coverable >= p.implicit_floor {
let ratio = cached as f64 / coverable.max(1) as f64;
let implicit = entry
.implicit
.map_or(ratio, |ewma| 0.5 * ewma + 0.5 * ratio);
entry.implicit = Some(implicit);
if implicit >= 0.5 {
missed = 0.0;
}
}
entry.premium += missed * (1.0 - p.cached_ratio);
}
pub(crate) fn created(
&self,
from_line: &str,
create: &Create,
name: String,
tokens: u64,
model: &str,
) -> Lease {
let now = self.now();
let lease = Lease {
name: name.clone(),
digest: create.digest.clone(),
covers: create.covers,
tokens,
expires_at: now + create.ttl_secs,
model: model.to_owned(),
};
let mut book = self.book();
if create.estimate > 0 && tokens > 0 {
let ratio = tokens as f64 / create.estimate as f64;
book.calibration = Some(book.calibration.map_or(ratio, |c| 0.5 * c + 0.5 * ratio));
}
let previous = book.lines.get(from_line).cloned().unwrap_or_default();
if create.covers > 0 {
book.lines.insert(
lease.digest.clone(),
Line {
premium: 0.0,
implicit: previous.implicit,
rolled_at: now,
last_call: previous.last_call,
longest_gap: previous.longest_gap,
failed_at: None,
},
);
}
book.lives.insert(
name.clone(),
Life {
tokens,
created: now,
expires_at: lease.expires_at,
ended: None,
last_read: now,
line_gap: previous.longest_gap,
},
);
book.events.push(CacheEvent::Created {
name: name.clone(),
tokens,
estimated: create.expected,
covers: create.covers,
ttl_secs: create.ttl_secs,
at: now,
});
book.report.created.push(CreatedCache {
name,
tokens,
at: now,
});
book.leases.insert(lease.digest.clone(), lease.clone());
tracing::info!(target: "gemini.cache.created", name = %lease.name, tokens, covers = create.covers);
lease
}
pub(crate) fn create_failed(
&self,
line: &str,
coverable: u64,
status: Option<u16>,
message: String,
) {
let now = self.now();
let mut book = self.book();
if let Some(entry) = book.lines.get_mut(line) {
entry.failed_at = Some(coverable);
}
book.events.push(CacheEvent::CreateFailed {
status,
message,
at: now,
});
}
pub(crate) fn extended(&self, lease: &Lease, ttl_secs: u64) {
let now = self.now();
let expires_at = now + ttl_secs;
let mut book = self.book();
if let Some(held) = book.leases.get_mut(&lease.digest) {
held.expires_at = expires_at;
}
if let Some(life) = book.lives.get_mut(&lease.name) {
life.expires_at = expires_at;
}
book.events.push(CacheEvent::Extended {
name: lease.name.clone(),
expires_at,
at: now,
});
tracing::info!(target: "gemini.cache.extended", name = %lease.name, expires_at);
}
pub(crate) fn touched(&self, lease: &Lease) {
let now = self.now();
let mut book = self.book();
let gap = book
.lines
.get(&lease.digest)
.map_or(0, |line| line.longest_gap);
if let Some(life) = book.lives.get_mut(&lease.name) {
life.last_read = now;
life.line_gap = life.line_gap.max(gap);
}
}
}
#[cfg(test)]
mod tests;