Skip to main content

rustbrain_core/
context.rs

1//! Graph-aware AI context assembly with token budgeting.
2//!
3//! Pipeline used by [`crate::Brain::context_for_prompt`]:
4//!
5//! 1. Ranked FTS seeds ([`crate::query`]) with stopword-aware query rewrite
6//! 2. Optional CSR k-hop expansion from `.brain/graph.mmap` (doc seeds preferred)
7//! 3. Score fusion (seed score × edge weight × hop decay) + symbol quality filters
8//! 4. Pack nodes with body excerpts until the approximate character/token budget is exhausted
9//!
10//! Token accounting is intentionally simple (`chars / 4`). It is good enough for
11//! agent prompt packing, not for billing-grade tokenizer parity.
12
13use crate::error::Result;
14use crate::fts::{is_generic_topic, prepare_search_query, tokenize_query};
15use crate::query::{QueryOptions, RankedHit};
16use crate::storage::Database;
17use crate::types::{ContextBundle, ContextNode, ContextRole, Node, NodeType};
18use serde::{Deserialize, Serialize};
19use std::collections::{HashMap, HashSet};
20use std::path::Path;
21use std::time::Instant;
22
23/// Max characters of body excerpt packed per note/seed.
24const MAX_EXCERPT_CHARS: usize = 900;
25/// Max characters for symbol neighbor excerpts (usually doc comments).
26const MAX_SYMBOL_EXCERPT_CHARS: usize = 280;
27/// Soft cap on packed symbol neighbors even when hop_to_symbols is on.
28const MAX_PACKED_SYMBOL_NEIGHBORS: usize = 8;
29
30/// Options controlling [`assemble_context`] and [`crate::Brain::context_for_prompt_with`].
31#[derive(Debug, Clone)]
32pub struct ContextOptions {
33    /// Approximate max tokens for the packed bundle (`chars ≈ tokens × 4`).
34    pub max_tokens: usize,
35    /// Graph expansion depth (`0` = seeds only; `1` is the usual default).
36    pub hop_depth: usize,
37    /// Max seed nodes taken from the ranked query before expansion.
38    pub max_seeds: usize,
39    /// Max total nodes packed into the bundle (seeds + neighbors).
40    pub max_nodes: usize,
41    /// Multiplicative decay applied per hop for neighbor scores.
42    pub hop_decay: f32,
43    /// Exclude pure code symbols from seeds (and from neighbors unless [`Self::hop_to_symbols`]).
44    pub no_symbols: bool,
45    /// Only include these node types when non-empty.
46    pub include_types: Vec<crate::types::NodeType>,
47    /// Always exclude these node types.
48    pub exclude_types: Vec<crate::types::NodeType>,
49    /// When true, allow graph hops *to* symbols even if `no_symbols` filters seeds.
50    /// Default true so ADR → `symbol:foo` remains useful for agents.
51    pub hop_to_symbols: bool,
52    /// Prefer expanding the graph from non-symbol seeds (docs/README). Default true.
53    pub hop_from_docs_only: bool,
54    /// Include body excerpts from FTS content (or summary fallback). Default true.
55    pub include_excerpts: bool,
56}
57
58impl Default for ContextOptions {
59    fn default() -> Self {
60        Self {
61            max_tokens: 2048,
62            hop_depth: 1,
63            max_seeds: 8,
64            max_nodes: 24,
65            hop_decay: 0.65,
66            // Agent-friendly: notes first; symbols arrive via hops when useful.
67            no_symbols: true,
68            include_types: Vec::new(),
69            exclude_types: Vec::new(),
70            hop_to_symbols: true,
71            hop_from_docs_only: true,
72            include_excerpts: true,
73        }
74    }
75}
76
77impl ContextOptions {
78    /// Agent-oriented defaults (same as [`Default`] as of 0.3.0).
79    pub fn agent() -> Self {
80        Self::default()
81    }
82
83    /// Include symbols as seeds and neighbors (power-user / debugging).
84    pub fn with_symbols(mut self) -> Self {
85        self.no_symbols = false;
86        self
87    }
88
89    fn seed_query_opts(&self) -> QueryOptions {
90        QueryOptions {
91            limit: self.max_seeds.saturating_mul(2).max(8),
92            no_symbols: self.no_symbols,
93            include_types: self.include_types.clone(),
94            exclude_types: self.exclude_types.clone(),
95            ..QueryOptions::default()
96        }
97    }
98
99    fn allows_pack(&self, ty: &crate::types::NodeType, role: ContextRole) -> bool {
100        use crate::types::NodeType;
101        if !self.include_types.is_empty() && !self.include_types.contains(ty) {
102            // Allow symbol neighbors when hop_to_symbols and role is Neighbor
103            if !(self.hop_to_symbols && role == ContextRole::Neighbor && *ty == NodeType::Symbol) {
104                return false;
105            }
106        }
107        if self.exclude_types.contains(ty) {
108            return false;
109        }
110        if *ty == NodeType::Symbol {
111            if role == ContextRole::Seed && self.no_symbols {
112                return false;
113            }
114            if role == ContextRole::Neighbor && self.no_symbols && !self.hop_to_symbols {
115                return false;
116            }
117        }
118        true
119    }
120}
121
122/// Intermediate scored candidate before token packing.
123#[derive(Debug, Clone)]
124struct Candidate {
125    node: Node,
126    score: f32,
127    role: ContextRole,
128    hop: u8,
129    excerpt: Option<String>,
130}
131
132/// Assemble a graph-aware context bundle for an agent prompt.
133///
134/// `brain_dir` is the `.brain` directory (for loading `graph.mmap`). When the
135/// mmap file is missing or the `mmap` feature is disabled, expansion is skipped
136/// and only ranked seeds are packed.
137///
138/// # Errors
139///
140/// Propagates FTS / database errors from the seed query.
141pub fn assemble_context(
142    db: &Database,
143    brain_dir: &Path,
144    prompt: &str,
145    opts: &ContextOptions,
146) -> Result<ContextBundle> {
147    let start = Instant::now();
148    let char_budget = opts.max_tokens.saturating_mul(4).max(256);
149    let query_tokens = prepare_search_query(prompt)
150        .map(|p| p.tokens)
151        .unwrap_or_else(|_| tokenize_query(prompt));
152
153    let qopts = opts.seed_query_opts();
154    let mut seeds = db.search_ranked(prompt, &qopts)?;
155    let generic = is_generic_topic(&query_tokens);
156    // Soft / empty retrieval → inject README hub + harvested goals (cold-start agents).
157    if seeds.is_empty() || generic {
158        inject_hub_seeds(db, &mut seeds, opts)?;
159    }
160
161    let mut candidates: Vec<Candidate> = Vec::new();
162    let mut seed_ids: HashSet<String> = HashSet::new();
163
164    for (i, hit) in seeds.into_iter().take(opts.max_seeds).enumerate() {
165        if !opts.allows_pack(&hit.node.node_type, ContextRole::Seed) {
166            continue;
167        }
168        // Skip pure ADR templates — they burn budget without knowledge.
169        if is_template_stub(&hit.node) {
170            continue;
171        }
172        seed_ids.insert(hit.node.id.clone());
173        // Slight rank-position prior
174        let mut score = hit.score * (1.0 / (1.0 + i as f32 * 0.05));
175        // Prefer hand-written / ADR over generated harvest when scores are close.
176        if is_generated_path(&hit.node) {
177            score *= 0.85;
178        }
179        if hit.node.node_type == NodeType::Adr {
180            score *= 1.08;
181        }
182        let excerpt = if opts.include_excerpts {
183            load_excerpt(db, &hit.node, MAX_EXCERPT_CHARS)
184        } else {
185            None
186        };
187        candidates.push(Candidate {
188            node: hit.node,
189            score,
190            role: ContextRole::Seed,
191            hop: 0,
192            excerpt,
193        });
194    }
195
196    // If filters wiped everything (e.g. only template matched), hub-inject again.
197    if candidates.is_empty() {
198        let mut hub_hits = Vec::new();
199        inject_hub_seeds(db, &mut hub_hits, opts)?;
200        for hit in hub_hits.into_iter().take(4) {
201            if !opts.allows_pack(&hit.node.node_type, ContextRole::Seed) {
202                continue;
203            }
204            if seed_ids.contains(&hit.node.id) {
205                continue;
206            }
207            seed_ids.insert(hit.node.id.clone());
208            let excerpt = if opts.include_excerpts {
209                load_excerpt(db, &hit.node, MAX_EXCERPT_CHARS)
210            } else {
211                None
212            };
213            candidates.push(Candidate {
214                node: hit.node,
215                score: hit.score,
216                role: ContextRole::Seed,
217                hop: 0,
218                excerpt,
219            });
220        }
221    }
222
223    let mut graph_nodes = 0usize;
224    let mut graph_edges = 0usize;
225    let mut neighbor_ids: Vec<String> = Vec::new();
226
227    #[cfg(feature = "mmap")]
228    if opts.hop_depth > 0 {
229        let mmap_path = brain_dir.join("graph.mmap");
230        if mmap_path.exists() {
231            if let Ok(graph) = crate::mmap::CsrMmapGraph::open(&mmap_path) {
232                graph_nodes = graph.node_count;
233                graph_edges = graph.edge_count;
234
235                let mut best_neighbor: HashMap<String, (f32, u8)> = HashMap::new();
236
237                for cand in candidates.iter().filter(|c| c.role == ContextRole::Seed) {
238                    if opts.hop_from_docs_only && cand.node.node_type == NodeType::Symbol {
239                        continue;
240                    }
241                    if let Some(idx) = graph.index_of(&cand.node.id) {
242                        for (nidx, edge_w) in
243                            graph.k_hop_neighborhood(idx as usize, opts.hop_depth)
244                        {
245                            if let Some(id) = graph.node_id(nidx as usize) {
246                                if seed_ids.contains(id) {
247                                    continue;
248                                }
249                                let hop: u8 = if opts.hop_depth <= 1 || edge_w >= 0.85 {
250                                    1
251                                } else {
252                                    2
253                                };
254                                let nscore = cand.score * edge_w * opts.hop_decay.powi(hop as i32);
255                                best_neighbor
256                                    .entry(id.to_string())
257                                    .and_modify(|(s, h)| {
258                                        if nscore > *s {
259                                            *s = nscore;
260                                            *h = hop;
261                                        }
262                                    })
263                                    .or_insert((nscore, hop));
264                            }
265                        }
266                    }
267                }
268
269                let mut neigh_sorted: Vec<_> = best_neighbor.into_iter().collect();
270                neigh_sorted.sort_by(|a, b| {
271                    b.1 .0
272                        .partial_cmp(&a.1 .0)
273                        .unwrap_or(std::cmp::Ordering::Equal)
274                });
275
276                for (id, (score, hop)) in neigh_sorted {
277                    neighbor_ids.push(id.clone());
278                    if let Some(node) = db.get_node(&id)? {
279                        if !opts.allows_pack(&node.node_type, ContextRole::Neighbor) {
280                            continue;
281                        }
282                        let mut nscore = score;
283                        if node.node_type == NodeType::Symbol {
284                            match symbol_neighbor_quality(&node, &query_tokens) {
285                                SymbolQuality::Drop => continue,
286                                SymbolQuality::Keep { boost } => nscore *= boost,
287                            }
288                        }
289                        let max_ex = if node.node_type == NodeType::Symbol {
290                            MAX_SYMBOL_EXCERPT_CHARS
291                        } else {
292                            MAX_EXCERPT_CHARS
293                        };
294                        let excerpt = if opts.include_excerpts {
295                            load_excerpt(db, &node, max_ex)
296                        } else {
297                            None
298                        };
299                        candidates.push(Candidate {
300                            node,
301                            score: nscore,
302                            role: ContextRole::Neighbor,
303                            hop,
304                            excerpt,
305                        });
306                    }
307                }
308            }
309        }
310    }
311
312    // Sort for agent-useful packs: seeds before neighbors, decisions before symbols,
313    // then raw score. Keeps ADR + README above opportunistic symbol hops.
314    candidates.sort_by(|a, b| {
315        pack_rank(b)
316            .partial_cmp(&pack_rank(a))
317            .unwrap_or(std::cmp::Ordering::Equal)
318    });
319
320    let mut packed: Vec<ContextNode> = Vec::new();
321    let mut used_chars = estimate_overhead(prompt);
322    let mut seen = HashSet::new();
323    let mut symbol_neighbors = 0usize;
324    let mut packed_excerpts: Vec<String> = Vec::new();
325
326    for cand in candidates {
327        if packed.len() >= opts.max_nodes {
328            break;
329        }
330        if !seen.insert(cand.node.id.clone()) {
331            continue;
332        }
333        if cand.role == ContextRole::Neighbor
334            && cand.node.node_type == NodeType::Symbol
335            && symbol_neighbors >= MAX_PACKED_SYMBOL_NEIGHBORS
336        {
337            continue;
338        }
339        // Prefer a single README-family hub (root readme XOR from-readme harvest).
340        if is_readme_family(&cand.node.id)
341            && packed.iter().any(|n| is_readme_family(&n.id))
342        {
343            continue;
344        }
345        // Drop near-duplicate body text (harvested clones of the same prose).
346        if let Some(ex) = &cand.excerpt {
347            if packed_excerpts
348                .iter()
349                .any(|prev| excerpt_jaccard(prev, ex) > 0.55)
350            {
351                continue;
352            }
353        }
354        let ctx_node = to_context_node(&cand);
355        let cost = estimate_node_chars(&ctx_node);
356        if used_chars + cost > char_budget && !packed.is_empty() {
357            // Always try to keep at least one seed
358            if cand.role == ContextRole::Neighbor {
359                continue;
360            }
361            if packed.iter().any(|n| n.role == ContextRole::Seed) {
362                continue;
363            }
364        }
365        if used_chars + cost > char_budget && !packed.is_empty() {
366            break;
367        }
368        used_chars += cost;
369        if cand.role == ContextRole::Neighbor && cand.node.node_type == NodeType::Symbol {
370            symbol_neighbors += 1;
371        }
372        if let Some(ex) = &ctx_node.excerpt {
373            packed_excerpts.push(ex.clone());
374        }
375        packed.push(ctx_node);
376    }
377
378    neighbor_ids.retain(|id| !seed_ids.contains(id));
379    neighbor_ids.truncate(48);
380
381    let latency_us = start.elapsed().as_micros() as u64;
382    let tokens_used = used_chars.div_ceil(4);
383
384    Ok(ContextBundle {
385        prompt: prompt.to_string(),
386        max_tokens: opts.max_tokens,
387        tokens_used,
388        nodes: packed,
389        neighbor_ids,
390        latency_us,
391        graph_nodes,
392        graph_edges,
393    })
394}
395
396enum SymbolQuality {
397    Drop,
398    Keep { boost: f32 },
399}
400
401/// Drop theme consts / short noise; boost query-matching or method-like symbols.
402fn symbol_neighbor_quality(node: &Node, query_tokens: &[String]) -> SymbolQuality {
403    let title = node.title.to_lowercase();
404    let id = node.id.to_lowercase();
405    let leaf = id.rsplit('/').next().unwrap_or(&id);
406
407    // Query term appears in symbol name → keep and boost.
408    for t in query_tokens {
409        if title.contains(t.as_str()) || id.contains(t.as_str()) {
410            return SymbolQuality::Keep { boost: 1.8 };
411        }
412    }
413
414    // Method / path-like symbols (Type::method or multi-segment).
415    if title.contains("::") || leaf.contains("::") {
416        return SymbolQuality::Keep { boost: 1.15 };
417    }
418
419    // Very short or SCREAMING_SNAKE theme constants without query match.
420    let name = title.split("::").last().unwrap_or(&title);
421    if name.len() <= 2 {
422        return SymbolQuality::Drop;
423    }
424    if name.chars().all(|c| c.is_ascii_uppercase() || c == '_') && name.len() <= 12 {
425        // OK, ERR, BG_*, ACCENT_* style noise
426        return SymbolQuality::Drop;
427    }
428    // Generic single-token modules without query match (derive, tree, …)
429    if !name.contains('_') && !title.contains("::") && name.len() <= 8 {
430        // Keep types that look like CamelCase API entry points.
431        let has_upper = name.chars().any(|c| c.is_ascii_uppercase());
432        let has_lower = name.chars().any(|c| c.is_ascii_lowercase());
433        if has_upper && has_lower {
434            return SymbolQuality::Keep { boost: 1.0 };
435        }
436        return SymbolQuality::Drop;
437    }
438
439    SymbolQuality::Keep { boost: 1.0 }
440}
441
442/// Soft-inject README hub + common bootstrap goals when FTS is empty or generic.
443fn inject_hub_seeds(
444    db: &Database,
445    seeds: &mut Vec<RankedHit>,
446    opts: &ContextOptions,
447) -> Result<()> {
448    let existing: HashSet<String> = seeds.iter().map(|h| h.node.id.clone()).collect();
449    let hubs = [
450        ("readme", 4.5_f32),
451        ("docs/goals/from-readme", 3.8),
452        ("docs/implementation/module-map.generated", 2.2),
453        ("docs/goals/readme", 1.5),
454    ];
455    for (id, score) in hubs {
456        if existing.contains(id) {
457            continue;
458        }
459        if let Some(node) = db.get_node(id)? {
460            if !opts.allows_pack(&node.node_type, ContextRole::Seed) {
461                continue;
462            }
463            if is_template_stub(&node) {
464                continue;
465            }
466            seeds.push(RankedHit {
467                node,
468                score,
469                reasons: vec!["hub-fallback".into()],
470            });
471        }
472    }
473    // Keep highest first for take(max_seeds).
474    seeds.sort_by(|a, b| {
475        b.score
476            .partial_cmp(&a.score)
477            .unwrap_or(std::cmp::Ordering::Equal)
478    });
479    Ok(())
480}
481
482fn is_template_stub(node: &Node) -> bool {
483    let id = node.id.to_lowercase();
484    let path = node.file_path.as_deref().unwrap_or("").to_lowercase();
485    let title = node.title.to_lowercase();
486    id.contains("template")
487        || path.ends_with("template.md")
488        || title.contains("template")
489        || title == "adr template"
490}
491
492fn is_generated_path(node: &Node) -> bool {
493    let path = node.file_path.as_deref().unwrap_or("");
494    path.contains("from-readme")
495        || path.contains(".generated.")
496        || path.ends_with("module-map.generated.md")
497}
498
499fn is_readme_family(id: &str) -> bool {
500    id == "readme" || id.contains("from-readme")
501}
502
503fn excerpt_jaccard(a: &str, b: &str) -> f32 {
504    let ta: HashSet<&str> = a
505        .split_whitespace()
506        .filter(|w| w.len() > 2)
507        .collect();
508    let tb: HashSet<&str> = b
509        .split_whitespace()
510        .filter(|w| w.len() > 2)
511        .collect();
512    if ta.is_empty() || tb.is_empty() {
513        return 0.0;
514    }
515    let inter = ta.intersection(&tb).count() as f32;
516    let uni = ta.union(&tb).count() as f32;
517    if uni == 0.0 {
518        0.0
519    } else {
520        inter / uni
521    }
522}
523
524fn pack_rank(c: &Candidate) -> f32 {
525    let role = match c.role {
526        ContextRole::Seed => 3.0,
527        ContextRole::Neighbor => 0.0,
528    };
529    let ty = match c.node.node_type {
530        NodeType::Adr => 2.5,
531        NodeType::Goal => 2.0,
532        NodeType::EdgeCase => 1.8,
533        NodeType::Concept => 1.4,
534        NodeType::Reference => 1.2,
535        NodeType::Alternative => 1.1,
536        NodeType::Symbol => 0.0,
537    };
538    // Score still dominates large gaps; bonuses break near-ties for agent packs.
539    c.score + role + ty
540}
541
542fn load_excerpt(db: &Database, node: &Node, max_chars: usize) -> Option<String> {
543    let raw = db
544        .get_fts_content(&node.id)
545        .ok()
546        .flatten()
547        .filter(|s| !s.trim().is_empty())
548        .or_else(|| node.summary.clone().filter(|s| !s.trim().is_empty()))?;
549    let cleaned = strip_yaml_frontmatter(&raw);
550    Some(truncate_excerpt(cleaned, max_chars))
551}
552
553/// Drop leading `---` YAML frontmatter so agent packs show body prose first.
554fn strip_yaml_frontmatter(s: &str) -> &str {
555    let t = s.trim_start();
556    if !t.starts_with("---") {
557        return s.trim();
558    }
559    let after_open = &t[3..];
560    // Allow optional newline after opening fence.
561    let body = after_open.strip_prefix('\n').unwrap_or(after_open);
562    if let Some(idx) = body.find("\n---") {
563        let rest = &body[idx + 4..];
564        return rest.trim_start_matches(['\r', '\n']).trim();
565    }
566    s.trim()
567}
568
569fn truncate_excerpt(s: &str, max_chars: usize) -> String {
570    let s = s.trim();
571    if s.chars().count() <= max_chars {
572        return s.to_string();
573    }
574    let mut out = String::new();
575    for (i, ch) in s.chars().enumerate() {
576        if i >= max_chars.saturating_sub(1) {
577            break;
578        }
579        out.push(ch);
580    }
581    out.push('…');
582    out
583}
584
585fn to_context_node(c: &Candidate) -> ContextNode {
586    ContextNode {
587        id: c.node.id.clone(),
588        node_type: c.node.node_type.clone(),
589        title: c.node.title.clone(),
590        summary: c.node.summary.clone(),
591        excerpt: c.excerpt.clone(),
592        file_path: c.node.file_path.clone(),
593        score_hint: c.score,
594        role: c.role,
595        hop: c.hop,
596    }
597}
598
599fn estimate_overhead(prompt: &str) -> usize {
600    120 + prompt.len()
601}
602
603fn estimate_node_chars(n: &ContextNode) -> usize {
604    let mut c = n.id.len() + n.title.len() + 48;
605    if let Some(s) = &n.summary {
606        c += s.len();
607    }
608    if let Some(ex) = &n.excerpt {
609        c += ex.len();
610    }
611    if let Some(p) = &n.file_path {
612        c += p.len();
613    }
614    c
615}
616
617/// Convert ranked hits to nodes (helper for CLI / adapters).
618pub fn hits_to_nodes(hits: Vec<RankedHit>) -> Vec<Node> {
619    hits.into_iter().map(|h| h.node).collect()
620}
621
622/// Multi-workspace ranked hit (workspace path + hit).
623///
624/// Prefer [`crate::GlobalRankedHit`] for the registry API; this type remains for
625/// ad-hoc tooling that reuses the same shape.
626#[derive(Debug, Clone, Serialize, Deserialize)]
627pub struct WorkspaceHit {
628    /// Absolute workspace path that produced the hit.
629    pub workspace: String,
630    /// Ranked hit within that workspace.
631    pub hit: RankedHit,
632}
633
634#[cfg(test)]
635mod tests {
636    use super::*;
637    use crate::brain::Brain;
638    use tempfile::tempdir;
639
640    #[test]
641    fn context_includes_graph_neighbors_within_budget() {
642        let dir = tempdir().unwrap();
643        let docs = dir.path().join("docs");
644        std::fs::create_dir_all(&docs).unwrap();
645        // Only seed A is FTS-matchable for "alphaunique"; B is only reachable via graph.
646        std::fs::write(
647            docs.join("alpha.md"),
648            "---\ntags: [alphaunique]\nnode_type: concept\n---\n# AlphaUnique\nLinks [[betaunique]].\nalphaunique body.\n",
649        )
650        .unwrap();
651        std::fs::write(
652            docs.join("beta.md"),
653            "---\ntags: [other]\nnode_type: concept\n---\n# BetaUnique\nNo alphaunique here.\n",
654        )
655        .unwrap();
656
657        let mut brain = Brain::create(dir.path()).unwrap();
658        brain.sync().unwrap();
659
660        let opts = ContextOptions {
661            max_tokens: 2048,
662            hop_depth: 1,
663            max_seeds: 4,
664            max_nodes: 12,
665            hop_decay: 0.7,
666            ..ContextOptions::default()
667        };
668        let ctx = assemble_context(brain.database(), brain.brain_dir(), "alphaunique", &opts)
669            .unwrap();
670        assert!(!ctx.nodes.is_empty());
671        assert!(ctx.tokens_used > 0);
672        assert!(ctx.tokens_used <= opts.max_tokens + 64); // soft bound with estimator slack
673
674        // Neighbor should be discovered via graph even if not an FTS seed.
675        let has_beta = ctx.nodes.iter().any(|n| n.id.contains("beta"))
676            || ctx.neighbor_ids.iter().any(|id| id.contains("beta"));
677        assert!(
678            has_beta,
679            "expected beta neighbor via graph; nodes={:?} neigh={:?}",
680            ctx.nodes.iter().map(|n| &n.id).collect::<Vec<_>>(),
681            ctx.neighbor_ids
682        );
683    }
684
685    #[test]
686    fn tight_budget_limits_nodes() {
687        let dir = tempdir().unwrap();
688        let docs = dir.path().join("docs");
689        std::fs::create_dir_all(&docs).unwrap();
690        for i in 0..8 {
691            std::fs::write(
692                docs.join(format!("n{i}.md")),
693                format!("---\ntags: [topic]\n---\n# Note {i}\nshared topic content {i}\n"),
694            )
695            .unwrap();
696        }
697        let mut brain = Brain::create(dir.path()).unwrap();
698        brain.sync().unwrap();
699        let opts = ContextOptions {
700            max_tokens: 40, // very tight
701            hop_depth: 0,
702            max_seeds: 8,
703            max_nodes: 20,
704            hop_decay: 0.5,
705            ..ContextOptions::default()
706        };
707        let ctx =
708            assemble_context(brain.database(), brain.brain_dir(), "topic", &opts).unwrap();
709        assert!(!ctx.nodes.is_empty());
710        assert!(ctx.nodes.len() < 8, "budget should clip nodes");
711    }
712
713    #[test]
714    fn natural_language_prompt_seeds_and_packs_excerpt() {
715        let dir = tempdir().unwrap();
716        std::fs::write(
717            dir.path().join("README.md"),
718            "# demotool\n\nLightweight **egui** explorer powered by DuckDB.\n\
719             Inspired by Duckling but avoids Tauri/WebView.\n",
720        )
721        .unwrap();
722        let mut brain = Brain::create(dir.path()).unwrap();
723        brain.sync().unwrap();
724
725        let ctx = assemble_context(
726            brain.database(),
727            brain.brain_dir(),
728            "why egui not tauri",
729            &ContextOptions::default(),
730        )
731        .unwrap();
732        assert!(
733            !ctx.nodes.is_empty(),
734            "expected seeds for NL prompt; packed=0"
735        );
736        let has_excerpt = ctx.nodes.iter().any(|n| {
737            n.excerpt
738                .as_ref()
739                .map(|e| e.to_lowercase().contains("egui") || e.to_lowercase().contains("tauri"))
740                .unwrap_or(false)
741        });
742        assert!(
743            has_excerpt,
744            "expected body excerpt mentioning egui/tauri; nodes={:?}",
745            ctx.nodes
746                .iter()
747                .map(|n| (&n.id, &n.excerpt))
748                .collect::<Vec<_>>()
749        );
750    }
751
752    #[test]
753    fn generic_overview_prompt_falls_back_to_hub() {
754        let dir = tempdir().unwrap();
755        std::fs::write(
756            dir.path().join("README.md"),
757            "# demotool\n\nNative egui + DuckDB CLI. Avoids Tauri.\n",
758        )
759        .unwrap();
760        let mut brain = Brain::create(dir.path()).unwrap();
761        brain.sync().unwrap();
762
763        for prompt in [
764            "summarize architecture",
765            "what is this project about",
766            "give an overview",
767        ] {
768            let ctx = assemble_context(
769                brain.database(),
770                brain.brain_dir(),
771                prompt,
772                &ContextOptions::default(),
773            )
774            .unwrap();
775            assert!(
776                !ctx.nodes.is_empty(),
777                "expected hub fallback for `{prompt}`; packed=0"
778            );
779            let has_hub = ctx.nodes.iter().any(|n| {
780                n.id == "readme"
781                    || n.id.contains("from-readme")
782                    || n.excerpt
783                        .as_ref()
784                        .map(|e| e.to_lowercase().contains("egui") || e.to_lowercase().contains("duckdb"))
785                        .unwrap_or(false)
786            });
787            assert!(
788                has_hub,
789                "expected README hub content for `{prompt}`; nodes={:?}",
790                ctx.nodes.iter().map(|n| &n.id).collect::<Vec<_>>()
791            );
792        }
793    }
794
795    #[test]
796    fn strips_frontmatter_from_excerpts() {
797        let dir = tempdir().unwrap();
798        std::fs::create_dir_all(dir.path().join("docs")).unwrap();
799        std::fs::write(
800            dir.path().join("docs/a.md"),
801            "---\ntags: [x]\nnode_type: concept\n---\n# Alpha\n\nBody about egui.\n",
802        )
803        .unwrap();
804        let mut brain = Brain::create(dir.path()).unwrap();
805        brain.sync().unwrap();
806        let ctx = assemble_context(
807            brain.database(),
808            brain.brain_dir(),
809            "egui",
810            &ContextOptions {
811                hop_depth: 0,
812                ..ContextOptions::default()
813            },
814        )
815        .unwrap();
816        let ex = ctx
817            .nodes
818            .iter()
819            .find_map(|n| n.excerpt.as_ref())
820            .expect("excerpt");
821        assert!(
822            !ex.trim_start().starts_with("---"),
823            "frontmatter leaked into excerpt: {ex}"
824        );
825        assert!(ex.to_lowercase().contains("egui") || ex.contains("Alpha"));
826    }
827
828    #[test]
829    fn prefers_adr_seed_over_symbol_neighbor() {
830        let dir = tempdir().unwrap();
831        std::fs::create_dir_all(dir.path().join("docs/adr")).unwrap();
832        std::fs::create_dir_all(dir.path().join("src")).unwrap();
833        std::fs::write(
834            dir.path().join("Cargo.toml"),
835            "[package]\nname = \"demo\"\nversion = \"0.1.0\"\nedition = \"2021\"\n",
836        )
837        .unwrap();
838        std::fs::write(
839            dir.path().join("src/lib.rs"),
840            "/// Run SQL via duckdb CLI.\npub fn run_sql() {}\n",
841        )
842        .unwrap();
843        std::fs::write(
844            dir.path().join("docs/adr/sql.md"),
845            "---\nnode_type: adr\n---\n# SQL via duckdb\n\nUse symbol:run_sql for SQL execution via duckdb CLI.\n",
846        )
847        .unwrap();
848        let mut brain = Brain::create(dir.path()).unwrap();
849        brain.sync().unwrap();
850        let _ = brain.sync().unwrap(); // resolve anchors
851        let ctx = assemble_context(
852            brain.database(),
853            brain.brain_dir(),
854            "how does SQL execution work",
855            &ContextOptions::default(),
856        )
857        .unwrap();
858        assert!(!ctx.nodes.is_empty());
859        let first = &ctx.nodes[0];
860        assert_eq!(
861            first.node_type,
862            NodeType::Adr,
863            "expected ADR first, got {:?} id={}",
864            first.node_type,
865            first.id
866        );
867    }
868
869    #[test]
870    fn dedups_near_identical_excerpts() {
871        let dir = tempdir().unwrap();
872        let body = "Lightweight egui explorer with DuckDB CLI. Avoids Tauri completely.\n".repeat(3);
873        std::fs::write(
874            dir.path().join("README.md"),
875            format!("# demotool\n\n{body}"),
876        )
877        .unwrap();
878        std::fs::create_dir_all(dir.path().join("docs/goals")).unwrap();
879        std::fs::write(
880            dir.path().join("docs/goals/from-readme.md"),
881            format!(
882                "---\nnode_type: goal\n---\n# Goals harvested from README\n\n{body}"
883            ),
884        )
885        .unwrap();
886        let mut brain = Brain::create(dir.path()).unwrap();
887        brain.sync().unwrap();
888        let ctx = assemble_context(
889            brain.database(),
890            brain.brain_dir(),
891            "egui duckdb tauri",
892            &ContextOptions {
893                max_tokens: 900,
894                hop_depth: 0,
895                ..ContextOptions::default()
896            },
897        )
898        .unwrap();
899        // Should not pack both near-identical README + harvest when budget is modest.
900        let ids: Vec<_> = ctx.nodes.iter().map(|n| n.id.as_str()).collect();
901        let both = ids.contains(&"readme") && ids.iter().any(|i| i.contains("from-readme"));
902        assert!(
903            !both || ctx.nodes.len() == 1,
904            "expected dedup of near-identical hubs; packed={ids:?}"
905        );
906    }
907}