Coverage for mcpgateway/services/root_service.py: 92%

69 statements  

« prev     ^ index     » next       coverage.py v7.9.1, created at 2025-06-22 12:53 +0100

1# -*- coding: utf-8 -*- 

2"""Root Service Implementation. 

3 

4Copyright 2025 

5SPDX-License-Identifier: Apache-2.0 

6Authors: Mihai Criveti 

7 

8This module implements root directory management according to the MCP specification. 

9It handles root registration, validation, and change notifications. 

10""" 

11 

12import asyncio 

13import logging 

14import os 

15from typing import AsyncGenerator, Dict, List, Optional 

16from urllib.parse import urlparse 

17 

18from mcpgateway.config import settings 

19from mcpgateway.types import Root 

20 

21logger = logging.getLogger(__name__) 

22 

23 

24class RootServiceError(Exception): 

25 """Base class for root service errors.""" 

26 

27 

28class RootService: 

29 """MCP root service. 

30 

31 Manages roots that can be exposed to MCP clients. 

32 Handles: 

33 - Root registration and validation 

34 - Change notifications 

35 - Root permissions and access control 

36 """ 

37 

38 def __init__(self): 

39 """Initialize root service.""" 

40 self._roots: Dict[str, Root] = {} 

41 self._subscribers: List[asyncio.Queue] = [] 

42 

43 async def initialize(self) -> None: 

44 """Initialize root service.""" 

45 logger.info("Initializing root service") 

46 # Add any configured default roots 

47 for root_uri in settings.default_roots: 

48 try: 

49 await self.add_root(root_uri) 

50 except RootServiceError as e: 

51 logger.error(f"Failed to add default root {root_uri}: {e}") 

52 

53 async def shutdown(self) -> None: 

54 """Shutdown root service.""" 

55 logger.info("Shutting down root service") 

56 # Clear all roots and subscribers 

57 self._roots.clear() 

58 self._subscribers.clear() 

59 

60 async def list_roots(self) -> List[Root]: 

61 """List available roots. 

62 

63 Returns: 

64 List of registered roots 

65 """ 

66 return list(self._roots.values()) 

67 

68 async def add_root(self, uri: str, name: Optional[str] = None) -> Root: 

69 """Add a new root. 

70 

71 Args: 

72 uri: Root URI 

73 name: Optional root name 

74 

75 Returns: 

76 Created root object 

77 

78 Raises: 

79 RootServiceError: If root is invalid or already exists 

80 """ 

81 try: 

82 root_uri = self._make_root_uri(uri) 

83 except ValueError as e: 

84 raise RootServiceError(f"Invalid root URI: {e}") 

85 

86 if root_uri in self._roots: 

87 raise RootServiceError(f"Root already exists: {root_uri}") 

88 

89 # Skip any access check; just store the key/value. 

90 root_obj = Root( 

91 uri=root_uri, 

92 name=name or os.path.basename(urlparse(root_uri).path) or root_uri, 

93 ) 

94 self._roots[root_uri] = root_obj 

95 

96 await self._notify_root_added(root_obj) 

97 logger.info(f"Added root: {root_uri}") 

98 return root_obj 

99 

100 async def remove_root(self, root_uri: str) -> None: 

101 """Remove a registered root. 

102 

103 Args: 

104 root_uri: Root URI to remove 

105 

106 Raises: 

107 RootServiceError: If root not found 

108 """ 

109 if root_uri not in self._roots: 

110 raise RootServiceError(f"Root not found: {root_uri}") 

111 root_obj = self._roots.pop(root_uri) 

112 await self._notify_root_removed(root_obj) 

113 logger.info(f"Removed root: {root_uri}") 

114 

115 async def subscribe_changes(self) -> AsyncGenerator[Dict, None]: 

116 """Subscribe to root changes. 

117 

118 Yields: 

119 Root change events 

120 """ 

121 queue: asyncio.Queue = asyncio.Queue() 

122 self._subscribers.append(queue) 

123 try: 

124 while True: 

125 event = await queue.get() 

126 yield event 

127 finally: 

128 self._subscribers.remove(queue) 

129 

130 def _make_root_uri(self, uri: str) -> str: 

131 """Convert input to a valid URI. 

132 

133 If no scheme is provided, assume a file URI and convert the path to an absolute path. 

134 

135 Args: 

136 uri: Input URI or filesystem path 

137 

138 Returns: 

139 A valid URI string 

140 """ 

141 parsed = urlparse(uri) 

142 if not parsed.scheme: 

143 # No scheme provided; assume a file URI. 

144 return f"file://{uri}" 

145 # If a scheme is present (e.g., http, https, ftp, etc.), return the URI as-is. 

146 return uri 

147 

148 async def _notify_root_added(self, root: Root) -> None: 

149 """Notify subscribers of root addition. 

150 

151 Args: 

152 root: Added root 

153 """ 

154 event = {"type": "root_added", "data": {"uri": root.uri, "name": root.name}} 

155 await self._notify_subscribers(event) 

156 

157 async def _notify_root_removed(self, root: Root) -> None: 

158 """Notify subscribers of root removal. 

159 

160 Args: 

161 root: Removed root 

162 """ 

163 event = {"type": "root_removed", "data": {"uri": root.uri}} 

164 await self._notify_subscribers(event) 

165 

166 async def _notify_subscribers(self, event: Dict) -> None: 

167 """Send event to all subscribers. 

168 

169 Args: 

170 event: Event to send 

171 """ 

172 for queue in self._subscribers: 

173 try: 

174 await queue.put(event) 

175 except Exception as e: 

176 logger.error(f"Failed to notify subscriber: {e}")