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