Coverage for src/quickmcp/server.py: 67%

238 statements  

« prev     ^ index     » next       coverage.py v7.10.4, created at 2025-08-20 20:56 +0200

1""" 

2QuickMCP Server - Simplified MCP server creation 

3""" 

4 

5import asyncio 

6import logging 

7from typing import Optional, Dict, Any, Callable, List, Union 

8from pathlib import Path 

9import sys 

10import json 

11import inspect 

12 

13from mcp.server import Server, InitializationOptions 

14from mcp.server.stdio import stdio_server 

15from pydantic import BaseModel 

16from mcp.types import Tool, Resource, Prompt, TextContent, ServerCapabilities 

17 

18from .autodiscovery import AutoDiscovery 

19 

20logger = logging.getLogger(__name__) 

21 

22 

23class QuickMCPServer: 

24 """ 

25 A simplified wrapper around the MCP Server class that makes it easy 

26 to create and run MCP servers with minimal boilerplate. 

27  

28 Example: 

29 ```python 

30 server = QuickMCPServer("my-server") 

31  

32 @server.tool() 

33 def add(a: int, b: int) -> int: 

34 return a + b 

35  

36 @server.resource("data://{key}") 

37 def get_data(key: str) -> str: 

38 return f"Data for {key}" 

39  

40 server.run() 

41 ``` 

42 """ 

43 

44 def __init__( 

45 self, 

46 name: str, 

47 version: str = "1.0.0", 

48 description: Optional[str] = None, 

49 log_level: str = "INFO", 

50 enable_autodiscovery: bool = True, 

51 discovery_metadata: Optional[Dict[str, Any]] = None, 

52 ): 

53 """ 

54 Initialize a QuickMCP server. 

55  

56 Args: 

57 name: The name of the server 

58 version: Server version 

59 description: Optional server description 

60 log_level: Logging level (DEBUG, INFO, WARNING, ERROR) 

61 enable_autodiscovery: Enable network autodiscovery 

62 discovery_metadata: Additional metadata for discovery 

63 """ 

64 self.name = name 

65 self.version = version 

66 self.description = description or f"{name} MCP Server" 

67 

68 # Set up logging - always use stderr to avoid interfering with stdio transport 

69 logging.basicConfig( 

70 level=getattr(logging, log_level.upper()), 

71 format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", 

72 stream=sys.stderr 

73 ) 

74 self.logger = logging.getLogger(name) 

75 

76 # Create the underlying MCP server 

77 self._server = Server(name, version=version) 

78 

79 # Track registered components 

80 self._tools: Dict[str, Callable] = {} 

81 self._resources: Dict[str, Callable] = {} 

82 self._prompts: Dict[str, Callable] = {} 

83 

84 # Store metadata 

85 self._tool_metadata: Dict[str, Dict[str, Any]] = {} 

86 self._resource_metadata: Dict[str, Dict[str, Any]] = {} 

87 self._prompt_metadata: Dict[str, Dict[str, Any]] = {} 

88 

89 # Register handlers with the MCP server 

90 self._register_handlers() 

91 

92 # Autodiscovery setup 

93 self.enable_autodiscovery = enable_autodiscovery 

94 self.discovery_metadata = discovery_metadata or {} 

95 self._autodiscovery: Optional[AutoDiscovery] = None 

96 

97 self.logger.info(f"Initialized QuickMCP server: {name} v{version}") 

98 

99 def _register_handlers(self): 

100 """Register the request handlers with the MCP server.""" 

101 

102 @self._server.list_tools() 

103 async def list_tools(): 

104 """List available tools.""" 

105 tools = [] 

106 for tool_name, func in self._tools.items(): 

107 metadata = self._tool_metadata.get(tool_name, {}) 

108 tools.append(Tool( 

109 name=tool_name, 

110 description=metadata.get("description", ""), 

111 inputSchema=metadata.get("schema", { 

112 "type": "object", 

113 "properties": {}, 

114 "required": [] 

115 }) 

116 )) 

117 return tools 

118 

119 @self._server.call_tool() 

120 async def call_tool(name: str, arguments: dict): 

121 """Execute a tool.""" 

122 if name not in self._tools: 

123 raise ValueError(f"Tool {name} not found") 

124 

125 func = self._tools[name] 

126 

127 # Execute the function 

128 if asyncio.iscoroutinefunction(func): 

129 result = await func(**arguments) 

130 else: 

131 result = func(**arguments) 

132 

133 # Return as TextContent 

134 return [TextContent( 

135 type="text", 

136 text=json.dumps(result) if not isinstance(result, str) else result 

137 )] 

