1use 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
18pub struct McpServer {
20 capabilities: ServerCapabilities,
22 auth_handler: Option<Box<dyn AuthHandler>>,
24 tool_provider: Option<Box<dyn ToolProvider>>,
26 server_info: ServerInfo,
28
29 #[cfg(feature = "sse-server")]
30 sse_transport: Option<SseTransport>,
32
33 #[cfg(feature = "sse-server")]
34 bind_address: Option<String>,
36}
37
38impl McpServer {
39 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 pub fn with_capabilities(mut self, capabilities: ServerCapabilities) -> Self {
61 self.capabilities = capabilities;
62 self
63 }
64
65 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 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 pub fn with_server_info(mut self, server_info: ServerInfo) -> Self {
79 self.server_info = server_info;
80 self
81 }
82
83 #[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 pub fn handle_request(
93 &self,
94 request: &UniversalRequest,
95 ) -> Result<UniversalResponse, ProtocolError> {
96 if let Some(ref auth_handler) = self.auth_handler {
98 auth_handler.validate_request(request)?;
99 }
100
101 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 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 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![], };
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 #[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 #[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
256pub 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 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 pub fn with_capabilities(mut self, capabilities: ServerCapabilities) -> Self {
287 self.capabilities = capabilities;
288 self
289 }
290
291 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 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 pub fn with_server_info(mut self, server_info: ServerInfo) -> Self {
305 self.server_info = server_info;
306 self
307 }
308
309 #[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 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}