1use std::collections::{BTreeMap, HashMap, HashSet};
85use std::sync::{Arc, Mutex, PoisonError};
86use std::time::Duration;
87
88use serde::{Deserialize, Serialize};
89use serde_json::value::RawValue;
90use sha2::{Digest, Sha256};
91
92use crate::completion::CacheCost;
93
94pub type Clock = Arc<dyn Fn() -> u64 + Send + Sync>;
96
97fn system_now() -> u64 {
102 #[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
103 {
104 std::time::SystemTime::now()
105 .duration_since(std::time::UNIX_EPOCH)
106 .map_or(0, |elapsed| elapsed.as_secs())
107 }
108 #[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
109 {
110 0
111 }
112}
113
114const MAX_TTL_SECS: u64 = 81 * 60;
117
118const EXPIRY_MARGIN_SECS: u64 = 5;
120
121const MIN_IDLE_SECS: u64 = 60;
124
125#[derive(Clone, Copy, Debug, PartialEq)]
130pub struct AutoCache {
131 pub ttl: Duration,
133 pub implicit_floor: u64,
135 pub implicit_ceiling: u64,
139 pub min_tokens: u64,
141 pub min_gain: u64,
143 pub cached_ratio: f64,
145 pub storage_ratio_per_hour: f64,
147}
148
149impl Default for AutoCache {
150 fn default() -> Self {
151 Self {
152 ttl: Duration::from_secs(60 * 60),
153 implicit_floor: 4_096,
154 implicit_ceiling: 16_000,
155 min_tokens: 1_024,
156 min_gain: 256,
157 cached_ratio: 0.1,
158 storage_ratio_per_hour: 0.5 / 0.75,
159 }
160 }
161}
162
163#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
165pub struct Lease {
166 pub name: String,
168 pub digest: String,
172 pub covers: usize,
174 pub tokens: u64,
176 pub expires_at: u64,
178 pub model: String,
180}
181
182#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
184pub enum CacheEvent {
185 Created {
187 name: String,
189 tokens: u64,
191 estimated: u64,
193 covers: usize,
195 ttl_secs: u64,
197 at: u64,
199 },
200 Retired {
202 name: String,
204 at: u64,
206 },
207 Lost {
209 name: String,
211 at: u64,
213 },
214 Extended {
216 name: String,
218 expires_at: u64,
220 at: u64,
222 },
223 CreateFailed {
225 status: Option<u16>,
227 message: String,
229 at: u64,
231 },
232}
233
234#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
236pub struct CreatedCache {
237 pub name: String,
239 pub tokens: u64,
241 pub at: u64,
243}
244
245#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
249pub struct CacheReport {
250 pub created: Vec<CreatedCache>,
252 pub reads: BTreeMap<String, u64>,
254 pub retired: Vec<String>,
257 pub lost: Vec<String>,
259 pub token_hours: f64,
262 pub live: Vec<Lease>,
264}
265
266impl From<&CacheReport> for CacheCost {
270 fn from(report: &CacheReport) -> Self {
271 Self {
272 cache_writes: report.created.iter().map(|created| created.tokens).sum(),
273 storage_token_hours: report.token_hours,
274 ..Self::default()
275 }
276 }
277}
278
279#[derive(Clone, Debug, Default)]
280struct Line {
281 premium: f64,
284 implicit: Option<f64>,
286 rolled_at: u64,
287 last_call: Option<u64>,
288 longest_gap: u64,
289 failed_at: Option<u64>,
292}
293
294#[derive(Clone, Debug)]
295struct Life {
296 tokens: u64,
297 created: u64,
298 expires_at: u64,
299 ended: Option<u64>,
300 last_read: u64,
301 line_gap: u64,
302}
303
304#[derive(Default)]
305pub(crate) struct Book {
306 pub(crate) leases: HashMap<String, Lease>,
307 lives: BTreeMap<String, Life>,
308 lineages: HashMap<String, HashSet<String>>,
310 lines: HashMap<String, Line>,
311 events: Vec<CacheEvent>,
312 report: CacheReport,
313 calibration: Option<f64>,
316}
317
318#[derive(Clone)]
322pub struct CacheBook {
323 inner: Arc<Mutex<Book>>,
324 pub(crate) create: Arc<futures::lock::Mutex<()>>,
326 policy: AutoCache,
327 clock: Clock,
328 pub(crate) display_prefix: Arc<str>,
329}
330
331impl std::fmt::Debug for CacheBook {
332 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
333 f.debug_struct("CacheBook")
334 .field("policy", &self.policy)
335 .field("display_prefix", &self.display_prefix)
336 .finish_non_exhaustive()
337 }
338}
339
340impl CacheBook {
341 pub fn new(policy: AutoCache) -> Self {
343 Self {
344 inner: Arc::default(),
345 create: Arc::default(),
346 policy,
347 clock: Arc::new(system_now),
348 display_prefix: Arc::from("rig-cache-"),
349 }
350 }
351
352 pub fn with_clock(mut self, clock: impl Fn() -> u64 + Send + Sync + 'static) -> Self {
354 self.clock = Arc::new(clock);
355 self
356 }
357
358 pub fn with_display_prefix(mut self, prefix: &str) -> Self {
361 self.display_prefix = Arc::from(prefix);
362 self
363 }
364
365 pub fn policy(&self) -> AutoCache {
367 self.policy
368 }
369
370 pub(crate) fn book(&self) -> std::sync::MutexGuard<'_, Book> {
371 self.inner.lock().unwrap_or_else(PoisonError::into_inner)
372 }
373
374 pub(crate) fn now(&self) -> u64 {
375 (self.clock)()
376 }
377
378 pub fn events(&self) -> Vec<CacheEvent> {
380 self.book().events.clone()
381 }
382
383 pub fn leases(&self) -> Vec<Lease> {
385 let mut leases: Vec<Lease> = self.book().leases.values().cloned().collect();
386 leases.sort_by(|a, b| a.name.cmp(&b.name));
387 leases
388 }
389
390 pub fn restore(&self, leases: Vec<Lease>) {
392 let now = self.now();
393 let mut book = self.book();
394 for lease in leases {
395 book.lives.entry(lease.name.clone()).or_insert(Life {
396 tokens: lease.tokens,
397 created: now,
398 expires_at: lease.expires_at,
399 ended: None,
400 last_read: now,
401 line_gap: 0,
402 });
403 book.leases.insert(lease.digest.clone(), lease);
404 }
405 }
406
407 pub fn report(&self) -> CacheReport {
409 let now = self.now();
410 let book = self.book();
411 let mut report = book.report.clone();
412 report.token_hours = book
413 .lives
414 .values()
415 .map(|life| {
416 let end = life.ended.unwrap_or(now).min(life.expires_at);
417 life.tokens as f64 * end.saturating_sub(life.created) as f64 / 3600.0
418 })
419 .sum();
420 report.live = book.leases.values().cloned().collect();
421 report.live.sort_by(|a, b| a.name.cmp(&b.name));
422 report
423 }
424
425 pub(crate) fn retired(&self, lease: &Lease) {
426 let now = self.now();
427 let mut book = self.book();
428 book.leases.remove(&lease.digest);
429 end_life(&mut book, &lease.name, now);
430 book.events.push(CacheEvent::Retired {
431 name: lease.name.clone(),
432 at: now,
433 });
434 book.report.retired.push(lease.name.clone());
435 tracing::info!(target: "gemini.cache.retired", name = %lease.name);
436 }
437
438 pub(crate) fn lost(&self, lease: &Lease) {
439 let now = self.now();
440 let mut book = self.book();
441 book.leases.remove(&lease.digest);
442 end_life(&mut book, &lease.name, now);
443 book.events.push(CacheEvent::Lost {
444 name: lease.name.clone(),
445 at: now,
446 });
447 book.report.lost.push(lease.name.clone());
448 tracing::info!(target: "gemini.cache.lost", name = %lease.name);
449 }
450}
451
452pub(crate) fn end_life(book: &mut Book, name: &str, now: u64) {
453 if let Some(life) = book.lives.get_mut(name)
454 && life.ended.is_none()
455 {
456 life.ended = Some(now);
457 }
458}
459
460pub(crate) fn short_digest(digest: &str) -> &str {
461 digest.get(..40).unwrap_or(digest)
462}
463
464const PREFIX_KEYS: [&str; 3] = ["systemInstruction", "tools", "toolConfig"];
468
469pub(crate) struct Parsed {
471 pub(crate) rest: Vec<(String, Box<RawValue>)>,
473 pub(crate) prefix: [Option<Box<RawValue>>; 3],
475 pub(crate) contents: Vec<Box<RawValue>>,
476 pub(crate) has_cached_content: bool,
477}
478
479pub(crate) fn parse(bytes: &[u8]) -> Option<Parsed> {
480 let fields: serde_json::Map<String, serde_json::Value> = serde_json::from_slice(bytes).ok()?;
481 let raw: HashMap<String, Box<RawValue>> = serde_json::from_slice(bytes).ok()?;
484 let mut parsed = Parsed {
485 rest: Vec::new(),
486 prefix: [None, None, None],
487 contents: Vec::new(),
488 has_cached_content: false,
489 };
490 for key in fields.keys() {
491 let value = raw.get(key)?.clone();
492 if key == "contents" {
493 parsed.contents = serde_json::from_str(value.get()).ok()?;
494 } else if let Some(slot) = PREFIX_KEYS.iter().position(|prefix| prefix == key) {
495 if value.get() != "null"
496 && let Some(field) = parsed.prefix.get_mut(slot)
497 {
498 *field = Some(value);
499 }
500 } else {
501 if key == "cachedContent" && value.get() != "null" {
502 parsed.has_cached_content = true;
503 }
504 parsed.rest.push((key.clone(), value));
505 }
506 }
507 Some(parsed)
508}
509
510fn hex(bytes: &[u8]) -> String {
511 use std::fmt::Write;
512 bytes.iter().fold(String::new(), |mut out, byte| {
513 let _ = write!(out, "{byte:02x}");
514 out
515 })
516}
517
518pub(crate) fn digests(model: &str, parsed: &Parsed) -> Vec<String> {
520 let mut hasher = Sha256::new();
521 hasher.update(model.as_bytes());
522 for field in &parsed.prefix {
523 hasher.update([0u8]);
524 if let Some(raw) = field {
525 hasher.update(raw.get().as_bytes());
526 }
527 }
528 let mut state = hasher.finalize().to_vec();
529 let mut out = vec![hex(&state)];
530 for content in &parsed.contents {
531 let mut hasher = Sha256::new();
532 hasher.update(&state);
533 hasher.update(content.get().as_bytes());
534 state = hasher.finalize().to_vec();
535 out.push(hex(&state));
536 }
537 out
538}
539
540fn estimate(json: &str) -> u64 {
544 let Ok(mut value) = serde_json::from_str::<serde_json::Value>(json) else {
545 return (json.len() / 4) as u64;
546 };
547 fn scrub(value: &mut serde_json::Value) {
548 match value {
549 serde_json::Value::Object(map) => {
550 map.remove("thoughtSignature");
551 map.values_mut().for_each(scrub);
552 }
553 serde_json::Value::Array(items) => items.iter_mut().for_each(scrub),
554 _ => {}
555 }
556 }
557 scrub(&mut value);
558 (value.to_string().len() / 4) as u64
559}
560
561pub(crate) fn is_user_text(raw: &RawValue) -> bool {
563 #[derive(Deserialize)]
564 struct Content {
565 role: Option<String>,
566 #[serde(default)]
567 parts: Vec<serde_json::Map<String, serde_json::Value>>,
568 }
569 serde_json::from_str::<Content>(raw.get()).is_ok_and(|content| {
570 content.role.as_deref() == Some("user")
571 && content.parts.iter().any(|part| part.contains_key("text"))
572 && !content
573 .parts
574 .iter()
575 .any(|part| part.contains_key("functionResponse"))
576 })
577}
578
579pub(crate) fn stripped(parsed: &Parsed, lease: &Lease) -> Option<Vec<u8>> {
581 let mut out: Vec<(String, &RawValue)> = Vec::new();
582 let name = serde_json::value::to_raw_value(&lease.name).ok()?;
583 let contents = serde_json::value::to_raw_value(&parsed.contents.get(lease.covers..)?).ok()?;
584 out.push(("cachedContent".to_owned(), &name));
585 out.push(("contents".to_owned(), &contents));
586 for (key, value) in &parsed.rest {
587 if key != "cachedContent" {
588 out.push((key.clone(), value));
589 }
590 }
591 let mut bytes = b"{".to_vec();
592 for (index, (key, value)) in out.iter().enumerate() {
593 if index > 0 {
594 bytes.push(b',');
595 }
596 bytes.extend(serde_json::to_vec(key).ok()?);
597 bytes.push(b':');
598 bytes.extend(value.get().as_bytes());
599 }
600 bytes.push(b'}');
601 Some(bytes)
602}
603
604pub(crate) fn cache_body(
607 model: &str,
608 parsed: &Parsed,
609 covers: usize,
610 display_name: &str,
611 ttl_secs: u64,
612) -> Option<Vec<u8>> {
613 let mut bytes = b"{".to_vec();
614 let mut field = |key: &str, value: &str| {
615 if bytes.len() > 1 {
616 bytes.push(b',');
617 }
618 bytes.extend(format!("\"{key}\":").as_bytes());
619 bytes.extend(value.as_bytes());
620 };
621 field(
622 "model",
623 &serde_json::to_string(&format!("models/{model}")).ok()?,
624 );
625 field("displayName", &serde_json::to_string(display_name).ok()?);
626 field("ttl", &format!("\"{ttl_secs}s\""));
627 for (key, value) in PREFIX_KEYS.iter().zip(&parsed.prefix) {
628 if let Some(value) = value {
629 field(key, value.get());
630 }
631 }
632 if covers > 0 {
633 let contents = serde_json::value::to_raw_value(&parsed.contents.get(..covers)?).ok()?;
634 field("contents", contents.get());
635 }
636 bytes.push(b'}');
637 Some(bytes)
638}
639
640pub(crate) struct Plan {
645 pub(crate) read: Option<Lease>,
646 pub(crate) create: Option<Create>,
647 pub(crate) retire: Vec<Lease>,
648 pub(crate) extend: Option<(Lease, u64)>,
649 pub(crate) line: String,
650 pub(crate) coverable: u64,
651}
652
653pub(crate) struct Create {
654 pub(crate) covers: usize,
655 pub(crate) digest: String,
656 pub(crate) ttl_secs: u64,
657 estimate: u64,
659 expected: u64,
661 pub(crate) replaces: Option<Lease>,
662}
663
664impl CacheBook {
665 pub(crate) fn plan(
666 &self,
667 model: &str,
668 d: &[String],
669 parsed: &Parsed,
670 roll_allowed: bool,
671 ) -> Plan {
672 let p = self.policy;
673 let now = self.now();
674 let n = parsed.contents.len();
675 let at = |k: usize| d.get(k).cloned().unwrap_or_default();
677 let mut book = self.book();
678 let calibration = book.calibration.unwrap_or(1.0);
679 let prefix_estimate: u64 = parsed
680 .prefix
681 .iter()
682 .flatten()
683 .map(|raw| estimate(raw.get()))
684 .sum();
685 let content_estimates: Vec<u64> = parsed
686 .contents
687 .iter()
688 .map(|raw| estimate(raw.get()))
689 .collect();
690
691 if n >= 1 {
694 book.lineages.entry(at(0)).or_default().insert(at(1));
695 }
696 let shared = book
697 .lineages
698 .get(&at(0))
699 .is_some_and(|seen| seen.len() >= 2);
700
701 let idle: Vec<Lease> = book
703 .leases
704 .values()
705 .filter(|lease| lease.covers > 0 && lease.model == model)
706 .filter(|lease| {
707 book.lives.get(&lease.name).is_some_and(|life| {
708 now.saturating_sub(life.last_read)
709 > MIN_IDLE_SECS.max(life.line_gap.saturating_mul(3))
710 })
711 })
712 .cloned()
713 .collect();
714
715 let read = (0..n.max(1))
717 .rev()
718 .find_map(|k| {
719 book.leases
720 .get(&at(k))
721 .filter(|lease| lease.model == model)
722 .filter(|lease| lease.expires_at > now + EXPIRY_MARGIN_SECS)
723 .cloned()
724 })
725 .filter(|lease| !idle.iter().any(|gone| gone.name == lease.name));
726 let covered = read.as_ref().map_or(0, |lease| lease.covers);
727 let tail: u64 = content_estimates
728 .get(covered..n.saturating_sub(1).max(covered))
729 .unwrap_or_default()
730 .iter()
731 .sum();
732 let coverable = match &read {
733 Some(lease) => lease.tokens + (tail as f64 * calibration) as u64,
734 None => ((prefix_estimate + tail) as f64 * calibration) as u64,
735 };
736
737 let line_key = read
739 .as_ref()
740 .filter(|lease| lease.covers > 0)
741 .map_or_else(|| at(1.min(n)), |lease| lease.digest.clone());
742 let prefix_leased = book.leases.contains_key(&at(0));
743 let line = book.lines.entry(line_key.clone()).or_insert_with(|| Line {
744 rolled_at: now,
745 ..Line::default()
746 });
747 if let Some(last) = line.last_call {
748 line.longest_gap = line.longest_gap.max(now.saturating_sub(last));
749 }
750 line.last_call = Some(now);
751 let longest_gap = line.longest_gap;
752 let ttl_secs = p
753 .ttl
754 .as_secs()
755 .max(longest_gap.saturating_mul(2).min(MAX_TTL_SECS))
756 .max(1);
757
758 let mut plan = Plan {
759 read: read.clone(),
760 create: None,
761 retire: idle,
762 extend: None,
763 line: line_key,
764 coverable,
765 };
766
767 let implicit_serving = line.implicit.is_none_or(|ratio| ratio >= 0.5);
769 if coverable >= p.implicit_ceiling && implicit_serving {
770 plan.read = None;
771 return plan;
772 }
773
774 let failed_below = line.failed_at.is_some_and(|at| coverable < at + p.min_gain);
775 if read.is_none()
776 && shared
777 && !prefix_leased
778 && (prefix_estimate as f64 * calibration) as u64 >= p.min_tokens
779 && !failed_below
780 {
781 plan.create = Some(Create {
782 covers: 0,
783 digest: at(0),
784 ttl_secs,
785 estimate: prefix_estimate,
786 expected: (prefix_estimate as f64 * calibration) as u64,
787 replaces: None,
788 });
789 } else if n >= 2 && roll_allowed && !failed_below {
790 let cached = read.as_ref().map_or(0, |lease| lease.tokens);
791 let gain = coverable.saturating_sub(cached);
792 let held_hours = now.saturating_sub(line.rolled_at).max(60) as f64 / 3600.0;
793 let cost = coverable as f64 * (1.0 + p.storage_ratio_per_hour * held_hours);
794 if gain >= p.min_gain && coverable >= p.min_tokens && line.premium >= cost {
795 plan.create = Some(Create {
796 covers: n - 1,
797 digest: at(n - 1),
798 ttl_secs,
799 estimate: prefix_estimate
800 + content_estimates
801 .get(..n - 1)
802 .unwrap_or_default()
803 .iter()
804 .sum::<u64>(),
805 expected: coverable,
806 replaces: read.clone().filter(|lease| lease.covers > 0),
807 });
808 }
809 }
810
811 if plan.create.is_none()
813 && let Some(lease) = &read
814 && longest_gap > 0
815 && lease.expires_at.saturating_sub(now) < longest_gap.saturating_mul(3) / 2
816 {
817 plan.extend = Some((lease.clone(), ttl_secs));
818 }
819 plan
820 }
821
822 pub(crate) fn observe(&self, line: &str, read: Option<&str>, coverable: u64, cached: u64) {
824 let p = self.policy;
825 let mut book = self.book();
826 if let Some(name) = read {
827 *book.report.reads.entry(name.to_owned()).or_default() += 1;
828 }
829 let entry = book.lines.entry(line.to_owned()).or_default();
830 let mut missed = coverable.saturating_sub(cached) as f64;
831 if read.is_none() && coverable >= p.implicit_floor {
832 let ratio = cached as f64 / coverable.max(1) as f64;
833 let implicit = entry
834 .implicit
835 .map_or(ratio, |ewma| 0.5 * ewma + 0.5 * ratio);
836 entry.implicit = Some(implicit);
837 if implicit >= 0.5 {
838 missed = 0.0;
839 }
840 }
841 entry.premium += missed * (1.0 - p.cached_ratio);
842 }
843
844 pub(crate) fn created(
845 &self,
846 from_line: &str,
847 create: &Create,
848 name: String,
849 tokens: u64,
850 model: &str,
851 ) -> Lease {
852 let now = self.now();
853 let lease = Lease {
854 name: name.clone(),
855 digest: create.digest.clone(),
856 covers: create.covers,
857 tokens,
858 expires_at: now + create.ttl_secs,
859 model: model.to_owned(),
860 };
861 let mut book = self.book();
862 if create.estimate > 0 && tokens > 0 {
863 let ratio = tokens as f64 / create.estimate as f64;
864 book.calibration = Some(book.calibration.map_or(ratio, |c| 0.5 * c + 0.5 * ratio));
865 }
866 let previous = book.lines.get(from_line).cloned().unwrap_or_default();
867 if create.covers > 0 {
868 book.lines.insert(
869 lease.digest.clone(),
870 Line {
871 premium: 0.0,
872 implicit: previous.implicit,
873 rolled_at: now,
874 last_call: previous.last_call,
875 longest_gap: previous.longest_gap,
876 failed_at: None,
877 },
878 );
879 }
880 book.lives.insert(
881 name.clone(),
882 Life {
883 tokens,
884 created: now,
885 expires_at: lease.expires_at,
886 ended: None,
887 last_read: now,
888 line_gap: previous.longest_gap,
889 },
890 );
891 book.events.push(CacheEvent::Created {
892 name: name.clone(),
893 tokens,
894 estimated: create.expected,
895 covers: create.covers,
896 ttl_secs: create.ttl_secs,
897 at: now,
898 });
899 book.report.created.push(CreatedCache {
900 name,
901 tokens,
902 at: now,
903 });
904 book.leases.insert(lease.digest.clone(), lease.clone());
905 tracing::info!(target: "gemini.cache.created", name = %lease.name, tokens, covers = create.covers);
906 lease
907 }
908
909 pub(crate) fn create_failed(
910 &self,
911 line: &str,
912 coverable: u64,
913 status: Option<u16>,
914 message: String,
915 ) {
916 let now = self.now();
917 let mut book = self.book();
918 if let Some(entry) = book.lines.get_mut(line) {
919 entry.failed_at = Some(coverable);
920 }
921 book.events.push(CacheEvent::CreateFailed {
922 status,
923 message,
924 at: now,
925 });
926 }
927
928 pub(crate) fn extended(&self, lease: &Lease, ttl_secs: u64) {
929 let now = self.now();
930 let expires_at = now + ttl_secs;
931 let mut book = self.book();
932 if let Some(held) = book.leases.get_mut(&lease.digest) {
933 held.expires_at = expires_at;
934 }
935 if let Some(life) = book.lives.get_mut(&lease.name) {
936 life.expires_at = expires_at;
937 }
938 book.events.push(CacheEvent::Extended {
939 name: lease.name.clone(),
940 expires_at,
941 at: now,
942 });
943 tracing::info!(target: "gemini.cache.extended", name = %lease.name, expires_at);
944 }
945
946 pub(crate) fn touched(&self, lease: &Lease) {
947 let now = self.now();
948 let mut book = self.book();
949 let gap = book
950 .lines
951 .get(&lease.digest)
952 .map_or(0, |line| line.longest_gap);
953 if let Some(life) = book.lives.get_mut(&lease.name) {
954 life.last_read = now;
955 life.line_gap = life.line_gap.max(gap);
956 }
957 }
958}