Skip to main content

lean_ctx/tools/
autonomy.rs

1use std::sync::atomic::{AtomicBool, Ordering};
2
3use crate::core::cache::SessionCache;
4use crate::core::config::AutonomyConfig;
5use crate::core::graph_index::ProjectIndex;
6use crate::core::protocol;
7use crate::core::tokens::count_tokens;
8use crate::tools::CrpMode;
9
10pub struct AutonomyState {
11    pub session_initialized: AtomicBool,
12    pub dedup_applied: AtomicBool,
13    pub config: AutonomyConfig,
14}
15
16impl Default for AutonomyState {
17    fn default() -> Self {
18        Self::new()
19    }
20}
21
22impl AutonomyState {
23    pub fn new() -> Self {
24        Self {
25            session_initialized: AtomicBool::new(false),
26            dedup_applied: AtomicBool::new(false),
27            config: AutonomyConfig::load(),
28        }
29    }
30
31    pub fn is_enabled(&self) -> bool {
32        self.config.enabled
33    }
34}
35
36pub fn session_lifecycle_pre_hook(
37    state: &AutonomyState,
38    tool_name: &str,
39    cache: &mut SessionCache,
40    task: Option<&str>,
41    project_root: Option<&str>,
42    crp_mode: CrpMode,
43) -> Option<String> {
44    if !state.is_enabled() || !state.config.auto_preload {
45        return None;
46    }
47
48    if tool_name == "ctx_overview" || tool_name == "ctx_preload" {
49        return None;
50    }
51
52    if state
53        .session_initialized
54        .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
55        .is_err()
56    {
57        return None;
58    }
59
60    let root = project_root
61        .map(|s| s.to_string())
62        .or_else(|| {
63            std::env::current_dir()
64                .ok()
65                .map(|p| p.to_string_lossy().to_string())
66        })
67        .unwrap_or_else(|| ".".to_string());
68
69    let result = if let Some(task_desc) = task {
70        crate::tools::ctx_preload::handle(cache, task_desc, Some(&root), crp_mode)
71    } else {
72        let cache_readonly = &*cache;
73        crate::tools::ctx_overview::handle(cache_readonly, None, Some(&root), crp_mode)
74    };
75
76    if result.contains("No directly relevant files") || result.trim().is_empty() {
77        return None;
78    }
79
80    Some(format!(
81        "--- AUTO CONTEXT ---\n{result}\n--- END AUTO CONTEXT ---"
82    ))
83}
84
85pub fn enrich_after_read(
86    state: &AutonomyState,
87    cache: &mut SessionCache,
88    file_path: &str,
89    project_root: Option<&str>,
90) -> EnrichResult {
91    let mut result = EnrichResult::default();
92
93    if !state.is_enabled() {
94        return result;
95    }
96
97    let root = project_root
98        .map(|s| s.to_string())
99        .or_else(|| {
100            std::env::current_dir()
101                .ok()
102                .map(|p| p.to_string_lossy().to_string())
103        })
104        .unwrap_or_else(|| ".".to_string());
105
106    let index = match ProjectIndex::load(&root) {
107        Some(idx) => idx,
108        None => return result,
109    };
110
111    if state.config.auto_related {
112        result.related_hint = build_related_hints(cache, file_path, &index);
113    }
114
115    if state.config.silent_preload {
116        silent_preload_imports(cache, file_path, &index, &root);
117    }
118
119    result
120}
121
122#[derive(Default)]
123pub struct EnrichResult {
124    pub related_hint: Option<String>,
125}
126
127fn build_related_hints(
128    cache: &SessionCache,
129    file_path: &str,
130    index: &ProjectIndex,
131) -> Option<String> {
132    let related: Vec<_> = index
133        .edges
134        .iter()
135        .filter(|e| e.from == file_path || e.to == file_path)
136        .map(|e| if e.from == file_path { &e.to } else { &e.from })
137        .filter(|path| cache.get(path).is_none())
138        .take(3)
139        .collect();
140
141    if related.is_empty() {
142        return None;
143    }
144
145    let hints: Vec<String> = related.iter().map(|p| protocol::shorten_path(p)).collect();
146
147    Some(format!("[related: {}]", hints.join(", ")))
148}
149
150fn silent_preload_imports(
151    cache: &mut SessionCache,
152    file_path: &str,
153    index: &ProjectIndex,
154    _project_root: &str,
155) {
156    let imports: Vec<String> = index
157        .edges
158        .iter()
159        .filter(|e| e.from == file_path)
160        .map(|e| e.to.clone())
161        .filter(|path| cache.get(path).is_none())
162        .take(2)
163        .collect();
164
165    for path in imports {
166        if let Ok(content) = std::fs::read_to_string(&path) {
167            let tokens = count_tokens(&content);
168            if tokens < 5000 {
169                cache.store(&path, content);
170            }
171        }
172    }
173}
174
175pub fn maybe_auto_dedup(state: &AutonomyState, cache: &mut SessionCache) {
176    if !state.is_enabled() || !state.config.auto_dedup {
177        return;
178    }
179
180    if state
181        .dedup_applied
182        .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
183        .is_err()
184    {
185        return;
186    }
187
188    let entries = cache.get_all_entries();
189    if entries.len() < state.config.dedup_threshold {
190        state.dedup_applied.store(false, Ordering::SeqCst);
191        return;
192    }
193
194    crate::tools::ctx_dedup::handle_action(cache, "apply");
195}
196
197pub fn shell_efficiency_hint(
198    state: &AutonomyState,
199    command: &str,
200    input_tokens: usize,
201    output_tokens: usize,
202) -> Option<String> {
203    if !state.is_enabled() {
204        return None;
205    }
206
207    if input_tokens == 0 {
208        return None;
209    }
210
211    let savings_pct = ((input_tokens - output_tokens) as f64 / input_tokens as f64) * 100.0;
212    if savings_pct >= 20.0 {
213        return None;
214    }
215
216    let cmd_lower = command.to_lowercase();
217    if cmd_lower.starts_with("grep ")
218        || cmd_lower.starts_with("rg ")
219        || cmd_lower.starts_with("find ")
220        || cmd_lower.starts_with("ag ")
221    {
222        return Some("[hint: ctx_search is more token-efficient for code search]".to_string());
223    }
224
225    if cmd_lower.starts_with("cat ") || cmd_lower.starts_with("head ") {
226        return Some("[hint: ctx_read provides cached, compressed file access]".to_string());
227    }
228
229    None
230}
231
232#[cfg(test)]
233mod tests {
234    use super::*;
235
236    #[test]
237    fn autonomy_state_starts_uninitialized() {
238        let state = AutonomyState::new();
239        assert!(!state.session_initialized.load(Ordering::SeqCst));
240        assert!(!state.dedup_applied.load(Ordering::SeqCst));
241    }
242
243    #[test]
244    fn session_initialized_fires_once() {
245        let state = AutonomyState::new();
246        let first = state.session_initialized.compare_exchange(
247            false,
248            true,
249            Ordering::SeqCst,
250            Ordering::SeqCst,
251        );
252        assert!(first.is_ok());
253        let second = state.session_initialized.compare_exchange(
254            false,
255            true,
256            Ordering::SeqCst,
257            Ordering::SeqCst,
258        );
259        assert!(second.is_err());
260    }
261
262    #[test]
263    fn shell_hint_for_grep() {
264        let state = AutonomyState::new();
265        let hint = shell_efficiency_hint(&state, "grep -rn foo .", 100, 95);
266        assert!(hint.is_some());
267        assert!(hint.unwrap().contains("ctx_search"));
268    }
269
270    #[test]
271    fn shell_hint_none_when_good_savings() {
272        let state = AutonomyState::new();
273        let hint = shell_efficiency_hint(&state, "grep -rn foo .", 100, 50);
274        assert!(hint.is_none());
275    }
276
277    #[test]
278    fn shell_hint_none_for_unknown_command() {
279        let state = AutonomyState::new();
280        let hint = shell_efficiency_hint(&state, "cargo build", 100, 95);
281        assert!(hint.is_none());
282    }
283
284    #[test]
285    fn disabled_state_blocks_all() {
286        let mut state = AutonomyState::new();
287        state.config.enabled = false;
288        assert!(!state.is_enabled());
289        let hint = shell_efficiency_hint(&state, "grep foo", 100, 95);
290        assert!(hint.is_none());
291    }
292}