138 

139 @self._server.list_resources() 

140 async def list_resources(): 

141 """List available resources.""" 

142 resources = [] 

143 for uri_template, func in self._resources.items(): 

144 metadata = self._resource_metadata.get(uri_template, {}) 

145 resources.append(Resource( 

146 uri=uri_template, 

147 name=metadata.get("name", func.__name__), 

148 description=metadata.get("description", ""), 

149 mimeType=metadata.get("mime_type", "text/plain") 

150 )) 

151 return resources 

152 

153 @self._server.read_resource() 

154 async def read_resource(uri: str): 

155 """Read a resource.""" 

156 # Find matching resource 

157 for uri_template, func in self._resources.items(): 

158 # Simple template matching (replace {param} with actual values) 

159 # This is a simplified version - production code would need proper URI template matching 

160 if self._uri_matches(uri, uri_template): 

161 params = self._extract_params(uri, uri_template) 

162 

163 if asyncio.iscoroutinefunction(func): 

164 content = await func(**params) 

165 else: 

166 content = func(**params) 

167 

168 metadata = self._resource_metadata.get(uri_template, {}) 

169 return TextContent( 

170 type="text", 

171 text=content, 

172 mimeType=metadata.get("mime_type", "text/plain") 

173 ) 

174 

175 raise ValueError(f"Resource {uri} not found") 

176 

177 @self._server.list_prompts() 

178 async def list_prompts(): 

179 """List available prompts.""" 

180 prompts = [] 

181 for prompt_name, func in self._prompts.items(): 

182 metadata = self._prompt_metadata.get(prompt_name, {}) 

183 prompts.append(Prompt( 

184 name=prompt_name, 

185 description=metadata.get("description", ""), 

186 arguments=metadata.get("arguments", []) 

187 )) 

188 return prompts 

189 

190 @self._server.get_prompt() 

191 async def get_prompt(name: str, arguments: dict): 

192 """Get a prompt.""" 

193 if name not in self._prompts: 

194 raise ValueError(f"Prompt {name} not found") 

195 

196 func = self._prompts[name] 

197 

198 if asyncio.iscoroutinefunction(func): 

199 content = await func(**arguments) 

200 else: 

201 content = func(**arguments) 

202 

203 return TextContent( 

204 type="text", 

205 text=content 

206 ) 

207 

208 def _uri_matches(self, uri: str, template: str) -> bool: 

209 """Check if a URI matches a template pattern.""" 

210 # Simple implementation - just check if the base matches 

211 # A full implementation would parse the template properly 

212 import re 

213 pattern = re.sub(r'\{[^}]+\}', r'[^/]+', template) 

214 return bool(re.match(f"^{pattern}$", uri)) 

215 

216 def _extract_params(self, uri: str, template: str) -> Dict[str, str]: 

217 """Extract parameters from a URI based on a template.""" 

218 # Simple implementation 

219 import re 

220 

221 # Find parameter names in template 

222 param_names = re.findall(r'\{([^}]+)\}', template) 

223 

224 # Create regex pattern from template 

225 pattern = template 

226 for param_name in param_names: 

227 pattern = pattern.replace(f"{{{param_name}}}", r'([^/]+)') 

228 

229 # Extract values 

230 match = re.match(f"^{pattern}$", uri) 

231 if match: 

232 return dict(zip(param_names, match.groups())) 

233 return {} 

234 

235 def tool( 

236 self, 

237 name: Optional[str] = None, 

238 description: Optional[str] = None, 

239 schema: Optional[Dict[str, Any]] = None, 

240 ) -> Callable: 

241 """ 

242 Decorator to register a function as an MCP tool. 

243  

244 Args: 

245 name: Optional tool name (defaults to function name) 

246 description: Optional tool description (defaults to docstring) 

247 schema: Optional JSON schema for parameters 

248  

249 Example: 

250 ```python 

251 @server.tool() 

252 def calculate(operation: str, a: float, b: float) -> float: 

253 '''Perform a calculation''' 

254 if operation == "add": 

255 return a + b 

256 elif operation == "multiply": 

257 return a * b 

258 ``` 

259 """ 

260 def decorator(func: Callable) -> Callable: 

261 tool_name = name or func.__name__ 

262 tool_desc = description or (func.__doc__ or "").strip() 

263 

264 # Generate schema from function signature if not provided 

265 tool_schema = schema 

266 if tool_schema is None: 

267 tool_schema = self._generate_schema_from_function(func) 

268 

269 # Track in our registry 

