hanzo_mcp/tools/
cloud_code.rs

1//! Cloud-backed code knowledge tools (HIP-0300).
2//!
3//! Cross-repo search, context bundling, cited RAG answers, and indexing over
4//! the live api.hanzo.ai code plane. These complement the local tree-sitter AST
5//! engine (`code` tool / `search/ast_search.rs`), which stays for instant
6//! single-file symbol/AST ops. Hybrid: local for one file, cloud for the corpus.
7//!
8//! - `code_search`  → GET  /v1/code/search  (symbol|text|semantic|hybrid)
9//! - `code_context` → POST /v1/code/context (token-budgeted span bundle)
10//! - `code_ask`     → GET  /v1/code/ask     (cited answer)
11//! - `code_index`   → POST /v1/code/index   (index a repo's files)
12
13use anyhow::Result;
14use async_trait::async_trait;
15use serde::Deserialize;
16use serde_json::{json, Value};
17use std::path::Path;
18use walkdir::WalkDir;
19
20use super::{envelope_err, envelope_ok};
21use crate::hanzo_api::HanzoApi;
22use crate::{MCPTool, ToolResult};
23
24const SKIP_DIRS: &[&str] = &[
25    ".git", "node_modules", "target", "dist", "build", "__pycache__", ".venv", "venv", ".next",
26];
27const MAX_INDEX_FILES: usize = 5000;
28const MAX_FILE_BYTES: u64 = 512 * 1024;
29
30/// Wrap a raw API result (or transport error) in the standard envelope.
31fn respond(tool: &'static str, action: &'static str, result: Result<Value>) -> Result<ToolResult> {
32    Ok(match result {
33        Ok(body) => ToolResult::ok(envelope_ok(tool, action, body)),
34        Err(e) => ToolResult::ok(envelope_err(tool, action, "UPSTREAM", e.to_string())),
35    })
36}
37
38fn require_key(api: &HanzoApi, tool: &'static str, action: &'static str) -> Option<Result<ToolResult>> {
39    if api.has_key() {
40        None
41    } else {
42        Some(Ok(ToolResult::ok(envelope_err(
43            tool,
44            action,
45            "NO_API_KEY",
46            crate::hanzo_api::NO_KEY,
47        ))))
48    }
49}
50
51// ---------------------------------------------------------------------------
52// code_search → GET /v1/code/search
53// ---------------------------------------------------------------------------
54
55#[derive(Debug, Default, Deserialize)]
56struct SearchArgs {
57    #[serde(alias = "q")]
58    query: Option<String>,
59    /// symbol | text | semantic | hybrid (default hybrid)
60    #[serde(alias = "type")]
61    kind: Option<String>,
62    repo: Option<String>,
63    limit: Option<u32>,
64}
65
66pub struct CodeSearchTool {
67    api: HanzoApi,
68}
69
70impl CodeSearchTool {
71    pub fn new() -> Self {
72        Self { api: HanzoApi::from_env() }
73    }
74    pub fn schema() -> Value {
75        json!({
76            "name": "code_search",
77            "description": "Hybrid cross-repo code search on api.hanzo.ai (symbol|text|semantic|hybrid). Cloud counterpart to the local `code` AST tool.",
78            "inputSchema": {
79                "type": "object",
80                "properties": {
81                    "query": { "type": "string", "description": "Search query" },
82                    "type": { "type": "string", "enum": ["symbol", "text", "semantic", "hybrid"], "default": "hybrid" },
83                    "repo": { "type": "string", "description": "Restrict to a repo" },
84                    "limit": { "type": "number", "default": 20 }
85                },
86                "required": ["query"]
87            }
88        })
89    }
90}
91
92impl Default for CodeSearchTool {
93    fn default() -> Self {
94        Self::new()
95    }
96}
97
98#[async_trait]
99impl MCPTool for CodeSearchTool {
100    fn name(&self) -> &str {
101        "code_search"
102    }
103    fn description(&self) -> &str {
104        "Hybrid cross-repo code search on api.hanzo.ai (symbol|text|semantic|hybrid)"
105    }
106    fn parameters(&self) -> Value {
107        Self::schema()["inputSchema"].clone()
108    }
109    async fn execute(&self, params: Value) -> Result<ToolResult> {
110        if let Some(r) = require_key(&self.api, "code_search", "search") {
111            return r;
112        }
113        let args: SearchArgs = serde_json::from_value(params).unwrap_or_default();
114        let query = match args.query.filter(|q| !q.trim().is_empty()) {
115            Some(q) => q,
116            None => return Ok(ToolResult::ok(envelope_err("code_search", "search", "INVALID_ARGS", "query required"))),
117        };
118
119        let mut q: Vec<(&str, String)> = vec![
120            ("q", query),
121            ("type", args.kind.unwrap_or_else(|| "hybrid".to_string())),
122            ("limit", args.limit.unwrap_or(20).to_string()),
123        ];
124        if let Some(repo) = args.repo.filter(|r| !r.is_empty()) {
125            q.push(("repo", repo));
126        }
127        respond("code_search", "search", self.api.get("/v1/code/search", &q).await)
128    }
129}
130
131// ---------------------------------------------------------------------------
132// code_context → POST /v1/code/context
133// ---------------------------------------------------------------------------
134
135#[derive(Debug, Default, Deserialize)]
136struct ContextArgs {
137    #[serde(alias = "q")]
138    query: Option<String>,
139    #[serde(alias = "budget_tokens", alias = "budget")]
140    budget_tokens: Option<u32>,
141    repo: Option<String>,
142}
143
144pub struct CodeContextTool {
145    api: HanzoApi,
146}
147
148impl CodeContextTool {
149    pub fn new() -> Self {
150        Self { api: HanzoApi::from_env() }
151    }
152    pub fn schema() -> Value {
153        json!({
154            "name": "code_context",
155            "description": "Token-budgeted code context bundle for a query from api.hanzo.ai.",
156            "inputSchema": {
157                "type": "object",
158                "properties": {
159                    "query": { "type": "string", "description": "What to gather context for" },
160                    "budgetTokens": { "type": "number", "default": 2000 },
161                    "repo": { "type": "string" }
162                },
163                "required": ["query"]
164            }
165        })
166    }
167}
168
169impl Default for CodeContextTool {
170    fn default() -> Self {
171        Self::new()
172    }
173}
174
175#[async_trait]
176impl MCPTool for CodeContextTool {
177    fn name(&self) -> &str {
178        "code_context"
179    }
180    fn description(&self) -> &str {
181        "Token-budgeted code context bundle for a query from api.hanzo.ai"
182    }
183    fn parameters(&self) -> Value {
184        Self::schema()["inputSchema"].clone()
185    }
186    async fn execute(&self, params: Value) -> Result<ToolResult> {
187        if let Some(r) = require_key(&self.api, "code_context", "context") {
188            return r;
189        }
190        let args: ContextArgs = serde_json::from_value(params).unwrap_or_default();
191        let query = match args.query.filter(|q| !q.trim().is_empty()) {
192            Some(q) => q,
193            None => return Ok(ToolResult::ok(envelope_err("code_context", "context", "INVALID_ARGS", "query required"))),
194        };
195        let mut body = json!({
196            "query": query,
197            "budgetTokens": args.budget_tokens.unwrap_or(2000),
198        });
199        if let Some(repo) = args.repo.filter(|r| !r.is_empty()) {
200            body["repo"] = json!(repo);
201        }
202        respond("code_context", "context", self.api.post("/v1/code/context", body).await)
203    }
204}
205
206// ---------------------------------------------------------------------------
207// code_ask → GET /v1/code/ask
208// ---------------------------------------------------------------------------
209
210#[derive(Debug, Default, Deserialize)]
211struct AskArgs {
212    #[serde(alias = "q", alias = "question")]
213    query: Option<String>,
214    repo: Option<String>,
215}
216
217pub struct CodeAskTool {
218    api: HanzoApi,
219}
220
221impl CodeAskTool {
222    pub fn new() -> Self {
223        Self { api: HanzoApi::from_env() }
224    }
225    pub fn schema() -> Value {
226        json!({
227            "name": "code_ask",
228            "description": "Ask a question about indexed code; returns a cited RAG answer from api.hanzo.ai.",
229            "inputSchema": {
230                "type": "object",
231                "properties": {
232                    "query": { "type": "string", "description": "Question about the code" },
233                    "repo": { "type": "string" }
234                },
235                "required": ["query"]
236            }
237        })
238    }
239}
240
241impl Default for CodeAskTool {
242    fn default() -> Self {
243        Self::new()
244    }
245}
246
247#[async_trait]
248impl MCPTool for CodeAskTool {
249    fn name(&self) -> &str {
250        "code_ask"
251    }
252    fn description(&self) -> &str {
253        "Ask a question about indexed code; returns a cited RAG answer from api.hanzo.ai"
254    }
255    fn parameters(&self) -> Value {
256        Self::schema()["inputSchema"].clone()
257    }
258    async fn execute(&self, params: Value) -> Result<ToolResult> {
259        if let Some(r) = require_key(&self.api, "code_ask", "ask") {
260            return r;
261        }
262        let args: AskArgs = serde_json::from_value(params).unwrap_or_default();
263        let query = match args.query.filter(|q| !q.trim().is_empty()) {
264            Some(q) => q,
265            None => return Ok(ToolResult::ok(envelope_err("code_ask", "ask", "INVALID_ARGS", "query required"))),
266        };
267        let mut q: Vec<(&str, String)> = vec![("q", query)];
268        if let Some(repo) = args.repo.filter(|r| !r.is_empty()) {
269            q.push(("repo", repo));
270        }
271        respond("code_ask", "ask", self.api.get("/v1/code/ask", &q).await)
272    }
273}
274
275// ---------------------------------------------------------------------------
276// code_index → POST /v1/code/index
277// ---------------------------------------------------------------------------
278
279#[derive(Debug, Deserialize)]
280struct IndexFile {
281    path: String,
282    content: String,
283}
284
285#[derive(Debug, Default, Deserialize)]
286struct IndexArgs {
287    repo: Option<String>,
288    /// Local directory to walk into `files` when `files` is omitted.
289    path: Option<String>,
290    files: Option<Vec<IndexFile>>,
291}
292
293pub struct CodeIndexTool {
294    api: HanzoApi,
295}
296
297impl CodeIndexTool {
298    pub fn new() -> Self {
299        Self { api: HanzoApi::from_env() }
300    }
301    pub fn schema() -> Value {
302        json!({
303            "name": "code_index",
304            "description": "Index a repo into api.hanzo.ai. Pass explicit `files`, or a local `path` to walk (text files, bounded).",
305            "inputSchema": {
306                "type": "object",
307                "properties": {
308                    "repo": { "type": "string", "description": "Repo name (defaults to the directory basename)" },
309                    "path": { "type": "string", "description": "Local directory to walk and index" },
310                    "files": {
311                        "type": "array",
312                        "description": "Explicit files to index",
313                        "items": {
314                            "type": "object",
315                            "properties": {
316                                "path": { "type": "string" },
317                                "content": { "type": "string" }
318                            },
319                            "required": ["path", "content"]
320                        }
321                    }
322                },
323                "required": []
324            }
325        })
326    }
327}
328
329impl Default for CodeIndexTool {
330    fn default() -> Self {
331        Self::new()
332    }
333}
334
335/// Walk a directory into indexable `{path, content}` files (bounded, text only).
336fn collect_files(dir: &Path) -> (Vec<Value>, usize) {
337    let mut files = Vec::new();
338    let mut skipped = 0usize;
339    for entry in WalkDir::new(dir)
340        .follow_links(false)
341        .into_iter()
342        .filter_entry(|e| {
343            !(e.file_type().is_dir()
344                && e.file_name().to_str().map_or(false, |n| SKIP_DIRS.contains(&n)))
345        })
346        .filter_map(|e| e.ok())
347        .filter(|e| e.file_type().is_file())
348    {
349        if files.len() >= MAX_INDEX_FILES {
350            break;
351        }
352        let path = entry.path();
353        match entry.metadata() {
354            Ok(m) if m.len() > MAX_FILE_BYTES => {
355                skipped += 1;
356                continue;
357            }
358            Ok(_) => {}
359            Err(_) => {
360                skipped += 1;
361                continue;
362            }
363        }
364        match std::fs::read_to_string(path) {
365            Ok(content) => {
366                let rel = path.strip_prefix(dir).unwrap_or(path).to_string_lossy().to_string();
367                files.push(json!({ "path": rel, "content": content }));
368            }
369            Err(_) => skipped += 1, // binary / non-utf8
370        }
371    }
372    (files, skipped)
373}
374
375#[async_trait]
376impl MCPTool for CodeIndexTool {
377    fn name(&self) -> &str {
378        "code_index"
379    }
380    fn description(&self) -> &str {
381        "Index a repo into api.hanzo.ai (explicit files or a local path to walk)"
382    }
383    fn parameters(&self) -> Value {
384        Self::schema()["inputSchema"].clone()
385    }
386    async fn execute(&self, params: Value) -> Result<ToolResult> {
387        if let Some(r) = require_key(&self.api, "code_index", "index") {
388            return r;
389        }
390        let args: IndexArgs = serde_json::from_value(params).unwrap_or_default();
391
392        // Resolve files: explicit list wins; else walk the local path.
393        let (files, repo, walked_skipped) = if let Some(list) = args.files {
394            let files: Vec<Value> = list
395                .into_iter()
396                .map(|f| json!({ "path": f.path, "content": f.content }))
397                .collect();
398            (files, args.repo.clone(), 0usize)
399        } else if let Some(path) = args.path.as_deref().filter(|p| !p.is_empty()) {
400            let dir = Path::new(path);
401            if !dir.is_dir() {
402                return Ok(ToolResult::ok(envelope_err("code_index", "index", "INVALID_ARGS", format!("not a directory: {}", path))));
403            }
404            let repo = args.repo.clone().or_else(|| {
405                dir.canonicalize()
406                    .ok()
407                    .and_then(|c| c.file_name().map(|n| n.to_string_lossy().to_string()))
408            });
409            let (files, skipped) = collect_files(dir);
410            (files, repo, skipped)
411        } else {
412            return Ok(ToolResult::ok(envelope_err("code_index", "index", "INVALID_ARGS", "provide `files` or a local `path`")));
413        };
414
415        if files.is_empty() {
416            return Ok(ToolResult::ok(envelope_err("code_index", "index", "EMPTY", "no indexable files found")));
417        }
418        let repo = match repo.filter(|r| !r.is_empty()) {
419            Some(r) => r,
420            None => return Ok(ToolResult::ok(envelope_err("code_index", "index", "INVALID_ARGS", "repo required"))),
421        };
422
423        let sent = files.len();
424        let body = json!({ "repo": repo, "files": files });
425        match self.api.post("/v1/code/index", body).await {
426            Ok(mut result) => {
427                if let Some(obj) = result.as_object_mut() {
428                    obj.insert("sent".into(), json!(sent));
429                    if walked_skipped > 0 {
430                        obj.insert("walkedSkipped".into(), json!(walked_skipped));
431                    }
432                }
433                Ok(ToolResult::ok(envelope_ok("code_index", "index", result)))
434            }
435            Err(e) => Ok(ToolResult::ok(envelope_err("code_index", "index", "UPSTREAM", e.to_string()))),
436        }
437    }
438}
439
440#[cfg(test)]
441mod tests {
442    use super::*;
443    use std::io::Write;
444
445    #[test]
446    fn search_schema_shape() {
447        let s = CodeSearchTool::schema();
448        assert_eq!(s["name"], "code_search");
449        assert_eq!(s["inputSchema"]["required"][0], "query");
450        let props = &s["inputSchema"]["properties"];
451        assert!(props.get("type").is_some());
452        assert!(props.get("repo").is_some());
453    }
454
455    #[test]
456    fn tool_names_are_stable() {
457        assert_eq!(CodeSearchTool::new().name(), "code_search");
458        assert_eq!(CodeContextTool::new().name(), "code_context");
459        assert_eq!(CodeAskTool::new().name(), "code_ask");
460        assert_eq!(CodeIndexTool::new().name(), "code_index");
461    }
462
463    #[test]
464    fn search_args_accept_q_and_type_aliases() {
465        let a: SearchArgs = serde_json::from_value(json!({ "q": "foo", "type": "symbol", "limit": 5 })).unwrap();
466        assert_eq!(a.query.as_deref(), Some("foo"));
467        assert_eq!(a.kind.as_deref(), Some("symbol"));
468        assert_eq!(a.limit, Some(5));
469    }
470
471    #[test]
472    fn collect_files_walks_and_skips_noise() {
473        let dir = tempfile::tempdir().unwrap();
474        let root = dir.path();
475        let mut f = std::fs::File::create(root.join("a.rs")).unwrap();
476        writeln!(f, "fn main() {{}}").unwrap();
477        std::fs::create_dir_all(root.join("node_modules")).unwrap();
478        std::fs::File::create(root.join("node_modules").join("junk.js")).unwrap();
479
480        let (files, _skipped) = collect_files(root);
481        assert_eq!(files.len(), 1, "node_modules must be pruned");
482        assert_eq!(files[0]["path"], "a.rs");
483    }
484
485    /// Live end-to-end via the registry: index a repo, then find a symbol in it.
486    /// Requires a live api.hanzo.ai + `sk-` key. Run: `cargo test -- --ignored`.
487    #[tokio::test]
488    #[ignore]
489    async fn live_index_then_search_via_registry() {
490        let registry = crate::ToolRegistry::with_defaults();
491        let repo = format!("mcp-rust-live-{}", std::process::id());
492
493        let indexed = registry
494            .execute("code_index", json!({
495                "repo": repo,
496                "files": [{ "path": "greet.rs", "content": "pub fn hanzo_greet() -> &'static str { \"hi\" }" }]
497            }))
498            .await
499            .unwrap();
500        assert_eq!(indexed.content["ok"], true, "index envelope: {}", indexed.content);
501        assert!(indexed.content["data"]["indexed"].as_u64().unwrap_or(0) >= 1);
502
503        let found = registry
504            .execute("code_search", json!({ "q": "hanzo_greet", "type": "hybrid", "repo": repo }))
505            .await
506            .unwrap();
507        assert_eq!(found.content["ok"], true, "search envelope: {}", found.content);
508        println!("code_search data: {}", found.content["data"]);
509    }
510}
hanzo_mcp/tools/cloud_code.rs

hanzo_mcp/tools/
cloud_code.rs
