lean_ctx/core/
graph_context.rs

1//! Graph-driven context loading — automatically includes related files
2//! based on Property Graph proximity and token budgeting.
3//!
4//! Used by `ctx_read` (task mode) to surface a small, budgeted set of
5//! related files (deterministic ordering; no output spam).
6
7use super::graph_provider::{self, GraphProviderSource};
8use super::tokens::count_tokens;
9
10#[derive(Debug)]
11pub struct GraphContext {
12    pub source: GraphProviderSource,
13    pub primary_file: String,
14    pub related_files: Vec<RelatedFile>,
15    pub total_tokens: usize,
16    pub budget_remaining: usize,
17}
18
19#[derive(Debug)]
20pub struct RelatedFile {
21    pub path: String,
22    pub relationship: Relationship,
23    pub token_count: usize,
24}
25
26#[derive(Debug, Clone)]
27pub enum Relationship {
28    DirectDependency,
29    DirectDependent,
30    TransitiveDependency,
31    TypeProvider,
32}
33
34impl Relationship {
35    pub fn label(&self) -> &'static str {
36        match self {
37            Relationship::DirectDependency => "imports",
38            Relationship::DirectDependent => "imported-by",
39            Relationship::TransitiveDependency => "transitive-dep",
40            Relationship::TypeProvider => "type-provider",
41        }
42    }
43
44    fn priority(&self) -> usize {
45        match self {
46            Relationship::DirectDependency => 0,
47            Relationship::TypeProvider => 1,
48            Relationship::DirectDependent => 2,
49            Relationship::TransitiveDependency => 3,
50        }
51    }
52}
53
54#[derive(Debug, Clone, Copy)]
55pub struct GraphContextOptions {
56    pub token_budget: usize,
57    pub max_files: usize,
58    pub max_edges: usize,
59    pub max_depth: usize,
60    pub allow_build: bool,
61}
62
63impl Default for GraphContextOptions {
64    fn default() -> Self {
65        Self {
66            token_budget: crate::core::budgets::GRAPH_CONTEXT_TOKEN_BUDGET,
67            max_files: crate::core::budgets::GRAPH_CONTEXT_MAX_FILES,
68            max_edges: crate::core::budgets::GRAPH_CONTEXT_MAX_EDGES,
69            max_depth: crate::core::budgets::GRAPH_CONTEXT_MAX_DEPTH,
70            allow_build: false,
71        }
72    }
73}
74
75pub fn build_graph_context(
76    file_path: &str,
77    project_root: &str,
78    options: Option<GraphContextOptions>,
79) -> Option<GraphContext> {
80    let opts = options.unwrap_or_default();
81
82    let rel_path = file_path
83        .strip_prefix(project_root)
84        .unwrap_or(file_path)
85        .trim_start_matches('/');
86
87    let provider_open = if opts.allow_build {
88        graph_provider::open_or_build(project_root)
89    } else {
90        graph_provider::open_best_effort(project_root)
91    }?;
92
93    let primary_content = std::fs::read_to_string(file_path).ok()?;
94    let primary_tokens = count_tokens(&primary_content);
95
96    let remaining = opts.token_budget.saturating_sub(primary_tokens);
97    if remaining < 200 {
98        return Some(GraphContext {
99            source: provider_open.source,
100            primary_file: rel_path.to_string(),
101            related_files: Vec::new(),
102            total_tokens: primary_tokens,
103            budget_remaining: 0,
104        });
105    }
106
107    let mut candidates = collect_candidates(&provider_open, rel_path, opts.max_depth);
108    candidates.sort_by(|a, b| {
109        a.relationship
110            .priority()
111            .cmp(&b.relationship.priority())
112            .then_with(|| a.path.cmp(&b.path))
113    });
114    if candidates.len() > opts.max_edges {
115        candidates.truncate(opts.max_edges);
116    }
117
118    let mut related: Vec<RelatedFile> = Vec::new();
119    let mut tokens_used = primary_tokens;
120    let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();
121    seen.insert(rel_path.to_string());
122
123    for candidate in candidates {
124        if related.len() >= opts.max_files {
125            break;
126        }
127        if seen.contains(&candidate.path) {
128            continue;
129        }
130
131        let abs_path = format!("{project_root}/{}", candidate.path);
132        if let Ok(content) = std::fs::read_to_string(&abs_path) {
133            let tokens = count_tokens(&content);
134            if tokens_used + tokens > opts.token_budget {
135                continue;
136            }
137            tokens_used += tokens;
138            seen.insert(candidate.path.clone());
139            related.push(RelatedFile {
140                path: candidate.path,
141                relationship: candidate.relationship,
142                token_count: tokens,
143            });
144        }
145    }
146
147    Some(GraphContext {
148        source: provider_open.source,
149        primary_file: rel_path.to_string(),
150        related_files: related,
151        total_tokens: tokens_used,
152        budget_remaining: opts.token_budget.saturating_sub(tokens_used),
153    })
154}
155
156struct Candidate {
157    path: String,
158    relationship: Relationship,
159}
160
161fn classify_dep(file: &str) -> Relationship {
162    if file.ends_with(".d.ts") {
163        Relationship::TypeProvider
164    } else {
165        Relationship::DirectDependency
166    }
167}
168
169fn collect_candidates(
170    open: &graph_provider::OpenGraphProvider,
171    file_path: &str,
172    max_depth: usize,
173) -> Vec<Candidate> {
174    let mut candidates: Vec<Candidate> = Vec::new();
175
176    for dep in open.provider.dependencies(file_path) {
177        let rel = classify_dep(&dep);
178        candidates.push(Candidate {
179            path: dep,
180            relationship: rel,
181        });
182    }
183
184    for dep in open.provider.dependents(file_path) {
185        candidates.push(Candidate {
186            path: dep,
187            relationship: Relationship::DirectDependent,
188        });
189    }
190
191    for affected in open.provider.related(file_path, max_depth.max(1)) {
192        let already = candidates.iter().any(|c| c.path == affected);
193        if !already {
194            candidates.push(Candidate {
195                path: affected,
196                relationship: Relationship::TransitiveDependency,
197            });
198        }
199    }
200
201    candidates
202}
203
204pub fn format_graph_context(ctx: &GraphContext) -> String {
205    if ctx.related_files.is_empty() {
206        return String::new();
207    }
208
209    let source = match ctx.source {
210        GraphProviderSource::PropertyGraph => "property_graph",
211        GraphProviderSource::GraphIndex => "graph_index",
212    };
213    let mut result = format!(
214        "\n--- GRAPH CONTEXT (source={source}, {} related files, {} tok) ---\n",
215        ctx.related_files.len(),
216        ctx.total_tokens
217    );
218
219    for rf in &ctx.related_files {
220        result.push_str(&format!(
221            "  {} [{}] ({} tok)\n",
222            rf.path,
223            rf.relationship.label(),
224            rf.token_count
225        ));
226    }
227
228    result.push_str("--- END GRAPH CONTEXT ---");
229    result
230}
231
232#[cfg(test)]
233mod tests {
234    use super::*;
235
236    #[test]
237    fn relationship_priorities() {
238        assert!(
239            Relationship::DirectDependency.priority() < Relationship::DirectDependent.priority()
240        );
241        assert!(
242            Relationship::DirectDependent.priority()
243                < Relationship::TransitiveDependency.priority()
244        );
245    }
246
247    #[test]
248    fn relationship_labels() {
249        assert_eq!(Relationship::DirectDependency.label(), "imports");
250        assert_eq!(Relationship::DirectDependent.label(), "imported-by");
251        assert_eq!(Relationship::TransitiveDependency.label(), "transitive-dep");
252        assert_eq!(Relationship::TypeProvider.label(), "type-provider");
253    }
254
255    #[test]
256    fn format_empty_context() {
257        let ctx = GraphContext {
258            source: GraphProviderSource::GraphIndex,
259            primary_file: "main.rs".to_string(),
260            related_files: vec![],
261            total_tokens: 100,
262            budget_remaining: 7900,
263        };
264        assert!(format_graph_context(&ctx).is_empty());
265    }
266
267    #[test]
268    fn format_with_related() {
269        let ctx = GraphContext {
270            source: GraphProviderSource::GraphIndex,
271            primary_file: "main.rs".to_string(),
272            related_files: vec![
273                RelatedFile {
274                    path: "lib.rs".to_string(),
275                    relationship: Relationship::DirectDependency,
276                    token_count: 500,
277                },
278                RelatedFile {
279                    path: "utils.rs".to_string(),
280                    relationship: Relationship::DirectDependent,
281                    token_count: 300,
282                },
283            ],
284            total_tokens: 900,
285            budget_remaining: 7100,
286        };
287        let output = format_graph_context(&ctx);
288        assert!(output.contains("2 related files"));
289        assert!(output.contains("lib.rs [imports]"));
290        assert!(output.contains("utils.rs [imported-by]"));
291    }
292
293    #[test]
294    fn nonexistent_root_returns_none() {
295        let result = build_graph_context("/nonexistent/file.rs", "/nonexistent", None);
296        assert!(result.is_none());
297    }
298}
lean_ctx/core/graph_context.rs

lean_ctx/core/
graph_context.rs