270 self._tools[tool_name] = func 

271 self._tool_metadata[tool_name] = { 

272 "description": tool_desc, 

273 "schema": tool_schema 

274 } 

275 self.logger.debug(f"Registered tool: {tool_name}") 

276 

277 return func 

278 

279 return decorator 

280 

281 def resource( 

282 self, 

283 uri_template: str, 

284 name: Optional[str] = None, 

285 description: Optional[str] = None, 

286 mime_type: str = "text/plain", 

287 ) -> Callable: 

288 """ 

289 Decorator to register a function as an MCP resource. 

290  

291 Args: 

292 uri_template: URI template for the resource (e.g., "file://{path}") 

293 name: Optional resource name 

294 description: Optional resource description 

295 mime_type: MIME type of the resource content 

296  

297 Example: 

298 ```python 

299 @server.resource("config://{section}") 

300 def get_config(section: str) -> str: 

301 '''Get configuration section''' 

302 return read_config_file(section) 

303 ``` 

304 """ 

305 def decorator(func: Callable) -> Callable: 

306 resource_name = name or func.__name__ 

307 resource_desc = description or (func.__doc__ or "").strip() 

308 

309 # Track in our registry 

310 self._resources[uri_template] = func 

311 self._resource_metadata[uri_template] = { 

312 "name": resource_name, 

313 "description": resource_desc, 

314 "mime_type": mime_type 

315 } 

316 self.logger.debug(f"Registered resource: {uri_template}") 

317 

318 return func 

319 

320 return decorator 

321 

322 def prompt( 

323 self, 

324 name: Optional[str] = None, 

325 description: Optional[str] = None, 

326 arguments: Optional[List[Dict[str, Any]]] = None, 

327 ) -> Callable: 

328 """ 

329 Decorator to register a function as an MCP prompt template. 

330  

331 Args: 

332 name: Optional prompt name (defaults to function name) 

333 description: Optional prompt description 

334 arguments: Optional list of argument schemas 

335  

336 Example: 

337 ```python 

338 @server.prompt() 

339 def analyze_code(language: str, code: str) -> str: 

340 '''Generate a code analysis prompt''' 

341 return f"Analyze this {language} code:\\n\\n{code}" 

342 ``` 

343 """ 

344 def decorator(func: Callable) -> Callable: 

345 prompt_name = name or func.__name__ 

346 prompt_desc = description or (func.__doc__ or "").strip() 

347 

348 # Generate arguments from function signature if not provided 

349 prompt_arguments = arguments 

350 if prompt_arguments is None: 

351 prompt_arguments = self._generate_arguments_from_function(func) 

352 

353 # Track in our registry 

354 self._prompts[prompt_name] = func 

355 self._prompt_metadata[prompt_name] = { 

356 "description": prompt_desc, 

357 "arguments": prompt_arguments 

358 } 

359 self.logger.debug(f"Registered prompt: {prompt_name}") 

360 

361 return func 

362 

363 return decorator 

364 

365 def _generate_schema_from_function(self, func: Callable) -> Dict[str, Any]: 

366 """Generate JSON schema from function signature.""" 

367 sig = inspect.signature(func) 

368 properties = {} 

369 required = [] 

370 

371 for param_name, param in sig.parameters.items(): 

372 if param_name == "self": 

373 continue 

374 

375 # Basic type mapping 

376 param_type = "string" # default 

377 if param.annotation != inspect.Parameter.empty: 

378 if param.annotation == int: 

379 param_type = "integer" 

380 elif param.annotation == float: 

381 param_type = "number" 

382 elif param.annotation == bool: 

383 param_type = "boolean" 

384 elif param.annotation == list: 

385 param_type = "array" 

386 elif param.annotation == dict: 

387 param_type = "object" 

388 

389 properties[param_name] = {"type": param_type} 

390 

391 if param.default == inspect.Parameter.empty: 

392 required.append(param_name) 

393 

394 return { 

395 "type": "object", 

396 "properties": properties, 

397 "required": required 

398 } 

399 

400 def _generate_arguments_from_function(self, func: Callable) -> List[Dict[str, Any]]: 

401 """Generate argument list from function signature.""" 

402 sig = inspect.signature(func) 

403 arguments = [] 

404 

405 for param_name, param in sig.parameters.items(): 

406 if param_name == "self": 

407 continue 

408 

409 arg = { 

410 "name": param_name, 

411 "required": param.default == inspect.Parameter.empty 

412 } 

413 arguments.append(arg) 

414 

415 return arguments 

416 

