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 view.metrics.quarantine_failures += 1;
1014 if qs.error_should_block() {
1015 tracing::warn!(
1016 error = %e,
1017 "quarantine failed for memory retrieval, fail_strategy=closed: blocking content"
1018 );
1019 view.security_events.push(
1020 zeph_common::SecurityEventCategory::Quarantine,
1021 "memory_retrieval",
1022 format!("Quarantine failed (fail-closed): {e}"),
1023 );
1024 msg.content = qs.blocked_fallback(&sanitized);
1025 return msg;
1026 }
1027 tracing::warn!(
1028 error = %e,
1029 "quarantine failed for memory retrieval, fail_strategy=open: using original sanitized content"
1030 );
1031 view.security_events.push(
1032 zeph_common::SecurityEventCategory::Quarantine,
1033 "memory_retrieval",
1034 format!("Quarantine failed: {e}"),
1035 );
1036 }
1037 }
1038 }
1039
1040 msg.content = sanitized.body;
1041 msg
1042 }
1043
1044 pub fn reset_conversation(
1054 &self,
1055 window: &mut MessageWindowView<'_>,
1056 _view: &mut ContextAssemblyView<'_>,
1057 ) -> Result<(), ContextError> {
1058 self.clear_history(window);
1059 Ok(())
1060 }
1061
1062 #[allow(
1079 clippy::cast_precision_loss,
1080 clippy::cast_possible_truncation,
1081 clippy::cast_sign_loss,
1082 clippy::too_many_lines
1083 )]
1084 #[tracing::instrument(name = "agent_context.service.maybe_compact", skip_all, err)]
1085 pub async fn maybe_compact(
1086 &self,
1087 summ: &mut ContextSummarizationView<'_>,
1088 status: &(impl StatusSink + ?Sized),
1089 ) -> Result<(), ContextError> {
1090 use zeph_context::manager::{CompactionState, CompactionTier};
1091
1092 if let Some(count) = summ.context_manager.turns_since_last_hard_compaction_mut() {
1094 *count += 1;
1095 }
1096
1097 if let CompactionState::Exhausted { warned } = summ.context_manager.compaction_state()
1099 && !warned
1100 {
1101 summ.context_manager
1102 .set_compaction_state(CompactionState::Exhausted { warned: true });
1103 tracing::warn!("compaction exhausted: context budget too tight for this session");
1104 }
1105 if summ.context_manager.compaction_state().is_exhausted() {
1106 return Ok(());
1107 }
1108
1109 if summ.server_compaction_active {
1111 let budget = summ
1112 .context_manager
1113 .budget
1114 .as_ref()
1115 .map_or(0, ContextBudget::max_tokens);
1116 if budget > 0 {
1117 let fallback = (budget * 95 / 100) as u64;
1118 if *summ.cached_prompt_tokens < fallback {
1119 return Ok(());
1120 }
1121 tracing::warn!(
1122 "server compaction active but context at 95%+ — falling back to client-side"
1123 );
1124 } else {
1125 return Ok(());
1126 }
1127 }
1128
1129 if summ
1131 .context_manager
1132 .compaction_state()
1133 .is_compacted_this_turn()
1134 {
1135 return Ok(());
1136 }
1137
1138 let in_cooldown = summ.context_manager.compaction_state().cooldown_remaining() > 0;
1140 if in_cooldown
1141 && let CompactionState::Cooling { turns_remaining } =
1142 summ.context_manager.compaction_state()
1143 {
1144 let next = turns_remaining - 1;
1145 summ.context_manager.set_compaction_state(if next == 0 {
1146 CompactionState::Ready
1147 } else {
1148 CompactionState::Cooling {
1149 turns_remaining: next,
1150 }
1151 });
1152 }
1153
1154 if let Some(ref fidelity_cfg) = summ.fidelity_config.clone()
1158 && fidelity_cfg.enabled
1159 && summ.context_manager.should_proactively_regrade(
1160 *summ.cached_prompt_tokens,
1161 fidelity_cfg.regrade_threshold,
1162 summ.server_compaction_active,
1163 )
1164 {
1165 use tracing::Instrument as _;
1166 let (regrade_embed_provider, regrade_compress_provider) = fidelity_provider_pair(
1167 summ.fidelity_semantic_provider.as_ref(),
1168 summ.fidelity_compress_provider.as_ref(),
1169 );
1170 FidelityScorer
1171 .score_and_apply(
1172 summ.messages,
1173 &summ.current_query,
1174 &[],
1175 fidelity_cfg,
1176 &*summ.token_counter,
1177 0,
1178 true, regrade_embed_provider,
1180 regrade_compress_provider,
1181 )
1182 .instrument(tracing::info_span!(
1183 "context.fidelity.regrade",
1184 budget_ratio = tracing::field::Empty,
1185 ))
1186 .await;
1187 persist_fidelity_tags(summ.messages, summ.memory.as_deref()).await;
1189 recompute_prompt_tokens_summ(summ);
1190 summ.context_manager.set_regraded_this_turn(true);
1191 tracing::debug!(
1192 cached_tokens = *summ.cached_prompt_tokens,
1193 "AgeMem proactive regrade complete"
1194 );
1195 }
1196
1197 match summ
1198 .context_manager
1199 .compaction_tier(*summ.cached_prompt_tokens)
1200 {
1201 CompactionTier::Soft => {
1202 self.do_soft_compaction(summ, status).await;
1203 Ok(())
1204 }
1205 CompactionTier::Hard => self.do_hard_compaction(summ, status, in_cooldown).await,
1206 _ => Ok(()),
1207 }
1208 }
1209
1210 #[tracing::instrument(name = "agent_context.service.do_soft_compaction", skip_all)]
1215 #[allow(
1216 clippy::cast_precision_loss,
1217 clippy::cast_possible_truncation,
1218 clippy::cast_sign_loss
1219 )]
1220 async fn do_soft_compaction(
1221 &self,
1222 summ: &mut ContextSummarizationView<'_>,
1223 status: &(impl StatusSink + ?Sized),
1224 ) {
1225 status.send_status("soft compacting context...").await;
1226
1227 match &summ.context_manager.compression.pruning_strategy {
1229 zeph_config::PruningStrategy::Subgoal | zeph_config::PruningStrategy::SubgoalMig => {
1230 crate::summarization::scheduling::maybe_refresh_subgoal(summ);
1231 }
1232 _ => crate::summarization::scheduling::maybe_refresh_task_goal(summ),
1233 }
1234
1235 let applied = crate::summarization::deferred::apply_deferred_summaries(summ);
1237
1238 if applied > 0
1240 && summ
1241 .context_manager
1242 .compression
1243 .pruning_strategy
1244 .is_subgoal()
1245 {
1246 summ.subgoal_registry
1247 .rebuild_after_compaction(summ.messages, 0);
1248 }
1249
1250 let budget = summ
1252 .context_manager
1253 .budget
1254 .as_ref()
1255 .map_or(0, ContextBudget::max_tokens);
1256 let soft_threshold =
1257 (budget as f32 * summ.context_manager.soft_compaction_threshold) as usize;
1258 let cached = usize::try_from(*summ.cached_prompt_tokens).unwrap_or(usize::MAX);
1259 let min_to_free = cached.saturating_sub(soft_threshold);
1260 if min_to_free > 0 {
1261 crate::summarization::pruning::prune_tool_outputs(summ, min_to_free);
1262 }
1263
1264 status.send_status("").await;
1265 tracing::info!(
1266 cached_tokens = *summ.cached_prompt_tokens,
1267 soft_threshold,
1268 "soft compaction complete"
1269 );
1270 }
1271
1272 #[tracing::instrument(name = "agent_context.service.do_hard_compaction", skip_all, err)]
1274 #[allow(
1275 clippy::cast_precision_loss,
1276 clippy::cast_possible_truncation,
1277 clippy::cast_sign_loss
1278 )]
1279 async fn do_hard_compaction(
1280 &self,
1281 summ: &mut ContextSummarizationView<'_>,
1282 status: &(impl StatusSink + ?Sized),
1283 in_cooldown: bool,
1284 ) -> Result<(), ContextError> {
1285 use zeph_context::manager::CompactionState;
1286
1287 let turns_since_last = summ
1289 .context_manager
1290 .turns_since_last_hard_compaction()
1291 .map(|t| u32::try_from(t).unwrap_or(u32::MAX));
1292 summ.context_manager
1293 .set_turns_since_last_hard_compaction(Some(0));
1294 if let Some(metrics) = summ.metrics {
1295 metrics.record_hard_compaction(turns_since_last);
1296 }
1297
1298 if in_cooldown {
1299 tracing::debug!(
1300 turns_remaining = summ.context_manager.compaction_state().cooldown_remaining(),
1301 "hard compaction skipped: cooldown active"
1302 );
1303 return Ok(());
1304 }
1305
1306 let budget = summ
1307 .context_manager
1308 .budget
1309 .as_ref()
1310 .map_or(0, ContextBudget::max_tokens);
1311 let hard_threshold =
1312 (budget as f32 * summ.context_manager.hard_compaction_threshold) as usize;
1313 let cached = usize::try_from(*summ.cached_prompt_tokens).unwrap_or(usize::MAX);
1314 let min_to_free = cached.saturating_sub(hard_threshold);
1315
1316 status.send_status("compacting context...").await;
1317
1318 crate::summarization::deferred::apply_deferred_summaries(summ);
1320
1321 let tokens_before = *summ.cached_prompt_tokens;
1329 let freed = crate::summarization::pruning::prune_tool_outputs(summ, min_to_free);
1330 if freed >= min_to_free {
1331 tracing::info!(freed, "hard compaction: pruning sufficient");
1332 summ.context_manager
1333 .set_compaction_state(CompactionState::CompactedThisTurn {
1334 cooldown: summ.context_manager.compaction_cooldown_turns(),
1335 });
1336 if let Err(e) = crate::summarization::deferred::flush_deferred_summaries(summ).await {
1337 tracing::warn!(%e, "flush_deferred_summaries failed after hard compaction");
1338 }
1339 status.send_status("").await;
1340 return Ok(());
1341 }
1342
1343 let preserve_tail = summ.context_manager.compaction_preserve_tail;
1345 let compactable = summ.messages.len().saturating_sub(preserve_tail + 1);
1346 if compactable <= 1 {
1347 tracing::warn!(
1348 compactable,
1349 "hard compaction: too few messages, marking exhausted"
1350 );
1351 summ.context_manager
1352 .set_compaction_state(CompactionState::Exhausted { warned: false });
1353 status.send_status("").await;
1354 return Ok(());
1355 }
1356
1357 tracing::info!(
1359 min_to_free,
1360 "hard compaction: falling back to LLM summarization"
1361 );
1362 let outcome = crate::summarization::compaction::compact_context(summ, None).await?;
1363
1364 let freed_tokens = tokens_before.saturating_sub(*summ.cached_prompt_tokens);
1365
1366 if !outcome.is_compacted() || freed_tokens == 0 {
1367 tracing::warn!("hard compaction: no net reduction, marking exhausted");
1368 summ.context_manager
1369 .set_compaction_state(CompactionState::Exhausted { warned: false });
1370 status.send_status("").await;
1371 return Ok(());
1372 }
1373
1374 if matches!(
1375 summ.context_manager
1376 .compaction_tier(*summ.cached_prompt_tokens),
1377 zeph_context::manager::CompactionTier::Hard
1378 ) {
1379 tracing::warn!(
1380 freed_tokens,
1381 "hard compaction: still above hard threshold after compaction, marking exhausted"
1382 );
1383 summ.context_manager
1384 .set_compaction_state(CompactionState::Exhausted { warned: false });
1385 status.send_status("").await;
1386 return Ok(());
1387 }
1388
1389 summ.context_manager
1390 .set_compaction_state(CompactionState::CompactedThisTurn {
1391 cooldown: summ.context_manager.compaction_cooldown_turns(),
1392 });
1393
1394 if tokens_before > *summ.cached_prompt_tokens {
1395 tracing::info!(
1396 tokens_before,
1397 tokens_after = *summ.cached_prompt_tokens,
1398 saved = freed_tokens,
1399 "context compaction complete"
1400 );
1401 }
1402
1403 status.send_status("").await;
1404 Ok(())
1405 }
1406
1407 #[tracing::instrument(name = "agent_context.service.maybe_summarize_tool_pair", skip_all)]
1413 pub async fn maybe_summarize_tool_pair(
1414 &self,
1415 summ: &mut ContextSummarizationView<'_>,
1416 providers: &ProviderHandles,
1417 ) {
1418 crate::summarization::deferred::maybe_summarize_tool_pair(
1419 summ,
1420 providers,
1421 &TxStatusSink(summ.status_tx.clone()),
1422 )
1423 .await;
1424 }
1425
1426 #[must_use]
1431 pub fn apply_deferred_summaries(&self, summ: &mut ContextSummarizationView<'_>) -> usize {
1432 crate::summarization::deferred::apply_deferred_summaries(summ)
1433 }
1434
1435 #[tracing::instrument(name = "agent_context.service.flush_deferred_summaries", skip_all)]
1440 pub async fn flush_deferred_summaries(&self, summ: &mut ContextSummarizationView<'_>) {
1441 if let Err(e) = crate::summarization::deferred::flush_deferred_summaries(summ).await {
1442 tracing::warn!(%e, "flush_deferred_summaries failed");
1443 }
1444 }
1445
1446 pub fn maybe_apply_deferred_summaries(&self, summ: &mut ContextSummarizationView<'_>) {
1452 crate::summarization::deferred::maybe_apply_deferred_summaries(summ);
1453 }
1454
1455 #[tracing::instrument(name = "agent_context.service.compact_context", skip_all, err)]
1470 pub async fn compact_context(
1471 &self,
1472 summ: &mut ContextSummarizationView<'_>,
1473 max_summary_tokens: Option<usize>,
1474 ) -> Result<crate::state::CompactionOutcome, crate::error::ContextError> {
1475 crate::summarization::compaction::compact_context(summ, max_summary_tokens).await
1476 }
1477
1478 pub fn maybe_soft_compact_mid_iteration(&self, summ: &mut ContextSummarizationView<'_>) {
1484 crate::summarization::scheduling::maybe_soft_compact_mid_iteration(summ);
1485 }
1486
1487 #[tracing::instrument(name = "agent_context.service.maybe_proactive_compress", skip_all)]
1493 pub async fn maybe_proactive_compress(
1494 &self,
1495 summ: &mut ContextSummarizationView<'_>,
1496 status: &(impl StatusSink + ?Sized),
1497 ) {
1498 let Some((_threshold, max_summary_tokens)) = summ
1499 .context_manager
1500 .should_proactively_compress(*summ.cached_prompt_tokens)
1501 else {
1502 return;
1503 };
1504
1505 if summ.server_compaction_active {
1506 let budget = summ
1507 .context_manager
1508 .budget
1509 .as_ref()
1510 .map_or(0, ContextBudget::max_tokens);
1511 if budget > 0 {
1512 let fallback = (budget * 95 / 100) as u64;
1513 if *summ.cached_prompt_tokens <= fallback {
1514 return;
1515 }
1516 tracing::warn!(
1517 cached_prompt_tokens = *summ.cached_prompt_tokens,
1518 fallback_threshold = fallback,
1519 "server compaction active but context at 95%+ — falling back to proactive"
1520 );
1521 } else {
1522 return;
1523 }
1524 }
1525
1526 status.send_status("compressing context...").await;
1527 tracing::info!(
1528 max_summary_tokens,
1529 cached_tokens = *summ.cached_prompt_tokens,
1530 "proactive compression triggered"
1531 );
1532
1533 match crate::summarization::compaction::compact_context(summ, Some(max_summary_tokens))
1534 .await
1535 {
1536 Ok(outcome) if outcome.is_compacted() => {
1537 summ.context_manager.set_compaction_state(
1538 zeph_context::manager::CompactionState::CompactedThisTurn { cooldown: 0 },
1539 );
1540 tracing::info!("proactive compression complete");
1541 }
1542 Ok(_) => {}
1543 Err(e) => tracing::warn!(%e, "proactive compression failed"),
1544 }
1545
1546 status.send_status("").await;
1547 }
1548
1549 pub fn maybe_refresh_task_goal(&self, summ: &mut ContextSummarizationView<'_>) {
1555 crate::summarization::scheduling::maybe_refresh_task_goal(summ);
1556 }
1557
1558 pub fn maybe_refresh_subgoal(&self, summ: &mut ContextSummarizationView<'_>) {
1563 crate::summarization::scheduling::maybe_refresh_subgoal(summ);
1564 }
1565}
1566
1567struct TxStatusSink(Option<tokio::sync::mpsc::UnboundedSender<String>>);
1573
1574impl StatusSink for TxStatusSink {
1575 fn send_status(&self, msg: &str) -> impl std::future::Future<Output = ()> + Send + '_ {
1576 if let Some(ref tx) = self.0 {
1577 let _ = tx.send(msg.to_owned());
1578 }
1579 std::future::ready(())
1580 }
1581}
1582
1583pub(crate) fn recompute_prompt_tokens(window: &mut MessageWindowView<'_>) {
1590 *window.cached_prompt_tokens = window
1591 .messages
1592 .iter()
1593 .map(|m| window.token_counter.count_message_tokens(m) as u64)
1594 .sum();
1595}
1596
1597fn fidelity_provider_pair<'a>(
1603 semantic: Option<&'a std::sync::Arc<zeph_llm::any::AnyProvider>>,
1604 compress: Option<&'a std::sync::Arc<zeph_llm::any::AnyProvider>>,
1605) -> (
1606 Option<&'a dyn zeph_llm::LlmProviderDyn>,
1607 Option<&'a dyn zeph_llm::LlmProviderDyn>,
1608) {
1609 let embed_provider = semantic
1610 .map(std::sync::Arc::as_ref)
1611 .map(|p| p as &dyn zeph_llm::LlmProviderDyn);
1612 let compress_provider = compress
1613 .map(std::sync::Arc::as_ref)
1614 .map(|p| p as &dyn zeph_llm::LlmProviderDyn);
1615 (embed_provider, compress_provider)
1616}
1617
1618#[tracing::instrument(name = "agent_context.service.persist_fidelity_tags", skip_all)]
1627async fn persist_fidelity_tags(
1628 messages: &[zeph_llm::provider::Message],
1629 memory: Option<&zeph_memory::semantic::SemanticMemory>,
1630) {
1631 let Some(mem) = memory else { return };
1632 let updates: Vec<(zeph_memory::MessageId, u8)> = messages
1633 .iter()
1634 .filter_map(|m| {
1635 let db_id = m.metadata.db_id?;
1636 let tag = m.metadata.fidelity_tag?;
1637 Some((zeph_memory::MessageId(db_id), tag as u8))
1638 })
1639 .collect();
1640 if updates.is_empty() {
1641 return;
1642 }
1643 if let Err(e) = mem.sqlite().update_fidelity_tags(&updates).await {
1644 tracing::warn!(
1645 count = updates.len(),
1646 error = %e,
1647 "failed to persist fidelity tags; floor invariant will not apply next turn"
1648 );
1649 }
1650}
1651
1652fn recompute_prompt_tokens_summ(summ: &mut crate::state::ContextSummarizationView<'_>) {
1656 *summ.cached_prompt_tokens = summ
1657 .messages
1658 .iter()
1659 .map(|m| summ.token_counter.count_message_tokens(m) as u64)
1660 .sum();
1661}
1662
1663pub(crate) fn remove_by_prefix(
1669 messages: &mut Vec<zeph_llm::provider::Message>,
1670 role: Role,
1671 prefix: &str,
1672) {
1673 messages.retain(|m| m.role != role || !m.content.starts_with(prefix));
1674}
1675
1676pub(crate) fn remove_by_part_or_prefix(
1685 messages: &mut Vec<zeph_llm::provider::Message>,
1686 prefix: &str,
1687 part_matches: impl Fn(&MessagePart) -> bool,
1688) {
1689 messages.retain(|m| {
1690 if m.role == Role::User {
1693 return !m.content.starts_with(prefix);
1694 }
1695 if m.role != Role::System {
1696 return true;
1697 }
1698 if m.parts.first().is_some_and(&part_matches) {
1699 return false;
1700 }
1701 !m.content.starts_with(prefix)
1702 });
1703}
1704
1705#[cfg(test)]
1706mod tests {
1707 use std::collections::HashSet;
1708 use std::sync::Arc;
1709
1710 use zeph_llm::provider::{Message, MessagePart, Role};
1711 use zeph_memory::TokenCounter;
1712
1713 use super::*;
1714 use crate::helpers::{GRAPH_FACTS_PREFIX, RECALL_PREFIX, SUMMARY_PREFIX};
1715 use crate::state::MessageWindowView;
1716
1717 fn make_counter() -> Arc<TokenCounter> {
1718 Arc::new(TokenCounter::default())
1719 }
1720
1721 fn make_window<'a>(
1722 messages: &'a mut Vec<Message>,
1723 cached: &'a mut u64,
1724 completed: &'a mut HashSet<String>,
1725 ) -> MessageWindowView<'a> {
1726 let last = Box::leak(Box::new(None::<i64>));
1727 let deferred_hide = Box::leak(Box::new(Vec::<i64>::new()));
1728 let deferred_summ = Box::leak(Box::new(Vec::<String>::new()));
1729 MessageWindowView {
1730 messages,
1731 last_persisted_message_id: last,
1732 deferred_db_hide_ids: deferred_hide,
1733 deferred_db_summaries: deferred_summ,
1734 cached_prompt_tokens: cached,
1735 token_counter: make_counter(),
1736 completed_tool_ids: completed,
1737 }
1738 }
1739
1740 fn sys(text: &str) -> Message {
1741 Message::from_legacy(Role::System, text)
1742 }
1743
1744 fn user(text: &str) -> Message {
1745 Message::from_legacy(Role::User, text)
1746 }
1747
1748 fn assistant(text: &str) -> Message {
1749 Message::from_legacy(Role::Assistant, text)
1750 }
1751
1752 #[test]
1753 fn clear_history_keeps_system_prompt() {
1754 let mut msgs = vec![sys("system"), user("hello"), assistant("hi")];
1755 let mut cached = 0u64;
1756 let mut completed = HashSet::new();
1757 completed.insert("tool_1".to_owned());
1758 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1759
1760 ContextService::new().clear_history(&mut window);
1761
1762 assert_eq!(window.messages.len(), 1);
1763 assert_eq!(window.messages[0].content, "system");
1764 assert!(
1765 window.completed_tool_ids.is_empty(),
1766 "completed_tool_ids must be cleared"
1767 );
1768 }
1769
1770 #[test]
1771 fn clear_history_empty_messages_is_noop() {
1772 let mut msgs: Vec<Message> = vec![];
1773 let mut cached = 0u64;
1774 let mut completed = HashSet::new();
1775 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1776
1777 ContextService::new().clear_history(&mut window);
1778
1779 assert!(window.messages.is_empty());
1780 }
1781
1782 #[test]
1783 fn remove_recall_messages_removes_by_prefix() {
1784 let mut msgs = vec![
1785 sys("system"),
1786 sys(&format!("{RECALL_PREFIX}some recalled text")),
1787 user("hello"),
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!(window.messages.len(), 2);
1796 assert!(
1797 window
1798 .messages
1799 .iter()
1800 .all(|m| !m.content.starts_with(RECALL_PREFIX))
1801 );
1802 }
1803
1804 #[test]
1807 fn remove_recall_messages_removes_user_role_recall() {
1808 let mut msgs = vec![
1809 sys("system"),
1810 user(&format!("{RECALL_PREFIX}recalled via tiered path")),
1811 user("real user message"),
1812 ];
1813 let mut cached = 0u64;
1814 let mut completed = HashSet::new();
1815 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1816
1817 ContextService::new().remove_recall_messages(&mut window);
1818
1819 assert_eq!(
1820 window.messages.len(),
1821 2,
1822 "Role::User recall message must be removed"
1823 );
1824 assert!(
1825 window
1826 .messages
1827 .iter()
1828 .all(|m| !m.content.starts_with(RECALL_PREFIX)),
1829 "no message with RECALL_PREFIX must remain"
1830 );
1831 assert!(
1832 window
1833 .messages
1834 .iter()
1835 .any(|m| m.content == "real user message"),
1836 "non-recall user message must survive"
1837 );
1838 }
1839
1840 #[test]
1841 fn remove_graph_facts_messages_removes_matching() {
1842 let mut msgs = vec![
1843 sys("system"),
1844 sys(&format!("{GRAPH_FACTS_PREFIX}fact1")),
1845 user("hello"),
1846 ];
1847 let mut cached = 0u64;
1848 let mut completed = HashSet::new();
1849 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1850
1851 ContextService::new().remove_graph_facts_messages(&mut window);
1852
1853 assert_eq!(window.messages.len(), 2);
1854 }
1855
1856 #[test]
1857 fn remove_summary_messages_removes_by_part() {
1858 let mut msgs = vec![
1859 sys("system"),
1860 Message::from_parts(
1861 Role::System,
1862 vec![MessagePart::Summary {
1863 text: format!("{SUMMARY_PREFIX}old summary"),
1864 }],
1865 ),
1866 user("hello"),
1867 ];
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().remove_summary_messages(&mut window);
1873
1874 assert_eq!(window.messages.len(), 2);
1875 }
1876
1877 #[test]
1878 fn trim_messages_to_budget_zero_is_noop() {
1879 let mut msgs = vec![sys("system"), user("a"), assistant("b"), user("c")];
1880 let original_len = msgs.len();
1881 let mut cached = 0u64;
1882 let mut completed = HashSet::new();
1883 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1884
1885 ContextService::new().trim_messages_to_budget(&mut window, 0);
1886
1887 assert_eq!(window.messages.len(), original_len);
1888 }
1889
1890 #[test]
1891 fn trim_messages_to_budget_keeps_recent() {
1892 let mut msgs = vec![
1894 sys("system"),
1895 user("message 1"),
1896 assistant("reply 1"),
1897 user("message 2"),
1898 ];
1899 let mut cached = 0u64;
1900 let mut completed = HashSet::new();
1901 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
1902
1903 ContextService::new().trim_messages_to_budget(&mut window, 1);
1905
1906 assert!(
1908 window.messages.len() < 4,
1909 "trim should remove some messages"
1910 );
1911 assert_eq!(
1912 window.messages[0].role,
1913 Role::System,
1914 "system prompt must survive trim"
1915 );
1916 }
1917
1918 mod inserted_count_tests {
1921 use parking_lot::RwLock;
1922 use std::borrow::Cow;
1923 use std::collections::HashSet;
1924 use std::sync::Arc;
1925
1926 use zeph_common::SecurityEventCategory;
1927 use zeph_config::memory::TieredRetrievalConfig;
1928 use zeph_config::{
1929 ContextFormat, ContextStrategy, DocumentConfig, GraphConfig, PersonaConfig,
1930 ReasoningConfig, TrajectoryConfig, TreeConfig,
1931 };
1932 use zeph_context::assembler::PreparedContext;
1933 use zeph_context::manager::ContextManager;
1934 use zeph_llm::provider::{Message, MessageMetadata, Role};
1935 use zeph_memory::TokenCounter;
1936 use zeph_sanitizer::ContentIsolationConfig;
1937 use zeph_sanitizer::ContentSanitizer;
1938 use zeph_skills::registry::SkillRegistry;
1939
1940 use super::super::*;
1941 use crate::state::{
1942 ContextAssemblyView, MessageWindowView, MetricsCounters, SecurityEventSink,
1943 };
1944
1945 fn make_task_supervisor() -> Arc<zeph_common::TaskSupervisor> {
1946 Arc::new(zeph_common::TaskSupervisor::new(
1947 tokio_util::sync::CancellationToken::new(),
1948 ))
1949 }
1950
1951 struct NoopSink;
1952 impl SecurityEventSink for NoopSink {
1953 fn push(&mut self, _: SecurityEventCategory, _: &'static str, _: String) {}
1954 }
1955
1956 fn make_counter() -> Arc<TokenCounter> {
1957 Arc::new(TokenCounter::default())
1958 }
1959
1960 fn make_window<'a>(
1961 messages: &'a mut Vec<Message>,
1962 cached: &'a mut u64,
1963 completed: &'a mut HashSet<String>,
1964 ) -> MessageWindowView<'a> {
1965 let last = Box::leak(Box::new(None::<i64>));
1966 let deferred_hide = Box::leak(Box::new(Vec::<i64>::new()));
1967 let deferred_summ = Box::leak(Box::new(Vec::<String>::new()));
1968 MessageWindowView {
1969 messages,
1970 last_persisted_message_id: last,
1971 deferred_db_hide_ids: deferred_hide,
1972 deferred_db_summaries: deferred_summ,
1973 cached_prompt_tokens: cached,
1974 token_counter: make_counter(),
1975 completed_tool_ids: completed,
1976 }
1977 }
1978
1979 fn mem_msg(content: &str) -> Message {
1980 Message {
1981 role: Role::User,
1982 content: content.to_string(),
1983 parts: vec![],
1984 metadata: MessageMetadata::default(),
1985 }
1986 }
1987
1988 fn scrub_noop(s: &str) -> Cow<'_, str> {
1989 Cow::Borrowed(s)
1990 }
1991
1992 #[tokio::test]
1993 async fn inserted_count_incremented_for_all_paths() {
1994 let mut msgs = vec![
1997 Message::from_legacy(Role::System, "system"),
1998 Message::from_legacy(Role::User, "user turn"),
1999 ];
2000 let mut cached = 0u64;
2001 let mut completed = HashSet::new();
2002 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2003
2004 let sanitizer = ContentSanitizer::new(&ContentIsolationConfig::default());
2005 let mut ctx_mgr = ContextManager::new();
2006 let mut sink = NoopSink;
2007 let mut last_confidence = None::<f32>;
2008 let mut last_skills_prompt = String::new();
2009 let mut active_skill_names = Vec::new();
2010 let registry = Arc::new(RwLock::new(SkillRegistry::default()));
2011
2012 let mut view = ContextAssemblyView {
2013 memory: None,
2014 conversation_id: None,
2015 recall_limit: 10,
2016 cross_session_score_threshold: 0.5,
2017 context_format: ContextFormat::default(),
2018 last_recall_confidence: &mut last_confidence,
2019 context_strategy: ContextStrategy::default(),
2020 crossover_turn_threshold: 0,
2021 cached_session_digest: None,
2022 digest_enabled: false, graph_config: GraphConfig::default(),
2024 document_config: DocumentConfig::default(),
2025 persona_config: PersonaConfig::default(),
2026 trajectory_config: TrajectoryConfig::default(),
2027 reasoning_config: ReasoningConfig::default(),
2028 memcot_config: zeph_config::MemCotConfig::default(),
2029 memcot_state: None,
2030 tree_config: TreeConfig::default(),
2031 last_skills_prompt: &mut last_skills_prompt,
2032 active_skill_names: &mut active_skill_names,
2033 skill_registry: registry,
2034 skill_paths: &[],
2035 correction_config: None,
2036 sidequest_turn_counter: 0,
2037 proactive_explorer: None,
2038 sanitizer: &sanitizer,
2039 quarantine_summarizer: None,
2040 context_manager: &mut ctx_mgr,
2041 token_counter: make_counter(),
2042 metrics: MetricsCounters::default(),
2043 security_events: &mut sink,
2044 cached_prompt_tokens: 0,
2045 redact_credentials: false,
2046 channel_skills: &[],
2047 scrub: scrub_noop,
2048 tiered_retrieval_config: TieredRetrievalConfig {
2049 enabled: false,
2050 ..TieredRetrievalConfig::default()
2051 },
2052 tiered_retrieval_classifier: None,
2053 tiered_retrieval_validator: None,
2054 type_aware_compose_config: zeph_config::memory::TypeAwareComposeConfig::default(),
2055 fidelity_config: None,
2056 fidelity_semantic_provider: None,
2057 fidelity_compress_provider: None,
2058 planned_next_tools: &[],
2059 status_tx: None,
2060 task_supervisor: make_task_supervisor(),
2061 };
2062
2063 let prepared = PreparedContext {
2065 graph_facts: Some(mem_msg("graph_facts")),
2066 doc_rag: Some(mem_msg("doc_rag")),
2067 corrections: Some(mem_msg("corrections")),
2068 recall: Some(mem_msg("recall")),
2069 recall_confidence: Some(0.9),
2070 cross_session: Some(mem_msg("cross_session")),
2071 summaries: Some(mem_msg("summaries")),
2072 code_context: None, persona_facts: Some(mem_msg("persona_facts")),
2074 trajectory_hints: Some(mem_msg("trajectory_hints")),
2075 tree_memory: Some(mem_msg("tree_memory")),
2076 reasoning_hints: Some(mem_msg("reasoning_hints")),
2077 memory_first: false,
2078 recent_history_budget: 100_000,
2079 background_tasks: vec![],
2080 };
2081
2082 let (_delta, inserted_count) = ContextService::new()
2083 .apply_prepared_context(&mut window, &mut view, prepared)
2084 .await;
2085
2086 assert_eq!(
2089 inserted_count, 10,
2090 "all 10 message-carrying PreparedContext fields must increment inserted_count"
2091 );
2092 }
2093 }
2094
2095 mod quarantine_fail_strategy_tests {
2099 use parking_lot::RwLock;
2100 use std::borrow::Cow;
2101 use std::sync::Arc;
2102
2103 use zeph_common::SecurityEventCategory;
2104 use zeph_config::memory::TieredRetrievalConfig;
2105 use zeph_config::{
2106 ContextFormat, ContextStrategy, DocumentConfig, GraphConfig, PersonaConfig,
2107 ReasoningConfig, TrajectoryConfig, TreeConfig,
2108 };
2109 use zeph_context::manager::ContextManager;
2110 use zeph_llm::any::AnyProvider;
2111 use zeph_llm::mock::MockProvider;
2112 use zeph_llm::provider::{Message, MessageMetadata, Role};
2113 use zeph_memory::TokenCounter;
2114 use zeph_sanitizer::quarantine::{
2115 GuardrailFailStrategy, QUARANTINE_BLOCKED_PLACEHOLDER, QuarantinedSummarizer,
2116 };
2117 use zeph_sanitizer::{ContentIsolationConfig, ContentSanitizer, MemorySourceHint};
2118 use zeph_skills::registry::SkillRegistry;
2119
2120 use super::super::*;
2121 use crate::state::{ContextAssemblyView, MetricsCounters, SecurityEventSink};
2122
2123 fn make_task_supervisor() -> Arc<zeph_common::TaskSupervisor> {
2124 Arc::new(zeph_common::TaskSupervisor::new(
2125 tokio_util::sync::CancellationToken::new(),
2126 ))
2127 }
2128
2129 struct RecordingSink {
2130 events: Vec<SecurityEventCategory>,
2131 }
2132 impl SecurityEventSink for RecordingSink {
2133 fn push(&mut self, category: SecurityEventCategory, _: &'static str, _: String) {
2134 self.events.push(category);
2135 }
2136 }
2137
2138 fn make_counter() -> Arc<TokenCounter> {
2139 Arc::new(TokenCounter::default())
2140 }
2141
2142 fn scrub_noop(s: &str) -> Cow<'_, str> {
2143 Cow::Borrowed(s)
2144 }
2145
2146 #[allow(clippy::too_many_arguments)]
2147 fn make_view<'a>(
2148 sanitizer: &'a ContentSanitizer,
2149 quarantine_summarizer: Option<&'a QuarantinedSummarizer>,
2150 last_confidence: &'a mut Option<f32>,
2151 last_skills_prompt: &'a mut String,
2152 active_skill_names: &'a mut Vec<String>,
2153 ctx_mgr: &'a mut ContextManager,
2154 sink: &'a mut RecordingSink,
2155 ) -> ContextAssemblyView<'a> {
2156 ContextAssemblyView {
2157 memory: None,
2158 conversation_id: None,
2159 recall_limit: 10,
2160 cross_session_score_threshold: 0.5,
2161 context_format: ContextFormat::default(),
2162 last_recall_confidence: last_confidence,
2163 context_strategy: ContextStrategy::default(),
2164 crossover_turn_threshold: 0,
2165 cached_session_digest: None,
2166 digest_enabled: false,
2167 graph_config: GraphConfig::default(),
2168 document_config: DocumentConfig::default(),
2169 persona_config: PersonaConfig::default(),
2170 trajectory_config: TrajectoryConfig::default(),
2171 reasoning_config: ReasoningConfig::default(),
2172 memcot_config: zeph_config::MemCotConfig::default(),
2173 memcot_state: None,
2174 tree_config: TreeConfig::default(),
2175 last_skills_prompt,
2176 active_skill_names,
2177 skill_registry: Arc::new(RwLock::new(SkillRegistry::default())),
2178 skill_paths: &[],
2179 correction_config: None,
2180 sidequest_turn_counter: 0,
2181 proactive_explorer: None,
2182 sanitizer,
2183 quarantine_summarizer,
2184 context_manager: ctx_mgr,
2185 token_counter: make_counter(),
2186 metrics: MetricsCounters::default(),
2187 security_events: sink,
2188 cached_prompt_tokens: 0,
2189 redact_credentials: false,
2190 channel_skills: &[],
2191 scrub: scrub_noop,
2192 tiered_retrieval_config: TieredRetrievalConfig {
2193 enabled: false,
2194 ..TieredRetrievalConfig::default()
2195 },
2196 tiered_retrieval_classifier: None,
2197 tiered_retrieval_validator: None,
2198 type_aware_compose_config: zeph_config::memory::TypeAwareComposeConfig::default(),
2199 fidelity_config: None,
2200 fidelity_semantic_provider: None,
2201 fidelity_compress_provider: None,
2202 planned_next_tools: &[],
2203 status_tx: None,
2204 task_supervisor: make_task_supervisor(),
2205 }
2206 }
2207
2208 fn quarantine_config(
2209 fail_strategy: GuardrailFailStrategy,
2210 ) -> zeph_config::QuarantineConfig {
2211 zeph_config::QuarantineConfig {
2212 enabled: true,
2213 sources: vec!["memory_retrieval".to_owned()],
2214 model: "mock".to_owned(),
2215 timeout_ms: 30_000,
2216 fail_strategy,
2217 }
2218 }
2219
2220 #[tokio::test]
2221 async fn fail_closed_default_blocks_content_on_quarantine_error() {
2222 let quarantine_provider = AnyProvider::Mock(MockProvider::failing());
2223 let qcfg = quarantine_config(GuardrailFailStrategy::Closed);
2224 let qs = QuarantinedSummarizer::new(quarantine_provider, &qcfg);
2225
2226 let sanitizer = ContentSanitizer::new(&ContentIsolationConfig::default());
2227 let mut ctx_mgr = ContextManager::new();
2228 let mut sink = RecordingSink { events: vec![] };
2229 let mut last_confidence = None::<f32>;
2230 let mut last_skills_prompt = String::new();
2231 let mut active_skill_names = Vec::new();
2232
2233 let mut view = make_view(
2234 &sanitizer,
2235 Some(&qs),
2236 &mut last_confidence,
2237 &mut last_skills_prompt,
2238 &mut active_skill_names,
2239 &mut ctx_mgr,
2240 &mut sink,
2241 );
2242
2243 let msg = Message {
2244 role: Role::User,
2245 content: "recalled secret memory".to_owned(),
2246 parts: vec![],
2247 metadata: MessageMetadata::default(),
2248 };
2249 let result = ContextService::new()
2250 .sanitize_memory_message(msg, MemorySourceHint::ExternalContent, &mut view)
2251 .await;
2252
2253 assert!(
2254 result.content.contains(QUARANTINE_BLOCKED_PLACEHOLDER),
2255 "fail-closed default must substitute the blocked placeholder: {}",
2256 result.content
2257 );
2258 assert!(
2259 !result.content.contains("recalled secret memory"),
2260 "fail-closed fallback must never leak the original untrusted content: {}",
2261 result.content
2262 );
2263 assert!(
2264 sink.events.contains(&SecurityEventCategory::Quarantine),
2265 "must emit a Quarantine security event for the fail-closed block"
2266 );
2267 }
2268
2269 #[tokio::test]
2270 async fn fail_open_preserves_legacy_fallback_behavior() {
2271 let quarantine_provider = AnyProvider::Mock(MockProvider::failing());
2272 let qcfg = quarantine_config(GuardrailFailStrategy::Open);
2273 let qs = QuarantinedSummarizer::new(quarantine_provider, &qcfg);
2274
2275 let sanitizer = ContentSanitizer::new(&ContentIsolationConfig::default());
2276 let mut ctx_mgr = ContextManager::new();
2277 let mut sink = RecordingSink { events: vec![] };
2278 let mut last_confidence = None::<f32>;
2279 let mut last_skills_prompt = String::new();
2280 let mut active_skill_names = Vec::new();
2281
2282 let mut view = make_view(
2283 &sanitizer,
2284 Some(&qs),
2285 &mut last_confidence,
2286 &mut last_skills_prompt,
2287 &mut active_skill_names,
2288 &mut ctx_mgr,
2289 &mut sink,
2290 );
2291
2292 let msg = Message {
2293 role: Role::User,
2294 content: "recalled memory content".to_owned(),
2295 parts: vec![],
2296 metadata: MessageMetadata::default(),
2297 };
2298 let result = ContextService::new()
2299 .sanitize_memory_message(msg, MemorySourceHint::ExternalContent, &mut view)
2300 .await;
2301
2302 assert!(
2303 !result.content.contains(QUARANTINE_BLOCKED_PLACEHOLDER),
2304 "fail-open must not use the blocked placeholder: {}",
2305 result.content
2306 );
2307 assert!(
2308 result.content.contains("recalled memory content"),
2309 "fail-open must fall back to the sanitized content: {}",
2310 result.content
2311 );
2312 }
2313 }
2314
2315 mod inject_semantic_recall_tests {
2316 use parking_lot::RwLock;
2317 use std::borrow::Cow;
2318 use std::collections::HashSet;
2319 use std::sync::Arc;
2320
2321 use zeph_config::memory::TieredRetrievalConfig;
2322 use zeph_config::{
2323 ContextFormat, ContextStrategy, DocumentConfig, GraphConfig, PersonaConfig,
2324 ReasoningConfig, TrajectoryConfig, TreeConfig,
2325 };
2326 use zeph_context::manager::ContextManager;
2327 use zeph_llm::provider::Message;
2328 use zeph_memory::TokenCounter;
2329 use zeph_sanitizer::ContentIsolationConfig;
2330 use zeph_sanitizer::ContentSanitizer;
2331 use zeph_skills::registry::SkillRegistry;
2332
2333 use zeph_common::SecurityEventCategory;
2334
2335 use super::super::*;
2336 use crate::helpers::RECALL_PREFIX;
2337 use crate::state::{
2338 ContextAssemblyView, MessageWindowView, MetricsCounters, SecurityEventSink,
2339 };
2340
2341 fn make_task_supervisor() -> Arc<zeph_common::TaskSupervisor> {
2342 Arc::new(zeph_common::TaskSupervisor::new(
2343 tokio_util::sync::CancellationToken::new(),
2344 ))
2345 }
2346
2347 struct NoopSink;
2348 impl SecurityEventSink for NoopSink {
2349 fn push(&mut self, _: SecurityEventCategory, _: &'static str, _: String) {}
2350 }
2351
2352 fn make_counter() -> Arc<TokenCounter> {
2353 Arc::new(TokenCounter::default())
2354 }
2355
2356 fn make_window<'a>(
2357 messages: &'a mut Vec<Message>,
2358 cached: &'a mut u64,
2359 completed: &'a mut HashSet<String>,
2360 ) -> MessageWindowView<'a> {
2361 let last = Box::leak(Box::new(None::<i64>));
2362 let deferred_hide = Box::leak(Box::new(Vec::<i64>::new()));
2363 let deferred_summ = Box::leak(Box::new(Vec::<String>::new()));
2364 MessageWindowView {
2365 messages,
2366 last_persisted_message_id: last,
2367 deferred_db_hide_ids: deferred_hide,
2368 deferred_db_summaries: deferred_summ,
2369 cached_prompt_tokens: cached,
2370 token_counter: make_counter(),
2371 completed_tool_ids: completed,
2372 }
2373 }
2374
2375 fn scrub_noop(s: &str) -> Cow<'_, str> {
2376 Cow::Borrowed(s)
2377 }
2378
2379 #[tokio::test]
2380 async fn tiered_recall_disabled_uses_flat_path() {
2381 let mut msgs: Vec<Message> = vec![];
2384 let mut cached = 0u64;
2385 let mut completed = HashSet::new();
2386 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2387
2388 let sanitizer = ContentSanitizer::new(&ContentIsolationConfig::default());
2389 let mut ctx_mgr = ContextManager::new();
2390 let mut sink = NoopSink;
2391 let mut last_confidence = None::<f32>;
2392 let mut last_skills_prompt = String::new();
2393 let mut active_skill_names = Vec::new();
2394 let registry = Arc::new(RwLock::new(SkillRegistry::default()));
2395
2396 let view = ContextAssemblyView {
2397 memory: None,
2398 conversation_id: None,
2399 recall_limit: 10,
2400 cross_session_score_threshold: 0.5,
2401 context_format: ContextFormat::default(),
2402 last_recall_confidence: &mut last_confidence,
2403 context_strategy: ContextStrategy::default(),
2404 crossover_turn_threshold: 0,
2405 cached_session_digest: None,
2406 digest_enabled: false,
2407 graph_config: GraphConfig::default(),
2408 document_config: DocumentConfig::default(),
2409 persona_config: PersonaConfig::default(),
2410 trajectory_config: TrajectoryConfig::default(),
2411 reasoning_config: ReasoningConfig::default(),
2412 memcot_config: zeph_config::MemCotConfig::default(),
2413 memcot_state: None,
2414 tree_config: TreeConfig::default(),
2415 last_skills_prompt: &mut last_skills_prompt,
2416 active_skill_names: &mut active_skill_names,
2417 skill_registry: registry,
2418 skill_paths: &[],
2419 correction_config: None,
2420 sidequest_turn_counter: 0,
2421 proactive_explorer: None,
2422 sanitizer: &sanitizer,
2423 quarantine_summarizer: None,
2424 context_manager: &mut ctx_mgr,
2425 token_counter: make_counter(),
2426 metrics: MetricsCounters::default(),
2427 security_events: &mut sink,
2428 cached_prompt_tokens: 0,
2429 redact_credentials: false,
2430 channel_skills: &[],
2431 scrub: scrub_noop,
2432 tiered_retrieval_config: TieredRetrievalConfig {
2433 enabled: false,
2434 ..TieredRetrievalConfig::default()
2435 },
2436 tiered_retrieval_classifier: None,
2437 tiered_retrieval_validator: None,
2438 type_aware_compose_config: zeph_config::memory::TypeAwareComposeConfig::default(),
2439 fidelity_config: None,
2440 fidelity_semantic_provider: None,
2441 fidelity_compress_provider: None,
2442 planned_next_tools: &[],
2443 status_tx: None,
2444 task_supervisor: make_task_supervisor(),
2445 };
2446
2447 let result = ContextService::new()
2448 .inject_semantic_recall("test query", 1000, &mut window, &view)
2449 .await;
2450
2451 assert!(result.is_ok(), "disabled tiered recall must return Ok(())");
2452 assert!(
2453 window
2454 .messages
2455 .iter()
2456 .all(|m| !m.content.starts_with(RECALL_PREFIX)),
2457 "no recall message must be injected when memory is None"
2458 );
2459 }
2460
2461 #[tokio::test]
2462 async fn tiered_recall_enabled_no_memory_returns_ok() {
2463 let mut msgs: Vec<Message> = vec![];
2466 let mut cached = 0u64;
2467 let mut completed = HashSet::new();
2468 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2469
2470 let sanitizer = ContentSanitizer::new(&ContentIsolationConfig::default());
2471 let mut ctx_mgr = ContextManager::new();
2472 let mut sink = NoopSink;
2473 let mut last_confidence = None::<f32>;
2474 let mut last_skills_prompt = String::new();
2475 let mut active_skill_names = Vec::new();
2476 let registry = Arc::new(RwLock::new(SkillRegistry::default()));
2477
2478 let view = ContextAssemblyView {
2479 memory: None,
2480 conversation_id: None,
2481 recall_limit: 10,
2482 cross_session_score_threshold: 0.5,
2483 context_format: ContextFormat::default(),
2484 last_recall_confidence: &mut last_confidence,
2485 context_strategy: ContextStrategy::default(),
2486 crossover_turn_threshold: 0,
2487 cached_session_digest: None,
2488 digest_enabled: false,
2489 graph_config: GraphConfig::default(),
2490 document_config: DocumentConfig::default(),
2491 persona_config: PersonaConfig::default(),
2492 trajectory_config: TrajectoryConfig::default(),
2493 reasoning_config: ReasoningConfig::default(),
2494 memcot_config: zeph_config::MemCotConfig::default(),
2495 memcot_state: None,
2496 tree_config: TreeConfig::default(),
2497 last_skills_prompt: &mut last_skills_prompt,
2498 active_skill_names: &mut active_skill_names,
2499 skill_registry: registry,
2500 skill_paths: &[],
2501 correction_config: None,
2502 sidequest_turn_counter: 0,
2503 proactive_explorer: None,
2504 sanitizer: &sanitizer,
2505 quarantine_summarizer: None,
2506 context_manager: &mut ctx_mgr,
2507 token_counter: make_counter(),
2508 metrics: MetricsCounters::default(),
2509 security_events: &mut sink,
2510 cached_prompt_tokens: 0,
2511 redact_credentials: false,
2512 channel_skills: &[],
2513 scrub: scrub_noop,
2514 tiered_retrieval_config: TieredRetrievalConfig {
2515 enabled: true,
2516 ..TieredRetrievalConfig::default()
2517 },
2518 tiered_retrieval_classifier: None,
2519 tiered_retrieval_validator: None,
2520 type_aware_compose_config: zeph_config::memory::TypeAwareComposeConfig::default(),
2521 fidelity_config: None,
2522 fidelity_semantic_provider: None,
2523 fidelity_compress_provider: None,
2524 planned_next_tools: &[],
2525 status_tx: None,
2526 task_supervisor: make_task_supervisor(),
2527 };
2528
2529 let result = ContextService::new()
2530 .inject_semantic_recall("test query", 1000, &mut window, &view)
2531 .await;
2532
2533 assert!(
2534 result.is_ok(),
2535 "enabled tiered recall with no memory must return Ok(())"
2536 );
2537 assert!(
2538 window.messages.is_empty(),
2539 "no recall message must be injected when memory is None"
2540 );
2541 }
2542
2543 #[tokio::test]
2548 async fn prepare_context_tiered_enabled_no_budget_returns_default() {
2549 let mut msgs: Vec<zeph_llm::provider::Message> = vec![];
2550 let mut cached = 0u64;
2551 let mut completed = HashSet::new();
2552 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2553
2554 let sanitizer = zeph_sanitizer::ContentSanitizer::new(
2555 &zeph_sanitizer::ContentIsolationConfig::default(),
2556 );
2557 let mut ctx_mgr = zeph_context::manager::ContextManager::new();
2558 assert!(ctx_mgr.budget.is_none());
2560
2561 let mut sink = NoopSink;
2562 let mut last_confidence = None::<f32>;
2563 let mut last_skills_prompt = String::new();
2564 let mut active_skill_names = Vec::new();
2565 let registry = Arc::new(RwLock::new(zeph_skills::registry::SkillRegistry::default()));
2566
2567 let mut view = ContextAssemblyView {
2568 memory: None,
2569 conversation_id: None,
2570 recall_limit: 10,
2571 cross_session_score_threshold: 0.5,
2572 context_format: ContextFormat::default(),
2573 last_recall_confidence: &mut last_confidence,
2574 context_strategy: ContextStrategy::default(),
2575 crossover_turn_threshold: 0,
2576 cached_session_digest: None,
2577 digest_enabled: false,
2578 graph_config: GraphConfig::default(),
2579 document_config: DocumentConfig::default(),
2580 persona_config: PersonaConfig::default(),
2581 trajectory_config: TrajectoryConfig::default(),
2582 reasoning_config: ReasoningConfig::default(),
2583 memcot_config: zeph_config::MemCotConfig::default(),
2584 memcot_state: None,
2585 tree_config: TreeConfig::default(),
2586 last_skills_prompt: &mut last_skills_prompt,
2587 active_skill_names: &mut active_skill_names,
2588 skill_registry: registry,
2589 skill_paths: &[],
2590 correction_config: None,
2591 sidequest_turn_counter: 0,
2592 proactive_explorer: None,
2593 sanitizer: &sanitizer,
2594 quarantine_summarizer: None,
2595 context_manager: &mut ctx_mgr,
2596 token_counter: make_counter(),
2597 metrics: MetricsCounters::default(),
2598 security_events: &mut sink,
2599 cached_prompt_tokens: 0,
2600 redact_credentials: false,
2601 channel_skills: &[],
2602 scrub: scrub_noop,
2603 tiered_retrieval_config: TieredRetrievalConfig {
2604 enabled: true,
2605 ..TieredRetrievalConfig::default()
2606 },
2607 tiered_retrieval_classifier: None,
2608 tiered_retrieval_validator: None,
2609 type_aware_compose_config: zeph_config::memory::TypeAwareComposeConfig::default(),
2610 fidelity_config: None,
2611 fidelity_semantic_provider: None,
2612 fidelity_compress_provider: None,
2613 planned_next_tools: &[],
2614 status_tx: None,
2615 task_supervisor: make_task_supervisor(),
2616 };
2617
2618 let result = ContextService::new()
2619 .prepare_context("test query", &mut window, &mut view)
2620 .await;
2621
2622 assert!(
2623 result.is_ok(),
2624 "prepare_context with tiered enabled and no budget must return Ok"
2625 );
2626 }
2627
2628 #[tokio::test]
2631 async fn inject_semantic_recall_bare_no_memory_returns_ok() {
2632 use zeph_config::memory::TieredRetrievalConfig;
2633
2634 let mut msgs: Vec<Message> = vec![];
2635 let mut cached = 0u64;
2636 let mut completed = HashSet::new();
2637 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2638
2639 let tiered_config = TieredRetrievalConfig {
2640 enabled: true,
2641 ..TieredRetrievalConfig::default()
2642 };
2643 let params = SemanticRecallParams {
2644 query: "test query",
2645 token_budget: 1000,
2646 recall_limit: 10,
2647 context_format: zeph_config::ContextFormat::default(),
2648 conversation_id: None,
2649 tiered_classifier: None,
2650 tiered_validator: None,
2651 tiered_config: &tiered_config,
2652 };
2653 let result = ContextService::new()
2654 .inject_semantic_recall_bare(params, &mut window, None)
2655 .await;
2656
2657 assert!(
2658 result.is_ok(),
2659 "inject_semantic_recall_bare with memory=None must return Ok(())"
2660 );
2661 assert!(
2662 window.messages.is_empty(),
2663 "no recall message must be injected when memory is None"
2664 );
2665 }
2666 }
2667
2668 mod prepare_context_proactive_explore_tests {
2675 use parking_lot::RwLock;
2676 use std::borrow::Cow;
2677 use std::collections::HashSet;
2678 use std::sync::Arc;
2679
2680 use zeph_config::memory::TieredRetrievalConfig;
2681 use zeph_config::{
2682 ContextFormat, ContextStrategy, DocumentConfig, GraphConfig, PersonaConfig,
2683 ReasoningConfig, TrajectoryConfig, TreeConfig,
2684 };
2685 use zeph_context::budget::ContextBudget;
2686 use zeph_context::manager::ContextManager;
2687 use zeph_llm::any::AnyProvider;
2688 use zeph_llm::mock::MockProvider;
2689 use zeph_llm::provider::Message;
2690 use zeph_memory::TokenCounter;
2691 use zeph_sanitizer::ContentIsolationConfig;
2692 use zeph_sanitizer::ContentSanitizer;
2693 use zeph_skills::generator::SkillGenerator;
2694 use zeph_skills::proactive::ProactiveExplorer;
2695 use zeph_skills::registry::SkillRegistry;
2696
2697 use zeph_common::SecurityEventCategory;
2698
2699 use super::super::*;
2700 use crate::state::{
2701 ContextAssemblyView, MessageWindowView, MetricsCounters, SecurityEventSink,
2702 };
2703
2704 fn make_task_supervisor() -> Arc<zeph_common::TaskSupervisor> {
2705 Arc::new(zeph_common::TaskSupervisor::new(
2706 tokio_util::sync::CancellationToken::new(),
2707 ))
2708 }
2709
2710 struct NoopSink;
2711 impl SecurityEventSink for NoopSink {
2712 fn push(&mut self, _: SecurityEventCategory, _: &'static str, _: String) {}
2713 }
2714
2715 fn make_counter() -> Arc<TokenCounter> {
2716 Arc::new(TokenCounter::default())
2717 }
2718
2719 fn make_window<'a>(
2720 messages: &'a mut Vec<Message>,
2721 cached: &'a mut u64,
2722 completed: &'a mut HashSet<String>,
2723 ) -> MessageWindowView<'a> {
2724 let last = Box::leak(Box::new(None::<i64>));
2725 let deferred_hide = Box::leak(Box::new(Vec::<i64>::new()));
2726 let deferred_summ = Box::leak(Box::new(Vec::<String>::new()));
2727 MessageWindowView {
2728 messages,
2729 last_persisted_message_id: last,
2730 deferred_db_hide_ids: deferred_hide,
2731 deferred_db_summaries: deferred_summ,
2732 cached_prompt_tokens: cached,
2733 token_counter: make_counter(),
2734 completed_tool_ids: completed,
2735 }
2736 }
2737
2738 fn scrub_noop(s: &str) -> Cow<'_, str> {
2739 Cow::Borrowed(s)
2740 }
2741
2742 fn mock_skill_content(name: &str) -> String {
2743 format!(
2744 "---\nname: {name}\ndescription: Test world-knowledge skill for {name}.\n---\n\n## Usage\n\nDetails.\n"
2745 )
2746 }
2747
2748 struct Fixture {
2753 registry: Arc<RwLock<SkillRegistry>>,
2754 skill_paths: Vec<std::path::PathBuf>,
2755 hub_dir: tempfile::TempDir,
2756 _skills_dir_guard: tempfile::TempDir,
2759 explorer: Arc<ProactiveExplorer>,
2760 }
2761
2762 fn setup_fixture(mock_provider: MockProvider) -> Fixture {
2763 let skills_dir = tempfile::tempdir().expect("create skills tempdir");
2764 let hub_dir = tempfile::tempdir().expect("create hub tempdir");
2765 let skill_paths = vec![skills_dir.path().to_path_buf()];
2766
2767 let registry = Arc::new(RwLock::new(
2768 SkillRegistry::load(&skill_paths).with_hub_dirs(vec![hub_dir.path().to_path_buf()]),
2769 ));
2770 assert!(
2771 registry.read().all_meta().is_empty(),
2772 "registry must start empty before any explore() reload"
2773 );
2774
2775 let generator = SkillGenerator::new(
2776 AnyProvider::Mock(mock_provider),
2777 skills_dir.path().to_path_buf(),
2778 );
2779 let explorer = Arc::new(ProactiveExplorer::new(
2780 generator,
2781 None,
2782 skills_dir.path().to_path_buf(),
2783 8_000,
2784 30_000,
2785 vec![],
2786 ));
2787
2788 Fixture {
2789 registry,
2790 skill_paths,
2791 hub_dir,
2792 _skills_dir_guard: skills_dir,
2793 explorer,
2794 }
2795 }
2796
2797 #[allow(clippy::too_many_arguments)]
2798 fn make_view<'a>(
2799 fixture: &'a Fixture,
2800 last_confidence: &'a mut Option<f32>,
2801 last_skills_prompt: &'a mut String,
2802 active_skill_names: &'a mut Vec<String>,
2803 sanitizer: &'a ContentSanitizer,
2804 ctx_mgr: &'a mut ContextManager,
2805 sink: &'a mut NoopSink,
2806 ) -> ContextAssemblyView<'a> {
2807 ContextAssemblyView {
2808 memory: None,
2809 conversation_id: None,
2810 recall_limit: 10,
2811 cross_session_score_threshold: 0.5,
2812 context_format: ContextFormat::default(),
2813 last_recall_confidence: last_confidence,
2814 context_strategy: ContextStrategy::default(),
2815 crossover_turn_threshold: 0,
2816 cached_session_digest: None,
2817 digest_enabled: false,
2818 graph_config: GraphConfig::default(),
2819 document_config: DocumentConfig::default(),
2820 persona_config: PersonaConfig::default(),
2821 trajectory_config: TrajectoryConfig::default(),
2822 reasoning_config: ReasoningConfig::default(),
2823 memcot_config: zeph_config::MemCotConfig::default(),
2824 memcot_state: None,
2825 tree_config: TreeConfig::default(),
2826 last_skills_prompt,
2827 active_skill_names,
2828 skill_registry: Arc::clone(&fixture.registry),
2829 skill_paths: &fixture.skill_paths,
2830 correction_config: None,
2831 sidequest_turn_counter: 0,
2832 proactive_explorer: Some(Arc::clone(&fixture.explorer)),
2833 sanitizer,
2834 quarantine_summarizer: None,
2835 context_manager: ctx_mgr,
2836 token_counter: make_counter(),
2837 metrics: MetricsCounters::default(),
2838 security_events: sink,
2839 cached_prompt_tokens: 0,
2840 redact_credentials: false,
2841 channel_skills: &[],
2842 scrub: scrub_noop,
2843 tiered_retrieval_config: TieredRetrievalConfig {
2844 enabled: false,
2845 ..TieredRetrievalConfig::default()
2846 },
2847 tiered_retrieval_classifier: None,
2848 tiered_retrieval_validator: None,
2849 type_aware_compose_config: zeph_config::memory::TypeAwareComposeConfig::default(),
2850 fidelity_config: None,
2851 fidelity_semantic_provider: None,
2852 fidelity_compress_provider: None,
2853 planned_next_tools: &[],
2854 status_tx: None,
2855 task_supervisor: make_task_supervisor(),
2856 }
2857 }
2858
2859 #[tokio::test]
2860 async fn prepare_context_proactive_explore_reloads_registry_off_lock() {
2861 let fixture = setup_fixture(MockProvider::with_responses(vec![mock_skill_content(
2865 "world-knowledge-git",
2866 )]));
2867 let expected_hub_dirs = vec![fixture.hub_dir.path().to_path_buf()];
2868
2869 let mut msgs: Vec<Message> = vec![];
2870 let mut cached = 0u64;
2871 let mut completed = HashSet::new();
2872 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2873
2874 let sanitizer = ContentSanitizer::new(&ContentIsolationConfig::default());
2875 let mut ctx_mgr = ContextManager::new();
2876 ctx_mgr.budget = Some(ContextBudget::new(100_000, 0.1));
2877 let mut sink = NoopSink;
2878 let mut last_confidence = None::<f32>;
2879 let mut last_skills_prompt = String::new();
2880 let mut active_skill_names = Vec::new();
2881
2882 let mut view = make_view(
2883 &fixture,
2884 &mut last_confidence,
2885 &mut last_skills_prompt,
2886 &mut active_skill_names,
2887 &sanitizer,
2888 &mut ctx_mgr,
2889 &mut sink,
2890 );
2891
2892 let result = ContextService::new()
2893 .prepare_context("please help me with a git rebase", &mut window, &mut view)
2894 .await;
2895
2896 assert!(result.is_ok(), "prepare_context must succeed: {result:?}");
2897
2898 let guard = fixture.registry.read();
2899 assert_eq!(
2900 guard.all_meta().len(),
2901 1,
2902 "reloaded registry must contain exactly the newly generated skill"
2903 );
2904 assert!(
2905 guard
2906 .all_meta()
2907 .iter()
2908 .any(|m| m.name == "world-knowledge-git"),
2909 "reloaded registry must expose the generated skill by name"
2910 );
2911 assert_eq!(
2912 guard.hub_dirs(),
2913 expected_hub_dirs.as_slice(),
2914 "hub_dirs read before spawn_blocking must survive the off-lock rebuild"
2915 );
2916 }
2917
2918 #[tokio::test]
2919 async fn prepare_context_proactive_explore_leaves_registry_unchanged_when_not_triggered() {
2920 let fixture = setup_fixture(MockProvider::with_responses(vec![mock_skill_content(
2924 "world-knowledge-unused",
2925 )]));
2926
2927 let existing_skill_dir = fixture.skill_paths[0].join("existing-skill");
2930 std::fs::create_dir_all(&existing_skill_dir).expect("create existing skill dir");
2931 std::fs::write(
2932 existing_skill_dir.join("SKILL.md"),
2933 mock_skill_content("existing-skill"),
2934 )
2935 .expect("write existing SKILL.md");
2936 {
2937 let mut guard = fixture.registry.write();
2938 *guard = SkillRegistry::load(&fixture.skill_paths)
2939 .with_hub_dirs(vec![fixture.hub_dir.path().to_path_buf()]);
2940 }
2941 assert_eq!(
2942 fixture.registry.read().all_meta().len(),
2943 1,
2944 "sanity: registry must be pre-populated before prepare_context runs"
2945 );
2946
2947 let mut msgs: Vec<Message> = vec![];
2948 let mut cached = 0u64;
2949 let mut completed = HashSet::new();
2950 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
2951
2952 let sanitizer = ContentSanitizer::new(&ContentIsolationConfig::default());
2953 let mut ctx_mgr = ContextManager::new();
2954 ctx_mgr.budget = Some(ContextBudget::new(100_000, 0.1));
2955 let mut sink = NoopSink;
2956 let mut last_confidence = None::<f32>;
2957 let mut last_skills_prompt = String::new();
2958 let mut active_skill_names = Vec::new();
2959
2960 let mut view = make_view(
2961 &fixture,
2962 &mut last_confidence,
2963 &mut last_skills_prompt,
2964 &mut active_skill_names,
2965 &sanitizer,
2966 &mut ctx_mgr,
2967 &mut sink,
2968 );
2969
2970 let result = ContextService::new()
2971 .prepare_context("how are you today", &mut window, &mut view)
2972 .await;
2973
2974 assert!(result.is_ok(), "prepare_context must succeed: {result:?}");
2975 let guard = fixture.registry.read();
2976 assert_eq!(
2977 guard.all_meta().len(),
2978 1,
2979 "registry must remain unchanged when no domain classifies and explore() never runs"
2980 );
2981 assert!(
2982 guard.all_meta().iter().any(|m| m.name == "existing-skill"),
2983 "the pre-existing skill must still be present, unreplaced by the mock's canned skill"
2984 );
2985 }
2986
2987 #[tokio::test]
2988 async fn prepare_context_proactive_explore_does_not_retrigger_for_known_domain() {
2989 let (mock_provider, recorder) =
2997 MockProvider::with_responses(vec![mock_skill_content("terraform-quickref")])
2998 .with_recording();
2999 let fixture = setup_fixture(mock_provider);
3000
3001 let sanitizer = ContentSanitizer::new(&ContentIsolationConfig::default());
3002
3003 {
3006 let mut msgs: Vec<Message> = vec![];
3007 let mut cached = 0u64;
3008 let mut completed = HashSet::new();
3009 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
3010 let mut ctx_mgr = ContextManager::new();
3011 ctx_mgr.budget = Some(ContextBudget::new(100_000, 0.1));
3012 let mut sink = NoopSink;
3013 let mut last_confidence = None::<f32>;
3014 let mut last_skills_prompt = String::new();
3015 let mut active_skill_names = Vec::new();
3016
3017 let mut view = make_view(
3018 &fixture,
3019 &mut last_confidence,
3020 &mut last_skills_prompt,
3021 &mut active_skill_names,
3022 &sanitizer,
3023 &mut ctx_mgr,
3024 &mut sink,
3025 );
3026
3027 let result = ContextService::new()
3028 .prepare_context("tell me about terraform modules", &mut window, &mut view)
3029 .await;
3030 assert!(
3031 result.is_ok(),
3032 "first prepare_context must succeed: {result:?}"
3033 );
3034 }
3035
3036 assert_eq!(
3037 recorder.lock().unwrap().len(),
3038 1,
3039 "first turn must trigger exactly one LLM generation call"
3040 );
3041 assert!(
3042 fixture
3043 .registry
3044 .read()
3045 .all_meta()
3046 .iter()
3047 .any(|m| m.proactive_domain.as_deref() == Some("terraform")),
3048 "reloaded registry must carry the stamped proactive_domain field"
3049 );
3050
3051 {
3054 let mut msgs: Vec<Message> = vec![];
3055 let mut cached = 0u64;
3056 let mut completed = HashSet::new();
3057 let mut window = make_window(&mut msgs, &mut cached, &mut completed);
3058 let mut ctx_mgr = ContextManager::new();
3059 ctx_mgr.budget = Some(ContextBudget::new(100_000, 0.1));
3060 let mut sink = NoopSink;
3061 let mut last_confidence = None::<f32>;
3062 let mut last_skills_prompt = String::new();
3063 let mut active_skill_names = Vec::new();
3064
3065 let mut view = make_view(
3066 &fixture,
3067 &mut last_confidence,
3068 &mut last_skills_prompt,
3069 &mut active_skill_names,
3070 &sanitizer,
3071 &mut ctx_mgr,
3072 &mut sink,
3073 );
3074
3075 let result = ContextService::new()
3076 .prepare_context("tell me more about terraform state", &mut window, &mut view)
3077 .await;
3078 assert!(
3079 result.is_ok(),
3080 "second prepare_context must succeed: {result:?}"
3081 );
3082 }
3083
3084 assert_eq!(
3085 recorder.lock().unwrap().len(),
3086 1,
3087 "second turn for an already-known domain must NOT trigger another LLM generation call"
3088 );
3089 }
3090 }
3091
3092 mod prune_token_bookkeeping_tests {
3100 use std::time::Duration;
3101
3102 use tokio_util::sync::CancellationToken;
3103 use zeph_common::task_supervisor::{BlockingHandle, TaskSupervisor};
3104 use zeph_context::budget::ContextBudget;
3105 use zeph_context::manager::{CompactionTier, ContextManager};
3106 use zeph_context::summarization::{MessageTokenCounter, SummarizationDeps};
3107 use zeph_llm::any::AnyProvider;
3108 use zeph_llm::mock::MockProvider;
3109 use zeph_llm::provider::MessageMetadata;
3110
3111 use super::*;
3112 use crate::compaction::{SubgoalExtractionResult, SubgoalRegistry};
3113 use crate::memory_backend::TokenCounterAdapter;
3114
3115 struct Fixture {
3118 messages: Vec<Message>,
3119 deferred_db_hide_ids: Vec<i64>,
3120 deferred_db_summaries: Vec<String>,
3121 deferred_db_trust_levels: Vec<Option<u8>>,
3122 cached_prompt_tokens: u64,
3123 context_manager: ContextManager,
3124 subgoal_registry: SubgoalRegistry,
3125 pending_task_goal: Option<BlockingHandle<Option<String>>>,
3126 pending_subgoal: Option<BlockingHandle<Option<SubgoalExtractionResult>>>,
3127 current_task_goal: Option<String>,
3128 task_goal_user_msg_hash: Option<u64>,
3129 subgoal_user_msg_hash: Option<u64>,
3130 token_counter: Arc<TokenCounter>,
3131 task_supervisor: Arc<TaskSupervisor>,
3132 provider_responses: Vec<String>,
3137 }
3138
3139 impl Fixture {
3140 fn new(messages: Vec<Message>, cached_prompt_tokens: u64, cm: ContextManager) -> Self {
3141 Self {
3142 messages,
3143 deferred_db_hide_ids: Vec::new(),
3144 deferred_db_summaries: Vec::new(),
3145 deferred_db_trust_levels: Vec::new(),
3146 cached_prompt_tokens,
3147 context_manager: cm,
3148 subgoal_registry: SubgoalRegistry::default(),
3149 pending_task_goal: None,
3150 pending_subgoal: None,
3151 current_task_goal: None,
3152 task_goal_user_msg_hash: None,
3153 subgoal_user_msg_hash: None,
3154 token_counter: make_counter(),
3155 task_supervisor: Arc::new(TaskSupervisor::new(CancellationToken::new())),
3156 provider_responses: Vec::new(),
3157 }
3158 }
3159
3160 fn view(&mut self) -> ContextSummarizationView<'_> {
3161 let token_counter_adapter: Arc<dyn MessageTokenCounter> =
3162 Arc::new(TokenCounterAdapter::new(Arc::clone(&self.token_counter)));
3163 ContextSummarizationView {
3164 messages: &mut self.messages,
3165 deferred_db_hide_ids: &mut self.deferred_db_hide_ids,
3166 deferred_db_summaries: &mut self.deferred_db_summaries,
3167 deferred_db_trust_levels: &mut self.deferred_db_trust_levels,
3168 cached_prompt_tokens: &mut self.cached_prompt_tokens,
3169 context_manager: &mut self.context_manager,
3170 server_compaction_active: false,
3171 token_counter: Arc::clone(&self.token_counter),
3172 summarization_deps: SummarizationDeps {
3173 provider: AnyProvider::Mock(MockProvider::with_responses(
3174 self.provider_responses.clone(),
3175 )),
3176 llm_timeout: Duration::from_secs(30),
3177 token_counter: token_counter_adapter,
3178 structured_summaries: false,
3179 on_anchored_summary: None,
3180 },
3181 task_supervisor: Arc::clone(&self.task_supervisor),
3182 memory: None,
3183 conversation_id: None,
3184 tool_call_cutoff: 100,
3185 subgoal_registry: &mut self.subgoal_registry,
3186 pending_task_goal: &mut self.pending_task_goal,
3187 pending_subgoal: &mut self.pending_subgoal,
3188 current_task_goal: &mut self.current_task_goal,
3189 task_goal_user_msg_hash: &mut self.task_goal_user_msg_hash,
3190 subgoal_user_msg_hash: &mut self.subgoal_user_msg_hash,
3191 status_tx: None,
3192 scrub: |s| std::borrow::Cow::Borrowed(s),
3193 compression_guidelines: None,
3194 probe: None,
3195 archive: None,
3196 persistence: None,
3197 metrics: None,
3198 typed_pages: None,
3199 fidelity_config: None,
3200 fidelity_semantic_provider: None,
3201 fidelity_compress_provider: None,
3202 current_query: String::new(),
3203 }
3204 }
3205 }
3206
3207 struct NoopStatus;
3208 impl StatusSink for NoopStatus {
3209 fn send_status(&self, _msg: &str) -> impl std::future::Future<Output = ()> + Send + '_ {
3210 std::future::ready(())
3211 }
3212 }
3213
3214 fn plain_msg(role: Role, content: &str) -> Message {
3215 Message {
3216 role,
3217 content: content.to_owned(),
3218 parts: vec![],
3219 metadata: MessageMetadata::default(),
3220 }
3221 }
3222
3223 fn tool_use_msg() -> Message {
3224 Message::from_parts(
3225 Role::Assistant,
3226 vec![MessagePart::ToolUse {
3227 id: "t1".into(),
3228 name: "shell".into(),
3229 input: serde_json::json!({}),
3230 }],
3231 )
3232 }
3233
3234 fn tool_output_msg(body: &str) -> Message {
3235 Message::from_parts(
3236 Role::User,
3237 vec![MessagePart::ToolOutput {
3238 tool_name: "shell".into(),
3239 body: body.to_owned(),
3240 compacted_at: None,
3241 }],
3242 )
3243 }
3244
3245 fn messages_with_one_tool_output(body: &str) -> Vec<Message> {
3248 vec![
3249 plain_msg(Role::System, "system"),
3250 tool_use_msg(),
3251 tool_output_msg(body),
3252 plain_msg(Role::User, "hello"),
3253 ]
3254 }
3255
3256 #[test]
3257 fn prune_tool_outputs_decrements_cached_tokens_by_exact_freed_amount() {
3258 let body = "large tool output ".repeat(200);
3259 let expected_freed = TokenCounter::default().count_tokens(&body);
3260 let mut ctx_mgr = ContextManager::new();
3261 ctx_mgr.prune_protect_tokens = 0;
3262
3263 let initial_tokens = 10_000u64;
3264 let mut fixture = Fixture::new(
3265 messages_with_one_tool_output(&body),
3266 initial_tokens,
3267 ctx_mgr,
3268 );
3269 let mut view = fixture.view();
3270
3271 let freed =
3272 crate::summarization::pruning::prune_tool_outputs(&mut view, expected_freed);
3273
3274 assert_eq!(
3275 freed, expected_freed,
3276 "returned freed amount must match the tool output's token count"
3277 );
3278 assert_eq!(
3279 *view.cached_prompt_tokens,
3280 initial_tokens - u64::try_from(freed).unwrap(),
3281 "cached_prompt_tokens must decrease by exactly the freed amount"
3282 );
3283 }
3284
3285 #[test]
3286 fn prune_tool_outputs_noop_leaves_cached_tokens_unchanged() {
3287 let mut ctx_mgr = ContextManager::new();
3288 ctx_mgr.prune_protect_tokens = 0;
3289
3290 let messages = vec![
3291 plain_msg(Role::System, "system"),
3292 plain_msg(Role::User, "hello"),
3293 ];
3294 let initial_tokens = 500u64;
3295 let mut fixture = Fixture::new(messages, initial_tokens, ctx_mgr);
3296 let mut view = fixture.view();
3297
3298 let freed = crate::summarization::pruning::prune_tool_outputs(&mut view, 100);
3299
3300 assert_eq!(freed, 0, "no ToolOutput parts exist to prune");
3301 assert_eq!(
3302 *view.cached_prompt_tokens, initial_tokens,
3303 "cached_prompt_tokens must be untouched when nothing is freed"
3304 );
3305 }
3306
3307 #[tokio::test]
3308 async fn do_soft_compaction_via_maybe_compact_reflects_real_freed_tokens() {
3309 let body = "large tool output ".repeat(200);
3310 let expected_freed = TokenCounter::default().count_tokens(&body);
3311
3312 let mut ctx_mgr = ContextManager::new();
3313 ctx_mgr.budget = Some(ContextBudget::new(1000, 0.2));
3314 ctx_mgr.soft_compaction_threshold = 0.5; ctx_mgr.hard_compaction_threshold = 0.9; ctx_mgr.prune_protect_tokens = 0;
3317
3318 let initial_tokens = 700u64; assert_eq!(
3320 ctx_mgr.compaction_tier(initial_tokens),
3321 CompactionTier::Soft
3322 );
3323
3324 let mut fixture = Fixture::new(
3325 messages_with_one_tool_output(&body),
3326 initial_tokens,
3327 ctx_mgr,
3328 );
3329 let mut view = fixture.view();
3330 let status = NoopStatus;
3331
3332 ContextService::new()
3333 .maybe_compact(&mut view, &status)
3334 .await
3335 .unwrap();
3336
3337 assert!(
3338 view.deferred_db_summaries.is_empty(),
3339 "no deferred summaries were queued in this scenario"
3340 );
3341 assert_eq!(
3342 *view.cached_prompt_tokens,
3343 initial_tokens - u64::try_from(expected_freed).unwrap(),
3344 "Soft-tier prune-only pass must decrement cached_prompt_tokens by exactly \
3345 the freed amount"
3346 );
3347 }
3348
3349 #[tokio::test]
3350 async fn do_hard_compaction_satisfied_by_pruning_alone_reflects_real_freed_tokens() {
3351 let body = "large tool output ".repeat(400);
3355 let expected_freed = TokenCounter::default().count_tokens(&body);
3356
3357 let mut ctx_mgr = ContextManager::new();
3358 ctx_mgr.budget = Some(ContextBudget::new(1_000_000, 0.2));
3359 ctx_mgr.soft_compaction_threshold = 0.5;
3360 ctx_mgr.hard_compaction_threshold = 0.6; ctx_mgr.prune_protect_tokens = 0;
3362
3363 let initial_tokens = 600_000u64 + (expected_freed as u64 / 2);
3366 assert_eq!(
3367 ctx_mgr.compaction_tier(initial_tokens),
3368 CompactionTier::Hard
3369 );
3370
3371 let mut fixture = Fixture::new(
3372 messages_with_one_tool_output(&body),
3373 initial_tokens,
3374 ctx_mgr,
3375 );
3376 let mut view = fixture.view();
3377 let status = NoopStatus;
3378
3379 ContextService::new()
3380 .maybe_compact(&mut view, &status)
3381 .await
3382 .unwrap();
3383
3384 assert!(
3385 view.deferred_db_summaries.is_empty(),
3386 "no deferred summaries were queued in this scenario"
3387 );
3388 assert_eq!(
3389 *view.cached_prompt_tokens,
3390 initial_tokens - u64::try_from(expected_freed).unwrap(),
3391 "Hard-tier pass satisfied by pruning alone must decrement cached_prompt_tokens \
3392 by exactly the freed amount"
3393 );
3394 assert!(
3395 view.context_manager
3396 .compaction_state()
3397 .is_compacted_this_turn(),
3398 "pruning-satisfied Hard tier must still mark the turn as compacted"
3399 );
3400 }
3401
3402 #[tokio::test]
3427 async fn do_hard_compaction_combined_prune_and_llm_savings_avoid_false_exhaustion() {
3428 let counter = TokenCounter::default();
3429
3430 let large_body = "large tool output content ".repeat(400);
3434 let summary_response = "ok".to_string();
3435
3436 let build_messages = || {
3437 vec![
3438 plain_msg(Role::System, "system"),
3439 tool_use_msg(),
3440 tool_output_msg(&large_body),
3441 plain_msg(Role::User, "t0"),
3442 plain_msg(Role::User, "t1"),
3443 plain_msg(Role::User, "t2"),
3444 ]
3445 };
3446
3447 let initial_tokens: u64 = build_messages()
3448 .iter()
3449 .map(|m| counter.count_message_tokens(m) as u64)
3450 .sum();
3451 let non_body_overhead =
3452 initial_tokens - u64::try_from(counter.count_tokens(&large_body)).unwrap();
3453 let budget_tokens = usize::try_from(initial_tokens).unwrap();
3454
3455 let make_ctx_mgr = |hard_ratio: f32| {
3456 let mut cm = ContextManager::new();
3457 cm.prune_protect_tokens = 0;
3458 cm.compaction_preserve_tail = 2;
3459 cm.budget = Some(ContextBudget::new(budget_tokens, 0.0));
3460 cm.hard_compaction_threshold = hard_ratio;
3461 cm
3462 };
3463
3464 let measured_final_tokens = {
3469 let mut fixture =
3470 Fixture::new(build_messages(), initial_tokens, make_ctx_mgr(0.0001));
3471 fixture.provider_responses = vec![summary_response.clone()];
3472 let mut view = fixture.view();
3473 ContextService::new()
3474 .do_hard_compaction(&mut view, &NoopStatus, false)
3475 .await
3476 .unwrap();
3477 *view.cached_prompt_tokens
3478 };
3479 assert!(
3480 measured_final_tokens < non_body_overhead,
3481 "test setup invariant: the LLM step must reduce tokens below the non-body \
3482 overhead ({non_body_overhead}) so a valid hard-threshold window exists; \
3483 measured {measured_final_tokens}"
3484 );
3485
3486 let hard_threshold_tokens = u64::midpoint(measured_final_tokens, non_body_overhead);
3491 #[allow(clippy::cast_precision_loss)]
3492 let hard_ratio = hard_threshold_tokens as f32 / budget_tokens as f32;
3493
3494 let mut fixture =
3495 Fixture::new(build_messages(), initial_tokens, make_ctx_mgr(hard_ratio));
3496 fixture.provider_responses = vec![summary_response];
3497 let mut view = fixture.view();
3498
3499 ContextService::new()
3500 .do_hard_compaction(&mut view, &NoopStatus, false)
3501 .await
3502 .unwrap();
3503
3504 assert!(
3505 !view.context_manager.compaction_state().is_exhausted(),
3506 "pruning's own savings must count toward the combined freed-tokens check, so \
3507 a Hard-tier pass where pruning alone was insufficient but prune+LLM combined \
3508 cross the hard threshold must not be marked Exhausted"
3509 );
3510 assert!(
3511 view.context_manager
3512 .compaction_state()
3513 .is_compacted_this_turn(),
3514 "the turn must be marked compacted given the combined prune + LLM reduction"
3515 );
3516 }
3517 }
3518}