1use super::compression::CompressionPipeline;
2use super::config::ContextConfig;
3use super::partitions::ContextPartitions;
4use super::pressure::{PressureAction, PressureMonitor};
5use super::renderer::RenderedContext;
6use super::renewal::RenewalPolicy;
7use super::skill_catalog::SkillCatalog;
8use super::task_state::{TaskState, TaskUpdate};
9use super::token_engine::ContextTokenEngine;
10use crate::mm::handle::{Handle, HandleId, HandleKind, HandleTable, Residency};
11use crate::types::capability::{CapabilityKind, CapabilityManifest};
12use crate::types::message::{Content, ContentPart, Message, ToolSchema};
13use crate::types::skill::SkillMetadata;
14use compact_str::CompactString;
15
16pub const MEMORY_TOOL_NAME: &str = "memory";
17pub const KNOWLEDGE_TOOL_NAME: &str = "knowledge";
18pub const READ_RESULT_TOOL_NAME: &str = "read_result";
20
21const META_TOOL_NAMES: &[&str] = &[
24 "pace",
25 "update_plan",
26 "skill",
27 MEMORY_TOOL_NAME,
28 KNOWLEDGE_TOOL_NAME,
29 READ_RESULT_TOOL_NAME,
30 "submit_workflow_nodes",
31 "start_workflow",
32];
33
34pub(crate) fn is_meta_tool(name: &str) -> bool {
37 META_TOOL_NAMES.contains(&name)
38}
39
40#[doc(hidden)]
45pub struct ContextManager {
46 pub partitions: ContextPartitions,
47 pub max_tokens: u32,
48 pub config: ContextConfig,
49 pub engine: ContextTokenEngine,
50 pub sprint: u32,
51 pub skills: SkillCatalog,
52 pub active_skills: std::collections::BTreeMap<CompactString, Option<u32>>,
60 pub stable_core_tools: std::collections::HashSet<CompactString>,
64 pub capabilities: CapabilityManifest,
65 pub memory_enabled: bool,
66 pub knowledge_enabled: bool,
67 pub plan_tool_enabled: bool,
68 last_observed_prompt_tokens: Option<u32>,
69 compression: CompressionPipeline,
70 pressure: PressureMonitor,
71 renewal: RenewalPolicy,
72
73 pub last_activity_ms: u64,
78
79 pub last_compact_ms: Option<u64>,
82
83 pub handles: HandleTable,
89 next_handle_id: HandleId,
91
92 frozen_history_len: usize,
98
99 pending_knowledge_sweeps: Vec<crate::context::partitions::KnowledgeSweep>,
102
103 knowledge_budget_warned: bool,
106}
107
108impl ContextManager {
109 pub fn new(max_tokens: u32) -> Self {
110 Self::with_config(max_tokens, ContextConfig::default(), ContextTokenEngine::char_approx())
111 }
112
113 pub fn with_config(max_tokens: u32, config: ContextConfig, engine: ContextTokenEngine) -> Self {
114 let compression = CompressionPipeline::new(&config);
115 let pressure = PressureMonitor::new(max_tokens, config.clone());
116 let renewal = RenewalPolicy::from_config(&config);
117 let partitions = ContextPartitions::new(&config);
118 Self {
119 partitions, max_tokens, config, engine,
120 sprint: 0,
121 skills: SkillCatalog::new(),
122 active_skills: std::collections::BTreeMap::new(),
123 stable_core_tools: std::collections::HashSet::new(),
124 capabilities: CapabilityManifest::new(),
125 memory_enabled: false, knowledge_enabled: false, plan_tool_enabled: false,
126 last_observed_prompt_tokens: None,
127 compression, pressure, renewal,
128 last_activity_ms: 0,
129 last_compact_ms: None,
130 handles: HandleTable::new(),
131 next_handle_id: 0,
132 frozen_history_len: 0,
133 pending_knowledge_sweeps: Vec::new(),
134 knowledge_budget_warned: false,
135 }
136 }
137
138 pub fn record_activity(&mut self, now_ms: u64) {
142 self.last_activity_ms = now_ms;
143 }
144
145 pub fn should_time_decay_compact(&self, now_ms: u64) -> bool {
148 let idle_ms = if let Some(last_compact) = self.last_compact_ms {
149 now_ms.saturating_sub(last_compact)
151 } else {
152 now_ms.saturating_sub(self.last_activity_ms)
154 };
155
156 let idle_minutes = idle_ms / 60_000;
157 idle_minutes >= self.config.micro_compact_idle_minutes as u64
158 }
159
160 pub fn recompute_handle_residency(&mut self) {
175 if self.rho() < self.config.collapse_threshold {
177 return;
178 }
179 let keep = self.config.preserve_recent_msgs;
180 let total = self
183 .handles
184 .all()
185 .iter()
186 .filter(|h| matches!(h.kind, HandleKind::ToolResult))
187 .count();
188 let cutoff = total.saturating_sub(keep);
189 for (i, handle) in self.handles.tool_result_handles_mut().enumerate() {
190 if i < cutoff && matches!(handle.residency, Residency::Resident) {
193 handle.residency = Residency::Collapsed;
194 }
195 }
196 }
197
198 pub fn reset_collapse_generation(&mut self) {
204 for handle in self.handles.all_mut() {
205 if matches!(handle.residency, Residency::Collapsed) {
206 handle.residency = Residency::Resident;
207 }
208 }
209 }
210
211 pub fn prune_orphaned_handles(&mut self) {
218 let live: std::collections::HashSet<CompactString> = self
219 .partitions
220 .history
221 .messages
222 .iter()
223 .flat_map(|m| match &m.content {
224 Content::Parts(parts) => parts
225 .iter()
226 .filter_map(|p| match p {
227 ContentPart::ToolResult { call_id, .. } => Some(call_id.clone()),
228 _ => None,
229 })
230 .collect::<Vec<_>>(),
231 _ => Vec::new(),
232 })
233 .collect();
234 self.handles
235 .retain(|h| h.source.as_ref().is_none_or(|s| live.contains(s)));
236 }
237
238 pub fn mark_spooled(&mut self, call_id: &str, spool_ref: impl Into<String>) {
242 let spool_ref = spool_ref.into();
243 if let Some(handle) = self
244 .handles
245 .all_mut()
246 .iter_mut()
247 .find(|h| h.source.as_deref() == Some(call_id))
248 {
249 handle.residency = Residency::SpooledOut { r: spool_ref };
250 }
251 }
252
253 pub fn rho(&self) -> f64 {
260 self.pressure
261 .pressure(&self.partitions, &self.engine, self.last_observed_prompt_tokens)
262 }
263
264 pub fn set_observed_prompt_tokens(&mut self, tokens: u32) {
265 self.last_observed_prompt_tokens = Some(tokens);
266 }
267
268 pub fn should_compress(&self) -> PressureAction {
269 self.pressure.recommend(self.rho())
275 }
276
277 pub fn compress(&mut self, action: PressureAction) -> (u32, Option<String>, Vec<Message>, Option<usize>) {
278 self.compress_with_time(action, None)
279 }
280
281 pub fn compress_with_time(
282 &mut self,
283 action: PressureAction,
284 now_ms: Option<u64>,
285 ) -> (u32, Option<String>, Vec<Message>, Option<usize>) {
286 let target = self.config.target_tokens(self.max_tokens);
287 self.compress_with_target(action, target, now_ms)
288 }
289
290 pub fn force_compress(&mut self) -> (u32, Option<String>, Vec<Message>, Option<usize>) {
291 self.compress_with_target(PressureAction::AutoCompact, 0, None)
292 }
293
294 pub fn compress_with_target(
301 &mut self,
302 action: PressureAction,
303 target_tokens: u32,
304 now_ms: Option<u64>,
305 ) -> (u32, Option<String>, Vec<Message>, Option<usize>) {
306 let result =
307 self.compression
308 .compress(&mut self.partitions, action, self.max_tokens, target_tokens, &self.engine);
309 if let Some(ts) = now_ms {
310 self.last_compact_ms = Some(ts);
311 }
312 if !result.2.is_empty() {
314 self.prune_orphaned_handles();
315 self.reset_collapse_generation();
318 self.sweep_knowledge_at_boundary();
321 }
322 if result.3.is_some() {
329 self.frozen_history_len = self.partitions.history.messages.len();
330 }
331 result
332 }
333
334 pub fn plan_compaction_params(&self) -> (u32, usize) {
338 (
339 self.config.target_tokens(self.max_tokens),
340 self.config.preserve_recent_turns,
341 )
342 }
343
344 pub fn should_renew(&self) -> bool {
347 self.renewal.should_renew(&self.pressure, &self.partitions, &self.engine)
348 }
349
350 pub fn renew(&mut self) {
351 self.partitions = self.renewal.renew(&self.partitions, self.max_tokens);
352 self.sprint += 1;
353 self.prune_orphaned_handles();
356 self.reset_collapse_generation();
357 self.sweep_knowledge_at_boundary();
359 self.frozen_history_len = self.partitions.history.messages.len();
361 }
362
363 pub fn render(&self) -> RenderedContext {
366 super::renderer::render_projected(
367 &self.partitions,
368 self.max_tokens,
369 &self.engine,
370 self.config.preserve_recent_msgs,
371 &self.handles,
372 self.frozen_history_len,
373 self.config.collapse_assistant_narration,
374 )
375 }
376
377 pub fn push_history(&mut self, msg: Message, tokens: u32) {
380 if let Content::Parts(parts) = &msg.content {
384 for part in parts {
385 if let ContentPart::ToolResult { call_id, output, .. } = part {
386 let id = self.alloc_handle_id();
387 let tok = self.engine.count(output).max(1);
388 self.handles.insert(Handle::resident_for(
389 id,
390 HandleKind::ToolResult,
391 tok,
392 call_id.clone(),
393 ));
394 }
395 }
396 }
397 self.partitions.history.push(msg, tokens);
398 }
399
400 fn alloc_handle_id(&mut self) -> HandleId {
401 let id = self.next_handle_id;
402 self.next_handle_id = self.next_handle_id.wrapping_add(1);
403 id
404 }
405
406 pub fn push_knowledge(&mut self, msg: Message, tokens: u32) {
408 self.partitions.knowledge.push(msg, tokens);
409 }
410
411 pub fn push_knowledge_entry(
415 &mut self,
416 key: Option<CompactString>,
417 msg: Message,
418 tokens: u32,
419 pinned: bool,
420 ) {
421 self.partitions.knowledge.push_entry(key, msg, tokens, pinned);
422 }
423
424 pub fn remove_knowledge(&mut self, key: &str) -> bool {
427 self.partitions.knowledge.remove(key)
428 }
429
430 fn sweep_knowledge_at_boundary(&mut self) {
434 let sweep = self.partitions.knowledge.sweep_at_boundary();
435 if sweep.changed {
436 if !sweep.removed_keys.is_empty() {
440 self.partitions.signals.push(format!(
441 "[KNOWLEDGE] entries removed at this boundary: {} — re-fetch via the memory tool if still needed.",
442 sweep.removed_keys.join(", ")
443 ));
444 }
445 self.pending_knowledge_sweeps.push(sweep);
446 }
447 self.knowledge_budget_warned = false;
449 }
450
451 pub fn enforce_knowledge_budget(&mut self) -> Option<(u32, u32)> {
460 let ratio = self.config.knowledge_budget_ratio;
461 if ratio <= 0.0 {
462 return None;
463 }
464 let budget = (self.max_tokens as f64 * ratio) as u32;
465 let used = self.partitions.knowledge.token_count;
466 if used <= budget {
467 return None;
468 }
469 let entries = &mut self.partitions.knowledge.entries;
470 let marked: u32 = entries.iter().filter(|e| e.evict_at_boundary).map(|e| e.tokens).sum();
471 let mut projected = used.saturating_sub(marked);
472 for entry in entries.iter_mut() {
473 if projected <= budget {
474 break;
475 }
476 if entry.evict_at_boundary || entry.pinned {
477 continue;
478 }
479 if entry.key.as_deref().is_some_and(|k| k.starts_with("skill:")) {
480 continue;
481 }
482 entry.evict_at_boundary = true;
483 projected = projected.saturating_sub(entry.tokens);
484 }
485 if self.knowledge_budget_warned {
486 return None;
487 }
488 self.knowledge_budget_warned = true;
489 Some((used, budget))
490 }
491
492 pub fn take_knowledge_sweeps(&mut self) -> Vec<crate::context::partitions::KnowledgeSweep> {
494 std::mem::take(&mut self.pending_knowledge_sweeps)
495 }
496
497 pub fn push_signal(&mut self, text: String) {
500 self.partitions.signals.push(text);
501 }
502
503 pub fn record_directive(&mut self, text: impl Into<String>) {
507 self.partitions.task_state.record_directive(text);
508 }
509
510 pub fn init_task(&mut self, goal: String, criteria: Vec<String>) {
513 self.partitions.task_state = TaskState { goal, criteria, ..Default::default() };
514 }
515
516 pub fn update_task(&mut self, update: TaskUpdate) {
517 self.partitions.task_state.apply(update);
518 }
519
520 pub fn note_tool_actions(&mut self, calls: &[(String, String)]) {
527 let summary = calls
528 .iter()
529 .filter(|(name, _)| !is_meta_tool(name))
530 .map(|(name, args)| {
531 if args.is_empty() { name.clone() } else { format!("{name}({args})") }
532 })
533 .collect::<Vec<_>>()
534 .join(", ");
535 self.partitions.task_state.note_actions(summary);
536 }
537
538 pub fn set_available_skills(&mut self, skills: Vec<SkillMetadata>) {
544 self.capabilities.remove_kind(CapabilityKind::Skill);
545 for skill in &skills { self.capabilities.add_skill(skill.clone()); }
546 self.skills.set_available(skills);
547 }
548
549 pub fn set_stable_core_tools(&mut self, ids: impl IntoIterator<Item = CompactString>) {
551 self.stable_core_tools = ids.into_iter().collect();
552 }
553
554 pub fn activate_skill(&mut self, name: impl Into<CompactString>) -> bool {
559 self.activate_skill_leased(name, None)
560 }
561
562 pub fn activate_skill_leased(
565 &mut self,
566 name: impl Into<CompactString>,
567 expires_at_turn: Option<u32>,
568 ) -> bool {
569 self.active_skills.insert(name.into(), expires_at_turn).is_none()
570 }
571
572 pub fn deactivate_skill(&mut self, name: &str) -> bool {
576 if self.active_skills.remove(name).is_none() {
577 return false;
578 }
579 self.partitions.knowledge.remove(&format!("skill:{name}"));
580 true
581 }
582
583 pub fn sweep_expired_skill_leases(&mut self, current_turn: u32) {
586 let expired: Vec<CompactString> = self
587 .active_skills
588 .iter()
589 .filter(|(_, lease)| lease.is_some_and(|t| current_turn >= t))
590 .map(|(name, _)| name.clone())
591 .collect();
592 for name in expired {
593 self.deactivate_skill(&name);
594 self.partitions.signals.push(format!(
596 "[SKILL] lease expired: {name} unloaded; the full toolset is restored."
597 ));
598 }
599 }
600
601 pub fn active_skill_tool_filter(&self) -> Option<std::collections::HashSet<CompactString>> {
607 if self.active_skills.is_empty() {
608 return None;
609 }
610 let mut union = std::collections::HashSet::new();
611 for name in self.active_skills.keys() {
612 let declared = self.skills.allowed_tools(name);
613 if declared.is_empty() {
614 return None; }
616 union.extend(declared.iter().cloned());
617 }
618 Some(union)
619 }
620
621 pub fn skill_tool_schema(&self) -> Option<ToolSchema> {
622 self.skills.build_tool_schema()
623 }
624
625 pub fn set_memory_enabled(&mut self, enabled: bool) {
628 self.memory_enabled = enabled;
629 if enabled {
630 self.capabilities.add_marker(CapabilityKind::Memory, MEMORY_TOOL_NAME,
631 "Search long-term memory through the memory meta-tool.");
632 } else {
633 self.capabilities.remove(CapabilityKind::Memory, MEMORY_TOOL_NAME);
634 }
635 }
636
637 pub fn set_knowledge_enabled(&mut self, enabled: bool) {
638 self.knowledge_enabled = enabled;
639 if enabled {
640 self.capabilities.add_marker(CapabilityKind::Knowledge, KNOWLEDGE_TOOL_NAME,
641 "Search external knowledge through the knowledge meta-tool.");
642 } else {
643 self.capabilities.remove(CapabilityKind::Knowledge, KNOWLEDGE_TOOL_NAME);
644 }
645 }
646
647 pub fn set_plan_tool_enabled(&mut self, enabled: bool) {
648 self.plan_tool_enabled = enabled;
649 if enabled {
650 self.capabilities.add_marker(CapabilityKind::Tool, "update_plan",
651 "Update task plan and progress through the planning meta-tool.");
652 } else {
653 self.capabilities.remove(CapabilityKind::Tool, "update_plan");
654 }
655 }
656
657 pub fn capability_inventory(&self) -> String { self.capabilities.format_inventory() }
658
659 pub fn meta_tool_schemas(&self) -> Vec<ToolSchema> {
660 let mut tools = Vec::new();
661 if let Some(t) = self.skill_tool_schema() { tools.push(t); }
662 if let Some(t) = self.memory_tool_schema() { tools.push(t); }
663 if let Some(t) = self.knowledge_tool_schema() { tools.push(t); }
664 if let Some(t) = self.plan_tool_schema() { tools.push(t); }
665 if let Some(t) = self.read_result_tool_schema() { tools.push(t); }
666 tools.sort_by(|a, b| a.name.cmp(&b.name));
667 tools
668 }
669
670 pub fn read_result_tool_schema(&self) -> Option<ToolSchema> {
676 let any_evicted = self
677 .handles
678 .all()
679 .iter()
680 .any(|h| !h.residency.occupies_context());
681 if !any_evicted { return None; }
682 Some(ToolSchema {
683 name: CompactString::new(READ_RESULT_TOOL_NAME),
684 description: "Re-read a tool result that was evicted from context (you see a \
685 placeholder like '[…tool result spooled…]' or a collapsed entry). \
686 Pass the tool call's call_id; use offset/max_bytes to page through \
687 large content."
688 .to_string(),
689 parameters: serde_json::json!({
690 "type": "object",
691 "properties": {
692 "call_id": { "type": "string" },
693 "offset": { "type": "integer", "description": "Byte offset to start from (default 0)." },
694 "max_bytes": { "type": "integer", "description": "Max bytes to return (default 4000)." }
695 },
696 "required": ["call_id"]
697 }),
698 })
699 }
700
701 pub fn plan_tool_schema(&self) -> Option<ToolSchema> {
702 if !self.plan_tool_enabled { return None; }
703 Some(ToolSchema {
704 name: CompactString::new("update_plan"),
705 description: "Update your task plan and progress. Call this after completing a step or when the plan changes.".to_string(),
706 parameters: serde_json::json!({
707 "type": "object",
708 "properties": {
709 "plan": { "type": "array", "items": { "type": "string" } },
710 "current_step": { "type": "integer" },
711 "progress": { "type": "string" },
712 "blocked_on": { "type": "array", "items": { "type": "string" } }
713 }
714 }),
715 })
716 }
717
718 pub fn memory_tool_schema(&self) -> Option<ToolSchema> {
719 if !self.memory_enabled { return None; }
720 Some(ToolSchema {
721 name: CompactString::new(MEMORY_TOOL_NAME),
722 description: "Search your long-term memory for relevant past experiences and knowledge.".to_string(),
723 parameters: serde_json::json!({
724 "type": "object",
725 "properties": {
726 "query": { "type": "string" },
727 "top_k": { "type": "integer" }
728 },
729 "required": ["query"]
730 }),
731 })
732 }
733
734 pub fn knowledge_tool_schema(&self) -> Option<ToolSchema> {
735 if !self.knowledge_enabled { return None; }
736 Some(ToolSchema {
737 name: CompactString::new(KNOWLEDGE_TOOL_NAME),
738 description: "Search the external knowledge base for facts, documentation, or reference data.".to_string(),
739 parameters: serde_json::json!({
740 "type": "object",
741 "properties": {
742 "query": { "type": "string" },
743 "top_k": { "type": "integer" }
744 },
745 "required": ["query"]
746 }),
747 })
748 }
749}
750
751#[cfg(test)]
752mod tests {
753 use super::*;
754 use crate::context::task_state::PlanStep;
755 use crate::types::message::Message;
756 use crate::types::skill::SkillMetadata;
757
758 #[test]
759 fn note_tool_actions_keys_on_name_and_args_so_legit_loops_dont_false_stop() {
760 let mut mgr = ContextManager::new(100_000);
763 mgr.init_task("process items".to_string(), vec![]);
764 mgr.note_tool_actions(&[("step".to_string(), "{\"n\":1}".to_string())]);
765 mgr.note_tool_actions(&[("step".to_string(), "{\"n\":2}".to_string())]);
766 mgr.note_tool_actions(&[("step".to_string(), "{\"n\":3}".to_string())]);
767 assert_eq!(
768 mgr.partitions.task_state.recent_actions,
769 ["step({\"n\":1})", "step({\"n\":2})", "step({\"n\":3})"]
770 );
771 let txt = mgr.render().state_turn.unwrap().content.as_text().unwrap().to_string();
772 assert!(!txt.contains("STOP:"), "same-tool/diff-args loop must not trip STOP: {txt}");
773
774 let mut mgr2 = ContextManager::new(100_000);
776 mgr2.init_task("g".to_string(), vec![]);
777 for _ in 0..3 {
778 mgr2.note_tool_actions(&[("document_read".to_string(), "{\"id\":\"x\"}".to_string())]);
779 }
780 let txt2 = mgr2.render().state_turn.unwrap().content.as_text().unwrap().to_string();
781 assert!(txt2.contains("STOP:"), "identical repeated call must trip STOP: {txt2}");
782
783 let mut mgr3 = ContextManager::new(100_000);
785 mgr3.init_task("g".to_string(), vec![]);
786 mgr3.note_tool_actions(&[("update_plan".to_string(), "{\"current_step\":1}".to_string())]);
787 assert!(mgr3.partitions.task_state.recent_actions.is_empty());
788 }
789
790 #[test]
791 fn manager_renew_advances_sprint_and_keeps_goal() {
792 let mut mgr = ContextManager::new(1_000);
793 mgr.init_task("test goal".to_string(), vec![]);
794 mgr.partitions.system.push(Message::system("rules"), 10);
795 for i in 0..10 { mgr.push_history(Message::user(format!("msg {i}")), 50); }
796 mgr.renew();
797 assert_eq!(mgr.partitions.task_state.goal, "test goal");
798 assert_eq!(mgr.sprint, 1);
799 }
800
801 #[test]
802 fn compress_only_touches_history() {
803 let mut mgr = ContextManager::new(1_000);
804 mgr.push_knowledge(Message::system("knowledge content"), 100);
805 for _ in 0..30 { mgr.push_history(Message::user("history msg"), 50); }
806 let knowledge_before = mgr.partitions.knowledge.token_count;
807 let history_before = mgr.partitions.history.token_count;
808 mgr.compress(PressureAction::AutoCompact);
809 assert_eq!(mgr.partitions.knowledge.token_count, knowledge_before);
810 assert!(mgr.partitions.history.token_count < history_before);
811 }
812
813 #[test]
814 fn init_task_sets_goal_and_criteria() {
815 let mut mgr = ContextManager::new(1_000);
816 mgr.init_task("analyse data".to_string(), vec!["criterion A".to_string()]);
817 assert_eq!(mgr.partitions.task_state.goal, "analyse data");
818 assert_eq!(mgr.partitions.task_state.criteria, ["criterion A"]);
819 }
820
821 #[test]
822 fn update_task_applies_plan() {
823 let mut mgr = ContextManager::new(1_000);
824 mgr.init_task("g".to_string(), vec![]);
825 mgr.update_task(TaskUpdate {
826 plan: Some(vec!["step 1".to_string(), "step 2".to_string()]),
827 current_step: Some(0),
828 ..Default::default()
829 });
830 assert_eq!(mgr.partitions.task_state.plan.len(), 2);
831 assert_eq!(mgr.partitions.task_state.current_step, Some(0));
832 }
833
834 #[test]
835 fn task_state_survives_autocompact() {
836 let mut mgr = ContextManager::new(1_000);
837 mgr.init_task("survive compression".to_string(), vec![]);
838 mgr.update_task(TaskUpdate {
839 plan: Some(vec!["fetch data".to_string(), "analyse".to_string()]),
840 ..Default::default()
841 });
842 for _ in 0..10 { mgr.push_history(Message::user("filler"), 50); }
843 mgr.compress(PressureAction::AutoCompact);
844 assert_eq!(mgr.partitions.task_state.goal, "survive compression");
845 assert_eq!(mgr.partitions.task_state.plan.len(), 2);
846 }
847
848 #[test]
849 fn render_includes_task_state_in_state_turn_not_system() {
850 let mut mgr = ContextManager::new(10_000);
851 mgr.init_task("find anomalies".to_string(), vec![]);
852 let rc = mgr.render();
853 assert!(!rc.system_text.contains("[TASK STATE]"), "task_state must not be in system_text");
854 let state = rc.state_turn.as_ref().expect("should have a state turn");
856 assert!(state.content.as_text().unwrap().contains("[TASK STATE] goal: find anomalies"));
857 }
858
859 #[test]
860 fn renewal_keeps_open_plan_steps_in_task_state() {
861 let mut mgr = ContextManager::new(1_000);
862 mgr.init_task("g".to_string(), vec![]);
863 mgr.partitions.task_state.plan = vec![
864 PlanStep { label: "done".to_string(), done: true },
865 PlanStep { label: "pending".to_string(), done: false },
866 ];
867 mgr.renew();
868 assert_eq!(mgr.partitions.task_state.open_steps(), vec!["pending"]);
869 }
870
871 #[test]
874 fn auto_compact_entry_logs_auto_compact_action() {
875 let mut mgr = ContextManager::new(1_000);
883 for i in 0..40 {
884 mgr.push_history(Message::user(format!("turn {i}: {}", "ctx ".repeat(40))), 200);
885 }
886 let (saved, summary, _, _) = mgr.force_compress();
887 assert!(saved > 0, "force_compress should compact a large history");
888 assert!(summary.is_some(), "auto-compact summarizes the archived turns");
889 let actions: Vec<&str> = mgr
890 .partitions
891 .task_state
892 .compression_log
893 .iter()
894 .map(|e| e.action.as_str())
895 .collect();
896 assert!(
897 actions.last() == Some(&"auto_compact"),
898 "auto-compact entry must log an auto_compact action; got {actions:?}"
899 );
900 }
901
902 #[test]
903 fn skill_tool_schema_empty_when_no_skills() {
904 let mgr = ContextManager::new(10_000);
905 assert!(mgr.skill_tool_schema().is_none());
906 }
907
908 #[test]
909 fn skill_tool_schema_present_when_registered() {
910 let mut mgr = ContextManager::new(10_000);
911 mgr.set_available_skills(vec![SkillMetadata::new("debug", "Debug helper")]);
912 assert!(mgr.skill_tool_schema().unwrap().description.contains("debug"));
913 }
914
915 #[test]
916 fn available_skills_are_reflected_in_capability_manifest() {
917 let mut mgr = ContextManager::new(1_000);
918 mgr.set_available_skills(vec![SkillMetadata::new("debug", "Debug helper")]);
919 let inventory = mgr.capability_inventory();
920 assert!(inventory.contains("debug"));
921 assert!(inventory.contains("Debug helper"));
922 }
923
924 #[test]
925 fn toggled_meta_tools_are_reflected_in_capability_manifest() {
926 let mut mgr = ContextManager::new(1_000);
927 mgr.set_memory_enabled(true);
928 assert!(mgr.capability_inventory().contains(MEMORY_TOOL_NAME));
929 mgr.set_memory_enabled(false);
930 assert!(!mgr.capability_inventory().contains(MEMORY_TOOL_NAME));
931 }
932
933 #[test]
934 fn meta_tool_schemas_are_sorted() {
935 let mut mgr = ContextManager::new(1_000);
936 mgr.set_available_skills(vec![SkillMetadata::new("debug", "Debug helper")]);
937 mgr.set_memory_enabled(true);
938 mgr.set_knowledge_enabled(true);
939 let names = mgr.meta_tool_schemas().into_iter().map(|s| s.name.to_string()).collect::<Vec<_>>();
940 assert_eq!(names, ["knowledge", "memory", "skill"]);
941 }
942
943 #[test]
944 fn b1_active_skill_state_and_tool_filter() {
945 let mut mgr = ContextManager::new(1_000);
946 let mut debug = SkillMetadata::new("debug", "Debug helper");
947 debug.allowed_tools = vec![CompactString::new("read"), CompactString::new("grep")];
948 let mut review = SkillMetadata::new("review", "Reviewer");
949 review.allowed_tools = vec![CompactString::new("git_diff")];
950 let plain = SkillMetadata::new("plain", "No tools declared"); mgr.set_available_skills(vec![debug, review, plain]);
952
953 assert!(mgr.active_skill_tool_filter().is_none());
955
956 assert!(mgr.activate_skill("debug"));
958 assert!(!mgr.activate_skill("debug")); let f = mgr.active_skill_tool_filter().unwrap();
962 assert_eq!(f.len(), 2);
963 assert!(f.contains(&CompactString::new("read")) && f.contains(&CompactString::new("grep")));
964
965 mgr.activate_skill("review");
967 let f = mgr.active_skill_tool_filter().unwrap();
968 assert_eq!(f.len(), 3);
969 assert!(f.contains(&CompactString::new("git_diff")));
970
971 mgr.activate_skill("plain");
973 assert!(mgr.active_skill_tool_filter().is_none());
974 }
975
976 #[test]
977 fn update_collapse_mode_collapses_old_tool_results_under_pressure() {
978 let mut mgr = ContextManager::new(1_000);
979 for i in 0..10 {
980 let m = Message::tool(vec![ContentPart::ToolResult {
981 call_id: format!("c{i}").into(),
982 output: "x".repeat(40),
983 is_error: false,
984 }]);
985 mgr.push_history(m, 40);
986 }
987 mgr.set_observed_prompt_tokens(950); assert!(mgr.rho() >= mgr.config.collapse_threshold);
990
991 mgr.recompute_handle_residency();
992 assert_eq!(mgr.handles.residency_for_source("c0"), Some(&Residency::Collapsed));
994 assert_eq!(mgr.handles.residency_for_source("c9"), Some(&Residency::Resident));
995
996 mgr.set_observed_prompt_tokens(100); mgr.recompute_handle_residency();
1000 assert_eq!(
1001 mgr.handles.residency_for_source("c0"),
1002 Some(&Residency::Collapsed),
1003 "collapse is sticky until a compaction boundary"
1004 );
1005
1006 mgr.reset_collapse_generation();
1008 assert_eq!(mgr.handles.residency_for_source("c0"), Some(&Residency::Resident));
1009 }
1010
1011 #[test]
1012 fn frozen_prefix_len_anchors_at_compaction_and_holds_across_appends() {
1013 let mut mgr = ContextManager::new(1_000);
1014 for i in 0..30 {
1016 mgr.push_history(Message::user(format!("turn {i}: {}", "ctx ".repeat(30))), 150);
1017 }
1018 assert!(mgr.render().frozen_prefix_len.is_none(), "no frozen region before any compaction");
1019
1020 let (saved, _, archived, _) = mgr.compress(PressureAction::AutoCompact);
1021 assert!(saved > 0 && !archived.is_empty(), "expected archival");
1022
1023 assert!(mgr.render().frozen_prefix_len.is_none(), "deep == tail right after compaction");
1025
1026 mgr.push_history(Message::user("new 1"), 5);
1028 let f1 = mgr.render().frozen_prefix_len.expect("frozen region exists once the tail grows");
1029 mgr.push_history(Message::assistant("reply 1"), 5);
1030 mgr.push_history(Message::user("new 2"), 5);
1031 let rc = mgr.render();
1032 let f2 = rc.frozen_prefix_len.expect("frozen region holds");
1033 assert_eq!(f1, f2, "the deep boundary is fixed between compactions; only the tail grows");
1034 assert!(f2 < rc.turns.len(), "deep boundary is distinct from the rolling tail");
1035 }
1036
1037 #[test]
1038 fn frozen_boundary_holds_through_a_prefix_safe_compaction() {
1039 let mut mgr = ContextManager::new(10_000);
1042 for i in 0..5 {
1043 mgr.push_history(Message::user(format!("m{i}")), 5);
1044 }
1045 mgr.frozen_history_len = 3; let (_, _, _, cache_at) = mgr.compress(PressureAction::None);
1050 assert!(cache_at.is_none(), "no-op compaction is prefix-safe");
1051 assert_eq!(mgr.frozen_history_len, 3, "prefix-safe compaction preserves the deep-cache anchor");
1052 }
1053
1054 #[test]
1055 fn collapse_generation_resets_on_autocompact() {
1056 let mut mgr = ContextManager::new(1_000);
1057 for i in 0..20 {
1060 mgr.push_history(tool_result_msg(&format!("c{i}"), &"x".repeat(120)), 60);
1061 }
1062 mgr.set_observed_prompt_tokens(980); mgr.recompute_handle_residency();
1064 assert_eq!(mgr.handles.residency_for_source("c0"), Some(&Residency::Collapsed));
1065
1066 let (saved, _, archived, _) = mgr.compress(PressureAction::AutoCompact);
1067 assert!(saved > 0 && !archived.is_empty(), "expected archival");
1068
1069 for h in mgr.handles.all() {
1072 if matches!(h.kind, HandleKind::ToolResult) {
1073 assert_eq!(h.residency, Residency::Resident, "generation reset un-collapses survivors");
1074 }
1075 }
1076 }
1077
1078 #[test]
1079 fn mark_spooled_sets_residency_and_survives_residency_recompute() {
1080 let mut mgr = ContextManager::new(1_000);
1081 mgr.push_history(
1082 Message::tool(vec![ContentPart::ToolResult {
1083 call_id: "big".into(),
1084 output: "preview only".to_string(),
1085 is_error: false,
1086 }]),
1087 10,
1088 );
1089 mgr.mark_spooled("big", "disk://big");
1090 assert_eq!(
1091 mgr.handles.residency_for_source("big"),
1092 Some(&Residency::SpooledOut { r: "disk://big".to_string() })
1093 );
1094
1095 mgr.set_observed_prompt_tokens(990);
1098 mgr.recompute_handle_residency();
1099 assert_eq!(
1100 mgr.handles.residency_for_source("big"),
1101 Some(&Residency::SpooledOut { r: "disk://big".to_string() })
1102 );
1103 }
1104
1105 #[test]
1106 fn push_history_indexes_tool_results_as_resident_handles() {
1107 let mut mgr = ContextManager::new(10_000);
1108 let msg = Message::tool(vec![ContentPart::ToolResult {
1109 call_id: "call_1".into(),
1110 output: "the tool output".to_string(),
1111 is_error: false,
1112 }]);
1113 mgr.push_history(msg, 20);
1114 assert_eq!(mgr.handles.all().len(), 1);
1116 assert_eq!(
1117 mgr.handles.residency_for_source("call_1"),
1118 Some(&Residency::Resident)
1119 );
1120 mgr.push_history(Message::user("hello"), 5);
1122 assert_eq!(mgr.handles.all().len(), 1);
1123 }
1124
1125 fn tool_result_msg(call_id: &str, output: &str) -> Message {
1128 Message::tool(vec![ContentPart::ToolResult {
1129 call_id: call_id.into(),
1130 output: output.to_string(),
1131 is_error: false,
1132 }])
1133 }
1134
1135 #[test]
1136 fn prune_orphaned_handles_drops_handles_whose_message_left_history() {
1137 let mut mgr = ContextManager::new(10_000);
1138 mgr.push_history(tool_result_msg("c0", "out 0"), 20);
1139 mgr.push_history(tool_result_msg("c1", "out 1"), 20);
1140 assert_eq!(mgr.handles.all().len(), 2);
1141
1142 mgr.partitions.history.messages.remove(0);
1144 mgr.prune_orphaned_handles();
1145
1146 assert_eq!(mgr.handles.all().len(), 1);
1148 assert!(mgr.handles.residency_for_source("c0").is_none());
1149 assert_eq!(
1150 mgr.handles.residency_for_source("c1"),
1151 Some(&Residency::Resident)
1152 );
1153 }
1154
1155 #[test]
1156 fn autocompact_prunes_handles_for_archived_tool_results() {
1157 let mut mgr = ContextManager::new(1_000);
1158 for i in 0..30 {
1160 mgr.push_history(tool_result_msg(&format!("c{i}"), &"x".repeat(200)), 80);
1161 }
1162 assert_eq!(mgr.handles.all().len(), 30);
1163
1164 let (saved, _, archived, _) = mgr.compress(PressureAction::AutoCompact);
1165 assert!(saved > 0 && !archived.is_empty(), "expected archival");
1166
1167 let live_tool_results = mgr
1170 .partitions
1171 .history
1172 .messages
1173 .iter()
1174 .filter(|m| matches!(&m.content, Content::Parts(p)
1175 if p.iter().any(|x| matches!(x, ContentPart::ToolResult { .. }))))
1176 .count();
1177 assert_eq!(mgr.handles.all().len(), live_tool_results);
1178 assert!(mgr.handles.all().len() < 30, "table must shrink with archival");
1179 }
1180
1181 #[test]
1182 fn renew_prunes_handles_for_dropped_history() {
1183 let mut mgr = ContextManager::new(1_000);
1184 mgr.init_task("g".to_string(), vec![]);
1185 for i in 0..20 {
1186 mgr.push_history(tool_result_msg(&format!("c{i}"), "data"), 60);
1187 }
1188 mgr.renew();
1189 for h in mgr.handles.all() {
1191 if let Some(src) = h.source.as_ref() {
1192 assert!(
1193 mgr.handles.residency_for_source(src).is_some(),
1194 "no dangling handle survives renewal"
1195 );
1196 }
1197 }
1198 assert!(mgr.handles.all().len() <= 20);
1199 }
1200
1201 #[test]
1202 fn recompute_residency_index_semantics_with_spooled_in_the_middle() {
1203 let mut mgr = ContextManager::new(1_000);
1206 for i in 0..6 {
1207 mgr.push_history(tool_result_msg(&format!("c{i}"), &"y".repeat(40)), 40);
1208 }
1209 mgr.mark_spooled("c2", "disk://c2");
1210
1211 mgr.set_observed_prompt_tokens(950); mgr.recompute_handle_residency();
1213
1214 assert_eq!(
1216 mgr.handles.residency_for_source("c2"),
1217 Some(&Residency::SpooledOut { r: "disk://c2".to_string() })
1218 );
1219 assert_eq!(mgr.handles.residency_for_source("c0"), Some(&Residency::Collapsed));
1220 assert_eq!(mgr.handles.residency_for_source("c5"), Some(&Residency::Resident));
1221 }
1222
1223 #[test]
1226 fn knowledge_budget_marks_oldest_unpinned_first_and_warns_once() {
1227 let mut mgr = ContextManager::new(100);
1230 mgr.push_knowledge(Message::system("oldest unkeyed"), 10);
1231 mgr.push_knowledge_entry(Some("a".into()), Message::system("keyed"), 10, false);
1232 mgr.push_knowledge_entry(Some("p".into()), Message::system("pinned"), 10, true);
1233 mgr.push_knowledge_entry(Some("skill:x".into()), Message::system("skill"), 10, false);
1234
1235 let warn = mgr.enforce_knowledge_budget();
1236 assert_eq!(warn, Some((40, 25)));
1237 let e = &mgr.partitions.knowledge.entries;
1239 assert!(e[0].evict_at_boundary);
1240 assert!(e[1].evict_at_boundary);
1241 assert!(!e[2].evict_at_boundary, "pinned exempt");
1242 assert!(!e[3].evict_at_boundary, "skill pin exempt");
1243
1244 assert_eq!(mgr.enforce_knowledge_budget(), None);
1246
1247 let sweep = mgr.partitions.knowledge.sweep_at_boundary();
1249 assert_eq!(sweep.tokens_freed, 20);
1250 assert_eq!(mgr.partitions.knowledge.token_count, 20);
1251 assert_eq!(mgr.enforce_knowledge_budget(), None);
1253 }
1254
1255 #[test]
1256 fn knowledge_budget_warning_stands_when_only_exempt_weight_remains() {
1257 let mut mgr = ContextManager::new(100);
1258 mgr.push_knowledge_entry(Some("p".into()), Message::system("pinned heavy"), 30, true);
1259 mgr.push_knowledge_entry(Some("skill:x".into()), Message::system("skill heavy"), 30, false);
1260
1261 assert_eq!(mgr.enforce_knowledge_budget(), Some((60, 25)));
1263 assert!(mgr.partitions.knowledge.entries.iter().all(|e| !e.evict_at_boundary));
1264 }
1265
1266 #[test]
1267 fn knowledge_budget_ratio_zero_disables() {
1268 let mut mgr = ContextManager::new(100);
1269 mgr.config.knowledge_budget_ratio = 0.0;
1270 mgr.push_knowledge(Message::system("huge"), 90);
1271 assert_eq!(mgr.enforce_knowledge_budget(), None);
1272 assert!(!mgr.partitions.knowledge.entries[0].evict_at_boundary);
1273 }
1274}