Skip to main content

task_graph_mcp/resources/
docs.rs

1//! Documentation resources - expose markdown files from docs/ directory via MCP resources.
2//!
3//! Provides access to project documentation through docs:// URIs:
4//! - `docs://index` - Lists all available documentation files
5//! - `docs://search/{query}` - Full-text search across all documentation files
6//! - `docs://FILENAME.md` - Returns content of specific doc file
7//! - `docs://subdir/FILENAME.md` - Supports recursive subdirectories
8
9use anyhow::Result;
10use serde_json::{Value, json};
11use std::path::Path;
12
13/// Metadata for a documentation file.
14#[derive(Debug, Clone)]
15pub struct DocInfo {
16    /// Relative path from docs/ directory (e.g., "GATES.md" or "diagrams/README.md")
17    pub relative_path: String,
18    /// File name only
19    pub name: String,
20    /// Size in bytes
21    pub size: u64,
22}
23
24/// Recursively find all markdown files in a directory.
25fn find_markdown_files(dir: &Path, base: &Path, files: &mut Vec<DocInfo>) -> Result<()> {
26    if !dir.exists() || !dir.is_dir() {
27        return Ok(());
28    }
29
30    for entry in std::fs::read_dir(dir)? {
31        let entry = entry?;
32        let path = entry.path();
33
34        if path.is_dir() {
35            // Recurse into subdirectories
36            find_markdown_files(&path, base, files)?;
37        } else if path.is_file() {
38            // Check if it's a markdown file
39            if let Some(ext) = path.extension()
40                && (ext == "md" || ext == "markdown")
41            {
42                let relative = path
43                    .strip_prefix(base)
44                    .map(|p| p.to_string_lossy().to_string().replace('\\', "/"))
45                    .unwrap_or_else(|_| path.file_name().unwrap().to_string_lossy().to_string());
46
47                let name = path
48                    .file_name()
49                    .map(|n| n.to_string_lossy().to_string())
50                    .unwrap_or_default();
51
52                let size = entry.metadata().map(|m| m.len()).unwrap_or(0);
53
54                files.push(DocInfo {
55                    relative_path: relative,
56                    name,
57                    size,
58                });
59            }
60        }
61    }
62
63    Ok(())
64}
65
66/// Validate a doc path to prevent path traversal attacks.
67/// Only allows alphanumeric characters, hyphens, underscores, dots, and forward slashes.
68fn validate_doc_path(path: &str) -> Result<()> {
69    if path.is_empty() {
70        return Err(anyhow::anyhow!("Doc path cannot be empty"));
71    }
72
73    if path.len() > 256 {
74        return Err(anyhow::anyhow!("Doc path too long (max 256 chars)"));
75    }
76
77    // Check for path traversal attempts
78    if path.contains("..") {
79        return Err(anyhow::anyhow!(
80            "Invalid doc path: path traversal not allowed"
81        ));
82    }
83
84    // Only allow safe characters
85    if !path
86        .chars()
87        .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.' || c == '/')
88    {
89        return Err(anyhow::anyhow!(
90            "Invalid doc path: only alphanumeric, hyphen, underscore, dot, and slash allowed"
91        ));
92    }
93
94    // Must end with .md or .markdown
95    if !path.ends_with(".md") && !path.ends_with(".markdown") {
96        return Err(anyhow::anyhow!(
97            "Invalid doc path: must end with .md or .markdown"
98        ));
99    }
100
101    Ok(())
102}
103
104/// List all documentation files as JSON.
105pub fn list_docs(docs_dir: Option<&Path>) -> Result<Value> {
106    let Some(dir) = docs_dir else {
107        return Ok(json!({
108            "docs": [],
109            "count": 0,
110            "docs_dir": null,
111            "error": "No docs directory configured"
112        }));
113    };
114
115    if !dir.exists() {
116        return Ok(json!({
117            "docs": [],
118            "count": 0,
119            "docs_dir": dir.display().to_string(),
120            "error": "Docs directory does not exist"
121        }));
122    }
123
124    let mut files = Vec::new();
125    find_markdown_files(dir, dir, &mut files)?;
126
127    // Sort by relative path for consistent ordering
128    files.sort_by(|a, b| a.relative_path.cmp(&b.relative_path));
129
130    let docs_list: Vec<Value> = files
131        .iter()
132        .map(|doc| {
133            json!({
134                "name": doc.name,
135                "path": doc.relative_path,
136                "uri": format!("docs://{}", doc.relative_path),
137                "size": doc.size,
138                "mime_type": "text/markdown",
139            })
140        })
141        .collect();
142
143    let count = docs_list.len();
144
145    Ok(json!({
146        "docs": docs_list,
147        "count": count,
148        "docs_dir": dir.display().to_string(),
149    }))
150}
151
152/// Get a specific documentation file's content as JSON.
153pub fn get_doc_resource(docs_dir: Option<&Path>, path: &str) -> Result<Value> {
154    validate_doc_path(path)?;
155
156    let Some(dir) = docs_dir else {
157        return Err(anyhow::anyhow!("No docs directory configured"));
158    };
159
160    // Construct the full path
161    let file_path = dir.join(path);
162
163    // Security check: ensure resolved path is within docs_dir
164    if let Ok(canonical_file) = file_path.canonicalize()
165        && let Ok(canonical_dir) = dir.canonicalize()
166        && !canonical_file.starts_with(&canonical_dir)
167    {
168        return Err(anyhow::anyhow!("Invalid doc path: outside docs directory"));
169    }
170
171    if !file_path.exists() {
172        return Err(anyhow::anyhow!("Documentation file not found: {}", path));
173    }
174
175    if !file_path.is_file() {
176        return Err(anyhow::anyhow!("Not a file: {}", path));
177    }
178
179    let content = std::fs::read_to_string(&file_path)
180        .map_err(|e| anyhow::anyhow!("Failed to read doc file: {}", e))?;
181
182    let size = std::fs::metadata(&file_path).map(|m| m.len()).unwrap_or(0);
183
184    let name = file_path
185        .file_name()
186        .map(|n| n.to_string_lossy().to_string())
187        .unwrap_or_default();
188
189    Ok(json!({
190        "name": name,
191        "path": path,
192        "uri": format!("docs://{}", path),
193        "content": content,
194        "size": size,
195        "mime_type": "text/markdown",
196    }))
197}
198
199/// A single match within a documentation file.
200#[derive(Debug, Clone)]
201struct DocMatch {
202    /// Line number (1-based) where the match was found
203    line_number: usize,
204    /// The matching line (trimmed)
205    line_text: String,
206    /// Context: a few lines around the match
207    context: String,
208}
209
210/// A search result for a documentation file.
211#[derive(Debug, Clone)]
212struct DocSearchResult {
213    /// Relative path from docs/ directory
214    relative_path: String,
215    /// File name only
216    name: String,
217    /// Size in bytes
218    size: u64,
219    /// Whether the filename itself matched the query
220    name_match: bool,
221    /// Matched lines within the file content
222    matches: Vec<DocMatch>,
223    /// Total number of matches in this file
224    match_count: usize,
225}
226
227/// Extract a context snippet around a line in the content.
228/// Returns up to `context_lines` lines before and after the matched line.
229fn extract_context(lines: &[&str], line_idx: usize, context_lines: usize) -> String {
230    let start = line_idx.saturating_sub(context_lines);
231    let end = (line_idx + context_lines + 1).min(lines.len());
232    lines[start..end].join("\n")
233}
234
235/// Search documentation files for a query string.
236///
237/// Performs case-insensitive substring matching across all markdown files.
238/// Searches both filenames and file content. Results are ranked by relevance:
239/// filename matches first, then by number of content matches.
240///
241/// The query is split into terms (space-separated) and all terms must appear
242/// in the file (either in the filename or content) for it to be a match.
243///
244/// Supports pagination via `limit` and `offset` parameters.
245pub fn search_docs(
246    docs_dir: Option<&Path>,
247    query: &str,
248    limit: Option<usize>,
249    offset: Option<usize>,
250) -> Result<Value> {
251    let Some(dir) = docs_dir else {
252        return Ok(json!({
253            "query": query,
254            "results": [],
255            "result_count": 0,
256            "total_matches": 0,
257            "has_more": false,
258            "error": "No docs directory configured"
259        }));
260    };
261
262    if !dir.exists() {
263        return Ok(json!({
264            "query": query,
265            "results": [],
266            "result_count": 0,
267            "total_matches": 0,
268            "has_more": false,
269            "docs_dir": dir.display().to_string(),
270            "error": "Docs directory does not exist"
271        }));
272    }
273
274    if query.trim().is_empty() {
275        return Err(anyhow::anyhow!("Search query cannot be empty"));
276    }
277
278    let limit = limit.unwrap_or(20).min(100);
279    let offset = offset.unwrap_or(0);
280    let context_lines = 2;
281    let max_matches_per_file = 5;
282
283    // Normalize query: split into lowercase terms
284    let query_lower = query.to_lowercase();
285    let terms: Vec<&str> = query_lower.split_whitespace().collect();
286
287    if terms.is_empty() {
288        return Err(anyhow::anyhow!("Search query cannot be empty"));
289    }
290
291    // Find all markdown files
292    let mut files = Vec::new();
293    find_markdown_files(dir, dir, &mut files)?;
294
295    // Search each file
296    let mut results: Vec<DocSearchResult> = Vec::new();
297
298    for doc in &files {
299        let file_path = dir.join(&doc.relative_path);
300        let content = match std::fs::read_to_string(&file_path) {
301            Ok(c) => c,
302            Err(_) => continue, // Skip unreadable files
303        };
304
305        let content_lower = content.to_lowercase();
306        let name_lower = doc.name.to_lowercase();
307        let path_lower = doc.relative_path.to_lowercase();
308
309        // Check if ALL terms appear somewhere in the file (name or content)
310        let all_terms_present = terms.iter().all(|term| {
311            name_lower.contains(term) || path_lower.contains(term) || content_lower.contains(term)
312        });
313
314        if !all_terms_present {
315            continue;
316        }
317
318        // Check if filename matches any term
319        let name_match = terms
320            .iter()
321            .any(|term| name_lower.contains(term) || path_lower.contains(term));
322
323        // Find matching lines in content
324        let lines: Vec<&str> = content.lines().collect();
325        let mut doc_matches = Vec::new();
326
327        for (idx, line) in lines.iter().enumerate() {
328            let line_lower = line.to_lowercase();
329            // A line matches if it contains any of the search terms
330            if terms.iter().any(|term| line_lower.contains(term)) {
331                let context = extract_context(&lines, idx, context_lines);
332                doc_matches.push(DocMatch {
333                    line_number: idx + 1,
334                    line_text: line.trim().to_string(),
335                    context,
336                });
337
338                if doc_matches.len() >= max_matches_per_file {
339                    break;
340                }
341            }
342        }
343
344        let match_count = if doc_matches.len() >= max_matches_per_file {
345            // Count all matches if we hit the per-file limit
346            lines
347                .iter()
348                .filter(|line| {
349                    let ll = line.to_lowercase();
350                    terms.iter().any(|term| ll.contains(term))
351                })
352                .count()
353        } else {
354            doc_matches.len()
355        };
356
357        // Only include files that have at least one match (name or content)
358        if name_match || !doc_matches.is_empty() {
359            results.push(DocSearchResult {
360                relative_path: doc.relative_path.clone(),
361                name: doc.name.clone(),
362                size: doc.size,
363                name_match,
364                matches: doc_matches,
365                match_count,
366            });
367        }
368    }
369
370    // Sort results: filename matches first, then by match count (descending)
371    results.sort_by(|a, b| {
372        b.name_match
373            .cmp(&a.name_match)
374            .then_with(|| b.match_count.cmp(&a.match_count))
375    });
376
377    let total_results = results.len();
378    let total_matches: usize = results.iter().map(|r| r.match_count).sum();
379
380    // Apply pagination
381    let paginated: Vec<&DocSearchResult> = results.iter().skip(offset).take(limit + 1).collect();
382    let has_more = paginated.len() > limit;
383    let paginated: Vec<&DocSearchResult> = paginated.into_iter().take(limit).collect();
384
385    // Convert to JSON
386    let results_json: Vec<Value> = paginated
387        .iter()
388        .map(|r| {
389            let matches_json: Vec<Value> = r
390                .matches
391                .iter()
392                .map(|m| {
393                    json!({
394                        "line": m.line_number,
395                        "text": m.line_text,
396                        "context": m.context,
397                    })
398                })
399                .collect();
400
401            json!({
402                "name": r.name,
403                "path": r.relative_path,
404                "uri": format!("docs://{}", r.relative_path),
405                "size": r.size,
406                "name_match": r.name_match,
407                "match_count": r.match_count,
408                "matches": matches_json,
409            })
410        })
411        .collect();
412
413    let result_count = results_json.len();
414
415    Ok(json!({
416        "query": query,
417        "results": results_json,
418        "result_count": result_count,
419        "total_files_matched": total_results,
420        "total_matches": total_matches,
421        "has_more": has_more,
422        "offset": offset,
423        "limit": limit,
424        "docs_dir": dir.display().to_string(),
425    }))
426}
427
428/// Get resources for all documentation files (for get_resources registration).
429pub fn get_doc_resources(docs_dir: Option<&Path>) -> Vec<(String, String, String)> {
430    let mut resources = Vec::new();
431
432    let Some(dir) = docs_dir else {
433        return resources;
434    };
435
436    if !dir.exists() {
437        return resources;
438    }
439
440    let mut files = Vec::new();
441    if find_markdown_files(dir, dir, &mut files).is_ok() {
442        for doc in files {
443            let uri = format!("docs://{}", doc.relative_path);
444            let name = doc.name.clone();
445            let description = format!("Documentation: {}", doc.relative_path);
446            resources.push((uri, name, description));
447        }
448    }
449
450    resources
451}
452
453#[cfg(test)]
454mod tests {
455    use super::*;
456    use std::fs;
457    use tempfile::TempDir;
458
459    #[test]
460    fn test_validate_doc_path() {
461        // Valid paths
462        assert!(validate_doc_path("README.md").is_ok());
463        assert!(validate_doc_path("GATES.md").is_ok());
464        assert!(validate_doc_path("diagrams/README.md").is_ok());
465        assert!(validate_doc_path("sub/dir/file.md").is_ok());
466        assert!(validate_doc_path("file_name-test.md").is_ok());
467
468        // Invalid paths
469        assert!(validate_doc_path("").is_err());
470        assert!(validate_doc_path("../etc/passwd").is_err());
471        assert!(validate_doc_path("..\\windows\\system32").is_err());
472        assert!(validate_doc_path("file.txt").is_err()); // wrong extension
473        assert!(validate_doc_path("file<script>.md").is_err()); // invalid chars
474    }
475
476    #[test]
477    fn test_list_docs_no_dir() {
478        let result = list_docs(None).unwrap();
479        assert_eq!(result["count"], 0);
480        assert!(result["error"].is_string());
481    }
482
483    #[test]
484    fn test_list_docs_with_files() {
485        let temp_dir = TempDir::new().unwrap();
486        let docs_path = temp_dir.path();
487
488        // Create test files
489        fs::write(docs_path.join("README.md"), "# Readme").unwrap();
490        fs::write(docs_path.join("GUIDE.md"), "# Guide").unwrap();
491
492        // Create subdirectory with file
493        fs::create_dir(docs_path.join("subdir")).unwrap();
494        fs::write(docs_path.join("subdir/NESTED.md"), "# Nested").unwrap();
495
496        let result = list_docs(Some(docs_path)).unwrap();
497        assert_eq!(result["count"], 3);
498
499        let docs = result["docs"].as_array().unwrap();
500        let paths: Vec<&str> = docs.iter().map(|d| d["path"].as_str().unwrap()).collect();
501
502        assert!(paths.contains(&"README.md"));
503        assert!(paths.contains(&"GUIDE.md"));
504        assert!(paths.contains(&"subdir/NESTED.md"));
505    }
506
507    #[test]
508    fn test_get_doc_resource() {
509        let temp_dir = TempDir::new().unwrap();
510        let docs_path = temp_dir.path();
511
512        let content = "# Test Document\n\nThis is a test.";
513        fs::write(docs_path.join("TEST.md"), content).unwrap();
514
515        let result = get_doc_resource(Some(docs_path), "TEST.md").unwrap();
516        assert_eq!(result["name"], "TEST.md");
517        assert_eq!(result["content"], content);
518        assert_eq!(result["mime_type"], "text/markdown");
519    }
520
521    #[test]
522    fn test_get_doc_resource_not_found() {
523        let temp_dir = TempDir::new().unwrap();
524        let result = get_doc_resource(Some(temp_dir.path()), "NONEXISTENT.md");
525        assert!(result.is_err());
526    }
527
528    #[test]
529    fn test_search_docs_no_dir() {
530        let result = search_docs(None, "test", None, None).unwrap();
531        assert_eq!(result["result_count"], 0);
532        assert!(result["error"].is_string());
533    }
534
535    #[test]
536    fn test_search_docs_empty_query() {
537        let temp_dir = TempDir::new().unwrap();
538        let result = search_docs(Some(temp_dir.path()), "", None, None);
539        assert!(result.is_err());
540    }
541
542    #[test]
543    fn test_search_docs_finds_content() {
544        let temp_dir = TempDir::new().unwrap();
545        let docs_path = temp_dir.path();
546
547        fs::write(
548            docs_path.join("GATES.md"),
549            "# Gates\n\nGates are quality checkpoints.\nThey verify task completion.",
550        )
551        .unwrap();
552        fs::write(
553            docs_path.join("DESIGN.md"),
554            "# Design\n\nArchitecture overview.\nSystem design document.",
555        )
556        .unwrap();
557
558        // Search for content in GATES.md
559        let result = search_docs(Some(docs_path), "checkpoints", None, None).unwrap();
560        assert_eq!(result["result_count"], 1);
561        let results = result["results"].as_array().unwrap();
562        assert_eq!(results[0]["name"], "GATES.md");
563        assert!(results[0]["match_count"].as_u64().unwrap() > 0);
564    }
565
566    #[test]
567    fn test_search_docs_filename_match() {
568        let temp_dir = TempDir::new().unwrap();
569        let docs_path = temp_dir.path();
570
571        fs::write(docs_path.join("GATES.md"), "# Gates\n\nContent here.").unwrap();
572        fs::write(docs_path.join("DESIGN.md"), "# Design\n\nOther content.").unwrap();
573
574        // Search for filename
575        let result = search_docs(Some(docs_path), "gates", None, None).unwrap();
576        assert_eq!(result["result_count"], 1);
577        let results = result["results"].as_array().unwrap();
578        assert_eq!(results[0]["name"], "GATES.md");
579        assert!(results[0]["name_match"].as_bool().unwrap());
580    }
581
582    #[test]
583    fn test_search_docs_case_insensitive() {
584        let temp_dir = TempDir::new().unwrap();
585        let docs_path = temp_dir.path();
586
587        fs::write(
588            docs_path.join("TEST.md"),
589            "# Test\n\nThis has UPPERCASE and lowercase content.",
590        )
591        .unwrap();
592
593        // Case-insensitive search
594        let result = search_docs(Some(docs_path), "uppercase", None, None).unwrap();
595        assert_eq!(result["result_count"], 1);
596
597        let result = search_docs(Some(docs_path), "UPPERCASE", None, None).unwrap();
598        assert_eq!(result["result_count"], 1);
599    }
600
601    #[test]
602    fn test_search_docs_multi_term() {
603        let temp_dir = TempDir::new().unwrap();
604        let docs_path = temp_dir.path();
605
606        fs::write(
607            docs_path.join("A.md"),
608            "# Alpha\n\nThis has alpha and beta content.",
609        )
610        .unwrap();
611        fs::write(
612            docs_path.join("B.md"),
613            "# Beta\n\nThis only has beta content.",
614        )
615        .unwrap();
616
617        // Multi-term: both must be present
618        let result = search_docs(Some(docs_path), "alpha beta", None, None).unwrap();
619        assert_eq!(result["result_count"], 1);
620        let results = result["results"].as_array().unwrap();
621        assert_eq!(results[0]["name"], "A.md");
622    }
623
624    #[test]
625    fn test_search_docs_pagination() {
626        let temp_dir = TempDir::new().unwrap();
627        let docs_path = temp_dir.path();
628
629        // Create several matching files
630        for i in 0..5 {
631            fs::write(
632                docs_path.join(format!("DOC{}.md", i)),
633                format!("# Doc {}\n\nSearchable content in doc {}.", i, i),
634            )
635            .unwrap();
636        }
637
638        // Search with limit
639        let result = search_docs(Some(docs_path), "searchable", Some(2), None).unwrap();
640        assert_eq!(result["result_count"], 2);
641        assert!(result["has_more"].as_bool().unwrap());
642
643        // Search with offset
644        let result = search_docs(Some(docs_path), "searchable", Some(2), Some(2)).unwrap();
645        assert_eq!(result["result_count"], 2);
646        assert!(result["has_more"].as_bool().unwrap());
647
648        // Search with offset past results
649        let result = search_docs(Some(docs_path), "searchable", Some(2), Some(4)).unwrap();
650        assert_eq!(result["result_count"], 1);
651        assert!(!result["has_more"].as_bool().unwrap());
652    }
653
654    #[test]
655    fn test_search_docs_with_subdirectories() {
656        let temp_dir = TempDir::new().unwrap();
657        let docs_path = temp_dir.path();
658
659        fs::write(docs_path.join("ROOT.md"), "# Root\n\nRoot level content.").unwrap();
660        fs::create_dir(docs_path.join("subdir")).unwrap();
661        fs::write(
662            docs_path.join("subdir/NESTED.md"),
663            "# Nested\n\nNested searchable content.",
664        )
665        .unwrap();
666
667        let result = search_docs(Some(docs_path), "searchable", None, None).unwrap();
668        assert_eq!(result["result_count"], 1);
669        let results = result["results"].as_array().unwrap();
670        assert_eq!(results[0]["path"], "subdir/NESTED.md");
671    }
672}