Skip to main content

task_graph_mcp/tools/
mod.rs

1//! MCP tool implementations.
2
3pub mod agents;
4pub mod attachments;
5pub mod claiming;
6pub mod context;
7pub mod deps;
8pub mod files;
9pub mod gates;
10pub mod query;
11pub mod schema;
12pub mod search;
13pub mod skills;
14pub mod tasks;
15pub mod tracking;
16pub mod workflows;
17
18pub use context::ToolContext;
19
20use crate::config::{AppConfig, Prompts, ServerPaths, workflows::WorkflowsConfig};
21use crate::db::Database;
22use crate::error::ToolError;
23use crate::format::{OutputFormat, ToolResult};
24use anyhow::Result;
25use rmcp::model::Tool;
26use serde_json::Value;
27use std::path::PathBuf;
28use std::sync::Arc;
29
30/// Tool handler that processes MCP tool calls.
31pub struct ToolHandler {
32    pub db: Arc<Database>,
33    pub media_dir: PathBuf,
34    pub skills_dir: PathBuf,
35    pub server_paths: Arc<ServerPaths>,
36    pub prompts: Arc<Prompts>,
37    /// Consolidated application configuration.
38    pub config: AppConfig,
39    pub default_format: OutputFormat,
40    pub default_page_size: i32,
41    pub path_mapper: Arc<crate::paths::PathMapper>,
42}
43
44impl ToolHandler {
45    #[allow(clippy::too_many_arguments)]
46    pub fn new(
47        db: Arc<Database>,
48        media_dir: PathBuf,
49        skills_dir: PathBuf,
50        server_paths: Arc<ServerPaths>,
51        prompts: Arc<Prompts>,
52        config: AppConfig,
53        default_format: OutputFormat,
54        default_page_size: i32,
55        path_mapper: Arc<crate::paths::PathMapper>,
56    ) -> Self {
57        Self {
58            db,
59            media_dir,
60            skills_dir,
61            server_paths,
62            prompts,
63            config,
64            default_format,
65            default_page_size,
66            path_mapper,
67        }
68    }
69
70    /// Get the workflow config for a worker.
71    /// Looks up the worker's workflow name and returns the corresponding config,
72    /// or falls back to the configured default workflow, or the base config.
73    pub fn get_workflow_for_worker(&self, worker_id: &str) -> Arc<WorkflowsConfig> {
74        // Look up worker's workflow name from database
75        if let Ok(Some(worker)) = self.db.get_worker(worker_id)
76            && let Some(ref workflow_name) = worker.workflow
77        {
78            // Try to get from named_workflows cache
79            if let Some(workflow_config) = self.config.workflows.get_named_workflow(workflow_name) {
80                return Arc::clone(workflow_config);
81            }
82        }
83        // Fall back to configured default workflow, or base config
84        if let Some(default_workflow) = self.config.workflows.get_default_workflow() {
85            Arc::clone(default_workflow)
86        } else {
87            Arc::clone(&self.config.workflows)
88        }
89    }
90
91    /// Get all available tools.
92    pub fn get_tools(&self) -> Vec<Tool> {
93        let mut tools = Vec::new();
94
95        // Worker tools
96        tools.extend(agents::get_tools(&self.prompts));
97
98        // Task tools (with dynamic state schema)
99        tools.extend(tasks::get_tools(&self.prompts, &self.config.states));
100
101        // Tracking tools
102        tools.extend(tracking::get_tools(&self.prompts, &self.config.states));
103
104        // Dependency tools
105        tools.extend(deps::get_tools(&self.prompts, &self.config.deps));
106
107        // Claiming tools (with dynamic state schema)
108        tools.extend(claiming::get_tools(&self.prompts, &self.config.states));
109
110        // File coordination tools
111        tools.extend(files::get_tools(&self.prompts));
112
113        // Attachment tools
114        tools.extend(attachments::get_tools(&self.prompts));
115
116        // Skill tools (no prompts needed, always available)
117        tools.extend(skills::get_tools());
118
119        // Schema introspection tools
120        tools.extend(schema::get_tools());
121
122        // Search tools
123        tools.extend(search::get_tools(&self.prompts));
124
125        // Query tools (read-only SQL)
126        tools.extend(query::get_tools());
127
128        // Gate checking tools
129        tools.extend(gates::get_tools(&self.prompts));
130
131        // Workflow discovery tools (no auth needed, callable before connect)
132        tools.extend(workflows::get_tools());
133
134        tools
135    }
136
137    /// Call a tool by name.
138    #[allow(unused_variables)]
139    pub async fn call_tool(
140        &self,
141        name: &str,
142        arguments: Value,
143        ctx: &ToolContext,
144    ) -> Result<ToolResult> {
145        // Helper to wrap JSON results
146        let json = |r: Result<Value>| r.map(ToolResult::Json);
147
148        match name {
149            // Worker tools
150            "connect" => {
151                // Resolve workflow from args (worker isn't registered yet during connect)
152                let workflow = arguments
153                    .get("workflow")
154                    .and_then(|v| v.as_str())
155                    .and_then(|name| self.config.workflows.get_named_workflow(name))
156                    .map(Arc::clone)
157                    .or_else(|| self.config.workflows.get_default_workflow().map(Arc::clone))
158                    .unwrap_or_else(|| Arc::clone(&self.config.workflows));
159                json(agents::connect(
160                    agents::ConnectOptions {
161                        db: &self.db,
162                        server_paths: &self.server_paths,
163                        config: &self.config,
164                        workflows: &workflow,
165                    },
166                    arguments,
167                ))
168            }
169            "disconnect" => json(agents::disconnect(&self.db, &self.config.states, arguments)),
170            "list_agents" => agents::list_agents(
171                &self.db,
172                &self.config.states,
173                self.default_format,
174                arguments,
175            ),
176            "cleanup_stale" => json(agents::cleanup_stale(
177                &self.db,
178                &self.config.states,
179                arguments,
180            )),
181
182            // Task tools
183            "create" => json(tasks::create(&self.db, &self.config, arguments)),
184            "create_tree" => json(tasks::create_tree(&self.db, &self.config, arguments)),
185            "get" => json(tasks::get(&self.db, self.default_format, arguments)),
186            "list_tasks" => json(tasks::list_tasks(
187                &self.db,
188                &self.config.states,
189                &self.config.deps,
190                self.default_format,
191                arguments,
192            )),
193            "update" => {
194                // Look up worker's workflow for prompts
195                let worker_id = arguments
196                    .get("worker_id")
197                    .and_then(|v| v.as_str())
198                    .unwrap_or("");
199                let workflow = self.get_workflow_for_worker(worker_id);
200                json(tasks::update(
201                    tasks::UpdateOptions {
202                        db: &self.db,
203                        config: &self.config,
204                        workflows: &workflow,
205                    },
206                    arguments,
207                ))
208            }
209            "delete" => json(tasks::delete(&self.db, arguments)),
210            "rename" => json(tasks::rename(&self.db, arguments)),
211            "scan" => json(tasks::scan(&self.db, self.default_format, arguments)),
212
213            // Tracking tools
214            "thinking" => json(tracking::thinking(&self.db, arguments)),
215            "task_history" => json(tracking::task_history(
216                &self.db,
217                &self.config.states,
218                self.default_format,
219                arguments,
220            )),
221            "log_metrics" => json(tracking::log_metrics(&self.db, arguments)),
222            "get_metrics" => json(tracking::get_metrics(&self.db, arguments)),
223            "project_history" => json(tracking::project_history(
224                &self.db,
225                self.default_format,
226                arguments,
227            )),
228
229            // Dependency tools
230            "link" => json(deps::link(&self.db, &self.config.deps, arguments)),
231            "unlink" => json(deps::unlink(&self.db, arguments)),
232            "relink" => json(deps::relink(&self.db, &self.config.deps, arguments)),
233
234            // Claiming tools
235            "claim" => {
236                // Look up worker's workflow for prompts
237                let worker_id = arguments
238                    .get("worker_id")
239                    .and_then(|v| v.as_str())
240                    .unwrap_or("");
241                let workflow = self.get_workflow_for_worker(worker_id);
242                json(claiming::claim(
243                    &self.db,
244                    &self.config,
245                    &workflow,
246                    arguments,
247                ))
248            }
249
250            // File coordination tools
251            "mark_file" => json(files::mark_file(&self.db, arguments)),
252            "unmark_file" => json(files::unmark_file(&self.db, arguments)),
253            "list_marks" => json(files::list_marks(&self.db, self.default_format, arguments)),
254            "mark_updates" => {
255                json(files::mark_updates_async(std::sync::Arc::clone(&self.db), arguments).await)
256            }
257
258            // Attachment tools
259            "attach" => json(attachments::attach(
260                &self.db,
261                &self.media_dir,
262                &self.config.attachments,
263                arguments,
264            )),
265            "attachments" => json(attachments::attachments(
266                &self.db,
267                &self.media_dir,
268                self.default_format,
269                arguments,
270            )),
271            "detach" => json(attachments::detach(&self.db, &self.media_dir, arguments)),
272
273            // Skill tools
274            name if skills::is_skill_tool(name) => {
275                json(skills::call_tool(&self.skills_dir, name, &arguments))
276            }
277
278            // Schema introspection tools
279            "get_schema" => json(schema::get_schema(&self.db, arguments)),
280
281            // Search tools
282            "search" => json(search::search(&self.db, self.default_page_size, arguments)),
283
284            // Query tools (read-only SQL)
285            "query" => query::query(&self.db, self.default_format, arguments),
286
287            // Gate checking tools
288            "check_gates" => {
289                // Look up worker's workflow for gate definitions
290                // Since check_gates doesn't require worker_id, use base workflow
291                json(gates::check_gates(
292                    &self.db,
293                    &self.config.workflows,
294                    arguments,
295                ))
296            }
297
298            // Workflow discovery tools (no connection required)
299            "list_workflows" => json(workflows::list_workflows(&self.config.workflows)),
300
301            _ => Err(ToolError::unknown_tool(name).into()),
302        }
303    }
304}
305
306/// Helper to create a tool definition.
307pub fn make_tool(name: &str, description: &str, properties: Value, required: Vec<&str>) -> Tool {
308    let input_schema = rmcp::model::JsonObject::from_iter([
309        ("type".to_string(), serde_json::json!("object")),
310        ("properties".to_string(), properties),
311        ("required".to_string(), serde_json::json!(required)),
312    ]);
313
314    Tool::new(name.to_string(), description.to_string(), input_schema)
315}
316
317/// Helper to create a tool definition with prompt overrides.
318/// Looks up the tool description in prompts, falls back to default_description.
319pub fn make_tool_with_prompts(
320    name: &str,
321    default_description: &str,
322    properties: Value,
323    required: Vec<&str>,
324    prompts: &Prompts,
325) -> Tool {
326    let description = prompts
327        .get_tool_description(name)
328        .unwrap_or(default_description);
329    make_tool(name, description, properties, required)
330}
331
332/// Helper to get a string from arguments.
333pub fn get_string(args: &Value, key: &str) -> Option<String> {
334    args.get(key).and_then(|v| v.as_str().map(String::from))
335}
336
337/// Helper to get an i32 from arguments.
338pub fn get_i32(args: &Value, key: &str) -> Option<i32> {
339    args.get(key).and_then(|v| v.as_i64().map(|n| n as i32))
340}
341
342/// Helper to get an i64 from arguments.
343pub fn get_i64(args: &Value, key: &str) -> Option<i64> {
344    args.get(key).and_then(|v| v.as_i64())
345}
346
347/// Helper to get an f64 from arguments.
348pub fn get_f64(args: &Value, key: &str) -> Option<f64> {
349    args.get(key).and_then(|v| v.as_f64())
350}
351
352/// Helper to get a bool from arguments.
353pub fn get_bool(args: &Value, key: &str) -> Option<bool> {
354    args.get(key).and_then(|v| v.as_bool())
355}
356
357/// Helper to get a string array from arguments.
358pub fn get_string_array(args: &Value, key: &str) -> Option<Vec<String>> {
359    args.get(key).and_then(|v| {
360        v.as_array().map(|arr| {
361            arr.iter()
362                .filter_map(|v| v.as_str().map(String::from))
363                .collect()
364        })
365    })
366}
367
368/// Helper to get either a single string or array of strings from arguments.
369/// Normalizes to a Vec<String>.
370pub fn get_string_or_array(args: &Value, key: &str) -> Option<Vec<String>> {
371    args.get(key).and_then(|v| {
372        if let Some(s) = v.as_str() {
373            // Single string - wrap in vec
374            Some(vec![s.to_string()])
375        } else {
376            v.as_array().map(|arr| {
377                arr.iter()
378                    .filter_map(|item| item.as_str().map(String::from))
379                    .collect()
380            })
381        }
382    })
383}
384
385/// Parsed result that may be a list of IDs or a wildcard "*".
386pub enum IdList {
387    Ids(Vec<String>),
388    Wildcard,
389}
390
391/// Like get_string_or_array, but recognizes "*" as a wildcard sentinel.
392pub fn get_string_or_array_or_wildcard(args: &Value, key: &str) -> Option<IdList> {
393    let vals = get_string_or_array(args, key)?;
394    if vals.len() == 1 && vals[0] == "*" {
395        Some(IdList::Wildcard)
396    } else {
397        Some(IdList::Ids(vals))
398    }
399}