1use zeph_context::budget::ContextBudget;
7use zeph_context::fidelity::FidelityScorer;
8use zeph_llm::LlmProvider;
9use zeph_llm::provider::{Message, MessagePart, Role};
10use zeph_skills::registry::SkillRegistry;
11
12use crate::error::ContextError;
13use crate::helpers::{
14 CODE_CONTEXT_PREFIX, CORRECTIONS_PREFIX, CROSS_SESSION_PREFIX, DOCUMENT_RAG_PREFIX,
15 GRAPH_FACTS_PREFIX, LSP_NOTE_PREFIX, PERSONA_PREFIX, REASONING_PREFIX, RECALL_PREFIX,
16 SESSION_DIGEST_PREFIX, SUMMARY_PREFIX, TRAJECTORY_PREFIX, TREE_MEMORY_PREFIX,
17};
18use crate::state::{
19 ContextAssemblyView, ContextDelta, ContextSummarizationView, MessageWindowView,
20 ProviderHandles, StatusSink,
21};
22
23pub struct SemanticRecallParams<'a> {
31 pub query: &'a str,
33 pub token_budget: usize,
35 pub recall_limit: usize,
37 pub context_format: zeph_config::ContextFormat,
39 pub conversation_id: Option<zeph_memory::ConversationId>,
41 pub tiered_classifier: Option<&'a std::sync::Arc<zeph_llm::any::AnyProvider>>,
43 pub tiered_validator: Option<&'a std::sync::Arc<zeph_llm::any::AnyProvider>>,
45 pub tiered_config: &'a zeph_config::memory::TieredRetrievalConfig,
47}
48
49#[derive(Debug, Default)]
67pub struct ContextService;
68
69impl ContextService {
70 #[must_use]
74 pub fn new() -> Self {
75 Self
76 }
77
78 pub fn clear_history(&self, window: &mut MessageWindowView<'_>) {
86 let system_prompt = window.messages.first().cloned();
87 window.messages.clear();
88 if let Some(sp) = system_prompt {
89 window.messages.push(sp);
90 }
91 window.completed_tool_ids.clear();
92 recompute_prompt_tokens(window);
93 }
94
95 pub fn remove_recall_messages(&self, window: &mut MessageWindowView<'_>) {
97 remove_by_part_or_prefix(window.messages, RECALL_PREFIX, |p| {
98 matches!(p, MessagePart::Recall { .. })
99 });
100 }
101
102 pub fn remove_correction_messages(&self, window: &mut MessageWindowView<'_>) {
104 remove_by_prefix(window.messages, Role::System, CORRECTIONS_PREFIX);
105 }
106
107 pub fn remove_graph_facts_messages(&self, window: &mut MessageWindowView<'_>) {
109 remove_by_prefix(window.messages, Role::System, GRAPH_FACTS_PREFIX);
110 }
111
112 pub fn remove_persona_facts_messages(&self, window: &mut MessageWindowView<'_>) {
114 remove_by_prefix(window.messages, Role::System, PERSONA_PREFIX);
115 }
116
117 pub fn remove_trajectory_hints_messages(&self, window: &mut MessageWindowView<'_>) {
119 remove_by_prefix(window.messages, Role::System, TRAJECTORY_PREFIX);
120 }
121
122 pub fn remove_tree_memory_messages(&self, window: &mut MessageWindowView<'_>) {
124 remove_by_prefix(window.messages, Role::System, TREE_MEMORY_PREFIX);
125 }
126
127 pub fn remove_reasoning_strategies_messages(&self, window: &mut MessageWindowView<'_>) {
129 remove_by_prefix(window.messages, Role::System, REASONING_PREFIX);
130 }
131
132 pub fn remove_lsp_messages(&self, window: &mut MessageWindowView<'_>) {
137 remove_by_prefix(window.messages, Role::System, LSP_NOTE_PREFIX);
138 }
139
140 pub fn remove_code_context_messages(&self, window: &mut MessageWindowView<'_>) {
142 remove_by_part_or_prefix(window.messages, CODE_CONTEXT_PREFIX, |p| {
143 matches!(p, MessagePart::CodeContext { .. })
144 });
145 }
146
147 pub fn remove_summary_messages(&self, window: &mut MessageWindowView<'_>) {
149 remove_by_part_or_prefix(window.messages, SUMMARY_PREFIX, |p| {
150 matches!(p, MessagePart::Summary { .. })
151 });
152 }
153
154 pub fn remove_cross_session_messages(&self, window: &mut MessageWindowView<'_>) {
156 remove_by_part_or_prefix(window.messages, CROSS_SESSION_PREFIX, |p| {
157 matches!(p, MessagePart::CrossSession { .. })
158 });
159 }
160
161 pub fn remove_session_digest_message(&self, window: &mut MessageWindowView<'_>) {
163 remove_by_prefix(window.messages, Role::User, SESSION_DIGEST_PREFIX);
164 }
165
166 pub fn remove_document_rag_messages(&self, window: &mut MessageWindowView<'_>) {
168 remove_by_prefix(window.messages, Role::System, DOCUMENT_RAG_PREFIX);
169 }
170
171 pub fn trim_messages_to_budget(&self, window: &mut MessageWindowView<'_>, token_budget: usize) {
179 if token_budget == 0 {
180 return;
181 }
182
183 let history_start = window
185 .messages
186 .iter()
187 .position(|m| m.role != Role::System)
188 .unwrap_or(window.messages.len());
189
190 if history_start >= window.messages.len() {
191 return;
192 }
193
194 let mut total = 0usize;
195 let mut keep_from = window.messages.len();
196
197 for i in (history_start..window.messages.len()).rev() {
198 let msg_tokens = window
199 .token_counter
200 .count_message_tokens(&window.messages[i]);
201 if total + msg_tokens > token_budget {
202 break;
203 }
204 total += msg_tokens;
205 keep_from = i;
206 }
207
208 if keep_from > history_start {
209 let removed = keep_from - history_start;
210 window.messages.drain(history_start..keep_from);
211 recompute_prompt_tokens(window);
212 tracing::info!(
213 removed,
214 token_budget,
215 "trimmed messages to fit context budget"
216 );
217 }
218 }
219
220 #[tracing::instrument(name = "agent_context.service.inject_semantic_recall", skip_all, err)]
232 pub async fn inject_semantic_recall(
233 &self,
234 query: &str,
235 token_budget: usize,
236 window: &mut MessageWindowView<'_>,
237 view: &ContextAssemblyView<'_>,
238 ) -> Result<(), ContextError> {
239 self.remove_recall_messages(window);
240
241 let params = SemanticRecallParams {
242 query,
243 token_budget,
244 recall_limit: view.recall_limit,
245 context_format: view.context_format,
246 conversation_id: view.conversation_id,
247 tiered_classifier: view.tiered_retrieval_classifier.as_ref(),
248 tiered_validator: view.tiered_retrieval_validator.as_ref(),
249 tiered_config: &view.tiered_retrieval_config,
250 };
251 let msg = self
252 .run_tiered_recall(¶ms, window, view.memory.as_deref())
253 .await?;
254
255 if let Some(msg) = msg
256 && window.messages.len() > 1
257 {
258 window.messages.insert(1, msg);
259 }
260
261 Ok(())
262 }
263
264 #[tracing::instrument(
276 name = "agent_context.service.inject_semantic_recall_bare",
277 skip_all,
278 err
279 )]
280 pub async fn inject_semantic_recall_bare(
281 &self,
282 params: SemanticRecallParams<'_>,
283 window: &mut MessageWindowView<'_>,
284 memory: Option<&zeph_memory::semantic::SemanticMemory>,
285 ) -> Result<(), ContextError> {
286 self.remove_recall_messages(window);
287
288 let msg = self.run_tiered_recall(¶ms, window, memory).await?;
289
290 if let Some(msg) = msg
291 && window.messages.len() > 1
292 {
293 window.messages.insert(1, msg);
294 }
295
296 Ok(())
297 }
298
299 #[tracing::instrument(name = "agent_context.service.run_tiered_recall", skip_all, err)]
304 async fn run_tiered_recall(
305 &self,
306 params: &SemanticRecallParams<'_>,
307 window: &MessageWindowView<'_>,
308 memory: Option<&zeph_memory::semantic::SemanticMemory>,
309 ) -> Result<Option<Message>, ContextError> {
310 if params.tiered_config.enabled {
311 use tracing::Instrument as _;
312 let Some(mem) = memory else {
313 return Ok(None);
314 };
315 let result = tokio::time::timeout(
316 std::time::Duration::from_secs(30),
317 zeph_memory::recall_tiered(
318 mem,
319 params.query,
320 params.conversation_id,
321 params.tiered_classifier,
322 params.tiered_validator,
323 params.tiered_config,
324 Some(params.token_budget),
325 )
326 .instrument(tracing::info_span!("agent_context.tiered_retrieval.recall")),
327 )
328 .await
329 .map_err(|_| {
330 tracing::warn!("tiered_retrieval: recall_tiered timed out after 30s");
331 ContextError::Memory(zeph_memory::MemoryError::Timeout(
332 "recall_tiered timed out".to_owned(),
333 ))
334 })?
335 .map_err(ContextError::Memory)?;
336
337 tracing::debug!(
338 intent = %result.intent,
339 tokens_used = result.tokens_used,
340 tier_escalated = result.tier_escalated,
341 count = result.messages.len(),
342 "tiered_retrieval: recall complete"
343 );
344
345 if result.messages.is_empty() {
346 return Ok(None);
347 }
348
349 let recalled_text = result
350 .messages
351 .iter()
352 .map(|m| m.message.content.as_str())
353 .collect::<Vec<_>>()
354 .join("\n---\n");
355 Ok(Some(Message::from_legacy(
356 Role::User,
357 format!("{RECALL_PREFIX}{recalled_text}"),
358 )))
359 } else {
360 let (msg, _score) = crate::helpers::fetch_semantic_recall_raw(
361 memory,
362 crate::helpers::SemanticRecallRawParams {
363 recall_limit: params.recall_limit,
364 context_format: params.context_format,
365 query: params.query,
366 token_budget: params.token_budget,
367 tc: &window.token_counter,
368 low_confidence_threshold: None,
369 },
370 None,
371 )
372 .await?;
373 Ok(msg)
374 }
375 }
376
377 #[tracing::instrument(
386 name = "agent_context.service.inject_cross_session_context",
387 skip_all,
388 err
389 )]
390 pub async fn inject_cross_session_context(
391 &self,
392 query: &str,
393 token_budget: usize,
394 window: &mut MessageWindowView<'_>,
395 view: &ContextAssemblyView<'_>,
396 ) -> Result<(), ContextError> {
397 self.remove_cross_session_messages(window);
398
399 if let Some(msg) = crate::helpers::fetch_cross_session_raw(
400 view.memory.as_deref(),
401 view.conversation_id,
402 view.cross_session_score_threshold,
403 query,
404 token_budget,
405 &view.token_counter,
406 )
407 .await?
408 && window.messages.len() > 1
409 {
410 window.messages.insert(1, msg);
411 tracing::debug!("injected cross-session context");
412 }
413
414 Ok(())
415 }
416
417 #[tracing::instrument(name = "agent_context.service.inject_summaries", skip_all, err)]
426 pub async fn inject_summaries(
427 &self,
428 token_budget: usize,
429 window: &mut MessageWindowView<'_>,
430 view: &ContextAssemblyView<'_>,
431 ) -> Result<(), ContextError> {
432 self.remove_summary_messages(window);
433
434 if let Some(msg) = crate::helpers::fetch_summaries_raw(
435 view.memory.as_deref(),
436 view.conversation_id,
437 token_budget,
438 &view.token_counter,
439 )
440 .await?
441 && window.messages.len() > 1
442 {
443 window.messages.insert(1, msg);
444 tracing::debug!("injected summaries into context");
445 }
446
447 Ok(())
448 }
449
450 #[tracing::instrument(name = "agent_context.service.disambiguate_skills", skip_all)]
455 pub async fn disambiguate_skills(
456 &self,
457 query: &str,
458 all_meta: &[&zeph_skills::loader::SkillMeta],
459 scored: &[zeph_skills::ScoredMatch],
460 providers: &ProviderHandles,
461 ) -> Option<Vec<usize>> {
462 use std::fmt::Write as _;
463
464 let mut candidates = String::new();
465 for sm in scored {
466 if let Some(meta) = all_meta.get(sm.index) {
467 let _ = writeln!(
468 candidates,
469 "- {} (score: {:.3}): {}",
470 meta.name, sm.score, meta.description
471 );
472 }
473 }
474
475 let prompt = format!(
476 "The user said: \"{query}\"\n\n\
477 These skills matched with similar scores:\n{candidates}\n\
478 Which skill best matches the user's intent? \
479 Return the skill_name, your confidence (0-1), and any extracted parameters."
480 );
481
482 let messages = vec![zeph_llm::provider::Message::from_legacy(
483 zeph_llm::provider::Role::User,
484 prompt,
485 )];
486 match providers
487 .disambiguate
488 .chat_typed::<zeph_skills::IntentClassification>(&messages)
489 .await
490 {
491 Ok(classification) => {
492 tracing::info!(
493 skill = %classification.skill_name,
494 confidence = classification.confidence,
495 "disambiguation selected skill"
496 );
497 let mut indices: Vec<usize> = scored.iter().map(|s| s.index).collect();
498 if let Some(pos) = indices.iter().position(|&i| {
499 all_meta
500 .get(i)
501 .is_some_and(|m| m.name == classification.skill_name)
502 }) {
503 indices.swap(0, pos);
504 }
505 Some(indices)
506 }
507 Err(e) => {
508 tracing::warn!("disambiguation failed, using original order: {e:#}");
509 None
510 }
511 }
512 }
513
514 #[allow(clippy::too_many_lines)] #[tracing::instrument(name = "agent_context.service.prepare_context", skip_all, err)]
527 pub async fn prepare_context(
528 &self,
529 query: &str,
530 window: &mut MessageWindowView<'_>,
531 view: &mut ContextAssemblyView<'_>,
532 ) -> Result<ContextDelta, ContextError> {
533 if view.context_manager.budget.is_none() {
534 return Ok(ContextDelta::default());
535 }
536
537 self.remove_session_digest_message(window);
539 self.remove_summary_messages(window);
540 self.remove_cross_session_messages(window);
541 self.remove_recall_messages(window);
542 self.remove_document_rag_messages(window);
543 self.remove_correction_messages(window);
544 self.remove_code_context_messages(window);
545 self.remove_graph_facts_messages(window);
546 self.remove_persona_facts_messages(window);
547 self.remove_trajectory_hints_messages(window);
548 self.remove_tree_memory_messages(window);
549 if view.reasoning_config.enabled {
550 self.remove_reasoning_strategies_messages(window);
551 }
552
553 if let Some(explorer) = view.proactive_explorer.clone()
555 && let Some(domain) = explorer.classify(query)
556 {
557 let already_known = {
558 let registry_guard = view.skill_registry.read();
559 explorer.has_knowledge(®istry_guard, &domain)
560 };
561 let excluded = explorer.is_excluded(&domain);
562
563 if !already_known && !excluded {
564 tracing::debug!(domain = %domain.0, query_len = query.len(), "proactive.explore triggered");
565 let timeout_ms = explorer.timeout_ms();
566 let result = tokio::time::timeout(
567 std::time::Duration::from_millis(timeout_ms),
568 explorer.explore(&domain),
569 )
570 .await;
571 match result {
572 Ok(Ok(())) => {
573 let hub_dirs = view.skill_registry.read().hub_dirs().to_vec();
574 let reload_paths = view.skill_paths.to_vec();
575 let span =
576 tracing::info_span!("agent_context.skills.registry.reload_blocking");
577 match tokio::task::spawn_blocking(move || {
578 let _enter = span.enter();
579 SkillRegistry::load(&reload_paths).with_hub_dirs(hub_dirs)
580 })
581 .await
582 {
583 Ok(new_registry) => {
584 *view.skill_registry.write() = new_registry;
585 tracing::debug!(domain = %domain.0, "proactive.explore complete, registry reloaded");
586 }
587 Err(e) => {
588 tracing::error!(
589 domain = %domain.0,
590 "proactive.explore: skill registry reload panicked, registry left unchanged: {e}"
591 );
592 }
593 }
594 }
595 Ok(Err(e)) => {
596 tracing::warn!(domain = %domain.0, error = %e, "proactive exploration failed");
597 }
598 Err(_) => {
599 tracing::warn!(domain = %domain.0, timeout_ms, "proactive exploration timed out");
600 }
601 }
602 }
603 }
604
605 let active_levels: &'static [zeph_memory::compression::CompressionLevel] =
607 if let Some(ref budget) = view.context_manager.budget {
608 let used = view.cached_prompt_tokens;
609 let max = budget.max_tokens();
610 #[allow(clippy::cast_precision_loss)]
611 let remaining_ratio = if max == 0 {
612 1.0_f32
613 } else {
614 1.0 - (used as f32 / max as f32).clamp(0.0, 1.0)
615 };
616 let levels =
617 zeph_memory::compression::RetrievalPolicy::default().select(remaining_ratio);
618 tracing::debug!(
619 remaining_ratio,
620 active_levels = ?levels,
621 "compression_spectrum: retrieval policy selected"
622 );
623 levels
624 } else {
625 &[]
626 };
627
628 let memory_backend: Option<std::sync::Arc<dyn zeph_common::memory::ContextMemoryBackend>> =
629 view.memory.clone().map(
630 |m| -> std::sync::Arc<dyn zeph_common::memory::ContextMemoryBackend> {
631 std::sync::Arc::new(crate::memory_backend::SemanticMemoryBackend::new(m))
632 },
633 );
634
635 let memory_view = zeph_context::input::ContextMemoryView {
636 memory: memory_backend,
637 conversation_id: view.conversation_id.map(|c| c.0),
638 recall_limit: view.recall_limit,
639 cross_session_score_threshold: view.cross_session_score_threshold,
640 context_strategy: view.context_strategy,
641 crossover_turn_threshold: view.crossover_turn_threshold,
642 cached_session_digest: view.cached_session_digest.clone(),
643 graph_config: view.graph_config.clone(),
644 document_config: view.document_config.clone(),
645 persona_config: view.persona_config.clone(),
646 trajectory_config: view.trajectory_config.clone(),
647 reasoning_config: view.reasoning_config.clone(),
648 memcot_config: view.memcot_config.clone(),
649 memcot_state: view.memcot_state.clone(),
650 tree_config: view.tree_config.clone(),
651 };
652
653 #[cfg(feature = "index")]
654 let index_access = view.index;
655 #[cfg(not(feature = "index"))]
656 let index_access: Option<&dyn zeph_context::input::IndexAccess> = None;
657
658 let router = crate::memory_backend::build_memory_router(view.context_manager);
659
660 let active_types = crate::type_aware_compose::resolve_active_functional_types(
664 &view.type_aware_compose_config,
665 query,
666 );
667
668 let input = zeph_context::input::ContextAssemblyInput {
669 memory: &memory_view,
670 context_manager: view.context_manager,
671 token_counter: &*view.token_counter,
672 skills_prompt: view.last_skills_prompt,
673 index: index_access,
674 correction_config: view.correction_config,
675 sidequest_turn_counter: view.sidequest_turn_counter,
676 messages: window.messages,
677 query,
678 scrub: view.scrub,
679 active_levels,
680 active_types: &active_types,
681 router,
682 planned_next_tools: view.planned_next_tools,
683 };
684
685 let mut prepared = zeph_context::assembler::ContextAssembler::gather(&input).await?;
686
687 if view.tiered_retrieval_config.enabled {
692 prepared.recall = None;
693 }
694
695 for handle in prepared.background_tasks.drain(..) {
699 let task_supervisor = std::sync::Arc::clone(&view.task_supervisor);
700 drop(task_supervisor.spawn_oneshot(
701 std::sync::Arc::from("context.assembly.background"),
702 move || async move {
703 let _ = handle.await;
704 },
705 ));
706 }
707
708 let (delta, inserted_count) = self.apply_prepared_context(window, view, prepared).await;
709
710 if view.tiered_retrieval_config.enabled {
711 self.inject_semantic_recall(query, usize::MAX, window, view)
712 .await?;
713 }
714
715 let memory_first_active =
720 view.context_strategy == zeph_config::ContextStrategy::MemoryFirst;
721 if let Some(fidelity_cfg) = view.fidelity_config
722 && fidelity_cfg.enabled
723 && !memory_first_active
724 {
725 use tracing::Instrument as _;
726 if let Some(ref tx) = view.status_tx {
727 let _ = tx.send("Scoring context fidelity\u{2026}".into());
728 }
729 let (embed_provider, compress_provider) = fidelity_provider_pair(
730 view.fidelity_semantic_provider.as_ref(),
731 view.fidelity_compress_provider.as_ref(),
732 );
733 let fidelity_span = tracing::info_span!(
734 "context.fidelity.score",
735 message_count = window.messages.len(),
736 query_len = query.len(),
737 );
738 FidelityScorer
739 .score_and_apply(
740 window.messages,
741 query,
742 view.planned_next_tools,
743 fidelity_cfg,
744 &*view.token_counter,
745 inserted_count,
746 false, embed_provider,
748 compress_provider,
749 )
750 .instrument(fidelity_span)
751 .await;
752 persist_fidelity_tags(window.messages, view.memory.as_deref()).await;
754 recompute_prompt_tokens(window);
755 if let Some(ref tx) = view.status_tx {
756 let _ = tx.send(String::new());
757 }
758 }
759
760 Ok(delta)
761 }
762
763 #[allow(clippy::too_many_lines)] #[tracing::instrument(name = "agent_context.service.apply_prepared_context", skip_all)]
773 async fn apply_prepared_context(
774 &self,
775 window: &mut MessageWindowView<'_>,
776 view: &mut ContextAssemblyView<'_>,
777 prepared: zeph_context::assembler::PreparedContext,
778 ) -> (ContextDelta, usize) {
779 use std::borrow::Cow;
780 use zeph_llm::provider::{Message, MessageMetadata, Role};
781 use zeph_sanitizer::{ContentSource, ContentSourceKind, MemorySourceHint};
782
783 *view.last_recall_confidence = prepared.recall_confidence;
785
786 if prepared.memory_first {
788 let history_start = 1usize;
789 let len = window.messages.len();
790 let keep_tail =
791 zeph_context::assembler::memory_first_keep_tail(window.messages, history_start);
792 if len > history_start + keep_tail {
793 window.messages.drain(history_start..len - keep_tail);
794 recompute_prompt_tokens(window);
795 tracing::debug!(
796 strategy = "memory_first",
797 keep_tail,
798 "dropped conversation history, kept last {keep_tail} messages"
799 );
800 }
801 }
802
803 let mut inserted_count: usize = 0;
806
807 let slots: &[(Option<Message>, MemorySourceHint, Option<&str>)] = &[
810 (
811 prepared.graph_facts,
812 MemorySourceHint::ExternalContent,
813 Some("injected knowledge graph facts into context"),
814 ),
815 (
816 prepared.doc_rag,
817 MemorySourceHint::ExternalContent,
818 Some("injected document RAG context"),
819 ),
820 (
821 prepared.corrections,
822 MemorySourceHint::ConversationHistory,
823 Some("injected past corrections into context"),
824 ),
825 (prepared.recall, MemorySourceHint::ConversationHistory, None),
826 (prepared.cross_session, MemorySourceHint::LlmSummary, None),
827 (
828 prepared.summaries,
829 MemorySourceHint::LlmSummary,
830 Some("injected summaries into context"),
831 ),
832 (
833 prepared.persona_facts,
834 MemorySourceHint::ExternalContent,
835 Some("injected persona facts into context"),
836 ),
837 (
838 prepared.trajectory_hints,
839 MemorySourceHint::ExternalContent,
840 Some("injected trajectory hints into context"),
841 ),
842 (
843 prepared.tree_memory,
844 MemorySourceHint::ExternalContent,
845 Some("injected tree memory summary into context"),
846 ),
847 (
848 prepared.reasoning_hints,
849 MemorySourceHint::ExternalContent,
850 Some("injected reasoning strategies into context"),
851 ),
852 ];
853 for (opt_msg, hint, label) in slots.iter().cloned() {
854 if let Some(msg) = opt_msg.filter(|_| window.messages.len() > 1) {
855 let sanitized = self.sanitize_memory_message(msg, hint, view).await;
856 window.messages.insert(1, sanitized);
857 inserted_count += 1;
858 if let Some(lbl) = label {
859 tracing::debug!("{lbl}");
860 }
861 }
862 }
863
864 let code_context = if let Some(text) = prepared.code_context {
866 let sanitized = view
867 .sanitizer
868 .sanitize(&text, ContentSource::new(ContentSourceKind::ToolResult));
869 view.metrics.sanitizer_runs += 1;
870 if !sanitized.injection_flags.is_empty() {
871 tracing::warn!(
872 flags = sanitized.injection_flags.len(),
873 "injection patterns detected in code RAG context"
874 );
875 view.metrics.sanitizer_injection_flags += sanitized.injection_flags.len() as u64;
876 let detail = sanitized
877 .injection_flags
878 .first()
879 .map_or_else(String::new, |f| {
880 format!("Detected pattern: {}", f.pattern_name)
881 });
882 view.security_events.push(
883 zeph_common::SecurityEventCategory::InjectionFlag,
884 "code_rag",
885 detail,
886 );
887 }
888 if sanitized.was_truncated {
889 view.metrics.sanitizer_truncations += 1;
890 view.security_events.push(
891 zeph_common::SecurityEventCategory::Truncation,
892 "code_rag",
893 "Content truncated to max_content_size".to_string(),
894 );
895 }
896 Some(sanitized.body)
897 } else {
898 None
899 };
900
901 if !prepared.memory_first {
902 self.trim_messages_to_budget(window, prepared.recent_history_budget);
903 }
904
905 if view.digest_enabled
907 && let Some((digest_text, _)) = view
908 .cached_session_digest
909 .clone()
910 .filter(|_| window.messages.len() > 1)
911 {
912 let digest_msg = Message {
913 role: Role::User,
914 content: format!("{}{digest_text}", crate::helpers::SESSION_DIGEST_PREFIX),
915 parts: vec![],
916 metadata: MessageMetadata::default(),
917 };
918 let sanitized = self
919 .sanitize_memory_message(digest_msg, MemorySourceHint::LlmSummary, view)
920 .await;
921 window.messages.insert(1, sanitized);
922 inserted_count += 1;
923 tracing::debug!("injected session digest into context");
924 }
925
926 if view.redact_credentials {
928 for msg in &mut *window.messages {
929 if msg.role == Role::System {
930 continue;
931 }
932 if let Cow::Owned(s) = (view.scrub)(&msg.content) {
933 msg.content = s;
934 }
935 }
936 }
937
938 recompute_prompt_tokens(window);
939
940 (ContextDelta { code_context }, inserted_count)
941 }
942
943 #[tracing::instrument(name = "agent_context.service.sanitize_memory_message", skip_all)]
953 async fn sanitize_memory_message(
954 &self,
955 mut msg: zeph_llm::provider::Message,
956 hint: zeph_sanitizer::MemorySourceHint,
957 view: &mut ContextAssemblyView<'_>,
958 ) -> zeph_llm::provider::Message {
959 use zeph_sanitizer::{ContentSource, ContentSourceKind};
960
961 let source = ContentSource::new(ContentSourceKind::MemoryRetrieval).with_memory_hint(hint);
962 let sanitized = view.sanitizer.sanitize(&msg.content, source);
963 view.metrics.sanitizer_runs += 1;
964 if !sanitized.injection_flags.is_empty() {
965 tracing::warn!(
966 flags = sanitized.injection_flags.len(),
967 "injection patterns detected in memory retrieval"
968 );
969 view.metrics.sanitizer_injection_flags += sanitized.injection_flags.len() as u64;
970 let detail = sanitized
971 .injection_flags
972 .first()
973 .map_or_else(String::new, |f| {
974 format!("Detected pattern: {}", f.pattern_name)
975 });
976 view.security_events.push(
977 zeph_common::SecurityEventCategory::InjectionFlag,
978 "memory_retrieval",
979 detail,
980 );
981 }
982 if sanitized.was_truncated {
983 view.metrics.sanitizer_truncations += 1;
984 view.security_events.push(
985 zeph_common::SecurityEventCategory::Truncation,
986 "memory_retrieval",
987 "Content truncated to max_content_size".to_string(),
988 );
989 }
990
991 if view.sanitizer.is_enabled()
993 && let Some(qs) = view.quarantine_summarizer
994 && qs.should_quarantine(ContentSourceKind::MemoryRetrieval)
995 {
996 match qs.extract_facts(&sanitized, view.sanitizer).await {
997 Ok((facts, flags)) => {
998 view.metrics.quarantine_invocations += 1;
999 view.security_events.push(
1000 zeph_common::SecurityEventCategory::Quarantine,
1001 "memory_retrieval",
1002 "Content quarantined, facts extracted".to_string(),
1003 );
1004 let escaped = zeph_sanitizer::ContentSanitizer::escape_delimiter_tags(&facts);
1005 msg.content = zeph_sanitizer::ContentSanitizer::apply_spotlight(
1006 &escaped,
1007 &sanitized.source,
1008 &flags,
1009 );
1010 return msg;
1011 }
1012 Err(e) => {
1013 tracing::warn!(
1014 error = %e,
1015 "quarantine failed for memory retrieval, using original sanitized content"
1016 );
1017 view.metrics.quarantine_failures += 1;
1018 view.security_events.push(
1019 zeph_common::SecurityEventCategory::Quarantine,
1020 "memory_retrieval",
1021 format!("Quarantine failed: {e}"),
1022 );
1023 }
1024 }
1025 }
1026
1027 msg.content = sanitized.body;
1028 msg
1029 }
1030
1031 pub fn reset_conversation(
1041 &self,
1042 window: &mut MessageWindowView<'_>,
1043 _view: &mut ContextAssemblyView<'_>,
1044 ) -> Result<(), ContextError> {
1045 self.clear_history(window);
1046 Ok(())
1047 }
1048
1049 #[allow(
1066 clippy::cast_precision_loss,
1067 clippy::cast_possible_truncation,
1068 clippy::cast_sign_loss,
1069 clippy::too_many_lines
1070 )]
1071 #[tracing::instrument(name = "agent_context.service.maybe_compact", skip_all, err)]
1072 pub async fn maybe_compact(
1073 &self,
1074 summ: &mut ContextSummarizationView<'_>,
1075 status: &(impl StatusSink + ?Sized),
1076 ) -> Result<(), ContextError> {
1077 use zeph_context::manager::{CompactionState, CompactionTier};
1078
1079 if let Some(count) = summ.context_manager.turns_since_last_hard_compaction_mut() {
1081 *count += 1;
1082 }
1083
1084 if let CompactionState::Exhausted { warned } = summ.context_manager.compaction_state()
1086 && !warned
1087 {
1088 summ.context_manager
1089 .set_compaction_state(CompactionState::Exhausted { warned: true });
1090 tracing::warn!("compaction exhausted: context budget too tight for this session");
1091 }
1092 if summ.context_manager.compaction_state().is_exhausted() {
1093 return Ok(());
1094 }
1095
1096 if summ.server_compaction_active {
1098 let budget = summ
1099 .context_manager
1100 .budget
1101 .as_ref()
1102 .map_or(0, ContextBudget::max_tokens);
1103 if budget > 0 {
1104 let fallback = (budget * 95 / 100) as u64;
1105 if *summ.cached_prompt_tokens < fallback {
1106 return Ok(());
1107 }
1108 tracing::warn!(
1109 "server compaction active but context at 95%+ — falling back to client-side"
1110 );
1111 } else {
1112 return Ok(());
1113 }
1114 }
1115
1116 if summ
1118 .context_manager
1119 .compaction_state()
1120 .is_compacted_this_turn()
1121 {
1122 return Ok(());
1123 }
1124
1125 let in_cooldown = summ.context_manager.compaction_state().cooldown_remaining() > 0;
1127 if in_cooldown
1128 && let CompactionState::Cooling { turns_remaining } =
1129 summ.context_manager.compaction_state()
1130 {
1131 let next = turns_remaining - 1;
1132 summ.context_manager.set_compaction_state(if next == 0 {
1133 CompactionState::Ready
1134 } else {
1135 CompactionState::Cooling {
1136 turns_remaining: next,
1137 }
1138 });
1139 }
1140
1141 if let Some(ref fidelity_cfg) = summ.fidelity_config.clone()
1145 && fidelity_cfg.enabled
1146 && summ.context_manager.should_proactively_regrade(
1147 *summ.cached_prompt_tokens,
1148 fidelity_cfg.regrade_threshold,
1149 summ.server_compaction_active,
1150 )
1151 {
1152 use tracing::Instrument as _;
1153 let (regrade_embed_provider, regrade_compress_provider) = fidelity_provider_pair(
1154 summ.fidelity_semantic_provider.as_ref(),
1155 summ.fidelity_compress_provider.as_ref(),
1156 );
1157 FidelityScorer
1158 .score_and_apply(
1159 summ.messages,
1160 &summ.current_query,
1161 &[],
1162 fidelity_cfg,
1163 &*summ.token_counter,
1164 0,
1165 true, regrade_embed_provider,
1167 regrade_compress_provider,
1168 )
1169 .instrument(tracing::info_span!(
1170 "context.fidelity.regrade",
1171 budget_ratio = tracing::field::Empty,
1172 ))
1173 .await;
1174 persist_fidelity_tags(summ.messages, summ.memory.as_deref()).await;
1176 recompute_prompt_tokens_summ(summ);
1177 summ.context_manager.set_regraded_this_turn(true);
1178 tracing::debug!(
1179 cached_tokens = *summ.cached_prompt_tokens,
1180 "AgeMem proactive regrade complete"
1181 );
1182 }
1183
1184 match summ
1185 .context_manager
1186 .compaction_tier(*summ.cached_prompt_tokens)
1187 {
1188 CompactionTier::Soft => {
1189 self.do_soft_compaction(summ, status).await;
1190 Ok(())
1191 }
1192 CompactionTier::Hard => self.do_hard_compaction(summ, status, in_cooldown).await,
1193 _ => Ok(()),
1194 }
1195 }
1196
1197 #[tracing::instrument(name = "agent_context.service.do_soft_compaction", skip_all)]
1202 #[allow(
1203 clippy::cast_precision_loss,
1204 clippy::cast_possible_truncation,
1205 clippy::cast_sign_loss
1206 )]
1207 async fn do_soft_compaction(
1208 &self,
1209 summ: &mut ContextSummarizationView<'_>,
1210 status: &(impl StatusSink + ?Sized),
1211 ) {
1212 status.send_status("soft compacting context...").await;
1213
1214 match &summ.context_manager.compression.pruning_strategy {
1216 zeph_config::PruningStrategy::Subgoal | zeph_config::PruningStrategy::SubgoalMig => {
1217 crate::summarization::scheduling::maybe_refresh_subgoal(summ);
1218 }
1219 _ => crate::summarization::scheduling::maybe_refresh_task_goal(summ),
1220 }
1221
1222 let applied = crate::summarization::deferred::apply_deferred_summaries(summ);
1224
1225 if applied > 0
1227 && summ
1228 .context_manager
1229 .compression
1230 .pruning_strategy
1231 .is_subgoal()
1232 {
1233 summ.subgoal_registry
1234 .rebuild_after_compaction(summ.messages, 0);
1235 }
1236
1237 let budget = summ
1239 .context_manager
1240 .budget
1241 .as_ref()
1242 .map_or(0, ContextBudget::max_tokens);
1243 let soft_threshold =
1244 (budget as f32 * summ.context_manager.soft_compaction_threshold) as usize;
1245 let cached = usize::try_from(*summ.cached_prompt_tokens).unwrap_or(usize::MAX);
1246 let min_to_free = cached.saturating_sub(soft_threshold);
1247 if min_to_free > 0 {
1248 crate::summarization::pruning::prune_tool_outputs(summ, min_to_free);
1249 }
1250
1251 status.send_status("").await;
1252 tracing::info!(
1253 cached_tokens = *summ.cached_prompt_tokens,
1254 soft_threshold,
1255 "soft compaction complete"
1256 );
1257 }
1258
1259 #[tracing::instrument(name = "agent_context.service.do_hard_compaction", skip_all, err)]
1261 #[allow(
1262 clippy::cast_precision_loss,
1263 clippy::cast_possible_truncation,
1264 clippy::cast_sign_loss
1265 )]
1266 async fn do_hard_compaction(
1267 &self,
1268 summ: &mut ContextSummarizationView<'_>,
1269 status: &(impl StatusSink + ?Sized),
1270 in_cooldown: bool,
1271 ) -> Result<(), ContextError> {
1272 use zeph_context::manager::CompactionState;
1273
1274 let turns_since_last = summ
1276 .context_manager
1277 .turns_since_last_hard_compaction()
1278 .map(|t| u32::try_from(t).unwrap_or(u32::MAX));
1279 summ.context_manager
1280 .set_turns_since_last_hard_compaction(Some(0));
1281 if let Some(metrics) = summ.metrics {
1282 metrics.record_hard_compaction(turns_since_last);
1283 }
1284
1285 if in_cooldown {
1286 tracing::debug!(
1287 turns_remaining = summ.context_manager.compaction_state().cooldown_remaining(),
1288 "hard compaction skipped: cooldown active"
1289 );
1290 return Ok(());
1291 }
1292
1293 let budget = summ
1294 .context_manager
1295 .budget
1296 .as_ref()
1297 .map_or(0, ContextBudget::max_tokens);
1298 let hard_threshold =
1299 (budget as f32 * summ.context_manager.hard_compaction_threshold) as usize;
1300 let cached = usize::try_from(*summ.cached_prompt_tokens).unwrap_or(usize::MAX);
1301 let min_to_free = cached.saturating_sub(hard_threshold);
1302
1303 status.send_status("compacting context...").await;
1304
1305 crate::summarization::deferred::apply_deferred_summaries(summ);
1307
1308 let tokens_before = *summ.cached_prompt_tokens;
1316 let freed = crate::summarization::pruning::prune_tool_outputs(summ, min_to_free);
1317 if freed >= min_to_free {
1318 tracing::info!(freed, "hard compaction: pruning sufficient");
1319 summ.context_manager
1320 .set_compaction_state(CompactionState::CompactedThisTurn {
1321 cooldown: summ.context_manager.compaction_cooldown_turns(),
1322 });
1323 if let Err(e) = crate::summarization::deferred::flush_deferred_summaries(summ).await {
1324 tracing::warn!(%e, "flush_deferred_summaries failed after hard compaction");
1325 }
1326 status.send_status("").await;
1327 return Ok(());
1328 }
1329
1330 let preserve_tail = summ.context_manager.compaction_preserve_tail;
1332 let compactable = summ.messages.len().saturating_sub(preserve_tail + 1);
1333 if compactable <= 1 {
1334 tracing::warn!(
1335 compactable,
1336 "hard compaction: too few messages, marking exhausted"
1337 );
1338 summ.context_manager
1339 .set_compaction_state(CompactionState::Exhausted { warned: false });
1340 status.send_status("").await;
1341 return Ok(());
1342 }
1343
1344 tracing::info!(
1346 min_to_free,
1347 "hard compaction: falling back to LLM summarization"
1348 );
1349 let outcome = crate::summarization::compaction::compact_context(summ, None).await?;
1350
1351 let freed_tokens = tokens_before.saturating_sub(*summ.cached_prompt_tokens);
1352
1353 if !outcome.is_compacted() || freed_tokens == 0 {
1354 tracing::warn!("hard compaction: no net reduction, marking exhausted");
1355 summ.context_manager
1356 .set_compaction_state(CompactionState::Exhausted { warned: false });
1357 status.send_status("").await;
1358 return Ok(());
1359 }
1360
1361 if matches!(
1362 summ.context_manager
1363 .compaction_tier(*summ.cached_prompt_tokens),
1364 zeph_context::manager::CompactionTier::Hard
1365 ) {
1366 tracing::warn!(
1367 freed_tokens,
1368 "hard compaction: still above hard threshold after compaction, marking exhausted"
1369 );
1370 summ.context_manager
1371 .set_compaction_state(CompactionState::Exhausted { warned: false });
1372 status.send_status("").await;
1373 return Ok(());
1374 }
1375
1376 summ.context_manager
1377 .set_compaction_state(CompactionState::CompactedThisTurn {
1378 cooldown: summ.context_manager.compaction_cooldown_turns(),
1379 });
1380
1381 if tokens_before > *summ.cached_prompt_tokens {
1382 tracing::info!(
1383 tokens_before,
1384 tokens_after = *summ.cached_prompt_tokens,
1385 saved = freed_tokens,
1386 "context compaction complete"
1387 );
1388 }
1389
1390 status.send_status("").await;
1391 Ok(())
1392 }
1393
1394 #[tracing::instrument(name = "agent_context.service.maybe_summarize_tool_pair", skip_all)]
1400 pub async fn maybe_summarize_tool_pair(
1401 &self,
1402 summ: &mut ContextSummarizationView<'_>,
1403 providers: &ProviderHandles,
1404 ) {
1405 crate::summarization::deferred::maybe_summarize_tool_pair(
1406 summ,
1407 providers,
1408 &TxStatusSink(summ.status_tx.clone()),
1409 )
1410 .await;
1411 }
1412
1413 #[must_use]
1418 pub fn apply_deferred_summaries(&self, summ: &mut ContextSummarizationView<'_>) -> usize {
1419 crate::summarization::deferred::apply_deferred_summaries(summ)
1420 }
1421
1422 #[tracing::instrument(name = "agent_context.service.flush_deferred_summaries", skip_all)]
1427 pub async fn flush_deferred_summaries(&self, summ: &mut ContextSummarizationView<'_>) {
1428 if let Err(e) = crate::summarization::deferred::flush_deferred_summaries(summ).await {
1429 tracing::warn!(%e, "flush_deferred_summaries failed");
1430 }
1431 }
1432
1433 pub fn maybe_apply_deferred_summaries(&self, summ: &mut ContextSummarizationView<'_>) {
1439 crate::summarization::deferred::maybe_apply_deferred_summaries(summ);
1440 }
1441
1442 #[tracing::instrument(name = "agent_context.service.compact_context", skip_all, err)]
1457 pub async fn compact_context(
1458 &self,
1459 summ: &mut ContextSummarizationView<'_>,
1460 max_summary_tokens: Option<usize>,
1461 ) -> Result<crate::state::CompactionOutcome, crate::error::ContextError> {
1462 crate::summarization::compaction::compact_context(summ, max_summary_tokens).await
1463 }
1464
1465 pub fn maybe_soft_compact_mid_iteration(&self, summ: &mut ContextSummarizationView<'_>) {
1471 crate::summarization::scheduling::maybe_soft_compact_mid_iteration(summ);
1472 }
1473
1474 #[tracing::instrument(name = "agent_context.service.maybe_proactive_compress", skip_all)]
1480 pub async fn maybe_proactive_compress(
1481 &self,
1482 summ: &mut ContextSummarizationView<'_>,
1483 status: &(impl StatusSink + ?Sized),
1484 ) {
1485 let Some((_threshold, max_summary_tokens)) = summ
1486 .context_manager
1487 .should_proactively_compress(*summ.cached_prompt_tokens)
1488 else {
1489 return;
1490 };
1491
1492 if summ.server_compaction_active {
1493 let budget = summ
1494 .context_manager
1495 .budget
1496 .as_ref()
1497 .map_or(0, ContextBudget::max_tokens);
1498 if budget > 0 {
1499 let fallback = (budget * 95 / 100) as u64;
1500 if *summ.cached_prompt_tokens <= fallback {
1501 return;
1502 }
1503 tracing::warn!(
1504 cached_prompt_tokens = *summ.cached_prompt_tokens,
1505 fallback_threshold = fallback,
1506 "server compaction active but context at 95%+ — falling back to proactive"
1507 );
1508 } else {
1509 return;
1510 }
1511 }
1512
1513 status.send_status("compressing context...").await;
1514 tracing::info!(
1515 max_summary_tokens,
1516 cached_tokens = *summ.cached_prompt_tokens,
1517 "proactive compression triggered"
1518 );
1519
1520 match crate::summarization::compaction::compact_context(summ, Some(max_summary_tokens))
1521 .await
1522 {
1523 Ok(outcome) if outcome.is_compacted() => {
1524 summ.context_manager.set_compaction_state(
1525 zeph_context::manager::CompactionState::CompactedThisTurn { cooldown: 0 },
1526 );
1527 tracing::info!("proactive compression complete");
1528 }
1529 Ok(_) => {}
1530 Err(e) => tracing::warn!(%e, "proactive compression failed"),
1531 }
1532
1533 status.send_status("").await;
1534 }
1535
1536 pub fn maybe_refresh_task_goal(&self, summ: &mut ContextSummarizationView<'_>) {
1542 crate::summarization::scheduling::maybe_refresh_task_goal(summ);
1543 }
1544
1545 pub fn maybe_refresh_subgoal(&self, summ: &mut ContextSummarizationView<'_>) {
1550 crate::summarization::scheduling::maybe_refresh_subgoal(summ);
1551 }
1552}
1553
1554struct TxStatusSink(Option<tokio::sync::mpsc::UnboundedSender<String>>);
1560
1561impl StatusSink for TxStatusSink {
1562 fn send_status(&self, msg: &str) -> impl std::future::Future<Output = ()> + Send + '_ {
1563 if let Some(ref tx) = self.0 {
1564 let _ = tx.send(msg.to_owned());
1565 }
1566 std::future::ready(())
1567 }
1568}
1569
1570pub(crate) fn recompute_prompt_tokens(window: &mut MessageWindowView<'_>) {
1577 *window.cached_prompt_tokens = window
1578 .messages
1579 .iter()
1580 .map(|m| window.token_counter.count_message_tokens(m) as u64)
1581 .sum();
1582}
1583
1584fn fidelity_provider_pair<'a>(
1590 semantic: Option<&'a std::sync::Arc<zeph_llm::any::AnyProvider>>,
1591 compress: Option<&'a std::sync::Arc<zeph_llm::any::AnyProvider>>,
1592) -> (
1593 Option<&'a dyn zeph_llm::LlmProviderDyn>,
1594 Option<&'a dyn zeph_llm::LlmProviderDyn>,
1595) {
1596 let embed_provider = semantic
1597 .map(std::sync::Arc::as_ref)
1598 .map(|p| p as &dyn zeph_llm::LlmProviderDyn);
1599 let compress_provider = compress
1600 .map(std::sync::Arc::as_ref)
1601 .map(|p| p as &dyn zeph_llm::LlmProviderDyn);
1602 (embed_provider, compress_provider)
1603}
1604
1605#[tracing::instrument(name = "agent_context.service.persist_fidelity_tags", skip_all)]
1614async fn persist_fidelity_tags(
1615 messages: &[zeph_llm::provider::Message],
1616 memory: Option<&zeph_memory::semantic::SemanticMemory>,
1617) {
1618 let Some(mem) = memory else { return };
1619 let updates: Vec<(zeph_memory::MessageId, u8)> = messages
1620 .iter()
1621 .filter_map(|m| {
1622 let db_id = m.metadata.db_id?;
1623 let tag = m.metadata.fidelity_tag?;
1624 Some((zeph_memory::MessageId(db_id), tag as u8))
1625 })
1626 .collect();
1627 if updates.is_empty() {
1628 return;
1629 }
1630 if let Err(e) = mem.sqlite().update_fidelity_tags(&updates).await {
1631 tracing::warn!(
1632 count = updates.len(),
1633 error = %e,
1634 "failed to persist fidelity tags; floor invariant will not apply next turn"
1635 );
1636 }
1637}
1638
1639fn recompute_prompt_tokens_summ(summ: &mut crate::state::ContextSummarizationView<'_>) {
1643 *summ.cached_prompt_tokens = summ
1644 .messages
1645 .iter()
1646 .map(|m| summ.token_counter.count_message_tokens(m) as u64)
1647 .sum();
1648}
1649
1650pub(crate) fn remove_by_prefix(
1656 messages: &mut Vec<zeph_llm::provider::Message>,
1657 role: Role,
1658 prefix: &str,
1659) {
1660 messages.retain(|m| m.role != role || !m.content.starts_with(prefix));
1661}
1662
1663pub(crate) fn remove_by_part_or_prefix(
1672 messages: &mut Vec<zeph_llm::provider::Message>,
1673 prefix: &str,
1674 part_matches: impl Fn(&MessagePart) -> bool,
1675) {
1676 messages.retain(|m| {
1677 if m.role == Role::User {
1680 return !m.content.starts_with(prefix);
1681 }
1682 if m.role != Role::System {
1683 return true;
1684 }
1685 if m.parts.first().is_some_and(&part_matches) {
1686 return false;
1687 }
1688 !m.content.starts_with(prefix)
1689 });
1690}
1691
1692#[cfg(test)]
1693mod tests {
1694 use std::collections::HashSet;
1695 use std::sync::Arc;
1696
1697 use zeph_llm::provider::{Message, MessagePart, Role};
1698 use zeph_memory::TokenCounter;
1699
1700 use super::*;
1701 use crate::helpers::{GRAPH_FACTS_PREFIX, RECALL_PREFIX, SUMMARY_PREFIX};
1702 use crate::state::MessageWindowView;
1703
1704 fn make_counter() -> Arc<TokenCounter> {
1705 Arc::new(TokenCounter::default())
1706 }
1707
1708 fn make_window<'a>(
1709 messages: &'a mut Vec<Message>,
1710 cached: &'a mut u64,
1711 completed: &'a mut HashSet<String>,
1712 ) -> MessageWindowView<'a> {
1713 let last = Box::leak(Box::new(None::<i64>));
1714 let deferred_hide = Box::leak(Box::new(Vec::<i64>::new()));
1715 let deferred_summ = Box::leak(Box::new(Vec::<String>::new()));
1716 MessageWindowView {
1717 messages,
1718 last_persisted_message_id: last,
1719 deferred_db_hide_ids: deferred_hide,
1720 deferred_db_summaries: deferred_summ,
1721 cached_prompt_tokens: cached,
1722 token_counter: make_counter(),
1723 completed_tool_ids: completed,
1724 }
1725 }
1726
1727 fn sys(text: &str) -> Message {
1728 Message::from_legacy(Role::System, text)
1729 }
1730
1731 fn user(text: &str) -> Message {
1732 Message::from_legacy(Role::User, text)
1733 }
1734
1735 fn assistant(text: &str) -> Message {
1736 Message::from_legacy(Role::Assistant, text)
1737 }
1738
1739 #[test]
1740 fn clear_history_keeps_system_prompt() {
1741 let mut msgs = vec![sys("system"), user("hello"), assistant("hi")];
1742 let mut cached = 0u64;
1743 let mut completed = HashSet::new();
1744 completed.insert("tool_1".to_owned());
1745 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1746
1747 ContextService::new().clear_history(&mut window);
1748
1749 assert_eq!(window.messages.len(), 1);
1750 assert_eq!(window.messages[0].content, "system");
1751 assert!(
1752 window.completed_tool_ids.is_empty(),
1753 "completed_tool_ids must be cleared"
1754 );
1755 }
1756
1757 #[test]
1758 fn clear_history_empty_messages_is_noop() {
1759 let mut msgs: Vec<Message> = vec![];
1760 let mut cached = 0u64;
1761 let mut completed = HashSet::new();
1762 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1763
1764 ContextService::new().clear_history(&mut window);
1765
1766 assert!(window.messages.is_empty());
1767 }
1768
1769 #[test]
1770 fn remove_recall_messages_removes_by_prefix() {
1771 let mut msgs = vec![
1772 sys("system"),
1773 sys(&format!("{RECALL_PREFIX}some recalled text")),
1774 user("hello"),
1775 ];
1776 let mut cached = 0u64;
1777 let mut completed = HashSet::new();
1778 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1779
1780 ContextService::new().remove_recall_messages(&mut window);
1781
1782 assert_eq!(window.messages.len(), 2);
1783 assert!(
1784 window
1785 .messages
1786 .iter()
1787 .all(|m| !m.content.starts_with(RECALL_PREFIX))
1788 );
1789 }
1790
1791 #[test]
1794 fn remove_recall_messages_removes_user_role_recall() {
1795 let mut msgs = vec![
1796 sys("system"),
1797 user(&format!("{RECALL_PREFIX}recalled via tiered path")),
1798 user("real user message"),
1799 ];
1800 let mut cached = 0u64;
1801 let mut completed = HashSet::new();
1802 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1803
1804 ContextService::new().remove_recall_messages(&mut window);
1805
1806 assert_eq!(
1807 window.messages.len(),
1808 2,
1809 "Role::User recall message must be removed"
1810 );
1811 assert!(
1812 window
1813 .messages
1814 .iter()
1815 .all(|m| !m.content.starts_with(RECALL_PREFIX)),
1816 "no message with RECALL_PREFIX must remain"
1817 );
1818 assert!(
1819 window
1820 .messages
1821 .iter()
1822 .any(|m| m.content == "real user message"),
1823 "non-recall user message must survive"
1824 );
1825 }
1826
1827 #[test]
1828 fn remove_graph_facts_messages_removes_matching() {
1829 let mut msgs = vec![
1830 sys("system"),
1831 sys(&format!("{GRAPH_FACTS_PREFIX}fact1")),
1832 user("hello"),
1833 ];
1834 let mut cached = 0u64;
1835 let mut completed = HashSet::new();
1836 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1837
1838 ContextService::new().remove_graph_facts_messages(&mut window);
1839
1840 assert_eq!(window.messages.len(), 2);
1841 }
1842
1843 #[test]
1844 fn remove_summary_messages_removes_by_part() {
1845 let mut msgs = vec![
1846 sys("system"),
1847 Message::from_parts(
1848 Role::System,
1849 vec![MessagePart::Summary {
1850 text: format!("{SUMMARY_PREFIX}old summary"),
1851 }],
1852 ),
1853 user("hello"),
1854 ];
1855 let mut cached = 0u64;
1856 let mut completed = HashSet::new();
1857 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1858
1859 ContextService::new().remove_summary_messages(&mut window);
1860
1861 assert_eq!(window.messages.len(), 2);
1862 }
1863
1864 #[test]
1865 fn trim_messages_to_budget_zero_is_noop() {
1866 let mut msgs = vec![sys("system"), user("a"), assistant("b"), user("c")];
1867 let original_len = msgs.len();
1868 let mut cached = 0u64;
1869 let mut completed = HashSet::new();
1870 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1871
1872 ContextService::new().trim_messages_to_budget(&mut window, 0);
1873
1874 assert_eq!(window.messages.len(), original_len);
1875 }
1876
1877 #[test]
1878 fn trim_messages_to_budget_keeps_recent() {
1879 let mut msgs = vec![
1881 sys("system"),
1882 user("message 1"),
1883 assistant("reply 1"),
1884 user("message 2"),
1885 ];
1886 let mut cached = 0u64;
1887 let mut completed = HashSet::new();
1888 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1889
1890 ContextService::new().trim_messages_to_budget(&mut window, 1);
1892
1893 assert!(
1895 window.messages.len() < 4,
1896 "trim should remove some messages"
1897 );
1898 assert_eq!(
1899 window.messages[0].role,
1900 Role::System,
1901 "system prompt must survive trim"
1902 );
1903 }
1904
1905 mod inserted_count_tests {
1908 use parking_lot::RwLock;
1909 use std::borrow::Cow;
1910 use std::collections::HashSet;
1911 use std::sync::Arc;
1912
1913 use zeph_common::SecurityEventCategory;
1914 use zeph_config::memory::TieredRetrievalConfig;
1915 use zeph_config::{
1916 ContextFormat, ContextStrategy, DocumentConfig, GraphConfig, PersonaConfig,
1917 ReasoningConfig, TrajectoryConfig, TreeConfig,
1918 };
1919 use zeph_context::assembler::PreparedContext;
1920 use zeph_context::manager::ContextManager;
1921 use zeph_llm::provider::{Message, MessageMetadata, Role};
1922 use zeph_memory::TokenCounter;
1923 use zeph_sanitizer::ContentIsolationConfig;
1924 use zeph_sanitizer::ContentSanitizer;
1925 use zeph_skills::registry::SkillRegistry;
1926
1927 use super::super::*;
1928 use crate::state::{
1929 ContextAssemblyView, MessageWindowView, MetricsCounters, SecurityEventSink,
1930 };
1931
1932 fn make_task_supervisor() -> Arc<zeph_common::TaskSupervisor> {
1933 Arc::new(zeph_common::TaskSupervisor::new(
1934 tokio_util::sync::CancellationToken::new(),
1935 ))
1936 }
1937
1938 struct NoopSink;
1939 impl SecurityEventSink for NoopSink {
1940 fn push(&mut self, _: SecurityEventCategory, _: &'static str, _: String) {}
1941 }
1942
1943 fn make_counter() -> Arc<TokenCounter> {
1944 Arc::new(TokenCounter::default())
1945 }
1946
1947 fn make_window<'a>(
1948 messages: &'a mut Vec<Message>,
1949 cached: &'a mut u64,
1950 completed: &'a mut HashSet<String>,
1951 ) -> MessageWindowView<'a> {
1952 let last = Box::leak(Box::new(None::<i64>));
1953 let deferred_hide = Box::leak(Box::new(Vec::<i64>::new()));
1954 let deferred_summ = Box::leak(Box::new(Vec::<String>::new()));
1955 MessageWindowView {
1956 messages,
1957 last_persisted_message_id: last,
1958 deferred_db_hide_ids: deferred_hide,
1959 deferred_db_summaries: deferred_summ,
1960 cached_prompt_tokens: cached,
1961 token_counter: make_counter(),
1962 completed_tool_ids: completed,
1963 }
1964 }
1965
1966 fn mem_msg(content: &str) -> Message {
1967 Message {
1968 role: Role::User,
1969 content: content.to_string(),
1970 parts: vec![],
1971 metadata: MessageMetadata::default(),
1972 }
1973 }
1974
1975 fn scrub_noop(s: &str) -> Cow<'_, str> {
1976 Cow::Borrowed(s)
1977 }
1978
1979 #[tokio::test]
1980 async fn inserted_count_incremented_for_all_paths() {
1981 let mut msgs = vec![
1984 Message::from_legacy(Role::System, "system"),
1985 Message::from_legacy(Role::User, "user turn"),
1986 ];
1987 let mut cached = 0u64;
1988 let mut completed = HashSet::new();
1989 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1990
1991 let sanitizer = ContentSanitizer::new(&ContentIsolationConfig::default());
1992 let mut ctx_mgr = ContextManager::new();
1993 let mut sink = NoopSink;
1994 let mut last_confidence = None::<f32>;
1995 let mut last_skills_prompt = String::new();
1996 let mut active_skill_names = Vec::new();
1997 let registry = Arc::new(RwLock::new(SkillRegistry::default()));
1998
1999 let mut view = ContextAssemblyView {
2000 memory: None,
2001 conversation_id: None,
2002 recall_limit: 10,
2003 cross_session_score_threshold: 0.5,
2004 context_format: ContextFormat::default(),
2005 last_recall_confidence: &mut last_confidence,
2006 context_strategy: ContextStrategy::default(),
2007 crossover_turn_threshold: 0,
2008 cached_session_digest: None,
2009 digest_enabled: false, graph_config: GraphConfig::default(),
2011 document_config: DocumentConfig::default(),
2012 persona_config: PersonaConfig::default(),
2013 trajectory_config: TrajectoryConfig::default(),
2014 reasoning_config: ReasoningConfig::default(),
2015 memcot_config: zeph_config::MemCotConfig::default(),
2016 memcot_state: None,
2017 tree_config: TreeConfig::default(),
2018 last_skills_prompt: &mut last_skills_prompt,
2019 active_skill_names: &mut active_skill_names,
2020 skill_registry: registry,
2021 skill_paths: &[],
2022 correction_config: None,
2023 sidequest_turn_counter: 0,
2024 proactive_explorer: None,
2025 sanitizer: &sanitizer,
2026 quarantine_summarizer: None,
2027 context_manager: &mut ctx_mgr,
2028 token_counter: make_counter(),
2029 metrics: MetricsCounters::default(),
2030 security_events: &mut sink,
2031 cached_prompt_tokens: 0,
2032 redact_credentials: false,
2033 channel_skills: &[],
2034 scrub: scrub_noop,
2035 tiered_retrieval_config: TieredRetrievalConfig {
2036 enabled: false,
2037 ..TieredRetrievalConfig::default()
2038 },
2039 tiered_retrieval_classifier: None,
2040 tiered_retrieval_validator: None,
2041 type_aware_compose_config: zeph_config::memory::TypeAwareComposeConfig::default(),
2042 fidelity_config: None,
2043 fidelity_semantic_provider: None,
2044 fidelity_compress_provider: None,
2045 planned_next_tools: &[],
2046 status_tx: None,
2047 task_supervisor: make_task_supervisor(),
2048 };
2049
2050 let prepared = PreparedContext {
2052 graph_facts: Some(mem_msg("graph_facts")),
2053 doc_rag: Some(mem_msg("doc_rag")),
2054 corrections: Some(mem_msg("corrections")),
2055 recall: Some(mem_msg("recall")),
2056 recall_confidence: Some(0.9),
2057 cross_session: Some(mem_msg("cross_session")),
2058 summaries: Some(mem_msg("summaries")),
2059 code_context: None, persona_facts: Some(mem_msg("persona_facts")),
2061 trajectory_hints: Some(mem_msg("trajectory_hints")),
2062 tree_memory: Some(mem_msg("tree_memory")),
2063 reasoning_hints: Some(mem_msg("reasoning_hints")),
2064 memory_first: false,
2065 recent_history_budget: 100_000,
2066 background_tasks: vec![],
2067 };
2068
2069 let (_delta, inserted_count) = ContextService::new()
2070 .apply_prepared_context(&mut window, &mut view, prepared)
2071 .await;
2072
2073 assert_eq!(
2076 inserted_count, 10,
2077 "all 10 message-carrying PreparedContext fields must increment inserted_count"
2078 );
2079 }
2080 }
2081
2082 mod inject_semantic_recall_tests {
2083 use parking_lot::RwLock;
2084 use std::borrow::Cow;
2085 use std::collections::HashSet;
2086 use std::sync::Arc;
2087
2088 use zeph_config::memory::TieredRetrievalConfig;
2089 use zeph_config::{
2090 ContextFormat, ContextStrategy, DocumentConfig, GraphConfig, PersonaConfig,
2091 ReasoningConfig, TrajectoryConfig, TreeConfig,
2092 };
2093 use zeph_context::manager::ContextManager;
2094 use zeph_llm::provider::Message;
2095 use zeph_memory::TokenCounter;
2096 use zeph_sanitizer::ContentIsolationConfig;
2097 use zeph_sanitizer::ContentSanitizer;
2098 use zeph_skills::registry::SkillRegistry;
2099
2100 use zeph_common::SecurityEventCategory;
2101
2102 use super::super::*;
2103 use crate::helpers::RECALL_PREFIX;
2104 use crate::state::{
2105 ContextAssemblyView, MessageWindowView, MetricsCounters, SecurityEventSink,
2106 };
2107
2108 fn make_task_supervisor() -> Arc<zeph_common::TaskSupervisor> {
2109 Arc::new(zeph_common::TaskSupervisor::new(
2110 tokio_util::sync::CancellationToken::new(),
2111 ))
2112 }
2113
2114 struct NoopSink;
2115 impl SecurityEventSink for NoopSink {
2116 fn push(&mut self, _: SecurityEventCategory, _: &'static str, _: String) {}
2117 }
2118
2119 fn make_counter() -> Arc<TokenCounter> {
2120 Arc::new(TokenCounter::default())
2121 }
2122
2123 fn make_window<'a>(
2124 messages: &'a mut Vec<Message>,
2125 cached: &'a mut u64,
2126 completed: &'a mut HashSet<String>,
2127 ) -> MessageWindowView<'a> {
2128 let last = Box::leak(Box::new(None::<i64>));
2129 let deferred_hide = Box::leak(Box::new(Vec::<i64>::new()));
2130 let deferred_summ = Box::leak(Box::new(Vec::<String>::new()));
2131 MessageWindowView {
2132 messages,
2133 last_persisted_message_id: last,
2134 deferred_db_hide_ids: deferred_hide,
2135 deferred_db_summaries: deferred_summ,
2136 cached_prompt_tokens: cached,
2137 token_counter: make_counter(),
2138 completed_tool_ids: completed,
2139 }
2140 }
2141
2142 fn scrub_noop(s: &str) -> Cow<'_, str> {
2143 Cow::Borrowed(s)
2144 }
2145
2146 #[tokio::test]
2147 async fn tiered_recall_disabled_uses_flat_path() {
2148 let mut msgs: Vec<Message> = vec![];
2151 let mut cached = 0u64;
2152 let mut completed = HashSet::new();
2153 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2154
2155 let sanitizer = ContentSanitizer::new(&ContentIsolationConfig::default());
2156 let mut ctx_mgr = ContextManager::new();
2157 let mut sink = NoopSink;
2158 let mut last_confidence = None::<f32>;
2159 let mut last_skills_prompt = String::new();
2160 let mut active_skill_names = Vec::new();
2161 let registry = Arc::new(RwLock::new(SkillRegistry::default()));
2162
2163 let view = ContextAssemblyView {
2164 memory: None,
2165 conversation_id: None,
2166 recall_limit: 10,
2167 cross_session_score_threshold: 0.5,
2168 context_format: ContextFormat::default(),
2169 last_recall_confidence: &mut last_confidence,
2170 context_strategy: ContextStrategy::default(),
2171 crossover_turn_threshold: 0,
2172 cached_session_digest: None,
2173 digest_enabled: false,
2174 graph_config: GraphConfig::default(),
2175 document_config: DocumentConfig::default(),
2176 persona_config: PersonaConfig::default(),
2177 trajectory_config: TrajectoryConfig::default(),
2178 reasoning_config: ReasoningConfig::default(),
2179 memcot_config: zeph_config::MemCotConfig::default(),
2180 memcot_state: None,
2181 tree_config: TreeConfig::default(),
2182 last_skills_prompt: &mut last_skills_prompt,
2183 active_skill_names: &mut active_skill_names,
2184 skill_registry: registry,
2185 skill_paths: &[],
2186 correction_config: None,
2187 sidequest_turn_counter: 0,
2188 proactive_explorer: None,
2189 sanitizer: &sanitizer,
2190 quarantine_summarizer: None,
2191 context_manager: &mut ctx_mgr,
2192 token_counter: make_counter(),
2193 metrics: MetricsCounters::default(),
2194 security_events: &mut sink,
2195 cached_prompt_tokens: 0,
2196 redact_credentials: false,
2197 channel_skills: &[],
2198 scrub: scrub_noop,
2199 tiered_retrieval_config: TieredRetrievalConfig {
2200 enabled: false,
2201 ..TieredRetrievalConfig::default()
2202 },
2203 tiered_retrieval_classifier: None,
2204 tiered_retrieval_validator: None,
2205 type_aware_compose_config: zeph_config::memory::TypeAwareComposeConfig::default(),
2206 fidelity_config: None,
2207 fidelity_semantic_provider: None,
2208 fidelity_compress_provider: None,
2209 planned_next_tools: &[],
2210 status_tx: None,
2211 task_supervisor: make_task_supervisor(),
2212 };
2213
2214 let result = ContextService::new()
2215 .inject_semantic_recall("test query", 1000, &mut window, &view)
2216 .await;
2217
2218 assert!(result.is_ok(), "disabled tiered recall must return Ok(())");
2219 assert!(
2220 window
2221 .messages
2222 .iter()
2223 .all(|m| !m.content.starts_with(RECALL_PREFIX)),
2224 "no recall message must be injected when memory is None"
2225 );
2226 }
2227
2228 #[tokio::test]
2229 async fn tiered_recall_enabled_no_memory_returns_ok() {
2230 let mut msgs: Vec<Message> = vec![];
2233 let mut cached = 0u64;
2234 let mut completed = HashSet::new();
2235 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2236
2237 let sanitizer = ContentSanitizer::new(&ContentIsolationConfig::default());
2238 let mut ctx_mgr = ContextManager::new();
2239 let mut sink = NoopSink;
2240 let mut last_confidence = None::<f32>;
2241 let mut last_skills_prompt = String::new();
2242 let mut active_skill_names = Vec::new();
2243 let registry = Arc::new(RwLock::new(SkillRegistry::default()));
2244
2245 let view = ContextAssemblyView {
2246 memory: None,
2247 conversation_id: None,
2248 recall_limit: 10,
2249 cross_session_score_threshold: 0.5,
2250 context_format: ContextFormat::default(),
2251 last_recall_confidence: &mut last_confidence,
2252 context_strategy: ContextStrategy::default(),
2253 crossover_turn_threshold: 0,
2254 cached_session_digest: None,
2255 digest_enabled: false,
2256 graph_config: GraphConfig::default(),
2257 document_config: DocumentConfig::default(),
2258 persona_config: PersonaConfig::default(),
2259 trajectory_config: TrajectoryConfig::default(),
2260 reasoning_config: ReasoningConfig::default(),
2261 memcot_config: zeph_config::MemCotConfig::default(),
2262 memcot_state: None,
2263 tree_config: TreeConfig::default(),
2264 last_skills_prompt: &mut last_skills_prompt,
2265 active_skill_names: &mut active_skill_names,
2266 skill_registry: registry,
2267 skill_paths: &[],
2268 correction_config: None,
2269 sidequest_turn_counter: 0,
2270 proactive_explorer: None,
2271 sanitizer: &sanitizer,
2272 quarantine_summarizer: None,
2273 context_manager: &mut ctx_mgr,
2274 token_counter: make_counter(),
2275 metrics: MetricsCounters::default(),
2276 security_events: &mut sink,
2277 cached_prompt_tokens: 0,
2278 redact_credentials: false,
2279 channel_skills: &[],
2280 scrub: scrub_noop,
2281 tiered_retrieval_config: TieredRetrievalConfig {
2282 enabled: true,
2283 ..TieredRetrievalConfig::default()
2284 },
2285 tiered_retrieval_classifier: None,
2286 tiered_retrieval_validator: None,
2287 type_aware_compose_config: zeph_config::memory::TypeAwareComposeConfig::default(),
2288 fidelity_config: None,
2289 fidelity_semantic_provider: None,
2290 fidelity_compress_provider: None,
2291 planned_next_tools: &[],
2292 status_tx: None,
2293 task_supervisor: make_task_supervisor(),
2294 };
2295
2296 let result = ContextService::new()
2297 .inject_semantic_recall("test query", 1000, &mut window, &view)
2298 .await;
2299
2300 assert!(
2301 result.is_ok(),
2302 "enabled tiered recall with no memory must return Ok(())"
2303 );
2304 assert!(
2305 window.messages.is_empty(),
2306 "no recall message must be injected when memory is None"
2307 );
2308 }
2309
2310 #[tokio::test]
2315 async fn prepare_context_tiered_enabled_no_budget_returns_default() {
2316 let mut msgs: Vec<zeph_llm::provider::Message> = vec![];
2317 let mut cached = 0u64;
2318 let mut completed = HashSet::new();
2319 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2320
2321 let sanitizer = zeph_sanitizer::ContentSanitizer::new(
2322 &zeph_sanitizer::ContentIsolationConfig::default(),
2323 );
2324 let mut ctx_mgr = zeph_context::manager::ContextManager::new();
2325 assert!(ctx_mgr.budget.is_none());
2327
2328 let mut sink = NoopSink;
2329 let mut last_confidence = None::<f32>;
2330 let mut last_skills_prompt = String::new();
2331 let mut active_skill_names = Vec::new();
2332 let registry = Arc::new(RwLock::new(zeph_skills::registry::SkillRegistry::default()));
2333
2334 let mut view = ContextAssemblyView {
2335 memory: None,
2336 conversation_id: None,
2337 recall_limit: 10,
2338 cross_session_score_threshold: 0.5,
2339 context_format: ContextFormat::default(),
2340 last_recall_confidence: &mut last_confidence,
2341 context_strategy: ContextStrategy::default(),
2342 crossover_turn_threshold: 0,
2343 cached_session_digest: None,
2344 digest_enabled: false,
2345 graph_config: GraphConfig::default(),
2346 document_config: DocumentConfig::default(),
2347 persona_config: PersonaConfig::default(),
2348 trajectory_config: TrajectoryConfig::default(),
2349 reasoning_config: ReasoningConfig::default(),
2350 memcot_config: zeph_config::MemCotConfig::default(),
2351 memcot_state: None,
2352 tree_config: TreeConfig::default(),
2353 last_skills_prompt: &mut last_skills_prompt,
2354 active_skill_names: &mut active_skill_names,
2355 skill_registry: registry,
2356 skill_paths: &[],
2357 correction_config: None,
2358 sidequest_turn_counter: 0,
2359 proactive_explorer: None,
2360 sanitizer: &sanitizer,
2361 quarantine_summarizer: None,
2362 context_manager: &mut ctx_mgr,
2363 token_counter: make_counter(),
2364 metrics: MetricsCounters::default(),
2365 security_events: &mut sink,
2366 cached_prompt_tokens: 0,
2367 redact_credentials: false,
2368 channel_skills: &[],
2369 scrub: scrub_noop,
2370 tiered_retrieval_config: TieredRetrievalConfig {
2371 enabled: true,
2372 ..TieredRetrievalConfig::default()
2373 },
2374 tiered_retrieval_classifier: None,
2375 tiered_retrieval_validator: None,
2376 type_aware_compose_config: zeph_config::memory::TypeAwareComposeConfig::default(),
2377 fidelity_config: None,
2378 fidelity_semantic_provider: None,
2379 fidelity_compress_provider: None,
2380 planned_next_tools: &[],
2381 status_tx: None,
2382 task_supervisor: make_task_supervisor(),
2383 };
2384
2385 let result = ContextService::new()
2386 .prepare_context("test query", &mut window, &mut view)
2387 .await;
2388
2389 assert!(
2390 result.is_ok(),
2391 "prepare_context with tiered enabled and no budget must return Ok"
2392 );
2393 }
2394
2395 #[tokio::test]
2398 async fn inject_semantic_recall_bare_no_memory_returns_ok() {
2399 use zeph_config::memory::TieredRetrievalConfig;
2400
2401 let mut msgs: Vec<Message> = vec![];
2402 let mut cached = 0u64;
2403 let mut completed = HashSet::new();
2404 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2405
2406 let tiered_config = TieredRetrievalConfig {
2407 enabled: true,
2408 ..TieredRetrievalConfig::default()
2409 };
2410 let params = SemanticRecallParams {
2411 query: "test query",
2412 token_budget: 1000,
2413 recall_limit: 10,
2414 context_format: zeph_config::ContextFormat::default(),
2415 conversation_id: None,
2416 tiered_classifier: None,
2417 tiered_validator: None,
2418 tiered_config: &tiered_config,
2419 };
2420 let result = ContextService::new()
2421 .inject_semantic_recall_bare(params, &mut window, None)
2422 .await;
2423
2424 assert!(
2425 result.is_ok(),
2426 "inject_semantic_recall_bare with memory=None must return Ok(())"
2427 );
2428 assert!(
2429 window.messages.is_empty(),
2430 "no recall message must be injected when memory is None"
2431 );
2432 }
2433 }
2434
2435 mod prepare_context_proactive_explore_tests {
2442 use parking_lot::RwLock;
2443 use std::borrow::Cow;
2444 use std::collections::HashSet;
2445 use std::sync::Arc;
2446
2447 use zeph_config::memory::TieredRetrievalConfig;
2448 use zeph_config::{
2449 ContextFormat, ContextStrategy, DocumentConfig, GraphConfig, PersonaConfig,
2450 ReasoningConfig, TrajectoryConfig, TreeConfig,
2451 };
2452 use zeph_context::budget::ContextBudget;
2453 use zeph_context::manager::ContextManager;
2454 use zeph_llm::any::AnyProvider;
2455 use zeph_llm::mock::MockProvider;
2456 use zeph_llm::provider::Message;
2457 use zeph_memory::TokenCounter;
2458 use zeph_sanitizer::ContentIsolationConfig;
2459 use zeph_sanitizer::ContentSanitizer;
2460 use zeph_skills::generator::SkillGenerator;
2461 use zeph_skills::proactive::ProactiveExplorer;
2462 use zeph_skills::registry::SkillRegistry;
2463
2464 use zeph_common::SecurityEventCategory;
2465
2466 use super::super::*;
2467 use crate::state::{
2468 ContextAssemblyView, MessageWindowView, MetricsCounters, SecurityEventSink,
2469 };
2470
2471 fn make_task_supervisor() -> Arc<zeph_common::TaskSupervisor> {
2472 Arc::new(zeph_common::TaskSupervisor::new(
2473 tokio_util::sync::CancellationToken::new(),
2474 ))
2475 }
2476
2477 struct NoopSink;
2478 impl SecurityEventSink for NoopSink {
2479 fn push(&mut self, _: SecurityEventCategory, _: &'static str, _: String) {}
2480 }
2481
2482 fn make_counter() -> Arc<TokenCounter> {
2483 Arc::new(TokenCounter::default())
2484 }
2485
2486 fn make_window<'a>(
2487 messages: &'a mut Vec<Message>,
2488 cached: &'a mut u64,
2489 completed: &'a mut HashSet<String>,
2490 ) -> MessageWindowView<'a> {
2491 let last = Box::leak(Box::new(None::<i64>));
2492 let deferred_hide = Box::leak(Box::new(Vec::<i64>::new()));
2493 let deferred_summ = Box::leak(Box::new(Vec::<String>::new()));
2494 MessageWindowView {
2495 messages,
2496 last_persisted_message_id: last,
2497 deferred_db_hide_ids: deferred_hide,
2498 deferred_db_summaries: deferred_summ,
2499 cached_prompt_tokens: cached,
2500 token_counter: make_counter(),
2501 completed_tool_ids: completed,
2502 }
2503 }
2504
2505 fn scrub_noop(s: &str) -> Cow<'_, str> {
2506 Cow::Borrowed(s)
2507 }
2508
2509 fn mock_skill_content(name: &str) -> String {
2510 format!(
2511 "---\nname: {name}\ndescription: Test world-knowledge skill for {name}.\n---\n\n## Usage\n\nDetails.\n"
2512 )
2513 }
2514
2515 struct Fixture {
2520 registry: Arc<RwLock<SkillRegistry>>,
2521 skill_paths: Vec<std::path::PathBuf>,
2522 hub_dir: tempfile::TempDir,
2523 _skills_dir_guard: tempfile::TempDir,
2526 explorer: Arc<ProactiveExplorer>,
2527 }
2528
2529 fn setup_fixture(mock_provider: MockProvider) -> Fixture {
2530 let skills_dir = tempfile::tempdir().expect("create skills tempdir");
2531 let hub_dir = tempfile::tempdir().expect("create hub tempdir");
2532 let skill_paths = vec![skills_dir.path().to_path_buf()];
2533
2534 let registry = Arc::new(RwLock::new(
2535 SkillRegistry::load(&skill_paths).with_hub_dirs(vec![hub_dir.path().to_path_buf()]),
2536 ));
2537 assert!(
2538 registry.read().all_meta().is_empty(),
2539 "registry must start empty before any explore() reload"
2540 );
2541
2542 let generator = SkillGenerator::new(
2543 AnyProvider::Mock(mock_provider),
2544 skills_dir.path().to_path_buf(),
2545 );
2546 let explorer = Arc::new(ProactiveExplorer::new(
2547 generator,
2548 None,
2549 skills_dir.path().to_path_buf(),
2550 8_000,
2551 30_000,
2552 vec![],
2553 ));
2554
2555 Fixture {
2556 registry,
2557 skill_paths,
2558 hub_dir,
2559 _skills_dir_guard: skills_dir,
2560 explorer,
2561 }
2562 }
2563
2564 #[allow(clippy::too_many_arguments)]
2565 fn make_view<'a>(
2566 fixture: &'a Fixture,
2567 last_confidence: &'a mut Option<f32>,
2568 last_skills_prompt: &'a mut String,
2569 active_skill_names: &'a mut Vec<String>,
2570 sanitizer: &'a ContentSanitizer,
2571 ctx_mgr: &'a mut ContextManager,
2572 sink: &'a mut NoopSink,
2573 ) -> ContextAssemblyView<'a> {
2574 ContextAssemblyView {
2575 memory: None,
2576 conversation_id: None,
2577 recall_limit: 10,
2578 cross_session_score_threshold: 0.5,
2579 context_format: ContextFormat::default(),
2580 last_recall_confidence: last_confidence,
2581 context_strategy: ContextStrategy::default(),
2582 crossover_turn_threshold: 0,
2583 cached_session_digest: None,
2584 digest_enabled: false,
2585 graph_config: GraphConfig::default(),
2586 document_config: DocumentConfig::default(),
2587 persona_config: PersonaConfig::default(),
2588 trajectory_config: TrajectoryConfig::default(),
2589 reasoning_config: ReasoningConfig::default(),
2590 memcot_config: zeph_config::MemCotConfig::default(),
2591 memcot_state: None,
2592 tree_config: TreeConfig::default(),
2593 last_skills_prompt,
2594 active_skill_names,
2595 skill_registry: Arc::clone(&fixture.registry),
2596 skill_paths: &fixture.skill_paths,
2597 correction_config: None,
2598 sidequest_turn_counter: 0,
2599 proactive_explorer: Some(Arc::clone(&fixture.explorer)),
2600 sanitizer,
2601 quarantine_summarizer: None,
2602 context_manager: ctx_mgr,
2603 token_counter: make_counter(),
2604 metrics: MetricsCounters::default(),
2605 security_events: sink,
2606 cached_prompt_tokens: 0,
2607 redact_credentials: false,
2608 channel_skills: &[],
2609 scrub: scrub_noop,
2610 tiered_retrieval_config: TieredRetrievalConfig {
2611 enabled: false,
2612 ..TieredRetrievalConfig::default()
2613 },
2614 tiered_retrieval_classifier: None,
2615 tiered_retrieval_validator: None,
2616 type_aware_compose_config: zeph_config::memory::TypeAwareComposeConfig::default(),
2617 fidelity_config: None,
2618 fidelity_semantic_provider: None,
2619 fidelity_compress_provider: None,
2620 planned_next_tools: &[],
2621 status_tx: None,
2622 task_supervisor: make_task_supervisor(),
2623 }
2624 }
2625
2626 #[tokio::test]
2627 async fn prepare_context_proactive_explore_reloads_registry_off_lock() {
2628 let fixture = setup_fixture(MockProvider::with_responses(vec![mock_skill_content(
2632 "world-knowledge-git",
2633 )]));
2634 let expected_hub_dirs = vec![fixture.hub_dir.path().to_path_buf()];
2635
2636 let mut msgs: Vec<Message> = vec![];
2637 let mut cached = 0u64;
2638 let mut completed = HashSet::new();
2639 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2640
2641 let sanitizer = ContentSanitizer::new(&ContentIsolationConfig::default());
2642 let mut ctx_mgr = ContextManager::new();
2643 ctx_mgr.budget = Some(ContextBudget::new(100_000, 0.1));
2644 let mut sink = NoopSink;
2645 let mut last_confidence = None::<f32>;
2646 let mut last_skills_prompt = String::new();
2647 let mut active_skill_names = Vec::new();
2648
2649 let mut view = make_view(
2650 &fixture,
2651 &mut last_confidence,
2652 &mut last_skills_prompt,
2653 &mut active_skill_names,
2654 &sanitizer,
2655 &mut ctx_mgr,
2656 &mut sink,
2657 );
2658
2659 let result = ContextService::new()
2660 .prepare_context("please help me with a git rebase", &mut window, &mut view)
2661 .await;
2662
2663 assert!(result.is_ok(), "prepare_context must succeed: {result:?}");
2664
2665 let guard = fixture.registry.read();
2666 assert_eq!(
2667 guard.all_meta().len(),
2668 1,
2669 "reloaded registry must contain exactly the newly generated skill"
2670 );
2671 assert!(
2672 guard
2673 .all_meta()
2674 .iter()
2675 .any(|m| m.name == "world-knowledge-git"),
2676 "reloaded registry must expose the generated skill by name"
2677 );
2678 assert_eq!(
2679 guard.hub_dirs(),
2680 expected_hub_dirs.as_slice(),
2681 "hub_dirs read before spawn_blocking must survive the off-lock rebuild"
2682 );
2683 }
2684
2685 #[tokio::test]
2686 async fn prepare_context_proactive_explore_leaves_registry_unchanged_when_not_triggered() {
2687 let fixture = setup_fixture(MockProvider::with_responses(vec![mock_skill_content(
2691 "world-knowledge-unused",
2692 )]));
2693
2694 let existing_skill_dir = fixture.skill_paths[0].join("existing-skill");
2697 std::fs::create_dir_all(&existing_skill_dir).expect("create existing skill dir");
2698 std::fs::write(
2699 existing_skill_dir.join("SKILL.md"),
2700 mock_skill_content("existing-skill"),
2701 )
2702 .expect("write existing SKILL.md");
2703 {
2704 let mut guard = fixture.registry.write();
2705 *guard = SkillRegistry::load(&fixture.skill_paths)
2706 .with_hub_dirs(vec![fixture.hub_dir.path().to_path_buf()]);
2707 }
2708 assert_eq!(
2709 fixture.registry.read().all_meta().len(),
2710 1,
2711 "sanity: registry must be pre-populated before prepare_context runs"
2712 );
2713
2714 let mut msgs: Vec<Message> = vec![];
2715 let mut cached = 0u64;
2716 let mut completed = HashSet::new();
2717 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2718
2719 let sanitizer = ContentSanitizer::new(&ContentIsolationConfig::default());
2720 let mut ctx_mgr = ContextManager::new();
2721 ctx_mgr.budget = Some(ContextBudget::new(100_000, 0.1));
2722 let mut sink = NoopSink;
2723 let mut last_confidence = None::<f32>;
2724 let mut last_skills_prompt = String::new();
2725 let mut active_skill_names = Vec::new();
2726
2727 let mut view = make_view(
2728 &fixture,
2729 &mut last_confidence,
2730 &mut last_skills_prompt,
2731 &mut active_skill_names,
2732 &sanitizer,
2733 &mut ctx_mgr,
2734 &mut sink,
2735 );
2736
2737 let result = ContextService::new()
2738 .prepare_context("how are you today", &mut window, &mut view)
2739 .await;
2740
2741 assert!(result.is_ok(), "prepare_context must succeed: {result:?}");
2742 let guard = fixture.registry.read();
2743 assert_eq!(
2744 guard.all_meta().len(),
2745 1,
2746 "registry must remain unchanged when no domain classifies and explore() never runs"
2747 );
2748 assert!(
2749 guard.all_meta().iter().any(|m| m.name == "existing-skill"),
2750 "the pre-existing skill must still be present, unreplaced by the mock's canned skill"
2751 );
2752 }
2753
2754 #[tokio::test]
2755 async fn prepare_context_proactive_explore_does_not_retrigger_for_known_domain() {
2756 let (mock_provider, recorder) =
2764 MockProvider::with_responses(vec![mock_skill_content("terraform-quickref")])
2765 .with_recording();
2766 let fixture = setup_fixture(mock_provider);
2767
2768 let sanitizer = ContentSanitizer::new(&ContentIsolationConfig::default());
2769
2770 {
2773 let mut msgs: Vec<Message> = vec![];
2774 let mut cached = 0u64;
2775 let mut completed = HashSet::new();
2776 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2777 let mut ctx_mgr = ContextManager::new();
2778 ctx_mgr.budget = Some(ContextBudget::new(100_000, 0.1));
2779 let mut sink = NoopSink;
2780 let mut last_confidence = None::<f32>;
2781 let mut last_skills_prompt = String::new();
2782 let mut active_skill_names = Vec::new();
2783
2784 let mut view = make_view(
2785 &fixture,
2786 &mut last_confidence,
2787 &mut last_skills_prompt,
2788 &mut active_skill_names,
2789 &sanitizer,
2790 &mut ctx_mgr,
2791 &mut sink,
2792 );
2793
2794 let result = ContextService::new()
2795 .prepare_context("tell me about terraform modules", &mut window, &mut view)
2796 .await;
2797 assert!(
2798 result.is_ok(),
2799 "first prepare_context must succeed: {result:?}"
2800 );
2801 }
2802
2803 assert_eq!(
2804 recorder.lock().unwrap().len(),
2805 1,
2806 "first turn must trigger exactly one LLM generation call"
2807 );
2808 assert!(
2809 fixture
2810 .registry
2811 .read()
2812 .all_meta()
2813 .iter()
2814 .any(|m| m.proactive_domain.as_deref() == Some("terraform")),
2815 "reloaded registry must carry the stamped proactive_domain field"
2816 );
2817
2818 {
2821 let mut msgs: Vec<Message> = vec![];
2822 let mut cached = 0u64;
2823 let mut completed = HashSet::new();
2824 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2825 let mut ctx_mgr = ContextManager::new();
2826 ctx_mgr.budget = Some(ContextBudget::new(100_000, 0.1));
2827 let mut sink = NoopSink;
2828 let mut last_confidence = None::<f32>;
2829 let mut last_skills_prompt = String::new();
2830 let mut active_skill_names = Vec::new();
2831
2832 let mut view = make_view(
2833 &fixture,
2834 &mut last_confidence,
2835 &mut last_skills_prompt,
2836 &mut active_skill_names,
2837 &sanitizer,
2838 &mut ctx_mgr,
2839 &mut sink,
2840 );
2841
2842 let result = ContextService::new()
2843 .prepare_context("tell me more about terraform state", &mut window, &mut view)
2844 .await;
2845 assert!(
2846 result.is_ok(),
2847 "second prepare_context must succeed: {result:?}"
2848 );
2849 }
2850
2851 assert_eq!(
2852 recorder.lock().unwrap().len(),
2853 1,
2854 "second turn for an already-known domain must NOT trigger another LLM generation call"
2855 );
2856 }
2857 }
2858
2859 mod prune_token_bookkeeping_tests {
2867 use std::time::Duration;
2868
2869 use tokio_util::sync::CancellationToken;
2870 use zeph_common::task_supervisor::{BlockingHandle, TaskSupervisor};
2871 use zeph_context::budget::ContextBudget;
2872 use zeph_context::manager::{CompactionTier, ContextManager};
2873 use zeph_context::summarization::{MessageTokenCounter, SummarizationDeps};
2874 use zeph_llm::any::AnyProvider;
2875 use zeph_llm::mock::MockProvider;
2876 use zeph_llm::provider::MessageMetadata;
2877
2878 use super::*;
2879 use crate::compaction::{SubgoalExtractionResult, SubgoalRegistry};
2880 use crate::memory_backend::TokenCounterAdapter;
2881
2882 struct Fixture {
2885 messages: Vec<Message>,
2886 deferred_db_hide_ids: Vec<i64>,
2887 deferred_db_summaries: Vec<String>,
2888 cached_prompt_tokens: u64,
2889 context_manager: ContextManager,
2890 subgoal_registry: SubgoalRegistry,
2891 pending_task_goal: Option<BlockingHandle<Option<String>>>,
2892 pending_subgoal: Option<BlockingHandle<Option<SubgoalExtractionResult>>>,
2893 current_task_goal: Option<String>,
2894 task_goal_user_msg_hash: Option<u64>,
2895 subgoal_user_msg_hash: Option<u64>,
2896 token_counter: Arc<TokenCounter>,
2897 task_supervisor: Arc<TaskSupervisor>,
2898 provider_responses: Vec<String>,
2903 }
2904
2905 impl Fixture {
2906 fn new(messages: Vec<Message>, cached_prompt_tokens: u64, cm: ContextManager) -> Self {
2907 Self {
2908 messages,
2909 deferred_db_hide_ids: Vec::new(),
2910 deferred_db_summaries: Vec::new(),
2911 cached_prompt_tokens,
2912 context_manager: cm,
2913 subgoal_registry: SubgoalRegistry::default(),
2914 pending_task_goal: None,
2915 pending_subgoal: None,
2916 current_task_goal: None,
2917 task_goal_user_msg_hash: None,
2918 subgoal_user_msg_hash: None,
2919 token_counter: make_counter(),
2920 task_supervisor: Arc::new(TaskSupervisor::new(CancellationToken::new())),
2921 provider_responses: Vec::new(),
2922 }
2923 }
2924
2925 fn view(&mut self) -> ContextSummarizationView<'_> {
2926 let token_counter_adapter: Arc<dyn MessageTokenCounter> =
2927 Arc::new(TokenCounterAdapter::new(Arc::clone(&self.token_counter)));
2928 ContextSummarizationView {
2929 messages: &mut self.messages,
2930 deferred_db_hide_ids: &mut self.deferred_db_hide_ids,
2931 deferred_db_summaries: &mut self.deferred_db_summaries,
2932 cached_prompt_tokens: &mut self.cached_prompt_tokens,
2933 context_manager: &mut self.context_manager,
2934 server_compaction_active: false,
2935 token_counter: Arc::clone(&self.token_counter),
2936 summarization_deps: SummarizationDeps {
2937 provider: AnyProvider::Mock(MockProvider::with_responses(
2938 self.provider_responses.clone(),
2939 )),
2940 llm_timeout: Duration::from_secs(30),
2941 token_counter: token_counter_adapter,
2942 structured_summaries: false,
2943 on_anchored_summary: None,
2944 },
2945 task_supervisor: Arc::clone(&self.task_supervisor),
2946 memory: None,
2947 conversation_id: None,
2948 tool_call_cutoff: 100,
2949 subgoal_registry: &mut self.subgoal_registry,
2950 pending_task_goal: &mut self.pending_task_goal,
2951 pending_subgoal: &mut self.pending_subgoal,
2952 current_task_goal: &mut self.current_task_goal,
2953 task_goal_user_msg_hash: &mut self.task_goal_user_msg_hash,
2954 subgoal_user_msg_hash: &mut self.subgoal_user_msg_hash,
2955 status_tx: None,
2956 scrub: |s| std::borrow::Cow::Borrowed(s),
2957 compression_guidelines: None,
2958 probe: None,
2959 archive: None,
2960 persistence: None,
2961 metrics: None,
2962 typed_pages: None,
2963 fidelity_config: None,
2964 fidelity_semantic_provider: None,
2965 fidelity_compress_provider: None,
2966 current_query: String::new(),
2967 }
2968 }
2969 }
2970
2971 struct NoopStatus;
2972 impl StatusSink for NoopStatus {
2973 fn send_status(&self, _msg: &str) -> impl std::future::Future<Output = ()> + Send + '_ {
2974 std::future::ready(())
2975 }
2976 }
2977
2978 fn plain_msg(role: Role, content: &str) -> Message {
2979 Message {
2980 role,
2981 content: content.to_owned(),
2982 parts: vec![],
2983 metadata: MessageMetadata::default(),
2984 }
2985 }
2986
2987 fn tool_use_msg() -> Message {
2988 Message::from_parts(
2989 Role::Assistant,
2990 vec![MessagePart::ToolUse {
2991 id: "t1".into(),
2992 name: "shell".into(),
2993 input: serde_json::json!({}),
2994 }],
2995 )
2996 }
2997
2998 fn tool_output_msg(body: &str) -> Message {
2999 Message::from_parts(
3000 Role::User,
3001 vec![MessagePart::ToolOutput {
3002 tool_name: "shell".into(),
3003 body: body.to_owned(),
3004 compacted_at: None,
3005 }],
3006 )
3007 }
3008
3009 fn messages_with_one_tool_output(body: &str) -> Vec<Message> {
3012 vec![
3013 plain_msg(Role::System, "system"),
3014 tool_use_msg(),
3015 tool_output_msg(body),
3016 plain_msg(Role::User, "hello"),
3017 ]
3018 }
3019
3020 #[test]
3021 fn prune_tool_outputs_decrements_cached_tokens_by_exact_freed_amount() {
3022 let body = "large tool output ".repeat(200);
3023 let expected_freed = TokenCounter::default().count_tokens(&body);
3024 let mut ctx_mgr = ContextManager::new();
3025 ctx_mgr.prune_protect_tokens = 0;
3026
3027 let initial_tokens = 10_000u64;
3028 let mut fixture = Fixture::new(
3029 messages_with_one_tool_output(&body),
3030 initial_tokens,
3031 ctx_mgr,
3032 );
3033 let mut view = fixture.view();
3034
3035 let freed =
3036 crate::summarization::pruning::prune_tool_outputs(&mut view, expected_freed);
3037
3038 assert_eq!(
3039 freed, expected_freed,
3040 "returned freed amount must match the tool output's token count"
3041 );
3042 assert_eq!(
3043 *view.cached_prompt_tokens,
3044 initial_tokens - u64::try_from(freed).unwrap(),
3045 "cached_prompt_tokens must decrease by exactly the freed amount"
3046 );
3047 }
3048
3049 #[test]
3050 fn prune_tool_outputs_noop_leaves_cached_tokens_unchanged() {
3051 let mut ctx_mgr = ContextManager::new();
3052 ctx_mgr.prune_protect_tokens = 0;
3053
3054 let messages = vec![
3055 plain_msg(Role::System, "system"),
3056 plain_msg(Role::User, "hello"),
3057 ];
3058 let initial_tokens = 500u64;
3059 let mut fixture = Fixture::new(messages, initial_tokens, ctx_mgr);
3060 let mut view = fixture.view();
3061
3062 let freed = crate::summarization::pruning::prune_tool_outputs(&mut view, 100);
3063
3064 assert_eq!(freed, 0, "no ToolOutput parts exist to prune");
3065 assert_eq!(
3066 *view.cached_prompt_tokens, initial_tokens,
3067 "cached_prompt_tokens must be untouched when nothing is freed"
3068 );
3069 }
3070
3071 #[tokio::test]
3072 async fn do_soft_compaction_via_maybe_compact_reflects_real_freed_tokens() {
3073 let body = "large tool output ".repeat(200);
3074 let expected_freed = TokenCounter::default().count_tokens(&body);
3075
3076 let mut ctx_mgr = ContextManager::new();
3077 ctx_mgr.budget = Some(ContextBudget::new(1000, 0.2));
3078 ctx_mgr.soft_compaction_threshold = 0.5; ctx_mgr.hard_compaction_threshold = 0.9; ctx_mgr.prune_protect_tokens = 0;
3081
3082 let initial_tokens = 700u64; assert_eq!(
3084 ctx_mgr.compaction_tier(initial_tokens),
3085 CompactionTier::Soft
3086 );
3087
3088 let mut fixture = Fixture::new(
3089 messages_with_one_tool_output(&body),
3090 initial_tokens,
3091 ctx_mgr,
3092 );
3093 let mut view = fixture.view();
3094 let status = NoopStatus;
3095
3096 ContextService::new()
3097 .maybe_compact(&mut view, &status)
3098 .await
3099 .unwrap();
3100
3101 assert!(
3102 view.deferred_db_summaries.is_empty(),
3103 "no deferred summaries were queued in this scenario"
3104 );
3105 assert_eq!(
3106 *view.cached_prompt_tokens,
3107 initial_tokens - u64::try_from(expected_freed).unwrap(),
3108 "Soft-tier prune-only pass must decrement cached_prompt_tokens by exactly \
3109 the freed amount"
3110 );
3111 }
3112
3113 #[tokio::test]
3114 async fn do_hard_compaction_satisfied_by_pruning_alone_reflects_real_freed_tokens() {
3115 let body = "large tool output ".repeat(400);
3119 let expected_freed = TokenCounter::default().count_tokens(&body);
3120
3121 let mut ctx_mgr = ContextManager::new();
3122 ctx_mgr.budget = Some(ContextBudget::new(1_000_000, 0.2));
3123 ctx_mgr.soft_compaction_threshold = 0.5;
3124 ctx_mgr.hard_compaction_threshold = 0.6; ctx_mgr.prune_protect_tokens = 0;
3126
3127 let initial_tokens = 600_000u64 + (expected_freed as u64 / 2);
3130 assert_eq!(
3131 ctx_mgr.compaction_tier(initial_tokens),
3132 CompactionTier::Hard
3133 );
3134
3135 let mut fixture = Fixture::new(
3136 messages_with_one_tool_output(&body),
3137 initial_tokens,
3138 ctx_mgr,
3139 );
3140 let mut view = fixture.view();
3141 let status = NoopStatus;
3142
3143 ContextService::new()
3144 .maybe_compact(&mut view, &status)
3145 .await
3146 .unwrap();
3147
3148 assert!(
3149 view.deferred_db_summaries.is_empty(),
3150 "no deferred summaries were queued in this scenario"
3151 );
3152 assert_eq!(
3153 *view.cached_prompt_tokens,
3154 initial_tokens - u64::try_from(expected_freed).unwrap(),
3155 "Hard-tier pass satisfied by pruning alone must decrement cached_prompt_tokens \
3156 by exactly the freed amount"
3157 );
3158 assert!(
3159 view.context_manager
3160 .compaction_state()
3161 .is_compacted_this_turn(),
3162 "pruning-satisfied Hard tier must still mark the turn as compacted"
3163 );
3164 }
3165
3166 #[tokio::test]
3191 async fn do_hard_compaction_combined_prune_and_llm_savings_avoid_false_exhaustion() {
3192 let counter = TokenCounter::default();
3193
3194 let large_body = "large tool output content ".repeat(400);
3198 let summary_response = "ok".to_string();
3199
3200 let build_messages = || {
3201 vec![
3202 plain_msg(Role::System, "system"),
3203 tool_use_msg(),
3204 tool_output_msg(&large_body),
3205 plain_msg(Role::User, "t0"),
3206 plain_msg(Role::User, "t1"),
3207 plain_msg(Role::User, "t2"),
3208 ]
3209 };
3210
3211 let initial_tokens: u64 = build_messages()
3212 .iter()
3213 .map(|m| counter.count_message_tokens(m) as u64)
3214 .sum();
3215 let non_body_overhead =
3216 initial_tokens - u64::try_from(counter.count_tokens(&large_body)).unwrap();
3217 let budget_tokens = usize::try_from(initial_tokens).unwrap();
3218
3219 let make_ctx_mgr = |hard_ratio: f32| {
3220 let mut cm = ContextManager::new();
3221 cm.prune_protect_tokens = 0;
3222 cm.compaction_preserve_tail = 2;
3223 cm.budget = Some(ContextBudget::new(budget_tokens, 0.0));
3224 cm.hard_compaction_threshold = hard_ratio;
3225 cm
3226 };
3227
3228 let measured_final_tokens = {
3233 let mut fixture =
3234 Fixture::new(build_messages(), initial_tokens, make_ctx_mgr(0.0001));
3235 fixture.provider_responses = vec![summary_response.clone()];
3236 let mut view = fixture.view();
3237 ContextService::new()
3238 .do_hard_compaction(&mut view, &NoopStatus, false)
3239 .await
3240 .unwrap();
3241 *view.cached_prompt_tokens
3242 };
3243 assert!(
3244 measured_final_tokens < non_body_overhead,
3245 "test setup invariant: the LLM step must reduce tokens below the non-body \
3246 overhead ({non_body_overhead}) so a valid hard-threshold window exists; \
3247 measured {measured_final_tokens}"
3248 );
3249
3250 let hard_threshold_tokens = u64::midpoint(measured_final_tokens, non_body_overhead);
3255 #[allow(clippy::cast_precision_loss)]
3256 let hard_ratio = hard_threshold_tokens as f32 / budget_tokens as f32;
3257
3258 let mut fixture =
3259 Fixture::new(build_messages(), initial_tokens, make_ctx_mgr(hard_ratio));
3260 fixture.provider_responses = vec![summary_response];
3261 let mut view = fixture.view();
3262
3263 ContextService::new()
3264 .do_hard_compaction(&mut view, &NoopStatus, false)
3265 .await
3266 .unwrap();
3267
3268 assert!(
3269 !view.context_manager.compaction_state().is_exhausted(),
3270 "pruning's own savings must count toward the combined freed-tokens check, so \
3271 a Hard-tier pass where pruning alone was insufficient but prune+LLM combined \
3272 cross the hard threshold must not be marked Exhausted"
3273 );
3274 assert!(
3275 view.context_manager
3276 .compaction_state()
3277 .is_compacted_this_turn(),
3278 "the turn must be marked compacted given the combined prune + LLM reduction"
3279 );
3280 }
3281 }
3282}