1use super::compression::CompressionPipeline;
2use super::config::{ContextConfig, PromptBudgetConfig};
3use super::partitions::ContextPartitions;
4use super::policy::ContextPolicyV1;
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::{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(
147 max_tokens,
148 ContextConfig::default(),
149 ContextTokenEngine::char_approx(),
150 )
151 }
152
153 pub fn with_config(max_tokens: u32, config: ContextConfig, engine: ContextTokenEngine) -> Self {
154 let compression = CompressionPipeline::new(&config);
155 let pressure = PressureMonitor::new(max_tokens, config.clone());
156 let renewal = RenewalPolicy::from_config(&config);
157 let partitions = ContextPartitions::new(&config);
158 Self {
159 partitions,
160 max_tokens,
161 config,
162 engine,
163 prompt_budget: PromptBudgetConfig::default(),
164 sprint: 0,
165 skills: SkillCatalog::new(),
166 active_skills: std::collections::BTreeMap::new(),
167 stable_core_tools: std::collections::HashSet::new(),
168 capabilities: CapabilityManifest::new(),
169 memory_enabled: false,
170 knowledge_enabled: false,
171 plan_tool_enabled: false,
172 last_observed_prompt_tokens: None,
173 compression,
174 pressure,
175 renewal,
176 last_activity_ms: 0,
177 last_compact_ms: None,
178 handles: HandleTable::new(),
179 next_handle_id: 0,
180 frozen_history_len: 0,
181 pending_knowledge_sweeps: Vec::new(),
182 knowledge_budget_warned: false,
183 knowledge_reference_step: 0,
184 }
185 }
186
187 pub fn apply_context_policy(&mut self, policy: &ContextPolicyV1) {
189 policy.apply_to(&mut self.config);
190 self.compression = CompressionPipeline::new(&self.config);
191 self.pressure = PressureMonitor::new(self.max_tokens, self.config.clone());
192 self.renewal = RenewalPolicy::from_config(&self.config);
193 }
194
195 pub fn record_activity(&mut self, now_ms: u64) {
199 self.last_activity_ms = now_ms;
200 }
201
202 pub fn should_time_decay_compact(&self, now_ms: u64) -> bool {
205 let idle_ms = if let Some(last_compact) = self.last_compact_ms {
206 now_ms.saturating_sub(last_compact)
208 } else {
209 now_ms.saturating_sub(self.last_activity_ms)
211 };
212
213 let idle_minutes = idle_ms / 60_000;
214 idle_minutes >= self.config.micro_compact_idle_minutes as u64
215 }
216
217 pub fn recompute_handle_residency(&mut self) {
232 if self.rho() < self.config.collapse_threshold {
234 return;
235 }
236 let keep = self.config.preserved_tool_results;
237 let total = self
240 .handles
241 .all()
242 .iter()
243 .filter(|h| matches!(h.kind, HandleKind::ToolResult))
244 .count();
245 let cutoff = total.saturating_sub(keep);
246 for (i, handle) in self.handles.tool_result_handles_mut().enumerate() {
247 if i < cutoff && matches!(handle.residency, Residency::Resident) {
250 handle.residency = Residency::Collapsed;
251 }
252 }
253 }
254
255 pub fn reset_collapse_generation(&mut self) {
261 for handle in self.handles.all_mut() {
262 if matches!(handle.residency, Residency::Collapsed) {
263 handle.residency = Residency::Resident;
264 }
265 }
266 }
267
268 pub fn prune_orphaned_handles(&mut self) {
280 let live: std::collections::HashSet<CompactString> = self
281 .partitions
282 .history
283 .messages
284 .iter()
285 .flat_map(|m| match &m.content {
286 Content::Parts(parts) => parts
287 .iter()
288 .filter_map(|p| match p {
289 ContentPart::ToolResult { call_id, .. } => Some(call_id.clone()),
290 _ => None,
291 })
292 .collect::<Vec<_>>(),
293 _ => Vec::new(),
294 })
295 .collect();
296 self.handles.retain(|h| {
297 h.residency.payload_ref().is_some()
298 || h.source.as_ref().is_none_or(|s| live.contains(s))
299 });
300 }
301
302 pub fn set_payload_residency(
317 &mut self,
318 source: &str,
319 kind: HandleKind,
320 tokens: u32,
321 residency: Residency,
322 ) -> Option<Residency> {
323 if let Some(handle) = self
324 .handles
325 .all_mut()
326 .iter_mut()
327 .find(|h| h.source.as_deref() == Some(source))
328 {
329 let previous = std::mem::replace(&mut handle.residency, residency);
330 handle.tokens = tokens;
331 return Some(previous);
332 }
333 let id = self.alloc_handle_id();
334 self.handles.insert(Handle {
335 id,
336 kind,
337 residency,
338 tokens,
339 source: Some(source.into()),
340 });
341 None
342 }
343
344 pub fn payload_residency(&self, source: &str) -> Option<&Residency> {
351 self.handles.residency_for_source(source)
352 }
353
354 pub fn rho(&self) -> f64 {
361 self.pressure.pressure(
362 &self.partitions,
363 &self.engine,
364 self.last_observed_prompt_tokens,
365 )
366 }
367
368 pub fn set_observed_prompt_tokens(&mut self, tokens: u32) {
369 self.last_observed_prompt_tokens = Some(tokens);
370 }
371
372 pub fn should_compress(&self) -> PressureAction {
373 self.pressure.recommend(self.rho())
379 }
380
381 pub fn compress(
382 &mut self,
383 action: PressureAction,
384 ) -> (u32, Option<String>, Vec<Message>, Option<usize>) {
385 self.compress_with_time(action, None)
386 }
387
388 pub fn compress_with_time(
389 &mut self,
390 action: PressureAction,
391 now_ms: Option<u64>,
392 ) -> (u32, Option<String>, Vec<Message>, Option<usize>) {
393 let target = self.config.target_tokens(self.max_tokens);
394 self.compress_with_target(action, target, now_ms)
395 }
396
397 pub fn force_compress(&mut self) -> (u32, Option<String>, Vec<Message>, Option<usize>) {
398 self.compress_with_target(PressureAction::AutoCompact, 0, None)
399 }
400
401 pub fn compress_with_target(
408 &mut self,
409 action: PressureAction,
410 target_tokens: u32,
411 now_ms: Option<u64>,
412 ) -> (u32, Option<String>, Vec<Message>, Option<usize>) {
413 let result = self.compression.compress(
414 &mut self.partitions,
415 action,
416 self.max_tokens,
417 target_tokens,
418 &self.engine,
419 );
420 if let Some(ts) = now_ms {
421 self.last_compact_ms = Some(ts);
422 }
423 if !result.2.is_empty() {
425 self.prune_orphaned_handles();
426 self.reset_collapse_generation();
429 self.sweep_knowledge_at_boundary();
432 }
433 if result.3.is_some() {
440 self.frozen_history_len = self.partitions.history.messages.len();
441 }
442 result
443 }
444
445 pub fn plan_compaction_params(&self) -> (u32, usize) {
449 (
450 self.config.target_tokens(self.max_tokens),
451 self.config.preserve_recent_turns,
452 )
453 }
454
455 pub fn should_renew(&self) -> bool {
458 self.renewal
459 .should_renew(&self.pressure, &self.partitions, &self.engine)
460 }
461
462 pub fn renew(&mut self) {
463 self.partitions = self.renewal.renew(&self.partitions, self.max_tokens);
464 self.sprint += 1;
465 self.prune_orphaned_handles();
468 self.reset_collapse_generation();
469 self.sweep_knowledge_at_boundary();
471 self.frozen_history_len = self.partitions.history.messages.len();
473 }
474
475 pub fn set_prompt_budget(&mut self, prompt_budget: PromptBudgetConfig) {
478 self.prompt_budget = prompt_budget;
479 }
480
481 pub fn available_input_tokens(&self) -> u32 {
482 self.max_tokens
483 .saturating_sub(self.prompt_budget.reserved_tokens())
484 }
485
486 pub fn render(&self) -> RenderedContext {
487 super::renderer::render_projected(
488 &self.partitions,
489 self.available_input_tokens(),
490 &self.engine,
491 self.config.preserve_recent_units,
492 &self.handles,
493 self.frozen_history_len,
494 self.config.collapse_assistant_narration,
495 )
496 }
497
498 pub fn push_history(&mut self, msg: Message, tokens: u32) {
501 self.knowledge_reference_step = self.knowledge_reference_step.saturating_add(1);
502 self.partitions
503 .knowledge
504 .observe_references(&msg, self.knowledge_reference_step);
505 if let Content::Parts(parts) = &msg.content {
509 for part in parts {
510 if let ContentPart::ToolResult {
511 call_id, output, ..
512 } = part
513 {
514 let id = self.alloc_handle_id();
515 let tok = self.engine.count(output).max(1);
516 self.handles.insert(Handle::resident_for(
517 id,
518 HandleKind::ToolResult,
519 tok,
520 call_id.clone(),
521 ));
522 }
523 }
524 }
525 self.partitions.history.push(msg, tokens);
526 }
527
528 fn alloc_handle_id(&mut self) -> HandleId {
529 let id = self.next_handle_id;
530 self.next_handle_id = self.next_handle_id.wrapping_add(1);
531 id
532 }
533
534 pub fn next_handle_id(&self) -> HandleId {
539 self.next_handle_id
540 }
541
542 pub fn restore_next_handle_id(&mut self, next: HandleId) {
548 self.next_handle_id = next;
549 }
550
551 pub fn frozen_history_len(&self) -> usize {
552 self.frozen_history_len
553 }
554
555 pub fn restore_frozen_history_len(&mut self, len: usize) -> bool {
556 if len > self.partitions.history.messages.len() {
557 return false;
558 }
559 self.frozen_history_len = len;
560 true
561 }
562
563 pub fn push_knowledge(&mut self, msg: Message, tokens: u32) {
565 self.partitions.knowledge.push(msg, tokens);
566 }
567
568 pub fn push_knowledge_entry(
572 &mut self,
573 key: Option<CompactString>,
574 msg: Message,
575 tokens: u32,
576 pinned: bool,
577 ) {
578 self.partitions
579 .knowledge
580 .push_entry(key, msg, tokens, pinned);
581 }
582
583 pub fn remove_knowledge(&mut self, key: &str) -> bool {
586 self.partitions.knowledge.remove(key)
587 }
588
589 fn sweep_knowledge_at_boundary(&mut self) {
593 let sweep = self.partitions.knowledge.sweep_at_boundary();
594 if sweep.changed {
595 if !sweep.removed_keys.is_empty() {
599 self.partitions.signals.push(format!(
600 "[KNOWLEDGE] entries removed at this boundary: {} — re-fetch via the memory tool if still needed.",
601 sweep.removed_keys.join(", ")
602 ));
603 }
604 self.pending_knowledge_sweeps.push(sweep);
605 }
606 self.knowledge_budget_warned = false;
608 }
609
610 pub fn enforce_knowledge_budget(&mut self) -> Option<(u32, u32)> {
619 let ratio = self.config.knowledge_budget_ratio;
620 if ratio <= 0.0 {
621 return None;
622 }
623 let budget = (self.max_tokens as f64 * ratio) as u32;
624 let used = self.partitions.knowledge.token_count;
625 if used <= budget {
626 return None;
627 }
628 let marked: u32 = self
629 .partitions
630 .knowledge
631 .entries
632 .iter()
633 .filter(|e| e.evict_at_boundary)
634 .map(|e| e.tokens)
635 .sum();
636 let mut projected = used.saturating_sub(marked);
637 let mut candidates = self
638 .partitions
639 .knowledge
640 .entries
641 .iter()
642 .enumerate()
643 .filter(|(_, entry)| {
644 !entry.evict_at_boundary
645 && !entry.pinned
646 && !entry
647 .key
648 .as_deref()
649 .is_some_and(|key| key.starts_with("skill:"))
650 })
651 .map(|(index, _)| {
652 let score = self
653 .partitions
654 .knowledge
655 .retention_score(index, self.knowledge_reference_step)
656 .unwrap_or(i64::MIN);
657 (score, index)
658 })
659 .collect::<Vec<_>>();
660 candidates.sort_by(|left, right| left.0.cmp(&right.0).then(left.1.cmp(&right.1)));
661 for (_, index) in candidates {
662 if projected <= budget {
663 break;
664 }
665 let entry = &mut self.partitions.knowledge.entries[index];
666 entry.evict_at_boundary = true;
667 projected = projected.saturating_sub(entry.tokens);
668 }
669 if self.knowledge_budget_warned {
670 return None;
671 }
672 self.knowledge_budget_warned = true;
673 Some((used, budget))
674 }
675
676 pub fn take_knowledge_sweeps(&mut self) -> Vec<crate::context::partitions::KnowledgeSweep> {
678 std::mem::take(&mut self.pending_knowledge_sweeps)
679 }
680
681 pub fn push_signal(&mut self, text: String) {
685 self.partitions.signals.push(text);
686 }
687
688 pub fn record_directive(&mut self, text: impl Into<String>) {
692 self.partitions.task_state.record_directive(text);
693 }
694
695 pub fn init_task(&mut self, goal: String, criteria: Vec<String>) {
698 self.partitions.task_state = TaskState {
699 goal,
700 criteria,
701 ..Default::default()
702 };
703 }
704
705 pub fn update_task(&mut self, update: TaskUpdate) {
706 self.partitions.task_state.apply(update);
707 }
708
709 pub fn note_tool_actions(&mut self, calls: &[(String, String)]) {
716 let summary = calls
717 .iter()
718 .filter(|(name, _)| !is_meta_tool(name))
719 .map(|(name, args)| {
720 if args.is_empty() {
721 name.clone()
722 } else {
723 format!("{name}({args})")
724 }
725 })
726 .collect::<Vec<_>>()
727 .join(", ");
728 self.partitions.task_state.note_actions(summary);
729 }
730
731 pub fn set_available_skills(&mut self, skills: Vec<SkillMetadata>) {
736 self.capabilities.remove_kind(CapabilityKind::Skill);
737 for skill in &skills {
738 self.capabilities.add_skill(skill.clone());
739 }
740 self.skills.set_available(skills);
741 }
742
743 pub fn set_stable_core_tools(&mut self, ids: impl IntoIterator<Item = CompactString>) {
745 self.stable_core_tools = ids.into_iter().collect();
746 }
747
748 pub fn activate_skill(&mut self, name: impl Into<CompactString>) -> bool {
753 self.activate_skill_leased(name, None)
754 }
755
756 pub fn activate_skill_leased(
759 &mut self,
760 name: impl Into<CompactString>,
761 expires_at_turn: Option<u32>,
762 ) -> bool {
763 self.active_skills
764 .insert(name.into(), expires_at_turn)
765 .is_none()
766 }
767
768 pub fn deactivate_skill(&mut self, name: &str) -> bool {
772 if self.active_skills.remove(name).is_none() {
773 return false;
774 }
775 self.partitions.knowledge.remove(&format!("skill:{name}"));
776 true
777 }
778
779 pub fn sweep_expired_skill_leases(&mut self, current_turn: u32) {
782 let expired: Vec<CompactString> = self
783 .active_skills
784 .iter()
785 .filter(|(_, lease)| lease.is_some_and(|t| current_turn >= t))
786 .map(|(name, _)| name.clone())
787 .collect();
788 for name in expired {
789 self.deactivate_skill(&name);
790 self.partitions.signals.push(format!(
792 "[SKILL] lease expired: {name} unloaded; the full toolset is restored."
793 ));
794 }
795 }
796
797 pub fn active_skill_tool_filter(&self) -> Option<std::collections::HashSet<CompactString>> {
803 if self.active_skills.is_empty() {
804 return None;
805 }
806 let mut union = std::collections::HashSet::new();
807 for name in self.active_skills.keys() {
808 let declared = self.skills.allowed_tools(name);
809 if declared.is_empty() {
810 return None; }
812 union.extend(declared.iter().cloned());
813 }
814 Some(union)
815 }
816
817 pub fn skill_tool_schema(&self) -> Option<ToolSchema> {
818 self.skills.build_tool_schema()
819 }
820
821 pub fn skill_available(&self, name: &str) -> bool {
824 self.skills.is_available(name)
825 }
826
827 pub fn set_memory_enabled(&mut self, enabled: bool) {
830 self.memory_enabled = enabled;
831 if enabled {
832 self.capabilities.add_marker(
833 CapabilityKind::Memory,
834 MEMORY_TOOL_NAME,
835 "Search long-term memory through the memory meta-tool.",
836 );
837 } else {
838 self.capabilities
839 .remove(CapabilityKind::Memory, MEMORY_TOOL_NAME);
840 }
841 }
842
843 pub fn set_knowledge_enabled(&mut self, enabled: bool) {
844 self.knowledge_enabled = enabled;
845 if enabled {
846 self.capabilities.add_marker(
847 CapabilityKind::Knowledge,
848 KNOWLEDGE_TOOL_NAME,
849 "Search external knowledge through the knowledge meta-tool.",
850 );
851 } else {
852 self.capabilities
853 .remove(CapabilityKind::Knowledge, KNOWLEDGE_TOOL_NAME);
854 }
855 }
856
857 pub fn set_plan_tool_enabled(&mut self, enabled: bool) {
858 self.plan_tool_enabled = enabled;
859 if enabled {
860 self.capabilities.add_marker(
861 CapabilityKind::Tool,
862 "update_plan",
863 "Update task plan and progress through the planning meta-tool.",
864 );
865 } else {
866 self.capabilities
867 .remove(CapabilityKind::Tool, "update_plan");
868 }
869 }
870
871 pub fn capability_inventory(&self) -> String {
872 self.capabilities.format_inventory()
873 }
874
875 pub fn meta_tool_schemas(&self) -> Vec<ToolSchema> {
876 let mut tools = Vec::new();
877 if let Some(t) = self.skill_tool_schema() {
878 tools.push(t);
879 }
880 if let Some(t) = self.memory_tool_schema() {
881 tools.push(t);
882 }
883 if let Some(t) = self.knowledge_tool_schema() {
884 tools.push(t);
885 }
886 if let Some(t) = self.plan_tool_schema() {
887 tools.push(t);
888 }
889 if let Some(t) = self.read_result_tool_schema() {
890 tools.push(t);
891 }
892 tools.sort_by(|a, b| a.name.cmp(&b.name));
893 tools
894 }
895
896 pub fn read_result_tool_schema(&self) -> Option<ToolSchema> {
902 let any_evicted = self
903 .handles
904 .all()
905 .iter()
906 .any(|h| !h.residency.occupies_context());
907 if !any_evicted {
908 return None;
909 }
910 Some(ToolSchema {
911 name: CompactString::new(READ_RESULT_TOOL_NAME),
912 description: "Re-read the full content of a tool result that was truncated to save \
913 context (marked '[Output truncated: … call the read_result tool …]'). \
914 Pass that marker's call_id; use offset/max_bytes to page through \
915 large content."
916 .to_string(),
917 parameters: serde_json::json!({
918 "type": "object",
919 "properties": {
920 "call_id": { "type": "string" },
921 "offset": { "type": "integer", "description": "Byte offset to start from (default 0)." },
922 "max_bytes": { "type": "integer", "description": "Max bytes to return (default 4000)." }
923 },
924 "required": ["call_id"]
925 }),
926 })
927 }
928
929 pub fn plan_tool_schema(&self) -> Option<ToolSchema> {
930 if !self.plan_tool_enabled {
931 return None;
932 }
933 Some(ToolSchema {
934 name: CompactString::new("update_plan"),
935 description: "Update your task plan and progress. Call this after completing a step or when the plan changes.".to_string(),
936 parameters: serde_json::json!({
937 "type": "object",
938 "properties": {
939 "plan": { "type": "array", "items": { "type": "string" } },
940 "current_step": { "type": "integer" },
941 "progress": { "type": "string" },
942 "blocked_on": { "type": "array", "items": { "type": "string" } }
943 }
944 }),
945 })
946 }
947
948 pub fn memory_tool_schema(&self) -> Option<ToolSchema> {
949 if !self.memory_enabled {
950 return None;
951 }
952 Some(ToolSchema {
953 name: CompactString::new(MEMORY_TOOL_NAME),
954 description:
955 "Search your long-term memory for relevant past experiences and knowledge."
956 .to_string(),
957 parameters: serde_json::json!({
958 "type": "object",
959 "properties": {
960 "query": { "type": "string" },
961 "top_k": { "type": "integer" }
962 },
963 "required": ["query"]
964 }),
965 })
966 }
967
968 pub fn knowledge_tool_schema(&self) -> Option<ToolSchema> {
969 if !self.knowledge_enabled {
970 return None;
971 }
972 Some(ToolSchema {
973 name: CompactString::new(KNOWLEDGE_TOOL_NAME),
974 description:
975 "Search the external knowledge base for facts, documentation, or reference data."
976 .to_string(),
977 parameters: serde_json::json!({
978 "type": "object",
979 "properties": {
980 "query": { "type": "string" },
981 "top_k": { "type": "integer" }
982 },
983 "required": ["query"]
984 }),
985 })
986 }
987}
988
989#[cfg(test)]
990mod tests {
991 use super::*;
992 use crate::context::task_state::PlanStep;
993 use crate::types::message::Message;
994 use crate::types::skill::SkillMetadata;
995
996 #[test]
997 fn note_tool_actions_keys_on_name_and_args_so_legit_loops_dont_false_stop() {
998 let mut mgr = ContextManager::new(100_000);
1001 mgr.init_task("process items".to_string(), vec![]);
1002 mgr.note_tool_actions(&[("step".to_string(), "{\"n\":1}".to_string())]);
1003 mgr.note_tool_actions(&[("step".to_string(), "{\"n\":2}".to_string())]);
1004 mgr.note_tool_actions(&[("step".to_string(), "{\"n\":3}".to_string())]);
1005 assert_eq!(
1006 mgr.partitions.task_state.recent_actions,
1007 ["step({\"n\":1})", "step({\"n\":2})", "step({\"n\":3})"]
1008 );
1009 let txt = mgr
1010 .render()
1011 .state_turn
1012 .unwrap()
1013 .content
1014 .as_text()
1015 .unwrap()
1016 .to_string();
1017 assert!(
1018 !txt.contains("STOP:"),
1019 "same-tool/diff-args loop must not trip STOP: {txt}"
1020 );
1021
1022 let mut mgr2 = ContextManager::new(100_000);
1024 mgr2.init_task("g".to_string(), vec![]);
1025 for _ in 0..3 {
1026 mgr2.note_tool_actions(&[("document_read".to_string(), "{\"id\":\"x\"}".to_string())]);
1027 }
1028 let txt2 = mgr2
1029 .render()
1030 .state_turn
1031 .unwrap()
1032 .content
1033 .as_text()
1034 .unwrap()
1035 .to_string();
1036 assert!(
1037 txt2.contains("STOP:"),
1038 "identical repeated call must trip STOP: {txt2}"
1039 );
1040
1041 let mut mgr3 = ContextManager::new(100_000);
1043 mgr3.init_task("g".to_string(), vec![]);
1044 mgr3.note_tool_actions(&[(
1045 "update_plan".to_string(),
1046 "{\"current_step\":1}".to_string(),
1047 )]);
1048 assert!(mgr3.partitions.task_state.recent_actions.is_empty());
1049 }
1050
1051 #[test]
1052 fn manager_renew_advances_sprint_and_keeps_goal() {
1053 let mut mgr = ContextManager::new(1_000);
1054 mgr.init_task("test goal".to_string(), vec![]);
1055 mgr.partitions.system.push(Message::system("rules"), 10);
1056 for i in 0..10 {
1057 mgr.push_history(Message::user(format!("msg {i}")), 50);
1058 }
1059 mgr.renew();
1060 assert_eq!(mgr.partitions.task_state.goal, "test goal");
1061 assert_eq!(mgr.sprint, 1);
1062 }
1063
1064 #[test]
1065 fn compress_only_touches_history() {
1066 let mut mgr = ContextManager::new(1_000);
1067 mgr.push_knowledge(Message::system("knowledge content"), 100);
1068 for _ in 0..30 {
1069 mgr.push_history(Message::user("history msg"), 50);
1070 }
1071 let knowledge_before = mgr.partitions.knowledge.token_count;
1072 let history_before = mgr.partitions.history.token_count;
1073 mgr.compress(PressureAction::AutoCompact);
1074 assert_eq!(mgr.partitions.knowledge.token_count, knowledge_before);
1075 assert!(mgr.partitions.history.token_count < history_before);
1076 }
1077
1078 #[test]
1079 fn init_task_sets_goal_and_criteria() {
1080 let mut mgr = ContextManager::new(1_000);
1081 mgr.init_task("analyse data".to_string(), vec!["criterion A".to_string()]);
1082 assert_eq!(mgr.partitions.task_state.goal, "analyse data");
1083 assert_eq!(mgr.partitions.task_state.criteria, ["criterion A"]);
1084 }
1085
1086 #[test]
1087 fn update_task_applies_plan() {
1088 let mut mgr = ContextManager::new(1_000);
1089 mgr.init_task("g".to_string(), vec![]);
1090 mgr.update_task(TaskUpdate {
1091 plan: Some(vec!["step 1".to_string(), "step 2".to_string()]),
1092 current_step: Some(0),
1093 ..Default::default()
1094 });
1095 assert_eq!(mgr.partitions.task_state.plan.len(), 2);
1096 assert_eq!(mgr.partitions.task_state.current_step, Some(0));
1097 }
1098
1099 #[test]
1100 fn task_state_survives_autocompact() {
1101 let mut mgr = ContextManager::new(1_000);
1102 mgr.init_task("survive compression".to_string(), vec![]);
1103 mgr.update_task(TaskUpdate {
1104 plan: Some(vec!["fetch data".to_string(), "analyse".to_string()]),
1105 ..Default::default()
1106 });
1107 for _ in 0..10 {
1108 mgr.push_history(Message::user("filler"), 50);
1109 }
1110 mgr.compress(PressureAction::AutoCompact);
1111 assert_eq!(mgr.partitions.task_state.goal, "survive compression");
1112 assert_eq!(mgr.partitions.task_state.plan.len(), 2);
1113 }
1114
1115 #[test]
1116 fn render_includes_task_state_in_state_turn_not_system() {
1117 let mut mgr = ContextManager::new(10_000);
1118 mgr.init_task("find anomalies".to_string(), vec![]);
1119 let rc = mgr.render();
1120 assert!(
1121 !rc.system_text.contains("[TASK STATE]"),
1122 "task_state must not be in system_text"
1123 );
1124 let state = rc.state_turn.as_ref().expect("should have a state turn");
1126 assert!(
1127 state
1128 .content
1129 .as_text()
1130 .unwrap()
1131 .contains("[TASK STATE] goal: find anomalies")
1132 );
1133 }
1134
1135 #[test]
1136 fn renewal_keeps_open_plan_steps_in_task_state() {
1137 let mut mgr = ContextManager::new(1_000);
1138 mgr.init_task("g".to_string(), vec![]);
1139 mgr.partitions.task_state.plan = vec![
1140 PlanStep {
1141 label: "done".to_string(),
1142 done: true,
1143 },
1144 PlanStep {
1145 label: "pending".to_string(),
1146 done: false,
1147 },
1148 ];
1149 mgr.renew();
1150 assert_eq!(mgr.partitions.task_state.open_steps(), vec!["pending"]);
1151 }
1152
1153 #[test]
1156 fn auto_compact_entry_logs_auto_compact_action() {
1157 let mut mgr = ContextManager::new(1_000);
1165 for i in 0..40 {
1166 mgr.push_history(
1167 Message::user(format!("turn {i}: {}", "ctx ".repeat(40))),
1168 200,
1169 );
1170 }
1171 let (saved, summary, _, _) = mgr.force_compress();
1172 assert!(saved > 0, "force_compress should compact a large history");
1173 assert!(
1174 summary.is_some(),
1175 "auto-compact summarizes the archived turns"
1176 );
1177 let actions: Vec<&str> = mgr
1178 .partitions
1179 .task_state
1180 .compression_log
1181 .iter()
1182 .map(|e| e.action.as_str())
1183 .collect();
1184 assert!(
1185 actions.last() == Some(&"auto_compact"),
1186 "auto-compact entry must log an auto_compact action; got {actions:?}"
1187 );
1188 }
1189
1190 #[test]
1191 fn skill_tool_schema_empty_when_no_skills() {
1192 let mgr = ContextManager::new(10_000);
1193 assert!(mgr.skill_tool_schema().is_none());
1194 }
1195
1196 #[test]
1197 fn skill_tool_schema_present_when_registered() {
1198 let mut mgr = ContextManager::new(10_000);
1199 mgr.set_available_skills(vec![SkillMetadata::new("debug", "Debug helper")]);
1200 assert!(
1201 mgr.skill_tool_schema()
1202 .unwrap()
1203 .description
1204 .contains("debug")
1205 );
1206 }
1207
1208 #[test]
1209 fn available_skills_are_reflected_in_capability_manifest() {
1210 let mut mgr = ContextManager::new(1_000);
1211 mgr.set_available_skills(vec![SkillMetadata::new("debug", "Debug helper")]);
1212 let inventory = mgr.capability_inventory();
1213 assert!(inventory.contains("debug"));
1214 assert!(inventory.contains("Debug helper"));
1215 }
1216
1217 #[test]
1218 fn toggled_meta_tools_are_reflected_in_capability_manifest() {
1219 let mut mgr = ContextManager::new(1_000);
1220 mgr.set_memory_enabled(true);
1221 assert!(mgr.capability_inventory().contains(MEMORY_TOOL_NAME));
1222 mgr.set_memory_enabled(false);
1223 assert!(!mgr.capability_inventory().contains(MEMORY_TOOL_NAME));
1224 }
1225
1226 #[test]
1227 fn meta_tool_schemas_are_sorted() {
1228 let mut mgr = ContextManager::new(1_000);
1229 mgr.set_available_skills(vec![SkillMetadata::new("debug", "Debug helper")]);
1230 mgr.set_memory_enabled(true);
1231 mgr.set_knowledge_enabled(true);
1232 let names = mgr
1233 .meta_tool_schemas()
1234 .into_iter()
1235 .map(|s| s.name.to_string())
1236 .collect::<Vec<_>>();
1237 assert_eq!(names, ["knowledge", "memory", "skill"]);
1238 }
1239
1240 #[test]
1241 fn b1_active_skill_state_and_tool_filter() {
1242 let mut mgr = ContextManager::new(1_000);
1243 let mut debug = SkillMetadata::new("debug", "Debug helper");
1244 debug.allowed_tools = vec![CompactString::new("read"), CompactString::new("grep")];
1245 let mut review = SkillMetadata::new("review", "Reviewer");
1246 review.allowed_tools = vec![CompactString::new("git_diff")];
1247 let plain = SkillMetadata::new("plain", "No tools declared"); mgr.set_available_skills(vec![debug, review, plain]);
1249
1250 assert!(mgr.active_skill_tool_filter().is_none());
1252
1253 assert!(mgr.activate_skill("debug"));
1255 assert!(!mgr.activate_skill("debug")); let f = mgr.active_skill_tool_filter().unwrap();
1259 assert_eq!(f.len(), 2);
1260 assert!(f.contains(&CompactString::new("read")) && f.contains(&CompactString::new("grep")));
1261
1262 mgr.activate_skill("review");
1264 let f = mgr.active_skill_tool_filter().unwrap();
1265 assert_eq!(f.len(), 3);
1266 assert!(f.contains(&CompactString::new("git_diff")));
1267
1268 mgr.activate_skill("plain");
1270 assert!(mgr.active_skill_tool_filter().is_none());
1271 }
1272
1273 #[test]
1274 fn update_collapse_mode_collapses_old_tool_results_under_pressure() {
1275 let mut mgr = ContextManager::new(1_000);
1276 for i in 0..10 {
1277 let m = Message::tool(vec![ContentPart::ToolResult {
1278 call_id: format!("c{i}").into(),
1279 output: "x".repeat(40),
1280 is_error: false,
1281 }]);
1282 mgr.push_history(m, 40);
1283 }
1284 mgr.set_observed_prompt_tokens(950); assert!(mgr.rho() >= mgr.config.collapse_threshold);
1287
1288 mgr.recompute_handle_residency();
1289 assert_eq!(
1291 mgr.handles.residency_for_source("c0"),
1292 Some(&Residency::Collapsed)
1293 );
1294 assert_eq!(
1295 mgr.handles.residency_for_source("c9"),
1296 Some(&Residency::Resident)
1297 );
1298
1299 mgr.set_observed_prompt_tokens(100); mgr.recompute_handle_residency();
1303 assert_eq!(
1304 mgr.handles.residency_for_source("c0"),
1305 Some(&Residency::Collapsed),
1306 "collapse is sticky until a compaction boundary"
1307 );
1308
1309 mgr.reset_collapse_generation();
1311 assert_eq!(
1312 mgr.handles.residency_for_source("c0"),
1313 Some(&Residency::Resident)
1314 );
1315 }
1316
1317 #[test]
1318 fn frozen_prefix_len_anchors_at_compaction_and_holds_across_appends() {
1319 let mut mgr = ContextManager::new(1_000);
1320 for i in 0..30 {
1322 mgr.push_history(
1323 Message::user(format!("turn {i}: {}", "ctx ".repeat(30))),
1324 150,
1325 );
1326 }
1327 assert!(
1328 mgr.render().frozen_prefix_len.is_none(),
1329 "no frozen region before any compaction"
1330 );
1331
1332 let (saved, _, archived, _) = mgr.compress(PressureAction::AutoCompact);
1333 assert!(saved > 0 && !archived.is_empty(), "expected archival");
1334
1335 assert!(
1337 mgr.render().frozen_prefix_len.is_none(),
1338 "deep == tail right after compaction"
1339 );
1340
1341 mgr.push_history(Message::user("new 1"), 5);
1343 let f1 = mgr
1344 .render()
1345 .frozen_prefix_len
1346 .expect("frozen region exists once the tail grows");
1347 mgr.push_history(Message::assistant("reply 1"), 5);
1348 mgr.push_history(Message::user("new 2"), 5);
1349 let rc = mgr.render();
1350 let f2 = rc.frozen_prefix_len.expect("frozen region holds");
1351 assert_eq!(
1352 f1, f2,
1353 "the deep boundary is fixed between compactions; only the tail grows"
1354 );
1355 assert!(
1356 f2 < rc.turns.len(),
1357 "deep boundary is distinct from the rolling tail"
1358 );
1359 }
1360
1361 #[test]
1362 fn frozen_boundary_holds_through_a_prefix_safe_compaction() {
1363 let mut mgr = ContextManager::new(10_000);
1366 for i in 0..5 {
1367 mgr.push_history(Message::user(format!("m{i}")), 5);
1368 }
1369 mgr.frozen_history_len = 3; let (_, _, _, cache_at) = mgr.compress(PressureAction::None);
1374 assert!(cache_at.is_none(), "no-op compaction is prefix-safe");
1375 assert_eq!(
1376 mgr.frozen_history_len, 3,
1377 "prefix-safe compaction preserves the deep-cache anchor"
1378 );
1379 }
1380
1381 #[test]
1382 fn collapse_generation_resets_on_autocompact() {
1383 let mut mgr = ContextManager::new(1_000);
1384 for i in 0..20 {
1387 mgr.push_history(tool_result_msg(&format!("c{i}"), &"x".repeat(120)), 60);
1388 }
1389 mgr.set_observed_prompt_tokens(980); mgr.recompute_handle_residency();
1391 assert_eq!(
1392 mgr.handles.residency_for_source("c0"),
1393 Some(&Residency::Collapsed)
1394 );
1395
1396 let (saved, _, archived, _) = mgr.compress(PressureAction::AutoCompact);
1397 assert!(saved > 0 && !archived.is_empty(), "expected archival");
1398
1399 for h in mgr.handles.all() {
1402 if matches!(h.kind, HandleKind::ToolResult) {
1403 assert_eq!(
1404 h.residency,
1405 Residency::Resident,
1406 "generation reset un-collapses survivors"
1407 );
1408 }
1409 }
1410 }
1411
1412 #[test]
1413 fn push_history_indexes_tool_results_as_resident_handles() {
1414 let mut mgr = ContextManager::new(10_000);
1415 let msg = Message::tool(vec![ContentPart::ToolResult {
1416 call_id: "call_1".into(),
1417 output: "the tool output".to_string(),
1418 is_error: false,
1419 }]);
1420 mgr.push_history(msg, 20);
1421 assert_eq!(mgr.handles.all().len(), 1);
1423 assert_eq!(
1424 mgr.handles.residency_for_source("call_1"),
1425 Some(&Residency::Resident)
1426 );
1427 mgr.push_history(Message::user("hello"), 5);
1429 assert_eq!(mgr.handles.all().len(), 1);
1430 }
1431
1432 fn tool_result_msg(call_id: &str, output: &str) -> Message {
1435 Message::tool(vec![ContentPart::ToolResult {
1436 call_id: call_id.into(),
1437 output: output.to_string(),
1438 is_error: false,
1439 }])
1440 }
1441
1442 #[test]
1443 fn prune_orphaned_handles_drops_handles_whose_message_left_history() {
1444 let mut mgr = ContextManager::new(10_000);
1445 mgr.push_history(tool_result_msg("c0", "out 0"), 20);
1446 mgr.push_history(tool_result_msg("c1", "out 1"), 20);
1447 assert_eq!(mgr.handles.all().len(), 2);
1448
1449 mgr.partitions.history.messages.remove(0);
1451 mgr.prune_orphaned_handles();
1452
1453 assert_eq!(mgr.handles.all().len(), 1);
1455 assert!(mgr.handles.residency_for_source("c0").is_none());
1456 assert_eq!(
1457 mgr.handles.residency_for_source("c1"),
1458 Some(&Residency::Resident)
1459 );
1460 }
1461
1462 #[test]
1463 fn autocompact_prunes_handles_for_archived_tool_results() {
1464 let mut mgr = ContextManager::new(1_000);
1465 for i in 0..30 {
1467 mgr.push_history(tool_result_msg(&format!("c{i}"), &"x".repeat(200)), 80);
1468 }
1469 assert_eq!(mgr.handles.all().len(), 30);
1470
1471 let (saved, _, archived, _) = mgr.compress(PressureAction::AutoCompact);
1472 assert!(saved > 0 && !archived.is_empty(), "expected archival");
1473
1474 let live_tool_results = mgr
1477 .partitions
1478 .history
1479 .messages
1480 .iter()
1481 .filter(|m| {
1482 matches!(&m.content, Content::Parts(p)
1483 if p.iter().any(|x| matches!(x, ContentPart::ToolResult { .. })))
1484 })
1485 .count();
1486 assert_eq!(mgr.handles.all().len(), live_tool_results);
1487 assert!(
1488 mgr.handles.all().len() < 30,
1489 "table must shrink with archival"
1490 );
1491 }
1492
1493 #[test]
1494 fn renew_prunes_handles_for_dropped_history() {
1495 let mut mgr = ContextManager::new(1_000);
1496 mgr.init_task("g".to_string(), vec![]);
1497 for i in 0..20 {
1498 mgr.push_history(tool_result_msg(&format!("c{i}"), "data"), 60);
1499 }
1500 mgr.renew();
1501 for h in mgr.handles.all() {
1503 if let Some(src) = h.source.as_ref() {
1504 assert!(
1505 mgr.handles.residency_for_source(src).is_some(),
1506 "no dangling handle survives renewal"
1507 );
1508 }
1509 }
1510 assert!(mgr.handles.all().len() <= 20);
1511 }
1512
1513 #[test]
1516 fn knowledge_budget_uses_stable_order_for_equal_value_and_warns_once() {
1517 let mut mgr = ContextManager::new(100);
1520 mgr.push_knowledge(Message::system("oldest unkeyed"), 10);
1521 mgr.push_knowledge_entry(Some("a".into()), Message::system("keyed"), 10, false);
1522 mgr.push_knowledge_entry(Some("p".into()), Message::system("pinned"), 10, true);
1523 mgr.push_knowledge_entry(Some("skill:x".into()), Message::system("skill"), 10, false);
1524
1525 let warn = mgr.enforce_knowledge_budget();
1526 assert_eq!(warn, Some((40, 25)));
1527 let e = &mgr.partitions.knowledge.entries;
1529 assert!(e[0].evict_at_boundary);
1530 assert!(e[1].evict_at_boundary);
1531 assert!(!e[2].evict_at_boundary, "pinned exempt");
1532 assert!(!e[3].evict_at_boundary, "skill pin exempt");
1533
1534 assert_eq!(mgr.enforce_knowledge_budget(), None);
1536
1537 let sweep = mgr.partitions.knowledge.sweep_at_boundary();
1539 assert_eq!(sweep.tokens_freed, 20);
1540 assert_eq!(mgr.partitions.knowledge.token_count, 20);
1541 assert_eq!(mgr.enforce_knowledge_budget(), None);
1543 }
1544
1545 #[test]
1546 fn knowledge_budget_warning_stands_when_only_exempt_weight_remains() {
1547 let mut mgr = ContextManager::new(100);
1548 mgr.push_knowledge_entry(Some("p".into()), Message::system("pinned heavy"), 30, true);
1549 mgr.push_knowledge_entry(
1550 Some("skill:x".into()),
1551 Message::system("skill heavy"),
1552 30,
1553 false,
1554 );
1555
1556 assert_eq!(mgr.enforce_knowledge_budget(), Some((60, 25)));
1558 assert!(
1559 mgr.partitions
1560 .knowledge
1561 .entries
1562 .iter()
1563 .all(|e| !e.evict_at_boundary)
1564 );
1565 }
1566
1567 #[test]
1568 fn knowledge_budget_retains_old_referenced_entry_over_new_irrelevant_entry() {
1569 let mut mgr = ContextManager::new(100);
1570 mgr.push_knowledge_entry(
1571 Some("project:orchid".into()),
1572 Message::system("ORCHID uses the Atlas storage engine"),
1573 10,
1574 false,
1575 );
1576 mgr.push_history(Message::user("For project:orchid keep using Atlas"), 5);
1578 mgr.push_knowledge_entry(
1579 Some("project:new".into()),
1580 Message::system("unrelated fresh material"),
1581 10,
1582 false,
1583 );
1584 mgr.push_knowledge_entry(
1585 Some("project:other".into()),
1586 Message::system("another unused reference"),
1587 10,
1588 false,
1589 );
1590
1591 assert_eq!(mgr.enforce_knowledge_budget(), Some((30, 25)));
1592 let entries = &mgr.partitions.knowledge.entries;
1593 assert!(
1594 !entries[0].evict_at_boundary,
1595 "a real reference must raise retention"
1596 );
1597 assert!(
1598 entries[1].evict_at_boundary,
1599 "lowest-value entry evicts first"
1600 );
1601 assert!(
1602 !entries[2].evict_at_boundary,
1603 "one eviction is enough to fit"
1604 );
1605 }
1606
1607 #[test]
1608 fn knowledge_budget_ratio_zero_disables() {
1609 let mut mgr = ContextManager::new(100);
1610 mgr.config.knowledge_budget_ratio = 0.0;
1611 mgr.push_knowledge(Message::system("huge"), 90);
1612 assert_eq!(mgr.enforce_knowledge_budget(), None);
1613 assert!(!mgr.partitions.knowledge.entries[0].evict_at_boundary);
1614 }
1615
1616 #[test]
1617 fn provider_and_output_reservations_reduce_the_hard_input_budget() {
1618 use crate::context::config::PromptBudgetConfig;
1619
1620 let mut mgr = ContextManager::new(100);
1621 mgr.set_prompt_budget(PromptBudgetConfig {
1622 prompt_overhead_tokens: 20,
1623 output_reserve_tokens: 20,
1624 safety_margin_tokens: 10,
1625 });
1626 mgr.partitions
1627 .system
1628 .push(Message::system("x".repeat(240)), 60);
1629
1630 let rendered = mgr.render();
1631 let overflow = rendered
1632 .budget_overflow
1633 .expect("fixed context exceeds input allowance");
1634 assert_eq!(overflow.max_tokens, 50);
1635 assert!(overflow.required_tokens > overflow.max_tokens);
1636 }
1637}