1use serde::{Deserialize, Serialize};
8
9#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
16#[serde(tag = "type")]
17pub enum EntryKind {
18 #[default]
20 Text,
21 UserMessage,
23 AssistantTurn {
25 tool_calls: Vec<SerializedToolCall>,
29 },
30 ToolResult {
32 tool_call_id: String,
35 tool_name: String,
37 is_error: bool,
40 },
41}
42
43#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
45pub struct SerializedToolCall {
46 pub id: String,
49 pub name: String,
51 pub arguments: serde_json::Value,
53 #[serde(default, skip_serializing_if = "Option::is_none")]
56 pub thought_signature: Option<String>,
57}
58
59#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
66#[serde(tag = "strategy", rename_all = "snake_case")]
67pub enum EvictionStrategy {
68 #[default]
70 PerItem,
71 Bulk {
74 overflow: usize,
77 },
78 Compact {
82 compact_count: usize,
84 },
85}
86
87#[derive(Debug, Clone, Serialize, Deserialize)]
94pub enum RegionKind {
95 Pinned,
102
103 SlidingWindow {
110 max_items: usize,
112 eviction_strategy: EvictionStrategy,
114 },
115
116 Temporary,
122
123 Compacting {
129 threshold_tokens: usize,
131 },
132
133 Clearable,
140
141 CompactHistory {
148 source_region: String,
150 },
151
152 HashMap {
156 max_entries: Option<usize>,
158 },
159
160 Checklist,
170
171 Custom {
187 script: String,
189 persistent: bool,
192 },
193}
194
195impl PartialEq for RegionKind {
196 #[inline(never)]
197 fn eq(&self, other: &Self) -> bool {
198 match (self, other) {
199 (Self::Pinned, Self::Pinned)
200 | (Self::Temporary, Self::Temporary)
201 | (Self::Clearable, Self::Clearable) => true,
202 (
203 Self::SlidingWindow {
204 max_items: a,
205 eviction_strategy: sa,
206 },
207 Self::SlidingWindow {
208 max_items: b,
209 eviction_strategy: sb,
210 },
211 ) => a == b && sa == sb,
212 (
213 Self::Compacting {
214 threshold_tokens: a,
215 },
216 Self::Compacting {
217 threshold_tokens: b,
218 },
219 ) => a == b,
220 (
221 Self::CompactHistory { source_region: a },
222 Self::CompactHistory { source_region: b },
223 ) => a == b,
224 (Self::HashMap { max_entries: a }, Self::HashMap { max_entries: b }) => a == b,
225 (Self::Checklist, Self::Checklist) => true,
226 (
227 Self::Custom {
228 script: a,
229 persistent: pa,
230 },
231 Self::Custom {
232 script: b,
233 persistent: pb,
234 },
235 ) => a == b && pa == pb,
236 _ => false,
237 }
238 }
239}
240impl Eq for RegionKind {}
241
242#[derive(Debug, Clone, PartialEq, Eq)]
249pub struct ChecklistItem {
250 pub id: usize,
252 pub text: String,
254 pub done: bool,
256 pub note: Option<String>,
258}
259
260const ITEM_ID: &str = "checklist_id";
262const ITEM_DONE: &str = "checklist_done";
264const ITEM_NOTE: &str = "checklist_note";
266
267impl RegionEntry {
268 pub fn as_checklist_item(&self) -> Option<ChecklistItem> {
270 let meta = self.metadata.as_ref()?;
271 Some(ChecklistItem {
272 id: meta.get(ITEM_ID)?.as_u64()? as usize,
273 text: self.content.clone(),
274 done: meta
275 .get(ITEM_DONE)
276 .and_then(|v| v.as_bool())
277 .unwrap_or(false),
278 note: meta
279 .get(ITEM_NOTE)
280 .and_then(|v| v.as_str())
281 .map(str::to_string),
282 })
283 }
284}
285
286impl Region {
287 pub fn checklist_items(&self) -> Vec<ChecklistItem> {
289 self.content
290 .iter()
291 .filter_map(RegionEntry::as_checklist_item)
292 .collect()
293 }
294
295 pub fn open_checklist_items(&self) -> Vec<ChecklistItem> {
297 self.checklist_items()
298 .into_iter()
299 .filter(|i| !i.done)
300 .collect()
301 }
302
303 pub fn add_checklist_item(
310 &mut self,
311 text: String,
312 tokens: usize,
313 ) -> crate::error::Result<usize> {
314 let id = self
315 .checklist_items()
316 .iter()
317 .map(|i| i.id)
318 .max()
319 .unwrap_or(0)
320 + 1;
321 self.add_entry_with_metadata(
322 text,
323 tokens,
324 serde_json::json!({ ITEM_ID: id, ITEM_DONE: false }),
325 )?;
326 Ok(id)
327 }
328
329 pub fn complete_checklist_item(&mut self, id: usize) -> bool {
331 self.set_item_field(id, ITEM_DONE, serde_json::Value::Bool(true))
332 }
333
334 pub fn note_checklist_item(&mut self, id: usize, note: &str) -> bool {
336 self.set_item_field(id, ITEM_NOTE, serde_json::Value::String(note.to_string()))
337 }
338
339 fn set_item_field(&mut self, id: usize, key: &str, value: serde_json::Value) -> bool {
341 for entry in &mut self.content {
342 let is_target = entry
343 .metadata
344 .as_ref()
345 .and_then(|m| m.get(ITEM_ID))
346 .and_then(serde_json::Value::as_u64)
347 .is_some_and(|found| found as usize == id);
348 if is_target && let Some(serde_json::Value::Object(meta)) = entry.metadata.as_mut() {
349 meta.insert(key.to_string(), value);
350 return true;
351 }
352 }
353 false
354 }
355
356 pub fn render_checklist(&self) -> String {
362 let items = self.checklist_items();
363 if items.is_empty() {
364 return String::new();
365 }
366 let (open, done): (Vec<_>, Vec<_>) = items.into_iter().partition(|i| !i.done);
367 let mut out = String::new();
368 for item in open.iter().chain(done.iter()) {
369 let box_ = match item.done {
370 true => "[x]",
371 false => "[ ]",
372 };
373 out.push_str(&format!("{box_} {} {}", item.id, item.text));
374 if let Some(note) = &item.note {
375 out.push_str(&format!("\n note: {note}"));
376 }
377 out.push('\n');
378 }
379 format!(
380 "Checklist ({} open, {} done):\n{}",
381 open.len(),
382 done.len(),
383 out.trim_end()
384 )
385 }
386}
387
388impl RegionKind {
389 pub fn cache_hint(&self) -> crate::cache::CacheHint {
391 match self {
392 RegionKind::Pinned | RegionKind::CompactHistory { .. } => {
393 crate::cache::CacheHint::Always
394 }
395 RegionKind::Compacting { .. } => crate::cache::CacheHint::UntilChanged,
396 RegionKind::SlidingWindow { .. } => crate::cache::CacheHint::SlidingPrefix {
397 stable_fraction: 0.75,
398 },
399 RegionKind::HashMap { .. } => crate::cache::CacheHint::UntilChanged,
400 RegionKind::Checklist => crate::cache::CacheHint::UntilChanged,
403 RegionKind::Temporary | RegionKind::Clearable => crate::cache::CacheHint::Never,
404 RegionKind::Custom { persistent, .. } => {
408 if *persistent {
409 crate::cache::CacheHint::Always
410 } else {
411 crate::cache::CacheHint::UntilChanged
412 }
413 }
414 }
415 }
416}
417
418#[derive(Debug, Clone, Serialize, Deserialize)]
423pub struct Region {
424 pub name: String,
426
427 pub kind: RegionKind,
429
430 pub content: Vec<RegionEntry>,
432
433 pub max_tokens: usize,
435
436 pub current_tokens: usize,
438
439 pub schema: Option<RegionSchema>,
441
442 #[serde(default, skip_serializing_if = "Option::is_none")]
444 pub taint: Option<crate::taint::RegionTaint>,
445
446 #[serde(default)]
450 pub needs_message_compaction: bool,
451
452 #[serde(default = "crate::region::default_true")]
458 pub summarizable: bool,
459}
460
461pub(crate) fn default_true() -> bool {
463 true
464}
465
466impl Region {
467 pub fn new(name: String, kind: RegionKind, max_tokens: usize) -> Self {
469 Self {
470 name,
471 kind,
472 content: Vec::new(),
473 max_tokens,
474 current_tokens: 0,
475 schema: None,
476 taint: None,
477 needs_message_compaction: false,
478 summarizable: true,
479 }
480 }
481
482 pub fn with_taint_tracking(mut self) -> Self {
484 self.taint = Some(crate::taint::RegionTaint::new());
485 self
486 }
487
488 pub fn enable_taint_tracking(&mut self) {
490 if self.taint.is_none() {
491 self.taint = Some(crate::taint::RegionTaint::new());
492 }
493 }
494
495 pub fn taint_level(&self) -> Option<crate::taint::TaintLevel> {
497 self.taint.as_ref().map(|t| t.level())
498 }
499
500 fn push_entry(
513 &mut self,
514 content: String,
515 tokens: usize,
516 metadata: Option<serde_json::Value>,
517 kind: EntryKind,
518 taint_level: crate::taint::TaintLevel,
519 ) -> crate::error::Result<()> {
520 if let Some(schema) = &self.schema {
521 schema.validate(&content)?;
522 }
523
524 if self.current_tokens + tokens > self.max_tokens {
525 return Err(crate::error::Error::TokenBudgetExceeded {
526 used: self.current_tokens + tokens,
527 max: self.max_tokens,
528 });
529 }
530
531 self.content.push(RegionEntry {
532 content,
533 tokens,
534 timestamp: chrono::Utc::now().timestamp(),
535 metadata,
536 kind,
537 key: None,
538 });
539 self.current_tokens += tokens;
540
541 if let Some(taint) = &mut self.taint {
545 taint.add_entry(taint_level);
546 }
547
548 self.enforce_sliding_window();
549
550 Ok(())
551 }
552
553 pub fn add_tainted_entry(
555 &mut self,
556 content: String,
557 tokens: usize,
558 taint_level: crate::taint::TaintLevel,
559 ) -> crate::error::Result<()> {
560 self.push_entry(content, tokens, None, EntryKind::default(), taint_level)
561 }
562
563 pub fn add_typed_tainted_entry(
572 &mut self,
573 content: String,
574 tokens: usize,
575 kind: EntryKind,
576 taint_level: crate::taint::TaintLevel,
577 ) -> crate::error::Result<()> {
578 self.push_entry(content, tokens, None, kind, taint_level)
579 }
580
581 pub fn with_schema(mut self, schema: RegionSchema) -> Self {
583 self.schema = Some(schema);
584 self
585 }
586
587 pub fn add_entry(&mut self, content: String, tokens: usize) -> crate::error::Result<()> {
592 self.push_entry(
593 content,
594 tokens,
595 None,
596 EntryKind::default(),
597 crate::taint::TaintLevel::Public,
598 )
599 }
600
601 pub fn add_entry_with_metadata(
603 &mut self,
604 content: String,
605 tokens: usize,
606 metadata: serde_json::Value,
607 ) -> crate::error::Result<()> {
608 self.push_entry(
609 content,
610 tokens,
611 Some(metadata),
612 EntryKind::default(),
613 crate::taint::TaintLevel::Public,
614 )
615 }
616
617 pub fn add_typed_entry(
623 &mut self,
624 content: String,
625 tokens: usize,
626 kind: EntryKind,
627 ) -> crate::error::Result<()> {
628 self.push_entry(
629 content,
630 tokens,
631 None,
632 kind,
633 crate::taint::TaintLevel::Public,
634 )
635 }
636
637 pub fn carry_entry(&mut self, entry: RegionEntry) -> crate::error::Result<()> {
651 if self.current_tokens + entry.tokens > self.max_tokens {
653 return Err(crate::error::Error::TokenBudgetExceeded {
654 used: self.current_tokens + entry.tokens,
655 max: self.max_tokens,
656 });
657 }
658
659 self.current_tokens += entry.tokens;
660 self.content.push(entry);
661
662 self.enforce_sliding_window();
664
665 Ok(())
666 }
667
668 pub fn upsert_by_key(
671 &mut self,
672 key: &str,
673 content: String,
674 tokens: usize,
675 ) -> Result<(), String> {
676 if let Some(pos) = self
678 .content
679 .iter()
680 .position(|e| e.key.as_deref() == Some(key))
681 {
682 let old_tokens = self.content[pos].tokens;
683 self.current_tokens -= old_tokens;
684 self.content[pos].content = content;
685 self.content[pos].tokens = tokens;
686 self.content[pos].timestamp = chrono::Utc::now().timestamp();
687 self.current_tokens += tokens;
688 return Ok(());
689 }
690
691 let max_entries = if let RegionKind::HashMap {
693 max_entries: Some(max),
694 } = &self.kind
695 {
696 Some(*max)
697 } else {
698 None
699 };
700 if let Some(max) = max_entries {
701 while self.content.len() >= max {
702 self.evict_lru_entry();
703 }
704 }
705
706 while self.current_tokens + tokens > self.max_tokens && !self.content.is_empty() {
708 self.evict_lru_entry();
709 }
710
711 if self.current_tokens + tokens > self.max_tokens {
712 return Err(format!(
713 "Entry ({} tokens) exceeds region budget ({} max)",
714 tokens, self.max_tokens
715 ));
716 }
717
718 self.content.push(RegionEntry {
719 content,
720 tokens,
721 timestamp: chrono::Utc::now().timestamp(),
722 metadata: None,
723 kind: EntryKind::default(),
724 key: Some(key.to_string()),
725 });
726 self.current_tokens += tokens;
727 Ok(())
728 }
729
730 pub fn get_by_key(&self, key: &str) -> Option<&RegionEntry> {
732 self.content.iter().find(|e| e.key.as_deref() == Some(key))
733 }
734
735 pub fn remove_by_key(&mut self, key: &str) -> bool {
737 if let Some(pos) = self
738 .content
739 .iter()
740 .position(|e| e.key.as_deref() == Some(key))
741 {
742 let tokens = self.content[pos].tokens;
743 self.content.remove(pos);
744 self.current_tokens -= tokens;
745 if let Some(taint) = &mut self.taint {
746 taint.remove_at(pos);
747 }
748 true
749 } else {
750 false
751 }
752 }
753
754 pub fn keys(&self) -> Vec<&str> {
756 self.content
757 .iter()
758 .filter_map(|e| e.key.as_deref())
759 .collect()
760 }
761
762 fn evict_lru_entry(&mut self) {
764 if self.content.is_empty() {
765 return;
766 }
767 let oldest_idx = self
768 .content
769 .iter()
770 .enumerate()
771 .min_by_key(|(_, e)| e.timestamp)
772 .map(|(i, _)| i)
773 .unwrap_or(0);
774 let tokens = self.content[oldest_idx].tokens;
775 self.content.remove(oldest_idx);
776 self.current_tokens -= tokens;
777 if let Some(taint) = &mut self.taint {
778 taint.remove_at(oldest_idx);
779 }
780 }
781
782 fn enforce_sliding_window(&mut self) {
793 if let RegionKind::SlidingWindow {
794 max_items,
795 eviction_strategy,
796 } = &self.kind
797 {
798 let max = *max_items;
799 match eviction_strategy.clone() {
800 EvictionStrategy::PerItem => {
801 while self.content.len() > max && self.remove_oldest().is_some() {}
805 }
806 EvictionStrategy::Bulk { overflow } => {
807 if self.content.len() > max + overflow {
808 while self.content.len() > max && self.remove_oldest().is_some() {}
809 }
810 }
811 EvictionStrategy::Compact { compact_count } => {
812 if self.content.len() > max + compact_count * 2 {
813 while self.content.len() > max && self.remove_oldest().is_some() {}
816 self.needs_message_compaction = false;
817 } else if self.content.len() > max + compact_count {
818 self.needs_message_compaction = true;
819 }
820 }
821 }
822 }
823 }
824
825 fn turn_group_size_at(&self, idx: usize) -> usize {
833 if idx >= self.content.len() {
834 return 0;
835 }
836 match &self.content[idx].kind {
837 EntryKind::AssistantTurn { .. } => {
838 let mut size = 1;
839 while idx + size < self.content.len() {
840 if matches!(self.content[idx + size].kind, EntryKind::ToolResult { .. }) {
841 size += 1;
842 } else {
843 break;
844 }
845 }
846 size
847 }
848 _ => 1,
849 }
850 }
851
852 pub fn clear(&mut self) {
854 self.content.clear();
855 self.current_tokens = 0;
856 if let Some(taint) = &mut self.taint {
857 taint.clear();
858 }
859 }
860
861 pub fn remove_oldest(&mut self) -> Option<RegionEntry> {
863 if self.content.is_empty() {
864 return None;
865 }
866 let group_size = self.turn_group_size_at(0);
870 let mut first = None;
871 let mut extra_tokens = 0usize;
872 let mut i = 0;
875 while i < group_size && !self.content.is_empty() {
876 let entry_tokens = self.content[0].tokens;
877 self.current_tokens -= entry_tokens;
878 let removed = self.content.remove(0);
879 if let Some(taint) = &mut self.taint {
880 taint.remove_oldest();
881 }
882 if i == 0 {
883 first = Some(removed);
884 } else {
885 extra_tokens += entry_tokens;
886 }
887 i += 1;
888 }
889 first.map(|mut entry| {
895 entry.tokens += extra_tokens;
896 entry
897 })
898 }
899
900 pub fn remove_entries_by_prefix(&mut self, prefix: &str) {
906 let mut i = 0;
907 while i < self.content.len() {
908 if self.content[i].content.starts_with(prefix) {
909 let tokens = self.content[i].tokens;
910 self.content.remove(i);
911 self.current_tokens -= tokens;
912 if let Some(taint) = &mut self.taint {
913 taint.remove_at(i);
914 }
915 } else {
916 i += 1;
917 }
918 }
919 }
920
921 pub fn entry_count(&self) -> usize {
923 self.content.len()
924 }
925
926 pub fn needs_compaction(&self) -> bool {
928 if let RegionKind::Compacting { threshold_tokens } = self.kind {
929 self.current_tokens > threshold_tokens
930 } else {
931 false
932 }
933 }
934}
935
936#[derive(Debug, Clone, Serialize, Deserialize)]
940pub struct RegionEntry {
941 pub content: String,
943
944 pub tokens: usize,
946
947 pub timestamp: i64,
949
950 pub metadata: Option<serde_json::Value>,
952
953 #[serde(default)]
957 pub kind: EntryKind,
958
959 #[serde(default, skip_serializing_if = "Option::is_none")]
961 pub key: Option<String>,
962}
963
964#[derive(Debug, Serialize, Deserialize)]
970pub struct RegionSchema {
971 pub format: ContentFormat,
973
974 #[serde(skip_serializing_if = "Option::is_none")]
976 pub custom_script: Option<String>,
977}
978
979impl Clone for RegionSchema {
980 fn clone(&self) -> Self {
981 Self {
982 format: self.format.clone(),
983 custom_script: self.custom_script.clone(),
984 }
985 }
986}
987
988impl RegionSchema {
989 pub fn new(format: ContentFormat) -> Self {
991 Self {
992 format,
993 custom_script: None,
994 }
995 }
996
997 pub fn with_custom_script(mut self, script: String) -> Self {
999 self.custom_script = Some(script);
1000 self
1001 }
1002
1003 pub fn validate(&self, content: &str) -> crate::error::Result<()> {
1005 match &self.format {
1006 ContentFormat::Json => {
1007 serde_json::from_str::<serde_json::Value>(content).map_err(|e| {
1008 crate::error::Error::ValidationFailed(format!("Invalid JSON: {}", e))
1009 })?;
1010 }
1011 ContentFormat::Mermaid => {
1012 if !content.contains("graph")
1014 && !content.contains("sequenceDiagram")
1015 && !content.contains("classDiagram")
1016 && !content.contains("stateDiagram")
1017 && !content.contains("erDiagram")
1018 && !content.contains("journey")
1019 && !content.contains("gantt")
1020 && !content.contains("pie")
1021 && !content.contains("flowchart")
1022 {
1023 return Err(crate::error::Error::ValidationFailed(
1024 "Mermaid diagrams must contain a valid diagram type (graph, sequenceDiagram, etc.)".to_string()
1025 ));
1026 }
1027 }
1028 ContentFormat::Code { .. } => {
1029 if content.trim().is_empty() {
1031 return Err(crate::error::Error::ValidationFailed(
1032 "Code cannot be empty".to_string(),
1033 ));
1034 }
1035 }
1036 ContentFormat::Markdown => {
1037 if content.trim().is_empty() {
1039 return Err(crate::error::Error::ValidationFailed(
1040 "Markdown content cannot be empty".to_string(),
1041 ));
1042 }
1043 }
1044 ContentFormat::Text | ContentFormat::Custom { .. } => {
1045 }
1047 }
1048
1049 Ok(())
1050 }
1051}
1052
1053#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
1055pub enum ContentFormat {
1056 Text,
1058
1059 Json,
1061
1062 Mermaid,
1064
1065 Code {
1067 language: String,
1070 },
1071
1072 Markdown,
1074
1075 Custom {
1077 format_name: String,
1080 },
1081}
1082
1083pub trait Validator: Send + Sync {
1089 fn validate(&self, content: &str) -> std::result::Result<(), crate::error::ValidationError>;
1091
1092 fn description(&self) -> &str;
1094}
1095
1096#[cfg(test)]
1097mod tests {
1098 use super::*;
1099
1100 fn checklist() -> Region {
1103 Region::new("todos".to_string(), RegionKind::Checklist, 10_000)
1104 }
1105
1106 #[test]
1112 fn a_malformed_entry_is_not_an_item() {
1113 let mut r = checklist();
1114 r.add_entry("a plain note".to_string(), 3).unwrap();
1116 r.add_entry_with_metadata(
1118 "something else".to_string(),
1119 3,
1120 serde_json::json!({ "unrelated": true }),
1121 )
1122 .unwrap();
1123 r.add_entry_with_metadata(
1125 "bad id".to_string(),
1126 3,
1127 serde_json::json!({ "checklist_id": "one" }),
1128 )
1129 .unwrap();
1130
1131 assert!(r.checklist_items().is_empty(), "none of those are items");
1132 assert!(r.open_checklist_items().is_empty());
1133 assert!(
1134 r.render_checklist().is_empty(),
1135 "and they do not render as a checklist"
1136 );
1137 }
1138
1139 #[test]
1142 fn a_checklist_caches_until_it_changes() {
1143 assert_eq!(
1144 RegionKind::Checklist.cache_hint(),
1145 crate::cache::CacheHint::UntilChanged
1146 );
1147 }
1148
1149 #[test]
1150 fn a_note_appears_in_the_render() {
1151 let mut r = checklist();
1152 let id = r.add_checklist_item("blocked".to_string(), 2).unwrap();
1153 r.note_checklist_item(id, "waiting on the manual");
1154 let rendered = r.render_checklist();
1155 assert!(
1156 rendered.contains("note: waiting on the manual"),
1157 "{rendered}"
1158 );
1159 }
1160
1161 #[test]
1164 fn an_item_over_budget_is_refused() {
1165 let mut r = Region::new("todos".to_string(), RegionKind::Checklist, 4);
1166 assert!(r.add_checklist_item("x".to_string(), 99).is_err());
1167 assert!(r.checklist_items().is_empty());
1168 }
1169
1170 #[test]
1171 fn an_added_item_starts_open_and_gets_an_id() {
1172 let mut r = checklist();
1173 let first = r
1174 .add_checklist_item("compute the fee table".to_string(), 5)
1175 .unwrap();
1176 let second = r
1177 .add_checklist_item("check the manual".to_string(), 5)
1178 .unwrap();
1179 assert_eq!((first, second), (1, 2), "ids are stable and sequential");
1180 assert_eq!(r.open_checklist_items().len(), 2);
1181 }
1182
1183 #[test]
1184 fn completing_an_item_closes_it_and_nothing_else() {
1185 let mut r = checklist();
1186 let id = r.add_checklist_item("one".to_string(), 2).unwrap();
1187 r.add_checklist_item("two".to_string(), 2).unwrap();
1188
1189 assert!(r.complete_checklist_item(id));
1190 let open = r.open_checklist_items();
1191 assert_eq!(open.len(), 1);
1192 assert_eq!(open[0].text, "two");
1193 assert_eq!(
1194 r.checklist_items().len(),
1195 2,
1196 "done items are kept, not deleted"
1197 );
1198 }
1199
1200 #[test]
1201 fn an_unknown_id_reports_failure_rather_than_ticking_something_else() {
1202 let mut r = checklist();
1205 r.add_checklist_item("one".to_string(), 2).unwrap();
1206 assert!(!r.complete_checklist_item(99));
1207 assert!(!r.note_checklist_item(99, "x"));
1208 assert_eq!(r.open_checklist_items().len(), 1);
1209 }
1210
1211 #[test]
1212 fn a_note_records_without_closing() {
1213 let mut r = checklist();
1214 let id = r
1215 .add_checklist_item("blocked thing".to_string(), 2)
1216 .unwrap();
1217 assert!(r.note_checklist_item(id, "waiting on the manual"));
1218 let item = &r.checklist_items()[0];
1219 assert!(!item.done, "a note is not a completion");
1220 assert_eq!(item.note.as_deref(), Some("waiting on the manual"));
1221 }
1222
1223 #[test]
1226 fn the_render_puts_open_items_first() {
1227 let mut r = checklist();
1228 let done = r
1229 .add_checklist_item("already finished".to_string(), 2)
1230 .unwrap();
1231 r.add_checklist_item("still to do".to_string(), 2).unwrap();
1232 r.complete_checklist_item(done);
1233
1234 let rendered = r.render_checklist();
1235 let open_at = rendered.find("still to do").expect("open item rendered");
1236 let done_at = rendered
1237 .find("already finished")
1238 .expect("done item rendered");
1239 assert!(open_at < done_at, "open before done:\n{rendered}");
1240 assert!(rendered.contains("1 open, 1 done"), "{rendered}");
1241 assert!(
1242 rendered.contains("[x]") && rendered.contains("[ ]"),
1243 "{rendered}"
1244 );
1245 }
1246
1247 #[test]
1248 fn an_empty_checklist_renders_nothing() {
1249 assert!(checklist().render_checklist().is_empty());
1251 }
1252
1253 #[test]
1256 fn ids_do_not_get_reused_after_a_drop() {
1257 let mut r = checklist();
1258 r.add_checklist_item("one".to_string(), 2).unwrap();
1259 let second = r.add_checklist_item("two".to_string(), 2).unwrap();
1260 r.content.remove(0);
1261 let third = r.add_checklist_item("three".to_string(), 2).unwrap();
1262 assert!(third > second, "a reused id would tick off the wrong item");
1263 }
1264
1265 #[test]
1266 fn test_region_creation() {
1267 let region = Region::new("test".to_string(), RegionKind::Pinned, 1000);
1268 assert_eq!(region.name, "test");
1269 assert_eq!(region.max_tokens, 1000);
1270 assert_eq!(region.current_tokens, 0);
1271 }
1272
1273 #[test]
1274 fn test_sliding_window_config() {
1275 let kind = RegionKind::SlidingWindow {
1276 max_items: 10,
1277 eviction_strategy: EvictionStrategy::PerItem,
1278 };
1279 let region = Region::new("history".to_string(), kind.clone(), 5000);
1280 assert_eq!(region.kind, kind);
1281 }
1282
1283 #[test]
1284 fn test_region_kind_equality() {
1285 assert_eq!(RegionKind::Clearable, RegionKind::Clearable);
1286 assert_eq!(
1287 RegionKind::Compacting {
1288 threshold_tokens: 500
1289 },
1290 RegionKind::Compacting {
1291 threshold_tokens: 500
1292 }
1293 );
1294 assert_eq!(
1295 RegionKind::CompactHistory {
1296 source_region: "conv".to_string()
1297 },
1298 RegionKind::CompactHistory {
1299 source_region: "conv".to_string()
1300 }
1301 );
1302 assert_ne!(RegionKind::Pinned, RegionKind::Temporary);
1303 }
1304
1305 #[test]
1306 fn custom_kind_equality_compares_script_and_persistent() {
1307 let a = RegionKind::Custom {
1308 script: "conv.rhai".to_string(),
1309 persistent: false,
1310 };
1311 assert_eq!(a, a.clone());
1312 assert_ne!(
1313 a,
1314 RegionKind::Custom {
1315 script: "other.rhai".to_string(),
1316 persistent: false,
1317 }
1318 );
1319 assert_ne!(
1320 a,
1321 RegionKind::Custom {
1322 script: "conv.rhai".to_string(),
1323 persistent: true,
1324 }
1325 );
1326 assert_ne!(a, RegionKind::Temporary);
1327 }
1328
1329 #[test]
1330 fn custom_kind_serde_round_trips() {
1331 let kind = RegionKind::Custom {
1332 script: "hooks/conv.rhai".to_string(),
1333 persistent: true,
1334 };
1335 let json = serde_json::to_string(&kind).unwrap();
1336 let back: RegionKind = serde_json::from_str(&json).unwrap();
1337 assert_eq!(kind, back);
1338 let old: RegionKind = serde_json::from_str("\"Pinned\"").unwrap();
1340 assert_eq!(old, RegionKind::Pinned);
1341 }
1342
1343 #[test]
1344 fn custom_kind_cache_hint_follows_persistent() {
1345 assert_eq!(
1346 RegionKind::Custom {
1347 script: "s.rhai".to_string(),
1348 persistent: true,
1349 }
1350 .cache_hint(),
1351 crate::cache::CacheHint::Always
1352 );
1353 assert_eq!(
1354 RegionKind::Custom {
1355 script: "s.rhai".to_string(),
1356 persistent: false,
1357 }
1358 .cache_hint(),
1359 crate::cache::CacheHint::UntilChanged
1360 );
1361 }
1362
1363 #[test]
1364 fn carry_entry_preserves_kind_metadata_key_and_timestamp() {
1365 let mut source = Region::new("conversation".to_string(), RegionKind::Temporary, 10_000);
1366 source
1367 .add_typed_entry(
1368 "result body".to_string(),
1369 10,
1370 EntryKind::ToolResult {
1371 tool_call_id: "call_1".to_string(),
1372 tool_name: "read_file".to_string(),
1373 is_error: false,
1374 },
1375 )
1376 .unwrap();
1377 let mut entry = source.content[0].clone();
1378 entry.metadata = Some(serde_json::json!({"origin": "test"}));
1379 entry.key = Some("k".to_string());
1380 let stamped = entry.timestamp;
1381
1382 let mut dest = Region::new("conversation".to_string(), RegionKind::Temporary, 10_000);
1383 dest.carry_entry(entry).unwrap();
1384
1385 let carried = &dest.content[0];
1386 assert!(matches!(
1387 &carried.kind,
1388 EntryKind::ToolResult { tool_call_id, .. } if tool_call_id == "call_1"
1389 ));
1390 assert_eq!(
1391 carried.metadata,
1392 Some(serde_json::json!({"origin": "test"}))
1393 );
1394 assert_eq!(carried.key.as_deref(), Some("k"));
1395 assert_eq!(carried.timestamp, stamped);
1396 assert_eq!(dest.current_tokens, 10);
1397 }
1398
1399 #[test]
1400 fn carry_entry_rejects_over_budget() {
1401 let mut dest = Region::new("small".to_string(), RegionKind::Temporary, 5);
1402 let mut source = Region::new("src".to_string(), RegionKind::Temporary, 100);
1403 source.add_entry("filler".to_string(), 10).unwrap();
1404 let err = dest.carry_entry(source.content[0].clone()).unwrap_err();
1405 assert_eq!(err.to_string(), "Content exceeds token budget: 10 > 5");
1406 assert!(dest.content.is_empty());
1407 assert_eq!(dest.current_tokens, 0);
1408 }
1409
1410 #[test]
1411 fn carry_entry_enforces_sliding_window_max_items() {
1412 let mut source = Region::new("src".to_string(), RegionKind::Temporary, 10_000);
1413 for i in 0..4 {
1414 source.add_entry(format!("msg{i}"), 10).unwrap();
1415 }
1416 let mut dest = Region::new(
1417 "conv".to_string(),
1418 RegionKind::SlidingWindow {
1419 max_items: 3,
1420 eviction_strategy: EvictionStrategy::PerItem,
1421 },
1422 10_000,
1423 );
1424 for entry in &source.content {
1425 dest.carry_entry(entry.clone()).unwrap();
1426 }
1427 assert_eq!(dest.content.len(), 3);
1428 assert_eq!(dest.content[0].content, "msg1");
1429 }
1430
1431 #[test]
1432 fn test_sliding_window_enforces_max_items() {
1433 let mut region = Region::new(
1434 "conv".to_string(),
1435 RegionKind::SlidingWindow {
1436 max_items: 3,
1437 eviction_strategy: EvictionStrategy::PerItem,
1438 },
1439 50000,
1440 );
1441
1442 region.add_entry("msg1".to_string(), 10).unwrap();
1443 region.add_entry("msg2".to_string(), 20).unwrap();
1444 region.add_entry("msg3".to_string(), 30).unwrap();
1445 assert_eq!(region.entry_count(), 3);
1446 assert_eq!(region.current_tokens, 60);
1447
1448 region.add_entry("msg4".to_string(), 40).unwrap();
1450 assert_eq!(region.entry_count(), 3);
1451 assert_eq!(region.content[0].content, "msg2");
1452 assert_eq!(region.content[2].content, "msg4");
1453 assert_eq!(region.current_tokens, 90); region.add_entry("msg5".to_string(), 50).unwrap();
1457 assert_eq!(region.entry_count(), 3);
1458 assert_eq!(region.content[0].content, "msg3");
1459 assert_eq!(region.current_tokens, 120); }
1461
1462 #[test]
1463 fn test_sliding_window_enforces_max_items_with_metadata() {
1464 let mut region = Region::new(
1465 "conv".to_string(),
1466 RegionKind::SlidingWindow {
1467 max_items: 2,
1468 eviction_strategy: EvictionStrategy::PerItem,
1469 },
1470 50000,
1471 );
1472
1473 region
1474 .add_entry_with_metadata("a".to_string(), 10, serde_json::json!({"idx": 1}))
1475 .unwrap();
1476 region
1477 .add_entry_with_metadata("b".to_string(), 20, serde_json::json!({"idx": 2}))
1478 .unwrap();
1479 region
1480 .add_entry_with_metadata("c".to_string(), 30, serde_json::json!({"idx": 3}))
1481 .unwrap();
1482
1483 assert_eq!(region.entry_count(), 2);
1484 assert_eq!(region.content[0].content, "b");
1485 assert_eq!(region.content[1].content, "c");
1486 assert_eq!(region.current_tokens, 50);
1487 }
1488
1489 #[test]
1490 fn test_cache_hint_pinned() {
1491 let kind = RegionKind::Pinned;
1492 assert_eq!(kind.cache_hint(), crate::cache::CacheHint::Always);
1493 }
1494
1495 #[test]
1496 fn test_cache_hint_compact_history() {
1497 let kind = RegionKind::CompactHistory {
1498 source_region: "conv".to_string(),
1499 };
1500 assert_eq!(kind.cache_hint(), crate::cache::CacheHint::Always);
1501 }
1502
1503 #[test]
1504 fn test_cache_hint_compacting() {
1505 let kind = RegionKind::Compacting {
1506 threshold_tokens: 1000,
1507 };
1508 assert_eq!(kind.cache_hint(), crate::cache::CacheHint::UntilChanged);
1509 }
1510
1511 #[test]
1512 fn test_cache_hint_sliding_window() {
1513 let kind = RegionKind::SlidingWindow {
1514 max_items: 10,
1515 eviction_strategy: EvictionStrategy::PerItem,
1516 };
1517 assert_eq!(
1518 kind.cache_hint(),
1519 crate::cache::CacheHint::SlidingPrefix {
1520 stable_fraction: 0.75
1521 }
1522 );
1523 }
1524
1525 #[test]
1526 fn test_cache_hint_temporary() {
1527 assert_eq!(
1528 RegionKind::Temporary.cache_hint(),
1529 crate::cache::CacheHint::Never
1530 );
1531 }
1532
1533 #[test]
1534 fn test_cache_hint_clearable() {
1535 assert_eq!(
1536 RegionKind::Clearable.cache_hint(),
1537 crate::cache::CacheHint::Never
1538 );
1539 }
1540
1541 #[test]
1544 fn test_with_schema_attaches_schema() {
1545 let schema = RegionSchema::new(ContentFormat::Json);
1546 let region =
1547 Region::new("data".to_string(), RegionKind::Temporary, 1000).with_schema(schema);
1548 assert!(region.schema.is_some());
1549 }
1550
1551 #[test]
1552 fn test_add_entry_rejects_content_failing_schema() {
1553 let schema = RegionSchema::new(ContentFormat::Json);
1554 let mut region =
1555 Region::new("data".to_string(), RegionKind::Temporary, 1000).with_schema(schema);
1556 let result = region.add_entry("not json".to_string(), 10);
1557 assert!(result.is_err());
1558 assert_eq!(region.entry_count(), 0);
1559 }
1560
1561 #[test]
1562 fn test_add_entry_accepts_content_passing_schema() {
1563 let schema = RegionSchema::new(ContentFormat::Json);
1564 let mut region =
1565 Region::new("data".to_string(), RegionKind::Temporary, 1000).with_schema(schema);
1566 let result = region.add_entry("{\"a\":1}".to_string(), 10);
1567 assert!(result.is_ok());
1568 assert_eq!(region.entry_count(), 1);
1569 }
1570
1571 #[test]
1572 fn test_add_entry_rejects_over_budget() {
1573 let mut region = Region::new("data".to_string(), RegionKind::Temporary, 10);
1574 let result = region.add_entry("too much".to_string(), 20);
1575 assert_eq!(
1576 result.unwrap_err().to_string(),
1577 "Content exceeds token budget: 20 > 10"
1578 );
1579 assert_eq!(region.entry_count(), 0);
1580 }
1581
1582 #[test]
1583 fn test_add_entry_with_metadata_rejects_content_failing_schema() {
1584 let schema = RegionSchema::new(ContentFormat::Json);
1585 let mut region =
1586 Region::new("data".to_string(), RegionKind::Temporary, 1000).with_schema(schema);
1587 let result =
1588 region.add_entry_with_metadata("not json".to_string(), 10, serde_json::json!({}));
1589 assert!(result.is_err());
1590 }
1591
1592 #[test]
1593 fn test_add_entry_with_metadata_rejects_over_budget() {
1594 let mut region = Region::new("data".to_string(), RegionKind::Temporary, 10);
1595 let result =
1596 region.add_entry_with_metadata("too much".to_string(), 20, serde_json::json!({}));
1597 assert_eq!(
1598 result.unwrap_err().to_string(),
1599 "Content exceeds token budget: 20 > 10"
1600 );
1601 }
1602
1603 #[test]
1604 fn test_add_entry_with_metadata_stores_metadata() {
1605 let mut region = Region::new("data".to_string(), RegionKind::Temporary, 1000);
1606 region
1607 .add_entry_with_metadata("hello".to_string(), 5, serde_json::json!({"k": "v"}))
1608 .unwrap();
1609 assert_eq!(
1610 region.content[0].metadata,
1611 Some(serde_json::json!({"k": "v"}))
1612 );
1613 }
1614
1615 #[test]
1618 fn test_clear_removes_all_content_and_resets_tokens() {
1619 let mut region = Region::new("data".to_string(), RegionKind::Temporary, 1000);
1620 region.add_entry("a".to_string(), 10).unwrap();
1621 region.add_entry("b".to_string(), 20).unwrap();
1622 assert_eq!(region.entry_count(), 2);
1623
1624 region.clear();
1625 assert_eq!(region.entry_count(), 0);
1626 assert_eq!(region.current_tokens, 0);
1627 }
1628
1629 #[test]
1630 fn test_remove_oldest_returns_and_removes_first_entry() {
1631 let mut region = Region::new("data".to_string(), RegionKind::Temporary, 1000);
1632 region.add_entry("first".to_string(), 10).unwrap();
1633 region.add_entry("second".to_string(), 20).unwrap();
1634
1635 let removed = region.remove_oldest().unwrap();
1636 assert_eq!(removed.content, "first");
1637 assert_eq!(region.entry_count(), 1);
1638 assert_eq!(region.current_tokens, 20);
1639 }
1640
1641 #[test]
1642 fn test_remove_oldest_returns_none_when_empty() {
1643 let mut region = Region::new("data".to_string(), RegionKind::Temporary, 1000);
1644 assert!(region.remove_oldest().is_none());
1645 }
1646
1647 #[test]
1648 fn test_needs_compaction_true_when_over_threshold() {
1649 let mut region = Region::new(
1650 "impl".to_string(),
1651 RegionKind::Compacting {
1652 threshold_tokens: 10,
1653 },
1654 1000,
1655 );
1656 region.add_entry("x".to_string(), 20).unwrap();
1657 assert!(region.needs_compaction());
1658 }
1659
1660 #[test]
1661 fn test_needs_compaction_false_when_under_threshold() {
1662 let mut region = Region::new(
1663 "impl".to_string(),
1664 RegionKind::Compacting {
1665 threshold_tokens: 100,
1666 },
1667 1000,
1668 );
1669 region.add_entry("x".to_string(), 20).unwrap();
1670 assert!(!region.needs_compaction());
1671 }
1672
1673 #[test]
1674 fn test_needs_compaction_false_for_non_compacting_kind() {
1675 let region = Region::new("data".to_string(), RegionKind::Temporary, 1000);
1676 assert!(!region.needs_compaction());
1677 }
1678
1679 #[test]
1682 fn test_region_schema_with_custom_script() {
1683 let schema = RegionSchema::new(ContentFormat::Custom {
1684 format_name: "special".to_string(),
1685 })
1686 .with_custom_script("validate_special()".to_string());
1687 assert_eq!(schema.custom_script.as_deref(), Some("validate_special()"));
1688 }
1689
1690 #[test]
1693 fn test_validate_json_valid() {
1694 let schema = RegionSchema::new(ContentFormat::Json);
1695 assert!(schema.validate("{\"a\": 1}").is_ok());
1696 }
1697
1698 #[test]
1699 fn test_validate_json_invalid() {
1700 let schema = RegionSchema::new(ContentFormat::Json);
1701 let err = schema.validate("not json").unwrap_err();
1702 assert!(err.to_string().starts_with("Region validation failed:"));
1703 }
1704
1705 #[test]
1706 fn test_validate_mermaid_valid() {
1707 let schema = RegionSchema::new(ContentFormat::Mermaid);
1708 assert!(schema.validate("graph TD\nA-->B").is_ok());
1709 }
1710
1711 #[test]
1712 fn test_validate_mermaid_all_recognized_diagram_types() {
1713 let schema = RegionSchema::new(ContentFormat::Mermaid);
1714 for kind in [
1715 "graph",
1716 "sequenceDiagram",
1717 "classDiagram",
1718 "stateDiagram",
1719 "erDiagram",
1720 "journey",
1721 "gantt",
1722 "pie",
1723 "flowchart",
1724 ] {
1725 assert!(schema.validate(&format!("{} content", kind)).is_ok());
1726 }
1727 }
1728
1729 #[test]
1730 fn test_validate_mermaid_invalid() {
1731 let schema = RegionSchema::new(ContentFormat::Mermaid);
1732 let err = schema.validate("just some text").unwrap_err();
1733 assert!(err.to_string().starts_with("Region validation failed:"));
1734 }
1735
1736 #[test]
1737 fn test_validate_code_non_empty_is_ok() {
1738 let schema = RegionSchema::new(ContentFormat::Code {
1739 language: "rust".to_string(),
1740 });
1741 assert!(schema.validate("fn main() {}").is_ok());
1742 }
1743
1744 #[test]
1745 fn test_validate_code_empty_is_error() {
1746 let schema = RegionSchema::new(ContentFormat::Code {
1747 language: "rust".to_string(),
1748 });
1749 let err = schema.validate(" ").unwrap_err();
1750 assert!(err.to_string().starts_with("Region validation failed:"));
1751 }
1752
1753 #[test]
1754 fn test_validate_markdown_non_empty_is_ok() {
1755 let schema = RegionSchema::new(ContentFormat::Markdown);
1756 assert!(schema.validate("# Heading").is_ok());
1757 }
1758
1759 #[test]
1760 fn test_validate_markdown_empty_is_error() {
1761 let schema = RegionSchema::new(ContentFormat::Markdown);
1762 let err = schema.validate("").unwrap_err();
1763 assert!(err.to_string().starts_with("Region validation failed:"));
1764 }
1765
1766 #[test]
1767 fn test_validate_text_has_no_restrictions() {
1768 let schema = RegionSchema::new(ContentFormat::Text);
1769 assert!(schema.validate("").is_ok());
1770 assert!(schema.validate("anything at all").is_ok());
1771 }
1772
1773 #[test]
1774 fn test_validate_custom_has_no_restrictions_here() {
1775 let schema = RegionSchema::new(ContentFormat::Custom {
1776 format_name: "special".to_string(),
1777 });
1778 assert!(schema.validate("").is_ok());
1781 assert!(schema.validate("whatever").is_ok());
1782 }
1783
1784 #[test]
1787 fn test_region_schema_clone_preserves_fields() {
1788 let schema = RegionSchema::new(ContentFormat::Text).with_custom_script("s".to_string());
1789 let cloned = schema.clone();
1790 assert_eq!(cloned.custom_script.as_deref(), Some("s"));
1791 assert_eq!(cloned.format, ContentFormat::Text);
1792 }
1793
1794 #[test]
1797 fn test_region_with_taint_tracking() {
1798 let region =
1799 Region::new("test".to_string(), RegionKind::Temporary, 1000).with_taint_tracking();
1800 assert!(region.taint.is_some());
1801 assert_eq!(region.taint_level(), Some(crate::taint::TaintLevel::Public));
1802 }
1803
1804 #[test]
1805 fn test_region_without_taint_tracking() {
1806 let region = Region::new("test".to_string(), RegionKind::Temporary, 1000);
1807 assert!(region.taint.is_none());
1808 assert_eq!(region.taint_level(), None);
1809 }
1810
1811 #[test]
1812 fn test_enable_taint_tracking() {
1813 let mut region = Region::new("test".to_string(), RegionKind::Temporary, 1000);
1814 assert!(region.taint.is_none());
1815 region.enable_taint_tracking();
1816 assert!(region.taint.is_some());
1817 region.enable_taint_tracking();
1819 assert!(region.taint.is_some());
1820 }
1821
1822 #[test]
1823 fn test_add_tainted_entry() {
1824 let mut region =
1825 Region::new("test".to_string(), RegionKind::Temporary, 1000).with_taint_tracking();
1826 region
1827 .add_tainted_entry(
1828 "secret data".to_string(),
1829 10,
1830 crate::taint::TaintLevel::Private,
1831 )
1832 .unwrap();
1833 assert_eq!(
1834 region.taint_level(),
1835 Some(crate::taint::TaintLevel::Private)
1836 );
1837 assert_eq!(region.entry_count(), 1);
1838 }
1839
1840 #[test]
1841 fn test_add_tainted_entry_validates_schema() {
1842 let mut region = Region::new("test".to_string(), RegionKind::Temporary, 1000)
1843 .with_taint_tracking()
1844 .with_schema(RegionSchema::new(ContentFormat::Json));
1845 let result = region.add_tainted_entry(
1846 "not json".to_string(),
1847 10,
1848 crate::taint::TaintLevel::Internal,
1849 );
1850 assert!(result.is_err());
1851 assert_eq!(region.entry_count(), 0);
1852 }
1853
1854 #[test]
1855 fn test_add_tainted_entry_checks_budget() {
1856 let mut region =
1857 Region::new("test".to_string(), RegionKind::Temporary, 10).with_taint_tracking();
1858 let result = region.add_tainted_entry(
1859 "too much".to_string(),
1860 20,
1861 crate::taint::TaintLevel::Internal,
1862 );
1863 assert!(result.is_err());
1864 }
1865
1866 #[test]
1867 fn test_add_entry_tracks_taint_as_public() {
1868 let mut region =
1869 Region::new("test".to_string(), RegionKind::Temporary, 1000).with_taint_tracking();
1870 region.add_entry("public data".to_string(), 10).unwrap();
1871 assert_eq!(region.taint_level(), Some(crate::taint::TaintLevel::Public));
1872 }
1873
1874 #[test]
1875 fn test_taint_recovery_on_remove_oldest() {
1876 let mut region =
1877 Region::new("test".to_string(), RegionKind::Temporary, 1000).with_taint_tracking();
1878 region
1879 .add_tainted_entry("private".to_string(), 10, crate::taint::TaintLevel::Private)
1880 .unwrap();
1881 region
1882 .add_tainted_entry("public".to_string(), 10, crate::taint::TaintLevel::Public)
1883 .unwrap();
1884 assert_eq!(
1885 region.taint_level(),
1886 Some(crate::taint::TaintLevel::Private)
1887 );
1888
1889 region.remove_oldest(); assert_eq!(region.taint_level(), Some(crate::taint::TaintLevel::Public));
1891 }
1892
1893 #[test]
1894 fn test_taint_recovery_on_clear() {
1895 let mut region =
1896 Region::new("test".to_string(), RegionKind::Temporary, 1000).with_taint_tracking();
1897 region
1898 .add_tainted_entry("private".to_string(), 10, crate::taint::TaintLevel::Private)
1899 .unwrap();
1900 region.clear();
1901 assert_eq!(region.taint_level(), Some(crate::taint::TaintLevel::Public));
1902 }
1903
1904 #[test]
1905 fn test_taint_recovery_on_sliding_window_eviction() {
1906 let mut region = Region::new(
1907 "conv".to_string(),
1908 RegionKind::SlidingWindow {
1909 max_items: 2,
1910 eviction_strategy: EvictionStrategy::PerItem,
1911 },
1912 50000,
1913 )
1914 .with_taint_tracking();
1915
1916 region
1917 .add_tainted_entry("private".to_string(), 10, crate::taint::TaintLevel::Private)
1918 .unwrap();
1919 region
1920 .add_tainted_entry("public1".to_string(), 10, crate::taint::TaintLevel::Public)
1921 .unwrap();
1922 assert_eq!(
1923 region.taint_level(),
1924 Some(crate::taint::TaintLevel::Private)
1925 );
1926
1927 region
1929 .add_tainted_entry("public2".to_string(), 10, crate::taint::TaintLevel::Public)
1930 .unwrap();
1931 assert_eq!(region.entry_count(), 2);
1932 assert_eq!(region.taint_level(), Some(crate::taint::TaintLevel::Public));
1933 }
1934
1935 #[test]
1936 fn test_taint_field_not_serialized_when_none() {
1937 let region = Region::new("test".to_string(), RegionKind::Temporary, 1000);
1938 let json = serde_json::to_string(®ion).unwrap();
1939 assert!(!json.contains("taint"));
1940 }
1941
1942 #[test]
1943 fn test_taint_field_deserialized_as_none_when_missing() {
1944 let json = r#"{"name":"test","kind":"Temporary","content":[],"max_tokens":1000,"current_tokens":0,"schema":null}"#;
1945 let region: Region = serde_json::from_str(json).unwrap();
1946 assert!(region.taint.is_none());
1947 }
1948
1949 #[test]
1950 fn test_add_typed_tainted_entry() {
1951 let mut region = Region::new(
1952 "conversation".to_string(),
1953 RegionKind::SlidingWindow {
1954 max_items: 100,
1955 eviction_strategy: EvictionStrategy::PerItem,
1956 },
1957 1000,
1958 )
1959 .with_taint_tracking();
1960
1961 region
1962 .add_typed_tainted_entry(
1963 "secret data".to_string(),
1964 10,
1965 EntryKind::ToolResult {
1966 tool_call_id: "tc_1".to_string(),
1967 tool_name: "calendar".to_string(),
1968 is_error: false,
1969 },
1970 crate::taint::TaintLevel::Private,
1971 )
1972 .unwrap();
1973
1974 assert_eq!(region.content.len(), 1);
1975 assert_eq!(
1976 region.content[0].kind,
1977 EntryKind::ToolResult {
1978 tool_call_id: "tc_1".to_string(),
1979 tool_name: "calendar".to_string(),
1980 is_error: false,
1981 }
1982 );
1983 assert_eq!(
1984 region.taint_level(),
1985 Some(crate::taint::TaintLevel::Private)
1986 );
1987 }
1988
1989 #[test]
1993 fn serialized_tool_call_round_trips_thought_signature_and_reads_old_json() {
1994 let with = SerializedToolCall {
1995 id: "c1".into(),
1996 name: "shell".into(),
1997 arguments: serde_json::json!({"command": "ls"}),
1998 thought_signature: Some("sig".into()),
1999 };
2000 let json = serde_json::to_string(&with).unwrap();
2001 let back: SerializedToolCall = serde_json::from_str(&json).unwrap();
2002 assert_eq!(back.thought_signature.as_deref(), Some("sig"));
2003
2004 let old = r#"{"id":"c2","name":"shell","arguments":{}}"#;
2006 let back: SerializedToolCall = serde_json::from_str(old).unwrap();
2007 assert_eq!(back.thought_signature, None);
2008
2009 let without = SerializedToolCall {
2012 id: "c3".into(),
2013 name: "shell".into(),
2014 arguments: serde_json::json!({}),
2015 thought_signature: None,
2016 };
2017 assert!(
2018 !serde_json::to_string(&without)
2019 .unwrap()
2020 .contains("thought_signature")
2021 );
2022 }
2023
2024 #[test]
2025 fn test_add_typed_tainted_entry_checks_budget() {
2026 let mut region = Region::new(
2027 "conversation".to_string(),
2028 RegionKind::SlidingWindow {
2029 max_items: 100,
2030 eviction_strategy: EvictionStrategy::PerItem,
2031 },
2032 5,
2033 )
2034 .with_taint_tracking();
2035
2036 let result = region.add_typed_tainted_entry(
2037 "too large".to_string(),
2038 100,
2039 EntryKind::ToolResult {
2040 tool_call_id: "tc_1".to_string(),
2041 tool_name: "tool".to_string(),
2042 is_error: false,
2043 },
2044 crate::taint::TaintLevel::Internal,
2045 );
2046 assert!(result.is_err());
2047 }
2048
2049 #[test]
2050 fn test_add_typed_tainted_entry_validates_schema() {
2051 let mut region = Region::new("test".to_string(), RegionKind::Pinned, 1000)
2052 .with_taint_tracking()
2053 .with_schema(RegionSchema::new(ContentFormat::Json));
2054
2055 let result = region.add_typed_tainted_entry(
2057 "not json".to_string(),
2058 5,
2059 EntryKind::Text,
2060 crate::taint::TaintLevel::Public,
2061 );
2062 assert!(result.is_err());
2063 }
2064
2065 #[test]
2066 fn test_add_typed_tainted_entry_without_taint_tracking() {
2067 let mut region = Region::new(
2070 "conversation".to_string(),
2071 RegionKind::SlidingWindow {
2072 max_items: 100,
2073 eviction_strategy: EvictionStrategy::PerItem,
2074 },
2075 1000,
2076 );
2077 region
2080 .add_typed_tainted_entry(
2081 "data".to_string(),
2082 10,
2083 EntryKind::Text,
2084 crate::taint::TaintLevel::Private,
2085 )
2086 .unwrap();
2087
2088 assert_eq!(region.content.len(), 1);
2089 assert_eq!(region.taint_level(), None); }
2091
2092 #[test]
2095 fn test_turn_group_size_at_assistant_with_tool_results() {
2096 let mut region = Region::new("conv".to_string(), RegionKind::Temporary, 50000);
2097 region
2098 .add_typed_entry(
2099 "assistant response".to_string(),
2100 10,
2101 EntryKind::AssistantTurn {
2102 tool_calls: vec![
2103 SerializedToolCall {
2104 id: "tc_1".to_string(),
2105 name: "read_file".to_string(),
2106 arguments: serde_json::json!({}),
2107 thought_signature: None,
2108 },
2109 SerializedToolCall {
2110 id: "tc_2".to_string(),
2111 name: "write_file".to_string(),
2112 arguments: serde_json::json!({}),
2113 thought_signature: None,
2114 },
2115 ],
2116 },
2117 )
2118 .unwrap();
2119 region
2120 .add_typed_entry(
2121 "result 1".to_string(),
2122 5,
2123 EntryKind::ToolResult {
2124 tool_call_id: "tc_1".to_string(),
2125 tool_name: "read_file".to_string(),
2126 is_error: false,
2127 },
2128 )
2129 .unwrap();
2130 region
2131 .add_typed_entry(
2132 "result 2".to_string(),
2133 5,
2134 EntryKind::ToolResult {
2135 tool_call_id: "tc_2".to_string(),
2136 tool_name: "write_file".to_string(),
2137 is_error: false,
2138 },
2139 )
2140 .unwrap();
2141
2142 assert_eq!(region.turn_group_size_at(0), 3);
2143 }
2144
2145 #[test]
2146 fn test_turn_group_size_at_assistant_at_end() {
2147 let mut region = Region::new("conv".to_string(), RegionKind::Temporary, 50000);
2148 region
2149 .add_typed_entry(
2150 "assistant with no tools".to_string(),
2151 10,
2152 EntryKind::AssistantTurn { tool_calls: vec![] },
2153 )
2154 .unwrap();
2155
2156 assert_eq!(region.turn_group_size_at(0), 1);
2157 }
2158
2159 #[test]
2160 fn test_turn_group_size_at_out_of_bounds() {
2161 let region = Region::new("conv".to_string(), RegionKind::Temporary, 50000);
2162 assert_eq!(region.turn_group_size_at(0), 0);
2163 assert_eq!(region.turn_group_size_at(99), 0);
2164 }
2165
2166 #[test]
2167 fn test_turn_group_size_at_non_assistant_entries() {
2168 let mut region = Region::new("conv".to_string(), RegionKind::Temporary, 50000);
2169 region
2170 .add_typed_entry("hello".to_string(), 5, EntryKind::Text)
2171 .unwrap();
2172 region
2173 .add_typed_entry("hi".to_string(), 5, EntryKind::UserMessage)
2174 .unwrap();
2175 region
2176 .add_typed_entry(
2177 "orphan result".to_string(),
2178 5,
2179 EntryKind::ToolResult {
2180 tool_call_id: "tc_x".to_string(),
2181 tool_name: "tool".to_string(),
2182 is_error: false,
2183 },
2184 )
2185 .unwrap();
2186
2187 assert_eq!(region.turn_group_size_at(0), 1); assert_eq!(region.turn_group_size_at(1), 1); assert_eq!(region.turn_group_size_at(2), 1); }
2191
2192 #[test]
2195 fn test_remove_oldest_evicts_entire_turn_group() {
2196 let mut region = Region::new("conv".to_string(), RegionKind::Temporary, 50000);
2197 region
2199 .add_typed_entry(
2200 "assistant".to_string(),
2201 100,
2202 EntryKind::AssistantTurn {
2203 tool_calls: vec![
2204 SerializedToolCall {
2205 id: "tc_1".to_string(),
2206 name: "read_file".to_string(),
2207 arguments: serde_json::json!({}),
2208 thought_signature: None,
2209 },
2210 SerializedToolCall {
2211 id: "tc_2".to_string(),
2212 name: "list_dir".to_string(),
2213 arguments: serde_json::json!({}),
2214 thought_signature: None,
2215 },
2216 ],
2217 },
2218 )
2219 .unwrap();
2220 region
2221 .add_typed_entry(
2222 "result 1".to_string(),
2223 30,
2224 EntryKind::ToolResult {
2225 tool_call_id: "tc_1".to_string(),
2226 tool_name: "read_file".to_string(),
2227 is_error: false,
2228 },
2229 )
2230 .unwrap();
2231 region
2232 .add_typed_entry(
2233 "result 2".to_string(),
2234 20,
2235 EntryKind::ToolResult {
2236 tool_call_id: "tc_2".to_string(),
2237 tool_name: "list_dir".to_string(),
2238 is_error: false,
2239 },
2240 )
2241 .unwrap();
2242 region
2244 .add_typed_entry("user msg".to_string(), 10, EntryKind::UserMessage)
2245 .unwrap();
2246
2247 assert_eq!(region.entry_count(), 4);
2248 assert_eq!(region.current_tokens, 160);
2249
2250 let removed = region.remove_oldest().unwrap();
2251 assert_eq!(removed.content, "assistant");
2254 assert_eq!(removed.tokens, 100 + 30 + 20); assert_eq!(region.entry_count(), 1);
2257 assert_eq!(region.content[0].content, "user msg");
2258 assert_eq!(region.current_tokens, 10);
2259 }
2260
2261 #[test]
2264 fn test_remove_oldest_turn_group_calls_taint_remove_for_each_entry() {
2265 let mut region =
2266 Region::new("conv".to_string(), RegionKind::Temporary, 50000).with_taint_tracking();
2267
2268 region
2270 .add_typed_tainted_entry(
2271 "assistant".to_string(),
2272 10,
2273 EntryKind::AssistantTurn {
2274 tool_calls: vec![SerializedToolCall {
2275 id: "tc_1".to_string(),
2276 name: "tool".to_string(),
2277 arguments: serde_json::json!({}),
2278 thought_signature: None,
2279 }],
2280 },
2281 crate::taint::TaintLevel::Private,
2282 )
2283 .unwrap();
2284 region
2285 .add_typed_tainted_entry(
2286 "result".to_string(),
2287 5,
2288 EntryKind::ToolResult {
2289 tool_call_id: "tc_1".to_string(),
2290 tool_name: "tool".to_string(),
2291 is_error: false,
2292 },
2293 crate::taint::TaintLevel::Internal,
2294 )
2295 .unwrap();
2296 region
2297 .add_tainted_entry(
2298 "public stuff".to_string(),
2299 5,
2300 crate::taint::TaintLevel::Public,
2301 )
2302 .unwrap();
2303
2304 assert_eq!(
2305 region.taint_level(),
2306 Some(crate::taint::TaintLevel::Private)
2307 );
2308 assert_eq!(region.taint.as_ref().unwrap().entry_count(), 3);
2309
2310 let removed = region.remove_oldest().unwrap();
2312 assert_eq!(removed.content, "assistant");
2313 assert_eq!(region.entry_count(), 1);
2314 assert_eq!(region.taint.as_ref().unwrap().entry_count(), 1);
2317 assert_eq!(region.taint_level(), Some(crate::taint::TaintLevel::Public));
2318 }
2319
2320 #[test]
2323 fn test_sliding_window_evicts_entire_turn_group() {
2324 let mut region = Region::new(
2325 "conv".to_string(),
2326 RegionKind::SlidingWindow {
2327 max_items: 3,
2328 eviction_strategy: EvictionStrategy::PerItem,
2329 },
2330 50000,
2331 );
2332
2333 region
2335 .add_typed_entry(
2336 "assistant".to_string(),
2337 10,
2338 EntryKind::AssistantTurn {
2339 tool_calls: vec![
2340 SerializedToolCall {
2341 id: "tc_1".to_string(),
2342 name: "t1".to_string(),
2343 arguments: serde_json::json!({}),
2344 thought_signature: None,
2345 },
2346 SerializedToolCall {
2347 id: "tc_2".to_string(),
2348 name: "t2".to_string(),
2349 arguments: serde_json::json!({}),
2350 thought_signature: None,
2351 },
2352 ],
2353 },
2354 )
2355 .unwrap();
2356 region
2357 .add_typed_entry(
2358 "r1".to_string(),
2359 5,
2360 EntryKind::ToolResult {
2361 tool_call_id: "tc_1".to_string(),
2362 tool_name: "t1".to_string(),
2363 is_error: false,
2364 },
2365 )
2366 .unwrap();
2367 region
2368 .add_typed_entry(
2369 "r2".to_string(),
2370 5,
2371 EntryKind::ToolResult {
2372 tool_call_id: "tc_2".to_string(),
2373 tool_name: "t2".to_string(),
2374 is_error: false,
2375 },
2376 )
2377 .unwrap();
2378
2379 assert_eq!(region.entry_count(), 3);
2380
2381 region
2384 .add_typed_entry("user msg".to_string(), 15, EntryKind::UserMessage)
2385 .unwrap();
2386
2387 assert_eq!(region.entry_count(), 1);
2389 assert_eq!(region.content[0].content, "user msg");
2390 assert_eq!(region.current_tokens, 15);
2391 }
2392
2393 #[test]
2396 fn test_add_entry_with_metadata_tracks_taint_as_public() {
2397 let mut region =
2398 Region::new("data".to_string(), RegionKind::Temporary, 1000).with_taint_tracking();
2399
2400 region
2401 .add_entry_with_metadata("content".to_string(), 10, serde_json::json!({"key": "val"}))
2402 .unwrap();
2403
2404 assert_eq!(region.taint_level(), Some(crate::taint::TaintLevel::Public));
2405 assert_eq!(region.taint.as_ref().unwrap().entry_count(), 1);
2406 assert_eq!(
2407 region.taint.as_ref().unwrap().entry_taint(0),
2408 Some(crate::taint::TaintLevel::Public)
2409 );
2410 }
2411
2412 #[test]
2415 fn test_add_typed_entry_tracks_taint_as_public() {
2416 let mut region =
2417 Region::new("conv".to_string(), RegionKind::Temporary, 1000).with_taint_tracking();
2418
2419 region
2420 .add_typed_entry(
2421 "assistant response".to_string(),
2422 10,
2423 EntryKind::AssistantTurn { tool_calls: vec![] },
2424 )
2425 .unwrap();
2426
2427 assert_eq!(region.taint_level(), Some(crate::taint::TaintLevel::Public));
2428 assert_eq!(region.taint.as_ref().unwrap().entry_count(), 1);
2429 assert_eq!(
2430 region.taint.as_ref().unwrap().entry_taint(0),
2431 Some(crate::taint::TaintLevel::Public)
2432 );
2433 }
2434
2435 #[test]
2438 fn test_per_item_strategy_evicts_one_at_a_time() {
2439 let mut region = Region::new(
2440 "conv".to_string(),
2441 RegionKind::SlidingWindow {
2442 max_items: 3,
2443 eviction_strategy: EvictionStrategy::PerItem,
2444 },
2445 50000,
2446 );
2447 for i in 0..5 {
2448 region.add_entry(format!("msg{}", i), 10).unwrap();
2449 }
2450 assert_eq!(region.entry_count(), 3);
2451 assert_eq!(region.content[0].content, "msg2");
2452 assert_eq!(region.content[1].content, "msg3");
2453 assert_eq!(region.content[2].content, "msg4");
2454 }
2455
2456 #[test]
2457 fn test_bulk_eviction_triggers_on_overflow() {
2458 let mut region = Region::new(
2459 "conv".to_string(),
2460 RegionKind::SlidingWindow {
2461 max_items: 5,
2462 eviction_strategy: EvictionStrategy::Bulk { overflow: 3 },
2463 },
2464 50000,
2465 );
2466 for i in 0..8 {
2469 region.add_entry(format!("msg{}", i), 10).unwrap();
2470 }
2471 assert_eq!(region.entry_count(), 8);
2472
2473 region.add_entry("msg8".to_string(), 10).unwrap();
2475 assert_eq!(region.entry_count(), 5);
2476 assert_eq!(region.content[0].content, "msg4");
2477 }
2478
2479 #[test]
2480 fn test_bulk_eviction_respects_turn_groups() {
2481 let mut region = Region::new(
2482 "conv".to_string(),
2483 RegionKind::SlidingWindow {
2484 max_items: 3,
2485 eviction_strategy: EvictionStrategy::Bulk { overflow: 2 },
2486 },
2487 50000,
2488 );
2489 region
2491 .add_typed_entry(
2492 "assistant".to_string(),
2493 10,
2494 EntryKind::AssistantTurn {
2495 tool_calls: vec![SerializedToolCall {
2496 id: "tc1".to_string(),
2497 name: "tool".to_string(),
2498 arguments: serde_json::json!({}),
2499 thought_signature: None,
2500 }],
2501 },
2502 )
2503 .unwrap();
2504 region
2505 .add_typed_entry(
2506 "result".to_string(),
2507 5,
2508 EntryKind::ToolResult {
2509 tool_call_id: "tc1".to_string(),
2510 tool_name: "tool".to_string(),
2511 is_error: false,
2512 },
2513 )
2514 .unwrap();
2515 region.add_entry("msg2".to_string(), 10).unwrap();
2517 region.add_entry("msg3".to_string(), 10).unwrap();
2518 region.add_entry("msg4".to_string(), 10).unwrap();
2519 assert_eq!(region.entry_count(), 5);
2521
2522 region.add_entry("msg5".to_string(), 10).unwrap();
2524 assert_eq!(region.entry_count(), 3);
2527 assert_eq!(region.content[0].content, "msg3");
2528 }
2529
2530 #[test]
2531 fn test_bulk_eviction_under_overflow_no_eviction() {
2532 let mut region = Region::new(
2533 "conv".to_string(),
2534 RegionKind::SlidingWindow {
2535 max_items: 5,
2536 eviction_strategy: EvictionStrategy::Bulk { overflow: 3 },
2537 },
2538 50000,
2539 );
2540 for i in 0..7 {
2542 region.add_entry(format!("msg{}", i), 10).unwrap();
2543 }
2544 assert_eq!(region.entry_count(), 7);
2546 }
2547
2548 #[test]
2549 fn test_compact_sets_needs_message_compaction_flag() {
2550 let mut region = Region::new(
2551 "conv".to_string(),
2552 RegionKind::SlidingWindow {
2553 max_items: 5,
2554 eviction_strategy: EvictionStrategy::Compact { compact_count: 3 },
2555 },
2556 50000,
2557 );
2558 assert!(!region.needs_message_compaction);
2559
2560 for i in 0..9 {
2562 region.add_entry(format!("msg{}", i), 10).unwrap();
2563 }
2564 assert!(region.needs_message_compaction);
2565 assert_eq!(region.entry_count(), 9);
2567 }
2568
2569 #[test]
2570 fn test_compact_fallback_to_bulk_eviction() {
2571 let mut region = Region::new(
2572 "conv".to_string(),
2573 RegionKind::SlidingWindow {
2574 max_items: 5,
2575 eviction_strategy: EvictionStrategy::Compact { compact_count: 3 },
2576 },
2577 50000,
2578 );
2579 for i in 0..12 {
2582 region.add_entry(format!("msg{}", i), 10).unwrap();
2583 }
2584 assert_eq!(region.entry_count(), 5);
2586 assert_eq!(region.content[0].content, "msg7");
2587 assert!(!region.needs_message_compaction);
2589 }
2590
2591 #[test]
2592 fn test_eviction_strategy_default_is_per_item() {
2593 assert_eq!(EvictionStrategy::default(), EvictionStrategy::PerItem);
2594 }
2595
2596 #[test]
2597 fn test_remove_entries_by_prefix() {
2598 let mut region = Region::new("system".to_string(), RegionKind::Pinned, 50000);
2599 region
2600 .add_entry("[Stage instructions: Be terse.]".to_string(), 10)
2601 .unwrap();
2602 region
2603 .add_entry("Core identity block".to_string(), 20)
2604 .unwrap();
2605 region
2606 .add_entry("[Stage instructions: Be verbose.]".to_string(), 15)
2607 .unwrap();
2608
2609 assert_eq!(region.entry_count(), 3);
2610 region.remove_entries_by_prefix("[Stage instructions:");
2611 assert_eq!(region.entry_count(), 1);
2612 assert_eq!(region.content[0].content, "Core identity block");
2613 assert_eq!(region.current_tokens, 20);
2614 }
2615
2616 #[test]
2617 fn test_remove_entries_by_prefix_with_taint_tracking() {
2618 let mut region =
2619 Region::new("system".to_string(), RegionKind::Pinned, 50000).with_taint_tracking();
2620 region
2621 .add_tainted_entry(
2622 "[Stage instructions: Be terse.]".to_string(),
2623 10,
2624 crate::taint::TaintLevel::Private,
2625 )
2626 .unwrap();
2627 region
2628 .add_tainted_entry(
2629 "Core identity block".to_string(),
2630 20,
2631 crate::taint::TaintLevel::Public,
2632 )
2633 .unwrap();
2634 region
2635 .add_tainted_entry(
2636 "[Stage instructions: Be verbose.]".to_string(),
2637 15,
2638 crate::taint::TaintLevel::Internal,
2639 )
2640 .unwrap();
2641
2642 assert_eq!(region.entry_count(), 3);
2643 assert_eq!(
2644 region.taint_level(),
2645 Some(crate::taint::TaintLevel::Private)
2646 );
2647
2648 region.remove_entries_by_prefix("[Stage instructions:");
2649 assert_eq!(region.entry_count(), 1);
2650 assert_eq!(region.content[0].content, "Core identity block");
2651 assert_eq!(region.current_tokens, 20);
2652 assert_eq!(region.taint_level(), Some(crate::taint::TaintLevel::Public));
2654 assert_eq!(region.taint.as_ref().unwrap().entry_count(), 1);
2655 }
2656
2657 #[test]
2658 fn test_compact_below_threshold_no_flag() {
2659 let mut region = Region::new(
2661 "conv".to_string(),
2662 RegionKind::SlidingWindow {
2663 max_items: 5,
2664 eviction_strategy: EvictionStrategy::Compact { compact_count: 3 },
2665 },
2666 50000,
2667 );
2668 for i in 0..8 {
2669 region.add_entry(format!("msg{}", i), 10).unwrap();
2670 }
2671 assert!(!region.needs_message_compaction);
2673 assert_eq!(region.entry_count(), 8);
2674 }
2675
2676 #[test]
2677 fn test_bulk_eviction_with_taint_tracking() {
2678 let mut region = Region::new(
2679 "conv".to_string(),
2680 RegionKind::SlidingWindow {
2681 max_items: 3,
2682 eviction_strategy: EvictionStrategy::Bulk { overflow: 2 },
2683 },
2684 50000,
2685 )
2686 .with_taint_tracking();
2687
2688 region
2690 .add_tainted_entry("private".to_string(), 10, crate::taint::TaintLevel::Private)
2691 .unwrap();
2692 for i in 1..5 {
2693 region
2694 .add_tainted_entry(format!("pub{}", i), 10, crate::taint::TaintLevel::Public)
2695 .unwrap();
2696 }
2697 assert_eq!(region.entry_count(), 5);
2698
2699 region
2701 .add_tainted_entry("pub5".to_string(), 10, crate::taint::TaintLevel::Public)
2702 .unwrap();
2703 assert_eq!(region.entry_count(), 3);
2704 assert_eq!(region.taint_level(), Some(crate::taint::TaintLevel::Public));
2706 }
2707
2708 #[test]
2709 fn test_eviction_strategy_serde_roundtrip() {
2710 let bulk = EvictionStrategy::Bulk { overflow: 5 };
2711 let json = serde_json::to_string(&bulk).unwrap();
2712 let parsed: EvictionStrategy = serde_json::from_str(&json).unwrap();
2713 assert_eq!(parsed, bulk);
2714
2715 let compact = EvictionStrategy::Compact { compact_count: 10 };
2716 let json = serde_json::to_string(&compact).unwrap();
2717 let parsed: EvictionStrategy = serde_json::from_str(&json).unwrap();
2718 assert_eq!(parsed, compact);
2719
2720 let per_item = EvictionStrategy::PerItem;
2721 let json = serde_json::to_string(&per_item).unwrap();
2722 let parsed: EvictionStrategy = serde_json::from_str(&json).unwrap();
2723 assert_eq!(parsed, per_item);
2724 }
2725
2726 #[test]
2727 fn test_sliding_window_kind_equality_with_eviction_strategy() {
2728 assert_eq!(
2729 RegionKind::SlidingWindow {
2730 max_items: 10,
2731 eviction_strategy: EvictionStrategy::Bulk { overflow: 3 },
2732 },
2733 RegionKind::SlidingWindow {
2734 max_items: 10,
2735 eviction_strategy: EvictionStrategy::Bulk { overflow: 3 },
2736 }
2737 );
2738 assert_ne!(
2739 RegionKind::SlidingWindow {
2740 max_items: 10,
2741 eviction_strategy: EvictionStrategy::PerItem,
2742 },
2743 RegionKind::SlidingWindow {
2744 max_items: 10,
2745 eviction_strategy: EvictionStrategy::Bulk { overflow: 3 },
2746 }
2747 );
2748 }
2749
2750 #[test]
2751 fn test_needs_message_compaction_default_false() {
2752 let region = Region::new("conv".to_string(), RegionKind::Temporary, 1000);
2753 assert!(!region.needs_message_compaction);
2754 }
2755
2756 #[test]
2759 fn test_add_typed_entry_validates_schema() {
2760 let mut region = Region::new("data".to_string(), RegionKind::Temporary, 1000)
2761 .with_schema(RegionSchema::new(ContentFormat::Json));
2762 let result = region.add_typed_entry("not json".to_string(), 5, EntryKind::Text);
2763 assert!(result.is_err());
2764 assert_eq!(region.entry_count(), 0);
2765 }
2766
2767 #[test]
2768 fn test_add_typed_entry_checks_budget() {
2769 let mut region = Region::new("data".to_string(), RegionKind::Temporary, 10);
2770 let result = region.add_typed_entry("too big".to_string(), 20, EntryKind::UserMessage);
2771 assert!(result.is_err());
2772 assert_eq!(region.entry_count(), 0);
2773 }
2774
2775 #[test]
2776 fn test_add_tainted_entry_without_taint_tracking() {
2777 let mut region = Region::new("data".to_string(), RegionKind::Temporary, 1000);
2779 region
2780 .add_tainted_entry("data".to_string(), 10, crate::taint::TaintLevel::Private)
2781 .unwrap();
2782 assert_eq!(region.entry_count(), 1);
2783 assert_eq!(region.taint_level(), None);
2784 }
2785
2786 #[test]
2787 fn test_remove_entries_by_prefix_no_match() {
2788 let mut region = Region::new("system".to_string(), RegionKind::Pinned, 50000);
2789 region.add_entry("Keep this".to_string(), 10).unwrap();
2790 region.add_entry("And this".to_string(), 20).unwrap();
2791 region.remove_entries_by_prefix("[Stage instructions:");
2792 assert_eq!(region.entry_count(), 2);
2793 assert_eq!(region.current_tokens, 30);
2794 }
2795
2796 #[test]
2799 fn test_hashmap_region_upsert_and_get() {
2800 let mut region = Region::new(
2801 "files".to_string(),
2802 RegionKind::HashMap { max_entries: None },
2803 10000,
2804 );
2805 region
2806 .upsert_by_key("src/main.rs", "fn main() {}".to_string(), 10)
2807 .unwrap();
2808 region
2809 .upsert_by_key("src/lib.rs", "pub mod foo;".to_string(), 8)
2810 .unwrap();
2811
2812 assert_eq!(region.entry_count(), 2);
2813 assert_eq!(region.current_tokens, 18);
2814
2815 let entry = region.get_by_key("src/main.rs").unwrap();
2816 assert_eq!(entry.content, "fn main() {}");
2817 assert_eq!(entry.key.as_deref(), Some("src/main.rs"));
2818 }
2819
2820 #[test]
2821 fn test_hashmap_region_upsert_replaces_existing() {
2822 let mut region = Region::new(
2823 "files".to_string(),
2824 RegionKind::HashMap { max_entries: None },
2825 10000,
2826 );
2827 region
2828 .upsert_by_key("file.rs", "version 1".to_string(), 10)
2829 .unwrap();
2830 assert_eq!(region.current_tokens, 10);
2831
2832 region
2833 .upsert_by_key("file.rs", "version 2".to_string(), 15)
2834 .unwrap();
2835 assert_eq!(region.entry_count(), 1);
2836 assert_eq!(region.current_tokens, 15);
2837 assert_eq!(region.get_by_key("file.rs").unwrap().content, "version 2");
2838 }
2839
2840 #[test]
2841 fn test_hashmap_region_remove_by_key() {
2842 let mut region = Region::new(
2843 "files".to_string(),
2844 RegionKind::HashMap { max_entries: None },
2845 10000,
2846 );
2847 region.upsert_by_key("a.rs", "aaa".to_string(), 10).unwrap();
2848 region.upsert_by_key("b.rs", "bbb".to_string(), 20).unwrap();
2849
2850 assert!(region.remove_by_key("a.rs"));
2851 assert_eq!(region.entry_count(), 1);
2852 assert_eq!(region.current_tokens, 20);
2853 assert!(region.get_by_key("a.rs").is_none());
2854 assert!(!region.remove_by_key("nonexistent"));
2855 }
2856
2857 #[test]
2858 fn test_hashmap_region_keys() {
2859 let mut region = Region::new(
2860 "files".to_string(),
2861 RegionKind::HashMap { max_entries: None },
2862 10000,
2863 );
2864 region.upsert_by_key("x.rs", "x".to_string(), 5).unwrap();
2865 region.upsert_by_key("y.rs", "y".to_string(), 5).unwrap();
2866
2867 let keys = region.keys();
2868 assert_eq!(keys.len(), 2);
2869 assert!(keys.contains(&"x.rs"));
2870 assert!(keys.contains(&"y.rs"));
2871 }
2872
2873 #[test]
2874 fn test_hashmap_region_lru_eviction_on_max_tokens() {
2875 let mut region = Region::new(
2876 "files".to_string(),
2877 RegionKind::HashMap { max_entries: None },
2878 30, );
2880 region.upsert_by_key("a.rs", "aaa".to_string(), 10).unwrap();
2881 region.content[0].timestamp -= 100;
2883 region.upsert_by_key("b.rs", "bbb".to_string(), 10).unwrap();
2884 region.upsert_by_key("c.rs", "ccc".to_string(), 10).unwrap();
2885 assert_eq!(region.entry_count(), 3);
2886 assert_eq!(region.current_tokens, 30);
2887
2888 region.upsert_by_key("d.rs", "ddd".to_string(), 10).unwrap();
2890 assert_eq!(region.entry_count(), 3);
2891 assert!(region.get_by_key("a.rs").is_none());
2892 assert!(region.get_by_key("d.rs").is_some());
2893 }
2894
2895 #[test]
2896 fn test_hashmap_region_max_entries_eviction() {
2897 let mut region = Region::new(
2898 "files".to_string(),
2899 RegionKind::HashMap {
2900 max_entries: Some(2),
2901 },
2902 10000,
2903 );
2904 region.upsert_by_key("a.rs", "aaa".to_string(), 10).unwrap();
2905 region.content[0].timestamp -= 100; region.upsert_by_key("b.rs", "bbb".to_string(), 10).unwrap();
2907 assert_eq!(region.entry_count(), 2);
2908
2909 region.upsert_by_key("c.rs", "ccc".to_string(), 10).unwrap();
2911 assert_eq!(region.entry_count(), 2);
2912 assert!(region.get_by_key("a.rs").is_none());
2913 assert!(region.get_by_key("c.rs").is_some());
2914 }
2915
2916 #[test]
2917 fn test_hashmap_region_upsert_too_large_for_budget() {
2918 let mut region = Region::new(
2919 "files".to_string(),
2920 RegionKind::HashMap { max_entries: None },
2921 5, );
2923 let result = region.upsert_by_key("big.rs", "huge content".to_string(), 100);
2924 assert!(result.is_err());
2925 }
2926
2927 #[test]
2928 fn test_hashmap_region_kind_equality() {
2929 assert_eq!(
2930 RegionKind::HashMap {
2931 max_entries: Some(10)
2932 },
2933 RegionKind::HashMap {
2934 max_entries: Some(10)
2935 }
2936 );
2937 assert_ne!(
2938 RegionKind::HashMap {
2939 max_entries: Some(10)
2940 },
2941 RegionKind::HashMap {
2942 max_entries: Some(20)
2943 }
2944 );
2945 assert_ne!(
2946 RegionKind::HashMap { max_entries: None },
2947 RegionKind::Pinned
2948 );
2949 }
2950
2951 #[test]
2952 fn test_hashmap_cache_hint() {
2953 let kind = RegionKind::HashMap { max_entries: None };
2954 assert_eq!(kind.cache_hint(), crate::cache::CacheHint::UntilChanged);
2955 }
2956
2957 #[test]
2958 fn test_region_entry_key_default_none() {
2959 let mut region = Region::new("test".to_string(), RegionKind::Temporary, 1000);
2960 region.add_entry("content".to_string(), 10).unwrap();
2961 assert!(region.content[0].key.is_none());
2962 }
2963
2964 #[test]
2965 fn test_region_entry_key_serde_skip_when_none() {
2966 let entry = RegionEntry {
2967 content: "test".to_string(),
2968 tokens: 5,
2969 timestamp: 0,
2970 metadata: None,
2971 kind: EntryKind::default(),
2972 key: None,
2973 };
2974 let json = serde_json::to_string(&entry).unwrap();
2975 assert!(!json.contains("key"));
2976 }
2977
2978 #[test]
2979 fn test_region_entry_key_serde_roundtrip() {
2980 let entry = RegionEntry {
2981 content: "test".to_string(),
2982 tokens: 5,
2983 timestamp: 0,
2984 metadata: None,
2985 kind: EntryKind::default(),
2986 key: Some("mykey".to_string()),
2987 };
2988 let json = serde_json::to_string(&entry).unwrap();
2989 assert!(json.contains("mykey"));
2990 let back: RegionEntry = serde_json::from_str(&json).unwrap();
2991 assert_eq!(back.key.as_deref(), Some("mykey"));
2992 }
2993
2994 #[test]
2997 fn test_hashmap_region_creation_and_basic_properties() {
2998 let region = Region::new(
2999 "lookup".to_string(),
3000 RegionKind::HashMap {
3001 max_entries: Some(5),
3002 },
3003 2000,
3004 );
3005 assert_eq!(region.name, "lookup");
3006 assert_eq!(
3007 region.kind,
3008 RegionKind::HashMap {
3009 max_entries: Some(5)
3010 }
3011 );
3012 assert_eq!(region.max_tokens, 2000);
3013 assert_eq!(region.current_tokens, 0);
3014 assert_eq!(region.entry_count(), 0);
3015 assert!(region.content.is_empty());
3016 }
3017
3018 #[test]
3019 fn test_hashmap_upsert_insert_new_entry() {
3020 let mut region = Region::new(
3021 "store".to_string(),
3022 RegionKind::HashMap {
3023 max_entries: Some(5),
3024 },
3025 5000,
3026 );
3027 region
3028 .upsert_by_key("config.toml", "[package]\nname = \"foo\"".to_string(), 12)
3029 .unwrap();
3030
3031 assert_eq!(region.entry_count(), 1);
3032 assert_eq!(region.current_tokens, 12);
3033
3034 let entry = region.get_by_key("config.toml").unwrap();
3035 assert_eq!(entry.content, "[package]\nname = \"foo\"");
3036 assert_eq!(entry.tokens, 12);
3037 assert_eq!(entry.key.as_deref(), Some("config.toml"));
3038 }
3039
3040 #[test]
3041 fn test_hashmap_upsert_update_existing_entry() {
3042 let mut region = Region::new(
3043 "store".to_string(),
3044 RegionKind::HashMap { max_entries: None },
3045 5000,
3046 );
3047 region
3048 .upsert_by_key("readme.md", "# Old".to_string(), 20)
3049 .unwrap();
3050 assert_eq!(region.current_tokens, 20);
3051
3052 region
3053 .upsert_by_key("readme.md", "# New and improved".to_string(), 35)
3054 .unwrap();
3055 assert_eq!(region.entry_count(), 1);
3056 assert_eq!(region.current_tokens, 35);
3057
3058 let entry = region.get_by_key("readme.md").unwrap();
3059 assert_eq!(entry.content, "# New and improved");
3060 assert_eq!(entry.tokens, 35);
3061 }
3062
3063 #[test]
3064 fn test_hashmap_upsert_lru_eviction_on_max_tokens() {
3065 let mut region = Region::new(
3066 "files".to_string(),
3067 RegionKind::HashMap { max_entries: None },
3068 100, );
3070
3071 region
3073 .upsert_by_key("first.rs", "first content".to_string(), 40)
3074 .unwrap();
3075 region.content[0].timestamp -= 200; region
3078 .upsert_by_key("second.rs", "second content".to_string(), 40)
3079 .unwrap();
3080 region.content[1].timestamp -= 100; region
3083 .upsert_by_key("third.rs", "third content".to_string(), 20)
3084 .unwrap();
3085 region
3089 .upsert_by_key("fourth.rs", "fourth content".to_string(), 30)
3090 .unwrap();
3091
3092 assert!(region.get_by_key("first.rs").is_none());
3094 assert!(region.get_by_key("fourth.rs").is_some());
3095 assert!(region.current_tokens <= 100);
3097 }
3098
3099 #[test]
3100 fn test_hashmap_upsert_max_entries_enforcement() {
3101 let mut region = Region::new(
3102 "cache".to_string(),
3103 RegionKind::HashMap {
3104 max_entries: Some(2),
3105 },
3106 50000,
3107 );
3108
3109 region
3110 .upsert_by_key("alpha", "aaa".to_string(), 10)
3111 .unwrap();
3112 region.content[0].timestamp -= 200; region.upsert_by_key("beta", "bbb".to_string(), 10).unwrap();
3115 region.content[1].timestamp -= 100;
3116
3117 region
3118 .upsert_by_key("gamma", "ccc".to_string(), 10)
3119 .unwrap();
3120
3121 assert_eq!(region.entry_count(), 2);
3123 assert!(region.get_by_key("alpha").is_none());
3124 assert!(region.get_by_key("beta").is_some());
3125 assert!(region.get_by_key("gamma").is_some());
3126 }
3127
3128 #[test]
3129 fn test_hashmap_get_by_key_found_and_not_found() {
3130 let mut region = Region::new(
3131 "data".to_string(),
3132 RegionKind::HashMap { max_entries: None },
3133 5000,
3134 );
3135 region
3136 .upsert_by_key("exists", "hello".to_string(), 5)
3137 .unwrap();
3138
3139 let found = region.get_by_key("exists");
3141 assert!(found.is_some());
3142 assert_eq!(found.unwrap().content, "hello");
3143
3144 let missing = region.get_by_key("does_not_exist");
3146 assert!(missing.is_none());
3147 }
3148
3149 #[test]
3150 fn test_hashmap_remove_by_key_exists() {
3151 let mut region = Region::new(
3152 "data".to_string(),
3153 RegionKind::HashMap { max_entries: None },
3154 5000,
3155 );
3156 region
3157 .upsert_by_key("target", "remove me".to_string(), 25)
3158 .unwrap();
3159 assert_eq!(region.current_tokens, 25);
3160
3161 let removed = region.remove_by_key("target");
3162 assert!(removed);
3163 assert_eq!(region.entry_count(), 0);
3164 assert_eq!(region.current_tokens, 0);
3165 assert!(region.get_by_key("target").is_none());
3166 }
3167
3168 #[test]
3169 fn test_hashmap_remove_by_key_does_not_exist() {
3170 let mut region = Region::new(
3171 "data".to_string(),
3172 RegionKind::HashMap { max_entries: None },
3173 5000,
3174 );
3175 let removed = region.remove_by_key("ghost");
3176 assert!(!removed);
3177 }
3178
3179 #[test]
3180 fn test_hashmap_keys_empty_populated_after_removal() {
3181 let mut region = Region::new(
3182 "data".to_string(),
3183 RegionKind::HashMap { max_entries: None },
3184 5000,
3185 );
3186
3187 assert!(region.keys().is_empty());
3189
3190 region.upsert_by_key("one", "1".to_string(), 5).unwrap();
3192 region.upsert_by_key("two", "2".to_string(), 5).unwrap();
3193 region.upsert_by_key("three", "3".to_string(), 5).unwrap();
3194
3195 let keys = region.keys();
3196 assert_eq!(keys.len(), 3);
3197 assert!(keys.contains(&"one"));
3198 assert!(keys.contains(&"two"));
3199 assert!(keys.contains(&"three"));
3200
3201 region.remove_by_key("two");
3203 let keys = region.keys();
3204 assert_eq!(keys.len(), 2);
3205 assert!(keys.contains(&"one"));
3206 assert!(!keys.contains(&"two"));
3207 assert!(keys.contains(&"three"));
3208 }
3209
3210 #[test]
3211 fn test_region_entry_serialization_with_key_field() {
3212 let entry_with_key = RegionEntry {
3214 content: "some data".to_string(),
3215 tokens: 10,
3216 timestamp: 1234567890,
3217 metadata: None,
3218 kind: EntryKind::default(),
3219 key: Some("mykey".to_string()),
3220 };
3221 let json = serde_json::to_string(&entry_with_key).unwrap();
3222 let deserialized: RegionEntry = serde_json::from_str(&json).unwrap();
3223 assert_eq!(deserialized.key.as_deref(), Some("mykey"));
3224 assert_eq!(deserialized.content, "some data");
3225 assert_eq!(deserialized.tokens, 10);
3226
3227 let entry_no_key = RegionEntry {
3229 content: "no key data".to_string(),
3230 tokens: 7,
3231 timestamp: 1234567890,
3232 metadata: None,
3233 kind: EntryKind::default(),
3234 key: None,
3235 };
3236 let json = serde_json::to_string(&entry_no_key).unwrap();
3237 assert!(!json.contains("\"key\""));
3238 let deserialized: RegionEntry = serde_json::from_str(&json).unwrap();
3239 assert!(deserialized.key.is_none());
3240 assert_eq!(deserialized.content, "no key data");
3241 }
3242
3243 #[test]
3244 fn test_hashmap_partial_eq() {
3245 let a = RegionKind::HashMap {
3246 max_entries: Some(5),
3247 };
3248 let b = RegionKind::HashMap {
3249 max_entries: Some(5),
3250 };
3251 let c = RegionKind::HashMap {
3252 max_entries: Some(10),
3253 };
3254 let d = RegionKind::HashMap { max_entries: None };
3255
3256 assert_eq!(a, b);
3257 assert_ne!(a, c);
3258 assert_ne!(a, d);
3259 assert_ne!(c, d);
3260 assert_ne!(a, RegionKind::Pinned);
3261 assert_ne!(a, RegionKind::Temporary);
3262 }
3263
3264 #[test]
3265 fn test_hashmap_cache_hint_returns_until_changed() {
3266 let kind = RegionKind::HashMap { max_entries: None };
3267 assert_eq!(kind.cache_hint(), crate::cache::CacheHint::UntilChanged);
3268
3269 let kind_with_max = RegionKind::HashMap {
3270 max_entries: Some(10),
3271 };
3272 assert_eq!(
3273 kind_with_max.cache_hint(),
3274 crate::cache::CacheHint::UntilChanged
3275 );
3276 }
3277
3278 #[test]
3281 fn test_remove_by_key_recomputes_taint_when_tracking_enabled() {
3282 let mut region = Region::new(
3285 "kv".to_string(),
3286 RegionKind::HashMap { max_entries: None },
3287 10_000,
3288 )
3289 .with_taint_tracking();
3290 region
3291 .upsert_by_key("k1", "value one".to_string(), 10)
3292 .unwrap();
3293 region
3294 .upsert_by_key("k2", "value two".to_string(), 10)
3295 .unwrap();
3296
3297 assert!(region.remove_by_key("k1"));
3298 assert!(!region.remove_by_key("missing"));
3299 assert_eq!(region.entry_count(), 1);
3300 assert_eq!(region.current_tokens, 10);
3301 }
3302
3303 #[test]
3304 fn test_evict_lru_entry_runs_taint_fixup() {
3305 let mut region = Region::new(
3309 "kv".to_string(),
3310 RegionKind::HashMap {
3311 max_entries: Some(1),
3312 },
3313 10_000,
3314 )
3315 .with_taint_tracking();
3316 region
3317 .upsert_by_key("first", "aaa".to_string(), 10)
3318 .unwrap();
3319 region
3320 .upsert_by_key("second", "bbb".to_string(), 10)
3321 .unwrap();
3322
3323 assert_eq!(region.entry_count(), 1);
3325 assert!(region.get_by_key("second").is_some());
3326 assert!(region.get_by_key("first").is_none());
3327 }
3328
3329 #[test]
3330 fn test_evict_lru_entry_on_empty_region_is_noop() {
3331 let mut region = Region::new(
3335 "kv".to_string(),
3336 RegionKind::HashMap {
3337 max_entries: Some(4),
3338 },
3339 1000,
3340 );
3341 assert_eq!(region.entry_count(), 0);
3342 region.evict_lru_entry();
3343 assert_eq!(region.entry_count(), 0);
3344 assert_eq!(region.current_tokens, 0);
3345 }
3346}