1use crate::tools::warm_guard::WarmGuard;
7use async_trait::async_trait;
8use limit_agent::AgentError;
9use limit_agent::Tool;
10use limit_tldr::{Config as TldrConfig, Language, TLDR};
11use serde::{Deserialize, Serialize};
12use serde_json::{json, Value};
13use std::path::{Path, PathBuf};
14use std::sync::atomic::{AtomicBool, Ordering};
15use std::sync::Arc;
16use tokio::sync::{Notify, OnceCell};
17use tracing::{debug, info, warn};
18
19#[derive(Debug, Clone, Serialize, Deserialize)]
21#[serde(rename_all = "snake_case")]
22pub enum AnalysisType {
23 Context,
25 Source,
27 Impact,
29 Cfg,
31 Dfg,
33 DeadCode,
35 Architecture,
37 Search,
39}
40
41#[derive(Debug, Clone, Serialize, Deserialize)]
43pub struct TldrParams {
44 pub analysis_type: AnalysisType,
46
47 pub function: Option<String>,
49
50 pub file: Option<String>,
52
53 #[serde(default = "default_depth")]
55 pub depth: usize,
56
57 #[serde(default = "default_entries")]
59 pub entries: Vec<String>,
60
61 pub query: Option<String>,
63
64 #[serde(default = "default_limit")]
66 pub limit: usize,
67
68 pub project_path: Option<String>,
70}
71
72fn default_depth() -> usize {
73 2
74}
75fn default_entries() -> Vec<String> {
76 vec!["main".to_string()]
77}
78fn default_limit() -> usize {
79 10
80}
81
82pub struct TldrTool {
84 cache: Arc<OnceCell<(PathBuf, Arc<TLDR>)>>,
86 default_project: PathBuf,
88 warm_notify: Arc<Notify>,
90 warm_started: Arc<AtomicBool>,
92}
93
94impl TldrTool {
95 pub fn new() -> Self {
97 let default_project = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."));
98 Self {
99 cache: Arc::new(OnceCell::new()),
100 default_project,
101 warm_notify: Arc::new(Notify::new()),
102 warm_started: Arc::new(AtomicBool::new(false)),
103 }
104 }
105
106 pub fn with_project<P: Into<PathBuf>>(project: P) -> Self {
108 Self {
109 cache: Arc::new(OnceCell::new()),
110 default_project: project.into(),
111 warm_notify: Arc::new(Notify::new()),
112 warm_started: Arc::new(AtomicBool::new(false)),
113 }
114 }
115
116 fn ensure_pre_warm_started(&self) {
118 if self
119 .warm_started
120 .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
121 .is_ok()
122 {
123 let project = self.default_project.clone();
124 let cache = Arc::clone(&self.cache);
125 let notify = Arc::clone(&self.warm_notify);
126
127 tokio::spawn(async move {
128 Self::pre_warm(project, cache, notify).await;
129 });
130 }
131 }
132
133 async fn pre_warm(
135 project_path: PathBuf,
136 cache: Arc<OnceCell<(PathBuf, Arc<TLDR>)>>,
137 notify: Arc<Notify>,
138 ) {
139 let cache_dir = match Self::get_cache_dir(&project_path) {
140 Ok(dir) => dir,
141 Err(e) => {
142 warn!("pre_warm: failed to get cache dir: {}", e);
143 notify.notify_waiters();
144 return;
145 }
146 };
147
148 let guard = WarmGuard::new(&cache_dir);
149 if guard.is_fresh(&project_path) {
150 info!("pre_warm: cache files fresh, loading without re-warming");
151 let config = TldrConfig {
155 language: Language::Auto,
156 max_depth: 3,
157 cache_dir: Some(cache_dir),
158 };
159 match TLDR::new(&project_path, config).await {
160 Ok(mut tldr) => match tldr.warm().await {
161 Ok(()) => {
162 let _ = cache.set((project_path, Arc::new(tldr))).map_err(|_| {
163 debug!("pre_warm: OnceCell already set (race with get_tldr)");
164 });
165 info!("pre_warm: warm from cache complete");
166 }
167 Err(e) => warn!("pre_warm: warm from cache failed: {}", e),
168 },
169 Err(e) => warn!("pre_warm: TLDR::new failed: {}", e),
170 }
171 notify.notify_waiters();
172 return;
173 }
174
175 info!("pre_warm: warming TLDR for {:?}", project_path);
176 let config = TldrConfig {
177 language: Language::Auto,
178 max_depth: 3,
179 cache_dir: Some(cache_dir),
180 };
181
182 match TLDR::new(&project_path, config).await {
183 Ok(mut tldr) => match tldr.warm().await {
184 Ok(()) => {
185 guard.save(&project_path);
186 info!("pre_warm: warm complete");
187 let _ = cache.set((project_path, Arc::new(tldr))).map_err(|_| {
188 debug!("pre_warm: OnceCell already set (race with get_tldr)");
189 });
190 notify.notify_waiters();
191 }
192 Err(e) => warn!("pre_warm: warm failed: {}", e),
193 },
194 Err(e) => warn!("pre_warm: TLDR::new failed: {}", e),
195 }
196 notify.notify_waiters();
197 }
198
199 async fn get_tldr(&self, project_path: &Path) -> Result<Arc<TLDR>, AgentError> {
204 let project_path = project_path.to_path_buf();
205 let project_path_for_check = project_path.clone();
206
207 self.ensure_pre_warm_started();
209
210 if let Some((cached_path, tldr)) = self.cache.get() {
212 if *cached_path == project_path_for_check {
213 debug!("TLDR cache hit for project: {:?}", project_path_for_check);
214 return Ok(Arc::clone(tldr));
215 }
216 warn!(
217 "get_tldr: ignoring project_path {:?}, using cached {:?}",
218 project_path_for_check, cached_path
219 );
220 return Ok(Arc::clone(tldr));
221 }
222
223 info!("get_tldr: waiting for pre_warm...");
225 tokio::select! {
226 _ = self.warm_notify.notified() => {
227 if let Some((cached_path, tldr)) = self.cache.get() {
229 if *cached_path == project_path_for_check {
230 debug!("TLDR cache hit after pre_warm for: {:?}", project_path_for_check);
231 return Ok(Arc::clone(tldr));
232 }
233 warn!(
234 "get_tldr: ignoring project_path {:?}, using cached {:?}",
235 project_path_for_check, cached_path
236 );
237 return Ok(Arc::clone(tldr));
238 }
239 warn!("get_tldr: pre_warm did not populate cache, falling back to lazy");
241 }
242 _ = tokio::time::sleep(std::time::Duration::from_secs(30)) => {
243 warn!("get_tldr: pre_warm timed out after 30s");
244 if let Some((cached_path, tldr)) = self.cache.get() {
246 if *cached_path == project_path_for_check {
247 info!("get_tldr: pre_warm completed during timeout, using cached result");
248 return Ok(Arc::clone(tldr));
249 }
250 warn!(
251 "get_tldr: ignoring project_path {:?}, using cached {:?}",
252 project_path_for_check, cached_path
253 );
254 return Ok(Arc::clone(tldr));
255 }
256 info!("get_tldr: falling back to lazy creation");
257 }
258 }
259
260 let cache = Arc::clone(&self.cache);
262 let result: Result<&(PathBuf, Arc<TLDR>), AgentError> = cache
263 .get_or_try_init(|| async {
264 info!(
265 "Lazy creating TLDR instance for project: {:?}",
266 project_path
267 );
268 let config = TldrConfig {
269 language: Language::Auto,
270 max_depth: 3,
271 cache_dir: Some(Self::get_cache_dir(&project_path)?),
272 };
273
274 let mut tldr = TLDR::new(&project_path, config)
275 .await
276 .map_err(|e| AgentError::ToolError(format!("Failed to create TLDR: {}", e)))?;
277
278 info!("Warming TLDR indexes...");
279 tldr.warm()
280 .await
281 .map_err(|e| AgentError::ToolError(format!("Failed to warm TLDR: {}", e)))?;
282
283 Ok((project_path, Arc::new(tldr)))
284 })
285 .await;
286
287 let (_cached_path, tldr) = result?;
288 debug!(
289 "TLDR cache hit (lazy) for project: {:?}",
290 project_path_for_check
291 );
292 Ok(Arc::clone(tldr))
293 }
294
295 fn get_cache_dir(project_path: &Path) -> Result<PathBuf, AgentError> {
297 let home = dirs::home_dir()
298 .ok_or_else(|| AgentError::ToolError("Cannot find home directory".into()))?;
299
300 let project_id = project_path
302 .canonicalize()
303 .map_err(|e| AgentError::ToolError(format!("Cannot canonicalize path: {}", e)))?
304 .to_string_lossy()
305 .to_string();
306
307 use std::collections::hash_map::DefaultHasher;
309 use std::hash::{Hash, Hasher};
310 let mut hasher = DefaultHasher::new();
311 project_id.hash(&mut hasher);
312 let hash = format!("{:x}", hasher.finish());
313
314 Ok(home
315 .join(".limit")
316 .join("projects")
317 .join(&hash)
318 .join("tldr"))
319 }
320
321 async fn build_source_result(
323 &self,
324 function: &str,
325 source_file: PathBuf,
326 start_line: usize,
327 end_line: usize,
328 project_path: &Path,
329 ) -> Result<Value, AgentError> {
330 let relative_file = source_file
331 .strip_prefix(project_path)
332 .unwrap_or(&source_file)
333 .to_path_buf();
334
335 let file_path = project_path.join(&source_file);
336 let source = tokio::fs::read_to_string(&file_path)
337 .await
338 .map_err(|e| AgentError::ToolError(format!("Failed to read file: {}", e)))?;
339
340 let lines: Vec<&str> = source.lines().collect();
341 let start = start_line.saturating_sub(1);
342 let end = end_line.min(lines.len());
343 let max_lines = 80;
344 let truncated = (end - start) > max_lines;
345 let actual_end = if truncated { start + max_lines } else { end };
346
347 let function_source = lines[start..actual_end].join("\n");
348
349 let mut result = json!({
350 "type": "source",
351 "function": function,
352 "file": relative_file.display().to_string(),
353 "line": start_line,
354 "end_line": actual_end,
355 "source": function_source
356 });
357 if truncated {
358 result["truncated"] = json!(true);
359 result["total_lines"] = json!(end - start);
360 }
361
362 Ok(result)
363 }
364
365 async fn analyze(&self, params: TldrParams) -> Result<Value, AgentError> {
367 let project_path = params
368 .project_path
369 .map(PathBuf::from)
370 .unwrap_or_else(|| self.default_project.clone());
371
372 let tldr = self.get_tldr(&project_path).await?;
373
374 let result = match params.analysis_type {
375 AnalysisType::Context => {
376 let function = params.function.ok_or_else(|| {
377 AgentError::ToolError("function parameter required for context analysis".into())
378 })?;
379
380 let context = tldr
381 .get_context(&function, params.depth)
382 .await
383 .map_err(|e| {
384 AgentError::ToolError(format!("Context analysis failed: {}", e))
385 })?;
386
387 Ok(json!({
388 "type": "context",
389 "function": function,
390 "depth": params.depth,
391 "context": context
392 }))
393 }
394
395 AnalysisType::Source => {
396 let function = params.function.ok_or_else(|| {
397 AgentError::ToolError("function parameter required for source analysis".into())
398 })?;
399
400 let (function, file_override) = if !function.starts_with("struct ") {
403 if let Some(pos) = function.find("::") {
404 let class_name = &function[..pos];
405 let method_name = &function[pos + 2..];
406 if !method_name.is_empty() {
407 let class_info = if let Some(ref file) = params.file {
408 let file_path = project_path.join(file);
409 tldr.find_class_in(class_name, &file_path).unwrap_or(None)
410 } else {
411 tldr.find_class(class_name).unwrap_or(None)
412 };
413 if let Some(info) = class_info {
414 let resolved_file = info
415 .file
416 .strip_prefix(&project_path)
417 .unwrap_or(&info.file)
418 .to_string_lossy()
419 .to_string();
420 (method_name.to_string(), Some(resolved_file))
421 } else {
422 (function, None)
423 }
424 } else {
425 (function, None)
426 }
427 } else {
428 (function, None)
429 }
430 } else {
431 (function, None)
432 };
433 let effective_file = file_override.or(params.file.clone());
435
436 let is_struct = function.starts_with("struct ");
437 let lookup_name = if is_struct {
438 function.strip_prefix("struct ").unwrap()
439 } else {
440 &function
441 };
442
443 let (source_file, start_line, end_line) = if is_struct {
444 let class_info = if let Some(ref file) = effective_file {
446 let file_path = project_path.join(file);
447 tldr.find_class_in(lookup_name, &file_path)
448 .map_err(|e| {
449 AgentError::ToolError(format!("Source analysis failed: {}", e))
450 })?
451 .ok_or_else(|| {
452 AgentError::ToolError(format!(
453 "Struct '{}' not found in '{}'",
454 lookup_name, file
455 ))
456 })?
457 } else {
458 tldr.find_class(lookup_name)
459 .map_err(|e| {
460 AgentError::ToolError(format!("Source analysis failed: {}", e))
461 })?
462 .ok_or_else(|| {
463 AgentError::ToolError(format!("Struct not found: {}", lookup_name))
464 })?
465 };
466 (class_info.file, class_info.line, class_info.end_line)
467 } else {
468 let func_info = if let Some(ref file) = effective_file {
470 let file_path = project_path.join(file);
471 if let Some(func) =
473 tldr.find_function_in(&function, &file_path).map_err(|e| {
474 AgentError::ToolError(format!("Source analysis failed: {}", e))
475 })?
476 {
477 func
478 } else if let Some(cls) =
479 tldr.find_class_in(&function, &file_path).map_err(|e| {
480 AgentError::ToolError(format!("Source analysis failed: {}", e))
481 })?
482 {
483 return self
485 .build_source_result(
486 &function,
487 cls.file,
488 cls.line,
489 cls.end_line,
490 &project_path,
491 )
492 .await;
493 } else {
494 return Err(AgentError::ToolError(format!(
495 "Function or struct '{}' not found in '{}'",
496 function, file
497 )));
498 }
499 } else {
500 let all_matches = tldr.find_all_functions(&function);
502 if all_matches.len() > 1 {
503 let match_list: Vec<String> = all_matches
504 .iter()
505 .take(5)
506 .map(|f| {
507 let relative =
508 f.file.strip_prefix(&project_path).unwrap_or(&f.file);
509 format!("{} ({}:{})", f.name, relative.display(), f.line)
510 })
511 .collect();
512 return Ok(json!({
513 "type": "disambiguation_needed",
514 "function": function,
515 "match_count": all_matches.len(),
516 "matches": match_list,
517 "hint": format!(
518 "Use file parameter to disambiguate, e.g.: {{\"analysis_type\": \"source\", \"function\": \"{}\", \"file\": \"path/to/file.rs\"}}",
519 function
520 )
521 }));
522 }
523 tldr.find_function(&function)
524 .await
525 .map_err(|e| {
526 AgentError::ToolError(format!("Source analysis failed: {}", e))
527 })?
528 .ok_or_else(|| {
529 AgentError::ToolError(format!("Function not found: {}", function))
530 })?
531 };
532 (func_info.file, func_info.line, func_info.end_line)
533 };
534
535 self.build_source_result(
536 &function,
537 source_file,
538 start_line,
539 end_line,
540 &project_path,
541 )
542 .await
543 }
544
545 AnalysisType::Impact => {
546 let function = params.function.ok_or_else(|| {
547 AgentError::ToolError("function parameter required for impact analysis".into())
548 })?;
549
550 let callers = tldr
551 .get_impact(&function)
552 .map_err(|e| AgentError::ToolError(format!("Impact analysis failed: {}", e)))?;
553
554 Ok(json!({
555 "type": "impact",
556 "function": function,
557 "callers": callers.iter().map(|c| json!({
558 "function": c.function,
559 "file": c.file.display().to_string(),
560 "line": c.line
561 })).collect::<Vec<_>>(),
562 "caller_count": callers.len()
563 }))
564 }
565
566 AnalysisType::Cfg => {
567 let file = params.file.ok_or_else(|| {
568 AgentError::ToolError("file parameter required for CFG analysis".into())
569 })?;
570 let function = params.function.ok_or_else(|| {
571 AgentError::ToolError("function parameter required for CFG analysis".into())
572 })?;
573
574 let file_path = project_path.join(&file);
575 let cfg = tldr
576 .get_cfg(&file_path, &function)
577 .map_err(|e| AgentError::ToolError(format!("CFG analysis failed: {}", e)))?;
578
579 Ok(json!({
580 "type": "cfg",
581 "function": function,
582 "file": file,
583 "complexity": cfg.complexity,
584 "blocks": cfg.blocks.len()
585 }))
586 }
587
588 AnalysisType::Dfg => {
589 let file = params.file.ok_or_else(|| {
590 AgentError::ToolError("file parameter required for DFG analysis".into())
591 })?;
592 let function = params.function.ok_or_else(|| {
593 AgentError::ToolError("function parameter required for DFG analysis".into())
594 })?;
595
596 let file_path = project_path.join(&file);
597 let dfg = tldr
598 .get_dfg(&file_path, &function)
599 .map_err(|e| AgentError::ToolError(format!("DFG analysis failed: {}", e)))?;
600
601 Ok(json!({
602 "type": "dfg",
603 "function": function,
604 "file": file,
605 "variables": dfg.variables,
606 "flows": dfg.flows.len()
607 }))
608 }
609
610 AnalysisType::DeadCode => {
611 let entries: Vec<&str> = params.entries.iter().map(|s| s.as_str()).collect();
612 let dead = tldr.find_dead_code(&entries).map_err(|e| {
613 AgentError::ToolError(format!("Dead code analysis failed: {}", e))
614 })?;
615
616 Ok(json!({
617 "type": "dead_code",
618 "entries": params.entries,
619 "dead_functions": dead.iter().map(|f| json!({
620 "name": f.name,
621 "file": f.file.display().to_string(),
622 "line": f.line
623 })).collect::<Vec<_>>(),
624 "dead_count": dead.len()
625 }))
626 }
627
628 AnalysisType::Architecture => {
629 let arch = tldr.detect_architecture().map_err(|e| {
630 AgentError::ToolError(format!("Architecture detection failed: {}", e))
631 })?;
632
633 let entry_sample: Vec<_> = arch.entry.iter().take(10).collect();
636 let middle_sample: Vec<_> = arch.middle.iter().take(10).collect();
637 let leaf_sample: Vec<_> = arch.leaf.iter().take(10).collect();
638
639 Ok(json!({
640 "type": "architecture",
641 "summary": {
642 "entry_points_count": arch.entry.len(),
643 "middle_layer_count": arch.middle.len(),
644 "leaf_functions_count": arch.leaf.len()
645 },
646 "sample_entry_points": entry_sample,
647 "sample_middle_layer": middle_sample,
648 "sample_leaf_functions": leaf_sample,
649 "note": "Showing top 10 of each category. Use Search analysis for specific functions."
650 }))
651 }
652
653 AnalysisType::Search => {
654 let query = params
655 .query
656 .unwrap_or_else(|| params.function.clone().unwrap_or_default());
657
658 let results = tldr
659 .semantic_search(&query, params.limit)
660 .await
661 .map_err(|e| AgentError::ToolError(format!("Search failed: {}", e)))?;
662
663 Ok(json!({
664 "type": "search",
665 "query": query,
666 "results": results.iter().map(|r| {
667 let relative = r
668 .file
669 .strip_prefix(&project_path)
670 .unwrap_or(&r.file);
671 json!({
672 "function": r.function,
673 "file": relative.display().to_string(),
674 "line": r.line,
675 "score": r.score,
676 "signature": r.signature
677 })
678 }).collect::<Vec<_>>()
679 }))
680 }
681 };
682
683 debug!("Analysis complete for: {:?}", params.analysis_type);
684 result
685 }
686}
687
688impl Default for TldrTool {
689 fn default() -> Self {
690 Self::new()
691 }
692}
693
694#[async_trait]
695impl Tool for TldrTool {
696 fn name(&self) -> &str {
697 "tldr_analyze"
698 }
699
700 async fn execute(&self, args: Value) -> Result<Value, AgentError> {
701 let params: TldrParams = serde_json::from_value(args)
703 .map_err(|e| AgentError::ToolError(format!("Invalid parameters: {}", e)))?;
704
705 info!("tldr_analyze invoked: type={:?}", params.analysis_type);
706 if let Some(ref f) = ¶ms.function {
707 debug!(" function: {}", f);
708 }
709 if let Some(ref q) = ¶ms.query {
710 debug!(" query: {}", q);
711 }
712
713 let result = match self.analyze(params).await {
714 Ok(r) => r,
715 Err(e) => {
716 tracing::warn!("tldr_analyze failed: {}", e);
717 return Err(e);
718 }
719 };
720 let result_str =
721 serde_json::to_string(&result).unwrap_or_else(|_| "serialize error".to_string());
722 info!(
723 "tldr_analyze result: {} chars, {} bytes",
724 result_str.chars().count(),
725 result_str.len()
726 );
727 Ok(result)
728 }
729}
730
731pub fn tldr_tool_definition() -> Value {
733 json!({
734 "name": "tldr_analyze",
735 "description": "Token-efficient code analysis. ALWAYS USE THIS when the user asks: 'what does X do', 'how does X work', 'explain X', 'tell me about X', 'what is X'. Saves 95% tokens vs reading raw code. Do NOT combine with file_read or bash — this tool provides all needed context. STRATEGY: (1) search to find functions, (2) source for 1-3 key functions only, (3) write answer. Do NOT read every function. Analysis types: search=find functions, context=dependencies, source=function code, impact=callers, architecture=layers.",
736 "parameters": {
737 "type": "object",
738 "properties": {
739 "analysis_type": {
740 "type": "string",
741 "enum": ["search", "context", "source", "impact", "cfg", "dfg", "dead_code", "architecture"],
742 "description": "Type: search=find by keyword, context=dependencies+callers, source=function code (use instead of file_read), impact=who calls this, cfg=control flow, dfg=data flow, dead_code=unreachable, architecture=module layers"
743 },
744 "function": {
745 "type": "string",
746 "description": "Function or struct name (required for context, source, impact, cfg, dfg). For structs, prefix with 'struct ' (e.g., 'struct AppConfig')"
747 },
748 "file": {
749 "type": "string",
750 "description": "File path relative to project root. Required for cfg, dfg. Optional for source (use to disambiguate when function name exists in multiple files)"
751 },
752 "depth": {
753 "type": "integer",
754 "description": "Depth for context traversal (default: 2)",
755 "default": 2
756 },
757 "entries": {
758 "type": "array",
759 "items": {"type": "string"},
760 "description": "Entry points for dead code detection (default: [\"main\"])",
761 "default": ["main"]
762 },
763 "query": {
764 "type": "string",
765 "description": "Search query for finding functions (supports patterns like 'daemon', 'auth', 'handle_*')"
766 },
767 "limit": {
768 "type": "integer",
769 "description": "Maximum results for search (default: 10)",
770 "default": 10
771 },
772 "project_path": {
773 "type": "string",
774 "description": "Project root directory (defaults to current directory). Do NOT use file paths here — use 'file' parameter for file paths."
775 }
776 },
777 "required": ["analysis_type"]
778 }
779 })
780}
781
782#[cfg(test)]
783mod tests {
784 use super::*;
785
786 #[test]
787 fn test_tool_definition() {
788 let def = tldr_tool_definition();
789 assert_eq!(def["name"], "tldr_analyze");
790 assert!(def["parameters"]["properties"]["analysis_type"]["enum"].is_array());
791 }
792
793 #[test]
794 fn test_params_deserialization() {
795 let json = json!({
796 "analysis_type": "context",
797 "function": "main",
798 "depth": 3
799 });
800
801 let params: TldrParams = serde_json::from_value(json).unwrap();
802 assert!(matches!(params.analysis_type, AnalysisType::Context));
803 assert_eq!(params.function, Some("main".to_string()));
804 assert_eq!(params.depth, 3);
805 }
806
807 #[tokio::test]
808 #[ignore = "requires fastembed model download — run with: cargo test -- --ignored test_cache_returns_cached_instance"]
809 async fn test_cache_returns_cached_instance() {
810 let tool = TldrTool::new();
811 let test_path = std::env::current_dir().unwrap();
812
813 let tldr1 = tool.get_tldr(&test_path).await.unwrap();
814 let tldr2 = tool.get_tldr(&test_path).await.unwrap();
815
816 let addr1 = Arc::as_ptr(&tldr1) as usize;
817 let addr2 = Arc::as_ptr(&tldr2) as usize;
818 assert_eq!(
819 addr1, addr2,
820 "Second call should return cached instance (same memory address)"
821 );
822 }
823}