1"""Resource cleanup handlers for graceful shutdown.

2

3 Provides concrete cleanup implementations for database connections,

4 file handles, HTTP clients, and other resources.

5

6 Phase 10.2: Production Hardening - Resource Cleanup

7"""

8

9 from __future__ import annotations

10

11 import importlib.util as import_util

12 import typing as t

13 from contextlib import suppress

14 from pathlib import Path

15

16

17 def _get_logger() -> t.Any:

18 """Get logger with lazy initialization."""

19 try:

20 from session_buddy.utils.logging import get_session_logger

21

22 return get_session_logger()

23 except Exception:

24 import logging

25

26 return logging.getLogger(__name__)

27

28

29 async def cleanup_database_connections() -> None:

30 """Cleanup DuckDB reflection database connections.

31

32 Closes all active database connections and flushes pending writes.

33 Safe to call even if database is not initialized.

34 """

35 logger = _get_logger()

36 logger.info("Cleaning up database connections")

37

38 try:

39 if import_util.find_spec("session_buddy.reflection_tools") is None: 39 ↛ 40 line 39 didn't jump to line 40 because the condition on line 39 was never true

40 logger.debug("Reflection database not available, skipping cleanup")

41 return

42

43 from session_buddy.reflection_tools import ReflectionDatabase

44

45 with suppress(Exception):

46 ReflectionDatabase().close()

47

48 logger.debug("Database cleanup completed successfully")

49 except Exception:

50 logger.exception("Error during database cleanup")

51 raise

52

53

54 async def _close_adapter_method(requests: t.Any, logger: t.Any) -> bool:

55 """Try to close using adapter's close method."""

56 if hasattr(requests, "close") and callable(requests.close):

57 maybe_await = requests.close()

58 if hasattr(maybe_await, "__await__"):

59 await maybe_await # type: ignore[func-returns-value]

60 logger.debug("Requests adapter cleanup completed successfully")

61 return True

62 return False

63

64

65 async def _close_underlying_client(requests: t.Any, logger: t.Any) -> bool:

66 """Try to close underlying HTTP client."""

67 if not hasattr(requests, "client"):

68 return False

69

70 client = requests.client

71 if hasattr(client, "aclose"):

72 await client.aclose()

73 logger.debug("HTTP client session closed (aclose)")

74 return True

75 if hasattr(client, "close"):

76 client.close()

77 logger.debug("HTTP client session closed (close)")

78 return True

79 return False

80

81

82 async def cleanup_http_clients() -> None:

83 """Cleanup HTTP client connections.

84

85 Closes HTTPClientAdapter instances and releases connection pools.

86 Safe to call even if HTTP clients are not initialized.

87 """

88 logger = _get_logger()

89 logger.info("Cleaning up HTTP client connections")

90

91 try:

92 from mcp_common.adapters.http.client import HTTPClientAdapter

93 from session_buddy.di.container import depends

94

95 try:

96 http_adapter = depends.get_sync(HTTPClientAdapter)

97 except Exception:

98 logger.debug("HTTPClientAdapter not available; skipping HTTP cleanup")

99 return

100

101 cleanup = getattr(http_adapter, "_cleanup_resources", None)

102 if callable(cleanup):

103 maybe_await = cleanup()

104 if hasattr(maybe_await, "__await__"):

105 await maybe_await # type: ignore[func-returns-value]

106 logger.debug("HTTP client adapter cleanup completed successfully")

107 return

108

109 if not await _close_adapter_method(http_adapter, logger):

110 await _close_underlying_client(http_adapter, logger)

111

112 except ModuleNotFoundError:

113 logger.debug("mcp_common.adapters module not available; skipping HTTP cleanup")

114 return

115 except Exception:

116 logger.exception("Error during HTTP client cleanup")

117 raise

118

119

120 async def cleanup_temp_files(temp_dir: Path | None = None) -> None:

121 """Cleanup temporary files created during session.

122

123 Args:

124 temp_dir: Optional temporary directory to clean (defaults to .claude/temp)

125

126 Removes temporary files but preserves important session data.

127

128 """

129 logger = _get_logger()

130

131 if temp_dir is None:

132 temp_dir = Path.home() / ".claude" / "temp"

133

134 if not temp_dir.exists():

