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 params.recall_limit,
363 params.context_format,
364 params.query,
365 params.token_budget,
366 &window.token_counter,
367 None,
368 None,
369 )
370 .await?;
371 Ok(msg)
372 }
373 }
374
375 #[tracing::instrument(
384 name = "agent_context.service.inject_cross_session_context",
385 skip_all,
386 err
387 )]
388 pub async fn inject_cross_session_context(
389 &self,
390 query: &str,
391 token_budget: usize,
392 window: &mut MessageWindowView<'_>,
393 view: &ContextAssemblyView<'_>,
394 ) -> Result<(), ContextError> {
395 self.remove_cross_session_messages(window);
396
397 if let Some(msg) = crate::helpers::fetch_cross_session_raw(
398 view.memory.as_deref(),
399 view.conversation_id,
400 view.cross_session_score_threshold,
401 query,
402 token_budget,
403 &view.token_counter,
404 )
405 .await?
406 && window.messages.len() > 1
407 {
408 window.messages.insert(1, msg);
409 tracing::debug!("injected cross-session context");
410 }
411
412 Ok(())
413 }
414
415 #[tracing::instrument(name = "agent_context.service.inject_summaries", skip_all, err)]
424 pub async fn inject_summaries(
425 &self,
426 token_budget: usize,
427 window: &mut MessageWindowView<'_>,
428 view: &ContextAssemblyView<'_>,
429 ) -> Result<(), ContextError> {
430 self.remove_summary_messages(window);
431
432 if let Some(msg) = crate::helpers::fetch_summaries_raw(
433 view.memory.as_deref(),
434 view.conversation_id,
435 token_budget,
436 &view.token_counter,
437 )
438 .await?
439 && window.messages.len() > 1
440 {
441 window.messages.insert(1, msg);
442 tracing::debug!("injected summaries into context");
443 }
444
445 Ok(())
446 }
447
448 #[tracing::instrument(name = "agent_context.service.disambiguate_skills", skip_all)]
453 pub async fn disambiguate_skills(
454 &self,
455 query: &str,
456 all_meta: &[&zeph_skills::loader::SkillMeta],
457 scored: &[zeph_skills::ScoredMatch],
458 providers: &ProviderHandles,
459 ) -> Option<Vec<usize>> {
460 use std::fmt::Write as _;
461
462 let mut candidates = String::new();
463 for sm in scored {
464 if let Some(meta) = all_meta.get(sm.index) {
465 let _ = writeln!(
466 candidates,
467 "- {} (score: {:.3}): {}",
468 meta.name, sm.score, meta.description
469 );
470 }
471 }
472
473 let prompt = format!(
474 "The user said: \"{query}\"\n\n\
475 These skills matched with similar scores:\n{candidates}\n\
476 Which skill best matches the user's intent? \
477 Return the skill_name, your confidence (0-1), and any extracted parameters."
478 );
479
480 let messages = vec![zeph_llm::provider::Message::from_legacy(
481 zeph_llm::provider::Role::User,
482 prompt,
483 )];
484 match providers
485 .disambiguate
486 .chat_typed::<zeph_skills::IntentClassification>(&messages)
487 .await
488 {
489 Ok(classification) => {
490 tracing::info!(
491 skill = %classification.skill_name,
492 confidence = classification.confidence,
493 "disambiguation selected skill"
494 );
495 let mut indices: Vec<usize> = scored.iter().map(|s| s.index).collect();
496 if let Some(pos) = indices.iter().position(|&i| {
497 all_meta
498 .get(i)
499 .is_some_and(|m| m.name == classification.skill_name)
500 }) {
501 indices.swap(0, pos);
502 }
503 Some(indices)
504 }
505 Err(e) => {
506 tracing::warn!("disambiguation failed, using original order: {e:#}");
507 None
508 }
509 }
510 }
511
512 #[allow(clippy::too_many_lines)] #[tracing::instrument(name = "agent_context.service.prepare_context", skip_all, err)]
525 pub async fn prepare_context(
526 &self,
527 query: &str,
528 window: &mut MessageWindowView<'_>,
529 view: &mut ContextAssemblyView<'_>,
530 ) -> Result<ContextDelta, ContextError> {
531 if view.context_manager.budget.is_none() {
532 return Ok(ContextDelta::default());
533 }
534
535 self.remove_session_digest_message(window);
537 self.remove_summary_messages(window);
538 self.remove_cross_session_messages(window);
539 self.remove_recall_messages(window);
540 self.remove_document_rag_messages(window);
541 self.remove_correction_messages(window);
542 self.remove_code_context_messages(window);
543 self.remove_graph_facts_messages(window);
544 self.remove_persona_facts_messages(window);
545 self.remove_trajectory_hints_messages(window);
546 self.remove_tree_memory_messages(window);
547 if view.reasoning_config.enabled {
548 self.remove_reasoning_strategies_messages(window);
549 }
550
551 if let Some(explorer) = view.proactive_explorer.clone()
553 && let Some(domain) = explorer.classify(query)
554 {
555 let already_known = {
556 let registry_guard = view.skill_registry.read();
557 explorer.has_knowledge(®istry_guard, &domain)
558 };
559 let excluded = explorer.is_excluded(&domain);
560
561 if !already_known && !excluded {
562 tracing::debug!(domain = %domain.0, query_len = query.len(), "proactive.explore triggered");
563 let timeout_ms = explorer.timeout_ms();
564 let result = tokio::time::timeout(
565 std::time::Duration::from_millis(timeout_ms),
566 explorer.explore(&domain),
567 )
568 .await;
569 match result {
570 Ok(Ok(())) => {
571 let hub_dirs = view.skill_registry.read().hub_dirs().to_vec();
572 let reload_paths = view.skill_paths.to_vec();
573 let span =
574 tracing::info_span!("agent_context.skills.registry.reload_blocking");
575 match tokio::task::spawn_blocking(move || {
576 let _enter = span.enter();
577 SkillRegistry::load(&reload_paths).with_hub_dirs(hub_dirs)
578 })
579 .await
580 {
581 Ok(new_registry) => {
582 *view.skill_registry.write() = new_registry;
583 tracing::debug!(domain = %domain.0, "proactive.explore complete, registry reloaded");
584 }
585 Err(e) => {
586 tracing::error!(
587 domain = %domain.0,
588 "proactive.explore: skill registry reload panicked, registry left unchanged: {e}"
589 );
590 }
591 }
592 }
593 Ok(Err(e)) => {
594 tracing::warn!(domain = %domain.0, error = %e, "proactive exploration failed");
595 }
596 Err(_) => {
597 tracing::warn!(domain = %domain.0, timeout_ms, "proactive exploration timed out");
598 }
599 }
600 }
601 }
602
603 let active_levels: &'static [zeph_memory::compression::CompressionLevel] =
605 if let Some(ref budget) = view.context_manager.budget {
606 let used = view.cached_prompt_tokens;
607 let max = budget.max_tokens();
608 #[allow(clippy::cast_precision_loss)]
609 let remaining_ratio = if max == 0 {
610 1.0_f32
611 } else {
612 1.0 - (used as f32 / max as f32).clamp(0.0, 1.0)
613 };
614 let levels =
615 zeph_memory::compression::RetrievalPolicy::default().select(remaining_ratio);
616 tracing::debug!(
617 remaining_ratio,
618 active_levels = ?levels,
619 "compression_spectrum: retrieval policy selected"
620 );
621 levels
622 } else {
623 &[]
624 };
625
626 let memory_backend: Option<std::sync::Arc<dyn zeph_common::memory::ContextMemoryBackend>> =
627 view.memory.clone().map(
628 |m| -> std::sync::Arc<dyn zeph_common::memory::ContextMemoryBackend> {
629 std::sync::Arc::new(crate::memory_backend::SemanticMemoryBackend::new(m))
630 },
631 );
632
633 let memory_view = zeph_context::input::ContextMemoryView {
634 memory: memory_backend,
635 conversation_id: view.conversation_id.map(|c| c.0),
636 recall_limit: view.recall_limit,
637 cross_session_score_threshold: view.cross_session_score_threshold,
638 context_strategy: view.context_strategy,
639 crossover_turn_threshold: view.crossover_turn_threshold,
640 cached_session_digest: view.cached_session_digest.clone(),
641 graph_config: view.graph_config.clone(),
642 document_config: view.document_config.clone(),
643 persona_config: view.persona_config.clone(),
644 trajectory_config: view.trajectory_config.clone(),
645 reasoning_config: view.reasoning_config.clone(),
646 memcot_config: view.memcot_config.clone(),
647 memcot_state: view.memcot_state.clone(),
648 tree_config: view.tree_config.clone(),
649 };
650
651 #[cfg(feature = "index")]
652 let index_access = view.index;
653 #[cfg(not(feature = "index"))]
654 let index_access: Option<&dyn zeph_context::input::IndexAccess> = None;
655
656 let router = crate::memory_backend::build_memory_router(view.context_manager);
657
658 let input = zeph_context::input::ContextAssemblyInput {
659 memory: &memory_view,
660 context_manager: view.context_manager,
661 token_counter: &*view.token_counter,
662 skills_prompt: view.last_skills_prompt,
663 index: index_access,
664 correction_config: view.correction_config,
665 sidequest_turn_counter: view.sidequest_turn_counter,
666 messages: window.messages,
667 query,
668 scrub: view.scrub,
669 active_levels,
670 router,
671 planned_next_tools: view.planned_next_tools,
672 };
673
674 let mut prepared = zeph_context::assembler::ContextAssembler::gather(&input).await?;
675
676 if view.tiered_retrieval_config.enabled {
681 prepared.recall = None;
682 }
683
684 for handle in prepared.background_tasks.drain(..) {
688 let task_supervisor = std::sync::Arc::clone(&view.task_supervisor);
689 drop(task_supervisor.spawn_oneshot(
690 std::sync::Arc::from("context.assembly.background"),
691 move || async move {
692 let _ = handle.await;
693 },
694 ));
695 }
696
697 let (delta, inserted_count) = self.apply_prepared_context(window, view, prepared).await;
698
699 if view.tiered_retrieval_config.enabled {
700 self.inject_semantic_recall(query, usize::MAX, window, view)
701 .await?;
702 }
703
704 let memory_first_active =
709 view.context_strategy == zeph_config::ContextStrategy::MemoryFirst;
710 if let Some(fidelity_cfg) = view.fidelity_config
711 && fidelity_cfg.enabled
712 && !memory_first_active
713 {
714 use tracing::Instrument as _;
715 if let Some(ref tx) = view.status_tx {
716 let _ = tx.send("Scoring context fidelity\u{2026}".into());
717 }
718 let (embed_provider, compress_provider) = fidelity_provider_pair(
719 view.fidelity_semantic_provider.as_ref(),
720 view.fidelity_compress_provider.as_ref(),
721 );
722 let fidelity_span = tracing::info_span!(
723 "context.fidelity.score",
724 message_count = window.messages.len(),
725 query_len = query.len(),
726 );
727 FidelityScorer
728 .score_and_apply(
729 window.messages,
730 query,
731 view.planned_next_tools,
732 fidelity_cfg,
733 &*view.token_counter,
734 inserted_count,
735 false, embed_provider,
737 compress_provider,
738 )
739 .instrument(fidelity_span)
740 .await;
741 persist_fidelity_tags(window.messages, view.memory.as_deref()).await;
743 recompute_prompt_tokens(window);
744 if let Some(ref tx) = view.status_tx {
745 let _ = tx.send(String::new());
746 }
747 }
748
749 Ok(delta)
750 }
751
752 #[allow(clippy::too_many_lines)] #[tracing::instrument(name = "agent_context.service.apply_prepared_context", skip_all)]
762 async fn apply_prepared_context(
763 &self,
764 window: &mut MessageWindowView<'_>,
765 view: &mut ContextAssemblyView<'_>,
766 prepared: zeph_context::assembler::PreparedContext,
767 ) -> (ContextDelta, usize) {
768 use std::borrow::Cow;
769 use zeph_llm::provider::{Message, MessageMetadata, Role};
770 use zeph_sanitizer::{ContentSource, ContentSourceKind, MemorySourceHint};
771
772 *view.last_recall_confidence = prepared.recall_confidence;
774
775 if prepared.memory_first {
777 let history_start = 1usize;
778 let len = window.messages.len();
779 let keep_tail =
780 zeph_context::assembler::memory_first_keep_tail(window.messages, history_start);
781 if len > history_start + keep_tail {
782 window.messages.drain(history_start..len - keep_tail);
783 recompute_prompt_tokens(window);
784 tracing::debug!(
785 strategy = "memory_first",
786 keep_tail,
787 "dropped conversation history, kept last {keep_tail} messages"
788 );
789 }
790 }
791
792 let mut inserted_count: usize = 0;
795
796 let slots: &[(Option<Message>, MemorySourceHint, Option<&str>)] = &[
799 (
800 prepared.graph_facts,
801 MemorySourceHint::ExternalContent,
802 Some("injected knowledge graph facts into context"),
803 ),
804 (
805 prepared.doc_rag,
806 MemorySourceHint::ExternalContent,
807 Some("injected document RAG context"),
808 ),
809 (
810 prepared.corrections,
811 MemorySourceHint::ConversationHistory,
812 Some("injected past corrections into context"),
813 ),
814 (prepared.recall, MemorySourceHint::ConversationHistory, None),
815 (prepared.cross_session, MemorySourceHint::LlmSummary, None),
816 (
817 prepared.summaries,
818 MemorySourceHint::LlmSummary,
819 Some("injected summaries into context"),
820 ),
821 (
822 prepared.persona_facts,
823 MemorySourceHint::ExternalContent,
824 Some("injected persona facts into context"),
825 ),
826 (
827 prepared.trajectory_hints,
828 MemorySourceHint::ExternalContent,
829 Some("injected trajectory hints into context"),
830 ),
831 (
832 prepared.tree_memory,
833 MemorySourceHint::ExternalContent,
834 Some("injected tree memory summary into context"),
835 ),
836 (
837 prepared.reasoning_hints,
838 MemorySourceHint::ExternalContent,
839 Some("injected reasoning strategies into context"),
840 ),
841 ];
842 for (opt_msg, hint, label) in slots.iter().cloned() {
843 if let Some(msg) = opt_msg.filter(|_| window.messages.len() > 1) {
844 let sanitized = self.sanitize_memory_message(msg, hint, view).await;
845 window.messages.insert(1, sanitized);
846 inserted_count += 1;
847 if let Some(lbl) = label {
848 tracing::debug!("{lbl}");
849 }
850 }
851 }
852
853 let code_context = if let Some(text) = prepared.code_context {
855 let sanitized = view
856 .sanitizer
857 .sanitize(&text, ContentSource::new(ContentSourceKind::ToolResult));
858 view.metrics.sanitizer_runs += 1;
859 if !sanitized.injection_flags.is_empty() {
860 tracing::warn!(
861 flags = sanitized.injection_flags.len(),
862 "injection patterns detected in code RAG context"
863 );
864 view.metrics.sanitizer_injection_flags += sanitized.injection_flags.len() as u64;
865 let detail = sanitized
866 .injection_flags
867 .first()
868 .map_or_else(String::new, |f| {
869 format!("Detected pattern: {}", f.pattern_name)
870 });
871 view.security_events.push(
872 zeph_common::SecurityEventCategory::InjectionFlag,
873 "code_rag",
874 detail,
875 );
876 }
877 if sanitized.was_truncated {
878 view.metrics.sanitizer_truncations += 1;
879 view.security_events.push(
880 zeph_common::SecurityEventCategory::Truncation,
881 "code_rag",
882 "Content truncated to max_content_size".to_string(),
883 );
884 }
885 Some(sanitized.body)
886 } else {
887 None
888 };
889
890 if !prepared.memory_first {
891 self.trim_messages_to_budget(window, prepared.recent_history_budget);
892 }
893
894 if view.digest_enabled
896 && let Some((digest_text, _)) = view
897 .cached_session_digest
898 .clone()
899 .filter(|_| window.messages.len() > 1)
900 {
901 let digest_msg = Message {
902 role: Role::User,
903 content: format!("{}{digest_text}", crate::helpers::SESSION_DIGEST_PREFIX),
904 parts: vec![],
905 metadata: MessageMetadata::default(),
906 };
907 let sanitized = self
908 .sanitize_memory_message(digest_msg, MemorySourceHint::LlmSummary, view)
909 .await;
910 window.messages.insert(1, sanitized);
911 inserted_count += 1;
912 tracing::debug!("injected session digest into context");
913 }
914
915 if view.redact_credentials {
917 for msg in &mut *window.messages {
918 if msg.role == Role::System {
919 continue;
920 }
921 if let Cow::Owned(s) = (view.scrub)(&msg.content) {
922 msg.content = s;
923 }
924 }
925 }
926
927 recompute_prompt_tokens(window);
928
929 (ContextDelta { code_context }, inserted_count)
930 }
931
932 #[tracing::instrument(name = "agent_context.service.sanitize_memory_message", skip_all)]
942 async fn sanitize_memory_message(
943 &self,
944 mut msg: zeph_llm::provider::Message,
945 hint: zeph_sanitizer::MemorySourceHint,
946 view: &mut ContextAssemblyView<'_>,
947 ) -> zeph_llm::provider::Message {
948 use zeph_sanitizer::{ContentSource, ContentSourceKind};
949
950 let source = ContentSource::new(ContentSourceKind::MemoryRetrieval).with_memory_hint(hint);
951 let sanitized = view.sanitizer.sanitize(&msg.content, source);
952 view.metrics.sanitizer_runs += 1;
953 if !sanitized.injection_flags.is_empty() {
954 tracing::warn!(
955 flags = sanitized.injection_flags.len(),
956 "injection patterns detected in memory retrieval"
957 );
958 view.metrics.sanitizer_injection_flags += sanitized.injection_flags.len() as u64;
959 let detail = sanitized
960 .injection_flags
961 .first()
962 .map_or_else(String::new, |f| {
963 format!("Detected pattern: {}", f.pattern_name)
964 });
965 view.security_events.push(
966 zeph_common::SecurityEventCategory::InjectionFlag,
967 "memory_retrieval",
968 detail,
969 );
970 }
971 if sanitized.was_truncated {
972 view.metrics.sanitizer_truncations += 1;
973 view.security_events.push(
974 zeph_common::SecurityEventCategory::Truncation,
975 "memory_retrieval",
976 "Content truncated to max_content_size".to_string(),
977 );
978 }
979
980 if view.sanitizer.is_enabled()
982 && let Some(qs) = view.quarantine_summarizer
983 && qs.should_quarantine(ContentSourceKind::MemoryRetrieval)
984 {
985 match qs.extract_facts(&sanitized, view.sanitizer).await {
986 Ok((facts, flags)) => {
987 view.metrics.quarantine_invocations += 1;
988 view.security_events.push(
989 zeph_common::SecurityEventCategory::Quarantine,
990 "memory_retrieval",
991 "Content quarantined, facts extracted".to_string(),
992 );
993 let escaped = zeph_sanitizer::ContentSanitizer::escape_delimiter_tags(&facts);
994 msg.content = zeph_sanitizer::ContentSanitizer::apply_spotlight(
995 &escaped,
996 &sanitized.source,
997 &flags,
998 );
999 return msg;
1000 }
1001 Err(e) => {
1002 tracing::warn!(
1003 error = %e,
1004 "quarantine failed for memory retrieval, using original sanitized content"
1005 );
1006 view.metrics.quarantine_failures += 1;
1007 view.security_events.push(
1008 zeph_common::SecurityEventCategory::Quarantine,
1009 "memory_retrieval",
1010 format!("Quarantine failed: {e}"),
1011 );
1012 }
1013 }
1014 }
1015
1016 msg.content = sanitized.body;
1017 msg
1018 }
1019
1020 pub fn reset_conversation(
1030 &self,
1031 window: &mut MessageWindowView<'_>,
1032 _view: &mut ContextAssemblyView<'_>,
1033 ) -> Result<(), ContextError> {
1034 self.clear_history(window);
1035 Ok(())
1036 }
1037
1038 #[allow(
1055 clippy::cast_precision_loss,
1056 clippy::cast_possible_truncation,
1057 clippy::cast_sign_loss,
1058 clippy::too_many_lines
1059 )]
1060 #[tracing::instrument(name = "agent_context.service.maybe_compact", skip_all, err)]
1061 pub async fn maybe_compact(
1062 &self,
1063 summ: &mut ContextSummarizationView<'_>,
1064 status: &(impl StatusSink + ?Sized),
1065 ) -> Result<(), ContextError> {
1066 use zeph_context::manager::{CompactionState, CompactionTier};
1067
1068 if let Some(count) = summ.context_manager.turns_since_last_hard_compaction_mut() {
1070 *count += 1;
1071 }
1072
1073 if let CompactionState::Exhausted { warned } = summ.context_manager.compaction_state()
1075 && !warned
1076 {
1077 summ.context_manager
1078 .set_compaction_state(CompactionState::Exhausted { warned: true });
1079 tracing::warn!("compaction exhausted: context budget too tight for this session");
1080 }
1081 if summ.context_manager.compaction_state().is_exhausted() {
1082 return Ok(());
1083 }
1084
1085 if summ.server_compaction_active {
1087 let budget = summ
1088 .context_manager
1089 .budget
1090 .as_ref()
1091 .map_or(0, ContextBudget::max_tokens);
1092 if budget > 0 {
1093 let fallback = (budget * 95 / 100) as u64;
1094 if *summ.cached_prompt_tokens < fallback {
1095 return Ok(());
1096 }
1097 tracing::warn!(
1098 "server compaction active but context at 95%+ — falling back to client-side"
1099 );
1100 } else {
1101 return Ok(());
1102 }
1103 }
1104
1105 if summ
1107 .context_manager
1108 .compaction_state()
1109 .is_compacted_this_turn()
1110 {
1111 return Ok(());
1112 }
1113
1114 let in_cooldown = summ.context_manager.compaction_state().cooldown_remaining() > 0;
1116 if in_cooldown
1117 && let CompactionState::Cooling { turns_remaining } =
1118 summ.context_manager.compaction_state()
1119 {
1120 let next = turns_remaining - 1;
1121 summ.context_manager.set_compaction_state(if next == 0 {
1122 CompactionState::Ready
1123 } else {
1124 CompactionState::Cooling {
1125 turns_remaining: next,
1126 }
1127 });
1128 }
1129
1130 if let Some(ref fidelity_cfg) = summ.fidelity_config.clone()
1134 && fidelity_cfg.enabled
1135 && summ.context_manager.should_proactively_regrade(
1136 *summ.cached_prompt_tokens,
1137 fidelity_cfg.regrade_threshold,
1138 summ.server_compaction_active,
1139 )
1140 {
1141 use tracing::Instrument as _;
1142 let (regrade_embed_provider, regrade_compress_provider) = fidelity_provider_pair(
1143 summ.fidelity_semantic_provider.as_ref(),
1144 summ.fidelity_compress_provider.as_ref(),
1145 );
1146 FidelityScorer
1147 .score_and_apply(
1148 summ.messages,
1149 &summ.current_query,
1150 &[],
1151 fidelity_cfg,
1152 &*summ.token_counter,
1153 0,
1154 true, regrade_embed_provider,
1156 regrade_compress_provider,
1157 )
1158 .instrument(tracing::info_span!(
1159 "context.fidelity.regrade",
1160 budget_ratio = tracing::field::Empty,
1161 ))
1162 .await;
1163 persist_fidelity_tags(summ.messages, summ.memory.as_deref()).await;
1165 recompute_prompt_tokens_summ(summ);
1166 summ.context_manager.set_regraded_this_turn(true);
1167 tracing::debug!(
1168 cached_tokens = *summ.cached_prompt_tokens,
1169 "AgeMem proactive regrade complete"
1170 );
1171 }
1172
1173 match summ
1174 .context_manager
1175 .compaction_tier(*summ.cached_prompt_tokens)
1176 {
1177 CompactionTier::Soft => {
1178 self.do_soft_compaction(summ, status).await;
1179 Ok(())
1180 }
1181 CompactionTier::Hard => self.do_hard_compaction(summ, status, in_cooldown).await,
1182 _ => Ok(()),
1183 }
1184 }
1185
1186 #[tracing::instrument(name = "agent_context.service.do_soft_compaction", skip_all)]
1191 #[allow(
1192 clippy::cast_precision_loss,
1193 clippy::cast_possible_truncation,
1194 clippy::cast_sign_loss
1195 )]
1196 async fn do_soft_compaction(
1197 &self,
1198 summ: &mut ContextSummarizationView<'_>,
1199 status: &(impl StatusSink + ?Sized),
1200 ) {
1201 status.send_status("soft compacting context...").await;
1202
1203 match &summ.context_manager.compression.pruning_strategy {
1205 zeph_config::PruningStrategy::Subgoal | zeph_config::PruningStrategy::SubgoalMig => {
1206 crate::summarization::scheduling::maybe_refresh_subgoal(summ);
1207 }
1208 _ => crate::summarization::scheduling::maybe_refresh_task_goal(summ),
1209 }
1210
1211 let applied = crate::summarization::deferred::apply_deferred_summaries(summ);
1213
1214 if applied > 0
1216 && summ
1217 .context_manager
1218 .compression
1219 .pruning_strategy
1220 .is_subgoal()
1221 {
1222 summ.subgoal_registry
1223 .rebuild_after_compaction(summ.messages, 0);
1224 }
1225
1226 let budget = summ
1228 .context_manager
1229 .budget
1230 .as_ref()
1231 .map_or(0, ContextBudget::max_tokens);
1232 let soft_threshold =
1233 (budget as f32 * summ.context_manager.soft_compaction_threshold) as usize;
1234 let cached = usize::try_from(*summ.cached_prompt_tokens).unwrap_or(usize::MAX);
1235 let min_to_free = cached.saturating_sub(soft_threshold);
1236 if min_to_free > 0 {
1237 crate::summarization::pruning::prune_tool_outputs(summ, min_to_free);
1238 }
1239
1240 status.send_status("").await;
1241 tracing::info!(
1242 cached_tokens = *summ.cached_prompt_tokens,
1243 soft_threshold,
1244 "soft compaction complete"
1245 );
1246 }
1247
1248 #[tracing::instrument(name = "agent_context.service.do_hard_compaction", skip_all, err)]
1250 #[allow(
1251 clippy::cast_precision_loss,
1252 clippy::cast_possible_truncation,
1253 clippy::cast_sign_loss
1254 )]
1255 async fn do_hard_compaction(
1256 &self,
1257 summ: &mut ContextSummarizationView<'_>,
1258 status: &(impl StatusSink + ?Sized),
1259 in_cooldown: bool,
1260 ) -> Result<(), ContextError> {
1261 use zeph_context::manager::CompactionState;
1262
1263 let turns_since_last = summ
1265 .context_manager
1266 .turns_since_last_hard_compaction()
1267 .map(|t| u32::try_from(t).unwrap_or(u32::MAX));
1268 summ.context_manager
1269 .set_turns_since_last_hard_compaction(Some(0));
1270 if let Some(metrics) = summ.metrics {
1271 metrics.record_hard_compaction(turns_since_last);
1272 }
1273
1274 if in_cooldown {
1275 tracing::debug!(
1276 turns_remaining = summ.context_manager.compaction_state().cooldown_remaining(),
1277 "hard compaction skipped: cooldown active"
1278 );
1279 return Ok(());
1280 }
1281
1282 let budget = summ
1283 .context_manager
1284 .budget
1285 .as_ref()
1286 .map_or(0, ContextBudget::max_tokens);
1287 let hard_threshold =
1288 (budget as f32 * summ.context_manager.hard_compaction_threshold) as usize;
1289 let cached = usize::try_from(*summ.cached_prompt_tokens).unwrap_or(usize::MAX);
1290 let min_to_free = cached.saturating_sub(hard_threshold);
1291
1292 status.send_status("compacting context...").await;
1293
1294 crate::summarization::deferred::apply_deferred_summaries(summ);
1296
1297 let tokens_before = *summ.cached_prompt_tokens;
1305 let freed = crate::summarization::pruning::prune_tool_outputs(summ, min_to_free);
1306 if freed >= min_to_free {
1307 tracing::info!(freed, "hard compaction: pruning sufficient");
1308 summ.context_manager
1309 .set_compaction_state(CompactionState::CompactedThisTurn {
1310 cooldown: summ.context_manager.compaction_cooldown_turns(),
1311 });
1312 if let Err(e) = crate::summarization::deferred::flush_deferred_summaries(summ).await {
1313 tracing::warn!(%e, "flush_deferred_summaries failed after hard compaction");
1314 }
1315 status.send_status("").await;
1316 return Ok(());
1317 }
1318
1319 let preserve_tail = summ.context_manager.compaction_preserve_tail;
1321 let compactable = summ.messages.len().saturating_sub(preserve_tail + 1);
1322 if compactable <= 1 {
1323 tracing::warn!(
1324 compactable,
1325 "hard compaction: too few messages, marking exhausted"
1326 );
1327 summ.context_manager
1328 .set_compaction_state(CompactionState::Exhausted { warned: false });
1329 status.send_status("").await;
1330 return Ok(());
1331 }
1332
1333 tracing::info!(
1335 min_to_free,
1336 "hard compaction: falling back to LLM summarization"
1337 );
1338 let outcome = crate::summarization::compaction::compact_context(summ, None).await?;
1339
1340 let freed_tokens = tokens_before.saturating_sub(*summ.cached_prompt_tokens);
1341
1342 if !outcome.is_compacted() || freed_tokens == 0 {
1343 tracing::warn!("hard compaction: no net reduction, marking exhausted");
1344 summ.context_manager
1345 .set_compaction_state(CompactionState::Exhausted { warned: false });
1346 status.send_status("").await;
1347 return Ok(());
1348 }
1349
1350 if matches!(
1351 summ.context_manager
1352 .compaction_tier(*summ.cached_prompt_tokens),
1353 zeph_context::manager::CompactionTier::Hard
1354 ) {
1355 tracing::warn!(
1356 freed_tokens,
1357 "hard compaction: still above hard threshold after compaction, marking exhausted"
1358 );
1359 summ.context_manager
1360 .set_compaction_state(CompactionState::Exhausted { warned: false });
1361 status.send_status("").await;
1362 return Ok(());
1363 }
1364
1365 summ.context_manager
1366 .set_compaction_state(CompactionState::CompactedThisTurn {
1367 cooldown: summ.context_manager.compaction_cooldown_turns(),
1368 });
1369
1370 if tokens_before > *summ.cached_prompt_tokens {
1371 tracing::info!(
1372 tokens_before,
1373 tokens_after = *summ.cached_prompt_tokens,
1374 saved = freed_tokens,
1375 "context compaction complete"
1376 );
1377 }
1378
1379 status.send_status("").await;
1380 Ok(())
1381 }
1382
1383 #[tracing::instrument(name = "agent_context.service.maybe_summarize_tool_pair", skip_all)]
1389 pub async fn maybe_summarize_tool_pair(
1390 &self,
1391 summ: &mut ContextSummarizationView<'_>,
1392 providers: &ProviderHandles,
1393 ) {
1394 crate::summarization::deferred::maybe_summarize_tool_pair(
1395 summ,
1396 providers,
1397 &TxStatusSink(summ.status_tx.clone()),
1398 )
1399 .await;
1400 }
1401
1402 #[must_use]
1407 pub fn apply_deferred_summaries(&self, summ: &mut ContextSummarizationView<'_>) -> usize {
1408 crate::summarization::deferred::apply_deferred_summaries(summ)
1409 }
1410
1411 #[tracing::instrument(name = "agent_context.service.flush_deferred_summaries", skip_all)]
1416 pub async fn flush_deferred_summaries(&self, summ: &mut ContextSummarizationView<'_>) {
1417 if let Err(e) = crate::summarization::deferred::flush_deferred_summaries(summ).await {
1418 tracing::warn!(%e, "flush_deferred_summaries failed");
1419 }
1420 }
1421
1422 pub fn maybe_apply_deferred_summaries(&self, summ: &mut ContextSummarizationView<'_>) {
1428 crate::summarization::deferred::maybe_apply_deferred_summaries(summ);
1429 }
1430
1431 #[tracing::instrument(name = "agent_context.service.compact_context", skip_all, err)]
1446 pub async fn compact_context(
1447 &self,
1448 summ: &mut ContextSummarizationView<'_>,
1449 max_summary_tokens: Option<usize>,
1450 ) -> Result<crate::state::CompactionOutcome, crate::error::ContextError> {
1451 crate::summarization::compaction::compact_context(summ, max_summary_tokens).await
1452 }
1453
1454 pub fn maybe_soft_compact_mid_iteration(&self, summ: &mut ContextSummarizationView<'_>) {
1460 crate::summarization::scheduling::maybe_soft_compact_mid_iteration(summ);
1461 }
1462
1463 #[tracing::instrument(name = "agent_context.service.maybe_proactive_compress", skip_all)]
1469 pub async fn maybe_proactive_compress(
1470 &self,
1471 summ: &mut ContextSummarizationView<'_>,
1472 status: &(impl StatusSink + ?Sized),
1473 ) {
1474 let Some((_threshold, max_summary_tokens)) = summ
1475 .context_manager
1476 .should_proactively_compress(*summ.cached_prompt_tokens)
1477 else {
1478 return;
1479 };
1480
1481 if summ.server_compaction_active {
1482 let budget = summ
1483 .context_manager
1484 .budget
1485 .as_ref()
1486 .map_or(0, ContextBudget::max_tokens);
1487 if budget > 0 {
1488 let fallback = (budget * 95 / 100) as u64;
1489 if *summ.cached_prompt_tokens <= fallback {
1490 return;
1491 }
1492 tracing::warn!(
1493 cached_prompt_tokens = *summ.cached_prompt_tokens,
1494 fallback_threshold = fallback,
1495 "server compaction active but context at 95%+ — falling back to proactive"
1496 );
1497 } else {
1498 return;
1499 }
1500 }
1501
1502 status.send_status("compressing context...").await;
1503 tracing::info!(
1504 max_summary_tokens,
1505 cached_tokens = *summ.cached_prompt_tokens,
1506 "proactive compression triggered"
1507 );
1508
1509 match crate::summarization::compaction::compact_context(summ, Some(max_summary_tokens))
1510 .await
1511 {
1512 Ok(outcome) if outcome.is_compacted() => {
1513 summ.context_manager.set_compaction_state(
1514 zeph_context::manager::CompactionState::CompactedThisTurn { cooldown: 0 },
1515 );
1516 tracing::info!("proactive compression complete");
1517 }
1518 Ok(_) => {}
1519 Err(e) => tracing::warn!(%e, "proactive compression failed"),
1520 }
1521
1522 status.send_status("").await;
1523 }
1524
1525 pub fn maybe_refresh_task_goal(&self, summ: &mut ContextSummarizationView<'_>) {
1531 crate::summarization::scheduling::maybe_refresh_task_goal(summ);
1532 }
1533
1534 pub fn maybe_refresh_subgoal(&self, summ: &mut ContextSummarizationView<'_>) {
1539 crate::summarization::scheduling::maybe_refresh_subgoal(summ);
1540 }
1541}
1542
1543struct TxStatusSink(Option<tokio::sync::mpsc::UnboundedSender<String>>);
1549
1550impl StatusSink for TxStatusSink {
1551 fn send_status(&self, msg: &str) -> impl std::future::Future<Output = ()> + Send + '_ {
1552 if let Some(ref tx) = self.0 {
1553 let _ = tx.send(msg.to_owned());
1554 }
1555 std::future::ready(())
1556 }
1557}
1558
1559pub(crate) fn recompute_prompt_tokens(window: &mut MessageWindowView<'_>) {
1566 *window.cached_prompt_tokens = window
1567 .messages
1568 .iter()
1569 .map(|m| window.token_counter.count_message_tokens(m) as u64)
1570 .sum();
1571}
1572
1573fn fidelity_provider_pair<'a>(
1579 semantic: Option<&'a std::sync::Arc<zeph_llm::any::AnyProvider>>,
1580 compress: Option<&'a std::sync::Arc<zeph_llm::any::AnyProvider>>,
1581) -> (
1582 Option<&'a dyn zeph_llm::LlmProviderDyn>,
1583 Option<&'a dyn zeph_llm::LlmProviderDyn>,
1584) {
1585 let embed_provider = semantic
1586 .map(std::sync::Arc::as_ref)
1587 .map(|p| p as &dyn zeph_llm::LlmProviderDyn);
1588 let compress_provider = compress
1589 .map(std::sync::Arc::as_ref)
1590 .map(|p| p as &dyn zeph_llm::LlmProviderDyn);
1591 (embed_provider, compress_provider)
1592}
1593
1594#[tracing::instrument(name = "agent_context.service.persist_fidelity_tags", skip_all)]
1603async fn persist_fidelity_tags(
1604 messages: &[zeph_llm::provider::Message],
1605 memory: Option<&zeph_memory::semantic::SemanticMemory>,
1606) {
1607 let Some(mem) = memory else { return };
1608 let updates: Vec<(zeph_memory::MessageId, u8)> = messages
1609 .iter()
1610 .filter_map(|m| {
1611 let db_id = m.metadata.db_id?;
1612 let tag = m.metadata.fidelity_tag?;
1613 Some((zeph_memory::MessageId(db_id), tag as u8))
1614 })
1615 .collect();
1616 if updates.is_empty() {
1617 return;
1618 }
1619 if let Err(e) = mem.sqlite().update_fidelity_tags(&updates).await {
1620 tracing::warn!(
1621 count = updates.len(),
1622 error = %e,
1623 "failed to persist fidelity tags; floor invariant will not apply next turn"
1624 );
1625 }
1626}
1627
1628fn recompute_prompt_tokens_summ(summ: &mut crate::state::ContextSummarizationView<'_>) {
1632 *summ.cached_prompt_tokens = summ
1633 .messages
1634 .iter()
1635 .map(|m| summ.token_counter.count_message_tokens(m) as u64)
1636 .sum();
1637}
1638
1639pub(crate) fn remove_by_prefix(
1645 messages: &mut Vec<zeph_llm::provider::Message>,
1646 role: Role,
1647 prefix: &str,
1648) {
1649 messages.retain(|m| m.role != role || !m.content.starts_with(prefix));
1650}
1651
1652pub(crate) fn remove_by_part_or_prefix(
1661 messages: &mut Vec<zeph_llm::provider::Message>,
1662 prefix: &str,
1663 part_matches: impl Fn(&MessagePart) -> bool,
1664) {
1665 messages.retain(|m| {
1666 if m.role == Role::User {
1669 return !m.content.starts_with(prefix);
1670 }
1671 if m.role != Role::System {
1672 return true;
1673 }
1674 if m.parts.first().is_some_and(&part_matches) {
1675 return false;
1676 }
1677 !m.content.starts_with(prefix)
1678 });
1679}
1680
1681#[cfg(test)]
1682mod tests {
1683 use std::collections::HashSet;
1684 use std::sync::Arc;
1685
1686 use zeph_llm::provider::{Message, MessagePart, Role};
1687 use zeph_memory::TokenCounter;
1688
1689 use super::*;
1690 use crate::helpers::{GRAPH_FACTS_PREFIX, RECALL_PREFIX, SUMMARY_PREFIX};
1691 use crate::state::MessageWindowView;
1692
1693 fn make_counter() -> Arc<TokenCounter> {
1694 Arc::new(TokenCounter::default())
1695 }
1696
1697 fn make_window<'a>(
1698 messages: &'a mut Vec<Message>,
1699 cached: &'a mut u64,
1700 completed: &'a mut HashSet<String>,
1701 ) -> MessageWindowView<'a> {
1702 let last = Box::leak(Box::new(None::<i64>));
1703 let deferred_hide = Box::leak(Box::new(Vec::<i64>::new()));
1704 let deferred_summ = Box::leak(Box::new(Vec::<String>::new()));
1705 MessageWindowView {
1706 messages,
1707 last_persisted_message_id: last,
1708 deferred_db_hide_ids: deferred_hide,
1709 deferred_db_summaries: deferred_summ,
1710 cached_prompt_tokens: cached,
1711 token_counter: make_counter(),
1712 completed_tool_ids: completed,
1713 }
1714 }
1715
1716 fn sys(text: &str) -> Message {
1717 Message::from_legacy(Role::System, text)
1718 }
1719
1720 fn user(text: &str) -> Message {
1721 Message::from_legacy(Role::User, text)
1722 }
1723
1724 fn assistant(text: &str) -> Message {
1725 Message::from_legacy(Role::Assistant, text)
1726 }
1727
1728 #[test]
1729 fn clear_history_keeps_system_prompt() {
1730 let mut msgs = vec![sys("system"), user("hello"), assistant("hi")];
1731 let mut cached = 0u64;
1732 let mut completed = HashSet::new();
1733 completed.insert("tool_1".to_owned());
1734 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1735
1736 ContextService::new().clear_history(&mut window);
1737
1738 assert_eq!(window.messages.len(), 1);
1739 assert_eq!(window.messages[0].content, "system");
1740 assert!(
1741 window.completed_tool_ids.is_empty(),
1742 "completed_tool_ids must be cleared"
1743 );
1744 }
1745
1746 #[test]
1747 fn clear_history_empty_messages_is_noop() {
1748 let mut msgs: Vec<Message> = vec![];
1749 let mut cached = 0u64;
1750 let mut completed = HashSet::new();
1751 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1752
1753 ContextService::new().clear_history(&mut window);
1754
1755 assert!(window.messages.is_empty());
1756 }
1757
1758 #[test]
1759 fn remove_recall_messages_removes_by_prefix() {
1760 let mut msgs = vec![
1761 sys("system"),
1762 sys(&format!("{RECALL_PREFIX}some recalled text")),
1763 user("hello"),
1764 ];
1765 let mut cached = 0u64;
1766 let mut completed = HashSet::new();
1767 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1768
1769 ContextService::new().remove_recall_messages(&mut window);
1770
1771 assert_eq!(window.messages.len(), 2);
1772 assert!(
1773 window
1774 .messages
1775 .iter()
1776 .all(|m| !m.content.starts_with(RECALL_PREFIX))
1777 );
1778 }
1779
1780 #[test]
1783 fn remove_recall_messages_removes_user_role_recall() {
1784 let mut msgs = vec![
1785 sys("system"),
1786 user(&format!("{RECALL_PREFIX}recalled via tiered path")),
1787 user("real user message"),
1788 ];
1789 let mut cached = 0u64;
1790 let mut completed = HashSet::new();
1791 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1792
1793 ContextService::new().remove_recall_messages(&mut window);
1794
1795 assert_eq!(
1796 window.messages.len(),
1797 2,
1798 "Role::User recall message must be removed"
1799 );
1800 assert!(
1801 window
1802 .messages
1803 .iter()
1804 .all(|m| !m.content.starts_with(RECALL_PREFIX)),
1805 "no message with RECALL_PREFIX must remain"
1806 );
1807 assert!(
1808 window
1809 .messages
1810 .iter()
1811 .any(|m| m.content == "real user message"),
1812 "non-recall user message must survive"
1813 );
1814 }
1815
1816 #[test]
1817 fn remove_graph_facts_messages_removes_matching() {
1818 let mut msgs = vec![
1819 sys("system"),
1820 sys(&format!("{GRAPH_FACTS_PREFIX}fact1")),
1821 user("hello"),
1822 ];
1823 let mut cached = 0u64;
1824 let mut completed = HashSet::new();
1825 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1826
1827 ContextService::new().remove_graph_facts_messages(&mut window);
1828
1829 assert_eq!(window.messages.len(), 2);
1830 }
1831
1832 #[test]
1833 fn remove_summary_messages_removes_by_part() {
1834 let mut msgs = vec![
1835 sys("system"),
1836 Message::from_parts(
1837 Role::System,
1838 vec![MessagePart::Summary {
1839 text: format!("{SUMMARY_PREFIX}old summary"),
1840 }],
1841 ),
1842 user("hello"),
1843 ];
1844 let mut cached = 0u64;
1845 let mut completed = HashSet::new();
1846 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1847
1848 ContextService::new().remove_summary_messages(&mut window);
1849
1850 assert_eq!(window.messages.len(), 2);
1851 }
1852
1853 #[test]
1854 fn trim_messages_to_budget_zero_is_noop() {
1855 let mut msgs = vec![sys("system"), user("a"), assistant("b"), user("c")];
1856 let original_len = msgs.len();
1857 let mut cached = 0u64;
1858 let mut completed = HashSet::new();
1859 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1860
1861 ContextService::new().trim_messages_to_budget(&mut window, 0);
1862
1863 assert_eq!(window.messages.len(), original_len);
1864 }
1865
1866 #[test]
1867 fn trim_messages_to_budget_keeps_recent() {
1868 let mut msgs = vec![
1870 sys("system"),
1871 user("message 1"),
1872 assistant("reply 1"),
1873 user("message 2"),
1874 ];
1875 let mut cached = 0u64;
1876 let mut completed = HashSet::new();
1877 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1878
1879 ContextService::new().trim_messages_to_budget(&mut window, 1);
1881
1882 assert!(
1884 window.messages.len() < 4,
1885 "trim should remove some messages"
1886 );
1887 assert_eq!(
1888 window.messages[0].role,
1889 Role::System,
1890 "system prompt must survive trim"
1891 );
1892 }
1893
1894 mod inserted_count_tests {
1897 use parking_lot::RwLock;
1898 use std::borrow::Cow;
1899 use std::collections::HashSet;
1900 use std::sync::Arc;
1901
1902 use zeph_common::SecurityEventCategory;
1903 use zeph_config::memory::TieredRetrievalConfig;
1904 use zeph_config::{
1905 ContextFormat, ContextStrategy, DocumentConfig, GraphConfig, PersonaConfig,
1906 ReasoningConfig, TrajectoryConfig, TreeConfig,
1907 };
1908 use zeph_context::assembler::PreparedContext;
1909 use zeph_context::manager::ContextManager;
1910 use zeph_llm::provider::{Message, MessageMetadata, Role};
1911 use zeph_memory::TokenCounter;
1912 use zeph_sanitizer::ContentIsolationConfig;
1913 use zeph_sanitizer::ContentSanitizer;
1914 use zeph_skills::registry::SkillRegistry;
1915
1916 use super::super::*;
1917 use crate::state::{
1918 ContextAssemblyView, MessageWindowView, MetricsCounters, SecurityEventSink,
1919 };
1920
1921 fn make_task_supervisor() -> Arc<zeph_common::TaskSupervisor> {
1922 Arc::new(zeph_common::TaskSupervisor::new(
1923 tokio_util::sync::CancellationToken::new(),
1924 ))
1925 }
1926
1927 struct NoopSink;
1928 impl SecurityEventSink for NoopSink {
1929 fn push(&mut self, _: SecurityEventCategory, _: &'static str, _: String) {}
1930 }
1931
1932 fn make_counter() -> Arc<TokenCounter> {
1933 Arc::new(TokenCounter::default())
1934 }
1935
1936 fn make_window<'a>(
1937 messages: &'a mut Vec<Message>,
1938 cached: &'a mut u64,
1939 completed: &'a mut HashSet<String>,
1940 ) -> MessageWindowView<'a> {
1941 let last = Box::leak(Box::new(None::<i64>));
1942 let deferred_hide = Box::leak(Box::new(Vec::<i64>::new()));
1943 let deferred_summ = Box::leak(Box::new(Vec::<String>::new()));
1944 MessageWindowView {
1945 messages,
1946 last_persisted_message_id: last,
1947 deferred_db_hide_ids: deferred_hide,
1948 deferred_db_summaries: deferred_summ,
1949 cached_prompt_tokens: cached,
1950 token_counter: make_counter(),
1951 completed_tool_ids: completed,
1952 }
1953 }
1954
1955 fn mem_msg(content: &str) -> Message {
1956 Message {
1957 role: Role::User,
1958 content: content.to_string(),
1959 parts: vec![],
1960 metadata: MessageMetadata::default(),
1961 }
1962 }
1963
1964 fn scrub_noop(s: &str) -> Cow<'_, str> {
1965 Cow::Borrowed(s)
1966 }
1967
1968 #[tokio::test]
1969 async fn inserted_count_incremented_for_all_paths() {
1970 let mut msgs = vec![
1973 Message::from_legacy(Role::System, "system"),
1974 Message::from_legacy(Role::User, "user turn"),
1975 ];
1976 let mut cached = 0u64;
1977 let mut completed = HashSet::new();
1978 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1979
1980 let sanitizer = ContentSanitizer::new(&ContentIsolationConfig::default());
1981 let mut ctx_mgr = ContextManager::new();
1982 let mut sink = NoopSink;
1983 let mut last_confidence = None::<f32>;
1984 let mut last_skills_prompt = String::new();
1985 let mut active_skill_names = Vec::new();
1986 let registry = Arc::new(RwLock::new(SkillRegistry::default()));
1987
1988 let mut view = ContextAssemblyView {
1989 memory: None,
1990 conversation_id: None,
1991 recall_limit: 10,
1992 cross_session_score_threshold: 0.5,
1993 context_format: ContextFormat::default(),
1994 last_recall_confidence: &mut last_confidence,
1995 context_strategy: ContextStrategy::default(),
1996 crossover_turn_threshold: 0,
1997 cached_session_digest: None,
1998 digest_enabled: false, graph_config: GraphConfig::default(),
2000 document_config: DocumentConfig::default(),
2001 persona_config: PersonaConfig::default(),
2002 trajectory_config: TrajectoryConfig::default(),
2003 reasoning_config: ReasoningConfig::default(),
2004 memcot_config: zeph_config::MemCotConfig::default(),
2005 memcot_state: None,
2006 tree_config: TreeConfig::default(),
2007 last_skills_prompt: &mut last_skills_prompt,
2008 active_skill_names: &mut active_skill_names,
2009 skill_registry: registry,
2010 skill_paths: &[],
2011 correction_config: None,
2012 sidequest_turn_counter: 0,
2013 proactive_explorer: None,
2014 sanitizer: &sanitizer,
2015 quarantine_summarizer: None,
2016 context_manager: &mut ctx_mgr,
2017 token_counter: make_counter(),
2018 metrics: MetricsCounters::default(),
2019 security_events: &mut sink,
2020 cached_prompt_tokens: 0,
2021 redact_credentials: false,
2022 channel_skills: &[],
2023 scrub: scrub_noop,
2024 tiered_retrieval_config: TieredRetrievalConfig {
2025 enabled: false,
2026 ..TieredRetrievalConfig::default()
2027 },
2028 tiered_retrieval_classifier: None,
2029 tiered_retrieval_validator: None,
2030 fidelity_config: None,
2031 fidelity_semantic_provider: None,
2032 fidelity_compress_provider: None,
2033 planned_next_tools: &[],
2034 status_tx: None,
2035 task_supervisor: make_task_supervisor(),
2036 };
2037
2038 let prepared = PreparedContext {
2040 graph_facts: Some(mem_msg("graph_facts")),
2041 doc_rag: Some(mem_msg("doc_rag")),
2042 corrections: Some(mem_msg("corrections")),
2043 recall: Some(mem_msg("recall")),
2044 recall_confidence: Some(0.9),
2045 cross_session: Some(mem_msg("cross_session")),
2046 summaries: Some(mem_msg("summaries")),
2047 code_context: None, persona_facts: Some(mem_msg("persona_facts")),
2049 trajectory_hints: Some(mem_msg("trajectory_hints")),
2050 tree_memory: Some(mem_msg("tree_memory")),
2051 reasoning_hints: Some(mem_msg("reasoning_hints")),
2052 memory_first: false,
2053 recent_history_budget: 100_000,
2054 background_tasks: vec![],
2055 };
2056
2057 let (_delta, inserted_count) = ContextService::new()
2058 .apply_prepared_context(&mut window, &mut view, prepared)
2059 .await;
2060
2061 assert_eq!(
2064 inserted_count, 10,
2065 "all 10 message-carrying PreparedContext fields must increment inserted_count"
2066 );
2067 }
2068 }
2069
2070 mod inject_semantic_recall_tests {
2071 use parking_lot::RwLock;
2072 use std::borrow::Cow;
2073 use std::collections::HashSet;
2074 use std::sync::Arc;
2075
2076 use zeph_config::memory::TieredRetrievalConfig;
2077 use zeph_config::{
2078 ContextFormat, ContextStrategy, DocumentConfig, GraphConfig, PersonaConfig,
2079 ReasoningConfig, TrajectoryConfig, TreeConfig,
2080 };
2081 use zeph_context::manager::ContextManager;
2082 use zeph_llm::provider::Message;
2083 use zeph_memory::TokenCounter;
2084 use zeph_sanitizer::ContentIsolationConfig;
2085 use zeph_sanitizer::ContentSanitizer;
2086 use zeph_skills::registry::SkillRegistry;
2087
2088 use zeph_common::SecurityEventCategory;
2089
2090 use super::super::*;
2091 use crate::helpers::RECALL_PREFIX;
2092 use crate::state::{
2093 ContextAssemblyView, MessageWindowView, MetricsCounters, SecurityEventSink,
2094 };
2095
2096 fn make_task_supervisor() -> Arc<zeph_common::TaskSupervisor> {
2097 Arc::new(zeph_common::TaskSupervisor::new(
2098 tokio_util::sync::CancellationToken::new(),
2099 ))
2100 }
2101
2102 struct NoopSink;
2103 impl SecurityEventSink for NoopSink {
2104 fn push(&mut self, _: SecurityEventCategory, _: &'static str, _: String) {}
2105 }
2106
2107 fn make_counter() -> Arc<TokenCounter> {
2108 Arc::new(TokenCounter::default())
2109 }
2110
2111 fn make_window<'a>(
2112 messages: &'a mut Vec<Message>,
2113 cached: &'a mut u64,
2114 completed: &'a mut HashSet<String>,
2115 ) -> MessageWindowView<'a> {
2116 let last = Box::leak(Box::new(None::<i64>));
2117 let deferred_hide = Box::leak(Box::new(Vec::<i64>::new()));
2118 let deferred_summ = Box::leak(Box::new(Vec::<String>::new()));
2119 MessageWindowView {
2120 messages,
2121 last_persisted_message_id: last,
2122 deferred_db_hide_ids: deferred_hide,
2123 deferred_db_summaries: deferred_summ,
2124 cached_prompt_tokens: cached,
2125 token_counter: make_counter(),
2126 completed_tool_ids: completed,
2127 }
2128 }
2129
2130 fn scrub_noop(s: &str) -> Cow<'_, str> {
2131 Cow::Borrowed(s)
2132 }
2133
2134 #[tokio::test]
2135 async fn tiered_recall_disabled_uses_flat_path() {
2136 let mut msgs: Vec<Message> = vec![];
2139 let mut cached = 0u64;
2140 let mut completed = HashSet::new();
2141 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2142
2143 let sanitizer = ContentSanitizer::new(&ContentIsolationConfig::default());
2144 let mut ctx_mgr = ContextManager::new();
2145 let mut sink = NoopSink;
2146 let mut last_confidence = None::<f32>;
2147 let mut last_skills_prompt = String::new();
2148 let mut active_skill_names = Vec::new();
2149 let registry = Arc::new(RwLock::new(SkillRegistry::default()));
2150
2151 let view = ContextAssemblyView {
2152 memory: None,
2153 conversation_id: None,
2154 recall_limit: 10,
2155 cross_session_score_threshold: 0.5,
2156 context_format: ContextFormat::default(),
2157 last_recall_confidence: &mut last_confidence,
2158 context_strategy: ContextStrategy::default(),
2159 crossover_turn_threshold: 0,
2160 cached_session_digest: None,
2161 digest_enabled: false,
2162 graph_config: GraphConfig::default(),
2163 document_config: DocumentConfig::default(),
2164 persona_config: PersonaConfig::default(),
2165 trajectory_config: TrajectoryConfig::default(),
2166 reasoning_config: ReasoningConfig::default(),
2167 memcot_config: zeph_config::MemCotConfig::default(),
2168 memcot_state: None,
2169 tree_config: TreeConfig::default(),
2170 last_skills_prompt: &mut last_skills_prompt,
2171 active_skill_names: &mut active_skill_names,
2172 skill_registry: registry,
2173 skill_paths: &[],
2174 correction_config: None,
2175 sidequest_turn_counter: 0,
2176 proactive_explorer: None,
2177 sanitizer: &sanitizer,
2178 quarantine_summarizer: None,
2179 context_manager: &mut ctx_mgr,
2180 token_counter: make_counter(),
2181 metrics: MetricsCounters::default(),
2182 security_events: &mut sink,
2183 cached_prompt_tokens: 0,
2184 redact_credentials: false,
2185 channel_skills: &[],
2186 scrub: scrub_noop,
2187 tiered_retrieval_config: TieredRetrievalConfig {
2188 enabled: false,
2189 ..TieredRetrievalConfig::default()
2190 },
2191 tiered_retrieval_classifier: None,
2192 tiered_retrieval_validator: None,
2193 fidelity_config: None,
2194 fidelity_semantic_provider: None,
2195 fidelity_compress_provider: None,
2196 planned_next_tools: &[],
2197 status_tx: None,
2198 task_supervisor: make_task_supervisor(),
2199 };
2200
2201 let result = ContextService::new()
2202 .inject_semantic_recall("test query", 1000, &mut window, &view)
2203 .await;
2204
2205 assert!(result.is_ok(), "disabled tiered recall must return Ok(())");
2206 assert!(
2207 window
2208 .messages
2209 .iter()
2210 .all(|m| !m.content.starts_with(RECALL_PREFIX)),
2211 "no recall message must be injected when memory is None"
2212 );
2213 }
2214
2215 #[tokio::test]
2216 async fn tiered_recall_enabled_no_memory_returns_ok() {
2217 let mut msgs: Vec<Message> = vec![];
2220 let mut cached = 0u64;
2221 let mut completed = HashSet::new();
2222 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2223
2224 let sanitizer = ContentSanitizer::new(&ContentIsolationConfig::default());
2225 let mut ctx_mgr = ContextManager::new();
2226 let mut sink = NoopSink;
2227 let mut last_confidence = None::<f32>;
2228 let mut last_skills_prompt = String::new();
2229 let mut active_skill_names = Vec::new();
2230 let registry = Arc::new(RwLock::new(SkillRegistry::default()));
2231
2232 let view = ContextAssemblyView {
2233 memory: None,
2234 conversation_id: None,
2235 recall_limit: 10,
2236 cross_session_score_threshold: 0.5,
2237 context_format: ContextFormat::default(),
2238 last_recall_confidence: &mut last_confidence,
2239 context_strategy: ContextStrategy::default(),
2240 crossover_turn_threshold: 0,
2241 cached_session_digest: None,
2242 digest_enabled: false,
2243 graph_config: GraphConfig::default(),
2244 document_config: DocumentConfig::default(),
2245 persona_config: PersonaConfig::default(),
2246 trajectory_config: TrajectoryConfig::default(),
2247 reasoning_config: ReasoningConfig::default(),
2248 memcot_config: zeph_config::MemCotConfig::default(),
2249 memcot_state: None,
2250 tree_config: TreeConfig::default(),
2251 last_skills_prompt: &mut last_skills_prompt,
2252 active_skill_names: &mut active_skill_names,
2253 skill_registry: registry,
2254 skill_paths: &[],
2255 correction_config: None,
2256 sidequest_turn_counter: 0,
2257 proactive_explorer: None,
2258 sanitizer: &sanitizer,
2259 quarantine_summarizer: None,
2260 context_manager: &mut ctx_mgr,
2261 token_counter: make_counter(),
2262 metrics: MetricsCounters::default(),
2263 security_events: &mut sink,
2264 cached_prompt_tokens: 0,
2265 redact_credentials: false,
2266 channel_skills: &[],
2267 scrub: scrub_noop,
2268 tiered_retrieval_config: TieredRetrievalConfig {
2269 enabled: true,
2270 ..TieredRetrievalConfig::default()
2271 },
2272 tiered_retrieval_classifier: None,
2273 tiered_retrieval_validator: None,
2274 fidelity_config: None,
2275 fidelity_semantic_provider: None,
2276 fidelity_compress_provider: None,
2277 planned_next_tools: &[],
2278 status_tx: None,
2279 task_supervisor: make_task_supervisor(),
2280 };
2281
2282 let result = ContextService::new()
2283 .inject_semantic_recall("test query", 1000, &mut window, &view)
2284 .await;
2285
2286 assert!(
2287 result.is_ok(),
2288 "enabled tiered recall with no memory must return Ok(())"
2289 );
2290 assert!(
2291 window.messages.is_empty(),
2292 "no recall message must be injected when memory is None"
2293 );
2294 }
2295
2296 #[tokio::test]
2301 async fn prepare_context_tiered_enabled_no_budget_returns_default() {
2302 let mut msgs: Vec<zeph_llm::provider::Message> = vec![];
2303 let mut cached = 0u64;
2304 let mut completed = HashSet::new();
2305 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2306
2307 let sanitizer = zeph_sanitizer::ContentSanitizer::new(
2308 &zeph_sanitizer::ContentIsolationConfig::default(),
2309 );
2310 let mut ctx_mgr = zeph_context::manager::ContextManager::new();
2311 assert!(ctx_mgr.budget.is_none());
2313
2314 let mut sink = NoopSink;
2315 let mut last_confidence = None::<f32>;
2316 let mut last_skills_prompt = String::new();
2317 let mut active_skill_names = Vec::new();
2318 let registry = Arc::new(RwLock::new(zeph_skills::registry::SkillRegistry::default()));
2319
2320 let mut view = ContextAssemblyView {
2321 memory: None,
2322 conversation_id: None,
2323 recall_limit: 10,
2324 cross_session_score_threshold: 0.5,
2325 context_format: ContextFormat::default(),
2326 last_recall_confidence: &mut last_confidence,
2327 context_strategy: ContextStrategy::default(),
2328 crossover_turn_threshold: 0,
2329 cached_session_digest: None,
2330 digest_enabled: false,
2331 graph_config: GraphConfig::default(),
2332 document_config: DocumentConfig::default(),
2333 persona_config: PersonaConfig::default(),
2334 trajectory_config: TrajectoryConfig::default(),
2335 reasoning_config: ReasoningConfig::default(),
2336 memcot_config: zeph_config::MemCotConfig::default(),
2337 memcot_state: None,
2338 tree_config: TreeConfig::default(),
2339 last_skills_prompt: &mut last_skills_prompt,
2340 active_skill_names: &mut active_skill_names,
2341 skill_registry: registry,
2342 skill_paths: &[],
2343 correction_config: None,
2344 sidequest_turn_counter: 0,
2345 proactive_explorer: None,
2346 sanitizer: &sanitizer,
2347 quarantine_summarizer: None,
2348 context_manager: &mut ctx_mgr,
2349 token_counter: make_counter(),
2350 metrics: MetricsCounters::default(),
2351 security_events: &mut sink,
2352 cached_prompt_tokens: 0,
2353 redact_credentials: false,
2354 channel_skills: &[],
2355 scrub: scrub_noop,
2356 tiered_retrieval_config: TieredRetrievalConfig {
2357 enabled: true,
2358 ..TieredRetrievalConfig::default()
2359 },
2360 tiered_retrieval_classifier: None,
2361 tiered_retrieval_validator: None,
2362 fidelity_config: None,
2363 fidelity_semantic_provider: None,
2364 fidelity_compress_provider: None,
2365 planned_next_tools: &[],
2366 status_tx: None,
2367 task_supervisor: make_task_supervisor(),
2368 };
2369
2370 let result = ContextService::new()
2371 .prepare_context("test query", &mut window, &mut view)
2372 .await;
2373
2374 assert!(
2375 result.is_ok(),
2376 "prepare_context with tiered enabled and no budget must return Ok"
2377 );
2378 }
2379
2380 #[tokio::test]
2383 async fn inject_semantic_recall_bare_no_memory_returns_ok() {
2384 use zeph_config::memory::TieredRetrievalConfig;
2385
2386 let mut msgs: Vec<Message> = vec![];
2387 let mut cached = 0u64;
2388 let mut completed = HashSet::new();
2389 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2390
2391 let tiered_config = TieredRetrievalConfig {
2392 enabled: true,
2393 ..TieredRetrievalConfig::default()
2394 };
2395 let params = SemanticRecallParams {
2396 query: "test query",
2397 token_budget: 1000,
2398 recall_limit: 10,
2399 context_format: zeph_config::ContextFormat::default(),
2400 conversation_id: None,
2401 tiered_classifier: None,
2402 tiered_validator: None,
2403 tiered_config: &tiered_config,
2404 };
2405 let result = ContextService::new()
2406 .inject_semantic_recall_bare(params, &mut window, None)
2407 .await;
2408
2409 assert!(
2410 result.is_ok(),
2411 "inject_semantic_recall_bare with memory=None must return Ok(())"
2412 );
2413 assert!(
2414 window.messages.is_empty(),
2415 "no recall message must be injected when memory is None"
2416 );
2417 }
2418 }
2419
2420 mod prepare_context_proactive_explore_tests {
2427 use parking_lot::RwLock;
2428 use std::borrow::Cow;
2429 use std::collections::HashSet;
2430 use std::sync::Arc;
2431
2432 use zeph_config::memory::TieredRetrievalConfig;
2433 use zeph_config::{
2434 ContextFormat, ContextStrategy, DocumentConfig, GraphConfig, PersonaConfig,
2435 ReasoningConfig, TrajectoryConfig, TreeConfig,
2436 };
2437 use zeph_context::budget::ContextBudget;
2438 use zeph_context::manager::ContextManager;
2439 use zeph_llm::any::AnyProvider;
2440 use zeph_llm::mock::MockProvider;
2441 use zeph_llm::provider::Message;
2442 use zeph_memory::TokenCounter;
2443 use zeph_sanitizer::ContentIsolationConfig;
2444 use zeph_sanitizer::ContentSanitizer;
2445 use zeph_skills::generator::SkillGenerator;
2446 use zeph_skills::proactive::ProactiveExplorer;
2447 use zeph_skills::registry::SkillRegistry;
2448
2449 use zeph_common::SecurityEventCategory;
2450
2451 use super::super::*;
2452 use crate::state::{
2453 ContextAssemblyView, MessageWindowView, MetricsCounters, SecurityEventSink,
2454 };
2455
2456 fn make_task_supervisor() -> Arc<zeph_common::TaskSupervisor> {
2457 Arc::new(zeph_common::TaskSupervisor::new(
2458 tokio_util::sync::CancellationToken::new(),
2459 ))
2460 }
2461
2462 struct NoopSink;
2463 impl SecurityEventSink for NoopSink {
2464 fn push(&mut self, _: SecurityEventCategory, _: &'static str, _: String) {}
2465 }
2466
2467 fn make_counter() -> Arc<TokenCounter> {
2468 Arc::new(TokenCounter::default())
2469 }
2470
2471 fn make_window<'a>(
2472 messages: &'a mut Vec<Message>,
2473 cached: &'a mut u64,
2474 completed: &'a mut HashSet<String>,
2475 ) -> MessageWindowView<'a> {
2476 let last = Box::leak(Box::new(None::<i64>));
2477 let deferred_hide = Box::leak(Box::new(Vec::<i64>::new()));
2478 let deferred_summ = Box::leak(Box::new(Vec::<String>::new()));
2479 MessageWindowView {
2480 messages,
2481 last_persisted_message_id: last,
2482 deferred_db_hide_ids: deferred_hide,
2483 deferred_db_summaries: deferred_summ,
2484 cached_prompt_tokens: cached,
2485 token_counter: make_counter(),
2486 completed_tool_ids: completed,
2487 }
2488 }
2489
2490 fn scrub_noop(s: &str) -> Cow<'_, str> {
2491 Cow::Borrowed(s)
2492 }
2493
2494 fn mock_skill_content(name: &str) -> String {
2495 format!(
2496 "---\nname: {name}\ndescription: Test world-knowledge skill for {name}.\n---\n\n## Usage\n\nDetails.\n"
2497 )
2498 }
2499
2500 struct Fixture {
2505 registry: Arc<RwLock<SkillRegistry>>,
2506 skill_paths: Vec<std::path::PathBuf>,
2507 hub_dir: tempfile::TempDir,
2508 _skills_dir_guard: tempfile::TempDir,
2511 explorer: Arc<ProactiveExplorer>,
2512 }
2513
2514 fn setup_fixture(mock_provider: MockProvider) -> Fixture {
2515 let skills_dir = tempfile::tempdir().expect("create skills tempdir");
2516 let hub_dir = tempfile::tempdir().expect("create hub tempdir");
2517 let skill_paths = vec![skills_dir.path().to_path_buf()];
2518
2519 let registry = Arc::new(RwLock::new(
2520 SkillRegistry::load(&skill_paths).with_hub_dirs(vec![hub_dir.path().to_path_buf()]),
2521 ));
2522 assert!(
2523 registry.read().all_meta().is_empty(),
2524 "registry must start empty before any explore() reload"
2525 );
2526
2527 let generator = SkillGenerator::new(
2528 AnyProvider::Mock(mock_provider),
2529 skills_dir.path().to_path_buf(),
2530 );
2531 let explorer = Arc::new(ProactiveExplorer::new(
2532 generator,
2533 None,
2534 skills_dir.path().to_path_buf(),
2535 8_000,
2536 30_000,
2537 vec![],
2538 ));
2539
2540 Fixture {
2541 registry,
2542 skill_paths,
2543 hub_dir,
2544 _skills_dir_guard: skills_dir,
2545 explorer,
2546 }
2547 }
2548
2549 #[allow(clippy::too_many_arguments)]
2550 fn make_view<'a>(
2551 fixture: &'a Fixture,
2552 last_confidence: &'a mut Option<f32>,
2553 last_skills_prompt: &'a mut String,
2554 active_skill_names: &'a mut Vec<String>,
2555 sanitizer: &'a ContentSanitizer,
2556 ctx_mgr: &'a mut ContextManager,
2557 sink: &'a mut NoopSink,
2558 ) -> ContextAssemblyView<'a> {
2559 ContextAssemblyView {
2560 memory: None,
2561 conversation_id: None,
2562 recall_limit: 10,
2563 cross_session_score_threshold: 0.5,
2564 context_format: ContextFormat::default(),
2565 last_recall_confidence: last_confidence,
2566 context_strategy: ContextStrategy::default(),
2567 crossover_turn_threshold: 0,
2568 cached_session_digest: None,
2569 digest_enabled: false,
2570 graph_config: GraphConfig::default(),
2571 document_config: DocumentConfig::default(),
2572 persona_config: PersonaConfig::default(),
2573 trajectory_config: TrajectoryConfig::default(),
2574 reasoning_config: ReasoningConfig::default(),
2575 memcot_config: zeph_config::MemCotConfig::default(),
2576 memcot_state: None,
2577 tree_config: TreeConfig::default(),
2578 last_skills_prompt,
2579 active_skill_names,
2580 skill_registry: Arc::clone(&fixture.registry),
2581 skill_paths: &fixture.skill_paths,
2582 correction_config: None,
2583 sidequest_turn_counter: 0,
2584 proactive_explorer: Some(Arc::clone(&fixture.explorer)),
2585 sanitizer,
2586 quarantine_summarizer: None,
2587 context_manager: ctx_mgr,
2588 token_counter: make_counter(),
2589 metrics: MetricsCounters::default(),
2590 security_events: sink,
2591 cached_prompt_tokens: 0,
2592 redact_credentials: false,
2593 channel_skills: &[],
2594 scrub: scrub_noop,
2595 tiered_retrieval_config: TieredRetrievalConfig {
2596 enabled: false,
2597 ..TieredRetrievalConfig::default()
2598 },
2599 tiered_retrieval_classifier: None,
2600 tiered_retrieval_validator: None,
2601 fidelity_config: None,
2602 fidelity_semantic_provider: None,
2603 fidelity_compress_provider: None,
2604 planned_next_tools: &[],
2605 status_tx: None,
2606 task_supervisor: make_task_supervisor(),
2607 }
2608 }
2609
2610 #[tokio::test]
2611 async fn prepare_context_proactive_explore_reloads_registry_off_lock() {
2612 let fixture = setup_fixture(MockProvider::with_responses(vec![mock_skill_content(
2616 "world-knowledge-git",
2617 )]));
2618 let expected_hub_dirs = vec![fixture.hub_dir.path().to_path_buf()];
2619
2620 let mut msgs: Vec<Message> = vec![];
2621 let mut cached = 0u64;
2622 let mut completed = HashSet::new();
2623 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2624
2625 let sanitizer = ContentSanitizer::new(&ContentIsolationConfig::default());
2626 let mut ctx_mgr = ContextManager::new();
2627 ctx_mgr.budget = Some(ContextBudget::new(100_000, 0.1));
2628 let mut sink = NoopSink;
2629 let mut last_confidence = None::<f32>;
2630 let mut last_skills_prompt = String::new();
2631 let mut active_skill_names = Vec::new();
2632
2633 let mut view = make_view(
2634 &fixture,
2635 &mut last_confidence,
2636 &mut last_skills_prompt,
2637 &mut active_skill_names,
2638 &sanitizer,
2639 &mut ctx_mgr,
2640 &mut sink,
2641 );
2642
2643 let result = ContextService::new()
2644 .prepare_context("please help me with a git rebase", &mut window, &mut view)
2645 .await;
2646
2647 assert!(result.is_ok(), "prepare_context must succeed: {result:?}");
2648
2649 let guard = fixture.registry.read();
2650 assert_eq!(
2651 guard.all_meta().len(),
2652 1,
2653 "reloaded registry must contain exactly the newly generated skill"
2654 );
2655 assert!(
2656 guard
2657 .all_meta()
2658 .iter()
2659 .any(|m| m.name == "world-knowledge-git"),
2660 "reloaded registry must expose the generated skill by name"
2661 );
2662 assert_eq!(
2663 guard.hub_dirs(),
2664 expected_hub_dirs.as_slice(),
2665 "hub_dirs read before spawn_blocking must survive the off-lock rebuild"
2666 );
2667 }
2668
2669 #[tokio::test]
2670 async fn prepare_context_proactive_explore_leaves_registry_unchanged_when_not_triggered() {
2671 let fixture = setup_fixture(MockProvider::with_responses(vec![mock_skill_content(
2675 "world-knowledge-unused",
2676 )]));
2677
2678 let existing_skill_dir = fixture.skill_paths[0].join("existing-skill");
2681 std::fs::create_dir_all(&existing_skill_dir).expect("create existing skill dir");
2682 std::fs::write(
2683 existing_skill_dir.join("SKILL.md"),
2684 mock_skill_content("existing-skill"),
2685 )
2686 .expect("write existing SKILL.md");
2687 {
2688 let mut guard = fixture.registry.write();
2689 *guard = SkillRegistry::load(&fixture.skill_paths)
2690 .with_hub_dirs(vec![fixture.hub_dir.path().to_path_buf()]);
2691 }
2692 assert_eq!(
2693 fixture.registry.read().all_meta().len(),
2694 1,
2695 "sanity: registry must be pre-populated before prepare_context runs"
2696 );
2697
2698 let mut msgs: Vec<Message> = vec![];
2699 let mut cached = 0u64;
2700 let mut completed = HashSet::new();
2701 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2702
2703 let sanitizer = ContentSanitizer::new(&ContentIsolationConfig::default());
2704 let mut ctx_mgr = ContextManager::new();
2705 ctx_mgr.budget = Some(ContextBudget::new(100_000, 0.1));
2706 let mut sink = NoopSink;
2707 let mut last_confidence = None::<f32>;
2708 let mut last_skills_prompt = String::new();
2709 let mut active_skill_names = Vec::new();
2710
2711 let mut view = make_view(
2712 &fixture,
2713 &mut last_confidence,
2714 &mut last_skills_prompt,
2715 &mut active_skill_names,
2716 &sanitizer,
2717 &mut ctx_mgr,
2718 &mut sink,
2719 );
2720
2721 let result = ContextService::new()
2722 .prepare_context("how are you today", &mut window, &mut view)
2723 .await;
2724
2725 assert!(result.is_ok(), "prepare_context must succeed: {result:?}");
2726 let guard = fixture.registry.read();
2727 assert_eq!(
2728 guard.all_meta().len(),
2729 1,
2730 "registry must remain unchanged when no domain classifies and explore() never runs"
2731 );
2732 assert!(
2733 guard.all_meta().iter().any(|m| m.name == "existing-skill"),
2734 "the pre-existing skill must still be present, unreplaced by the mock's canned skill"
2735 );
2736 }
2737
2738 #[tokio::test]
2739 async fn prepare_context_proactive_explore_does_not_retrigger_for_known_domain() {
2740 let (mock_provider, recorder) =
2748 MockProvider::with_responses(vec![mock_skill_content("terraform-quickref")])
2749 .with_recording();
2750 let fixture = setup_fixture(mock_provider);
2751
2752 let sanitizer = ContentSanitizer::new(&ContentIsolationConfig::default());
2753
2754 {
2757 let mut msgs: Vec<Message> = vec![];
2758 let mut cached = 0u64;
2759 let mut completed = HashSet::new();
2760 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2761 let mut ctx_mgr = ContextManager::new();
2762 ctx_mgr.budget = Some(ContextBudget::new(100_000, 0.1));
2763 let mut sink = NoopSink;
2764 let mut last_confidence = None::<f32>;
2765 let mut last_skills_prompt = String::new();
2766 let mut active_skill_names = Vec::new();
2767
2768 let mut view = make_view(
2769 &fixture,
2770 &mut last_confidence,
2771 &mut last_skills_prompt,
2772 &mut active_skill_names,
2773 &sanitizer,
2774 &mut ctx_mgr,
2775 &mut sink,
2776 );
2777
2778 let result = ContextService::new()
2779 .prepare_context("tell me about terraform modules", &mut window, &mut view)
2780 .await;
2781 assert!(
2782 result.is_ok(),
2783 "first prepare_context must succeed: {result:?}"
2784 );
2785 }
2786
2787 assert_eq!(
2788 recorder.lock().unwrap().len(),
2789 1,
2790 "first turn must trigger exactly one LLM generation call"
2791 );
2792 assert!(
2793 fixture
2794 .registry
2795 .read()
2796 .all_meta()
2797 .iter()
2798 .any(|m| m.proactive_domain.as_deref() == Some("terraform")),
2799 "reloaded registry must carry the stamped proactive_domain field"
2800 );
2801
2802 {
2805 let mut msgs: Vec<Message> = vec![];
2806 let mut cached = 0u64;
2807 let mut completed = HashSet::new();
2808 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2809 let mut ctx_mgr = ContextManager::new();
2810 ctx_mgr.budget = Some(ContextBudget::new(100_000, 0.1));
2811 let mut sink = NoopSink;
2812 let mut last_confidence = None::<f32>;
2813 let mut last_skills_prompt = String::new();
2814 let mut active_skill_names = Vec::new();
2815
2816 let mut view = make_view(
2817 &fixture,
2818 &mut last_confidence,
2819 &mut last_skills_prompt,
2820 &mut active_skill_names,
2821 &sanitizer,
2822 &mut ctx_mgr,
2823 &mut sink,
2824 );
2825
2826 let result = ContextService::new()
2827 .prepare_context("tell me more about terraform state", &mut window, &mut view)
2828 .await;
2829 assert!(
2830 result.is_ok(),
2831 "second prepare_context must succeed: {result:?}"
2832 );
2833 }
2834
2835 assert_eq!(
2836 recorder.lock().unwrap().len(),
2837 1,
2838 "second turn for an already-known domain must NOT trigger another LLM generation call"
2839 );
2840 }
2841 }
2842
2843 mod prune_token_bookkeeping_tests {
2851 use std::time::Duration;
2852
2853 use tokio_util::sync::CancellationToken;
2854 use zeph_common::task_supervisor::{BlockingHandle, TaskSupervisor};
2855 use zeph_context::budget::ContextBudget;
2856 use zeph_context::manager::{CompactionTier, ContextManager};
2857 use zeph_context::summarization::{MessageTokenCounter, SummarizationDeps};
2858 use zeph_llm::any::AnyProvider;
2859 use zeph_llm::mock::MockProvider;
2860 use zeph_llm::provider::MessageMetadata;
2861
2862 use super::*;
2863 use crate::compaction::{SubgoalExtractionResult, SubgoalRegistry};
2864 use crate::memory_backend::TokenCounterAdapter;
2865
2866 struct Fixture {
2869 messages: Vec<Message>,
2870 deferred_db_hide_ids: Vec<i64>,
2871 deferred_db_summaries: Vec<String>,
2872 cached_prompt_tokens: u64,
2873 context_manager: ContextManager,
2874 subgoal_registry: SubgoalRegistry,
2875 pending_task_goal: Option<BlockingHandle<Option<String>>>,
2876 pending_subgoal: Option<BlockingHandle<Option<SubgoalExtractionResult>>>,
2877 current_task_goal: Option<String>,
2878 task_goal_user_msg_hash: Option<u64>,
2879 subgoal_user_msg_hash: Option<u64>,
2880 token_counter: Arc<TokenCounter>,
2881 task_supervisor: Arc<TaskSupervisor>,
2882 provider_responses: Vec<String>,
2887 }
2888
2889 impl Fixture {
2890 fn new(messages: Vec<Message>, cached_prompt_tokens: u64, cm: ContextManager) -> Self {
2891 Self {
2892 messages,
2893 deferred_db_hide_ids: Vec::new(),
2894 deferred_db_summaries: Vec::new(),
2895 cached_prompt_tokens,
2896 context_manager: cm,
2897 subgoal_registry: SubgoalRegistry::default(),
2898 pending_task_goal: None,
2899 pending_subgoal: None,
2900 current_task_goal: None,
2901 task_goal_user_msg_hash: None,
2902 subgoal_user_msg_hash: None,
2903 token_counter: make_counter(),
2904 task_supervisor: Arc::new(TaskSupervisor::new(CancellationToken::new())),
2905 provider_responses: Vec::new(),
2906 }
2907 }
2908
2909 fn view(&mut self) -> ContextSummarizationView<'_> {
2910 let token_counter_adapter: Arc<dyn MessageTokenCounter> =
2911 Arc::new(TokenCounterAdapter::new(Arc::clone(&self.token_counter)));
2912 ContextSummarizationView {
2913 messages: &mut self.messages,
2914 deferred_db_hide_ids: &mut self.deferred_db_hide_ids,
2915 deferred_db_summaries: &mut self.deferred_db_summaries,
2916 cached_prompt_tokens: &mut self.cached_prompt_tokens,
2917 context_manager: &mut self.context_manager,
2918 server_compaction_active: false,
2919 token_counter: Arc::clone(&self.token_counter),
2920 summarization_deps: SummarizationDeps {
2921 provider: AnyProvider::Mock(MockProvider::with_responses(
2922 self.provider_responses.clone(),
2923 )),
2924 llm_timeout: Duration::from_secs(30),
2925 token_counter: token_counter_adapter,
2926 structured_summaries: false,
2927 on_anchored_summary: None,
2928 },
2929 task_supervisor: Arc::clone(&self.task_supervisor),
2930 memory: None,
2931 conversation_id: None,
2932 tool_call_cutoff: 100,
2933 subgoal_registry: &mut self.subgoal_registry,
2934 pending_task_goal: &mut self.pending_task_goal,
2935 pending_subgoal: &mut self.pending_subgoal,
2936 current_task_goal: &mut self.current_task_goal,
2937 task_goal_user_msg_hash: &mut self.task_goal_user_msg_hash,
2938 subgoal_user_msg_hash: &mut self.subgoal_user_msg_hash,
2939 status_tx: None,
2940 scrub: |s| std::borrow::Cow::Borrowed(s),
2941 compression_guidelines: None,
2942 probe: None,
2943 archive: None,
2944 persistence: None,
2945 metrics: None,
2946 typed_pages: None,
2947 fidelity_config: None,
2948 fidelity_semantic_provider: None,
2949 fidelity_compress_provider: None,
2950 current_query: String::new(),
2951 }
2952 }
2953 }
2954
2955 struct NoopStatus;
2956 impl StatusSink for NoopStatus {
2957 fn send_status(&self, _msg: &str) -> impl std::future::Future<Output = ()> + Send + '_ {
2958 std::future::ready(())
2959 }
2960 }
2961
2962 fn plain_msg(role: Role, content: &str) -> Message {
2963 Message {
2964 role,
2965 content: content.to_owned(),
2966 parts: vec![],
2967 metadata: MessageMetadata::default(),
2968 }
2969 }
2970
2971 fn tool_use_msg() -> Message {
2972 Message::from_parts(
2973 Role::Assistant,
2974 vec![MessagePart::ToolUse {
2975 id: "t1".into(),
2976 name: "shell".into(),
2977 input: serde_json::json!({}),
2978 }],
2979 )
2980 }
2981
2982 fn tool_output_msg(body: &str) -> Message {
2983 Message::from_parts(
2984 Role::User,
2985 vec![MessagePart::ToolOutput {
2986 tool_name: "shell".into(),
2987 body: body.to_owned(),
2988 compacted_at: None,
2989 }],
2990 )
2991 }
2992
2993 fn messages_with_one_tool_output(body: &str) -> Vec<Message> {
2996 vec![
2997 plain_msg(Role::System, "system"),
2998 tool_use_msg(),
2999 tool_output_msg(body),
3000 plain_msg(Role::User, "hello"),
3001 ]
3002 }
3003
3004 #[test]
3005 fn prune_tool_outputs_decrements_cached_tokens_by_exact_freed_amount() {
3006 let body = "large tool output ".repeat(200);
3007 let expected_freed = TokenCounter::default().count_tokens(&body);
3008 let mut ctx_mgr = ContextManager::new();
3009 ctx_mgr.prune_protect_tokens = 0;
3010
3011 let initial_tokens = 10_000u64;
3012 let mut fixture = Fixture::new(
3013 messages_with_one_tool_output(&body),
3014 initial_tokens,
3015 ctx_mgr,
3016 );
3017 let mut view = fixture.view();
3018
3019 let freed =
3020 crate::summarization::pruning::prune_tool_outputs(&mut view, expected_freed);
3021
3022 assert_eq!(
3023 freed, expected_freed,
3024 "returned freed amount must match the tool output's token count"
3025 );
3026 assert_eq!(
3027 *view.cached_prompt_tokens,
3028 initial_tokens - u64::try_from(freed).unwrap(),
3029 "cached_prompt_tokens must decrease by exactly the freed amount"
3030 );
3031 }
3032
3033 #[test]
3034 fn prune_tool_outputs_noop_leaves_cached_tokens_unchanged() {
3035 let mut ctx_mgr = ContextManager::new();
3036 ctx_mgr.prune_protect_tokens = 0;
3037
3038 let messages = vec![
3039 plain_msg(Role::System, "system"),
3040 plain_msg(Role::User, "hello"),
3041 ];
3042 let initial_tokens = 500u64;
3043 let mut fixture = Fixture::new(messages, initial_tokens, ctx_mgr);
3044 let mut view = fixture.view();
3045
3046 let freed = crate::summarization::pruning::prune_tool_outputs(&mut view, 100);
3047
3048 assert_eq!(freed, 0, "no ToolOutput parts exist to prune");
3049 assert_eq!(
3050 *view.cached_prompt_tokens, initial_tokens,
3051 "cached_prompt_tokens must be untouched when nothing is freed"
3052 );
3053 }
3054
3055 #[tokio::test]
3056 async fn do_soft_compaction_via_maybe_compact_reflects_real_freed_tokens() {
3057 let body = "large tool output ".repeat(200);
3058 let expected_freed = TokenCounter::default().count_tokens(&body);
3059
3060 let mut ctx_mgr = ContextManager::new();
3061 ctx_mgr.budget = Some(ContextBudget::new(1000, 0.2));
3062 ctx_mgr.soft_compaction_threshold = 0.5; ctx_mgr.hard_compaction_threshold = 0.9; ctx_mgr.prune_protect_tokens = 0;
3065
3066 let initial_tokens = 700u64; assert_eq!(
3068 ctx_mgr.compaction_tier(initial_tokens),
3069 CompactionTier::Soft
3070 );
3071
3072 let mut fixture = Fixture::new(
3073 messages_with_one_tool_output(&body),
3074 initial_tokens,
3075 ctx_mgr,
3076 );
3077 let mut view = fixture.view();
3078 let status = NoopStatus;
3079
3080 ContextService::new()
3081 .maybe_compact(&mut view, &status)
3082 .await
3083 .unwrap();
3084
3085 assert!(
3086 view.deferred_db_summaries.is_empty(),
3087 "no deferred summaries were queued in this scenario"
3088 );
3089 assert_eq!(
3090 *view.cached_prompt_tokens,
3091 initial_tokens - u64::try_from(expected_freed).unwrap(),
3092 "Soft-tier prune-only pass must decrement cached_prompt_tokens by exactly \
3093 the freed amount"
3094 );
3095 }
3096
3097 #[tokio::test]
3098 async fn do_hard_compaction_satisfied_by_pruning_alone_reflects_real_freed_tokens() {
3099 let body = "large tool output ".repeat(400);
3103 let expected_freed = TokenCounter::default().count_tokens(&body);
3104
3105 let mut ctx_mgr = ContextManager::new();
3106 ctx_mgr.budget = Some(ContextBudget::new(1_000_000, 0.2));
3107 ctx_mgr.soft_compaction_threshold = 0.5;
3108 ctx_mgr.hard_compaction_threshold = 0.6; ctx_mgr.prune_protect_tokens = 0;
3110
3111 let initial_tokens = 600_000u64 + (expected_freed as u64 / 2);
3114 assert_eq!(
3115 ctx_mgr.compaction_tier(initial_tokens),
3116 CompactionTier::Hard
3117 );
3118
3119 let mut fixture = Fixture::new(
3120 messages_with_one_tool_output(&body),
3121 initial_tokens,
3122 ctx_mgr,
3123 );
3124 let mut view = fixture.view();
3125 let status = NoopStatus;
3126
3127 ContextService::new()
3128 .maybe_compact(&mut view, &status)
3129 .await
3130 .unwrap();
3131
3132 assert!(
3133 view.deferred_db_summaries.is_empty(),
3134 "no deferred summaries were queued in this scenario"
3135 );
3136 assert_eq!(
3137 *view.cached_prompt_tokens,
3138 initial_tokens - u64::try_from(expected_freed).unwrap(),
3139 "Hard-tier pass satisfied by pruning alone must decrement cached_prompt_tokens \
3140 by exactly the freed amount"
3141 );
3142 assert!(
3143 view.context_manager
3144 .compaction_state()
3145 .is_compacted_this_turn(),
3146 "pruning-satisfied Hard tier must still mark the turn as compacted"
3147 );
3148 }
3149
3150 #[tokio::test]
3175 async fn do_hard_compaction_combined_prune_and_llm_savings_avoid_false_exhaustion() {
3176 let counter = TokenCounter::default();
3177
3178 let large_body = "large tool output content ".repeat(400);
3182 let summary_response = "ok".to_string();
3183
3184 let build_messages = || {
3185 vec![
3186 plain_msg(Role::System, "system"),
3187 tool_use_msg(),
3188 tool_output_msg(&large_body),
3189 plain_msg(Role::User, "t0"),
3190 plain_msg(Role::User, "t1"),
3191 plain_msg(Role::User, "t2"),
3192 ]
3193 };
3194
3195 let initial_tokens: u64 = build_messages()
3196 .iter()
3197 .map(|m| counter.count_message_tokens(m) as u64)
3198 .sum();
3199 let non_body_overhead =
3200 initial_tokens - u64::try_from(counter.count_tokens(&large_body)).unwrap();
3201 let budget_tokens = usize::try_from(initial_tokens).unwrap();
3202
3203 let make_ctx_mgr = |hard_ratio: f32| {
3204 let mut cm = ContextManager::new();
3205 cm.prune_protect_tokens = 0;
3206 cm.compaction_preserve_tail = 2;
3207 cm.budget = Some(ContextBudget::new(budget_tokens, 0.0));
3208 cm.hard_compaction_threshold = hard_ratio;
3209 cm
3210 };
3211
3212 let measured_final_tokens = {
3217 let mut fixture =
3218 Fixture::new(build_messages(), initial_tokens, make_ctx_mgr(0.0001));
3219 fixture.provider_responses = vec![summary_response.clone()];
3220 let mut view = fixture.view();
3221 ContextService::new()
3222 .do_hard_compaction(&mut view, &NoopStatus, false)
3223 .await
3224 .unwrap();
3225 *view.cached_prompt_tokens
3226 };
3227 assert!(
3228 measured_final_tokens < non_body_overhead,
3229 "test setup invariant: the LLM step must reduce tokens below the non-body \
3230 overhead ({non_body_overhead}) so a valid hard-threshold window exists; \
3231 measured {measured_final_tokens}"
3232 );
3233
3234 let hard_threshold_tokens = u64::midpoint(measured_final_tokens, non_body_overhead);
3239 #[allow(clippy::cast_precision_loss)]
3240 let hard_ratio = hard_threshold_tokens as f32 / budget_tokens as f32;
3241
3242 let mut fixture =
3243 Fixture::new(build_messages(), initial_tokens, make_ctx_mgr(hard_ratio));
3244 fixture.provider_responses = vec![summary_response];
3245 let mut view = fixture.view();
3246
3247 ContextService::new()
3248 .do_hard_compaction(&mut view, &NoopStatus, false)
3249 .await
3250 .unwrap();
3251
3252 assert!(
3253 !view.context_manager.compaction_state().is_exhausted(),
3254 "pruning's own savings must count toward the combined freed-tokens check, so \
3255 a Hard-tier pass where pruning alone was insufficient but prune+LLM combined \
3256 cross the hard threshold must not be marked Exhausted"
3257 );
3258 assert!(
3259 view.context_manager
3260 .compaction_state()
3261 .is_compacted_this_turn(),
3262 "the turn must be marked compacted given the combined prune + LLM reduction"
3263 );
3264 }
3265 }
3266}