1use std::future::Future;
17use std::pin::Pin;
18
19use futures::StreamExt as _;
20use futures::stream::FuturesUnordered;
21
22use zeph_common::memory::{
23 AsyncMemoryRouter, CompressionLevel, FunctionalType, GraphRecallParams, TokenCounting,
24};
25use zeph_llm::provider::{Message, MessageMetadata, MessagePart, Role};
26
27use crate::error::AssemblerError;
28use crate::input::ContextAssemblyInput;
29use crate::slot::ContextSlot;
30
31pub(crate) fn levels_to_flags(levels: &[CompressionLevel]) -> (bool, bool, bool) {
39 if levels.is_empty() {
40 return (true, true, true);
41 }
42 let episodic = levels.contains(&CompressionLevel::Episodic);
43 let procedural = levels.contains(&CompressionLevel::Procedural);
44 let declarative = levels.contains(&CompressionLevel::Declarative);
45 (episodic, procedural, declarative)
46}
47
48pub(crate) fn type_active(active: &[FunctionalType], t: FunctionalType) -> bool {
57 active.is_empty() || active.contains(&t)
58}
59
60pub const SUMMARY_PREFIX: &str = "[conversation summaries]\n";
62pub const CROSS_SESSION_PREFIX: &str = "[cross-session context]\n";
64pub const RECALL_PREFIX: &str = "[semantic recall]\n";
66pub const CORRECTIONS_PREFIX: &str = "[past corrections]\n";
68pub const DOCUMENT_RAG_PREFIX: &str = "## Relevant documents\n";
70pub const GRAPH_FACTS_PREFIX: &str = "[known facts]\n";
72
73const MEMORY_FETCH_TIMEOUT_MS: u64 = 1000;
81
82#[derive(Default)]
87pub struct PreparedContext {
88 pub graph_facts: Option<Message>,
90 pub doc_rag: Option<Message>,
92 pub corrections: Option<Message>,
94 pub recall: Option<Message>,
96 pub recall_confidence: Option<f32>,
98 pub cross_session: Option<Message>,
100 pub summaries: Option<Message>,
102 pub code_context: Option<String>,
104 pub persona_facts: Option<Message>,
106 pub trajectory_hints: Option<Message>,
108 pub tree_memory: Option<Message>,
110 pub reasoning_hints: Option<Message>,
112 pub memory_first: bool,
114 pub recent_history_budget: usize,
116 pub background_tasks: Vec<tokio::task::JoinHandle<()>>,
121}
122
123pub struct ContextAssembler;
127
128type CtxFuture<'a> = Pin<Box<dyn Future<Output = Result<ContextSlot, AssemblerError>> + Send + 'a>>;
129
130fn empty_prepared_context() -> PreparedContext {
131 PreparedContext::default()
132}
133
134fn resolve_effective_strategy(
135 memory: &crate::input::ContextMemoryView,
136 sidequest_turn_counter: u64,
137) -> zeph_config::ContextStrategy {
138 match memory.context_strategy {
139 zeph_config::ContextStrategy::MemoryFirst => zeph_config::ContextStrategy::MemoryFirst,
140 zeph_config::ContextStrategy::Adaptive => {
141 if sidequest_turn_counter >= u64::from(memory.crossover_turn_threshold) {
142 zeph_config::ContextStrategy::MemoryFirst
143 } else {
144 zeph_config::ContextStrategy::FullHistory
145 }
146 }
147 _ => zeph_config::ContextStrategy::FullHistory,
148 }
149}
150
151fn correction_params(cfg: Option<&crate::input::CorrectionConfig>) -> (usize, f32) {
152 cfg.filter(|c| c.correction_detection)
153 .map_or((3, 0.75), |c| {
154 (
155 c.correction_recall_limit as usize,
156 c.correction_min_similarity,
157 )
158 })
159}
160
161#[allow(clippy::too_many_arguments, clippy::too_many_lines)]
168fn schedule_context_fetchers<'r>(
169 memory: &'r crate::input::ContextMemoryView,
170 tc: &'r dyn TokenCounting,
171 query: &'r str,
172 scrub: fn(&str) -> std::borrow::Cow<'_, str>,
173 index: Option<&'r dyn crate::input::IndexAccess>,
174 router_ref: &'r dyn AsyncMemoryRouter,
175 summaries_budget: usize,
176 cross_session_budget: usize,
177 semantic_recall_budget: usize,
178 code_context_budget: usize,
179 graph_facts_budget: usize,
180 recall_limit: usize,
181 min_sim: f32,
182 active_levels: &[CompressionLevel],
183 active_types: &[FunctionalType],
184) -> FuturesUnordered<CtxFuture<'r>> {
185 let (episodic_active, procedural_active, declarative_active) = levels_to_flags(active_levels);
195
196 let fetchers: FuturesUnordered<CtxFuture<'r>> = FuturesUnordered::new();
197
198 if episodic_active
199 && summaries_budget > 0
200 && type_active(active_types, FunctionalType::CrossSessionSummary)
201 {
202 fetchers.push(Box::pin(async move {
203 fetch_summaries(memory, summaries_budget, tc)
204 .await
205 .map(ContextSlot::Summaries)
206 }));
207 }
208 if episodic_active
209 && cross_session_budget > 0
210 && type_active(active_types, FunctionalType::CrossSessionSummary)
211 {
212 fetchers.push(Box::pin(async move {
213 fetch_cross_session(memory, query, cross_session_budget, tc)
214 .await
215 .map(ContextSlot::CrossSession)
216 }));
217 }
218 if episodic_active
219 && semantic_recall_budget > 0
220 && type_active(active_types, FunctionalType::Episodic)
221 {
222 fetchers.push(Box::pin(async move {
223 fetch_semantic_recall(memory, query, semantic_recall_budget, tc, Some(router_ref))
224 .await
225 .map(|(msg, score)| ContextSlot::SemanticRecall(msg, score))
226 }));
227 }
228 if episodic_active && semantic_recall_budget > 0 {
233 fetchers.push(Box::pin(async move {
234 fetch_document_rag(memory, query, semantic_recall_budget, tc)
235 .await
236 .map(ContextSlot::DocumentRag)
237 }));
238 }
239 fetchers.push(Box::pin(async move {
241 fetch_corrections(memory, query, recall_limit, min_sim, scrub)
242 .await
243 .map(ContextSlot::Corrections)
244 }));
245 if code_context_budget > 0
247 && let Some(idx) = index
248 {
249 fetchers.push(Box::pin(async move {
250 let result: Result<Option<String>, AssemblerError> = if let Ok(r) =
251 tokio::time::timeout(
252 std::time::Duration::from_millis(MEMORY_FETCH_TIMEOUT_MS),
253 idx.fetch_code_rag(query, code_context_budget),
254 )
255 .await
256 {
257 r
258 } else {
259 tracing::warn!("code RAG fetch timed out ({MEMORY_FETCH_TIMEOUT_MS}ms)");
260 Ok(None)
261 };
262 result.map(ContextSlot::CodeContext)
263 }));
264 }
265 if declarative_active
266 && graph_facts_budget > 0
267 && type_active(active_types, FunctionalType::GraphFact)
268 {
269 fetchers.push(Box::pin(async move {
270 fetch_graph_facts(memory, query, graph_facts_budget, tc)
271 .await
272 .map(ContextSlot::GraphFacts)
273 }));
274 }
275 if declarative_active
276 && memory.persona_config.context_budget_tokens > 0
277 && type_active(active_types, FunctionalType::UserFact)
278 {
279 fetchers.push(Box::pin(async move {
280 let persona_budget = memory.persona_config.context_budget_tokens;
281 fetch_persona_facts(memory, persona_budget, tc)
282 .await
283 .map(ContextSlot::PersonaFacts)
284 }));
285 }
286 if procedural_active && memory.trajectory_config.context_budget_tokens > 0 {
289 fetchers.push(Box::pin(async move {
290 let tbudget = memory.trajectory_config.context_budget_tokens;
291 fetch_trajectory_hints(memory, tbudget, tc)
292 .await
293 .map(ContextSlot::TrajectoryHints)
294 }));
295 }
296 if declarative_active && memory.tree_config.context_budget_tokens > 0 {
299 fetchers.push(Box::pin(async move {
300 let tbudget = memory.tree_config.context_budget_tokens;
301 fetch_tree_memory(memory, tbudget, tc)
302 .await
303 .map(ContextSlot::TreeMemory)
304 }));
305 }
306 if procedural_active
307 && memory.reasoning_config.enabled
308 && memory.reasoning_config.context_budget_tokens > 0
309 && type_active(active_types, FunctionalType::ReasoningStrategy)
310 {
311 fetchers.push(Box::pin(async move {
312 let rbudget = memory.reasoning_config.context_budget_tokens;
313 let top_k = memory.reasoning_config.top_k;
314 fetch_reasoning_strategies(memory, query, rbudget, top_k, tc)
315 .await
316 .map(|(msg, handle)| ContextSlot::ReasoningStrategies(msg, handle))
317 }));
318 }
319
320 fetchers
321}
322
323async fn drive_fetchers(
324 mut fetchers: FuturesUnordered<CtxFuture<'_>>,
325 prepared: &mut PreparedContext,
326) -> Result<(), AssemblerError> {
327 while let Some(result) = fetchers.next().await {
328 match result {
329 Ok(slot) => match slot {
330 ContextSlot::Summaries(msg) => prepared.summaries = msg,
331 ContextSlot::CrossSession(msg) => prepared.cross_session = msg,
332 ContextSlot::SemanticRecall(msg, score) => {
333 prepared.recall = msg;
334 prepared.recall_confidence = score;
335 }
336 ContextSlot::DocumentRag(msg) => prepared.doc_rag = msg,
337 ContextSlot::Corrections(msg) => prepared.corrections = msg,
338 ContextSlot::CodeContext(text) => prepared.code_context = text,
339 ContextSlot::GraphFacts(msg) => prepared.graph_facts = msg,
340 ContextSlot::PersonaFacts(msg) => prepared.persona_facts = msg,
341 ContextSlot::TrajectoryHints(msg) => prepared.trajectory_hints = msg,
342 ContextSlot::TreeMemory(msg) => prepared.tree_memory = msg,
343 ContextSlot::ReasoningStrategies(msg, handle) => {
344 prepared.reasoning_hints = msg;
345 if let Some(h) = handle {
346 prepared.background_tasks.push(h);
347 }
348 }
349 },
350 Err(e) => return Err(e),
351 }
352 }
353 Ok(())
354}
355
356impl ContextAssembler {
357 #[tracing::instrument(
365 name = "context.assembler.gather",
366 skip_all,
367 fields(active_types = ?input.active_types)
368 )]
369 pub async fn gather(
370 input: &ContextAssemblyInput<'_>,
371 ) -> Result<PreparedContext, AssemblerError> {
372 let Some(ref budget) = input.context_manager.budget else {
373 return Ok(empty_prepared_context());
374 };
375
376 let memory = input.memory;
377 let tc = input.token_counter;
378
379 let effective_strategy = resolve_effective_strategy(memory, input.sidequest_turn_counter);
380 let memory_first = effective_strategy == zeph_config::ContextStrategy::MemoryFirst;
381
382 let system_prompt = input
383 .messages
384 .first()
385 .filter(|m| m.role == Role::System)
386 .map_or("", |m| m.content.as_str());
387
388 let digest_tokens = memory
389 .cached_session_digest
390 .as_ref()
391 .map_or(0, |(_, tokens)| *tokens);
392
393 let alloc = budget.allocate_with_opts(
394 system_prompt,
395 input.skills_prompt,
396 tc,
397 memory.graph_config.enabled,
398 digest_tokens,
399 memory_first,
400 );
401
402 let (recall_limit, min_sim) = correction_params(input.correction_config.as_ref());
403
404 let router_ref: &dyn AsyncMemoryRouter = input.router.as_ref();
405
406 tracing::debug!(
407 active_sources = alloc.active_sources(),
408 active_levels = ?input.active_levels,
409 "context budget allocated"
410 );
411
412 let fetchers = schedule_context_fetchers(
413 memory,
414 tc,
415 input.query,
416 input.scrub,
417 input.index,
418 router_ref,
419 alloc.summaries,
420 alloc.cross_session,
421 alloc.semantic_recall,
422 alloc.code_context,
423 alloc.graph_facts,
424 recall_limit,
425 min_sim,
426 input.active_levels,
427 input.active_types,
428 );
429
430 let mut prepared = empty_prepared_context();
431 prepared.memory_first = memory_first;
432 prepared.recent_history_budget = alloc.recent_history;
433
434 drive_fetchers(fetchers, &mut prepared).await?;
435 Ok(prepared)
436 }
437}
438
439pub fn effective_recall_timeout_ms(configured: u64) -> u64 {
444 if configured == 0 {
445 tracing::warn!(
446 "recall_timeout_ms is 0, which would disable spreading activation recall; \
447 clamping to 100ms"
448 );
449 100
450 } else {
451 configured
452 }
453}
454
455use crate::input::ContextMemoryView;
456
457#[tracing::instrument(name = "context.graph_facts", skip_all)]
458#[allow(clippy::too_many_lines)] pub(crate) async fn fetch_graph_facts(
460 memory: &ContextMemoryView,
461 query: &str,
462 budget_tokens: usize,
463 tc: &dyn TokenCounting,
464) -> Result<Option<Message>, AssemblerError> {
465 use zeph_common::memory::{RecallView, SpreadingActivationParams, classify_graph_subgraph};
466
467 if budget_tokens == 0 || !memory.graph_config.enabled {
468 return Ok(None);
469 }
470 let Some(ref mem) = memory.memory else {
471 return Ok(None);
472 };
473 let recall_limit = memory.graph_config.recall_limit;
474 let temporal_decay_rate = memory.graph_config.temporal_decay_rate;
475 let sa_config = &memory.graph_config.spreading_activation;
476
477 let fused_query;
479 let effective_query = if let Some(ref state) = memory.memcot_state {
480 let max_state_chars = 2 * query.len();
481 let state_slice = if state.len() > max_state_chars {
482 let boundary = state.floor_char_boundary(max_state_chars);
483 &state[..boundary]
484 } else {
485 state.as_str()
486 };
487 fused_query = format!("[state] {state_slice}\n{query}");
488 &fused_query as &str
489 } else {
490 query
491 };
492
493 let edge_types = classify_graph_subgraph(effective_query);
494
495 let view = match memory.memcot_config.recall_view {
496 zeph_config::RecallViewConfig::ZoomIn => RecallView::ZoomIn,
497 zeph_config::RecallViewConfig::ZoomOut => RecallView::ZoomOut,
498 _ => RecallView::Head,
499 };
500
501 let sa_params = if sa_config.enabled {
502 Some(SpreadingActivationParams {
503 decay_lambda: sa_config.decay_lambda,
504 max_hops: sa_config.max_hops,
505 activation_threshold: sa_config.activation_threshold,
506 inhibition_threshold: sa_config.inhibition_threshold,
507 max_activated_nodes: sa_config.max_activated_nodes,
508 temporal_decay_rate,
509 seed_structural_weight: sa_config.seed_structural_weight,
510 seed_community_cap: sa_config.seed_community_cap,
511 alpha: sa_config.alpha,
512 })
513 } else {
514 None
515 };
516
517 let timeout_ms = effective_recall_timeout_ms(sa_config.recall_timeout_ms);
518 let recall_fut = mem.recall_graph_facts(
519 effective_query,
520 GraphRecallParams {
521 limit: recall_limit,
522 view,
523 zoom_out_neighbor_cap: memory.memcot_config.zoom_out_neighbor_cap,
524 max_hops: memory.graph_config.max_hops,
525 temporal_decay_rate,
526 edge_types: &edge_types,
527 spreading_activation: sa_params,
528 },
529 );
530 let recalled = match tokio::time::timeout(
531 std::time::Duration::from_millis(timeout_ms),
532 recall_fut,
533 )
534 .await
535 {
536 Ok(Ok(facts)) => facts,
537 Ok(Err(e)) => {
538 tracing::warn!("graph recall failed: {e:#}");
539 Vec::new()
540 }
541 Err(_) => {
542 tracing::warn!("graph recall timed out ({timeout_ms}ms)");
543 Vec::new()
544 }
545 };
546
547 if recalled.is_empty() {
548 return Ok(None);
549 }
550
551 let mut body = String::from(GRAPH_FACTS_PREFIX);
552 let mut tokens_so_far = tc.count_tokens(&body);
553
554 for rf in &recalled {
555 let fact_text = rf.fact.replace(['\n', '\r', '<', '>'], " ");
556 let line = if let Some(score) = rf.activation_score {
557 format!(
558 "- {} (confidence: {:.2}, activation: {:.2})\n",
559 fact_text, rf.confidence, score
560 )
561 } else {
562 format!("- {} (confidence: {:.2})\n", fact_text, rf.confidence)
563 };
564 let line_tokens = tc.count_tokens(&line);
565 if tokens_so_far + line_tokens > budget_tokens {
566 break;
567 }
568 body.push_str(&line);
569 tokens_so_far += line_tokens;
570
571 for nb in &rf.neighbors {
573 let nb_text = nb.fact.replace(['\n', '\r', '<', '>'], " ");
574 let nb_line = format!(" ~ {} (confidence: {:.2})\n", nb_text, nb.confidence);
575 let nb_tokens = tc.count_tokens(&nb_line);
576 if tokens_so_far + nb_tokens > budget_tokens {
577 break;
578 }
579 body.push_str(&nb_line);
580 tokens_so_far += nb_tokens;
581 }
582
583 if let Some(ref snippet) = rf.provenance_snippet {
585 let snip_line = format!(
586 " [source: {}]\n",
587 snippet.replace(['\n', '\r', '<', '>'], " ")
588 );
589 let snip_tokens = tc.count_tokens(&snip_line);
590 if tokens_so_far + snip_tokens <= budget_tokens {
591 body.push_str(&snip_line);
592 tokens_so_far += snip_tokens;
593 }
594 }
595 }
596
597 if body == GRAPH_FACTS_PREFIX {
598 return Ok(None);
599 }
600
601 Ok(Some(Message::from_legacy(Role::System, body)))
602}
603
604fn append_budgeted_lines(
610 prefix: &str,
611 lines: impl Iterator<Item = String>,
612 budget_tokens: usize,
613 tc: &dyn TokenCounting,
614) -> Option<String> {
615 let mut body = String::from(prefix);
616 let mut tokens_so_far = tc.count_tokens(&body);
617
618 for line in lines {
619 let line_tokens = tc.count_tokens(&line);
620 if tokens_so_far + line_tokens > budget_tokens {
621 break;
622 }
623 body.push_str(&line);
624 tokens_so_far += line_tokens;
625 }
626
627 if body == prefix { None } else { Some(body) }
628}
629
630#[tracing::instrument(name = "context.persona_facts", skip_all)]
631pub(crate) async fn fetch_persona_facts(
632 memory: &ContextMemoryView,
633 budget_tokens: usize,
634 tc: &dyn TokenCounting,
635) -> Result<Option<Message>, AssemblerError> {
636 if budget_tokens == 0 || !memory.persona_config.enabled {
637 return Ok(None);
638 }
639 let Some(ref mem) = memory.memory else {
640 return Ok(None);
641 };
642
643 let min_confidence = memory.persona_config.min_confidence;
644 let facts = if let Ok(result) = tokio::time::timeout(
645 std::time::Duration::from_millis(MEMORY_FETCH_TIMEOUT_MS),
646 mem.load_persona_facts(min_confidence),
647 )
648 .await
649 {
650 result.map_err(AssemblerError::Memory)?
651 } else {
652 tracing::warn!("persona facts load timed out ({MEMORY_FETCH_TIMEOUT_MS}ms)");
653 Vec::new()
654 };
655
656 if facts.is_empty() {
657 return Ok(None);
658 }
659
660 let lines = facts
661 .iter()
662 .map(|fact| format!("[{}] {}\n", fact.category, fact.content));
663 Ok(
664 append_budgeted_lines(crate::slot::PERSONA_PREFIX, lines, budget_tokens, tc)
665 .map(|body| Message::from_legacy(Role::System, body)),
666 )
667}
668
669#[tracing::instrument(name = "context.trajectory_hints", skip_all)]
670pub(crate) async fn fetch_trajectory_hints(
671 memory: &ContextMemoryView,
672 budget_tokens: usize,
673 tc: &dyn TokenCounting,
674) -> Result<Option<Message>, AssemblerError> {
675 if budget_tokens == 0 || !memory.trajectory_config.enabled {
676 return Ok(None);
677 }
678 let Some(ref mem) = memory.memory else {
679 return Ok(None);
680 };
681
682 let top_k = memory.trajectory_config.recall_top_k;
683 let min_conf = memory.trajectory_config.min_confidence;
684 let entries = if let Ok(result) = tokio::time::timeout(
688 std::time::Duration::from_millis(MEMORY_FETCH_TIMEOUT_MS),
689 mem.load_trajectory_entries(Some("procedural"), top_k),
690 )
691 .await
692 {
693 result.map_err(AssemblerError::Memory)?
694 } else {
695 tracing::warn!("trajectory entries load timed out ({MEMORY_FETCH_TIMEOUT_MS}ms)");
696 Vec::new()
697 };
698
699 if entries.is_empty() {
700 return Ok(None);
701 }
702
703 let lines = entries
704 .iter()
705 .filter(|e| e.confidence >= min_conf)
706 .take(top_k)
707 .map(|entry| format!("- {}: {}\n", entry.intent, entry.outcome));
708 Ok(
709 append_budgeted_lines(crate::slot::TRAJECTORY_PREFIX, lines, budget_tokens, tc)
710 .map(|body| Message::from_legacy(Role::System, body)),
711 )
712}
713
714#[tracing::instrument(name = "context.tree_memory", skip_all)]
715pub(crate) async fn fetch_tree_memory(
716 memory: &ContextMemoryView,
717 budget_tokens: usize,
718 tc: &dyn TokenCounting,
719) -> Result<Option<Message>, AssemblerError> {
720 if budget_tokens == 0 || !memory.tree_config.enabled {
721 return Ok(None);
722 }
723 let Some(ref mem) = memory.memory else {
724 return Ok(None);
725 };
726
727 let top_k = memory.tree_config.recall_top_k;
728 let nodes = if let Ok(result) = tokio::time::timeout(
729 std::time::Duration::from_millis(MEMORY_FETCH_TIMEOUT_MS),
730 mem.load_tree_nodes(1, top_k),
731 )
732 .await
733 {
734 result.map_err(AssemblerError::Memory)?
735 } else {
736 tracing::warn!("tree nodes load timed out ({MEMORY_FETCH_TIMEOUT_MS}ms)");
737 Vec::new()
738 };
739
740 if nodes.is_empty() {
741 return Ok(None);
742 }
743
744 let lines = nodes
745 .iter()
746 .take(top_k)
747 .map(|node| format!("- {}\n", node.content));
748 Ok(
749 append_budgeted_lines(crate::slot::TREE_MEMORY_PREFIX, lines, budget_tokens, tc)
750 .map(|body| Message::from_legacy(Role::System, body)),
751 )
752}
753
754#[tracing::instrument(name = "context.reasoning_strategies", skip_all)]
755pub(crate) async fn fetch_reasoning_strategies(
756 memory: &ContextMemoryView,
757 query: &str,
758 budget_tokens: usize,
759 top_k: usize,
760 tc: &dyn TokenCounting,
761) -> Result<(Option<Message>, Option<tokio::task::JoinHandle<()>>), AssemblerError> {
762 let budget_tokens = budget_tokens.min(500);
764 if budget_tokens == 0 {
765 return Ok((None, None));
766 }
767 let Some(ref mem) = memory.memory else {
768 return Ok((None, None));
769 };
770
771 let strategies = if let Ok(result) = tokio::time::timeout(
772 std::time::Duration::from_millis(MEMORY_FETCH_TIMEOUT_MS),
773 mem.retrieve_reasoning_strategies(query, top_k),
774 )
775 .await
776 {
777 result.map_err(AssemblerError::Memory)?
778 } else {
779 tracing::warn!("reasoning strategies retrieval timed out ({MEMORY_FETCH_TIMEOUT_MS}ms)");
780 Vec::new()
781 };
782
783 if strategies.is_empty() {
784 return Ok((None, None));
785 }
786
787 let mut body = String::from(crate::slot::REASONING_PREFIX);
788 let mut tokens_so_far = tc.count_tokens(&body);
789 let mut injected_ids: Vec<String> = Vec::new();
790
791 for s in strategies.iter().take(top_k) {
792 let safe_summary = s.summary.replace(['\n', '\r', '<', '>'], " ");
795 let line = format!("- [{}] {}\n", s.outcome, safe_summary);
796 let line_tokens = tc.count_tokens(&line);
797 if tokens_so_far + line_tokens > budget_tokens {
798 break;
799 }
800 body.push_str(&line);
801 tokens_so_far += line_tokens;
802 injected_ids.push(s.id.clone());
803 }
804
805 if body == crate::slot::REASONING_PREFIX {
806 return Ok((None, None));
807 }
808
809 let handle = if injected_ids.is_empty() {
813 None
814 } else {
815 let mem_clone = mem.clone();
816 let mark_used = async move {
817 if let Err(e) = mem_clone.mark_reasoning_used(&injected_ids).await {
818 tracing::warn!(error = %e, "reasoning: mark_used failed");
819 }
820 };
821 Some(tokio::spawn(mark_used)) };
823
824 Ok((Some(Message::from_legacy(Role::System, body)), handle))
825}
826
827#[tracing::instrument(name = "context.corrections", skip_all)]
828pub(crate) async fn fetch_corrections(
829 memory: &ContextMemoryView,
830 query: &str,
831 limit: usize,
832 min_score: f32,
833 scrub: fn(&str) -> std::borrow::Cow<'_, str>,
834) -> Result<Option<Message>, AssemblerError> {
835 let Some(ref mem) = memory.memory else {
836 return Ok(None);
837 };
838 let corrections = if let Ok(result) = tokio::time::timeout(
839 std::time::Duration::from_millis(MEMORY_FETCH_TIMEOUT_MS),
840 mem.retrieve_corrections(query, limit, min_score),
841 )
842 .await
843 {
844 result.map_err(AssemblerError::Memory)?
845 } else {
846 tracing::warn!("corrections retrieval timed out ({MEMORY_FETCH_TIMEOUT_MS}ms)");
847 Vec::new()
848 };
849 if corrections.is_empty() {
850 return Ok(None);
851 }
852 let mut text = String::from(CORRECTIONS_PREFIX);
853 for c in &corrections {
854 text.push_str("- Past user correction: \"");
855 text.push_str(&scrub(&c.correction_text));
856 text.push_str("\"\n");
857 }
858 Ok(Some(Message::from_legacy(Role::System, text)))
859}
860
861#[tracing::instrument(name = "context.semantic_recall", skip_all)]
862pub(crate) async fn fetch_semantic_recall(
863 memory: &ContextMemoryView,
864 query: &str,
865 token_budget: usize,
866 tc: &dyn TokenCounting,
867 router: Option<&dyn AsyncMemoryRouter>,
868) -> Result<(Option<Message>, Option<f32>), AssemblerError> {
869 let Some(ref mem) = memory.memory else {
870 return Ok((None, None));
871 };
872 if memory.recall_limit == 0 || token_budget == 0 {
873 return Ok((None, None));
874 }
875
876 let recalled = if let Ok(result) = tokio::time::timeout(
877 std::time::Duration::from_millis(MEMORY_FETCH_TIMEOUT_MS),
878 mem.recall(query, memory.recall_limit, router),
879 )
880 .await
881 {
882 result.map_err(AssemblerError::Memory)?
883 } else {
884 tracing::warn!("semantic recall timed out ({MEMORY_FETCH_TIMEOUT_MS}ms)");
885 Vec::new()
886 };
887 if recalled.is_empty() {
888 return Ok((None, None));
889 }
890
891 let top_score = recalled.first().map(|r| r.score);
892
893 let mut recall_text = String::with_capacity(token_budget * 3);
894 recall_text.push_str(RECALL_PREFIX);
895 let mut tokens_used = tc.count_tokens(&recall_text);
896
897 for item in &recalled {
898 if item.content.starts_with("[skipped]") || item.content.starts_with("[stopped]") {
899 continue;
900 }
901 let entry = format!("- [{}] {}\n", item.role, item.content);
902 let entry_tokens = tc.count_tokens(&entry);
903 if tokens_used + entry_tokens > token_budget {
904 break;
905 }
906 recall_text.push_str(&entry);
907 tokens_used += entry_tokens;
908 }
909
910 if tokens_used > tc.count_tokens(RECALL_PREFIX) {
911 Ok((
912 Some(Message::from_parts(
913 Role::System,
914 vec![MessagePart::Recall { text: recall_text }],
915 )),
916 top_score,
917 ))
918 } else {
919 Ok((None, None))
920 }
921}
922
923#[tracing::instrument(name = "context.document_rag", skip_all)]
924pub(crate) async fn fetch_document_rag(
925 memory: &ContextMemoryView,
926 query: &str,
927 token_budget: usize,
928 tc: &dyn TokenCounting,
929) -> Result<Option<Message>, AssemblerError> {
930 if !memory.document_config.rag_enabled || token_budget == 0 {
931 return Ok(None);
932 }
933 let Some(ref mem) = memory.memory else {
934 return Ok(None);
935 };
936
937 let collection = &memory.document_config.collection;
938 let top_k = memory.document_config.top_k;
939 let chunks = if let Ok(result) = tokio::time::timeout(
940 std::time::Duration::from_millis(MEMORY_FETCH_TIMEOUT_MS),
941 mem.search_document_collection(collection, query, top_k),
942 )
943 .await
944 {
945 result.map_err(AssemblerError::Memory)?
946 } else {
947 tracing::warn!("document RAG search timed out ({MEMORY_FETCH_TIMEOUT_MS}ms)");
948 Vec::new()
949 };
950 if chunks.is_empty() {
951 return Ok(None);
952 }
953
954 let mut text = String::from(DOCUMENT_RAG_PREFIX);
955 let mut tokens_used = tc.count_tokens(&text);
956
957 for chunk in &chunks {
958 if chunk.text.is_empty() {
959 continue;
960 }
961 let entry = format!("{}\n", chunk.text);
962 let cost = tc.count_tokens(&entry);
963 if tokens_used + cost > token_budget {
964 break;
965 }
966 text.push_str(&entry);
967 tokens_used += cost;
968 }
969
970 if tokens_used > tc.count_tokens(DOCUMENT_RAG_PREFIX) {
971 Ok(Some(Message {
972 role: Role::System,
973 content: text,
974 parts: vec![],
975 metadata: MessageMetadata::default(),
976 }))
977 } else {
978 Ok(None)
979 }
980}
981
982#[tracing::instrument(name = "context.summaries", skip_all)]
983pub(crate) async fn fetch_summaries(
984 memory: &ContextMemoryView,
985 token_budget: usize,
986 tc: &dyn TokenCounting,
987) -> Result<Option<Message>, AssemblerError> {
988 let (Some(mem), Some(cid)) = (&memory.memory, memory.conversation_id) else {
989 return Ok(None);
990 };
991 if token_budget == 0 {
992 return Ok(None);
993 }
994
995 let summaries = if let Ok(result) = tokio::time::timeout(
996 std::time::Duration::from_millis(MEMORY_FETCH_TIMEOUT_MS),
997 mem.load_summaries(cid),
998 )
999 .await
1000 {
1001 result.map_err(AssemblerError::Memory)?
1002 } else {
1003 tracing::warn!("summaries load timed out ({MEMORY_FETCH_TIMEOUT_MS}ms)");
1004 Vec::new()
1005 };
1006 if summaries.is_empty() {
1007 return Ok(None);
1008 }
1009
1010 let mut summary_text = String::from(SUMMARY_PREFIX);
1011 let mut tokens_used = tc.count_tokens(&summary_text);
1012
1013 for summary in summaries.iter().rev() {
1014 let first = summary.first_message_id.unwrap_or(0);
1015 let last = summary.last_message_id.unwrap_or(0);
1016 let entry = format!("- Messages {first}-{last}: {}\n", summary.content);
1017 let cost = tc.count_tokens(&entry);
1018 if tokens_used + cost > token_budget {
1019 break;
1020 }
1021 summary_text.push_str(&entry);
1022 tokens_used += cost;
1023 }
1024
1025 if tokens_used > tc.count_tokens(SUMMARY_PREFIX) {
1026 Ok(Some(Message::from_parts(
1027 Role::System,
1028 vec![MessagePart::Summary { text: summary_text }],
1029 )))
1030 } else {
1031 Ok(None)
1032 }
1033}
1034
1035#[tracing::instrument(name = "context.cross_session", skip_all)]
1036pub(crate) async fn fetch_cross_session(
1037 memory: &ContextMemoryView,
1038 query: &str,
1039 token_budget: usize,
1040 tc: &dyn TokenCounting,
1041) -> Result<Option<Message>, AssemblerError> {
1042 let (Some(mem), Some(cid)) = (&memory.memory, memory.conversation_id) else {
1043 return Ok(None);
1044 };
1045 if token_budget == 0 {
1046 return Ok(None);
1047 }
1048
1049 let threshold = memory.cross_session_score_threshold;
1050 let summaries = if let Ok(result) = tokio::time::timeout(
1051 std::time::Duration::from_millis(MEMORY_FETCH_TIMEOUT_MS),
1052 mem.search_session_summaries(query, 5, Some(cid)),
1053 )
1054 .await
1055 {
1056 result.map_err(AssemblerError::Memory)?
1057 } else {
1058 tracing::warn!("cross-session search timed out ({MEMORY_FETCH_TIMEOUT_MS}ms)");
1059 Vec::new()
1060 };
1061 let results: Vec<_> = summaries
1062 .into_iter()
1063 .filter(|r| r.score >= threshold)
1064 .collect();
1065 if results.is_empty() {
1066 return Ok(None);
1067 }
1068
1069 let mut text = String::from(CROSS_SESSION_PREFIX);
1070 let mut tokens_used = tc.count_tokens(&text);
1071
1072 for item in &results {
1073 let entry = format!("- {}\n", item.summary_text);
1074 let cost = tc.count_tokens(&entry);
1075 if tokens_used + cost > token_budget {
1076 break;
1077 }
1078 text.push_str(&entry);
1079 tokens_used += cost;
1080 }
1081
1082 if tokens_used > tc.count_tokens(CROSS_SESSION_PREFIX) {
1083 Ok(Some(Message::from_parts(
1084 Role::System,
1085 vec![MessagePart::CrossSession { text }],
1086 )))
1087 } else {
1088 Ok(None)
1089 }
1090}
1091
1092pub const MAX_KEEP_TAIL_SCAN: usize = 50;
1095
1096#[must_use]
1104pub fn memory_first_keep_tail(messages: &[Message], history_start: usize) -> usize {
1105 use zeph_llm::provider::MessagePart;
1106
1107 let mut keep_tail = 2usize;
1108 let len = messages.len();
1109 let max = len.saturating_sub(history_start);
1110
1111 while keep_tail < max {
1112 let first_retained = &messages[len - keep_tail];
1113 let is_tool_result = first_retained.role == Role::User
1114 && first_retained
1115 .parts
1116 .iter()
1117 .any(|p| matches!(p, MessagePart::ToolResult { .. }));
1118
1119 if is_tool_result {
1120 keep_tail += 1;
1121 } else {
1122 break;
1123 }
1124
1125 if keep_tail >= MAX_KEEP_TAIL_SCAN {
1126 let preceding_idx = len.saturating_sub(keep_tail + 1);
1127 if preceding_idx >= history_start {
1128 let preceding = &messages[preceding_idx];
1129 let is_tool_use = preceding.role == Role::Assistant
1130 && preceding
1131 .parts
1132 .iter()
1133 .any(|p| matches!(p, MessagePart::ToolUse { .. }));
1134 if is_tool_use {
1135 keep_tail += 1;
1136 }
1137 }
1138 break;
1139 }
1140 }
1141
1142 keep_tail
1143}
1144
1145#[cfg(test)]
1146mod tests {
1147 use super::*;
1148 use crate::input::ContextMemoryView;
1149 use zeph_common::memory::CompressionLevel;
1150 use zeph_config::{
1151 ContextStrategy, DocumentConfig, GraphConfig, PersonaConfig, ReasoningConfig,
1152 TrajectoryConfig, TreeConfig,
1153 };
1154
1155 struct NaiveTokenCounter;
1156 impl zeph_common::memory::TokenCounting for NaiveTokenCounter {
1157 fn count_tokens(&self, text: &str) -> usize {
1158 text.split_whitespace().count()
1159 }
1160 fn count_tool_schema_tokens(&self, schema: &serde_json::Value) -> usize {
1161 schema.to_string().split_whitespace().count()
1162 }
1163 }
1164
1165 fn empty_view() -> ContextMemoryView {
1166 ContextMemoryView {
1167 memory: None,
1168 conversation_id: None,
1169 recall_limit: 10,
1170 cross_session_score_threshold: 0.5,
1171 context_strategy: ContextStrategy::default(),
1172 crossover_turn_threshold: 5,
1173 cached_session_digest: None,
1174 graph_config: GraphConfig::default(),
1175 document_config: DocumentConfig::default(),
1176 persona_config: PersonaConfig::default(),
1177 trajectory_config: TrajectoryConfig::default(),
1178 reasoning_config: ReasoningConfig::default(),
1179 memcot_config: zeph_config::MemCotConfig::default(),
1180 memcot_state: None,
1181 tree_config: TreeConfig::default(),
1182 }
1183 }
1184
1185 #[tokio::test]
1188 async fn fetch_graph_facts_returns_none_when_memory_is_none() {
1189 let view = empty_view();
1190 let tc = NaiveTokenCounter;
1191 let result = fetch_graph_facts(&view, "test", 1000, &tc).await.unwrap();
1192 assert!(result.is_none());
1193 }
1194
1195 #[tokio::test]
1196 async fn fetch_graph_facts_returns_none_when_budget_zero() {
1197 let mut view = empty_view();
1198 view.graph_config.enabled = true;
1199 let tc = NaiveTokenCounter;
1200 let result = fetch_graph_facts(&view, "test", 0, &tc).await.unwrap();
1201 assert!(result.is_none());
1202 }
1203
1204 #[tokio::test]
1205 async fn fetch_graph_facts_returns_none_when_graph_disabled() {
1206 let mut view = empty_view();
1207 view.graph_config.enabled = false;
1208 let tc = NaiveTokenCounter;
1209 let result = fetch_graph_facts(&view, "test", 1000, &tc).await.unwrap();
1210 assert!(result.is_none());
1211 }
1212
1213 #[tokio::test]
1216 async fn fetch_persona_facts_returns_none_when_memory_is_none() {
1217 let view = empty_view();
1218 let tc = NaiveTokenCounter;
1219 let result = fetch_persona_facts(&view, 1000, &tc).await.unwrap();
1220 assert!(result.is_none());
1221 }
1222
1223 #[tokio::test]
1224 async fn fetch_persona_facts_returns_none_when_budget_zero() {
1225 let mut view = empty_view();
1226 view.persona_config.enabled = true;
1227 let tc = NaiveTokenCounter;
1228 let result = fetch_persona_facts(&view, 0, &tc).await.unwrap();
1229 assert!(result.is_none());
1230 }
1231
1232 #[tokio::test]
1235 async fn fetch_trajectory_hints_returns_none_when_memory_is_none() {
1236 let view = empty_view();
1237 let tc = NaiveTokenCounter;
1238 let result = fetch_trajectory_hints(&view, 1000, &tc).await.unwrap();
1239 assert!(result.is_none());
1240 }
1241
1242 #[tokio::test]
1243 async fn fetch_trajectory_hints_returns_none_when_budget_zero() {
1244 let mut view = empty_view();
1245 view.trajectory_config.enabled = true;
1246 let tc = NaiveTokenCounter;
1247 let result = fetch_trajectory_hints(&view, 0, &tc).await.unwrap();
1248 assert!(result.is_none());
1249 }
1250
1251 #[tokio::test]
1254 async fn fetch_tree_memory_returns_none_when_memory_is_none() {
1255 let view = empty_view();
1256 let tc = NaiveTokenCounter;
1257 let result = fetch_tree_memory(&view, 1000, &tc).await.unwrap();
1258 assert!(result.is_none());
1259 }
1260
1261 #[tokio::test]
1262 async fn fetch_tree_memory_returns_none_when_budget_zero() {
1263 let mut view = empty_view();
1264 view.tree_config.enabled = true;
1265 let tc = NaiveTokenCounter;
1266 let result = fetch_tree_memory(&view, 0, &tc).await.unwrap();
1267 assert!(result.is_none());
1268 }
1269
1270 #[tokio::test]
1273 async fn fetch_corrections_returns_none_when_memory_is_none() {
1274 let view = empty_view();
1275 let result = fetch_corrections(&view, "test", 10, 0.5, |s| s.into())
1276 .await
1277 .unwrap();
1278 assert!(result.is_none());
1279 }
1280
1281 #[tokio::test]
1284 async fn fetch_semantic_recall_returns_none_when_memory_is_none() {
1285 let view = empty_view();
1286 let tc = NaiveTokenCounter;
1287 let result = fetch_semantic_recall(&view, "test", 1000, &tc, None)
1288 .await
1289 .unwrap();
1290 assert!(result.0.is_none() && result.1.is_none());
1291 }
1292
1293 #[tokio::test]
1294 async fn fetch_semantic_recall_returns_none_when_budget_zero() {
1295 let view = empty_view();
1296 let tc = NaiveTokenCounter;
1297 let result = fetch_semantic_recall(&view, "test", 0, &tc, None)
1298 .await
1299 .unwrap();
1300 assert!(result.0.is_none() && result.1.is_none());
1301 }
1302
1303 #[tokio::test]
1306 async fn fetch_document_rag_returns_none_when_memory_is_none() {
1307 let mut view = empty_view();
1308 view.document_config.rag_enabled = true;
1309 let tc = NaiveTokenCounter;
1310 let result = fetch_document_rag(&view, "test", 1000, &tc).await.unwrap();
1311 assert!(result.is_none());
1312 }
1313
1314 #[tokio::test]
1315 async fn fetch_document_rag_returns_none_when_rag_disabled() {
1316 let view = empty_view();
1317 let tc = NaiveTokenCounter;
1318 let result = fetch_document_rag(&view, "test", 1000, &tc).await.unwrap();
1319 assert!(result.is_none());
1320 }
1321
1322 #[tokio::test]
1325 async fn fetch_summaries_returns_none_when_memory_is_none() {
1326 let view = empty_view();
1327 let tc = NaiveTokenCounter;
1328 let result = fetch_summaries(&view, 1000, &tc).await.unwrap();
1329 assert!(result.is_none());
1330 }
1331
1332 #[tokio::test]
1335 async fn fetch_cross_session_returns_none_when_memory_is_none() {
1336 let view = empty_view();
1337 let tc = NaiveTokenCounter;
1338 let result = fetch_cross_session(&view, "test", 1000, &tc).await.unwrap();
1339 assert!(result.is_none());
1340 }
1341
1342 #[test]
1345 fn levels_to_flags_empty_slice_enables_all_tiers() {
1346 let (e, p, d) = levels_to_flags(&[]);
1347 assert!(e, "episodic should be active for empty slice");
1348 assert!(p, "procedural should be active for empty slice");
1349 assert!(d, "declarative should be active for empty slice");
1350 }
1351
1352 #[test]
1353 fn levels_to_flags_full_set_enables_all_tiers() {
1354 let all = &[
1355 CompressionLevel::Episodic,
1356 CompressionLevel::Procedural,
1357 CompressionLevel::Declarative,
1358 ];
1359 let (e, p, d) = levels_to_flags(all);
1360 assert!(e);
1361 assert!(p);
1362 assert!(d);
1363 }
1364
1365 #[test]
1366 fn levels_to_flags_episodic_only() {
1367 let (e, p, d) = levels_to_flags(&[CompressionLevel::Episodic]);
1368 assert!(e);
1369 assert!(!p, "procedural should be inactive");
1370 assert!(!d, "declarative should be inactive");
1371 }
1372
1373 #[test]
1374 fn levels_to_flags_episodic_and_procedural() {
1375 let (e, p, d) =
1376 levels_to_flags(&[CompressionLevel::Episodic, CompressionLevel::Procedural]);
1377 assert!(e);
1378 assert!(p);
1379 assert!(!d, "declarative should be inactive");
1380 }
1381
1382 #[test]
1383 fn levels_to_flags_declarative_only() {
1384 let (e, p, d) = levels_to_flags(&[CompressionLevel::Declarative]);
1385 assert!(!e, "episodic should be inactive");
1386 assert!(!p, "procedural should be inactive");
1387 assert!(d);
1388 }
1389
1390 #[test]
1393 fn type_active_empty_active_set_means_all_types() {
1394 assert!(type_active(&[], FunctionalType::Episodic));
1395 assert!(type_active(&[], FunctionalType::GraphFact));
1396 assert!(type_active(&[], FunctionalType::BehavioralRule));
1397 }
1398
1399 #[test]
1400 fn type_active_nonempty_set_gates_by_membership() {
1401 let active = [FunctionalType::UserFact];
1402 assert!(type_active(&active, FunctionalType::UserFact));
1403 assert!(!type_active(&active, FunctionalType::Episodic));
1404 assert!(!type_active(&active, FunctionalType::GraphFact));
1405 }
1406
1407 struct NoopRouter;
1410 impl zeph_common::memory::MemoryRouter for NoopRouter {
1411 fn route(&self, _query: &str) -> zeph_common::memory::MemoryRoute {
1412 zeph_common::memory::MemoryRoute::default()
1413 }
1414 }
1415 impl AsyncMemoryRouter for NoopRouter {
1416 fn route_async<'a>(
1417 &'a self,
1418 _query: &'a str,
1419 ) -> std::pin::Pin<
1420 Box<dyn std::future::Future<Output = zeph_common::memory::RoutingDecision> + Send + 'a>,
1421 > {
1422 Box::pin(async move {
1423 zeph_common::memory::RoutingDecision {
1424 route: zeph_common::memory::MemoryRoute::default(),
1425 confidence: 1.0,
1426 reasoning: None,
1427 }
1428 })
1429 }
1430 }
1431
1432 fn full_active_view() -> ContextMemoryView {
1436 let mut view = empty_view();
1437 view.persona_config.context_budget_tokens = 100;
1438 view.trajectory_config.context_budget_tokens = 100;
1439 view.tree_config.context_budget_tokens = 100;
1440 view.reasoning_config.enabled = true;
1441 view.reasoning_config.context_budget_tokens = 100;
1442 view.document_config.rag_enabled = true;
1443 view
1444 }
1445
1446 #[allow(clippy::too_many_arguments)]
1447 fn schedule_all_budgeted<'r>(
1448 view: &'r ContextMemoryView,
1449 tc: &'r NaiveTokenCounter,
1450 router: &'r NoopRouter,
1451 active_types: &'r [FunctionalType],
1452 ) -> FuturesUnordered<CtxFuture<'r>> {
1453 schedule_context_fetchers(
1454 view,
1455 tc,
1456 "query",
1457 |s| s.into(),
1458 None,
1459 router,
1460 100,
1461 100,
1462 100,
1463 100,
1464 100,
1465 10,
1466 0.5,
1467 &[],
1468 active_types,
1469 )
1470 }
1471
1472 #[test]
1473 fn schedule_context_fetchers_schedules_everything_when_active_types_empty() {
1474 let view = full_active_view();
1475 let tc = NaiveTokenCounter;
1476 let router = NoopRouter;
1477 let fetchers = schedule_all_budgeted(&view, &tc, &router, &[]);
1478 assert_eq!(fetchers.len(), 10);
1482 }
1483
1484 #[test]
1485 fn schedule_context_fetchers_gates_to_user_fact_only_sc1() {
1486 let view = full_active_view();
1491 let tc = NaiveTokenCounter;
1492 let router = NoopRouter;
1493 let active = [FunctionalType::UserFact];
1494 let fetchers = schedule_all_budgeted(&view, &tc, &router, &active);
1495 assert_eq!(fetchers.len(), 5);
1497 }
1498
1499 #[test]
1500 fn schedule_context_fetchers_document_rag_survives_episodic_exclusion_n2() {
1501 let view = full_active_view();
1505 let tc = NaiveTokenCounter;
1506 let router = NoopRouter;
1507 let active = [FunctionalType::GraphFact];
1508 let fetchers = schedule_all_budgeted(&view, &tc, &router, &active);
1509 assert_eq!(fetchers.len(), 5);
1512 }
1513
1514 #[test]
1515 fn schedule_context_fetchers_gates_cross_session_summary_both_slots() {
1516 let view = full_active_view();
1518 let tc = NaiveTokenCounter;
1519 let router = NoopRouter;
1520 let active = [FunctionalType::CrossSessionSummary];
1521 let fetchers = schedule_all_budgeted(&view, &tc, &router, &active);
1522 assert_eq!(fetchers.len(), 6);
1524 }
1525
1526 #[tokio::test]
1535 async fn gather_with_user_fact_active_type_excludes_other_slots_sc4() {
1536 let mock = MockMemoryBackend {
1537 persona_facts: vec![MemPersonaFact {
1538 category: "preference".to_string(),
1539 content: "prefers concise answers".to_string(),
1540 }],
1541 ..Default::default()
1542 };
1543 let mut memory = mock_view(mock);
1544 memory.persona_config.enabled = true;
1545 memory.persona_config.context_budget_tokens = 1000;
1546 memory.graph_config.enabled = true;
1547 memory.reasoning_config.enabled = true;
1548 memory.reasoning_config.context_budget_tokens = 500;
1549 memory.document_config.rag_enabled = false;
1550
1551 let mut context_manager = crate::manager::ContextManager::new();
1552 context_manager.budget = Some(crate::budget::ContextBudget::new(128_000, 0.1));
1553
1554 let tc = NaiveTokenCounter;
1555 let active_types = [FunctionalType::UserFact];
1556
1557 let input = crate::input::ContextAssemblyInput {
1558 memory: &memory,
1559 context_manager: &context_manager,
1560 token_counter: &tc,
1561 skills_prompt: "",
1562 index: None,
1563 correction_config: None,
1564 sidequest_turn_counter: 0,
1565 messages: &[],
1566 query: "what do you know about me?",
1567 scrub: |s| s.into(),
1568 active_levels: &[],
1569 active_types: &active_types,
1570 router: Box::new(NoopRouter),
1571 planned_next_tools: &[],
1572 };
1573
1574 let prepared = ContextAssembler::gather(&input).await.unwrap();
1575
1576 assert!(
1577 prepared.recall.is_none(),
1578 "Episodic excluded from active set: recall must be None"
1579 );
1580 assert!(
1581 prepared.reasoning_hints.is_none(),
1582 "ReasoningStrategy excluded from active set: reasoning_hints must be None"
1583 );
1584 assert!(
1585 prepared.graph_facts.is_none(),
1586 "GraphFact excluded from active set: graph_facts must be None"
1587 );
1588 assert!(
1589 prepared.summaries.is_none(),
1590 "CrossSessionSummary excluded from active set: summaries must be None"
1591 );
1592 assert!(
1593 prepared.persona_facts.is_some(),
1594 "UserFact is in the active set: persona_facts must be Some"
1595 );
1596 }
1597
1598 #[tokio::test]
1601 async fn fetch_reasoning_strategies_returns_none_when_memory_is_none() {
1602 let mut view = empty_view();
1603 view.reasoning_config.enabled = true;
1604 let tc = NaiveTokenCounter;
1605 let (result, handle) = fetch_reasoning_strategies(&view, "query", 1000, 3, &tc)
1606 .await
1607 .unwrap();
1608 assert!(result.is_none());
1609 assert!(handle.is_none());
1610 }
1611
1612 #[tokio::test]
1613 async fn fetch_reasoning_strategies_returns_none_when_budget_zero() {
1614 let mut view = empty_view();
1615 view.reasoning_config.enabled = true;
1616 let tc = NaiveTokenCounter;
1617 let (result, handle) = fetch_reasoning_strategies(&view, "query", 0, 3, &tc)
1618 .await
1619 .unwrap();
1620 assert!(result.is_none());
1621 assert!(handle.is_none());
1622 }
1623
1624 use std::sync::{Arc, Mutex};
1627 use zeph_common::memory::{
1628 ContextMemoryBackend, GraphRecallParams, MemCorrection, MemDocumentChunk, MemGraphFact,
1629 MemPersonaFact, MemReasoningStrategy, MemRecalledMessage, MemSessionSummary, MemSummary,
1630 MemTrajectoryEntry, MemTreeNode,
1631 };
1632
1633 const KNOWN_FAIL_ON: &[&str] = &[
1635 "load_persona_facts",
1636 "load_trajectory_entries",
1637 "load_tree_nodes",
1638 "load_summaries",
1639 "retrieve_reasoning_strategies",
1640 "mark_reasoning_used",
1641 "retrieve_corrections",
1642 "recall",
1643 "recall_graph_facts",
1644 "search_session_summaries",
1645 "search_document_collection",
1646 ];
1647
1648 #[derive(Default)]
1649 struct MockMemoryBackend {
1650 persona_facts: Vec<MemPersonaFact>,
1651 trajectory_entries: Vec<MemTrajectoryEntry>,
1652 tree_nodes: Vec<MemTreeNode>,
1653 summaries: Vec<MemSummary>,
1654 reasoning_strategies: Vec<MemReasoningStrategy>,
1655 corrections: Vec<MemCorrection>,
1656 recalled: Vec<MemRecalledMessage>,
1657 graph_facts: Vec<MemGraphFact>,
1658 session_summaries: Vec<MemSessionSummary>,
1659 document_chunks: Vec<MemDocumentChunk>,
1660 fail_on: Option<&'static str>,
1662 delay: Option<std::time::Duration>,
1665 marked_ids: Mutex<Vec<String>>,
1667 }
1668
1669 impl MockMemoryBackend {
1670 fn with_fail_on(method: &'static str) -> Self {
1671 debug_assert!(
1672 KNOWN_FAIL_ON.contains(&method),
1673 "unknown fail_on method name: {method}"
1674 );
1675 Self {
1676 fail_on: Some(method),
1677 ..Default::default()
1678 }
1679 }
1680
1681 fn fail_err(method: &str) -> Box<dyn std::error::Error + Send + Sync> {
1682 format!("mock error in {method}").into()
1683 }
1684 }
1685
1686 impl ContextMemoryBackend for MockMemoryBackend {
1687 fn load_persona_facts<'a>(
1688 &'a self,
1689 _min_confidence: f64,
1690 ) -> std::pin::Pin<
1691 Box<
1692 dyn std::future::Future<
1693 Output = Result<
1694 Vec<MemPersonaFact>,
1695 Box<dyn std::error::Error + Send + Sync>,
1696 >,
1697 > + Send
1698 + 'a,
1699 >,
1700 > {
1701 let result = if self.fail_on == Some("load_persona_facts") {
1702 Err(Self::fail_err("load_persona_facts"))
1703 } else {
1704 Ok(self.persona_facts.clone())
1705 };
1706 let delay = self.delay;
1707 Box::pin(async move {
1708 if let Some(d) = delay {
1709 tokio::time::sleep(d).await;
1710 }
1711 result
1712 })
1713 }
1714
1715 fn load_trajectory_entries<'a>(
1716 &'a self,
1717 _tier: Option<&'a str>,
1718 _top_k: usize,
1719 ) -> std::pin::Pin<
1720 Box<
1721 dyn std::future::Future<
1722 Output = Result<
1723 Vec<MemTrajectoryEntry>,
1724 Box<dyn std::error::Error + Send + Sync>,
1725 >,
1726 > + Send
1727 + 'a,
1728 >,
1729 > {
1730 let result = if self.fail_on == Some("load_trajectory_entries") {
1731 Err(Self::fail_err("load_trajectory_entries"))
1732 } else {
1733 Ok(self.trajectory_entries.clone())
1734 };
1735 Box::pin(async move { result })
1736 }
1737
1738 fn load_tree_nodes<'a>(
1739 &'a self,
1740 _level: u32,
1741 _top_k: usize,
1742 ) -> std::pin::Pin<
1743 Box<
1744 dyn std::future::Future<
1745 Output = Result<Vec<MemTreeNode>, Box<dyn std::error::Error + Send + Sync>>,
1746 > + Send
1747 + 'a,
1748 >,
1749 > {
1750 let result = if self.fail_on == Some("load_tree_nodes") {
1751 Err(Self::fail_err("load_tree_nodes"))
1752 } else {
1753 Ok(self.tree_nodes.clone())
1754 };
1755 Box::pin(async move { result })
1756 }
1757
1758 fn load_summaries<'a>(
1759 &'a self,
1760 _conversation_id: i64,
1761 ) -> std::pin::Pin<
1762 Box<
1763 dyn std::future::Future<
1764 Output = Result<Vec<MemSummary>, Box<dyn std::error::Error + Send + Sync>>,
1765 > + Send
1766 + 'a,
1767 >,
1768 > {
1769 let result = if self.fail_on == Some("load_summaries") {
1770 Err(Self::fail_err("load_summaries"))
1771 } else {
1772 Ok(self.summaries.clone())
1773 };
1774 Box::pin(async move { result })
1775 }
1776
1777 fn retrieve_reasoning_strategies<'a>(
1778 &'a self,
1779 _query: &'a str,
1780 _top_k: usize,
1781 ) -> std::pin::Pin<
1782 Box<
1783 dyn std::future::Future<
1784 Output = Result<
1785 Vec<MemReasoningStrategy>,
1786 Box<dyn std::error::Error + Send + Sync>,
1787 >,
1788 > + Send
1789 + 'a,
1790 >,
1791 > {
1792 let result = if self.fail_on == Some("retrieve_reasoning_strategies") {
1793 Err(Self::fail_err("retrieve_reasoning_strategies"))
1794 } else {
1795 Ok(self.reasoning_strategies.clone())
1796 };
1797 Box::pin(async move { result })
1798 }
1799
1800 fn mark_reasoning_used<'a>(
1801 &'a self,
1802 ids: &'a [String],
1803 ) -> std::pin::Pin<
1804 Box<
1805 dyn std::future::Future<
1806 Output = Result<(), Box<dyn std::error::Error + Send + Sync>>,
1807 > + Send
1808 + 'a,
1809 >,
1810 > {
1811 if self.fail_on == Some("mark_reasoning_used") {
1812 return Box::pin(async move { Err(Self::fail_err("mark_reasoning_used")) });
1813 }
1814 let mut guard = self.marked_ids.lock().expect("marked_ids poisoned");
1815 guard.extend_from_slice(ids);
1816 Box::pin(async move { Ok(()) })
1817 }
1818
1819 fn retrieve_corrections<'a>(
1820 &'a self,
1821 _query: &'a str,
1822 _limit: usize,
1823 _min_score: f32,
1824 ) -> std::pin::Pin<
1825 Box<
1826 dyn std::future::Future<
1827 Output = Result<
1828 Vec<MemCorrection>,
1829 Box<dyn std::error::Error + Send + Sync>,
1830 >,
1831 > + Send
1832 + 'a,
1833 >,
1834 > {
1835 let result = if self.fail_on == Some("retrieve_corrections") {
1836 Err(Self::fail_err("retrieve_corrections"))
1837 } else {
1838 Ok(self.corrections.clone())
1839 };
1840 Box::pin(async move { result })
1841 }
1842
1843 fn recall<'a>(
1844 &'a self,
1845 _query: &'a str,
1846 _limit: usize,
1847 _router: Option<&'a dyn zeph_common::memory::AsyncMemoryRouter>,
1848 ) -> std::pin::Pin<
1849 Box<
1850 dyn std::future::Future<
1851 Output = Result<
1852 Vec<MemRecalledMessage>,
1853 Box<dyn std::error::Error + Send + Sync>,
1854 >,
1855 > + Send
1856 + 'a,
1857 >,
1858 > {
1859 let result = if self.fail_on == Some("recall") {
1860 Err(Self::fail_err("recall"))
1861 } else {
1862 Ok(self.recalled.clone())
1863 };
1864 let delay = self.delay;
1865 Box::pin(async move {
1866 if let Some(d) = delay {
1867 tokio::time::sleep(d).await;
1868 }
1869 result
1870 })
1871 }
1872
1873 fn recall_graph_facts<'a>(
1874 &'a self,
1875 _query: &'a str,
1876 _params: GraphRecallParams<'a>,
1877 ) -> std::pin::Pin<
1878 Box<
1879 dyn std::future::Future<
1880 Output = Result<
1881 Vec<MemGraphFact>,
1882 Box<dyn std::error::Error + Send + Sync>,
1883 >,
1884 > + Send
1885 + 'a,
1886 >,
1887 > {
1888 let result = if self.fail_on == Some("recall_graph_facts") {
1889 Err(Self::fail_err("recall_graph_facts"))
1890 } else {
1891 Ok(self.graph_facts.clone())
1892 };
1893 Box::pin(async move { result })
1894 }
1895
1896 fn search_session_summaries<'a>(
1897 &'a self,
1898 _query: &'a str,
1899 _limit: usize,
1900 _current_conversation_id: Option<i64>,
1901 ) -> std::pin::Pin<
1902 Box<
1903 dyn std::future::Future<
1904 Output = Result<
1905 Vec<MemSessionSummary>,
1906 Box<dyn std::error::Error + Send + Sync>,
1907 >,
1908 > + Send
1909 + 'a,
1910 >,
1911 > {
1912 let result = if self.fail_on == Some("search_session_summaries") {
1913 Err(Self::fail_err("search_session_summaries"))
1914 } else {
1915 Ok(self.session_summaries.clone())
1916 };
1917 Box::pin(async move { result })
1918 }
1919
1920 fn search_document_collection<'a>(
1921 &'a self,
1922 _collection: &'a str,
1923 _query: &'a str,
1924 _top_k: usize,
1925 ) -> std::pin::Pin<
1926 Box<
1927 dyn std::future::Future<
1928 Output = Result<
1929 Vec<MemDocumentChunk>,
1930 Box<dyn std::error::Error + Send + Sync>,
1931 >,
1932 > + Send
1933 + 'a,
1934 >,
1935 > {
1936 let result = if self.fail_on == Some("search_document_collection") {
1937 Err(Self::fail_err("search_document_collection"))
1938 } else {
1939 Ok(self.document_chunks.clone())
1940 };
1941 Box::pin(async move { result })
1942 }
1943 }
1944
1945 fn mock_view(mock: MockMemoryBackend) -> ContextMemoryView {
1946 let mut v = empty_view();
1947 v.memory = Some(Arc::new(mock));
1948 v
1949 }
1950
1951 #[tokio::test]
1954 async fn fetch_graph_facts_returns_message_when_memory_present() {
1955 let mock = MockMemoryBackend {
1956 graph_facts: vec![zeph_common::memory::MemGraphFact {
1957 fact: "Rust is fast".to_string(),
1958 confidence: 0.9,
1959 activation_score: None,
1960 neighbors: vec![],
1961 provenance_snippet: None,
1962 }],
1963 ..Default::default()
1964 };
1965 let mut view = mock_view(mock);
1966 view.graph_config.enabled = true;
1967 view.graph_config.spreading_activation.recall_timeout_ms = 5000;
1969 let tc = NaiveTokenCounter;
1970 let result = fetch_graph_facts(&view, "test", 1000, &tc).await.unwrap();
1971 assert!(result.is_some(), "expected Some message");
1972 let msg = result.unwrap();
1973 assert!(
1974 msg.content.contains("Rust is fast"),
1975 "expected fact text in output, got: {}",
1976 msg.content
1977 );
1978 assert!(
1979 msg.content.starts_with(GRAPH_FACTS_PREFIX),
1980 "expected GRAPH_FACTS_PREFIX"
1981 );
1982 }
1983
1984 #[tokio::test]
1985 async fn fetch_graph_facts_swallows_error_and_returns_none() {
1986 let mock = MockMemoryBackend::with_fail_on("recall_graph_facts");
1987 let mut view = mock_view(mock);
1988 view.graph_config.enabled = true;
1989 view.graph_config.spreading_activation.recall_timeout_ms = 5000;
1990 let tc = NaiveTokenCounter;
1991 let result = fetch_graph_facts(&view, "test", 1000, &tc).await.unwrap();
1993 assert!(
1994 result.is_none(),
1995 "expected None when recall_graph_facts errors"
1996 );
1997 }
1998
1999 #[tokio::test]
2000 async fn fetch_graph_facts_returns_none_when_facts_empty() {
2001 let mock = MockMemoryBackend::default(); let mut view = mock_view(mock);
2003 view.graph_config.enabled = true;
2004 view.graph_config.spreading_activation.recall_timeout_ms = 5000;
2005 let tc = NaiveTokenCounter;
2006 let result = fetch_graph_facts(&view, "test", 1000, &tc).await.unwrap();
2007 assert!(result.is_none());
2008 }
2009
2010 #[tokio::test]
2013 async fn fetch_persona_facts_returns_message_when_persona_enabled() {
2014 let mock = MockMemoryBackend {
2015 persona_facts: vec![MemPersonaFact {
2016 category: "preference".to_string(),
2017 content: "prefers concise answers".to_string(),
2018 }],
2019 ..Default::default()
2020 };
2021 let mut view = mock_view(mock);
2022 view.persona_config.enabled = true;
2023 view.persona_config.context_budget_tokens = 1000;
2024 let tc = NaiveTokenCounter;
2025 let result = fetch_persona_facts(&view, 1000, &tc).await.unwrap();
2026 assert!(result.is_some());
2027 let msg = result.unwrap();
2028 assert!(msg.content.contains("preference"));
2029 assert!(msg.content.contains("prefers concise answers"));
2030 assert!(msg.content.starts_with(crate::slot::PERSONA_PREFIX));
2031 }
2032
2033 #[tokio::test]
2034 async fn fetch_persona_facts_propagates_error() {
2035 let mock = MockMemoryBackend::with_fail_on("load_persona_facts");
2036 let mut view = mock_view(mock);
2037 view.persona_config.enabled = true;
2038 let tc = NaiveTokenCounter;
2039 let result = fetch_persona_facts(&view, 1000, &tc).await;
2040 assert!(
2041 result.is_err(),
2042 "expected Err from load_persona_facts failure"
2043 );
2044 }
2045
2046 #[tokio::test]
2049 async fn fetch_trajectory_hints_returns_message_when_trajectory_enabled() {
2050 let mock = MockMemoryBackend {
2051 trajectory_entries: vec![MemTrajectoryEntry {
2052 intent: "summarize code".to_string(),
2053 outcome: "produced concise summary".to_string(),
2054 confidence: 0.9,
2055 }],
2056 ..Default::default()
2057 };
2058 let mut view = mock_view(mock);
2059 view.trajectory_config.enabled = true;
2060 view.trajectory_config.context_budget_tokens = 1000;
2061 view.trajectory_config.min_confidence = 0.5;
2062 let tc = NaiveTokenCounter;
2063 let result = fetch_trajectory_hints(&view, 1000, &tc).await.unwrap();
2064 assert!(result.is_some());
2065 let msg = result.unwrap();
2066 assert!(msg.content.contains("summarize code"));
2067 assert!(msg.content.starts_with(crate::slot::TRAJECTORY_PREFIX));
2068 }
2069
2070 #[tokio::test]
2071 async fn fetch_trajectory_hints_passes_tier_filter() {
2072 let mock = MockMemoryBackend {
2075 trajectory_entries: vec![
2076 MemTrajectoryEntry {
2077 intent: "debug async code".to_string(),
2078 outcome: "fixed deadlock".to_string(),
2079 confidence: 0.85,
2080 },
2081 MemTrajectoryEntry {
2082 intent: "low confidence task".to_string(),
2083 outcome: "irrelevant".to_string(),
2084 confidence: 0.3,
2085 },
2086 ],
2087 ..Default::default()
2088 };
2089 let mut view = mock_view(mock);
2090 view.trajectory_config.enabled = true;
2091 view.trajectory_config.context_budget_tokens = 1000;
2092 view.trajectory_config.min_confidence = 0.5;
2093 let tc = NaiveTokenCounter;
2094 let result = fetch_trajectory_hints(&view, 1000, &tc).await.unwrap();
2095 assert!(result.is_some(), "expected Some message");
2096 let msg = result.unwrap();
2097 assert!(
2098 msg.content.contains("debug async code"),
2099 "high-confidence entry must be included"
2100 );
2101 assert!(
2102 !msg.content.contains("low confidence task"),
2103 "entry below min_confidence must be filtered out"
2104 );
2105 }
2106
2107 #[tokio::test]
2108 async fn fetch_trajectory_hints_propagates_error() {
2109 let mock = MockMemoryBackend::with_fail_on("load_trajectory_entries");
2110 let mut view = mock_view(mock);
2111 view.trajectory_config.enabled = true;
2112 let tc = NaiveTokenCounter;
2113 let result = fetch_trajectory_hints(&view, 1000, &tc).await;
2114 assert!(result.is_err());
2115 }
2116
2117 #[tokio::test]
2120 async fn fetch_tree_memory_returns_message_when_tree_enabled() {
2121 let mock = MockMemoryBackend {
2122 tree_nodes: vec![MemTreeNode {
2123 content: "Topic: async Rust patterns".to_string(),
2124 }],
2125 ..Default::default()
2126 };
2127 let mut view = mock_view(mock);
2128 view.tree_config.enabled = true;
2129 view.tree_config.context_budget_tokens = 1000;
2130 let tc = NaiveTokenCounter;
2131 let result = fetch_tree_memory(&view, 1000, &tc).await.unwrap();
2132 assert!(result.is_some());
2133 let msg = result.unwrap();
2134 assert!(msg.content.contains("async Rust patterns"));
2135 assert!(msg.content.starts_with(crate::slot::TREE_MEMORY_PREFIX));
2136 }
2137
2138 #[tokio::test]
2139 async fn fetch_tree_memory_propagates_error() {
2140 let mock = MockMemoryBackend::with_fail_on("load_tree_nodes");
2141 let mut view = mock_view(mock);
2142 view.tree_config.enabled = true;
2143 let tc = NaiveTokenCounter;
2144 let result = fetch_tree_memory(&view, 1000, &tc).await;
2145 assert!(result.is_err());
2146 }
2147
2148 #[tokio::test]
2151 async fn fetch_corrections_returns_message_when_corrections_present() {
2152 let mock = MockMemoryBackend {
2153 corrections: vec![MemCorrection {
2154 correction_text: "use snake_case not camelCase".to_string(),
2155 }],
2156 ..Default::default()
2157 };
2158 let view = mock_view(mock);
2159 let result = fetch_corrections(&view, "query", 10, 0.5, |s| s.into())
2160 .await
2161 .unwrap();
2162 assert!(result.is_some());
2163 let msg = result.unwrap();
2164 assert!(msg.content.contains("snake_case"));
2165 assert!(msg.content.starts_with(CORRECTIONS_PREFIX));
2166 }
2167
2168 #[tokio::test]
2169 async fn fetch_corrections_propagates_error() {
2170 let mock = MockMemoryBackend::with_fail_on("retrieve_corrections");
2173 let view = mock_view(mock);
2174 let result = fetch_corrections(&view, "query", 10, 0.5, |s| s.into()).await;
2175 assert!(result.is_err(), "expected Err, got {result:?}");
2176 }
2177
2178 #[tokio::test]
2181 async fn fetch_semantic_recall_returns_message_with_content() {
2182 let mock = MockMemoryBackend {
2183 recalled: vec![
2184 MemRecalledMessage {
2185 role: "user".to_string(),
2186 content: "how does tokio work".to_string(),
2187 score: 0.95,
2188 },
2189 MemRecalledMessage {
2190 role: "assistant".to_string(),
2191 content: "tokio is an async runtime".to_string(),
2192 score: 0.88,
2193 },
2194 ],
2195 ..Default::default()
2196 };
2197 let mut view = mock_view(mock);
2198 view.recall_limit = 10;
2199 let tc = NaiveTokenCounter;
2200 let (msg, score) = fetch_semantic_recall(&view, "tokio", 1000, &tc, None)
2201 .await
2202 .unwrap();
2203 assert!(msg.is_some(), "expected Some message");
2204 assert!(score.is_some_and(|s| (s - 0.95_f32).abs() < f32::EPSILON));
2206 let msg = msg.unwrap();
2207 let has_recall_part = msg.parts.iter().any(|p| {
2209 if let zeph_llm::provider::MessagePart::Recall { text } = p {
2210 text.contains("how does tokio work")
2211 } else {
2212 false
2213 }
2214 });
2215 assert!(has_recall_part, "expected recalled content in Recall part");
2216 }
2217
2218 #[tokio::test]
2219 async fn fetch_semantic_recall_returns_none_when_recalled_empty() {
2220 let mock = MockMemoryBackend::default();
2221 let mut view = mock_view(mock);
2222 view.recall_limit = 10;
2223 let tc = NaiveTokenCounter;
2224 let (msg, score) = fetch_semantic_recall(&view, "query", 1000, &tc, None)
2225 .await
2226 .unwrap();
2227 assert!(msg.is_none());
2228 assert!(score.is_none());
2229 }
2230
2231 #[tokio::test]
2232 async fn fetch_semantic_recall_propagates_error() {
2233 let mock = MockMemoryBackend::with_fail_on("recall");
2234 let mut view = mock_view(mock);
2235 view.recall_limit = 10;
2236 let tc = NaiveTokenCounter;
2237 let result = fetch_semantic_recall(&view, "query", 1000, &tc, None).await;
2238 assert!(result.is_err());
2239 }
2240
2241 #[tokio::test]
2244 async fn fetch_document_rag_returns_message_when_rag_enabled() {
2245 let mock = MockMemoryBackend {
2246 document_chunks: vec![MemDocumentChunk {
2247 text: "Rust ownership rules prevent data races".to_string(),
2248 }],
2249 ..Default::default()
2250 };
2251 let mut view = mock_view(mock);
2252 view.document_config.rag_enabled = true;
2253 let tc = NaiveTokenCounter;
2254 let result = fetch_document_rag(&view, "ownership", 1000, &tc)
2255 .await
2256 .unwrap();
2257 assert!(result.is_some());
2258 let msg = result.unwrap();
2259 assert!(msg.content.contains("ownership rules"));
2260 assert!(msg.content.starts_with(DOCUMENT_RAG_PREFIX));
2261 }
2262
2263 #[tokio::test]
2264 async fn fetch_document_rag_propagates_error() {
2265 let mock = MockMemoryBackend::with_fail_on("search_document_collection");
2266 let mut view = mock_view(mock);
2267 view.document_config.rag_enabled = true;
2268 let tc = NaiveTokenCounter;
2269 let result = fetch_document_rag(&view, "query", 1000, &tc).await;
2270 assert!(result.is_err());
2271 }
2272
2273 #[tokio::test]
2276 async fn fetch_summaries_returns_message_when_summaries_present() {
2277 let mock = MockMemoryBackend {
2278 summaries: vec![MemSummary {
2279 first_message_id: Some(1),
2280 last_message_id: Some(5),
2281 content: "User asked about async Rust".to_string(),
2282 }],
2283 ..Default::default()
2284 };
2285 let mut view = mock_view(mock);
2286 view.conversation_id = Some(42);
2287 let tc = NaiveTokenCounter;
2288 let result = fetch_summaries(&view, 1000, &tc).await.unwrap();
2289 assert!(result.is_some());
2290 let msg = result.unwrap();
2291 let has_summary_part = msg.parts.iter().any(|p| {
2292 if let zeph_llm::provider::MessagePart::Summary { text } = p {
2293 text.contains("Messages 1-5") && text.contains("async Rust")
2294 } else {
2295 false
2296 }
2297 });
2298 assert!(
2299 has_summary_part,
2300 "expected Summary part with messages range"
2301 );
2302 }
2303
2304 #[tokio::test]
2305 async fn fetch_summaries_returns_none_without_conversation_id() {
2306 let mock = MockMemoryBackend {
2307 summaries: vec![MemSummary {
2308 first_message_id: Some(1),
2309 last_message_id: Some(5),
2310 content: "some content".to_string(),
2311 }],
2312 ..Default::default()
2313 };
2314 let mut view = mock_view(mock);
2315 view.conversation_id = None; let tc = NaiveTokenCounter;
2317 let result = fetch_summaries(&view, 1000, &tc).await.unwrap();
2318 assert!(result.is_none());
2319 }
2320
2321 #[tokio::test]
2322 async fn fetch_summaries_propagates_error() {
2323 let mock = MockMemoryBackend::with_fail_on("load_summaries");
2324 let mut view = mock_view(mock);
2325 view.conversation_id = Some(42);
2326 let tc = NaiveTokenCounter;
2327 let result = fetch_summaries(&view, 1000, &tc).await;
2328 assert!(result.is_err());
2329 }
2330
2331 #[tokio::test]
2334 async fn fetch_cross_session_returns_message_when_results_present() {
2335 let mock = MockMemoryBackend {
2336 session_summaries: vec![MemSessionSummary {
2337 summary_text: "Previous session: debugging tokio deadlock".to_string(),
2338 score: 0.9,
2339 }],
2340 ..Default::default()
2341 };
2342 let mut view = mock_view(mock);
2343 view.conversation_id = Some(1);
2344 view.cross_session_score_threshold = 0.5;
2345 let tc = NaiveTokenCounter;
2346 let result = fetch_cross_session(&view, "async", 1000, &tc)
2347 .await
2348 .unwrap();
2349 assert!(result.is_some());
2350 let msg = result.unwrap();
2351 let has_cross_session_part = msg.parts.iter().any(|p| {
2352 if let zeph_llm::provider::MessagePart::CrossSession { text } = p {
2353 text.contains("tokio deadlock")
2354 } else {
2355 false
2356 }
2357 });
2358 assert!(has_cross_session_part);
2359 }
2360
2361 #[tokio::test]
2362 async fn fetch_cross_session_propagates_error() {
2363 let mock = MockMemoryBackend::with_fail_on("search_session_summaries");
2364 let mut view = mock_view(mock);
2365 view.conversation_id = Some(1);
2366 let tc = NaiveTokenCounter;
2367 let result = fetch_cross_session(&view, "query", 1000, &tc).await;
2368 assert!(result.is_err());
2369 }
2370
2371 #[tokio::test]
2374 async fn fetch_reasoning_strategies_returns_message_and_marks_used() {
2375 let mock = Arc::new(MockMemoryBackend {
2376 reasoning_strategies: vec![
2377 MemReasoningStrategy {
2378 id: "strat-1".to_string(),
2379 outcome: "success".to_string(),
2380 summary: "break the problem into small steps".to_string(),
2381 },
2382 MemReasoningStrategy {
2383 id: "strat-2".to_string(),
2384 outcome: "success".to_string(),
2385 summary: "use tracing spans for debugging".to_string(),
2386 },
2387 ],
2388 ..Default::default()
2389 });
2390 let marked_ids = Arc::clone(&mock);
2391 let mut view = empty_view();
2392 view.memory = Some(mock);
2393 view.reasoning_config.enabled = true;
2394 view.reasoning_config.context_budget_tokens = 1000;
2395 let tc = NaiveTokenCounter;
2396 let (result, handle) = fetch_reasoning_strategies(&view, "debug", 1000, 5, &tc)
2397 .await
2398 .unwrap();
2399 assert!(result.is_some());
2400 let msg = result.unwrap();
2401 assert!(msg.content.starts_with(crate::slot::REASONING_PREFIX));
2402 assert!(msg.content.contains("break the problem"));
2403
2404 if let Some(h) = handle {
2406 h.await.expect("mark_reasoning_used task panicked");
2407 }
2408
2409 let ids = marked_ids.marked_ids.lock().expect("marked_ids poisoned");
2410 assert!(
2411 ids.contains(&"strat-1".to_string()),
2412 "expected strat-1 marked"
2413 );
2414 assert!(
2415 ids.contains(&"strat-2".to_string()),
2416 "expected strat-2 marked"
2417 );
2418 }
2419
2420 #[tokio::test]
2421 async fn fetch_reasoning_strategies_propagates_error() {
2422 let mock = MockMemoryBackend::with_fail_on("retrieve_reasoning_strategies");
2423 let mut view = mock_view(mock);
2424 view.reasoning_config.enabled = true;
2425 let tc = NaiveTokenCounter;
2426 let result = fetch_reasoning_strategies(&view, "query", 1000, 3, &tc).await;
2427 assert!(result.is_err());
2428 }
2429
2430 #[tokio::test]
2433 async fn fetch_semantic_recall_skips_skipped_and_stopped_messages() {
2434 let mock = MockMemoryBackend {
2435 recalled: vec![
2436 MemRecalledMessage {
2437 role: "user".to_string(),
2438 content: "[skipped] some content".to_string(),
2439 score: 0.95,
2440 },
2441 MemRecalledMessage {
2442 role: "user".to_string(),
2443 content: "[stopped] other content".to_string(),
2444 score: 0.90,
2445 },
2446 MemRecalledMessage {
2447 role: "user".to_string(),
2448 content: "valid content to recall".to_string(),
2449 score: 0.85,
2450 },
2451 ],
2452 ..Default::default()
2453 };
2454 let mut view = mock_view(mock);
2455 view.recall_limit = 10;
2456 let tc = NaiveTokenCounter;
2457 let (msg, _) = fetch_semantic_recall(&view, "query", 1000, &tc, None)
2458 .await
2459 .unwrap();
2460 assert!(msg.is_some());
2461 let msg = msg.unwrap();
2462 let full_text = msg.parts.iter().find_map(|p| {
2463 if let zeph_llm::provider::MessagePart::Recall { text } = p {
2464 Some(text.clone())
2465 } else {
2466 None
2467 }
2468 });
2469 let text = full_text.unwrap_or_default();
2470 assert!(
2471 !text.contains("[skipped]"),
2472 "skipped messages must be excluded"
2473 );
2474 assert!(
2475 !text.contains("[stopped]"),
2476 "stopped messages must be excluded"
2477 );
2478 assert!(
2479 text.contains("valid content to recall"),
2480 "valid messages must be included"
2481 );
2482 }
2483
2484 #[tokio::test]
2485 async fn fetch_cross_session_filters_below_threshold() {
2486 let mock = MockMemoryBackend {
2487 session_summaries: vec![
2488 MemSessionSummary {
2489 summary_text: "high relevance session".to_string(),
2490 score: 0.9,
2491 },
2492 MemSessionSummary {
2493 summary_text: "low relevance session".to_string(),
2494 score: 0.2,
2495 },
2496 ],
2497 ..Default::default()
2498 };
2499 let mut view = mock_view(mock);
2500 view.conversation_id = Some(1);
2501 view.cross_session_score_threshold = 0.5;
2502 let tc = NaiveTokenCounter;
2503 let result = fetch_cross_session(&view, "query", 1000, &tc)
2504 .await
2505 .unwrap();
2506 assert!(result.is_some());
2507 let msg = result.unwrap();
2508 let text = msg
2509 .parts
2510 .iter()
2511 .find_map(|p| {
2512 if let zeph_llm::provider::MessagePart::CrossSession { text } = p {
2513 Some(text.clone())
2514 } else {
2515 None
2516 }
2517 })
2518 .unwrap_or_default();
2519 assert!(
2520 text.contains("high relevance"),
2521 "high score must be included"
2522 );
2523 assert!(
2524 !text.contains("low relevance"),
2525 "low score must be filtered out"
2526 );
2527 }
2528
2529 #[tokio::test]
2530 async fn fetch_document_rag_skips_empty_chunks() {
2531 let mock = MockMemoryBackend {
2532 document_chunks: vec![
2533 MemDocumentChunk {
2534 text: String::new(),
2535 }, MemDocumentChunk {
2537 text: "real content here".to_string(),
2538 },
2539 ],
2540 ..Default::default()
2541 };
2542 let mut view = mock_view(mock);
2543 view.document_config.rag_enabled = true;
2544 let tc = NaiveTokenCounter;
2545 let result = fetch_document_rag(&view, "query", 1000, &tc).await.unwrap();
2546 assert!(result.is_some());
2547 let msg = result.unwrap();
2548 assert!(msg.content.contains("real content here"));
2549 assert!(!msg.content.contains("\n\n\n"));
2551 }
2552
2553 #[tokio::test]
2554 async fn fetch_graph_facts_sanitizes_injection_payloads() {
2555 let mock = MockMemoryBackend {
2557 graph_facts: vec![zeph_common::memory::MemGraphFact {
2558 fact: "fact with <script>alert(1)</script> and\nnewline".to_string(),
2559 confidence: 0.8,
2560 activation_score: None,
2561 neighbors: vec![],
2562 provenance_snippet: None,
2563 }],
2564 ..Default::default()
2565 };
2566 let mut view = mock_view(mock);
2567 view.graph_config.enabled = true;
2568 view.graph_config.spreading_activation.recall_timeout_ms = 5000;
2569 let tc = NaiveTokenCounter;
2570 let result = fetch_graph_facts(&view, "test", 1000, &tc).await.unwrap();
2571 assert!(result.is_some());
2572 let msg = result.unwrap();
2573 assert!(
2574 !msg.content.contains('<'),
2575 "angle brackets must be sanitized"
2576 );
2577 assert!(
2580 !msg.content.contains("\n\n"),
2581 "embedded newlines must be sanitized, no double-newline sequences expected"
2582 );
2583 }
2584
2585 #[tokio::test]
2586 async fn fetch_reasoning_strategies_sanitizes_injection_payloads() {
2587 let mock = MockMemoryBackend {
2589 reasoning_strategies: vec![MemReasoningStrategy {
2590 id: "s1".to_string(),
2591 outcome: "success".to_string(),
2592 summary: "strategy with <b>bold</b> and\nnewline".to_string(),
2593 }],
2594 ..Default::default()
2595 };
2596 let mut view = mock_view(mock);
2597 view.reasoning_config.enabled = true;
2598 let tc = NaiveTokenCounter;
2599 let (result, _handle) = fetch_reasoning_strategies(&view, "query", 1000, 3, &tc)
2600 .await
2601 .unwrap();
2602 assert!(result.is_some());
2603 let msg = result.unwrap();
2604 assert!(
2605 !msg.content.contains('<'),
2606 "angle brackets must be sanitized in strategy summaries"
2607 );
2608 }
2609
2610 #[tokio::test]
2613 async fn fetch_persona_facts_truncates_at_budget() {
2614 let tc = NaiveTokenCounter;
2615 let first_line = "[pref] brief\n";
2617 let budget = tc.count_tokens(crate::slot::PERSONA_PREFIX) + tc.count_tokens(first_line);
2618 let mock = MockMemoryBackend {
2619 persona_facts: vec![
2620 MemPersonaFact {
2621 category: "pref".to_string(),
2622 content: "brief".to_string(),
2623 },
2624 MemPersonaFact {
2625 category: "lang".to_string(),
2626 content: "english".to_string(),
2627 },
2628 ],
2629 ..Default::default()
2630 };
2631 let mut view = mock_view(mock);
2632 view.persona_config.enabled = true;
2633 let result = fetch_persona_facts(&view, budget, &tc).await.unwrap();
2634 let msg = result.unwrap();
2635 assert!(msg.content.contains("brief"), "first fact must be included");
2636 assert!(
2637 !msg.content.contains("english"),
2638 "second fact must be truncated by budget"
2639 );
2640 }
2641
2642 #[tokio::test]
2643 async fn fetch_semantic_recall_truncates_at_budget() {
2644 let tc = NaiveTokenCounter;
2645 let first_entry = "- [user] first message\n";
2647 let budget = tc.count_tokens(RECALL_PREFIX) + tc.count_tokens(first_entry);
2648 let mock = MockMemoryBackend {
2649 recalled: vec![
2650 MemRecalledMessage {
2651 role: "user".to_string(),
2652 content: "first message".to_string(),
2653 score: 0.95,
2654 },
2655 MemRecalledMessage {
2656 role: "user".to_string(),
2657 content: "second message that should be truncated".to_string(),
2658 score: 0.80,
2659 },
2660 ],
2661 ..Default::default()
2662 };
2663 let mut view = mock_view(mock);
2664 view.recall_limit = 10;
2665 let (msg, _) = fetch_semantic_recall(&view, "query", budget, &tc, None)
2666 .await
2667 .unwrap();
2668 assert!(msg.is_some());
2669 let text = msg
2670 .unwrap()
2671 .parts
2672 .iter()
2673 .find_map(|p| {
2674 if let zeph_llm::provider::MessagePart::Recall { text } = p {
2675 Some(text.clone())
2676 } else {
2677 None
2678 }
2679 })
2680 .unwrap_or_default();
2681 assert!(
2682 text.contains("first message"),
2683 "first entry must be included"
2684 );
2685 assert!(
2686 !text.contains("second message"),
2687 "second entry must be truncated by budget"
2688 );
2689 }
2690
2691 #[tokio::test]
2694 async fn fetch_graph_facts_sanitizes_provenance_snippet() {
2695 use zeph_common::memory::MemGraphNeighbor;
2696 let mock = MockMemoryBackend {
2697 graph_facts: vec![zeph_common::memory::MemGraphFact {
2698 fact: "safe fact".to_string(),
2699 confidence: 0.9,
2700 activation_score: None,
2701 neighbors: vec![MemGraphNeighbor {
2702 fact: "neighbor".to_string(),
2703 confidence: 0.7,
2704 }],
2705 provenance_snippet: Some("source with <injection>\nand newline".to_string()),
2706 }],
2707 ..Default::default()
2708 };
2709 let mut view = mock_view(mock);
2710 view.graph_config.enabled = true;
2711 view.graph_config.spreading_activation.recall_timeout_ms = 5000;
2712 let tc = NaiveTokenCounter;
2713 let result = fetch_graph_facts(&view, "test", 1000, &tc).await.unwrap();
2714 assert!(result.is_some());
2715 let msg = result.unwrap();
2716 assert!(
2717 !msg.content.contains('<'),
2718 "angle brackets in provenance_snippet must be sanitized"
2719 );
2720 assert!(
2721 !msg.content.contains("\n\n"),
2722 "newlines in provenance_snippet must be sanitized"
2723 );
2724 assert!(
2725 msg.content.contains("[source:"),
2726 "provenance snippet must be rendered"
2727 );
2728 }
2729
2730 #[tokio::test(start_paused = true)]
2738 async fn fetch_persona_facts_degrades_to_empty_on_timeout() {
2739 let mock = MockMemoryBackend {
2740 persona_facts: vec![MemPersonaFact {
2741 category: "pref".to_string(),
2742 content: "would have been returned".to_string(),
2743 }],
2744 delay: Some(std::time::Duration::from_millis(
2745 MEMORY_FETCH_TIMEOUT_MS + 1000,
2746 )),
2747 ..Default::default()
2748 };
2749 let mut view = mock_view(mock);
2750 view.persona_config.enabled = true;
2751 let tc = NaiveTokenCounter;
2752 let result = fetch_persona_facts(&view, 1000, &tc).await;
2753 assert!(
2754 result.is_ok(),
2755 "timeout must degrade gracefully, not propagate as an error: {result:?}"
2756 );
2757 assert!(
2758 result.unwrap().is_none(),
2759 "timed-out fetch must yield no message, not the stale backend data"
2760 );
2761 }
2762
2763 #[tokio::test(start_paused = true)]
2764 async fn fetch_semantic_recall_degrades_to_empty_on_timeout() {
2765 let mock = MockMemoryBackend {
2766 recalled: vec![MemRecalledMessage {
2767 role: "user".to_string(),
2768 content: "would have been returned".to_string(),
2769 score: 0.95,
2770 }],
2771 delay: Some(std::time::Duration::from_millis(
2772 MEMORY_FETCH_TIMEOUT_MS + 1000,
2773 )),
2774 ..Default::default()
2775 };
2776 let mut view = mock_view(mock);
2777 view.recall_limit = 10;
2778 let tc = NaiveTokenCounter;
2779 let result = fetch_semantic_recall(&view, "query", 1000, &tc, None).await;
2780 assert!(
2781 result.is_ok(),
2782 "timeout must degrade gracefully, not propagate as an error: {result:?}"
2783 );
2784 let (msg, score) = result.unwrap();
2785 assert!(msg.is_none(), "timed-out recall must yield no message");
2786 assert!(score.is_none(), "timed-out recall must yield no score");
2787 }
2788
2789 #[test]
2792 fn append_budgeted_lines_empty_input_returns_none() {
2793 let tc = NaiveTokenCounter;
2794 let result = append_budgeted_lines("prefix\n", std::iter::empty(), 1000, &tc);
2795 assert!(result.is_none());
2796 }
2797
2798 #[test]
2799 fn append_budgeted_lines_all_items_fit() {
2800 let tc = NaiveTokenCounter;
2801 let lines = vec![
2802 "one\n".to_string(),
2803 "two\n".to_string(),
2804 "three\n".to_string(),
2805 ];
2806 let result = append_budgeted_lines("prefix\n", lines.into_iter(), 1000, &tc).unwrap();
2807 assert!(result.starts_with("prefix\n"));
2808 assert!(result.contains("one"));
2809 assert!(result.contains("two"));
2810 assert!(result.contains("three"));
2811 }
2812
2813 #[test]
2814 fn append_budgeted_lines_truncates_at_budget() {
2815 let tc = NaiveTokenCounter;
2816 let prefix = "prefix\n";
2817 let first = "one\n";
2818 let budget = tc.count_tokens(prefix) + tc.count_tokens(first);
2820 let lines = vec![first.to_string(), "two extra words here\n".to_string()];
2821 let result = append_budgeted_lines(prefix, lines.into_iter(), budget, &tc).unwrap();
2822 assert!(result.contains("one"), "first line must fit in budget");
2823 assert!(
2824 !result.contains("two extra words"),
2825 "second line must be truncated by budget"
2826 );
2827 }
2828
2829 #[test]
2830 fn append_budgeted_lines_zero_budget_returns_none() {
2831 let tc = NaiveTokenCounter;
2832 let lines = vec!["one\n".to_string()];
2833 let result = append_budgeted_lines("prefix\n", lines.into_iter(), 0, &tc);
2834 assert!(result.is_none(), "no line can fit within a zero budget");
2835 }
2836}