Coverage for mcp_bridge/tools/lsp/manager.py: 0%

1""" 

2Persistent LSP Server Manager 

3 

4Manages persistent Language Server Protocol (LSP) servers for improved performance. 

5Implements lazy initialization, JSON-RPC communication, and graceful shutdown. 

6 

7Architecture: 

8- Servers start on first use (lazy initialization) 

9- JSON-RPC over stdio using pygls BaseLanguageClient 

10- Supports Python (jedi-language-server) and TypeScript (typescript-language-server) 

11- Graceful shutdown on MCP server exit 

12- Health checks and idle timeout management 

13""" 

14 

15import asyncio 

16import logging 

17import threading 

18import time 

19from dataclasses import dataclass, field 

20from typing import Optional 

21 

22from lsprotocol.types import ( 

23 ClientCapabilities, 

24 InitializedParams, 

25 InitializeParams, 

26) 

27from pygls.client import JsonRPCClient 

28 

29logger = logging.getLogger(__name__) 

30 

31# Configuration for LSP server lifecycle management 

32LSP_CONFIG = { 

33 "idle_timeout": 1800, # 30 minutes 

34 "health_check_interval": 300, # 5 minutes 

35 "health_check_timeout": 5.0, 

36} 

37 

38 

39@dataclass 

40class LSPServer: 

41 """Metadata for a persistent LSP server.""" 

42 

43 name: str 

44 command: list[str] 

45 client: JsonRPCClient | None = None 

46 initialized: bool = False 

47 process: asyncio.subprocess.Process | None = None 

48 pid: int | None = None # Track subprocess PID for explicit cleanup 

49 last_used: float = field(default_factory=time.time) # Timestamp of last usage 

50 created_at: float = field(default_factory=time.time) # Timestamp of server creation 

51 

52 

53class LSPManager: 

54 """ 

55 Singleton manager for persistent LSP servers. 

56 

57 Implements: 

58 - Lazy server initialization (start on first use) 

59 - Process lifecycle management with GC protection 

60 - Exponential backoff for crash recovery 

61 - Graceful shutdown with signal handling 

62 - Health checks and idle server shutdown 

63 """ 

64 

65 _instance: Optional["LSPManager"] = None 

66 

67 def __new__(cls): 

68 if cls._instance is None: 

69 cls._instance = super().__new__(cls) 

70 return cls._instance 

71 

72 def __init__(self): 

73 if hasattr(self, "_initialized"): 

74 return 

75 self._initialized = True 

76 self._servers: dict[str, LSPServer] = {} 

77 self._lock = asyncio.Lock() 

78 self._restart_attempts: dict[str, int] = {} 

79 self._health_monitor_task: asyncio.Task | None = None 

80 

81 # Register available LSP servers 

82 self._register_servers() 

83 

84 def _register_servers(self): 

85 """Register available LSP server configurations.""" 

86 self._servers["python"] = LSPServer(name="python", command=["jedi-language-server"]) 

87 self._servers["typescript"] = LSPServer( 

88 name="typescript", command=["typescript-language-server", "--stdio"] 

89 ) 

90 

91 async def get_server(self, language: str) -> JsonRPCClient | None: 

92 """ 

93 Get or start a persistent LSP server for the given language. 

94 

95 Args: 

96 language: Language identifier (e.g., "python", "typescript") 

97 

98 Returns: 

99 JsonRPCClient instance or None if server unavailable 

100 """ 

101 if language not in self._servers: 

102 logger.warning(f"No LSP server configured for language: {language}") 

103 return None 

104 

105 server = self._servers[language] 

106 

107 # Return existing initialized server 

108 if server.initialized and server.client: 

109 # Update last_used timestamp 

110 server.last_used = time.time() 

111 # Start health monitor on first use 

112 if self._health_monitor_task is None or self._health_monitor_task.done(): 

113 self._health_monitor_task = asyncio.create_task(self._background_health_monitor()) 

114 return server.client 

115 

116 # Start server with lock to prevent race conditions 

117 async with self._lock: 

118 # Double-check after acquiring lock 

119 if server.initialized and server.client: 

120 server.last_used = time.time() 

121 if self._health_monitor_task is None or self._health_monitor_task.done(): 

122 self._health_monitor_task = asyncio.create_task(self._background_health_monitor()) 

123 return server.client 

124 

125 try: 

126 await self._start_server(server) 

127 # Start health monitor on first server creation 

128 if self._health_monitor_task is None or self._health_monitor_task.done(): 

129 self._health_monitor_task = asyncio.create_task(self._background_health_monitor()) 

130 return server.client 

131 except Exception as e: 

132 logger.error(f"Failed to start {language} LSP server: {e}") 

133 return None 

134 

135 async def _start_server(self, server: LSPServer): 

136 """ 

137 Start a persistent LSP server process. 

138 

139 Implements: 

140 - Process health validation after start 

141 - LSP initialization handshake 

142 - GC protection via persistent reference 

143 - Timestamp tracking for idle detection 

144 

145 Args: 

146 server: LSPServer metadata object 

147 """ 

