1use super::compression::CompressionPipeline;
2use super::config::{ContextConfig, PromptBudgetConfig};
3use super::partitions::ContextPartitions;
4use super::policy::ContextPolicy;
5use super::pressure::{PressureAction, PressureMonitor};
6use super::renderer::RenderedContext;
7use super::renewal::RenewalPolicy;
8use super::skill_catalog::SkillCatalog;
9use super::task_state::{TaskState, TaskUpdate};
10use super::token_engine::ContextTokenEngine;
11use crate::mm::handle::{Handle, HandleId, HandleKind, HandleTable, Residency};
12use crate::types::capability::{Capability, CapabilityKind, CapabilityManifest};
13use crate::types::message::{Content, ContentPart, Message, ToolSchema};
14use crate::types::skill::SkillMetadata;
15use compact_str::CompactString;
16
17pub const MEMORY_TOOL_NAME: &str = "memory";
18pub const KNOWLEDGE_TOOL_NAME: &str = "knowledge";
19pub const READ_RESULT_TOOL_NAME: &str = "read_result";
21
22const META_TOOL_NAMES: &[&str] = &[
25 "pace",
26 "update_plan",
27 "skill",
28 MEMORY_TOOL_NAME,
29 KNOWLEDGE_TOOL_NAME,
30 READ_RESULT_TOOL_NAME,
31 "submit_workflow_nodes",
32 "start_workflow",
33];
34
35pub(crate) fn is_meta_tool(name: &str) -> bool {
38 META_TOOL_NAMES.contains(&name)
39}
40
41pub(crate) const EXPOSURE_EXEMPT_META_TOOLS: &[&str] = &[
59 "skill",
60 MEMORY_TOOL_NAME,
61 KNOWLEDGE_TOOL_NAME,
62 "update_plan",
63 READ_RESULT_TOOL_NAME,
64];
65
66pub(crate) fn is_exposure_exempt_meta_tool(name: &str) -> bool {
69 EXPOSURE_EXEMPT_META_TOOLS.contains(&name)
70}
71
72#[doc(hidden)]
77pub struct ContextManager {
78 pub partitions: ContextPartitions,
79 pub max_tokens: u32,
80 pub config: ContextConfig,
81 pub engine: ContextTokenEngine,
82 pub prompt_budget: PromptBudgetConfig,
85 pub sprint: u32,
86 pub skills: SkillCatalog,
87 pub active_skills: std::collections::BTreeMap<CompactString, Option<u32>>,
95 pub stable_core_tools: std::collections::HashSet<CompactString>,
99 pub capabilities: CapabilityManifest,
100 pub memory_enabled: bool,
101 pub knowledge_enabled: bool,
102 pub plan_tool_enabled: bool,
103 last_observed_prompt_tokens: Option<u32>,
104 compression: CompressionPipeline,
105 pressure: PressureMonitor,
106 renewal: RenewalPolicy,
107
108 pub last_activity_ms: u64,
112
113 pub last_compact_ms: Option<u64>,
116
117 pub handles: HandleTable,
122 next_handle_id: HandleId,
124
125 frozen_history_len: usize,
132
133 pending_knowledge_sweeps: Vec<crate::context::partitions::KnowledgeSweep>,
136
137 knowledge_budget_warned: bool,
140 knowledge_reference_step: u64,
142}
143
144impl ContextManager {
145 pub fn new(max_tokens: u32) -> Self {
146 Self::with_config(
149 max_tokens,
150 ContextConfig::default(),
151 ContextTokenEngine::fallback_estimator(),
152 )
153 }
154
155 pub fn with_config(max_tokens: u32, config: ContextConfig, engine: ContextTokenEngine) -> Self {
156 let compression = CompressionPipeline::new(&config);
157 let pressure = PressureMonitor::new(max_tokens, config.clone());
158 let renewal = RenewalPolicy::from_config(&config);
159 let partitions = ContextPartitions::new(&config);
160 Self {
161 partitions,
162 max_tokens,
163 config,
164 engine,
165 prompt_budget: PromptBudgetConfig::default(),
166 sprint: 0,
167 skills: SkillCatalog::new(),
168 active_skills: std::collections::BTreeMap::new(),
169 stable_core_tools: std::collections::HashSet::new(),
170 capabilities: CapabilityManifest::new(),
171 memory_enabled: false,
172 knowledge_enabled: false,
173 plan_tool_enabled: false,
174 last_observed_prompt_tokens: None,
175 compression,
176 pressure,
177 renewal,
178 last_activity_ms: 0,
179 last_compact_ms: None,
180 handles: HandleTable::new(),
181 next_handle_id: 0,
182 frozen_history_len: 0,
183 pending_knowledge_sweeps: Vec::new(),
184 knowledge_budget_warned: false,
185 knowledge_reference_step: 0,
186 }
187 }
188
189 pub fn apply_context_policy(&mut self, policy: &ContextPolicy) {
191 policy.apply_to(&mut self.config);
192 self.compression = CompressionPipeline::new(&self.config);
193 self.pressure = PressureMonitor::new(self.max_tokens, self.config.clone());
194 self.renewal = RenewalPolicy::from_config(&self.config);
195 }
196
197 pub fn record_activity(&mut self, now_ms: u64) {
201 self.last_activity_ms = now_ms;
202 }
203
204 pub fn should_time_decay_compact(&self, now_ms: u64) -> bool {
207 let idle_ms = if let Some(last_compact) = self.last_compact_ms {
208 now_ms.saturating_sub(last_compact)
210 } else {
211 now_ms.saturating_sub(self.last_activity_ms)
213 };
214
215 let idle_minutes = idle_ms / 60_000;
216 idle_minutes >= self.config.micro_compact_idle_minutes as u64
217 }
218
219 pub fn recompute_handle_residency(&mut self) {
234 if self.rho() < self.config.collapse_threshold {
236 return;
237 }
238 let keep = self.config.preserved_tool_results;
239 let total = self
242 .handles
243 .all()
244 .iter()
245 .filter(|h| matches!(h.kind, HandleKind::ToolResult))
246 .count();
247 let cutoff = total.saturating_sub(keep);
248 for (i, handle) in self.handles.tool_result_handles_mut().enumerate() {
249 if i < cutoff && matches!(handle.residency, Residency::Resident) {
252 handle.residency = Residency::Collapsed;
253 }
254 }
255 }
256
257 pub fn reset_collapse_generation(&mut self) {
263 for handle in self.handles.all_mut() {
264 if matches!(handle.residency, Residency::Collapsed) {
265 handle.residency = Residency::Resident;
266 }
267 }
268 }
269
270 pub fn prune_orphaned_handles(&mut self) {
282 let live: std::collections::HashSet<CompactString> = self
283 .partitions
284 .history
285 .messages
286 .iter()
287 .flat_map(|m| match &m.content {
288 Content::Parts(parts) => parts
289 .iter()
290 .filter_map(|p| match p {
291 ContentPart::ToolResult { call_id, .. } => Some(call_id.clone()),
292 _ => None,
293 })
294 .collect::<Vec<_>>(),
295 _ => Vec::new(),
296 })
297 .collect();
298 self.handles.retain(|h| {
299 h.residency.payload_ref().is_some()
300 || h.source.as_ref().is_none_or(|s| live.contains(s))
301 });
302 }
303
304 pub fn set_payload_residency(
319 &mut self,
320 source: &str,
321 kind: HandleKind,
322 tokens: u32,
323 residency: Residency,
324 ) -> Option<Residency> {
325 if let Some(handle) = self
326 .handles
327 .all_mut()
328 .iter_mut()
329 .find(|h| h.source.as_deref() == Some(source))
330 {
331 let previous = std::mem::replace(&mut handle.residency, residency);
332 handle.tokens = tokens;
333 return Some(previous);
334 }
335 let id = self.alloc_handle_id();
336 self.handles.insert(Handle {
337 id,
338 kind,
339 residency,
340 tokens,
341 source: Some(source.into()),
342 });
343 None
344 }
345
346 pub fn payload_residency(&self, source: &str) -> Option<&Residency> {
353 self.handles.residency_for_source(source)
354 }
355
356 pub fn rho(&self) -> f64 {
363 self.pressure.pressure(
364 &self.partitions,
365 &self.engine,
366 self.last_observed_prompt_tokens,
367 )
368 }
369
370 pub fn set_observed_prompt_tokens(&mut self, tokens: u32) {
371 self.last_observed_prompt_tokens = Some(tokens);
372 }
373
374 pub fn should_compress(&self) -> PressureAction {
375 self.pressure.recommend(self.rho())
381 }
382
383 pub fn compress(
384 &mut self,
385 action: PressureAction,
386 ) -> (u32, Option<String>, Vec<Message>, Option<usize>) {
387 self.compress_with_time(action, None)
388 }
389
390 pub fn compress_with_time(
391 &mut self,
392 action: PressureAction,
393 now_ms: Option<u64>,
394 ) -> (u32, Option<String>, Vec<Message>, Option<usize>) {
395 let target = self.config.target_tokens(self.max_tokens);
396 self.compress_with_target(action, target, now_ms)
397 }
398
399 pub fn force_compress(&mut self) -> (u32, Option<String>, Vec<Message>, Option<usize>) {
400 self.compress_with_target(PressureAction::AutoCompact, 0, None)
401 }
402
403 pub fn compress_with_target(
410 &mut self,
411 action: PressureAction,
412 target_tokens: u32,
413 now_ms: Option<u64>,
414 ) -> (u32, Option<String>, Vec<Message>, Option<usize>) {
415 let result = self.compression.compress(
416 &mut self.partitions,
417 action,
418 self.max_tokens,
419 target_tokens,
420 &self.engine,
421 );
422 if let Some(ts) = now_ms {
423 self.last_compact_ms = Some(ts);
424 }
425 if !result.2.is_empty() {
427 self.prune_orphaned_handles();
428 self.reset_collapse_generation();
431 self.sweep_knowledge_at_boundary();
434 }
435 if result.3.is_some() {
442 self.frozen_history_len = self.partitions.history.messages.len();
443 }
444 result
445 }
446
447 pub fn plan_compaction_params(&self) -> (u32, usize) {
451 (
452 self.config.target_tokens(self.max_tokens),
453 self.config.preserve_recent_turns,
454 )
455 }
456
457 pub fn should_renew(&self) -> bool {
460 self.renewal
461 .should_renew(&self.pressure, &self.partitions, &self.engine)
462 }
463
464 pub fn renew(&mut self) {
465 self.partitions = self.renewal.renew(&self.partitions, self.max_tokens);
466 self.sprint += 1;
467 self.prune_orphaned_handles();
470 self.reset_collapse_generation();
471 self.sweep_knowledge_at_boundary();
473 self.frozen_history_len = self.partitions.history.messages.len();
475 }
476
477 pub fn set_prompt_budget(&mut self, prompt_budget: PromptBudgetConfig) {
480 self.prompt_budget = prompt_budget;
481 }
482
483 pub fn available_input_tokens(&self) -> u32 {
484 self.max_tokens
485 .saturating_sub(self.prompt_budget.reserved_tokens())
486 }
487
488 pub fn render(&self) -> RenderedContext {
489 super::renderer::render_projected(
490 &self.partitions,
491 self.available_input_tokens(),
492 &self.engine,
493 self.config.preserve_recent_units,
494 &self.handles,
495 self.frozen_history_len,
496 self.config.collapse_assistant_narration,
497 )
498 }
499
500 pub fn push_history(&mut self, msg: Message, tokens: u32) {
503 self.knowledge_reference_step = self.knowledge_reference_step.saturating_add(1);
504 self.partitions
505 .knowledge
506 .observe_references(&msg, self.knowledge_reference_step);
507 if let Content::Parts(parts) = &msg.content {
511 for part in parts {
512 if let ContentPart::ToolResult {
513 call_id, output, ..
514 } = part
515 {
516 let id = self.alloc_handle_id();
517 let tok = self.engine.count(output).max(1);
518 self.handles.insert(Handle::resident_for(
519 id,
520 HandleKind::ToolResult,
521 tok,
522 call_id.clone(),
523 ));
524 }
525 }
526 }
527 self.partitions.history.push(msg, tokens);
528 }
529
530 fn alloc_handle_id(&mut self) -> HandleId {
531 let id = self.next_handle_id;
532 self.next_handle_id = self.next_handle_id.wrapping_add(1);
533 id
534 }
535
536 pub fn next_handle_id(&self) -> HandleId {
541 self.next_handle_id
542 }
543
544 pub fn restore_next_handle_id(&mut self, next: HandleId) {
550 self.next_handle_id = next;
551 }
552
553 pub fn frozen_history_len(&self) -> usize {
554 self.frozen_history_len
555 }
556
557 pub fn restore_frozen_history_len(&mut self, len: usize) -> bool {
558 if len > self.partitions.history.messages.len() {
559 return false;
560 }
561 self.frozen_history_len = len;
562 true
563 }
564
565 pub fn push_knowledge(&mut self, msg: Message, tokens: u32) {
567 self.partitions.knowledge.push(msg, tokens);
568 }
569
570 pub fn push_knowledge_entry(
574 &mut self,
575 key: Option<CompactString>,
576 msg: Message,
577 tokens: u32,
578 pinned: bool,
579 ) {
580 self.partitions
581 .knowledge
582 .push_entry(key, msg, tokens, pinned);
583 }
584
585 pub fn remove_knowledge(&mut self, key: &str) -> bool {
588 self.partitions.knowledge.remove(key)
589 }
590
591 fn sweep_knowledge_at_boundary(&mut self) {
595 let sweep = self.partitions.knowledge.sweep_at_boundary();
596 if sweep.changed {
597 if !sweep.removed_keys.is_empty() {
601 self.partitions.signals.push(format!(
602 "[KNOWLEDGE] entries removed at this boundary: {} — re-fetch via the memory tool if still needed.",
603 sweep.removed_keys.join(", ")
604 ));
605 }
606 self.pending_knowledge_sweeps.push(sweep);
607 }
608 self.knowledge_budget_warned = false;
610 }
611
612 pub fn enforce_knowledge_budget(&mut self) -> Option<(u32, u32)> {
621 let ratio = self.config.knowledge_budget_ratio;
622 if ratio <= 0.0 {
623 return None;
624 }
625 let budget = (self.max_tokens as f64 * ratio) as u32;
626 let used = self.partitions.knowledge.token_count;
627 if used <= budget {
628 return None;
629 }
630 let marked: u32 = self
631 .partitions
632 .knowledge
633 .entries
634 .iter()
635 .filter(|e| e.evict_at_boundary)
636 .map(|e| e.tokens)
637 .sum();
638 let mut projected = used.saturating_sub(marked);
639 let mut candidates = self
640 .partitions
641 .knowledge
642 .entries
643 .iter()
644 .enumerate()
645 .filter(|(_, entry)| {
646 !entry.evict_at_boundary
647 && !entry.pinned
648 && !entry
649 .key
650 .as_deref()
651 .is_some_and(|key| key.starts_with("skill:"))
652 })
653 .map(|(index, _)| {
654 let score = self
655 .partitions
656 .knowledge
657 .retention_score(index, self.knowledge_reference_step)
658 .unwrap_or(i64::MIN);
659 (score, index)
660 })
661 .collect::<Vec<_>>();
662 candidates.sort_by(|left, right| left.0.cmp(&right.0).then(left.1.cmp(&right.1)));
663 for (_, index) in candidates {
664 if projected <= budget {
665 break;
666 }
667 let entry = &mut self.partitions.knowledge.entries[index];
668 entry.evict_at_boundary = true;
669 projected = projected.saturating_sub(entry.tokens);
670 }
671 if self.knowledge_budget_warned {
672 return None;
673 }
674 self.knowledge_budget_warned = true;
675 Some((used, budget))
676 }
677
678 pub fn take_knowledge_sweeps(&mut self) -> Vec<crate::context::partitions::KnowledgeSweep> {
680 std::mem::take(&mut self.pending_knowledge_sweeps)
681 }
682
683 pub fn push_signal(&mut self, text: String) {
687 self.partitions.signals.push(text);
688 }
689
690 pub fn record_directive(&mut self, text: impl Into<String>) {
694 self.partitions.task_state.record_directive(text);
695 }
696
697 pub fn init_task(&mut self, goal: String, criteria: Vec<String>) {
700 self.partitions.task_state = TaskState {
701 goal,
702 criteria,
703 ..Default::default()
704 };
705 }
706
707 pub fn update_task(&mut self, update: TaskUpdate) {
708 self.partitions.task_state.apply(update);
709 }
710
711 pub fn note_tool_actions(&mut self, calls: &[(String, String)]) {
718 let summary = calls
719 .iter()
720 .filter(|(name, _)| !is_meta_tool(name))
721 .map(|(name, args)| {
722 if args.is_empty() {
723 name.clone()
724 } else {
725 format!("{name}({args})")
726 }
727 })
728 .collect::<Vec<_>>()
729 .join(", ");
730 self.partitions.task_state.note_actions(summary);
731 }
732
733 pub fn set_available_skills(&mut self, skills: Vec<SkillMetadata>) {
738 self.capabilities.remove_kind(CapabilityKind::Skill);
739 for skill in &skills {
740 self.capabilities.add_skill(skill.clone());
741 }
742 self.skills.set_available(skills);
743 }
744
745 pub fn set_stable_core_tools(&mut self, ids: impl IntoIterator<Item = CompactString>) {
747 self.stable_core_tools = ids.into_iter().collect();
748 }
749
750 pub fn activate_skill(&mut self, name: impl Into<CompactString>) -> bool {
755 self.activate_skill_leased(name, None)
756 }
757
758 pub fn activate_skill_leased(
761 &mut self,
762 name: impl Into<CompactString>,
763 expires_at_turn: Option<u32>,
764 ) -> bool {
765 self.active_skills
766 .insert(name.into(), expires_at_turn)
767 .is_none()
768 }
769
770 pub fn deactivate_skill(&mut self, name: &str) -> bool {
774 if self.active_skills.remove(name).is_none() {
775 return false;
776 }
777 self.partitions.knowledge.remove(&format!("skill:{name}"));
778 true
779 }
780
781 pub fn sweep_expired_skill_leases(&mut self, current_turn: u32) {
784 let expired: Vec<CompactString> = self
785 .active_skills
786 .iter()
787 .filter(|(_, lease)| lease.is_some_and(|t| current_turn >= t))
788 .map(|(name, _)| name.clone())
789 .collect();
790 for name in expired {
791 self.deactivate_skill(&name);
792 self.partitions.signals.push(format!(
794 "[SKILL] lease expired: {name} unloaded; the full toolset is restored."
795 ));
796 }
797 }
798
799 pub fn active_skill_tool_filter(&self) -> Option<std::collections::HashSet<CompactString>> {
805 if self.active_skills.is_empty() {
806 return None;
807 }
808 let mut union = std::collections::HashSet::new();
809 for name in self.active_skills.keys() {
810 let declared = self.skills.allowed_tools(name);
811 if declared.is_empty() {
812 return None; }
814 union.extend(declared.iter().cloned());
815 }
816 Some(union)
817 }
818
819 pub fn active_skill_capabilities(&self) -> Vec<Capability> {
823 self.active_skills
824 .keys()
825 .flat_map(|name| self.skills.capability_grants(name).iter().cloned())
826 .collect()
827 }
828
829 pub fn skill_capability_grants(&self, name: &str) -> &[Capability] {
830 self.skills.capability_grants(name)
831 }
832
833 pub fn skill_tool_schema(&self) -> Option<ToolSchema> {
834 self.skills.build_tool_schema()
835 }
836
837 pub fn skill_available(&self, name: &str) -> bool {
840 self.skills.is_available(name)
841 }
842
843 pub fn set_memory_enabled(&mut self, enabled: bool) {
846 self.memory_enabled = enabled;
847 if enabled {
848 self.capabilities.add_marker(
849 CapabilityKind::Memory,
850 MEMORY_TOOL_NAME,
851 "Search long-term memory through the memory meta-tool.",
852 );
853 } else {
854 self.capabilities
855 .remove(CapabilityKind::Memory, MEMORY_TOOL_NAME);
856 }
857 }
858
859 pub fn set_knowledge_enabled(&mut self, enabled: bool) {
860 self.knowledge_enabled = enabled;
861 if enabled {
862 self.capabilities.add_marker(
863 CapabilityKind::Knowledge,
864 KNOWLEDGE_TOOL_NAME,
865 "Search external knowledge through the knowledge meta-tool.",
866 );
867 } else {
868 self.capabilities
869 .remove(CapabilityKind::Knowledge, KNOWLEDGE_TOOL_NAME);
870 }
871 }
872
873 pub fn set_plan_tool_enabled(&mut self, enabled: bool) {
874 self.plan_tool_enabled = enabled;
875 if enabled {
876 self.capabilities.add_marker(
877 CapabilityKind::Tool,
878 "update_plan",
879 "Update task plan and progress through the planning meta-tool.",
880 );
881 } else {
882 self.capabilities
883 .remove(CapabilityKind::Tool, "update_plan");
884 }
885 }
886
887 pub fn capability_inventory(&self) -> String {
888 self.capabilities.format_inventory()
889 }
890
891 pub fn meta_tool_schemas(&self) -> Vec<ToolSchema> {
892 let mut tools = Vec::new();
893 if let Some(t) = self.skill_tool_schema() {
894 tools.push(t);
895 }
896 if let Some(t) = self.memory_tool_schema() {
897 tools.push(t);
898 }
899 if let Some(t) = self.knowledge_tool_schema() {
900 tools.push(t);
901 }
902 if let Some(t) = self.plan_tool_schema() {
903 tools.push(t);
904 }
905 if let Some(t) = self.read_result_tool_schema() {
906 tools.push(t);
907 }
908 tools.sort_by(|a, b| a.name.cmp(&b.name));
909 tools
910 }
911
912 pub fn read_result_tool_schema(&self) -> Option<ToolSchema> {
918 let any_evicted = self
919 .handles
920 .all()
921 .iter()
922 .any(|h| !h.residency.occupies_context());
923 if !any_evicted {
924 return None;
925 }
926 Some(ToolSchema {
927 name: CompactString::new(READ_RESULT_TOOL_NAME),
928 description: "Re-read the full content of a tool result that was truncated to save \
929 context (marked '[Output truncated: … call the read_result tool …]'). \
930 Pass that marker's call_id; use offset/max_bytes to page through \
931 large content."
932 .to_string(),
933 parameters: serde_json::json!({
934 "type": "object",
935 "properties": {
936 "call_id": { "type": "string" },
937 "offset": { "type": "integer", "description": "Byte offset to start from (default 0)." },
938 "max_bytes": { "type": "integer", "description": "Max bytes to return (default 4000)." }
939 },
940 "required": ["call_id"]
941 }),
942 })
943 }
944
945 pub fn plan_tool_schema(&self) -> Option<ToolSchema> {
946 if !self.plan_tool_enabled {
947 return None;
948 }
949 Some(ToolSchema {
950 name: CompactString::new("update_plan"),
951 description: "Update your task plan and progress. Call this after completing a step or when the plan changes.".to_string(),
952 parameters: serde_json::json!({
953 "type": "object",
954 "properties": {
955 "plan": { "type": "array", "items": { "type": "string" } },
956 "current_step": { "type": "integer" },
957 "progress": { "type": "string" },
958 "blocked_on": { "type": "array", "items": { "type": "string" } }
959 }
960 }),
961 })
962 }
963
964 pub fn memory_tool_schema(&self) -> Option<ToolSchema> {
965 if !self.memory_enabled {
966 return None;
967 }
968 Some(ToolSchema {
969 name: CompactString::new(MEMORY_TOOL_NAME),
970 description:
971 "Search your long-term memory for relevant past experiences and knowledge."
972 .to_string(),
973 parameters: serde_json::json!({
974 "type": "object",
975 "properties": {
976 "query": { "type": "string" },
977 "top_k": { "type": "integer" }
978 },
979 "required": ["query"]
980 }),
981 })
982 }
983
984 pub fn knowledge_tool_schema(&self) -> Option<ToolSchema> {
985 if !self.knowledge_enabled {
986 return None;
987 }
988 Some(ToolSchema {
989 name: CompactString::new(KNOWLEDGE_TOOL_NAME),
990 description:
991 "Search the external knowledge base for facts, documentation, or reference data."
992 .to_string(),
993 parameters: serde_json::json!({
994 "type": "object",
995 "properties": {
996 "query": { "type": "string" },
997 "top_k": { "type": "integer" }
998 },
999 "required": ["query"]
1000 }),
1001 })
1002 }
1003}
1004
1005#[cfg(test)]
1006mod tests {
1007 use super::*;
1008 use crate::context::task_state::PlanStep;
1009 use crate::types::message::Message;
1010 use crate::types::skill::SkillMetadata;
1011
1012 #[test]
1013 fn note_tool_actions_keys_on_name_and_args_so_legit_loops_dont_false_stop() {
1014 let mut mgr = ContextManager::new(100_000);
1017 mgr.init_task("process items".to_string(), vec![]);
1018 mgr.note_tool_actions(&[("step".to_string(), "{\"n\":1}".to_string())]);
1019 mgr.note_tool_actions(&[("step".to_string(), "{\"n\":2}".to_string())]);
1020 mgr.note_tool_actions(&[("step".to_string(), "{\"n\":3}".to_string())]);
1021 assert_eq!(
1022 mgr.partitions.task_state.recent_actions,
1023 ["step({\"n\":1})", "step({\"n\":2})", "step({\"n\":3})"]
1024 );
1025 let txt = mgr
1026 .render()
1027 .state_turn
1028 .unwrap()
1029 .content
1030 .as_text()
1031 .unwrap()
1032 .to_string();
1033 assert!(
1034 !txt.contains("STOP:"),
1035 "same-tool/diff-args loop must not trip STOP: {txt}"
1036 );
1037
1038 let mut mgr2 = ContextManager::new(100_000);
1040 mgr2.init_task("g".to_string(), vec![]);
1041 for _ in 0..3 {
1042 mgr2.note_tool_actions(&[("document_read".to_string(), "{\"id\":\"x\"}".to_string())]);
1043 }
1044 let txt2 = mgr2
1045 .render()
1046 .state_turn
1047 .unwrap()
1048 .content
1049 .as_text()
1050 .unwrap()
1051 .to_string();
1052 assert!(
1053 txt2.contains("STOP:"),
1054 "identical repeated call must trip STOP: {txt2}"
1055 );
1056
1057 let mut mgr3 = ContextManager::new(100_000);
1059 mgr3.init_task("g".to_string(), vec![]);
1060 mgr3.note_tool_actions(&[(
1061 "update_plan".to_string(),
1062 "{\"current_step\":1}".to_string(),
1063 )]);
1064 assert!(mgr3.partitions.task_state.recent_actions.is_empty());
1065 }
1066
1067 #[test]
1068 fn manager_renew_advances_sprint_and_keeps_goal() {
1069 let mut mgr = ContextManager::new(1_000);
1070 mgr.init_task("test goal".to_string(), vec![]);
1071 mgr.partitions.system.push(Message::system("rules"), 10);
1072 for i in 0..10 {
1073 mgr.push_history(Message::user(format!("msg {i}")), 50);
1074 }
1075 mgr.renew();
1076 assert_eq!(mgr.partitions.task_state.goal, "test goal");
1077 assert_eq!(mgr.sprint, 1);
1078 }
1079
1080 #[test]
1081 fn compress_only_touches_history() {
1082 let mut mgr = ContextManager::new(1_000);
1083 mgr.push_knowledge(Message::system("knowledge content"), 100);
1084 for _ in 0..30 {
1085 mgr.push_history(Message::user("history msg"), 50);
1086 }
1087 let knowledge_before = mgr.partitions.knowledge.token_count;
1088 let history_before = mgr.partitions.history.token_count;
1089 mgr.compress(PressureAction::AutoCompact);
1090 assert_eq!(mgr.partitions.knowledge.token_count, knowledge_before);
1091 assert!(mgr.partitions.history.token_count < history_before);
1092 }
1093
1094 #[test]
1095 fn init_task_sets_goal_and_criteria() {
1096 let mut mgr = ContextManager::new(1_000);
1097 mgr.init_task("analyse data".to_string(), vec!["criterion A".to_string()]);
1098 assert_eq!(mgr.partitions.task_state.goal, "analyse data");
1099 assert_eq!(mgr.partitions.task_state.criteria, ["criterion A"]);
1100 }
1101
1102 #[test]
1103 fn update_task_applies_plan() {
1104 let mut mgr = ContextManager::new(1_000);
1105 mgr.init_task("g".to_string(), vec![]);
1106 mgr.update_task(TaskUpdate {
1107 plan: Some(vec!["step 1".to_string(), "step 2".to_string()]),
1108 current_step: Some(0),
1109 ..Default::default()
1110 });
1111 assert_eq!(mgr.partitions.task_state.plan.len(), 2);
1112 assert_eq!(mgr.partitions.task_state.current_step, Some(0));
1113 }
1114
1115 #[test]
1116 fn task_state_survives_autocompact() {
1117 let mut mgr = ContextManager::new(1_000);
1118 mgr.init_task("survive compression".to_string(), vec![]);
1119 mgr.update_task(TaskUpdate {
1120 plan: Some(vec!["fetch data".to_string(), "analyse".to_string()]),
1121 ..Default::default()
1122 });
1123 for _ in 0..10 {
1124 mgr.push_history(Message::user("filler"), 50);
1125 }
1126 mgr.compress(PressureAction::AutoCompact);
1127 assert_eq!(mgr.partitions.task_state.goal, "survive compression");
1128 assert_eq!(mgr.partitions.task_state.plan.len(), 2);
1129 }
1130
1131 #[test]
1132 fn render_includes_task_state_in_state_turn_not_system() {
1133 let mut mgr = ContextManager::new(10_000);
1134 mgr.init_task("find anomalies".to_string(), vec![]);
1135 let rc = mgr.render();
1136 assert!(
1137 !rc.system_text.contains("[TASK STATE]"),
1138 "task_state must not be in system_text"
1139 );
1140 let state = rc.state_turn.as_ref().expect("should have a state turn");
1142 assert!(
1143 state
1144 .content
1145 .as_text()
1146 .unwrap()
1147 .contains("[TASK STATE] goal: find anomalies")
1148 );
1149 }
1150
1151 #[test]
1152 fn renewal_keeps_open_plan_steps_in_task_state() {
1153 let mut mgr = ContextManager::new(1_000);
1154 mgr.init_task("g".to_string(), vec![]);
1155 mgr.partitions.task_state.plan = vec![
1156 PlanStep {
1157 label: "done".to_string(),
1158 done: true,
1159 },
1160 PlanStep {
1161 label: "pending".to_string(),
1162 done: false,
1163 },
1164 ];
1165 mgr.renew();
1166 assert_eq!(mgr.partitions.task_state.open_steps(), vec!["pending"]);
1167 }
1168
1169 #[test]
1172 fn auto_compact_entry_logs_auto_compact_action() {
1173 let mut mgr = ContextManager::new(1_000);
1181 for i in 0..40 {
1182 mgr.push_history(
1183 Message::user(format!("turn {i}: {}", "ctx ".repeat(40))),
1184 200,
1185 );
1186 }
1187 let (saved, summary, _, _) = mgr.force_compress();
1188 assert!(saved > 0, "force_compress should compact a large history");
1189 assert!(
1190 summary.is_some(),
1191 "auto-compact summarizes the archived turns"
1192 );
1193 let actions: Vec<&str> = mgr
1194 .partitions
1195 .task_state
1196 .compression_log
1197 .iter()
1198 .map(|e| e.action.as_str())
1199 .collect();
1200 assert!(
1201 actions.last() == Some(&"auto_compact"),
1202 "auto-compact entry must log an auto_compact action; got {actions:?}"
1203 );
1204 }
1205
1206 #[test]
1207 fn skill_tool_schema_empty_when_no_skills() {
1208 let mgr = ContextManager::new(10_000);
1209 assert!(mgr.skill_tool_schema().is_none());
1210 }
1211
1212 #[test]
1213 fn skill_tool_schema_present_when_registered() {
1214 let mut mgr = ContextManager::new(10_000);
1215 mgr.set_available_skills(vec![SkillMetadata::new("debug", "Debug helper")]);
1216 assert!(
1217 mgr.skill_tool_schema()
1218 .unwrap()
1219 .description
1220 .contains("debug")
1221 );
1222 }
1223
1224 #[test]
1225 fn available_skills_are_reflected_in_capability_manifest() {
1226 let mut mgr = ContextManager::new(1_000);
1227 mgr.set_available_skills(vec![SkillMetadata::new("debug", "Debug helper")]);
1228 let inventory = mgr.capability_inventory();
1229 assert!(inventory.contains("debug"));
1230 assert!(inventory.contains("Debug helper"));
1231 }
1232
1233 #[test]
1234 fn toggled_meta_tools_are_reflected_in_capability_manifest() {
1235 let mut mgr = ContextManager::new(1_000);
1236 mgr.set_memory_enabled(true);
1237 assert!(mgr.capability_inventory().contains(MEMORY_TOOL_NAME));
1238 mgr.set_memory_enabled(false);
1239 assert!(!mgr.capability_inventory().contains(MEMORY_TOOL_NAME));
1240 }
1241
1242 #[test]
1243 fn meta_tool_schemas_are_sorted() {
1244 let mut mgr = ContextManager::new(1_000);
1245 mgr.set_available_skills(vec![SkillMetadata::new("debug", "Debug helper")]);
1246 mgr.set_memory_enabled(true);
1247 mgr.set_knowledge_enabled(true);
1248 let names = mgr
1249 .meta_tool_schemas()
1250 .into_iter()
1251 .map(|s| s.name.to_string())
1252 .collect::<Vec<_>>();
1253 assert_eq!(names, ["knowledge", "memory", "skill"]);
1254 }
1255
1256 #[test]
1257 fn b1_active_skill_state_and_tool_filter() {
1258 let mut mgr = ContextManager::new(1_000);
1259 let mut debug = SkillMetadata::new("debug", "Debug helper");
1260 debug.allowed_tools = vec![CompactString::new("read"), CompactString::new("grep")];
1261 let mut review = SkillMetadata::new("review", "Reviewer");
1262 review.allowed_tools = vec![CompactString::new("git_diff")];
1263 let plain = SkillMetadata::new("plain", "No tools declared"); mgr.set_available_skills(vec![debug, review, plain]);
1265
1266 assert!(mgr.active_skill_tool_filter().is_none());
1268
1269 assert!(mgr.activate_skill("debug"));
1271 assert!(!mgr.activate_skill("debug")); let f = mgr.active_skill_tool_filter().unwrap();
1275 assert_eq!(f.len(), 2);
1276 assert!(f.contains(&CompactString::new("read")) && f.contains(&CompactString::new("grep")));
1277
1278 mgr.activate_skill("review");
1280 let f = mgr.active_skill_tool_filter().unwrap();
1281 assert_eq!(f.len(), 3);
1282 assert!(f.contains(&CompactString::new("git_diff")));
1283
1284 mgr.activate_skill("plain");
1286 assert!(mgr.active_skill_tool_filter().is_none());
1287 }
1288
1289 #[test]
1290 fn active_skill_capability_grants_follow_activation_deactivation_and_lease_expiry() {
1291 use crate::types::capability::{
1292 ActionSet, Capability, CapabilityId, ConstraintSet, Principal, ResourceSelector,
1293 };
1294
1295 let grant = Capability {
1296 id: CapabilityId("read-src".into()),
1297 kind: CapabilityKind::Tool,
1298 resource: ResourceSelector("/repo/src/**".into()),
1299 actions: ActionSet(["read".into()].into_iter().collect()),
1300 constraints: ConstraintSet::default(),
1301 lease: None,
1302 delegatable: false,
1303 issuer: Principal("root".into()),
1304 };
1305 let mut review = SkillMetadata::new("review", "Review source files");
1306 review.capability_grants = vec![grant.clone()];
1307
1308 let mut mgr = ContextManager::new(1_000);
1309 mgr.set_available_skills(vec![review]);
1310 assert!(mgr.active_skill_capabilities().is_empty());
1311
1312 mgr.activate_skill("review");
1313 assert_eq!(mgr.active_skill_capabilities(), vec![grant.clone()]);
1314
1315 mgr.deactivate_skill("review");
1316 assert!(mgr.active_skill_capabilities().is_empty());
1317
1318 mgr.activate_skill_leased("review", Some(3));
1319 mgr.sweep_expired_skill_leases(2);
1320 assert_eq!(mgr.active_skill_capabilities(), vec![grant]);
1321
1322 mgr.sweep_expired_skill_leases(3);
1323 assert!(mgr.active_skill_capabilities().is_empty());
1324 }
1325
1326 #[test]
1327 fn update_collapse_mode_collapses_old_tool_results_under_pressure() {
1328 let mut mgr = ContextManager::new(1_000);
1329 for i in 0..10 {
1330 let m = Message::tool(vec![ContentPart::ToolResult {
1331 call_id: format!("c{i}").into(),
1332 output: "x".repeat(40),
1333 is_error: false,
1334 durable_content: None,
1335 }]);
1336 mgr.push_history(m, 40);
1337 }
1338 mgr.set_observed_prompt_tokens(950); assert!(mgr.rho() >= mgr.config.collapse_threshold);
1341
1342 mgr.recompute_handle_residency();
1343 assert_eq!(
1345 mgr.handles.residency_for_source("c0"),
1346 Some(&Residency::Collapsed)
1347 );
1348 assert_eq!(
1349 mgr.handles.residency_for_source("c9"),
1350 Some(&Residency::Resident)
1351 );
1352
1353 mgr.set_observed_prompt_tokens(100); mgr.recompute_handle_residency();
1357 assert_eq!(
1358 mgr.handles.residency_for_source("c0"),
1359 Some(&Residency::Collapsed),
1360 "collapse is sticky until a compaction boundary"
1361 );
1362
1363 mgr.reset_collapse_generation();
1365 assert_eq!(
1366 mgr.handles.residency_for_source("c0"),
1367 Some(&Residency::Resident)
1368 );
1369 }
1370
1371 #[test]
1372 fn frozen_prefix_len_anchors_at_compaction_and_holds_across_appends() {
1373 let mut mgr = ContextManager::new(1_000);
1374 for i in 0..30 {
1376 mgr.push_history(
1377 Message::user(format!("turn {i}: {}", "ctx ".repeat(30))),
1378 150,
1379 );
1380 }
1381 assert!(
1382 mgr.render().frozen_prefix_len.is_none(),
1383 "no frozen region before any compaction"
1384 );
1385
1386 let (saved, _, archived, _) = mgr.compress(PressureAction::AutoCompact);
1387 assert!(saved > 0 && !archived.is_empty(), "expected archival");
1388
1389 assert!(
1391 mgr.render().frozen_prefix_len.is_none(),
1392 "deep == tail right after compaction"
1393 );
1394
1395 mgr.push_history(Message::user("new 1"), 5);
1397 let f1 = mgr
1398 .render()
1399 .frozen_prefix_len
1400 .expect("frozen region exists once the tail grows");
1401 mgr.push_history(Message::assistant("reply 1"), 5);
1402 mgr.push_history(Message::user("new 2"), 5);
1403 let rc = mgr.render();
1404 let f2 = rc.frozen_prefix_len.expect("frozen region holds");
1405 assert_eq!(
1406 f1, f2,
1407 "the deep boundary is fixed between compactions; only the tail grows"
1408 );
1409 assert!(
1410 f2 < rc.turns.len(),
1411 "deep boundary is distinct from the rolling tail"
1412 );
1413 }
1414
1415 #[test]
1416 fn frozen_boundary_holds_through_a_prefix_safe_compaction() {
1417 let mut mgr = ContextManager::new(10_000);
1420 for i in 0..5 {
1421 mgr.push_history(Message::user(format!("m{i}")), 5);
1422 }
1423 mgr.frozen_history_len = 3; let (_, _, _, cache_at) = mgr.compress(PressureAction::None);
1428 assert!(cache_at.is_none(), "no-op compaction is prefix-safe");
1429 assert_eq!(
1430 mgr.frozen_history_len, 3,
1431 "prefix-safe compaction preserves the deep-cache anchor"
1432 );
1433 }
1434
1435 #[test]
1436 fn collapse_generation_resets_on_autocompact() {
1437 let mut mgr = ContextManager::new(1_000);
1438 for i in 0..20 {
1441 mgr.push_history(tool_result_msg(&format!("c{i}"), &"x".repeat(120)), 60);
1442 }
1443 mgr.set_observed_prompt_tokens(980); mgr.recompute_handle_residency();
1445 assert_eq!(
1446 mgr.handles.residency_for_source("c0"),
1447 Some(&Residency::Collapsed)
1448 );
1449
1450 let (saved, _, archived, _) = mgr.compress(PressureAction::AutoCompact);
1451 assert!(saved > 0 && !archived.is_empty(), "expected archival");
1452
1453 for h in mgr.handles.all() {
1456 if matches!(h.kind, HandleKind::ToolResult) {
1457 assert_eq!(
1458 h.residency,
1459 Residency::Resident,
1460 "generation reset un-collapses survivors"
1461 );
1462 }
1463 }
1464 }
1465
1466 #[test]
1467 fn push_history_indexes_tool_results_as_resident_handles() {
1468 let mut mgr = ContextManager::new(10_000);
1469 let msg = Message::tool(vec![ContentPart::ToolResult {
1470 call_id: "call_1".into(),
1471 output: "the tool output".to_string(),
1472 is_error: false,
1473 durable_content: None,
1474 }]);
1475 mgr.push_history(msg, 20);
1476 assert_eq!(mgr.handles.all().len(), 1);
1478 assert_eq!(
1479 mgr.handles.residency_for_source("call_1"),
1480 Some(&Residency::Resident)
1481 );
1482 mgr.push_history(Message::user("hello"), 5);
1484 assert_eq!(mgr.handles.all().len(), 1);
1485 }
1486
1487 fn tool_result_msg(call_id: &str, output: &str) -> Message {
1490 Message::tool(vec![ContentPart::ToolResult {
1491 call_id: call_id.into(),
1492 output: output.to_string(),
1493 is_error: false,
1494 durable_content: None,
1495 }])
1496 }
1497
1498 #[test]
1499 fn prune_orphaned_handles_drops_handles_whose_message_left_history() {
1500 let mut mgr = ContextManager::new(10_000);
1501 mgr.push_history(tool_result_msg("c0", "out 0"), 20);
1502 mgr.push_history(tool_result_msg("c1", "out 1"), 20);
1503 assert_eq!(mgr.handles.all().len(), 2);
1504
1505 mgr.partitions.history.messages.remove(0);
1507 mgr.prune_orphaned_handles();
1508
1509 assert_eq!(mgr.handles.all().len(), 1);
1511 assert!(mgr.handles.residency_for_source("c0").is_none());
1512 assert_eq!(
1513 mgr.handles.residency_for_source("c1"),
1514 Some(&Residency::Resident)
1515 );
1516 }
1517
1518 #[test]
1519 fn autocompact_prunes_handles_for_archived_tool_results() {
1520 let mut mgr = ContextManager::new(1_000);
1521 for i in 0..30 {
1523 mgr.push_history(tool_result_msg(&format!("c{i}"), &"x".repeat(200)), 80);
1524 }
1525 assert_eq!(mgr.handles.all().len(), 30);
1526
1527 let (saved, _, archived, _) = mgr.compress(PressureAction::AutoCompact);
1528 assert!(saved > 0 && !archived.is_empty(), "expected archival");
1529
1530 let live_tool_results = mgr
1533 .partitions
1534 .history
1535 .messages
1536 .iter()
1537 .filter(|m| {
1538 matches!(&m.content, Content::Parts(p)
1539 if p.iter().any(|x| matches!(x, ContentPart::ToolResult { .. })))
1540 })
1541 .count();
1542 assert_eq!(mgr.handles.all().len(), live_tool_results);
1543 assert!(
1544 mgr.handles.all().len() < 30,
1545 "table must shrink with archival"
1546 );
1547 }
1548
1549 #[test]
1550 fn renew_prunes_handles_for_dropped_history() {
1551 let mut mgr = ContextManager::new(1_000);
1552 mgr.init_task("g".to_string(), vec![]);
1553 for i in 0..20 {
1554 mgr.push_history(tool_result_msg(&format!("c{i}"), "data"), 60);
1555 }
1556 mgr.renew();
1557 for h in mgr.handles.all() {
1559 if let Some(src) = h.source.as_ref() {
1560 assert!(
1561 mgr.handles.residency_for_source(src).is_some(),
1562 "no dangling handle survives renewal"
1563 );
1564 }
1565 }
1566 assert!(mgr.handles.all().len() <= 20);
1567 }
1568
1569 #[test]
1572 fn knowledge_budget_uses_stable_order_for_equal_value_and_warns_once() {
1573 let mut mgr = ContextManager::new(100);
1576 mgr.push_knowledge(Message::system("oldest unkeyed"), 10);
1577 mgr.push_knowledge_entry(Some("a".into()), Message::system("keyed"), 10, false);
1578 mgr.push_knowledge_entry(Some("p".into()), Message::system("pinned"), 10, true);
1579 mgr.push_knowledge_entry(Some("skill:x".into()), Message::system("skill"), 10, false);
1580
1581 let warn = mgr.enforce_knowledge_budget();
1582 assert_eq!(warn, Some((40, 25)));
1583 let e = &mgr.partitions.knowledge.entries;
1585 assert!(e[0].evict_at_boundary);
1586 assert!(e[1].evict_at_boundary);
1587 assert!(!e[2].evict_at_boundary, "pinned exempt");
1588 assert!(!e[3].evict_at_boundary, "skill pin exempt");
1589
1590 assert_eq!(mgr.enforce_knowledge_budget(), None);
1592
1593 let sweep = mgr.partitions.knowledge.sweep_at_boundary();
1595 assert_eq!(sweep.tokens_freed, 20);
1596 assert_eq!(mgr.partitions.knowledge.token_count, 20);
1597 assert_eq!(mgr.enforce_knowledge_budget(), None);
1599 }
1600
1601 #[test]
1602 fn knowledge_budget_warning_stands_when_only_exempt_weight_remains() {
1603 let mut mgr = ContextManager::new(100);
1604 mgr.push_knowledge_entry(Some("p".into()), Message::system("pinned heavy"), 30, true);
1605 mgr.push_knowledge_entry(
1606 Some("skill:x".into()),
1607 Message::system("skill heavy"),
1608 30,
1609 false,
1610 );
1611
1612 assert_eq!(mgr.enforce_knowledge_budget(), Some((60, 25)));
1614 assert!(
1615 mgr.partitions
1616 .knowledge
1617 .entries
1618 .iter()
1619 .all(|e| !e.evict_at_boundary)
1620 );
1621 }
1622
1623 #[test]
1624 fn knowledge_budget_retains_old_referenced_entry_over_new_irrelevant_entry() {
1625 let mut mgr = ContextManager::new(100);
1626 mgr.push_knowledge_entry(
1627 Some("project:orchid".into()),
1628 Message::system("ORCHID uses the Atlas storage engine"),
1629 10,
1630 false,
1631 );
1632 mgr.push_history(Message::user("For project:orchid keep using Atlas"), 5);
1634 mgr.push_knowledge_entry(
1635 Some("project:new".into()),
1636 Message::system("unrelated fresh material"),
1637 10,
1638 false,
1639 );
1640 mgr.push_knowledge_entry(
1641 Some("project:other".into()),
1642 Message::system("another unused reference"),
1643 10,
1644 false,
1645 );
1646
1647 assert_eq!(mgr.enforce_knowledge_budget(), Some((30, 25)));
1648 let entries = &mgr.partitions.knowledge.entries;
1649 assert!(
1650 !entries[0].evict_at_boundary,
1651 "a real reference must raise retention"
1652 );
1653 assert!(
1654 entries[1].evict_at_boundary,
1655 "lowest-value entry evicts first"
1656 );
1657 assert!(
1658 !entries[2].evict_at_boundary,
1659 "one eviction is enough to fit"
1660 );
1661 }
1662
1663 #[test]
1664 fn knowledge_budget_ratio_zero_disables() {
1665 let mut mgr = ContextManager::new(100);
1666 mgr.config.knowledge_budget_ratio = 0.0;
1667 mgr.push_knowledge(Message::system("huge"), 90);
1668 assert_eq!(mgr.enforce_knowledge_budget(), None);
1669 assert!(!mgr.partitions.knowledge.entries[0].evict_at_boundary);
1670 }
1671
1672 #[test]
1673 fn provider_and_output_reservations_reduce_the_hard_input_budget() {
1674 use crate::context::config::PromptBudgetConfig;
1675
1676 let mut mgr = ContextManager::new(100);
1677 mgr.set_prompt_budget(PromptBudgetConfig {
1678 prompt_overhead_tokens: 20,
1679 output_reserve_tokens: 20,
1680 safety_margin_tokens: 10,
1681 });
1682 mgr.partitions.system.push(
1690 Message::system(
1691 "System policy directive number seven requires strict adherence. ".repeat(10),
1692 ),
1693 60,
1694 );
1695
1696 let rendered = mgr.render();
1697 let overflow = rendered
1698 .budget_overflow
1699 .expect("fixed context exceeds input allowance");
1700 assert_eq!(overflow.max_tokens, 50);
1701 assert!(overflow.required_tokens > overflow.max_tokens);
1702 }
1703}