1use std::collections::BTreeSet;
14
15use sha2::{Digest, Sha256};
16
17use super::ArtifactRecord;
18use crate::stdlib::xml::escape_xml_text;
19
20#[derive(Clone, Copy, Debug, Eq, PartialEq)]
22pub enum AssembleStrategy {
23 Recency,
25 Relevance,
28 RoundRobin,
31}
32
33impl AssembleStrategy {
34 pub fn parse(value: &str) -> Result<Self, String> {
35 match value {
36 "recency" => Ok(Self::Recency),
37 "relevance" => Ok(Self::Relevance),
38 "round_robin" => Ok(Self::RoundRobin),
39 other => Err(format!(
40 "assemble_context: strategy must be one of recency | relevance | round_robin (got {other:?})"
41 )),
42 }
43 }
44
45 pub fn as_str(&self) -> &'static str {
46 match self {
47 Self::Recency => "recency",
48 Self::Relevance => "relevance",
49 Self::RoundRobin => "round_robin",
50 }
51 }
52}
53
54#[derive(Clone, Copy, Debug, Eq, PartialEq)]
55pub enum AssembleDedup {
56 None,
57 Chunked,
59 Semantic,
63}
64
65impl AssembleDedup {
66 pub fn parse(value: &str) -> Result<Self, String> {
67 match value {
68 "none" => Ok(Self::None),
69 "chunked" => Ok(Self::Chunked),
70 "semantic" => Ok(Self::Semantic),
71 other => Err(format!(
72 "assemble_context: dedup must be one of none | chunked | semantic (got {other:?})"
73 )),
74 }
75 }
76
77 pub fn as_str(&self) -> &'static str {
78 match self {
79 Self::None => "none",
80 Self::Chunked => "chunked",
81 Self::Semantic => "semantic",
82 }
83 }
84}
85
86#[derive(Clone, Debug)]
87pub struct AssembleOptions {
88 pub budget_tokens: usize,
89 pub dedup: AssembleDedup,
90 pub strategy: AssembleStrategy,
91 pub query: Option<String>,
92 pub microcompact_threshold: usize,
94 pub semantic_overlap: f64,
96}
97
98impl Default for AssembleOptions {
99 fn default() -> Self {
100 Self {
101 budget_tokens: 8_000,
102 dedup: AssembleDedup::Chunked,
103 strategy: AssembleStrategy::Relevance,
104 query: None,
105 microcompact_threshold: 2_000,
106 semantic_overlap: 0.85,
107 }
108 }
109}
110
111#[derive(Clone, Debug)]
113pub struct AssembledChunk {
114 pub id: String,
115 pub artifact_id: String,
116 pub artifact_kind: String,
117 pub title: Option<String>,
118 pub source: Option<String>,
119 pub text: String,
120 pub estimated_tokens: usize,
121 pub chunk_index: usize,
122 pub chunk_count: usize,
123 pub score: f64,
124}
125
126#[derive(Clone, Debug)]
128pub struct AssembledArtifactSummary {
129 pub artifact_id: String,
130 pub artifact_kind: String,
131 pub chunks_included: usize,
132 pub chunks_total: usize,
133 pub tokens_included: usize,
134}
135
136#[derive(Clone, Debug)]
138pub struct AssembledExclusion {
139 pub artifact_id: String,
140 pub chunk_id: Option<String>,
141 pub reason: &'static str,
142 pub detail: Option<String>,
143}
144
145#[derive(Clone, Debug)]
147pub struct AssembledReason {
148 pub chunk_id: String,
149 pub artifact_id: String,
150 pub strategy: &'static str,
151 pub score: f64,
152 pub included: bool,
153 pub reason: &'static str,
154}
155
156#[derive(Clone, Debug)]
157pub struct AssembledContext {
158 pub chunks: Vec<AssembledChunk>,
159 pub included: Vec<AssembledArtifactSummary>,
160 pub dropped: Vec<AssembledExclusion>,
161 pub reasons: Vec<AssembledReason>,
162 pub total_tokens: usize,
163 pub budget_tokens: usize,
164 pub strategy: AssembleStrategy,
165 pub dedup: AssembleDedup,
166}
167
168pub fn stable_chunk_id(artifact_id: &str, text: &str) -> String {
172 let mut hasher = Sha256::new();
173 hasher.update(text.as_bytes());
174 let digest = hasher.finalize();
175 let hex = digest
176 .iter()
177 .take(8)
178 .map(|byte| format!("{byte:02x}"))
179 .collect::<String>();
180 format!("{artifact_id}#{hex}")
181}
182
183pub fn estimate_chunk_tokens(text: &str) -> usize {
186 text.len().div_ceil(4)
187}
188
189pub fn chunk_text(text: &str, target_tokens: usize) -> Vec<String> {
194 if text.is_empty() {
195 return Vec::new();
196 }
197 let target_chars = (target_tokens.max(1)).saturating_mul(4);
198 if text.len() <= target_chars {
199 return vec![text.to_string()];
200 }
201
202 let mut chunks = Vec::new();
203 let mut current = String::new();
204 let push_current = |current: &mut String, chunks: &mut Vec<String>| {
205 if !current.is_empty() {
206 chunks.push(std::mem::take(current));
207 }
208 };
209
210 for paragraph in split_paragraphs(text) {
211 if current.len() + paragraph.len() + 2 > target_chars && !current.is_empty() {
212 push_current(&mut current, &mut chunks);
213 }
214 if paragraph.len() > target_chars {
215 push_current(&mut current, &mut chunks);
217 let mut inner = String::new();
218 for line in paragraph.split_inclusive('\n') {
219 if inner.len() + line.len() > target_chars && !inner.is_empty() {
220 chunks.push(std::mem::take(&mut inner));
221 }
222 if line.len() > target_chars {
223 let mut i = 0;
225 let bytes = line.as_bytes();
226 while i < line.len() {
227 let mut end = (i + target_chars).min(line.len());
228 while end < line.len() && (bytes[end] & 0b1100_0000) == 0b1000_0000 {
229 end += 1;
230 }
231 if !inner.is_empty() {
232 chunks.push(std::mem::take(&mut inner));
233 }
234 #[expect(
235 clippy::string_slice,
236 reason = "end skips continuation bytes, so i/end are char boundaries"
237 )]
238 chunks.push(line[i..end].to_string());
239 i = end;
240 }
241 } else {
242 inner.push_str(line);
243 }
244 }
245 if !inner.is_empty() {
246 chunks.push(inner);
247 }
248 } else {
249 if !current.is_empty() {
250 current.push_str("\n\n");
251 }
252 current.push_str(paragraph);
253 }
254 }
255 push_current(&mut current, &mut chunks);
256 chunks
257}
258
259#[expect(
260 clippy::string_slice,
261 reason = "start and i index ASCII newline bytes found by the byte scan"
262)]
263fn split_paragraphs(text: &str) -> Vec<&str> {
264 let mut out = Vec::new();
265 let mut start = 0;
266 let bytes = text.as_bytes();
267 let mut i = 0;
268 while i + 1 < bytes.len() {
269 if bytes[i] == b'\n' && bytes[i + 1] == b'\n' {
270 let segment = text[start..i].trim_matches('\n');
271 if !segment.is_empty() {
272 out.push(segment);
273 }
274 let mut j = i;
276 while j < bytes.len() && bytes[j] == b'\n' {
277 j += 1;
278 }
279 start = j;
280 i = j;
281 } else {
282 i += 1;
283 }
284 }
285 let tail = text[start..].trim_matches('\n');
286 if !tail.is_empty() {
287 out.push(tail);
288 }
289 if out.is_empty() && !text.is_empty() {
290 out.push(text);
291 }
292 out
293}
294
295fn trigrams(text: &str) -> BTreeSet<[u8; 3]> {
299 let normalized: Vec<u8> = text
300 .chars()
301 .filter_map(|c| {
302 if c.is_alphanumeric() {
303 Some(c.to_ascii_lowercase() as u8)
304 } else if c.is_whitespace() {
305 Some(b' ')
306 } else {
307 None
308 }
309 })
310 .collect();
311 let mut out = BTreeSet::new();
312 if normalized.len() < 3 {
313 return out;
314 }
315 for window in normalized.windows(3) {
316 out.insert([window[0], window[1], window[2]]);
317 }
318 out
319}
320
321fn jaccard(a: &BTreeSet<[u8; 3]>, b: &BTreeSet<[u8; 3]>) -> f64 {
322 if a.is_empty() && b.is_empty() {
323 return 1.0;
324 }
325 let intersection = a.intersection(b).count() as f64;
326 let union = a.union(b).count() as f64;
327 if union == 0.0 {
328 0.0
329 } else {
330 intersection / union
331 }
332}
333
334fn keyword_overlap_score(text: &str, query: &str) -> f64 {
335 if query.trim().is_empty() {
336 return 0.0;
337 }
338 let query_terms: BTreeSet<String> = query
339 .split_whitespace()
340 .filter(|term| term.len() > 2)
341 .map(|term| term.to_ascii_lowercase())
342 .collect();
343 if query_terms.is_empty() {
344 return 0.0;
345 }
346 let mut matches = 0usize;
347 let lower = text.to_ascii_lowercase();
348 for term in &query_terms {
349 if lower.contains(term.as_str()) {
350 matches += 1;
351 }
352 }
353 let base = matches as f64 / query_terms.len() as f64;
354 let density = (matches as f64) / (text.len() as f64 / 400.0 + 1.0);
358 base * 0.7 + density.min(1.0) * 0.3
359}
360
361pub fn build_candidate_chunks(
365 artifacts: &[ArtifactRecord],
366 options: &AssembleOptions,
367 dropped: &mut Vec<AssembledExclusion>,
368) -> Vec<AssembledChunk> {
369 let mut candidates = Vec::new();
370 for artifact in artifacts {
371 let Some(text) = artifact.text.as_ref() else {
372 dropped.push(AssembledExclusion {
373 artifact_id: artifact.id.clone(),
374 chunk_id: None,
375 reason: "no_text",
376 detail: None,
377 });
378 continue;
379 };
380 let trimmed = text.trim();
381 if trimmed.is_empty() {
382 dropped.push(AssembledExclusion {
383 artifact_id: artifact.id.clone(),
384 chunk_id: None,
385 reason: "empty_text",
386 detail: None,
387 });
388 continue;
389 }
390 let estimated = artifact
391 .estimated_tokens
392 .unwrap_or_else(|| estimate_chunk_tokens(text));
393 let pieces: Vec<String> = if estimated > options.microcompact_threshold {
394 chunk_text(text, options.microcompact_threshold)
395 } else {
396 vec![text.clone()]
397 };
398 let count = pieces.len();
399 for (idx, piece) in pieces.into_iter().enumerate() {
400 let id = stable_chunk_id(&artifact.id, &piece);
401 let tokens = estimate_chunk_tokens(&piece);
402 candidates.push(AssembledChunk {
403 id,
404 artifact_id: artifact.id.clone(),
405 artifact_kind: artifact.kind.clone(),
406 title: artifact.title.clone(),
407 source: artifact.source.clone(),
408 text: piece,
409 estimated_tokens: tokens,
410 chunk_index: idx,
411 chunk_count: count,
412 score: 0.0,
413 });
414 }
415 }
416 candidates
417}
418
419pub fn dedup_chunks(
423 mut chunks: Vec<AssembledChunk>,
424 mode: AssembleDedup,
425 semantic_overlap: f64,
426) -> (Vec<AssembledChunk>, Vec<AssembledExclusion>) {
427 let mut dropped = Vec::new();
428 match mode {
429 AssembleDedup::None => (chunks, dropped),
430 AssembleDedup::Chunked => {
431 let mut seen: BTreeSet<String> = BTreeSet::new();
432 chunks.retain(|chunk| {
433 let key = normalized_text_key(&chunk.text);
434 if seen.insert(key) {
435 true
436 } else {
437 dropped.push(AssembledExclusion {
438 artifact_id: chunk.artifact_id.clone(),
439 chunk_id: Some(chunk.id.clone()),
440 reason: "duplicate",
441 detail: Some("chunked".to_string()),
442 });
443 false
444 }
445 });
446 (chunks, dropped)
447 }
448 AssembleDedup::Semantic => {
449 let mut kept: Vec<(AssembledChunk, BTreeSet<[u8; 3]>)> = Vec::new();
450 for chunk in chunks.drain(..) {
451 let trigrams_new = trigrams(&chunk.text);
452 let mut duplicate = false;
453 for (existing, existing_trigrams) in &kept {
454 if jaccard(&trigrams_new, existing_trigrams) >= semantic_overlap {
455 dropped.push(AssembledExclusion {
456 artifact_id: chunk.artifact_id.clone(),
457 chunk_id: Some(chunk.id.clone()),
458 reason: "duplicate",
459 detail: Some(format!("semantic≈{}", existing.id)),
460 });
461 duplicate = true;
462 break;
463 }
464 }
465 if !duplicate {
466 kept.push((chunk, trigrams_new));
467 }
468 }
469 (kept.into_iter().map(|(chunk, _)| chunk).collect(), dropped)
470 }
471 }
472}
473
474fn normalized_text_key(text: &str) -> String {
475 text.split_whitespace().collect::<Vec<_>>().join(" ")
476}
477
478pub fn score_chunks(
483 chunks: &mut [AssembledChunk],
484 artifacts: &[ArtifactRecord],
485 options: &AssembleOptions,
486 custom_scores: Option<&[f64]>,
487) {
488 match options.strategy {
489 AssembleStrategy::Recency => {
490 let order: std::collections::BTreeMap<&str, (String, usize)> = artifacts
492 .iter()
493 .enumerate()
494 .map(|(idx, artifact)| (artifact.id.as_str(), (artifact.created_at.clone(), idx)))
495 .collect();
496 for chunk in chunks.iter_mut() {
497 let (created_at, input_idx) = order
498 .get(chunk.artifact_id.as_str())
499 .cloned()
500 .unwrap_or_else(|| (String::new(), 0));
501 let recency_rank = created_at
505 .chars()
506 .fold(0u64, |acc, c| acc.wrapping_mul(131).wrapping_add(c as u64));
507 chunk.score = recency_rank as f64 / u64::MAX as f64
508 - (input_idx as f64) * 1e-9
509 - (chunk.chunk_index as f64) * 1e-12;
510 }
511 }
512 AssembleStrategy::Relevance => {
513 if let Some(scores) = custom_scores {
514 for (chunk, score) in chunks.iter_mut().zip(scores.iter()) {
515 chunk.score = *score;
516 }
517 } else {
519 let query = options.query.as_deref().unwrap_or("");
520 for chunk in chunks.iter_mut() {
521 chunk.score = keyword_overlap_score(&chunk.text, query);
522 }
523 }
524 }
525 AssembleStrategy::RoundRobin => {
526 for (idx, chunk) in chunks.iter_mut().enumerate() {
529 chunk.score = 1.0 - (idx as f64) * 1e-6;
530 }
531 }
532 }
533}
534
535pub fn pack_budget(
537 chunks: Vec<AssembledChunk>,
538 options: &AssembleOptions,
539) -> (Vec<AssembledChunk>, Vec<AssembledChunk>) {
540 let mut sorted = chunks;
541 match options.strategy {
542 AssembleStrategy::RoundRobin => {
543 let mut groups: Vec<Vec<AssembledChunk>> = Vec::new();
545 let mut group_index: std::collections::BTreeMap<String, usize> =
546 std::collections::BTreeMap::new();
547 for chunk in sorted.drain(..) {
549 let key = chunk.artifact_id.clone();
550 let idx = match group_index.get(&key) {
551 Some(idx) => *idx,
552 None => {
553 let idx = groups.len();
554 group_index.insert(key.clone(), idx);
555 groups.push(Vec::new());
556 idx
557 }
558 };
559 groups[idx].push(chunk);
560 }
561 for group in &mut groups {
563 group.sort_by_key(|chunk| chunk.chunk_index);
564 }
565 let mut interleaved = Vec::new();
566 let max_len = groups.iter().map(Vec::len).max().unwrap_or(0);
567 for i in 0..max_len {
568 for group in &mut groups {
569 if i < group.len() {
570 interleaved.push(group[i].clone());
571 }
572 }
573 }
574 sorted = interleaved;
575 }
576 _ => {
577 sorted.sort_by(|a, b| {
578 b.score
579 .partial_cmp(&a.score)
580 .unwrap_or(std::cmp::Ordering::Equal)
581 .then_with(|| a.artifact_id.cmp(&b.artifact_id))
582 .then_with(|| a.chunk_index.cmp(&b.chunk_index))
583 });
584 }
585 }
586
587 let mut selected = Vec::new();
588 let mut rejected = Vec::new();
589 let mut used = 0usize;
590 for chunk in sorted {
591 if used + chunk.estimated_tokens > options.budget_tokens {
592 rejected.push(chunk);
593 continue;
594 }
595 used += chunk.estimated_tokens;
596 selected.push(chunk);
597 }
598 (selected, rejected)
599}
600
601pub fn assemble_context(
604 artifacts: &[ArtifactRecord],
605 options: &AssembleOptions,
606 custom_scores: Option<&[f64]>,
607) -> AssembledContext {
608 let mut dropped = Vec::new();
609 let candidates = build_candidate_chunks(artifacts, options, &mut dropped);
610 let custom_map: Option<std::collections::BTreeMap<String, f64>> = custom_scores.map(|scores| {
614 candidates
615 .iter()
616 .zip(scores.iter().copied())
617 .map(|(chunk, score)| (chunk.id.clone(), score))
618 .collect()
619 });
620 let (mut deduped, dedup_dropped) =
621 dedup_chunks(candidates, options.dedup, options.semantic_overlap);
622 dropped.extend(dedup_dropped);
623
624 if let Some(map) = custom_map.as_ref() {
625 for chunk in deduped.iter_mut() {
626 chunk.score = map.get(&chunk.id).copied().unwrap_or(0.0);
627 }
628 } else {
629 score_chunks(&mut deduped, artifacts, options, None);
630 }
631
632 let (selected, rejected) = pack_budget(deduped, options);
633
634 let mut reasons = Vec::new();
635 let mut included_tokens: std::collections::BTreeMap<String, (String, usize, usize, usize)> =
636 std::collections::BTreeMap::new();
637 let mut total_counts: std::collections::BTreeMap<String, usize> =
639 std::collections::BTreeMap::new();
640 for chunk in selected.iter().chain(rejected.iter()) {
641 *total_counts.entry(chunk.artifact_id.clone()).or_insert(0) += 1;
642 }
643
644 for chunk in &selected {
645 reasons.push(AssembledReason {
646 chunk_id: chunk.id.clone(),
647 artifact_id: chunk.artifact_id.clone(),
648 strategy: options.strategy.as_str(),
649 score: chunk.score,
650 included: true,
651 reason: "selected",
652 });
653 let entry = included_tokens
654 .entry(chunk.artifact_id.clone())
655 .or_insert_with(|| {
656 (
657 chunk.artifact_kind.clone(),
658 0,
659 *total_counts.get(&chunk.artifact_id).unwrap_or(&0),
660 0,
661 )
662 });
663 entry.1 += 1;
664 entry.3 += chunk.estimated_tokens;
665 }
666 for chunk in &rejected {
667 reasons.push(AssembledReason {
668 chunk_id: chunk.id.clone(),
669 artifact_id: chunk.artifact_id.clone(),
670 strategy: options.strategy.as_str(),
671 score: chunk.score,
672 included: false,
673 reason: "budget_exceeded",
674 });
675 dropped.push(AssembledExclusion {
676 artifact_id: chunk.artifact_id.clone(),
677 chunk_id: Some(chunk.id.clone()),
678 reason: "budget_exceeded",
679 detail: None,
680 });
681 }
682
683 let total_tokens = selected.iter().map(|chunk| chunk.estimated_tokens).sum();
684 let included: Vec<AssembledArtifactSummary> = included_tokens
685 .into_iter()
686 .map(
687 |(artifact_id, (kind, included, total, tokens))| AssembledArtifactSummary {
688 artifact_id,
689 artifact_kind: kind,
690 chunks_included: included,
691 chunks_total: total,
692 tokens_included: tokens,
693 },
694 )
695 .collect();
696
697 AssembledContext {
698 chunks: selected,
699 included,
700 dropped,
701 reasons,
702 total_tokens,
703 budget_tokens: options.budget_tokens,
704 strategy: options.strategy,
705 dedup: options.dedup,
706 }
707}
708
709pub fn render_assembled_chunks(assembled: &AssembledContext) -> String {
715 let mut parts = Vec::with_capacity(assembled.chunks.len() + 1);
716 for chunk in &assembled.chunks {
717 let title = chunk
718 .title
719 .clone()
720 .unwrap_or_else(|| format!("{} {}", chunk.artifact_kind, chunk.artifact_id));
721 parts.push(format!(
722 "<artifact>\n<title>{}</title>\n<kind>{}</kind>\n<source>{}</source>\n\
723<chunk_id>{}</chunk_id>\n<chunk_index>{} of {}</chunk_index>\n<body>\n{}\n</body>\n</artifact>",
724 escape_xml_text(&title),
725 escape_xml_text(&chunk.artifact_kind),
726 escape_xml_text(chunk.source.as_deref().unwrap_or("unknown")),
727 escape_xml_text(&chunk.id),
728 chunk.chunk_index + 1,
729 chunk.chunk_count,
730 chunk.text,
731 ));
732 }
733 parts.push(format!(
734 "<context_budget>\n<used_tokens>{}</used_tokens>\n<budget_tokens>{}</budget_tokens>\n<strategy>{}</strategy>\n<dedup>{}</dedup>\n</context_budget>",
735 assembled.total_tokens,
736 assembled.budget_tokens,
737 assembled.strategy.as_str(),
738 assembled.dedup.as_str(),
739 ));
740 parts.join("\n\n")
741}
742
743#[cfg(test)]
744mod tests {
745 use super::*;
746
747 fn artifact(id: &str, text: &str) -> ArtifactRecord {
748 ArtifactRecord {
749 type_name: "artifact".to_string(),
750 id: id.to_string(),
751 kind: "resource".to_string(),
752 title: Some(id.to_string()),
753 text: Some(text.to_string()),
754 data: None,
755 source: None,
756 created_at: format!("2026-04-{id:0>2}T00:00:00Z"),
757 freshness: None,
758 priority: Some(50),
759 lineage: Vec::new(),
760 relevance: None,
761 estimated_tokens: None,
762 stage: None,
763 metadata: Default::default(),
764 }
765 .normalize()
766 }
767
768 #[test]
769 fn chunk_ids_are_stable_and_content_addressed() {
770 let a = artifact("01", "alpha bravo charlie");
771 let options = AssembleOptions::default();
772 let mut dropped = Vec::new();
773 let first = build_candidate_chunks(&[a.clone()], &options, &mut dropped);
774 let second = build_candidate_chunks(&[a], &options, &mut dropped);
775 assert_eq!(first[0].id, second[0].id);
776 assert!(first[0].id.starts_with("01#"));
777 let different = artifact("01", "delta echo foxtrot");
779 let different_chunks = build_candidate_chunks(&[different], &options, &mut dropped);
780 assert_ne!(first[0].id, different_chunks[0].id);
781 }
782
783 #[test]
784 fn chunked_dedup_drops_exact_duplicates() {
785 let a = artifact("01", "shared body");
786 let b = artifact("02", "shared body");
787 let options = AssembleOptions {
788 budget_tokens: 10_000,
789 dedup: AssembleDedup::Chunked,
790 strategy: AssembleStrategy::Recency,
791 ..AssembleOptions::default()
792 };
793 let result = assemble_context(&[a, b], &options, None);
794 assert_eq!(result.chunks.len(), 1);
795 assert!(result.dropped.iter().any(|d| d.reason == "duplicate"));
796 }
797
798 #[test]
799 fn semantic_dedup_catches_near_duplicates() {
800 let a = artifact(
801 "01",
802 "The parser drift issue was diagnosed by tracing token spans.",
803 );
804 let b = artifact(
805 "02",
806 "The parser drift issue, diagnosed by tracing token spans, appeared in the tokenizer.",
807 );
808 let options = AssembleOptions {
809 dedup: AssembleDedup::Semantic,
810 strategy: AssembleStrategy::Recency,
811 semantic_overlap: 0.5,
812 ..AssembleOptions::default()
813 };
814 let result = assemble_context(&[a, b], &options, None);
815 assert_eq!(result.chunks.len(), 1);
817 assert!(result.dropped.iter().any(|d| d.reason == "duplicate"
818 && d.detail
819 .as_deref()
820 .is_some_and(|s| s.starts_with("semantic"))));
821 }
822
823 #[test]
824 fn budget_enforcement_trims_excess_chunks() {
825 let text = "word ".repeat(5_000); let a = artifact("01", &text);
827 let options = AssembleOptions {
828 budget_tokens: 500,
829 dedup: AssembleDedup::None,
830 strategy: AssembleStrategy::Recency,
831 microcompact_threshold: 200,
832 ..AssembleOptions::default()
833 };
834 let result = assemble_context(&[a], &options, None);
835 assert!(result.total_tokens <= options.budget_tokens);
836 assert!(result
837 .reasons
838 .iter()
839 .any(|r| !r.included && r.reason == "budget_exceeded"));
840 }
841
842 #[test]
843 fn relevance_strategy_prefers_query_matches() {
844 let a = artifact("01", "completely unrelated content about weather");
845 let b = artifact("02", "parser drift diagnostics token spans hotspot");
846 let options = AssembleOptions {
847 budget_tokens: 12,
849 dedup: AssembleDedup::None,
850 strategy: AssembleStrategy::Relevance,
851 query: Some("parser drift diagnostics".to_string()),
852 microcompact_threshold: 10_000,
853 ..AssembleOptions::default()
854 };
855 let result = assemble_context(&[a, b], &options, None);
856 assert_eq!(result.chunks.len(), 1);
857 assert_eq!(result.chunks[0].artifact_id, "02");
858 }
859
860 #[test]
861 fn round_robin_interleaves_artifacts() {
862 let a = artifact("01", "alpha aaaa\n\nbeta bbbb\n\ngamma ccc");
866 let b = artifact("02", "delta dddd\n\nepsilon ee\n\nzeta ff");
867 let options = AssembleOptions {
868 budget_tokens: 10_000,
869 dedup: AssembleDedup::None,
870 strategy: AssembleStrategy::RoundRobin,
871 microcompact_threshold: 3,
872 ..AssembleOptions::default()
873 };
874 let result = assemble_context(&[a, b], &options, None);
875 let order: Vec<&str> = result
876 .chunks
877 .iter()
878 .map(|c| c.artifact_id.as_str())
879 .collect();
880 assert!(order.len() >= 4);
883 assert_eq!(order[0], "01");
884 assert_eq!(order[1], "02");
885 assert_eq!(order[2], "01");
886 assert_eq!(order[3], "02");
887 }
888
889 #[test]
890 fn custom_scores_override_default_ranker() {
891 let a = artifact("01", "first body content");
892 let b = artifact("02", "second body content");
893 let options = AssembleOptions {
894 budget_tokens: 6,
896 dedup: AssembleDedup::None,
897 strategy: AssembleStrategy::Relevance,
898 query: Some("first".to_string()),
899 microcompact_threshold: 10_000,
900 ..AssembleOptions::default()
901 };
902 let mut dropped = Vec::new();
903 let candidates = build_candidate_chunks(&[a.clone(), b.clone()], &options, &mut dropped);
904 assert_eq!(candidates.len(), 2);
905 let scores = vec![0.1, 0.9];
908 let result = assemble_context(&[a, b], &options, Some(&scores));
909 assert_eq!(result.chunks.len(), 1);
910 assert_eq!(result.chunks[0].artifact_id, "02");
911 }
912
913 #[test]
914 fn reasons_name_strategy_and_inclusion() {
915 let a = artifact("01", "included body");
916 let b = artifact("02", "dropped body because budget");
917 let options = AssembleOptions {
918 budget_tokens: 5,
919 dedup: AssembleDedup::None,
920 strategy: AssembleStrategy::Recency,
921 microcompact_threshold: 10_000,
922 ..AssembleOptions::default()
923 };
924 let result = assemble_context(&[a, b], &options, None);
925 assert!(result.reasons.iter().any(|r| r.included));
926 assert!(result.reasons.iter().any(|r| !r.included));
927 for reason in &result.reasons {
928 assert_eq!(reason.strategy, "recency");
929 }
930 }
931
932 #[test]
933 fn empty_artifact_reports_dropped() {
934 let mut empty = artifact("01", "");
935 empty.text = Some(String::new());
936 let options = AssembleOptions::default();
937 let result = assemble_context(&[empty], &options, None);
938 assert!(result.chunks.is_empty());
939 assert!(result
940 .dropped
941 .iter()
942 .any(|d| d.reason == "empty_text" || d.reason == "no_text"));
943 }
944}