Skip to main content

harn_vm/orchestration/
assemble.rs

1//! Adaptive context assembly with cross-artifact microcompaction.
2//!
3//! `assemble_context` packs a set of artifacts into a token-budgeted slice
4//! of chunks, deduplicating overlap across artifacts, snipping oversized
5//! entries into chunked form, and returning an observability record that
6//! names why each chunk was included or dropped.
7//!
8//! The core is intentionally deterministic: given the same input artifacts
9//! and options, it produces the same chunk ids and ordering. A host-supplied
10//! ranker callback is the only non-deterministic hook; the VM-side binding
11//! invokes it via the same pattern as `compress_callback`.
12
13use std::collections::BTreeSet;
14
15use sha2::{Digest, Sha256};
16
17use super::ArtifactRecord;
18use crate::stdlib::xml::escape_xml_text;
19
20/// Strategy used to order chunks when packing into the budget.
21#[derive(Clone, Copy, Debug, Eq, PartialEq)]
22pub enum AssembleStrategy {
23    /// Sort by artifact `created_at` (newest first), then by chunk index.
24    Recency,
25    /// Sort by ranker score (highest first). Default ranker is token-overlap
26    /// against `query`; a host callback can supply a custom one.
27    Relevance,
28    /// Interleave chunks one-per-artifact in artifact input order, cycling
29    /// until the budget fills. Gives every artifact a chance to contribute.
30    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    /// Hash each chunk's normalized text; drop later duplicates.
58    Chunked,
59    /// Shingle-based overlap detection. Treat chunks whose trigram
60    /// Jaccard similarity exceeds 0.85 as duplicates. Still fully
61    /// deterministic — no embeddings or callback required.
62    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    /// Artifacts larger than this many tokens are split into chunks.
93    pub microcompact_threshold: usize,
94    /// Minimum overlap ratio (0.0-1.0) that counts as a semantic duplicate.
95    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/// One unit of packed context with a stable content-addressed id.
112#[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/// Per-artifact summary of what made it into the pack.
127#[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/// Reason an artifact or chunk was excluded from the final pack.
137#[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/// Per-chunk rationale — the "why" field the issue calls out.
146#[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
168/// Content-addressed chunk id — stable across runs for the same text.
169/// The leading `artifact_id` prefix keeps chunks from one artifact visually
170/// grouped in transcripts and replay diffs.
171pub 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
183/// Approximate token count using the same chars-per-token heuristic as
184/// `estimate_message_tokens`. One token ~= 4 characters.
185pub fn estimate_chunk_tokens(text: &str) -> usize {
186    text.len().div_ceil(4)
187}
188
189/// Split a text into chunks of roughly `target_tokens` each, snapping to
190/// paragraph (`\n\n`) boundaries when possible, then falling back to line
191/// breaks, then raw character boundaries. Chunks never exceed the target
192/// more than transiently — the final chunk may be short.
193pub 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            // Oversized paragraph: split by lines.
216            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                    // Still too big: fall back to char-boundary splits.
224                    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            // skip all consecutive newlines
275            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
295/// Trigram set for shingle-based semantic dedup. Lowercases and strips
296/// non-alphanumeric characters; produces trigrams over UTF-8 byte windows
297/// so behavior is stable across platforms.
298fn 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    // Length penalty: prefer chunks where matched query terms are dense,
355    // so a 100-char chunk that mentions "parser" scores higher than a
356    // 10k-char chunk that mentions "parser" once.
357    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
361/// Build every candidate chunk from the input artifacts. Artifacts that
362/// exceed `microcompact_threshold` tokens get split; smaller ones become
363/// a single chunk each. Skips artifacts with no text body.
364pub 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
419/// Apply dedup. Returns (kept, dropped-by-dedup). `dropped` is the slice
420/// needed for the caller's observability record; the reason is always
421/// `"duplicate"`.
422pub 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
478/// Score chunks under `strategy`. With `Relevance` and no callback scores,
479/// falls back to keyword overlap against `options.query`. The `custom_scores`
480/// Option is the hook where a host-supplied ranker slots in — the caller
481/// (stdlib binding) invokes the closure and passes the resulting Vec here.
482pub 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            // Newer artifacts first; within an artifact, earlier chunks first.
491            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                // Score is a synthetic "recency score" in [0, 1] based on
502                // lexicographic order of created_at with the input index as
503                // a stable tiebreaker. Newer created_at → higher score.
504                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                // Any trailing chunks without a custom score keep 0.0.
518            } 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            // Round-robin is handled at pack time; score just reflects
527            // input order so ties break deterministically.
528            for (idx, chunk) in chunks.iter_mut().enumerate() {
529                chunk.score = 1.0 - (idx as f64) * 1e-6;
530            }
531        }
532    }
533}
534
535/// Pack chunks under `budget_tokens`. Returns (selected, rejected-for-budget).
536pub 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            // Group by artifact_id preserving first-appearance order, then interleave.
544            let mut groups: Vec<Vec<AssembledChunk>> = Vec::new();
545            let mut group_index: std::collections::BTreeMap<String, usize> =
546                std::collections::BTreeMap::new();
547            // Preserve input order of artifacts by walking sorted-as-given.
548            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            // Within each group, keep by chunk_index ascending.
562            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
601/// Core assembly pass. The caller is responsible for supplying
602/// `custom_scores` (Some when a host ranker produced them, None otherwise).
603pub 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    // When a custom_scores slice was supplied, it's indexed over the
611    // *pre-dedup* candidate list — the caller saw those chunk ids. Build
612    // a score map keyed by chunk id so dedup doesn't misalign the slice.
613    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    // chunk_count per artifact from selected + rejected, for "of X chunks" observability.
638    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
709/// Render assembled chunks into the XML-ish `<artifact>` format used by
710/// `render_artifacts_context`, so swapping in `assemble_context` at the
711/// workflow stage layer produces the same prompt shape agents already
712/// expect. Appends a trailing `<context_budget>` summary so the agent
713/// (and replay diff) can see how much of the budget got used.
714pub 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        // Different text → different id.
778        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        // One of the two should dedup out.
816        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); // ~25_000 chars → ~6_250 tokens
826        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            // Tight budget: only one ~11-token chunk fits.
848            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        // Each paragraph is ~10 chars (~3 tokens); microcompact threshold of
863        // 3 tokens (=12 chars) forces one chunk per paragraph without
864        // fragmenting the final one mid-word.
865        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        // First four positions must alternate even if counts don't match
881        // exactly — interleaving is the invariant, not total chunk count.
882        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            // Tight budget: only one ~5-token chunk fits.
895            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        // Host-supplied ranker deliberately inverts the default order: it
906        // scores the second artifact higher even though query says "first".
907        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}