kernel/keep_warm.rs
1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at https://mozilla.org/MPL/2.0/.
4// Copyright (c) 2026 2youg1 and the sprawling contributors
5
6//! Keep-warm: whether, and when, a prompt cache still in use is renewed
7//! before the provider lets it expire.
8//! The part `crates/kernel/spec/KeepWarm.lean` specifies this module.
9//!
10//! **Off unless the city enables it.** A renewal is a request the person
11//! pays for, so [`KeepWarm::Off`] is the default and plans nothing. Spend
12//! is observed where every request's usage is recorded; this decision
13//! reads no balance and refuses nothing.
14
15use crate::consts_external::PROMPT_CACHE_TTL_SECS;
16use serde::{Deserialize, Serialize};
17
18/// Whether this city renews a warm prompt cache before it expires.
19#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
20#[serde(rename_all = "snake_case")]
21#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
22pub enum KeepWarm {
23 /// No renewal is ever planned: the city sends no request of its own.
24 #[default]
25 Off,
26 /// A prefix a real request carried within the last cache lifetime is
27 /// renewed once, just before that lifetime ends.
28 FiveMinute,
29}
30
31/// One prefix's cache as last seen, in milliseconds: when a real request
32/// carried it, and when the provider last refreshed it. A renewal moves
33/// only the second, which is what bounds renewals to one per real use.
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub struct CacheUse {
36 used_ms: u64,
37 refreshed_ms: u64,
38}
39
40impl CacheUse {
41 /// A real request carrying the prefix was sent at `at_ms`.
42 pub fn sent(at_ms: u64) -> CacheUse {
43 CacheUse {
44 used_ms: at_ms,
45 refreshed_ms: at_ms,
46 }
47 }
48
49 /// A renewal request refreshed the cache at `at_ms`; the last real
50 /// use stays where it was.
51 #[must_use]
52 pub fn renewed(self, at_ms: u64) -> CacheUse {
53 CacheUse {
54 refreshed_ms: self.refreshed_ms.max(at_ms),
55 ..self
56 }
57 }
58}
59
60/// The instant the next renewal request is sent, or `None` when none is.
61///
62/// `lead_ms` is the round trip to the provider as the caller measured it:
63/// the renewal leaves that long before the cache expires, so the provider
64/// reads it while the entry is still alive. A renewal whose instant lies
65/// more than one cache lifetime after the last real use is not sent.
66pub fn renewal_due(setting: KeepWarm, cache: CacheUse, lead_ms: u64) -> Option<u64> {
67 match setting {
68 KeepWarm::Off => None,
69 KeepWarm::FiveMinute => {
70 let ttl_ms = PROMPT_CACHE_TTL_SECS.checked_mul(1000)?;
71 let due = cache
72 .refreshed_ms
73 .checked_add(ttl_ms)?
74 .saturating_sub(lead_ms);
75 (due.saturating_sub(cache.used_ms) <= ttl_ms).then_some(due)
76 }
77 }
78}
79
80#[cfg(test)]
81mod tests {
82 use super::*;
83
84 const TTL_MS: u64 = PROMPT_CACHE_TTL_SECS * 1000;
85 const LEAD_MS: u64 = 2_000;
86
87 #[test]
88 fn the_default_setting_plans_no_request_of_its_own() {
89 let use_at = 1_000_000;
90 let cases = [
91 CacheUse::sent(use_at),
92 CacheUse::sent(use_at).renewed(use_at + TTL_MS - LEAD_MS),
93 CacheUse::sent(0),
94 ];
95 for cache in cases {
96 assert_eq!(renewal_due(KeepWarm::default(), cache, LEAD_MS), None);
97 }
98 }
99
100 #[test]
101 fn five_minute_renews_a_used_prefix_once_just_before_it_expires() {
102 let use_at = 1_000_000;
103 let first = renewal_due(KeepWarm::FiveMinute, CacheUse::sent(use_at), LEAD_MS);
104 assert_eq!(first, Some(use_at + TTL_MS - LEAD_MS));
105 let renewed = CacheUse::sent(use_at).renewed(use_at + TTL_MS - LEAD_MS);
106 assert_eq!(renewal_due(KeepWarm::FiveMinute, renewed, LEAD_MS), None);
107 }
108}