1use std::time::Duration;
25
26use futures::StreamExt as _;
27use futures::future::BoxFuture;
28use tracing::Instrument as _;
29use tracing::info_span;
30use zeph_common::memory::TokenCounting;
31use zeph_common::{ContextFidelity, PlannedToolHint};
32use zeph_llm::LlmError;
33use zeph_llm::LlmProviderDyn;
34use zeph_llm::provider::{EmbedFuture, Message, MessageMetadata, MessagePart, Role};
35
36use crate::assembler::CORRECTIONS_PREFIX;
37
38pub use zeph_config::FidelityConfig;
40
41trait EmbedCall: Send + Sync {
47 fn call<'a>(&'a self, text: &'a str) -> BoxFuture<'a, Result<Vec<f32>, LlmError>>;
48}
49
50struct ClosureEmbed<F>(F);
52
53impl<F> EmbedCall for ClosureEmbed<F>
54where
55 F: Fn(&str) -> EmbedFuture + Send + Sync,
56{
57 fn call<'a>(&'a self, text: &'a str) -> BoxFuture<'a, Result<Vec<f32>, LlmError>> {
58 Box::pin((self.0)(text))
61 }
62}
63
64struct ProviderEmbed<'p>(&'p dyn LlmProviderDyn);
66
67impl EmbedCall for ProviderEmbed<'_> {
68 fn call<'a>(&'a self, text: &'a str) -> BoxFuture<'a, Result<Vec<f32>, LlmError>> {
69 self.0.embed(text)
70 }
71}
72
73#[tracing::instrument(name = "context.fidelity.embed_prepass_dyn", skip_all)]
79async fn embed_prepass_dyn(
80 messages: &[Message],
81 embed: &dyn EmbedCall,
82 config: &FidelityConfig,
83 inserted_count: usize,
84) -> std::collections::HashMap<usize, Vec<f32>> {
85 let concurrency = if config.embed_concurrency == 0 {
86 tracing::warn!(
87 "embed_concurrency is 0, clamping to 1; set a positive value in [memory.fidelity]"
88 );
89 1
90 } else {
91 config.embed_concurrency
92 };
93
94 let tasks = messages.iter().enumerate().filter_map(|(i, msg)| {
95 if is_exempt(msg, i, inserted_count)
96 || msg.content.is_empty()
97 || msg.metadata.embedding.is_some()
98 {
99 return None;
100 }
101 let content = match config.max_embed_input_tokens {
102 Some(n) => truncate_to_byte_limit(&msg.content, n.saturating_mul(4)),
103 None => msg.content.clone(),
104 };
105 Some((i, content))
106 });
107
108 let timeout = Duration::from_secs(config.embed_timeout_secs);
109 futures::stream::iter(tasks)
110 .map(|(i, content)| async move {
111 let result = tokio::time::timeout(timeout, embed.call(&content)).await;
112 match result {
113 Ok(Ok(vec)) => Some((i, vec)),
114 Ok(Err(e)) => {
115 tracing::warn!(idx = i, err = %e, "embed_prepass: embed failed, skipping");
116 None
117 }
118 Err(_) => {
119 tracing::warn!(idx = i, "embed_prepass: embed timed out, skipping");
120 None
121 }
122 }
123 })
124 .buffer_unordered(concurrency)
125 .filter_map(|opt| async move { opt })
126 .collect()
127 .await
128}
129
130#[tracing::instrument(name = "context.fidelity.embed_prepass", skip_all)]
160pub async fn embed_prepass<F>(
161 messages: &[Message],
162 embed: &F,
163 config: &FidelityConfig,
164 inserted_count: usize,
165) -> std::collections::HashMap<usize, Vec<f32>>
166where
167 F: Fn(&str) -> EmbedFuture + Send + Sync,
168{
169 embed_prepass_dyn(messages, &ClosureEmbed(embed), config, inserted_count).await
170}
171
172fn truncate_to_byte_limit(s: &str, max_bytes: usize) -> String {
181 if s.len() <= max_bytes {
182 return s.to_string();
183 }
184 let boundary = s.floor_char_boundary(max_bytes);
185 s[..boundary].to_string()
186}
187
188struct FidelityScore {
189 score: f32,
190 level: ContextFidelity,
191 original_tokens: u32,
192}
193
194pub struct FidelityScorer;
220
221impl FidelityScorer {
222 #[tracing::instrument(name = "context.fidelity.score_and_apply", skip_all)]
254 #[allow(clippy::too_many_arguments, clippy::too_many_lines)]
255 pub async fn score_and_apply(
256 &self,
257 messages: &mut Vec<Message>,
258 query: &str,
259 planned_tools: &[PlannedToolHint],
260 config: &FidelityConfig,
261 tc: &dyn TokenCounting,
262 inserted_count: usize,
263 allow_upgrade: bool,
264 embed_provider: Option<&dyn LlmProviderDyn>,
265 compress_provider: Option<&dyn LlmProviderDyn>,
266 ) {
267 if !config.enabled || messages.is_empty() {
268 return;
269 }
270
271 let query_embedding: Option<Vec<f32>> = if let (true, Some(p)) =
273 (config.semantic_scoring_provider.is_some(), embed_provider)
274 && p.supports_embeddings()
275 {
276 match tokio::time::timeout(
277 Duration::from_secs(config.embed_timeout_secs),
278 p.embed(query),
279 )
280 .instrument(info_span!("context.fidelity.embed_query"))
281 .await
282 {
283 Ok(Ok(v)) => Some(v),
284 Ok(Err(e)) => {
285 tracing::warn!(error = %e, "semantic scoring provider unavailable, falling back to keyword");
286 None
287 }
288 Err(_) => {
289 tracing::warn!("fidelity query embed timed out, falling back to keyword");
290 None
291 }
292 }
293 } else {
294 None
295 };
296
297 if let (Some(q_emb), Some(p)) = (&query_embedding, embed_provider) {
301 let n = messages.len();
302 let score_end = if n > config.max_scored_messages {
303 n.saturating_sub(config.exempt_tail_messages)
304 } else {
305 n
306 };
307 let embeddings = embed_prepass_dyn(
308 &messages[..score_end],
309 &ProviderEmbed(p),
310 config,
311 inserted_count,
312 )
313 .instrument(info_span!("context.fidelity.embed_prepass"))
314 .await;
315 for (i, emb) in embeddings {
316 messages[i].metadata.embedding = Some(emb);
317 }
318 let _ = q_emb; }
320
321 let scores = compute_scores(
322 messages,
323 query,
324 planned_tools,
325 config,
326 tc,
327 inserted_count,
328 allow_upgrade,
329 query_embedding.as_deref(),
330 );
331 apply_scores(messages, &scores, config, tc, compress_provider).await;
332
333 let _merge_span = info_span!("context.fidelity.merge").entered();
334 let merged_count = merge_consecutive_placeholders(messages);
335 tracing::debug!(merged_count, "fidelity merge complete");
336 }
337}
338
339#[allow(clippy::too_many_arguments)]
340fn compute_scores(
341 messages: &[Message],
342 query: &str,
343 planned_tools: &[PlannedToolHint],
344 config: &FidelityConfig,
345 tc: &dyn TokenCounting,
346 inserted_count: usize,
347 allow_upgrade: bool,
348 query_embedding: Option<&[f32]>,
349) -> Vec<Option<FidelityScore>> {
350 let n = messages.len();
351
352 let score_end = if n > config.max_scored_messages {
354 n.saturating_sub(config.exempt_tail_messages)
355 } else {
356 n
357 };
358
359 let semantic_active = query.len() >= config.min_query_length;
360 let plan_active = !planned_tools.is_empty();
361 let query_words: std::collections::HashSet<&str> = if semantic_active {
363 query.split_whitespace().collect()
364 } else {
365 std::collections::HashSet::default()
366 };
367
368 let mut weight_sum = config.w_temporal + config.w_importance;
370 if semantic_active {
371 weight_sum += config.w_semantic;
372 }
373 if plan_active {
374 weight_sum += config.w_plan;
375 }
376 if weight_sum <= 0.0 {
377 weight_sum = 1.0;
378 }
379
380 #[allow(clippy::cast_precision_loss)]
381 let max_dist = score_end.saturating_sub(1) as f32;
382
383 let mut scores: Vec<Option<FidelityScore>> = (0..n).map(|_| None).collect();
384
385 for (i, msg) in messages.iter().enumerate().take(score_end) {
386 if is_exempt(msg, i, inserted_count) {
387 continue;
388 }
389
390 #[allow(clippy::cast_possible_truncation)]
391 let original_tokens = tc.count_tokens(&msg.content) as u32;
392
393 #[allow(clippy::cast_precision_loss)]
396 let temporal = if max_dist > 0.0 {
397 let distance_from_end = (score_end - 1 - i) as f32;
398 1.0 - distance_from_end / max_dist
399 } else {
400 1.0
401 };
402 let importance = if msg
404 .parts
405 .iter()
406 .any(|p| matches!(p, MessagePart::ToolResult { .. }))
407 {
408 0.4
409 } else {
410 role_weight(msg.role)
411 };
412 let semantic = if semantic_active {
413 match (query_embedding, msg.metadata.embedding.as_deref()) {
414 (Some(q_emb), Some(m_emb)) => semantic_overlap(m_emb, q_emb),
415 _ => keyword_overlap(&msg.content, &query_words),
416 }
417 } else {
418 0.0
419 };
420 let plan = if plan_active {
421 plan_relevance(&msg.content, planned_tools)
422 } else {
423 0.0
424 };
425
426 let raw = config.w_temporal * temporal
427 + config.w_importance * importance
428 + if semantic_active {
429 config.w_semantic * semantic
430 } else {
431 0.0
432 }
433 + if plan_active {
434 config.w_plan * plan
435 } else {
436 0.0
437 };
438
439 let score = (raw / weight_sum).clamp(0.0, 1.0);
440 let candidate_level = score_to_level(score, config);
441
442 let level = if allow_upgrade {
447 candidate_level
448 } else {
449 match msg.metadata.fidelity_tag {
450 Some(ContextFidelity::Placeholder) => ContextFidelity::Placeholder,
451 Some(ContextFidelity::Compressed) => {
452 if candidate_level == ContextFidelity::Full {
453 ContextFidelity::Compressed
454 } else {
455 candidate_level
456 }
457 }
458 _ => candidate_level,
459 }
460 };
461
462 scores[i] = Some(FidelityScore {
463 score,
464 level,
465 original_tokens,
466 });
467 }
468
469 apply_tool_pair_atomicity(messages, &mut scores, config);
470 scores
471}
472
473#[tracing::instrument(name = "context.fidelity.apply", skip_all)]
474async fn apply_scores(
475 messages: &mut [Message],
476 scores: &[Option<FidelityScore>],
477 config: &FidelityConfig,
478 tc: &dyn TokenCounting,
479 provider: Option<&dyn LlmProviderDyn>,
480) {
481 let (mut full_count, mut compressed_count, mut placeholder_count, mut tokens_saved) =
482 (0u32, 0u32, 0u32, 0u32);
483
484 for (i, msg) in messages.iter_mut().enumerate() {
485 let Some(ref fs) = scores[i] else { continue };
486 match fs.level {
487 ContextFidelity::Compressed => {
488 #[allow(clippy::cast_possible_truncation)]
489 let original_tokens = fs.original_tokens;
490 render_compressed(msg, config, tc, provider).await;
491 #[allow(clippy::cast_possible_truncation)]
492 let new_tokens = tc.count_tokens(&msg.content) as u32;
493 tokens_saved += original_tokens.saturating_sub(new_tokens);
494 compressed_count += 1;
495 }
496 ContextFidelity::Placeholder => {
497 render_placeholder(msg, fs.score, fs.original_tokens);
498 placeholder_count += 1;
499 }
500 _ => {
502 msg.metadata.fidelity_tag = Some(ContextFidelity::Full);
503 full_count += 1;
504 }
505 }
506 }
507
508 tracing::debug!(
509 full_count,
510 compressed_count,
511 placeholder_count,
512 tokens_saved,
513 "fidelity apply complete"
514 );
515}
516
517fn is_exempt(msg: &Message, idx: usize, inserted_count: usize) -> bool {
518 (idx == 0 && msg.role == Role::System)
523 || msg.metadata.focus_pinned
524 || msg.content.starts_with(CORRECTIONS_PREFIX)
525 || (idx >= 1 && idx < 1 + inserted_count)
526}
527
528fn role_weight(role: Role) -> f32 {
529 match role {
530 Role::System => 1.0,
531 Role::User => 0.8,
532 Role::Assistant | _ => 0.6,
533 }
534}
535
536fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
537 if a.len() != b.len() || a.is_empty() {
538 return 0.0;
539 }
540 let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
541 let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
542 let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
543 if norm_a == 0.0 || norm_b == 0.0 {
544 return 0.0;
545 }
546 (dot / (norm_a * norm_b)).clamp(0.0, 1.0)
547}
548
549fn semantic_overlap(msg_embedding: &[f32], query_embedding: &[f32]) -> f32 {
550 cosine_similarity(msg_embedding, query_embedding)
551}
552
553fn keyword_overlap(content: &str, query_words: &std::collections::HashSet<&str>) -> f32 {
557 let content_words: std::collections::HashSet<&str> = content.split_whitespace().collect();
558 let min_len = content_words.len().min(query_words.len());
559 if min_len == 0 {
560 return 0.0;
561 }
562 #[allow(clippy::cast_precision_loss)]
563 let result = content_words.intersection(query_words).count() as f32 / min_len as f32;
564 result.clamp(0.0, 1.0)
565}
566
567fn plan_relevance(content: &str, planned_tools: &[PlannedToolHint]) -> f32 {
571 if planned_tools.is_empty() {
572 return 0.0;
573 }
574 let content_words: std::collections::HashSet<&str> = content.split_whitespace().collect();
575 let mut weighted_sum = 0.0f32;
576 let mut weight_total = 0.0f32;
577 for hint in planned_tools {
578 let dist = f32::from(hint.distance_from_current.max(1));
579 let weight = 1.0 / dist;
580 weight_total += weight;
581 let hint_words: std::collections::HashSet<&str> =
582 hint.keywords.iter().map(String::as_str).collect();
583 let min_len = content_words.len().min(hint_words.len());
584 if min_len == 0 {
585 continue;
586 }
587 #[allow(clippy::cast_precision_loss)]
588 let overlap = content_words.intersection(&hint_words).count() as f32 / min_len as f32;
589 weighted_sum += weight * overlap.clamp(0.0, 1.0);
590 }
591 if weight_total <= 0.0 {
592 return 0.0;
593 }
594 (weighted_sum / weight_total).clamp(0.0, 1.0)
595}
596
597fn apply_tool_pair_atomicity(
602 messages: &[Message],
603 scores: &mut [Option<FidelityScore>],
604 config: &FidelityConfig,
605) {
606 let mut tool_result_map: std::collections::HashMap<&str, usize> =
608 std::collections::HashMap::new();
609 for (i, msg) in messages.iter().enumerate() {
610 for part in &msg.parts {
611 if let MessagePart::ToolResult { tool_use_id, .. } = part {
612 tool_result_map.insert(tool_use_id.as_str(), i);
613 }
614 }
615 }
616
617 for (i, msg) in messages.iter().enumerate().rev() {
619 for part in &msg.parts {
620 if let MessagePart::ToolUse { id, .. } = part
621 && let Some(&result_idx) = tool_result_map.get(id.as_str())
622 {
623 let score_a = scores[i].as_ref().map_or(1.0, |s| s.score);
624 let score_b = scores[result_idx].as_ref().map_or(1.0, |s| s.score);
625 let min_score = score_a.min(score_b);
626
627 let level_a = scores[i]
631 .as_ref()
632 .map_or(ContextFidelity::Full, |s| s.level);
633 let level_b = scores[result_idx]
634 .as_ref()
635 .map_or(ContextFidelity::Full, |s| s.level);
636 let float_level = score_to_level(min_score, config);
637 let min_level = more_restrictive(more_restrictive(level_a, level_b), float_level);
638
639 let tokens_a = scores[i].as_ref().map_or(0, |s| s.original_tokens);
640 let tokens_b = scores[result_idx].as_ref().map_or(0, |s| s.original_tokens);
641 scores[i] = Some(FidelityScore {
642 score: min_score,
643 level: min_level,
644 original_tokens: tokens_a,
645 });
646 scores[result_idx] = Some(FidelityScore {
647 score: min_score,
648 level: min_level,
649 original_tokens: tokens_b,
650 });
651 }
652 }
653 }
654}
655
656fn more_restrictive(a: ContextFidelity, b: ContextFidelity) -> ContextFidelity {
660 use ContextFidelity::{Compressed, Full, Placeholder};
661 match (a, b) {
662 (Placeholder, _) | (_, Placeholder) => Placeholder,
663 (Compressed, _) | (_, Compressed) => Compressed,
664 _ => Full,
665 }
666}
667
668fn score_to_level(score: f32, config: &FidelityConfig) -> ContextFidelity {
669 if score >= config.full_threshold {
670 ContextFidelity::Full
671 } else if score >= config.compressed_threshold {
672 ContextFidelity::Compressed
673 } else {
674 ContextFidelity::Placeholder
675 }
676}
677
678#[tracing::instrument(name = "context.fidelity.render_compressed", skip_all)]
679async fn render_compressed(
680 msg: &mut Message,
681 config: &FidelityConfig,
682 tc: &dyn TokenCounting,
683 provider: Option<&dyn LlmProviderDyn>,
684) {
685 if let Some(summary) = msg.metadata.deferred_summary.take() {
687 msg.content = summary;
688 } else if config.compress_provider.is_some()
689 && let Some(p) = provider
690 {
691 let input_tokens = tc.count_tokens(&msg.content);
693 if input_tokens > config.compressed_max_tokens * 2 && input_tokens > 0 {
695 if let Some(max_in) = config.max_compress_input_tokens {
697 apply_input_cap(&mut msg.content, max_in);
698 }
699
700 let prompt = format!(
701 "Summarize in {} tokens or fewer: {}",
702 config.compressed_max_tokens, msg.content
703 );
704 let req = vec![Message {
705 role: Role::User,
706 content: prompt,
707 parts: vec![],
708 metadata: MessageMetadata::default(),
709 }];
710
711 let span = info_span!(
712 "context.fidelity.compress_llm",
713 input_tokens,
714 cached = false,
715 );
716 let result = tokio::time::timeout(
717 Duration::from_secs(config.compress_timeout_secs),
718 p.chat(&req),
719 )
720 .instrument(span)
721 .await;
722
723 match result {
724 Ok(Ok(summary)) => {
725 msg.metadata.deferred_summary = Some(summary.clone());
726 msg.content = summary;
727 }
728 Ok(Err(e)) => {
729 tracing::debug!(error = %e, "compress_llm failed, falling back to truncation");
730 }
731 Err(_) => {
732 tracing::warn!("compress_llm timed out, falling back to truncation");
733 }
734 }
735 }
736 } else if let Some(max_in) = config.max_compress_input_tokens {
737 apply_input_cap(&mut msg.content, max_in);
739 }
740
741 truncate_to_tokens(&mut msg.content, config.compressed_max_tokens, tc);
744 msg.parts.clear();
745 msg.metadata.fidelity_tag = Some(ContextFidelity::Compressed);
746}
747
748pub fn apply_input_cap(content: &mut String, max_tokens: usize) {
756 let max_bytes = max_tokens.saturating_mul(4);
757 if content.len() > max_bytes {
758 let boundary = content.floor_char_boundary(max_bytes);
759 content.truncate(boundary);
760 }
761}
762
763fn truncate_to_tokens(content: &mut String, max_tokens: usize, tc: &dyn TokenCounting) {
764 if tc.count_tokens(content) <= max_tokens {
765 return;
766 }
767 let mut lo: usize = 0;
772 let mut hi: usize = content.len();
773 while hi - lo > 1 {
774 let mid = content.floor_char_boundary(usize::midpoint(lo, hi));
775 if mid == lo {
776 hi = mid;
779 } else if tc.count_tokens(&content[..mid]) <= max_tokens {
780 lo = mid;
781 } else {
782 hi = mid;
783 }
784 }
785 content.truncate(lo);
786}
787
788fn render_placeholder(msg: &mut Message, score: f32, original_tokens: u32) {
789 let role_str = match msg.role {
790 Role::System => "system",
791 Role::Assistant => "assistant",
792 Role::User | _ => "user",
793 };
794 msg.content = format!(
795 "[placeholder: role={role_str}, original_tokens={original_tokens}, importance={score:.2}]"
796 );
797 msg.parts.clear();
798 msg.metadata.fidelity_tag = Some(ContextFidelity::Placeholder);
799}
800
801fn merge_consecutive_placeholders(messages: &mut Vec<Message>) -> usize {
805 let mut merged_count = 0usize;
806 let mut i = 0;
807 while i < messages.len() {
808 if messages[i].metadata.fidelity_tag != Some(ContextFidelity::Placeholder)
809 || messages[i].role == Role::System
810 {
811 i += 1;
812 continue;
813 }
814 let role = messages[i].role;
815 let mut j = i + 1;
816 while j < messages.len()
817 && messages[j].metadata.fidelity_tag == Some(ContextFidelity::Placeholder)
818 && messages[j].role == role
819 {
820 j += 1;
821 }
822 if j - i <= 1 {
823 i += 1;
824 continue;
825 }
826 let count = j - i;
827 let mut total_tokens = 0u32;
828 let mut importance_sum = 0.0f32;
829 for msg in &messages[i..j] {
830 total_tokens += parse_placeholder_tokens(&msg.content);
831 importance_sum += parse_placeholder_importance(&msg.content);
832 }
833 debug_assert!(count >= 2, "placeholder merge triggered with count={count}");
834 #[allow(clippy::cast_precision_loss)]
835 let avg_importance = if count > 0 {
836 importance_sum / count as f32
837 } else {
838 0.0
839 };
840 let role_str = match role {
841 Role::System => "system",
842 Role::Assistant => "assistant",
843 Role::User | _ => "user",
844 };
845 let merged_content = format!(
846 "[placeholder: {count} messages, role={role_str}, total_tokens={total_tokens}, avg_importance={avg_importance:.2}]"
847 );
848 let first = messages[i].clone();
849 messages.drain(i..j);
850 messages.insert(
851 i,
852 Message {
853 role: first.role,
854 content: merged_content,
855 parts: vec![],
856 metadata: {
857 let mut m = first.metadata;
858 m.fidelity_tag = Some(ContextFidelity::Placeholder);
859 m
860 },
861 },
862 );
863 merged_count += count - 1;
864 i += 1;
865 }
866 merged_count
867}
868
869fn parse_placeholder_tokens(content: &str) -> u32 {
870 for part in content.split(',') {
871 let part = part.trim();
872 for prefix in &["original_tokens=", "total_tokens="] {
873 if let Some(rest) = part.strip_prefix(prefix)
874 && let Ok(n) = rest.trim_end_matches(']').trim().parse::<u32>()
875 {
876 return n;
877 }
878 }
879 }
880 0
881}
882
883fn parse_placeholder_importance(content: &str) -> f32 {
884 for part in content.split(',') {
885 let part = part.trim();
886 for prefix in &["importance=", "avg_importance="] {
887 if let Some(rest) = part.strip_prefix(prefix)
888 && let Ok(v) = rest.trim_end_matches(']').trim().parse::<f32>()
889 {
890 return v;
891 }
892 }
893 }
894 0.0
895}
896
897#[cfg(test)]
898mod tests {
899 use super::*;
900 use zeph_common::ProviderName;
901 use zeph_llm::provider::{Message, MessageMetadata, MessagePart, Role};
902
903 struct FixedTc(usize);
904 impl TokenCounting for FixedTc {
905 fn count_tokens(&self, text: &str) -> usize {
906 text.len() / self.0.max(1)
907 }
908
909 fn count_tool_schema_tokens(&self, _schema: &serde_json::Value) -> usize {
910 0
911 }
912 }
913
914 fn make_msg(role: Role, content: &str) -> Message {
915 Message {
916 role,
917 content: content.to_string(),
918 parts: vec![],
919 metadata: MessageMetadata::default(),
920 }
921 }
922
923 fn make_cfg() -> FidelityConfig {
924 FidelityConfig {
925 enabled: true,
926 w_semantic: 0.3,
927 w_temporal: 0.3,
928 w_importance: 0.2,
929 w_plan: 0.2,
930 full_threshold: 0.7,
931 compressed_threshold: 0.3,
932 compressed_max_tokens: 50,
933 regrade_threshold: 0.6,
934 min_query_length: 8,
935 max_scored_messages: 500,
936 exempt_tail_messages: 0,
937 compress_provider: None,
938 semantic_scoring_provider: None,
939 lookahead_depth: 3,
940 embed_concurrency: 32,
941 max_embed_input_tokens: None,
942 max_compress_input_tokens: None,
943 embed_timeout_secs: 30,
944 compress_timeout_secs: 30,
945 }
946 }
947
948 #[tokio::test]
950 async fn empty_window_no_change() {
951 let scorer = FidelityScorer;
952 let cfg = make_cfg();
953 let tc = FixedTc(4);
954 let mut messages: Vec<Message> = vec![];
955 scorer
956 .score_and_apply(
957 &mut messages,
958 "query text",
959 &[],
960 &cfg,
961 &tc,
962 0,
963 false,
964 None,
965 None,
966 )
967 .await;
968 assert!(messages.is_empty());
969 }
970
971 #[tokio::test]
973 async fn all_exempt_no_downgrade() {
974 let scorer = FidelityScorer;
975 let cfg = make_cfg();
976 let tc = FixedTc(4);
977 let mut messages = vec![
978 make_msg(Role::System, "system prompt"),
979 make_msg(Role::User, "memory context"),
981 ];
982 scorer
983 .score_and_apply(&mut messages, "short", &[], &cfg, &tc, 1, false, None, None)
984 .await;
985 for msg in &messages {
986 assert!(
987 msg.metadata.fidelity_tag.is_none()
988 || msg.metadata.fidelity_tag == Some(ContextFidelity::Full)
989 );
990 }
991 }
992
993 #[tokio::test]
995 async fn tool_pair_atomicity() {
996 let scorer = FidelityScorer;
997 let cfg = FidelityConfig {
999 full_threshold: 0.9,
1000 compressed_threshold: 0.5,
1001 ..make_cfg()
1002 };
1003 let tc = FixedTc(4);
1004 let tool_use_id = "abc123".to_string();
1005 let mut tool_use_msg = make_msg(Role::Assistant, "calling tool");
1006 tool_use_msg.parts = vec![MessagePart::ToolUse {
1007 id: tool_use_id.clone(),
1008 name: "shell".to_string(),
1009 input: serde_json::json!({}),
1010 }];
1011 let mut tool_result_msg = make_msg(Role::User, "tool result body");
1012 tool_result_msg.parts = vec![MessagePart::ToolResult {
1013 tool_use_id: tool_use_id.clone(),
1014 content: "result".to_string(),
1015 is_error: false,
1016 }];
1017 let mut messages = vec![
1018 make_msg(Role::System, "system"),
1019 tool_use_msg,
1020 tool_result_msg,
1021 ];
1022 scorer
1023 .score_and_apply(
1024 &mut messages,
1025 "completely unrelated query blah",
1026 &[],
1027 &cfg,
1028 &tc,
1029 0,
1030 false,
1031 None,
1032 None,
1033 )
1034 .await;
1035 let tag_a = messages[1].metadata.fidelity_tag;
1036 let tag_b = messages[2].metadata.fidelity_tag;
1037 assert_eq!(tag_a, tag_b, "tool pair must share fidelity level");
1038 }
1039
1040 #[tokio::test]
1042 async fn same_role_placeholder_merge() {
1043 let scorer = FidelityScorer;
1044 let cfg = FidelityConfig {
1046 full_threshold: 2.0, compressed_threshold: 1.5, ..make_cfg()
1049 };
1050 let tc = FixedTc(4);
1051 let mut messages: Vec<Message> = std::iter::once(make_msg(Role::System, "system"))
1052 .chain((0..5).map(|i| make_msg(Role::Assistant, &format!("msg {i}"))))
1053 .collect();
1054 scorer
1055 .score_and_apply(
1056 &mut messages,
1057 "some query here",
1058 &[],
1059 &cfg,
1060 &tc,
1061 0,
1062 false,
1063 None,
1064 None,
1065 )
1066 .await;
1067 assert_eq!(
1069 messages.len(),
1070 2,
1071 "5 assistant placeholders must merge to 1"
1072 );
1073 assert!(messages[1].content.contains("5 messages"));
1074 }
1075
1076 #[tokio::test]
1078 async fn score_normalization_no_panic() {
1079 let scorer = FidelityScorer;
1080 let cfg = make_cfg();
1081 let tc = FixedTc(4);
1082 let mut messages = vec![
1083 make_msg(Role::System, "system"),
1084 make_msg(Role::User, "hello"),
1085 make_msg(Role::Assistant, "world response"),
1086 ];
1087 scorer
1088 .score_and_apply(
1089 &mut messages,
1090 "hello world signal",
1091 &[],
1092 &cfg,
1093 &tc,
1094 0,
1095 false,
1096 None,
1097 None,
1098 )
1099 .await;
1100 for msg in &messages {
1101 let _ = msg.metadata.fidelity_tag;
1102 }
1103 }
1104
1105 #[tokio::test]
1107 async fn short_query_fallback() {
1108 let scorer = FidelityScorer;
1109 let cfg = FidelityConfig {
1110 min_query_length: 8,
1111 ..make_cfg()
1112 };
1113 let tc = FixedTc(4);
1114 let mut messages = vec![
1115 make_msg(Role::System, "system"),
1116 make_msg(Role::User, "test"),
1117 ];
1118 scorer
1120 .score_and_apply(&mut messages, "short", &[], &cfg, &tc, 0, false, None, None)
1121 .await;
1122 }
1123
1124 #[tokio::test]
1130 async fn memory_first_bypass_is_callers_responsibility() {
1131 let scorer = FidelityScorer;
1132 let cfg = FidelityConfig {
1135 enabled: false,
1136 ..make_cfg()
1137 };
1138 let tc = FixedTc(4);
1139 let mut messages = vec![
1140 make_msg(Role::System, "system prompt"),
1141 make_msg(Role::User, "memory-injected context"),
1142 make_msg(Role::Assistant, "response"),
1143 ];
1144 let before: Vec<_> = messages.iter().map(|m| m.content.clone()).collect();
1145 scorer
1147 .score_and_apply(
1148 &mut messages,
1149 "some user query text here",
1150 &[],
1151 &cfg,
1152 &tc,
1153 2,
1154 false,
1155 None,
1156 None,
1157 )
1158 .await;
1159 for (msg, orig) in messages.iter().zip(&before) {
1160 assert_eq!(msg.content, *orig, "content must be unchanged");
1161 assert!(
1162 msg.metadata.fidelity_tag.is_none(),
1163 "no fidelity tag must be set"
1164 );
1165 }
1166 }
1167
1168 #[tokio::test]
1170 async fn enabled_false_guard() {
1171 let scorer = FidelityScorer;
1172 let cfg = FidelityConfig {
1173 enabled: false,
1174 ..make_cfg()
1175 };
1176 let tc = FixedTc(4);
1177 let mut messages = vec![
1178 make_msg(Role::System, "system"),
1179 make_msg(Role::User, "user message that would normally be scored"),
1180 ];
1181 let original_contents: Vec<String> = messages.iter().map(|m| m.content.clone()).collect();
1182 scorer
1183 .score_and_apply(
1184 &mut messages,
1185 "query text here",
1186 &[],
1187 &cfg,
1188 &tc,
1189 0,
1190 false,
1191 None,
1192 None,
1193 )
1194 .await;
1195 for (msg, orig) in messages.iter().zip(&original_contents) {
1196 assert_eq!(msg.content, *orig);
1197 assert!(msg.metadata.fidelity_tag.is_none());
1198 }
1199 }
1200
1201 #[tokio::test]
1203 async fn score_always_in_range() {
1204 let scorer = FidelityScorer;
1205 let cfg = FidelityConfig {
1206 enabled: true,
1207 w_semantic: 0.0,
1208 w_temporal: 0.0,
1209 w_importance: 0.0,
1210 w_plan: 0.0,
1211 full_threshold: 0.7,
1212 compressed_threshold: 0.3,
1213 compressed_max_tokens: 50,
1214 regrade_threshold: 0.6,
1215 min_query_length: 0,
1216 max_scored_messages: 500,
1217 exempt_tail_messages: 0,
1218 compress_provider: None,
1219 semantic_scoring_provider: None,
1220 lookahead_depth: 3,
1221 embed_concurrency: 32,
1222 max_embed_input_tokens: None,
1223 max_compress_input_tokens: None,
1224 embed_timeout_secs: 30,
1225 compress_timeout_secs: 30,
1226 };
1227 let tc = FixedTc(4);
1228 let mut messages = vec![make_msg(Role::System, ""), make_msg(Role::User, "")];
1229 scorer
1231 .score_and_apply(&mut messages, "", &[], &cfg, &tc, 0, false, None, None)
1232 .await;
1233 }
1234
1235 #[tokio::test]
1237 async fn placeholder_uses_tc_count_tokens() {
1238 let scorer = FidelityScorer;
1239 let cfg = FidelityConfig {
1240 full_threshold: 2.0,
1241 compressed_threshold: 1.5,
1242 ..make_cfg()
1243 };
1244 let tc = FixedTc(1); let mut messages = vec![
1246 make_msg(Role::System, "system"),
1247 make_msg(Role::User, "user message content for placeholder rendering"),
1248 ];
1249 scorer
1250 .score_and_apply(
1251 &mut messages,
1252 "some query text here",
1253 &[],
1254 &cfg,
1255 &tc,
1256 0,
1257 false,
1258 None,
1259 None,
1260 )
1261 .await;
1262 assert_eq!(
1263 messages[1].metadata.fidelity_tag,
1264 Some(ContextFidelity::Placeholder)
1265 );
1266 assert!(messages[1].content.starts_with("[placeholder:"));
1267 }
1268
1269 #[tokio::test]
1276 async fn exempt_tail_messages_large_window() {
1277 let scorer = FidelityScorer;
1278 let cfg = FidelityConfig {
1279 full_threshold: 2.0,
1281 compressed_threshold: 1.5,
1282 max_scored_messages: 10,
1283 exempt_tail_messages: 5,
1284 ..make_cfg()
1285 };
1286 let tc = FixedTc(4);
1287
1288 let mut messages: Vec<Message> = std::iter::once(make_msg(Role::System, "system prompt"))
1295 .chain((1..15).map(|i| make_msg(Role::Assistant, &format!("assistant message {i}"))))
1296 .chain((15..20).map(|i| {
1297 let mut m = make_msg(Role::User, &format!("tail message {i}"));
1298 m.metadata.focus_pinned = true;
1299 m
1300 }))
1301 .collect();
1302
1303 scorer
1304 .score_and_apply(
1305 &mut messages,
1306 "query text here long",
1307 &[],
1308 &cfg,
1309 &tc,
1310 0,
1311 false,
1312 None,
1313 None,
1314 )
1315 .await;
1316
1317 let tail: Vec<_> = messages
1319 .iter()
1320 .filter(|m| m.metadata.focus_pinned)
1321 .collect();
1322 assert_eq!(
1323 tail.len(),
1324 5,
1325 "all 5 tail messages must survive the merge pass"
1326 );
1327 for msg in &tail {
1328 assert!(
1329 msg.metadata.fidelity_tag.is_none(),
1330 "tail message must have no fidelity_tag, got {:?}",
1331 msg.metadata.fidelity_tag
1332 );
1333 }
1334 }
1335
1336 #[tokio::test]
1340 async fn exempt_tail_messages_small_window_no_effect() {
1341 let scorer = FidelityScorer;
1342 let cfg = FidelityConfig {
1343 full_threshold: 2.0,
1344 compressed_threshold: 1.5,
1345 max_scored_messages: 10,
1346 exempt_tail_messages: 5,
1347 ..make_cfg()
1348 };
1349 let tc = FixedTc(4);
1350 let roles = [Role::User, Role::Assistant];
1353 let mut messages: Vec<Message> = std::iter::once(make_msg(Role::System, "system prompt"))
1354 .chain((1..8usize).map(|i| make_msg(roles[i % 2], &format!("message {i}"))))
1355 .collect();
1356 scorer
1357 .score_and_apply(
1358 &mut messages,
1359 "query text here long",
1360 &[],
1361 &cfg,
1362 &tc,
1363 0,
1364 false,
1365 None,
1366 None,
1367 )
1368 .await;
1369 let untagged_count = messages[1..]
1372 .iter()
1373 .filter(|m| m.metadata.fidelity_tag.is_none())
1374 .count();
1375 assert_eq!(
1376 untagged_count, 0,
1377 "all non-system messages must be scored when n <= max_scored_messages"
1378 );
1379 }
1380
1381 #[tokio::test]
1383 async fn compressed_uses_deferred_summary() {
1384 let scorer = FidelityScorer;
1385 let cfg = FidelityConfig {
1386 full_threshold: 2.0, compressed_threshold: 0.0, compressed_max_tokens: 5,
1389 ..make_cfg()
1390 };
1391 let tc = FixedTc(4);
1392 let mut msg_with_summary =
1393 make_msg(Role::User, "original long content that would be truncated");
1394 msg_with_summary.metadata.deferred_summary = Some("short summary".to_string());
1395 let mut messages = vec![make_msg(Role::System, "system"), msg_with_summary];
1396 scorer
1397 .score_and_apply(
1398 &mut messages,
1399 "query text here long",
1400 &[],
1401 &cfg,
1402 &tc,
1403 0,
1404 false,
1405 None,
1406 None,
1407 )
1408 .await;
1409 assert_eq!(
1410 messages[1].metadata.fidelity_tag,
1411 Some(ContextFidelity::Compressed)
1412 );
1413 assert_eq!(messages[1].content, "short summary");
1414 }
1415
1416 fn make_msg_with_fidelity(role: Role, content: &str, tag: Option<ContextFidelity>) -> Message {
1419 let mut m = make_msg(role, content);
1420 m.metadata.fidelity_tag = tag;
1421 m
1422 }
1423
1424 #[tokio::test]
1426 async fn floor_prevents_compressed_upgrade_to_full() {
1427 let scorer = FidelityScorer;
1428 let cfg = FidelityConfig {
1429 full_threshold: 0.0,
1431 compressed_threshold: -1.0,
1432 ..make_cfg()
1433 };
1434 let tc = FixedTc(4);
1435 let mut messages = vec![
1436 make_msg(Role::System, "system"),
1437 make_msg_with_fidelity(
1438 Role::User,
1439 "query text here long keyword",
1440 Some(ContextFidelity::Compressed),
1441 ),
1442 ];
1443 scorer
1444 .score_and_apply(
1445 &mut messages,
1446 "query text here long keyword",
1447 &[],
1448 &cfg,
1449 &tc,
1450 0,
1451 false,
1452 None,
1453 None,
1454 )
1455 .await;
1456 assert_eq!(
1457 messages[1].metadata.fidelity_tag,
1458 Some(ContextFidelity::Compressed),
1459 "Compressed floor must block upgrade to Full"
1460 );
1461 }
1462
1463 #[tokio::test]
1465 async fn floor_prevents_placeholder_upgrade_to_full() {
1466 let scorer = FidelityScorer;
1467 let cfg = FidelityConfig {
1468 full_threshold: 0.0,
1469 compressed_threshold: -1.0,
1470 ..make_cfg()
1471 };
1472 let tc = FixedTc(4);
1473 let mut messages = vec![
1474 make_msg(Role::System, "system"),
1475 make_msg_with_fidelity(
1476 Role::User,
1477 "query text here long keyword",
1478 Some(ContextFidelity::Placeholder),
1479 ),
1480 ];
1481 scorer
1482 .score_and_apply(
1483 &mut messages,
1484 "query text here long keyword",
1485 &[],
1486 &cfg,
1487 &tc,
1488 0,
1489 false,
1490 None,
1491 None,
1492 )
1493 .await;
1494 assert_eq!(
1495 messages[1].metadata.fidelity_tag,
1496 Some(ContextFidelity::Placeholder),
1497 "Placeholder floor must block upgrade to Full"
1498 );
1499 }
1500
1501 #[tokio::test]
1503 async fn floor_prevents_placeholder_upgrade_to_compressed() {
1504 let scorer = FidelityScorer;
1505 let cfg = FidelityConfig {
1507 full_threshold: 2.0,
1508 compressed_threshold: 0.0,
1509 ..make_cfg()
1510 };
1511 let tc = FixedTc(4);
1512 let mut messages = vec![
1513 make_msg(Role::System, "system"),
1514 make_msg_with_fidelity(
1515 Role::User,
1516 "message content",
1517 Some(ContextFidelity::Placeholder),
1518 ),
1519 ];
1520 scorer
1521 .score_and_apply(
1522 &mut messages,
1523 "query text here long",
1524 &[],
1525 &cfg,
1526 &tc,
1527 0,
1528 false,
1529 None,
1530 None,
1531 )
1532 .await;
1533 assert_eq!(
1534 messages[1].metadata.fidelity_tag,
1535 Some(ContextFidelity::Placeholder),
1536 "Placeholder floor must block upgrade to Compressed"
1537 );
1538 }
1539
1540 #[tokio::test]
1542 async fn floor_allows_further_downgrade() {
1543 let scorer = FidelityScorer;
1544 let cfg = FidelityConfig {
1545 full_threshold: 2.0,
1546 compressed_threshold: 2.0, ..make_cfg()
1548 };
1549 let tc = FixedTc(4);
1550 let mut messages = vec![
1551 make_msg(Role::System, "system"),
1552 make_msg_with_fidelity(
1553 Role::User,
1554 "some content",
1555 Some(ContextFidelity::Compressed),
1556 ),
1557 ];
1558 scorer
1559 .score_and_apply(
1560 &mut messages,
1561 "query text here long",
1562 &[],
1563 &cfg,
1564 &tc,
1565 0,
1566 false,
1567 None,
1568 None,
1569 )
1570 .await;
1571 assert_eq!(
1572 messages[1].metadata.fidelity_tag,
1573 Some(ContextFidelity::Placeholder),
1574 "downgrade from Compressed to Placeholder must be allowed"
1575 );
1576 }
1577
1578 #[tokio::test]
1580 async fn floor_no_constraint_when_none() {
1581 let scorer = FidelityScorer;
1582 let cfg = FidelityConfig {
1584 full_threshold: 0.0,
1585 compressed_threshold: -1.0,
1586 ..make_cfg()
1587 };
1588 let tc = FixedTc(4);
1589 let mut messages = vec![
1590 make_msg(Role::System, "system"),
1591 make_msg_with_fidelity(Role::User, "query text here long keyword", None),
1592 ];
1593 scorer
1594 .score_and_apply(
1595 &mut messages,
1596 "query text here long keyword",
1597 &[],
1598 &cfg,
1599 &tc,
1600 0,
1601 false,
1602 None,
1603 None,
1604 )
1605 .await;
1606 assert_eq!(
1607 messages[1].metadata.fidelity_tag,
1608 Some(ContextFidelity::Full),
1609 "None tag must not constrain scoring"
1610 );
1611 }
1612
1613 #[tokio::test]
1615 async fn allow_upgrade_bypasses_floor() {
1616 let scorer = FidelityScorer;
1617 let cfg = FidelityConfig {
1618 full_threshold: 0.0,
1619 compressed_threshold: -1.0,
1620 ..make_cfg()
1621 };
1622 let tc = FixedTc(4);
1623 let mut messages = vec![
1624 make_msg(Role::System, "system"),
1625 make_msg_with_fidelity(
1626 Role::User,
1627 "query text here long keyword",
1628 Some(ContextFidelity::Placeholder),
1629 ),
1630 ];
1631 scorer
1632 .score_and_apply(
1633 &mut messages,
1634 "query text here long keyword",
1635 &[],
1636 &cfg,
1637 &tc,
1638 0,
1639 true,
1640 None,
1641 None,
1642 )
1643 .await;
1644 assert_eq!(
1645 messages[1].metadata.fidelity_tag,
1646 Some(ContextFidelity::Full),
1647 "allow_upgrade=true must bypass the Placeholder floor"
1648 );
1649 }
1650
1651 #[test]
1655 fn truncate_no_op_below_limit() {
1656 let tc = FixedTc(1); let mut s = "hello".to_string(); truncate_to_tokens(&mut s, 10, &tc);
1659 assert_eq!(s, "hello");
1660 }
1661
1662 #[test]
1664 fn truncate_no_op_at_limit() {
1665 let tc = FixedTc(1);
1666 let mut s = "hello".to_string(); truncate_to_tokens(&mut s, 5, &tc);
1668 assert_eq!(s, "hello");
1669 }
1670
1671 #[test]
1673 fn truncate_minimal_one_over_limit() {
1674 let tc = FixedTc(1); let mut s = "abcdef".to_string(); truncate_to_tokens(&mut s, 5, &tc);
1677 assert!(
1678 tc.count_tokens(&s) <= 5,
1679 "result must fit in 5 tokens, got {}",
1680 tc.count_tokens(&s)
1681 );
1682 assert!(!s.is_empty(), "must keep prefix, not empty");
1683 }
1684
1685 #[test]
1687 fn truncate_preserves_90pct_of_limit() {
1688 let tc = FixedTc(1);
1690 let s_orig = "a".repeat(90);
1691 let mut s = s_orig.clone();
1692 truncate_to_tokens(&mut s, 100, &tc);
1693 assert_eq!(s, s_orig, "90% of limit must not be truncated");
1694 }
1695
1696 #[test]
1698 fn truncate_empty_string_no_op() {
1699 let tc = FixedTc(1);
1700 let mut s = String::new();
1701 truncate_to_tokens(&mut s, 5, &tc);
1702 assert!(s.is_empty());
1703 }
1704
1705 #[test]
1707 fn truncate_max_tokens_zero_clears_content() {
1708 let tc = FixedTc(1);
1709 let mut s = "hello world".to_string();
1710 truncate_to_tokens(&mut s, 0, &tc);
1711 assert!(s.is_empty(), "max_tokens=0 must clear content");
1712 }
1713
1714 #[test]
1716 fn truncate_multibyte_stays_on_char_boundary() {
1717 let tc = FixedTc(3);
1721 let mut s = "日本語".to_string();
1722 truncate_to_tokens(&mut s, 2, &tc);
1723 assert!(
1724 s.is_char_boundary(s.len()),
1725 "result must be on a valid char boundary"
1726 );
1727 assert!(tc.count_tokens(&s) <= 2);
1728 assert_eq!(s, "日本");
1729 }
1730
1731 #[tokio::test]
1733 async fn mixed_fidelity_tool_pair_floor_plus_atomicity() {
1734 let scorer = FidelityScorer;
1735 let cfg = FidelityConfig {
1737 full_threshold: 0.0,
1738 compressed_threshold: -1.0,
1739 ..make_cfg()
1740 };
1741 let tc = FixedTc(4);
1742 let tool_id = "tool-42".to_string();
1743
1744 let mut tool_use_msg = make_msg_with_fidelity(Role::Assistant, "call tool", None);
1745 tool_use_msg.parts = vec![MessagePart::ToolUse {
1746 id: tool_id.clone(),
1747 name: "shell".to_string(),
1748 input: serde_json::json!({}),
1749 }];
1750
1751 let mut tool_result_msg =
1752 make_msg_with_fidelity(Role::User, "result body", Some(ContextFidelity::Compressed));
1753 tool_result_msg.parts = vec![MessagePart::ToolResult {
1754 tool_use_id: tool_id.clone(),
1755 content: "output".to_string(),
1756 is_error: false,
1757 }];
1758
1759 let mut messages = vec![
1760 make_msg(Role::System, "system"),
1761 tool_use_msg,
1762 tool_result_msg,
1763 ];
1764
1765 scorer
1766 .score_and_apply(
1767 &mut messages,
1768 "query text here long",
1769 &[],
1770 &cfg,
1771 &tc,
1772 0,
1773 false,
1774 None,
1775 None,
1776 )
1777 .await;
1778
1779 let tag_use = messages[1].metadata.fidelity_tag;
1783 let tag_result = messages[2].metadata.fidelity_tag;
1784 assert_eq!(
1785 tag_use, tag_result,
1786 "tool pair must share the same fidelity level"
1787 );
1788 assert_eq!(
1789 tag_use,
1790 Some(ContextFidelity::Compressed),
1791 "atomicity must bring the tool-use down to the tool-result floor"
1792 );
1793 }
1794
1795 #[tokio::test]
1797 async fn compress_llm_path_stores_deferred_summary() {
1798 use zeph_llm::LlmError;
1799 use zeph_llm::provider::ChatStream;
1800
1801 #[derive(Debug)]
1802 struct MockProvider;
1803
1804 impl zeph_llm::provider::LlmProvider for MockProvider {
1805 async fn chat(&self, _messages: &[Message]) -> Result<String, LlmError> {
1806 Ok("summary text".to_string())
1807 }
1808
1809 async fn chat_stream(&self, _messages: &[Message]) -> Result<ChatStream, LlmError> {
1810 Err(LlmError::Unavailable)
1811 }
1812
1813 fn supports_streaming(&self) -> bool {
1814 false
1815 }
1816
1817 async fn embed(&self, _text: &str) -> Result<Vec<f32>, LlmError> {
1818 Err(LlmError::EmbedUnsupported {
1819 provider: "mock".into(),
1820 })
1821 }
1822
1823 fn supports_embeddings(&self) -> bool {
1824 false
1825 }
1826
1827 fn name(&self) -> &'static str {
1828 "mock"
1829 }
1830 }
1831
1832 let scorer = FidelityScorer;
1833 let cfg = FidelityConfig {
1834 enabled: true,
1835 full_threshold: 2.0,
1837 compressed_threshold: 0.0,
1838 compressed_max_tokens: 5,
1839 compress_provider: Some(ProviderName::new("mock")),
1840 ..make_cfg()
1841 };
1842 let tc = FixedTc(1);
1844 let content = "a".repeat(50); let mut messages = vec![
1846 make_msg(Role::System, "system"),
1847 make_msg(Role::User, &content),
1848 ];
1849
1850 let provider = MockProvider;
1851 scorer
1852 .score_and_apply(
1853 &mut messages,
1854 "some query text here",
1855 &[],
1856 &cfg,
1857 &tc,
1858 0,
1859 false,
1860 None,
1861 Some(&provider),
1862 )
1863 .await;
1864
1865 assert_eq!(
1866 messages[1].metadata.fidelity_tag,
1867 Some(ContextFidelity::Compressed),
1868 );
1869 assert!(
1871 tc.count_tokens(&messages[1].content) <= 5,
1872 "content must be capped to compressed_max_tokens after LLM summary"
1873 );
1874 assert_eq!(
1876 messages[1].metadata.deferred_summary,
1877 Some("summary text".to_string()),
1878 );
1879 }
1880
1881 #[tokio::test]
1883 async fn compress_llm_skipped_when_provider_none() {
1884 let scorer = FidelityScorer;
1885 let cfg = FidelityConfig {
1886 enabled: true,
1887 full_threshold: 2.0,
1888 compressed_threshold: 0.0,
1889 compressed_max_tokens: 5,
1890 compress_provider: Some(ProviderName::new("mock")),
1891 ..make_cfg()
1892 };
1893 let tc = FixedTc(1);
1894 let content = "a".repeat(50);
1895 let mut messages = vec![
1896 make_msg(Role::System, "system"),
1897 make_msg(Role::User, &content),
1898 ];
1899
1900 scorer
1901 .score_and_apply(
1902 &mut messages,
1903 "some query text here",
1904 &[],
1905 &cfg,
1906 &tc,
1907 0,
1908 false,
1909 None,
1910 None,
1911 )
1912 .await;
1913
1914 assert_eq!(
1915 messages[1].metadata.fidelity_tag,
1916 Some(ContextFidelity::Compressed),
1917 );
1918 assert!(
1920 messages[1].metadata.deferred_summary.is_none(),
1921 "deferred_summary must not be populated via truncation path"
1922 );
1923 assert!(
1925 messages[1].content.len() <= 5,
1926 "content must be truncated, got len={}",
1927 messages[1].content.len()
1928 );
1929 }
1930
1931 #[test]
1933 fn cosine_similarity_identical() {
1934 let v = vec![1.0f32, 0.0, 0.0];
1935 assert!((cosine_similarity(&v, &v) - 1.0).abs() < 1e-6);
1936 }
1937
1938 #[test]
1939 fn cosine_similarity_orthogonal() {
1940 let a = vec![1.0f32, 0.0, 0.0];
1941 let b = vec![0.0f32, 1.0, 0.0];
1942 assert!(cosine_similarity(&a, &b).abs() < 1e-6);
1943 }
1944
1945 #[test]
1946 fn cosine_similarity_zero_vector() {
1947 let a = vec![0.0f32, 0.0, 0.0];
1948 let b = vec![1.0f32, 0.0, 0.0];
1949 assert!(cosine_similarity(&a, &b).abs() < f32::EPSILON);
1950 }
1951
1952 #[test]
1953 fn cosine_similarity_empty() {
1954 assert!(cosine_similarity(&[], &[]).abs() < f32::EPSILON);
1955 }
1956
1957 #[test]
1958 fn cosine_similarity_dimension_mismatch() {
1959 let a = vec![1.0f32, 0.0];
1960 let b = vec![1.0f32, 0.0, 0.0];
1961 assert!(cosine_similarity(&a, &b).abs() < f32::EPSILON);
1962 }
1963
1964 #[tokio::test]
1971 async fn semantic_scoring_higher_for_similar_messages() {
1972 use zeph_llm::LlmError;
1973 use zeph_llm::provider::ChatStream;
1974
1975 #[derive(Debug)]
1976 struct EmbedMockProvider;
1977
1978 impl zeph_llm::provider::LlmProvider for EmbedMockProvider {
1979 async fn chat(&self, _: &[Message]) -> Result<String, LlmError> {
1980 Err(LlmError::Unavailable)
1981 }
1982 async fn chat_stream(&self, _: &[Message]) -> Result<ChatStream, LlmError> {
1983 Err(LlmError::Unavailable)
1984 }
1985 fn supports_streaming(&self) -> bool {
1986 false
1987 }
1988 async fn embed(&self, text: &str) -> Result<Vec<f32>, LlmError> {
1989 let v = if text.contains("cat")
1990 || text.contains("mat")
1991 || text.contains("feline")
1992 || text.contains("rug")
1993 {
1994 if text.contains("feline") || text.contains("rug") {
1995 vec![0.9f32, 0.1, 0.0]
1996 } else {
1997 vec![1.0f32, 0.0, 0.0]
1998 }
1999 } else {
2000 vec![0.0f32, 0.0, 1.0]
2001 };
2002 Ok(v)
2003 }
2004 fn supports_embeddings(&self) -> bool {
2005 true
2006 }
2007 fn name(&self) -> &'static str {
2008 "embed-mock"
2009 }
2010 }
2011
2012 let provider = EmbedMockProvider;
2013 let scorer = FidelityScorer;
2014 let cfg = FidelityConfig {
2015 enabled: true,
2016 semantic_scoring_provider: Some(ProviderName::new("embed-mock")),
2017 full_threshold: 0.0,
2021 compressed_threshold: 0.0,
2022 w_semantic: 1.0,
2023 w_temporal: 0.0,
2024 w_importance: 0.0,
2025 w_plan: 0.0,
2026 ..make_cfg()
2027 };
2028 let tc = FixedTc(4);
2029 let cat_msg = make_msg(Role::User, "The cat is on the mat");
2030 let feline_msg = make_msg(Role::User, "A feline rests on the rug");
2031 let stock_msg = make_msg(Role::User, "Stock prices fell today");
2032 let mut messages = vec![
2033 make_msg(Role::System, "system"),
2034 cat_msg,
2035 feline_msg,
2036 stock_msg,
2037 ];
2038
2039 scorer
2040 .score_and_apply(
2041 &mut messages,
2042 "cat mat",
2043 &[],
2044 &cfg,
2045 &tc,
2046 0,
2047 false,
2048 Some(&provider),
2049 None,
2050 )
2051 .await;
2052
2053 assert!(
2056 messages[1].metadata.embedding.is_some(),
2057 "cat message must have embedding"
2058 );
2059 assert!(
2060 messages[2].metadata.embedding.is_some(),
2061 "feline message must have embedding"
2062 );
2063 assert!(
2064 messages[3].metadata.embedding.is_some(),
2065 "stock message must have embedding"
2066 );
2067
2068 let query_emb = [1.0f32, 0.0, 0.0];
2070 let cat_emb = messages[1].metadata.embedding.as_ref().unwrap();
2071 let feline_emb = messages[2].metadata.embedding.as_ref().unwrap();
2072 let stock_emb = messages[3].metadata.embedding.as_ref().unwrap();
2073 assert!(
2074 cosine_similarity(cat_emb, &query_emb) > cosine_similarity(stock_emb, &query_emb),
2075 "cat message must be more similar to query than stock message"
2076 );
2077 assert!(
2078 cosine_similarity(feline_emb, &query_emb) > cosine_similarity(stock_emb, &query_emb),
2079 "feline message must be more similar to query than stock message"
2080 );
2081 }
2082
2083 #[tokio::test]
2085 async fn semantic_scoring_falls_back_to_keyword_when_provider_none() {
2086 let scorer = FidelityScorer;
2087 let cfg = FidelityConfig {
2088 enabled: true,
2089 semantic_scoring_provider: None,
2090 ..make_cfg()
2091 };
2092 let tc = FixedTc(4);
2093 let mut messages = vec![
2094 make_msg(Role::System, "system"),
2095 make_msg(Role::User, "cat mat keyword test"),
2096 make_msg(Role::User, "something unrelated here"),
2097 ];
2098
2099 scorer
2100 .score_and_apply(
2101 &mut messages,
2102 "query text here long",
2103 &[],
2104 &cfg,
2105 &tc,
2106 0,
2107 false,
2108 None,
2109 None,
2110 )
2111 .await;
2112
2113 for msg in &messages {
2115 assert!(
2116 msg.metadata.embedding.is_none(),
2117 "no embedding must be computed when provider is None"
2118 );
2119 }
2120 for msg in &messages[1..] {
2122 assert!(
2123 msg.metadata.fidelity_tag.is_some(),
2124 "all non-system messages must be scored via keyword path"
2125 );
2126 }
2127 }
2128
2129 #[tokio::test]
2131 async fn w_plan_produces_nonzero_score_for_matching_message() {
2132 use zeph_common::PlannedToolHint;
2133
2134 let scorer = FidelityScorer;
2135 let cfg = FidelityConfig {
2138 w_semantic: 0.0,
2139 w_temporal: 0.0,
2140 w_importance: 0.0,
2141 w_plan: 1.0,
2142 full_threshold: 0.5,
2143 compressed_threshold: 0.1,
2144 min_query_length: 100, ..make_cfg()
2146 };
2147 let tc = FixedTc(4);
2148
2149 let hint = PlannedToolHint::new("shell", vec!["cargo".to_string(), "build".to_string()], 1);
2151
2152 let mut messages = vec![
2155 make_msg(Role::System, "system prompt"),
2156 make_msg(Role::User, "run cargo build to compile"),
2157 make_msg(Role::User, "what is the weather today"),
2158 ];
2159
2160 scorer
2161 .score_and_apply(
2162 &mut messages,
2163 "q", &[hint],
2165 &cfg,
2166 &tc,
2167 0,
2168 false,
2169 None,
2170 None,
2171 )
2172 .await;
2173
2174 assert_eq!(
2176 messages[1].metadata.fidelity_tag,
2177 Some(ContextFidelity::Full),
2178 "message matching planned tool keywords must reach Full fidelity"
2179 );
2180
2181 assert_ne!(
2183 messages[2].metadata.fidelity_tag,
2184 Some(ContextFidelity::Full),
2185 "message with no keyword overlap must not reach Full fidelity via w_plan"
2186 );
2187 }
2188
2189 #[test]
2192 fn truncate_to_byte_limit_no_op_when_short() {
2193 assert_eq!(truncate_to_byte_limit("hello", 10), "hello");
2194 }
2195
2196 #[test]
2197 fn truncate_to_byte_limit_exact_limit_no_op() {
2198 assert_eq!(truncate_to_byte_limit("hello", 5), "hello");
2199 }
2200
2201 #[test]
2202 fn truncate_to_byte_limit_over_limit() {
2203 let s = truncate_to_byte_limit("abcdefgh", 5);
2204 assert_eq!(s.len(), 5);
2205 assert_eq!(s, "abcde");
2206 }
2207
2208 #[test]
2209 fn truncate_to_byte_limit_multibyte_boundary() {
2210 let s = truncate_to_byte_limit("日本語", 6);
2212 assert!(s.is_char_boundary(s.len()));
2213 assert_eq!(s, "日本");
2214 }
2215
2216 #[test]
2219 fn apply_input_cap_no_op_below_limit() {
2220 let mut s = "hello".to_string();
2221 apply_input_cap(&mut s, 10); assert_eq!(s, "hello");
2223 }
2224
2225 #[test]
2226 fn apply_input_cap_truncates_over_limit() {
2227 let mut s = "abcdefgh".to_string();
2229 apply_input_cap(&mut s, 1);
2230 assert_eq!(s, "abcd");
2231 }
2232
2233 #[test]
2234 fn apply_input_cap_multibyte() {
2235 let mut s = "日本語".to_string();
2238 apply_input_cap(&mut s, 1);
2239 assert!(s.is_char_boundary(s.len()));
2240 assert_eq!(s, "日");
2241 }
2242
2243 #[tokio::test]
2246 async fn embed_prepass_returns_embeddings_for_non_exempt() {
2247 let messages = vec![
2248 make_msg(Role::System, "system prompt"), make_msg(Role::User, "user message"),
2250 make_msg(Role::Assistant, "assistant reply"),
2251 ];
2252 let cfg = FidelityConfig::default();
2253 let embed = |_text: &str| -> EmbedFuture { Box::pin(async { Ok(vec![1.0f32, 2.0, 3.0]) }) };
2254 let result = embed_prepass(&messages, &embed, &cfg, 0).await;
2255 assert!(!result.contains_key(&0));
2257 assert_eq!(result[&1], vec![1.0, 2.0, 3.0]);
2258 assert_eq!(result[&2], vec![1.0, 2.0, 3.0]);
2259 }
2260
2261 #[tokio::test]
2262 async fn embed_prepass_skips_empty_content() {
2263 let messages = vec![make_msg(Role::System, "system"), make_msg(Role::User, "")];
2264 let cfg = FidelityConfig::default();
2265 let embed = |_text: &str| -> EmbedFuture { Box::pin(async { Ok(vec![1.0f32]) }) };
2266 let result = embed_prepass(&messages, &embed, &cfg, 0).await;
2267 assert!(!result.contains_key(&1), "empty content must be skipped");
2268 }
2269
2270 #[tokio::test]
2271 async fn embed_prepass_skips_inserted_memory() {
2272 let messages = vec![
2273 make_msg(Role::System, "system"),
2274 make_msg(Role::User, "injected memory"),
2275 make_msg(Role::User, "real user message"),
2276 ];
2277 let cfg = FidelityConfig::default();
2278 let embed = |_text: &str| -> EmbedFuture { Box::pin(async { Ok(vec![1.0f32]) }) };
2279 let result = embed_prepass(&messages, &embed, &cfg, 1).await;
2281 assert!(!result.contains_key(&0), "system is exempt");
2282 assert!(!result.contains_key(&1), "inserted memory is exempt");
2283 assert!(
2284 result.contains_key(&2),
2285 "real user message must be embedded"
2286 );
2287 }
2288
2289 #[tokio::test]
2290 async fn embed_prepass_silently_skips_errors() {
2291 let messages = vec![
2292 make_msg(Role::System, "system"),
2293 make_msg(Role::User, "user"),
2294 ];
2295 let cfg = FidelityConfig::default();
2296 let embed = |_text: &str| -> EmbedFuture {
2297 Box::pin(async {
2298 Err(zeph_llm::LlmError::EmbedUnsupported {
2299 provider: "mock".to_string(),
2300 })
2301 })
2302 };
2303 let result = embed_prepass(&messages, &embed, &cfg, 0).await;
2304 assert!(result.is_empty(), "errors must be silently skipped");
2305 }
2306
2307 #[tokio::test]
2308 async fn embed_prepass_truncates_content_when_cap_set() {
2309 let long_content = "a".repeat(100);
2311 let messages = vec![
2312 make_msg(Role::System, "system"),
2313 make_msg(Role::User, &long_content),
2314 ];
2315 let cfg = FidelityConfig {
2316 max_embed_input_tokens: Some(1), ..FidelityConfig::default()
2318 };
2319 let seen_len = std::sync::Arc::new(std::sync::Mutex::new(0usize));
2320 let seen_len_clone = seen_len.clone();
2321 let embed = move |text: &str| -> EmbedFuture {
2322 let len = text.len();
2323 let seen = seen_len_clone.clone();
2324 Box::pin(async move {
2325 *seen.lock().unwrap() = len;
2326 Ok(vec![1.0f32])
2327 })
2328 };
2329 embed_prepass(&messages, &embed, &cfg, 0).await;
2330 assert_eq!(
2331 *seen_len.lock().unwrap(),
2332 4,
2333 "content must be truncated to max_embed_input_tokens * 4 bytes"
2334 );
2335 }
2336
2337 #[tokio::test]
2338 async fn embed_prepass_concurrency_zero_clamped_to_one() {
2339 let messages = vec![
2341 make_msg(Role::System, "system"),
2342 make_msg(Role::User, "user message"),
2343 ];
2344 let cfg = FidelityConfig {
2345 embed_concurrency: 0,
2346 ..FidelityConfig::default()
2347 };
2348 let embed = |_text: &str| -> EmbedFuture { Box::pin(async { Ok(vec![1.0f32]) }) };
2349 let result = embed_prepass(&messages, &embed, &cfg, 0).await;
2350 assert!(
2351 result.contains_key(&1),
2352 "result must be produced even with concurrency=0"
2353 );
2354 }
2355
2356 #[tokio::test]
2357 async fn embed_prepass_skips_cached_embeddings() {
2358 let mut msg_with_cache = make_msg(Role::User, "already embedded");
2359 msg_with_cache.metadata.embedding = Some(vec![9.0f32]);
2360 let messages = vec![
2361 make_msg(Role::System, "system"),
2362 msg_with_cache,
2363 make_msg(Role::User, "needs embedding"),
2364 ];
2365 let cfg = FidelityConfig::default();
2366 let embed = |_text: &str| -> EmbedFuture { Box::pin(async { Ok(vec![1.0f32]) }) };
2367 let result = embed_prepass(&messages, &embed, &cfg, 0).await;
2368 assert!(
2369 !result.contains_key(&1),
2370 "message with cached embedding must be skipped"
2371 );
2372 assert!(
2373 result.contains_key(&2),
2374 "message without embedding must be processed"
2375 );
2376 }
2377
2378 #[tokio::test(start_paused = true)]
2379 async fn embed_prepass_timeout_skips_message() {
2380 let messages = vec![
2381 make_msg(Role::System, "system"),
2382 make_msg(Role::User, "user"),
2383 ];
2384 let cfg = FidelityConfig::default();
2385 let embed = |_text: &str| -> EmbedFuture {
2386 Box::pin(async {
2387 tokio::time::sleep(Duration::from_secs(45)).await;
2389 Ok(vec![1.0f32])
2390 })
2391 };
2392 let result = embed_prepass(&messages, &embed, &cfg, 0).await;
2394 assert!(result.is_empty(), "timed-out embed must be skipped");
2395 }
2396
2397 #[tokio::test(start_paused = true)]
2398 async fn embed_prepass_custom_timeout_respected() {
2399 let messages = vec![
2400 make_msg(Role::System, "system"),
2401 make_msg(Role::User, "user"),
2402 ];
2403 let cfg = FidelityConfig {
2405 embed_timeout_secs: 5,
2406 ..FidelityConfig::default()
2407 };
2408 let embed = |_text: &str| -> EmbedFuture {
2409 Box::pin(async {
2410 tokio::time::sleep(Duration::from_secs(10)).await;
2411 Ok(vec![1.0f32])
2412 })
2413 };
2414 let result = embed_prepass(&messages, &embed, &cfg, 0).await;
2415 assert!(
2416 result.is_empty(),
2417 "custom embed_timeout_secs must be honored"
2418 );
2419 }
2420
2421 #[test]
2424 fn fidelity_config_new_fields_defaults() {
2425 let cfg = FidelityConfig::default();
2426 assert_eq!(cfg.embed_concurrency, 32);
2427 assert!(cfg.max_embed_input_tokens.is_none());
2428 assert!(cfg.max_compress_input_tokens.is_none());
2429 }
2430
2431 #[test]
2432 fn fidelity_config_timeout_defaults() {
2433 let cfg = FidelityConfig::default();
2434 assert_eq!(cfg.embed_timeout_secs, 30);
2435 assert_eq!(cfg.compress_timeout_secs, 30);
2436 }
2437
2438 #[test]
2439 fn fidelity_config_new_fields_custom() {
2440 let cfg = FidelityConfig {
2441 embed_concurrency: 8,
2442 max_embed_input_tokens: Some(512),
2443 max_compress_input_tokens: Some(1024),
2444 ..FidelityConfig::default()
2445 };
2446 assert_eq!(cfg.embed_concurrency, 8);
2447 assert_eq!(cfg.max_embed_input_tokens, Some(512));
2448 assert_eq!(cfg.max_compress_input_tokens, Some(1024));
2449 }
2450
2451 #[tokio::test]
2453 async fn render_compressed_truncates_oversized_deferred_summary() {
2454 let scorer = FidelityScorer;
2455 let cfg = FidelityConfig {
2458 full_threshold: 2.0,
2459 compressed_threshold: 0.0,
2460 compressed_max_tokens: 3,
2461 ..make_cfg()
2462 };
2463 let tc = FixedTc(1);
2464 let mut msg = make_msg(Role::User, "original long content");
2465 msg.metadata.deferred_summary = Some("ten chars!".to_string());
2467 let mut messages = vec![make_msg(Role::System, "sys"), msg];
2468 scorer
2469 .score_and_apply(
2470 &mut messages,
2471 "query text here long",
2472 &[],
2473 &cfg,
2474 &tc,
2475 0,
2476 false,
2477 None,
2478 None,
2479 )
2480 .await;
2481 let compressed = &messages[1];
2482 assert_eq!(
2483 compressed.metadata.fidelity_tag,
2484 Some(ContextFidelity::Compressed)
2485 );
2486 assert!(
2487 tc.count_tokens(&compressed.content) <= 3,
2488 "deferred_summary result must be truncated to compressed_max_tokens"
2489 );
2490 }
2491
2492 #[tokio::test]
2494 async fn render_compressed_applies_max_compress_input_tokens() {
2495 let scorer = FidelityScorer;
2496 let cfg = FidelityConfig {
2499 full_threshold: 2.0,
2500 compressed_threshold: 0.0,
2501 compressed_max_tokens: 100,
2502 max_compress_input_tokens: Some(2), ..make_cfg()
2504 };
2505 let tc = FixedTc(1);
2506 let content_20 = "a".repeat(20); let mut msg = make_msg(Role::User, &content_20);
2508 let mut messages = vec![make_msg(Role::System, "sys"), msg.clone()];
2510 scorer
2511 .score_and_apply(
2512 &mut messages,
2513 "query text here long",
2514 &[],
2515 &cfg,
2516 &tc,
2517 0,
2518 false,
2519 None,
2520 None,
2521 )
2522 .await;
2523 let compressed = &messages[1];
2524 assert_eq!(
2525 compressed.metadata.fidelity_tag,
2526 Some(ContextFidelity::Compressed)
2527 );
2528 assert_eq!(
2530 compressed.content.len(),
2531 8,
2532 "content must be capped to max_compress_input_tokens * 4 bytes"
2533 );
2534 msg.metadata.deferred_summary = Some("short".to_string());
2536 let mut messages2 = vec![make_msg(Role::System, "sys"), msg];
2537 scorer
2538 .score_and_apply(
2539 &mut messages2,
2540 "query text here long",
2541 &[],
2542 &cfg,
2543 &tc,
2544 0,
2545 false,
2546 None,
2547 None,
2548 )
2549 .await;
2550 assert_eq!(
2551 messages2[1].content, "short",
2552 "deferred_summary must bypass input cap"
2553 );
2554 }
2555}