Coverage for session_buddy / resource_cleanup.py: 54.35%
182 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-04 00:43 -0800
1"""Resource cleanup handlers for graceful shutdown.
3Provides concrete cleanup implementations for database connections,
4file handles, HTTP clients, and other resources.
6Phase 10.2: Production Hardening - Resource Cleanup
7"""
9from __future__ import annotations
11import importlib.util as import_util
12import typing as t
13from contextlib import suppress
14from pathlib import Path
17def _get_logger() -> t.Any:
18 """Get logger with lazy initialization."""
19 try:
20 from session_buddy.utils.logging import get_session_logger
22 return get_session_logger()
23 except Exception:
24 import logging
26 return logging.getLogger(__name__)
29async def cleanup_database_connections() -> None:
30 """Cleanup DuckDB reflection database connections.
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")
38 try:
39 if import_util.find_spec("session_buddy.reflection_tools") is None: 39 ↛ 40line 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
43 from session_buddy.reflection_tools import ReflectionDatabase
45 with suppress(Exception):
46 ReflectionDatabase().close()
48 logger.debug("Database cleanup completed successfully")
49 except Exception:
50 logger.exception("Error during database cleanup")
51 raise
54async 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
65async 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
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
82async def cleanup_http_clients() -> None:
83 """Cleanup HTTP client connections.
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")
91 try:
92 from mcp_common.adapters.http.client import HTTPClientAdapter
93 from session_buddy.di.container import depends
95 try:
96 http_adapter = depends.get_sync(HTTPClientAdapter)
97 except Exception:
98 logger.debug("HTTPClientAdapter not available; skipping HTTP cleanup")
99 return
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
109 if not await _close_adapter_method(http_adapter, logger):
110 await _close_underlying_client(http_adapter, logger)
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
120async def cleanup_temp_files(temp_dir: Path | None = None) -> None:
121 """Cleanup temporary files created during session.
123 Args:
124 temp_dir: Optional temporary directory to clean (defaults to .claude/temp)
126 Removes temporary files but preserves important session data.
128 """
129 logger = _get_logger()
131 if temp_dir is None:
132 temp_dir = Path.home() / ".claude" / "temp"
134 if not temp_dir.exists():
135 logger.debug(f"Temp directory does not exist: {temp_dir}")
136 return
138 logger.info(f"Cleaning up temporary files in {temp_dir}")
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 ↛ 143line 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}")
151 logger.debug(f"Removed {files_removed} temporary files")
153 except Exception:
154 logger.exception("Error during temp file cleanup")
155 raise
158async def cleanup_file_handles() -> None:
159 """Cleanup open file handles and flush buffers.
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")
167 try:
168 # Flush all open file descriptors
169 import sys
171 if hasattr(sys.stdout, "flush"): 171 ↛ 173line 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 ↛ 176line 173 didn't jump to line 176 because the condition on line 173 was always true
174 sys.stderr.flush()
176 logger.debug("File handle cleanup completed successfully")
178 except Exception:
179 logger.exception("Error during file handle cleanup")
180 raise
183async def cleanup_session_state() -> None:
184 """Cleanup session state and persistence.
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")
192 try:
193 # Try to save session state if session manager exists
194 from session_buddy.di.container import depends
196 with suppress(Exception):
197 from session_buddy.core import SessionLifecycleManager
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")
204 except ImportError:
205 logger.debug("Session manager not available")
206 except Exception:
207 logger.exception("Error during session state cleanup")
208 raise
211async def cleanup_background_tasks() -> None:
212 """Cleanup background tasks and async operations.
214 Cancels or waits for background tasks to complete gracefully.
215 """
216 logger = _get_logger()
217 logger.info("Cleaning up background tasks")
219 try:
220 import asyncio
222 # Get current event loop
223 try:
224 loop = asyncio.get_running_loop()
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 ]
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()
241 # Wait for tasks to cancel
242 await asyncio.gather(*pending_tasks, return_exceptions=True)
244 logger.debug("Background task cleanup completed successfully")
246 except RuntimeError:
247 logger.debug("No running event loop, skipping task cleanup")
249 except Exception:
250 logger.exception("Error during background task cleanup")
251 raise
254async def cleanup_logging_handlers() -> None:
255 """Cleanup logging handlers and flush log buffers.
257 Ensures all log messages are written before shutdown.
258 """
259 logger = _get_logger()
260 logger.info("Cleaning up logging handlers")
262 try:
263 import logging
265 for handler in logging.root.handlers.copy():
266 _cleanup_handler(handler)
268 logger.debug("Logging handler cleanup completed successfully")
270 except Exception:
271 logger.exception("Error during logging handler cleanup")
272 raise
275def _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 ↛ 279line 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 ↛ 283line 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)
295def 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.
301 Args:
302 shutdown_manager: ShutdownManager instance
303 temp_dir: Optional temporary directory for file cleanup
305 This is the main entry point for registering cleanup handlers.
306 Called during server initialization.
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()
316 """
317 logger = _get_logger()
318 logger.info("Registering all resource cleanup handlers")
320 # Register cleanup tasks in priority order (highest first)
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 )
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 )
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 )
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 )
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 )
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)
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 )
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 )
388 logger.info(
389 f"Registered {len(shutdown_manager._cleanup_tasks)} resource cleanup handlers",
390 )
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]