Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-session/src/lexigram/ai/session/stores/in_memory.py: 57%

1"""In-memory session store — for development and testing.""" 

2 

3from __future__ import annotations 

4 

5from datetime import UTC, datetime 

6from typing import TYPE_CHECKING 

7 

8if TYPE_CHECKING: 

9 from lexigram.contracts.ai.session import ( 

10 SessionCheckpoint, 

11 SessionState, 

12 ) 

13 

14 

15class InMemorySessionStore: 

16 """In-process session and checkpoint storage. 

17 

18 All data is stored in plain Python dicts and is lost when the process 

19 exits. Suitable for development, testing, and single-process deployments. 

20 """ 

21 

22 def __init__(self) -> None: 

23 self._sessions: dict[str, SessionState] = {} 

24 self._checkpoints: dict[str, SessionCheckpoint] = {} 

25 

26 # ------------------------------------------------------------------ 

27 # Session CRUD 

28 # ------------------------------------------------------------------ 

29 

30 async def save(self, state: SessionState) -> None: 

31 """Persist or overwrite a session state. 

32 

33 Args: 

34 state: The session state to save. 

35 """ 

36 self._sessions[state.session_id] = state 

37 

38 async def load(self, session_id: str) -> SessionState | None: 

39 """Return the session state for *session_id*, or ``None``. 

40 

41 Args: 

42 session_id: The session to load. 

43 

44 Returns: 

45 The session state, or ``None`` if not found. 

46 """ 

47 return self._sessions.get(session_id) 

48 

49 async def delete(self, session_id: str) -> None: 

50 """Remove a session from the store. 

51 

52 Args: 

53 session_id: The session to delete. 

54 """ 

55 self._sessions.pop(session_id, None) 

56 

57 async def list_sessions(self, user_id: str) -> list[SessionState]: 

58 """List all sessions belonging to *user_id*. 

59 

60 Args: 

61 user_id: The user to filter by. 

62 

63 Returns: 

64 All sessions owned by that user. 

65 """ 

66 return [s for s in self._sessions.values() if s.user_id == user_id] 

67 

68 # ------------------------------------------------------------------ 

69 # Checkpoint CRUD 

70 # ------------------------------------------------------------------ 

71 

72 async def save_checkpoint(self, checkpoint: SessionCheckpoint) -> None: 

73 """Persist an immutable session checkpoint. 

74 

75 Args: 

76 checkpoint: The checkpoint to store. 

77 """ 

78 self._checkpoints[checkpoint.checkpoint_id] = checkpoint 

79 

80 async def load_checkpoint(self, checkpoint_id: str) -> SessionCheckpoint | None: 

81 """Return the checkpoint for *checkpoint_id*, or ``None``. 

82 

83 Args: 

84 checkpoint_id: The checkpoint to load. 

85 

86 Returns: 

87 The checkpoint, or ``None`` if not found. 

88 """ 

89 return self._checkpoints.get(checkpoint_id) 

90 

91 async def list_checkpoints(self, session_id: str) -> list[SessionCheckpoint]: 

92 """List all checkpoints for *session_id*, sorted oldest-first. 

93 

94 Args: 

95 session_id: The session to query. 

96 

97 Returns: 

98 Checkpoints in chronological order. 

99 """ 

100 return sorted( 

101 (c for c in self._checkpoints.values() if c.session_id == session_id), 

102 key=lambda c: c.created_at, 

103 ) 

104 

105 async def delete_checkpoint(self, checkpoint_id: str) -> None: 

106 """Remove a checkpoint from the store. 

107 

108 Args: 

109 checkpoint_id: The checkpoint to delete. 

110 """ 

111 self._checkpoints.pop(checkpoint_id, None) 

112 

113 # ------------------------------------------------------------------ 

114 # Housekeeping 

115 # ------------------------------------------------------------------ 

116 

117 async def expire_old_sessions(self, ttl_seconds: int) -> int: 

118 """Delete sessions whose last update is older than *ttl_seconds*. 

119 

120 Args: 

121 ttl_seconds: Maximum allowed age in seconds. 

122 

123 Returns: 

124 Number of sessions expired. 

125 """ 

126 now = datetime.now(UTC) 

127 expired_ids = [ 

128 sid 

129 for sid, state in self._sessions.items() 

130 if (now - state.updated_at).total_seconds() > ttl_seconds 

131 ] 

132 for sid in expired_ids: 

133 del self._sessions[sid] 

134 return len(expired_ids) 

135 

136 

137__all__ = ["InMemorySessionStore"]