| 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. **Detects drift** by comparing the current prefix against the |
| 13 | //! pinned fingerprint before every request. |
| 14 | //! 3. **Diagnoses** the cause of drift — did the system prompt change? |
| 15 | //! Did the tool set change? Both? |
| 16 | //! 4. **Emits events** so the TUI can surface stability to the user. |
| 17 | //! |
| 18 | //! ## Three-region model (from Reasonix) |
| 19 | //! |
| 20 | //! ```text |
| 21 | //! ┌─────────────────────────────────────────┐ |
| 22 | //! │ IMMUTABLE PREFIX │ ← fixed for session |
| 23 | //! │ system + tool_specs │ cache hit candidate |
| 24 | //! ├─────────────────────────────────────────┤ |
| 25 | //! │ APPEND-ONLY HISTORY │ ← grows monotonically |
| 26 | //! │ [assistant₁][tool₁][assistant₂]... │ preserves prefix of prior turns |
| 27 | //! ├─────────────────────────────────────────┤ |
| 28 | //! │ LATEST USER TURN │ ← the only new content per request |
| 29 | //! └─────────────────────────────────────────┘ |
| 30 | //! ``` |
| 31 | |
| 32 | use std::collections::hash_map::DefaultHasher; |
| 33 | use std::collections::{HashMap, VecDeque}; |
| 34 | use std::hash::{Hash, Hasher}; |
| 35 | |
| 36 | use serde::{Deserialize, Serialize}; |
| 37 | |
| 38 | use crate::models::{SystemPrompt, Tool}; |
| 39 | |
| 40 | /// A snapshot of the immutable prefix's fingerprint. |
| 41 | /// |
| 42 | /// Two snapshots with the same `combined` hash are guaranteed to |
| 43 | /// produce the same byte prefix when serialized for the API. |
| 44 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 45 | pub struct PrefixFingerprint { |
| 46 | /// SHA-256 of the system prompt text. |
| 47 | pub system_sha256: String, |
| 48 | /// SHA-256 of the full tool catalog JSON (names, descriptions, schemas). |
| 49 | pub tools_sha256: String, |
| 50 | /// SHA-256 of system_sha256 ++ tools_sha256 (combined). |
| 51 | pub combined_sha256: String, |
| 52 | } |
| 53 | |
| 54 | impl PrefixFingerprint { |
| 55 | /// Compute a fingerprint from system prompt text and tool list. |
| 56 | /// |
| 57 | /// Tools are serialized to the same JSON shape the chat API receives |
| 58 | /// (`type`, `name`, `description`, `parameters`, `strict`), sorted |
| 59 | /// lexicographically by JSON text, then SHA-256 hashed. This catches |
| 60 | /// schema/description drift that actually affects the API prefix, |
| 61 | /// while ignoring internal-only fields like `allowed_callers` (#2264). |
| 62 | /// |
| 63 | /// This entry point shares a process-local [`ToolCatalogCache`] with |
| 64 | /// every other call, so a stable tool set (the common case after the |
| 65 | /// first turn of a session) avoids the per-tool JSON serialization |
| 66 | /// and sort/join entirely. Callers that hold their own cache — e.g. |
| 67 | /// [`PrefixStabilityManager`] — should use |
| 68 | /// [`Self::compute_with_tool_cache`] to share *that* cache instead |
| 69 | /// and avoid the thread-local lookup. |
| 70 | #[cfg(test)] |
| 71 | pub fn compute(system_text: &str, tools: Option<&[Tool]>) -> Self { |
| 72 | let mut cache = ToolCatalogCache::new(); |
| 73 | Self::compute_with_tool_cache(system_text, tools, &mut cache) |
| 74 | } |
| 75 | |
| 76 | /// Compute a fingerprint while reusing a [`ToolCatalogCache`] for the |
| 77 | /// tool-side work. The cache holds the joined+sorted+SHA-256'd catalog |
| 78 | /// under a content-derived identity so the per-tool JSON serialization |
| 79 | /// and the sort/join only run on the first call for a given tool set. |
| 80 | /// |
| 81 | /// On a cache hit this function avoids the entire tool serialization |
| 82 | /// path, which can be 100+ microseconds for a 60-tool catalog. |
| 83 | pub fn compute_with_tool_cache( |
| 84 | system_text: &str, |
| 85 | tools: Option<&[Tool]>, |
| 86 | cache: &mut ToolCatalogCache, |
| 87 | ) -> Self { |
| 88 | let system_sha256 = sha256_hex(system_text.as_bytes()); |
| 89 | |
| 90 | let tools_sha256 = match tools { |
| 91 | Some(tools) if !tools.is_empty() => { |
| 92 | // `fingerprint_for` consults the cache first; on a hit |
| 93 | // it returns the pre-computed hex digest directly. |
| 94 | cache.fingerprint_for(tools).sha256_hex |
| 95 | } |
| 96 | _ => sha256_hex(b""), |
| 97 | }; |
| 98 | |
| 99 | let combined = format!("{system_sha256}:{tools_sha256}"); |
| 100 | let combined_sha256 = sha256_hex(combined.as_bytes()); |
| 101 | Self { |
| 102 | system_sha256, |
| 103 | tools_sha256, |
| 104 | combined_sha256, |
| 105 | } |
| 106 | } |
| 107 | } |
| 108 | |
| 109 | /// A change record describing what drifted in the prefix. |
| 110 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 111 | pub struct PrefixChange { |
| 112 | /// The old fingerprint (before the change). |
| 113 | pub old: PrefixFingerprint, |
| 114 | /// The new fingerprint (after the change). |
| 115 | pub new: PrefixFingerprint, |
| 116 | /// Whether the system prompt component changed. |
| 117 | pub system_changed: bool, |
| 118 | /// Whether the tool set component changed. |
| 119 | pub tools_changed: bool, |
| 120 | } |
| 121 | |
| 122 | #[allow(dead_code)] |
| 123 | impl PrefixChange { |
| 124 | /// Returns a human-readable description of what changed. |
| 125 | pub fn description(&self) -> String { |
| 126 | let mut parts = Vec::new(); |
| 127 | if self.system_changed { |
| 128 | parts.push("system prompt"); |
| 129 | } |
| 130 | if self.tools_changed { |
| 131 | parts.push("tool set"); |
| 132 | } |
| 133 | if parts.is_empty() { |
| 134 | return "unknown (fingerprint mismatch but no component detected)".to_string(); |
| 135 | } |
| 136 | format!("prefix cache invalidated: {} changed", parts.join(" and ")) |
| 137 | } |
| 138 | |
| 139 | /// Returns a short label for TUI chip display. |
| 140 | pub fn label(&self) -> &'static str { |
| 141 | if self.system_changed && self.tools_changed { |
| 142 | "sys+tools" |
| 143 | } else if self.system_changed { |
| 144 | "sys" |
| 145 | } else if self.tools_changed { |
| 146 | "tools" |
| 147 | } else { |
| 148 | "prefix" |
| 149 | } |
| 150 | } |
| 151 | } |
| 152 | |
| 153 | /// Monitors and manages prefix-cache stability across turns. |
| 154 | /// |
| 155 | /// This is the core abstraction, mirroring Reasonix's `ImmutablePrefix` |
| 156 | /// concept but adapted to CodeWhale's existing architecture where the |
| 157 | /// system prompt is rebuilt each turn and tools are registered at startup. |
| 158 | /// |
| 159 | /// Usage: |
| 160 | /// ```ignore |
| 161 | /// let mgr = PrefixStabilityManager::new(system_text, tools); |
| 162 | /// if mgr.check_and_update(system_text, tools) { |
| 163 | /// println!("Prefix is stable (cache-friendly)"); |
| 164 | /// } else { |
| 165 | /// let change = mgr.last_change().unwrap(); |
| 166 | /// println!("Prefix drifted: {}", change.description()); |
| 167 | /// } |
| 168 | /// ``` |
| 169 | #[derive(Debug, Clone)] |
| 170 | pub struct PrefixStabilityManager { |
| 171 | /// The pinned fingerprint from session start or last stabilization. |
| 172 | pinned: Option<PrefixFingerprint>, |
| 173 | /// The most recent fingerprint (computed during last check). |
| 174 | current: Option<PrefixFingerprint>, |
| 175 | /// The last detected change, if any. |
| 176 | last_change: Option<PrefixChange>, |
| 177 | /// Total number of prefix changes detected this session. |
| 178 | change_count: u64, |
| 179 | /// Total number of stability checks performed. |
| 180 | check_count: u64, |
| 181 | /// Process-local cache for the tool-catalog JSON serialization. Avoids |
| 182 | /// re-running `tool_to_api_json` + sort + join on every `check_and_update` |
| 183 | /// when the tool set is unchanged (the common case once tools are |
| 184 | /// registered at session start). |
| 185 | tool_catalog_cache: ToolCatalogCache, |
| 186 | } |
| 187 | |
| 188 | /// Default capacity for the tool-catalog serialization cache. Sized for |
| 189 | /// "session + 1 or 2 forked subagent catalogs" without unbounded growth. |
| 190 | const TOOL_CATALOG_CACHE_CAPACITY: usize = 8; |
| 191 | |
| 192 | /// Bounded LRU cache of `(tool_set_identity) -> sha256_hex`. |
| 193 | /// |
| 194 | /// The cache key is a content-derived `u64` hash of the tool list (length + |
| 195 | /// per-tool `name` + `description` + serialized `input_schema`). On a hit, |
| 196 | /// `PrefixFingerprint::compute` skips the per-tool JSON serialization, the |
| 197 | /// sort, and the join — a workload that can be 100+ microseconds for a |
| 198 | /// 60-tool catalog. On a miss, the work runs once and only the digest is |
| 199 | /// retained (#3854); the joined catalog string is ephemeral. |
| 200 | #[derive(Debug, Default, Clone)] |
| 201 | pub struct ToolCatalogCache { |
| 202 | by_identity: HashMap<u64, CachedCatalog>, |
| 203 | insertion_order: VecDeque<u64>, |
| 204 | capacity: usize, |
| 205 | } |
| 206 | |
| 207 | /// One entry in [`ToolCatalogCache`]. Production only needs the pre-computed |
| 208 | /// SHA-256 digest of the sorted joined catalog. |
| 209 | #[derive(Debug, Clone)] |
| 210 | pub struct CachedCatalog { |
| 211 | /// SHA-256 hex digest of the newline-joined, sorted tool-catalog JSON. |
| 212 | pub sha256_hex: String, |
| 213 | } |
| 214 | |
| 215 | impl ToolCatalogCache { |
| 216 | /// Create a cache with the default capacity. |
| 217 | #[must_use] |
| 218 | pub fn new() -> Self { |
| 219 | Self::with_capacity(TOOL_CATALOG_CACHE_CAPACITY) |
| 220 | } |
| 221 | |
| 222 | /// Create a cache that holds at most `capacity` tool-set entries. |
| 223 | /// Smaller values save memory at the cost of more cache misses. |
| 224 | #[must_use] |
| 225 | pub fn with_capacity(capacity: usize) -> Self { |
| 226 | let cap = capacity.max(1); |
| 227 | Self { |
| 228 | by_identity: HashMap::with_capacity(cap), |
| 229 | insertion_order: VecDeque::with_capacity(cap), |
| 230 | capacity: cap, |
| 231 | } |
| 232 | } |
| 233 | |
| 234 | /// Compute (or recall) the joined-and-hashed tool catalog for `tools`. |
| 235 | /// The cache is keyed on a content-derived `u64` identity so two `&[Tool]` |
| 236 | /// slices with the same payloads — in the same order — hit the same entry. |
| 237 | pub fn fingerprint_for(&mut self, tools: &[Tool]) -> CachedCatalog { |
| 238 | let identity = tool_set_identity(tools); |
| 239 | if let Some(cached) = self.by_identity.get(&identity) { |
| 240 | return cached.clone(); |
| 241 | } |
| 242 | |
| 243 | // Miss: serialize, sort, join, hash. Keep only the digest in the |
| 244 | // cache — the joined string is not needed on the hot path (#3854). |
| 245 | let mut serialized: Vec<String> = tools.iter().filter_map(tool_to_api_json).collect(); |
| 246 | serialized.sort(); |
| 247 | let joined = serialized.join("\n"); |
| 248 | let entry = CachedCatalog { |
| 249 | sha256_hex: sha256_hex(joined.as_bytes()), |
| 250 | }; |
| 251 | |
| 252 | if self.by_identity.len() >= self.capacity |
| 253 | && let Some(oldest) = self.insertion_order.pop_front() |
| 254 | { |
| 255 | self.by_identity.remove(&oldest); |
| 256 | } |
| 257 | self.by_identity.insert(identity, entry.clone()); |
| 258 | self.insertion_order.push_back(identity); |
| 259 | entry |
| 260 | } |
| 261 | |
| 262 | /// Drop every cached entry. Used by tool-registry mutation paths |
| 263 | /// (e.g. plugin hot-reload, MCP attach) when the caller cannot |
| 264 | /// easily prove the tool set is unchanged. |
| 265 | #[allow(dead_code)] // observability; called by /cache flush and tests |
| 266 | pub fn invalidate(&mut self) { |
| 267 | self.by_identity.clear(); |
| 268 | self.insertion_order.clear(); |
| 269 | } |
| 270 | |
| 271 | /// Returns the number of cached entries. |
| 272 | #[must_use] |
| 273 | pub fn len(&self) -> usize { |
| 274 | self.by_identity.len() |
| 275 | } |
| 276 | |
| 277 | /// Returns `true` if the cache has no entries. |
| 278 | #[allow(dead_code)] // observability; surfaced via /status |
| 279 | #[must_use] |
| 280 | pub fn is_empty(&self) -> bool { |
| 281 | self.by_identity.is_empty() |
| 282 | } |
| 283 | |
| 284 | /// Returns `(current_entries, capacity)` for observability. Surfaced via |
| 285 | /// the `/status` chip in a follow-up; tests exercise the path. |
| 286 | #[allow(dead_code)] // surfaced via /status in a follow-up; tests exercise it |
| 287 | #[must_use] |
| 288 | pub fn stats(&self) -> (usize, usize) { |
| 289 | (self.len(), self.capacity) |
| 290 | } |
| 291 | } |
| 292 | |
| 293 | /// Content-derived identity for a tool slice. Order-sensitive: two slices |
| 294 | /// with the same tools in different orders produce different identities. |
| 295 | /// (The downstream fingerprint itself is order-insensitive — the sort in |
| 296 | /// `fingerprint_for` takes care of that — but the cache key matches the |
| 297 | /// input order so re-registration of the same set in the same order hits.) |
| 298 | fn tool_set_identity(tools: &[Tool]) -> u64 { |
| 299 | let mut hasher = DefaultHasher::new(); |
| 300 | tools.len().hash(&mut hasher); |
| 301 | for tool in tools { |
| 302 | tool.name.hash(&mut hasher); |
| 303 | tool.description.hash(&mut hasher); |
| 304 | // `strict` participates in `tool_to_api_json` output (it is part of |
| 305 | // the wire-format the chat API receives), so it MUST be part of the |
| 306 | // identity. Omitting it lets two semantically different catalogs |
| 307 | // collide and serve a stale fingerprint. |
| 308 | tool.strict.hash(&mut hasher); |
| 309 | // Walk the schema JSON directly instead of materializing it as a |
| 310 | // String. For a 60-tool catalog this saves ~25-40 KB of allocation |
| 311 | // on every cache miss. |
| 312 | hash_json_value(&tool.input_schema, &mut hasher); |
| 313 | } |
| 314 | hasher.finish() |
| 315 | } |
| 316 | |
| 317 | /// Fold a `serde_json::Value` into the hasher without allocating a |
| 318 | /// `String`. Numeric variants are hashed via their bit pattern so `1` and |
| 319 | /// `1.0` produce distinct identities (matching the JSON spec). |
| 320 | fn hash_json_value<H: Hasher>(value: &serde_json::Value, state: &mut H) { |
| 321 | match value { |
| 322 | serde_json::Value::Null => 0u8.hash(state), |
| 323 | serde_json::Value::Bool(b) => { |
| 324 | 1u8.hash(state); |
| 325 | b.hash(state); |
| 326 | } |
| 327 | serde_json::Value::Number(n) => { |
| 328 | 2u8.hash(state); |
| 329 | if let Some(i) = n.as_i64() { |
| 330 | i.hash(state); |
| 331 | } else if let Some(u) = n.as_u64() { |
| 332 | u.hash(state); |
| 333 | } else if let Some(f) = n.as_f64() { |
| 334 | f.to_bits().hash(state); |
| 335 | } |
| 336 | } |
| 337 | serde_json::Value::String(s) => { |
| 338 | 3u8.hash(state); |
| 339 | s.hash(state); |
| 340 | } |
| 341 | serde_json::Value::Array(arr) => { |
| 342 | 4u8.hash(state); |
| 343 | arr.len().hash(state); |
| 344 | for v in arr { |
| 345 | hash_json_value(v, state); |
| 346 | } |
| 347 | } |
| 348 | serde_json::Value::Object(obj) => { |
| 349 | 5u8.hash(state); |
| 350 | obj.len().hash(state); |
| 351 | // Iterate by sorted key so `{"a":1,"b":2}` and `{"b":2,"a":1}` |
| 352 | // collide — the wire format already canonicalizes via the |
| 353 | // `serde_json` Map ordering, but a defensively-sorted view |
| 354 | // future-proofs against schema serializers that emit |
| 355 | // declaration order. |
| 356 | let mut entries: Vec<(&String, &serde_json::Value)> = obj.iter().collect(); |
| 357 | entries.sort_by(|a, b| a.0.cmp(b.0)); |
| 358 | for (k, v) in entries { |
| 359 | k.hash(state); |
| 360 | hash_json_value(v, state); |
| 361 | } |
| 362 | } |
| 363 | } |
| 364 | } |
| 365 | |
| 366 | /// Process-local fallback cache used by `PrefixFingerprint::compute` |
| 367 | /// (when available). Callers that maintain their own cache (e.g. |
| 368 | /// [`PrefixStabilityManager`]) should prefer |
| 369 | /// [`PrefixFingerprint::compute_with_tool_cache`] and pass the cache in |
| 370 | /// directly, both to share state and to avoid the thread-local lookup |
| 371 | /// on the hot path. |
| 372 | #[allow(dead_code)] |
| 373 | impl PrefixStabilityManager { |
| 374 | /// Create a new manager and immediately pin the first fingerprint. |
| 375 | pub fn new(system_text: &str, tools: Option<&[Tool]>) -> Self { |
| 376 | let mut cache = ToolCatalogCache::new(); |
| 377 | let fp = PrefixFingerprint::compute_with_tool_cache(system_text, tools, &mut cache); |
| 378 | Self { |
| 379 | pinned: Some(fp.clone()), |
| 380 | current: Some(fp), |
| 381 | last_change: None, |
| 382 | change_count: 0, |
| 383 | check_count: 0, |
| 384 | tool_catalog_cache: cache, |
| 385 | } |
| 386 | } |
| 387 | |
| 388 | /// Create a manager in "unpinned" state — no initial fingerprint. |
| 389 | /// Call `pin()` or `check_and_update()` to establish the baseline. |
| 390 | pub fn new_unpinned() -> Self { |
| 391 | Self { |
| 392 | pinned: None, |
| 393 | current: None, |
| 394 | last_change: None, |
| 395 | change_count: 0, |
| 396 | check_count: 0, |
| 397 | tool_catalog_cache: ToolCatalogCache::new(), |
| 398 | } |
| 399 | } |
| 400 | |
| 401 | /// Explicitly pin a fingerprint, replacing any prior pinned state. |
| 402 | /// Returns `true` if this is the first pin, or `false` if replacing. |
| 403 | /// Note: does NOT increment `check_count` — that counter is reserved |
| 404 | /// for `check_and_update` calls so `stability_ratio()` stays accurate. |
| 405 | pub fn pin(&mut self, system_text: &str, tools: Option<&[Tool]>) -> bool { |
| 406 | let fp = PrefixFingerprint::compute_with_tool_cache( |
| 407 | system_text, |
| 408 | tools, |
| 409 | &mut self.tool_catalog_cache, |
| 410 | ); |
| 411 | let was_unpinned = self.pinned.is_none(); |
| 412 | self.pinned = Some(fp.clone()); |
| 413 | self.current = Some(fp); |
| 414 | was_unpinned |
| 415 | } |
| 416 | |
| 417 | /// Check whether the current prefix matches the pinned fingerprint. |
| 418 | /// Updates internal state and returns: |
| 419 | /// - `Ok(true)` if the prefix is stable (fingerprint matches pinned). |
| 420 | /// - `Ok(false)` if the prefix changed but was automatically re-pinned. |
| 421 | /// - `Err(change)` if the prefix changed; caller should surface this. |
| 422 | /// |
| 423 | /// After calling this, `last_change()` returns the detected change. |
| 424 | pub fn check_and_update( |
| 425 | &mut self, |
| 426 | system_text: &str, |
| 427 | tools: Option<&[Tool]>, |
| 428 | ) -> Result<bool, Box<PrefixChange>> { |
| 429 | // Use the cached tool-catalog fingerprint path so a stable tool set |
| 430 | // (the common case after the first turn) does not re-serialize the |
| 431 | // full tool list. The system-prompt side is hashed on every call |
| 432 | // because the system prompt changes more often (mode flips, |
| 433 | // project-context refreshes, canonical state overlays). |
| 434 | let fp = PrefixFingerprint::compute_with_tool_cache( |
| 435 | system_text, |
| 436 | tools, |
| 437 | &mut self.tool_catalog_cache, |
| 438 | ); |
| 439 | let old_fp = self.current.replace(fp.clone()); |
| 440 | self.check_count += 1; |
| 441 | |
| 442 | let pinned = match &self.pinned { |
| 443 | Some(p) => p, |
| 444 | None => { |
| 445 | // First check: pin now. |
| 446 | self.pinned = Some(fp); |
| 447 | self.last_change = None; |
| 448 | return Ok(true); |
| 449 | } |
| 450 | }; |
| 451 | |
| 452 | if fp.combined_sha256 == pinned.combined_sha256 { |
| 453 | // Stable — no change. |
| 454 | Ok(true) |
| 455 | } else { |
| 456 | // Change detected. |
| 457 | let old = old_fp.unwrap_or_else(|| pinned.clone()); |
| 458 | let system_changed = fp.system_sha256 != pinned.system_sha256; |
| 459 | let tools_changed = fp.tools_sha256 != pinned.tools_sha256; |
| 460 | |
| 461 | let change = PrefixChange { |
| 462 | old, |
| 463 | new: fp.clone(), |
| 464 | system_changed, |
| 465 | tools_changed, |
| 466 | }; |
| 467 | |
| 468 | self.last_change = Some(change.clone()); |
| 469 | self.change_count += 1; |
| 470 | |
| 471 | // Re-pin to the new prefix so subsequent checks are |
| 472 | // against the latest baseline. Use the original fp |
| 473 | // (avoid recomputing the hash — clone was for the change record). |
| 474 | self.pinned = Some(fp); |
| 475 | |
| 476 | Err(Box::new(change)) |
| 477 | } |
| 478 | } |
| 479 | |
| 480 | /// Returns the most recent prefix change, if any. |
| 481 | pub fn last_change(&self) -> Option<&PrefixChange> { |
| 482 | self.last_change.as_ref() |
| 483 | } |
| 484 | |
| 485 | /// Returns the pinned fingerprint. |
| 486 | pub fn pinned_fingerprint(&self) -> Option<&PrefixFingerprint> { |
| 487 | self.pinned.as_ref() |
| 488 | } |
| 489 | |
| 490 | /// Returns the current (most recently computed) fingerprint. |
| 491 | pub fn current_fingerprint(&self) -> Option<&PrefixFingerprint> { |
| 492 | self.current.as_ref() |
| 493 | } |
| 494 | |
| 495 | /// Returns the total number of prefix changes detected. |
| 496 | pub fn change_count(&self) -> u64 { |
| 497 | self.change_count |
| 498 | } |
| 499 | |
| 500 | /// Returns the total number of stability checks performed. |
| 501 | pub fn check_count(&self) -> u64 { |
| 502 | self.check_count |
| 503 | } |
| 504 | |
| 505 | /// Returns the prefix stability rate as a fraction (0.0 – 1.0). |
| 506 | /// 1.0 means the prefix has never changed. Returns 1.0 when no |
| 507 | /// checks have been performed (to avoid division by zero). |
| 508 | pub fn stability_ratio(&self) -> f64 { |
| 509 | if self.check_count == 0 { |
| 510 | 1.0 |
| 511 | } else { |
| 512 | let stable_checks = self.check_count - self.change_count; |
| 513 | stable_checks as f64 / self.check_count as f64 |
| 514 | } |
| 515 | } |
| 516 | |
| 517 | /// Returns a human-readable stability summary. |
| 518 | pub fn summary(&self) -> String { |
| 519 | let pct = self.stability_ratio() * 100.0; |
| 520 | let pinned_short = self |
| 521 | .pinned |
| 522 | .as_ref() |
| 523 | .map(|fp| { |
| 524 | if fp.combined_sha256.len() >= 12 { |
| 525 | &fp.combined_sha256[..12] |
| 526 | } else { |
| 527 | &fp.combined_sha256 |
| 528 | } |
| 529 | }) |
| 530 | .unwrap_or("none"); |
| 531 | |
| 532 | format!( |
| 533 | "Prefix stability: {pct:.1}% ({stable}/{total} checks stable) | fingerprint: {pinned_short} | changes: {changes}", |
| 534 | pct = pct, |
| 535 | stable = self.check_count.saturating_sub(self.change_count), |
| 536 | total = self.check_count, |
| 537 | pinned_short = pinned_short, |
| 538 | changes = self.change_count, |
| 539 | ) |
| 540 | } |
| 541 | } |
| 542 | |
| 543 | /// Serialize a tool to the same JSON shape the chat API receives, |
| 544 | /// excluding internal-only fields like `allowed_callers`, `defer_loading`, |
| 545 | /// `input_examples`, and `cache_control` that are never sent to DeepSeek. |
| 546 | fn tool_to_api_json(tool: &Tool) -> Option<String> { |
| 547 | let mut value = serde_json::json!({ |
| 548 | "type": "function", |
| 549 | "function": { |
| 550 | "name": tool.name, |
| 551 | "description": tool.description, |
| 552 | "parameters": tool.input_schema, |
| 553 | } |
| 554 | }); |
| 555 | if let Some(strict) = tool.strict |
| 556 | && let Some(function) = value.get_mut("function") |
| 557 | { |
| 558 | function["strict"] = serde_json::json!(strict); |
| 559 | } |
| 560 | serde_json::to_string(&value).ok() |
| 561 | } |
| 562 | |
| 563 | /// Compute the SHA-256 hex digest of a byte slice. |
| 564 | fn sha256_hex(bytes: &[u8]) -> String { |
| 565 | crate::hashing::sha256_hex(bytes) |
| 566 | } |
| 567 | |
| 568 | /// Extract the system prompt text from an optional SystemPrompt, |
| 569 | /// returning an owned String. This is used for prefix fingerprinting |
| 570 | /// and avoids lifetime/leak issues with the rare SystemPrompt::Blocks case. |
| 571 | pub fn system_prompt_text(system: Option<&SystemPrompt>) -> String { |
| 572 | match system { |
| 573 | Some(SystemPrompt::Text(text)) => text.clone(), |
| 574 | Some(SystemPrompt::Blocks(blocks)) => { |
| 575 | let mut text = String::new(); |
| 576 | for block in blocks { |
| 577 | text.push_str(&block.text); |
| 578 | text.push('\n'); |
| 579 | } |
| 580 | text |
| 581 | } |
| 582 | None => String::new(), |
| 583 | } |
| 584 | } |
| 585 | |
| 586 | #[cfg(test)] |
| 587 | mod tests { |
| 588 | use super::*; |
| 589 | |
| 590 | fn make_tool(name: &str) -> Tool { |
| 591 | Tool { |
| 592 | name: name.to_string(), |
| 593 | description: String::new(), |
| 594 | input_schema: serde_json::Value::Null, |
| 595 | tool_type: None, |
| 596 | allowed_callers: None, |
| 597 | defer_loading: None, |
| 598 | input_examples: None, |
| 599 | strict: None, |
| 600 | cache_control: None, |
| 601 | } |
| 602 | } |
| 603 | |
| 604 | #[test] |
| 605 | fn same_prefix_produces_same_fingerprint() { |
| 606 | let a = PrefixFingerprint::compute("hello world", None); |
| 607 | let b = PrefixFingerprint::compute("hello world", None); |
| 608 | assert_eq!(a.combined_sha256, b.combined_sha256); |
| 609 | } |
| 610 | |
| 611 | #[test] |
| 612 | fn different_system_produces_different_fingerprint() { |
| 613 | let a = PrefixFingerprint::compute("hello", None); |
| 614 | let b = PrefixFingerprint::compute("world", None); |
| 615 | assert_ne!(a.combined_sha256, b.combined_sha256); |
| 616 | } |
| 617 | |
| 618 | #[test] |
| 619 | fn tool_order_does_not_affect_fingerprint() { |
| 620 | let tools_a = vec![make_tool("read_file"), make_tool("write_file")]; |
| 621 | let tools_b = vec![make_tool("write_file"), make_tool("read_file")]; |
| 622 | let a = PrefixFingerprint::compute("system", Some(&tools_a)); |
| 623 | let b = PrefixFingerprint::compute("system", Some(&tools_b)); |
| 624 | assert_eq!(a.combined_sha256, b.combined_sha256); |
| 625 | } |
| 626 | |
| 627 | #[test] |
| 628 | fn different_tools_produce_different_fingerprint() { |
| 629 | let tools_a = vec![make_tool("read_file")]; |
| 630 | let tools_b = vec![make_tool("write_file")]; |
| 631 | let a = PrefixFingerprint::compute("system", Some(&tools_a)); |
| 632 | let b = PrefixFingerprint::compute("system", Some(&tools_b)); |
| 633 | assert_ne!(a.combined_sha256, b.combined_sha256); |
| 634 | } |
| 635 | |
| 636 | #[test] |
| 637 | fn manager_starts_stable() { |
| 638 | let mut mgr = PrefixStabilityManager::new("system prompt", None); |
| 639 | assert!(mgr.check_and_update("system prompt", None).unwrap()); |
| 640 | assert_eq!(mgr.change_count(), 0); |
| 641 | assert_eq!(mgr.check_count(), 1); |
| 642 | } |
| 643 | |
| 644 | #[test] |
| 645 | fn manager_detects_change() { |
| 646 | let mut mgr = PrefixStabilityManager::new("system prompt", None); |
| 647 | let result = mgr.check_and_update("different prompt", None); |
| 648 | assert!(result.is_err()); |
| 649 | assert_eq!(mgr.change_count(), 1); |
| 650 | let change = mgr.last_change().unwrap(); |
| 651 | assert!(change.system_changed); |
| 652 | assert!(!change.tools_changed); |
| 653 | } |
| 654 | |
| 655 | #[test] |
| 656 | fn manager_detects_tool_change() { |
| 657 | let tools_a = vec![make_tool("read_file")]; |
| 658 | let tools_b = vec![make_tool("write_file")]; |
| 659 | let mut mgr = PrefixStabilityManager::new("system", Some(&tools_a)); |
| 660 | let result = mgr.check_and_update("system", Some(&tools_b)); |
| 661 | assert!(result.is_err()); |
| 662 | let change = mgr.last_change().unwrap(); |
| 663 | assert!(!change.system_changed); |
| 664 | assert!(change.tools_changed); |
| 665 | } |
| 666 | |
| 667 | #[test] |
| 668 | fn manager_re_pins_after_change() { |
| 669 | let mut mgr = PrefixStabilityManager::new("old", None); |
| 670 | let _ = mgr.check_and_update("new", None); |
| 671 | // After re-pin, the new "new" should be stable. |
| 672 | assert!(mgr.check_and_update("new", None).unwrap()); |
| 673 | assert_eq!(mgr.change_count(), 1); |
| 674 | } |
| 675 | |
| 676 | #[test] |
| 677 | fn stability_ratio_is_one_for_no_changes() { |
| 678 | let mut mgr = PrefixStabilityManager::new("hello", None); |
| 679 | mgr.check_and_update("hello", None).unwrap(); |
| 680 | mgr.check_and_update("hello", None).unwrap(); |
| 681 | assert!((mgr.stability_ratio() - 1.0).abs() < f64::EPSILON); |
| 682 | assert_eq!(mgr.check_count(), 2); |
| 683 | assert_eq!(mgr.change_count(), 0); |
| 684 | } |
| 685 | |
| 686 | #[test] |
| 687 | fn stability_ratio_reflects_change_rate() { |
| 688 | let mut mgr = PrefixStabilityManager::new("hello", None); |
| 689 | mgr.check_and_update("hello", None).unwrap(); // check 1: stable |
| 690 | let _ = mgr.check_and_update("world", None); // check 2: changed |
| 691 | mgr.check_and_update("world", None).unwrap(); // check 3: stable |
| 692 | // 2 stable out of 3 checks = 0.666... |
| 693 | // (check_count=0 at start, so 3 checks: 3 checks - 1 change = 2 stable) |
| 694 | assert!((mgr.stability_ratio() - 2.0 / 3.0).abs() < 0.01); |
| 695 | assert_eq!(mgr.check_count(), 3); |
| 696 | assert_eq!(mgr.change_count(), 1); |
| 697 | } |
| 698 | |
| 699 | #[test] |
| 700 | fn empty_tools_and_none_tools_produce_same_hash() { |
| 701 | let empty = PrefixFingerprint::compute("system", Some(&[])); |
| 702 | let none = PrefixFingerprint::compute("system", None); |
| 703 | // Both should produce sha256(b"") for the tool component |
| 704 | assert_eq!(empty.tools_sha256, none.tools_sha256); |
| 705 | } |
| 706 | |
| 707 | #[test] |
| 708 | fn empty_system_produces_sha256_of_empty_string() { |
| 709 | let fp = PrefixFingerprint::compute("", None); |
| 710 | let expected = sha256_hex(b""); |
| 711 | assert_eq!(fp.system_sha256, expected); |
| 712 | } |
| 713 | |
| 714 | #[test] |
| 715 | fn prefix_change_description_is_informative() { |
| 716 | let old = PrefixFingerprint::compute("old", None); |
| 717 | let new = PrefixFingerprint::compute("new", None); |
| 718 | let change = PrefixChange { |
| 719 | old, |
| 720 | new, |
| 721 | system_changed: true, |
| 722 | tools_changed: false, |
| 723 | }; |
| 724 | assert_eq!( |
| 725 | change.description(), |
| 726 | "prefix cache invalidated: system prompt changed" |
| 727 | ); |
| 728 | assert_eq!(change.label(), "sys"); |
| 729 | } |
| 730 | |
| 731 | #[test] |
| 732 | fn new_unpinned_has_no_change_history() { |
| 733 | let mut mgr = PrefixStabilityManager::new_unpinned(); |
| 734 | assert!(mgr.pinned_fingerprint().is_none()); |
| 735 | assert!(mgr.current_fingerprint().is_none()); |
| 736 | assert!(mgr.last_change().is_none()); |
| 737 | assert_eq!(mgr.change_count(), 0); |
| 738 | assert_eq!(mgr.check_count(), 0); |
| 739 | // First check should pin automatically and count as a check. |
| 740 | assert!(mgr.check_and_update("hello", None).unwrap()); |
| 741 | assert!(mgr.pinned_fingerprint().is_some()); |
| 742 | assert_eq!(mgr.check_count(), 1); |
| 743 | } |
| 744 | |
| 745 | #[test] |
| 746 | fn fingerprint_detects_schema_change_not_just_name_change() { |
| 747 | let tool_a = make_tool("my_tool"); |
| 748 | let mut tool_a_v2 = make_tool("my_tool"); |
| 749 | tool_a_v2.description = "updated description".to_string(); |
| 750 | |
| 751 | let a = PrefixFingerprint::compute("system", Some(&[tool_a])); |
| 752 | let b = PrefixFingerprint::compute("system", Some(&[tool_a_v2])); |
| 753 | // Same name, different description — must produce different hash. |
| 754 | assert_ne!(a.tools_sha256, b.tools_sha256); |
| 755 | assert_ne!(a.combined_sha256, b.combined_sha256); |
| 756 | } |
| 757 | |
| 758 | #[test] |
| 759 | fn system_prompt_text_returns_empty_for_none() { |
| 760 | assert_eq!(system_prompt_text(None), ""); |
| 761 | } |
| 762 | |
| 763 | // ── ToolCatalogCache tests ────────────────────────────────── |
| 764 | |
| 765 | #[test] |
| 766 | fn tool_catalog_cache_miss_then_hit_returns_same_digest() { |
| 767 | let mut cache = ToolCatalogCache::new(); |
| 768 | let tools = vec![make_tool("read_file"), make_tool("write_file")]; |
| 769 | |
| 770 | let first = cache.fingerprint_for(&tools); |
| 771 | assert_eq!(cache.len(), 1); |
| 772 | |
| 773 | let second = cache.fingerprint_for(&tools); |
| 774 | assert_eq!(cache.len(), 1, "second call should be a cache hit"); |
| 775 | assert_eq!(first.sha256_hex, second.sha256_hex); |
| 776 | } |
| 777 | |
| 778 | #[test] |
| 779 | fn tool_catalog_cache_different_tool_sets_dont_collide() { |
| 780 | let mut cache = ToolCatalogCache::new(); |
| 781 | let a = vec![make_tool("read_file")]; |
| 782 | let b = vec![make_tool("write_file")]; |
| 783 | |
| 784 | let entry_a = cache.fingerprint_for(&a); |
| 785 | let entry_b = cache.fingerprint_for(&b); |
| 786 | assert_eq!(cache.len(), 2); |
| 787 | assert_ne!(entry_a.sha256_hex, entry_b.sha256_hex); |
| 788 | } |
| 789 | |
| 790 | #[test] |
| 791 | fn tool_catalog_cache_pinned_by_input_order() { |
| 792 | // The identity hash includes the input order so re-registering the |
| 793 | // same set with a different permutation produces a separate cache |
| 794 | // entry. The sorted-and-joined digest still matches the order- |
| 795 | // independent fingerprint that the chat API sees. |
| 796 | let mut cache = ToolCatalogCache::new(); |
| 797 | let a = vec![make_tool("read_file"), make_tool("write_file")]; |
| 798 | let b = vec![make_tool("write_file"), make_tool("read_file")]; |
| 799 | let entry_a = cache.fingerprint_for(&a); |
| 800 | let entry_b = cache.fingerprint_for(&b); |
| 801 | // Digests match (sorted join) but the two cache entries are distinct |
| 802 | // because their identities differ. |
| 803 | assert_eq!(entry_a.sha256_hex, entry_b.sha256_hex); |
| 804 | assert_eq!(cache.len(), 2); |
| 805 | } |
| 806 | |
| 807 | #[test] |
| 808 | fn tool_catalog_cache_detects_schema_change() { |
| 809 | let mut cache = ToolCatalogCache::new(); |
| 810 | let tool_v1 = make_tool("t"); |
| 811 | let mut tool_v2 = make_tool("t"); |
| 812 | tool_v2.description = "updated".to_string(); |
| 813 | |
| 814 | let entry_v1 = cache.fingerprint_for(&[tool_v1]); |
| 815 | let entry_v2 = cache.fingerprint_for(&[tool_v2]); |
| 816 | assert_ne!(entry_v1.sha256_hex, entry_v2.sha256_hex); |
| 817 | assert_eq!(cache.len(), 2); |
| 818 | } |
| 819 | |
| 820 | #[test] |
| 821 | fn tool_catalog_cache_respects_capacity() { |
| 822 | let mut cache = ToolCatalogCache::with_capacity(2); |
| 823 | cache.fingerprint_for(&[make_tool("a")]); |
| 824 | cache.fingerprint_for(&[make_tool("b")]); |
| 825 | cache.fingerprint_for(&[make_tool("c")]); |
| 826 | assert_eq!(cache.len(), 2); |
| 827 | // The first entry was evicted; a re-query for it should miss. |
| 828 | let re_entry = cache.fingerprint_for(&[make_tool("a")]); |
| 829 | // After the re-query, the cache has [b, c, a] — 3 entries? No, |
| 830 | // capacity 2 means oldest is evicted when we insert the 3rd unique. |
| 831 | // After inserting a, the cache holds the most recent 2: {c, a}. |
| 832 | assert_eq!(cache.len(), 2); |
| 833 | // The returned digest should match a fresh fingerprint of the same set. |
| 834 | let fresh = cache.fingerprint_for(&[make_tool("a")]); |
| 835 | assert_eq!(re_entry.sha256_hex, fresh.sha256_hex); |
| 836 | } |
| 837 | |
| 838 | #[test] |
| 839 | fn tool_catalog_cache_invalidate_clears_all() { |
| 840 | let mut cache = ToolCatalogCache::new(); |
| 841 | cache.fingerprint_for(&[make_tool("a")]); |
| 842 | cache.fingerprint_for(&[make_tool("b")]); |
| 843 | cache.invalidate(); |
| 844 | assert!(cache.is_empty()); |
| 845 | assert_eq!(cache.len(), 0); |
| 846 | } |
| 847 | |
| 848 | #[test] |
| 849 | fn tool_catalog_cache_empty_slice_uses_zero_capacity_path() { |
| 850 | // Empty input is fine — should produce a stable, non-empty digest. |
| 851 | let mut cache = ToolCatalogCache::new(); |
| 852 | let entry = cache.fingerprint_for(&[]); |
| 853 | assert!(!entry.sha256_hex.is_empty()); |
| 854 | let again = cache.fingerprint_for(&[]); |
| 855 | assert_eq!(entry.sha256_hex, again.sha256_hex); |
| 856 | } |
| 857 | |
| 858 | #[test] |
| 859 | fn compute_with_tool_cache_matches_compute_uncached() { |
| 860 | // The cached and uncached paths must produce identical fingerprints |
| 861 | // for the same inputs — otherwise we'd silently corrupt the prefix |
| 862 | // cache and invalidate every request. |
| 863 | let mut cache = ToolCatalogCache::new(); |
| 864 | let tools = vec![make_tool("alpha"), make_tool("beta")]; |
| 865 | |
| 866 | let cached = PrefixFingerprint::compute_with_tool_cache("sys", Some(&tools), &mut cache); |
| 867 | let uncached = PrefixFingerprint::compute("sys", Some(&tools)); |
| 868 | assert_eq!(cached.combined_sha256, uncached.combined_sha256); |
| 869 | assert_eq!(cached.tools_sha256, uncached.tools_sha256); |
| 870 | } |
| 871 | |
| 872 | #[test] |
| 873 | fn manager_check_and_update_uses_cached_tool_fingerprint() { |
| 874 | // After the first call populates the cache, subsequent calls with |
| 875 | // the same tool list should not invalidate the prefix. |
| 876 | let tools = vec![make_tool("t1")]; |
| 877 | let mut mgr = PrefixStabilityManager::new("sys", Some(&tools)); |
| 878 | assert!(mgr.check_and_update("sys", Some(&tools)).is_ok()); |
| 879 | assert!(mgr.check_and_update("sys", Some(&tools)).is_ok()); |
| 880 | assert_eq!(mgr.change_count(), 0); |
| 881 | } |
| 882 | } |
| 883 |