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

1"""

2 Persistent LSP Server Manager

3

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

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

6

7 Architecture:

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

15 import asyncio

16 import logging

17 import threading

18 import time

19 from dataclasses import dataclass, field

20 from typing import Optional

21

22 from lsprotocol.types import (

23 ClientCapabilities,

24 InitializedParams,

25 InitializeParams,

26)

27 from pygls.client import JsonRPCClient

28

29 logger = logging.getLogger(__name__)

30

31# Configuration for LSP server lifecycle management

32 LSP_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

40 class 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

53 class 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

434 def 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