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