Coverage for src / lexigram / contracts / ai / session.py: 2%

1"""Session management contracts for stateful conversations.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Sequence 

6from dataclasses import dataclass, field 

7from datetime import UTC, datetime 

8from enum import StrEnum 

9from typing import TYPE_CHECKING, Any 

10 

11from typing_extensions import Protocol, runtime_checkable 

12 

13from lexigram.contracts.exceptions import LexigramError 

14 

15if TYPE_CHECKING: 

16 from lexigram.contracts.ai.llm import ChatMessageProtocol 

17 

18 

19# Session Errors 

20class SessionError(LexigramError): 

21 """Base class for session-related errors.""" 

22 

23 _code = "LEX_ERR_SES_001" 

24 

25 

26class TaskCancelledError(LexigramError): 

27 """Error raised when a task is cancelled.""" 

28 

29 _code = "LEX_ERR_SES_002" 

30 

31 

32class TaskError(LexigramError): 

33 """Base class for task execution errors.""" 

34 

35 _code = "LEX_ERR_SES_003" 

36 

37 

38class TaskTimeoutError(TaskError): 

39 """Error raised when a task times out.""" 

40 

41 _code = "LEX_ERR_SES_004" 

42 

43 

44class TaskValidationError(TaskError): 

45 """Error raised when task input validation fails.""" 

46 

47 _code = "LEX_ERR_SES_005" 

48 

49 

50class SessionStatus(StrEnum): 

51 """Status of a conversation session. 

52 

53 Attributes: 

54 ACTIVE: Session is currently active 

55 SUSPENDED: Session is paused but can be resumed 

56 CLOSED: Session has been terminated 

57 EXPIRED: Session has expired due to TTL 

58 """ 

59 

60 ACTIVE = "active" 

61 SUSPENDED = "suspended" 

62 CLOSED = "closed" 

63 EXPIRED = "expired" 

64 

65 

66@dataclass(frozen=True) 

67class SessionTurn: 

68 """A single turn (exchange) in a conversation session. 

69 

70 Attributes: 

71 turn_id: Unique ID for this turn 

72 role: The role (e.g., "user", "assistant", "system", "tool") 

73 content: The content of this turn 

74 timestamp: When this turn occurred 

75 tool_calls: List of tool calls made during this turn 

76 skill_results: Results from skill executions in this turn 

77 metadata: Additional metadata about the turn 

78 tokens_used: Tokens consumed by this turn 

79 cost: If metered, cost of this turn 

80 model: Model that produced this turn (if applicable) 

81 provider: Provider that served this turn (if applicable) 

82 """ 

83 

84 turn_id: str 

85 role: str 

86 content: str 

87 timestamp: datetime 

88 tool_calls: list[dict[str, Any]] = field(default_factory=list) 

89 skill_results: list[dict[str, Any]] = field(default_factory=list) 

90 metadata: dict[str, Any] = field(default_factory=dict) 

91 tokens_used: int = 0 

92 cost: float = 0.0 

93 model: str | None = None 

94 provider: str | None = None 

95 

96 

97@dataclass(frozen=True) 

98class SessionState: 

99 """Complete state of a conversation session. 

100 

101 Frozen dataclass representing the current state of a session, 

102 including all turns and metadata. List and dict fields are frozen 

103 by reference — their contents remain mutable where needed. 

104 

105 Attributes: 

106 session_id: Unique session identifier 

107 user_id: User who owns this session 

108 status: Current session status 

109 turns: All turns in this session 

110 metadata: Session-level metadata 

111 active_tools: Tools currently active in this session 

112 active_skills: Skills currently active in this session 

113 system_prompt: Optional system prompt for this session 

114 variables: User-defined session-level variables 

115 created_at: When session was created 

116 updated_at: When session was last modified 

117 checkpoint_id: ID of last checkpoint (if any) 

118 total_tokens: Cumulative tokens used across all turns 

119 total_cost: Cumulative cost across all turns 

120 turn_count: Number of turns recorded 

121 parent_session_id: Parent session ID if this is a branch 

