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::Changelog => 1.95,
567        NodeType::Plan => 1.9,
568        NodeType::EdgeCase => 1.8,
569        NodeType::Analysis => 1.65,
570        NodeType::Concept => 1.4,
571        NodeType::Reference => 1.2,
572        NodeType::Alternative => 1.1,
573        NodeType::Symbol => 0.0,
574    };
575    // Score still dominates large gaps; bonuses break near-ties for agent packs.
576    c.score + role + ty
577}
578
579fn load_excerpt(db: &Database, node: &Node, max_chars: usize) -> Option<String> {
580    let raw = db
581        .get_fts_content(&node.id)
582        .ok()
583        .flatten()
584        .filter(|s| !s.trim().is_empty())
585        .or_else(|| node.summary.clone().filter(|s| !s.trim().is_empty()))?;
586    let cleaned = strip_yaml_frontmatter(&raw);
587    Some(truncate_excerpt(cleaned, max_chars))
588}
589
590/// Drop leading `---` YAML frontmatter so agent packs show body prose first.
591fn strip_yaml_frontmatter(s: &str) -> &str {
592    let t = s.trim_start();
593    if !t.starts_with("---") {
594        return s.trim();
595    }
596    let after_open = &t[3..];
597    // Allow optional newline after opening fence.
598    let body = after_open.strip_prefix('\n').unwrap_or(after_open);
599    if let Some(idx) = body.find("\n---") {
600        let rest = &body[idx + 4..];
601        return rest.trim_start_matches(['\r', '\n']).trim();
602    }
603    s.trim()
604}
605
606fn truncate_excerpt(s: &str, max_chars: usize) -> String {
607    let s = s.trim();
608    if s.chars().count() <= max_chars {
609        return s.to_string();
610    }
611    let mut out = String::new();
612    for (i, ch) in s.chars().enumerate() {
613        if i >= max_chars.saturating_sub(1) {
614            break;
615        }
616        out.push(ch);
617    }
618    out.push('…');
619    out
620}
621
622fn to_context_node(c: &Candidate) -> ContextNode {
623    ContextNode {
624        id: c.node.id.clone(),
625        node_type: c.node.node_type.clone(),
626        title: c.node.title.clone(),
627        summary: c.node.summary.clone(),
628        excerpt: c.excerpt.clone(),
629        file_path: c.node.file_path.clone(),
630        score_hint: c.score,
631        role: c.role,
632        hop: c.hop,
633    }
634}
635
636fn estimate_overhead(prompt: &str) -> usize {
637    120 + prompt.len()
638}
639
640fn estimate_node_chars(n: &ContextNode) -> usize {
641    let mut c = n.id.len() + n.title.len() + 48;
642    if let Some(s) = &n.summary {
643        c += s.len();
644    }
645    if let Some(ex) = &n.excerpt {
646        c += ex.len();
647    }
648    if let Some(p) = &n.file_path {
649        c += p.len();
650    }
651    c
652}
653
654/// Convert ranked hits to nodes (helper for CLI / adapters).
655pub fn hits_to_nodes(hits: Vec<RankedHit>) -> Vec<Node> {
656    hits.into_iter().map(|h| h.node).collect()
657}
658
659/// Multi-workspace ranked hit (workspace path + hit).
660///
661/// Prefer [`crate::GlobalRankedHit`] for the registry API; this type remains for
662/// ad-hoc tooling that reuses the same shape.
663#[derive(Debug, Clone, Serialize, Deserialize)]
664pub struct WorkspaceHit {
665    /// Absolute workspace path that produced the hit.
666    pub workspace: String,
667    /// Ranked hit within that workspace.
668    pub hit: RankedHit,
669}
670
671#[cfg(test)]
672mod tests {
673    use super::*;
674    use crate::brain::Brain;
675    use crate::indexer::WorkspaceIndexer;
676    use crate::storage::Database;
677    use tempfile::tempdir;
678
679    #[test]
680    fn context_includes_graph_neighbors_within_budget() {
681        let dir = tempdir().unwrap();
682        let docs = dir.path().join("docs");
683        std::fs::create_dir_all(&docs).unwrap();
684        // Only seed A is FTS-matchable for "alphaunique"; B is only reachable via graph.
685        std::fs::write(
686            docs.join("alpha.md"),
687            "---\ntags: [alphaunique]\nnode_type: concept\n---\n# AlphaUnique\nLinks [[betaunique]].\nalphaunique body.\n",
688        )
689        .unwrap();
690        std::fs::write(
691            docs.join("beta.md"),
692            "---\ntags: [other]\nnode_type: concept\n---\n# BetaUnique\nNo alphaunique here.\n",
693        )
694        .unwrap();
695
696        let mut brain = Brain::create(dir.path()).unwrap();
697        brain.sync().unwrap();
698
699        let opts = ContextOptions {
700            max_tokens: 2048,
701            hop_depth: 1,
702            max_seeds: 4,
703            max_nodes: 12,
704            hop_decay: 0.7,
705            ..ContextOptions::default()
706        };
707        let ctx = assemble_context(brain.database(), brain.brain_dir(), "alphaunique", &opts)
708            .unwrap();
709        assert!(!ctx.nodes.is_empty());
710        assert!(ctx.tokens_used > 0);
711        assert!(ctx.tokens_used <= opts.max_tokens + 64); // soft bound with estimator slack
712
713        // Neighbor should be discovered via graph even if not an FTS seed.
714        let has_beta = ctx.nodes.iter().any(|n| n.id.contains("beta"))
715            || ctx.neighbor_ids.iter().any(|id| id.contains("beta"));
716        assert!(
717            has_beta,
718            "expected beta neighbor via graph; nodes={:?} neigh={:?}",
719            ctx.nodes.iter().map(|n| &n.id).collect::<Vec<_>>(),
720            ctx.neighbor_ids
721        );
722    }
723
724    #[test]
725    fn tight_budget_limits_nodes() {
726        let dir = tempdir().unwrap();
727        let docs = dir.path().join("docs");
728        std::fs::create_dir_all(&docs).unwrap();
729        for i in 0..8 {
730            std::fs::write(
731                docs.join(format!("n{i}.md")),
732                format!("---\ntags: [topic]\n---\n# Note {i}\nshared topic content {i}\n"),
733            )
734            .unwrap();
735        }
736        let mut brain = Brain::create(dir.path()).unwrap();
737        brain.sync().unwrap();
738        let opts = ContextOptions {
739            max_tokens: 40, // very tight
740            hop_depth: 0,
741            max_seeds: 8,
742            max_nodes: 20,
743            hop_decay: 0.5,
744            ..ContextOptions::default()
745        };
746        let ctx =
747            assemble_context(brain.database(), brain.brain_dir(), "topic", &opts).unwrap();
748        assert!(!ctx.nodes.is_empty());
749        assert!(ctx.nodes.len() < 8, "budget should clip nodes");
750    }
751
752    #[test]
753    fn natural_language_prompt_seeds_and_packs_excerpt() {
754        let dir = tempdir().unwrap();
755        std::fs::write(
756            dir.path().join("README.md"),
757            "# demotool\n\nLightweight **egui** explorer powered by DuckDB.\n\
758             Inspired by Duckling but avoids Tauri/WebView.\n",
759        )
760        .unwrap();
761        let mut brain = Brain::create(dir.path()).unwrap();
762        brain.sync().unwrap();
763
764        let ctx = assemble_context(
765            brain.database(),
766            brain.brain_dir(),
767            "why egui not tauri",
768            &ContextOptions::default(),
769        )
770        .unwrap();
771        assert!(
772            !ctx.nodes.is_empty(),
773            "expected seeds for NL prompt; packed=0"
774        );
775        let has_excerpt = ctx.nodes.iter().any(|n| {
776            n.excerpt
777                .as_ref()
778                .map(|e| e.to_lowercase().contains("egui") || e.to_lowercase().contains("tauri"))
779                .unwrap_or(false)
780        });
781        assert!(
782            has_excerpt,
783            "expected body excerpt mentioning egui/tauri; nodes={:?}",
784            ctx.nodes
785                .iter()
786                .map(|n| (&n.id, &n.excerpt))
787                .collect::<Vec<_>>()
788        );
789    }
790
791    #[test]
792    fn release_prompt_prefers_changelog_hub() {
793        let dir = tempdir().unwrap();
794        std::fs::write(
795            dir.path().join("README.md"),
796            "# Demo\n\nShip notes live in the changelog.\n",
797        )
798        .unwrap();
799        std::fs::write(
800            dir.path().join("CHANGELOG.md"),
801            "# Changelog\n\n## [1.2.3] - 2026-01-01\n\n### Added\n- important feature alpha\n",
802        )
803        .unwrap();
804        let brain = dir.path().join(".brain");
805        std::fs::create_dir_all(&brain).unwrap();
806        let db = Database::open(brain.join("db.sqlite")).unwrap();
807        let indexer = WorkspaceIndexer::new(db, dir.path());
808        indexer.index_workspace().unwrap();
809        let db = Database::open(brain.join("db.sqlite")).unwrap();
810        let ctx = assemble_context(
811            &db,
812            &brain,
813            "what shipped in the changelog release",
814            &ContextOptions {
815                max_tokens: 800,
816                hop_depth: 0,
817                ..ContextOptions::default()
818            },
819        )
820        .unwrap();
821        assert!(
822            ctx.nodes.iter().any(|n| n.id == "changelog"),
823            "expected changelog hub packed; nodes={:?}",
824            ctx.nodes.iter().map(|n| &n.id).collect::<Vec<_>>()
825        );
826    }
827
828    #[test]
829    fn generic_overview_prompt_falls_back_to_hub() {
830        let dir = tempdir().unwrap();
831        std::fs::write(
832            dir.path().join("README.md"),
833            "# demotool\n\nNative egui + DuckDB CLI. Avoids Tauri.\n",
834        )
835        .unwrap();
836        let mut brain = Brain::create(dir.path()).unwrap();
837        brain.sync().unwrap();
838
839        for prompt in [
840            "summarize architecture",
841            "what is this project about",
842            "give an overview",
843        ] {
844            let ctx = assemble_context(
845                brain.database(),
846                brain.brain_dir(),
847                prompt,
848                &ContextOptions::default(),
849            )
850            .unwrap();
851            assert!(
852                !ctx.nodes.is_empty(),
853                "expected hub fallback for `{prompt}`; packed=0"
854            );
855            let has_hub = ctx.nodes.iter().any(|n| {
856                n.id == "readme"
857                    || n.id.contains("from-readme")
858                    || n.excerpt
859                        .as_ref()
860                        .map(|e| e.to_lowercase().contains("egui") || e.to_lowercase().contains("duckdb"))
861                        .unwrap_or(false)
862            });
863            assert!(
864                has_hub,
865                "expected README hub content for `{prompt}`; nodes={:?}",
866                ctx.nodes.iter().map(|n| &n.id).collect::<Vec<_>>()
867            );
868        }
869    }
870
871    #[test]
872    fn strips_frontmatter_from_excerpts() {
873        let dir = tempdir().unwrap();
874        std::fs::create_dir_all(dir.path().join("docs")).unwrap();
875        std::fs::write(
876            dir.path().join("docs/a.md"),
877            "---\ntags: [x]\nnode_type: concept\n---\n# Alpha\n\nBody about egui.\n",
878        )
879        .unwrap();
880        let mut brain = Brain::create(dir.path()).unwrap();
881        brain.sync().unwrap();
882        let ctx = assemble_context(
883            brain.database(),
884            brain.brain_dir(),
885            "egui",
886            &ContextOptions {
887                hop_depth: 0,
888                ..ContextOptions::default()
889            },
890        )
891        .unwrap();
892        let ex = ctx
893            .nodes
894            .iter()
895            .find_map(|n| n.excerpt.as_ref())
896            .expect("excerpt");
897        assert!(
898            !ex.trim_start().starts_with("---"),
899            "frontmatter leaked into excerpt: {ex}"
900        );
901        assert!(ex.to_lowercase().contains("egui") || ex.contains("Alpha"));
902    }
903
904    #[test]
905    fn prefers_adr_seed_over_symbol_neighbor() {
906        let dir = tempdir().unwrap();
907        std::fs::create_dir_all(dir.path().join("docs/adr")).unwrap();
908        std::fs::create_dir_all(dir.path().join("src")).unwrap();
909        std::fs::write(
910            dir.path().join("Cargo.toml"),
911            "[package]\nname = \"demo\"\nversion = \"0.1.0\"\nedition = \"2021\"\n",
912        )
913        .unwrap();
914        std::fs::write(
915            dir.path().join("src/lib.rs"),
916            "/// Run SQL via duckdb CLI.\npub fn run_sql() {}\n",
917        )
918        .unwrap();
919        std::fs::write(
920            dir.path().join("docs/adr/sql.md"),
921            "---\nnode_type: adr\n---\n# SQL via duckdb\n\nUse symbol:run_sql for SQL execution via duckdb CLI.\n",
922        )
923        .unwrap();
924        let mut brain = Brain::create(dir.path()).unwrap();
925        brain.sync().unwrap();
926        let _ = brain.sync().unwrap(); // resolve anchors
927        let ctx = assemble_context(
928            brain.database(),
929            brain.brain_dir(),
930            "how does SQL execution work",
931            &ContextOptions::default(),
932        )
933        .unwrap();
934        assert!(!ctx.nodes.is_empty());
935        let first = &ctx.nodes[0];
936        assert_eq!(
937            first.node_type,
938            NodeType::Adr,
939            "expected ADR first, got {:?} id={}",
940            first.node_type,
941            first.id
942        );
943    }
944
945    #[test]
946    fn dedups_near_identical_excerpts() {
947        let dir = tempdir().unwrap();
948        let body = "Lightweight egui explorer with DuckDB CLI. Avoids Tauri completely.\n".repeat(3);
949        std::fs::write(
950            dir.path().join("README.md"),
951            format!("# demotool\n\n{body}"),
952        )
953        .unwrap();
954        std::fs::create_dir_all(dir.path().join("docs/goals")).unwrap();
955        std::fs::write(
956            dir.path().join("docs/goals/from-readme.md"),
957            format!(
958                "---\nnode_type: goal\n---\n# Goals harvested from README\n\n{body}"
959            ),
960        )
961        .unwrap();
962        let mut brain = Brain::create(dir.path()).unwrap();
963        brain.sync().unwrap();
964        let ctx = assemble_context(
965            brain.database(),
966            brain.brain_dir(),
967            "egui duckdb tauri",
968            &ContextOptions {
969                max_tokens: 900,
970                hop_depth: 0,
971                ..ContextOptions::default()
972            },
973        )
974        .unwrap();
975        // Should not pack both near-identical README + harvest when budget is modest.
976        let ids: Vec<_> = ctx.nodes.iter().map(|n| n.id.as_str()).collect();
977        let both = ids.contains(&"readme") && ids.iter().any(|i| i.contains("from-readme"));
978        assert!(
979            !both || ctx.nodes.len() == 1,
980            "expected dedup of near-identical hubs; packed={ids:?}"
981        );
982    }
983}