1use std::{
10 collections::BTreeMap,
11 fs::{self, OpenOptions},
12 io::Write,
13 path::PathBuf,
14};
15
16use http::{HeaderMap, HeaderValue, Response};
19use serde::{Deserialize, Serialize};
20use serde_json::Value;
21
22#[derive(Clone, Debug, Deserialize, Serialize)]
24pub struct DeterminismProof {
25 pub request_id: String,
27 pub prefix_hash_before: String,
29 pub prefix_hash_after: String,
31 pub is_stable: bool,
34 pub frozen_bytes: usize,
36 pub modification_start_byte: usize,
38}
39
40#[derive(Clone, Debug)]
42pub struct DeterminismGuard {
43 request_id: String,
44}
45
46impl DeterminismGuard {
47 #[must_use]
48 pub fn new(request_id: impl Into<String>) -> Self {
49 Self {
50 request_id: request_id.into(),
51 }
52 }
53
54 #[must_use]
56 pub fn verify(&self, before: &[Value], after: &[Value]) -> DeterminismProof {
57 verify_with_request_id(&self.request_id, before, after)
58 }
59
60 #[must_use]
84 pub fn verify_forwarded(
85 &self,
86 before: &[Value],
87 after: &[Value],
88 sent: Option<&crate::prefix_replay::SentPrefix>,
89 ) -> DeterminismProof {
90 let empty = crate::prefix_replay::SentPrefix::default();
91 let sent = sent.unwrap_or(&empty);
92 let client = message_hashes(before);
93 let outgoing = message_hashes(after);
94 let protected = sent
95 .client
96 .iter()
97 .zip(&client)
98 .take_while(|(then, now)| then == now)
99 .count()
100 .min(sent.forwarded.len());
101 let then = &sent.forwarded[..protected];
102 let now = outgoing.get(..protected).unwrap_or(&outgoing);
103 let unchanged = then
104 .iter()
105 .zip(now)
106 .take_while(|(then, now)| then == now)
107 .count();
108 let prefix_hash_before = hash_list(then);
109 let prefix_hash_after = hash_list(now);
110 let elision_step = unchanged < now.len()
115 && after
116 .get(unchanged)
117 .is_some_and(crate::session_elide::message_is_elided)
118 && !before
119 .get(unchanged)
120 .is_some_and(crate::session_elide::message_is_elided);
121 DeterminismProof {
122 request_id: self.request_id.clone(),
123 is_stable: prefix_hash_before == prefix_hash_after || elision_step,
124 prefix_hash_before,
125 prefix_hash_after,
126 frozen_bytes: canonical_message_prefix_len(after, unchanged),
127 modification_start_byte: common_prefix_len(
128 &canonical_messages(before),
129 &canonical_messages(after),
130 ),
131 }
132 }
133}
134
135#[must_use]
146pub fn message_hashes(messages: &[Value]) -> Vec<u64> {
147 let last_marked = cache_breakpoint_len(messages).checked_sub(1);
148 messages
149 .iter()
150 .enumerate()
151 .map(|(index, message)| {
152 let cached = if Some(index) == last_marked {
153 cached_part(message)
154 } else {
155 without_cache_control(message)
156 };
157 let bytes = lean_ctx::core::canonical::canonical_json_bytes(&cached);
158 let digest = blake3::hash(&bytes);
159 let mut head = [0u8; 8];
160 head.copy_from_slice(&digest.as_bytes()[..8]);
161 u64::from_le_bytes(head)
162 })
163 .collect()
164}
165
166fn cached_part(message: &Value) -> Value {
169 let mut cached = without_cache_control(message);
170 if message.get("cache_control").is_none()
171 && let Some(blocks) = message.get("content").and_then(Value::as_array)
172 && let Some(last) = blocks.iter().rposition(block_has_cache_control)
173 && let Some(content) = cached.get_mut("content").and_then(Value::as_array_mut)
174 {
175 content.truncate(last + 1);
176 }
177 cached
178}
179
180fn block_has_cache_control(block: &Value) -> bool {
181 block.get("cache_control").is_some()
182 || block
183 .get("content")
184 .and_then(Value::as_array)
185 .is_some_and(|items| items.iter().any(|item| item.get("cache_control").is_some()))
186}
187
188fn hash_list(hashes: &[u64]) -> String {
189 let mut hasher = blake3::Hasher::new();
190 for hash in hashes {
191 hasher.update(&hash.to_le_bytes());
192 }
193 hasher.finalize().to_hex().to_string()
194}
195
196#[must_use]
204pub fn verify_determinism(before: &[Value], after: &[Value]) -> DeterminismProof {
205 let request_id = blake3::hash(&canonical_messages(before))
206 .to_hex()
207 .to_string();
208 verify_with_request_id(&request_id, before, after)
209}
210
211fn verify_with_request_id(request_id: &str, before: &[Value], after: &[Value]) -> DeterminismProof {
212 let before_bytes = canonical_messages(before);
213 let after_bytes = canonical_messages(after);
214 let frozen_bytes = common_prefix_len(&before_bytes, &after_bytes);
215 let cached_messages = cache_breakpoint_len(before);
216 let proof_prefix_bytes = if cached_messages > 0 {
217 canonical_message_prefix_len(before, cached_messages)
218 } else {
219 0
225 };
226 let before_prefix = &before_bytes[..proof_prefix_bytes];
227 let after_prefix = after_bytes
228 .get(..proof_prefix_bytes)
229 .unwrap_or(&after_bytes);
230 let prefix_hash_before = blake3::hash(before_prefix).to_hex().to_string();
231 let prefix_hash_after = blake3::hash(after_prefix).to_hex().to_string();
232
233 DeterminismProof {
234 request_id: request_id.to_owned(),
235 is_stable: prefix_hash_before == prefix_hash_after,
236 prefix_hash_before,
237 prefix_hash_after,
238 frozen_bytes,
239 modification_start_byte: frozen_bytes,
240 }
241}
242
243#[must_use]
248pub fn cache_relevant_messages(request: &Value) -> Vec<Value> {
249 let mut messages = Vec::new();
250 if let Some(system) = request.get("system") {
251 messages.push(serde_json::json!({"role": "system", "content": system}));
252 }
253 if let Some(items) = request
254 .get("messages")
255 .or_else(|| request.get("input"))
256 .and_then(Value::as_array)
257 {
258 messages.extend(items.iter().cloned());
259 }
260 messages
261}
262
263#[must_use]
266#[cfg(feature = "http-server")]
267pub fn parse_request_body(
268 body: &[u8],
269 parts: &http::request::Parts,
270 limit: usize,
271) -> Option<Value> {
272 use crate::codec::{
273 RequestBodyEncoding, decode_gzip_bounded, decode_zstd_bounded, request_body_encoding,
274 };
275
276 let decoded = match request_body_encoding(parts) {
277 RequestBodyEncoding::Identity => body.to_vec(),
278 RequestBodyEncoding::Gzip => decode_gzip_bounded(body, limit).ok()?,
279 RequestBodyEncoding::Zstd => decode_zstd_bounded(body, limit).ok()?,
280 RequestBodyEncoding::Passthrough => return None,
281 };
282 serde_json::from_slice(&decoded).ok()
283}
284
285pub fn apply_response_headers<B>(response: &mut Response<B>, proof: &DeterminismProof) {
287 let headers = response.headers_mut();
288 headers.insert(
289 "x-leanctx-determinism",
290 HeaderValue::from_static(if proof.is_stable {
291 "stable"
292 } else {
293 "unstable"
294 }),
295 );
296 insert_header(
297 headers,
298 "x-leanctx-frozen-bytes",
299 &proof.frozen_bytes.to_string(),
300 );
301 headers.insert(
302 "x-leanctx-cache-safe",
303 HeaderValue::from_static(if proof.is_stable { "true" } else { "false" }),
304 );
305 insert_header(headers, "x-leanctx-proof-hash", &proof.prefix_hash_after);
306}
307
308pub fn ensure_response_headers<B>(response: &mut Response<B>) {
311 if !response.headers().contains_key("x-leanctx-determinism") {
312 apply_response_headers(response, &verify_determinism(&[], &[]));
313 }
314}
315
316fn insert_header(headers: &mut HeaderMap, name: &str, value: &str) {
317 if let (Ok(header_name), Ok(header_value)) = (
318 http::HeaderName::from_bytes(name.as_bytes()),
319 HeaderValue::from_str(value),
320 ) {
321 headers.insert(header_name, header_value);
322 }
323}
324
325#[derive(Clone, Debug, Deserialize, Serialize)]
326struct AuditRecord {
327 proof: DeterminismProof,
328 violation: bool,
329}
330
331#[derive(Clone, Debug, Deserialize, Serialize)]
333pub struct DeterminismAuditReport {
334 pub total_requests: usize,
335 pub stable_count: usize,
336 pub unstable_count: usize,
337 pub stability_rate: f64,
338 pub violations: Vec<DeterminismProof>,
339}
340
341const AUDIT_FILE: &str = "proxy-determinism-audit.jsonl";
342
343fn audit_path() -> Option<PathBuf> {
344 lean_ctx::core::data_dir::lean_ctx_data_dir()
345 .ok()
346 .map(|dir| dir.join(AUDIT_FILE))
347}
348
349pub fn record_audit(proof: &DeterminismProof, violation: bool) {
351 let Some(path) = audit_path() else {
352 return;
353 };
354 let Some(parent) = path.parent() else {
355 return;
356 };
357 if let Err(error) = fs::create_dir_all(parent) {
358 tracing::debug!(%error, "determinism audit directory unavailable");
359 return;
360 }
361 let record = AuditRecord {
362 proof: proof.clone(),
363 violation,
364 };
365 let Ok(line) = serde_json::to_string(&record) else {
366 return;
367 };
368 match OpenOptions::new().create(true).append(true).open(path) {
369 Ok(mut file) => {
370 if let Err(error) = writeln!(file, "{line}") {
371 tracing::debug!(%error, "determinism audit write failed");
372 }
373 }
374 Err(error) => tracing::debug!(%error, "determinism audit file unavailable"),
375 }
376}
377
378#[must_use]
380pub fn audit_report() -> DeterminismAuditReport {
381 let records = audit_path()
382 .and_then(|path| fs::read_to_string(path).ok())
383 .map(|contents| {
384 contents
385 .lines()
386 .filter_map(|line| serde_json::from_str::<AuditRecord>(line).ok())
387 .collect::<Vec<_>>()
388 })
389 .unwrap_or_default();
390 let total_requests = records.len();
391 let violations = records
392 .iter()
393 .filter(|record| record.violation || !record.proof.is_stable)
394 .map(|record| record.proof.clone())
395 .collect::<Vec<_>>();
396 let unstable_count = violations.len();
397 let stable_count = total_requests.saturating_sub(unstable_count);
398 let stability_rate = if total_requests == 0 {
399 1.0
400 } else {
401 stable_count as f64 / total_requests as f64
402 };
403 DeterminismAuditReport {
404 total_requests,
405 stable_count,
406 unstable_count,
407 stability_rate,
408 violations,
409 }
410}
411
412fn canonical_messages(messages: &[Value]) -> Vec<u8> {
417 let content: Vec<Value> = messages.iter().map(without_cache_control).collect();
418 lean_ctx::core::canonical::canonical_json_bytes(&Value::Array(content))
419}
420
421fn canonical_message_prefix_len(messages: &[Value], count: usize) -> usize {
422 if count == 0 {
423 return 0;
424 }
425 let mut bytes = Vec::from(b"[".as_slice());
426 for (index, message) in messages.iter().take(count).enumerate() {
427 if index > 0 {
428 bytes.push(b',');
429 }
430 bytes.extend(lean_ctx::core::canonical::canonical_json_bytes(
431 &without_cache_control(message),
432 ));
433 }
434 bytes.len()
435}
436
437fn without_cache_control(value: &Value) -> Value {
438 match value {
439 Value::Object(map) => Value::Object(
440 map.iter()
441 .filter(|(key, _)| key.as_str() != "cache_control")
442 .map(|(key, item)| (key.clone(), without_cache_control(item)))
443 .collect(),
444 ),
445 Value::Array(items) => Value::Array(items.iter().map(without_cache_control).collect()),
446 other => other.clone(),
447 }
448}
449
450fn common_prefix_len(before: &[u8], after: &[u8]) -> usize {
451 before
452 .iter()
453 .zip(after)
454 .take_while(|(left, right)| left == right)
455 .count()
456}
457
458fn cache_breakpoint_len(messages: &[Value]) -> usize {
459 messages
460 .iter()
461 .enumerate()
462 .filter_map(|(index, message)| message_has_cache_control(message).then_some(index + 1))
463 .last()
464 .unwrap_or_default()
465}
466
467fn message_has_cache_control(message: &Value) -> bool {
468 if message.get("cache_control").is_some() {
469 return true;
470 }
471 message
472 .get("content")
473 .and_then(Value::as_array)
474 .is_some_and(|blocks| blocks.iter().any(block_has_cache_control))
475}
476
477#[cfg(test)]
478mod tests {
479 use http::Response;
480 use serde_json::json;
481
482 use super::{
483 DeterminismGuard, apply_response_headers, canonical_messages, message_hashes,
484 verify_determinism,
485 };
486 use crate::prefix_replay::SentPrefix;
487
488 #[test]
494 fn guard_protects_what_the_proxy_forwarded_not_the_client_breakpoint() {
495 let marker = json!({"type": "ephemeral"});
496 let result = |id: &str, text: &str, marked: bool| {
497 let mut block = json!({"type": "tool_result", "tool_use_id": id, "content": text});
498 if marked {
499 block["cache_control"] = marker.clone();
500 }
501 json!({"role": "user", "content": [block]})
502 };
503 let opening = vec![
504 json!({"role": "system", "content": "rules"}),
505 json!({"role": "user", "content": "fix the bug"}),
506 json!({"role": "assistant", "content": "reading"}),
507 ];
508 let client_two = [opening.clone(), vec![result("t1", "raw one", true)]].concat();
509 let sent_two = [opening.clone(), vec![result("t1", "short one", true)]].concat();
510 let guard = DeterminismGuard::new("trace");
511
512 let first_sight = guard.verify_forwarded(&client_two, &sent_two, None);
513 assert!(first_sight.is_stable, "nothing was forwarded yet");
514
515 let sent = SentPrefix {
516 client: message_hashes(&client_two),
517 forwarded: message_hashes(&sent_two),
518 };
519 let mut client_three = client_two.clone();
520 client_three[3] = result("t1", "raw one", false);
521 client_three.push(json!({"role": "assistant", "content": "next"}));
522 client_three.push(result("t2", "raw two", true));
523 let mut compressed = client_three.clone();
524 compressed[3] = result("t1", "short one", false);
525 compressed[5] = result("t2", "short two", true);
526 assert!(!verify_determinism(&client_three, &compressed).is_stable);
527 let proof = guard.verify_forwarded(&client_three, &compressed, Some(&sent));
528 assert!(proof.is_stable, "the newest tool result is new content");
529 assert!(proof.frozen_bytes > 0);
530
531 let busted = guard.verify_forwarded(&client_three, &client_three, Some(&sent));
532 assert!(
533 !busted.is_stable,
534 "t1 left compressed before: sending it raw changes the cached prefix"
535 );
536
537 let mut steered = compressed.clone();
540 steered[5]["content"]
541 .as_array_mut()
542 .unwrap()
543 .push(json!({"type": "text", "text": "Be brief."}));
544 let sent = SentPrefix {
545 client: message_hashes(&client_three),
546 forwarded: message_hashes(&steered),
547 };
548 let mut client_four = client_three.clone();
549 client_four[5] = result("t2", "raw two", false);
550 client_four.push(json!({"role": "assistant", "content": "done"}));
551 client_four.push(result("t3", "raw three", true));
552 let mut unsteered = compressed.clone();
553 unsteered[5] = result("t2", "short two", false);
554 unsteered.extend_from_slice(&client_four[6..]);
555 assert!(
556 guard
557 .verify_forwarded(&client_four, &unsteered, Some(&sent))
558 .is_stable,
559 "a block after the last marker was never cached"
560 );
561 unsteered[5] = result("t2", "shorter two", false);
562 assert!(
563 !guard
564 .verify_forwarded(&client_four, &unsteered, Some(&sent))
565 .is_stable,
566 "the marked block itself was cached"
567 );
568
569 let stub = format!(
572 "{} (9000 chars); full text: cat x]",
573 crate::session_elide::STUB_OPEN
574 );
575 let mut elided = unsteered.clone();
576 elided[3] = result("t1", &stub, false);
577 elided[5] = result("t2", "short two", false);
578 assert!(
579 guard
580 .verify_forwarded(&client_four, &elided, Some(&sent))
581 .is_stable,
582 "an elision step breaks the prefix once, on purpose"
583 );
584 }
585
586 #[test]
587 fn stable_cached_prefix_produces_a_stable_proof() {
588 let before = vec![
589 json!({"role": "user", "content": "frozen", "cache_control": {"type": "ephemeral"}}),
590 json!({"role": "user", "content": "live"}),
591 ];
592 let after = vec![
593 before[0].clone(),
594 json!({"role": "user", "content": "rewritten live"}),
595 ];
596
597 let proof = verify_determinism(&before, &after);
598
599 assert!(proof.is_stable);
600 assert_eq!(proof.prefix_hash_before, proof.prefix_hash_after);
601 }
602
603 #[test]
604 fn modified_cached_prefix_is_unstable() {
605 let before = vec![json!({"role": "user", "content": "frozen", "cache_control": {}})];
606 let after = vec![json!({"role": "user", "content": "changed", "cache_control": {}})];
607
608 assert!(!verify_determinism(&before, &after).is_stable);
609 }
610
611 #[test]
612 fn single_message_without_cache_control_allows_compression() {
613 let before = vec![json!({"role": "user", "content": "frozen"})];
614 let after = vec![json!({"role": "user", "content": "changed"})];
615 assert!(verify_determinism(&before, &after).is_stable);
616 }
617
618 #[test]
619 fn last_message_modification_without_cache_control_is_stable() {
620 let before = vec![
621 json!({"role": "system", "content": "rules"}),
622 json!({"role": "user", "content": "original"}),
623 ];
624 let after = vec![
625 json!({"role": "system", "content": "rules"}),
626 json!({"role": "user", "content": "compressed"}),
627 ];
628 assert!(verify_determinism(&before, &after).is_stable);
629 }
630
631 #[test]
632 fn any_modification_without_cache_control_is_stable() {
633 let before = vec![
636 json!({"role": "system", "content": "rules"}),
637 json!({"role": "user", "content": "q1"}),
638 json!({"role": "assistant", "content": "a1"}),
639 json!({"role": "user", "content": "q2"}),
640 ];
641 let after = vec![
642 json!({"role": "system", "content": "MODIFIED"}),
643 json!({"role": "user", "content": "q1"}),
644 json!({"role": "assistant", "content": "a1"}),
645 json!({"role": "user", "content": "q2"}),
646 ];
647 assert!(verify_determinism(&before, &after).is_stable);
648 }
649
650 #[test]
651 fn explicit_cache_control_still_catches_violations() {
652 let before = vec![
655 json!({"role": "system", "content": [{"type": "text", "text": "rules", "cache_control": {}}]}),
656 json!({"role": "user", "content": "q1"}),
657 ];
658 let after = vec![
659 json!({"role": "system", "content": [{"type": "text", "text": "MODIFIED", "cache_control": {}}]}),
660 json!({"role": "user", "content": "q1"}),
661 ];
662 assert!(!verify_determinism(&before, &after).is_stable);
663 }
664
665 #[test]
668 fn changing_a_cache_lifetime_is_not_a_violation() {
669 let marked = |ttl: &str| {
670 vec![
671 json!({"role": "user", "content": [{"type": "text", "text": "q1"}]}),
672 json!({"role": "user", "content": [{"type": "tool_result", "tool_use_id": "t",
673 "content": "ok", "cache_control": {"type": "ephemeral", "ttl": ttl}}]}),
674 ]
675 };
676 assert!(verify_determinism(&marked("1h"), &marked("5m")).is_stable);
677 }
678
679 #[test]
680 fn cached_system_message_modification_is_detected() {
681 let before = vec![
682 json!({"role": "system", "content": [{"type": "text", "text": "rules", "cache_control": {}}]}),
683 ];
684 let after = vec![
685 json!({"role": "system", "content": [{"type": "text", "text": "changed", "cache_control": {}}]}),
686 ];
687
688 assert!(!verify_determinism(&before, &after).is_stable);
689 }
690
691 #[test]
692 fn frozen_byte_count_matches_the_common_canonical_prefix() {
693 let before = vec![json!({"role": "user", "content": "same", "cache_control": {}})];
694 let after = vec![json!({"role": "user", "content": "different", "cache_control": {}})];
695 let before_bytes = canonical_messages(&before);
696 let after_bytes = canonical_messages(&after);
697 let expected = before_bytes
698 .iter()
699 .zip(&after_bytes)
700 .take_while(|(left, right)| left == right)
701 .count();
702
703 assert_eq!(verify_determinism(&before, &after).frozen_bytes, expected);
704 }
705
706 #[test]
707 fn proof_headers_are_set() {
708 let proof = verify_determinism(&[], &[]);
709 let mut response = Response::new(());
710
711 apply_response_headers(&mut response, &proof);
712
713 let headers = response.headers();
714 assert_eq!(headers["x-leanctx-determinism"], "stable");
715 assert_eq!(headers["x-leanctx-frozen-bytes"], "2");
716 assert_eq!(headers["x-leanctx-cache-safe"], "true");
717 assert_eq!(headers["x-leanctx-proof-hash"], proof.prefix_hash_after);
718 }
719}