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 all candidates by score and pack into token budget.
313    candidates.sort_by(|a, b| {
314        // Prefer non-symbols slightly when scores are close (stable agent packs).
315        let a_bonus = if a.node.node_type == NodeType::Symbol {
316            0.0
317        } else {
318            0.05
319        };
320        let b_bonus = if b.node.node_type == NodeType::Symbol {
321            0.0
322        } else {
323            0.05
324        };
325        (b.score + b_bonus)
326            .partial_cmp(&(a.score + a_bonus))
327            .unwrap_or(std::cmp::Ordering::Equal)
328    });
329
330    let mut packed: Vec<ContextNode> = Vec::new();
331    let mut used_chars = estimate_overhead(prompt);
332    let mut seen = HashSet::new();
333    let mut symbol_neighbors = 0usize;
334    let mut packed_excerpts: Vec<String> = Vec::new();
335
336    for cand in candidates {
337        if packed.len() >= opts.max_nodes {
338            break;
339        }
340        if !seen.insert(cand.node.id.clone()) {
341            continue;
342        }
343        if cand.role == ContextRole::Neighbor
344            && cand.node.node_type == NodeType::Symbol
345            && symbol_neighbors >= MAX_PACKED_SYMBOL_NEIGHBORS
346        {
347            continue;
348        }
349        // Prefer a single README-family hub (root readme XOR from-readme harvest).
350        if is_readme_family(&cand.node.id)
351            && packed.iter().any(|n| is_readme_family(&n.id))
352        {
353            continue;
354        }
355        // Drop near-duplicate body text (harvested clones of the same prose).
356        if let Some(ex) = &cand.excerpt {
357            if packed_excerpts
358                .iter()
359                .any(|prev| excerpt_jaccard(prev, ex) > 0.55)
360            {
361                continue;
362            }
363        }
364        let ctx_node = to_context_node(&cand);
365        let cost = estimate_node_chars(&ctx_node);
366        if used_chars + cost > char_budget && !packed.is_empty() {
367            // Always try to keep at least one seed
368            if cand.role == ContextRole::Neighbor {
369                continue;
370            }
371            if packed.iter().any(|n| n.role == ContextRole::Seed) {
372                continue;
373            }
374        }
375        if used_chars + cost > char_budget && !packed.is_empty() {
376            break;
377        }
378        used_chars += cost;
379        if cand.role == ContextRole::Neighbor && cand.node.node_type == NodeType::Symbol {
380            symbol_neighbors += 1;
381        }
382        if let Some(ex) = &ctx_node.excerpt {
383            packed_excerpts.push(ex.clone());
384        }
385        packed.push(ctx_node);
386    }
387
388    neighbor_ids.retain(|id| !seed_ids.contains(id));
389    neighbor_ids.truncate(48);
390
391    let latency_us = start.elapsed().as_micros() as u64;
392    let tokens_used = used_chars.div_ceil(4);
393
394    Ok(ContextBundle {
395        prompt: prompt.to_string(),
396        max_tokens: opts.max_tokens,
397        tokens_used,
398        nodes: packed,
399        neighbor_ids,
400        latency_us,
401        graph_nodes,
402        graph_edges,
403    })
404}
405
406enum SymbolQuality {
407    Drop,
408    Keep { boost: f32 },
409}
410
411/// Drop theme consts / short noise; boost query-matching or method-like symbols.
412fn symbol_neighbor_quality(node: &Node, query_tokens: &[String]) -> SymbolQuality {
413    let title = node.title.to_lowercase();
414    let id = node.id.to_lowercase();
415    let leaf = id.rsplit('/').next().unwrap_or(&id);
416
417    // Query term appears in symbol name → keep and boost.
418    for t in query_tokens {
419        if title.contains(t.as_str()) || id.contains(t.as_str()) {
420            return SymbolQuality::Keep { boost: 1.8 };
421        }
422    }
423
424    // Method / path-like symbols (Type::method or multi-segment).
425    if title.contains("::") || leaf.contains("::") {
426        return SymbolQuality::Keep { boost: 1.15 };
427    }
428
429    // Very short or SCREAMING_SNAKE theme constants without query match.
430    let name = title.split("::").last().unwrap_or(&title);
431    if name.len() <= 2 {
432        return SymbolQuality::Drop;
433    }
434    if name.chars().all(|c| c.is_ascii_uppercase() || c == '_') && name.len() <= 12 {
435        // OK, ERR, BG_*, ACCENT_* style noise
436        return SymbolQuality::Drop;
437    }
438    // Generic single-token modules without query match (derive, tree, …)
439    if !name.contains('_') && !title.contains("::") && name.len() <= 8 {
440        // Keep types that look like CamelCase API entry points.
441        let has_upper = name.chars().any(|c| c.is_ascii_uppercase());
442        let has_lower = name.chars().any(|c| c.is_ascii_lowercase());
443        if has_upper && has_lower {
444            return SymbolQuality::Keep { boost: 1.0 };
445        }
446        return SymbolQuality::Drop;
447    }
448
449    SymbolQuality::Keep { boost: 1.0 }
450}
451
452/// Soft-inject README hub + common bootstrap goals when FTS is empty or generic.
453fn inject_hub_seeds(
454    db: &Database,
455    seeds: &mut Vec<RankedHit>,
456    opts: &ContextOptions,
457) -> Result<()> {
458    let existing: HashSet<String> = seeds.iter().map(|h| h.node.id.clone()).collect();
459    let hubs = [
460        ("readme", 4.5_f32),
461        ("docs/goals/from-readme", 3.8),
462        ("docs/implementation/module-map.generated", 2.2),
463        ("docs/goals/readme", 1.5),
464    ];
465    for (id, score) in hubs {
466        if existing.contains(id) {
467            continue;
468        }
469        if let Some(node) = db.get_node(id)? {
470            if !opts.allows_pack(&node.node_type, ContextRole::Seed) {
471                continue;
472            }
473            if is_template_stub(&node) {
474                continue;
475            }
476            seeds.push(RankedHit {
477                node,
478                score,
479                reasons: vec!["hub-fallback".into()],
480            });
481        }
482    }
483    // Keep highest first for take(max_seeds).
484    seeds.sort_by(|a, b| {
485        b.score
486            .partial_cmp(&a.score)
487            .unwrap_or(std::cmp::Ordering::Equal)
488    });
489    Ok(())
490}
491
492fn is_template_stub(node: &Node) -> bool {
493    let id = node.id.to_lowercase();
494    let path = node.file_path.as_deref().unwrap_or("").to_lowercase();
495    let title = node.title.to_lowercase();
496    id.contains("template")
497        || path.ends_with("template.md")
498        || title.contains("template")
499        || title == "adr template"
500}
501
502fn is_generated_path(node: &Node) -> bool {
503    let path = node.file_path.as_deref().unwrap_or("");
504    path.contains("from-readme")
505        || path.contains(".generated.")
506        || path.ends_with("module-map.generated.md")
507}
508
509fn is_readme_family(id: &str) -> bool {
510    id == "readme" || id.contains("from-readme")
511}
512
513fn excerpt_jaccard(a: &str, b: &str) -> f32 {
514    let ta: HashSet<&str> = a
515        .split_whitespace()
516        .filter(|w| w.len() > 2)
517        .collect();
518    let tb: HashSet<&str> = b
519        .split_whitespace()
520        .filter(|w| w.len() > 2)
521        .collect();
522    if ta.is_empty() || tb.is_empty() {
523        return 0.0;
524    }
525    let inter = ta.intersection(&tb).count() as f32;
526    let uni = ta.union(&tb).count() as f32;
527    if uni == 0.0 {
528        0.0
529    } else {
530        inter / uni
531    }
532}
533
534fn load_excerpt(db: &Database, node: &Node, max_chars: usize) -> Option<String> {
535    let raw = db
536        .get_fts_content(&node.id)
537        .ok()
538        .flatten()
539        .filter(|s| !s.trim().is_empty())
540        .or_else(|| node.summary.clone().filter(|s| !s.trim().is_empty()))?;
541    Some(truncate_excerpt(&raw, max_chars))
542}
543
544fn truncate_excerpt(s: &str, max_chars: usize) -> String {
545    let s = s.trim();
546    if s.chars().count() <= max_chars {
547        return s.to_string();
548    }
549    let mut out = String::new();
550    for (i, ch) in s.chars().enumerate() {
551        if i >= max_chars.saturating_sub(1) {
552            break;
553        }
554        out.push(ch);
555    }
556    out.push('…');
557    out
558}
559
560fn to_context_node(c: &Candidate) -> ContextNode {
561    ContextNode {
562        id: c.node.id.clone(),
563        node_type: c.node.node_type.clone(),
564        title: c.node.title.clone(),
565        summary: c.node.summary.clone(),
566        excerpt: c.excerpt.clone(),
567        file_path: c.node.file_path.clone(),
568        score_hint: c.score,
569        role: c.role,
570        hop: c.hop,
571    }
572}
573
574fn estimate_overhead(prompt: &str) -> usize {
575    120 + prompt.len()
576}
577
578fn estimate_node_chars(n: &ContextNode) -> usize {
579    let mut c = n.id.len() + n.title.len() + 48;
580    if let Some(s) = &n.summary {
581        c += s.len();
582    }
583    if let Some(ex) = &n.excerpt {
584        c += ex.len();
585    }
586    if let Some(p) = &n.file_path {
587        c += p.len();
588    }
589    c
590}
591
592/// Convert ranked hits to nodes (helper for CLI / adapters).
593pub fn hits_to_nodes(hits: Vec<RankedHit>) -> Vec<Node> {
594    hits.into_iter().map(|h| h.node).collect()
595}
596
597/// Multi-workspace ranked hit (workspace path + hit).
598///
599/// Prefer [`crate::GlobalRankedHit`] for the registry API; this type remains for
600/// ad-hoc tooling that reuses the same shape.
601#[derive(Debug, Clone, Serialize, Deserialize)]
602pub struct WorkspaceHit {
603    /// Absolute workspace path that produced the hit.
604    pub workspace: String,
605    /// Ranked hit within that workspace.
606    pub hit: RankedHit,
607}
608
609#[cfg(test)]
610mod tests {
611    use super::*;
612    use crate::brain::Brain;
613    use tempfile::tempdir;
614
615    #[test]
616    fn context_includes_graph_neighbors_within_budget() {
617        let dir = tempdir().unwrap();
618        let docs = dir.path().join("docs");
619        std::fs::create_dir_all(&docs).unwrap();
620        // Only seed A is FTS-matchable for "alphaunique"; B is only reachable via graph.
621        std::fs::write(
622            docs.join("alpha.md"),
623            "---\ntags: [alphaunique]\nnode_type: concept\n---\n# AlphaUnique\nLinks [[betaunique]].\nalphaunique body.\n",
624        )
625        .unwrap();
626        std::fs::write(
627            docs.join("beta.md"),
628            "---\ntags: [other]\nnode_type: concept\n---\n# BetaUnique\nNo alphaunique here.\n",
629        )
630        .unwrap();
631
632        let mut brain = Brain::create(dir.path()).unwrap();
633        brain.sync().unwrap();
634
635        let opts = ContextOptions {
636            max_tokens: 2048,
637            hop_depth: 1,
638            max_seeds: 4,
639            max_nodes: 12,
640            hop_decay: 0.7,
641            ..ContextOptions::default()
642        };
643        let ctx = assemble_context(brain.database(), brain.brain_dir(), "alphaunique", &opts)
644            .unwrap();
645        assert!(!ctx.nodes.is_empty());
646        assert!(ctx.tokens_used > 0);
647        assert!(ctx.tokens_used <= opts.max_tokens + 64); // soft bound with estimator slack
648
649        // Neighbor should be discovered via graph even if not an FTS seed.
650        let has_beta = ctx.nodes.iter().any(|n| n.id.contains("beta"))
651            || ctx.neighbor_ids.iter().any(|id| id.contains("beta"));
652        assert!(
653            has_beta,
654            "expected beta neighbor via graph; nodes={:?} neigh={:?}",
655            ctx.nodes.iter().map(|n| &n.id).collect::<Vec<_>>(),
656            ctx.neighbor_ids
657        );
658    }
659
660    #[test]
661    fn tight_budget_limits_nodes() {
662        let dir = tempdir().unwrap();
663        let docs = dir.path().join("docs");
664        std::fs::create_dir_all(&docs).unwrap();
665        for i in 0..8 {
666            std::fs::write(
667                docs.join(format!("n{i}.md")),
668                format!("---\ntags: [topic]\n---\n# Note {i}\nshared topic content {i}\n"),
669            )
670            .unwrap();
671        }
672        let mut brain = Brain::create(dir.path()).unwrap();
673        brain.sync().unwrap();
674        let opts = ContextOptions {
675            max_tokens: 40, // very tight
676            hop_depth: 0,
677            max_seeds: 8,
678            max_nodes: 20,
679            hop_decay: 0.5,
680            ..ContextOptions::default()
681        };
682        let ctx =
683            assemble_context(brain.database(), brain.brain_dir(), "topic", &opts).unwrap();
684        assert!(!ctx.nodes.is_empty());
685        assert!(ctx.nodes.len() < 8, "budget should clip nodes");
686    }
687
688    #[test]
689    fn natural_language_prompt_seeds_and_packs_excerpt() {
690        let dir = tempdir().unwrap();
691        std::fs::write(
692            dir.path().join("README.md"),
693            "# demotool\n\nLightweight **egui** explorer powered by DuckDB.\n\
694             Inspired by Duckling but avoids Tauri/WebView.\n",
695        )
696        .unwrap();
697        let mut brain = Brain::create(dir.path()).unwrap();
698        brain.sync().unwrap();
699
700        let ctx = assemble_context(
701            brain.database(),
702            brain.brain_dir(),
703            "why egui not tauri",
704            &ContextOptions::default(),
705        )
706        .unwrap();
707        assert!(
708            !ctx.nodes.is_empty(),
709            "expected seeds for NL prompt; packed=0"
710        );
711        let has_excerpt = ctx.nodes.iter().any(|n| {
712            n.excerpt
713                .as_ref()
714                .map(|e| e.to_lowercase().contains("egui") || e.to_lowercase().contains("tauri"))
715                .unwrap_or(false)
716        });
717        assert!(
718            has_excerpt,
719            "expected body excerpt mentioning egui/tauri; nodes={:?}",
720            ctx.nodes
721                .iter()
722                .map(|n| (&n.id, &n.excerpt))
723                .collect::<Vec<_>>()
724        );
725    }
726
727    #[test]
728    fn generic_overview_prompt_falls_back_to_hub() {
729        let dir = tempdir().unwrap();
730        std::fs::write(
731            dir.path().join("README.md"),
732            "# demotool\n\nNative egui + DuckDB CLI. Avoids Tauri.\n",
733        )
734        .unwrap();
735        let mut brain = Brain::create(dir.path()).unwrap();
736        brain.sync().unwrap();
737
738        for prompt in [
739            "summarize architecture",
740            "what is this project about",
741            "give an overview",
742        ] {
743            let ctx = assemble_context(
744                brain.database(),
745                brain.brain_dir(),
746                prompt,
747                &ContextOptions::default(),
748            )
749            .unwrap();
750            assert!(
751                !ctx.nodes.is_empty(),
752                "expected hub fallback for `{prompt}`; packed=0"
753            );
754            let has_hub = ctx.nodes.iter().any(|n| {
755                n.id == "readme"
756                    || n.id.contains("from-readme")
757                    || n.excerpt
758                        .as_ref()
759                        .map(|e| e.to_lowercase().contains("egui") || e.to_lowercase().contains("duckdb"))
760                        .unwrap_or(false)
761            });
762            assert!(
763                has_hub,
764                "expected README hub content for `{prompt}`; nodes={:?}",
765                ctx.nodes.iter().map(|n| &n.id).collect::<Vec<_>>()
766            );
767        }
768    }
769
770    #[test]
771    fn dedups_near_identical_excerpts() {
772        let dir = tempdir().unwrap();
773        let body = "Lightweight egui explorer with DuckDB CLI. Avoids Tauri completely.\n".repeat(3);
774        std::fs::write(
775            dir.path().join("README.md"),
776            format!("# demotool\n\n{body}"),
777        )
778        .unwrap();
779        std::fs::create_dir_all(dir.path().join("docs/goals")).unwrap();
780        std::fs::write(
781            dir.path().join("docs/goals/from-readme.md"),
782            format!(
783                "---\nnode_type: goal\n---\n# Goals harvested from README\n\n{body}"
784            ),
785        )
786        .unwrap();
787        let mut brain = Brain::create(dir.path()).unwrap();
788        brain.sync().unwrap();
789        let ctx = assemble_context(
790            brain.database(),
791            brain.brain_dir(),
792            "egui duckdb tauri",
793            &ContextOptions {
794                max_tokens: 900,
795                hop_depth: 0,
796                ..ContextOptions::default()
797            },
798        )
799        .unwrap();
800        // Should not pack both near-identical README + harvest when budget is modest.
801        let ids: Vec<_> = ctx.nodes.iter().map(|n| n.id.as_str()).collect();
802        let both = ids.contains(&"readme") && ids.iter().any(|i| i.contains("from-readme"));
803        assert!(
804            !both || ctx.nodes.len() == 1,
805            "expected dedup of near-identical hubs; packed={ids:?}"
806        );
807    }
808}