|
| 1 | +use rmcp::{ |
| 2 | + ErrorData as McpError, ServerHandler, |
| 3 | + handler::server::{router::tool::ToolRouter, wrapper::Parameters}, |
| 4 | + model::*, |
| 5 | + schemars, tool, tool_handler, tool_router, |
| 6 | +}; |
| 7 | +use serde::{Deserialize, Serialize}; |
| 8 | +use std::sync::Arc; |
| 9 | +use tokio::sync::Mutex; |
| 10 | + |
| 11 | +#[derive(Debug, Clone, Serialize, Deserialize)] |
| 12 | +struct Task { |
| 13 | + id: usize, |
| 14 | + title: String, |
| 15 | + description: String, |
| 16 | + completed: bool, |
| 17 | +} |
| 18 | + |
| 19 | +#[derive(Debug, Clone)] |
| 20 | +struct TaskManager { |
| 21 | + tasks: Arc<Mutex<Vec<Task>>>, |
| 22 | + next_id: Arc<Mutex<usize>>, |
| 23 | + tool_router: ToolRouter<TaskManager>, |
| 24 | +} |
| 25 | + |
| 26 | +#[derive(Debug, Deserialize, schemars::JsonSchema)] |
| 27 | +struct AddTaskRequest { |
| 28 | + #[schemars(description = "The title of the task")] |
| 29 | + title: String, |
| 30 | + #[schemars(description = "A detailed description of the task")] |
| 31 | + description: String, |
| 32 | +} |
| 33 | + |
| 34 | +#[derive(Debug, Deserialize, schemars::JsonSchema)] |
| 35 | +struct CompleteTaskRequest { |
| 36 | + #[schemars(description = "The ID of the task to mark as completed")] |
| 37 | + id: usize, |
| 38 | +} |
| 39 | + |
| 40 | +#[derive(Debug, Deserialize, schemars::JsonSchema)] |
| 41 | +struct GetTaskRequest { |
| 42 | + #[schemars(description = "The ID of the task to retrieve")] |
| 43 | + id: usize, |
| 44 | +} |
| 45 | + |
| 46 | +#[tool_router] |
| 47 | +impl TaskManager { |
| 48 | + fn new() -> Self { |
| 49 | + Self { |
| 50 | + tasks: Arc::new(Mutex::new(Vec::new())), |
| 51 | + next_id: Arc::new(Mutex::new(1)), |
| 52 | + tool_router: Self::tool_router(), |
| 53 | + } |
| 54 | + } |
| 55 | + |
| 56 | + #[tool(description = "Add a new task to the task manager")] |
| 57 | + async fn add_task( |
| 58 | + &self, |
| 59 | + Parameters(AddTaskRequest { title, description }): Parameters<AddTaskRequest>, |
| 60 | + ) -> Result<CallToolResult, McpError> { |
| 61 | + let mut tasks = self.tasks.lock().await; |
| 62 | + let mut next_id = self.next_id.lock().await; |
| 63 | + |
| 64 | + let task = Task { |
| 65 | + id: *next_id, |
| 66 | + title, |
| 67 | + description, |
| 68 | + completed: false, |
| 69 | + }; |
| 70 | + |
| 71 | + *next_id += 1; |
| 72 | + tasks.push(task.clone()); |
| 73 | + |
| 74 | + let response = serde_json::json!({ |
| 75 | + "success": true, |
| 76 | + "task": task, |
| 77 | + "message": format!("Task '{}' added successfully with ID {}", task.title, task.id) |
| 78 | + }); |
| 79 | + |
| 80 | + Ok(CallToolResult::success(vec![Content::text( |
| 81 | + serde_json::to_string_pretty(&response).unwrap(), |
| 82 | + )])) |
| 83 | + } |
| 84 | + |
| 85 | + #[tool(description = "Mark a task as completed")] |
| 86 | + async fn complete_task( |
| 87 | + &self, |
| 88 | + Parameters(CompleteTaskRequest { id }): Parameters<CompleteTaskRequest>, |
| 89 | + ) -> Result<CallToolResult, McpError> { |
| 90 | + let mut tasks = self.tasks.lock().await; |
| 91 | + |
| 92 | + if let Some(task) = tasks.iter_mut().find(|t| t.id == id) { |
| 93 | + task.completed = true; |
| 94 | + let response = serde_json::json!({ |
| 95 | + "success": true, |
| 96 | + "task": task, |
| 97 | + "message": format!("Task '{}' marked as completed", task.title) |
| 98 | + }); |
| 99 | + |
| 100 | + Ok(CallToolResult::success(vec![Content::text( |
| 101 | + serde_json::to_string_pretty(&response).unwrap(), |
| 102 | + )])) |
| 103 | + } else { |
| 104 | + Err(McpError::invalid_params( |
| 105 | + format!("Task with ID {} not found", id), |
| 106 | + None, |
| 107 | + )) |
| 108 | + } |
| 109 | + } |
| 110 | + |
| 111 | + #[tool(description = "List all tasks in the task manager")] |
| 112 | + async fn list_tasks(&self) -> Result<CallToolResult, McpError> { |
| 113 | + let tasks = self.tasks.lock().await; |
| 114 | + |
| 115 | + let response = serde_json::json!({ |
| 116 | + "total": tasks.len(), |
| 117 | + "tasks": tasks.clone() |
| 118 | + }); |
| 119 | + |
| 120 | + Ok(CallToolResult::success(vec![Content::text( |
| 121 | + serde_json::to_string_pretty(&response).unwrap(), |
| 122 | + )])) |
| 123 | + } |
| 124 | + |
| 125 | + #[tool(description = "Get a specific task by ID")] |
| 126 | + async fn get_task( |
| 127 | + &self, |
| 128 | + Parameters(GetTaskRequest { id }): Parameters<GetTaskRequest>, |
| 129 | + ) -> Result<CallToolResult, McpError> { |
| 130 | + let tasks = self.tasks.lock().await; |
| 131 | + |
| 132 | + if let Some(task) = tasks.iter().find(|t| t.id == id) { |
| 133 | + let response = serde_json::json!({ |
| 134 | + "success": true, |
| 135 | + "task": task |
| 136 | + }); |
| 137 | + |
| 138 | + Ok(CallToolResult::success(vec![Content::text( |
| 139 | + serde_json::to_string_pretty(&response).unwrap(), |
| 140 | + )])) |
| 141 | + } else { |
| 142 | + Err(McpError::invalid_params( |
| 143 | + format!("Task with ID {} not found", id), |
| 144 | + None, |
| 145 | + )) |
| 146 | + } |
| 147 | + } |
| 148 | +} |
| 149 | + |
| 150 | +#[tool_handler] |
| 151 | +impl ServerHandler for TaskManager { |
| 152 | + fn get_info(&self) -> ServerInfo { |
| 153 | + ServerInfo { |
| 154 | + protocol_version: ProtocolVersion::V_2024_11_05, |
| 155 | + capabilities: ServerCapabilities::builder().enable_tools().build(), |
| 156 | + server_info: Implementation { |
| 157 | + name: "task-manager".to_string(), |
| 158 | + version: "0.1.0".to_string(), |
| 159 | + title: None, |
| 160 | + website_url: None, |
| 161 | + icons: None, |
| 162 | + }, |
| 163 | + instructions: Some( |
| 164 | + "A task manager MCP server that allows you to add, complete, list, and retrieve tasks with real-time updates." |
| 165 | + .to_string(), |
| 166 | + ), |
| 167 | + } |
| 168 | + } |
| 169 | +} |
| 170 | + |
| 171 | +#[shuttle_runtime::main] |
| 172 | +async fn main() -> shuttle_axum::ShuttleAxum { |
| 173 | + tracing::info!("Starting Task Manager MCP Server"); |
| 174 | + |
| 175 | + let service = rmcp::transport::streamable_http_server::StreamableHttpService::new( |
| 176 | + || Ok(TaskManager::new()), |
| 177 | + rmcp::transport::streamable_http_server::session::local::LocalSessionManager::default() |
| 178 | + .into(), |
| 179 | + Default::default(), |
| 180 | + ); |
| 181 | + |
| 182 | + let router = axum::Router::new().nest_service("/mcp", service); |
| 183 | + |
| 184 | + Ok(router.into()) |
| 185 | +} |
0 commit comments