Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-session/src/lexigram/ai/session/manager/cleanup.py: 22%

1"""Session cleanup scheduler — background task for expiring old sessions.""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6from datetime import UTC, datetime 

7from typing import Any 

8 

9from lexigram.ai.session.config import SessionConfig 

10from lexigram.contracts.ai.session import SessionManagerProtocol 

11from lexigram.logging import ( 

12 get_logger, 

13) 

14 

15logger = get_logger(__name__) 

16 

17 

18class SessionCleanupScheduler: 

19 """Runs a background loop to close sessions that have exceeded their TTL. 

20 

21 Session TTL allows the system to automatically close abandoned or 

22 stale sessions, which in turn triggers memory consolidation and 

23 frees up active session counts for users. 

24 

25 Args: 

26 manager: Session manager instance to list and close sessions. 

27 config: Session configuration containing TTL settings. 

28 """ 

29 

30 def __init__( 

31 self, 

32 manager: SessionManagerProtocol, 

33 config: SessionConfig, 

34 ) -> None: 

35 """Initialise the cleanup scheduler. 

36 

37 Args: 

38 manager: The configured session manager instance. 

39 config: Session configuration. 

40 """ 

41 self._manager = manager 

42 self._config = config 

43 self._running = False 

44 self._task: asyncio.Task[Any] | None = None 

45 

46 async def start(self) -> None: 

47 """Start the background cleanup loop.""" 

48 if self._running: 

49 return 

50 self._running = True 

51 self._task = asyncio.create_task(self._cleanup_loop()) 

52 logger.info( 

53 "session_cleanup_started", 

54 interval=self._config.cleanup_interval_s, 

55 ttl=self._config.session_ttl, 

56 ) 

57 

58 async def stop(self) -> None: 

59 """Stop the background cleanup loop gracefully.""" 

60 if not self._running: 

61 return 

62 self._running = False 

63 if self._task is not None: 

64 self._task.cancel() 

65 try: 

66 await self._task 

67 except asyncio.CancelledError: 

68 pass 

69 self._task = None 

70 logger.info("session_cleanup_stopped") 

71 

72 async def run_cleanup_pass(self) -> int: 

73 """Run a single cleanup pass to find and close expired sessions. 

74 

75 Returns: 

76 Number of sessions closed during this pass. 

77 """ 

78 # Session TTL is disabled if set to 0 

79 if self._config.session_ttl <= 0: 

80 return 0 

81 

82 closed_count = 0 

83 now = datetime.now(UTC) 

84 

85 # In a real distributed system, we would query the store directly 

86 # for expired sessions. For the protocol, we list users/sessions. 

87 # Since this is an abstract implementation, we rely on the store's 

88 # underlying querying capability if exposed, otherwise we only clean up 

89 # sessions we actively know about. If working with bounded stores, 

90 # we iterate. 

91 try: 

92 # Assume store is injected into manager and accessible via private getattr 

93 # or the manager implements an undocumented list_all_active() 

94 active_sessions = [] 

95 store = getattr(self._manager, "_store", None) 

96 if store and hasattr(store, "list_all_active"): 

97 active_sessions = await store.list_all_active() 

98 

99 for session in active_sessions: 

100 age_s = (now - session.updated_at).total_seconds() 

101 if age_s > self._config.session_ttl: 

102 try: 

103 await self._manager.close(session.session_id) 

104 closed_count += 1 

105 logger.debug( 

106 "session_ttl_expired", session_id=session.session_id 

107 ) 

108 except Exception as e: # noqa: BLE001 # per-session cleanup must not stop other sessions from being cleaned 

109 logger.warning( 

110 "session_cleanup_failed", 

111 session_id=session.session_id, 

112 error=str(e), 

113 ) 

114 except Exception as e: # noqa: BLE001 # cleanup pass error boundary; logged and swallowed 

115 logger.error("session_cleanup_pass_error", error=str(e)) 

116 

117 if closed_count > 0: 

118 logger.info("session_cleanup_pass_completed", closed=closed_count) 

119 

120 return closed_count 

121 

122 async def _cleanup_loop(self) -> None: 

123 """Background loop continuously running cleanup passes.""" 

124 while self._running: 

125 try: 

126 await asyncio.sleep(self._config.cleanup_interval_s) 

127 if not self._running: 

128 break 

129 await self.run_cleanup_pass() 

130 except asyncio.CancelledError: 

131 break 

132 except Exception as e: # noqa: BLE001 

133 logger.error("session_cleanup_loop_error", error=str(e)) 

134 # Prevent tight loop on persistent error 

135 await asyncio.sleep(60.0) 

136 

137 

138__all__ = ["SessionCleanupScheduler"]