runtime/prefix/
warmth.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//! One session's keep-warm account: the last real request each prefix
7//! carried, and the renewal sent for it when `kernel::keep_warm` says it
8//! is due (`crates/runtime/spec/Prefix.lean` §8-4-2).
9//!
10//! **Off keeps nothing.** Under [`KeepWarm::Off`] a request is not even
11//! cloned, so the default setting costs no memory per session and has no
12//! renewal it could send.
13
14use std::collections::BTreeMap;
15use std::sync::Arc;
16use std::sync::atomic::{AtomicU64, Ordering};
17
18use kernel::keep_warm::{CacheUse, KeepWarm, renewal_due};
19use kernel::{
20    AxError, B3Hash, Ceiling, EarlyCalls, Increments, Model, ModelRequest, ModelReturn, TimeMs,
21};
22
23/// Keep-warm bookkeeping for the prefixes one session has sent.
24#[derive(Debug, Clone)]
25pub struct Warmth {
26    setting: KeepWarm,
27    kept: BTreeMap<[B3Hash; 4], Kept>,
28}
29
30/// A prefix's last real request, and its cache as last seen.
31#[derive(Debug, Clone)]
32struct Kept {
33    request: ModelRequest<'static>,
34    cache: CacheUse,
35}
36
37impl Warmth {
38    pub fn new(setting: KeepWarm) -> Warmth {
39        Warmth {
40            setting,
41            kept: BTreeMap::new(),
42        }
43    }
44
45    /// A real request carrying `request.segments` was sent at `at_ms`.
46    pub fn sent(&mut self, request: &ModelRequest, at_ms: u64) {
47        match self.setting {
48            KeepWarm::Off => return,
49            KeepWarm::FiveMinute => {}
50        }
51        self.kept.insert(
52            request.segments,
53            Kept {
54                request: request.clone().into_owned(),
55                cache: CacheUse::sent(at_ms),
56            },
57        );
58    }
59
60    /// The instant the earliest renewal is sent with the lead `lead`
61    /// reads now, or `None` when no prefix is to be renewed and a timer
62    /// need not wake.
63    pub fn next_due(&self, lead: &Lead) -> Option<u64> {
64        self.kept
65            .values()
66            .filter_map(|kept| renewal_due(self.setting, kept.cache, lead.ms()))
67            .min()
68    }
69
70    /// Hands every renewal due by `now_ms` to `send`, one at a time.
71    /// Every answer comes back in the round, a refusal beside the
72    /// answers that preceded it rather than instead of them: their usage
73    /// the caller owes the ledger all the same. Nothing more is sent
74    /// after a refusal, and a prefix that owes no renewal is let go
75    /// here.
76    pub fn renew_due(
77        &mut self,
78        lead: &Lead,
79        now_ms: u64,
80        mut send: impl FnMut(&ModelRequest) -> Result<ModelReturn, AxError>,
81    ) -> RenewalRound {
82        let (setting, lead_ms) = (self.setting, lead.ms());
83        self.kept
84            .retain(|_, kept| renewal_due(setting, kept.cache, lead_ms).is_some());
85        let mut round = RenewalRound::default();
86        for kept in self.kept.values_mut() {
87            if renewal_due(setting, kept.cache, lead_ms).is_some_and(|due| due <= now_ms) {
88                match send(&renewal_of(&kept.request)) {
89                    Ok(answer) => {
90                        round.answers.push(answer);
91                        kept.cache = kept.cache.renewed();
92                    }
93                    Err(refused) => {
94                        round.refused = Some(refused);
95                        break;
96                    }
97                }
98            }
99        }
100        round
101    }
102}
103
104/// What one wake's renewals came back with: every answer the provider
105/// gave, and the first refusal when one came. The answers survive a
106/// later refusal, because each one's usage is already spent.
107#[derive(Debug, Default)]
108pub struct RenewalRound {
109    /// Each renewal the provider answered, in the order sent.
110    pub answers: Vec<ModelReturn>,
111    /// The refusal that ended the round, if one did.
112    pub refused: Option<AxError>,
113}
114
115/// The last real request, asking for one output token: the cache hits on
116/// the prefix bytes whatever the ceiling, so one token is the lowest
117/// price a renewal can pay.
118fn renewal_of(request: &ModelRequest<'static>) -> ModelRequest<'static> {
119    let mut renewal = request.clone();
120    renewal.chat.max_tokens = Ceiling::new(1);
121    renewal
122}
123
124/// How long a request to one endpoint takes from leaving the city to the
125/// provider starting it, in milliseconds, as last measured; 0 before the
126/// first reading.
127///
128/// A clone shares the reading, so every door that reaches the endpoint
129/// renews by the same lead: the network path and the permit gate belong
130/// to the endpoint, and one door sees one renewal per real use, too few
131/// to keep a reading of its own (`crates/runtime/spec/Prefix.lean` D100).
132#[derive(Debug, Clone, Default)]
133pub struct Lead(Arc<AtomicU64>);
134
135impl Lead {
136    /// The latest reading.
137    pub fn ms(&self) -> u64 {
138        self.0.load(Ordering::Relaxed)
139    }
140
141    fn measured(&self, ms: u64) {
142        self.0.store(ms, Ordering::Relaxed);
143    }
144}
145
146/// The door a run's requests enter the keep-warm account by: it owns the
147/// adapter a run calls and is itself the [`Model`] the turn loop is
148/// handed, so every request that reached the provider is recorded
149/// without the turn loop knowing keep-warm exists.
150///
151/// It forwards every call method of [`Model`] rather than inheriting a
152/// default, because a default is correct and only slower: an inherited
153/// `call_speculating` drops the early calls, and no read then starts
154/// while the model writes (`crates/runtime/spec/Prefix.lean` D97).
155///
156/// The clock is the caller's, injected from the one sampling point. The
157/// door measures its endpoint's [`Lead`] from a renewal's round trip and
158/// from a real call's first increment, never from a whole call, whose
159/// generation the provider's lifetime does not wait for (D100).
160pub struct Warmed<C> {
161    model: Box<dyn Model + Send>,
162    warmth: Warmth,
163    clock: C,
164    lead: Lead,
165}
166
167impl<C: FnMut() -> Result<TimeMs, AxError>> Warmed<C> {
168    /// Wraps `model`, which reaches the endpoint `lead` is read for;
169    /// under [`KeepWarm::Off`] it only forwards.
170    pub fn new(model: Box<dyn Model + Send>, setting: KeepWarm, clock: C, lead: Lead) -> Warmed<C> {
171        Warmed {
172            model,
173            warmth: Warmth::new(setting),
174            clock,
175            lead,
176        }
177    }
178
179    /// See [`Warmth::next_due`]; the lead is the endpoint's reading now.
180    pub fn next_due(&self) -> Option<u64> {
181        self.warmth.next_due(&self.lead)
182    }
183
184    /// Sends every renewal due by `now_ms` through the adapter this door
185    /// owns, and takes each one's round trip as the endpoint's lead: a
186    /// renewal asks for one token, so its answer follows the prefill.
187    /// A clock or adapter failure ends the round as its refusal.
188    pub fn renew_due(&mut self, now_ms: u64) -> RenewalRound {
189        let Warmed {
190            model,
191            warmth,
192            clock,
193            lead,
194        } = self;
195        let lead: &Lead = lead;
196        warmth.renew_due(lead, now_ms, |renewal| {
197            let left = clock()?;
198            let answer = model.call(renewal)?;
199            let back = clock()?;
200            lead.measured(back.value().saturating_sub(left.value()));
201            Ok(answer)
202        })
203    }
204
205    /// Makes one real call through `send`, which reports each increment
206    /// to the sink it is handed, and records it when the provider
207    /// answered. The time to the first increment becomes the endpoint's
208    /// lead; it includes the prefill and the first token, so it errs
209    /// early, never late. A call with no increment leaves the lead.
210    fn recorded(
211        &mut self,
212        request: &ModelRequest,
213        send: impl FnOnce(&mut dyn Model, Increments<'_>) -> Result<ModelReturn, AxError>,
214    ) -> Result<ModelReturn, AxError> {
215        let Warmed {
216            model,
217            warmth,
218            clock,
219            lead,
220        } = self;
221        let left = clock()?;
222        let mut first = None;
223        let answer = send(model.as_mut(), &mut |_| {
224            if first.is_none() {
225                first = Some(clock());
226            }
227        })?;
228        if let Some(first) = first {
229            lead.measured(first?.value().saturating_sub(left.value()));
230        }
231        warmth.sent(request, left.value());
232        Ok(answer)
233    }
234}
235
236impl<C: FnMut() -> Result<TimeMs, AxError>> Model for Warmed<C> {
237    fn provider_account(&self) -> Option<kernel::event::record::ProviderAccountBinding> {
238        self.model.provider_account()
239    }
240
241    fn account_roster(&self) -> Option<kernel::model::AccountRoster> {
242        self.model.account_roster()
243    }
244
245    fn select_account(&mut self, account: &kernel::ServerLabel) {
246        self.model.select_account(account);
247    }
248
249    fn prefix_held(&self) -> Option<kernel::event::record::PrefixHeld> {
250        self.model.prefix_held()
251    }
252
253    fn call(&mut self, req: &ModelRequest) -> Result<ModelReturn, AxError> {
254        self.recorded(req, |model, _| model.call(req))
255    }
256
257    fn call_streaming(
258        &mut self,
259        req: &ModelRequest,
260        onto: Increments<'_>,
261    ) -> Result<ModelReturn, AxError> {
262        self.recorded(req, |model, first| {
263            model.call_streaming(req, &mut |increment| {
264                first(increment);
265                onto(increment);
266            })
267        })
268    }
269
270    fn call_speculating(
271        &mut self,
272        req: &ModelRequest,
273        onto: Increments<'_>,
274        early: EarlyCalls<'_>,
275    ) -> Result<ModelReturn, AxError> {
276        self.recorded(req, |model, first| {
277            model.call_speculating(
278                req,
279                &mut |increment| {
280                    first(increment);
281                    onto(increment);
282                },
283                early,
284            )
285        })
286    }
287}
288
289#[cfg(test)]
290#[allow(
291    clippy::unwrap_used,
292    clippy::expect_used,
293    clippy::indexing_slicing,
294    clippy::arithmetic_side_effects
295)]
296mod tests {
297    use super::*;
298    use kernel::consts_external::{PROMPT_CACHE_RENEWAL_MARGIN_SECS, PROMPT_CACHE_TTL_SECS};
299    use kernel::{BuildingPolicy, ChatRequest, ContentBlock, Increment};
300    use std::sync::Arc;
301    use std::sync::atomic::{AtomicU64, Ordering};
302
303    const TTL_MS: u64 = PROMPT_CACHE_TTL_SECS * 1000;
304    const MARGIN_MS: u64 = PROMPT_CACHE_RENEWAL_MARGIN_SECS * 1000;
305    const LEAD_MS: u64 = 2_000;
306    const USED_AT: u64 = 1_000_000;
307
308    /// Counts what reaches the provider, and keeps each request it saw.
309    #[derive(Default)]
310    struct Provider {
311        seen: Vec<ModelRequest<'static>>,
312    }
313
314    impl Model for Provider {
315        fn call(&mut self, req: &ModelRequest) -> Result<ModelReturn, AxError> {
316            self.seen.push(req.clone().into_owned());
317            Ok(ModelReturn::bare(
318                kernel::model::message_payload(&[ContentBlock::Text {
319                    text: "warm".to_owned(),
320                }])
321                .unwrap(),
322                Vec::new(),
323            ))
324        }
325    }
326
327    fn request() -> ModelRequest<'static> {
328        ModelRequest {
329            policy: BuildingPolicy::default(),
330            segments: [B3Hash::digest(b"prefix"); 4],
331            chat: ChatRequest::empty("script", Ceiling::new(512).unwrap()),
332        }
333    }
334
335    #[test]
336    fn a_later_refusal_keeps_the_answers_before_it() {
337        let mut warmth = Warmth::new(KeepWarm::FiveMinute);
338        let mut first = request();
339        first.segments = [B3Hash::digest(b"one"); 4];
340        let mut second = request();
341        second.segments = [B3Hash::digest(b"two"); 4];
342        warmth.sent(&first, USED_AT);
343        warmth.sent(&second, USED_AT);
344        let due = warmth.next_due(&lead()).unwrap();
345        let mut calls = 0;
346        let round = warmth.renew_due(&lead(), due, |_| {
347            calls += 1;
348            if calls == 1 {
349                Ok(ModelReturn::bare(
350                    kernel::model::message_payload(&[ContentBlock::Text {
351                        text: "warm".to_owned(),
352                    }])
353                    .unwrap(),
354                    Vec::new(),
355                ))
356            } else {
357                Err(AxError::failure(
358                    kernel::AxCode::Provider,
359                    "renew a cached prefix",
360                    "rate limited",
361                )
362                .with_recovery("renewal is skipped until the next real request"))
363            }
364        });
365        assert_eq!(round.answers.len(), 1, "the answered renewal survives");
366        assert!(round.refused.is_some(), "the refusal ends the round");
367        assert_eq!(calls, 2, "nothing more is sent after the refusal");
368        assert!(
369            warmth.next_due(&lead()).is_some(),
370            "the refused prefix stays"
371        );
372    }
373
374    /// An endpoint's lead that has read [`LEAD_MS`].
375    fn lead() -> Lead {
376        let lead = Lead::default();
377        lead.measured(LEAD_MS);
378        lead
379    }
380
381    /// Asks `warmth` for renewals at every instant a timer could wake,
382    /// from the real use to well past one cache lifetime.
383    fn wake_through_two_lifetimes(warmth: &mut Warmth, provider: &mut Provider) {
384        for now in (USED_AT..=USED_AT + 2 * TTL_MS).step_by(1_000) {
385            let round = warmth.renew_due(&lead(), now, |renewal| provider.call(renewal));
386            assert!(round.refused.is_none(), "{:?}", round.refused);
387        }
388    }
389
390    #[test]
391    fn the_default_setting_sends_no_request_of_its_own() {
392        let mut warmth = Warmth::new(KeepWarm::default());
393        let mut provider = Provider::default();
394        warmth.sent(&request(), USED_AT);
395        assert_eq!(warmth.next_due(&lead()), None);
396        wake_through_two_lifetimes(&mut warmth, &mut provider);
397        assert!(
398            provider.seen.is_empty(),
399            "{} renewals sent",
400            provider.seen.len()
401        );
402    }
403
404    #[test]
405    fn five_minute_resends_the_last_request_once_with_one_output_token() {
406        let mut warmth = Warmth::new(KeepWarm::FiveMinute);
407        let mut provider = Provider::default();
408        warmth.sent(&request(), USED_AT);
409        assert_eq!(
410            warmth.next_due(&lead()),
411            Some(USED_AT + TTL_MS - MARGIN_MS - LEAD_MS)
412        );
413        wake_through_two_lifetimes(&mut warmth, &mut provider);
414        let mut renewal = request();
415        renewal.chat.max_tokens = Ceiling::new(1);
416        assert_eq!(provider.seen, vec![renewal]);
417        assert_eq!(warmth.next_due(&lead()), None);
418    }
419
420    /// What reached the provider behind a [`Warmed`] door, read after
421    /// the door owns the provider.
422    type Seen = std::sync::Arc<std::sync::Mutex<Vec<ModelRequest<'static>>>>;
423
424    struct Behind(Seen);
425
426    impl Model for Behind {
427        fn call(&mut self, req: &ModelRequest) -> Result<ModelReturn, AxError> {
428            let mut provider = Provider::default();
429            let answer = provider.call(req);
430            self.0.lock().unwrap().push(req.clone().into_owned());
431            answer
432        }
433    }
434
435    fn door(setting: KeepWarm) -> (Warmed<impl FnMut() -> Result<TimeMs, AxError>>, Seen) {
436        let seen = Seen::default();
437        let clock = || Ok(TimeMs::new(USED_AT));
438        let warmed = Warmed::new(
439            Box::new(Behind(seen.clone())),
440            setting,
441            clock,
442            Lead::default(),
443        );
444        (warmed, seen)
445    }
446
447    fn wake_door_through_two_lifetimes(
448        warmed: &mut Warmed<impl FnMut() -> Result<TimeMs, AxError>>,
449    ) {
450        for now in (USED_AT..=USED_AT + 2 * TTL_MS).step_by(1_000) {
451            let round = warmed.renew_due(now);
452            assert!(round.refused.is_none(), "{:?}", round.refused);
453        }
454    }
455
456    #[test]
457    fn the_default_door_forwards_a_run_and_sends_nothing_of_its_own() {
458        let (mut warmed, seen) = door(KeepWarm::default());
459        warmed.call(&request()).unwrap();
460        wake_door_through_two_lifetimes(&mut warmed);
461        assert_eq!(*seen.lock().unwrap(), vec![request()]);
462    }
463
464    #[test]
465    fn a_five_minute_door_renews_what_a_run_sent_through_it_once() {
466        let (mut warmed, seen) = door(KeepWarm::FiveMinute);
467        warmed.call_streaming(&request(), &mut |_| {}).unwrap();
468        assert_eq!(warmed.next_due(), Some(USED_AT + TTL_MS - MARGIN_MS));
469        wake_door_through_two_lifetimes(&mut warmed);
470        let mut renewal = request();
471        renewal.chat.max_tokens = Ceiling::new(1);
472        assert_eq!(*seen.lock().unwrap(), vec![request(), renewal]);
473    }
474
475    /// How long a real call takes from sending to settling.
476    const CALL_MS: u64 = 200_000;
477    /// How long a streamed real call takes to say its first word.
478    const FIRST_WORD_MS: u64 = 3_000;
479    /// How long a renewal takes from sending to settling.
480    const RENEWAL_MS: u64 = 2_000;
481
482    /// A provider that moves the door's clock itself: a real call says
483    /// its first word [`FIRST_WORD_MS`] after it was sent, when it is
484    /// streamed, and settles [`CALL_MS`] after; a renewal settles
485    /// [`RENEWAL_MS`] after.
486    struct Slow {
487        now: Arc<AtomicU64>,
488        seen: Seen,
489    }
490
491    impl Slow {
492        fn answer(
493            &mut self,
494            req: &ModelRequest,
495            onto: Increments<'_>,
496        ) -> Result<ModelReturn, AxError> {
497            self.seen.lock().unwrap().push(req.clone().into_owned());
498            if req.chat.max_tokens == Ceiling::new(1) {
499                self.now.fetch_add(RENEWAL_MS, Ordering::SeqCst);
500            } else {
501                self.now.fetch_add(FIRST_WORD_MS, Ordering::SeqCst);
502                onto(&Increment::Said("warm".to_owned()));
503                self.now
504                    .fetch_add(CALL_MS - FIRST_WORD_MS, Ordering::SeqCst);
505            }
506            Provider::default().call(req)
507        }
508    }
509
510    impl Model for Slow {
511        fn call(&mut self, req: &ModelRequest) -> Result<ModelReturn, AxError> {
512            self.answer(req, &mut |_| {})
513        }
514
515        fn call_streaming(
516            &mut self,
517            req: &ModelRequest,
518            onto: Increments<'_>,
519        ) -> Result<ModelReturn, AxError> {
520            self.answer(req, onto)
521        }
522    }
523
524    /// A five-minute door in front of a [`Slow`] provider, on the clock
525    /// the provider moves, reading the endpoint's lead through `lead`.
526    fn slow_door(
527        lead: Lead,
528    ) -> (
529        Warmed<impl FnMut() -> Result<TimeMs, AxError>>,
530        Arc<AtomicU64>,
531        Seen,
532    ) {
533        let (now, seen) = (Arc::new(AtomicU64::new(USED_AT)), Seen::default());
534        let provider = Slow {
535            now: now.clone(),
536            seen: seen.clone(),
537        };
538        let read = now.clone();
539        let clock = move || Ok(TimeMs::new(read.load(Ordering::SeqCst)));
540        let warmed = Warmed::new(Box::new(provider), KeepWarm::FiveMinute, clock, lead);
541        (warmed, now, seen)
542    }
543
544    #[test]
545    fn a_long_call_leaves_its_renewal_due_a_lifetime_after_it_was_sent() {
546        let (mut warmed, now, seen) = slow_door(Lead::default());
547        warmed.call(&request()).unwrap();
548        let landed = now.load(Ordering::SeqCst);
549        assert_eq!(warmed.next_due(), Some(USED_AT + TTL_MS - MARGIN_MS));
550        assert!(warmed.renew_due(landed).refused.is_none());
551        assert_eq!(
552            *seen.lock().unwrap(),
553            vec![request()],
554            "renewed at the landing"
555        );
556
557        let due = warmed.next_due().unwrap();
558        now.store(due, Ordering::SeqCst);
559        assert!(warmed.renew_due(due).refused.is_none());
560        let sent_again = now.load(Ordering::SeqCst);
561        warmed.call(&request()).unwrap();
562        assert_eq!(
563            warmed.next_due(),
564            Some(sent_again + TTL_MS - RENEWAL_MS - MARGIN_MS)
565        );
566    }
567
568    #[test]
569    fn a_streamed_call_reads_the_lead_at_its_first_word() {
570        let (mut warmed, _, _) = slow_door(Lead::default());
571        let mut words = 0;
572        warmed
573            .call_streaming(&request(), &mut |_| words += 1)
574            .unwrap();
575        assert_eq!(
576            (words, warmed.next_due()),
577            (1, Some(USED_AT + TTL_MS - FIRST_WORD_MS - MARGIN_MS))
578        );
579    }
580
581    #[test]
582    fn doors_to_one_endpoint_share_its_lead() {
583        let lead = Lead::default();
584        let (mut renewing, now, _) = slow_door(lead.clone());
585        let (mut waiting, _, _) = slow_door(lead);
586        waiting.call(&request()).unwrap();
587        renewing.call(&request()).unwrap();
588        let due = renewing.next_due().unwrap();
589        now.store(due, Ordering::SeqCst);
590        assert!(renewing.renew_due(due).refused.is_none());
591        assert_eq!(
592            waiting.next_due(),
593            Some(USED_AT + TTL_MS - MARGIN_MS - RENEWAL_MS)
594        );
595    }
596}
runtime/prefix/warmth.rs

runtime/prefix/
warmth.rs