122 branch_name: Name of this branch if forked 

123 """ 

124 

125 session_id: str 

126 user_id: str 

127 status: SessionStatus 

128 turns: list[SessionTurn] = field(default_factory=list) 

129 metadata: dict[str, Any] = field(default_factory=dict) 

130 active_tools: list[str] = field(default_factory=list) 

131 active_skills: list[str] = field(default_factory=list) 

132 system_prompt: str | None = None 

133 variables: dict[str, Any] = field(default_factory=dict) 

134 created_at: datetime = field(default_factory=lambda: datetime.now(UTC)) 

135 updated_at: datetime = field(default_factory=lambda: datetime.now(UTC)) 

136 checkpoint_id: str | None = None 

137 total_tokens: int = 0 

138 total_cost: float = 0.0 

139 turn_count: int = 0 

140 parent_session_id: str | None = None 

141 branch_name: str | None = None 

142 

143 

144@dataclass(frozen=True) 

145class SessionCheckpoint: 

146 """A snapshot of session state at a point in time. 

147 

148 Immutable checkpoint that can be used for restoring session state 

149 or implementing branching/versioning. 

150 

151 Attributes: 

152 checkpoint_id: Unique identifier for this checkpoint 

153 session_id: ID of the session this checkpoint belongs to 

154 state: The complete session state at checkpoint time 

155 created_at: When this checkpoint was created 

156 parent_checkpoint_id: ID of the parent checkpoint (for DAG tracking) 

157 metadata: Optional checkpoint metadata 

158 """ 

159 

160 checkpoint_id: str 

161 session_id: str 

162 state: SessionState 

163 created_at: datetime 

164 parent_checkpoint_id: str | None = None 

165 metadata: dict[str, Any] = field(default_factory=dict) 

166 

167 

168@runtime_checkable 

169class SessionStoreProtocol(Protocol): 

170 """Protocol for storing and retrieving session state. 

171 

172 Implementations provide persistence for session state, supporting 

173 creation, loading, deletion, and enumeration of sessions. 

174 """ 

175 

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

177 """Save or update a session state. 

178 

179 Args: 

180 state: The session state to save 

181 """ 

182 ... 

183 

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

185 """Load a session state by ID. 

186 

187 Args: 

188 session_id: The session ID to load 

189 

190 Returns: 

191 The session state, or None if not found 

192 """ 

193 ... 

194 

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

196 """Delete a session. 

197 

198 Args: 

199 session_id: The session ID to delete 

200 """ 

201 ... 

202 

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

204 """List all sessions for a user. 

205 

206 Args: 

207 user_id: The user ID to query 

208 

209 Returns: 

210 List of sessions owned by this user 

211 """ 

212 ... 

213 

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

215 """Persist an immutable session checkpoint. 

216 

217 Args: 

218 checkpoint: The checkpoint to save. 

219 """ 

220 ... 

221 

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

223 """Load a checkpoint by ID. 

224 

225 Args: 

226 checkpoint_id: The checkpoint ID to load. 

227 

228 Returns: 

229 The checkpoint, or None if not found. 

230 """ 

231 ... 

232 

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

234 """List all checkpoints for a session. 

235 

236 Args: 

237 session_id: The session to list checkpoints for. 

238 

239 Returns: 

240 All checkpoints in chronological order. 

241 """ 

242 ... 

243 

244 

245@runtime_checkable 

246class SessionManagerProtocol(Protocol): 

247 """Protocol for managing session lifecycle. 

248 

249 Handles session creation, resumption, state management, 

250 checkpointing, and restoration. 

251 """ 

252 

253 async def create( 

254 self, user_id: str, metadata: dict[str, Any] | None = None 

255 ) -> SessionState: 

256 """Create a new session. 

257 

258 Args: 

259 user_id: The user ID 

260 metadata: Optional session metadata 

261 

262 Returns: 

263 The newly created session 

264 """ 

265 ... 

266 

267 async def resume(self, session_id: str) -> SessionState | None: 

268 """Resume an existing session. 

269 

270 Args: 

