Skip to main content

ctx/
smart.rs

1//! Smart context selection for LLM workflows.
2//!
3//! This module provides intelligent file selection based on:
4//! - Semantic search using embeddings
5//! - Call graph expansion (callers and callees)
6//! - Token budget management
7//!
8//! # Usage
9//!
10//! ```ignore
11//! let config = SmartConfig::default();
12//! let result = smart_context(&db, &analytics, &provider, "add caching", config)?;
13//! for file in result.selected_files {
14//!     println!("{}: {} tokens", file.path, file.token_count);
15//! }
16//! ```
17
18use std::collections::{BTreeSet, HashMap};
19use std::path::Path;
20
21use crate::analytics::{Analytics, CallGraphNode, ImpactNode};
22use crate::db::Database;
23use crate::embeddings::{semantic_search, Embedding, EmbeddingProvider, SearchResult};
24use crate::error::{CtxError, Result};
25use crate::tokens::{count_file_tokens, select_by_token_budget, Encoding, HasTokenCount};
26use crate::utils::lexical_tokens;
27
28/// Configuration for smart context selection.
29#[derive(Debug, Clone)]
30pub struct SmartConfig {
31    /// Maximum tokens in output
32    pub max_tokens: usize,
33    /// Call graph expansion depth
34    pub depth: i32,
35    /// Number of initial semantic matches to find
36    pub top: usize,
37    /// Tokenizer encoding to use
38    pub encoding: Encoding,
39}
40
41impl Default for SmartConfig {
42    fn default() -> Self {
43        Self {
44            max_tokens: 8000,
45            depth: 2,
46            top: 10,
47            encoding: Encoding::default(),
48        }
49    }
50}
51
52/// Reason why a file was selected.
53#[derive(Debug, Clone)]
54pub enum SelectionReason {
55    /// Direct semantic match to task embedding
56    SemanticMatch {
57        /// The symbol that matched
58        symbol: String,
59        /// Similarity score (0.0-1.0)
60        score: f32,
61    },
62    /// Called by a matched symbol
63    CalledBy {
64        /// The symbol that calls this file's code
65        caller: String,
66        /// Depth in call graph
67        depth: i32,
68    },
69    /// Calls a matched symbol
70    Calls {
71        /// The symbol that this file calls
72        callee: String,
73        /// Depth in call graph
74        depth: i32,
75    },
76    /// In the same module as a matched symbol
77    #[allow(dead_code)] // Used in tests and pattern matching
78    SameModule {
79        /// The related symbol
80        symbol: String,
81    },
82    /// User explicitly requested this file
83    #[allow(dead_code)] // Used in tests and pattern matching
84    Explicit,
85}
86
87impl SelectionReason {
88    /// Get a human-readable description of the reason.
89    pub fn description(&self) -> String {
90        match self {
91            SelectionReason::SemanticMatch { symbol, score } => {
92                format!("SemanticMatch: \"{}\" (score: {:.2})", symbol, score)
93            }
94            SelectionReason::CalledBy { caller, depth } => {
95                format!("CalledBy: {} (depth {})", caller, depth)
96            }
97            SelectionReason::Calls { callee, depth } => {
98                format!("Calls: {} (depth {})", callee, depth)
99            }
100            SelectionReason::SameModule { symbol } => {
101                format!("SameModule: shares context with {}", symbol)
102            }
103            SelectionReason::Explicit => "Explicit: user requested".to_string(),
104        }
105    }
106
107    /// Get the relevance weight for this reason type.
108    fn weight(&self) -> f32 {
109        match self {
110            SelectionReason::SemanticMatch { score, .. } => *score,
111            SelectionReason::CalledBy { depth, .. } => 0.8 / (*depth as f32).max(1.0),
112            SelectionReason::Calls { depth, .. } => 0.7 / (*depth as f32).max(1.0),
113            SelectionReason::SameModule { .. } => 0.5,
114            SelectionReason::Explicit => 1.0,
115        }
116    }
117}
118
119/// A selected file with relevance information.
120#[derive(Debug, Clone)]
121pub struct FileSelection {
122    /// File path (relative to project root)
123    pub path: String,
124    /// Relevance score (0.0-1.0)
125    pub relevance_score: f32,
126    /// Why this file was selected
127    pub reasons: Vec<SelectionReason>,
128    /// Token count for this file
129    pub token_count: usize,
130    /// Number of distinct task tokens that appear in this file's path or in the
131    /// names of the symbols that selected it (see [`FileSelection::compute_lexical_score`]).
132    /// A high-precision relevance signal layered on top of the semantic/graph tiers.
133    pub lexical_score: f32,
134}
135
136impl HasTokenCount for FileSelection {
137    fn token_count(&self) -> usize {
138        self.token_count
139    }
140}
141
142impl FileSelection {
143    /// Create a new file selection.
144    fn new(path: String, reason: SelectionReason) -> Self {
145        let score = reason.weight();
146        Self {
147            path,
148            relevance_score: score,
149            reasons: vec![reason],
150            token_count: 0,
151            lexical_score: 0.0,
152        }
153    }
154
155    /// Add another reason for selection, updating the relevance score.
156    fn add_reason(&mut self, reason: SelectionReason) {
157        let new_weight = reason.weight();
158        // Take the maximum weight as the relevance score
159        if new_weight > self.relevance_score {
160            self.relevance_score = new_weight;
161        }
162        self.reasons.push(reason);
163    }
164
165    /// The best (maximum) direct semantic-match score among this file's reasons,
166    /// or `None` if the file was pulled in only through call-graph expansion.
167    ///
168    /// Files that directly match the task embedding are ranked ahead of graph-only
169    /// files (see [`rank_files`]), so this must be read independently of
170    /// `relevance_score`, which `add_reason` collapses to the max weight across
171    /// reasons (and which the flat 0.7/0.8 call-graph constants would otherwise
172    /// dominate over a compressed semantic score).
173    fn best_semantic_score(&self) -> Option<f32> {
174        self.reasons
175            .iter()
176            .filter_map(|r| match r {
177                SelectionReason::SemanticMatch { score, .. } => Some(*score),
178                _ => None,
179            })
180            .reduce(f32::max)
181    }
182
183    /// Count the distinct `task_tokens` that appear in this file's **path**.
184    ///
185    /// This is the lexical relevance signal: when a task word ("openai",
186    /// "solidity", "sql") is literally present in a candidate's path, that is
187    /// strong, high-precision evidence the file is on-topic even if the embedding
188    /// ranked it low. Returns 0.0 when there is no overlap.
189    ///
190    /// Path only — deliberately *not* symbol names. Matching symbol names is
191    /// low-precision: ubiquitous identifiers (notably `ctx`, the tool's own name,
192    /// which appears as a test helper across `tests/*_cli.rs`) would score a hit
193    /// on nearly every file and drown out the genuinely on-topic one.
194    fn compute_lexical_score(&self, task_tokens: &BTreeSet<String>) -> f32 {
195        if task_tokens.is_empty() {
196            return 0.0;
197        }
198        let path_tokens = lexical_tokens(&self.path);
199        task_tokens.intersection(&path_tokens).count() as f32
200    }
201}
202
203/// Result of smart context selection.
204#[derive(Debug, Clone)]
205pub struct SmartContext {
206    /// The task description used for selection
207    #[allow(dead_code)] // Part of public API
208    pub task: String,
209    /// Selected files with their relevance information
210    pub selected_files: Vec<FileSelection>,
211    /// Total token count of selected files
212    pub total_tokens: usize,
213    /// Whether selection was truncated by token limit
214    pub truncated: bool,
215    /// Number of files omitted due to token limit
216    pub omitted_count: usize,
217}
218
219/// Select files relevant to a task using semantic search and call graph analysis.
220///
221/// # Algorithm
222///
223/// 1. Embed the task description
224/// 2. Find top-N symbols matching the embedding
225/// 3. For each matched symbol:
226///    - Add the file containing the symbol
227///    - Expand call graph (callers + callees) to specified depth
228///    - Add files from expanded symbols
229/// 4. Deduplicate files, keeping highest relevance
230/// 5. Count tokens for each file
231/// 6. Sort by relevance score descending
232/// 7. Include files until token limit is reached
233pub fn smart_context(
234    db: &Database,
235    analytics: &Analytics,
236    provider: &dyn EmbeddingProvider,
237    task: &str,
238    config: SmartConfig,
239) -> Result<SmartContext> {
240    // 1. Embed the task description
241    let task_embedding = provider.embed(task)?;
242
243    smart_context_with_embedding(db, analytics, task, &task_embedding, config)
244}
245
246/// Select files relevant to a task using a pre-computed embedding.
247///
248/// This variant is useful when the embedding has been computed asynchronously
249/// (e.g., in the MCP server with OpenAI's async API) to avoid blocking the async runtime.
250///
251/// See [`smart_context`] for the full algorithm description.
252pub fn smart_context_with_embedding(
253    db: &Database,
254    analytics: &Analytics,
255    task: &str,
256    task_embedding: &Embedding,
257    config: SmartConfig,
258) -> Result<SmartContext> {
259    // 2. Semantic search for matching symbols
260    let matches = semantic_search(db, task_embedding, config.top)?;
261
262    if matches.is_empty() {
263        return Err(CtxError::NoMatches);
264    }
265
266    // 3. Build file selection map
267    let mut files: HashMap<String, FileSelection> = HashMap::new();
268
269    for result in &matches {
270        // Add the file containing the matched symbol
271        add_file(
272            &mut files,
273            &result.file_path,
274            SelectionReason::SemanticMatch {
275                symbol: result.name.clone(),
276                score: result.score,
277            },
278        );
279
280        // Expand call graph
281        expand_symbol(&mut files, analytics, result, config.depth)?;
282    }
283
284    // 4. Convert to vector and count tokens
285    let mut selections: Vec<FileSelection> = files.into_values().collect();
286
287    // Count tokens for each file
288    for selection in &mut selections {
289        selection.token_count = count_file_token_safe(&selection.path, config.encoding);
290    }
291
292    // 4b. Lexical relevance: reward candidates whose path or matched symbol names
293    // contain the task's own tokens. This rescues on-topic files the embedding
294    // ranked low (e.g. `embeddings/openai.rs` for "…openai") without any tunable
295    // float weight — the score is a plain distinct-token-hit count.
296    let task_tokens = lexical_tokens(task);
297    for selection in &mut selections {
298        selection.lexical_score = selection.compute_lexical_score(&task_tokens);
299    }
300
301    // 5. Rank files by relevance
302    rank_files(&mut selections);
303
304    // 6. Apply token limit, but never silently drop the single most-relevant file.
305    let (selected, total_tokens, omitted) =
306        select_with_guaranteed_top(selections, config.max_tokens);
307
308    Ok(SmartContext {
309        task: task.to_string(),
310        selected_files: selected,
311        total_tokens,
312        truncated: omitted > 0,
313        omitted_count: omitted,
314    })
315}
316
317/// Apply the token budget while guaranteeing the top-ranked file is always
318/// included — even when that file alone exceeds the budget.
319///
320/// The greedy first-fit selector would otherwise skip an oversized rank-1 file
321/// (e.g. a 9k-token parser larger than the whole 8k budget) and backfill with
322/// smaller, less-relevant files, so the command returns everything *except* the
323/// file the task is really about. Here the top file is force-included, then the
324/// remainder is first-fit against the leftover budget via the shared
325/// [`select_by_token_budget`].
326fn select_with_guaranteed_top(
327    ranked: Vec<FileSelection>,
328    max_tokens: usize,
329) -> (Vec<FileSelection>, usize, usize) {
330    let mut iter = ranked.into_iter();
331    let Some(top) = iter.next() else {
332        return (Vec::new(), 0, 0);
333    };
334    let top_tokens = top.token_count;
335    let rest: Vec<FileSelection> = iter.collect();
336    // Budget the remainder against whatever is left after the guaranteed top file
337    // (saturating: an oversized top file leaves a zero budget for the rest).
338    let remaining_budget = max_tokens.saturating_sub(top_tokens);
339    let (mut selected_rest, rest_tokens, omitted) = select_by_token_budget(rest, remaining_budget);
340
341    let mut selected = Vec::with_capacity(selected_rest.len() + 1);
342    selected.push(top);
343    selected.append(&mut selected_rest);
344    (selected, top_tokens + rest_tokens, omitted)
345}
346
347/// Add a file to the selection map, merging reasons if already present.
348fn add_file(files: &mut HashMap<String, FileSelection>, path: &str, reason: SelectionReason) {
349    if let Some(existing) = files.get_mut(path) {
350        existing.add_reason(reason);
351    } else {
352        files.insert(
353            path.to_string(),
354            FileSelection::new(path.to_string(), reason),
355        );
356    }
357}
358
359/// Expand a symbol's call graph and add related files.
360fn expand_symbol(
361    files: &mut HashMap<String, FileSelection>,
362    analytics: &Analytics,
363    result: &SearchResult,
364    depth: i32,
365) -> Result<()> {
366    // Expand callers (who calls this symbol - impact analysis)
367    if let Ok(callers) = analytics.impact_analysis(&result.symbol_id, depth) {
368        for caller in callers {
369            add_file_from_impact(files, &caller, &result.name);
370        }
371    }
372
373    // Expand callees (what does this symbol call - call graph)
374    if let Ok(callees) = analytics.call_graph(&result.symbol_id, depth) {
375        for callee in callees {
376            add_file_from_call_graph(files, &callee, &result.name);
377        }
378    }
379
380    Ok(())
381}
382
383/// Add a file from impact analysis result.
384fn add_file_from_impact(
385    files: &mut HashMap<String, FileSelection>,
386    node: &ImpactNode,
387    callee_name: &str,
388) {
389    add_file(
390        files,
391        &node.file_path,
392        SelectionReason::CalledBy {
393            caller: node.name.clone(),
394            depth: node.distance,
395        },
396    );
397
398    // If in the same file as the callee, also add SameModule reason
399    if let Some(existing) = files.get(&node.file_path) {
400        if existing
401            .reasons
402            .iter()
403            .any(|r| matches!(r, SelectionReason::SemanticMatch { .. }))
404        {
405            // Already has a semantic match, skip SameModule
406        } else {
407            // Check if this is the same module (we'll add SameModule reason separately if needed)
408            let _ = callee_name; // Suppress unused warning
409        }
410    }
411}
412
413/// Add a file from call graph result.
414fn add_file_from_call_graph(
415    files: &mut HashMap<String, FileSelection>,
416    node: &CallGraphNode,
417    caller_name: &str,
418) {
419    add_file(
420        files,
421        &node.file_path,
422        SelectionReason::Calls {
423            callee: node.name.clone(),
424            depth: node.depth,
425        },
426    );
427
428    let _ = caller_name; // Suppress unused warning
429}
430
431/// Ranking key for a file: `(is_low_relevance, lexical_score, tier_score)`.
432///
433/// - **Relevance tier** (`is_low_relevance`): a file is in the top tier if it has
434///   a direct semantic match **or** a lexical hit (a task token in its path/name).
435///   Graph-only files with no lexical hit sort last, so a generic call-graph hub
436///   (flat weight 0.8) never displaces an on-topic file.
437/// - **`lexical_score`**: within a tier, more task-token hits rank first — this is
438///   what surfaces `embeddings/openai.rs` for "…openai" over a higher-scored but
439///   off-topic semantic match.
440/// - **`tier_score`**: the semantic score for semantic matches, else
441///   `relevance_score`; scores are only ever compared within the same scale.
442fn rank_key(f: &FileSelection) -> (bool, f32, f32) {
443    let sem = f.best_semantic_score();
444    let is_low_relevance = sem.is_none() && f.lexical_score == 0.0;
445    let tier_score = sem.unwrap_or(f.relevance_score);
446    (is_low_relevance, f.lexical_score, tier_score)
447}
448
449/// Rank files for selection. The ordering is tiered and fully deterministic:
450///
451/// 1. Relevant files (semantic match or lexical hit) rank above files pulled in
452///    only through call-graph expansion.
453/// 2. Within a tier: more lexical hits first, then higher tier score.
454/// 3. Ties break by `path` ascending. Because the input is collected from a
455///    `HashMap`, this tie-break is what makes the command deterministic across runs
456///    (many files share the same 0.5/0.7/0.8 weight and the same lexical score).
457fn rank_files(files: &mut [FileSelection]) {
458    files.sort_by(|a, b| {
459        let a_key = rank_key(a);
460        let b_key = rank_key(b);
461        a_key
462            .0
463            .cmp(&b_key.0) // false (relevant) sorts before true (low-relevance)
464            .then_with(|| {
465                b_key
466                    .1
467                    .partial_cmp(&a_key.1) // lexical_score descending
468                    .unwrap_or(std::cmp::Ordering::Equal)
469            })
470            .then_with(|| {
471                b_key
472                    .2
473                    .partial_cmp(&a_key.2) // tier_score descending
474                    .unwrap_or(std::cmp::Ordering::Equal)
475            })
476            .then_with(|| a.path.cmp(&b.path))
477    });
478}
479
480/// Count tokens in a file, returning 0 on error.
481fn count_file_token_safe(path: &str, encoding: Encoding) -> usize {
482    count_file_tokens(Path::new(path), encoding)
483        .map(|tc| tc.count)
484        .unwrap_or(0)
485}
486
487/// Format the smart context result for display.
488pub fn format_explain(result: &SmartContext) -> String {
489    let mut output = String::new();
490
491    output.push_str(&format!(
492        "Selected {} files ({} tokens):\n\n",
493        result.selected_files.len(),
494        result.total_tokens
495    ));
496
497    for (i, file) in result.selected_files.iter().enumerate() {
498        output.push_str(&format!(
499            "{}. {} ({} tokens)\n",
500            i + 1,
501            file.path,
502            file.token_count
503        ));
504
505        for reason in &file.reasons {
506            output.push_str(&format!("   - {}\n", reason.description()));
507        }
508        output.push('\n');
509    }
510
511    if result.truncated {
512        output.push_str(&format!(
513            "({} files omitted due to token limit)\n",
514            result.omitted_count
515        ));
516    }
517
518    output
519}
520
521/// Format the smart context result for dry-run mode.
522pub fn format_dry_run(result: &SmartContext) -> String {
523    let mut output = String::new();
524
525    output.push_str(&format!(
526        "Would select {} files ({} tokens):\n",
527        result.selected_files.len(),
528        result.total_tokens
529    ));
530
531    for file in &result.selected_files {
532        let primary_reason = file
533            .reasons
534            .first()
535            .map(|r| match r {
536                SelectionReason::SemanticMatch { .. } => "SemanticMatch",
537                SelectionReason::CalledBy { .. } => "CalledBy",
538                SelectionReason::Calls { .. } => "Calls",
539                SelectionReason::SameModule { .. } => "SameModule",
540                SelectionReason::Explicit => "Explicit",
541            })
542            .unwrap_or("Unknown");
543
544        output.push_str(&format!(
545            "  {} ({} tokens) - {}\n",
546            file.path, file.token_count, primary_reason
547        ));
548    }
549
550    if result.truncated {
551        output.push_str(&format!(
552            "\n({} files would be omitted due to token limit)\n",
553            result.omitted_count
554        ));
555    }
556
557    output
558}
559
560#[cfg(test)]
561mod tests {
562    use super::*;
563
564    #[test]
565    fn test_selection_reason_weight() {
566        let semantic = SelectionReason::SemanticMatch {
567            symbol: "test".to_string(),
568            score: 0.9,
569        };
570        assert!((semantic.weight() - 0.9).abs() < 0.001);
571
572        let called_by = SelectionReason::CalledBy {
573            caller: "main".to_string(),
574            depth: 1,
575        };
576        assert!((called_by.weight() - 0.8).abs() < 0.001);
577
578        let called_by_depth2 = SelectionReason::CalledBy {
579            caller: "main".to_string(),
580            depth: 2,
581        };
582        assert!((called_by_depth2.weight() - 0.4).abs() < 0.001);
583
584        let calls = SelectionReason::Calls {
585            callee: "helper".to_string(),
586            depth: 1,
587        };
588        assert!((calls.weight() - 0.7).abs() < 0.001);
589
590        let same_module = SelectionReason::SameModule {
591            symbol: "related".to_string(),
592        };
593        assert!((same_module.weight() - 0.5).abs() < 0.001);
594
595        let explicit = SelectionReason::Explicit;
596        assert!((explicit.weight() - 1.0).abs() < 0.001);
597    }
598
599    #[test]
600    fn test_file_selection_add_reason() {
601        let mut selection = FileSelection::new(
602            "src/main.rs".to_string(),
603            SelectionReason::Calls {
604                callee: "helper".to_string(),
605                depth: 1,
606            },
607        );
608
609        assert!((selection.relevance_score - 0.7).abs() < 0.001);
610        assert_eq!(selection.reasons.len(), 1);
611
612        // Add a higher-weight reason
613        selection.add_reason(SelectionReason::SemanticMatch {
614            symbol: "main".to_string(),
615            score: 0.95,
616        });
617
618        assert!((selection.relevance_score - 0.95).abs() < 0.001);
619        assert_eq!(selection.reasons.len(), 2);
620    }
621
622    #[test]
623    fn test_select_by_token_budget() {
624        let files = vec![
625            FileSelection {
626                path: "a.rs".to_string(),
627                relevance_score: 1.0,
628                reasons: vec![SelectionReason::Explicit],
629                token_count: 100,
630                lexical_score: 0.0,
631            },
632            FileSelection {
633                path: "b.rs".to_string(),
634                relevance_score: 0.8,
635                reasons: vec![SelectionReason::Explicit],
636                token_count: 200,
637                lexical_score: 0.0,
638            },
639            FileSelection {
640                path: "c.rs".to_string(),
641                relevance_score: 0.5,
642                reasons: vec![SelectionReason::Explicit],
643                token_count: 150,
644                lexical_score: 0.0,
645            },
646        ];
647
648        // All fit
649        let (selected, total, omitted) = select_by_token_budget(files.clone(), 500);
650        assert_eq!(selected.len(), 3);
651        assert_eq!(total, 450);
652        assert_eq!(omitted, 0);
653
654        // Only first two fit
655        let (selected, total, omitted) = select_by_token_budget(files.clone(), 300);
656        assert_eq!(selected.len(), 2);
657        assert_eq!(total, 300);
658        assert_eq!(omitted, 1);
659
660        // Only first fits
661        let (selected, total, omitted) = select_by_token_budget(files, 150);
662        assert_eq!(selected.len(), 1);
663        assert_eq!(total, 100);
664        assert_eq!(omitted, 2);
665    }
666
667    #[test]
668    fn test_format_dry_run() {
669        let result = SmartContext {
670            task: "add caching".to_string(),
671            selected_files: vec![
672                FileSelection {
673                    path: "src/main.rs".to_string(),
674                    relevance_score: 0.9,
675                    reasons: vec![SelectionReason::SemanticMatch {
676                        symbol: "main".to_string(),
677                        score: 0.9,
678                    }],
679                    token_count: 500,
680                    lexical_score: 0.0,
681                },
682                FileSelection {
683                    path: "src/lib.rs".to_string(),
684                    relevance_score: 0.7,
685                    reasons: vec![SelectionReason::Calls {
686                        callee: "helper".to_string(),
687                        depth: 1,
688                    }],
689                    token_count: 300,
690                    lexical_score: 0.0,
691                },
692            ],
693            total_tokens: 800,
694            truncated: false,
695            omitted_count: 0,
696        };
697
698        let output = format_dry_run(&result);
699        assert!(output.contains("Would select 2 files"));
700        assert!(output.contains("src/main.rs"));
701        assert!(output.contains("SemanticMatch"));
702    }
703
704    #[test]
705    fn test_smart_config_default() {
706        let config = SmartConfig::default();
707        assert_eq!(config.max_tokens, 8000);
708        assert_eq!(config.depth, 2);
709        assert_eq!(config.top, 10);
710    }
711
712    #[test]
713    fn test_add_file_merges_reasons() {
714        let mut files: HashMap<String, FileSelection> = HashMap::new();
715
716        // First add
717        add_file(
718            &mut files,
719            "src/main.rs",
720            SelectionReason::SemanticMatch {
721                symbol: "main".to_string(),
722                score: 0.9,
723            },
724        );
725        assert_eq!(files.len(), 1);
726        assert!((files.get("src/main.rs").unwrap().relevance_score - 0.9).abs() < 0.001);
727
728        // Second add to same file - should merge
729        add_file(
730            &mut files,
731            "src/main.rs",
732            SelectionReason::CalledBy {
733                caller: "run".to_string(),
734                depth: 1,
735            },
736        );
737        assert_eq!(files.len(), 1); // Still only 1 file
738        assert_eq!(files.get("src/main.rs").unwrap().reasons.len(), 2);
739        // Score should be max of the two
740        assert!((files.get("src/main.rs").unwrap().relevance_score - 0.9).abs() < 0.001);
741    }
742
743    #[test]
744    fn test_rank_files_sorts_by_relevance() {
745        let mut files = vec![
746            FileSelection {
747                path: "low.rs".to_string(),
748                relevance_score: 0.3,
749                reasons: vec![SelectionReason::Explicit],
750                token_count: 100,
751                lexical_score: 0.0,
752            },
753            FileSelection {
754                path: "high.rs".to_string(),
755                relevance_score: 0.9,
756                reasons: vec![SelectionReason::Explicit],
757                token_count: 100,
758                lexical_score: 0.0,
759            },
760            FileSelection {
761                path: "mid.rs".to_string(),
762                relevance_score: 0.5,
763                reasons: vec![SelectionReason::Explicit],
764                token_count: 100,
765                lexical_score: 0.0,
766            },
767        ];
768
769        rank_files(&mut files);
770
771        assert_eq!(files[0].path, "high.rs");
772        assert_eq!(files[1].path, "mid.rs");
773        assert_eq!(files[2].path, "low.rs");
774    }
775
776    /// A file that directly matches the task embedding must rank above a file that
777    /// was only pulled in through call-graph expansion, even though the graph
778    /// neighbour's flat weight (0.8) is numerically larger than the compressed
779    /// semantic score (0.5). This is the core relevance regression.
780    #[test]
781    fn test_semantic_match_ranks_above_graph_only() {
782        let mut files = vec![
783            FileSelection {
784                path: "graph_hub.rs".to_string(),
785                relevance_score: 0.8, // CalledBy depth 1
786                reasons: vec![SelectionReason::CalledBy {
787                    caller: "run".to_string(),
788                    depth: 1,
789                }],
790                token_count: 100,
791                lexical_score: 0.0,
792            },
793            FileSelection {
794                path: "on_topic.rs".to_string(),
795                relevance_score: 0.5,
796                reasons: vec![SelectionReason::SemanticMatch {
797                    symbol: "run_sql".to_string(),
798                    score: 0.5,
799                }],
800                token_count: 100,
801                lexical_score: 0.0,
802            },
803        ];
804
805        rank_files(&mut files);
806
807        assert_eq!(
808            files[0].path, "on_topic.rs",
809            "semantic match must outrank a higher-weighted graph-only file"
810        );
811        assert_eq!(files[1].path, "graph_hub.rs");
812    }
813
814    /// Within the semantic tier, higher score wins; equal scores break by path so
815    /// the order is stable.
816    #[test]
817    fn test_semantic_tier_orders_by_score_then_path() {
818        let mk = |path: &str, score: f32| FileSelection {
819            path: path.to_string(),
820            relevance_score: score,
821            reasons: vec![SelectionReason::SemanticMatch {
822                symbol: "s".to_string(),
823                score,
824            }],
825            token_count: 100,
826            lexical_score: 0.0,
827        };
828        // Two files tie at 0.6; "a.rs" must precede "b.rs" by path.
829        let mut files = vec![mk("b.rs", 0.6), mk("c.rs", 0.4), mk("a.rs", 0.6)];
830
831        rank_files(&mut files);
832
833        assert_eq!(
834            files.iter().map(|f| f.path.as_str()).collect::<Vec<_>>(),
835            vec!["a.rs", "b.rs", "c.rs"]
836        );
837    }
838
839    /// Ranking must not depend on the (HashMap-derived) input order: the same set of
840    /// selections presented in any order yields the same ranked paths. This is what
841    /// makes `ctx smart` deterministic across runs.
842    #[test]
843    fn test_rank_files_is_deterministic() {
844        let make_set = || {
845            vec![
846                FileSelection {
847                    path: "z_graph.rs".to_string(),
848                    relevance_score: 0.8,
849                    reasons: vec![SelectionReason::CalledBy {
850                        caller: "run".to_string(),
851                        depth: 1,
852                    }],
853                    token_count: 100,
854                    lexical_score: 0.0,
855                },
856                FileSelection {
857                    path: "a_graph.rs".to_string(),
858                    relevance_score: 0.8, // ties with z_graph.rs -> path breaks it
859                    reasons: vec![SelectionReason::CalledBy {
860                        caller: "run".to_string(),
861                        depth: 1,
862                    }],
863                    token_count: 100,
864                    lexical_score: 0.0,
865                },
866                FileSelection {
867                    path: "seed.rs".to_string(),
868                    relevance_score: 0.5,
869                    reasons: vec![SelectionReason::SemanticMatch {
870                        symbol: "seed".to_string(),
871                        score: 0.5,
872                    }],
873                    token_count: 100,
874                    lexical_score: 0.0,
875                },
876            ]
877        };
878
879        let mut a = make_set();
880        rank_files(&mut a);
881        let order_a: Vec<String> = a.iter().map(|f| f.path.clone()).collect();
882
883        // Present the same set reversed; ranking must produce the identical order.
884        let mut b = make_set();
885        b.reverse();
886        rank_files(&mut b);
887        let order_b: Vec<String> = b.iter().map(|f| f.path.clone()).collect();
888
889        assert_eq!(order_a, order_b);
890        // Semantic seed first, then graph files tie-broken by path.
891        assert_eq!(order_a, vec!["seed.rs", "a_graph.rs", "z_graph.rs"]);
892    }
893
894    /// A graph-only file whose path/name matches a task token (lexical_score > 0)
895    /// is promoted into the relevant tier, above a semantic match that shares no
896    /// task token. This is what surfaces `embeddings/openai.rs` for "…openai".
897    #[test]
898    fn test_lexical_promotes_graph_only_file() {
899        let mut files = vec![
900            FileSelection {
901                path: "off_topic_semantic.rs".to_string(),
902                relevance_score: 0.6,
903                reasons: vec![SelectionReason::SemanticMatch {
904                    symbol: "unrelated".to_string(),
905                    score: 0.6,
906                }],
907                token_count: 100,
908                lexical_score: 0.0,
909            },
910            FileSelection {
911                path: "embeddings/openai.rs".to_string(),
912                relevance_score: 0.7, // Calls, graph-only
913                reasons: vec![SelectionReason::Calls {
914                    callee: "embed".to_string(),
915                    depth: 1,
916                }],
917                token_count: 100,
918                lexical_score: 2.0, // "embeddings" + "openai"
919            },
920        ];
921
922        rank_files(&mut files);
923
924        assert_eq!(
925            files[0].path, "embeddings/openai.rs",
926            "a lexical hit must outrank a semantic match with no task-token overlap"
927        );
928    }
929
930    /// Within the relevant tier, more lexical hits rank first; equal lexical scores
931    /// fall back to tier score, then path.
932    #[test]
933    fn test_lexical_orders_within_tier() {
934        let mk = |path: &str, score: f32, lex: f32| FileSelection {
935            path: path.to_string(),
936            relevance_score: score,
937            reasons: vec![SelectionReason::SemanticMatch {
938                symbol: "s".to_string(),
939                score,
940            }],
941            token_count: 100,
942            lexical_score: lex,
943        };
944        // two_hits has the most lexical overlap and must lead despite a lower score.
945        let mut files = vec![
946            mk("high_score_no_lexical.rs", 0.9, 0.0),
947            mk("two_hits.rs", 0.4, 2.0),
948            mk("one_hit.rs", 0.4, 1.0),
949        ];
950
951        rank_files(&mut files);
952
953        assert_eq!(
954            files.iter().map(|f| f.path.as_str()).collect::<Vec<_>>(),
955            vec!["two_hits.rs", "one_hit.rs", "high_score_no_lexical.rs"]
956        );
957    }
958
959    /// `compute_lexical_score` counts distinct task tokens present in a file's
960    /// path (only), and does not count symbol names — so ubiquitous identifiers
961    /// like `ctx` cannot inflate the score.
962    #[test]
963    fn test_compute_lexical_score() {
964        let sel = FileSelection {
965            path: "src/embeddings/openai.rs".to_string(),
966            relevance_score: 0.5,
967            // Symbol name deliberately contains "ctx"; it must NOT be counted.
968            reasons: vec![SelectionReason::SemanticMatch {
969                symbol: "ctx_embed".to_string(),
970                score: 0.5,
971            }],
972            token_count: 100,
973            lexical_score: 0.0,
974        };
975        // "embeddings" + "openai" from the path overlap the task; "with" is a stopword.
976        let task = lexical_tokens("generate embeddings with openai");
977        assert!((sel.compute_lexical_score(&task) - 2.0).abs() < 0.001);
978        // "ctx" appears in the symbol but not the path, so it does not match.
979        let ctx_task = lexical_tokens("ctx tooling");
980        assert!((sel.compute_lexical_score(&ctx_task)).abs() < 0.001);
981        // No overlap -> 0.
982        let other = lexical_tokens("parse solidity contracts");
983        assert!((sel.compute_lexical_score(&other)).abs() < 0.001);
984    }
985
986    /// The rank-1 file is always included even when it alone exceeds the budget,
987    /// rather than being silently dropped in favour of smaller lower-ranked files.
988    #[test]
989    fn test_budget_includes_oversized_top_file() {
990        let ranked = vec![
991            FileSelection {
992                path: "huge_top.rs".to_string(),
993                relevance_score: 0.9,
994                reasons: vec![SelectionReason::SemanticMatch {
995                    symbol: "s".to_string(),
996                    score: 0.9,
997                }],
998                token_count: 9000, // larger than the whole budget
999                lexical_score: 1.0,
1000            },
1001            FileSelection {
1002                path: "small_next.rs".to_string(),
1003                relevance_score: 0.5,
1004                reasons: vec![SelectionReason::SemanticMatch {
1005                    symbol: "s".to_string(),
1006                    score: 0.5,
1007                }],
1008                token_count: 500,
1009                lexical_score: 0.0,
1010            },
1011        ];
1012
1013        let (selected, total, omitted) = select_with_guaranteed_top(ranked, 8000);
1014
1015        assert_eq!(selected.len(), 1, "only the oversized top file is included");
1016        assert_eq!(selected[0].path, "huge_top.rs");
1017        assert_eq!(total, 9000);
1018        assert_eq!(omitted, 1, "the smaller file is omitted for lack of budget");
1019    }
1020
1021    /// When the top file fits, the remainder is first-fit against the leftover
1022    /// budget (same behaviour as before for the non-oversized case).
1023    #[test]
1024    fn test_budget_first_fits_remainder() {
1025        let ranked = vec![
1026            FileSelection {
1027                path: "top.rs".to_string(),
1028                relevance_score: 0.9,
1029                reasons: vec![SelectionReason::Explicit],
1030                token_count: 3000,
1031                lexical_score: 0.0,
1032            },
1033            FileSelection {
1034                path: "mid.rs".to_string(),
1035                relevance_score: 0.8,
1036                reasons: vec![SelectionReason::Explicit],
1037                token_count: 4000,
1038                lexical_score: 0.0,
1039            },
1040            FileSelection {
1041                path: "tail.rs".to_string(),
1042                relevance_score: 0.7,
1043                reasons: vec![SelectionReason::Explicit],
1044                token_count: 2000, // would overflow the 8000 budget after top+mid
1045                lexical_score: 0.0,
1046            },
1047        ];
1048
1049        let (selected, total, omitted) = select_with_guaranteed_top(ranked, 8000);
1050
1051        assert_eq!(
1052            selected.iter().map(|f| f.path.as_str()).collect::<Vec<_>>(),
1053            vec!["top.rs", "mid.rs"]
1054        );
1055        assert_eq!(total, 7000);
1056        assert_eq!(omitted, 1);
1057    }
1058
1059    #[test]
1060    fn test_selection_reason_description() {
1061        let semantic = SelectionReason::SemanticMatch {
1062            symbol: "test".to_string(),
1063            score: 0.9,
1064        };
1065        assert!(semantic.description().contains("SemanticMatch"));
1066        assert!(semantic.description().contains("test"));
1067
1068        let called_by = SelectionReason::CalledBy {
1069            caller: "main".to_string(),
1070            depth: 2,
1071        };
1072        assert!(called_by.description().contains("CalledBy"));
1073        assert!(called_by.description().contains("main"));
1074        assert!(called_by.description().contains("2"));
1075    }
1076}