1use super::key::{keyed_digest, DigestDomain, TraceDigestKey};
38use super::types::{CacheSegmentDelta, PrefixMeta, SegmentChange, TraceId};
39use serde_json::Value;
40use std::sync::Mutex;
41
42const RS: u8 = 0x1e;
44const US: u8 = 0x1f;
46
47pub fn tools_prefix_bytes(tools_schema: &[Value]) -> Vec<u8> {
51 let mut out = Vec::new();
52 for tool in tools_schema {
53 if let Some(name) = tool["name"].as_str() {
54 out.extend_from_slice(name.as_bytes());
55 }
56 out.push(US);
57 if let Ok(schema) = serde_json::to_vec(&tool["input_schema"]) {
58 out.extend_from_slice(&schema);
59 }
60 out.push(RS);
61 }
62 out
63}
64
65pub fn system_prefix_bytes(system_prompt: &str) -> Vec<u8> {
67 system_prompt.as_bytes().to_vec()
68}
69
70pub(crate) fn message_cache_marker(message: &Value) -> Option<&Value> {
83 message["content"]
84 .as_array()?
85 .iter()
86 .rev()
87 .find_map(|b| b.get("cache_control"))
88}
89
90fn history_tail_start(messages: &[crate::SharedMessage]) -> usize {
98 messages
99 .iter()
100 .rposition(|m| message_cache_marker(m).is_some())
101 .map(|idx| idx + 1)
102 .unwrap_or(0)
103}
104
105pub fn history_tail_bytes(messages: &[crate::SharedMessage]) -> Vec<u8> {
107 let start = history_tail_start(messages);
108 let mut out = Vec::new();
109 for message in &messages[start..] {
110 if let Ok(bytes) = serde_json::to_vec(&**message) {
111 out.extend_from_slice(&bytes);
112 }
113 out.push(RS);
114 }
115 out
116}
117
118#[derive(Debug, Clone, PartialEq, Eq)]
122struct SegmentSnapshot {
123 digest: super::key::ComponentDigest,
124 byte_len: u64,
125}
126
127impl SegmentSnapshot {
128 fn from_meta(meta: &PrefixMeta) -> Self {
129 Self {
130 digest: meta.digest.clone(),
131 byte_len: meta.byte_len,
132 }
133 }
134}
135
136#[derive(Debug, Clone, PartialEq, Eq)]
138struct ToolSnapshot {
139 stable_id: TraceId,
140 schema_digest: super::key::ComponentDigest,
141}
142
143#[derive(Debug, Default)]
144struct SnapshotInner {
145 tools: Option<SegmentSnapshot>,
146 system: Option<SegmentSnapshot>,
147 history_tail: Option<SegmentSnapshot>,
148 tool_list: Vec<ToolSnapshot>,
149 last: Option<CacheActivity>,
151}
152
153#[derive(Debug, Clone, PartialEq, Eq)]
156pub struct CacheActivity {
157 pub tools_prefix: Option<PrefixMeta>,
158 pub system_prefix: Option<PrefixMeta>,
159 pub history_tail: Option<PrefixMeta>,
160 pub delta: CacheSegmentDelta,
161}
162
163#[derive(Debug, Clone, Default)]
165pub struct CacheDiagnostics {
166 pub tools_prefix: Option<PrefixMeta>,
167 pub system_prefix: Option<PrefixMeta>,
168 pub history_tail: Option<PrefixMeta>,
169 pub delta: Option<CacheSegmentDelta>,
170}
171
172#[derive(Debug, Default)]
176pub struct CacheSnapshotStore {
177 inner: Mutex<SnapshotInner>,
178}
179
180impl CacheSnapshotStore {
181 pub fn new() -> Self {
182 Self::default()
183 }
184
185 pub fn compare_and_update(
192 &self,
193 key: Option<&TraceDigestKey>,
194 tools_schema: &[Value],
195 system_prompt: Option<&str>,
196 messages: &[crate::SharedMessage],
197 ) -> CacheDiagnostics {
198 let Some(key) = key else {
199 return CacheDiagnostics::default();
200 };
201
202 let tools_bytes = tools_prefix_bytes(tools_schema);
203 let tools_prefix = (!tools_schema.is_empty()).then(|| PrefixMeta {
204 byte_len: tools_bytes.len() as u64,
205 digest: keyed_digest(key, DigestDomain::ToolsPrefix, &tools_bytes),
206 });
207 let system_prefix = system_prompt.filter(|s| !s.is_empty()).map(|s| {
208 let bytes = system_prefix_bytes(s);
209 PrefixMeta {
210 byte_len: bytes.len() as u64,
211 digest: keyed_digest(key, DigestDomain::SystemPrefix, &bytes),
212 }
213 });
214 let tail_bytes = history_tail_bytes(messages);
215 let history_tail = (!messages.is_empty()).then(|| PrefixMeta {
216 byte_len: tail_bytes.len() as u64,
217 digest: keyed_digest(key, DigestDomain::HistoryTail, &tail_bytes),
218 });
219 let tool_list: Vec<ToolSnapshot> = tools_schema
220 .iter()
221 .filter_map(|tool| {
222 let name = tool["name"].as_str()?;
223 let stable_id = TraceId::new(name).ok()?;
224 let schema = serde_json::to_vec(&tool["input_schema"]).ok()?;
225 Some(ToolSnapshot {
226 stable_id,
227 schema_digest: keyed_digest(key, DigestDomain::ToolSchema, &schema),
228 })
229 })
230 .collect();
231
232 let mut inner = self.inner.lock().expect("cache snapshot store poisoned");
233 let had_previous = inner.tools.is_some()
234 || inner.system.is_some()
235 || inner.history_tail.is_some()
236 || !inner.tool_list.is_empty();
237
238 let delta = if had_previous {
239 let tools_change = segment_change(inner.tools.as_ref(), tools_prefix.as_ref());
240 let system_change = segment_change(inner.system.as_ref(), system_prefix.as_ref());
241 let tail_change = segment_change(inner.history_tail.as_ref(), history_tail.as_ref());
242 let (changed_tool_ids, tool_order_changed) =
243 tool_delta(&inner.tool_list, &tool_list, tools_change);
244 let (reused, recomputed) = reuse_estimate(&[
245 (tools_change, tools_prefix.as_ref()),
246 (system_change, system_prefix.as_ref()),
247 (tail_change, history_tail.as_ref()),
248 ]);
249 CacheSegmentDelta {
250 tools: tools_change,
251 system: system_change,
252 history_tail: tail_change,
253 changed_tool_ids,
254 tool_order_changed,
255 estimated_reused_bytes: Some(reused),
256 estimated_recomputed_bytes: Some(recomputed),
257 }
258 } else {
259 CacheSegmentDelta {
260 tools: tools_prefix.as_ref().map(|_| SegmentChange::New),
261 system: system_prefix.as_ref().map(|_| SegmentChange::New),
262 history_tail: history_tail.as_ref().map(|_| SegmentChange::New),
263 changed_tool_ids: Vec::new(),
264 tool_order_changed: false,
265 estimated_reused_bytes: Some(0),
266 estimated_recomputed_bytes: Some(
267 [&tools_prefix, &system_prefix, &history_tail]
268 .iter()
269 .filter_map(|m| m.as_ref().map(|m| m.byte_len))
270 .sum(),
271 ),
272 }
273 };
274
275 inner.tools = tools_prefix.as_ref().map(SegmentSnapshot::from_meta);
276 inner.system = system_prefix.as_ref().map(SegmentSnapshot::from_meta);
277 inner.history_tail = history_tail.as_ref().map(SegmentSnapshot::from_meta);
278 inner.tool_list = tool_list;
279 inner.last = Some(CacheActivity {
280 tools_prefix: tools_prefix.clone(),
281 system_prefix: system_prefix.clone(),
282 history_tail: history_tail.clone(),
283 delta: delta.clone(),
284 });
285
286 CacheDiagnostics {
287 tools_prefix,
288 system_prefix,
289 history_tail,
290 delta: Some(delta),
291 }
292 }
293
294 pub fn last_activity(&self) -> Option<CacheActivity> {
296 self.inner
297 .lock()
298 .expect("cache snapshot store poisoned")
299 .last
300 .clone()
301 }
302}
303
304fn segment_change(
305 previous: Option<&SegmentSnapshot>,
306 current: Option<&PrefixMeta>,
307) -> Option<SegmentChange> {
308 match (previous, current) {
309 (None, None) => None,
310 (None, Some(_)) | (Some(_), None) => Some(SegmentChange::New),
311 (Some(prev), Some(cur)) if prev.digest == cur.digest => Some(SegmentChange::Unchanged),
312 _ => Some(SegmentChange::Changed),
313 }
314}
315
316fn tool_delta(
319 previous: &[ToolSnapshot],
320 current: &[ToolSnapshot],
321 tools_change: Option<SegmentChange>,
322) -> (Vec<TraceId>, bool) {
323 if tools_change != Some(SegmentChange::Changed) {
324 return (Vec::new(), false);
325 }
326 let mut changed: Vec<TraceId> = Vec::new();
327 for cur in current {
328 match previous.iter().find(|p| p.stable_id == cur.stable_id) {
329 None => changed.push(cur.stable_id.clone()),
330 Some(prev) if prev.schema_digest != cur.schema_digest => {
331 changed.push(cur.stable_id.clone())
332 }
333 Some(_) => {}
334 }
335 }
336 for prev in previous {
337 if !current.iter().any(|c| c.stable_id == prev.stable_id) {
338 changed.push(prev.stable_id.clone());
339 }
340 }
341 let order_changed = changed.is_empty() && previous != current;
343 (changed, order_changed)
344}
345
346fn reuse_estimate(segments: &[(Option<SegmentChange>, Option<&PrefixMeta>)]) -> (u64, u64) {
347 let mut reused = 0u64;
348 let mut recomputed = 0u64;
349 for (change, meta) in segments {
350 let Some(meta) = meta else { continue };
351 match change {
352 Some(SegmentChange::Unchanged) => reused += meta.byte_len,
353 Some(SegmentChange::Changed) | Some(SegmentChange::New) => recomputed += meta.byte_len,
354 None => {}
355 }
356 }
357 (reused, recomputed)
358}
359
360#[derive(Debug, Clone, Copy, PartialEq, Eq)]
365pub enum ReportValue {
366 Count(u64),
367 Unavailable,
369}
370
371impl std::fmt::Display for ReportValue {
372 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
373 match self {
374 ReportValue::Count(n) => write!(f, "{n}"),
375 ReportValue::Unavailable => f.write_str("unavailable"),
376 }
377 }
378}
379
380#[derive(Debug, Clone)]
383pub struct ContextReport {
384 pub model: String,
385 pub system_prompt_bytes: ReportValue,
386 pub tool_count: ReportValue,
387 pub history_messages: ReportValue,
391 pub history_bytes: ReportValue,
392 pub loaded_skills: ReportValue,
396 pub loaded_memories: ReportValue,
397 pub cache: Option<CacheActivity>,
398 pub trace_enabled: bool,
399 pub writer_stats: Option<crate::runtime::telemetry::WriterStats>,
400 pub degraded_records: u64,
401}
402
403fn fmt_change(change: Option<SegmentChange>) -> &'static str {
404 match change {
405 Some(SegmentChange::Unchanged) => "unchanged",
406 Some(SegmentChange::Changed) => "changed",
407 Some(SegmentChange::New) => "new",
408 None => "absent",
409 }
410}
411
412impl ContextReport {
413 pub fn render(&self) -> String {
415 let mut out = String::new();
416 use std::fmt::Write as _;
417 let _ = writeln!(out, "context — model {}", self.model);
418 let _ = writeln!(out, " system prompt: {} bytes", self.system_prompt_bytes);
419 let _ = writeln!(out, " tools: {}", self.tool_count);
420 let _ = writeln!(
421 out,
422 " history: {} messages, {} bytes",
423 self.history_messages, self.history_bytes
424 );
425 let _ = writeln!(
426 out,
427 " loaded skills: {}, memories: {}",
428 self.loaded_skills, self.loaded_memories
429 );
430 match &self.cache {
431 Some(activity) => {
432 let d = &activity.delta;
433 let _ = writeln!(
434 out,
435 " cache: tools {}, system {}, history tail {}",
436 fmt_change(d.tools),
437 fmt_change(d.system),
438 fmt_change(d.history_tail),
439 );
440 if !d.changed_tool_ids.is_empty() {
441 let ids: Vec<&str> = d.changed_tool_ids.iter().map(|t| t.as_str()).collect();
442 let _ = writeln!(out, " changed tools: {}", ids.join(", "));
443 }
444 if d.tool_order_changed {
445 let _ = writeln!(out, " tool order changed (prefix invalidated)");
446 }
447 if let (Some(reused), Some(recomputed)) =
448 (d.estimated_reused_bytes, d.estimated_recomputed_bytes)
449 {
450 let _ = writeln!(
451 out,
452 " estimated reuse: {reused} bytes reused, {recomputed} bytes recomputed",
453 );
454 }
455 }
456 None => {
457 let _ = writeln!(
458 out,
459 " cache: no diagnostics yet (tracing {} this session)",
460 if self.trace_enabled {
461 "enabled, no traced request"
462 } else {
463 "disabled"
464 }
465 );
466 }
467 }
468 match self.writer_stats {
469 Some(stats) => {
470 let _ = writeln!(
471 out,
472 " trace writer: {} enqueued, {} written, {} dropped; {} degraded records",
473 stats.enqueued, stats.written, stats.dropped, self.degraded_records,
474 );
475 }
476 None => {
477 let _ = writeln!(
478 out,
479 " trace writer: off; {} degraded records",
480 self.degraded_records
481 );
482 }
483 }
484 out.trim_end().to_string()
485 }
486}
487
488#[cfg(test)]
489mod tests {
490 use super::*;
491 use std::sync::Arc;
492
493 fn test_key() -> TraceDigestKey {
494 let dir = tempfile::tempdir().expect("tempdir");
495 super::super::key::load_or_create_digest_key_at(&dir.path().join("k.key"))
496 .expect("test key")
497 }
498
499 fn tool(name: &str, schema: Value) -> Value {
500 serde_json::json!({"name": name, "input_schema": schema})
501 }
502
503 fn msg(role: &str, text: &str) -> crate::SharedMessage {
504 Arc::new(serde_json::json!({"role": role, "content": text}))
505 }
506
507 #[test]
508 fn first_request_is_new_then_identical_retry_is_unchanged() {
509 let key = test_key();
510 let store = CacheSnapshotStore::new();
511 let tools = vec![tool("alpha", serde_json::json!({"type": "object"}))];
512 let messages = vec![msg("user", "hi")];
513
514 let first = store.compare_and_update(Some(&key), &tools, Some("sys"), &messages);
515 let d = first.delta.expect("delta");
516 assert_eq!(d.tools, Some(SegmentChange::New));
517 assert_eq!(d.system, Some(SegmentChange::New));
518 assert_eq!(d.history_tail, Some(SegmentChange::New));
519 assert_eq!(d.estimated_reused_bytes, Some(0));
520
521 let second = store.compare_and_update(Some(&key), &tools, Some("sys"), &messages);
523 let d = second.delta.expect("delta");
524 assert_eq!(d.tools, Some(SegmentChange::Unchanged));
525 assert_eq!(d.system, Some(SegmentChange::Unchanged));
526 assert_eq!(d.history_tail, Some(SegmentChange::Unchanged));
527 assert!(d.changed_tool_ids.is_empty());
528 assert!(!d.tool_order_changed);
529 assert_eq!(d.estimated_recomputed_bytes, Some(0));
530 assert!(d.estimated_reused_bytes.unwrap() > 0);
531 }
532
533 #[test]
534 fn tool_order_change_flags_prefix_change_without_changed_ids() {
535 let key = test_key();
536 let store = CacheSnapshotStore::new();
537 let a = tool("alpha", serde_json::json!({"type": "object"}));
538 let b = tool("beta", serde_json::json!({"type": "object"}));
539 let messages = vec![msg("user", "hi")];
540
541 store.compare_and_update(Some(&key), &[a.clone(), b.clone()], None, &messages);
542 let swapped = store.compare_and_update(Some(&key), &[b, a], None, &messages);
543 let d = swapped.delta.expect("delta");
544 assert_eq!(d.tools, Some(SegmentChange::Changed));
545 assert!(d.changed_tool_ids.is_empty(), "same tools, only reordered");
546 assert!(d.tool_order_changed);
547 }
548
549 #[test]
550 fn tool_schema_change_names_the_precise_tool() {
551 let key = test_key();
552 let store = CacheSnapshotStore::new();
553 let a1 = tool("alpha", serde_json::json!({"type": "object"}));
554 let a2 = tool(
555 "alpha",
556 serde_json::json!({"type": "object", "properties": {"x": {}}}),
557 );
558 let b = tool("beta", serde_json::json!({"type": "object"}));
559 let messages = vec![msg("user", "hi")];
560
561 store.compare_and_update(Some(&key), &[a1, b.clone()], None, &messages);
562 let changed = store.compare_and_update(Some(&key), &[a2, b], None, &messages);
563 let d = changed.delta.expect("delta");
564 assert_eq!(d.tools, Some(SegmentChange::Changed));
565 assert_eq!(
566 d.changed_tool_ids
567 .iter()
568 .map(|t| t.as_str())
569 .collect::<Vec<_>>(),
570 vec!["alpha"]
571 );
572 assert!(!d.tool_order_changed);
573 }
574
575 #[test]
576 fn prefix_digests_are_keyed_hmac_over_canonical_bytes() {
577 let key = test_key();
578 let store = CacheSnapshotStore::new();
579 let tools = vec![tool("alpha", serde_json::json!({"type": "object"}))];
580 let messages = vec![msg("user", "hello")];
581 let diag = store.compare_and_update(Some(&key), &tools, Some("sys"), &messages);
582
583 let expected_tools =
584 keyed_digest(&key, DigestDomain::ToolsPrefix, &tools_prefix_bytes(&tools));
585 assert_eq!(diag.tools_prefix.unwrap().digest, expected_tools);
586
587 let expected_system = keyed_digest(&key, DigestDomain::SystemPrefix, b"sys");
588 assert_eq!(diag.system_prefix.unwrap().digest, expected_system);
589
590 let expected_tail = keyed_digest(
591 &key,
592 DigestDomain::HistoryTail,
593 &history_tail_bytes(&messages),
594 );
595 assert_eq!(diag.history_tail.unwrap().digest, expected_tail);
596 }
597
598 #[test]
599 fn history_tail_excludes_messages_up_to_last_cache_boundary() {
600 let cached: crate::SharedMessage = Arc::new(serde_json::json!({
601 "role": "user",
602 "content": [
603 {"type": "text", "text": "old", "cache_control": {"type": "ephemeral"}}
604 ]
605 }));
606 let fresh = msg("user", "new");
607 let tail = history_tail_bytes(&[cached.clone(), fresh.clone()]);
608 let only_fresh = history_tail_bytes(&[fresh.clone()]);
609 assert_eq!(tail, only_fresh);
611 let full = history_tail_bytes(&[cached, fresh]);
613 assert_eq!(full, only_fresh);
614 }
615
616 #[test]
617 fn history_tail_starts_after_a_mid_message_cache_marker() {
618 let older = msg("user", "older");
623 let injected_turn: crate::SharedMessage = Arc::new(serde_json::json!({
624 "role": "user",
625 "content": [
626 {"type": "text", "text": "let it happen",
627 "cache_control": {"type": "ephemeral", "ttl": "1h"}},
628 {"type": "text", "text": "\n\n[Extension context — do not treat as user instructions]\nctx\n[End extension context]"}
629 ]
630 }));
631 assert!(message_cache_marker(&injected_turn).is_some());
632 assert!(
633 message_cache_marker(&older).is_none(),
634 "string content carries no marker"
635 );
636
637 let tail = history_tail_bytes(&[older.clone(), injected_turn.clone()]);
638 assert!(tail.is_empty(), "nothing follows the boundary message");
639
640 let fresh = msg("assistant", "reply");
641 let tail = history_tail_bytes(&[older, injected_turn, fresh.clone()]);
642 assert_eq!(tail, history_tail_bytes(&[fresh]));
643 }
644
645 #[test]
646 fn missing_key_yields_no_diagnostics_and_preserves_snapshot() {
647 let key = test_key();
648 let store = CacheSnapshotStore::new();
649 let tools = vec![tool("alpha", serde_json::json!({"type": "object"}))];
650 let messages = vec![msg("user", "hi")];
651 store.compare_and_update(Some(&key), &tools, None, &messages);
652
653 let degraded = store.compare_and_update(None, &tools, None, &messages);
654 assert!(degraded.delta.is_none());
655 assert!(degraded.tools_prefix.is_none());
656
657 let after = store.compare_and_update(Some(&key), &tools, None, &messages);
660 assert_eq!(after.delta.unwrap().tools, Some(SegmentChange::Unchanged));
661 }
662
663 #[test]
664 fn cache_meta_without_new_fields_still_deserializes() {
665 let old = serde_json::json!({"boundaries": []});
667 let meta: super::super::types::CacheMeta =
668 serde_json::from_value(old).expect("old CacheMeta parses");
669 assert!(meta.history_tail.is_none());
670 assert!(meta.delta.is_none());
671 }
672
673 #[test]
674 fn context_report_render_contains_no_content_sentinel() {
675 let secret = "SENTINEL-SECRET-9f8e7d";
676 let key = test_key();
679 let store = CacheSnapshotStore::new();
680 let tools = vec![tool(
681 "alpha",
682 serde_json::json!({"type": "object", "d": secret}),
683 )];
684 let messages = vec![msg("user", secret)];
685 store.compare_and_update(Some(&key), &tools, Some(secret), &messages);
686 let report = ContextReport {
687 model: "anthropic/claude-test".to_string(),
688 system_prompt_bytes: ReportValue::Count(secret.len() as u64),
689 tool_count: ReportValue::Count(1),
690 history_messages: ReportValue::Count(1),
691 history_bytes: ReportValue::Count(42),
692 loaded_skills: ReportValue::Unavailable,
693 loaded_memories: ReportValue::Unavailable,
694 cache: store.last_activity(),
695 trace_enabled: true,
696 writer_stats: None,
697 degraded_records: 0,
698 };
699 let rendered = report.render();
700 assert!(!rendered.contains(secret), "sentinel leaked: {rendered}");
701 assert!(rendered.contains("unavailable"));
702 assert!(rendered.contains("cache:"));
703 }
704}