271 session_id: The session ID to resume 

272 

273 Returns: 

274 The session state, or None if not found or closed 

275 """ 

276 ... 

277 

278 async def add_turn(self, session_id: str, turn: SessionTurn) -> None: 

279 """Add a turn to a session. 

280 

281 Args: 

282 session_id: The session to update 

283 turn: The turn to add 

284 """ 

285 ... 

286 

287 async def get_state(self, session_id: str) -> SessionState | None: 

288 """Get the current state of a session. 

289 

290 Args: 

291 session_id: The session ID 

292 

293 Returns: 

294 The session state, or None if not found 

295 """ 

296 ... 

297 

298 async def checkpoint(self, session_id: str) -> SessionCheckpoint: 

299 """Create a checkpoint of the session. 

300 

301 Args: 

302 session_id: The session to checkpoint 

303 

304 Returns: 

305 The created checkpoint 

306 """ 

307 ... 

308 

309 async def restore(self, checkpoint_id: str) -> SessionState: 

310 """Restore session state from a checkpoint. 

311 

312 Args: 

313 checkpoint_id: The checkpoint ID to restore from 

314 

315 Returns: 

316 The restored session state 

317 """ 

318 ... 

319 

320 async def close(self, session_id: str) -> None: 

321 """Close/terminate a session. 

322 

323 Args: 

324 session_id: The session to close 

325 """ 

326 ... 

327 

328 async def suspend(self, session_id: str) -> SessionState: 

329 """Suspend an active session without closing it. 

330 

331 Args: 

332 session_id: The session to suspend. 

333 

334 Returns: 

335 Updated session state with status ``suspended``. 

336 """ 

337 ... 

338 

339 

340@runtime_checkable 

341class SessionContextProtocol(Protocol): 

342 """Protocol for per-request session context. 

343 

344 Provides access to the current session within a request scope, 

345 with ability to bind and retrieve session state. 

346 """ 

347 

348 @property 

349 def session_id(self) -> str: 

350 """Get the current session ID. 

351 

352 Returns: 

353 The current session ID 

354 

355 Raises: 

356 SessionError: If no session is currently bound 

357 """ 

358 ... 

359 

360 @property 

361 def state(self) -> SessionState: 

362 """Get the current session state. 

363 

364 Returns: 

365 The current session state 

366 

367 Raises: 

368 SessionError: If no session is currently bound 

369 """ 

370 ... 

371 

372 async def get_or_create(self, user_id: str) -> SessionState: 

373 """Get an existing session or create a new one. 

374 

375 Args: 

376 user_id: The user ID 

377 

378 Returns: 

379 An existing or newly created session 

380 """ 

381 ... 

382 

383 

384@runtime_checkable 

385class ContextPrunerProtocol(Protocol): 

386 """Relevance-based conversation history pruning. 

387 

388 Unlike sliding-window truncation, implementations score each turn for 

389 relevance to the current query and preserve high-value turns regardless 

390 of their position in the history. 

391 """ 

392 

393 async def prune( 

394 self, 

395 history: Sequence[ChatMessageProtocol], 

396 current_query: str, 

397 max_turns: int, 

398 ) -> list[ChatMessageProtocol]: 

399 """Prune history to at most max_turns, preserving relevant turns. 

400 

401 Args: 

402 history: Full conversation history as ChatMessageProtocol instances. 

403 current_query: The current user query (used for relevance scoring). 

404 max_turns: Maximum number of turns to retain. 

405 

406 Returns: 

407 Pruned history in chronological order. 

408 """ 

409 ... 

410 

411 

412__all__ = [ 

413 "ContextPrunerProtocol", 

414 "SessionCheckpoint", 

415 "SessionContextProtocol", 

416 "SessionError", 

417 "SessionManagerProtocol", 

418 "SessionState", 

419 "SessionStatus", 

420 "SessionStoreProtocol", 

421 "SessionTurn", 

422 "TaskCancelledError", 

423 "TaskError", 

424 "TaskTimeoutError", 

425 "TaskValidationError", 

426]