135 logger.debug(f"Temp directory does not exist: {temp_dir}")

136 return

137

138 logger.info(f"Cleaning up temporary files in {temp_dir}")

139

140 try:

141 # Remove temporary files

142 files_removed = 0

143 for temp_file in temp_dir.glob("*"):

144 if temp_file.is_file(): 144 ↛ 143 line 144 didn't jump to line 143 because the condition on line 144 was always true

145 try:

146 temp_file.unlink()

147 files_removed += 1

148 except (OSError, PermissionError) as e:

149 logger.warning(f"Could not remove temp file {temp_file}: {e}")

150

151 logger.debug(f"Removed {files_removed} temporary files")

152

153 except Exception:

154 logger.exception("Error during temp file cleanup")

155 raise

156

157

158 async def cleanup_file_handles() -> None:

159 """Cleanup open file handles and flush buffers.

160

161 Ensures all file handles are properly closed and data is flushed.

162 Safe to call multiple times.

163 """

164 logger = _get_logger()

165 logger.info("Cleaning up file handles and flushing buffers")

166

167 try:

168 # Flush all open file descriptors

169 import sys

170

171 if hasattr(sys.stdout, "flush"): 171 ↛ 173 line 171 didn't jump to line 173 because the condition on line 171 was always true

172 sys.stdout.flush()

173 if hasattr(sys.stderr, "flush"): 173 ↛ 176 line 173 didn't jump to line 176 because the condition on line 173 was always true

174 sys.stderr.flush()

175

176 logger.debug("File handle cleanup completed successfully")

177

178 except Exception:

179 logger.exception("Error during file handle cleanup")

180 raise

181

182

183 async def cleanup_session_state() -> None:

184 """Cleanup session state and persistence.

185

186 Saves current session state and cleans up any runtime data.

187 Safe to call even if session management is not active.

188 """

189 logger = _get_logger()

190 logger.info("Cleaning up session state")

191

192 try:

193 # Try to save session state if session manager exists

194 from session_buddy.di.container import depends

195

196 with suppress(Exception):

197 from session_buddy.core import SessionLifecycleManager

198

199 session_mgr = depends.get_sync(SessionLifecycleManager)

200 if session_mgr and hasattr(session_mgr, "_save_state"):

201 # Save any pending state

202 logger.debug("Session state cleanup completed successfully")

203

204 except ImportError:

205 logger.debug("Session manager not available")

206 except Exception:

207 logger.exception("Error during session state cleanup")

208 raise

209

210

211 async def cleanup_background_tasks() -> None:

212 """Cleanup background tasks and async operations.

213

214 Cancels or waits for background tasks to complete gracefully.

215 """

216 logger = _get_logger()

217 logger.info("Cleaning up background tasks")

218

219 try:

220 import asyncio

221

222 # Get current event loop

223 try:

224 loop = asyncio.get_running_loop()

225

226 # Cancel pending tasks (except current task)

227 current_task = asyncio.current_task(loop)

228 pending_tasks = [

229 task

230 for task in asyncio.all_tasks(loop)

231 if task != current_task and not task.done()

232 ]

233

234 if pending_tasks:

235 logger.debug(

236 f"Cancelling {len(pending_tasks)} pending background tasks",

237 )

238 for task in pending_tasks:

239 task.cancel()

240

241 # Wait for tasks to cancel

242 await asyncio.gather(*pending_tasks, return_exceptions=True)

243

244 logger.debug("Background task cleanup completed successfully")

245

246 except RuntimeError:

247 logger.debug("No running event loop, skipping task cleanup")

248

249 except Exception:

250 logger.exception("Error during background task cleanup")

251 raise

252

253

254 async def cleanup_logging_handlers() -> None:

255 """Cleanup logging handlers and flush log buffers.

256

257 Ensures all log messages are written before shutdown.

258 """

259 logger = _get_logger()

260 logger.info("Cleaning up logging handlers")

261

262 try:

263 import logging

264

265 for handler in logging.root.handlers.copy():

266 _cleanup_handler(handler)

267

268 logger.debug("Logging handler cleanup completed successfully")

269

270 except Exception:

271 logger.exception("Error during logging handler cleanup")

272 raise

273

274

275 def _cleanup_handler(handler: t.Any) -> None:

