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
41#[doc(hidden)]
46pub struct ContextManager {
47 pub partitions: ContextPartitions,
48 pub max_tokens: u32,
49 pub config: ContextConfig,
50 pub engine: ContextTokenEngine,
51 pub prompt_budget: PromptBudgetConfig,
54 pub sprint: u32,
55 pub skills: SkillCatalog,
56 pub active_skills: std::collections::BTreeMap<CompactString, Option<u32>>,
64 pub stable_core_tools: std::collections::HashSet<CompactString>,
68 pub capabilities: CapabilityManifest,
69 pub memory_enabled: bool,
70 pub knowledge_enabled: bool,
71 pub plan_tool_enabled: bool,
72 last_observed_prompt_tokens: Option<u32>,
73 compression: CompressionPipeline,
74 pressure: PressureMonitor,
75 renewal: RenewalPolicy,
76
77 pub last_activity_ms: u64,
81
82 pub last_compact_ms: Option<u64>,
85
86 pub handles: HandleTable,
91 next_handle_id: HandleId,
93
94 frozen_history_len: usize,
100
101 pending_knowledge_sweeps: Vec<crate::context::partitions::KnowledgeSweep>,
104
105 knowledge_budget_warned: bool,
108 knowledge_reference_step: u64,
110}
111
112impl ContextManager {
113 pub fn new(max_tokens: u32) -> Self {
114 Self::with_config(
115 max_tokens,
116 ContextConfig::default(),
117 ContextTokenEngine::char_approx(),
118 )
119 }
120
121 pub fn with_config(max_tokens: u32, config: ContextConfig, engine: ContextTokenEngine) -> Self {
122 let compression = CompressionPipeline::new(&config);
123 let pressure = PressureMonitor::new(max_tokens, config.clone());
124 let renewal = RenewalPolicy::from_config(&config);
125 let partitions = ContextPartitions::new(&config);
126 Self {
127 partitions,
128 max_tokens,
129 config,
130 engine,
131 prompt_budget: PromptBudgetConfig::default(),
132 sprint: 0,
133 skills: SkillCatalog::new(),
134 active_skills: std::collections::BTreeMap::new(),
135 stable_core_tools: std::collections::HashSet::new(),
136 capabilities: CapabilityManifest::new(),
137 memory_enabled: false,
138 knowledge_enabled: false,
139 plan_tool_enabled: false,
140 last_observed_prompt_tokens: None,
141 compression,
142 pressure,
143 renewal,
144 last_activity_ms: 0,
145 last_compact_ms: None,
146 handles: HandleTable::new(),
147 next_handle_id: 0,
148 frozen_history_len: 0,
149 pending_knowledge_sweeps: Vec::new(),
150 knowledge_budget_warned: false,
151 knowledge_reference_step: 0,
152 }
153 }
154
155 pub fn apply_context_policy(&mut self, policy: &ContextPolicyV1) {
157 policy.apply_to(&mut self.config);
158 self.compression = CompressionPipeline::new(&self.config);
159 self.pressure = PressureMonitor::new(self.max_tokens, self.config.clone());
160 self.renewal = RenewalPolicy::from_config(&self.config);
161 }
162
163 pub fn record_activity(&mut self, now_ms: u64) {
167 self.last_activity_ms = now_ms;
168 }
169
170 pub fn should_time_decay_compact(&self, now_ms: u64) -> bool {
173 let idle_ms = if let Some(last_compact) = self.last_compact_ms {
174 now_ms.saturating_sub(last_compact)
176 } else {
177 now_ms.saturating_sub(self.last_activity_ms)
179 };
180
181 let idle_minutes = idle_ms / 60_000;
182 idle_minutes >= self.config.micro_compact_idle_minutes as u64
183 }
184
185 pub fn recompute_handle_residency(&mut self) {
200 if self.rho() < self.config.collapse_threshold {
202 return;
203 }
204 let keep = self.config.preserved_tool_results;
205 let total = self
208 .handles
209 .all()
210 .iter()
211 .filter(|h| matches!(h.kind, HandleKind::ToolResult))
212 .count();
213 let cutoff = total.saturating_sub(keep);
214 for (i, handle) in self.handles.tool_result_handles_mut().enumerate() {
215 if i < cutoff && matches!(handle.residency, Residency::Resident) {
218 handle.residency = Residency::Collapsed;
219 }
220 }
221 }
222
223 pub fn reset_collapse_generation(&mut self) {
229 for handle in self.handles.all_mut() {
230 if matches!(handle.residency, Residency::Collapsed) {
231 handle.residency = Residency::Resident;
232 }
233 }
234 }
235
236 pub fn prune_orphaned_handles(&mut self) {
243 let live: std::collections::HashSet<CompactString> = self
244 .partitions
245 .history
246 .messages
247 .iter()
248 .flat_map(|m| match &m.content {
249 Content::Parts(parts) => parts
250 .iter()
251 .filter_map(|p| match p {
252 ContentPart::ToolResult { call_id, .. } => Some(call_id.clone()),
253 _ => None,
254 })
255 .collect::<Vec<_>>(),
256 _ => Vec::new(),
257 })
258 .collect();
259 self.handles
260 .retain(|h| h.source.as_ref().is_none_or(|s| live.contains(s)));
261 }
262
263 pub fn mark_spooled(&mut self, call_id: &str, spool_ref: impl Into<String>) {
267 let spool_ref = spool_ref.into();
268 if let Some(handle) = self
269 .handles
270 .all_mut()
271 .iter_mut()
272 .find(|h| h.source.as_deref() == Some(call_id))
273 {
274 handle.residency = Residency::SpooledOut { r: spool_ref };
275 }
276 }
277
278 pub fn rho(&self) -> f64 {
285 self.pressure.pressure(
286 &self.partitions,
287 &self.engine,
288 self.last_observed_prompt_tokens,
289 )
290 }
291
292 pub fn set_observed_prompt_tokens(&mut self, tokens: u32) {
293 self.last_observed_prompt_tokens = Some(tokens);
294 }
295
296 pub fn should_compress(&self) -> PressureAction {
297 self.pressure.recommend(self.rho())
303 }
304
305 pub fn compress(
306 &mut self,
307 action: PressureAction,
308 ) -> (u32, Option<String>, Vec<Message>, Option<usize>) {
309 self.compress_with_time(action, None)
310 }
311
312 pub fn compress_with_time(
313 &mut self,
314 action: PressureAction,
315 now_ms: Option<u64>,
316 ) -> (u32, Option<String>, Vec<Message>, Option<usize>) {
317 let target = self.config.target_tokens(self.max_tokens);
318 self.compress_with_target(action, target, now_ms)
319 }
320
321 pub fn force_compress(&mut self) -> (u32, Option<String>, Vec<Message>, Option<usize>) {
322 self.compress_with_target(PressureAction::AutoCompact, 0, None)
323 }
324
325 pub fn compress_with_target(
332 &mut self,
333 action: PressureAction,
334 target_tokens: u32,
335 now_ms: Option<u64>,
336 ) -> (u32, Option<String>, Vec<Message>, Option<usize>) {
337 let result = self.compression.compress(
338 &mut self.partitions,
339 action,
340 self.max_tokens,
341 target_tokens,
342 &self.engine,
343 );
344 if let Some(ts) = now_ms {
345 self.last_compact_ms = Some(ts);
346 }
347 if !result.2.is_empty() {
349 self.prune_orphaned_handles();
350 self.reset_collapse_generation();
353 self.sweep_knowledge_at_boundary();
356 }
357 if result.3.is_some() {
364 self.frozen_history_len = self.partitions.history.messages.len();
365 }
366 result
367 }
368
369 pub fn plan_compaction_params(&self) -> (u32, usize) {
373 (
374 self.config.target_tokens(self.max_tokens),
375 self.config.preserve_recent_turns,
376 )
377 }
378
379 pub fn should_renew(&self) -> bool {
382 self.renewal
383 .should_renew(&self.pressure, &self.partitions, &self.engine)
384 }
385
386 pub fn renew(&mut self) {
387 self.partitions = self.renewal.renew(&self.partitions, self.max_tokens);
388 self.sprint += 1;
389 self.prune_orphaned_handles();
392 self.reset_collapse_generation();
393 self.sweep_knowledge_at_boundary();
395 self.frozen_history_len = self.partitions.history.messages.len();
397 }
398
399 pub fn set_prompt_budget(&mut self, prompt_budget: PromptBudgetConfig) {
402 self.prompt_budget = prompt_budget;
403 }
404
405 pub fn available_input_tokens(&self) -> u32 {
406 self.max_tokens
407 .saturating_sub(self.prompt_budget.reserved_tokens())
408 }
409
410 pub fn render(&self) -> RenderedContext {
411 super::renderer::render_projected(
412 &self.partitions,
413 self.available_input_tokens(),
414 &self.engine,
415 self.config.preserve_recent_units,
416 &self.handles,
417 self.frozen_history_len,
418 self.config.collapse_assistant_narration,
419 )
420 }
421
422 pub fn push_history(&mut self, msg: Message, tokens: u32) {
425 self.knowledge_reference_step = self.knowledge_reference_step.saturating_add(1);
426 self.partitions
427 .knowledge
428 .observe_references(&msg, self.knowledge_reference_step);
429 if let Content::Parts(parts) = &msg.content {
433 for part in parts {
434 if let ContentPart::ToolResult {
435 call_id, output, ..
436 } = part
437 {
438 let id = self.alloc_handle_id();
439 let tok = self.engine.count(output).max(1);
440 self.handles.insert(Handle::resident_for(
441 id,
442 HandleKind::ToolResult,
443 tok,
444 call_id.clone(),
445 ));
446 }
447 }
448 }
449 self.partitions.history.push(msg, tokens);
450 }
451
452 fn alloc_handle_id(&mut self) -> HandleId {
453 let id = self.next_handle_id;
454 self.next_handle_id = self.next_handle_id.wrapping_add(1);
455 id
456 }
457
458 pub fn push_knowledge(&mut self, msg: Message, tokens: u32) {
460 self.partitions.knowledge.push(msg, tokens);
461 }
462
463 pub fn push_knowledge_entry(
467 &mut self,
468 key: Option<CompactString>,
469 msg: Message,
470 tokens: u32,
471 pinned: bool,
472 ) {
473 self.partitions
474 .knowledge
475 .push_entry(key, msg, tokens, pinned);
476 }
477
478 pub fn remove_knowledge(&mut self, key: &str) -> bool {
481 self.partitions.knowledge.remove(key)
482 }
483
484 fn sweep_knowledge_at_boundary(&mut self) {
488 let sweep = self.partitions.knowledge.sweep_at_boundary();
489 if sweep.changed {
490 if !sweep.removed_keys.is_empty() {
494 self.partitions.signals.push(format!(
495 "[KNOWLEDGE] entries removed at this boundary: {} — re-fetch via the memory tool if still needed.",
496 sweep.removed_keys.join(", ")
497 ));
498 }
499 self.pending_knowledge_sweeps.push(sweep);
500 }
501 self.knowledge_budget_warned = false;
503 }
504
505 pub fn enforce_knowledge_budget(&mut self) -> Option<(u32, u32)> {
514 let ratio = self.config.knowledge_budget_ratio;
515 if ratio <= 0.0 {
516 return None;
517 }
518 let budget = (self.max_tokens as f64 * ratio) as u32;
519 let used = self.partitions.knowledge.token_count;
520 if used <= budget {
521 return None;
522 }
523 let marked: u32 = self
524 .partitions
525 .knowledge
526 .entries
527 .iter()
528 .filter(|e| e.evict_at_boundary)
529 .map(|e| e.tokens)
530 .sum();
531 let mut projected = used.saturating_sub(marked);
532 let mut candidates = self
533 .partitions
534 .knowledge
535 .entries
536 .iter()
537 .enumerate()
538 .filter(|(_, entry)| {
539 !entry.evict_at_boundary
540 && !entry.pinned
541 && !entry
542 .key
543 .as_deref()
544 .is_some_and(|key| key.starts_with("skill:"))
545 })
546 .map(|(index, _)| {
547 let score = self
548 .partitions
549 .knowledge
550 .retention_score(index, self.knowledge_reference_step)
551 .unwrap_or(i64::MIN);
552 (score, index)
553 })
554 .collect::<Vec<_>>();
555 candidates.sort_by(|left, right| left.0.cmp(&right.0).then(left.1.cmp(&right.1)));
556 for (_, index) in candidates {
557 if projected <= budget {
558 break;
559 }
560 let entry = &mut self.partitions.knowledge.entries[index];
561 entry.evict_at_boundary = true;
562 projected = projected.saturating_sub(entry.tokens);
563 }
564 if self.knowledge_budget_warned {
565 return None;
566 }
567 self.knowledge_budget_warned = true;
568 Some((used, budget))
569 }
570
571 pub fn take_knowledge_sweeps(&mut self) -> Vec<crate::context::partitions::KnowledgeSweep> {
573 std::mem::take(&mut self.pending_knowledge_sweeps)
574 }
575
576 pub fn push_signal(&mut self, text: String) {
580 self.partitions.signals.push(text);
581 }
582
583 pub fn record_directive(&mut self, text: impl Into<String>) {
587 self.partitions.task_state.record_directive(text);
588 }
589
590 pub fn init_task(&mut self, goal: String, criteria: Vec<String>) {
593 self.partitions.task_state = TaskState {
594 goal,
595 criteria,
596 ..Default::default()
597 };
598 }
599
600 pub fn update_task(&mut self, update: TaskUpdate) {
601 self.partitions.task_state.apply(update);
602 }
603
604 pub fn note_tool_actions(&mut self, calls: &[(String, String)]) {
611 let summary = calls
612 .iter()
613 .filter(|(name, _)| !is_meta_tool(name))
614 .map(|(name, args)| {
615 if args.is_empty() {
616 name.clone()
617 } else {
618 format!("{name}({args})")
619 }
620 })
621 .collect::<Vec<_>>()
622 .join(", ");
623 self.partitions.task_state.note_actions(summary);
624 }
625
626 pub fn set_available_skills(&mut self, skills: Vec<SkillMetadata>) {
631 self.capabilities.remove_kind(CapabilityKind::Skill);
632 for skill in &skills {
633 self.capabilities.add_skill(skill.clone());
634 }
635 self.skills.set_available(skills);
636 }
637
638 pub fn set_stable_core_tools(&mut self, ids: impl IntoIterator<Item = CompactString>) {
640 self.stable_core_tools = ids.into_iter().collect();
641 }
642
643 pub fn activate_skill(&mut self, name: impl Into<CompactString>) -> bool {
648 self.activate_skill_leased(name, None)
649 }
650
651 pub fn activate_skill_leased(
654 &mut self,
655 name: impl Into<CompactString>,
656 expires_at_turn: Option<u32>,
657 ) -> bool {
658 self.active_skills
659 .insert(name.into(), expires_at_turn)
660 .is_none()
661 }
662
663 pub fn deactivate_skill(&mut self, name: &str) -> bool {
667 if self.active_skills.remove(name).is_none() {
668 return false;
669 }
670 self.partitions.knowledge.remove(&format!("skill:{name}"));
671 true
672 }
673
674 pub fn sweep_expired_skill_leases(&mut self, current_turn: u32) {
677 let expired: Vec<CompactString> = self
678 .active_skills
679 .iter()
680 .filter(|(_, lease)| lease.is_some_and(|t| current_turn >= t))
681 .map(|(name, _)| name.clone())
682 .collect();
683 for name in expired {
684 self.deactivate_skill(&name);
685 self.partitions.signals.push(format!(
687 "[SKILL] lease expired: {name} unloaded; the full toolset is restored."
688 ));
689 }
690 }
691
692 pub fn active_skill_tool_filter(&self) -> Option<std::collections::HashSet<CompactString>> {
698 if self.active_skills.is_empty() {
699 return None;
700 }
701 let mut union = std::collections::HashSet::new();
702 for name in self.active_skills.keys() {
703 let declared = self.skills.allowed_tools(name);
704 if declared.is_empty() {
705 return None; }
707 union.extend(declared.iter().cloned());
708 }
709 Some(union)
710 }
711
712 pub fn skill_tool_schema(&self) -> Option<ToolSchema> {
713 self.skills.build_tool_schema()
714 }
715
716 pub fn set_memory_enabled(&mut self, enabled: bool) {
719 self.memory_enabled = enabled;
720 if enabled {
721 self.capabilities.add_marker(
722 CapabilityKind::Memory,
723 MEMORY_TOOL_NAME,
724 "Search long-term memory through the memory meta-tool.",
725 );
726 } else {
727 self.capabilities
728 .remove(CapabilityKind::Memory, MEMORY_TOOL_NAME);
729 }
730 }
731
732 pub fn set_knowledge_enabled(&mut self, enabled: bool) {
733 self.knowledge_enabled = enabled;
734 if enabled {
735 self.capabilities.add_marker(
736 CapabilityKind::Knowledge,
737 KNOWLEDGE_TOOL_NAME,
738 "Search external knowledge through the knowledge meta-tool.",
739 );
740 } else {
741 self.capabilities
742 .remove(CapabilityKind::Knowledge, KNOWLEDGE_TOOL_NAME);
743 }
744 }
745
746 pub fn set_plan_tool_enabled(&mut self, enabled: bool) {
747 self.plan_tool_enabled = enabled;
748 if enabled {
749 self.capabilities.add_marker(
750 CapabilityKind::Tool,
751 "update_plan",
752 "Update task plan and progress through the planning meta-tool.",
753 );
754 } else {
755 self.capabilities
756 .remove(CapabilityKind::Tool, "update_plan");
757 }
758 }
759
760 pub fn capability_inventory(&self) -> String {
761 self.capabilities.format_inventory()
762 }
763
764 pub fn meta_tool_schemas(&self) -> Vec<ToolSchema> {
765 let mut tools = Vec::new();
766 if let Some(t) = self.skill_tool_schema() {
767 tools.push(t);
768 }
769 if let Some(t) = self.memory_tool_schema() {
770 tools.push(t);
771 }
772 if let Some(t) = self.knowledge_tool_schema() {
773 tools.push(t);
774 }
775 if let Some(t) = self.plan_tool_schema() {
776 tools.push(t);
777 }
778 if let Some(t) = self.read_result_tool_schema() {
779 tools.push(t);
780 }
781 tools.sort_by(|a, b| a.name.cmp(&b.name));
782 tools
783 }
784
785 pub fn read_result_tool_schema(&self) -> Option<ToolSchema> {
791 let any_evicted = self
792 .handles
793 .all()
794 .iter()
795 .any(|h| !h.residency.occupies_context());
796 if !any_evicted {
797 return None;
798 }
799 Some(ToolSchema {
800 name: CompactString::new(READ_RESULT_TOOL_NAME),
801 description: "Re-read a tool result that was evicted from context (you see a \
802 placeholder like '[…tool result spooled…]' or a collapsed entry). \
803 Pass the tool call's call_id; use offset/max_bytes to page through \
804 large content."
805 .to_string(),
806 parameters: serde_json::json!({
807 "type": "object",
808 "properties": {
809 "call_id": { "type": "string" },
810 "offset": { "type": "integer", "description": "Byte offset to start from (default 0)." },
811 "max_bytes": { "type": "integer", "description": "Max bytes to return (default 4000)." }
812 },
813 "required": ["call_id"]
814 }),
815 })
816 }
817
818 pub fn plan_tool_schema(&self) -> Option<ToolSchema> {
819 if !self.plan_tool_enabled {
820 return None;
821 }
822 Some(ToolSchema {
823 name: CompactString::new("update_plan"),
824 description: "Update your task plan and progress. Call this after completing a step or when the plan changes.".to_string(),
825 parameters: serde_json::json!({
826 "type": "object",
827 "properties": {
828 "plan": { "type": "array", "items": { "type": "string" } },
829 "current_step": { "type": "integer" },
830 "progress": { "type": "string" },
831 "blocked_on": { "type": "array", "items": { "type": "string" } }
832 }
833 }),
834 })
835 }
836
837 pub fn memory_tool_schema(&self) -> Option<ToolSchema> {
838 if !self.memory_enabled {
839 return None;
840 }
841 Some(ToolSchema {
842 name: CompactString::new(MEMORY_TOOL_NAME),
843 description:
844 "Search your long-term memory for relevant past experiences and knowledge."
845 .to_string(),
846 parameters: serde_json::json!({
847 "type": "object",
848 "properties": {
849 "query": { "type": "string" },
850 "top_k": { "type": "integer" }
851 },
852 "required": ["query"]
853 }),
854 })
855 }
856
857 pub fn knowledge_tool_schema(&self) -> Option<ToolSchema> {
858 if !self.knowledge_enabled {
859 return None;
860 }
861 Some(ToolSchema {
862 name: CompactString::new(KNOWLEDGE_TOOL_NAME),
863 description:
864 "Search the external knowledge base for facts, documentation, or reference data."
865 .to_string(),
866 parameters: serde_json::json!({
867 "type": "object",
868 "properties": {
869 "query": { "type": "string" },
870 "top_k": { "type": "integer" }
871 },
872 "required": ["query"]
873 }),
874 })
875 }
876}
877
878#[cfg(test)]
879mod tests {
880 use super::*;
881 use crate::context::task_state::PlanStep;
882 use crate::types::message::Message;
883 use crate::types::skill::SkillMetadata;
884
885 #[test]
886 fn note_tool_actions_keys_on_name_and_args_so_legit_loops_dont_false_stop() {
887 let mut mgr = ContextManager::new(100_000);
890 mgr.init_task("process items".to_string(), vec![]);
891 mgr.note_tool_actions(&[("step".to_string(), "{\"n\":1}".to_string())]);
892 mgr.note_tool_actions(&[("step".to_string(), "{\"n\":2}".to_string())]);
893 mgr.note_tool_actions(&[("step".to_string(), "{\"n\":3}".to_string())]);
894 assert_eq!(
895 mgr.partitions.task_state.recent_actions,
896 ["step({\"n\":1})", "step({\"n\":2})", "step({\"n\":3})"]
897 );
898 let txt = mgr
899 .render()
900 .state_turn
901 .unwrap()
902 .content
903 .as_text()
904 .unwrap()
905 .to_string();
906 assert!(
907 !txt.contains("STOP:"),
908 "same-tool/diff-args loop must not trip STOP: {txt}"
909 );
910
911 let mut mgr2 = ContextManager::new(100_000);
913 mgr2.init_task("g".to_string(), vec![]);
914 for _ in 0..3 {
915 mgr2.note_tool_actions(&[("document_read".to_string(), "{\"id\":\"x\"}".to_string())]);
916 }
917 let txt2 = mgr2
918 .render()
919 .state_turn
920 .unwrap()
921 .content
922 .as_text()
923 .unwrap()
924 .to_string();
925 assert!(
926 txt2.contains("STOP:"),
927 "identical repeated call must trip STOP: {txt2}"
928 );
929
930 let mut mgr3 = ContextManager::new(100_000);
932 mgr3.init_task("g".to_string(), vec![]);
933 mgr3.note_tool_actions(&[(
934 "update_plan".to_string(),
935 "{\"current_step\":1}".to_string(),
936 )]);
937 assert!(mgr3.partitions.task_state.recent_actions.is_empty());
938 }
939
940 #[test]
941 fn manager_renew_advances_sprint_and_keeps_goal() {
942 let mut mgr = ContextManager::new(1_000);
943 mgr.init_task("test goal".to_string(), vec![]);
944 mgr.partitions.system.push(Message::system("rules"), 10);
945 for i in 0..10 {
946 mgr.push_history(Message::user(format!("msg {i}")), 50);
947 }
948 mgr.renew();
949 assert_eq!(mgr.partitions.task_state.goal, "test goal");
950 assert_eq!(mgr.sprint, 1);
951 }
952
953 #[test]
954 fn compress_only_touches_history() {
955 let mut mgr = ContextManager::new(1_000);
956 mgr.push_knowledge(Message::system("knowledge content"), 100);
957 for _ in 0..30 {
958 mgr.push_history(Message::user("history msg"), 50);
959 }
960 let knowledge_before = mgr.partitions.knowledge.token_count;
961 let history_before = mgr.partitions.history.token_count;
962 mgr.compress(PressureAction::AutoCompact);
963 assert_eq!(mgr.partitions.knowledge.token_count, knowledge_before);
964 assert!(mgr.partitions.history.token_count < history_before);
965 }
966
967 #[test]
968 fn init_task_sets_goal_and_criteria() {
969 let mut mgr = ContextManager::new(1_000);
970 mgr.init_task("analyse data".to_string(), vec!["criterion A".to_string()]);
971 assert_eq!(mgr.partitions.task_state.goal, "analyse data");
972 assert_eq!(mgr.partitions.task_state.criteria, ["criterion A"]);
973 }
974
975 #[test]
976 fn update_task_applies_plan() {
977 let mut mgr = ContextManager::new(1_000);
978 mgr.init_task("g".to_string(), vec![]);
979 mgr.update_task(TaskUpdate {
980 plan: Some(vec!["step 1".to_string(), "step 2".to_string()]),
981 current_step: Some(0),
982 ..Default::default()
983 });
984 assert_eq!(mgr.partitions.task_state.plan.len(), 2);
985 assert_eq!(mgr.partitions.task_state.current_step, Some(0));
986 }
987
988 #[test]
989 fn task_state_survives_autocompact() {
990 let mut mgr = ContextManager::new(1_000);
991 mgr.init_task("survive compression".to_string(), vec![]);
992 mgr.update_task(TaskUpdate {
993 plan: Some(vec!["fetch data".to_string(), "analyse".to_string()]),
994 ..Default::default()
995 });
996 for _ in 0..10 {
997 mgr.push_history(Message::user("filler"), 50);
998 }
999 mgr.compress(PressureAction::AutoCompact);
1000 assert_eq!(mgr.partitions.task_state.goal, "survive compression");
1001 assert_eq!(mgr.partitions.task_state.plan.len(), 2);
1002 }
1003
1004 #[test]
1005 fn render_includes_task_state_in_state_turn_not_system() {
1006 let mut mgr = ContextManager::new(10_000);
1007 mgr.init_task("find anomalies".to_string(), vec![]);
1008 let rc = mgr.render();
1009 assert!(
1010 !rc.system_text.contains("[TASK STATE]"),
1011 "task_state must not be in system_text"
1012 );
1013 let state = rc.state_turn.as_ref().expect("should have a state turn");
1015 assert!(
1016 state
1017 .content
1018 .as_text()
1019 .unwrap()
1020 .contains("[TASK STATE] goal: find anomalies")
1021 );
1022 }
1023
1024 #[test]
1025 fn renewal_keeps_open_plan_steps_in_task_state() {
1026 let mut mgr = ContextManager::new(1_000);
1027 mgr.init_task("g".to_string(), vec![]);
1028 mgr.partitions.task_state.plan = vec![
1029 PlanStep {
1030 label: "done".to_string(),
1031 done: true,
1032 },
1033 PlanStep {
1034 label: "pending".to_string(),
1035 done: false,
1036 },
1037 ];
1038 mgr.renew();
1039 assert_eq!(mgr.partitions.task_state.open_steps(), vec!["pending"]);
1040 }
1041
1042 #[test]
1045 fn auto_compact_entry_logs_auto_compact_action() {
1046 let mut mgr = ContextManager::new(1_000);
1054 for i in 0..40 {
1055 mgr.push_history(
1056 Message::user(format!("turn {i}: {}", "ctx ".repeat(40))),
1057 200,
1058 );
1059 }
1060 let (saved, summary, _, _) = mgr.force_compress();
1061 assert!(saved > 0, "force_compress should compact a large history");
1062 assert!(
1063 summary.is_some(),
1064 "auto-compact summarizes the archived turns"
1065 );
1066 let actions: Vec<&str> = mgr
1067 .partitions
1068 .task_state
1069 .compression_log
1070 .iter()
1071 .map(|e| e.action.as_str())
1072 .collect();
1073 assert!(
1074 actions.last() == Some(&"auto_compact"),
1075 "auto-compact entry must log an auto_compact action; got {actions:?}"
1076 );
1077 }
1078
1079 #[test]
1080 fn skill_tool_schema_empty_when_no_skills() {
1081 let mgr = ContextManager::new(10_000);
1082 assert!(mgr.skill_tool_schema().is_none());
1083 }
1084
1085 #[test]
1086 fn skill_tool_schema_present_when_registered() {
1087 let mut mgr = ContextManager::new(10_000);
1088 mgr.set_available_skills(vec![SkillMetadata::new("debug", "Debug helper")]);
1089 assert!(
1090 mgr.skill_tool_schema()
1091 .unwrap()
1092 .description
1093 .contains("debug")
1094 );
1095 }
1096
1097 #[test]
1098 fn available_skills_are_reflected_in_capability_manifest() {
1099 let mut mgr = ContextManager::new(1_000);
1100 mgr.set_available_skills(vec![SkillMetadata::new("debug", "Debug helper")]);
1101 let inventory = mgr.capability_inventory();
1102 assert!(inventory.contains("debug"));
1103 assert!(inventory.contains("Debug helper"));
1104 }
1105
1106 #[test]
1107 fn toggled_meta_tools_are_reflected_in_capability_manifest() {
1108 let mut mgr = ContextManager::new(1_000);
1109 mgr.set_memory_enabled(true);
1110 assert!(mgr.capability_inventory().contains(MEMORY_TOOL_NAME));
1111 mgr.set_memory_enabled(false);
1112 assert!(!mgr.capability_inventory().contains(MEMORY_TOOL_NAME));
1113 }
1114
1115 #[test]
1116 fn meta_tool_schemas_are_sorted() {
1117 let mut mgr = ContextManager::new(1_000);
1118 mgr.set_available_skills(vec![SkillMetadata::new("debug", "Debug helper")]);
1119 mgr.set_memory_enabled(true);
1120 mgr.set_knowledge_enabled(true);
1121 let names = mgr
1122 .meta_tool_schemas()
1123 .into_iter()
1124 .map(|s| s.name.to_string())
1125 .collect::<Vec<_>>();
1126 assert_eq!(names, ["knowledge", "memory", "skill"]);
1127 }
1128
1129 #[test]
1130 fn b1_active_skill_state_and_tool_filter() {
1131 let mut mgr = ContextManager::new(1_000);
1132 let mut debug = SkillMetadata::new("debug", "Debug helper");
1133 debug.allowed_tools = vec![CompactString::new("read"), CompactString::new("grep")];
1134 let mut review = SkillMetadata::new("review", "Reviewer");
1135 review.allowed_tools = vec![CompactString::new("git_diff")];
1136 let plain = SkillMetadata::new("plain", "No tools declared"); mgr.set_available_skills(vec![debug, review, plain]);
1138
1139 assert!(mgr.active_skill_tool_filter().is_none());
1141
1142 assert!(mgr.activate_skill("debug"));
1144 assert!(!mgr.activate_skill("debug")); let f = mgr.active_skill_tool_filter().unwrap();
1148 assert_eq!(f.len(), 2);
1149 assert!(f.contains(&CompactString::new("read")) && f.contains(&CompactString::new("grep")));
1150
1151 mgr.activate_skill("review");
1153 let f = mgr.active_skill_tool_filter().unwrap();
1154 assert_eq!(f.len(), 3);
1155 assert!(f.contains(&CompactString::new("git_diff")));
1156
1157 mgr.activate_skill("plain");
1159 assert!(mgr.active_skill_tool_filter().is_none());
1160 }
1161
1162 #[test]
1163 fn update_collapse_mode_collapses_old_tool_results_under_pressure() {
1164 let mut mgr = ContextManager::new(1_000);
1165 for i in 0..10 {
1166 let m = Message::tool(vec![ContentPart::ToolResult {
1167 call_id: format!("c{i}").into(),
1168 output: "x".repeat(40),
1169 is_error: false,
1170 }]);
1171 mgr.push_history(m, 40);
1172 }
1173 mgr.set_observed_prompt_tokens(950); assert!(mgr.rho() >= mgr.config.collapse_threshold);
1176
1177 mgr.recompute_handle_residency();
1178 assert_eq!(
1180 mgr.handles.residency_for_source("c0"),
1181 Some(&Residency::Collapsed)
1182 );
1183 assert_eq!(
1184 mgr.handles.residency_for_source("c9"),
1185 Some(&Residency::Resident)
1186 );
1187
1188 mgr.set_observed_prompt_tokens(100); mgr.recompute_handle_residency();
1192 assert_eq!(
1193 mgr.handles.residency_for_source("c0"),
1194 Some(&Residency::Collapsed),
1195 "collapse is sticky until a compaction boundary"
1196 );
1197
1198 mgr.reset_collapse_generation();
1200 assert_eq!(
1201 mgr.handles.residency_for_source("c0"),
1202 Some(&Residency::Resident)
1203 );
1204 }
1205
1206 #[test]
1207 fn frozen_prefix_len_anchors_at_compaction_and_holds_across_appends() {
1208 let mut mgr = ContextManager::new(1_000);
1209 for i in 0..30 {
1211 mgr.push_history(
1212 Message::user(format!("turn {i}: {}", "ctx ".repeat(30))),
1213 150,
1214 );
1215 }
1216 assert!(
1217 mgr.render().frozen_prefix_len.is_none(),
1218 "no frozen region before any compaction"
1219 );
1220
1221 let (saved, _, archived, _) = mgr.compress(PressureAction::AutoCompact);
1222 assert!(saved > 0 && !archived.is_empty(), "expected archival");
1223
1224 assert!(
1226 mgr.render().frozen_prefix_len.is_none(),
1227 "deep == tail right after compaction"
1228 );
1229
1230 mgr.push_history(Message::user("new 1"), 5);
1232 let f1 = mgr
1233 .render()
1234 .frozen_prefix_len
1235 .expect("frozen region exists once the tail grows");
1236 mgr.push_history(Message::assistant("reply 1"), 5);
1237 mgr.push_history(Message::user("new 2"), 5);
1238 let rc = mgr.render();
1239 let f2 = rc.frozen_prefix_len.expect("frozen region holds");
1240 assert_eq!(
1241 f1, f2,
1242 "the deep boundary is fixed between compactions; only the tail grows"
1243 );
1244 assert!(
1245 f2 < rc.turns.len(),
1246 "deep boundary is distinct from the rolling tail"
1247 );
1248 }
1249
1250 #[test]
1251 fn frozen_boundary_holds_through_a_prefix_safe_compaction() {
1252 let mut mgr = ContextManager::new(10_000);
1255 for i in 0..5 {
1256 mgr.push_history(Message::user(format!("m{i}")), 5);
1257 }
1258 mgr.frozen_history_len = 3; let (_, _, _, cache_at) = mgr.compress(PressureAction::None);
1263 assert!(cache_at.is_none(), "no-op compaction is prefix-safe");
1264 assert_eq!(
1265 mgr.frozen_history_len, 3,
1266 "prefix-safe compaction preserves the deep-cache anchor"
1267 );
1268 }
1269
1270 #[test]
1271 fn collapse_generation_resets_on_autocompact() {
1272 let mut mgr = ContextManager::new(1_000);
1273 for i in 0..20 {
1276 mgr.push_history(tool_result_msg(&format!("c{i}"), &"x".repeat(120)), 60);
1277 }
1278 mgr.set_observed_prompt_tokens(980); mgr.recompute_handle_residency();
1280 assert_eq!(
1281 mgr.handles.residency_for_source("c0"),
1282 Some(&Residency::Collapsed)
1283 );
1284
1285 let (saved, _, archived, _) = mgr.compress(PressureAction::AutoCompact);
1286 assert!(saved > 0 && !archived.is_empty(), "expected archival");
1287
1288 for h in mgr.handles.all() {
1291 if matches!(h.kind, HandleKind::ToolResult) {
1292 assert_eq!(
1293 h.residency,
1294 Residency::Resident,
1295 "generation reset un-collapses survivors"
1296 );
1297 }
1298 }
1299 }
1300
1301 #[test]
1302 fn mark_spooled_sets_residency_and_survives_residency_recompute() {
1303 let mut mgr = ContextManager::new(1_000);
1304 mgr.push_history(
1305 Message::tool(vec![ContentPart::ToolResult {
1306 call_id: "big".into(),
1307 output: "preview only".to_string(),
1308 is_error: false,
1309 }]),
1310 10,
1311 );
1312 mgr.mark_spooled("big", "disk://big");
1313 assert_eq!(
1314 mgr.handles.residency_for_source("big"),
1315 Some(&Residency::SpooledOut {
1316 r: "disk://big".to_string()
1317 })
1318 );
1319
1320 mgr.set_observed_prompt_tokens(990);
1323 mgr.recompute_handle_residency();
1324 assert_eq!(
1325 mgr.handles.residency_for_source("big"),
1326 Some(&Residency::SpooledOut {
1327 r: "disk://big".to_string()
1328 })
1329 );
1330 }
1331
1332 #[test]
1333 fn push_history_indexes_tool_results_as_resident_handles() {
1334 let mut mgr = ContextManager::new(10_000);
1335 let msg = Message::tool(vec![ContentPart::ToolResult {
1336 call_id: "call_1".into(),
1337 output: "the tool output".to_string(),
1338 is_error: false,
1339 }]);
1340 mgr.push_history(msg, 20);
1341 assert_eq!(mgr.handles.all().len(), 1);
1343 assert_eq!(
1344 mgr.handles.residency_for_source("call_1"),
1345 Some(&Residency::Resident)
1346 );
1347 mgr.push_history(Message::user("hello"), 5);
1349 assert_eq!(mgr.handles.all().len(), 1);
1350 }
1351
1352 fn tool_result_msg(call_id: &str, output: &str) -> Message {
1355 Message::tool(vec![ContentPart::ToolResult {
1356 call_id: call_id.into(),
1357 output: output.to_string(),
1358 is_error: false,
1359 }])
1360 }
1361
1362 #[test]
1363 fn prune_orphaned_handles_drops_handles_whose_message_left_history() {
1364 let mut mgr = ContextManager::new(10_000);
1365 mgr.push_history(tool_result_msg("c0", "out 0"), 20);
1366 mgr.push_history(tool_result_msg("c1", "out 1"), 20);
1367 assert_eq!(mgr.handles.all().len(), 2);
1368
1369 mgr.partitions.history.messages.remove(0);
1371 mgr.prune_orphaned_handles();
1372
1373 assert_eq!(mgr.handles.all().len(), 1);
1375 assert!(mgr.handles.residency_for_source("c0").is_none());
1376 assert_eq!(
1377 mgr.handles.residency_for_source("c1"),
1378 Some(&Residency::Resident)
1379 );
1380 }
1381
1382 #[test]
1383 fn autocompact_prunes_handles_for_archived_tool_results() {
1384 let mut mgr = ContextManager::new(1_000);
1385 for i in 0..30 {
1387 mgr.push_history(tool_result_msg(&format!("c{i}"), &"x".repeat(200)), 80);
1388 }
1389 assert_eq!(mgr.handles.all().len(), 30);
1390
1391 let (saved, _, archived, _) = mgr.compress(PressureAction::AutoCompact);
1392 assert!(saved > 0 && !archived.is_empty(), "expected archival");
1393
1394 let live_tool_results = mgr
1397 .partitions
1398 .history
1399 .messages
1400 .iter()
1401 .filter(|m| {
1402 matches!(&m.content, Content::Parts(p)
1403 if p.iter().any(|x| matches!(x, ContentPart::ToolResult { .. })))
1404 })
1405 .count();
1406 assert_eq!(mgr.handles.all().len(), live_tool_results);
1407 assert!(
1408 mgr.handles.all().len() < 30,
1409 "table must shrink with archival"
1410 );
1411 }
1412
1413 #[test]
1414 fn renew_prunes_handles_for_dropped_history() {
1415 let mut mgr = ContextManager::new(1_000);
1416 mgr.init_task("g".to_string(), vec![]);
1417 for i in 0..20 {
1418 mgr.push_history(tool_result_msg(&format!("c{i}"), "data"), 60);
1419 }
1420 mgr.renew();
1421 for h in mgr.handles.all() {
1423 if let Some(src) = h.source.as_ref() {
1424 assert!(
1425 mgr.handles.residency_for_source(src).is_some(),
1426 "no dangling handle survives renewal"
1427 );
1428 }
1429 }
1430 assert!(mgr.handles.all().len() <= 20);
1431 }
1432
1433 #[test]
1434 fn recompute_residency_index_semantics_with_spooled_in_the_middle() {
1435 let mut mgr = ContextManager::new(1_000);
1438 for i in 0..6 {
1439 mgr.push_history(tool_result_msg(&format!("c{i}"), &"y".repeat(40)), 40);
1440 }
1441 mgr.mark_spooled("c2", "disk://c2");
1442
1443 mgr.set_observed_prompt_tokens(950); mgr.recompute_handle_residency();
1445
1446 assert_eq!(
1448 mgr.handles.residency_for_source("c2"),
1449 Some(&Residency::SpooledOut {
1450 r: "disk://c2".to_string()
1451 })
1452 );
1453 assert_eq!(
1454 mgr.handles.residency_for_source("c0"),
1455 Some(&Residency::Collapsed)
1456 );
1457 assert_eq!(
1458 mgr.handles.residency_for_source("c5"),
1459 Some(&Residency::Resident)
1460 );
1461 }
1462
1463 #[test]
1466 fn knowledge_budget_uses_stable_order_for_equal_value_and_warns_once() {
1467 let mut mgr = ContextManager::new(100);
1470 mgr.push_knowledge(Message::system("oldest unkeyed"), 10);
1471 mgr.push_knowledge_entry(Some("a".into()), Message::system("keyed"), 10, false);
1472 mgr.push_knowledge_entry(Some("p".into()), Message::system("pinned"), 10, true);
1473 mgr.push_knowledge_entry(Some("skill:x".into()), Message::system("skill"), 10, false);
1474
1475 let warn = mgr.enforce_knowledge_budget();
1476 assert_eq!(warn, Some((40, 25)));
1477 let e = &mgr.partitions.knowledge.entries;
1479 assert!(e[0].evict_at_boundary);
1480 assert!(e[1].evict_at_boundary);
1481 assert!(!e[2].evict_at_boundary, "pinned exempt");
1482 assert!(!e[3].evict_at_boundary, "skill pin exempt");
1483
1484 assert_eq!(mgr.enforce_knowledge_budget(), None);
1486
1487 let sweep = mgr.partitions.knowledge.sweep_at_boundary();
1489 assert_eq!(sweep.tokens_freed, 20);
1490 assert_eq!(mgr.partitions.knowledge.token_count, 20);
1491 assert_eq!(mgr.enforce_knowledge_budget(), None);
1493 }
1494
1495 #[test]
1496 fn knowledge_budget_warning_stands_when_only_exempt_weight_remains() {
1497 let mut mgr = ContextManager::new(100);
1498 mgr.push_knowledge_entry(Some("p".into()), Message::system("pinned heavy"), 30, true);
1499 mgr.push_knowledge_entry(
1500 Some("skill:x".into()),
1501 Message::system("skill heavy"),
1502 30,
1503 false,
1504 );
1505
1506 assert_eq!(mgr.enforce_knowledge_budget(), Some((60, 25)));
1508 assert!(
1509 mgr.partitions
1510 .knowledge
1511 .entries
1512 .iter()
1513 .all(|e| !e.evict_at_boundary)
1514 );
1515 }
1516
1517 #[test]
1518 fn knowledge_budget_retains_old_referenced_entry_over_new_irrelevant_entry() {
1519 let mut mgr = ContextManager::new(100);
1520 mgr.push_knowledge_entry(
1521 Some("project:orchid".into()),
1522 Message::system("ORCHID uses the Atlas storage engine"),
1523 10,
1524 false,
1525 );
1526 mgr.push_history(Message::user("For project:orchid keep using Atlas"), 5);
1528 mgr.push_knowledge_entry(
1529 Some("project:new".into()),
1530 Message::system("unrelated fresh material"),
1531 10,
1532 false,
1533 );
1534 mgr.push_knowledge_entry(
1535 Some("project:other".into()),
1536 Message::system("another unused reference"),
1537 10,
1538 false,
1539 );
1540
1541 assert_eq!(mgr.enforce_knowledge_budget(), Some((30, 25)));
1542 let entries = &mgr.partitions.knowledge.entries;
1543 assert!(
1544 !entries[0].evict_at_boundary,
1545 "a real reference must raise retention"
1546 );
1547 assert!(
1548 entries[1].evict_at_boundary,
1549 "lowest-value entry evicts first"
1550 );
1551 assert!(
1552 !entries[2].evict_at_boundary,
1553 "one eviction is enough to fit"
1554 );
1555 }
1556
1557 #[test]
1558 fn knowledge_budget_ratio_zero_disables() {
1559 let mut mgr = ContextManager::new(100);
1560 mgr.config.knowledge_budget_ratio = 0.0;
1561 mgr.push_knowledge(Message::system("huge"), 90);
1562 assert_eq!(mgr.enforce_knowledge_budget(), None);
1563 assert!(!mgr.partitions.knowledge.entries[0].evict_at_boundary);
1564 }
1565
1566 #[test]
1567 fn provider_and_output_reservations_reduce_the_hard_input_budget() {
1568 use crate::context::config::PromptBudgetConfig;
1569
1570 let mut mgr = ContextManager::new(100);
1571 mgr.set_prompt_budget(PromptBudgetConfig {
1572 prompt_overhead_tokens: 20,
1573 output_reserve_tokens: 20,
1574 safety_margin_tokens: 10,
1575 });
1576 mgr.partitions
1577 .system
1578 .push(Message::system("x".repeat(240)), 60);
1579
1580 let rendered = mgr.render();
1581 let overflow = rendered
1582 .budget_overflow
1583 .expect("fixed context exceeds input allowance");
1584 assert_eq!(overflow.max_tokens, 50);
1585 assert!(overflow.required_tokens > overflow.max_tokens);
1586 }
1587}