Skip to main content

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}