Skip to main content

mcp_protocol/
server.rs

1//! # MCP Server
2//!
3//! **MCP server implementation** for exposing tools and memory operations
4//! - **HTTP Server**: Standard HTTP/JSON-RPC (default)
5//! - **SSE Server**: Direct SSE connections (feature: "sse-server")
6
7use crate::{
8    AuthHandler, CallToolRequest, CallToolResult, InitializeRequest, InitializeResult,
9    ListToolsResult, ServerCapabilities, ServerInfo, ToolProvider,
10};
11use protocol_transport_core::{ProtocolError, UniversalRequest, UniversalResponse};
12use serde_json::json;
13use std::collections::HashMap;
14
15#[cfg(feature = "sse-server")]
16use protocol_transport_core::{SseTransport, Transport, TransportFactory};
17
18/// **MCP Server** - Expose tools and capabilities via MCP protocol
19pub struct McpServer {
20    /// Server capabilities
21    capabilities: ServerCapabilities,
22    /// Authentication handler
23    auth_handler: Option<Box<dyn AuthHandler>>,
24    /// Tool provider
25    tool_provider: Option<Box<dyn ToolProvider>>,
26    /// Server info
27    server_info: ServerInfo,
28
29    #[cfg(feature = "sse-server")]
30    /// SSE transport for direct client connections
31    sse_transport: Option<SseTransport>,
32
33    #[cfg(feature = "sse-server")]
34    /// Server bind address
35    bind_address: Option<String>,
36}
37
38impl McpServer {
39    /// Create new MCP server
40    pub fn new() -> Self {
41        Self {
42            capabilities: ServerCapabilities::default(),
43            auth_handler: None,
44            tool_provider: None,
45            server_info: ServerInfo {
46                name: "promptfleet-mcp-server".to_string(),
47                version: env!("CARGO_PKG_VERSION").to_string(),
48                description: Some("PromptFleet MCP Server".to_string()),
49            },
50
51            #[cfg(feature = "sse-server")]
52            sse_transport: None,
53
54            #[cfg(feature = "sse-server")]
55            bind_address: None,
56        }
57    }
58
59    /// Configure server capabilities
60    pub fn with_capabilities(mut self, capabilities: ServerCapabilities) -> Self {
61        self.capabilities = capabilities;
62        self
63    }
64
65    /// Configure authentication handler
66    pub fn with_auth_handler<H: AuthHandler + 'static>(mut self, handler: H) -> Self {
67        self.auth_handler = Some(Box::new(handler));
68        self
69    }
70
71    /// Configure tool provider
72    pub fn with_tool_provider<P: ToolProvider + 'static>(mut self, provider: P) -> Self {
73        self.tool_provider = Some(Box::new(provider));
74        self
75    }
76
77    /// Configure server info
78    pub fn with_server_info(mut self, server_info: ServerInfo) -> Self {
79        self.server_info = server_info;
80        self
81    }
82
83    /// Bind SSE server to address (feature: "sse-server")
84    #[cfg(feature = "sse-server")]
85    pub fn with_sse_server(mut self, bind_address: &str) -> Self {
86        self.sse_transport = Some(TransportFactory::mcp_sse(bind_address));
87        self.bind_address = Some(bind_address.to_string());
88        self
89    }
90
91    /// Handle incoming MCP request
92    pub fn handle_request(
93        &self,
94        request: &UniversalRequest,
95    ) -> Result<UniversalResponse, ProtocolError> {
96        // Check authentication if configured
97        if let Some(ref auth_handler) = self.auth_handler {
98            auth_handler.validate_request(request)?;
99        }
100
101        // Parse JSON-RPC request
102        let request_body = String::from_utf8(request.body.clone())
103            .map_err(|e| ProtocolError::Parsing(format!("Invalid UTF-8 request: {}", e)))?;
104
105        let request_json: serde_json::Value = serde_json::from_str(&request_body)
106            .map_err(|e| ProtocolError::Parsing(format!("Invalid JSON request: {}", e)))?;
107
108        let method = request_json
109            .get("method")
110            .and_then(|m| m.as_str())
111            .ok_or_else(|| ProtocolError::Parsing("Missing 'method' field".to_string()))?;
112
113        let params = request_json.get("params").cloned().unwrap_or(json!({}));
114
115        let id = request_json.get("id").cloned();
116
117        // Handle MCP methods
118        let response_json = match method {
119            "initialize" => self.handle_initialize(params, id)?,
120            "tools/list" => self.handle_list_tools(params, id)?,
121            "tools/call" => {
122                #[cfg(not(target_arch = "wasm32"))]
123                {
124                    tokio::runtime::Handle::current().block_on(self.handle_call_tool(params, id))?
125                }
126                #[cfg(target_arch = "wasm32")]
127                {
128                    return Err(ProtocolError::internal_error(
129                        "Sync MCP server handle_request is not supported in WASM; use async handler",
130                    ));
131                }
132            }
133            _ => json!({
134                "jsonrpc": "2.0",
135                "id": id,
136                "error": {
137                    "code": -32601,
138                    "message": format!("Method '{}' not found", method)
139                }
140            }),
141        };
142
143        // Build response
144        let response_body = response_json.to_string().into_bytes();
145        let mut headers = HashMap::new();
146        headers.insert("content-type".to_string(), "application/json".to_string());
147        headers.insert("x-protocol".to_string(), "MCP".to_string());
148
149        Ok(UniversalResponse {
150            status: 200,
151            headers,
152            body: response_body,
153            protocol: "MCP".to_string(),
154            correlation_id: request.correlation_id.clone(),
155        })
156    }
157
158    fn handle_initialize(
159        &self,
160        params: serde_json::Value,
161        id: Option<serde_json::Value>,
162    ) -> Result<serde_json::Value, ProtocolError> {
163        let _init_request: InitializeRequest = serde_json::from_value(params)
164            .map_err(|e| ProtocolError::Parsing(format!("Invalid initialize request: {}", e)))?;
165
166        let result = InitializeResult {
167            protocol_version: crate::MCP_PROTOCOL_VERSION.to_string(),
168            capabilities: self.capabilities.clone(),
169            server_info: self.server_info.clone(),
170        };
171
172        Ok(json!({
173            "jsonrpc": "2.0",
174            "id": id,
175            "result": result
176        }))
177    }
178
179    fn handle_list_tools(
180        &self,
181        _params: serde_json::Value,
182        id: Option<serde_json::Value>,
183    ) -> Result<serde_json::Value, ProtocolError> {
184        let tools = match &self.tool_provider {
185            Some(provider) => provider.list_tools()?,
186            None => vec![], // No provider configured
187        };
188
189        let result = ListToolsResult { tools };
190
191        Ok(json!({
192            "jsonrpc": "2.0",
193            "id": id,
194            "result": result
195        }))
196    }
197
198    async fn handle_call_tool(
199        &self,
200        params: serde_json::Value,
201        id: Option<serde_json::Value>,
202    ) -> Result<serde_json::Value, ProtocolError> {
203        let call_request: CallToolRequest = serde_json::from_value(params)
204            .map_err(|e| ProtocolError::Parsing(format!("Invalid call_tool request: {}", e)))?;
205
206        let result = match &self.tool_provider {
207            Some(provider) => {
208                provider
209                    .call_tool(&call_request.name, call_request.arguments)
210                    .await?
211            }
212            None => CallToolResult {
213                content: vec![crate::Content::text("No tool provider configured")],
214                is_error: Some(true),
215            },
216        };
217
218        Ok(json!({
219            "jsonrpc": "2.0",
220            "id": id,
221            "result": result
222        }))
223    }
224
225    /// Start SSE server (feature: "sse-server")
226    #[cfg(feature = "sse-server")]
227    pub async fn start_sse_server(&self) -> Result<(), ProtocolError> {
228        if let Some(ref transport) = self.sse_transport {
229            log::info!(
230                "Starting MCP SSE Server on {}",
231                self.bind_address.as_deref().unwrap_or("unknown")
232            );
233
234            transport.health_check().await.map_err(|e| {
235                ProtocolError::internal_error(&format!("Failed to start SSE server: {:?}", e))
236            })
237        } else {
238            Err(ProtocolError::internal_error("No SSE transport configured"))
239        }
240    }
241
242    /// Stop SSE server (feature: "sse-server")
243    #[cfg(feature = "sse-server")]
244    pub async fn stop_sse_server(&self) -> Result<(), ProtocolError> {
245        log::info!("Stopping MCP SSE Server");
246        Ok(())
247    }
248}
249
250impl Default for McpServer {
251    fn default() -> Self {
252        Self::new()
253    }
254}
255
256/// **MCP Server Builder** - Convenient server configuration
257pub struct McpServerBuilder {
258    capabilities: ServerCapabilities,
259    auth_handler: Option<Box<dyn AuthHandler>>,
260    tool_provider: Option<Box<dyn ToolProvider>>,
261    server_info: ServerInfo,
262
263    #[cfg(feature = "sse-server")]
264    bind_address: Option<String>,
265}
266
267impl McpServerBuilder {
268    /// Create new server builder
269    pub fn new() -> Self {
270        Self {
271            capabilities: ServerCapabilities::default(),
272            auth_handler: None,
273            tool_provider: None,
274            server_info: ServerInfo {
275                name: "promptfleet-mcp-server".to_string(),
276                version: env!("CARGO_PKG_VERSION").to_string(),
277                description: Some("PromptFleet MCP Server".to_string()),
278            },
279
280            #[cfg(feature = "sse-server")]
281            bind_address: None,
282        }
283    }
284
285    /// Set server capabilities
286    pub fn with_capabilities(mut self, capabilities: ServerCapabilities) -> Self {
287        self.capabilities = capabilities;
288        self
289    }
290
291    /// Set authentication handler
292    pub fn with_auth_handler<H: AuthHandler + 'static>(mut self, handler: H) -> Self {
293        self.auth_handler = Some(Box::new(handler));
294        self
295    }
296
297    /// Set tool provider
298    pub fn with_tool_provider<P: ToolProvider + 'static>(mut self, provider: P) -> Self {
299        self.tool_provider = Some(Box::new(provider));
300        self
301    }
302
303    /// Set server info
304    pub fn with_server_info(mut self, server_info: ServerInfo) -> Self {
305        self.server_info = server_info;
306        self
307    }
308
309    /// Bind SSE server to address (feature: "sse-server")
310    #[cfg(feature = "sse-server")]
311    pub fn with_sse_server(mut self, bind_address: &str) -> Self {
312        self.bind_address = Some(bind_address.to_string());
313        self
314    }
315
316    /// Build the MCP server
317    pub fn build(self) -> McpServer {
318        let mut server = McpServer::new()
319            .with_capabilities(self.capabilities)
320            .with_server_info(self.server_info);
321
322        if let Some(handler) = self.auth_handler {
323            server.auth_handler = Some(handler);
324        }
325
326        if let Some(provider) = self.tool_provider {
327            server.tool_provider = Some(provider);
328        }
329
330        #[cfg(feature = "sse-server")]
331        {
332            if let Some(ref address) = self.bind_address {
333                server = server.with_sse_server(address);
334            }
335        }
336
337        server
338    }
339}
340
341impl Default for McpServerBuilder {
342    fn default() -> Self {
343        Self::new()
344    }
345}