417 def add_tool_from_function( 

418 self, 

419 func: Callable, 

420 name: Optional[str] = None, 

421 description: Optional[str] = None, 

422 ) -> None: 

423 """ 

424 Add a tool from an existing function. 

425  

426 Args: 

427 func: The function to add as a tool 

428 name: Optional tool name 

429 description: Optional tool description 

430 """ 

431 tool_name = name or func.__name__ 

432 tool_desc = description or (func.__doc__ or "").strip() 

433 

434 self._tools[tool_name] = func 

435 self._tool_metadata[tool_name] = { 

436 "description": tool_desc, 

437 "schema": self._generate_schema_from_function(func) 

438 } 

439 

440 self.logger.info(f"Added tool from function: {tool_name}") 

441 

442 def list_tools(self) -> List[str]: 

443 """Get a list of registered tool names.""" 

444 return list(self._tools.keys()) 

445 

446 def list_resources(self) -> List[str]: 

447 """Get a list of registered resource URIs.""" 

448 return list(self._resources.keys()) 

449 

450 def list_prompts(self) -> List[str]: 

451 """Get a list of registered prompt names.""" 

452 return list(self._prompts.keys()) 

453 

454 def get_info(self) -> Dict[str, Any]: 

455 """ 

456 Get server information including registered components. 

457  

458 Returns: 

459 Dictionary with server info 

460 """ 

461 return { 

462 "name": self.name, 

463 "version": self.version, 

464 "description": self.description, 

465 "tools": self.list_tools(), 

466 "resources": self.list_resources(), 

467 "prompts": self.list_prompts(), 

468 } 

469 

470 def start_autodiscovery(self, transport: str = "stdio", host: str = "localhost", port: int = 8000) -> None: 

471 """ 

472 Start autodiscovery broadcasting. 

473  

474 Args: 

475 transport: Transport type 

476 host: Host for network transports 

477 port: Port for network transports 

478 """ 

479 if not self.enable_autodiscovery: 

480 return 

481 

482 if self._autodiscovery: 

483 self.logger.warning("Autodiscovery already started") 

484 return 

485 

486 # Create autodiscovery instance 

487 self._autodiscovery = AutoDiscovery( 

488 server_name=self.name, 

489 server_version=self.version, 

490 server_description=self.description, 

491 transport=transport, 

492 host=host, 

493 port=port, 

494 metadata=self.discovery_metadata 

495 ) 

496 

497 # Update capabilities 

498 self._autodiscovery.update_capabilities( 

499 tools=self.list_tools(), 

500 resources=self.list_resources(), 

501 prompts=self.list_prompts() 

502 ) 

503 

504 # Start broadcasting 

505 self._autodiscovery.start() 

506 self.logger.info("Autodiscovery broadcasting started") 

507 

508 def stop_autodiscovery(self) -> None: 

509 """Stop autodiscovery broadcasting.""" 

510 if self._autodiscovery: 

511 self._autodiscovery.stop() 

512 self._autodiscovery = None 

513 self.logger.info("Autodiscovery broadcasting stopped") 

514 

515 def run( 

516 self, 

517 transport: str = "stdio", 

518 host: str = "localhost", 

519 port: int = 8000, 

520 **kwargs 

521 ) -> None: 

522 """ 

523 Run the MCP server. 

524  

525 Args: 

526 transport: Transport type ("stdio", "sse", "http") 

527 host: Host for network transports 

528 port: Port for network transports 

529 **kwargs: Additional transport-specific arguments 

530  

531 Example: 

532 ```python 

533 # Run as stdio server (default) 

534 server.run() 

535  

536 # Run as SSE server 

537 server.run(transport="sse", port=8080) 

538 ``` 

539 """ 

540 # Check for --info flag (for discovery) 

541 if "--info" in sys.argv: 

542 info = { 

543 "name": self.name, 

544 "version": self.version, 

545 "description": self.description, 

546 "capabilities": { 

547 "tools": self.list_tools(), 

548 "resources": self.list_resources(), 

549 "prompts": self.list_prompts() 

550 }, 

551 "metadata": self.discovery_metadata 

552 } 

553 print(json.dumps(info)) 

554 sys.exit(0) 

555 

556 self.logger.info(f"Starting {self.name} server with {transport} transport") 

557 self.logger.info(f"Registered {len(self._tools)} tools, " 

558 f"{len(self._resources)} resources, " 

559 f"{len(self._prompts)} prompts") 

560 

561 # Start autodiscovery if enabled (but not for stdio transport) 

562 if transport != "stdio": 

563 self.start_autodiscovery(transport=transport, host=host, port=port) 