148 try: 

149 # Create pygls client 

150 client = JsonRPCClient() 

151 

152 logger.info(f"Starting {server.name} LSP server: {' '.join(server.command)}") 

153 

154 # Start server process (start_io expects cmd as first arg, then *args) 

155 await client.start_io(server.command[0], *server.command[1:]) 

156 

157 # Brief delay for process startup 

158 await asyncio.sleep(0.2) 

159 

160 # Capture subprocess from client 

161 if not hasattr(client, "_server") or client._server is None: 

162 raise ConnectionError( 

163 f"{server.name} LSP server process not accessible after start_io()" 

164 ) 

165 

166 server.process = client._server 

167 server.pid = server.process.pid 

168 logger.debug(f"{server.name} LSP server started with PID: {server.pid}") 

169 

170 # Validate process is still running 

171 if server.process.returncode is not None: 

172 raise ConnectionError( 

173 f"{server.name} LSP server exited immediately (code {server.process.returncode})" 

174 ) 

175 

176 # Perform LSP initialization handshake 

177 init_params = InitializeParams( 

178 process_id=None, root_uri=None, capabilities=ClientCapabilities() 

179 ) 

180 

181 try: 

182 # Send initialize request via protocol 

183 response = await asyncio.wait_for( 

184 client.protocol.send_request_async("initialize", init_params), timeout=10.0 

185 ) 

186 

187 # Send initialized notification 

188 client.protocol.notify("initialized", InitializedParams()) 

189 

190 logger.info(f"{server.name} LSP server initialized: {response}") 

191 

192 except TimeoutError: 

193 raise ConnectionError(f"{server.name} LSP server initialization timed out") 

194 

195 # Store client reference (GC protection) 

196 server.client = client 

197 server.initialized = True 

198 server.created_at = time.time() 

199 server.last_used = time.time() 

200 

201 # Reset restart attempts on successful start 

202 self._restart_attempts[server.name] = 0 

203 

204 logger.info(f"{server.name} LSP server started successfully") 

205 

206 except Exception as e: 

207 logger.error(f"Failed to start {server.name} LSP server: {e}", exc_info=True) 

208 # Cleanup on failure 

209 if server.client: 

210 try: 

211 await server.client.stop() 

212 except: 

213 pass 

214 server.client = None 

215 server.initialized = False 

216 server.process = None 

217 server.pid = None 

218 raise 

219 

220 async def _restart_with_backoff(self, server: LSPServer): 

221 """ 

222 Restart a crashed LSP server with exponential backoff. 

223 

224 Strategy: delay = 2^attempt + jitter (max 60s) 

225 

226 Args: 

227 server: LSPServer to restart 

228 """ 

229 import random 

230 

231 attempt = self._restart_attempts.get(server.name, 0) 

232 self._restart_attempts[server.name] = attempt + 1 

233 

234 # Exponential backoff with jitter (max 60s) 

235 delay = min(60, (2**attempt) + random.uniform(0, 1)) 

236 

237 logger.warning( 

238 f"{server.name} LSP server crashed. Restarting in {delay:.2f}s (attempt {attempt + 1})" 

239 ) 

240 await asyncio.sleep(delay) 

241 

242 # Reset state before restart 

243 server.initialized = False 

244 server.client = None 

245 server.process = None 

246 server.pid = None 

247 

248 try: 

249 await self._start_server(server) 

250 except Exception as e: 

251 logger.error(f"Restart failed for {server.name}: {e}") 

252 

253 async def _health_check_server(self, server: LSPServer) -> bool: 

254 """ 

255 Perform health check on an LSP server. 

256 

257 Args: 

258 server: LSPServer to check 

259 

260 Returns: 

261 True if server is healthy, False otherwise 

262 """ 

263 if not server.initialized or not server.client: 

264 return False 

265 

266 try: 

267 # Simple health check: send initialize request 

268 # Most servers respond to repeated initialize calls 

269 init_params = InitializeParams( 

270 process_id=None, root_uri=None, capabilities=ClientCapabilities() 

271 ) 

272 response = await asyncio.wait_for( 

273 server.client.protocol.send_request_async("initialize", init_params), 

274 timeout=LSP_CONFIG["health_check_timeout"], 

275 ) 

276 logger.debug(f"{server.name} LSP server health check passed") 

277 return True 

278 except TimeoutError: 

279 logger.warning(f"{server.name} LSP server health check timed out") 

280 return False 

281 except Exception as e: 

282 logger.warning(f"{server.name} LSP server health check failed: {e}") 

283 return False 

284 

285 async def _shutdown_single_server(self, name: str, server: LSPServer): 

286 """ 

287 Gracefully shutdown a single LSP server. 

288 

289 Args: 

290 name: Server name (key) 

291 server: LSPServer instance 

292 """ 

293 if not server.initialized or not server.client: 

294 return 

295 

296 try: 

297 logger.info(f"Shutting down {name} LSP server") 

