返回 CodeWhale
prefix_cache.rs
根目录 / crates / core / src / prefix_cache.rs
1 //! Prefix-cache stability manager (inspired by Reasonix's Pillar 1).
2 //!
3 //! DeepSeek's automatic prefix caching activates only when the *exact*
4 //! byte prefix of a request matches the prior request. Any system-prompt
5 //! drift, tool-list reordering, or message-rewriting busts the cache
6 //! for every token after the changed byte.
7 //!
8 //! This module provides a `PrefixStabilityManager` that:
9 //!
10 //! 1. **Fingerprints** the immutable prefix (system prompt + tool specs)
11 //! at session start, using SHA-256 for strong collision resistance.
12 //! 2. **Verifies** the current prefix against the pinned fingerprint before
13 //! every request.
14 //! 3. **Attributes** every change: a header change the engine declared
15 //! (`/model`, `/mode`, goal edits, MCP or deferred-tool activation,
16 //! session sync) re-pins under a logged reason; an undeclared change is
17 //! *drift* — it is recorded and reported, and the original pin stays so
18 //! later checks keep counting the miss instead of quietly adopting it.
19 //! 4. **Emits events** so the TUI can surface stability to the user.
20 //!
21 //! The invariant this guards: after session start, system and tools are
22 //! frozen bytes; history only grows; a miss is allowed only when we log why.
23 //!
24 //! ## Three-region model (from Reasonix)
25 //!
26 //! ```text
27 //! ┌─────────────────────────────────────────┐
28 //! │ IMMUTABLE PREFIX │ ← fixed for session
29 //! │ system + tool_specs │ cache hit candidate
30 //! ├─────────────────────────────────────────┤
31 //! │ APPEND-ONLY HISTORY │ ← grows monotonically
32 //! │ [assistant₁][tool₁][assistant₂]... │ preserves prefix of prior turns
33 //! ├─────────────────────────────────────────┤
34 //! │ LATEST USER TURN │ ← the only new content per request
35 //! └─────────────────────────────────────────┘
36 //! ```
37
38 use std::collections::hash_map::DefaultHasher;
39 use std::collections::{HashMap, VecDeque};
40 use std::hash::{Hash, Hasher};
41
42 use serde::{Deserialize, Serialize};
43
44 use crate::request::{SystemPrompt, Tool};
45
46 /// A snapshot of the immutable prefix's fingerprint.
47 ///
48 /// Matching hashes show stable system text and the normalized OpenAI-style
49 /// tool catalog. They do not measure provider cache hits or fingerprint every
50 /// provider-specific wire transformation; request replay tests cover those.
51 #[derive(Debug, Clone, Serialize, Deserialize)]
52 pub struct PrefixFingerprint {
53 /// SHA-256 of the system prompt text.
54 pub system_sha256: String,
55 /// SHA-256 of the full tool catalog JSON (names, descriptions, schemas).
56 pub tools_sha256: String,
57 /// SHA-256 of system_sha256 ++ tools_sha256 (combined).
58 pub combined_sha256: String,
59 }
60
61 impl PrefixFingerprint {
62 /// Compute a fingerprint from system prompt text and tool list.
63 ///
64 /// Tools are serialized to the same JSON shape the chat API receives
65 /// (`type`, `name`, `description`, `parameters`, `strict`) **in provider
66 /// wire order**, then SHA-256 hashed. Order is load-bearing: the provider
67 /// KV cache is order-sensitive, so sorting before hashing could read a
68 /// false-green while the real prefix cache misses every turn (ops C1,
69 /// DSH invariant). This also catches schema/description drift that
70 /// actually affects the API prefix, while ignoring internal-only fields
71 /// like `allowed_callers` (#2264).
72 ///
73 /// This entry point shares a process-local [`ToolCatalogCache`] with
74 /// every other call, so a stable tool set (the common case after the
75 /// first turn of a session) avoids the per-tool JSON serialization
76 /// and join entirely. Callers that hold their own cache — e.g.
77 /// [`PrefixStabilityManager`] — should use
78 /// [`Self::compute_with_tool_cache`] to share *that* cache instead
79 /// and avoid the thread-local lookup.
80 #[cfg(test)]
81 pub fn compute(system_text: &str, tools: Option<&[Tool]>) -> Self {
82 let mut cache = ToolCatalogCache::new();
83 Self::compute_with_tool_cache(system_text, tools, &mut cache)
84 }
85
86 /// Compute a fingerprint while reusing a [`ToolCatalogCache`] for the
87 /// tool-side work. The cache holds the joined+SHA-256'd catalog
88 /// under a content-derived identity so the per-tool JSON serialization
89 /// and the join only run on the first call for a given tool set.
90 ///
91 /// On a cache hit this function avoids the entire tool serialization
92 /// path, which can be 100+ microseconds for a 60-tool catalog.
93 pub fn compute_with_tool_cache(
94 system_text: &str,
95 tools: Option<&[Tool]>,
96 cache: &mut ToolCatalogCache,
97 ) -> Self {
98 let system_sha256 = sha256_hex(system_text.as_bytes());
99
100 let tools_sha256 = match tools {
101 Some(tools) if !tools.is_empty() => {
102 // `fingerprint_for` consults the cache first; on a hit
103 // it returns the pre-computed hex digest directly.
104 cache.fingerprint_for(tools).sha256_hex
105 }
106 _ => sha256_hex(b""),
107 };
108
109 let combined = format!("{system_sha256}:{tools_sha256}");
110 let combined_sha256 = sha256_hex(combined.as_bytes());
111 Self {
112 system_sha256,
113 tools_sha256,
114 combined_sha256,
115 }
116 }
117 }
118
119 /// A change record describing what drifted in the prefix.
120 #[derive(Debug, Clone, Serialize, Deserialize)]
121 pub struct PrefixChange {
122 /// The old fingerprint (before the change).
123 pub old: PrefixFingerprint,
124 /// The new fingerprint (after the change).
125 pub new: PrefixFingerprint,
126 /// Whether the system prompt component changed.
127 pub system_changed: bool,
128 /// Whether the tool set component changed.
129 pub tools_changed: bool,
130 }
131
132 #[allow(dead_code)]
133 impl PrefixChange {
134 /// Returns a human-readable description of what changed.
135 pub fn description(&self) -> String {
136 let mut parts = Vec::new();
137 if self.system_changed {
138 parts.push("system prompt");
139 }
140 if self.tools_changed {
141 parts.push("tool set");
142 }
143 if parts.is_empty() {
144 return "unknown (fingerprint mismatch but no component detected)".to_string();
145 }
146 format!("prefix cache invalidated: {} changed", parts.join(" and "))
147 }
148
149 /// Returns a short label for TUI chip display.
150 pub fn label(&self) -> &'static str {
151 if self.system_changed && self.tools_changed {
152 "sys+tools"
153 } else if self.system_changed {
154 "sys"
155 } else if self.tools_changed {
156 "tools"
157 } else {
158 "prefix"
159 }
160 }
161 }
162
163 /// Monitors and manages prefix-cache stability across turns.
164 ///
165 /// This is the core abstraction, mirroring Reasonix's `ImmutablePrefix`
166 /// concept but adapted to CodeWhale's existing architecture where the
167 /// system prompt is rebuilt each turn and tools are registered at startup.
168 ///
169 /// Usage:
170 /// ```ignore
171 /// let mgr = PrefixStabilityManager::new(system_text, tools);
172 /// if mgr.check_and_update(system_text, tools) {
173 /// println!("Prefix is stable (cache-friendly)");
174 /// } else {
175 /// let change = mgr.last_change().unwrap();
176 /// println!("Prefix drifted: {}", change.description());
177 /// }
178 /// ```
179 #[derive(Debug, Clone)]
180 pub struct PrefixStabilityManager {
181 /// The pinned fingerprint from session start or last stabilization.
182 pinned: Option<PrefixFingerprint>,
183 /// The most recent fingerprint (computed during last check).
184 current: Option<PrefixFingerprint>,
185 /// The last detected change, if any.
186 last_change: Option<PrefixChange>,
187 /// Total number of prefix changes detected this session.
188 change_count: u64,
189 /// Total number of stability checks performed.
190 check_count: u64,
191 /// Why the current pin exists: `initial`, `resume`, or `change:<what>`.
192 pin_reason: Option<String>,
193 /// Bounded log of every attributed change and every undeclared drift.
194 history: VecDeque<PrefixHistoryEntry>,
195 /// Explanation of the most recent expected cache miss (a declared header
196 /// change, a history reset such as compaction, or undeclared drift).
197 last_miss_reason: Option<String>,
198 /// `<context_update>` snapshots appended this session (workspace drift
199 /// delivered as history, with the pinned header untouched).
200 context_update_count: u64,
201 /// Process-local cache for the tool-catalog JSON serialization. Avoids
202 /// re-running `tool_to_api_json` + join on every `check_and_update`
203 /// when the tool set is unchanged (the common case once tools are
204 /// registered at session start).
205 tool_catalog_cache: ToolCatalogCache,
206 }
207
208 /// Maximum retained [`PrefixHistoryEntry`] records per session.
209 const PREFIX_HISTORY_CAP: usize = 32;
210
211 /// One attributed prefix event: a declared header change (re-pinned) or an
212 /// undeclared drift (pin kept).
213 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
214 pub struct PrefixHistoryEntry {
215 /// `change:<what>` for declared header changes, `drift:<component>` for
216 /// undeclared changes, `reset:<what>` for history resets.
217 pub reason: String,
218 /// Whether the pin was replaced by this event.
219 pub repinned: bool,
220 /// Combined SHA-256 before the event.
221 pub from_sha256: String,
222 /// Combined SHA-256 the request actually carried.
223 pub to_sha256: String,
224 }
225
226 /// Outcome of [`PrefixStabilityManager::check`].
227 #[derive(Debug, Clone)]
228 pub enum PrefixCheck {
229 /// The request prefix matches the pin byte-for-byte.
230 Stable,
231 /// The prefix changed and the engine declared why; the pin moved.
232 Repinned {
233 reason: String,
234 change: PrefixChange,
235 },
236 /// The prefix changed with no declared reason. The pin did NOT move.
237 Drift { change: PrefixChange },
238 }
239
240 /// Default capacity for the tool-catalog serialization cache. Sized for
241 /// "session + 1 or 2 forked subagent catalogs" without unbounded growth.
242 const TOOL_CATALOG_CACHE_CAPACITY: usize = 8;
243
244 /// Bounded LRU cache of `(tool_set_identity) -> sha256_hex`.
245 ///
246 /// The cache key is a content-derived `u64` hash of the tool list (length +
247 /// per-tool `name` + `description` + serialized `input_schema`). On a hit,
248 /// `PrefixFingerprint::compute` skips the per-tool JSON serialization and
249 /// the join — a workload that can be 100+ microseconds for a
250 /// 60-tool catalog. On a miss, the work runs once and only the digest is
251 /// retained (#3854); the joined catalog string is ephemeral.
252 #[derive(Debug, Default, Clone)]
253 pub struct ToolCatalogCache {
254 by_identity: HashMap<u64, CachedCatalog>,
255 insertion_order: VecDeque<u64>,
256 capacity: usize,
257 }
258
259 /// One entry in [`ToolCatalogCache`]. Production only needs the pre-computed
260 /// SHA-256 digest of the in-order joined catalog.
261 #[derive(Debug, Clone)]
262 pub struct CachedCatalog {
263 /// SHA-256 hex digest of the newline-joined, in-order tool-catalog JSON.
264 pub sha256_hex: String,
265 }
266
267 impl ToolCatalogCache {
268 /// Create a cache with the default capacity.
269 #[must_use]
270 pub fn new() -> Self {
271 Self::with_capacity(TOOL_CATALOG_CACHE_CAPACITY)
272 }
273
274 /// Create a cache that holds at most `capacity` tool-set entries.
275 /// Smaller values save memory at the cost of more cache misses.
276 #[must_use]
277 pub fn with_capacity(capacity: usize) -> Self {
278 let cap = capacity.max(1);
279 Self {
280 by_identity: HashMap::with_capacity(cap),
281 insertion_order: VecDeque::with_capacity(cap),
282 capacity: cap,
283 }
284 }
285
286 /// Compute (or recall) the joined-and-hashed tool catalog for `tools`.
287 /// The cache is keyed on a content-derived `u64` identity so two `&[Tool]`
288 /// slices with the same payloads — in the same order — hit the same entry.
289 pub fn fingerprint_for(&mut self, tools: &[Tool]) -> CachedCatalog {
290 let identity = tool_set_identity(tools);
291 if let Some(cached) = self.by_identity.get(&identity) {
292 return cached.clone();
293 }
294
295 // Miss: serialize, join, hash — in wire order, never sorted. The
296 // provider cache is order-sensitive; a sorted fingerprint can say
297 // "stable" while the real prefix misses every turn (ops C1). Keep
298 // only the digest in the cache — the joined string is not needed on
299 // the hot path (#3854).
300 let serialized: Vec<String> = tools.iter().filter_map(tool_to_api_json).collect();
301 let joined = serialized.join("\n");
302 let entry = CachedCatalog {
303 sha256_hex: sha256_hex(joined.as_bytes()),
304 };
305
306 if self.by_identity.len() >= self.capacity
307 && let Some(oldest) = self.insertion_order.pop_front()
308 {
309 self.by_identity.remove(&oldest);
310 }
311 self.by_identity.insert(identity, entry.clone());
312 self.insertion_order.push_back(identity);
313 entry
314 }
315
316 /// Drop every cached entry. Used by tool-registry mutation paths
317 /// (e.g. plugin hot-reload, MCP attach) when the caller cannot
318 /// easily prove the tool set is unchanged.
319 #[allow(dead_code)] // observability; called by /cache flush and tests
320 pub fn invalidate(&mut self) {
321 self.by_identity.clear();
322 self.insertion_order.clear();
323 }
324
325 /// Returns the number of cached entries.
326 #[must_use]
327 pub fn len(&self) -> usize {
328 self.by_identity.len()
329 }
330
331 /// Returns `true` if the cache has no entries.
332 #[allow(dead_code)] // observability; surfaced via /status
333 #[must_use]
334 pub fn is_empty(&self) -> bool {
335 self.by_identity.is_empty()
336 }
337
338 /// Returns `(current_entries, capacity)` for observability. Surfaced via
339 /// the `/status` chip in a follow-up; tests exercise the path.
340 #[allow(dead_code)] // surfaced via /status in a follow-up; tests exercise it
341 #[must_use]
342 pub fn stats(&self) -> (usize, usize) {
343 (self.len(), self.capacity)
344 }
345 }
346
347 /// Content-derived identity for a tool slice. Order-sensitive: two slices
348 /// with the same tools in different orders produce different identities.
349 /// (The downstream fingerprint itself is order-insensitive — the sort in
350 /// `fingerprint_for` takes care of that — but the cache key matches the
351 /// input order so re-registration of the same set in the same order hits.)
352 fn tool_set_identity(tools: &[Tool]) -> u64 {
353 let mut hasher = DefaultHasher::new();
354 tools.len().hash(&mut hasher);
355 for tool in tools {
356 tool.name.hash(&mut hasher);
357 tool.description.hash(&mut hasher);
358 // `strict` participates in `tool_to_api_json` output (it is part of
359 // the wire-format the chat API receives), so it MUST be part of the
360 // identity. Omitting it lets two semantically different catalogs
361 // collide and serve a stale fingerprint.
362 tool.strict.hash(&mut hasher);
363 // Walk the schema JSON directly instead of materializing it as a
364 // String. For a 60-tool catalog this saves ~25-40 KB of allocation
365 // on every cache miss.
366 hash_json_value(&tool.input_schema, &mut hasher);
367 }
368 hasher.finish()
369 }
370
371 /// Fold a `serde_json::Value` into the hasher without allocating a
372 /// `String`. Numeric variants are hashed via their bit pattern so `1` and
373 /// `1.0` produce distinct identities (matching the JSON spec).
374 fn hash_json_value<H: Hasher>(value: &serde_json::Value, state: &mut H) {
375 match value {
376 serde_json::Value::Null => 0u8.hash(state),
377 serde_json::Value::Bool(b) => {
378 1u8.hash(state);
379 b.hash(state);
380 }
381 serde_json::Value::Number(n) => {
382 2u8.hash(state);
383 if let Some(i) = n.as_i64() {
384 i.hash(state);
385 } else if let Some(u) = n.as_u64() {
386 u.hash(state);
387 } else if let Some(f) = n.as_f64() {
388 f.to_bits().hash(state);
389 }
390 }
391 serde_json::Value::String(s) => {
392 3u8.hash(state);
393 s.hash(state);
394 }
395 serde_json::Value::Array(arr) => {
396 4u8.hash(state);
397 arr.len().hash(state);
398 for v in arr {
399 hash_json_value(v, state);
400 }
401 }
402 serde_json::Value::Object(obj) => {
403 5u8.hash(state);
404 obj.len().hash(state);
405 // Iterate by sorted key so `{"a":1,"b":2}` and `{"b":2,"a":1}`
406 // collide — the wire format already canonicalizes via the
407 // `serde_json` Map ordering, but a defensively-sorted view
408 // future-proofs against schema serializers that emit
409 // declaration order.
410 let mut entries: Vec<(&String, &serde_json::Value)> = obj.iter().collect();
411 entries.sort_by(|a, b| a.0.cmp(b.0));
412 for (k, v) in entries {
413 k.hash(state);
414 hash_json_value(v, state);
415 }
416 }
417 }
418 }
419
420 /// Process-local fallback cache used by `PrefixFingerprint::compute`
421 /// (when available). Callers that maintain their own cache (e.g.
422 /// [`PrefixStabilityManager`]) should prefer
423 /// [`PrefixFingerprint::compute_with_tool_cache`] and pass the cache in
424 /// directly, both to share state and to avoid the thread-local lookup
425 /// on the hot path.
426 #[allow(dead_code)]
427 impl PrefixStabilityManager {
428 /// Create a new manager and immediately pin the first fingerprint.
429 pub fn new(system_text: &str, tools: Option<&[Tool]>) -> Self {
430 let mut cache = ToolCatalogCache::new();
431 let fp = PrefixFingerprint::compute_with_tool_cache(system_text, tools, &mut cache);
432 Self {
433 pinned: Some(fp.clone()),
434 current: Some(fp),
435 last_change: None,
436 change_count: 0,
437 check_count: 0,
438 pin_reason: Some("initial".to_string()),
439 history: VecDeque::new(),
440 last_miss_reason: None,
441 context_update_count: 0,
442 tool_catalog_cache: cache,
443 }
444 }
445
446 /// Create a manager in "unpinned" state — no initial fingerprint.
447 /// Call `pin()` or `check_and_update()` to establish the baseline.
448 pub fn new_unpinned() -> Self {
449 Self {
450 pinned: None,
451 current: None,
452 last_change: None,
453 change_count: 0,
454 check_count: 0,
455 pin_reason: None,
456 history: VecDeque::new(),
457 last_miss_reason: None,
458 context_update_count: 0,
459 tool_catalog_cache: ToolCatalogCache::new(),
460 }
461 }
462
463 /// Explicitly pin a fingerprint, replacing any prior pinned state.
464 /// Returns `true` if this is the first pin, or `false` if replacing.
465 /// Note: does NOT increment `check_count` — that counter is reserved
466 /// for `check_and_update` calls so `stability_ratio()` stays accurate.
467 pub fn pin(&mut self, system_text: &str, tools: Option<&[Tool]>) -> bool {
468 self.pin_with_reason(system_text, tools, "initial")
469 }
470
471 /// Pin under an explicit reason (`initial`, `resume`, `change:<what>`).
472 pub fn pin_with_reason(
473 &mut self,
474 system_text: &str,
475 tools: Option<&[Tool]>,
476 reason: &str,
477 ) -> bool {
478 let fp = PrefixFingerprint::compute_with_tool_cache(
479 system_text,
480 tools,
481 &mut self.tool_catalog_cache,
482 );
483 let was_unpinned = self.pinned.is_none();
484 self.pinned = Some(fp.clone());
485 self.current = Some(fp);
486 self.pin_reason = Some(reason.to_string());
487 was_unpinned
488 }
489
490 /// Record an expected miss that is not a header change (compaction,
491 /// `/clear`, an edited turn). The pin is untouched; the reason is kept so
492 /// `/cache stats` can explain the next low hit-rate turn.
493 pub fn note_history_reset(&mut self, what: &str) {
494 let hash = self
495 .pinned
496 .as_ref()
497 .map(|fp| fp.combined_sha256.clone())
498 .unwrap_or_default();
499 self.push_history(PrefixHistoryEntry {
500 reason: format!("reset:{what}"),
501 repinned: false,
502 from_sha256: hash.clone(),
503 to_sha256: hash,
504 });
505 self.last_miss_reason = Some(format!("reset:{what}"));
506 }
507
508 /// Record that workspace drift was delivered as a `<context_update>`
509 /// history append. Not a miss: the pin and the prefix are unchanged.
510 pub fn note_context_update(&mut self) {
511 self.context_update_count = self.context_update_count.saturating_add(1);
512 let hash = self
513 .pinned
514 .as_ref()
515 .map(|fp| fp.combined_sha256.clone())
516 .unwrap_or_default();
517 self.push_history(PrefixHistoryEntry {
518 reason: "context_update".to_string(),
519 repinned: false,
520 from_sha256: hash.clone(),
521 to_sha256: hash,
522 });
523 }
524
525 /// Number of `<context_update>` snapshots appended this session.
526 pub fn context_update_count(&self) -> u64 {
527 self.context_update_count
528 }
529
530 fn push_history(&mut self, entry: PrefixHistoryEntry) {
531 if self.history.len() >= PREFIX_HISTORY_CAP {
532 self.history.pop_front();
533 }
534 self.history.push_back(entry);
535 }
536
537 /// Verify the request prefix against the pin, attributing any change.
538 ///
539 /// `declared_change` names a header change the engine performed on
540 /// purpose (`model`, `mode`, `goal`, `tools:+web_fetch`, `mcp`, …). When
541 /// the prefix changed and a reason is declared, the pin moves and the
542 /// change is logged as `change:<reason>`. When it changed with no
543 /// declared reason, that is drift: it is logged as `drift:<component>`
544 /// and the pin stays put, so the same undeclared prefix keeps counting as
545 /// a miss instead of becoming the new baseline.
546 pub fn check(
547 &mut self,
548 system_text: &str,
549 tools: Option<&[Tool]>,
550 declared_change: Option<&str>,
551 ) -> PrefixCheck {
552 let fp = PrefixFingerprint::compute_with_tool_cache(
553 system_text,
554 tools,
555 &mut self.tool_catalog_cache,
556 );
557 let old_fp = self.current.replace(fp.clone());
558 self.check_count += 1;
559
560 let pinned = match &self.pinned {
561 Some(p) => p.clone(),
562 None => {
563 self.pinned = Some(fp);
564 self.pin_reason = Some(declared_change.unwrap_or("initial").to_string());
565 self.last_change = None;
566 return PrefixCheck::Stable;
567 }
568 };
569
570 if fp.combined_sha256 == pinned.combined_sha256 {
571 return PrefixCheck::Stable;
572 }
573
574 let old = old_fp.unwrap_or_else(|| pinned.clone());
575 let system_changed = fp.system_sha256 != pinned.system_sha256;
576 let tools_changed = fp.tools_sha256 != pinned.tools_sha256;
577 let change = PrefixChange {
578 old,
579 new: fp.clone(),
580 system_changed,
581 tools_changed,
582 };
583 self.last_change = Some(change.clone());
584 self.change_count += 1;
585
586 match declared_change {
587 Some(reason) => {
588 let reason = format!("change:{reason}");
589 self.push_history(PrefixHistoryEntry {
590 reason: reason.clone(),
591 repinned: true,
592 from_sha256: pinned.combined_sha256.clone(),
593 to_sha256: fp.combined_sha256.clone(),
594 });
595 self.last_miss_reason = Some(reason.clone());
596 self.pinned = Some(fp);
597 self.pin_reason = Some(reason.clone());
598 PrefixCheck::Repinned { reason, change }
599 }
600 None => {
601 let reason = format!("drift:{}", change.label());
602 self.push_history(PrefixHistoryEntry {
603 reason: reason.clone(),
604 repinned: false,
605 from_sha256: pinned.combined_sha256.clone(),
606 to_sha256: fp.combined_sha256.clone(),
607 });
608 self.last_miss_reason = Some(reason);
609 PrefixCheck::Drift { change }
610 }
611 }
612 }
613
614 /// Why the current pin exists.
615 pub fn pin_reason(&self) -> Option<&str> {
616 self.pin_reason.as_deref()
617 }
618
619 /// Attributed change/drift/reset history, oldest first.
620 pub fn history(&self) -> impl Iterator<Item = &PrefixHistoryEntry> {
621 self.history.iter()
622 }
623
624 /// Explanation of the most recent expected miss, if any.
625 pub fn last_miss_reason(&self) -> Option<&str> {
626 self.last_miss_reason.as_deref()
627 }
628
629 /// Check whether the current prefix matches the pinned fingerprint
630 /// without a declared header change.
631 ///
632 /// - `Ok(true)` when the prefix is stable (or this was the first pin).
633 /// - `Err(change)` when the prefix drifted. The pin is **kept**: an
634 /// undeclared change never becomes the new baseline. Use [`Self::check`]
635 /// with a declared reason to move the pin on purpose.
636 pub fn check_and_update(
637 &mut self,
638 system_text: &str,
639 tools: Option<&[Tool]>,
640 ) -> Result<bool, Box<PrefixChange>> {
641 match self.check(system_text, tools, None) {
642 PrefixCheck::Stable => Ok(true),
643 PrefixCheck::Repinned { change, .. } | PrefixCheck::Drift { change } => {
644 Err(Box::new(change))
645 }
646 }
647 }
648
649 /// Returns the most recent prefix change, if any.
650 pub fn last_change(&self) -> Option<&PrefixChange> {
651 self.last_change.as_ref()
652 }
653
654 /// Returns the pinned fingerprint.
655 pub fn pinned_fingerprint(&self) -> Option<&PrefixFingerprint> {
656 self.pinned.as_ref()
657 }
658
659 /// Returns the current (most recently computed) fingerprint.
660 pub fn current_fingerprint(&self) -> Option<&PrefixFingerprint> {
661 self.current.as_ref()
662 }
663
664 /// Returns the total number of prefix changes detected.
665 pub fn change_count(&self) -> u64 {
666 self.change_count
667 }
668
669 /// Returns the total number of stability checks performed.
670 pub fn check_count(&self) -> u64 {
671 self.check_count
672 }
673
674 /// Returns the prefix stability rate as a fraction (0.0 – 1.0).
675 /// 1.0 means the prefix has never changed. Returns 1.0 when no
676 /// checks have been performed (to avoid division by zero).
677 pub fn stability_ratio(&self) -> f64 {
678 if self.check_count == 0 {
679 1.0
680 } else {
681 let stable_checks = self.check_count - self.change_count;
682 stable_checks as f64 / self.check_count as f64
683 }
684 }
685
686 /// Returns a human-readable stability summary.
687 pub fn summary(&self) -> String {
688 let pct = self.stability_ratio() * 100.0;
689 let pinned_short = self
690 .pinned
691 .as_ref()
692 .map(|fp| {
693 if fp.combined_sha256.len() >= 12 {
694 &fp.combined_sha256[..12]
695 } else {
696 &fp.combined_sha256
697 }
698 })
699 .unwrap_or("none");
700
701 format!(
702 "Prefix stability: {pct:.1}% ({stable}/{total} checks stable) | fingerprint: {pinned_short} | changes: {changes}",
703 pct = pct,
704 stable = self.check_count.saturating_sub(self.change_count),
705 total = self.check_count,
706 pinned_short = pinned_short,
707 changes = self.change_count,
708 )
709 }
710 }
711
712 /// Serialize a tool to the same JSON shape the chat API receives,
713 /// excluding internal-only fields like `allowed_callers`, `defer_loading`,
714 /// `input_examples`, and `cache_control` that are never sent to DeepSeek.
715 fn tool_to_api_json(tool: &Tool) -> Option<String> {
716 let mut value = serde_json::json!({
717 "type": "function",
718 "function": {
719 "name": tool.name,
720 "description": tool.description,
721 "parameters": tool.input_schema,
722 }
723 });
724 if let Some(strict) = tool.strict
725 && let Some(function) = value.get_mut("function")
726 {
727 function["strict"] = serde_json::json!(strict);
728 }
729 serde_json::to_string(&value).ok()
730 }
731
732 /// Compute the SHA-256 hex digest of a byte slice.
733 fn sha256_hex(bytes: &[u8]) -> String {
734 use sha2::{Digest, Sha256};
735 use std::fmt::Write;
736 let mut hex = String::with_capacity(64);
737 for byte in Sha256::digest(bytes) {
738 let _ = write!(&mut hex, "{byte:02x}");
739 }
740 hex
741 }
742
743 /// Bounded line delta between the session context the model last saw and a
744 /// fresh composition, rendered as a `<context_update>` user-role message.
745 ///
746 /// The pinned system prompt is never rewritten; this is how workspace,
747 /// instruction, skills, memory, and goal drift reaches the model as a normal
748 /// history append. Returns `None` when the two texts are line-identical (only
749 /// whitespace/ordering noise), so no empty update is ever sent.
750 pub const CONTEXT_UPDATE_MAX_LINES: usize = 80;
751 pub const CONTEXT_UPDATE_MAX_BYTES: usize = 6_000;
752
753 pub fn context_update_message(known: &str, current: &str) -> Option<String> {
754 use std::collections::HashMap;
755 let mut counts: HashMap<&str, i64> = HashMap::new();
756 for line in known.lines() {
757 *counts.entry(line).or_insert(0) -= 1;
758 }
759 for line in current.lines() {
760 *counts.entry(line).or_insert(0) += 1;
761 }
762 let mut added: Vec<&str> = Vec::new();
763 for line in current.lines() {
764 if let Some(count) = counts.get_mut(line)
765 && *count > 0
766 {
767 *count -= 1;
768 if !line.trim().is_empty() {
769 added.push(line);
770 }
771 }
772 }
773 let mut removed: Vec<&str> = Vec::new();
774 for line in known.lines() {
775 if let Some(count) = counts.get_mut(line)
776 && *count < 0
777 {
778 *count += 1;
779 if !line.trim().is_empty() {
780 removed.push(line);
781 }
782 }
783 }
784 if added.is_empty() && removed.is_empty() {
785 return None;
786 }
787
788 let mut out = String::from(
789 "<context_update>\nSession context changed since it was pinned; the pinned system \
790 prompt is unchanged. Delta (+ added, - removed):\n",
791 );
792 let mut lines_written = 0usize;
793 let mut truncated = 0usize;
794 for (sign, group) in [("+ ", &added), ("- ", &removed)] {
795 for line in group {
796 if lines_written >= CONTEXT_UPDATE_MAX_LINES
797 || out.len() + sign.len() + line.len() + 1 > CONTEXT_UPDATE_MAX_BYTES
798 {
799 truncated += 1;
800 continue;
801 }
802 out.push_str(sign);
803 out.push_str(line);
804 out.push('\n');
805 lines_written += 1;
806 }
807 }
808 if truncated > 0 {
809 out.push_str(&format!(
810 "(+{truncated} more changed lines; re-read the files you need)\n"
811 ));
812 }
813 out.push_str("</context_update>");
814 Some(out)
815 }
816
817 /// Extract the system prompt text from an optional SystemPrompt,
818 /// returning an owned String. This is used for prefix fingerprinting
819 /// and avoids lifetime/leak issues with the rare SystemPrompt::Blocks case.
820 pub fn system_prompt_text(system: Option<&SystemPrompt>) -> String {
821 match system {
822 Some(SystemPrompt::Text(text)) => text.clone(),
823 Some(SystemPrompt::Blocks(blocks)) => {
824 let mut text = String::new();
825 for block in blocks {
826 text.push_str(&block.text);
827 text.push('\n');
828 }
829 text
830 }
831 None => String::new(),
832 }
833 }
834
835 #[cfg(test)]
836 mod tests {
837 use super::*;
838
839 fn make_tool(name: &str) -> Tool {
840 Tool {
841 name: name.to_string(),
842 description: String::new(),
843 input_schema: serde_json::Value::Null,
844 tool_type: None,
845 allowed_callers: None,
846 defer_loading: None,
847 input_examples: None,
848 strict: None,
849 cache_control: None,
850 }
851 }
852
853 #[test]
854 fn same_prefix_produces_same_fingerprint() {
855 let a = PrefixFingerprint::compute("hello world", None);
856 let b = PrefixFingerprint::compute("hello world", None);
857 assert_eq!(a.combined_sha256, b.combined_sha256);
858 }
859
860 #[test]
861 fn different_system_produces_different_fingerprint() {
862 let a = PrefixFingerprint::compute("hello", None);
863 let b = PrefixFingerprint::compute("world", None);
864 assert_ne!(a.combined_sha256, b.combined_sha256);
865 }
866
867 #[test]
868 fn tool_order_affects_fingerprint() {
869 // The provider KV cache is order-sensitive; sorting before hashing
870 // could read a false-green while the real prefix cache missed every
871 // turn (ops C1).
872 let tools_a = vec![make_tool("read_file"), make_tool("write_file")];
873 let tools_b = vec![make_tool("write_file"), make_tool("read_file")];
874 let a = PrefixFingerprint::compute("system", Some(&tools_a));
875 let b = PrefixFingerprint::compute("system", Some(&tools_b));
876 assert_ne!(a.combined_sha256, b.combined_sha256);
877 }
878
879 #[test]
880 fn different_tools_produce_different_fingerprint() {
881 let tools_a = vec![make_tool("read_file")];
882 let tools_b = vec![make_tool("write_file")];
883 let a = PrefixFingerprint::compute("system", Some(&tools_a));
884 let b = PrefixFingerprint::compute("system", Some(&tools_b));
885 assert_ne!(a.combined_sha256, b.combined_sha256);
886 }
887
888 #[test]
889 fn manager_starts_stable() {
890 let mut mgr = PrefixStabilityManager::new("system prompt", None);
891 assert!(mgr.check_and_update("system prompt", None).unwrap());
892 assert_eq!(mgr.change_count(), 0);
893 assert_eq!(mgr.check_count(), 1);
894 }
895
896 #[test]
897 fn manager_detects_change() {
898 let mut mgr = PrefixStabilityManager::new("system prompt", None);
899 let result = mgr.check_and_update("different prompt", None);
900 assert!(result.is_err());
901 assert_eq!(mgr.change_count(), 1);
902 let change = mgr.last_change().unwrap();
903 assert!(change.system_changed);
904 assert!(!change.tools_changed);
905 }
906
907 #[test]
908 fn manager_detects_tool_change() {
909 let tools_a = vec![make_tool("read_file")];
910 let tools_b = vec![make_tool("write_file")];
911 let mut mgr = PrefixStabilityManager::new("system", Some(&tools_a));
912 let result = mgr.check_and_update("system", Some(&tools_b));
913 assert!(result.is_err());
914 let change = mgr.last_change().unwrap();
915 assert!(!change.system_changed);
916 assert!(change.tools_changed);
917 }
918
919 #[test]
920 fn undeclared_drift_never_moves_the_pin() {
921 let mut mgr = PrefixStabilityManager::new("old", None);
922 assert_eq!(mgr.pin_reason(), Some("initial"));
923 assert!(mgr.check_and_update("new", None).is_err());
924 // The pin stays on "old": the same undeclared prefix is still a miss.
925 assert!(mgr.check_and_update("new", None).is_err());
926 assert!(mgr.check_and_update("old", None).unwrap());
927 assert_eq!(mgr.change_count(), 2);
928 assert_eq!(mgr.last_miss_reason(), Some("drift:sys"));
929 let history: Vec<_> = mgr.history().collect();
930 assert_eq!(history.len(), 2);
931 assert!(history.iter().all(|entry| !entry.repinned));
932 assert_eq!(mgr.pin_reason(), Some("initial"));
933 }
934
935 #[test]
936 fn declared_header_change_repins_under_a_logged_reason() {
937 let mut mgr = PrefixStabilityManager::new("old", None);
938 match mgr.check("new", None, Some("model")) {
939 PrefixCheck::Repinned { reason, change } => {
940 assert_eq!(reason, "change:model");
941 assert!(change.system_changed);
942 }
943 other => panic!("expected repin, got {other:?}"),
944 }
945 assert_eq!(mgr.pin_reason(), Some("change:model"));
946 assert!(matches!(mgr.check("new", None, None), PrefixCheck::Stable));
947 assert!(matches!(
948 mgr.check("newer", None, None),
949 PrefixCheck::Drift { .. }
950 ));
951 assert_eq!(mgr.pin_reason(), Some("change:model"));
952 let reasons: Vec<&str> = mgr.history().map(|e| e.reason.as_str()).collect();
953 assert_eq!(reasons, vec!["change:model", "drift:sys"]);
954 }
955
956 #[test]
957 fn declared_reason_on_a_stable_prefix_is_a_noop() {
958 let mut mgr = PrefixStabilityManager::new("same", None);
959 assert!(matches!(
960 mgr.check("same", None, Some("mode")),
961 PrefixCheck::Stable
962 ));
963 assert_eq!(mgr.change_count(), 0);
964 assert_eq!(mgr.pin_reason(), Some("initial"));
965 }
966
967 #[test]
968 fn context_update_message_reports_added_and_removed_lines_bounded() {
969 let known = "## Files\nsrc/a.rs\nsrc/b.rs\n## Instructions\nbe kind\n";
970 let current = "## Files\nsrc/a.rs\nsrc/b.rs\nsrc/c.rs\n## Instructions\nbe precise\n";
971 let update = context_update_message(known, current).expect("delta");
972 assert!(update.starts_with("<context_update>"));
973 assert!(update.ends_with("</context_update>"));
974 assert!(update.contains("+ src/c.rs"));
975 assert!(update.contains("+ be precise"));
976 assert!(update.contains("- be kind"));
977 assert!(!update.contains("- src/a.rs"));
978 assert!(context_update_message(known, known).is_none());
979
980 let mut big = String::new();
981 for i in 0..500 {
982 big.push_str(&format!("line {i}\n"));
983 }
984 let bounded = context_update_message("", &big).expect("delta");
985 assert!(
986 bounded.len() <= CONTEXT_UPDATE_MAX_BYTES + 128,
987 "{}",
988 bounded.len()
989 );
990 assert!(bounded.contains("more changed lines"));
991 }
992
993 #[test]
994 fn context_update_is_logged_but_is_not_a_miss() {
995 let mut mgr = PrefixStabilityManager::new("sys", None);
996 mgr.note_context_update();
997 assert_eq!(mgr.context_update_count(), 1);
998 assert_eq!(mgr.last_miss_reason(), None);
999 assert!(matches!(mgr.check("sys", None, None), PrefixCheck::Stable));
1000 }
1001
1002 #[test]
1003 fn history_reset_is_logged_without_moving_the_pin() {
1004 let mut mgr = PrefixStabilityManager::new("sys", None);
1005 mgr.note_history_reset("compaction");
1006 assert_eq!(mgr.last_miss_reason(), Some("reset:compaction"));
1007 assert!(matches!(mgr.check("sys", None, None), PrefixCheck::Stable));
1008 assert_eq!(mgr.history().count(), 1);
1009 }
1010
1011 #[test]
1012 fn stability_ratio_is_one_for_no_changes() {
1013 let mut mgr = PrefixStabilityManager::new("hello", None);
1014 mgr.check_and_update("hello", None).unwrap();
1015 mgr.check_and_update("hello", None).unwrap();
1016 assert!((mgr.stability_ratio() - 1.0).abs() < f64::EPSILON);
1017 assert_eq!(mgr.check_count(), 2);
1018 assert_eq!(mgr.change_count(), 0);
1019 }
1020
1021 #[test]
1022 fn stability_ratio_reflects_change_rate() {
1023 let mut mgr = PrefixStabilityManager::new("hello", None);
1024 mgr.check_and_update("hello", None).unwrap(); // check 1: stable
1025 let _ = mgr.check("world", None, Some("model")); // check 2: declared change
1026 mgr.check_and_update("world", None).unwrap(); // check 3: stable
1027 // 2 stable out of 3 checks = 0.666...
1028 // (check_count=0 at start, so 3 checks: 3 checks - 1 change = 2 stable)
1029 assert!((mgr.stability_ratio() - 2.0 / 3.0).abs() < 0.01);
1030 assert_eq!(mgr.check_count(), 3);
1031 assert_eq!(mgr.change_count(), 1);
1032 }
1033
1034 #[test]
1035 fn empty_tools_and_none_tools_produce_same_hash() {
1036 let empty = PrefixFingerprint::compute("system", Some(&[]));
1037 let none = PrefixFingerprint::compute("system", None);
1038 // Both should produce sha256(b"") for the tool component
1039 assert_eq!(empty.tools_sha256, none.tools_sha256);
1040 }
1041
1042 #[test]
1043 fn empty_system_produces_sha256_of_empty_string() {
1044 let fp = PrefixFingerprint::compute("", None);
1045 let expected = sha256_hex(b"");
1046 assert_eq!(fp.system_sha256, expected);
1047 }
1048
1049 #[test]
1050 fn prefix_change_description_is_informative() {
1051 let old = PrefixFingerprint::compute("old", None);
1052 let new = PrefixFingerprint::compute("new", None);
1053 let change = PrefixChange {
1054 old,
1055 new,
1056 system_changed: true,
1057 tools_changed: false,
1058 };
1059 assert_eq!(
1060 change.description(),
1061 "prefix cache invalidated: system prompt changed"
1062 );
1063 assert_eq!(change.label(), "sys");
1064 }
1065
1066 #[test]
1067 fn new_unpinned_has_no_change_history() {
1068 let mut mgr = PrefixStabilityManager::new_unpinned();
1069 assert!(mgr.pinned_fingerprint().is_none());
1070 assert!(mgr.current_fingerprint().is_none());
1071 assert!(mgr.last_change().is_none());
1072 assert_eq!(mgr.change_count(), 0);
1073 assert_eq!(mgr.check_count(), 0);
1074 // First check should pin automatically and count as a check.
1075 assert!(mgr.check_and_update("hello", None).unwrap());
1076 assert!(mgr.pinned_fingerprint().is_some());
1077 assert_eq!(mgr.check_count(), 1);
1078 }
1079
1080 #[test]
1081 fn fingerprint_detects_schema_change_not_just_name_change() {
1082 let tool_a = make_tool("my_tool");
1083 let mut tool_a_v2 = make_tool("my_tool");
1084 tool_a_v2.description = "updated description".to_string();
1085
1086 let a = PrefixFingerprint::compute("system", Some(&[tool_a]));
1087 let b = PrefixFingerprint::compute("system", Some(&[tool_a_v2]));
1088 // Same name, different description — must produce different hash.
1089 assert_ne!(a.tools_sha256, b.tools_sha256);
1090 assert_ne!(a.combined_sha256, b.combined_sha256);
1091 }
1092
1093 #[test]
1094 fn system_prompt_text_returns_empty_for_none() {
1095 assert_eq!(system_prompt_text(None), "");
1096 }
1097
1098 // ── ToolCatalogCache tests ──────────────────────────────────
1099
1100 #[test]
1101 fn tool_catalog_cache_miss_then_hit_returns_same_digest() {
1102 let mut cache = ToolCatalogCache::new();
1103 let tools = vec![make_tool("read_file"), make_tool("write_file")];
1104
1105 let first = cache.fingerprint_for(&tools);
1106 assert_eq!(cache.len(), 1);
1107
1108 let second = cache.fingerprint_for(&tools);
1109 assert_eq!(cache.len(), 1, "second call should be a cache hit");
1110 assert_eq!(first.sha256_hex, second.sha256_hex);
1111 }
1112
1113 #[test]
1114 fn tool_catalog_cache_different_tool_sets_dont_collide() {
1115 let mut cache = ToolCatalogCache::new();
1116 let a = vec![make_tool("read_file")];
1117 let b = vec![make_tool("write_file")];
1118
1119 let entry_a = cache.fingerprint_for(&a);
1120 let entry_b = cache.fingerprint_for(&b);
1121 assert_eq!(cache.len(), 2);
1122 assert_ne!(entry_a.sha256_hex, entry_b.sha256_hex);
1123 }
1124
1125 #[test]
1126 fn tool_catalog_cache_fingerprint_is_order_sensitive_like_the_provider_cache() {
1127 // The identity hash includes the input order so re-registering the
1128 // same set with a different permutation produces a separate cache
1129 // entry. The digest is now order-sensitive too: the provider KV cache
1130 // is order-sensitive, so a sorted fingerprint read a false-green while
1131 // the real prefix missed every turn (ops C1). Different wire order =
1132 // different prefix = different fingerprint.
1133 let mut cache = ToolCatalogCache::new();
1134 let a = vec![make_tool("read_file"), make_tool("write_file")];
1135 let b = vec![make_tool("write_file"), make_tool("read_file")];
1136 let entry_a = cache.fingerprint_for(&a);
1137 let entry_b = cache.fingerprint_for(&b);
1138 assert_ne!(entry_a.sha256_hex, entry_b.sha256_hex);
1139 assert_eq!(cache.len(), 2);
1140 // Re-requesting the original order returns the cached digest.
1141 let again = cache.fingerprint_for(&a);
1142 assert_eq!(again.sha256_hex, entry_a.sha256_hex);
1143 }
1144
1145 #[test]
1146 fn tool_catalog_cache_detects_schema_change() {
1147 let mut cache = ToolCatalogCache::new();
1148 let tool_v1 = make_tool("t");
1149 let mut tool_v2 = make_tool("t");
1150 tool_v2.description = "updated".to_string();
1151
1152 let entry_v1 = cache.fingerprint_for(&[tool_v1]);
1153 let entry_v2 = cache.fingerprint_for(&[tool_v2]);
1154 assert_ne!(entry_v1.sha256_hex, entry_v2.sha256_hex);
1155 assert_eq!(cache.len(), 2);
1156 }
1157
1158 #[test]
1159 fn tool_catalog_cache_respects_capacity() {
1160 let mut cache = ToolCatalogCache::with_capacity(2);
1161 cache.fingerprint_for(&[make_tool("a")]);
1162 cache.fingerprint_for(&[make_tool("b")]);
1163 cache.fingerprint_for(&[make_tool("c")]);
1164 assert_eq!(cache.len(), 2);
1165 // The first entry was evicted; a re-query for it should miss.
1166 let re_entry = cache.fingerprint_for(&[make_tool("a")]);
1167 // After the re-query, the cache has [b, c, a] — 3 entries? No,
1168 // capacity 2 means oldest is evicted when we insert the 3rd unique.
1169 // After inserting a, the cache holds the most recent 2: {c, a}.
1170 assert_eq!(cache.len(), 2);
1171 // The returned digest should match a fresh fingerprint of the same set.
1172 let fresh = cache.fingerprint_for(&[make_tool("a")]);
1173 assert_eq!(re_entry.sha256_hex, fresh.sha256_hex);
1174 }
1175
1176 #[test]
1177 fn tool_catalog_cache_invalidate_clears_all() {
1178 let mut cache = ToolCatalogCache::new();
1179 cache.fingerprint_for(&[make_tool("a")]);
1180 cache.fingerprint_for(&[make_tool("b")]);
1181 cache.invalidate();
1182 assert!(cache.is_empty());
1183 assert_eq!(cache.len(), 0);
1184 }
1185
1186 #[test]
1187 fn tool_catalog_cache_empty_slice_uses_zero_capacity_path() {
1188 // Empty input is fine — should produce a stable, non-empty digest.
1189 let mut cache = ToolCatalogCache::new();
1190 let entry = cache.fingerprint_for(&[]);
1191 assert!(!entry.sha256_hex.is_empty());
1192 let again = cache.fingerprint_for(&[]);
1193 assert_eq!(entry.sha256_hex, again.sha256_hex);
1194 }
1195
1196 #[test]
1197 fn compute_with_tool_cache_matches_compute_uncached() {
1198 // The cached and uncached paths must produce identical fingerprints
1199 // for the same inputs — otherwise we'd silently corrupt the prefix
1200 // cache and invalidate every request.
1201 let mut cache = ToolCatalogCache::new();
1202 let tools = vec![make_tool("alpha"), make_tool("beta")];
1203
1204 let cached = PrefixFingerprint::compute_with_tool_cache("sys", Some(&tools), &mut cache);
1205 let uncached = PrefixFingerprint::compute("sys", Some(&tools));
1206 assert_eq!(cached.combined_sha256, uncached.combined_sha256);
1207 assert_eq!(cached.tools_sha256, uncached.tools_sha256);
1208 }
1209
1210 #[test]
1211 fn manager_check_and_update_uses_cached_tool_fingerprint() {
1212 // After the first call populates the cache, subsequent calls with
1213 // the same tool list should not invalidate the prefix.
1214 let tools = vec![make_tool("t1")];
1215 let mut mgr = PrefixStabilityManager::new("sys", Some(&tools));
1216 assert!(mgr.check_and_update("sys", Some(&tools)).is_ok());
1217 assert!(mgr.check_and_update("sys", Some(&tools)).is_ok());
1218 assert_eq!(mgr.change_count(), 0);
1219 }
1220 }
1221
1221 lines RUST