564 

565 try: 

566 if transport == "stdio": 

567 self.run_stdio() 

568 elif transport == "sse": 

569 self.run_sse(host=host, port=port, **kwargs) 

570 elif transport == "http": 

571 self.run_http(host=host, port=port, **kwargs) 

572 else: 

573 raise ValueError(f"Unknown transport: {transport}") 

574 finally: 

575 # Stop autodiscovery when server stops 

576 self.stop_autodiscovery() 

577 

578 def run_stdio(self) -> None: 

579 """Run the server with stdio transport.""" 

580 # Disable logging for stdio transport to avoid interfering with JSON-RPC 

581 logging.getLogger().setLevel(logging.CRITICAL) 

582 

583 async def run_async(): 

584 # Create initialization options 

585 init_options = InitializationOptions( 

586 server_name=self.name, 

587 server_version=self.version, 

588 capabilities=ServerCapabilities( 

589 tools={"listTools": True, "callTool": True} if self._tools else {}, 

590 resources={"listResources": True, "readResource": True} if self._resources else {}, 

591 prompts={"listPrompts": True, "getPrompt": True} if self._prompts else {} 

592 ), 

593 instructions=self.description 

594 ) 

595 

596 async with stdio_server() as (read_stream, write_stream): 

597 await self._server.run( 

598 read_stream=read_stream, 

599 write_stream=write_stream, 

600 initialization_options=init_options 

601 ) 

602 

603 asyncio.run(run_async()) 

604 

605 def run_sse(self, host: str = "localhost", port: int = 8000, **kwargs) -> None: 

606 """ 

607 Run the server with SSE (Server-Sent Events) transport. 

608  

609 Args: 

610 host: Host to bind to 

611 port: Port to bind to 

612 """ 

613 try: 

614 from mcp.server.sse import SseServerTransport 

615 from starlette.applications import Starlette 

616 from starlette.routing import Route 

617 import uvicorn 

618 except ImportError: 

619 self.logger.error("SSE transport requires 'http' extras: pip install quickmcp[http]") 

620 raise 

621 

622 self.logger.info(f"Running SSE server on http://{host}:{port}/sse") 

623 

624 # Create SSE transport 

625 transport = SseServerTransport(self._server) 

626 

627 # Create Starlette app 

628 app = Starlette( 

629 routes=[ 

630 Route("/sse", endpoint=transport.handle_sse, methods=["GET"]), 

631 Route("/", endpoint=self._health_check, methods=["GET"]), 

632 ] 

633 ) 

634 

635 # Run with uvicorn 

636 uvicorn.run(app, host=host, port=port, **kwargs) 

637 

638 def run_http(self, host: str = "localhost", port: int = 8000, **kwargs) -> None: 

639 """ 

640 Run the server with HTTP transport. 

641  

642 Args: 

643 host: Host to bind to 

644 port: Port to bind to 

645 """ 

646 # Similar to SSE but with regular HTTP endpoints 

647 raise NotImplementedError("HTTP transport not yet implemented") 

648 

649 async def _health_check(self, request) -> Dict[str, Any]: 

650 """Health check endpoint for network transports.""" 

651 from starlette.responses import JSONResponse 

652 return JSONResponse({ 

653 "status": "healthy", 

654 "server": self.name, 

655 "version": self.version, 

656 "tools": len(self._tools), 

657 "resources": len(self._resources), 

658 "prompts": len(self._prompts), 

659 }) 

660 

661 def export_openapi(self) -> Dict[str, Any]: 

662 """ 

663 Export server specification as OpenAPI schema. 

664  

665 Returns: 

666 OpenAPI specification dictionary 

667 """ 

668 return { 

669 "openapi": "3.0.0", 

670 "info": { 

671 "title": self.name, 

672 "version": self.version, 

673 "description": self.description, 

674 }, 

675 "paths": { 

676 "/tools": { 

677 "get": { 

678 "summary": "List available tools", 

679 "responses": { 

680 "200": { 

681 "description": "List of tools", 

682 "content": { 

683 "application/json": { 

684 "schema": { 

685 "type": "array", 

686 "items": {"type": "string"} 

687 } 

688 } 

689 } 

690 } 

691 } 

692 } 

693 }, 

694 # Add more endpoints as needed 

695 } 

696 } 

697 

698 def __repr__(self) -> str: 

699 return (f"QuickMCPServer(name='{self.name}', version='{self.version}', " 

700 f"tools={len(self._tools)}, resources={len(self._resources)}, " 

701 f"prompts={len(self._prompts)})")