lean_ctx_proxy/
cache_attribution.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Prompt-cache miss attribution (#986, cache-economics telemetry).
3//!
4//! The proxy already keeps the client-cached prefix byte-stable (#448) and can
5//! re-seed a leaner one on a cold resume (#480). What it could not yet *measure*
6//! is **why** a turn fails to hit the provider prompt-cache — the single most
7//! actionable cache signal. There are only two causes, and they want opposite
8//! fixes:
9//!
10//! - **TTL lapse** — the cacheable prefix is byte-identical to last turn, but the
11//!   idle gap exceeded the provider's cache TTL, so the entry expired. The fix is
12//!   the cold-prefix repack (#480) / longer TTL, never a prefix change.
13//! - **Prefix change** — the cacheable prefix is *different* from last turn, so
14//!   the provider re-writes from the first changed byte regardless of timing. The
15//!   fix is to stop mutating the prefix (a moving system prompt, an edited earlier
16//!   turn, volatile fields — see the cache-aligner #940/#974).
17//!
18//! This module classifies every anchored turn (`cached > 0`) into one of four
19//! outcomes by comparing the `cached_prefix_hash` and idle time against the
20//! conversation's previous turn, and exposes cumulative gauges on `/status`. It
21//! is **measurement-only** — the request body is never touched — and gated behind
22//! the opt-in `proxy.cache_policy`, so a default proxy pays nothing.
23
24use std::collections::HashMap;
25use std::sync::{Mutex, OnceLock};
26use std::time::{SystemTime, UNIX_EPOCH};
27
28use serde_json::Value;
29
30use super::cold_prefix;
31
32pub use lean_ctx::core::proxy_cache_attribution::{CacheAttribution, CacheOutcome, snapshot};
33
34/// Hard cap on tracked conversations so a long-lived proxy can't grow the
35/// last-prefix map without bound; the oldest entry is evicted past this.
36const MAX_TRACKED: usize = 4096;
37
38/// Previous-turn record for one conversation: the cacheable-prefix hash and the
39/// Unix-seconds timestamp it was last seen.
40#[derive(Debug, Clone, Copy)]
41struct PrefixState {
42    prefix_hash: u64,
43    last_touch: u64,
44}
45
46fn store() -> &'static Mutex<HashMap<u64, PrefixState>> {
47    static STORE: OnceLock<Mutex<HashMap<u64, PrefixState>>> = OnceLock::new();
48    STORE.get_or_init(|| Mutex::new(HashMap::new()))
49}
50
51fn now_secs() -> u64 {
52    SystemTime::now()
53        .duration_since(UNIX_EPOCH)
54        .map_or(0, |d| d.as_secs())
55}
56
57/// Pure classification of a turn's cache outcome. `prev` is the conversation's
58/// previous `(prefix_hash, last_touch)`, `curr_hash` this turn's cacheable-prefix
59/// hash, `now`/`ttl_secs` the idle clock. Pure (no globals, no I/O) so the
60/// TTL-vs-prefix decision is unit-tested independently of the live store.
61#[must_use]
62pub fn classify(prev: Option<(u64, u64)>, curr_hash: u64, now: u64, ttl_secs: u64) -> CacheOutcome {
63    match prev {
64        None => CacheOutcome::ColdStart,
65        Some((prev_hash, last_touch)) => {
66            if prev_hash != curr_hash {
67                CacheOutcome::PrefixChange
68            } else if now.saturating_sub(last_touch) > ttl_secs {
69                CacheOutcome::TtlLapse
70            } else {
71                CacheOutcome::WarmReuse
72            }
73        }
74    }
75}
76
77fn evict_oldest(map: &mut HashMap<u64, PrefixState>) {
78    if let Some(oldest) = map
79        .iter()
80        .min_by_key(|(_, s)| s.last_touch)
81        .map(|(k, _)| *k)
82    {
83        map.remove(&oldest);
84    }
85}
86
87/// Attribute this request's cache outcome and record it, updating the
88/// conversation's last-seen prefix baseline for the next turn. Only anchored
89/// turns (`cached > 0`) are attributable; an unanchored turn returns `None`
90/// (the cache-aligner telemetry #940 covers "client never anchors"). The caller
91/// owns the opt-in gate — this only runs when `proxy.cache_policy` is enabled.
92pub fn record_request(messages: &[Value], cached: usize) -> Option<CacheOutcome> {
93    let conv_key = cold_prefix::conversation_key(messages)?;
94    let curr_hash = cold_prefix::cached_prefix_hash(messages, cached)?;
95    let ttl = cold_prefix::resolved_ttl_secs(messages, cached).unwrap_or(0);
96    let now = now_secs();
97
98    let outcome = {
99        let mut map = store()
100            .lock()
101            .unwrap_or_else(std::sync::PoisonError::into_inner);
102        let prev = map.get(&conv_key).map(|s| (s.prefix_hash, s.last_touch));
103        let outcome = classify(prev, curr_hash, now, ttl);
104        map.insert(
105            conv_key,
106            PrefixState {
107                prefix_hash: curr_hash,
108                last_touch: now,
109            },
110        );
111        if map.len() > MAX_TRACKED {
112            evict_oldest(&mut map);
113        }
114        outcome
115    };
116    lean_ctx::core::proxy_cache_attribution::record_outcome(outcome);
117    Some(outcome)
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123    use serde_json::json;
124
125    #[test]
126    fn classify_distinguishes_ttl_lapse_from_prefix_change() {
127        // No prior turn → cold start.
128        assert_eq!(classify(None, 1, 100, 300), CacheOutcome::ColdStart);
129        // Same prefix, within TTL → warm reuse (should hit).
130        assert_eq!(
131            classify(Some((1, 100)), 1, 350, 300),
132            CacheOutcome::WarmReuse
133        );
134        // Same prefix, idle past TTL → ttl lapse.
135        assert_eq!(
136            classify(Some((1, 100)), 1, 500, 300),
137            CacheOutcome::TtlLapse
138        );
139        // Different prefix → prefix change, regardless of timing.
140        assert_eq!(
141            classify(Some((1, 100)), 2, 110, 300),
142            CacheOutcome::PrefixChange
143        );
144        // Different prefix wins even past TTL (the mutation is the root cause).
145        assert_eq!(
146            classify(Some((1, 100)), 2, 9999, 300),
147            CacheOutcome::PrefixChange
148        );
149    }
150
151    fn anchored(first_text: &str) -> Vec<Value> {
152        vec![
153            json!({"role": "user", "content": [
154                {"type": "text", "text": first_text, "cache_control": {"type": "ephemeral"}}
155            ]}),
156            json!({"role": "assistant", "content": "ok"}),
157        ]
158    }
159
160    #[test]
161    fn unanchored_turn_is_not_attributed() {
162        let msgs = anchored("unanchored-attribution-test");
163        // cached == 0: nothing anchored to attribute.
164        assert_eq!(record_request(&msgs, 0), None);
165    }
166
167    #[test]
168    fn first_anchored_turn_is_cold_start_then_warm() {
169        let msgs = anchored("cold-then-warm-attribution-test");
170        assert_eq!(record_request(&msgs, 1), Some(CacheOutcome::ColdStart));
171        // Immediately again (idle ~0, prefix identical) → warm reuse.
172        assert_eq!(record_request(&msgs, 1), Some(CacheOutcome::WarmReuse));
173    }
174
175    #[test]
176    fn prefix_change_detected_with_stable_head() {
177        // Stable head (conversation key), but the cached prefix spans two messages
178        // and the second one changes — a true mid-prefix mutation.
179        let head = json!({"role": "user", "content": [
180            {"type": "text", "text": "stable-head-attribution", "cache_control": {"type": "ephemeral"}}
181        ]});
182        let v1 = vec![
183            head.clone(),
184            json!({"role": "assistant", "content": "answer one"}),
185        ];
186        let v2 = vec![
187            head,
188            json!({"role": "assistant", "content": "answer two CHANGED"}),
189        ];
190
191        assert_eq!(record_request(&v1, 2), Some(CacheOutcome::ColdStart));
192        assert_eq!(record_request(&v2, 2), Some(CacheOutcome::PrefixChange));
193    }
194}
lean_ctx_proxy/cache_attribution.rs

lean_ctx_proxy/
cache_attribution.rs