298 

299 # LSP protocol shutdown request 

300 try: 

301 await asyncio.wait_for( 

302 server.client.protocol.send_request_async("shutdown", None), timeout=5.0 

303 ) 

304 except TimeoutError: 

305 logger.warning(f"{name} LSP server shutdown request timed out") 

306 

307 # Send exit notification 

308 server.client.protocol.notify("exit", None) 

309 

310 # Stop the client 

311 await server.client.stop() 

312 

313 # Terminate subprocess using stored process reference 

314 if server.process is not None: 

315 try: 

316 if server.process.returncode is not None: 

317 logger.debug(f"{name} already exited (code {server.process.returncode})") 

318 else: 

319 server.process.terminate() 

320 try: 

321 await asyncio.wait_for(server.process.wait(), timeout=2.0) 

322 except TimeoutError: 

323 server.process.kill() 

324 await asyncio.wait_for(server.process.wait(), timeout=1.0) 

325 except Exception as e: 

326 logger.warning(f"Error terminating {name}: {e}") 

327 

328 # Mark as uninitialized 

329 server.initialized = False 

330 server.client = None 

331 server.process = None 

332 server.pid = None 

333 

334 except Exception as e: 

335 logger.error(f"Error shutting down {name} LSP server: {e}") 

336 

337 async def _background_health_monitor(self): 

338 """ 

339 Background task for health checking and idle server shutdown. 

340 

341 Runs periodically to: 

342 - Check health of running servers 

343 - Shutdown idle servers 

344 - Restart crashed servers 

345 """ 

346 logger.info("LSP health monitor task started") 

347 try: 

348 while True: 

349 await asyncio.sleep(LSP_CONFIG["health_check_interval"]) 

350 

351 current_time = time.time() 

352 idle_threshold = current_time - LSP_CONFIG["idle_timeout"] 

353 

354 async with self._lock: 

355 for name, server in self._servers.items(): 

356 if not server.initialized or not server.client: 

357 continue 

358 

359 # Check if server is idle 

360 if server.last_used < idle_threshold: 

361 logger.info( 

362 f"{name} LSP server idle for {(current_time - server.last_used) / 60:.1f} minutes, shutting down" 

363 ) 

364 await self._shutdown_single_server(name, server) 

365 continue 

366 

367 # Health check for active servers 

368 is_healthy = await self._health_check_server(server) 

369 if not is_healthy: 

370 logger.warning(f"{name} LSP server health check failed, restarting") 

371 await self._shutdown_single_server(name, server) 

372 try: 

373 await self._start_server(server) 

374 except Exception as e: 

375 logger.error(f"Failed to restart {name} LSP server: {e}") 

376 

377 except asyncio.CancelledError: 

378 logger.info("LSP health monitor task cancelled") 

379 raise 

380 except Exception as e: 

381 logger.error(f"LSP health monitor task error: {e}", exc_info=True) 

382 

383 def get_status(self) -> dict: 

384 """Get status of managed servers including idle information.""" 

385 current_time = time.time() 

386 status = {} 

387 for name, server in self._servers.items(): 

388 idle_seconds = current_time - server.last_used 

389 uptime_seconds = current_time - server.created_at if server.created_at else 0 

390 status[name] = { 

391 "running": server.initialized and server.client is not None, 

392 "pid": server.pid, 

393 "command": " ".join(server.command), 

394 "restarts": self._restart_attempts.get(name, 0), 

395 "idle_seconds": idle_seconds, 

396 "idle_minutes": idle_seconds / 60.0, 

397 "uptime_seconds": uptime_seconds, 

398 } 

399 return status 

400 

401 async def shutdown(self): 

402 """ 

403 Gracefully shutdown all LSP servers. 

404 

405 Implements: 

406 - Health monitor cancellation 

407 - LSP protocol shutdown (shutdown request + exit notification) 

408 - Pending task cancellation 

409 - Process cleanup with timeout 

410 """ 

411 logger.info("Shutting down LSP manager...") 

412 

413 # Cancel health monitor task 

414 if self._health_monitor_task and not self._health_monitor_task.done(): 

415 logger.info("Cancelling health monitor task") 

416 self._health_monitor_task.cancel() 

417 try: 

418 await self._health_monitor_task 

419 except asyncio.CancelledError: 

420 pass 

421 

422 async with self._lock: 

423 for name, server in self._servers.items(): 

424 await self._shutdown_single_server(name, server) 

425 

426 logger.info("LSP manager shutdown complete") 

427 

428 

429# Singleton accessor 

430_manager_instance: LSPManager | None = None 

431_manager_lock = threading.Lock() 

432 

433 

434def get_lsp_manager() -> LSPManager: 

435 """Get the global LSP manager singleton.""" 

436 global _manager_instance 

437 if _manager_instance is None: 

438 with _manager_lock: 

439 # Double-check pattern to avoid race condition 

440 if _manager_instance is None: 

441 _manager_instance = LSPManager() 

442 return _manager_instance