276 """Detach, flush, and close a single logging handler safely."""

277 try:

278 if hasattr(handler, "remove") and not hasattr(handler, "flush"): 278 ↛ 279 line 278 didn't jump to line 279 because the condition on line 278 was never true

279 return

280 if hasattr(handler, "flush") and hasattr(handler, "close"): 280 ↛ 283 line 280 didn't jump to line 283 because the condition on line 280 was always true

281 handler.flush()

282 handler.close()

283 elif hasattr(handler, "close"):

284 handler.close()

285 elif hasattr(handler, "flush"):

286 handler.flush()

287 except TypeError as e:

288 if "_LoggerProxy.remove()" in str(e) and "handler_id" in str(e):

289 return

290 raise

291 except Exception as e:

292 print(f"Error closing log handler: {e}", file=__import__("sys").stderr)

293

294

295 def register_all_cleanup_handlers(

296 shutdown_manager: t.Any,

297 temp_dir: Path | None = None,

298) -> None:

299 """Register all resource cleanup handlers with shutdown manager.

300

301 Args:

302 shutdown_manager: ShutdownManager instance

303 temp_dir: Optional temporary directory for file cleanup

304

305 This is the main entry point for registering cleanup handlers.

306 Called during server initialization.

307

308 Example:

309 >>> from session_buddy.shutdown_manager import get_shutdown_manager

310 >>> from session_buddy.resource_cleanup import register_all_cleanup_handlers

311 >>>

312 >>> shutdown_mgr = get_shutdown_manager()

313 >>> register_all_cleanup_handlers(shutdown_mgr)

314 >>> shutdown_mgr.setup_signal_handlers()

315

316 """

317 logger = _get_logger()

318 logger.info("Registering all resource cleanup handlers")

319

320 # Register cleanup tasks in priority order (highest first)

321

322 # Priority 100: Critical database and connection cleanup

323 shutdown_manager.register_cleanup(

324 name="database_connections",

325 callback=cleanup_database_connections,

326 priority=100,

327 timeout_seconds=10.0,

328 critical=False, # Don't stop other cleanups if this fails

329 )

330

331 shutdown_manager.register_cleanup(

332 name="http_clients",

333 callback=cleanup_http_clients,

334 priority=100,

335 timeout_seconds=10.0,

336 critical=False,

337 )

338

339 # Priority 80: Background tasks

340 shutdown_manager.register_cleanup(

341 name="background_tasks",

342 callback=cleanup_background_tasks,

343 priority=80,

344 timeout_seconds=15.0,

345 critical=False,

346 )

347

348 # Priority 60: Session state

349 shutdown_manager.register_cleanup(

350 name="session_state",

351 callback=cleanup_session_state,

352 priority=60,

353 timeout_seconds=10.0,

354 critical=False,

355 )

356

357 # Priority 40: File handles

358 shutdown_manager.register_cleanup(

359 name="file_handles",

360 callback=cleanup_file_handles,

361 priority=40,

362 timeout_seconds=5.0,

363 critical=False,

364 )

365

366 # Priority 20: Temp files

367 async def _cleanup_temp_files_wrapper() -> None:

368 """Wrapper to properly await cleanup_temp_files coroutine."""

369 await cleanup_temp_files(temp_dir)

370

371 shutdown_manager.register_cleanup(

372 name="temp_files",

373 callback=_cleanup_temp_files_wrapper,

374 priority=20,

375 timeout_seconds=10.0,

376 critical=False,

377 )

378

379 # Priority 10: Logging (last, so we can log other cleanups)

380 shutdown_manager.register_cleanup(

381 name="logging_handlers",

382 callback=cleanup_logging_handlers,

383 priority=10,

384 timeout_seconds=5.0,

385 critical=False,

386 )

387

388 logger.info(

389 f"Registered {len(shutdown_manager._cleanup_tasks)} resource cleanup handlers",

390 )

391

392

393__all__ = [

394 "cleanup_background_tasks",

395 "cleanup_database_connections",

396 "cleanup_file_handles",

397 "cleanup_http_clients",

398 "cleanup_logging_handlers",

399 "cleanup_session_state",

400 "cleanup_temp_files",

401 "register_all_cleanup_handlers",

402]