Skip to main content

lean_ctx/core/
context_deficit.rs

1use super::context_ledger::{ContextLedger, PressureAction};
2use super::intent_engine::{IntentScope, StructuredIntent, TaskType};
3
4#[derive(Debug, Clone)]
5pub struct ContextDeficit {
6    pub missing_targets: Vec<String>,
7    pub suggested_files: Vec<SuggestedFile>,
8    pub pressure_action: PressureAction,
9    pub budget_remaining: usize,
10}
11
12#[derive(Debug, Clone)]
13pub struct SuggestedFile {
14    pub path: String,
15    pub reason: DeficitReason,
16    pub estimated_tokens: usize,
17    pub recommended_mode: String,
18}
19
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub enum DeficitReason {
22    TargetNotLoaded,
23    DependencyOfTarget,
24    TestFileForTarget,
25    ConfigForTask,
26}
27
28impl DeficitReason {
29    fn priority(&self) -> u8 {
30        match self {
31            Self::TargetNotLoaded => 0,
32            Self::DependencyOfTarget => 1,
33            Self::TestFileForTarget => 2,
34            Self::ConfigForTask => 3,
35        }
36    }
37}
38
39pub fn detect_deficit(
40    ledger: &ContextLedger,
41    intent: &StructuredIntent,
42    known_files: &[String],
43) -> ContextDeficit {
44    let loaded_paths: Vec<&str> = ledger.entries.iter().map(|e| e.path.as_str()).collect();
45    let pressure = ledger.pressure();
46
47    let mut missing_targets = Vec::new();
48    let mut suggestions = Vec::new();
49
50    for target in &intent.targets {
51        if target.contains('.') || target.contains('/') {
52            let is_loaded = loaded_paths
53                .iter()
54                .any(|p| p.ends_with(target) || p.contains(target));
55            if !is_loaded {
56                missing_targets.push(target.clone());
57
58                let matching: Vec<&String> = known_files
59                    .iter()
60                    .filter(|f| f.ends_with(target) || f.contains(target))
61                    .collect();
62
63                for file in matching {
64                    let mode = mode_for_pressure(&pressure.recommendation, &intent.scope);
65                    suggestions.push(SuggestedFile {
66                        path: file.clone(),
67                        reason: DeficitReason::TargetNotLoaded,
68                        estimated_tokens: estimate_tokens_for_mode(&mode),
69                        recommended_mode: mode,
70                    });
71                }
72            }
73        }
74    }
75
76    if intent.task_type == TaskType::FixBug || intent.task_type == TaskType::Test {
77        for target in &intent.targets {
78            if target.contains('.') || target.contains('/') {
79                let test_patterns = derive_test_paths(target);
80                for test_path in &test_patterns {
81                    let matching: Vec<&String> = known_files
82                        .iter()
83                        .filter(|f| f.contains(test_path))
84                        .collect();
85                    for file in matching {
86                        if !loaded_paths.contains(&file.as_str())
87                            && !suggestions.iter().any(|s| s.path == *file)
88                        {
89                            let mode = mode_for_pressure(&pressure.recommendation, &intent.scope);
90                            suggestions.push(SuggestedFile {
91                                path: file.clone(),
92                                reason: DeficitReason::TestFileForTarget,
93                                estimated_tokens: estimate_tokens_for_mode(&mode),
94                                recommended_mode: mode,
95                            });
96                        }
97                    }
98                }
99            }
100        }
101    }
102
103    if intent.task_type == TaskType::Config || intent.task_type == TaskType::Deploy {
104        let config_patterns = [
105            "Cargo.toml",
106            "package.json",
107            "tsconfig.json",
108            "pyproject.toml",
109            ".env",
110            "Dockerfile",
111        ];
112        for pattern in &config_patterns {
113            let matching: Vec<&String> = known_files
114                .iter()
115                .filter(|f| f.ends_with(pattern))
116                .collect();
117            for file in matching {
118                if !loaded_paths.contains(&file.as_str())
119                    && !suggestions.iter().any(|s| s.path == *file)
120                {
121                    suggestions.push(SuggestedFile {
122                        path: file.clone(),
123                        reason: DeficitReason::ConfigForTask,
124                        estimated_tokens: estimate_tokens_for_mode("full"),
125                        recommended_mode: "full".to_string(),
126                    });
127                }
128            }
129        }
130    }
131
132    suggestions.sort_by_key(|s| s.reason.priority());
133
134    let budget_remaining = pressure.remaining_tokens;
135    let mut cumulative = 0usize;
136    suggestions.retain(|s| {
137        cumulative += s.estimated_tokens;
138        cumulative <= budget_remaining
139    });
140
141    ContextDeficit {
142        missing_targets,
143        suggested_files: suggestions,
144        pressure_action: pressure.recommendation,
145        budget_remaining,
146    }
147}
148
149fn mode_for_pressure(action: &PressureAction, scope: &IntentScope) -> String {
150    match action {
151        PressureAction::EvictLeastRelevant => "reference".to_string(),
152        PressureAction::ForceCompression => "signatures".to_string(),
153        PressureAction::SuggestCompression => match scope {
154            IntentScope::SingleFile => "full".to_string(),
155            IntentScope::MultiFile => "signatures".to_string(),
156            IntentScope::CrossModule | IntentScope::ProjectWide => "map".to_string(),
157        },
158        PressureAction::NoAction => match scope {
159            IntentScope::SingleFile | IntentScope::MultiFile => "full".to_string(),
160            IntentScope::CrossModule => "signatures".to_string(),
161            IntentScope::ProjectWide => "map".to_string(),
162        },
163    }
164}
165
166fn estimate_tokens_for_mode(mode: &str) -> usize {
167    match mode {
168        "full" => 2000,
169        "signatures" => 400,
170        "map" => 200,
171        "reference" => 50,
172        "aggressive" => 800,
173        _ => 1000,
174    }
175}
176
177fn derive_test_paths(file_path: &str) -> Vec<String> {
178    let stem = std::path::Path::new(file_path)
179        .file_stem()
180        .and_then(|s| s.to_str())
181        .unwrap_or("");
182    if stem.is_empty() {
183        return Vec::new();
184    }
185    vec![
186        format!("{stem}_test"),
187        format!("test_{stem}"),
188        format!("{stem}.test"),
189        format!("{stem}.spec"),
190        format!("{stem}_spec"),
191    ]
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197
198    #[test]
199    fn detects_missing_target() {
200        let ledger = ContextLedger::new();
201        let intent = StructuredIntent::from_query("fix bug in auth.rs");
202        let known = vec!["src/auth.rs".to_string(), "src/db.rs".to_string()];
203        let deficit = detect_deficit(&ledger, &intent, &known);
204        assert!(!deficit.missing_targets.is_empty());
205        assert!(deficit
206            .suggested_files
207            .iter()
208            .any(|s| s.path.contains("auth")));
209    }
210
211    #[test]
212    fn no_deficit_when_loaded() {
213        let mut ledger = ContextLedger::new();
214        ledger.record("src/auth.rs", "full", 500, 500);
215        let intent = StructuredIntent::from_query("fix bug in auth.rs");
216        let known = vec!["src/auth.rs".to_string()];
217        let deficit = detect_deficit(&ledger, &intent, &known);
218        assert!(deficit.missing_targets.is_empty());
219    }
220
221    #[test]
222    fn suggests_test_files_for_fixbug() {
223        let mut ledger = ContextLedger::new();
224        ledger.record("src/auth.rs", "full", 500, 500);
225        let intent = StructuredIntent::from_query("fix bug in auth.rs");
226        let known = vec!["src/auth.rs".to_string(), "tests/auth_test.rs".to_string()];
227        let deficit = detect_deficit(&ledger, &intent, &known);
228        let test_suggestions: Vec<_> = deficit
229            .suggested_files
230            .iter()
231            .filter(|s| s.reason == DeficitReason::TestFileForTarget)
232            .collect();
233        assert!(
234            !test_suggestions.is_empty(),
235            "should suggest test files for FixBug"
236        );
237    }
238
239    #[test]
240    fn respects_budget() {
241        let mut ledger = ContextLedger::with_window_size(1000);
242        ledger.record("existing.rs", "full", 900, 900);
243        let intent = StructuredIntent::from_query("fix bug in big_file.rs");
244        let known = vec!["src/big_file.rs".to_string()];
245        let deficit = detect_deficit(&ledger, &intent, &known);
246        assert!(
247            deficit.suggested_files.is_empty() || deficit.budget_remaining < 200,
248            "should respect budget constraints"
249        );
250    }
251
252    #[test]
253    fn config_task_suggests_config_files() {
254        let ledger = ContextLedger::new();
255        let intent = StructuredIntent::from_query("configure env settings for the project");
256        let known = vec![
257            "src/main.rs".to_string(),
258            "Cargo.toml".to_string(),
259            "package.json".to_string(),
260        ];
261        let deficit = detect_deficit(&ledger, &intent, &known);
262        let config_suggestions: Vec<_> = deficit
263            .suggested_files
264            .iter()
265            .filter(|s| s.reason == DeficitReason::ConfigForTask)
266            .collect();
267        assert!(!config_suggestions.is_empty());
268    }
269
270    #[test]
271    fn mode_adapts_to_pressure() {
272        let mode_low = mode_for_pressure(&PressureAction::NoAction, &IntentScope::SingleFile);
273        let mode_high = mode_for_pressure(
274            &PressureAction::EvictLeastRelevant,
275            &IntentScope::SingleFile,
276        );
277        assert_eq!(mode_low, "full");
278        assert_eq!(mode_high, "reference");
279    }
280
281    #[test]
282    fn derive_test_paths_generates_variants() {
283        let paths = derive_test_paths("src/auth.rs");
284        assert!(paths.contains(&"auth_test".to_string()));
285        assert!(paths.contains(&"test_auth".to_string()));
286        assert!(paths.contains(&"auth.test".to_string()));
287        assert!(paths.contains(&"auth.spec".to_string()));
288    }
289}