Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-session/src/lexigram/ai/session/manager/core.py: 22%
117 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Session manager — central orchestrator for session lifecycle."""
3from __future__ import annotations
5import asyncio
6from dataclasses import replace
7from datetime import UTC, datetime
8from typing import Any
9from uuid import uuid4
11from lexigram.ai.session.config import SessionConfig
12from lexigram.ai.session.exceptions import (
13 CheckpointNotFoundError,
14 SessionCapacityError,
15 SessionClosedError,
16 SessionNotFoundError,
17)
18from lexigram.ai.session.state import SessionStateMachine
19from lexigram.contracts.ai.memory import MemoryEntry, WorkingMemoryProtocol
20from lexigram.contracts.ai.session import (
21 SessionCheckpoint,
22 SessionManagerProtocol,
23 SessionState,
24 SessionStatus,
25 SessionStoreProtocol,
26 SessionTurn,
27)
28from lexigram.logging import (
29 get_logger,
30)
32logger = get_logger(__name__)
35class SessionManagerImpl(SessionManagerProtocol):
36 """Manages session lifecycle with persistence and optional memory integration.
38 Implements ``SessionManagerProtocol``. Orchestrates create, resume,
39 suspend, close, add_turn, checkpoint, and restore operations. Publishes
40 domain events via an optional event bus and triggers memory consolidation
41 on session close when a consolidator is available.
43 Constructor parameters are injected by the DI container; every dependency
44 except *config* and *store* is optional to allow graceful degradation.
46 Args:
47 config: Session configuration.
48 store: Persistence backend for session state and checkpoints.
49 memory: Optional working memory for recording turns.
50 consolidator: Optional memory consolidator triggered on close.
51 event_bus: Optional event bus for lifecycle events.
52 """
54 def __init__(
55 self,
56 config: SessionConfig,
57 store: SessionStoreProtocol,
58 memory: WorkingMemoryProtocol | None = None,
59 consolidator: Any | None = None,
60 event_bus: Any | None = None,
61 ) -> None:
62 self._config = config
63 self._store = store
64 self._memory = memory
65 self._consolidator = consolidator
66 self._event_bus = event_bus
67 self._fsm = SessionStateMachine()
68 self._background_tasks: set[asyncio.Task[Any]] = set()
70 # ------------------------------------------------------------------
71 # SessionManagerProtocol
72 # ------------------------------------------------------------------
74 async def create(
75 self,
76 user_id: str,
77 metadata: dict[str, Any] | None = None,
78 ) -> SessionState:
79 """Create a new session.
81 Args:
82 user_id: The owning user.
83 metadata: Optional session metadata.
85 Returns:
86 Newly created ``SessionState`` with status ``active``.
88 Raises:
89 SessionCapacityError: If the per-user session limit is reached.
90 """
91 existing = await self._store.list_sessions(user_id)
92 active_count = sum(1 for s in existing if s.status != SessionStatus.CLOSED)
93 if active_count >= self._config.max_sessions_per_user:
94 raise SessionCapacityError(
95 f"user {user_id!r} already has {active_count} active sessions "
96 f"(max {self._config.max_sessions_per_user})"
97 )
99 now = datetime.now(UTC)
100 state = SessionState(
101 session_id=str(uuid4()),
102 user_id=user_id,
103 status=SessionStatus.ACTIVE,
104 metadata=metadata or {},
105 created_at=now,
106 updated_at=now,
107 )
108 await self._store.save(state)
109 await self._emit("session.created", state)
110 logger.info("session_created", session_id=state.session_id, user_id=user_id)
111 return state
113 async def resume(self, session_id: str) -> SessionState | None:
114 """Resume an existing session.
116 Transitions SUSPENDED → ACTIVE. Returns the updated state or ``None``
117 if the session does not exist.
119 Args:
120 session_id: ID of the session to resume.
122 Returns:
123 Updated ``SessionState``, or ``None`` if not found.
125 Raises:
126 SessionClosedError: If the session is already closed.
127 """
128 state = await self._store.load(session_id)
129 if state is None:
130 return None
131 if state.status == SessionStatus.CLOSED:
132 raise SessionClosedError(session_id)
133 self._fsm.validate(session_id, state.status.value, SessionStatus.ACTIVE.value)
134 state = replace(
135 state,
136 status=SessionStatus.ACTIVE,
137 updated_at=datetime.now(UTC),
138 )
139 await self._store.save(state)
140 await self._emit("session.resumed", state)
141 logger.info("session_resumed", session_id=session_id)
142 return state
144 async def add_turn(self, session_id: str, turn: SessionTurn) -> None:
145 """Add a conversation turn and update session metrics.
147 Also records the turn in working memory if available, and triggers
148 an auto-checkpoint when the configured interval is reached.
150 Args:
151 session_id: Target session ID.
152 turn: The turn to append.
154 Raises:
155 SessionNotFoundError: If the session does not exist.
156 SessionClosedError: If the session is closed.
157 SessionCapacityError: If the per-session turn limit is reached.
158 """
159 state = await self._store.load(session_id)
160 if state is None:
161 raise SessionNotFoundError(session_id)
162 if state.status == SessionStatus.CLOSED:
163 raise SessionClosedError(session_id)
164 if len(state.turns) >= self._config.max_turns_per_session:
165 raise SessionCapacityError(
166 f"session {session_id!r} has reached the turn limit "
167 f"({self._config.max_turns_per_session})"
168 )
170 state.turns.append(turn)
171 state = replace(
172 state,
173 turn_count=state.turn_count + 1,
174 total_tokens=state.total_tokens + turn.tokens_used,
175 total_cost=state.total_cost + turn.cost,
176 updated_at=datetime.now(UTC),
177 )
179 if self._memory is not None:
180 entry = MemoryEntry(
181 id=turn.turn_id,
182 owner_id=state.user_id or session_id,
183 content=turn.content,
184 role=turn.role,
185 timestamp=turn.timestamp,
186 metadata=turn.metadata,
187 )
188 await self._memory.add(entry)
190 await self._store.save(state)
192 if (
193 self._config.auto_checkpoint_interval
194 and state.turn_count % self._config.auto_checkpoint_interval == 0
195 ):
196 await self.checkpoint(session_id)
198 async def get_state(self, session_id: str) -> SessionState | None:
199 """Return the current state of a session, or ``None`` if not found.
201 Args:
202 session_id: The session ID to look up.
204 Returns:
205 ``SessionState`` or ``None``.
206 """
207 return await self._store.load(session_id)
209 async def checkpoint(self, session_id: str) -> SessionCheckpoint:
210 """Create an immutable snapshot of the current session state.
212 Args:
213 session_id: The session to snapshot.
215 Returns:
216 The created ``SessionCheckpoint``.
218 Raises:
219 SessionNotFoundError: If the session does not exist.
220 """
221 state = await self._store.load(session_id)
222 if state is None:
223 raise SessionNotFoundError(session_id)
225 import copy
227 checkpoint = SessionCheckpoint(
228 checkpoint_id=str(uuid4()),
229 session_id=session_id,
230 state=copy.deepcopy(state),
231 created_at=datetime.now(UTC),
232 parent_checkpoint_id=state.checkpoint_id,
233 )
234 state = replace(state, checkpoint_id=checkpoint.checkpoint_id)
235 await self._store.save(state)
236 await self._store.save_checkpoint(checkpoint)
237 logger.info(
238 "session_checkpointed",
239 session_id=session_id,
240 checkpoint_id=checkpoint.checkpoint_id,
241 )
242 return checkpoint
244 async def restore(self, checkpoint_id: str) -> SessionState:
245 """Restore a session to a previous checkpoint.
247 Args:
248 checkpoint_id: The checkpoint to restore from.
250 Returns:
251 The restored ``SessionState``.
253 Raises:
254 CheckpointNotFoundError: If the checkpoint does not exist.
255 """
256 checkpoint = await self._store.load_checkpoint(checkpoint_id)
257 if checkpoint is None:
258 raise CheckpointNotFoundError(checkpoint_id)
260 import copy
262 restored = copy.deepcopy(checkpoint.state)
263 restored = replace(
264 restored,
265 updated_at=datetime.now(UTC),
266 checkpoint_id=checkpoint_id,
267 )
268 await self._store.save(restored)
269 await self._emit("session.restored", restored)
270 logger.info(
271 "session_restored",
272 session_id=restored.session_id,
273 checkpoint_id=checkpoint_id,
274 )
275 return restored
277 async def close(self, session_id: str) -> None:
278 """Close a session and optionally trigger memory consolidation.
280 Args:
281 session_id: The session to close.
283 Raises:
284 SessionNotFoundError: If the session does not exist.
285 """
286 state = await self._store.load(session_id)
287 if state is None:
288 raise SessionNotFoundError(session_id)
290 self._fsm.validate(session_id, state.status.value, SessionStatus.CLOSED.value)
291 state = replace(
292 state,
293 status=SessionStatus.CLOSED,
294 updated_at=datetime.now(UTC),
295 )
296 await self._store.save(state)
298 if self._consolidator is not None and self._config.consolidate_on_close:
299 task = asyncio.create_task(
300 self._consolidator.consolidate(session_id=session_id)
301 )
302 self._background_tasks.add(task)
303 task.add_done_callback(self._background_tasks.discard)
305 await self._emit("session.closed", state)
306 logger.info("session_closed", session_id=session_id)
308 async def suspend(self, session_id: str) -> SessionState:
309 """Suspend an active session without closing it.
311 Args:
312 session_id: The session to suspend.
314 Returns:
315 Updated ``SessionState`` with status ``suspended``.
317 Raises:
318 SessionNotFoundError: If the session does not exist.
319 SessionTransitionError: If the session cannot be suspended.
320 """
321 state = await self._store.load(session_id)
322 if state is None:
323 raise SessionNotFoundError(session_id)
324 self._fsm.validate(
325 session_id, state.status.value, SessionStatus.SUSPENDED.value
326 )
327 state = replace(
328 state,
329 status=SessionStatus.SUSPENDED,
330 updated_at=datetime.now(UTC),
331 )
332 await self._store.save(state)
333 await self._emit("session.suspended", state)
334 logger.info("session_suspended", session_id=session_id)
335 return state
337 # ------------------------------------------------------------------
338 # Internal helpers
339 # ------------------------------------------------------------------
341 async def _emit(self, event_type: str, state: SessionState) -> None:
342 """Publish a session lifecycle event if an event bus is configured.
344 Args:
345 event_type: Dot-separated event name (e.g. ``session.created``).
346 state: The session state associated with the event.
347 """
348 if self._event_bus is None:
349 return
350 try:
351 await self._event_bus.publish(
352 event_type,
353 {"session_id": state.session_id, "status": state.status.value},
354 )
355 except Exception as exc: # noqa: BLE001
356 logger.warning(
357 "session_event_publish_failed", event_type=event_type, error=str(exc)
358 )
360 async def dispose(self) -> None:
361 """Container dispose hook. No-op — shutdown is handled by SessionProvider."""
364__all__ = ["SessionManagerImpl"]