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