1use bevy_ecs::prelude::*;
4use leviath_core::Region;
5use serde::{Deserialize, Serialize};
6
7#[derive(Component, Debug, Clone)]
12pub struct AgentState {
13 pub agent_id: String,
15
16 pub current_stage: String,
18
19 pub iteration: usize,
21
22 pub status: AgentStatus,
24
25 pub spawned_children_ids: Vec<String>,
27
28 pub pending_wait: Option<String>,
30
31 pub accepts_messages: bool,
34}
35
36#[derive(Component, Debug, Clone)]
38pub struct ParentRef {
39 pub parent_entity: Entity,
41
42 pub parent_agent_id: String,
44
45 pub depth: usize,
47}
48
49#[derive(Component, Debug, Clone)]
51pub struct SubAgentChildren {
52 pub children: Vec<Entity>,
54
55 pub max_child_depth: usize,
57}
58
59#[derive(Component, Debug, Clone)]
69pub struct AwaitingInteraction;
70
71#[derive(Component, Debug, Clone, Copy, Default)]
85pub struct GateAutoApprove;
86
87#[derive(Component, Debug, Clone, Copy, Default)]
98pub struct InteractionAutoApprove;
99
100#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Hash)]
105pub enum AgentStatus {
106 Idle,
108
109 Active,
111
112 Waiting,
114
115 Paused,
119
120 Complete,
122
123 Error { message: String },
125
126 Cancelled,
128}
129
130impl AgentStatus {
131 pub fn label(&self) -> &'static str {
139 match self {
140 Self::Idle => "idle",
141 Self::Active => "active",
142 Self::Waiting => "waiting",
143 Self::Paused => "paused",
144 Self::Complete => "complete",
145 Self::Error { .. } => "error",
146 Self::Cancelled => "cancelled",
147 }
148 }
149}
150
151impl std::fmt::Display for AgentStatus {
152 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
156 match self {
157 Self::Error { message } => write!(f, "error: {message}"),
158 other => f.write_str(other.label()),
159 }
160 }
161}
162
163#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
177#[serde(rename_all = "snake_case", tag = "reason")]
178pub enum WaitReason {
179 ToolApproval,
181
182 UserPrompt,
185
186 TaintGate,
188
189 InteractionPoint,
191
192 FanOutWorkers {
194 outstanding: usize,
196 },
197
198 Children {
201 outstanding: usize,
203 },
204}
205
206impl WaitReason {
207 pub fn needs_a_person(&self) -> bool {
210 !matches!(self, Self::FanOutWorkers { .. } | Self::Children { .. })
211 }
212}
213
214impl std::fmt::Display for WaitReason {
215 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
216 match self {
217 Self::ToolApproval => f.write_str("tool approval"),
218 Self::UserPrompt => f.write_str("user prompt"),
219 Self::TaintGate => f.write_str("taint gate"),
220 Self::InteractionPoint => f.write_str("checkpoint"),
221 Self::FanOutWorkers { outstanding } => write!(f, "workers({outstanding})"),
222 Self::Children { outstanding } => write!(f, "children({outstanding})"),
223 }
224 }
225}
226
227#[derive(Debug, Clone)]
229pub struct EvictionResult {
230 pub tokens_freed: usize,
232 pub needs_compaction: Vec<String>,
234}
235
236#[derive(Component, Debug, Clone, Default)]
242pub struct InferenceConfig {
243 pub temperature: Option<f32>,
245 pub max_output_tokens: Option<usize>,
247 pub extra_params: serde_json::Map<String, serde_json::Value>,
252 pub batch_tool_hint: bool,
257 pub shell_hint: bool,
263 pub request_timeout_secs: Option<u64>,
268}
269
270#[derive(Component, Debug, Clone)]
275pub struct ToolResultRoutingComponent {
276 pub routing: leviath_core::ToolResultRouting,
278}
279
280#[derive(Debug, Clone)]
286pub struct AssembledContext {
287 pub system_blocks: Vec<leviath_providers::SystemBlock>,
289 pub messages: Vec<leviath_providers::Message>,
291}
292
293fn cache_hint_sort_priority(hint: leviath_core::CacheHint) -> u8 {
298 use leviath_core::CacheHint;
299 match hint {
300 CacheHint::Always => 0, CacheHint::SlidingPrefix { .. } => 1, CacheHint::UntilChanged => 2, CacheHint::Never => 3, }
305}
306
307#[derive(Component, Debug, Clone)]
309pub struct ContextWindow {
310 pub regions: Vec<Region>,
312
313 pub current_tokens: usize,
315
316 pub max_tokens: usize,
318
319 pub region_scripts: std::collections::HashMap<
327 String,
328 std::sync::Arc<leviath_scripting::region_hook::RegionScript>,
329 >,
330}
331
332impl ContextWindow {
333 pub fn new(max_tokens: usize) -> Self {
335 Self {
336 regions: Vec::new(),
337 current_tokens: 0,
338 max_tokens,
339 region_scripts: std::collections::HashMap::new(),
340 }
341 }
342
343 fn custom_script_for(
346 &self,
347 region_name: &str,
348 ) -> Option<std::sync::Arc<leviath_scripting::region_hook::RegionScript>> {
349 let region = self.get_region(region_name)?;
350 let leviath_core::RegionKind::Custom { script, .. } = ®ion.kind else {
351 return None;
352 };
353 self.region_scripts.get(script).cloned()
354 }
355
356 fn on_write_outcome(
364 &self,
365 region_name: &str,
366 content: String,
367 tokens: usize,
368 kind: &leviath_core::EntryKind,
369 ) -> Option<(String, usize)> {
370 let Some(script) = self.custom_script_for(region_name) else {
371 return Some((content, tokens));
372 };
373 if !script.has_on_write() {
374 return Some((content, tokens));
375 }
376 let region = self
378 .get_region(region_name)
379 .expect("custom_script_for resolved through this region");
380 match crate::custom_region::apply_on_write(&script, region, content, tokens, kind) {
381 crate::custom_region::OnWriteOutcome::Accept(content, tokens) => {
382 Some((content, tokens))
383 }
384 crate::custom_region::OnWriteOutcome::Drop => None,
385 }
386 }
387
388 fn try_custom_overflow(&mut self, region_name: &str, incoming_tokens: usize) -> bool {
394 let Some(script) = self.custom_script_for(region_name) else {
395 return false;
396 };
397 if !script.has_on_overflow() {
398 return false;
399 }
400 let region = self
401 .get_region_mut(region_name)
402 .expect("custom_script_for resolved through this region");
403 let needed = (region.current_tokens + incoming_tokens).saturating_sub(region.max_tokens);
404 let freed = crate::custom_region::apply_overflow(&script, region, needed);
405 self.current_tokens = self.calculate_tokens();
406 freed >= needed && needed > 0
407 }
408
409 pub fn get_region(&self, name: &str) -> Option<&Region> {
411 self.regions.iter().find(|r| r.name == name)
412 }
413
414 pub fn get_region_mut(&mut self, name: &str) -> Option<&mut Region> {
416 self.regions.iter_mut().find(|r| r.name == name)
417 }
418
419 pub fn add_region(&mut self, region: Region) {
421 self.regions.push(region);
422 self.current_tokens = self.calculate_tokens();
423 }
424
425 pub fn add_to_region(
427 &mut self,
428 region_name: &str,
429 content: String,
430 tokens: usize,
431 ) -> leviath_core::Result<()> {
432 let Some((content, tokens)) =
433 self.on_write_outcome(region_name, content, tokens, &leviath_core::EntryKind::Text)
434 else {
435 return Ok(()); };
437 self.write_to_region(region_name, tokens, &mut |region, tokens| {
438 region.add_entry(content.clone(), tokens)
439 })
440 }
441
442 pub fn replace_region(&mut self, region_name: &str, content: String, tokens: usize) -> bool {
447 let Some((content, tokens)) =
450 self.on_write_outcome(region_name, content, tokens, &leviath_core::EntryKind::Text)
451 else {
452 return self.get_region(region_name).is_some();
454 };
455 if let Some(region) = self.get_region_mut(region_name) {
456 region.clear();
457 let _ = region.add_entry(content, tokens);
458 self.current_tokens = self.calculate_tokens();
459 true
460 } else {
461 false
462 }
463 }
464
465 pub fn add_typed_entry(
471 &mut self,
472 region_name: &str,
473 kind: leviath_core::EntryKind,
474 content: String,
475 tokens: usize,
476 ) -> leviath_core::Result<()> {
477 let Some((content, tokens)) = self.on_write_outcome(region_name, content, tokens, &kind)
478 else {
479 return Ok(());
480 };
481 self.write_to_region(region_name, tokens, &mut |region, tokens| {
482 region.add_typed_entry(content.clone(), tokens, kind.clone())
483 })
484 }
485
486 fn write_to_region(
491 &mut self,
492 region_name: &str,
493 tokens: usize,
494 insert: &mut dyn FnMut(&mut Region, usize) -> leviath_core::Result<()>,
495 ) -> leviath_core::Result<()> {
496 if self.get_region(region_name).is_none() {
497 return Err(leviath_core::Error::RegionNotFound(region_name.to_string()));
498 }
499 let first = {
500 let region = self.get_region_mut(region_name).expect("checked above");
501 insert(region, tokens)
502 };
503 match first {
504 Ok(()) => {
505 self.current_tokens = self.calculate_tokens();
506 Ok(())
507 }
508 Err(leviath_core::Error::TokenBudgetExceeded { .. })
509 if self.try_custom_overflow(region_name, tokens) =>
510 {
511 let region = self.get_region_mut(region_name).expect("checked above");
512 let retried = insert(region, tokens);
513 self.current_tokens = self.calculate_tokens();
514 retried
515 }
516 Err(e) => Err(e),
517 }
518 }
519
520 pub fn calculate_tokens(&self) -> usize {
522 self.regions.iter().map(|r| r.current_tokens).sum()
523 }
524
525 pub fn needs_eviction(&self, threshold: f32) -> bool {
527 let usage_ratio = self.current_tokens as f32 / self.max_tokens as f32;
528 usage_ratio >= threshold
529 }
530
531 pub fn try_evict(&mut self, target_free_tokens: usize) -> leviath_core::Result<EvictionResult> {
537 use leviath_core::RegionKind;
538
539 let initial_tokens = self.current_tokens;
540
541 let has_evictable = self.regions.iter().any(|r| {
543 matches!(
544 r.kind,
545 RegionKind::Clearable
546 | RegionKind::Temporary
547 | RegionKind::Custom {
548 persistent: false,
549 ..
550 }
551 )
552 });
553
554 if !has_evictable {
555 tracing::warn!(
556 "Context window has no Clearable or Temporary regions. \
557 This may be intentional, but usually indicates a configuration error."
558 );
559 }
560
561 for region in &mut self.regions {
563 if matches!(region.kind, RegionKind::Clearable) && !region.content.is_empty() {
564 let freed = region.current_tokens;
565 region.clear();
566 self.current_tokens -= freed;
567 tracing::debug!(
568 region = %region.name,
569 tokens_freed = freed,
570 "Cleared Clearable region (all-or-nothing)"
571 );
572
573 if self.max_tokens.saturating_sub(self.current_tokens) >= target_free_tokens {
574 return Ok(EvictionResult {
575 tokens_freed: initial_tokens - self.current_tokens,
576 needs_compaction: Vec::new(),
577 });
578 }
579 }
580 }
581
582 let mut custom_freed = 0usize;
589 for i in 0..self.regions.len() {
590 let needed = target_free_tokens
591 .saturating_sub(self.max_tokens.saturating_sub(self.current_tokens));
592 if needed == 0 {
593 break;
594 }
595 let region = &self.regions[i];
596 if !matches!(
597 region.kind,
598 RegionKind::Custom {
599 persistent: false,
600 ..
601 }
602 ) || region.content.is_empty()
603 {
604 continue;
605 }
606 let Some(script) = self.custom_script_for(®ion.name.clone()) else {
607 continue;
608 };
609 if !script.has_on_overflow() {
610 continue;
611 }
612 let freed = crate::custom_region::apply_overflow(&script, &mut self.regions[i], needed);
613 self.current_tokens = self.current_tokens.saturating_sub(freed);
614 custom_freed += freed;
615 if freed > 0 {
616 tracing::debug!(
617 region = %self.regions[i].name,
618 tokens_freed = freed,
619 "custom region's on_overflow chose its own evictions"
620 );
621 }
622 }
623 if custom_freed > 0
629 && self.max_tokens.saturating_sub(self.current_tokens) >= target_free_tokens
630 {
631 return Ok(EvictionResult {
632 tokens_freed: initial_tokens - self.current_tokens,
633 needs_compaction: Vec::new(),
634 });
635 }
636
637 loop {
642 let mut evicted_any = false;
643
644 for region in &mut self.regions {
645 if matches!(
646 region.kind,
647 RegionKind::Temporary
648 | RegionKind::Custom {
649 persistent: false,
650 ..
651 }
652 ) && let Some(entry) = region.remove_oldest()
653 {
654 let freed = entry.tokens;
655 self.current_tokens -= freed;
656 evicted_any = true;
657
658 tracing::debug!(
659 region = %region.name,
660 tokens_freed = freed,
661 "Evicted temporary region entry (oldest first)"
662 );
663
664 if self.max_tokens.saturating_sub(self.current_tokens) >= target_free_tokens {
665 return Ok(EvictionResult {
666 tokens_freed: initial_tokens - self.current_tokens,
667 needs_compaction: Vec::new(),
668 });
669 }
670 }
671 }
672
673 if !evicted_any {
674 break;
675 }
676 }
677
678 let mut needs_compaction = Vec::new();
680 if self.max_tokens.saturating_sub(self.current_tokens) < target_free_tokens {
681 for region in &self.regions {
682 if region.needs_compaction() {
683 needs_compaction.push(region.name.clone());
684 }
685 }
686 }
687
688 let pinned_tokens: usize = self
693 .regions
694 .iter()
695 .filter(|r| {
696 matches!(
697 r.kind,
698 RegionKind::Pinned
699 | RegionKind::CompactHistory { .. }
700 | RegionKind::Custom {
701 persistent: true,
702 ..
703 }
704 )
705 })
706 .map(|r| r.current_tokens)
707 .sum();
708
709 if pinned_tokens > self.max_tokens {
710 return Err(leviath_core::Error::PinnedRegionsOverBudget {
711 pinned_tokens,
712 total_budget: self.max_tokens,
713 });
714 }
715
716 Ok(EvictionResult {
717 tokens_freed: initial_tokens - self.current_tokens,
718 needs_compaction,
719 })
720 }
721
722 pub fn assemble(&self) -> AssembledContext {
730 self.assemble_with_meta(&crate::custom_region::AssembleMeta::default())
731 }
732
733 pub fn assemble_with_meta(
738 &self,
739 meta: &crate::custom_region::AssembleMeta,
740 ) -> AssembledContext {
741 use leviath_core::{CacheHint, EntryKind};
742
743 let mut system_blocks = Vec::new();
744 let mut messages: Vec<leviath_providers::Message> = Vec::new();
745
746 for region in &self.regions {
747 let is_custom = matches!(region.kind, leviath_core::RegionKind::Custom { .. });
750 if region.content.is_empty() && !is_custom {
751 continue;
752 }
753
754 match ®ion.kind {
755 leviath_core::RegionKind::Pinned => {
757 let text = region
758 .content
759 .iter()
760 .map(|e| e.content.as_str())
761 .collect::<Vec<_>>()
762 .join("\n\n");
763 system_blocks.push(leviath_providers::SystemBlock {
764 text,
765 cache_hint: CacheHint::Always,
766 });
767 }
768 leviath_core::RegionKind::CompactHistory { .. } => {
769 let text = region
770 .content
771 .iter()
772 .map(|e| e.content.as_str())
773 .collect::<Vec<_>>()
774 .join("\n\n");
775 system_blocks.push(leviath_providers::SystemBlock {
776 text,
777 cache_hint: CacheHint::Always,
778 });
779 }
780
781 leviath_core::RegionKind::SlidingWindow { .. } => {
786 let mut pending_tool_results: Vec<leviath_providers::ContentBlock> = Vec::new();
787
788 for entry in ®ion.content {
789 if !matches!(entry.kind, EntryKind::ToolResult { .. })
791 && !pending_tool_results.is_empty()
792 {
793 messages.push(leviath_providers::Message {
794 role: "user".to_string(),
795 content: leviath_providers::MessageContent::Blocks(std::mem::take(
796 &mut pending_tool_results,
797 )),
798 cache_breakpoint: false,
799 });
800 }
801
802 match &entry.kind {
803 EntryKind::UserMessage => {
804 messages.push(leviath_providers::Message {
805 role: "user".to_string(),
806 content: entry.content.clone().into(),
807 cache_breakpoint: false,
808 });
809 }
810 EntryKind::AssistantTurn { tool_calls } => {
811 if tool_calls.is_empty() {
812 messages.push(leviath_providers::Message {
813 role: "assistant".to_string(),
814 content: entry.content.clone().into(),
815 cache_breakpoint: false,
816 });
817 } else {
818 let mut blocks = Vec::new();
819 if !entry.content.is_empty() {
820 blocks.push(leviath_providers::ContentBlock::Text {
821 text: entry.content.clone(),
822 });
823 }
824 for tc in tool_calls {
825 blocks.push(leviath_providers::ContentBlock::ToolUse {
826 id: tc.id.clone(),
827 name: tc.name.clone(),
828 input: tc.arguments.clone(),
829 thought_signature: tc.thought_signature.clone(),
830 });
831 }
832 messages.push(leviath_providers::Message {
833 role: "assistant".to_string(),
834 content: leviath_providers::MessageContent::Blocks(blocks),
835 cache_breakpoint: false,
836 });
837 }
838 }
839 EntryKind::ToolResult {
840 tool_call_id,
841 is_error,
842 ..
843 } => {
844 pending_tool_results.push(
846 leviath_providers::ContentBlock::ToolResult {
847 tool_use_id: tool_call_id.clone(),
848 content: entry.content.clone(),
849 is_error: *is_error,
850 },
851 );
852 }
853 EntryKind::Text => {
854 let trimmed = entry.content.trim();
855 if let Some(rest) = trimmed.strip_prefix("Assistant: ") {
856 messages.push(leviath_providers::Message {
857 role: "assistant".to_string(),
858 content: rest.to_string().into(),
859 cache_breakpoint: false,
860 });
861 } else if let Some(rest) = trimmed.strip_prefix("User: ") {
862 messages.push(leviath_providers::Message {
863 role: "user".to_string(),
864 content: rest.to_string().into(),
865 cache_breakpoint: false,
866 });
867 } else {
868 messages.push(leviath_providers::Message {
869 role: "user".to_string(),
870 content: entry.content.clone().into(),
871 cache_breakpoint: false,
872 });
873 }
874 }
875 }
876 }
877
878 if !pending_tool_results.is_empty() {
880 messages.push(leviath_providers::Message {
881 role: "user".to_string(),
882 content: leviath_providers::MessageContent::Blocks(std::mem::take(
883 &mut pending_tool_results,
884 )),
885 cache_breakpoint: false,
886 });
887 }
888 }
889
890 leviath_core::RegionKind::Compacting { .. } => {
892 let text = region
893 .content
894 .iter()
895 .map(|e| e.content.as_str())
896 .collect::<Vec<_>>()
897 .join("\n\n");
898 system_blocks.push(leviath_providers::SystemBlock {
899 text: format!("[{}]:\n{}", region.name, text),
900 cache_hint: CacheHint::UntilChanged,
901 });
902 }
903 leviath_core::RegionKind::Temporary => {
904 let text = region
905 .content
906 .iter()
907 .map(|e| e.content.as_str())
908 .collect::<Vec<_>>()
909 .join("\n\n");
910 system_blocks.push(leviath_providers::SystemBlock {
911 text: format!("[{}]:\n{}", region.name, text),
912 cache_hint: CacheHint::Never,
913 });
914 }
915 leviath_core::RegionKind::Clearable => {
916 let text = region
917 .content
918 .iter()
919 .map(|e| e.content.as_str())
920 .collect::<Vec<_>>()
921 .join("\n\n");
922 system_blocks.push(leviath_providers::SystemBlock {
923 text: format!("[{}]:\n{}", region.name, text),
924 cache_hint: CacheHint::Never,
925 });
926 }
927
928 leviath_core::RegionKind::Custom { script, persistent } => {
933 crate::custom_region::render_custom_region(
934 region,
935 self.region_scripts.get(script),
936 *persistent,
937 meta,
938 self.current_tokens,
939 self.max_tokens,
940 &mut system_blocks,
941 &mut messages,
942 );
943 }
944
945 leviath_core::RegionKind::HashMap { .. } => {
947 let text = region
948 .content
949 .iter()
950 .map(|e| {
951 if let Some(key) = &e.key {
952 format!("### [{}]\n{}", key, e.content)
953 } else {
954 e.content.clone()
955 }
956 })
957 .collect::<Vec<_>>()
958 .join("\n\n");
959 system_blocks.push(leviath_providers::SystemBlock {
960 text: format!("[{}]:\n{}", region.name, text),
961 cache_hint: CacheHint::UntilChanged,
962 });
963 }
964 }
965 }
966
967 system_blocks.sort_by_key(|block| cache_hint_sort_priority(block.cache_hint));
974
975 let mut tool_use_ids = std::collections::HashSet::new();
980 let mut tool_result_ids = std::collections::HashSet::new();
981
982 for msg in &messages {
983 if let leviath_providers::MessageContent::Blocks(blocks) = &msg.content {
984 for block in blocks {
985 match block {
986 leviath_providers::ContentBlock::ToolUse { id, .. } => {
987 tool_use_ids.insert(id.clone());
988 }
989 leviath_providers::ContentBlock::ToolResult { tool_use_id, .. } => {
990 tool_result_ids.insert(tool_use_id.clone());
991 }
992 _ => {}
993 }
994 }
995 }
996 }
997
998 let orphaned_tool_uses: std::collections::HashSet<_> =
999 tool_use_ids.difference(&tool_result_ids).cloned().collect();
1000 let orphaned_tool_results: std::collections::HashSet<_> =
1001 tool_result_ids.difference(&tool_use_ids).cloned().collect();
1002
1003 if !orphaned_tool_uses.is_empty() || !orphaned_tool_results.is_empty() {
1004 tracing::warn!(
1005 orphaned_tool_uses = orphaned_tool_uses.len(),
1006 orphaned_tool_results = orphaned_tool_results.len(),
1007 "Stripping orphaned tool_use/tool_result blocks from assembled context"
1008 );
1009
1010 messages = messages
1011 .into_iter()
1012 .filter_map(|msg| {
1013 if let leviath_providers::MessageContent::Blocks(blocks) = &msg.content {
1014 let filtered: Vec<_> = blocks
1015 .iter()
1016 .filter(|block| match block {
1017 leviath_providers::ContentBlock::ToolUse { id, .. } => {
1018 !orphaned_tool_uses.contains(id)
1019 }
1020 leviath_providers::ContentBlock::ToolResult {
1021 tool_use_id, ..
1022 } => !orphaned_tool_results.contains(tool_use_id),
1023 _ => true,
1024 })
1025 .cloned()
1026 .collect();
1027
1028 if filtered.is_empty() {
1029 None
1031 } else {
1032 Some(leviath_providers::Message {
1033 role: msg.role.clone(),
1034 content: leviath_providers::MessageContent::Blocks(filtered),
1035 cache_breakpoint: msg.cache_breakpoint,
1036 })
1037 }
1038 } else {
1039 Some(msg)
1040 }
1041 })
1042 .collect();
1043 }
1044
1045 if messages.len() >= 5 {
1058 let bp_idx = messages.len() - 4;
1059 messages[bp_idx].cache_breakpoint = true;
1060 } else if messages.len() >= 2 {
1061 messages[0].cache_breakpoint = true;
1063 }
1064
1065 if !messages.iter().any(|m| m.role == "user") {
1067 messages.push(leviath_providers::Message {
1068 role: "user".to_string(),
1069 content: "Begin.".into(),
1070 cache_breakpoint: false,
1071 });
1072 }
1073
1074 if messages.last().map(|m| m.role.as_str()) == Some("assistant") {
1081 messages.push(leviath_providers::Message {
1082 role: "user".to_string(),
1083 content: "Continue.".into(),
1084 cache_breakpoint: false,
1085 });
1086 }
1087
1088 AssembledContext {
1089 system_blocks,
1090 messages,
1091 }
1092 }
1093
1094 pub fn enable_taint_tracking(&mut self) {
1096 for region in &mut self.regions {
1097 region.enable_taint_tracking();
1098 }
1099 }
1100
1101 pub fn add_tainted_to_region(
1103 &mut self,
1104 region_name: &str,
1105 content: String,
1106 tokens: usize,
1107 taint_level: leviath_core::TaintLevel,
1108 ) -> leviath_core::Result<()> {
1109 let Some((content, tokens)) =
1110 self.on_write_outcome(region_name, content, tokens, &leviath_core::EntryKind::Text)
1111 else {
1112 return Ok(());
1113 };
1114 self.write_to_region(region_name, tokens, &mut |region, tokens| {
1115 region.add_tainted_entry(content.clone(), tokens, taint_level)
1116 })
1117 }
1118
1119 pub fn add_typed_tainted_to_region(
1126 &mut self,
1127 region_name: &str,
1128 kind: leviath_core::EntryKind,
1129 content: String,
1130 tokens: usize,
1131 taint_level: leviath_core::TaintLevel,
1132 ) -> leviath_core::Result<()> {
1133 let Some((content, tokens)) = self.on_write_outcome(region_name, content, tokens, &kind)
1134 else {
1135 return Ok(());
1136 };
1137 self.write_to_region(region_name, tokens, &mut |region, tokens| {
1138 region.add_typed_tainted_entry(content.clone(), tokens, kind.clone(), taint_level)
1139 })
1140 }
1141
1142 pub fn overall_taint(&self) -> Option<leviath_core::TaintLevel> {
1145 let mut max_taint = None;
1146 for region in &self.regions {
1147 if let Some(level) = region.taint_level() {
1148 max_taint = Some(match max_taint {
1149 Some(current) => level.max(current),
1150 None => level,
1151 });
1152 }
1153 }
1154 max_taint
1155 }
1156
1157 pub fn taint_summary(&self) -> Vec<(String, leviath_core::TaintLevel)> {
1159 self.regions
1160 .iter()
1161 .filter_map(|r| r.taint_level().map(|t| (r.name.clone(), t)))
1162 .collect()
1163 }
1164}
1165
1166#[derive(Component, Debug, Clone)]
1171pub struct InferenceResult {
1172 pub response: String,
1174
1175 pub tool_calls: Vec<ToolCall>,
1177
1178 pub tokens_used: usize,
1180
1181 pub timestamp: i64,
1183}
1184
1185#[derive(Debug, Clone, Serialize, Deserialize)]
1187pub struct ToolCall {
1188 pub tool_id: String,
1190
1191 pub name: String,
1193
1194 pub arguments: serde_json::Value,
1196 #[serde(default, skip_serializing_if = "Option::is_none")]
1199 pub thought_signature: Option<String>,
1200}
1201
1202#[derive(Debug, Clone)]
1204pub struct AgentMessage {
1205 pub agent_id: String,
1207 pub content: String,
1209 pub target_region: Option<String>,
1211}
1212
1213#[derive(Component, Debug, Clone)]
1215pub struct MessageInbox {
1216 pub messages: Vec<AgentMessage>,
1218}
1219
1220impl MessageInbox {
1221 pub fn new() -> Self {
1223 Self {
1224 messages: Vec::new(),
1225 }
1226 }
1227
1228 pub fn push(&mut self, msg: AgentMessage) {
1232 self.messages.push(msg);
1233 }
1234
1235 pub fn drain_all(&mut self) -> Vec<AgentMessage> {
1237 std::mem::take(&mut self.messages)
1238 }
1239}
1240
1241impl Default for MessageInbox {
1242 fn default() -> Self {
1243 Self::new()
1244 }
1245}
1246
1247#[cfg(test)]
1248mod tests {
1249 use super::*;
1250 use crate::test_support::with_tracing;
1251 use leviath_core::{EvictionStrategy, Region, RegionKind};
1252
1253 #[test]
1254 fn test_context_window_creation() {
1255 let window = ContextWindow::new(10000);
1256 assert_eq!(window.max_tokens, 10000);
1257 assert_eq!(window.current_tokens, 0);
1258 }
1259
1260 #[test]
1261 fn test_needs_eviction() {
1262 let mut window = ContextWindow::new(10000);
1263 window.current_tokens = 9500;
1264 assert!(window.needs_eviction(0.9));
1265
1266 window.current_tokens = 5000;
1267 assert!(!window.needs_eviction(0.9));
1268 }
1269
1270 #[test]
1271 fn test_add_region() {
1272 let mut window = ContextWindow::new(10000);
1273 let region = Region::new("test".to_string(), RegionKind::Pinned, 1000);
1274 window.add_region(region);
1275 assert_eq!(window.regions.len(), 1);
1276 }
1277
1278 #[test]
1279 fn replace_region_overwrites_existing_and_reports_missing() {
1280 let mut window = ContextWindow::new(10000);
1281 let mut region = Region::new("plan".to_string(), RegionKind::Pinned, 6000);
1282 region.add_entry("old plan".to_string(), 3).unwrap();
1283 window.add_region(region);
1284
1285 assert!(window.replace_region("plan", "new plan".to_string(), 3));
1287 let plan = window.get_region("plan").unwrap();
1288 assert_eq!(plan.content.len(), 1);
1289 assert_eq!(plan.content[0].content, "new plan");
1290
1291 assert!(!window.replace_region("nope", "x".to_string(), 1));
1293 }
1294
1295 #[test]
1296 fn test_clearable_eviction() {
1297 let mut window = ContextWindow::new(10000);
1298 let mut region = Region::new("scratch".to_string(), RegionKind::Clearable, 5000);
1299 region
1300 .add_entry("test content 1".to_string(), 1000)
1301 .unwrap();
1302 region
1303 .add_entry("test content 2".to_string(), 1000)
1304 .unwrap();
1305 window.add_region(region);
1306
1307 assert_eq!(window.current_tokens, 2000);
1308
1309 let result = with_tracing(|| window.try_evict(1000)).unwrap();
1311 assert_eq!(result.tokens_freed, 2000);
1312 assert!(result.needs_compaction.is_empty());
1313 assert_eq!(window.current_tokens, 0);
1314 }
1315
1316 #[test]
1317 fn test_temporary_eviction_oldest_first() {
1318 let mut window = ContextWindow::new(10000);
1319 let mut region = Region::new("temp".to_string(), RegionKind::Temporary, 5000);
1320 region.add_entry("old content".to_string(), 1000).unwrap();
1321 region
1322 .add_entry("middle content".to_string(), 1000)
1323 .unwrap();
1324 region.add_entry("new content".to_string(), 1000).unwrap();
1325 window.add_region(region);
1326
1327 assert_eq!(window.current_tokens, 3000);
1328
1329 let result = with_tracing(|| window.try_evict(500)).unwrap();
1331 assert!(result.tokens_freed >= 1000); assert!(result.needs_compaction.is_empty());
1333
1334 let region = window.get_region("temp").unwrap();
1336 assert_eq!(region.content.len(), 2);
1337 assert_eq!(region.content[0].content, "middle content");
1338 }
1339
1340 fn assert_sliding_window_unreduced(initial_count: usize, after_count: usize) {
1341 assert_eq!(
1342 initial_count, after_count,
1343 "SlidingWindow should never be reduced during eviction"
1344 );
1345 }
1346
1347 #[test]
1348 fn test_sliding_window_never_reduced() {
1349 let mut window = ContextWindow::new(10000);
1350 let mut region = Region::new(
1351 "conversation".to_string(),
1352 RegionKind::SlidingWindow {
1353 max_items: 5,
1354 eviction_strategy: EvictionStrategy::PerItem,
1355 },
1356 5000,
1357 );
1358 region.add_entry("msg 1".to_string(), 1000).unwrap();
1359 region.add_entry("msg 2".to_string(), 1000).unwrap();
1360 region.add_entry("msg 3".to_string(), 1000).unwrap();
1361 window.add_region(region);
1362
1363 let initial_count = window.get_region("conversation").unwrap().content.len();
1364
1365 window.try_evict(1000).ok();
1367
1368 let after_count = window.get_region("conversation").unwrap().content.len();
1369 assert_sliding_window_unreduced(initial_count, after_count);
1370 }
1371
1372 #[test]
1373 #[should_panic(expected = "SlidingWindow should never be reduced during eviction")]
1374 fn test_sliding_window_never_reduced_panics_on_mismatch() {
1375 assert_sliding_window_unreduced(3, 2);
1376 }
1377
1378 fn assert_pinned_unevicted(initial_tokens: usize, after_tokens: usize) {
1379 assert_eq!(
1380 initial_tokens, after_tokens,
1381 "Pinned region should never be evicted"
1382 );
1383 }
1384
1385 #[test]
1386 fn test_pinned_never_touched() {
1387 let mut window = ContextWindow::new(10000);
1388 let mut region = Region::new("architecture".to_string(), RegionKind::Pinned, 3000);
1389 region
1390 .add_entry("architecture diagram".to_string(), 2000)
1391 .unwrap();
1392 window.add_region(region);
1393
1394 let initial_tokens = window.get_region("architecture").unwrap().current_tokens;
1395
1396 window.try_evict(1000).ok();
1398
1399 let after_tokens = window.get_region("architecture").unwrap().current_tokens;
1400 assert_pinned_unevicted(initial_tokens, after_tokens);
1401 }
1402
1403 #[test]
1404 #[should_panic(expected = "Pinned region should never be evicted")]
1405 fn test_pinned_never_touched_panics_on_mismatch() {
1406 assert_pinned_unevicted(2000, 1000);
1407 }
1408
1409 #[test]
1410 fn test_eviction_cascade_order() {
1411 let mut window = ContextWindow::new(10000);
1412
1413 let mut clearable = Region::new("scratch".to_string(), RegionKind::Clearable, 2000);
1415 clearable
1416 .add_entry("scratch data".to_string(), 1000)
1417 .unwrap();
1418 window.add_region(clearable);
1419
1420 let mut temporary = Region::new("temp".to_string(), RegionKind::Temporary, 3000);
1422 temporary
1423 .add_entry("temp data 1".to_string(), 1000)
1424 .unwrap();
1425 temporary
1426 .add_entry("temp data 2".to_string(), 1000)
1427 .unwrap();
1428 window.add_region(temporary);
1429
1430 assert_eq!(window.current_tokens, 3000);
1431
1432 window.try_evict(500).unwrap();
1434
1435 assert_eq!(window.get_region("scratch").unwrap().current_tokens, 0);
1437
1438 assert!(window.get_region("temp").unwrap().current_tokens > 0);
1440 }
1441
1442 #[test]
1443 fn test_message_inbox() {
1444 let mut inbox = MessageInbox::new();
1445 assert!(inbox.messages.is_empty());
1446
1447 inbox.push(AgentMessage {
1448 agent_id: "agent-1".to_string(),
1449 content: "hello".to_string(),
1450 target_region: None,
1451 });
1452 assert_eq!(inbox.messages.len(), 1);
1453
1454 let drained = inbox.drain_all();
1455 assert_eq!(drained.len(), 1);
1456 assert!(inbox.messages.is_empty());
1457 }
1458
1459 #[test]
1460 fn message_inbox_preserves_fifo_order() {
1461 let mut inbox = MessageInbox::new();
1462 for content in ["first", "second", "third"] {
1463 inbox.push(AgentMessage {
1464 agent_id: "a".to_string(),
1465 content: content.to_string(),
1466 target_region: None,
1467 });
1468 }
1469
1470 let msgs = inbox.drain_all();
1471 assert_eq!(msgs[0].content, "first");
1472 assert_eq!(msgs[1].content, "second");
1473 assert_eq!(msgs[2].content, "third");
1474 }
1475
1476 #[test]
1477 fn test_eviction_result_identifies_compaction_regions() {
1478 let mut window = ContextWindow::new(1000);
1480 let mut compacting = Region::new(
1482 "impl".to_string(),
1483 RegionKind::Compacting {
1484 threshold_tokens: 500,
1485 },
1486 900,
1487 );
1488 compacting
1489 .add_entry("lots of content".to_string(), 600)
1490 .unwrap();
1491 window.add_region(compacting);
1492
1493 assert_eq!(window.current_tokens, 600);
1494
1495 let result = window.try_evict(500).unwrap();
1497 assert_eq!(result.tokens_freed, 0);
1498 assert_eq!(result.needs_compaction, vec!["impl".to_string()]);
1499 }
1500
1501 #[test]
1502 fn test_try_evict_returns_needs_compaction_when_full() {
1503 let mut window = ContextWindow::new(1200);
1504
1505 let mut compacting = Region::new(
1507 "analysis".to_string(),
1508 RegionKind::Compacting {
1509 threshold_tokens: 800,
1510 },
1511 1100,
1512 );
1513 compacting.add_entry("data 1".to_string(), 500).unwrap();
1514 compacting.add_entry("data 2".to_string(), 500).unwrap();
1515 window.add_region(compacting);
1516
1517 let result = window.try_evict(500).unwrap();
1519 assert_eq!(result.tokens_freed, 0);
1520 assert!(result.needs_compaction.contains(&"analysis".to_string()));
1521 }
1522
1523 #[test]
1524 fn test_try_evict_errors_when_pinned_regions_exceed_budget() {
1525 let mut window = ContextWindow::new(1000);
1529 let mut pinned = Region::new("architecture".to_string(), RegionKind::Pinned, 2000);
1530 pinned
1531 .add_entry("huge pinned doc".to_string(), 1500)
1532 .unwrap();
1533 window.add_region(pinned);
1534
1535 let result = window.try_evict(100);
1536 assert!(result.is_err());
1537 let err_str = result.unwrap_err().to_string();
1538 assert!(err_str.contains("Pinned regions"));
1539 }
1540
1541 #[test]
1542 fn test_clearable_eviction_continues_past_insufficient_first_region() {
1543 let mut window = ContextWindow::new(2000);
1548
1549 let mut region_a = Region::new("a".to_string(), RegionKind::Clearable, 1000);
1550 region_a.add_entry("small".to_string(), 500).unwrap();
1551 window.add_region(region_a);
1552
1553 let mut region_b = Region::new("b".to_string(), RegionKind::Clearable, 1000);
1554 region_b.add_entry("large".to_string(), 1000).unwrap();
1555 window.add_region(region_b);
1556
1557 assert_eq!(window.current_tokens, 1500);
1558
1559 let result = with_tracing(|| window.try_evict(1400)).unwrap();
1563 assert_eq!(result.tokens_freed, 1500);
1564 assert_eq!(window.current_tokens, 0);
1565 assert_eq!(window.get_region("a").unwrap().current_tokens, 0);
1566 assert_eq!(window.get_region("b").unwrap().current_tokens, 0);
1567 }
1568
1569 #[test]
1570 fn test_agent_status_cancelled() {
1571 assert_eq!(AgentStatus::Cancelled, AgentStatus::Cancelled);
1572 }
1573
1574 #[test]
1575 fn test_parent_ref_component() {
1576 let parent_ref = super::ParentRef {
1577 parent_entity: Entity::from_raw_u32(42)
1578 .expect("a small literal index is always a valid entity id"),
1579 parent_agent_id: "coder-01".to_string(),
1580 depth: 1,
1581 };
1582 assert_eq!(parent_ref.parent_agent_id, "coder-01");
1583 assert_eq!(parent_ref.depth, 1);
1584 }
1585
1586 #[test]
1587 fn test_children_component() {
1588 let children = super::SubAgentChildren {
1589 children: vec![
1590 Entity::from_raw_u32(1).expect("a small literal index is always a valid entity id"),
1591 Entity::from_raw_u32(2).expect("a small literal index is always a valid entity id"),
1592 ],
1593 max_child_depth: 3,
1594 };
1595 assert_eq!(children.children.len(), 2);
1596 assert_eq!(children.max_child_depth, 3);
1597 }
1598
1599 #[test]
1600 fn test_agent_state_with_children_fields() {
1601 let state = AgentState {
1602 agent_id: "test-01".to_string(),
1603 current_stage: "analyze".to_string(),
1604 iteration: 0,
1605 status: AgentStatus::Active,
1606 spawned_children_ids: vec!["child-01".to_string(), "child-02".to_string()],
1607 pending_wait: Some("child-01".to_string()),
1608 accepts_messages: true,
1609 };
1610 assert_eq!(state.spawned_children_ids.len(), 2);
1611 assert_eq!(state.pending_wait, Some("child-01".to_string()));
1612 }
1613
1614 #[test]
1617 fn test_context_window_get_region() {
1618 let mut window = ContextWindow::new(10000);
1619 let region = Region::new("test".to_string(), RegionKind::Pinned, 1000);
1620 window.add_region(region);
1621
1622 assert!(window.get_region("test").is_some());
1623 assert!(window.get_region("nonexistent").is_none());
1624 }
1625
1626 #[test]
1627 fn test_context_window_get_region_mut() {
1628 let mut window = ContextWindow::new(10000);
1629 let region = Region::new("test".to_string(), RegionKind::Temporary, 1000);
1630 window.add_region(region);
1631
1632 let region = window.get_region_mut("test").unwrap();
1633 region.add_entry("new content".to_string(), 50).unwrap();
1634 assert_eq!(region.content.len(), 1);
1635
1636 assert!(window.get_region_mut("nonexistent").is_none());
1637 }
1638
1639 #[test]
1640 fn test_context_window_add_to_region_success() {
1641 let mut window = ContextWindow::new(10000);
1642 let region = Region::new("conv".to_string(), RegionKind::Temporary, 5000);
1643 window.add_region(region);
1644
1645 let result = window.add_to_region("conv", "Hello".to_string(), 10);
1646 assert!(result.is_ok());
1647 assert_eq!(window.current_tokens, 10);
1648 }
1649
1650 #[test]
1651 fn test_context_window_add_to_region_not_found() {
1652 let mut window = ContextWindow::new(10000);
1653 let result = window.add_to_region("nonexistent", "Hello".to_string(), 10);
1654 assert!(result.is_err());
1655 }
1656
1657 #[test]
1658 fn test_context_window_calculate_tokens() {
1659 let mut window = ContextWindow::new(10000);
1660 let mut r1 = Region::new("a".to_string(), RegionKind::Pinned, 5000);
1661 r1.add_entry("x".to_string(), 100).unwrap();
1662 let mut r2 = Region::new("b".to_string(), RegionKind::Temporary, 5000);
1663 r2.add_entry("y".to_string(), 200).unwrap();
1664 window.add_region(r1);
1665 window.add_region(r2);
1666
1667 assert_eq!(window.calculate_tokens(), 300);
1668 }
1669
1670 #[test]
1671 fn test_context_window_needs_eviction_boundary() {
1672 let mut window = ContextWindow::new(100);
1673 window.current_tokens = 90;
1675 assert!(window.needs_eviction(0.9));
1676
1677 window.current_tokens = 89;
1679 assert!(!window.needs_eviction(0.9));
1680 }
1681
1682 #[test]
1683 fn test_eviction_result_default_fields() {
1684 let result = EvictionResult {
1685 tokens_freed: 0,
1686 needs_compaction: Vec::new(),
1687 };
1688 assert_eq!(result.tokens_freed, 0);
1689 assert!(result.needs_compaction.is_empty());
1690 }
1691
1692 #[test]
1693 fn test_message_inbox_default() {
1694 let inbox = MessageInbox::default();
1695 assert!(inbox.messages.is_empty());
1696 }
1697
1698 #[test]
1699 fn test_message_inbox_drain_all_empties() {
1700 let mut inbox = MessageInbox::new();
1701 inbox.push(AgentMessage {
1702 agent_id: "a".to_string(),
1703 content: "msg".to_string(),
1704 target_region: None,
1705 });
1706 let _ = inbox.drain_all();
1707 assert!(inbox.messages.is_empty());
1708 let result = inbox.drain_all();
1710 assert!(result.is_empty());
1711 }
1712
1713 #[test]
1714 fn test_agent_message_clone() {
1715 let msg = AgentMessage {
1716 agent_id: "agent-1".to_string(),
1717 content: "hello".to_string(),
1718 target_region: Some("conv".to_string()),
1719 };
1720 let cloned = msg.clone();
1721 assert_eq!(cloned.agent_id, "agent-1");
1722 assert_eq!(cloned.content, "hello");
1723 assert_eq!(cloned.target_region, Some("conv".to_string()));
1724 }
1725
1726 #[test]
1727 fn test_agent_status_serialization() {
1728 let status = AgentStatus::Active;
1729 let json = serde_json::to_string(&status).unwrap();
1730 assert!(json.contains("Active"));
1731
1732 let error_status = AgentStatus::Error {
1733 message: "boom".to_string(),
1734 };
1735 let json = serde_json::to_string(&error_status).unwrap();
1736 assert!(json.contains("boom"));
1737 }
1738
1739 #[test]
1740 fn test_tool_call_serialization() {
1741 let tc = ToolCall {
1742 tool_id: "tool-1".to_string(),
1743 name: "search".to_string(),
1744 arguments: serde_json::json!({"query": "rust"}),
1745 thought_signature: None,
1746 };
1747 let json = serde_json::to_string(&tc).unwrap();
1748 assert!(json.contains("search"));
1749 assert!(json.contains("rust"));
1750 }
1751
1752 #[test]
1753 fn test_eviction_with_only_pinned_region_frees_nothing() {
1754 let mut window = ContextWindow::new(10000);
1756 let mut pinned = Region::new("pinned".to_string(), RegionKind::Pinned, 5000);
1757 pinned
1758 .add_entry("important data".to_string(), 2000)
1759 .unwrap();
1760 window.add_region(pinned);
1761
1762 let result = with_tracing(|| window.try_evict(500)).unwrap();
1763 assert_eq!(result.tokens_freed, 0);
1764 assert!(result.needs_compaction.is_empty());
1765 }
1766
1767 #[test]
1768 fn test_inference_result_fields() {
1769 let ir = InferenceResult {
1770 response: "Hello".to_string(),
1771 tool_calls: vec![ToolCall {
1772 tool_id: "t1".to_string(),
1773 name: "search".to_string(),
1774 arguments: serde_json::json!({}),
1775 thought_signature: None,
1776 }],
1777 tokens_used: 100,
1778 timestamp: 99999,
1779 };
1780 assert_eq!(ir.response, "Hello");
1781 assert_eq!(ir.tool_calls.len(), 1);
1782 assert_eq!(ir.tokens_used, 100);
1783 }
1784
1785 #[test]
1786 fn test_sub_agent_children_clone() {
1787 let children = SubAgentChildren {
1788 children: vec![
1789 Entity::from_raw_u32(1).expect("a small literal index is always a valid entity id"),
1790 ],
1791 max_child_depth: 2,
1792 };
1793 let cloned = children.clone();
1794 assert_eq!(cloned.children.len(), 1);
1795 assert_eq!(cloned.max_child_depth, 2);
1796 }
1797
1798 #[test]
1811 fn try_evict_continues_loop_when_each_entry_removal_is_insufficient() {
1812 let mut window = ContextWindow::new(1000);
1813 let mut temp = Region::new("cache".to_string(), RegionKind::Temporary, 800);
1814 temp.add_entry("entry1".to_string(), 50).unwrap();
1815 temp.add_entry("entry2".to_string(), 50).unwrap();
1816 window.add_region(temp);
1817 window.current_tokens = 950; let result = window.try_evict(200).unwrap();
1821 assert_eq!(result.tokens_freed, 100); }
1823
1824 #[test]
1827 fn test_enable_taint_tracking_on_context_window() {
1828 let mut window = ContextWindow::new(10000);
1829 window.add_region(Region::new(
1830 "conv".to_string(),
1831 RegionKind::SlidingWindow {
1832 max_items: 10,
1833 eviction_strategy: EvictionStrategy::PerItem,
1834 },
1835 5000,
1836 ));
1837 window.add_region(Region::new(
1838 "tools".to_string(),
1839 RegionKind::Temporary,
1840 3000,
1841 ));
1842
1843 assert!(window.overall_taint().is_none());
1844 window.enable_taint_tracking();
1845 assert_eq!(
1846 window.overall_taint(),
1847 Some(leviath_core::TaintLevel::Public)
1848 );
1849 }
1850
1851 #[test]
1852 fn test_add_tainted_to_region() {
1853 let mut window = ContextWindow::new(10000);
1854 let region =
1855 Region::new("tools".to_string(), RegionKind::Temporary, 5000).with_taint_tracking();
1856 window.add_region(region);
1857
1858 window
1859 .add_tainted_to_region(
1860 "tools",
1861 "secret data".to_string(),
1862 10,
1863 leviath_core::TaintLevel::Private,
1864 )
1865 .unwrap();
1866
1867 assert_eq!(
1868 window.get_region("tools").and_then(|r| r.taint_level()),
1869 Some(leviath_core::TaintLevel::Private)
1870 );
1871 assert_eq!(
1872 window.overall_taint(),
1873 Some(leviath_core::TaintLevel::Private)
1874 );
1875 }
1876
1877 #[test]
1878 fn test_add_tainted_to_nonexistent_region() {
1879 let mut window = ContextWindow::new(10000);
1880 let result = window.add_tainted_to_region(
1881 "nope",
1882 "data".to_string(),
1883 10,
1884 leviath_core::TaintLevel::Public,
1885 );
1886 assert!(result.is_err());
1887 }
1888
1889 #[test]
1890 fn test_overall_taint_is_max_across_regions() {
1891 let mut window = ContextWindow::new(10000);
1892 let r1 = Region::new("a".to_string(), RegionKind::Temporary, 5000).with_taint_tracking();
1893 let r2 = Region::new("b".to_string(), RegionKind::Temporary, 5000).with_taint_tracking();
1894 window.add_region(r1);
1895 window.add_region(r2);
1896
1897 window
1898 .add_tainted_to_region("a", "x".to_string(), 5, leviath_core::TaintLevel::Internal)
1899 .unwrap();
1900 window
1901 .add_tainted_to_region("b", "y".to_string(), 5, leviath_core::TaintLevel::Public)
1902 .unwrap();
1903
1904 assert_eq!(
1905 window.overall_taint(),
1906 Some(leviath_core::TaintLevel::Internal)
1907 );
1908 }
1909
1910 #[test]
1911 fn test_taint_summary() {
1912 let mut window = ContextWindow::new(10000);
1913 let r1 = Region::new("conv".to_string(), RegionKind::Temporary, 5000).with_taint_tracking();
1914 let r2 =
1915 Region::new("tools".to_string(), RegionKind::Temporary, 5000).with_taint_tracking();
1916 window.add_region(r1);
1917 window.add_region(r2);
1918
1919 window
1920 .add_tainted_to_region(
1921 "conv",
1922 "x".to_string(),
1923 5,
1924 leviath_core::TaintLevel::Private,
1925 )
1926 .unwrap();
1927
1928 let summary = window.taint_summary();
1929 assert_eq!(summary.len(), 2);
1930 assert!(
1931 summary
1932 .iter()
1933 .any(|(name, level)| name == "conv" && *level == leviath_core::TaintLevel::Private)
1934 );
1935 assert!(
1936 summary
1937 .iter()
1938 .any(|(name, level)| name == "tools" && *level == leviath_core::TaintLevel::Public)
1939 );
1940 }
1941
1942 #[test]
1943 fn test_taint_recovery_through_eviction() {
1944 with_tracing(|| {});
1945 let mut window = ContextWindow::new(100);
1946 let r = Region::new("temp".to_string(), RegionKind::Temporary, 100).with_taint_tracking();
1947 window.add_region(r);
1948
1949 window
1950 .add_tainted_to_region(
1951 "temp",
1952 "private".to_string(),
1953 30,
1954 leviath_core::TaintLevel::Private,
1955 )
1956 .unwrap();
1957 window
1958 .add_tainted_to_region(
1959 "temp",
1960 "public".to_string(),
1961 30,
1962 leviath_core::TaintLevel::Public,
1963 )
1964 .unwrap();
1965
1966 assert_eq!(
1967 window.get_region("temp").and_then(|r| r.taint_level()),
1968 Some(leviath_core::TaintLevel::Private)
1969 );
1970
1971 window.current_tokens = 96; let result = window.try_evict(10).unwrap();
1974 assert!(result.tokens_freed > 0);
1975
1976 assert_eq!(
1978 window.get_region("temp").and_then(|r| r.taint_level()),
1979 Some(leviath_core::TaintLevel::Public)
1980 );
1981 }
1982
1983 #[test]
1986 fn test_assemble_appends_user_nudge_when_conversation_ends_with_assistant() {
1987 let mut window = ContextWindow::new(100_000);
1991 window.add_region(Region::new(
1992 "conversation".to_string(),
1993 RegionKind::SlidingWindow {
1994 max_items: 100,
1995 eviction_strategy: EvictionStrategy::PerItem,
1996 },
1997 50_000,
1998 ));
1999 window
2000 .add_typed_entry(
2001 "conversation",
2002 leviath_core::EntryKind::UserMessage,
2003 "do the task".to_string(),
2004 10,
2005 )
2006 .unwrap();
2007 window
2008 .add_typed_entry(
2009 "conversation",
2010 leviath_core::EntryKind::AssistantTurn { tool_calls: vec![] },
2011 "All done with stage one.".to_string(),
2012 10,
2013 )
2014 .unwrap();
2015
2016 let assembled = window.assemble();
2017 assert_eq!(
2018 assembled.messages.last().map(|m| m.role.as_str()),
2019 Some("user"),
2020 "the assembled conversation must end with a user message"
2021 );
2022 }
2023
2024 #[test]
2025 fn test_assemble_strips_orphaned_tool_use() {
2026 let mut window = ContextWindow::new(100_000);
2027 let region = Region::new(
2028 "conversation".to_string(),
2029 RegionKind::SlidingWindow {
2030 max_items: 100,
2031 eviction_strategy: EvictionStrategy::PerItem,
2032 },
2033 50_000,
2034 );
2035 window.add_region(region);
2036
2037 window
2039 .add_typed_entry(
2040 "conversation",
2041 leviath_core::EntryKind::AssistantTurn {
2042 tool_calls: vec![leviath_core::SerializedToolCall {
2043 id: "tc_orphan".to_string(),
2044 name: "read_file".to_string(),
2045 arguments: serde_json::json!({"path": "foo.rs"}),
2046 thought_signature: None,
2047 }],
2048 },
2049 "Let me read that file.".to_string(),
2050 50,
2051 )
2052 .unwrap();
2053
2054 let assembled = with_tracing(|| window.assemble());
2055
2056 for msg in &assembled.messages {
2058 if let leviath_providers::MessageContent::Blocks(blocks) = &msg.content {
2059 for block in blocks {
2060 assert!(
2061 !matches!(block, leviath_providers::ContentBlock::ToolUse { .. }),
2062 "Orphaned tool_use should have been stripped"
2063 );
2064 }
2065 }
2066 }
2067 assert!(
2069 assembled
2070 .messages
2071 .iter()
2072 .any(|m| m.role == "assistant" && m.content.as_text().contains("read that file"))
2073 );
2074 }
2075
2076 #[test]
2077 fn test_assemble_strips_orphaned_tool_result() {
2078 let mut window = ContextWindow::new(100_000);
2079 let region = Region::new(
2080 "conversation".to_string(),
2081 RegionKind::SlidingWindow {
2082 max_items: 100,
2083 eviction_strategy: EvictionStrategy::PerItem,
2084 },
2085 50_000,
2086 );
2087 window.add_region(region);
2088
2089 window
2091 .add_typed_entry(
2092 "conversation",
2093 leviath_core::EntryKind::UserMessage,
2094 "Hello".to_string(),
2095 10,
2096 )
2097 .unwrap();
2098
2099 window
2101 .add_typed_entry(
2102 "conversation",
2103 leviath_core::EntryKind::ToolResult {
2104 tool_call_id: "tc_missing".to_string(),
2105 tool_name: "read_file".to_string(),
2106 is_error: false,
2107 },
2108 "file contents here".to_string(),
2109 20,
2110 )
2111 .unwrap();
2112
2113 let assembled = with_tracing(|| window.assemble());
2114
2115 assert_eq!(assembled.messages.len(), 1);
2118 assert_eq!(assembled.messages[0].role, "user");
2119 assert_eq!(
2120 assembled.messages[0].content,
2121 leviath_providers::MessageContent::Text("Hello".to_string())
2122 );
2123 }
2124
2125 #[test]
2126 fn test_assemble_paired_tool_use_result_passes_through() {
2127 let mut window = ContextWindow::new(100_000);
2128 let region = Region::new(
2129 "conversation".to_string(),
2130 RegionKind::SlidingWindow {
2131 max_items: 100,
2132 eviction_strategy: EvictionStrategy::PerItem,
2133 },
2134 50_000,
2135 );
2136 window.add_region(region);
2137
2138 window
2140 .add_typed_entry(
2141 "conversation",
2142 leviath_core::EntryKind::UserMessage,
2143 "Fix the bug".to_string(),
2144 10,
2145 )
2146 .unwrap();
2147
2148 window
2150 .add_typed_entry(
2151 "conversation",
2152 leviath_core::EntryKind::AssistantTurn {
2153 tool_calls: vec![leviath_core::SerializedToolCall {
2154 id: "tc_1".to_string(),
2155 name: "read_file".to_string(),
2156 arguments: serde_json::json!({"path": "main.rs"}),
2157 thought_signature: None,
2158 }],
2159 },
2160 "".to_string(),
2161 10,
2162 )
2163 .unwrap();
2164
2165 window
2167 .add_typed_entry(
2168 "conversation",
2169 leviath_core::EntryKind::ToolResult {
2170 tool_call_id: "tc_1".to_string(),
2171 tool_name: "read_file".to_string(),
2172 is_error: false,
2173 },
2174 "fn main() {}".to_string(),
2175 10,
2176 )
2177 .unwrap();
2178
2179 let assembled = window.assemble();
2180
2181 let has_tool_use = assembled.messages.iter().any(|m| {
2183 if let leviath_providers::MessageContent::Blocks(blocks) = &m.content {
2184 blocks
2185 .iter()
2186 .any(|b| matches!(b, leviath_providers::ContentBlock::ToolUse { id, .. } if id == "tc_1"))
2187 } else {
2188 false
2189 }
2190 });
2191 let has_tool_result = assembled.messages.iter().any(|m| {
2192 if let leviath_providers::MessageContent::Blocks(blocks) = &m.content {
2193 blocks
2194 .iter()
2195 .any(|b| matches!(b, leviath_providers::ContentBlock::ToolResult { tool_use_id, .. } if tool_use_id == "tc_1"))
2196 } else {
2197 false
2198 }
2199 });
2200 assert!(has_tool_use, "Paired tool_use should remain");
2201 assert!(has_tool_result, "Paired tool_result should remain");
2202 }
2203
2204 #[test]
2205 fn test_assemble_removes_empty_assistant_after_stripping() {
2206 let mut window = ContextWindow::new(100_000);
2207 let region = Region::new(
2208 "conversation".to_string(),
2209 RegionKind::SlidingWindow {
2210 max_items: 100,
2211 eviction_strategy: EvictionStrategy::PerItem,
2212 },
2213 50_000,
2214 );
2215 window.add_region(region);
2216
2217 window
2219 .add_typed_entry(
2220 "conversation",
2221 leviath_core::EntryKind::UserMessage,
2222 "Do something".to_string(),
2223 10,
2224 )
2225 .unwrap();
2226
2227 window
2229 .add_typed_entry(
2230 "conversation",
2231 leviath_core::EntryKind::AssistantTurn {
2232 tool_calls: vec![leviath_core::SerializedToolCall {
2233 id: "tc_gone".to_string(),
2234 name: "bash".to_string(),
2235 arguments: serde_json::json!({"command": "ls"}),
2236 thought_signature: None,
2237 }],
2238 },
2239 "".to_string(),
2240 10,
2241 )
2242 .unwrap();
2243
2244 let assembled = with_tracing(|| window.assemble());
2245
2246 let assistant_msgs: Vec<_> = assembled
2248 .messages
2249 .iter()
2250 .filter(|m| m.role == "assistant")
2251 .collect();
2252 assert!(
2253 assistant_msgs.is_empty(),
2254 "Assistant message with only orphaned tool_use should be removed entirely"
2255 );
2256 }
2257
2258 #[test]
2259 fn test_assemble_strips_multiple_orphaned_tool_uses_in_one_message() {
2260 let mut window = ContextWindow::new(100_000);
2261 let region = Region::new(
2262 "conversation".to_string(),
2263 RegionKind::SlidingWindow {
2264 max_items: 100,
2265 eviction_strategy: EvictionStrategy::PerItem,
2266 },
2267 50_000,
2268 );
2269 window.add_region(region);
2270
2271 window
2273 .add_typed_entry(
2274 "conversation",
2275 leviath_core::EntryKind::UserMessage,
2276 "Do two things".to_string(),
2277 10,
2278 )
2279 .unwrap();
2280
2281 window
2283 .add_typed_entry(
2284 "conversation",
2285 leviath_core::EntryKind::AssistantTurn {
2286 tool_calls: vec![
2287 leviath_core::SerializedToolCall {
2288 id: "tc_orphan_1".to_string(),
2289 name: "read_file".to_string(),
2290 arguments: serde_json::json!({"path": "a.rs"}),
2291 thought_signature: None,
2292 },
2293 leviath_core::SerializedToolCall {
2294 id: "tc_orphan_2".to_string(),
2295 name: "bash".to_string(),
2296 arguments: serde_json::json!({"cmd": "ls"}),
2297 thought_signature: None,
2298 },
2299 ],
2300 },
2301 "Let me do both.".to_string(),
2302 50,
2303 )
2304 .unwrap();
2305
2306 let assembled = with_tracing(|| window.assemble());
2307
2308 for msg in &assembled.messages {
2310 if let leviath_providers::MessageContent::Blocks(blocks) = &msg.content {
2311 for block in blocks {
2312 assert!(
2313 !matches!(block, leviath_providers::ContentBlock::ToolUse { .. }),
2314 "All orphaned tool_uses should have been stripped"
2315 );
2316 }
2317 }
2318 }
2319 assert!(
2321 assembled
2322 .messages
2323 .iter()
2324 .any(|m| m.role == "assistant" && m.content.as_text().contains("do both"))
2325 );
2326 }
2327
2328 #[test]
2329 fn test_assemble_mixed_valid_and_orphaned_in_same_message() {
2330 let mut window = ContextWindow::new(100_000);
2331 let region = Region::new(
2332 "conversation".to_string(),
2333 RegionKind::SlidingWindow {
2334 max_items: 100,
2335 eviction_strategy: EvictionStrategy::PerItem,
2336 },
2337 50_000,
2338 );
2339 window.add_region(region);
2340
2341 window
2343 .add_typed_entry(
2344 "conversation",
2345 leviath_core::EntryKind::UserMessage,
2346 "Do stuff".to_string(),
2347 10,
2348 )
2349 .unwrap();
2350
2351 window
2353 .add_typed_entry(
2354 "conversation",
2355 leviath_core::EntryKind::AssistantTurn {
2356 tool_calls: vec![
2357 leviath_core::SerializedToolCall {
2358 id: "tc_valid".to_string(),
2359 name: "read_file".to_string(),
2360 arguments: serde_json::json!({"path": "main.rs"}),
2361 thought_signature: None,
2362 },
2363 leviath_core::SerializedToolCall {
2364 id: "tc_orphan".to_string(),
2365 name: "bash".to_string(),
2366 arguments: serde_json::json!({"cmd": "ls"}),
2367 thought_signature: None,
2368 },
2369 ],
2370 },
2371 "".to_string(),
2372 10,
2373 )
2374 .unwrap();
2375
2376 window
2378 .add_typed_entry(
2379 "conversation",
2380 leviath_core::EntryKind::ToolResult {
2381 tool_call_id: "tc_valid".to_string(),
2382 tool_name: "read_file".to_string(),
2383 is_error: false,
2384 },
2385 "fn main() {}".to_string(),
2386 10,
2387 )
2388 .unwrap();
2389
2390 let assembled = with_tracing(|| window.assemble());
2391
2392 let tool_use_ids: Vec<&str> = assembled
2394 .messages
2395 .iter()
2396 .filter_map(|m| match &m.content {
2397 leviath_providers::MessageContent::Blocks(blocks) => Some(blocks),
2398 _ => None,
2399 })
2400 .flatten()
2401 .filter_map(|b| match b {
2402 leviath_providers::ContentBlock::ToolUse { id, .. } => Some(id.as_str()),
2403 _ => None,
2404 })
2405 .collect();
2406 assert!(
2408 tool_use_ids.contains(&"tc_valid"),
2409 "Valid tool_use should remain"
2410 );
2411 assert!(
2412 !tool_use_ids.contains(&"tc_orphan"),
2413 "Orphaned tool_use should be stripped"
2414 );
2415
2416 let has_result = assembled.messages.iter().any(|m| {
2418 if let leviath_providers::MessageContent::Blocks(blocks) = &m.content {
2419 blocks.iter().any(|b| {
2420 matches!(b, leviath_providers::ContentBlock::ToolResult { tool_use_id, .. } if tool_use_id == "tc_valid")
2421 })
2422 } else {
2423 false
2424 }
2425 });
2426 assert!(has_result, "Valid tool_result should remain");
2427 }
2428
2429 #[test]
2432 fn test_assemble_compact_history_region_produces_system_block_always() {
2433 let mut window = ContextWindow::new(100_000);
2434 let mut region = Region::new(
2435 "history".to_string(),
2436 RegionKind::CompactHistory {
2437 source_region: "conv".to_string(),
2438 },
2439 10_000,
2440 );
2441 region
2442 .add_entry("summary of earlier conversation".to_string(), 50)
2443 .unwrap();
2444 window.add_region(region);
2445
2446 let assembled = window.assemble();
2447
2448 assert_eq!(assembled.system_blocks.len(), 1);
2449 assert_eq!(
2450 assembled.system_blocks[0].text,
2451 "summary of earlier conversation"
2452 );
2453 assert_eq!(
2454 assembled.system_blocks[0].cache_hint,
2455 leviath_core::CacheHint::Always
2456 );
2457 }
2458
2459 #[test]
2460 fn test_assemble_compacting_region_produces_system_block_until_changed() {
2461 let mut window = ContextWindow::new(100_000);
2462 let mut region = Region::new(
2463 "impl".to_string(),
2464 RegionKind::Compacting {
2465 threshold_tokens: 500,
2466 },
2467 10_000,
2468 );
2469 region
2470 .add_entry("implementation details".to_string(), 50)
2471 .unwrap();
2472 window.add_region(region);
2473
2474 let assembled = window.assemble();
2475
2476 assert_eq!(assembled.system_blocks.len(), 1);
2477 assert_eq!(
2478 assembled.system_blocks[0].text,
2479 "[impl]:\nimplementation details"
2480 );
2481 assert_eq!(
2482 assembled.system_blocks[0].cache_hint,
2483 leviath_core::CacheHint::UntilChanged
2484 );
2485 }
2486
2487 #[test]
2488 fn test_assemble_temporary_region_produces_system_block_never() {
2489 let mut window = ContextWindow::new(100_000);
2490 let mut region = Region::new("scratch".to_string(), RegionKind::Temporary, 10_000);
2491 region.add_entry("temp data".to_string(), 20).unwrap();
2492 window.add_region(region);
2493
2494 let assembled = window.assemble();
2495
2496 assert_eq!(assembled.system_blocks.len(), 1);
2497 assert_eq!(assembled.system_blocks[0].text, "[scratch]:\ntemp data");
2498 assert_eq!(
2499 assembled.system_blocks[0].cache_hint,
2500 leviath_core::CacheHint::Never
2501 );
2502 }
2503
2504 #[test]
2505 fn test_assemble_clearable_region_produces_system_block_never() {
2506 let mut window = ContextWindow::new(100_000);
2507 let mut region = Region::new("cache".to_string(), RegionKind::Clearable, 10_000);
2508 region.add_entry("clearable data".to_string(), 20).unwrap();
2509 window.add_region(region);
2510
2511 let assembled = window.assemble();
2512
2513 assert_eq!(assembled.system_blocks.len(), 1);
2514 assert_eq!(assembled.system_blocks[0].text, "[cache]:\nclearable data");
2515 assert_eq!(
2516 assembled.system_blocks[0].cache_hint,
2517 leviath_core::CacheHint::Never
2518 );
2519 }
2520
2521 fn custom_kind(script: &str, persistent: bool) -> RegionKind {
2522 RegionKind::Custom {
2523 script: script.to_string(),
2524 persistent,
2525 }
2526 }
2527
2528 #[test]
2529 fn test_assemble_custom_region_falls_back_to_temporary_style_block() {
2530 let mut window = ContextWindow::new(100_000);
2534 let mut region = Region::new("brain".to_string(), custom_kind("b.rhai", false), 10_000);
2535 region.add_entry("thought one".to_string(), 10).unwrap();
2536 region.add_entry("thought two".to_string(), 10).unwrap();
2537 window.add_region(region);
2538
2539 let assembled = window.assemble();
2540
2541 assert_eq!(assembled.system_blocks.len(), 1);
2542 assert_eq!(
2543 assembled.system_blocks[0].text,
2544 "[brain]:\nthought one\n\nthought two"
2545 );
2546 assert_eq!(
2547 assembled.system_blocks[0].cache_hint,
2548 leviath_core::CacheHint::Never
2549 );
2550 }
2551
2552 #[test]
2553 fn try_evict_evicts_non_persistent_custom_regions_oldest_first() {
2554 let mut window = ContextWindow::new(100);
2555 let mut region = Region::new("brain".to_string(), custom_kind("b.rhai", false), 100);
2556 region.add_entry("old".to_string(), 40).unwrap();
2557 region.add_entry("new".to_string(), 40).unwrap();
2558 window.add_region(region);
2559 window.current_tokens = 80;
2560
2561 let result = with_tracing(|| window.try_evict(30).unwrap());
2562 assert!(result.tokens_freed >= 40);
2563 let brain = window.get_region("brain").unwrap();
2564 assert_eq!(brain.content.len(), 1);
2565 assert_eq!(brain.content[0].content, "new");
2566 }
2567
2568 #[test]
2569 fn try_evict_never_touches_persistent_custom_and_counts_it_as_pinned() {
2570 let mut window = ContextWindow::new(50);
2573 let mut vault = Region::new("vault".to_string(), custom_kind("v.rhai", true), 100);
2574 vault.add_entry("precious".to_string(), 60).unwrap();
2575 window.add_region(vault);
2576 window.current_tokens = 60;
2577
2578 let err = with_tracing(|| window.try_evict(10).unwrap_err());
2579 assert_eq!(
2580 err.to_string(),
2581 "Pinned regions (60) exceed total budget (50)"
2582 );
2583 assert_eq!(window.get_region("vault").unwrap().content.len(), 1);
2584 }
2585
2586 fn custom_window(src: &str, persistent: bool) -> ContextWindow {
2589 let mut window = ContextWindow::new(10_000);
2590 window.add_region(Region::new(
2591 "brain".to_string(),
2592 RegionKind::Custom {
2593 script: "s.rhai".to_string(),
2594 persistent,
2595 },
2596 100,
2597 ));
2598 window.region_scripts.insert(
2599 "s.rhai".to_string(),
2600 std::sync::Arc::new(leviath_scripting::region_hook::compile("s.rhai", src).unwrap()),
2601 );
2602 window
2603 }
2604
2605 #[test]
2606 fn custom_region_on_write_fires_across_all_write_methods() {
2607 let src = r#"
2608 fn render(ctx) { "" }
2609 fn on_write(ctx) { `${ctx.entry.kind}:${ctx.entry.content}` }
2610 "#;
2611 let mut window = custom_window(src, false);
2612
2613 window.add_to_region("brain", "a".to_string(), 1).unwrap();
2614 window
2615 .add_typed_entry(
2616 "brain",
2617 leviath_core::EntryKind::UserMessage,
2618 "b".to_string(),
2619 1,
2620 )
2621 .unwrap();
2622 window
2623 .add_tainted_to_region(
2624 "brain",
2625 "c".to_string(),
2626 1,
2627 leviath_core::TaintLevel::Public,
2628 )
2629 .unwrap();
2630 window
2631 .add_typed_tainted_to_region(
2632 "brain",
2633 leviath_core::EntryKind::UserMessage,
2634 "d".to_string(),
2635 1,
2636 leviath_core::TaintLevel::Public,
2637 )
2638 .unwrap();
2639
2640 let contents: Vec<_> = window
2641 .get_region("brain")
2642 .unwrap()
2643 .content
2644 .iter()
2645 .map(|e| e.content.as_str())
2646 .collect();
2647 assert_eq!(
2648 contents,
2649 vec!["text:a", "user_message:b", "text:c", "user_message:d"],
2650 "every write method passes through on_write with the entry kind visible"
2651 );
2652 assert_eq!(window.current_tokens, window.calculate_tokens());
2654
2655 assert!(window.replace_region("brain", "e".to_string(), 1));
2656 let region = window.get_region("brain").unwrap();
2657 assert_eq!(region.content.len(), 1);
2658 assert_eq!(region.content[0].content, "text:e");
2659 }
2660
2661 #[test]
2662 fn custom_region_on_write_drop_reports_success_without_storing() {
2663 let src = r#"
2664 fn render(ctx) { "" }
2665 fn on_write(ctx) { false }
2666 "#;
2667 let mut window = custom_window(src, false);
2668 window
2669 .add_to_region("brain", "spam".to_string(), 1)
2670 .unwrap();
2671 assert!(window.get_region("brain").unwrap().content.is_empty());
2672
2673 assert!(window.replace_region("brain", "more spam".to_string(), 1));
2675 assert!(window.get_region("brain").unwrap().content.is_empty());
2676 }
2677
2678 #[test]
2679 fn custom_region_on_write_drop_covers_typed_and_tainted_methods() {
2680 let src = r#"
2682 fn render(ctx) { "" }
2683 fn on_write(ctx) { false }
2684 "#;
2685 let mut window = custom_window(src, false);
2686 window
2687 .add_typed_entry(
2688 "brain",
2689 leviath_core::EntryKind::UserMessage,
2690 "a".to_string(),
2691 1,
2692 )
2693 .unwrap();
2694 window
2695 .add_tainted_to_region(
2696 "brain",
2697 "b".to_string(),
2698 1,
2699 leviath_core::TaintLevel::Public,
2700 )
2701 .unwrap();
2702 window
2703 .add_typed_tainted_to_region(
2704 "brain",
2705 leviath_core::EntryKind::UserMessage,
2706 "c".to_string(),
2707 1,
2708 leviath_core::TaintLevel::Public,
2709 )
2710 .unwrap();
2711 assert!(window.get_region("brain").unwrap().content.is_empty());
2712 }
2713
2714 #[test]
2715 fn try_evict_skips_custom_region_whose_script_has_no_on_overflow() {
2716 let mut window = ContextWindow::new(100);
2719 window.add_region(Region::new(
2720 "brain".to_string(),
2721 RegionKind::Custom {
2722 script: "s.rhai".to_string(),
2723 persistent: false,
2724 },
2725 100,
2726 ));
2727 window.region_scripts.insert(
2728 "s.rhai".to_string(),
2729 std::sync::Arc::new(
2730 leviath_scripting::region_hook::compile("s.rhai", "fn render(ctx) { \"\" }")
2731 .unwrap(),
2732 ),
2733 );
2734 window
2735 .add_to_region("brain", "old".to_string(), 40)
2736 .unwrap();
2737 window
2738 .add_to_region("brain", "new".to_string(), 40)
2739 .unwrap();
2740
2741 let result = with_tracing(|| window.try_evict(30).unwrap());
2742 assert!(result.tokens_freed >= 40);
2743 let brain = window.get_region("brain").unwrap();
2744 assert_eq!(brain.content.len(), 1);
2745 assert_eq!(brain.content[0].content, "new", "oldest-first fallback ran");
2746 }
2747
2748 #[test]
2749 fn try_evict_falls_to_oldest_first_when_script_frees_nothing() {
2750 let src = r#"
2753 fn render(ctx) { "" }
2754 fn on_overflow(ctx) { [] }
2755 "#;
2756 let mut window = ContextWindow::new(100);
2757 window.add_region(Region::new(
2758 "brain".to_string(),
2759 RegionKind::Custom {
2760 script: "s.rhai".to_string(),
2761 persistent: false,
2762 },
2763 100,
2764 ));
2765 window.region_scripts.insert(
2766 "s.rhai".to_string(),
2767 std::sync::Arc::new(leviath_scripting::region_hook::compile("s.rhai", src).unwrap()),
2768 );
2769 window
2770 .add_to_region("brain", "old".to_string(), 40)
2771 .unwrap();
2772 window
2773 .add_to_region("brain", "new".to_string(), 40)
2774 .unwrap();
2775
2776 let result = with_tracing(|| window.try_evict(30).unwrap());
2777 assert!(result.tokens_freed >= 40);
2778 assert_eq!(window.get_region("brain").unwrap().content.len(), 1);
2779 }
2780
2781 #[test]
2782 fn non_custom_regions_bypass_the_on_write_seam() {
2783 let mut window = custom_window(
2786 "fn render(ctx) { \"\" }\nfn on_write(ctx) { \"MANGLED\" }",
2787 false,
2788 );
2789 window.add_region(Region::new("plain".to_string(), RegionKind::Temporary, 100));
2790 window
2791 .add_to_region("plain", "untouched".to_string(), 2)
2792 .unwrap();
2793 assert_eq!(
2794 window.get_region("plain").unwrap().content[0].content,
2795 "untouched"
2796 );
2797 }
2798
2799 #[test]
2800 fn write_to_missing_region_still_errors() {
2801 let mut window = custom_window("fn render(ctx) { \"\" }", false);
2802 let err = window
2803 .add_to_region("ghost", "x".to_string(), 1)
2804 .unwrap_err();
2805 assert!(err.to_string().contains("ghost"), "{err}");
2806 }
2807
2808 #[test]
2809 fn custom_region_add_time_overflow_retries_after_script_drops() {
2810 let src = r#"
2813 fn render(ctx) { "" }
2814 fn on_overflow(ctx) { [0] }
2815 "#;
2816 let mut window = custom_window(src, false);
2817 window
2818 .add_to_region("brain", "big".to_string(), 90)
2819 .unwrap();
2820 window
2821 .add_to_region("brain", "next".to_string(), 20)
2822 .unwrap();
2823
2824 let region = window.get_region("brain").unwrap();
2825 assert_eq!(region.content.len(), 1);
2826 assert_eq!(region.content[0].content, "next");
2827 assert_eq!(window.current_tokens, 20);
2828 }
2829
2830 #[test]
2831 fn custom_region_add_time_overflow_propagates_when_still_too_big() {
2832 let src = r#"
2835 fn render(ctx) { "" }
2836 fn on_overflow(ctx) { [] }
2837 "#;
2838 let mut window = custom_window(src, false);
2839 window
2840 .add_to_region("brain", "big".to_string(), 90)
2841 .unwrap();
2842 let err =
2843 with_tracing(|| window.add_to_region("brain", "too much".to_string(), 50)).unwrap_err();
2844 assert_eq!(err.to_string(), "Content exceeds token budget: 140 > 100");
2845 assert_eq!(window.get_region("brain").unwrap().content.len(), 1);
2846 }
2847
2848 #[test]
2849 fn custom_region_without_on_overflow_gets_no_retry() {
2850 let mut window = custom_window("fn render(ctx) { \"\" }", false);
2851 window
2852 .add_to_region("brain", "big".to_string(), 90)
2853 .unwrap();
2854 let err = window
2855 .add_to_region("brain", "too much".to_string(), 50)
2856 .unwrap_err();
2857 assert_eq!(err.to_string(), "Content exceeds token budget: 140 > 100");
2858 }
2859
2860 #[test]
2861 fn try_evict_lets_custom_script_choose_what_to_drop() {
2862 let src = r#"
2866 fn render(ctx) { "" }
2867 fn on_overflow(ctx) {
2868 let drops = [];
2869 for (entry, i) in ctx.entries {
2870 if !entry.content.contains("ERROR") { drops.push(i); }
2871 }
2872 drops
2873 }
2874 "#;
2875 let mut window = ContextWindow::new(100);
2876 window.add_region(Region::new(
2877 "brain".to_string(),
2878 RegionKind::Custom {
2879 script: "s.rhai".to_string(),
2880 persistent: false,
2881 },
2882 100,
2883 ));
2884 window.region_scripts.insert(
2885 "s.rhai".to_string(),
2886 std::sync::Arc::new(leviath_scripting::region_hook::compile("s.rhai", src).unwrap()),
2887 );
2888 window
2889 .add_to_region("brain", "ok one".to_string(), 30)
2890 .unwrap();
2891 window
2892 .add_to_region("brain", "ERROR two".to_string(), 30)
2893 .unwrap();
2894 window
2895 .add_to_region("brain", "ok three".to_string(), 30)
2896 .unwrap();
2897
2898 let result = with_tracing(|| window.try_evict(40)).unwrap();
2899 assert!(result.tokens_freed >= 40);
2900 let contents: Vec<_> = window
2901 .get_region("brain")
2902 .unwrap()
2903 .content
2904 .iter()
2905 .map(|e| e.content.as_str())
2906 .collect();
2907 assert_eq!(
2908 contents,
2909 vec!["ERROR two"],
2910 "script retention choice honored"
2911 );
2912 }
2913
2914 #[test]
2917 fn assemble_custom_region_renders_through_script() {
2918 let src = r#"fn render(ctx) { `<brain iter=${ctx.stage_iterations}>` }"#;
2919 let mut window = custom_window(src, false);
2920 window
2921 .add_to_region("brain", "note".to_string(), 2)
2922 .unwrap();
2923
2924 let assembled = window.assemble();
2926 assert_eq!(assembled.system_blocks.len(), 1);
2927 assert_eq!(assembled.system_blocks[0].text, "<brain iter=0>");
2928 assert_eq!(
2929 assembled.system_blocks[0].cache_hint,
2930 leviath_core::CacheHint::UntilChanged
2931 );
2932
2933 let assembled = window.assemble_with_meta(&crate::custom_region::AssembleMeta {
2935 stage_name: "plan".to_string(),
2936 stage_iterations: 7,
2937 model: "m".to_string(),
2938 });
2939 assert_eq!(assembled.system_blocks[0].text, "<brain iter=7>");
2940 }
2941
2942 #[test]
2943 fn assemble_custom_region_renders_even_when_empty() {
2944 let src = r#"fn render(ctx) { `<empty count=${ctx.entries.len()}>` }"#;
2946 let window = custom_window(src, false);
2947 let assembled = window.assemble();
2948 assert_eq!(assembled.system_blocks.len(), 1);
2949 assert_eq!(assembled.system_blocks[0].text, "<empty count=0>");
2950 }
2951
2952 #[test]
2953 fn assemble_custom_conversation_takeover_renders_single_user_message() {
2954 let src = r#"
2958 fn render(ctx) {
2959 let xml = "<context>";
2960 for entry in ctx.entries {
2961 xml += `<event kind="${entry.kind}">${entry.content}</event>`;
2962 }
2963 xml += "</context>";
2964 #{ messages: [ #{ role: "user", content: xml } ] }
2965 }
2966 "#;
2967 let mut window = ContextWindow::new(10_000);
2968 window.add_region(Region::new(
2969 "conversation".to_string(),
2970 RegionKind::Custom {
2971 script: "conv.rhai".to_string(),
2972 persistent: false,
2973 },
2974 5_000,
2975 ));
2976 window.region_scripts.insert(
2977 "conv.rhai".to_string(),
2978 std::sync::Arc::new(leviath_scripting::region_hook::compile("conv.rhai", src).unwrap()),
2979 );
2980 window
2981 .add_typed_entry(
2982 "conversation",
2983 leviath_core::EntryKind::UserMessage,
2984 "do the task".to_string(),
2985 4,
2986 )
2987 .unwrap();
2988 window
2989 .add_typed_entry(
2990 "conversation",
2991 leviath_core::EntryKind::ToolResult {
2992 tool_call_id: "c1".to_string(),
2993 tool_name: "shell".to_string(),
2994 is_error: false,
2995 },
2996 "output".to_string(),
2997 2,
2998 )
2999 .unwrap();
3000
3001 let assembled = window.assemble();
3002 assert!(assembled.system_blocks.is_empty());
3003 assert_eq!(assembled.messages.len(), 1);
3004 assert_eq!(assembled.messages[0].role, "user");
3005 assert_eq!(
3006 assembled.messages[0].content.as_text(),
3007 "<context><event kind=\"user_message\">do the task</event>\
3008 <event kind=\"tool_result\">output</event></context>"
3009 );
3010 }
3011
3012 #[test]
3013 fn assemble_custom_script_emitting_nothing_gets_begin_fallback() {
3014 let window = custom_window("fn render(ctx) { \"\" }", false);
3017 let assembled = window.assemble();
3018 assert!(assembled.system_blocks.is_empty());
3019 assert_eq!(assembled.messages.len(), 1);
3020 assert_eq!(assembled.messages[0].content.as_text(), "Begin.");
3021 }
3022
3023 #[test]
3024 fn assemble_custom_unpaired_tool_blocks_are_sanitized() {
3025 let src = r#"
3029 fn render(ctx) {
3030 #{ messages: [
3031 #{ role: "user", content: "hello" },
3032 #{ role: "user", tool_results: [
3033 #{ tool_call_id: "ghost", content: "orphan" },
3034 ] },
3035 ] }
3036 }
3037 "#;
3038 let window = custom_window(src, false);
3039 let assembled = window.assemble();
3040 assert_eq!(assembled.messages.len(), 1, "orphan tool_result stripped");
3041 assert_eq!(assembled.messages[0].content.as_text(), "hello");
3042 }
3043
3044 #[test]
3047 fn test_assemble_text_entry_with_assistant_prefix() {
3048 let mut window = ContextWindow::new(100_000);
3049 let region = Region::new(
3050 "conv".to_string(),
3051 RegionKind::SlidingWindow {
3052 max_items: 100,
3053 eviction_strategy: EvictionStrategy::PerItem,
3054 },
3055 50_000,
3056 );
3057 window.add_region(region);
3058
3059 window
3060 .add_typed_entry(
3061 "conv",
3062 leviath_core::EntryKind::Text,
3063 "Assistant: I can help with that.".to_string(),
3064 10,
3065 )
3066 .unwrap();
3067
3068 let assembled = window.assemble();
3069
3070 let assistant_msgs: Vec<_> = assembled
3071 .messages
3072 .iter()
3073 .filter(|m| m.role == "assistant")
3074 .collect();
3075 assert_eq!(assistant_msgs.len(), 1);
3076 assert_eq!(assistant_msgs[0].content.as_text(), "I can help with that.");
3077 }
3078
3079 #[test]
3080 fn test_assemble_text_entry_with_user_prefix() {
3081 let mut window = ContextWindow::new(100_000);
3082 let region = Region::new(
3083 "conv".to_string(),
3084 RegionKind::SlidingWindow {
3085 max_items: 100,
3086 eviction_strategy: EvictionStrategy::PerItem,
3087 },
3088 50_000,
3089 );
3090 window.add_region(region);
3091
3092 window
3093 .add_typed_entry(
3094 "conv",
3095 leviath_core::EntryKind::Text,
3096 "User: What is Rust?".to_string(),
3097 10,
3098 )
3099 .unwrap();
3100
3101 let assembled = window.assemble();
3102
3103 let user_msgs: Vec<_> = assembled
3104 .messages
3105 .iter()
3106 .filter(|m| m.role == "user")
3107 .collect();
3108 assert_eq!(user_msgs.len(), 1);
3109 assert_eq!(user_msgs[0].content.as_text(), "What is Rust?");
3110 }
3111
3112 #[test]
3113 fn test_assemble_text_entry_without_prefix_defaults_to_user() {
3114 let mut window = ContextWindow::new(100_000);
3115 let region = Region::new(
3116 "conv".to_string(),
3117 RegionKind::SlidingWindow {
3118 max_items: 100,
3119 eviction_strategy: EvictionStrategy::PerItem,
3120 },
3121 50_000,
3122 );
3123 window.add_region(region);
3124
3125 window
3126 .add_typed_entry(
3127 "conv",
3128 leviath_core::EntryKind::Text,
3129 "some plain text".to_string(),
3130 10,
3131 )
3132 .unwrap();
3133
3134 let assembled = window.assemble();
3135
3136 let user_msgs: Vec<_> = assembled
3137 .messages
3138 .iter()
3139 .filter(|m| m.role == "user")
3140 .collect();
3141 assert_eq!(user_msgs.len(), 1);
3142 assert_eq!(user_msgs[0].content.as_text(), "some plain text");
3143 }
3144
3145 #[test]
3148 fn test_assemble_assistant_turn_with_text_and_tool_calls() {
3149 let mut window = ContextWindow::new(100_000);
3150 let region = Region::new(
3151 "conv".to_string(),
3152 RegionKind::SlidingWindow {
3153 max_items: 100,
3154 eviction_strategy: EvictionStrategy::PerItem,
3155 },
3156 50_000,
3157 );
3158 window.add_region(region);
3159
3160 window
3162 .add_typed_entry(
3163 "conv",
3164 leviath_core::EntryKind::UserMessage,
3165 "Read my file".to_string(),
3166 10,
3167 )
3168 .unwrap();
3169
3170 window
3172 .add_typed_entry(
3173 "conv",
3174 leviath_core::EntryKind::AssistantTurn {
3175 tool_calls: vec![leviath_core::SerializedToolCall {
3176 id: "tc_a".to_string(),
3177 name: "read_file".to_string(),
3178 arguments: serde_json::json!({"path": "foo.rs"}),
3179 thought_signature: None,
3180 }],
3181 },
3182 "Sure, let me read it.".to_string(),
3183 20,
3184 )
3185 .unwrap();
3186
3187 window
3189 .add_typed_entry(
3190 "conv",
3191 leviath_core::EntryKind::ToolResult {
3192 tool_call_id: "tc_a".to_string(),
3193 tool_name: "read_file".to_string(),
3194 is_error: false,
3195 },
3196 "fn main() {}".to_string(),
3197 10,
3198 )
3199 .unwrap();
3200
3201 let assembled = window.assemble();
3202
3203 let assistant_msg = assembled
3205 .messages
3206 .iter()
3207 .find(|m| m.role == "assistant")
3208 .expect("should have assistant message");
3209
3210 assert_eq!(
3213 assistant_msg.content,
3214 leviath_providers::MessageContent::Blocks(vec![
3215 leviath_providers::ContentBlock::Text {
3216 text: "Sure, let me read it.".to_string(),
3217 },
3218 leviath_providers::ContentBlock::ToolUse {
3219 id: "tc_a".to_string(),
3220 name: "read_file".to_string(),
3221 input: serde_json::json!({"path": "foo.rs"}),
3222 thought_signature: None,
3223 },
3224 ])
3225 );
3226 }
3227
3228 #[test]
3229 fn test_assemble_assistant_turn_no_text_only_tool_calls() {
3230 let mut window = ContextWindow::new(100_000);
3231 let region = Region::new(
3232 "conv".to_string(),
3233 RegionKind::SlidingWindow {
3234 max_items: 100,
3235 eviction_strategy: EvictionStrategy::PerItem,
3236 },
3237 50_000,
3238 );
3239 window.add_region(region);
3240
3241 window
3243 .add_typed_entry(
3244 "conv",
3245 leviath_core::EntryKind::UserMessage,
3246 "Do it".to_string(),
3247 10,
3248 )
3249 .unwrap();
3250
3251 window
3253 .add_typed_entry(
3254 "conv",
3255 leviath_core::EntryKind::AssistantTurn {
3256 tool_calls: vec![leviath_core::SerializedToolCall {
3257 id: "tc_b".to_string(),
3258 name: "bash".to_string(),
3259 arguments: serde_json::json!({"cmd": "ls"}),
3260 thought_signature: None,
3261 }],
3262 },
3263 "".to_string(),
3264 10,
3265 )
3266 .unwrap();
3267
3268 window
3270 .add_typed_entry(
3271 "conv",
3272 leviath_core::EntryKind::ToolResult {
3273 tool_call_id: "tc_b".to_string(),
3274 tool_name: "bash".to_string(),
3275 is_error: false,
3276 },
3277 "file1.rs\nfile2.rs".to_string(),
3278 10,
3279 )
3280 .unwrap();
3281
3282 let assembled = window.assemble();
3283
3284 let assistant_msg = assembled
3285 .messages
3286 .iter()
3287 .find(|m| m.role == "assistant")
3288 .expect("should have assistant message");
3289
3290 assert_eq!(
3292 assistant_msg.content,
3293 leviath_providers::MessageContent::Blocks(vec![
3294 leviath_providers::ContentBlock::ToolUse {
3295 id: "tc_b".to_string(),
3296 name: "bash".to_string(),
3297 input: serde_json::json!({"cmd": "ls"}),
3298 thought_signature: None,
3299 },
3300 ])
3301 );
3302 }
3303
3304 #[test]
3307 fn test_assemble_consecutive_tool_results_flushed_on_non_tool_result() {
3308 let mut window = ContextWindow::new(100_000);
3309 let region = Region::new(
3310 "conv".to_string(),
3311 RegionKind::SlidingWindow {
3312 max_items: 100,
3313 eviction_strategy: EvictionStrategy::PerItem,
3314 },
3315 50_000,
3316 );
3317 window.add_region(region);
3318
3319 window
3321 .add_typed_entry(
3322 "conv",
3323 leviath_core::EntryKind::UserMessage,
3324 "Run two tools".to_string(),
3325 10,
3326 )
3327 .unwrap();
3328
3329 window
3331 .add_typed_entry(
3332 "conv",
3333 leviath_core::EntryKind::AssistantTurn {
3334 tool_calls: vec![
3335 leviath_core::SerializedToolCall {
3336 id: "tc_1".to_string(),
3337 name: "read_file".to_string(),
3338 arguments: serde_json::json!({"path": "a.rs"}),
3339 thought_signature: None,
3340 },
3341 leviath_core::SerializedToolCall {
3342 id: "tc_2".to_string(),
3343 name: "read_file".to_string(),
3344 arguments: serde_json::json!({"path": "b.rs"}),
3345 thought_signature: None,
3346 },
3347 ],
3348 },
3349 "".to_string(),
3350 10,
3351 )
3352 .unwrap();
3353
3354 window
3356 .add_typed_entry(
3357 "conv",
3358 leviath_core::EntryKind::ToolResult {
3359 tool_call_id: "tc_1".to_string(),
3360 tool_name: "read_file".to_string(),
3361 is_error: false,
3362 },
3363 "content of a.rs".to_string(),
3364 10,
3365 )
3366 .unwrap();
3367 window
3368 .add_typed_entry(
3369 "conv",
3370 leviath_core::EntryKind::ToolResult {
3371 tool_call_id: "tc_2".to_string(),
3372 tool_name: "read_file".to_string(),
3373 is_error: false,
3374 },
3375 "content of b.rs".to_string(),
3376 10,
3377 )
3378 .unwrap();
3379
3380 window
3382 .add_typed_entry(
3383 "conv",
3384 leviath_core::EntryKind::UserMessage,
3385 "Now fix the bug".to_string(),
3386 10,
3387 )
3388 .unwrap();
3389
3390 let assembled = window.assemble();
3391
3392 assert_eq!(assembled.messages.len(), 4);
3395
3396 let tool_result_msg = &assembled.messages[2];
3398 assert_eq!(tool_result_msg.role, "user");
3399 assert_eq!(
3402 tool_result_msg.content,
3403 leviath_providers::MessageContent::Blocks(vec![
3404 leviath_providers::ContentBlock::ToolResult {
3405 tool_use_id: "tc_1".to_string(),
3406 content: "content of a.rs".to_string(),
3407 is_error: false,
3408 },
3409 leviath_providers::ContentBlock::ToolResult {
3410 tool_use_id: "tc_2".to_string(),
3411 content: "content of b.rs".to_string(),
3412 is_error: false,
3413 },
3414 ])
3415 );
3416
3417 assert_eq!(assembled.messages[3].role, "user");
3419 assert_eq!(assembled.messages[3].content.as_text(), "Now fix the bug");
3420 }
3421
3422 #[test]
3425 fn test_assemble_injects_begin_when_no_user_messages() {
3426 let mut window = ContextWindow::new(100_000);
3427 let mut pinned = Region::new("system".to_string(), RegionKind::Pinned, 10_000);
3429 pinned
3430 .add_entry("You are a helpful assistant.".to_string(), 20)
3431 .unwrap();
3432 window.add_region(pinned);
3433
3434 let assembled = window.assemble();
3435
3436 assert_eq!(assembled.messages.len(), 1);
3438 assert_eq!(assembled.messages[0].role, "user");
3439 assert_eq!(assembled.messages[0].content.as_text(), "Begin.");
3440 }
3441
3442 #[test]
3445 fn test_add_typed_entry_to_nonexistent_region() {
3446 let mut window = ContextWindow::new(10000);
3447 let result = window.add_typed_entry(
3448 "nonexistent",
3449 leviath_core::EntryKind::UserMessage,
3450 "hello".to_string(),
3451 10,
3452 );
3453 assert!(result.is_err());
3454 let err_str = result.unwrap_err().to_string();
3455 assert!(
3456 err_str.contains("nonexistent"),
3457 "Error should mention the missing region name"
3458 );
3459 }
3460
3461 #[test]
3462 fn test_add_typed_tainted_to_nonexistent_region() {
3463 let mut window = ContextWindow::new(10000);
3464 let result = window.add_typed_tainted_to_region(
3465 "ghost",
3466 leviath_core::EntryKind::UserMessage,
3467 "data".to_string(),
3468 10,
3469 leviath_core::TaintLevel::Public,
3470 );
3471 assert!(result.is_err());
3472 let err_str = result.unwrap_err().to_string();
3473 assert!(
3474 err_str.contains("ghost"),
3475 "Error should mention the missing region name"
3476 );
3477 }
3478
3479 #[test]
3480 fn test_assemble_tool_result_before_any_tool_use() {
3481 let mut window = ContextWindow::new(100_000);
3483 let region = Region::new(
3484 "conversation".to_string(),
3485 RegionKind::SlidingWindow {
3486 max_items: 100,
3487 eviction_strategy: EvictionStrategy::PerItem,
3488 },
3489 50_000,
3490 );
3491 window.add_region(region);
3492
3493 window
3495 .add_typed_entry(
3496 "conversation",
3497 leviath_core::EntryKind::ToolResult {
3498 tool_call_id: "tc_nowhere".to_string(),
3499 tool_name: "read_file".to_string(),
3500 is_error: false,
3501 },
3502 "orphan result".to_string(),
3503 10,
3504 )
3505 .unwrap();
3506
3507 let assembled = with_tracing(|| window.assemble());
3508
3509 assert_eq!(assembled.messages.len(), 1);
3512 assert_eq!(assembled.messages[0].role, "user");
3513 assert_eq!(
3514 assembled.messages[0].content,
3515 leviath_providers::MessageContent::Text("Begin.".to_string())
3516 );
3517 }
3518
3519 #[test]
3522 fn test_assemble_sets_cache_breakpoint_on_stable_prefix() {
3523 let mut window = ContextWindow::new(100_000);
3524 let region = Region::new(
3525 "conv".to_string(),
3526 RegionKind::SlidingWindow {
3527 max_items: 100,
3528 eviction_strategy: EvictionStrategy::PerItem,
3529 },
3530 50_000,
3531 );
3532 window.add_region(region);
3533
3534 for i in 0..10 {
3536 let kind = if i % 2 == 0 {
3537 leviath_core::EntryKind::UserMessage
3538 } else {
3539 leviath_core::EntryKind::AssistantTurn { tool_calls: vec![] }
3540 };
3541 window
3542 .add_typed_entry("conv", kind, format!("message {i}"), 10)
3543 .unwrap();
3544 }
3545
3546 let assembled = window.assemble();
3547 assert_eq!(assembled.messages.len(), 11);
3550 assert_eq!(assembled.messages.last().unwrap().role, "user");
3551
3552 let bp_idx = 6;
3555 for (i, msg) in assembled.messages.iter().enumerate() {
3556 if i == bp_idx {
3557 assert!(
3558 msg.cache_breakpoint,
3559 "Message at index {i} should have cache_breakpoint = true"
3560 );
3561 } else {
3562 assert!(
3563 !msg.cache_breakpoint,
3564 "Message at index {i} should have cache_breakpoint = false"
3565 );
3566 }
3567 }
3568 }
3569
3570 #[test]
3571 fn test_assemble_cache_breakpoint_small_conversation() {
3572 let mut window = ContextWindow::new(100_000);
3573 let region = Region::new(
3574 "conv".to_string(),
3575 RegionKind::SlidingWindow {
3576 max_items: 100,
3577 eviction_strategy: EvictionStrategy::PerItem,
3578 },
3579 50_000,
3580 );
3581 window.add_region(region);
3582
3583 window
3585 .add_typed_entry(
3586 "conv",
3587 leviath_core::EntryKind::UserMessage,
3588 "Hello".to_string(),
3589 10,
3590 )
3591 .unwrap();
3592 window
3593 .add_typed_entry(
3594 "conv",
3595 leviath_core::EntryKind::AssistantTurn { tool_calls: vec![] },
3596 "Hi there".to_string(),
3597 10,
3598 )
3599 .unwrap();
3600 window
3601 .add_typed_entry(
3602 "conv",
3603 leviath_core::EntryKind::UserMessage,
3604 "How are you?".to_string(),
3605 10,
3606 )
3607 .unwrap();
3608
3609 let assembled = window.assemble();
3610 assert_eq!(assembled.messages.len(), 3);
3611
3612 assert!(
3614 assembled.messages[0].cache_breakpoint,
3615 "First message should have cache_breakpoint in small conversation"
3616 );
3617 assert!(!assembled.messages[1].cache_breakpoint);
3618 assert!(!assembled.messages[2].cache_breakpoint);
3619 }
3620
3621 #[test]
3622 fn test_assemble_cache_breakpoint_too_few_messages() {
3623 let mut window = ContextWindow::new(100_000);
3624 let region = Region::new(
3625 "conv".to_string(),
3626 RegionKind::SlidingWindow {
3627 max_items: 100,
3628 eviction_strategy: EvictionStrategy::PerItem,
3629 },
3630 50_000,
3631 );
3632 window.add_region(region);
3633
3634 window
3636 .add_typed_entry(
3637 "conv",
3638 leviath_core::EntryKind::UserMessage,
3639 "Solo message".to_string(),
3640 10,
3641 )
3642 .unwrap();
3643
3644 let assembled = window.assemble();
3645 assert_eq!(assembled.messages.len(), 1);
3646
3647 assert!(
3649 !assembled.messages[0].cache_breakpoint,
3650 "Single message should not get a cache breakpoint"
3651 );
3652 }
3653
3654 #[test]
3655 fn test_assemble_system_blocks_sorted_by_cache_stability() {
3656 use leviath_core::CacheHint;
3657
3658 let mut window = ContextWindow::new(100_000);
3659
3660 let mut clearable = Region::new("scratch".to_string(), RegionKind::Clearable, 10_000);
3662 clearable
3663 .add_entry("clearable data".to_string(), 20)
3664 .unwrap();
3665 window.add_region(clearable);
3666
3667 let mut temporary = Region::new("temp".to_string(), RegionKind::Temporary, 10_000);
3668 temporary
3669 .add_entry("temporary data".to_string(), 20)
3670 .unwrap();
3671 window.add_region(temporary);
3672
3673 let mut compacting = Region::new(
3674 "impl".to_string(),
3675 RegionKind::Compacting {
3676 threshold_tokens: 500,
3677 },
3678 10_000,
3679 );
3680 compacting
3681 .add_entry("compacting data".to_string(), 20)
3682 .unwrap();
3683 window.add_region(compacting);
3684
3685 let mut pinned = Region::new("system".to_string(), RegionKind::Pinned, 10_000);
3686 pinned
3687 .add_entry("pinned system prompt".to_string(), 20)
3688 .unwrap();
3689 window.add_region(pinned);
3690
3691 let assembled = window.assemble();
3692
3693 assert_eq!(assembled.system_blocks.len(), 4);
3694
3695 assert_eq!(
3698 assembled.system_blocks[0].cache_hint,
3699 CacheHint::Always,
3700 "First system block should be Always (Pinned)"
3701 );
3702 assert_eq!(
3703 assembled.system_blocks[1].cache_hint,
3704 CacheHint::UntilChanged,
3705 "Second system block should be UntilChanged (Compacting)"
3706 );
3707 assert_eq!(
3708 assembled.system_blocks[2].cache_hint,
3709 CacheHint::Never,
3710 "Third system block should be Never"
3711 );
3712 assert_eq!(
3713 assembled.system_blocks[3].cache_hint,
3714 CacheHint::Never,
3715 "Fourth system block should be Never"
3716 );
3717 }
3718
3719 #[test]
3722 fn test_add_typed_tainted_to_region_success() {
3723 let mut window = ContextWindow::new(10000);
3724 let mut region = Region::new(
3725 "conv".to_string(),
3726 RegionKind::SlidingWindow {
3727 max_items: 50,
3728 eviction_strategy: EvictionStrategy::PerItem,
3729 },
3730 5000,
3731 );
3732 region.enable_taint_tracking();
3733 window.add_region(region);
3734
3735 window
3736 .add_typed_tainted_to_region(
3737 "conv",
3738 leviath_core::EntryKind::ToolResult {
3739 tool_call_id: "tc_1".to_string(),
3740 tool_name: "read_file".to_string(),
3741 is_error: false,
3742 },
3743 "secret data".to_string(),
3744 100,
3745 leviath_core::TaintLevel::Private,
3746 )
3747 .unwrap();
3748
3749 assert_eq!(window.current_tokens, 100);
3750 assert_eq!(
3751 window.get_region("conv").and_then(|r| r.taint_level()),
3752 Some(leviath_core::TaintLevel::Private)
3753 );
3754 }
3755
3756 #[test]
3757 fn test_add_typed_tainted_to_region_not_found() {
3758 let mut window = ContextWindow::new(10000);
3759 let result = window.add_typed_tainted_to_region(
3760 "nonexistent",
3761 leviath_core::EntryKind::Text,
3762 "data".to_string(),
3763 10,
3764 leviath_core::TaintLevel::Public,
3765 );
3766 assert!(result.is_err());
3767 }
3768
3769 #[test]
3770 fn test_assemble_consecutive_tool_results_flushed_at_end() {
3771 let mut window = ContextWindow::new(100_000);
3774 let region = Region::new(
3775 "conv".to_string(),
3776 RegionKind::SlidingWindow {
3777 max_items: 100,
3778 eviction_strategy: EvictionStrategy::PerItem,
3779 },
3780 50_000,
3781 );
3782 window.add_region(region);
3783
3784 window
3786 .add_typed_entry(
3787 "conv",
3788 leviath_core::EntryKind::UserMessage,
3789 "do something".to_string(),
3790 10,
3791 )
3792 .unwrap();
3793 window
3794 .add_typed_entry(
3795 "conv",
3796 leviath_core::EntryKind::AssistantTurn {
3797 tool_calls: vec![leviath_core::SerializedToolCall {
3798 id: "tc_1".to_string(),
3799 name: "read_file".to_string(),
3800 arguments: serde_json::json!({"path": "foo.rs"}),
3801 thought_signature: None,
3802 }],
3803 },
3804 "Let me read that".to_string(),
3805 10,
3806 )
3807 .unwrap();
3808 window
3809 .add_typed_entry(
3810 "conv",
3811 leviath_core::EntryKind::ToolResult {
3812 tool_call_id: "tc_1".to_string(),
3813 tool_name: "read_file".to_string(),
3814 is_error: false,
3815 },
3816 "fn main() {}".to_string(),
3817 10,
3818 )
3819 .unwrap();
3820
3821 let assembled = window.assemble();
3822 assert_eq!(assembled.messages.len(), 3);
3824 assert_eq!(assembled.messages[2].role, "user");
3825 assert_eq!(
3827 assembled.messages[2].content,
3828 leviath_providers::MessageContent::Blocks(vec![
3829 leviath_providers::ContentBlock::ToolResult {
3830 tool_use_id: "tc_1".to_string(),
3831 content: "fn main() {}".to_string(),
3832 is_error: false,
3833 },
3834 ])
3835 );
3836 }
3837
3838 #[test]
3839 fn test_assemble_compact_history_with_sliding_prefix_sorting() {
3840 use leviath_core::CacheHint;
3842
3843 let mut window = ContextWindow::new(100_000);
3844
3845 let mut temp = Region::new("temp".to_string(), RegionKind::Temporary, 10_000);
3846 temp.add_entry("temp data".to_string(), 10).unwrap();
3847 window.add_region(temp);
3848
3849 let mut history = Region::new(
3850 "history".to_string(),
3851 RegionKind::CompactHistory {
3852 source_region: "impl".to_string(),
3853 },
3854 10_000,
3855 );
3856 history.add_entry("summary data".to_string(), 10).unwrap();
3857 window.add_region(history);
3858
3859 let assembled = window.assemble();
3860 assert_eq!(assembled.system_blocks.len(), 2);
3861 assert_eq!(assembled.system_blocks[0].cache_hint, CacheHint::Always);
3863 assert_eq!(assembled.system_blocks[1].cache_hint, CacheHint::Never);
3864 }
3865
3866 #[test]
3867 fn cache_hint_sort_priority_orders_by_stability() {
3868 use leviath_core::CacheHint;
3869 assert_eq!(cache_hint_sort_priority(CacheHint::Always), 0);
3871 assert_eq!(
3872 cache_hint_sort_priority(CacheHint::SlidingPrefix {
3873 stable_fraction: 0.75
3874 }),
3875 1
3876 );
3877 assert_eq!(cache_hint_sort_priority(CacheHint::UntilChanged), 2);
3878 assert_eq!(cache_hint_sort_priority(CacheHint::Never), 3);
3879 assert!(
3881 cache_hint_sort_priority(CacheHint::Always)
3882 < cache_hint_sort_priority(CacheHint::SlidingPrefix {
3883 stable_fraction: 0.5
3884 })
3885 );
3886 }
3887
3888 #[test]
3889 fn test_assemble_empty_regions_skipped() {
3890 let mut window = ContextWindow::new(100_000);
3891 window.add_region(Region::new(
3892 "system".to_string(),
3893 RegionKind::Pinned,
3894 10_000,
3895 ));
3896 let assembled = window.assemble();
3898 assert!(assembled.system_blocks.is_empty());
3899 }
3900
3901 #[test]
3902 fn test_assemble_hashmap_region_with_keys() {
3903 let mut window = ContextWindow::new(100_000);
3904 let mut region = Region::new(
3905 "files".to_string(),
3906 RegionKind::HashMap { max_entries: None },
3907 10_000,
3908 );
3909 region
3910 .upsert_by_key("src/main.rs", "fn main() {}".to_string(), 10)
3911 .unwrap();
3912 region
3913 .upsert_by_key("src/lib.rs", "pub mod foo;".to_string(), 8)
3914 .unwrap();
3915 window.add_region(region);
3916
3917 let assembled = window.assemble();
3918 assert_eq!(assembled.system_blocks.len(), 1);
3919 let block_text = &assembled.system_blocks[0].text;
3920 assert!(block_text.contains("[files]:"));
3921 assert!(block_text.contains("### [src/main.rs]"));
3922 assert!(block_text.contains("fn main() {}"));
3923 assert!(block_text.contains("### [src/lib.rs]"));
3924 assert!(block_text.contains("pub mod foo;"));
3925 }
3926
3927 #[test]
3928 fn test_assemble_hashmap_region_cache_hint() {
3929 let mut window = ContextWindow::new(100_000);
3930 let mut region = Region::new(
3931 "files".to_string(),
3932 RegionKind::HashMap { max_entries: None },
3933 10_000,
3934 );
3935 region
3936 .upsert_by_key("a.rs", "content".to_string(), 5)
3937 .unwrap();
3938 window.add_region(region);
3939
3940 let assembled = window.assemble();
3941 assert_eq!(assembled.system_blocks.len(), 1);
3942 assert_eq!(
3943 assembled.system_blocks[0].cache_hint,
3944 leviath_core::CacheHint::UntilChanged
3945 );
3946 }
3947
3948 #[test]
3951 fn test_assemble_hashmap_single_keyed_entry() {
3952 let mut window = ContextWindow::new(100_000);
3953 let mut region = Region::new(
3954 "context".to_string(),
3955 RegionKind::HashMap { max_entries: None },
3956 10_000,
3957 );
3958 region
3959 .upsert_by_key("config.toml", "key = \"value\"".to_string(), 10)
3960 .unwrap();
3961 window.add_region(region);
3962
3963 let assembled = window.assemble();
3964
3965 assert_eq!(assembled.system_blocks.len(), 1);
3966 let block_text = &assembled.system_blocks[0].text;
3967 assert!(
3968 block_text.starts_with("[context]:"),
3969 "System block should start with [region_name]: prefix"
3970 );
3971 assert!(
3972 block_text.contains("### [config.toml]"),
3973 "Entry should have ### [key] header"
3974 );
3975 assert!(
3976 block_text.contains("key = \"value\""),
3977 "Entry content should be present"
3978 );
3979 }
3980
3981 #[test]
3982 fn test_assemble_hashmap_multiple_keyed_entries() {
3983 let mut window = ContextWindow::new(100_000);
3984 let mut region = Region::new(
3985 "tracked_files".to_string(),
3986 RegionKind::HashMap { max_entries: None },
3987 10_000,
3988 );
3989 region
3990 .upsert_by_key("alpha.rs", "fn alpha() {}".to_string(), 10)
3991 .unwrap();
3992 region
3993 .upsert_by_key("beta.rs", "fn beta() {}".to_string(), 10)
3994 .unwrap();
3995 region
3996 .upsert_by_key("gamma.rs", "fn gamma() {}".to_string(), 10)
3997 .unwrap();
3998 window.add_region(region);
3999
4000 let assembled = window.assemble();
4001
4002 assert_eq!(assembled.system_blocks.len(), 1);
4003 let block_text = &assembled.system_blocks[0].text;
4004 assert!(block_text.starts_with("[tracked_files]:"));
4005 assert!(block_text.contains("### [alpha.rs]"));
4006 assert!(block_text.contains("fn alpha() {}"));
4007 assert!(block_text.contains("### [beta.rs]"));
4008 assert!(block_text.contains("fn beta() {}"));
4009 assert!(block_text.contains("### [gamma.rs]"));
4010 assert!(block_text.contains("fn gamma() {}"));
4011 }
4012
4013 #[test]
4014 fn test_assemble_hashmap_empty_region_skipped() {
4015 let mut window = ContextWindow::new(100_000);
4016 let region = Region::new(
4017 "empty_map".to_string(),
4018 RegionKind::HashMap { max_entries: None },
4019 10_000,
4020 );
4021 window.add_region(region);
4023
4024 let assembled = window.assemble();
4025
4026 assert!(
4027 assembled.system_blocks.is_empty(),
4028 "Empty HashMap region should not produce a system block"
4029 );
4030 }
4031
4032 #[test]
4033 fn test_assemble_mixed_pinned_hashmap_sliding_window() {
4034 use leviath_core::CacheHint;
4035
4036 let mut window = ContextWindow::new(100_000);
4037
4038 let mut pinned = Region::new("system".to_string(), RegionKind::Pinned, 10_000);
4040 pinned
4041 .add_entry("You are a helpful assistant.".to_string(), 20)
4042 .unwrap();
4043 window.add_region(pinned);
4044
4045 let mut hashmap = Region::new(
4047 "files".to_string(),
4048 RegionKind::HashMap { max_entries: None },
4049 10_000,
4050 );
4051 hashmap
4052 .upsert_by_key("main.rs", "fn main() {}".to_string(), 10)
4053 .unwrap();
4054 window.add_region(hashmap);
4055
4056 let mut sliding = Region::new(
4058 "conv".to_string(),
4059 RegionKind::SlidingWindow {
4060 max_items: 100,
4061 eviction_strategy: EvictionStrategy::PerItem,
4062 },
4063 50_000,
4064 );
4065 sliding
4066 .add_typed_entry(
4067 "Hello there".to_string(),
4068 10,
4069 leviath_core::EntryKind::UserMessage,
4070 )
4071 .unwrap();
4072 window.add_region(sliding);
4073
4074 let assembled = window.assemble();
4075
4076 assert_eq!(assembled.system_blocks.len(), 2);
4078
4079 assert_eq!(
4081 assembled.system_blocks[0].cache_hint,
4082 CacheHint::Always,
4083 "Pinned region should sort first (Always cache hint)"
4084 );
4085 assert!(
4086 assembled.system_blocks[0]
4087 .text
4088 .contains("You are a helpful assistant."),
4089 "First system block should be the pinned content"
4090 );
4091
4092 assert_eq!(
4093 assembled.system_blocks[1].cache_hint,
4094 CacheHint::UntilChanged,
4095 "HashMap region should sort second (UntilChanged cache hint)"
4096 );
4097 assert!(
4098 assembled.system_blocks[1].text.starts_with("[files]:"),
4099 "HashMap system block should have [region_name]: prefix"
4100 );
4101 assert!(
4102 assembled.system_blocks[1].text.contains("### [main.rs]"),
4103 "HashMap system block should contain ### [key] header"
4104 );
4105
4106 assert!(
4108 assembled
4109 .messages
4110 .iter()
4111 .any(|m| m.role == "user" && m.content.as_text().contains("Hello there")),
4112 "SlidingWindow entries should appear as messages"
4113 );
4114 }
4115
4116 #[test]
4117 fn test_add_tainted_to_region_propagates_budget_error() {
4118 let mut window = ContextWindow::new(10_000);
4121 let mut region = Region::new("conv".to_string(), RegionKind::Temporary, 10);
4122 region.enable_taint_tracking();
4123 window.add_region(region);
4124
4125 let result = window.add_tainted_to_region(
4126 "conv",
4127 "far too many tokens".to_string(),
4128 100,
4129 leviath_core::TaintLevel::Private,
4130 );
4131 assert!(result.is_err());
4132 }
4133
4134 #[test]
4135 fn test_add_typed_tainted_to_region_propagates_budget_error() {
4136 let mut window = ContextWindow::new(10_000);
4139 let region = Region::new("conv".to_string(), RegionKind::Temporary, 10);
4140 window.add_region(region);
4141
4142 let result = window.add_typed_tainted_to_region(
4143 "conv",
4144 leviath_core::EntryKind::Text,
4145 "far too many tokens".to_string(),
4146 100,
4147 leviath_core::TaintLevel::Public,
4148 );
4149 assert!(result.is_err());
4150 }
4151
4152 #[test]
4153 fn test_assemble_hashmap_region_entry_without_key() {
4154 let mut window = ContextWindow::new(10_000);
4157 let region = Region::new(
4158 "kv".to_string(),
4159 RegionKind::HashMap { max_entries: None },
4160 5000,
4161 );
4162 window.add_region(region);
4163 window
4165 .add_to_region("kv", "keyless content".to_string(), 10)
4166 .unwrap();
4167
4168 let assembled = window.assemble();
4169 assert!(
4170 assembled
4171 .system_blocks
4172 .iter()
4173 .any(|b| b.text.contains("keyless content")),
4174 "keyless HashMap entry should appear verbatim in a system block"
4175 );
4176 }
4177
4178 #[test]
4184 fn status_labels_are_fixed() {
4185 assert_eq!(AgentStatus::Idle.label(), "idle");
4186 assert_eq!(AgentStatus::Active.label(), "active");
4187 assert_eq!(AgentStatus::Waiting.label(), "waiting");
4188 assert_eq!(AgentStatus::Paused.label(), "paused");
4189 assert_eq!(AgentStatus::Complete.label(), "complete");
4190 assert_eq!(AgentStatus::Cancelled.label(), "cancelled");
4191 assert_eq!(
4192 AgentStatus::Error {
4193 message: "boom".to_string()
4194 }
4195 .label(),
4196 "error"
4197 );
4198 }
4199
4200 #[test]
4203 fn display_matches_label_except_for_an_error() {
4204 for status in [
4205 AgentStatus::Idle,
4206 AgentStatus::Active,
4207 AgentStatus::Waiting,
4208 AgentStatus::Paused,
4209 AgentStatus::Complete,
4210 AgentStatus::Cancelled,
4211 ] {
4212 assert_eq!(status.to_string(), status.label());
4213 }
4214 assert_eq!(
4215 AgentStatus::Error {
4216 message: "disk full".to_string()
4217 }
4218 .to_string(),
4219 "error: disk full"
4220 );
4221 }
4222
4223 #[test]
4224 fn wait_reasons_read_as_short_phrases() {
4225 assert_eq!(WaitReason::ToolApproval.to_string(), "tool approval");
4226 assert_eq!(WaitReason::UserPrompt.to_string(), "user prompt");
4227 assert_eq!(WaitReason::TaintGate.to_string(), "taint gate");
4228 assert_eq!(WaitReason::InteractionPoint.to_string(), "checkpoint");
4229 assert_eq!(
4230 WaitReason::FanOutWorkers { outstanding: 4 }.to_string(),
4231 "workers(4)"
4232 );
4233 assert_eq!(
4234 WaitReason::Children { outstanding: 1 }.to_string(),
4235 "children(1)"
4236 );
4237 }
4238
4239 #[test]
4242 fn only_prompts_need_a_person() {
4243 for reason in [
4244 WaitReason::ToolApproval,
4245 WaitReason::UserPrompt,
4246 WaitReason::TaintGate,
4247 WaitReason::InteractionPoint,
4248 ] {
4249 assert!(reason.needs_a_person(), "{reason} is blocked on someone");
4250 }
4251 for reason in [
4252 WaitReason::FanOutWorkers { outstanding: 2 },
4253 WaitReason::Children { outstanding: 2 },
4254 ] {
4255 assert!(!reason.needs_a_person(), "{reason} resolves on its own");
4256 }
4257 }
4258}